diff --git a/CHANGELOG.md b/CHANGELOG.md index 6e008666..480b0dfc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,28 +1,45 @@ # NVIDIA CUTLASS Changelog + +## 3.5 (2024-03-18) + +- Implicit GEMM Convolutions targeting Hopper SM90A via WGMMA + [TMA im2col](./include/cute/atom/copy_traits_sm90_im2col.hpp) + + Native implementation in CUTLASS 3.x using CuTe, mirroring the [same design hierarchy as that of GEMMs](./media/docs/gemm_api_3x.md). + + Support for 1D, 2D, and 3D convolutions in a [rank-agnostic fashion](./include/cutlass/conv/convnd_problem_shape.hpp). + + Support for [Fprop](./test/unit/conv/device_3x/fprop/sm90_conv3d_fprop_implicit_gemm_s8_s8_s32_tensorop_s32.cu), [Dgrad](./test/unit/conv/device_3x/dgrad/sm90_conv2d_dgrad_implicit_gemm_f16_f16_f32_tensorop_f16.cu), and [Wgrad](./test/unit/conv/device_3x/wgrad/sm90_conv1d_wgrad_implicit_gemm_f16_f16_f32_tensorop_f16.cu) algorithms + + [CUTLASS profiler support](./python/cutlass_library/conv3x_emitter.py) for 2D and 3D convolutions implemented via the 3.x API. + + NOTE: this is a beta release. Further updates to CUTLASS will include major performance improvements, feature enablement, and possible breaking changes to the API until 3.7 release. Your feedback is welcome on the design! +- Support for [Ada (SM89) FP8 tensor cores via the 2.x API](./examples/58_ada_fp8_gemm/ada_fp8_gemm.cu). Requires CUDA 12.4 or newer. +- [Ampere gather/scatter convolution example](./examples/59_ampere_gather_scatter_gemm/README.md) in CuTe and CUTLASS 3.x + + Showcasing how custom kernels can be written and optimized using CUTLASS 3.x and CuTe and the general strategy for implementing convolutions as specializations of GETTs. + + Implementation of a coarse grained sparse gather/scatter kernel achieving peak performance on Ampere class tensor cores. +- Updates to CuTe documentation for [`cute::Tensor<>`](./media/docs/cute/03_tensor.md), [MMA atoms](./media/docs/cute/0t_mma_atom.md), and an overhauled [CuTe GEMM tutorial series](./examples/cute/tutorial). +- Extensions to CuTe to support [L2 prefetching](./include/cute/algorithm/prefetch.hpp) and [TMA store+reductions](./include/cute/arch/copy_sm90_tma.hpp#L1337). +- Updates and bugfixes from the community (thanks!) + ## [3.4.1](https://github.com/NVIDIA/cutlass/releases/tag/v3.4.1) (2024-02-14) -- Statically available [CUTLASS Version macros](/include/cutlass/version.h) that allow for handling API changes between CUTLASS releases on the users' side. -- Improvements for Hopper [Group-GEMMs](/examples/57_hopper_grouped_gemm) and [Pointer-Array Batched GEMMs](/examples/56_hopper_ptr_array_batched_gemm). +- Statically available [CUTLASS Version macros](./include/cutlass/version.h) that allow for handling API changes between CUTLASS releases on the users' side. +- Improvements for Hopper [Group-GEMMs](./examples/57_hopper_grouped_gemm) and [Pointer-Array Batched GEMMs](./examples/56_hopper_ptr_array_batched_gemm). - Updates and bugfixes from the community (thanks!). ## [3.4.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.4.0) (2024-01-12) -* Expanded [Mixed-input Hopper GEMMs](/examples/55_hopper_mixed_dtype_gemm) support covering {16-bit, 8-bit} x {8-bit, 4-bit} input types with fast numerical converters and group scaling factors. -* Performance improvements to [Mixed-input Hopper GEMMs](/examples/55_hopper_mixed_dtype_gemm) -* Beta release of [Pointer-Array Batched GEMMs](/examples/56_hopper_ptr_array_batched_gemm) now available on Hopper GPUs utilizing TMA and WGMMA (requires CUDA 12.3 or above). -* Beta release of [Group-GEMM](/examples/57_hopper_grouped_gemm) utilizing TMA and WGMMA (requires CUDA 12.3 or above). -* [Ampere Sparse GEMM](/examples/15_ampere_sparse_tensorop_gemm/ampere_sparse_tensorop_gemm_with_visitor.cu) supports Epilogue Visitor Tree (EVT) now. -* NamedBarriers usability improvement and list of [ReservedNamedBarriers](/include/cutlass/arch/barrier.h) has been officially released. -* Improved [CuTe documentation](/media/docs/cute/) including improved clarity and depth of [Quickstart](/media/docs/cute/00_quickstart.md), [CuTe Layout](/media/docs/cute/01_layout.md), and [CuTe Layout Algebra](/media/docs/cute/02_layout_algebra.md). Associated code comments, post-conditions, and details in [CuTe Core Unit Tests](/test/unit/cute/core/) also improved. +* Expanded [Mixed-input Hopper GEMMs](./examples/55_hopper_mixed_dtype_gemm) support covering {16-bit, 8-bit} x {8-bit, 4-bit} input types with fast numerical converters and group scaling factors. +* Performance improvements to [Mixed-input Hopper GEMMs](./examples/55_hopper_mixed_dtype_gemm) +* Beta release of [Pointer-Array Batched GEMMs](./examples/56_hopper_ptr_array_batched_gemm) now available on Hopper GPUs utilizing TMA and WGMMA (requires CUDA 12.3 or above). +* Beta release of [Group-GEMM](./examples/57_hopper_grouped_gemm) utilizing TMA and WGMMA (requires CUDA 12.3 or above). +* [Ampere Sparse GEMM](./examples/15_ampere_sparse_tensorop_gemm/ampere_sparse_tensorop_gemm_with_visitor.cu) supports Epilogue Visitor Tree (EVT) now. +* NamedBarriers usability improvement and list of [ReservedNamedBarriers](./include/cutlass/arch/barrier.h) has been officially released. +* Improved [CuTe documentation](./media/docs/cute/) including improved clarity and depth of [Quickstart](./media/docs/cute/00_quickstart.md), [CuTe Layout](./media/docs/cute/01_layout.md), and [CuTe Layout Algebra](./media/docs/cute/02_layout_algebra.md). Associated code comments, post-conditions, and details in [CuTe Core Unit Tests](./test/unit/cute/core/) also improved. ## [3.3](https://github.com/NVIDIA/cutlass/releases/tag/v3.3.0) (2023-10-31) -* [Mixed-input Hopper GEMMs](/examples/55_hopper_mixed_dtype_gemm) support covering 16-bit x 8-bit input operand types. +* [Mixed-input Hopper GEMMs](./examples/55_hopper_mixed_dtype_gemm) support covering 16-bit x 8-bit input operand types. * [Mixed-input Ampere GEMMs](https://github.com/NVIDIA/cutlass/pull/1084) with support for canonical layouts (TN). The implementation supports upcast on operandB {fp16, bf16} x {s8, u8}, and upcast on operandA {s8, u8} x {fp16, bf16}. -* [Copy Async based Hopper GEMMs](/test/unit/gemm/device/sm90_gemm_bf16_bf16_bf16_alignx_tensor_op_f32_warpspecialized_cooperative.cu) - which support lower than 16B aligned input tensors. +* [Copy Async based Hopper GEMMs](./test/unit/gemm/device/sm90_gemm_bf16_bf16_bf16_alignx_tensor_op_f32_warpspecialized_cooperative.cu) - which support lower than 16B aligned input tensors. * Kernel schedules and Builder support for mixed precision and Copy Async GEMMs with < 16B aligned input tensors. * Profiler support for lower-aligned Hopper GEMMs. -* Performance Improvements to [Scatter-Gather Hopper Example](/examples/52_hopper_gather_scatter_fusion). +* Performance Improvements to [Scatter-Gather Hopper Example](./examples/52_hopper_gather_scatter_fusion). * Sub-Byte type fixes and improvements. -* EVT Support for RELU with Aux bitmap tensor store (used in dRELU). See [SM90 EVT fusions](/include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp) for details. +* EVT Support for RELU with Aux bitmap tensor store (used in dRELU). See [SM90 EVT fusions](./include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp) for details. * Fusion support for backprop fusions including drelu, dgelu, and dbias. * Support for void-C kernels and SM80 mixed-input GEMMs in the CUTLASS Python interface @@ -34,7 +51,7 @@ * SM80 EVT support in C++ and Python. * Other SM90 epilogue improvements. * Splitting CUTLASS library into smaller units based on operation, arch and datatypes. See [1105](https://github.com/NVIDIA/cutlass/discussions/1105) for details. -* Making `tools/library/scripts` packageable - `tools/library/scripts` is now moving to `python/cutlass_library`. See the Python [README](/python/README.md) for details. +* Making `tools/library/scripts` packageable - `tools/library/scripts` is now moving to `python/cutlass_library`. See the Python [README](./python/README.md) for details. * SM90 TF32 kernel improvements for all layouts. * SM90 rasterization direction support in the CUTLASS profiler. * Improvement for CUTLASS profiler build times. @@ -42,65 +59,65 @@ ## [3.2.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.2.0) (2023-08-03) -* New warp-specialized persistent FP8 GEMM kernel [kernel schedules](/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp) and [mainloops](/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp) targeting Hopper architecture that achieve great performance with TMA, WGMMA, and threadblock clusters. An example showcasing [Hopper warp-specialized FP8 GEMMs](/examples/54_hopper_fp8_warp_specialized_gemm). FP8 GEMMs come with a fast accumulation mode. When enabled, problem execution might be faster but at the cost of lower accuracy because intermediate results will not periodically be promoted to a higher precision. -* New [Epilogue Visitor Tree (EVT)](/examples/49_hopper_gemm_with_collective_builder/49_collective_builder.cu) support for Hopper TMA epilogues. EVTs allows for user-defined customized epilogue fusion patterns without having to write a new epilogue. -* [Stream-K](/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp) feature for Hopper. Note that this is only a functional implementation of stream-K, and should not be used for performance comparison. Optimizations are expected in a future release. -* Improved CTA rasterization and support for CTA swizzling for Hopper kernels using the [Tile Scheduler](/include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp). +* New warp-specialized persistent FP8 GEMM kernel [kernel schedules](./include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp) and [mainloops](./include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp) targeting Hopper architecture that achieve great performance with TMA, WGMMA, and threadblock clusters. An example showcasing [Hopper warp-specialized FP8 GEMMs](./examples/54_hopper_fp8_warp_specialized_gemm). FP8 GEMMs come with a fast accumulation mode. When enabled, problem execution might be faster but at the cost of lower accuracy because intermediate results will not periodically be promoted to a higher precision. +* New [Epilogue Visitor Tree (EVT)](./examples/49_hopper_gemm_with_collective_builder/49_collective_builder.cu) support for Hopper TMA epilogues. EVTs allows for user-defined customized epilogue fusion patterns without having to write a new epilogue. +* [Stream-K](./include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp) feature for Hopper. Note that this is only a functional implementation of stream-K, and should not be used for performance comparison. Optimizations are expected in a future release. +* Improved CTA rasterization and support for CTA swizzling for Hopper kernels using the [Tile Scheduler](./include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp). * Improved performance for [warp-specialized TensorFloat-32 (TF32) GEMM kernels](test/unit/gemm/device/sm90_gemm_tf32_tf32_f32_tensor_op_f32_gmma_rs_cluster_warpspecialized.cu) targeting Hopper TMA. -* [Hopper GEMM+Permute](/examples/53_hopper_gemm_permute/53_hopper_gemm_permute.cu), an example of fusing tensor reordering (permutation) with GEMM mainloop or epilogue. -* New CUTLASS 2D Convolution Python interface. New [example](/examples/python/03_basic_conv2d.ipynb) here. +* [Hopper GEMM+Permute](./examples/53_hopper_gemm_permute/53_hopper_gemm_permute.cu), an example of fusing tensor reordering (permutation) with GEMM mainloop or epilogue. +* New CUTLASS 2D Convolution Python interface. New [example](./examples/python/03_basic_conv2d.ipynb) here. * Support for Windows (MSVC) builds. Tested with Visual Studio 2019 v16.11.27 on Windows 10.0. * Optimal performance using [**CUDA 12.2u1**](https://developer.nvidia.com/cuda-downloads) * Updates and bugfixes from the community (thanks!) ## [3.1.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.1.0) (2023-04-14) -* New CUTLASS Python interface that aims to provide an ease-of-use interface for instantiating, emitting, compiling, and running CUTLASS kernels via Python. More details [here](/python/README.md) and new [examples](/examples/python). +* New CUTLASS Python interface that aims to provide an ease-of-use interface for instantiating, emitting, compiling, and running CUTLASS kernels via Python. More details [here](./python/README.md) and new [examples](./examples/python). * New [efficient epilogues](test/unit/gemm/device/sm90_gemm_f16_f16_f16_tensor_op_f32_cluster_warpspecialized_cooperative.cu#L783) using TMA for Hopper. * Support for [fused epilogues](test/unit/gemm/device/sm90_gemm_f16_f16_f16_tensor_op_f32_cluster_warpspecialized_cooperative_bias_elementwise.cu), such Bias, ReLU and GELU, using the new efficient epilogues. * New [warp-specialized TensorFloat-32 (TF32) GEMM kernels](test/unit/gemm/device/sm90_gemm_tf32_tf32_f32_tensor_op_f32_gmma_rs_cluster_warpspecialized.cu) targeting Hopper TMA. -* New [*warp-specialized persistent cooperative*](include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp) kernel design that allows for larger tile sizes and improves performance on Hopper. -* An [example](examples/51_hopper_gett) showcasing GEMM-Like Tensor-Tensor Contraction (GETT) capability on Hopper. -* Epilogue builders. Similar to mainloop builders (see [example 49](/examples/49_hopper_gemm_with_collective_builder/49_collective_builder.cu)), epilogue builders aim to generate the best-possible epilogue while exposing incremental opt-ins for greater customization. +* New [*warp-specialized persistent cooperative*](./include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp) kernel design that allows for larger tile sizes and improves performance on Hopper. +* An [example](./examples/51_hopper_gett) showcasing GEMM-Like Tensor-Tensor Contraction (GETT) capability on Hopper. +* Epilogue builders. Similar to mainloop builders (see [example 49](./examples/49_hopper_gemm_with_collective_builder/49_collective_builder.cu)), epilogue builders aim to generate the best-possible epilogue while exposing incremental opt-ins for greater customization. * Profiler support for overriding kernel and epilogue builder auto schedules for 3.x API kernels, allowing specific policies to be run in the CUTLASS profiler. -* Performance optimizations for the [*warp-specialized persistent ping-pong*](include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp) kernel. -* Changes to the [GEMM API 3.x](media/docs/gemm_api_3x.md), involving the host-facing arguments and the underlying `Params` structs. -* [FMHA Backward Pass](examples/41_fused_multi_head_attention/fused_multi_head_attention_backward.cu) from Meta xFormers. -* [Streamk GEMM with Broadcast](examples/47_ampere_gemm_universal_streamk/ampere_gemm_universal_streamk_broadcast.cu) enables epilogue broadcast with StreamK GEMM. -* [Batched B2B GEMM](examples/13_two_tensor_op_fusion) now can run multiple Back-to-Back GEMM with the same problem size in parallel. +* Performance optimizations for the [*warp-specialized persistent ping-pong*](./include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp) kernel. +* Changes to the [GEMM API 3.x](./media/docs/gemm_api_3x.md), involving the host-facing arguments and the underlying `Params` structs. +* [FMHA Backward Pass](./examples/41_fused_multi_head_attention/fused_multi_head_attention_backward.cu) from Meta xFormers. +* [Streamk GEMM with Broadcast](./examples/47_ampere_gemm_universal_streamk/ampere_gemm_universal_streamk_broadcast.cu) enables epilogue broadcast with StreamK GEMM. +* [Batched B2B GEMM](./examples/13_two_tensor_op_fusion) now can run multiple Back-to-Back GEMM with the same problem size in parallel. * [Batched Strided GEMV](test/unit/gemm/device/gemv.cu) support both row major and column major input matrix. -* [Permute + GEMM fusion](examples/39_gemm_permute) can fuse Permute with following GEMM now. Before, we only support fusing GEMM with Permute in the epilogue. -* [Row Broadcast](include/cutlass/epilogue/threadblock/predicated_tile_iterator_row_broadcast.h) can be fused in the epilogue. +* [Permute + GEMM fusion](./examples/39_gemm_permute) can fuse Permute with following GEMM now. Before, we only support fusing GEMM with Permute in the epilogue. +* [Row Broadcast](./include/cutlass/epilogue/threadblock/predicated_tile_iterator_row_broadcast.h) can be fused in the epilogue. * The GitHub branch is renamed from `master` to `main` in this release. * Optimal performance using [**CUDA 12.1**](https://developer.nvidia.com/cuda-downloads) * Updates and bugfixes from the community (thanks!) ## [3.0.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.0.0) (2023-01-23) -* [CuTe](/media/docs/cute/00_quickstart.md), a [new core library and backend](/include/cute) for CUTLASS 3.0 that defines a single Layout vocabulary type and an associated algebra of layouts for a much more expressive and composable abstraction for tensors, sets of parallel agents, and operations by said agents on tensors. -* [A new conceptual operation hierarchy](media/docs/cutlass_3x_design.md) that replaces the architecture-centric hierarchy of CUTLASS 2.x and [documentation for CUTLASS 3.0's GEMM API changes](/media/docs/gemm_api_3x.md). -* Strict API backwards compatibility that exposes both 2.x and 3.x API kernels through the same [`device::GemmUniversalAdapter`](include/cutlass/gemm/device/gemm_universal_adapter.h) and [`kernel::GemmUniversal`](include/cutlass/gemm/kernel/gemm_universal.hpp) types, allowing users to include both APIs in the same translation units. More information can be found in the [3.x backwards compatibility section](media/docs/cutlass_3x_backwards_compatibility.md). -* Updates to [Functionality](media/docs/functionality.md) which directs users on which kernels are supported via CUTLASS-2 and CUTLASS-3. -* Updates to [Compatibility](/README.md#compatibility) Section regarding supported compilers, operating systems, CUDA Toolkits, Hardware Architectures and [Target Architecture](/README.md#Target-Architecture). -* New warp-specialized GEMM [kernel schedules](include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp) and [mainloops](include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp) targeting Hopper architecture that achieve great performance with TMA, WGMMA, and threadblock clusters. +* [CuTe](./media/docs/cute/00_quickstart.md), a [new core library and backend](./include/cute) for CUTLASS 3.0 that defines a single Layout vocabulary type and an associated algebra of layouts for a much more expressive and composable abstraction for tensors, sets of parallel agents, and operations by said agents on tensors. +* [A new conceptual operation hierarchy](./media/docs/cutlass_3x_design.md) that replaces the architecture-centric hierarchy of CUTLASS 2.x and [documentation for CUTLASS 3.0's GEMM API changes](./media/docs/gemm_api_3x.md). +* Strict API backwards compatibility that exposes both 2.x and 3.x API kernels through the same [`device::GemmUniversalAdapter`](./include/cutlass/gemm/device/gemm_universal_adapter.h) and [`kernel::GemmUniversal`](./include/cutlass/gemm/kernel/gemm_universal.hpp) types, allowing users to include both APIs in the same translation units. More information can be found in the [3.x backwards compatibility section](./media/docs/cutlass_3x_backwards_compatibility.md). +* Updates to [Functionality](./media/docs/functionality.md) which directs users on which kernels are supported via CUTLASS-2 and CUTLASS-3. +* Updates to [Compatibility](./README.md#compatibility) Section regarding supported compilers, operating systems, CUDA Toolkits, Hardware Architectures and [Target Architecture](./README.md#Target-Architecture). +* New warp-specialized GEMM [kernel schedules](./include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp) and [mainloops](./include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp) targeting Hopper architecture that achieve great performance with TMA, WGMMA, and threadblock clusters. * Extensions to CUTLASS profiler to support threadblock cluster shapes in library and profiler tile configurations. -* [CUTLASS library integration](/tools/library/src/gemm_operation_3x.hpp) for 3.x API kernels built through the new `CollectiveBuilder` API, enabling CUTLASS profiler. -* Support for [Hopper GEMMs](examples/48_hopper_warp_specialized_gemm) through the new 3.0 API with CuTe-based exposure of the Hopper [Tensor Memory Accelerator](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor) and [WGMMA Tensor Core](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-instructions) features. -* Set of examples that demonstrate the usage of the new 3.0 API to easily build GEMM kernels targeting Hopper: examples [48](examples/48_hopper_warp_specialized_gemm), [49](examples/49_hopper_gemm_schedules_with_collective_builder), and [50](examples/50_hopper_gemm_with_epilogue_swizzle). +* [CUTLASS library integration](./tools/library/src/gemm_operation_3x.hpp) for 3.x API kernels built through the new `CollectiveBuilder` API, enabling CUTLASS profiler. +* Support for [Hopper GEMMs](./examples/48_hopper_warp_specialized_gemm) through the new 3.0 API with CuTe-based exposure of the Hopper [Tensor Memory Accelerator](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor) and [WGMMA Tensor Core](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-instructions) features. +* Set of examples that demonstrate the usage of the new 3.0 API to easily build GEMM kernels targeting Hopper: examples [48](./examples/48_hopper_warp_specialized_gemm), [49](./examples/49_hopper_gemm_schedules_with_collective_builder), and [50](./examples/50_hopper_gemm_with_epilogue_swizzle). ## [2.11.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.11.0) (2022-11-19) -* [Stream-K](/examples/47_ampere_gemm_universal_streamk), which is a new general way to do split-K. It can not only improve performance, but can also significantly reduce the number of tile sizes that need to be profiled to find the best one. -* [Fused multi-head attention Kernel](/examples/41_fused_multi_head_attention). It has two variants: one uses batched GEMM for the fixed sequence length, and the other one uses group GEMM for the variable sequence length. Both versions just need one kernel. -* [Dual GEMM](/examples/45_dual_gemm), which can fuse A x B and A x C into one kernel. Two GEMMs has no producer-consumer dependency. -* Hopper improves [double precision matrix multiplication](/test/unit/gemm/device/gemm_f64n_f64t_f64t_tensor_op_f64_sm90.cu) by 2x compared to Ampere at iso-clocks. It is supported since CUDA 11.8. -* [BLAS3](/test/unit/gemm/device/hemm_cf64_cf64_cf64_tensor_op_f64_sm90.cu) functions with Hoppers new double precision matrix multiplication instructions. -* [ELL Block Sparse GEMM](/examples/43_ell_block_sparse_gemm), which uses an [ELL matrix](https://developer.nvidia.com/blog/accelerating-matrix-multiplication-with-block-sparse-format-and-nvidia-tensor-cores/) to describe the sparsity of A matrix. B and output matrices are still dense. The block size can be arbitary. -* Optimized [Group Conv](/examples/42_ampere_tensorop_group_conv) for SingleGroup mode, which requires that the output channel per group is a multiple of Threadblock tile N. -* [Optimized DepthWise Conv](/examples/46_depthwise_simt_conv2dfprop/depthwise_simt_conv2dfprop.cu). Two new modes are added - * [kOptimized](/test/unit/conv/device/depthwise_conv2d_fprop_direct_conv_f16nhwc_f16nhwc_f16nhwc_simt_f16_sm60.cu) - use direct conv to compute instead of implicit GEMM. +* [Stream-K](./examples/47_ampere_gemm_universal_streamk), which is a new general way to do split-K. It can not only improve performance, but can also significantly reduce the number of tile sizes that need to be profiled to find the best one. +* [Fused multi-head attention Kernel](./examples/41_fused_multi_head_attention). It has two variants: one uses batched GEMM for the fixed sequence length, and the other one uses group GEMM for the variable sequence length. Both versions just need one kernel. +* [Dual GEMM](./examples/45_dual_gemm), which can fuse A x B and A x C into one kernel. Two GEMMs has no producer-consumer dependency. +* Hopper improves [double precision matrix multiplication](./test/unit/gemm/device/gemm_f64n_f64t_f64t_tensor_op_f64_sm90.cu) by 2x compared to Ampere at iso-clocks. It is supported since CUDA 11.8. +* [BLAS3](./test/unit/gemm/device/hemm_cf64_cf64_cf64_tensor_op_f64_sm90.cu) functions with Hoppers new double precision matrix multiplication instructions. +* [ELL Block Sparse GEMM](./examples/43_ell_block_sparse_gemm), which uses an [ELL matrix](https://developer.nvidia.com/blog/accelerating-matrix-multiplication-with-block-sparse-format-and-nvidia-tensor-cores/) to describe the sparsity of A matrix. B and output matrices are still dense. The block size can be arbitary. +* Optimized [Group Conv](./examples/42_ampere_tensorop_group_conv) for SingleGroup mode, which requires that the output channel per group is a multiple of Threadblock tile N. +* [Optimized DepthWise Conv](./examples/46_depthwise_simt_conv2dfprop/depthwise_simt_conv2dfprop.cu). Two new modes are added + * [kOptimized](./test/unit/conv/device/depthwise_conv2d_fprop_direct_conv_f16nhwc_f16nhwc_f16nhwc_simt_f16_sm60.cu) - use direct conv to compute instead of implicit GEMM. * The restrictions are: 1) input ,output channel and group number should be multiple of (128 / sizeof(input element)). 2) The input filter size should be the same as the template parameter configuration. - * [kFixedStrideDilation](/test/unit/conv/device/depthwise_conv2d_fprop_direct_conv_fixed_stride_dilation_f16nhwc_f16nhwc_f16nhwc_simt_f16_sm60.cu) - which puts stride and dilation into templates to further improve the performance. In this mode, kernel persistents some inputs into register to squeeze more performance, so large filter/stride/dilation is not recommanded. + * [kFixedStrideDilation](./test/unit/conv/device/depthwise_conv2d_fprop_direct_conv_fixed_stride_dilation_f16nhwc_f16nhwc_f16nhwc_simt_f16_sm60.cu) - which puts stride and dilation into templates to further improve the performance. In this mode, kernel persistents some inputs into register to squeeze more performance, so large filter/stride/dilation is not recommanded. * The restrictions are: 1) input, output channel and group number should be multiple of (128 / sizeof(input element)). 2) input filter size, stride, dilation should same as the template parameter configuration. -* [Scripts](/examples/44_multi_gemm_ir_and_codegen) to fuse multiple back-to-back GEMM. Its implementation was discussed in a GTC'22 Spring [talk](https://www.nvidia.com/en-us/on-demand/session/gtcspring22-s41606/). -* [FP8 data type definition](/include/cutlass/float8.h) and [conversion routines](/include/cutlass/numeric_conversion.h#L1274-2115). +* [Scripts](./examples/44_multi_gemm_ir_and_codegen) to fuse multiple back-to-back GEMM. Its implementation was discussed in a GTC'22 Spring [talk](https://www.nvidia.com/en-us/on-demand/session/gtcspring22-s41606/). +* [FP8 data type definition](./include/cutlass/float8.h) and [conversion routines](./include/cutlass/numeric_conversion.h#L1274-2115). * Updates and bugfixes from the community (thanks!). Big shout out to Meta's [xFormers](https://github.com/facebookresearch/xformers). * **Deprecation announcement:** CUTLASS plans to deprecate the following: @@ -109,54 +126,54 @@ * CUDA 10.2 ## [2.10.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.10.0) (2022-08-23) -* [CUTLASS Python](/examples/40_cutlass_py) now supports GEMM, CONV, Group GEMM for different data types as well as different epilogue flavours. -* Optimizations for CUTLASS's [Grouped GEMM](examples/24_gemm_grouped/gemm_grouped.cu) kernel. Threadblock scheduling part is improved. Some computation can be moved to the host side if applicable. [Grouped Syr2k](examples/38_syr2k_grouped/syr2k_grouped.cu) kernels are added, too. -* Optimizations for [GEMM+Softmax](examples/35_gemm_softmax). All the reduction computation is fused into the previous GEMM. More template arguments are provided to fine tune the performance. -* [Grouped GEMM for Multihead Attention](examples/41_multi_head_attention). This general group gemm based MHA does not require the sequence length of all GEMMs to be the same which makes it most useful for natural language processing. -* [GEMM + Layer norm fusion for Ampere](examples/37_gemm_layernorm_gemm_fusion/) splits the layernorm into two parts and both of them can be fused into the GEMMs before and after separately. In addition to use square sum to compute variance of layernorm, [Shift-K](https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Computing_shifted_data) is provided if square sum raise numerical issues. -* [GEMM Epilogue Permutation Fusion](examples/39_gemm_permute) can apply user provided permutation layout mapping in the GEMM epilogue. +* [CUTLASS Python](./examples/40_cutlass_py) now supports GEMM, CONV, Group GEMM for different data types as well as different epilogue flavours. +* Optimizations for CUTLASS's [Grouped GEMM](./examples/24_gemm_grouped/gemm_grouped.cu) kernel. Threadblock scheduling part is improved. Some computation can be moved to the host side if applicable. [Grouped Syr2k](./examples/38_syr2k_grouped/syr2k_grouped.cu) kernels are added, too. +* Optimizations for [GEMM+Softmax](./examples/35_gemm_softmax). All the reduction computation is fused into the previous GEMM. More template arguments are provided to fine tune the performance. +* [Grouped GEMM for Multihead Attention](./examples/41_multi_head_attention). This general group gemm based MHA does not require the sequence length of all GEMMs to be the same which makes it most useful for natural language processing. +* [GEMM + Layer norm fusion for Ampere](./examples/37_gemm_layernorm_gemm_fusion/) splits the layernorm into two parts and both of them can be fused into the GEMMs before and after separately. In addition to use square sum to compute variance of layernorm, [Shift-K](https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Computing_shifted_data) is provided if square sum raise numerical issues. +* [GEMM Epilogue Permutation Fusion](./examples/39_gemm_permute) can apply user provided permutation layout mapping in the GEMM epilogue. * [Grouped convolution targeting implicit GEMM](test/unit/conv/device/group_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu) introduces the first group convolution implementation to CUTLASS. It is an Analytical implementation, not an Optimized. The restrictions are: 1) input and output channel number should be multiple of group number. 2) split-K is not supported. The implementation has 2 modes: * kSingleGroup: output channel per group is multiple of Threadblock tile N. * kMultipleGroup: Threadblock tile N is multiple of output channel per group. * [Depthwise separable convolution](test/unit/conv/device/depthwise_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_simt_f16_sm60.cu) introduces the first depthwise convolution which is also Analytical for now. The restrictions are: 1) SIMT only 2) No split-K 3) input channel equals to output channel equals to group number. -* Standalone [Layernorm](/tools/util/include/cutlass/util/device_layernorm.h) and [Pooling](/tools/util/include/cutlass/util/device_nhwc_pooling.h) kernels. -* [Back-to-back GEMM/CONV](examples/13_two_tensor_op_fusion) relaxes the requirement that the first GEMM K dimension needs to be the multiple of Threadblock Tile K dimension. +* Standalone [Layernorm](./tools/util/include/cutlass/util/device_layernorm.h) and [Pooling](./tools/util/include/cutlass/util/device_nhwc_pooling.h) kernels. +* [Back-to-back GEMM/CONV](./examples/13_two_tensor_op_fusion) relaxes the requirement that the first GEMM K dimension needs to be the multiple of Threadblock Tile K dimension. * Optimal performance using [**CUDA 11.6u2**](https://developer.nvidia.com/cuda-downloads) * Updates and bugfixes from the community (thanks!) ## [2.9.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.9.0) (2022-04-21) -* [First layer Convolution kernels](/test/unit/conv/device/conv2d_fprop_fixed_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu) specialized for small channel counts and reduced alignment - * [Few channels](/include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_few_channels.h) specialization for reduced alignment capabilities - * [Fixed channels](/include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_fixed_channels.h) further specialized when channel count perfectly matches the access vector size - * [Unit tests](/test/unit/conv/device/conv2d_fprop_few_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu) - * [Python-based instance emitter](/python/cutlass_library/generator.py) in the CUTLASS Library and support in the Profiler +* [First layer Convolution kernels](./test/unit/conv/device/conv2d_fprop_fixed_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu) specialized for small channel counts and reduced alignment + * [Few channels](./include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_few_channels.h) specialization for reduced alignment capabilities + * [Fixed channels](./include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_fixed_channels.h) further specialized when channel count perfectly matches the access vector size + * [Unit tests](./test/unit/conv/device/conv2d_fprop_few_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu) + * [Python-based instance emitter](./python/cutlass_library/generator.py) in the CUTLASS Library and support in the Profiler * [BLAS3](https://docs.nvidia.com/cuda/cublas/index.html#cublas-level-3-function-reference) operators accelerated by Tensor Cores * Supported types: f32, cf32, f64, cf64, tf32x3, complex tf32x3 - * [HERK](/test/unit/gemm/device/her2k_cf32h_cf32n_tensor_op_fast_f32_sm80.cu) with [emitter](/python/cutlass_library/rank_k_operation.py) - * [SYRK](/test/unit/gemm/device/syrk_f32n_f32t_tensor_op_fast_f32_sm80.cu) with [emitter](/python/cutlass_library/rank_k_operation.py) - * [SYMM](/test/unit/gemm/device/symm_f32n_f32n_tensor_op_fast_f32_ls_sm80.cu) with [emitter](/python/cutlass_library/symm_operation.py) - * [TRMM](/test/unit/gemm/device/trmm_f32n_f32t_f32t_tensor_op_fast_f32_ls_sm80.cu) with [emitter](/python/cutlass_library/trmm_operation.py) - * [Unit tests](/test/unit/gemm/device/testbed_rank_k_universal.h) -* [CUTLASS Python](/examples/40_cutlass_py) demonstrating JIT compilation of CUTLASS kernels and a Python-based runtime using [CUDA Python](https://developer.nvidia.com/cuda-python) - * [Python-based runtime](/tools/library/scripts/rt.py) interoperable with existing emitters -* [GEMM + Softmax example](/examples/35_gemm_softmax) -* [Gather and Scatter Fusion with GEMM](/examples/36_gather_scatter_fusion) can gather inputs and scatters outputs based on indices vectors in the same GEMM kernel. + * [HERK](./test/unit/gemm/device/her2k_cf32h_cf32n_tensor_op_fast_f32_sm80.cu) with [emitter](./python/cutlass_library/rank_k_operation.py) + * [SYRK](./test/unit/gemm/device/syrk_f32n_f32t_tensor_op_fast_f32_sm80.cu) with [emitter](./python/cutlass_library/rank_k_operation.py) + * [SYMM](./test/unit/gemm/device/symm_f32n_f32n_tensor_op_fast_f32_ls_sm80.cu) with [emitter](./python/cutlass_library/symm_operation.py) + * [TRMM](./test/unit/gemm/device/trmm_f32n_f32t_f32t_tensor_op_fast_f32_ls_sm80.cu) with [emitter](./python/cutlass_library/trmm_operation.py) + * [Unit tests](./test/unit/gemm/device/testbed_rank_k_universal.h) +* [CUTLASS Python](./examples/40_cutlass_py) demonstrating JIT compilation of CUTLASS kernels and a Python-based runtime using [CUDA Python](https://developer.nvidia.com/cuda-python) + * [Python-based runtime](./tools/library/scripts/rt.py) interoperable with existing emitters +* [GEMM + Softmax example](./examples/35_gemm_softmax) +* [Gather and Scatter Fusion with GEMM](./examples/36_gather_scatter_fusion) can gather inputs and scatters outputs based on indices vectors in the same GEMM kernel. * It can select random rows in a row major matrix. * It can select random columns in a column major matrix. -* [Back-to-back GEMM/CONV](examples/13_two_tensor_op_fusion) fully supports buffering the first GEMM/CONV results in the shared memory for the latter one to use. It can eliminate register spill when the tile size is big. Additionally, bias vector add is supported in the first GEMM/CONV. +* [Back-to-back GEMM/CONV](./examples/13_two_tensor_op_fusion) fully supports buffering the first GEMM/CONV results in the shared memory for the latter one to use. It can eliminate register spill when the tile size is big. Additionally, bias vector add is supported in the first GEMM/CONV. * Supported kernels: GEMM and CONV. * Supported types: fp16 and int8. * Supported architectures: Turing and Ampere. -* [Transposed Convolution](/examples/34_transposed_conv2d) (a.k.a Deconvolution) support which reuses Dgrad implementation. -* [Utility functions](/tools/util/include/cutlass/util) that can pad NHWC and convert between NCHW and NHWC. +* [Transposed Convolution](./examples/34_transposed_conv2d) (a.k.a Deconvolution) support which reuses Dgrad implementation. +* [Utility functions](./tools/util/include/cutlass/util) that can pad NHWC and convert between NCHW and NHWC. * [Small alignment implicit gemm](https://github.com/NVIDIA/cutlass/issues/242) support for Fprop/Dgrad/Wgrad so that padding is no longer mandated to use tensor cores in these kernels. * Epilogue enhancement: * Eliminate bank conflicts in int8 tensor core kernels. * Half2 usage if epilogue compute type is fp16. * More activation functions: Silu, Hardswish, Leaky Relu. - * New elementwise fusion pattern for [residual block](/include/cutlass/epilogue/thread/linear_combination_residual_block.h). -* [Group GEMM](/examples/24_gemm_grouped) thread block number calculation fix which helps to launch the intended number of threadblocks to fully occupy the GPUs. + * New elementwise fusion pattern for [residual block](./include/cutlass/epilogue/thread/linear_combination_residual_block.h). +* [Group GEMM](./examples/24_gemm_grouped) thread block number calculation fix which helps to launch the intended number of threadblocks to fully occupy the GPUs. * [Parallel GEMM splitk](https://github.com/NVIDIA/cutlass/pull/277) support in the CUTLASS profiler. * Optimal performance using [**CUDA 11.6u2**](https://developer.nvidia.com/cuda-downloads) * Updates and bugfixes from the community (thanks!) @@ -166,17 +183,17 @@ * **TF32x3:** emulated single-precision using Tensor Cores * 45+ TFLOPs on NVIDIA A100 - * [GEMM SDK example](/examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/27_ampere_3xtf32_fast_accurate_tensorop_gemm.cu) (real) - * [COMPLEX GEMM SDK example](/examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_3xtf32_complex_gemm.cu) (complex) - * [Implicit GEMM Convolution SDK example](/examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/ampere_3xtf32_fast_accurate_tensorop_fprop.cu) + * [GEMM SDK example](./examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/27_ampere_3xtf32_fast_accurate_tensorop_gemm.cu) (real) + * [COMPLEX GEMM SDK example](./examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_3xtf32_complex_gemm.cu) (complex) + * [Implicit GEMM Convolution SDK example](./examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/ampere_3xtf32_fast_accurate_tensorop_fprop.cu) * **Mainloop fusion for Convolution:** convolution with fused per-channel scale-bias-relu - * [Conv Fprop SDK example](/examples/25_ampere_fprop_mainloop_fusion/ampere_fprop_mainloop_fusion.cu) - * [Conv WGrad SDK example](/examples/26_ampere_wgrad_mainloop_fusion/ampere_wgrad_mainloop_fusion.cu) - * [cutlass::conv::device::ImplicitGemmConvolutionFusion](/include/cutlass/conv/device/implicit_gemm_convolution_fusion.h) + * [Conv Fprop SDK example](./examples/25_ampere_fprop_mainloop_fusion/ampere_fprop_mainloop_fusion.cu) + * [Conv WGrad SDK example](./examples/26_ampere_wgrad_mainloop_fusion/ampere_wgrad_mainloop_fusion.cu) + * [cutlass::conv::device::ImplicitGemmConvolutionFusion](./include/cutlass/conv/device/implicit_gemm_convolution_fusion.h) * **Grouped GEMM:** similar to batched GEMM with distinct problem size per group - * [SDK example](/examples/24_gemm_grouped) with performance comparison with Batched Strided GEMM - * [cutlass::gemm::device::GemmGrouped](/include/cutlass/gemm/device/gemm_grouped.h) -* [Implicit GEMM Convolution fusion](/examples/13_two_tensor_op_fusion/) supports staging 1st convolution's output accumulator in the shared memory on Turing. This allows more flexible warp tile sizes and less regsiter pressue. + * [SDK example](./examples/24_gemm_grouped) with performance comparison with Batched Strided GEMM + * [cutlass::gemm::device::GemmGrouped](./include/cutlass/gemm/device/gemm_grouped.h) +* [Implicit GEMM Convolution fusion](./examples/13_two_tensor_op_fusion/) supports staging 1st convolution's output accumulator in the shared memory on Turing. This allows more flexible warp tile sizes and less regsiter pressue. * Optimal performance using [**CUDA 11.5**](https://developer.nvidia.com/cuda-downloads) * Updates from the community (thanks!) @@ -186,11 +203,11 @@ * CUDA 10.2 ## [2.7.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.7.0) (2021-09-24) - * Mainloop fusion for GEMM: [summation over A or B](/examples/23_ampere_gemm_operand_reduction_fusion/ampere_gemm_operand_reduction_fusion.cu) - * [Strided DGRAD (optimized iterators)](/include/cutlass/conv/kernel/default_conv2d_dgrad.h) - * [Half-precision GELU_taylor activation functions](/include/cutlass/epilogue/thread/activation.h#L196) + * Mainloop fusion for GEMM: [summation over A or B](./examples/23_ampere_gemm_operand_reduction_fusion/ampere_gemm_operand_reduction_fusion.cu) + * [Strided DGRAD (optimized iterators)](./include/cutlass/conv/kernel/default_conv2d_dgrad.h) + * [Half-precision GELU_taylor activation functions](./include/cutlass/epilogue/thread/activation.h#L196) * Use these when accumulation and epilogue compute types are all `cutlass::half_t` - * Tuning and bug fixes to [fused GEMM + GEMM example](/examples/13_two_tensor_op_fusion/) + * Tuning and bug fixes to [fused GEMM + GEMM example](./examples/13_two_tensor_op_fusion/) * Support for smaller than 128b aligned Convolutions: [see examples](test/unit/conv/device/conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.cu#L272) * Caching of results to accelerate Convolution [unit tests](test/unit/conv/device/cache_testbed_output.h) * Can be enabled or disabled by running `cmake .. -DCUTLASS_TEST_ENABLE_CACHED_RESULTS=OFF` @@ -205,24 +222,24 @@ ## [2.6.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.6.0) (2021-07-22) * Optimal performance when compiled with the [CUDA 11.4 Toolkit](https://developer.nvidia.com/cuda-toolkit) - * Adopt the new L2 prefetch feature in [cp.async](/include/cutlass/arch/memory.h) and [global load](/include/cutlass/arch/memory_sm80.h) + * Adopt the new L2 prefetch feature in [cp.async](./include/cutlass/arch/memory.h) and [global load](./include/cutlass/arch/memory_sm80.h) * Fused operators with GEMM and Convolution * [Fused broadcast in epilogue](test/unit/gemm/device/gemm_with_broadcast_f16n_f16n_f16n_tensorop_f32_sm75.cu) - * [Fused partial reduction in epilogue](/test/unit/gemm/device/gemm_with_reduction_f16n_f16n_f16n_tensorop_f32_sm75.cu) + * [Fused partial reduction in epilogue](./test/unit/gemm/device/gemm_with_reduction_f16n_f16n_f16n_tensorop_f32_sm75.cu) * 64b tensor strides and leading dimensions support for GEMMs * Affine rank=2 matrix layouts - * Row stride and column stride for matrices using [cutlass::layout::AffineRank2](/include/cutlass/layout/matrix.h) - * Support [FP64 tensor core](/examples/18_ampere_fp64_tensorop_affine2_gemm/ampere_fp64_tensorop_affine2_gemm.cu) and SIMT GEMM. - * [Batched GEMV](/test/unit/gemm/device/gemv.cu) preview implementation + * Row stride and column stride for matrices using [cutlass::layout::AffineRank2](./include/cutlass/layout/matrix.h) + * Support [FP64 tensor core](./examples/18_ampere_fp64_tensorop_affine2_gemm/ampere_fp64_tensorop_affine2_gemm.cu) and SIMT GEMM. + * [Batched GEMV](./test/unit/gemm/device/gemv.cu) preview implementation * [New strided Dgrad](test/unit/conv/device/conv2d_strided_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.cu) implementation * Accelerates over previous implementation by cutting down redundant math by 4x * Support using new `Dy` and `w` analytic iterators and existing `cutlass::conv::device::ImplicitGemmConvolution` interface * Quaternion-valued GEMM and Convolution in single- and double-precision (targeting CUDA Cores) - * Updates to [quaternion.h](/include/cutlass/quaternion.h) and [functional.h](/include/cutlass/functional.h) - * SDK Example for [GEMM](/examples/21_quaternion_gemm/quaternion_gemm.cu) and [Convolution](/examples/22_quaternion_conv/quaternion_conv.cu) - * [Unit tests for GEMM](/test/unit/gemm/device/simt_qgemm_nn_sm50.cu) and [Convolution](/test/unit/conv/device/conv2d_fprop_implicit_gemm_qf32nhwc_qf32nhwc_qf32nhwc_simt_f32_sm50.cu) + * Updates to [quaternion.h](./include/cutlass/quaternion.h) and [functional.h](./include/cutlass/functional.h) + * SDK Example for [GEMM](./examples/21_quaternion_gemm/quaternion_gemm.cu) and [Convolution](./examples/22_quaternion_conv/quaternion_conv.cu) + * [Unit tests for GEMM](./test/unit/gemm/device/simt_qgemm_nn_sm50.cu) and [Convolution](./test/unit/conv/device/conv2d_fprop_implicit_gemm_qf32nhwc_qf32nhwc_qf32nhwc_simt_f32_sm50.cu) * Many improvements to the epilogue. - * Provide an [option](/include/cutlass/epilogue/threadblock/epilogue.h) to not fully unroll the epilogue to reduce the code size and improve the performance when using complicated elementwise operations + * Provide an [option](./include/cutlass/epilogue/threadblock/epilogue.h) to not fully unroll the epilogue to reduce the code size and improve the performance when using complicated elementwise operations * Performance improvement for FP16 tensor core kernels * Bug fixes * Enhanced Clang support and the combination of Clang 13 and CUDA 11.4 can build and run kernels from Pascal and Ampere. @@ -234,14 +251,14 @@ ## [2.5.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.5.0) (2021-02-26) * Tensor reductions * _m_-to-_n_ reductions of tensors with affine layout - * [Specializations](/test/unit/reduction/device/tensor_reduce_contiguous.cu) for reductions including contiguous dimension - * [Specializations](/test/unit/reduction/device/tensor_reduce_strided.cu) for reductions excluding contiguous dimension + * [Specializations](./test/unit/reduction/device/tensor_reduce_contiguous.cu) for reductions including contiguous dimension + * [Specializations](./test/unit/reduction/device/tensor_reduce_strided.cu) for reductions excluding contiguous dimension * Custom reduction functors such as `cutlass::logical_and` * Large tensor support, up to 2^63 elements (however, each dimension is limited to an extent of 2^31) * Optimizations for 3-D convolution - * [Optimized tile iterators](include/cutlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h) using precomputed delta table for 3-D convolution + * [Optimized tile iterators](./include/cutlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h) using precomputed delta table for 3-D convolution * Full coverage of [forward](test/unit/conv/device/conv3d_fprop_implicit_gemm_f16ndhwc_f16ndhwc_f32ndhwc_tensor_op_f32_sm80.cu) and [backwards](test/unit/conv/device/conv3d_dgrad_implicit_gemm_f16ndhwc_f16ndhwc_f32ndhwc_tensor_op_f32_sm80.cu) passes for 3D convolution - * [Fused Convolution+Convolution example](/examples/13_two_tensor_op_fusion/README.md) + * [Fused Convolution+Convolution example](./examples/13_two_tensor_op_fusion/README.md) * Corrections and bug fixes reported by the CUTLASS community * Thank you for filing these issues! @@ -256,7 +273,7 @@ * Global memory iterators supporting Fprop, Dgrad, and Wgrad * `MmaMultistage` for implicit GEMM convolution for NVIDIA Ampere architecture * `MmaPipeline` for implicit GEMM convolution for NVIDIA Volta and Turing architectures - * [Documentation](/media/docs/implicit_gemm_convolution.md) describing Implicit GEMM Convolution algorithm and implementation + * [Documentation](./media/docs/implicit_gemm_convolution.md) describing Implicit GEMM Convolution algorithm and implementation ## [2.3.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.3.0) (2020-09-23) * [NVIDIA Ampere Architecture features](https://devblogs.nvidia.com/nvidia-ampere-architecture-in-depth/) @@ -264,13 +281,13 @@ * Direct access to Sparse Tensor Cores and maximum performance via [`mma.sp.sync`](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma-and-friends) * Fast SGEMM targeting GeForce RTX 30-series CUDA Cores * Minor Features: - * [Activation functions](/include/cutlass/epilogue/thread/activation.h) such as [GeLU](/include/cutlass/epilogue/thread/linear_combination_gelu.h) and [Sigmoid](/include/cutlass/epilogue/thread/linear_combination_sigmoid.h) - * Small [matrix](/include/cutlass/matrix.h) and [quaternion](/include/cutlass/quaternion.h) template classes in device code - * [Floating-point constants](/include/cutlass/constants.h) + * [Activation functions](./include/cutlass/epilogue/thread/activation.h) such as [GeLU](./include/cutlass/epilogue/thread/linear_combination_gelu.h) and [Sigmoid](./include/cutlass/epilogue/thread/linear_combination_sigmoid.h) + * Small [matrix](./include/cutlass/matrix.h) and [quaternion](./include/cutlass/quaternion.h) template classes in device code + * [Floating-point constants](./include/cutlass/constants.h) * NVIDIA Ampere GPU Architecture examples and documentation: - * [Tensor Float 32](/examples/14_ampere_tf32_tensorop_gemm/ampere_tf32_tensorop_gemm.cu) and - * [Sparse Tensor Cores](/examples/15_ampere_sparse_tensorop_gemm/ampere_sparse_tensorop_gemm.cu) - * Documentation added on CUTLASS [efficient row-major epilogue](/media/docs/gemm_api.md#efficient-epilogue) + * [Tensor Float 32](./examples/14_ampere_tf32_tensorop_gemm/ampere_tf32_tensorop_gemm.cu) and + * [Sparse Tensor Cores](./examples/15_ampere_sparse_tensorop_gemm/ampere_sparse_tensorop_gemm.cu) + * Documentation added on CUTLASS [efficient row-major epilogue](./media/docs/gemm_api.md#efficient-epilogue) ## [2.2.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.2.0) (2020-06-08) * [NVIDIA Ampere Architecture features](https://devblogs.nvidia.com/nvidia-ampere-architecture-in-depth/) @@ -290,11 +307,11 @@ * Disabled F16C by default for compatibility - enable on cmake command line with `-DCUTLASS_ENABLE_F16C=ON` ## [2.1.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.1.0) (2020-04-06) - * BLAS-style host-side API added to [CUTLASS Library](/media/docs/quickstart.md#cutlass-library) + * BLAS-style host-side API added to [CUTLASS Library](./media/docs/quickstart.md#cutlass-library) * API to launch compiled kernel instances for GEMM and planar complex GEMM * Planar Complex GEMM kernels targeting Volta and Turing Tensor Cores * Computes complex matrix products on matrices stored as disjoint real and imaginary parts - * [SDK Examples of Planar Complex GEMMs](/examples/10_planar_complex/planar_complex.cu) + * [SDK Examples of Planar Complex GEMMs](./examples/10_planar_complex/planar_complex.cu) * Minor enhancements and bug fixes ## [2.0.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.0.0) (2019-11-19) @@ -304,10 +321,10 @@ * Encapsulated functionality embodying modern C++11 programming techniques * Optimized containers and data types for efficient, generic, portable device code * Updates to: - * [Quick start guide](/media/docs/quickstart.md) - * [Documentation](/README.md#documentation) - * [Utilities](/media/docs/utilities.md) - * [CUTLASS Profiler](/media/docs/profiler.md) + * [Quick start guide](./media/docs/quickstart.md) + * [Documentation](./README.md#documentation) + * [Utilities](./media/docs/utilities.md) + * [CUTLASS Profiler](./media/docs/profiler.md) * Native Turing Tensor Cores * Efficient GEMM kernels targeting Turing Tensor Cores * Mixed-precision floating point, 8-bit integer, 4-bit integer, and binarized operands diff --git a/CMakeLists.txt b/CMakeLists.txt index ed759073..dd06a605 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -67,14 +67,13 @@ elseif (CUDA_VERSION VERSION_LESS 11.4) message(WARNING "CUTLASS ${CUTLASS_VERSION} support for CUDA ${CUDA_VERSION} is deprecated, please use CUDA 11.8 or higher.") endif() -if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS 7.5) - message(FATAL_ERROR "GCC version must be at least 7.5!") +if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS 7.3) + message(FATAL_ERROR "GCC version must be at least 7.3!") endif() if (CUDA_COMPILER MATCHES "[Cc]lang" AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS 7.0) message(FATAL_ERROR "Clang 7.0+ required for GPU compilation") endif() - find_package(Doxygen QUIET) ################################################################################ @@ -168,6 +167,7 @@ endif() include(GNUInstallDirs) link_directories(${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs) +link_directories(${CUDA_TOOLKIT_ROOT_DIR}/lib64) ################################################################################################### # diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 91537ea7..538bb658 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -1,6 +1,6 @@ -![ALT](/media/images/gemm-hierarchy-with-epilogue-no-labels.png "CUTLASS") +![ALT](./media/images/gemm-hierarchy-with-epilogue-no-labels.png "CUTLASS") -[README](/README.md#documentation) > **Contributors** +[README](./README.md#documentation) > **Contributors** # CUTLASS Developers and Contributors diff --git a/CUDA.cmake b/CUDA.cmake index de176f74..755b7476 100644 --- a/CUDA.cmake +++ b/CUDA.cmake @@ -326,6 +326,14 @@ function(cutlass_add_library NAME) cxx_std_11 ) + get_target_property(TARGET_TYPE ${NAME} TYPE) + + if (TARGET_TYPE MATCHES "SHARED") + set_target_properties(${NAME} PROPERTIES CUDA_RUNTIME_LIBRARY Shared) + elseif(TARGET_TYPE MATCHES "STATIC") + set_target_properties(${NAME} PROPERTIES CUDA_RUNTIME_LIBRARY Static) + endif() + if(__EXPORT_NAME) add_library(nvidia::cutlass::${__EXPORT_NAME} ALIAS ${NAME}) set_target_properties(${NAME} PROPERTIES EXPORT_NAME ${__EXPORT_NAME}) @@ -336,10 +344,19 @@ endfunction() function(cutlass_add_executable NAME) set(options) - set(oneValueArgs) + set(oneValueArgs CUDA_RUNTIME_LIBRARY) set(multiValueArgs) cmake_parse_arguments(_ "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + if (NOT DEFINED __CUDA_RUNTIME_LIBRARY) + set(__CUDA_RUNTIME_LIBRARY Shared) + endif() + + set(__CUDA_RUNTIME_LIBRARY_ALLOWED None Shared Static) + if (NOT __CUDA_RUNTIME_LIBRARY IN_LIST __CUDA_RUNTIME_LIBRARY_ALLOWED) + message(FATAL_ERROR "CUDA_RUNTIME_LIBRARY value '${__CUDA_RUNTIME_LIBRARY}' is not in allowed list of '${__CUDA_RUNTIME_LIBRARY_ALLOWED}'") + endif() + cutlass_unify_source_files(TARGET_SOURCE_ARGS ${__UNPARSED_ARGUMENTS}) if(CUTLASS_NATIVE_CUDA OR CUDA_COMPILER MATCHES "clang") @@ -359,6 +376,8 @@ function(cutlass_add_executable NAME) cxx_std_11 ) + set_target_properties(${NAME} PROPERTIES CUDA_RUNTIME_LIBRARY ${__CUDA_RUNTIME_LIBRARY}) + endfunction() function(cutlass_target_sources NAME) diff --git a/PUBLICATIONS.md b/PUBLICATIONS.md index 487309f7..32b76e5f 100644 --- a/PUBLICATIONS.md +++ b/PUBLICATIONS.md @@ -4,6 +4,7 @@ - ["A Case Study in CUDA Kernel Fusion: Implementing FlashAttention-2 on NVIDIA Hopper Architecture using the CUTLASS Library"](https://arxiv.org/abs/2312.11918). Ganesh Bikshandi, Jay Shah. _arXiv_, December 2023. +- ["Benchmarking GPU Tensor Cores on General Matrix Multiplication Kernels through CUTLASS"](https://www.mdpi.com/2076-3417/13/24/13022). Xuanteng Huang, Xianwei Zhang, Panfei Yang, Nong Xiao. _Journal of Applied Sciences_, December 2023. - ["A Speed Odyssey for Deployable Quantization of LLMs"](https://arxiv.org/abs/2311.09550). Qingyuan Li, Ran Meng, Yiduo Li, Bo Zhang, Liang Li, Yifan Lu, Xiangxiang Chu, Yerui Sun, Yuchen Xie. _arXiv_, November 2023. diff --git a/README.md b/README.md index d642ee72..98ddbb01 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,8 @@ -![ALT](/media/images/gemm-hierarchy-with-epilogue-no-labels.png "Complete CUDA GEMM decomposition") +![ALT](./media/images/gemm-hierarchy-with-epilogue-no-labels.png "Complete CUDA GEMM decomposition") -# CUTLASS 3.4 +# CUTLASS 3.5 -_CUTLASS 3.4 - February 2024_ +_CUTLASS 3.5 - March 2024_ CUTLASS is a collection of CUDA C++ template abstractions for implementing high-performance matrix-matrix multiplication (GEMM) and related computations at all levels @@ -19,16 +19,16 @@ mixed-precision computations, providing specialized data-movement and multiply-accumulate abstractions for half-precision floating point (FP16), BFloat16 (BF16), Tensor Float 32 (TF32), single-precision floating point (FP32), -[FP32 emulation via tensor core instruction](/examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm), +[FP32 emulation via tensor core instruction](./examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm), double-precision floating point (FP64) types, integer data types (4b and 8b), and binary data types (1b). CUTLASS demonstrates warp-synchronous matrix multiply operations targeting the programmable, high-throughput _Tensor Cores_ implemented by NVIDIA's Volta, Turing, Ampere, and Hopper architectures. -See the [Quick Start Guide](/media/docs/quickstart.md) to get started quickly. +See the [Quick Start Guide](./media/docs/quickstart.md) to get started quickly. -See the [functionality listing](/media/docs/functionality.md) for the list of operations +See the [functionality listing](./media/docs/functionality.md) for the list of operations supported at each level of the execution model hierarchy. CUTLASS 3.0 introduced a new core library, CuTe, to describe and manipulate tensors of threads and data. @@ -37,25 +37,27 @@ CuTe is a collection of C++ CUDA template abstractions for defining and operatin The core abstractions of CuTe are hierarchically multidimensional layouts which can be composed with data arrays to represent tensors. The representation of layouts is powerful enough to represent nearly everything we need to implement efficient dense linear algebra. Layouts can also be combined and manipulated via functional composition, on which we build a large set of common operations such as tiling and partitioning. CUTLASS 3.0 and beyond adopts CuTe throughout the GEMM hierarchy in its templates. This greatly simplifies the design -and improves code composability and readability. More documentation specific to CuTe can be found in its [dedicated documentation directory](/media/docs/cute/00_quickstart.md). +and improves code composability and readability. More documentation specific to CuTe can be found in its [dedicated documentation directory](./media/docs/cute/00_quickstart.md). In addition to GEMMs, CUTLASS implements high-performance convolution via the implicit GEMM algorithm. Implicit GEMM is the formulation of a convolution operation as a GEMM thereby taking advantage of CUTLASS's modular GEMM pipeline. This allows CUTLASS to build convolutions by reusing highly-optimized GEMM components. -# What's New in CUTLASS 3.4 +# What's New in CUTLASS 3.5 -CUTLASS 3.4.1 is an update to CUTLASS adding: -- Statically available [CUTLASS Version macros](/include/cutlass/version.h) that allow for handling API changes between CUTLASS releases on the users' side. -- Improvements for Hopper [Group-GEMM](/examples/57_hopper_grouped_gemm) and [Pointer-Array Batched GEMM](/examples/56_hopper_ptr_array_batched_gemm). -- Updates and bugfixes from the community (thanks!). +CUTLASS 3.5 is an update to CUTLASS adding: -CUTLASS 3.4.0 is an update to CUTLASS adding: - -- Improved [Mixed-input Hopper GEMMs](/examples/55_hopper_mixed_dtype_gemm) supporting {16-bit, 8-bit} x {8-bit, 4-bit} input types with fast numerical converters and group scaling factors tuned for optimal performance on Hopper H100. -- Beta release of [Pointer-Array Batched GEMMs](/examples/56_hopper_ptr_array_batched_gemm) utilizing TMA and Hopper H100 tensor cores now available. (Requires CUDA 12.3 or above) -- Beta release of [Group-GEMM](/examples/57_hopper_grouped_gemm) - commonly used in optimization of Mixture-Of-Expert models, is now available on Hopper GPUs taking advantage of TMA and Hopper H100 tensor cores. (Requires CUDA 12.3 or above) -- [Ampere Sparse GEMM](/examples/15_ampere_sparse_tensorop_gemm/ampere_sparse_tensorop_gemm_with_visitor.cu) supports Epilogue Visitor Tree (EVT) now. -- Improvements to NamedBarriers including details of [ReservedNamedBarriers](/include/cutlass/arch/barrier.h) used within the CUTLASS library. -- Improved [CuTe documentation](/media/docs/cute/) including improved clarity and depth of [Quickstart](/media/docs/cute/00_quickstart.md), [CuTe Layout](/media/docs/cute/01_layout.md), and [CuTe Layout Algebra](/media/docs/cute/02_layout_algebra.md). Associated code comments, post-conditions, and details in [CuTe Core Unit Tests](/test/unit/cute/core/) also improved. +- Implicit GEMM Convolutions targeting Hopper SM90A via WGMMA + [TMA im2col](./include/cute/atom/copy_traits_sm90_im2col.hpp) + + Native implementation in CUTLASS 3.x using CuTe, mirroring the [same design hierarchy as that of GEMMs](./media/docs/gemm_api_3x.md). + + Support for 1D, 2D, and 3D convolutions in a [rank-agnostic fashion](./include/cutlass/conv/convnd_problem_shape.hpp). + + Support for [Fprop](./test/unit/conv/device_3x/fprop/sm90_conv3d_fprop_implicit_gemm_s8_s8_s32_tensorop_s32.cu), [Dgrad](./test/unit/conv/device_3x/dgrad/sm90_conv2d_dgrad_implicit_gemm_f16_f16_f32_tensorop_f16.cu), and [Wgrad](./test/unit/conv/device_3x/wgrad/sm90_conv1d_wgrad_implicit_gemm_f16_f16_f32_tensorop_f16.cu) algorithms + + [CUTLASS profiler support](./python/cutlass_library/conv3x_emitter.py) for 2D and 3D convolutions implemented via the 3.x API. + + NOTE: this is a beta release. Further updates to CUTLASS will include major performance improvements, feature enablement, and possible breaking changes to the API until 3.7 release. Your feedback is welcome on the design! +- Support for [Ada (SM89) FP8 tensor cores via the 2.x API](./examples/58_ada_fp8_gemm/ada_fp8_gemm.cu). Requires CUDA 12.4 or newer. +- [Ampere gather/scatter convolution example](./examples/59_ampere_gather_scatter_gemm/README.md) in CuTe and CUTLASS 3.x + + Showcasing how custom kernels can be written and optimized using CUTLASS 3.x and CuTe and the general strategy for implementing convolutions as specializations of GETTs. + + Implementation of a coarse grained sparse gather/scatter kernel achieving peak performance on Ampere class tensor cores. +- Updates to CuTe documentation for [`cute::Tensor<>`](./media/docs/cute/03_tensor.md), [MMA atoms](./media/docs/cute/0t_mma_atom.md), and an overhauled [CuTe GEMM tutorial series](./examples/cute/tutorial). +- Extensions to CuTe to support [L2 prefetching](./include/cute/algorithm/prefetch.hpp) and [TMA store+reductions](./include/cute/arch/copy_sm90_tma.hpp#L1337). +- Updates and bugfixes from the community (thanks!) Minimum requirements: @@ -98,7 +100,7 @@ as shown in the above figure. Tensor Core operations are implemented using CUDA # Compatibility CUTLASS requires a C++17 host compiler and -performs best when built with the [**CUDA 12.3.2 Toolkit**](https://developer.nvidia.com/cuda-downloads). +performs best when built with the [**CUDA 12.4 Toolkit**](https://developer.nvidia.com/cuda-downloads). It is also compatible with CUDA 11.4, CUDA 11.5, CUDA 11.6, CUDA 11.7, CUDA 11.8, CUDA 12.0, CUDA 12.1, CUDA 12.2.2, CUDA 12.3.1 and CUDA 12.3.2. ## Operating Systems @@ -142,28 +144,28 @@ The target architecture information is passed on to CUTLASS via the cmake flag ` cmake .. -DCUTLASS_NVCC_ARCHS="90a" ``` -Please refer to the [functionality documentation](media/docs/functionality.md) for details on which kernels require which target architectures. +Please refer to the [functionality documentation](./media/docs/functionality.md) for details on which kernels require which target architectures. # Documentation CUTLASS is described in the following documents and the accompanying [Doxygen documentation](https://nvidia.github.io/cutlass). -- [Quick Start Guide](/media/docs/quickstart.md) - build and run CUTLASS -- [Functionality](/media/docs/functionality.md) - summarizes functionality available in CUTLASS -- [Efficient GEMM in CUDA](media/docs/efficient_gemm.md) - describes how GEMM kernels may be implemented efficiently in CUDA -- [CUTLASS 3.x Design](media/docs/cutlass_3x_design.md) - describes the CUTLASS 3.x design, its benefits, and how CuTe enables us to write much more composable components -- [GEMM API 3.x](media/docs/gemm_api_3x.md) - describes the CUTLASS 3.x GEMM model and C++ template concepts -- [GEMM API 2.x](media/docs/gemm_api.md) - describes the CUTLASS 2.x GEMM model and C++ template concepts -- [Implicit GEMM Convolution](media/docs/implicit_gemm_convolution.md) - describes 2-D and 3-D convolution in CUTLASS -- [Code Organization](media/docs/code_organization.md) - describes the organization and contents of the CUTLASS project -- [Terminology](media/docs/terminology.md) - describes terms used in the code -- [Programming Guidelines](media/docs/programming_guidelines.md) - guidelines for writing efficient modern CUDA C++ -- [Fundamental types](media/docs/fundamental_types.md) - describes basic C++ classes used in CUTLASS to represent numeric quantities and arrays -- [Layouts](media/docs/layout.md) - describes layouts of matrices and tensors in memory -- [Tile Iterators](media/docs/tile_iterator_concept.md) - describes C++ concepts for iterating over tiles of matrices in memory -- [CUTLASS Profiler](media/docs/profiler.md) - command-line driven profiling application -- [CUTLASS Utilities](media/docs/utilities.md) - additional templates used to facilate rapid development +- [Quick Start Guide](./media/docs/quickstart.md) - build and run CUTLASS +- [Functionality](./media/docs/functionality.md) - summarizes functionality available in CUTLASS +- [Efficient GEMM in CUDA](./media/docs/efficient_gemm.md) - describes how GEMM kernels may be implemented efficiently in CUDA +- [CUTLASS 3.x Design](./media/docs/cutlass_3x_design.md) - describes the CUTLASS 3.x design, its benefits, and how CuTe enables us to write much more composable components +- [GEMM API 3.x](./media/docs/gemm_api_3x.md) - describes the CUTLASS 3.x GEMM model and C++ template concepts +- [GEMM API 2.x](./media/docs/gemm_api.md) - describes the CUTLASS 2.x GEMM model and C++ template concepts +- [Implicit GEMM Convolution](./media/docs/implicit_gemm_convolution.md) - describes 2-D and 3-D convolution in CUTLASS +- [Code Organization](./media/docs/code_organization.md) - describes the organization and contents of the CUTLASS project +- [Terminology](./media/docs/terminology.md) - describes terms used in the code +- [Programming Guidelines](./media/docs/programming_guidelines.md) - guidelines for writing efficient modern CUDA C++ +- [Fundamental types](./media/docs/fundamental_types.md) - describes basic C++ classes used in CUTLASS to represent numeric quantities and arrays +- [Layouts](./media/docs/layout.md) - describes layouts of matrices and tensors in memory +- [Tile Iterators](./media/docs/tile_iterator_concept.md) - describes C++ concepts for iterating over tiles of matrices in memory +- [CUTLASS Profiler](./media/docs/profiler.md) - command-line driven profiling application +- [CUTLASS Utilities](./media/docs/utilities.md) - additional templates used to facilate rapid development # Resources We have also described the structure of an efficient GEMM in our talk at the @@ -182,7 +184,7 @@ projects. Client applications should target CUTLASS's `include/` directory in th paths. CUTLASS unit tests, examples, and utilities can be build with CMake. -The minimum version of CMake is given in the [Quickstart guide](media/docs/quickstart.md). +The minimum version of CMake is given in the [Quickstart guide](./media/docs/quickstart.md). Make sure the `CUDACXX` environment variable points to NVCC in the CUDA Toolkit installed on your system. @@ -227,7 +229,7 @@ CUTLASS is arranged as a header-only library along with Utilities, Tools, Exampl and template concepts defined in the CUTLASS project. A detailed explanation of the source code organization may be found in the -[CUTLASS documentation](media/docs/code_organization.md), but several main components are summarized below. +[CUTLASS documentation](./media/docs/code_organization.md), but several main components are summarized below. ## CUTLASS Template Library @@ -276,7 +278,7 @@ include/ # client applications should target this directory ### CUTLASS SDK Examples -[CUTLASS SDK examples](/examples) apply CUTLASS templates to implement basic computations. +[CUTLASS SDK examples](./examples) apply CUTLASS templates to implement basic computations. ### Tools @@ -301,7 +303,7 @@ tools/ The `test/unit/` directory consist of unit tests implemented with Google Test that demonstrate basic usage of Core API components and complete tests of the CUTLASS GEMM computations. -Instructions for building and running the Unit tests are described in the [Quickstart guide](media/docs/quickstart.md). +Instructions for building and running the Unit tests are described in the [Quickstart guide](./media/docs/quickstart.md). # Performance Profiling @@ -517,9 +519,9 @@ reference_device: Passed ## More Details on Compiling CUTLASS Kernels and CUTLASS Profiler - Please follow the links for more CMake examples on selectively compiling CUTLASS kernels: - - [GEMM CMake Examples](media/docs/quickstart.md#gemm-cmake-examples) - - [Implicit GEMM convolution CMake Examples](media/docs/quickstart.md#convolution-cmake-examples) -- [Further details about the CUTLASS Profiler are described here.](media/docs/profiler.md) + - [GEMM CMake Examples](./media/docs/quickstart.md#gemm-cmake-examples) + - [Implicit GEMM convolution CMake Examples](./media/docs/quickstart.md#convolution-cmake-examples) +- [Further details about the CUTLASS Profiler are described here.](./media/docs/profiler.md) # About diff --git a/cmake/googletest.cmake b/cmake/googletest.cmake index 89d64e6d..0350fb2d 100644 --- a/cmake/googletest.cmake +++ b/cmake/googletest.cmake @@ -37,7 +37,7 @@ endif() FetchContent_Declare( googletest GIT_REPOSITORY https://github.com/google/googletest.git - GIT_TAG v1.13.0 + GIT_TAG v1.14.0 ) FetchContent_GetProperties(googletest) diff --git a/examples/03_visualize_layout/visualize_layout.h b/examples/03_visualize_layout/visualize_layout.h index 2244f60b..f070bad2 100644 --- a/examples/03_visualize_layout/visualize_layout.h +++ b/examples/03_visualize_layout/visualize_layout.h @@ -260,7 +260,7 @@ private: if (options.vectorize <= 2) return std::make_pair(false, -1); // Boundary check. - if (i > elements.size() || (i + options.vectorize - 1) > elements.size()) + if (i > int(elements.size()) || (i + options.vectorize - 1) > int(elements.size())) return std::make_pair(false, -1); // Check if either all elements are valid or invalid. diff --git a/examples/04_tile_iterator/tile_iterator.cu b/examples/04_tile_iterator/tile_iterator.cu index e6eb380d..b9441a56 100644 --- a/examples/04_tile_iterator/tile_iterator.cu +++ b/examples/04_tile_iterator/tile_iterator.cu @@ -94,7 +94,7 @@ __global__ void copy( typename Iterator::Fragment fragment; - for(int i = 0; i < fragment.size(); ++i) { + for(size_t i = 0; i < fragment.size(); ++i) { fragment[i] = 0; } diff --git a/examples/05_batched_gemm/batched_gemm.cu b/examples/05_batched_gemm/batched_gemm.cu index be1ca4a9..5fb7518f 100644 --- a/examples/05_batched_gemm/batched_gemm.cu +++ b/examples/05_batched_gemm/batched_gemm.cu @@ -207,15 +207,15 @@ cudaError_t strided_batched_gemm_nn_reference( cudaError_t result = cudaSuccess; - if (A.size() < lda * k * batch_count) { + if (A.size() < size_t(lda * k * batch_count)) { std::cout << "the size of A is too small" << std::endl; return cudaErrorInvalidValue; } - if (B.size() < ldb * n) { + if (B.size() < size_t(ldb * n)) { std::cout << "the size of B is too small" << std::endl; return cudaErrorInvalidValue; } - if (C.size() < ldc * n * batch_count) { + if (C.size() < size_t(ldc * n * batch_count)) { std::cout << "the size of C is too small" << std::endl; return cudaErrorInvalidValue; } diff --git a/examples/13_two_tensor_op_fusion/b2b_grouped_gemm_run.h b/examples/13_two_tensor_op_fusion/b2b_grouped_gemm_run.h index 63d5802d..eb04105f 100644 --- a/examples/13_two_tensor_op_fusion/b2b_grouped_gemm_run.h +++ b/examples/13_two_tensor_op_fusion/b2b_grouped_gemm_run.h @@ -102,7 +102,7 @@ struct B2bFusedGroupedGemmRun if (dist_kind == cutlass::Distribution::Uniform) { cutlass::reference::host::TensorFillRandomUniform( - view, seed, 2, -2, 0); + view, seed, 1, -1, 0); } else if (dist_kind == cutlass::Distribution::Identity) { diff --git a/examples/13_two_tensor_op_fusion/kernel/b2b_gemm.h b/examples/13_two_tensor_op_fusion/kernel/b2b_gemm.h index add21f12..fca87a1d 100644 --- a/examples/13_two_tensor_op_fusion/kernel/b2b_gemm.h +++ b/examples/13_two_tensor_op_fusion/kernel/b2b_gemm.h @@ -157,35 +157,34 @@ struct B2bGemm { // Data members // - GemmUniversalMode mode; - GemmCoord problem_size_0; - GemmCoord problem_size_1; - typename B2bMma::IteratorA0::TensorRef ref_A0; - typename B2bMma::IteratorB0::TensorRef ref_B0; - typename Epilogue::OutputTileIterator::TensorRef ref_C0; - typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Scale0; - typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Bias0; - typename B2bMma::IteratorB1::TensorRef ref_B1; - typename Epilogue::OutputTileIterator::TensorRef ref_C1; - typename Epilogue::OutputTileIterator::TensorRef ref_D1; - int64_t batch_stride_A0; - int64_t batch_stride_B0; - int64_t batch_stride_B1; - int64_t batch_stride_C1; - int64_t batch_stride_D1; - int64_t batch_stride_Bias0; - int64_t batch_stride_Scale0; - typename OutputOp0::Params epilogue0; - typename OutputOp1::Params epilogue1; - int batch_count; + GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm; + GemmCoord problem_size_0{0,0,0}; + GemmCoord problem_size_1{0,0,0}; + typename B2bMma::IteratorA0::TensorRef ref_A0{}; + typename B2bMma::IteratorB0::TensorRef ref_B0{}; + typename Epilogue::OutputTileIterator::TensorRef ref_C0{}; + typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Scale0{}; + typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Bias0{}; + typename B2bMma::IteratorB1::TensorRef ref_B1{}; + typename Epilogue::OutputTileIterator::TensorRef ref_C1{}; + typename Epilogue::OutputTileIterator::TensorRef ref_D1{}; + int64_t batch_stride_A0{0}; + int64_t batch_stride_B0{0}; + int64_t batch_stride_B1{0}; + int64_t batch_stride_C1{0}; + int64_t batch_stride_D1{0}; + int64_t batch_stride_Bias0{0}; + int64_t batch_stride_Scale0{0}; + typename OutputOp0::Params epilogue0 {}; + typename OutputOp1::Params epilogue1 {}; + int batch_count{1}; // // Methods // /// Default ctor - CUTLASS_HOST_DEVICE - Arguments() : mode(mode), problem_size_0(0, 0, 0), problem_size_1(0, 0, 0), batch_count(1) {} + Arguments() = default; /// Constructs an Arguments structure CUTLASS_HOST_DEVICE @@ -285,47 +284,45 @@ struct B2bGemm { /// Parameters structure struct Params { - cutlass::gemm::GemmUniversalMode mode; - cutlass::gemm::GemmCoord problem_size_0; - cutlass::gemm::GemmCoord problem_size_1; - cutlass::gemm::GemmCoord grid_tiled_shape; - int swizzle_log_tile; - typename B2bMma::IteratorA0::Params params_A0; - typename B2bMma::IteratorA0::TensorRef ref_A0; - typename B2bMma::IteratorB0::Params params_B0; - typename B2bMma::IteratorB0::TensorRef ref_B0; - typename Epilogue::OutputTileIterator::Params params_C0; - typename Epilogue::OutputTileIterator::TensorRef ref_C0; - typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Scale0; - typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Bias0; - typename B2bMma::IteratorB1::Params params_B1; - typename B2bMma::IteratorB1::TensorRef ref_B1; - typename Epilogue::OutputTileIterator::Params params_C1; - typename Epilogue::OutputTileIterator::TensorRef ref_C1; - typename Epilogue::OutputTileIterator::Params params_D1; - typename Epilogue::OutputTileIterator::TensorRef ref_D1; - typename OutputOp0::Params output_op_0; - typename OutputOp1::Params output_op_1; - int64_t batch_stride_A0; - int64_t batch_stride_B0; - int64_t batch_stride_B1; - int64_t batch_stride_C1; - int64_t batch_stride_D1; - int64_t batch_stride_Bias0; - int64_t batch_stride_Scale0; - int *semaphore; - int gemm_k_iterations_0; - int gemm_k_size_0; - int gemm_k_iterations_1; - int gemm_k_size_1; + cutlass::gemm::GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm; + cutlass::gemm::GemmCoord problem_size_0{}; + cutlass::gemm::GemmCoord problem_size_1{}; + cutlass::gemm::GemmCoord grid_tiled_shape{}; + int swizzle_log_tile{0}; + typename B2bMma::IteratorA0::Params params_A0{}; + typename B2bMma::IteratorA0::TensorRef ref_A0{}; + typename B2bMma::IteratorB0::Params params_B0{}; + typename B2bMma::IteratorB0::TensorRef ref_B0{}; + typename Epilogue::OutputTileIterator::Params params_C0{}; + typename Epilogue::OutputTileIterator::TensorRef ref_C0{}; + typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Scale0{}; + typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Bias0{}; + typename B2bMma::IteratorB1::Params params_B1{}; + typename B2bMma::IteratorB1::TensorRef ref_B1{}; + typename Epilogue::OutputTileIterator::Params params_C1{}; + typename Epilogue::OutputTileIterator::TensorRef ref_C1{}; + typename Epilogue::OutputTileIterator::Params params_D1{}; + typename Epilogue::OutputTileIterator::TensorRef ref_D1{}; + typename OutputOp0::Params output_op_0{}; + typename OutputOp1::Params output_op_1{}; + int64_t batch_stride_A0{0}; + int64_t batch_stride_B0{0}; + int64_t batch_stride_B1{0}; + int64_t batch_stride_C1{0}; + int64_t batch_stride_D1{0}; + int64_t batch_stride_Bias0{0}; + int64_t batch_stride_Scale0{0}; + int *semaphore = nullptr; + int gemm_k_iterations_0{0}; + int gemm_k_size_0{0}; + int gemm_k_iterations_1{0}; + int gemm_k_size_1{0}; // // Methods // - CUTLASS_HOST_DEVICE - Params(): mode(mode), swizzle_log_tile(0), semaphore(0), gemm_k_iterations_0(0), gemm_k_size_0(0), - gemm_k_iterations_1(0), gemm_k_size_1(0) { } + Params() = default; CUTLASS_HOST_DEVICE Params( diff --git a/examples/23_ampere_gemm_operand_reduction_fusion/CMakeLists.txt b/examples/23_ampere_gemm_operand_reduction_fusion/CMakeLists.txt index df22da5b..e5b4ec03 100644 --- a/examples/23_ampere_gemm_operand_reduction_fusion/CMakeLists.txt +++ b/examples/23_ampere_gemm_operand_reduction_fusion/CMakeLists.txt @@ -27,10 +27,14 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - +set(TEST_STANDARD --m=1024 --n=1024 --k=1024) +set(TEST_LARGE_PERFCHECK --m=4096 --n=3456 --k=4096 --perf-check) cutlass_example_add_executable( 23_ampere_gemm_operand_reduction_fusion ampere_gemm_operand_reduction_fusion.cu + TEST_COMMAND_OPTIONS + TEST_STANDARD + TEST_LARGE_PERFCHECK ) diff --git a/examples/23_ampere_gemm_operand_reduction_fusion/ampere_gemm_operand_reduction_fusion.cu b/examples/23_ampere_gemm_operand_reduction_fusion/ampere_gemm_operand_reduction_fusion.cu index 43a1e734..23fd6c18 100644 --- a/examples/23_ampere_gemm_operand_reduction_fusion/ampere_gemm_operand_reduction_fusion.cu +++ b/examples/23_ampere_gemm_operand_reduction_fusion/ampere_gemm_operand_reduction_fusion.cu @@ -377,22 +377,22 @@ Result profile(Options const &options) { cutlass::reference::host::TensorFillRandomUniform( tensor_a.host_view(), 1997, - ElementInputA(2), - ElementInputA(-2), + ElementInputA(1), + ElementInputA(-1), 0); // <- Fill tensor A on host with uniform-distribution random data cutlass::reference::host::TensorFillRandomUniform( tensor_b.host_view(), 2003, - ElementInputB(2), - ElementInputB(-2), + ElementInputB(1), + ElementInputB(-1), 0); // <- Fill tensor B on host with uniform-distribution random data cutlass::reference::host::TensorFillRandomUniform( tensor_c.host_view(), 2017, - ElementOutput(2), - ElementOutput(-2), + ElementOutput(1), + ElementOutput(-1), 0); // <- Fill matrix C on host with uniform-distribution random data cutlass::reference::host::TensorFill( tensor_d.host_view()); // <- fill matrix D on host with zeros diff --git a/examples/24_gemm_grouped/gemm_grouped.cu b/examples/24_gemm_grouped/gemm_grouped.cu index 9ebd9d1c..993d554f 100644 --- a/examples/24_gemm_grouped/gemm_grouped.cu +++ b/examples/24_gemm_grouped/gemm_grouped.cu @@ -789,7 +789,7 @@ public: problem_count_check += bin.second.size(); } - if (problem_count_check != this->problem_count()) { + if (problem_count_check != size_t(this->problem_count())) { std::cout << "\n***\nERROR in BINNING LOGIC!\n***\n" << std::endl; } diff --git a/examples/31_basic_syrk/basic_syrk.cu b/examples/31_basic_syrk/basic_syrk.cu index 2a5c0c7d..9f9cd93a 100644 --- a/examples/31_basic_syrk/basic_syrk.cu +++ b/examples/31_basic_syrk/basic_syrk.cu @@ -113,10 +113,10 @@ cudaError_t CutlassSsyrkNN( >, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>, 5, // Stages - 1, // AligmentA + 1, // AlignmentA false, // SplitKSerail - cutlass::arch::OpMultiplyAdd, - cutlass::ComplexTransform::kNone, + cutlass::arch::OpMultiplyAdd, + cutlass::ComplexTransform::kNone, cutlass::BlasMode::kSymmetric >; @@ -149,7 +149,7 @@ cudaError_t CutlassSsyrkNN( // // Launch the CUTLASS SYRK kernel. // - + cutlass::Status status = syrk_operator(args); // diff --git a/examples/35_gemm_softmax/gemm_softmax.cu b/examples/35_gemm_softmax/gemm_softmax.cu index 1942d49d..27156ea0 100644 --- a/examples/35_gemm_softmax/gemm_softmax.cu +++ b/examples/35_gemm_softmax/gemm_softmax.cu @@ -456,7 +456,7 @@ struct Testbed { bool verify_tensor(std::vector vector_Input, \ std::vector vector_Input_Ref) { - int64_t size = (vector_Input.size() < vector_Input_Ref.size()) ? vector_Input.size() : vector_Input_Ref.size(); + auto size = int64_t((vector_Input.size() < vector_Input_Ref.size()) ? vector_Input.size() : vector_Input_Ref.size()); float abs_tol = options.tolerance; float rel_tol = options.tolerance; diff --git a/examples/37_gemm_layernorm_gemm_fusion/gemm_layernorm.cu b/examples/37_gemm_layernorm_gemm_fusion/gemm_layernorm.cu index 8bd906c3..b5a0a1dc 100644 --- a/examples/37_gemm_layernorm_gemm_fusion/gemm_layernorm.cu +++ b/examples/37_gemm_layernorm_gemm_fusion/gemm_layernorm.cu @@ -454,48 +454,48 @@ struct Testbed { cutlass::reference::host::TensorFillRandomUniform( tensor_A0.host_view(), options.seed, - ElementInputA0(5), - ElementInputA0(-5), + ElementInputA0(4), + ElementInputA0(-4), 0 ); cutlass::reference::host::TensorFillRandomUniform( tensor_B0.host_view(), options.seed + 1, - ElementInputB0(5), - ElementInputB0(-5), + ElementInputB0(4), + ElementInputB0(-4), 0 ); cutlass::reference::host::TensorFillRandomUniform( tensor_A1.host_view(), options.seed + 2, - ElementInputA1(5), - ElementInputA1(-5), + ElementInputA1(4), + ElementInputA1(-4), 0 ); cutlass::reference::host::TensorFillRandomUniform( tensor_Beta.host_view(), options.seed + 3, - ElementInputScaleBias(5), - ElementInputScaleBias(-5), + ElementInputScaleBias(4), + ElementInputScaleBias(-4), 0 ); cutlass::reference::host::TensorFillRandomUniform( tensor_Gamma.host_view(), options.seed + 4, - ElementInputScaleBias(5), - ElementInputScaleBias(-5), + ElementInputScaleBias(4), + ElementInputScaleBias(-4), 0 ); cutlass::reference::host::TensorFillRandomUniform( tensor_Shifted_K.host_view(), options.seed + 5, - ElementOutput(5), - ElementOutput(-6), + ElementOutput(4), + ElementOutput(-5), 0 ); diff --git a/examples/38_syr2k_grouped/syr2k_grouped.cu b/examples/38_syr2k_grouped/syr2k_grouped.cu index 2342b1ef..c1fb82e8 100644 --- a/examples/38_syr2k_grouped/syr2k_grouped.cu +++ b/examples/38_syr2k_grouped/syr2k_grouped.cu @@ -803,7 +803,7 @@ public: // Use 'D' for the in/out workspace this->block_D.copy_from_device(this->block_C.get()); - for (int i = 0; i < this->options.problem_sizes.size(); ++i) { + for (size_t i = 0; i < this->options.problem_sizes.size(); ++i) { cutlass::gemm::GemmCoord const & problem = this->options.problem_sizes[i]; int32_t batch_count = 1; int64_t lda = this->lda_host.at(i); @@ -904,10 +904,10 @@ public: // Run profiling loop // - int last_stream_idx = 0; + size_t last_stream_idx = 0; for (int iter = 0; iter < this->options.iterations; ++iter) { - for (int i = 0; i < this->options.problem_sizes.size(); ++i) { + for (size_t i = 0; i < this->options.problem_sizes.size(); ++i) { cutlass::gemm::GemmCoord const & problem = this->options.problem_sizes[i]; int32_t batch_count = 1; int64_t lda = this->lda_host.at(i); @@ -1146,7 +1146,7 @@ public: ); // Initialize the Rank2K object - Rank2K rank2k; + Rank2K rank2k{}; size_t workspace_size = rank2k.get_workspace_size(args); cutlass::DeviceAllocation workspace(workspace_size); diff --git a/examples/41_fused_multi_head_attention/debug_utils.h b/examples/41_fused_multi_head_attention/debug_utils.h index 040f793b..90c0a69b 100644 --- a/examples/41_fused_multi_head_attention/debug_utils.h +++ b/examples/41_fused_multi_head_attention/debug_utils.h @@ -40,7 +40,7 @@ // Nans & inf detection #define NANCHECK(frag) \ { \ - for (int _i = 0; _i < frag.size(); ++_i) { \ + for (size_t _i = 0; _i < frag.size(); ++_i) { \ assert(std::isfinite(float(frag[_i]))); \ assert(!std::isnan(float(frag[_i]))); \ } \ @@ -147,7 +147,7 @@ constexpr __string_view __get_type_name() { { \ auto typeStr = __get_type_name(); \ PRINT_B0_T0("printing %s (%s)", name, typeStr.data); \ - for (int _start = 0; _start < frag.size(); _start += 8) { \ + for (size_t _start = 0; _start < frag.size(); _start += 8) { \ PRINT_ACCUM8_T0_L0_START(" ", frag, _start); \ } \ /*__syncthreads(); \ diff --git a/examples/41_fused_multi_head_attention/fmha_grouped.h b/examples/41_fused_multi_head_attention/fmha_grouped.h index a0d49d33..22779b59 100644 --- a/examples/41_fused_multi_head_attention/fmha_grouped.h +++ b/examples/41_fused_multi_head_attention/fmha_grouped.h @@ -167,58 +167,39 @@ public: // Data members // - GemmCoord *problem_sizes0; - GemmCoord *problem_sizes1; + GemmCoord *problem_sizes0{nullptr}; + GemmCoord *problem_sizes1{nullptr}; - int problem_count; - int threadblock_count; + int problem_count{0}; + int threadblock_count{0}; - ElementQ ** ptr_Q; - ElementK ** ptr_K; - ElementP ** ptr_P; - ElementV ** ptr_V; - ElementO ** ptr_O; - ElementOAccum ** ptr_O_accum; + ElementQ ** ptr_Q{nullptr}; + ElementK ** ptr_K{nullptr}; + ElementP ** ptr_P{nullptr}; + ElementV ** ptr_V{nullptr}; + ElementO ** ptr_O{nullptr}; + ElementOAccum ** ptr_O_accum{nullptr}; - typename LayoutQ::Stride::LongIndex *ldq; - typename LayoutK::Stride::LongIndex *ldk; - typename LayoutP::Stride::LongIndex *ldv; - typename LayoutO::Stride::LongIndex *ldo; - - // Scale - ElementAccumulator scale; + typename LayoutQ::Stride::LongIndex *ldq{nullptr}; + typename LayoutK::Stride::LongIndex *ldk{nullptr}; + typename LayoutP::Stride::LongIndex *ldv{nullptr}; + typename LayoutO::Stride::LongIndex *ldo{nullptr}; // Whether causal masking is to be performed - bool causal; + bool causal{false}; + + // Scale + ElementAccumulator scale{0}; // Only used by device-level operator - GemmCoord *host_problem_sizes; + GemmCoord *host_problem_sizes{nullptr}; // // Methods // - - /// Default ctor - CUTLASS_HOST_DEVICE - Arguments(): - problem_count(0), - threadblock_count(0), - ptr_Q(nullptr), - ptr_K(nullptr), - ptr_P(nullptr), - ptr_V(nullptr), - ptr_O(nullptr), - ptr_O_accum(nullptr), - ldq(nullptr), - ldk(nullptr), - ldv(nullptr), - ldo(nullptr), - scale(0), - causal(false), - host_problem_sizes(nullptr) - { - - } + + /// Default ctor + Arguments() = default; /// Ctor CUTLASS_HOST_DEVICE diff --git a/examples/41_fused_multi_head_attention/fused_multihead_attention_fixed_seqlen.cu b/examples/41_fused_multi_head_attention/fused_multihead_attention_fixed_seqlen.cu index a26de16e..cf02a7b9 100644 --- a/examples/41_fused_multi_head_attention/fused_multihead_attention_fixed_seqlen.cu +++ b/examples/41_fused_multi_head_attention/fused_multihead_attention_fixed_seqlen.cu @@ -286,7 +286,7 @@ struct Options { // Number of real-valued multiply-adds int64_t fops = int64_t(); - for (int i = 0; i < problem_sizes0.size(); ++i) { + for (size_t i = 0; i < problem_sizes0.size(); ++i) { auto const& problem0 = problem_sizes0[i]; auto const& problem1 = problem_sizes1[i]; for (int row = 0; row < problem0.m(); ++row) { diff --git a/examples/41_fused_multi_head_attention/fused_multihead_attention_variable_seqlen.cu b/examples/41_fused_multi_head_attention/fused_multihead_attention_variable_seqlen.cu index a07ad507..49d8699a 100644 --- a/examples/41_fused_multi_head_attention/fused_multihead_attention_variable_seqlen.cu +++ b/examples/41_fused_multi_head_attention/fused_multihead_attention_variable_seqlen.cu @@ -340,7 +340,7 @@ struct Options { // Number of real-valued multiply-adds int64_t fops = int64_t(); - for (int i = 0; i < problem_sizes0.size(); ++i) { + for (size_t i = 0; i < problem_sizes0.size(); ++i) { auto const& problem0 = problem_sizes0[i]; auto const& problem1 = problem_sizes1[i]; diff --git a/examples/41_fused_multi_head_attention/gemm/custom_mma_multistage.h b/examples/41_fused_multi_head_attention/gemm/custom_mma_multistage.h index 893f765c..eedcb637 100644 --- a/examples/41_fused_multi_head_attention/gemm/custom_mma_multistage.h +++ b/examples/41_fused_multi_head_attention/gemm/custom_mma_multistage.h @@ -244,11 +244,13 @@ class CustomMmaMultistage : public CustomMmaBase { CUTLASS_DEVICE bool set_prologue_done(bool value) { prologue_done_ = value; + return true; } CUTLASS_DEVICE bool set_zero_outside_bounds(bool value) { zero_outside_bounds_ = value; + return true; } template diff --git a/examples/41_fused_multi_head_attention/kernel_backward.h b/examples/41_fused_multi_head_attention/kernel_backward.h index 1132af69..e7372f13 100644 --- a/examples/41_fused_multi_head_attention/kernel_backward.h +++ b/examples/41_fused_multi_head_attention/kernel_backward.h @@ -1446,7 +1446,7 @@ struct AttentionBackwardKernel { uint8_t lane_id) { cutlass::Array dropout_keep_mask_doivj; - dropout_keep_mask_doivj.fill(1); + dropout_keep_mask_doivj.fill(cutlass::uint1b_t{1}); const float dropout_scale = kApplyDropout ? 1.0 / (1.0 - p.dropout_prob) : 1.0f; @@ -1744,7 +1744,7 @@ struct AttentionBackwardKernel { [&](int accum_m) {}, [&](int accum_m /*q*/, int accum_n /*k*/, int idx) { if (zij.at({accum_n, accum_m}) == scalar_t(0)) { - dropout_keep_mask_doivj[idx] = cutlass::uint1b_t(0); + dropout_keep_mask_doivj[idx] = cutlass::uint1b_t{0}; } }, [&](int accum_m) {}); diff --git a/examples/41_fused_multi_head_attention/kernel_forward.h b/examples/41_fused_multi_head_attention/kernel_forward.h index b5e989fd..4c80f549 100644 --- a/examples/41_fused_multi_head_attention/kernel_forward.h +++ b/examples/41_fused_multi_head_attention/kernel_forward.h @@ -40,7 +40,6 @@ #include #include -#include "cutlass/bfloat16.h" #include "cutlass/fast_math.h" #include "cutlass/gemm/gemm.h" #include "cutlass/layout/matrix.h" diff --git a/examples/43_ell_block_sparse_gemm/ell_block_sparse_gemm.cu b/examples/43_ell_block_sparse_gemm/ell_block_sparse_gemm.cu index 582e61fb..52d2d0cb 100644 --- a/examples/43_ell_block_sparse_gemm/ell_block_sparse_gemm.cu +++ b/examples/43_ell_block_sparse_gemm/ell_block_sparse_gemm.cu @@ -452,7 +452,7 @@ public: // Determine SMEM requirements and waive if not satisfied // - int smem_size = int(sizeof(typename Gemm::GemmKernel::SharedStorage)); + size_t smem_size = sizeof(typename Gemm::GemmKernel::SharedStorage); cudaDeviceProp properties; int device_idx; @@ -509,7 +509,7 @@ public: ); // Initialize the GEMM object - Gemm gemm; + Gemm gemm{}; result.status = gemm.initialize(args); diff --git a/examples/51_hopper_gett/gett_kernel.cuh b/examples/51_hopper_gett/gett_kernel.cuh index c2c94068..6a775d13 100644 --- a/examples/51_hopper_gett/gett_kernel.cuh +++ b/examples/51_hopper_gett/gett_kernel.cuh @@ -102,7 +102,8 @@ gett_kernel( ElementB, StrideB, 128 / cutlass::sizeof_bits::value, ElementAccumulator, TileShape, Shape<_1,_2,_1>, - cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, cutlass::gemm::collective::KernelScheduleAuto >::CollectiveOp; diff --git a/examples/52_hopper_gather_scatter_fusion/52_hopper_gather_scatter_fusion.cu b/examples/52_hopper_gather_scatter_fusion/52_hopper_gather_scatter_fusion.cu index 6f6826ad..52a8c19c 100644 --- a/examples/52_hopper_gather_scatter_fusion/52_hopper_gather_scatter_fusion.cu +++ b/examples/52_hopper_gather_scatter_fusion/52_hopper_gather_scatter_fusion.cu @@ -289,7 +289,8 @@ struct ExampleRunner ElementAccumulator, Shape<_128,_128,_64>, Shape<_2,_2,_1>, - cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename EpilogueOpt::SharedStorage))>, cutlass::gemm::collective::KernelScheduleAuto >::CollectiveOp; diff --git a/examples/52_hopper_gather_scatter_fusion/gather_gemm.hpp b/examples/52_hopper_gather_scatter_fusion/gather_gemm.hpp index 09b0a9b8..04777d8a 100644 --- a/examples/52_hopper_gather_scatter_fusion/gather_gemm.hpp +++ b/examples/52_hopper_gather_scatter_fusion/gather_gemm.hpp @@ -39,6 +39,11 @@ #include "gather_tensor.hpp" +namespace cutlass { + ///Forward declaration + struct CudaHostAdapter; +} + namespace cutlass::gemm::kernel { /////////////////////////////////////////////////////////////////////////////// @@ -143,10 +148,10 @@ public: // Kernel entry point API struct Params { - GemmUniversalMode mode; - ProblemShape problem_shape; - MainloopParams mainloop; - EpilogueParams epilogue; + GemmUniversalMode mode{}; + ProblemShape problem_shape{}; + MainloopParams mainloop{}; + EpilogueParams epilogue{}; GatherA gather_A{}; GatherB gather_B{}; }; @@ -191,14 +196,15 @@ public: } static - int + size_t get_workspace_size(Arguments const& args) { return 0; } static cutlass::Status - initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { + initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr, + CudaHostAdapter* cuda_adapter = nullptr) { return Status::kSuccess; } diff --git a/examples/52_hopper_gather_scatter_fusion/scatter_epilogue.hpp b/examples/52_hopper_gather_scatter_fusion/scatter_epilogue.hpp index 3b8acd20..323e9606 100644 --- a/examples/52_hopper_gather_scatter_fusion/scatter_epilogue.hpp +++ b/examples/52_hopper_gather_scatter_fusion/scatter_epilogue.hpp @@ -39,7 +39,7 @@ #include "cutlass/epilogue/collective/detail.hpp" #include "cute/tensor.hpp" -#include "cute/numeric/int.hpp" +#include "cute/numeric/numeric_types.hpp" #include "gather_tensor.hpp" diff --git a/examples/53_hopper_gemm_permute/53_hopper_gemm_permute.cu b/examples/53_hopper_gemm_permute/53_hopper_gemm_permute.cu index 23b8a39c..20a282b0 100644 --- a/examples/53_hopper_gemm_permute/53_hopper_gemm_permute.cu +++ b/examples/53_hopper_gemm_permute/53_hopper_gemm_permute.cu @@ -393,7 +393,8 @@ private: ElementB, StrideB, 128 / cutlass::sizeof_bits::value, ElementAccumulator, TileShape, ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, cutlass::gemm::collective::KernelScheduleAuto >::CollectiveOp; @@ -403,7 +404,8 @@ private: ElementB, StrideBPermute, 128 / cutlass::sizeof_bits::value, ElementAccumulator, TileShapePermute, ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpiloguePermute::SharedStorage))>, cutlass::gemm::collective::KernelScheduleAuto >::CollectiveOp; diff --git a/examples/53_hopper_gemm_permute/permute_kernel.cuh b/examples/53_hopper_gemm_permute/permute_kernel.cuh index 11578085..8abe7010 100644 --- a/examples/53_hopper_gemm_permute/permute_kernel.cuh +++ b/examples/53_hopper_gemm_permute/permute_kernel.cuh @@ -37,7 +37,7 @@ #include "cutlass/layout/matrix.h" #include "cutlass/tensor_view.h" #include "cutlass/fast_math.h" -#include "cute/numeric/uint128.hpp" +#include "cute/numeric/numeric_types.hpp" namespace example { diff --git a/examples/58_ada_fp8_gemm/CMakeLists.txt b/examples/58_ada_fp8_gemm/CMakeLists.txt new file mode 100644 index 00000000..2af32542 --- /dev/null +++ b/examples/58_ada_fp8_gemm/CMakeLists.txt @@ -0,0 +1,34 @@ + +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +cutlass_example_add_executable( + 58_ada_fp8_gemm + ada_fp8_gemm.cu + ) diff --git a/examples/58_ada_fp8_gemm/ada_fp8_gemm.cu b/examples/58_ada_fp8_gemm/ada_fp8_gemm.cu new file mode 100644 index 00000000..0a6c034d --- /dev/null +++ b/examples/58_ada_fp8_gemm/ada_fp8_gemm.cu @@ -0,0 +1,826 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Example of running an Ada FP8 GEMM. + + In addition to using FP8 Tensor Core instructions, the Ada FP8 GEMM uses a distinct epilogue + that enables additional scaling of operands/outputs, storing a pre-activation-function output + tensor (called the "auxiliary" output), and computing the absolute maximum value of the + outputs. + + Pseudocode for this epilogue is as follows: + + Aux = ((alpha * scale_a * scale_b) * accumulator) + ((beta * scale_c) * source) + bias + D = activation(Aux) + + if Aux is fp8 type: + abs_max_output = max( abs(aux) | (for every aux in Aux)) + Aux = scale_aux * Aux + endif + + if D is fp8 type: + abs_max_output = max( abs(d) | (for every d in D)) + D = scale_d * D + endif + + Parameter Aux is optionally stored to global memory +*/ + +#include +#include +#include + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/util/command_line.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/reference/host/gemm_complex.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/distribution.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_norm.h" +#include "cutlass/util/reference/host/gemm.h" + +#include "cutlass/epilogue/thread/activation.h" +#include "cutlass/epilogue/thread/linear_combination_generic_with_scaling.h" +#include "cutlass/gemm/device/gemm_universal_with_absmax.h" + +#include "cutlass/layout/matrix.h" +#include "cutlass/matrix_coord.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" + + +using ElementA = cutlass::float_e4m3_t; +using ElementB = cutlass::float_e4m3_t; +using ElementOutput = cutlass::float_e4m3_t; +using ElementAuxOutput = ElementOutput; +using ElementAccumulator = float; +using LayoutA = cutlass::layout::RowMajor; +using LayoutB = cutlass::layout::ColumnMajor; +using LayoutC = cutlass::layout::RowMajor; +static int const kStages = 3; +static int const kAlignmentA = 16; +static int const kAlignmentB = 16; + +using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax< + cutlass::epilogue::thread::ReLu, + ElementOutput, + ElementAuxOutput, + 128 / cutlass::sizeof_bits::value, + ElementAccumulator, + ElementAccumulator + >; + +template +using Gemm_ = cutlass::gemm::device::GemmUniversalWithAbsMax< + ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC, + ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89, + cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>, + EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages, + kAlignmentA, kAlignmentB, MathOperator + >; + +using ElementAbsmax = typename EpilogueOutputOp::ElementAbsmax; + + +// Command line options parsing +struct Options { + + bool help; + bool error; + bool reference_check; + cutlass::gemm::GemmCoord problem_size; + + int iterations; + int warmup_iterations; + + bool scale_A; + bool scale_B; + bool scale_C; + + float alpha; + float beta; + + Options(): + help(false), + error(false), + reference_check(false), + iterations(20), + warmup_iterations(5), + scale_A(true), + scale_B(true), + scale_C(true), + alpha(1.f), + beta(0.f) + { } + + // Parses the command line + void parse(int argc, char const **args) { + cutlass::CommandLine cmd(argc, args); + + if (cmd.check_cmd_line_flag("help")) { + help = true; + return; + } + + cmd.get_cmd_line_argument("iterations", iterations, 20); + cmd.get_cmd_line_argument("warmup_iterations", warmup_iterations, 5); + cmd.get_cmd_line_argument("reference-check", reference_check, false); + cmd.get_cmd_line_argument("scale-A", scale_A, true); + cmd.get_cmd_line_argument("scale-B", scale_B, true); + cmd.get_cmd_line_argument("scale-C", scale_C, true); + cmd.get_cmd_line_argument("alpha", alpha, 1.f); + cmd.get_cmd_line_argument("beta", beta, 0.f); + + int m, n, k; + cmd.get_cmd_line_argument("m", m, 1024); + cmd.get_cmd_line_argument("n", n, 1024); + cmd.get_cmd_line_argument("k", k, 1024); + + problem_size = cutlass::gemm::GemmCoord{m, n, k}; + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "58_ada_fp8_gemm\n\n" + << " This example executes a GEMM using Ada FP8 Tensor Core operations. In addition to performing\n" + << " a normal GEMM, the kernel performs the following operations:\n" + << " Aux = ((alpha * scale_a * scale_b) * accumulator) + ((beta * scale_c) * source) + bias\n" + << " D = activation(Aux)\n\n" + << " if Aux is fp8:\n" + << " abs_max_output = max( abs(aux) | (for every aux in Aux) )\n" + << " Aux = scale_aux * Aux\n\n" + << " if D is fp8 type:\n" + << " abs_max_output = max( abs(d) | (for every d in D) )\n" + << " D = scale_d * D\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement\n\n" + << " --m= Sets the M dimension of the GEMM\n" + << " --n= Sets the N dimension of the GEMM\n" + << " --k= Sets the K dimension of the GEMM\n" + << " --scale-A= Whether to apply a scaling factor to operand A (default: true)\n" + << " --scale-B= Whether to apply a scaling factor to operand B (default: true)\n" + << " --scale-C= Whether to apply a scaling factor to operand C (default: true)\n" + << " --iterations= Number of profiling iterations to perform\n" + << " --warmup-iterations= Number of warmup iterations to perform\n" + << " --reference-check= If true, performs reference check\n"; + + return out; + } + + /// Compute performance in GFLOP/s + float gflops(float runtime_s) const { + // Two flops per multiply-add + return 2.0f * float(problem_size.product()) / float(1.0e9) / runtime_s; + } +}; + +/// Helper class to run the kernel +template +struct TestbedRunner { + + using ElementAccumulator = typename Gemm::ElementAccumulator; + using ElementCompute = typename Gemm::GemmKernel::Epilogue::OutputOp::ElementCompute; + using ElementScalingFactor = typename Gemm::EpilogueOutputOp::ElementScalingFactor; + + static bool const kScaleAux = Gemm::EpilogueOutputOp::kIsScalingAndAmaxAuxOutputNeeded; + static bool const kScaleOutput = Gemm::EpilogueOutputOp::kIsScalingAndAmaxOutputNeeded; + + /// Initialization + cutlass::Distribution::Kind init_A; + cutlass::Distribution::Kind init_B; + cutlass::Distribution::Kind init_C; + uint64_t seed; + + cutlass::HostTensor tensor_A; + cutlass::HostTensor tensor_B; + cutlass::HostTensor tensor_C; + cutlass::HostTensor tensor_Aux; + cutlass::HostTensor tensor_D; + cutlass::HostTensor tensor_Vector; + cutlass::HostTensor tmp_D; + cutlass::HostTensor reference_D; + cutlass::HostTensor reference_Aux; + cutlass::HostTensor scale_A; + cutlass::HostTensor scale_B; + cutlass::HostTensor scale_C; + cutlass::HostTensor scale_D; + cutlass::HostTensor scale_Aux; + cutlass::HostTensor abs_max_Aux; + cutlass::HostTensor abs_max_D; + cutlass::HostTensor reference_abs_max_Aux; + cutlass::HostTensor reference_abs_max_D; + + // + // Methods + // + + TestbedRunner( + bool scaleA = true, + bool scaleB = true, + bool scaleC = true, + cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_C_ = cutlass::Distribution::Uniform, + uint64_t seed_ = 2080 + ): + init_A(init_A_), init_B(init_B_), init_C(init_C_), seed(seed_) { } + + /// Helper to initialize scaling factors + template + bool initialize_scale_factor(cutlass::TensorView view, uint64_t seed, int bits=0) { + cutlass::reference::host::TensorFillRandomUniform(view, seed, double(1.), double(0.), bits); + return true; + } + + /// Helper to initialize a tensor view + template + bool initialize_tensor( + cutlass::TensorView view, + cutlass::Distribution::Kind dist_kind, + uint64_t seed) { + + if (dist_kind == cutlass::Distribution::Uniform) { + + double scope_max, scope_min; + int bits_input = cutlass::sizeof_bits::value; + int bits_output = cutlass::sizeof_bits::value; + + if (bits_input == 1) { + scope_max = 2; + scope_min = 0; + } else if (bits_input <= 8) { + scope_max = 2; + scope_min = -2; + } else if (bits_output == 16) { + scope_max = 5; + scope_min = -5; + } else { + scope_max = 8; + scope_min = -8; + } + + cutlass::reference::host::TensorFillRandomUniform( + view, seed, scope_max, scope_min, 0); + } + else if (dist_kind == cutlass::Distribution::Identity) { + + cutlass::reference::host::TensorFillIdentity(view); + } + else if (dist_kind == cutlass::Distribution::Gaussian) { + + cutlass::reference::host::TensorFillRandomGaussian(view, seed, 0, 0.5); + } + else if (dist_kind == cutlass::Distribution::Sequential) { + + cutlass::reference::host::BlockFillSequential( + view.data(), view.capacity()); + } + else { + std::cerr << "Not implemented"; + return false; + } + + return true; + } + + /// Initializes data structures + void initialize(const Options& options) { + // + // Allocate the GEMM workspace + // + + tensor_A.resize(options.problem_size.mk()); + tensor_B.resize(options.problem_size.kn()); + tensor_C.resize(options.problem_size.mn()); + tensor_D.resize(options.problem_size.mn()); + tensor_Vector.resize({1, options.problem_size.n()}); + reference_D.resize(options.problem_size.mn(), false); + tmp_D.resize(options.problem_size.mn(), false); + + initialize_tensor(tensor_A.host_view(), init_A, seed + 2019); + initialize_tensor(tensor_B.host_view(), init_B, seed + 2018); + initialize_tensor(tensor_C.host_view(), init_C, seed + 2017); + initialize_tensor(tensor_Vector.host_view(), init_C, seed + 2020); + + // It is possible to randomly initialize to all zeros, so override this with non-zeros + // in the upper left corner of each operand. + cutlass::Coord<2> origin(0); + tensor_A.host_view().at(origin) = typename Gemm::ElementA(1); + tensor_B.host_view().at(origin) = typename Gemm::ElementB(1); + tensor_C.host_view().at(origin) = typename Gemm::ElementC(1); + tensor_Vector.host_view().at(origin) = typename Gemm::ElementC(1); + + cutlass::reference::host::TensorFill(tensor_D.host_view()); + cutlass::reference::host::TensorCopy(reference_D.host_view(), tensor_C.host_view()); + + tensor_A.sync_device(); + tensor_B.sync_device(); + tensor_C.sync_device(); + tensor_D.sync_device(); + tensor_Vector.sync_device(); + + int scale_bits = 2; + if (options.scale_A) { + scale_A.resize({1, 1}); + initialize_scale_factor(scale_A.host_view(), seed + 2021, scale_bits); + scale_A.sync_device(); + } + + if (options.scale_B) { + scale_B.resize({1, 1}); + initialize_scale_factor(scale_B.host_view(), seed + 2022, scale_bits); + scale_B.sync_device(); + } + + if (options.scale_C) { + scale_C.resize({1, 1}); + initialize_scale_factor(scale_C.host_view(), seed + 2023, scale_bits); + scale_C.sync_device(); + } + + if (kScaleOutput) { + scale_D.resize({1, 1}); + initialize_scale_factor(scale_D.host_view(), seed + 2024, scale_bits); + scale_D.sync_device(); + + abs_max_D.resize({1, 1}); + cutlass::reference::host::TensorFill(abs_max_D.host_view()); + abs_max_D.sync_device(); + + reference_abs_max_D.resize({1, 1}); + } + + if (kScaleAux) { + tensor_Aux.resize(options.problem_size.mn()); + cutlass::reference::host::TensorFill(tensor_Aux.host_view()); + tensor_Aux.sync_device(); + + scale_Aux.resize({1, 1}); + initialize_scale_factor(scale_Aux.host_view(), seed + 2025, scale_bits); + scale_Aux.sync_device(); + + abs_max_Aux.resize({1, 1}); + cutlass::reference::host::TensorFill(abs_max_Aux.host_view()); + abs_max_Aux.sync_device(); + + reference_Aux.resize(options.problem_size.mn(), false); + reference_abs_max_Aux.resize({1, 1}); + } + } + + /// Compares computed reference with device reference and outputs to a file if incorrect + bool compare_reference(const Options& options) { + + tensor_D.sync_host(); + + bool passed = cutlass::reference::host::TensorEquals(reference_D.host_view(), tensor_D.host_view()); + + if (kScaleAux) { + tensor_Aux.sync_host(); + abs_max_Aux.sync_host(); + passed &= cutlass::reference::host::TensorEquals(reference_Aux.host_view(), tensor_Aux.host_view()); + passed &= cutlass::reference::host::TensorEquals(abs_max_Aux.host_view(), reference_abs_max_Aux.host_view()); + } + + if (kScaleOutput) { + abs_max_D.sync_host(); + passed &= cutlass::reference::host::TensorEquals(abs_max_D.host_view(), reference_abs_max_D.host_view()); + } + + if (!passed) { + std::cerr << "Reference check failed" << std::endl; + + std::string output_file = "testbed_with_amax_errors.txt"; + std::ofstream file(output_file); + + file + << "problem: " << options.problem_size + << ", alpha: " << options.alpha << ", beta: " << options.beta << "\n\n"; + + file + << "A =\n" << tensor_A.host_view() + << "\nB =\n" << tensor_B.host_view() + << "\nC =\n" << tensor_C.host_view() + << "\nVector =\n" << tensor_Vector.host_view() + << "\nScaleA = " << scale_A.host_view() + << "\nScaleB = " << scale_B.host_view() + << "\nScaleC = " << scale_C.host_view() + << "\nScaleD = " << scale_D.host_view() + << "\nScaleAux = " << scale_Aux.host_view() + << "\n\nReference D =\n" << reference_D.host_view() + << "\nComputed D =\n" << tensor_D.host_view(); + if (kScaleAux) { + file + << "\n\nReference Aux =\n" << reference_Aux.host_view() + << "\nComputed Aux =\n" << tensor_Aux.host_view() + << "\n\nReference Absmax Aux = " << reference_abs_max_Aux.host_view() + << "\nComputed Absmax Aux = " << abs_max_Aux.host_view(); + } + if (kScaleOutput) { + file + << "\n\nReference Absmax D = " << reference_abs_max_D.host_view() + << "\nComputed Absmax D = " << abs_max_D.host_view(); + } + + std::cerr << "Dumped results to " << output_file << std::endl; + + } + + return passed; + } + + /// Verifies the result is a GEMM + bool verify(const Options& options) { + + cutlass::Coord<2> origin(0); + ElementCompute scaled_alpha = options.alpha; + if (options.scale_A) { + scaled_alpha *= scale_A.host_view().at(origin); + } + if (options.scale_B) { + scaled_alpha *= scale_B.host_view().at(origin); + } + + ElementCompute scaled_beta = options.beta; + if (options.scale_C) { + scaled_beta *= scale_C.host_view().at(origin); + } + + // + // Verify + // + + cutlass::reference::host::GemmComplex< + typename Gemm::ElementA, typename Gemm::LayoutA, + typename Gemm::ElementB, typename Gemm::LayoutB, + typename Gemm::ElementC, typename Gemm::LayoutC, + ElementCompute, ElementAccumulator, ElementAccumulator + >( + options.problem_size, + scaled_alpha, + tensor_A.host_ref(), + Gemm::kTransformA, + tensor_B.host_ref(), + Gemm::kTransformB, + scaled_beta, + tensor_C.host_ref(), + tmp_D.host_ref(), + ElementAccumulator(0) + ); + + ElementCompute tmp_abs_max_Aux(0.); + ElementCompute tmp_abs_max_D(0.); + + cutlass::NumericConverter cvt_c_to_compute; + cutlass::NumericConverter cvt_accum_to_compute; + cutlass::NumericConverter cvt_compute_to_accum; + cutlass::NumericConverter cvt_compute_to_d; + cutlass::NumericConverter cvt_compute_to_aux; + + cutlass::absolute_value_op abs; + cutlass::maximum_with_nan_propogation max; + cutlass::epilogue::thread::ReLu act; + + ElementScalingFactor d_scale = kScaleOutput ? scale_D.host_view().at(origin) : ElementScalingFactor(1.); + + for (int m = 0; m < options.problem_size.m(); ++m) { + for (int n = 0; n < options.problem_size.n(); ++n) { + ElementCompute intermediate = cvt_accum_to_compute(tmp_D.host_view().at({m, n})); + ElementCompute bias = cvt_c_to_compute(tensor_Vector.host_view().at({0, n})); + ElementCompute aux = intermediate + bias; + ElementCompute d = act(aux); + tmp_abs_max_Aux = max(abs(aux), tmp_abs_max_Aux); + tmp_abs_max_D = max(abs(d), tmp_abs_max_D); + reference_D.host_view().at({m, n}) = cvt_compute_to_d(d * d_scale); + + if (kScaleAux) { + reference_Aux.host_view().at({m, n}) = cvt_compute_to_aux(aux * scale_Aux.host_view().at(origin)); + } + } + } + + if (kScaleAux) { + reference_abs_max_Aux.host_view().at(origin) = cvt_compute_to_accum(tmp_abs_max_Aux); + } + + if (kScaleOutput) { + reference_abs_max_D.host_view().at(origin) = cvt_compute_to_accum(tmp_abs_max_D); + } + + return compare_reference(options); + } + + /// Returns true if the CUDA device is sufficient to execute the kernel. + bool sufficient() const { + + if (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 4)) { + std::cerr << "This example requires CUDA 12.4 or greater." << std::endl; + return false; + } + + size_t smem_size = sizeof(typename Gemm::GemmKernel::SharedStorage); + + cudaDeviceProp properties; + int device_idx; + cudaError_t result = cudaGetDevice(&device_idx); + + if (result != cudaSuccess) { + std::cerr << "cudaGetDevice() failed with error: " << cudaGetErrorString(result) << std::endl; + return false; + } + + result = cudaGetDeviceProperties(&properties, device_idx); + + if (result != cudaSuccess) { + std::cerr << "cudaGetDeviceProperties() failed with error: " << cudaGetErrorString(result) << std::endl; + return false; + } + + if (properties.major < 8 || (properties.major == 8 && properties.minor < 9)) { + std::cerr << "CUTLASS's Ada FP8 GEMM example requires a device of compute capability 89 or higher.\n" << std::endl; + return false; + } + + if (properties.sharedMemPerBlockOptin < smem_size) { + std::cerr << "Insufficient shared memory. Need " << smem_size + << ", but device only has " << properties.sharedMemPerBlockOptin << std::endl; + return false; + } + + return true; + } + + /// Executes one test + bool run(Options& options) + { + + // Waive test if insufficient CUDA device + if (!sufficient()) { + std::cerr << "Insufficient resources to run the kernel." << std::endl; + return false; + } + + this->initialize(options); + + // + // Initialize the GEMM operator + // + + typename Gemm::EpilogueOutputOp::Params::ActivationParams activation_params{ + ElementCompute(options.alpha), + ElementCompute(options.beta) + }; + typename Gemm::EpilogueOutputOp::Params epilogue_params{ + activation_params, + scale_A.device_data(), + scale_B.device_data(), + scale_C.device_data(), + scale_D.device_data(), + scale_Aux.device_data(), + abs_max_Aux.device_data(), + abs_max_D.device_data() + }; + + typename Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + options.problem_size, + /* batch_count = */ 1, + epilogue_params, + tensor_A.device_data(), + tensor_B.device_data(), + tensor_C.device_data(), + tensor_D.device_data(), + tensor_Aux.device_data(), + tensor_Vector.device_data(), + options.problem_size.m() * options.problem_size.k(), + options.problem_size.n() * options.problem_size.k(), + options.problem_size.m() * options.problem_size.n(), + options.problem_size.m() * options.problem_size.n(), + (int)options.problem_size.m(), // Batch stride vector + tensor_A.layout().stride(0), + tensor_B.layout().stride(0), + tensor_C.layout().stride(0), + tensor_D.layout().stride(0), + (int64_t)0 // Leading dimension of vector. This must be 0 + }; + + Gemm gemm_op; + + cutlass::Status status = gemm_op.can_implement(arguments); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Gemm::can_implement() failed" << std::endl; + return false; + } + + size_t workspace_size = Gemm::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + status = gemm_op.initialize(arguments, workspace.get()); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Gemm::initialize() failed" << std::endl; + return false; + } + + // + // Run the GEMM + // + + status = gemm_op(); + + if (status != cutlass::Status::kSuccess) { + std::cerr << "Gemm::run() failed" << std::endl; + return false; + } + + cudaError_t cuda_error = cudaDeviceSynchronize(); + if (cuda_error != cudaSuccess) { + std::cerr << "CUDA error: " << cudaGetErrorString(cuda_error) << std::endl; + return false; + } + + // + // Verify + // + + bool passed = true; + if (options.reference_check) { + passed &= this->verify(options); + } else { + std::cout << "Skipped reference check" << std::endl; + } + + // + // Warm up + // + + for (int i = 0; i < options.warmup_iterations; ++i) { + gemm_op(); + } + + // + // Profile + // + + cudaEvent_t events[2]; + cudaError_t error; + for (auto & event : events) { + error = cudaEventCreate(&event); + if (error != cudaSuccess) { + std::cerr << "cudaEventCreate() failed: " << cudaGetErrorString(error) << std::endl; + return false; + } + } + + // Record an event at the start of a series of GEMM operations + error = cudaEventRecord(events[0]); + if (error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(error) << std::endl; + return false; + } + + // Run profiling loop + for (int iter = 0; iter < options.iterations; ++iter) { + gemm_op(); + } + + // Record an event when the GEMM operations have been launched. + error = cudaEventRecord(events[1]); + if (error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(error) << std::endl; + return false; + } + + // Wait for work on the device to complete. + error = cudaEventSynchronize(events[1]); + if (error != cudaSuccess) { + std::cerr << "cudaEventSynchronize() failed: " << cudaGetErrorString(error) << std::endl; + return false; + } + + // Measure elapsed runtime + float runtime_ms = 0; + error = cudaEventElapsedTime(&runtime_ms, events[0], events[1]); + if (error != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(error) << std::endl; + return false; + } + + // Compute average runtime and GFLOPs. + runtime_ms = runtime_ms / float(options.iterations); + float gflops = options.gflops(runtime_ms / 1000.0f); + + std::cout << "Problem size: " << options.problem_size.m() << 'x' << options.problem_size.n() << 'x' << options.problem_size.k() << std::endl; + std::cout << "Runtime (ms): " << runtime_ms << std::endl; + std::cout << "GFLOPs/sec: " << gflops << std::endl; + + // Cleanup + for (auto event : events) { + (void)cudaEventDestroy(event); + } + + return passed; + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +int main(int argc, char const** argv) { + + cudaDeviceProp props; + + cudaError_t error = cudaGetDeviceProperties(&props, 0); + if (error != cudaSuccess) { + std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl; + return -1; + } + + if (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 4) || + (props.major != 8 && props.minor != 9)) { + + // + // This example requires an NVIDIA Ada-architecture GPU. + // + + std::cout + << "CUTLASS's FP8 SM89 example requires a GPU of NVIDIA's Ada architecture " + << "and CUDA toolkit version 12.4 or later.\n"; + + return 0; + } + + // + // Parse options + // + + Options options; + + options.parse(argc, argv); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + if (options.error) { + std::cerr << "Aborting execution." << std::endl; + return -1; + } + + std::cout << "Running GEMM with staged accumulation (OpMultiplyAdd)" << std::endl; + std::cout << "=====================================================" << std::endl; + TestbedRunner> testbed_staged_accum; + bool passed = testbed_staged_accum.run(options); + + if (passed) { + std::cout << "Passed" << std::endl; + } else { + std::cout << "Failed" << std::endl; + } + + std::cout << "\nRunning GEMM with fast accumulation (OpMultiplyAddFastAccum)" << std::endl; + std::cout << "============================================================" << std::endl; + TestbedRunner> testbed_fast_accum; + passed = testbed_fast_accum.run(options); + + if (passed) { + std::cout << "Passed" << std::endl; + } else { + std::cout << "Failed" << std::endl; + } + + return 0; +} diff --git a/examples/59_ampere_gather_scatter_conv/CMakeLists.txt b/examples/59_ampere_gather_scatter_conv/CMakeLists.txt new file mode 100644 index 00000000..e7f16400 --- /dev/null +++ b/examples/59_ampere_gather_scatter_conv/CMakeLists.txt @@ -0,0 +1,36 @@ +# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +cutlass_example_add_executable( + 59_ampere_gather_scatter_conv + ampere_gather_scatter_conv.cu +) + +if (CUTLASS_ENABLE_OPENMP_TESTS AND OpenMP_CXX_FOUND) + target_link_libraries(59_ampere_gather_scatter_conv PRIVATE OpenMP::OpenMP_CXX) +endif() diff --git a/examples/59_ampere_gather_scatter_conv/README.md b/examples/59_ampere_gather_scatter_conv/README.md new file mode 100644 index 00000000..4aac0536 --- /dev/null +++ b/examples/59_ampere_gather_scatter_conv/README.md @@ -0,0 +1,209 @@ +# Example 59: Ampere gather/scatter convolution + +CuTe and CUTLASS 3.x based Ampere convolution forward propagation kernel capable of operating on both affine and gather/scatter tensors. + +Example executions: +```sh +./59_ampere_gather_scatter_conv +./59_ampere_gather_scatter_conv --n=108 +./59_ampere_gather_scatter_conv --n=4096 --i=1 +./59_ampere_gather_scatter_conv --n=1080 --i=1000 +./59_ampere_gather_scatter_conv --n=131072 --i=1000 --no-check +``` + +This example demonstrates a few super cool features of CUTLASS and CuTe. It shows off +1. A dense conv 3D fprop kernel written as a single file ... +2. ... that leverages off-the-shelf CUTLASS collectives to show how custom kernels can use collectives ... +3. ... and uses the exact same templated kernel to also stamp out a gather/scatter 3D fprop conv ... +4. ... while getting near peak performance of the Ampere class tensor core on Ampere and Ada GPUs ... +5. ... by using static cute shapes and strides in case problem shapes are known at compile time. + +## A dense conv 3D fprop kernel written in CUTLASS 3.x and CuTe + +The most common strategy for implementing high performance convolution kernels on the GPU is to transform +the activation tensor in such a way that we can perform the computation as a GEMM. This is called the +image to column (im2col) transformation. [CUTLASS 2.x implementation of im2col based convolutions is +documented separately](../../media/docs/implicit_gemm_convolution.md), and here we consider a fresh approach for CuTe. + +A 3D convolution has the following input tensors: +- Activation tensor (Act): `((N,(D,H,W)), (C,(1,1,1)))` +- Filter tensor (Flt): `( K, (C,(T,R,S)))` +- Output tensor (Out): `((N,(Z,P,Q)), K )` + +Where +- N := number of images +- DHW := spatial dimensions of the activation tensor +- C := channel dimension of the activation tensor +- K := channel dimension of the filter and output tensor +- TRS := spoke dimensions of the filter tensor +- ZPQ := spatial dimensions of the output tensor + +As is evident in the tensor shapes, these cannot be issued to a GEMM just yet, since there is no +logical M, N, and K modes we can group the tensor modes into. + +Notice that every spoke of the filter tensor (TRS) will be applied to some (offset) view of the +activation tensor, thus expanding the logical size of the activation tensor. +Additionally, a similar logical transform of the spatial dimensions can be encoded as a function of the +padding, dilations, traversal strides, and filter spokes. This gets us to our im2col transform: + +im2col transform affects the component shapes/strides of the activation tensor in the following way: +- ZPQ Shape : changes DHW domain with formula `(1 + (DHW + pad - (((TRS-1) * dilation) + 1)) / traversal_stride)` +- TRS Shape : TRS domain instead of `(1,1,1)` +- ZPQ Strides : Original DHW strides get `elem_scale()`-ed by traversal strides DHW +- TRS Strides : Original DHW strides get `elem_scale()`-ed by dilation DHW + +With this transform applied, we end up with a set of input and output tensors that +are logically consistent in their MNK dimensions, thus allowing us to dispatch to a GEMM. +im2col activation layout: ((N,(Z,P,Q)), (C,(T,R,S))) // logical (M,K) +filter layout : ( K, (C,(T,R,S))) // logical (N,K) +output layout : ((N,(Z,P,Q)), K ) // logical (M,N) + +CuTe's layout representation and algebra make these folded tensors easy to represent and manipulate. +This is most evident in the reference check code used in this example: + +```cpp +for (size_t logical_m = 0; logical_m < size<0>(mOutputRef); ++logical_m) { + for (size_t logical_n = 0; logical_n < size<1>(mOutputRef); ++logical_n) { + auto accumulator = float(0); + for (size_t logical_k = 0; logical_k < size<1>(mStencil); ++logical_k) { + accumulator += mStencil(logical_m, logical_k) * mActivation(logical_n, logical_k); + } + mOutputRef(logical_m, logical_n) = accumulator; + } +} +``` + +Which succinctly demonstrates how im2col transform allows us to implement convolutions +as GEMMs with special layout transformations on the input tensor. + +Note: in the example kernel's implementation we treat activations as the B tensor +and filter as the A tensor, thus making their logical dimensions NK and MK respectively. + +## Leveraging CUTLASS collectives off the shelf in a custom kernel + +Now that we have transformed our problem in such a way that allows us to dispatch to a GEMM, +we can reuse much of the machinery CUTLASS offers to implement this forward pass convolution +operator. CUTLASS decomposes these "moving parts" of GPU linear algebra into reusable, +modular software components abstracted by C++ template classes. This example +demonstrates how some of the lower layers of the hierarchy can be re-used for custom kernels +by writing a custom kernel for convolution that re-uses the Ampere/Ada GEMM collectives +from CUTLASS 3. + +A kernel author is free to compose their custom components with any of the existing templates +in the CUTLASS hierarchy to leverage existing high performance implementations from the CUTLASS +team. In this example, we write a custom kernel layer and compose with an existing collective. +However, any of the CUTLASS kernels can be composed with bespoke collectives if the desired +customization is a mainloop or epilogue fusion without changes to the grid planning, +tile scheduling, load balancing, or thread marshalling. + +## Implementing gather/scatter and dense convolution with the same kernel + +Functionality and correctness of the implemented kernel, as a virtue of using +CuTe and off the shelf CUTLASS collectives, only relies on the logical consistency of +the layouts of input and output tensors. This means that we can freely change how +the logical coordinates of the tensors map into the index space, and even how these dereferences +happen. [CUTLASS example 52](../52_hopper_gather_scatter_fusion/) demonstrates this by implementing a custom stride that +supports indexed indirection for tensor data accesses. This allows for example 52 +to implement a GEMM where inputs are gathered and output is scattered based on an index buffer. + +We re-use the same custom stride utilities in this example to implement a convolution kernel +that gathers along the NDHW dimensions of the activation tensor and scatters the output along the +NZPQ dimensions of the output tensor, treating the channel dimensions as the dense vectors. + +Our dense affine im2col transformed activation tensor: + +```cpp +// im2col transformed activation layout: ((nzpq), (ctrs)) => idx +auto xformed_act_layout = make_layout( + make_shape (make_shape ( N, Z, P, Q), make_shape ( C, T, R, S)), + make_stride(make_stride(D*H*W*C, H*W*C, W*C, C), make_stride(_1{}, H*W*C, W*C, C))); +``` + +now becomes a composed layout that uses `IndexedGather`: + +```cpp +// Inner layout of the composition: +// ((nzpq), (csrt)) => (idx_buffer_idx, dense_offset) +auto EG = E<0>{}; // Gather basis (1,0) (idx_buffer_idx) +auto EC = E<1>{}; // Contiguous basis (0,1) (dense_offset) +auto xformed_act_logical_inner = make_layout( + make_shape (make_shape ( N, Z, P, Q), make_shape ( C, T, R, S)), + make_stride(make_stride(D*H*W*EG, H*W*EG, W*EG, EG), make_stride(EC, H*W*EG, W*EG, EG))); + +// Outer layout of the composition: +// (idx_buffer_idx, dense_offset) => idx +// IndexedGather obtains idx by applying (gmem_base_ptr + gather_idx_buf[idx_buffer_idx] + dense_offset) +auto xformed_act_gather_outer = make_layout( + make_shape(_1{},_1{}), + make_stride(CustomStride{IndexedGather{gather_idx_buf}, C}, _1{})); + +// Compose the inner and outer layouts +// ((nzpq), (ctrs)) => idx +auto xformed_act_composed_layout = composition( + xformed_act_gather_outer, + make_arithmetic_tuple(_0{}, _0{}), + xformed_act_logical_inner); +``` + +Here, we create a composed layout whose inner layout has the same logical MK shape as earlier, +but with an outer layout that uses the custom strides with an index buffer to access memory with +indirections. A custom stride requires two inputs to compute the index that a certain coordinate maps to: +the index buffer offset and the dense offset into the vector. This entails that our inner layout +(the one with the logical MK shape) has a rank-2 codomain `(idx_buffer_idx, dense_offset)`. +We can set up such a layout with scaled basis strides, which allow us to map a domain onto a +codomain with multiple orthogonal bases. The two codomain basis are the +index buffer offsets (rank 0 basis), and the dense vector offsets (rank 1 basis). +A similar composed layout is set up for the output scatter tensor. + +This tensor still has a logical MK shape and is backed by a CuTe layout, which means we can still +tile, partition, and otherwise manipulate it with CuTe's layout algebra in the same way we would any +other tensor. Substituting the activation tensor's affine layout for this gather layout requires +no changes to the implementation of the kernel whatsoever. Everything composes. This example +stamps out a dense 3D convolution as well as gather/scatter 3D convolution using the same kernel template, +with the only difference between them being the layouts of the input and output tensors. + +Convolutions are just a special case of tensor contractions, and as [example 51](../51_hopper_gett) +demonstrates, the exact same collective used in this example can also be used to implement arbitrary GETTs. +Of course, this also means that the same kernel can implement gather/scatter GETTs as well! + +This demonstrates the composition power of not just CuTe, but also CUTLASS 3's two level +micro kernel abstraction. A single highly tuned temporal micro-kernel (collective) can be implemented once +and applied to compute dense GETTs, gather/scatter GETTs, dense convolutions, and gather/scatter convolutions. + +## Peak performance on Ampere and Ada GPUs by leveraging domain specific knowledge + +Often, when implementing custom kernels, a user has more knowledge of the problem domain that can be +exploited to deliver higher performance than otherwise could be through general kernels. In this example +we presume that the shape of each of the images (DHWC dimensions) as well as the filter (TRS) are available +a-priori and that the tile shape evenly divides the problem. Number of images (N) is still left as a runtime +parameter. + +Knowing the extents of our tensors at compile time allows us to encode them as static cute shapes rather than +a dynamic problem shape, resulting in the elimination of most of the index computation instructions such as +expensive div/mods. Knowing that the problem shape is divisible by the tile shape allows us to use the +Ampere collective that does not perform predication on global memory loads, further reducing overheads +and allowing us to achieve near peak performance on RTX Ampere and Ada GPUs. + +Running this example on an RTX 3080Ti prints the following performance numbers (some output culled for brevity): + +``` +$> ./examples/59_ampere_gather_scatter_conv/59_ampere_gather_scatter_conv --n=131072 --i=128 --no-check +Ampere convolution forward propogation kernel supporting both affine and gather/scatter tensors. + +Allocating tensors ... done. +Initializing data ... done. +Initializing gather/scatter index buffers ... done. + +Running dense fprop kernel +Conv TFLOP count = 0.927713 +Conv dense perf: 31.027376ms | TFLOP/s = 29.899819 + +Running gather/scatter fprop kernel +Conv TFLOP count = 0.927713 +Conv gather/scatter perf: 28.973721ms | TFLOP/s = 32.019117 +``` + +With this in mind, this example kernel has the following limitations: +- This example kernel only supports dynamic image count, all other conv problem shape must be defined as `cute::Constant<>`s +- Problem shapes (including dynamic image count `N`) must be evenly divisible by the tile shape +- It does not perform fp32->tf32 numeric conversion, gmem inputs must be rounded to tf32 already diff --git a/examples/59_ampere_gather_scatter_conv/ampere_conv_kernel.h b/examples/59_ampere_gather_scatter_conv/ampere_conv_kernel.h new file mode 100644 index 00000000..cc00cced --- /dev/null +++ b/examples/59_ampere_gather_scatter_conv/ampere_conv_kernel.h @@ -0,0 +1,320 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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/tensor.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/atom/copy_atom.hpp" +#include + +#include "cutlass/util/print_error.hpp" + +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/collective/collective_mma.hpp" + +using namespace cute; + +struct AmpereUnpredicatedFprop { + // + // Static config for conv problem shape + // + using D = _6; + using H = _4; + using W = _4; + + using T = _3; + using R = _3; + using S = _3; + + using Z = _4; + using P = _2; + using Q = _2; + + using C = _64; + using K = _128; + + // Tiler config + using Tiler_K = decltype(cute::min(K{}, _128{})); + using Tiler_C = decltype(cute::min(C{}, _32{})); + using Tiler_N = _4; + using TileM = Tiler_K; + using TileN = Shape; + using TileK = Shape; + using PIPE = _3; + using TilerFlt = Shape; + using TilerAct = Shape; + using TilerOut = Shape; + + using TileSizeM = Int; + using TileSizeN = Int; + using TileSizeK = Int; + static constexpr int Stages = PIPE::value; + + using ElementFlt = tfloat32_t; + using ElementAct = tfloat32_t; + using ElementOut = float; + + using TiledMma = TiledMMA< + MMA_Atom, + Layout>, + Tile<_32,_32,Underscore>>; + + static constexpr int MaxThreadsPerBlock = size(TiledMma{}); + static constexpr int MinBlocksPerMultiprocessor = 1; + + union SharedStorage { + struct { + ElementFlt sAMatrix[size(TileM{}) * size(TileK{}) * size(PIPE{})]; + ElementAct sBMatrix[size(TileN{}) * size(TileK{}) * size(PIPE{})]; + } mainloop; + + struct { + ElementOut sCMatrix[size(TileM{}) * size(TileN{})]; + } epilogue; + }; + + // + // Stencil tensor + // + + using GmemLayoutFlt = decltype(make_ordered_layout( + Shape< K, Shape< C, T, R, S>>{}, + tuple<_4, tuple<_0,_3,_2,_1>>{})); + + // We have 64 elements * 32b each in the major mode that we can vectorize + // Max vector size is 128b, so lay 16 threads along the major mode with a vector size of 4 + // Rest along the minor mode + using GmemTiledCopyFlt = decltype(make_tiled_copy( + Copy_Atom, ElementFlt>{}, + Layout, + Stride< _8, _1>>{}, + Layout>{})); + + // Following layout is also correct, but trades off dynamic strides in the slice for bank conflict free accesses + // using SmemLayoutFlt = decltype( + // composition(Swizzle<3,2,3>{}, + // make_ordered_layout( + // Shape{}, + // tuple< _1, _0, _2>{}))); + + using SmemLayoutAtomFlt = decltype( + composition(Swizzle<1,2,3>{}, + Layout>, + Stride<_4,Stride<_1,_32>>>{})); + + using SmemCopyAtomFlt = Copy_Atom; + + // + // Activation tensor + // + + // Activation tensor is major in the contraction mode, so vectorize that mode first + // Then lay out the rest of the threads along the other mode + using GmemTiledCopyAct = decltype(make_tiled_copy( + Copy_Atom, ElementAct>{}, + Layout, + Stride< _8, _1>>{}, + Layout>{})); + + // Following layout is also correct, but trades off dynamic strides in the slice for bank conflict free accesses + // using SmemLayoutAct = decltype( + // composition(Swizzle<3,2,3>{}, + // make_ordered_layout( + // Shape{}, + // tuple< _1, _0, _2>{}))); + + using SmemLayoutAtomAct = decltype( + composition(Swizzle<1,2,3>{}, + Layout>, + Stride<_4,Stride<_1,_32>>>{})); + + using SmemCopyAtomAct = Copy_Atom; + + // + // Output tensor + // + + using GmemTiledCopyOut = decltype(make_tiled_copy( + Copy_Atom, ElementAct>{}, + Layout, + Stride<_1, _8>>{}, + Layout>{})); + + using SmemCopyAtomOut = Copy_Atom, ElementOut>; + + // This can be optimized to make accesses BCF, but we use a col-major layout here to show off composability + using SmemLayoutOut = Layout>; + + // + // Conv functor + // + template + void __device__ + operator()(cute::Tensor mFlt, // ( K, (C,T,R,S)) + TensorActivation mAct, // ((N,Z,P,Q), (C,T,R,S)) + TensorOutput mOut, // ( K, (N,Z,P,Q)) + char* smem_buf) const { + using namespace cute; + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveMma< + cutlass::gemm::MainloopSm80CpAsyncUnpredicated, + Shape, + ElementFlt, + Underscore, // Ignore the stride, we are passing full cute::Tensor to operator() + ElementAct, + Underscore, // Ignore the stride, we are passing full cute::Tensor to operator() + TiledMma, + GmemTiledCopyFlt, + SmemLayoutAtomFlt, + SmemCopyAtomFlt, + cute::identity, + GmemTiledCopyAct, + SmemLayoutAtomAct, + SmemCopyAtomAct, + cute::identity>; + + TiledMma tiled_mma; + Tensor accum = partition_fragment_C(tiled_mma, TilerOut{}); + clear(accum); + + // Set up tensors + // NOTE: blockIdx.x projects onto act-NDHW mode, y along the flt-K mode for the sake of higher dynamic range in NDHW + Tensor gA_mk = local_tile(mFlt, TilerFlt{}, make_coord(_,_)); // (BLK_M,BLK_K,m',k') + Tensor gB_nk = local_tile(mAct, TilerAct{}, make_coord(_,_)); // (BLK_N,BLK_K,n',_1) + Tensor gC_mn = local_tile(mOut, TilerOut{}, make_coord(_,_)); // (BLK_M,BLK_N,m',n') + + // Compute m_coord and n_coord with their post-tiled shapes + auto m_coord = idx2crd(int(blockIdx.y), shape<2>(gA_mk)); + auto n_coord = idx2crd(int(blockIdx.x), shape<2>(gB_nk)); + Tensor gA = gA_mk(_,_,m_coord,_); // (BLK_M,BLK_K,k') + Tensor gB = gB_nk(_,_,n_coord,_); // (BLK_N,BLK_K,_1) + Tensor gC = gC_mn(_,_,m_coord,n_coord); // (BLK_M,BLK_N) + + auto k_tile_iter = cute::make_coord_iterator(size<2>(gA)); + int k_tile_count = size<2>(gA); + + CollectiveMainloop collective_mma; + collective_mma( + accum, + gA, + gB, + accum, + k_tile_iter, k_tile_count, + Underscore{}, // no residue since we do not support predication + threadIdx.x, + smem_buf); + + // + // Epilogue + // + SharedStorage& storage = *reinterpret_cast(smem_buf); + Tensor sC = make_tensor(make_smem_ptr(&storage.epilogue.sCMatrix[0]), SmemLayoutOut{}); + + auto smem_tiled_copy_C = make_tiled_copy_C(SmemCopyAtomOut{}, tiled_mma); + auto smem_thr_copy_C = smem_tiled_copy_C.get_slice(threadIdx.x); + auto tCrC = smem_thr_copy_C.retile_S(accum); + auto tCsC = smem_thr_copy_C.partition_D(sC); + copy(smem_tiled_copy_C, tCrC, tCsC); + + __syncthreads(); + + GmemTiledCopyOut gmem_tiled_copy_C; + auto gmem_thr_copy_C = gmem_tiled_copy_C.get_slice(threadIdx.x); + auto tDsC = gmem_thr_copy_C.partition_S(sC); + auto tDgC = gmem_thr_copy_C.partition_D(gC); + copy(gmem_tiled_copy_C, tDsC, tDgC); + + #if 0 + if (thread0()) { + print("mAct = "); print(mAct); print('\n'); + print("mFlt = "); print(mFlt); print('\n'); + print("mOut = "); print(mOut); print('\n'); + print("gA = "); print(gA); print('\n'); + print("gB = "); print(gB); print('\n'); + print("gC = "); print(gC); print('\n'); + print("sA = "); print(sA.layout()); print('\n'); + print("sB = "); print(sB.layout()); print('\n'); + print("sC = "); print(sC.layout()); print('\n'); + print("tAgA = "); print(tAgA.layout()); print('\n'); + print("tBgB = "); print(tBgB.layout()); print('\n'); + print("tAsA = "); print(tAsA.layout()); print('\n'); + print("tBsB = "); print(tBsB.layout()); print('\n'); + print("tCsA = "); print(tCsA.layout()); print('\n'); + print("tCsB = "); print(tCsB.layout()); print('\n'); + print("tCrC = "); print(tCrC.layout()); print('\n'); + print("tCsC = "); print(tCsC.layout()); print('\n'); + print("tDsC = "); print(tDsC.layout()); print('\n'); + print("tDgC = "); print(tDgC.layout()); print('\n'); + print("gmem tiled copy A = "); print(gmem_tiled_copy_A); print('\n'); + print("gmem tiled copy B = "); print(gmem_tiled_copy_B); print('\n'); + print("gmem tiled copy C = "); print(gmem_tiled_copy_C); print('\n'); + print("k_tile_count = "); print(size<2>(gA)); print('\n'); + print("k_tile_iter = "); print(*k_tile_iter); print('\n'); + print("K_BLOCK_MAX = "); print(K_BLOCK_MAX); print('\n'); + } + #endif + } +}; + +template +inline int +fprop_reference( + TensorFlt mStencil, // Logical MK: ( K, (C,T,R,S)) + TensorAct mActivation, // Logical NK: ((N,Z,P,Q), (C,T,R,S)) + TensorOut mOutput, // Logical MN: ( K, (N,Z,P,Q)) + TensorOut mOutputRef) { + int32_t N = size<1,0>(mOutputRef); + int32_t Z = size<1,1>(mOutputRef); + int32_t P = size<1,2>(mOutputRef); + int32_t Q = size<1,3>(mOutputRef); + int32_t T = size<1,3>(mStencil); + int32_t R = size<1,2>(mStencil); + int32_t S = size<1,1>(mStencil); + int32_t C = size<1,0>(mStencil); + + size_t K = static_cast(size<0>(mOutputRef)); + size_t NZPQ = static_cast(size<1>(mOutputRef)); + size_t CTRS = static_cast(size<1>(mStencil)); + +#if defined(_OPENMP) + #pragma omp parallel for +#endif + for (size_t logical_m = 0; logical_m < K; ++logical_m) { + for (size_t logical_n = 0; logical_n < NZPQ; ++logical_n) { + auto accumulator = float(0); + for (size_t logical_k = 0; logical_k < CTRS; ++logical_k) { + accumulator += mStencil(logical_m, logical_k) * mActivation(logical_n, logical_k); + } + mOutputRef(logical_m, logical_n) = accumulator; + } + } + + return print_relative_error(mOutput, mOutputRef, /*print_verbose*/ false, /*print_error*/ true, /*error_margin*/ 0.01); +} diff --git a/examples/59_ampere_gather_scatter_conv/ampere_gather_scatter_conv.cu b/examples/59_ampere_gather_scatter_conv/ampere_gather_scatter_conv.cu new file mode 100644 index 00000000..341d1e9f --- /dev/null +++ b/examples/59_ampere_gather_scatter_conv/ampere_gather_scatter_conv.cu @@ -0,0 +1,392 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Example demonstrating CuTe and CUTLASS 3.x based Ampere convolution forward propogation kernel + capable of operating on both affine and gather/scatter tensors. + + This example demonstartes a few super cool features of CUTLASS and CuTe. It shows off + 1. A dense conv 3D fprop kernel written as a single file ... + 2. ... that leverages off the shelf CUTLASS collectives to show how custom kernels can use collectives ... + 3. ... and uses the exact same templated kernel to also stamp out a gather/scatter 3D fprop conv ... + 4. ... while getting near peak performance of the Ampere class tensor core on Ampere and Ada GPUs ... + 5. ... by using static cute shapes and strides in case problem shapes are known at compile time. + + Full documentation for this example can be found within the README.md file in this directory. + + Example executions: + ./59_ampere_gather_scatter_conv + ./59_ampere_gather_scatter_conv --n=108 + ./59_ampere_gather_scatter_conv --n=4096 --i=1 + ./59_ampere_gather_scatter_conv --n=1080 --i=1000 + ./59_ampere_gather_scatter_conv --n=131072 --i=1000 --no-check +*/ + +#include +#include + +#include "ampere_conv_kernel.h" +#include "gather_tensor.hpp" + +#include "cutlass/util/command_line.h" + +bool check_cuda_result(cudaError_t code, const char* file, int line) { + if (code == cudaSuccess) { + return true; + } + + std::cerr << "CUDA error at (" << file << "," << line << ")\n\t" << unsigned(code) << " -- " << cudaGetErrorString(code) << "\n"; + return false; +} + +#define CHECK_CUDA(code) (check_cuda_result(code, __FILE__, __LINE__)) + +using namespace cute; +using example::IndexedGather; +using example::CustomStride; + +template +__global__ +__launch_bounds__(Operator::MaxThreadsPerBlock, Operator::MinBlocksPerMultiprocessor) +void kernel_entrypoint(FilterTensor mFlt, ActivationTensor mAct, OutputTensor mOut) { + extern __shared__ char smem_buf[]; + Operator op; + op(mFlt, mAct, mOut, smem_buf); +} + +int ampere_dense_conv_fprop( + int num_images, + float* activations, + float* filter, + float* output, + float* output_ref, + int num_iterations = 1, + bool do_ref_check = true) { + auto D = typename AmpereUnpredicatedFprop::D{}; + auto H = typename AmpereUnpredicatedFprop::H{}; + auto W = typename AmpereUnpredicatedFprop::W{}; + auto Z = typename AmpereUnpredicatedFprop::Z{}; + auto P = typename AmpereUnpredicatedFprop::P{}; + auto Q = typename AmpereUnpredicatedFprop::Q{}; + auto C = typename AmpereUnpredicatedFprop::C{}; + auto K = typename AmpereUnpredicatedFprop::K{}; + auto S = typename AmpereUnpredicatedFprop::S{}; + auto R = typename AmpereUnpredicatedFprop::R{}; + auto T = typename AmpereUnpredicatedFprop::T{}; + + int N = num_images; // dynamic + if (num_images % int(typename AmpereUnpredicatedFprop::Tiler_N{}) != 0) { + printf("ERROR: Input image count must be evenly divisible by CTA tiler N.\n"); + return 1; + } + + // Tensor Activation: (n,d,h,w,c)::(?,6,4,4,64):(6144,1536,384,64,1) + auto activation_layout = make_layout( + make_shape (make_shape ( N, D, H, W), make_shape ( C, _1{},_1{},_1{})), + make_stride(make_stride(D*H*W*C, H*W*C, W*C, C), make_stride(_1{}, _0{},_0{},_0{}))); + + auto xformed_act_layout = make_layout( + make_shape (make_shape(N, Z, P, Q), make_shape ( C, T, R, S)), + make_stride(stride<0>(activation_layout), make_stride(_1{}, H*W*C, W*C, C))); + + // Tensor Filter : (k,c,s,r,t)::(128,3,3,3,64):(1728,576,192,64,1) + auto filter_layout = AmpereUnpredicatedFprop::GmemLayoutFlt{}; + + // Tensor Output : (n,z,p,q,k)::(?,4,2,2,128):(2048,1024,512,128,1) + auto output_layout = make_ordered_layout( + make_shape( K, make_shape( N, Z, P, Q)), + make_tuple(_0{}, make_tuple(_4{},_3{},_2{},_1{}))); + + Tensor mActivation = make_tensor(make_gmem_ptr(activations), activation_layout); + Tensor mXformedAct = make_tensor(make_gmem_ptr(activations), xformed_act_layout); + Tensor mFilter = make_tensor(make_gmem_ptr(filter), filter_layout); + Tensor mOutput = make_tensor(make_gmem_ptr(output), output_layout); // (K, (N,Z,P,Q)) + Tensor mOutputRef = make_tensor(make_gmem_ptr(output_ref), output_layout); + + print("xformed act layout ((N,Z,P,Q), (C,T,R,S)) = "); print(xformed_act_layout); print("\n"); + + cudaEvent_t start, stop; + CHECK_CUDA(cudaEventCreate(&start)); + CHECK_CUDA(cudaEventCreate(&stop)); + + constexpr size_t smem_size = sizeof(typename AmpereUnpredicatedFprop::SharedStorage); + Tensor gOutput_mn = zipped_divide(mOutput, typename AmpereUnpredicatedFprop::TilerOut{}); // ((BLK_M, BLK_N), (m', n')) + dim3 lauch_grid {static_cast(size<1,1>(gOutput_mn)), static_cast(size<1,0>(gOutput_mn)), 1}; + + CHECK_CUDA(cudaFuncSetAttribute( + kernel_entrypoint, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size)); + + CHECK_CUDA(cudaEventRecord(start)); + for (int i = 0; i < num_iterations; ++i) { + kernel_entrypoint + <<>>( + mFilter, mXformedAct, mOutput); + } + CHECK_CUDA(cudaEventRecord(stop)); + CHECK_CUDA(cudaEventSynchronize(stop)); + + float milliseconds = 0; + cudaEventElapsedTime(&milliseconds, start, stop); + milliseconds /= float(num_iterations); + + double tflop_count = (2 * double(size<0>(xformed_act_layout)) * double(size(filter_layout))) / double(1e12); + double tflops = tflop_count / (double(milliseconds) / double(1e3)); + + printf("Conv TFLOP count = %f\n", tflop_count); + printf("Conv dense perf: %fms | TFLOP/s = %f\n", milliseconds, tflops); + + if (do_ref_check) { + printf("Running host reference check ...\n"); + return fprop_reference(mFilter, mXformedAct, mOutput, mOutputRef); + } + else { + return 0; + } +} + +int ampere_gather_scatter_conv_fprop( + int num_images, + float* activations, + uint32_t *gather_idx_buf, + float* filter, + float* output, + uint32_t *scatter_idx_buf, + int num_iterations = 1) { + auto D = typename AmpereUnpredicatedFprop::D{}; + auto H = typename AmpereUnpredicatedFprop::H{}; + auto W = typename AmpereUnpredicatedFprop::W{}; + auto Z = typename AmpereUnpredicatedFprop::Z{}; + auto P = typename AmpereUnpredicatedFprop::P{}; + auto Q = typename AmpereUnpredicatedFprop::Q{}; + auto C = typename AmpereUnpredicatedFprop::C{}; + auto K = typename AmpereUnpredicatedFprop::K{}; + auto S = typename AmpereUnpredicatedFprop::S{}; + auto R = typename AmpereUnpredicatedFprop::R{}; + auto T = typename AmpereUnpredicatedFprop::T{}; + + int N = num_images; // dynamic + if (N % int(typename AmpereUnpredicatedFprop::Tiler_N{}) != 0) { + printf("ERROR: Input image count must be evenly divisible by CTA tiler N. Got num_images = %d\n", N); + return 1; + } + + // Tensor Filter : (k,c,s,r,t)::(128,3,3,3,64):(1728,576,192,64,1) + auto filter_layout = AmpereUnpredicatedFprop::GmemLayoutFlt{}; + + // Tensor Output : (n,z,p,q,k)::(?,4,2,2,128):(2048,1024,512,128,1) + auto output_layout = make_ordered_layout( + make_shape( K, make_shape( N, Z, P, Q)), + make_tuple(_0{}, make_tuple(_4{},_3{},_2{},_1{}))); + + // Input gather layout + // inner_layout(make_coord((nzpq), (csrt))) => (idx_buffer_idx, dense_c_idx) + auto EG = E<0>{}; // Gather basis (1,0) (idx_buffer_idx) + auto EC = E<1>{}; // Contiguous basis (0,1) (dense_offset) + auto xformed_act_logical_inner = make_layout( + make_shape (make_shape ( N, Z, P, Q), make_shape ( C, T, R, S)), + make_stride(make_stride(D*H*W*EG, H*W*EG, W*EG, EG), make_stride(EC, H*W*EG, W*EG, EG))); + + // outer_layout(make_coord(idx_buffer_idx, dense_c_idx)) => idx + // IndexedGather obtains idx by applying (gmem_base_ptr + gather_idx_buf[idx_buffer_idx] + dense_offset) + auto xformed_act_gather_outer = make_layout( + make_shape(_1{},_1{}), + make_stride(CustomStride{IndexedGather{gather_idx_buf}, C}, _1{})); + + // Compose the inner and outer layouts + // gather_composed(make_coord((nzpq), (csrt))) => idx + auto xformed_act_composed_layout = composition( + xformed_act_gather_outer, + make_arithmetic_tuple(_0{}, _0{}), + xformed_act_logical_inner); + + // Output scatter layout + auto out_basis_stride = make_stride( + E<1>{}, + make_stride(Z*P*Q*E<0>{}, P*Q*E<0>{}, Q*E<0>{}, _1{}*E<0>{})); // -> (crd0, crd1) + auto out_basis_layout = make_layout(shape(output_layout), out_basis_stride); + auto out_scatter_layout = make_layout( + make_shape(_1{},_1{}), + make_stride(CustomStride{IndexedGather{scatter_idx_buf}, K}, _1{})); + auto out_composed_layout = composition( + out_scatter_layout, + make_arithmetic_tuple(_0{},_0{}), + out_basis_layout); + + Tensor mXformedActGather = make_tensor(make_gmem_ptr(activations), xformed_act_composed_layout); + Tensor mFilter = make_tensor(make_gmem_ptr(filter), filter_layout); + Tensor mOutputScatter = make_tensor(make_gmem_ptr(output), out_composed_layout); // (K, (N,Z,P,Q)) + + Tensor gOutput_mn = zipped_divide(mOutputScatter, typename AmpereUnpredicatedFprop::TilerOut{}); // ((BLK_M, BLK_N), (m', n')) + dim3 lauch_grid {static_cast(size<1,1>(gOutput_mn)), static_cast(size<1,0>(gOutput_mn)), 1}; + constexpr size_t smem_size = sizeof(typename AmpereUnpredicatedFprop::SharedStorage); + + print("xforemed gather layout ((N,Z,P,Q), (C,T,R,S)) = "); print(xformed_act_composed_layout); print("\n"); + print("Output scatter layout ( K, (N,Z,P,Q)) = "); print(out_composed_layout); print("\n"); + print("Filter layout ( K, (C,T,R,S)) = "); print(filter_layout); print("\n"); + + CHECK_CUDA(cudaFuncSetAttribute( + kernel_entrypoint, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size)); + + cudaEvent_t start, stop; + CHECK_CUDA(cudaEventCreate(&start)); + CHECK_CUDA(cudaEventCreate(&stop)); + CHECK_CUDA(cudaEventRecord(start)); + for (int i = 0; i < num_iterations; ++i) { + kernel_entrypoint + <<>>( + mFilter, mXformedActGather, mOutputScatter); + } + CHECK_CUDA(cudaEventRecord(stop)); + CHECK_CUDA(cudaEventSynchronize(stop)); + float milliseconds = 0; + cudaEventElapsedTime(&milliseconds, start, stop); + milliseconds /= float(num_iterations); + + double tflop_count = (2 * double(size<0>(xformed_act_logical_inner)) * double(size(filter_layout))) / double(1e12); + double tflops = tflop_count / (double(milliseconds) / double(1e3)); + printf("Conv TFLOP count = %f\n", tflop_count); + printf("Conv gather/scatter perf: %fms | TFLOP/s = %f\n", milliseconds, tflops); + + return 0; +} + +int +main(int argc, char const** argv) { + cutlass::CommandLine cmd(argc, argv); + std::cout << "Ampere convolution forward propogation kernel supporting both affine and gather/scatter tensors.\n\n"; + if (cmd.check_cmd_line_flag("help")) { + std::cout + << "Options:\n" + "\t--n= Sets the number of images for the input activation tensor (dataset size). Default = 131072.\n" + "\t--i= Sets the benchmarking repetitions. Default = 128.\n" + "\t--nocheck If specified, skips the reference check for dense kernel.\n" + "\t--help Displays this help message and exits.\n"; + return 0; + } + + + cudaDeviceProp props; + cudaError_t error = cudaGetDeviceProperties(&props, 0); + if (error != cudaSuccess) { + std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl; + return -1; + } + if (props.major < 8) { + std::cerr << "This example requires an Ampere GPU or newer.\n"; + return 0; + } + + int num_images = 4320; + cmd.get_cmd_line_argument("n", num_images, 4320); + int num_iterations = 128; + cmd.get_cmd_line_argument("i", num_iterations, 128); + bool do_host_ref_check = not cmd.check_cmd_line_flag("no-check"); + + auto D = typename AmpereUnpredicatedFprop::D{}; + auto H = typename AmpereUnpredicatedFprop::H{}; + auto W = typename AmpereUnpredicatedFprop::W{}; + auto Z = typename AmpereUnpredicatedFprop::Z{}; + auto P = typename AmpereUnpredicatedFprop::P{}; + auto Q = typename AmpereUnpredicatedFprop::Q{}; + auto C = typename AmpereUnpredicatedFprop::C{}; + auto K = typename AmpereUnpredicatedFprop::K{}; + + auto activation_layout = make_layout( + make_shape (make_shape (num_images, D, H, W), make_shape ( C, _1{},_1{},_1{})), + make_stride(make_stride( D*H*W*C, H*W*C, W*C, C), make_stride(_1{}, _0{},_0{},_0{}))); + + auto filter_layout = typename AmpereUnpredicatedFprop::GmemLayoutFlt{}; + + auto output_layout = make_ordered_layout( + make_shape( K, make_shape(num_images, Z, P, Q)), + make_step (_0{}, make_step ( _4{},_3{},_2{},_1{}))); + + print("Filter layout ( K, (C,T,R,S)) = "); print(filter_layout); print("\n"); + print("Activation layout ((N,D,H,W), (C,1,1,1)) = "); print(activation_layout); print("\n"); + print("Output layout ( K, (N,Z,P,Q)) = "); print(output_layout); print("\n"); + + // allocate tensors + std::cout << "Allocating tensors ... "; + thrust::universal_vector activation_data(size_t(cute::size(activation_layout)), float(0)); + thrust::universal_vector filter_data(size_t(cute::size(filter_layout)), float(0)); + thrust::universal_vector output_data(size_t(cute::size(output_layout)), float(0)); + thrust::universal_vector output_data_ref(size_t(cute::size(output_layout)), float(0)); + std::cout << "done.\n"; + + // init tensors + std::cout << "Initializing data ... " << std::flush; + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_real_distribution uniform_dist(-1.0, 1.0); + for (std::size_t i = 0; i < size_t(cute::size(activation_layout)); ++i) { + activation_data[i] = uniform_dist(gen); + } + + for (std::size_t i = 0; i < size_t(cute::size(filter_layout)); ++i) { + filter_data[i] = uniform_dist(gen); + } + std::cout << "done.\n"; + + // set up index buffers for gather/scatter, fill with indireciton indices in reversed order + std::cout << "Initializing gather/scatter index buffers ... "; + thrust::universal_vector gather_idx_buf(size_t(size<0>(activation_layout))); + thrust::universal_vector scatter_idx_buf(size_t(size<1>(output_layout))); + thrust::sequence(gather_idx_buf.rbegin(), gather_idx_buf.rend()); + thrust::sequence(scatter_idx_buf.rbegin(), scatter_idx_buf.rend()); + std::cout << "done.\n"; + + // launch dense + std::cout << "\nRunning dense fprop kernel\n"; + int passed = ampere_dense_conv_fprop( + num_images, + activation_data.data().get(), + filter_data.data().get(), + output_data.data().get(), + output_data_ref.data().get(), + num_iterations, + do_host_ref_check); + + // launch gather/scatter + std::cout << "\nRunning gather/scatter fprop kernel\n"; + ampere_gather_scatter_conv_fprop( + num_images, + activation_data.data().get(), + gather_idx_buf.data().get(), + filter_data.data().get(), + output_data.data().get(), + scatter_idx_buf.data().get(), + num_iterations); + + return passed; +} diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index a2daa07a..9cb125d9 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -138,6 +138,8 @@ foreach(EXAMPLE 55_hopper_mixed_dtype_gemm 56_hopper_ptr_array_batched_gemm 57_hopper_grouped_gemm + 58_ada_fp8_gemm + 59_ampere_gather_scatter_conv ) add_subdirectory(${EXAMPLE}) diff --git a/examples/52_hopper_gather_scatter_fusion/gather_tensor.hpp b/examples/common/gather_tensor.hpp similarity index 97% rename from examples/52_hopper_gather_scatter_fusion/gather_tensor.hpp rename to examples/common/gather_tensor.hpp index 5dc13941..62616e00 100644 --- a/examples/52_hopper_gather_scatter_fusion/gather_tensor.hpp +++ b/examples/common/gather_tensor.hpp @@ -32,6 +32,7 @@ #include "cute/layout.hpp" #include "cute/tensor.hpp" +#include "cute/util/print.hpp" namespace example { @@ -59,7 +60,7 @@ struct IndexedGather CUTE_HOST_DEVICE friend void print(IndexedGather const &s) { - print("Indexed"); + cute::print("Indexed"); } Index const *indices_; @@ -81,9 +82,9 @@ struct StridedGather CUTE_HOST_DEVICE friend void print(StridedGather const &s) { - print("Strided{"); + cute::print("Strided{"); print(s.stride_); - print("}"); + cute::print("}"); } Stride stride_; @@ -109,11 +110,11 @@ struct CustomStride CUTE_HOST_DEVICE friend void print(CustomStride const & s) { - print("Custom{"); + cute::print("Custom{"); print(s.func_); - print(","); + cute::print(","); print(s.stride_); - print("}"); + cute::print("}"); } template diff --git a/examples/cute/tutorial/CMakeLists.txt b/examples/cute/tutorial/CMakeLists.txt index dfe4fccb..1e4dad5f 100644 --- a/examples/cute/tutorial/CMakeLists.txt +++ b/examples/cute/tutorial/CMakeLists.txt @@ -29,8 +29,23 @@ cutlass_example_add_executable( - sgemm_nt_1 - sgemm_nt_1.cu + sgemm_1 + sgemm_1.cu +) + +cutlass_example_add_executable( + sgemm_2 + sgemm_2.cu +) + +cutlass_example_add_executable( + sgemm_sm70 + sgemm_sm70.cu +) + +cutlass_example_add_executable( + sgemm_sm80 + sgemm_sm80.cu ) cutlass_example_add_executable( diff --git a/examples/cute/tutorial/sgemm_1.cu b/examples/cute/tutorial/sgemm_1.cu new file mode 100644 index 00000000..46bf537a --- /dev/null +++ b/examples/cute/tutorial/sgemm_1.cu @@ -0,0 +1,469 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#include +#include +#include + +#include +#include + +#include + +#include "cutlass/util/print_error.hpp" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/helper_cuda.hpp" + +template +__global__ static +__launch_bounds__(decltype(size(CThreadLayout{}))::value) +void +gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler, + TA const* A, AStride dA, ASmemLayout sA_layout, AThreadLayout tA, + TB const* B, BStride dB, BSmemLayout sB_layout, BThreadLayout tB, + TC * C, CStride dC, CSmemLayout , CThreadLayout tC, + Alpha alpha, Beta beta) +{ + using namespace cute; + + // Preconditions + CUTE_STATIC_ASSERT_V(rank(shape_MNK) == Int<3>{}); // (M, N, K) + CUTE_STATIC_ASSERT_V(rank(cta_tiler) == Int<3>{}); // (BLK_M, BLK_N, BLK_K) + + static_assert(is_static::value); + static_assert(is_static::value); + static_assert(is_static::value); + + CUTE_STATIC_ASSERT_V(size(tA) == size(tB)); // NumThreads + CUTE_STATIC_ASSERT_V(size(tC) == size(tA)); // NumThreads + + CUTE_STATIC_ASSERT_V(size<0>(cta_tiler) % size<0>(tA) == Int<0>{}); // BLK_M / THR_M + CUTE_STATIC_ASSERT_V(size<2>(cta_tiler) % size<1>(tA) == Int<0>{}); // BLK_K / THR_K + CUTE_STATIC_ASSERT_V(size<1>(cta_tiler) % size<0>(tB) == Int<0>{}); // BLK_N / THR_N + CUTE_STATIC_ASSERT_V(size<2>(cta_tiler) % size<1>(tB) == Int<0>{}); // BLK_K / THR_K + CUTE_STATIC_ASSERT_V(size<0>(cta_tiler) % size<0>(tC) == Int<0>{}); // BLK_M / THR_M + CUTE_STATIC_ASSERT_V(size<1>(cta_tiler) % size<1>(tC) == Int<0>{}); // BLK_N / THR_N + + static_assert(is_static::value); + static_assert(is_static::value); + static_assert(is_static::value); + + CUTE_STATIC_ASSERT_V(size<0>(ASmemLayout{}) == size<0>(cta_tiler)); // BLK_M + CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<0>(cta_tiler)); // BLK_M + CUTE_STATIC_ASSERT_V(size<0>(BSmemLayout{}) == size<1>(cta_tiler)); // BLK_N + CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<1>(cta_tiler)); // BLK_N + CUTE_STATIC_ASSERT_V(size<1>(ASmemLayout{}) == size<2>(cta_tiler)); // BLK_K + CUTE_STATIC_ASSERT_V(size<1>(BSmemLayout{}) == size<2>(cta_tiler)); // BLK_K + + CUTE_STATIC_ASSERT_V(congruent(select<0,2>(shape_MNK), dA)); // dA strides for shape MK + CUTE_STATIC_ASSERT_V(congruent(select<1,2>(shape_MNK), dB)); // dB strides for shape NK + CUTE_STATIC_ASSERT_V(congruent(select<0,1>(shape_MNK), dC)); // dC strides for shape MN + + // + // Full and Tiled Tensors + // + + // Represent the full tensors + Tensor mA = make_tensor(make_gmem_ptr(A), select<0,2>(shape_MNK), dA); // (M,K) + Tensor mB = make_tensor(make_gmem_ptr(B), select<1,2>(shape_MNK), dB); // (N,K) + Tensor mC = make_tensor(make_gmem_ptr(C), select<0,1>(shape_MNK), dC); // (M,N) + + // Get the appropriate blocks for this thread block + auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k) + Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) + + // Shared memory buffers + __shared__ TA smemA[cosize_v]; + __shared__ TB smemB[cosize_v]; + Tensor sA = make_tensor(make_smem_ptr(smemA), sA_layout); // (BLK_M,BLK_K) + Tensor sB = make_tensor(make_smem_ptr(smemB), sB_layout); // (BLK_N,BLK_K) + + // + // Partition the copying of A and B tiles across the threads + // + + // TUTORIAL: Example of simple raked partitioning of ThreadLayouts tA|tB over data A|B tiles + + Tensor tAgA = local_partition(gA, tA, threadIdx.x); // (THR_M,THR_K,k) + Tensor tAsA = local_partition(sA, tA, threadIdx.x); // (THR_M,THR_K) + + Tensor tBgB = local_partition(gB, tB, threadIdx.x); // (THR_N,THR_K,k) + Tensor tBsB = local_partition(sB, tB, threadIdx.x); // (THR_N,THR_K) + + CUTE_STATIC_ASSERT_V(size<0>(tAgA) == size<0>(tAsA)); // THR_M + CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tAsA)); // THR_K + CUTE_STATIC_ASSERT_V(size<0>(tBgB) == size<0>(tBsB)); // THR_N + CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBsB)); // THR_K + + // + // Define A/B partitioning and C accumulators + // + + // TUTORIAL: Example of partitioning via projections of a ThreadLayout tC + + // Partition sA (M,K) by the rows of tC + Tensor tCsA = local_partition(sA, tC, threadIdx.x, Step<_1, X>{}); // (THR_M,BLK_K) + // Partition sB (N,K) by the cols of tC + Tensor tCsB = local_partition(sB, tC, threadIdx.x, Step< X,_1>{}); // (THR_N,BLK_K) + // Partition gC (M,N) by the tile of tC + Tensor tCgC = local_partition(gC, tC, threadIdx.x, Step<_1,_1>{}); // (THR_M,THR_N) + + // Allocate the accumulators -- same shape/layout as the partitioned data + Tensor tCrC = make_tensor_like(tCgC); // (THR_M,THR_N) + + CUTE_STATIC_ASSERT_V(size<0>(tCrC) == size<0>(tCgC)); // THR_M + CUTE_STATIC_ASSERT_V(size<0>(tCrC) == size<0>(tCsA)); // THR_M + CUTE_STATIC_ASSERT_V(size<1>(tCrC) == size<1>(tCgC)); // THR_N + CUTE_STATIC_ASSERT_V(size<1>(tCrC) == size<0>(tCsB)); // THR_N + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCsB)); // BLK_K + + // Clear the accumulators + clear(tCrC); + +#if 0 + if(thread0()) { + print(" mA : "); print( mA); print("\n"); + print(" gA : "); print( gA); print("\n"); + print(" sA : "); print( sA); print("\n"); + print("tAgA : "); print(tAgA); print("\n"); + print("tAsA : "); print(tAsA); print("\n"); + } +#endif + +#if 0 + if(thread0()) { + print(" mB : "); print( mB); print("\n"); + print(" gB : "); print( gB); print("\n"); + print(" sB : "); print( sB); print("\n"); + print("tBgB : "); print(tBgB); print("\n"); + print("tBsB : "); print(tBsB); print("\n"); + } +#endif + +#if 0 + if(thread0()) { + print(" mC : "); print( mC); print("\n"); + print(" gC : "); print( gC); print("\n"); + print("tCsA : "); print(tCsA); print("\n"); + print("tCsB : "); print(tCsB); print("\n"); + print("tCgC : "); print(tCgC); print("\n"); + print("tCrC : "); print(tCrC); print("\n"); + } +#endif + +#if 1 + + // TUTORIAL: Example of a simple mainloop that read tiles of data into shared memory, + // and then computes on those tiles. + // copy(.) operates on the global and shared memory via the tA|tB partitioning + // gemm(.) operates on the shared and register memory via the tC partitioning + + auto K_TILE_MAX = size<2>(tAgA); + + for (int k_tile = 0; k_tile < K_TILE_MAX; ++k_tile) + { + // Copy gmem to smem with tA|tB thread-partitioned tensors + copy(tAgA(_,_,k_tile), tAsA); // A (THR_M,THR_K) -> (THR_M,THR_K) + copy(tBgB(_,_,k_tile), tBsB); // B (THR_N,THR_K) -> (THR_N,THR_K) + + // TUTORIAL: The above call to copy(tAgA(_,_,k_tile), tAsA) is equivalent to + // Tensor tAgAk = tAgA(_,_,k_tile); + // CUTE_UNROLL + // for (int i = 0; i < size(tAsA); ++i) { + // tAsA(i) = tAgAk(i); + // } + + cp_async_fence(); // Label the end of (potential) cp.async instructions + cp_async_wait<0>(); // Sync on all (potential) cp.async instructions + __syncthreads(); // Wait for all threads to write to smem + + // Compute gemm on tC thread-partitioned smem + gemm(tCsA, tCsB, tCrC); // (THR_M,THR_N) += (THR_M,BLK_K) * (THR_N,BLK_K) + + // TUTORIAL: The above call to gemm(tCsA, tCsB, tCrC) is equivalent to + // CUTE_UNROLL + // for (int k = 0; k < size<1>(tCsA); ++k) { + // CUTE_UNROLL + // for (int m = 0; m < size<0>(tCrC); ++m) { + // CUTE_UNROLL + // for (int n = 0; n < size<1>(tCrC); ++n) { + // tCrC(m,n) += tCsA(m,k) * tCsB(n,k); + // } + // } + // } + + __syncthreads(); // Wait for all threads to read from smem + } + +#endif + + // + // Epilogue + // + + axpby(alpha, tCrC, beta, tCgC); + + // TUTORIAL: The above call to axpby(alpha, tCrC, beta, tCgC) is equivalent to + // CUTE_UNROLL + // for (int i = 0; i < size(tCsA); ++i) { + // tCgC(i) = alpha * tCrC(i) + beta * tCgC(i); + // } +} + +// Setup params for an NT GEMM +// Use m-major smem sA, n-major smem sB, and mn-major threads tA|tB +template +void +gemm_nt(int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + using namespace cute; + + // Define shapes (dynamic) + auto M = int(m); + auto N = int(n); + auto K = int(k); + auto prob_shape = make_shape(M, N, K); // (M, N, K) + + // Define NT strides (mixed) + auto dA = make_stride(Int<1>{}, ldA); // (dM, dK) + auto dB = make_stride(Int<1>{}, ldB); // (dN, dK) + auto dC = make_stride(Int<1>{}, ldC); // (dM, dN) + + // Define CTA tile sizes (static) + auto bM = Int<128>{}; + auto bN = Int<128>{}; + auto bK = Int< 8>{}; + auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K) + + // Define the smem layouts (static) + auto sA = make_layout(make_shape(bM, bK)); // (m,k) -> smem_idx; m-major + auto sB = make_layout(make_shape(bN, bK)); // (n,k) -> smem_idx; n-major + auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx; m-major + + // Define the thread layouts (static) + auto tA = make_layout(make_shape(Int<32>{}, Int< 8>{})); // (m,k) -> thr_idx + auto tB = make_layout(make_shape(Int<32>{}, Int< 8>{})); // (n,k) -> thr_idx + auto tC = make_layout(make_shape(Int<16>{}, Int<16>{})); // (m,n) -> thr_idx + + dim3 dimBlock(size(tC)); + dim3 dimGrid(size(ceil_div(M, bM)), + size(ceil_div(N, bN))); + gemm_device<<>> + (prob_shape, cta_tiler, + A, dA, sA, tA, + B, dB, sB, tB, + C, dC, sC, tC, + alpha, beta); +} + +// Setup params for a TN GEMM +// Use padded m-major smem sA, padded n-major smem sB, and k-major threads tA|tB +template +void +gemm_tn(int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + using namespace cute; + + // Define shapes (dynamic) + auto M = int(m); + auto N = int(n); + auto K = int(k); + auto prob_shape = make_shape(M, N, K); // (M, N, K) + + // Define TN strides (mixed) + auto dA = make_stride(ldA, Int<1>{}); // (dM, dK) + auto dB = make_stride(ldB, Int<1>{}); // (dN, dK) + auto dC = make_stride(Int<1>{}, ldC); // (dM, dN) + + // Define CTA tile sizes (static) + auto bM = Int<128>{}; + auto bN = Int<128>{}; + auto bK = Int< 8>{}; + auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K) + + // Define the smem layouts (static) + auto sA = make_layout(make_shape(bM,bK), LayoutRight{}); // (m,k) -> smem_idx; k-major + auto sB = make_layout(make_shape(bN,bK), LayoutRight{}); // (n,k) -> smem_idx; k-major + auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx; m-major + + // Define the thread layouts (static) + auto tA = make_layout(make_shape(Int<32>{}, Int< 8>{}), LayoutRight{}); // (m,k) -> thr_idx; k-major + auto tB = make_layout(make_shape(Int<32>{}, Int< 8>{}), LayoutRight{}); // (n,k) -> thr_idx; k-major + auto tC = make_layout(make_shape(Int<16>{}, Int<16>{})); // (m,n) -> thr_idx; m-major + + dim3 dimBlock(size(tC)); + dim3 dimGrid(size(ceil_div(M, bM)), + size(ceil_div(N, bN))); + gemm_device<<>> + (prob_shape, cta_tiler, + A, dA, sA, tA, + B, dB, sB, tB, + C, dC, sC, tC, + alpha, beta); +} + +template +void +gemm(char transA, char transB, int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + if (transA == 'N' && transB == 'T') { + return gemm_nt(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream); + } else + if (transA == 'T' && transB == 'N') { + return gemm_tn(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream); + } + assert(false && "Not implemented"); +} + + +int main(int argc, char** argv) +{ + int m = 5120; + if (argc >= 2) + sscanf(argv[1], "%d", &m); + + int n = 5120; + if (argc >= 3) + sscanf(argv[2], "%d", &n); + + int k = 4096; + if (argc >= 4) + sscanf(argv[3], "%d", &k); + + char transA = 'N'; + if (argc >= 5) + sscanf(argv[4], "%c", &transA); + + char transB = 'T'; + if (argc >= 6) + sscanf(argv[5], "%c", &transB); + + using TA = float; + using TB = float; + using TC = float; + using TI = float; + + TI alpha = 1.0; + TI beta = 0.0; + + std::cout << "M = " << m << std::endl; + std::cout << "N = " << n << std::endl; + std::cout << "K = " << k << std::endl; + std::cout << "C = A^" << transA << " B^" << transB << std::endl; + + cute::device_init(0); + + thrust::host_vector h_A(m*k); + thrust::host_vector h_B(n*k); + thrust::host_vector h_C(m*n); + + for (int j = 0; j < m*k; ++j) h_A[j] = static_cast( 2*(rand() / double(RAND_MAX)) - 1 ); + for (int j = 0; j < n*k; ++j) h_B[j] = static_cast( 2*(rand() / double(RAND_MAX)) - 1 ); + for (int j = 0; j < m*n; ++j) h_C[j] = static_cast(-1); + + thrust::device_vector d_A = h_A; + thrust::device_vector d_B = h_B; + thrust::device_vector d_C = h_C; + + double gflops = (2.0*m*n*k) * 1e-9; + + const int timing_iterations = 100; + GPU_Clock timer; + + int ldA = 0, ldB = 0, ldC = m; + + if (transA == 'N') { + ldA = m; + } else if (transA == 'T') { + ldA = k; + } else { + assert(false); + } + + if (transB == 'N') { + ldB = k; + } else if (transB == 'T') { + ldB = n; + } else { + assert(false); + } + // Run once + d_C = h_C; + gemm(transA, transB, m, n, k, + alpha, + d_A.data().get(), ldA, + d_B.data().get(), ldB, + beta, + d_C.data().get(), ldC); + CUTE_CHECK_LAST(); + thrust::host_vector cute_result = d_C; + + // Timing iterations + timer.start(); + for (int i = 0; i < timing_iterations; ++i) { + gemm(transA, transB, m, n, k, + alpha, + d_A.data().get(), ldA, + d_B.data().get(), ldB, + beta, + d_C.data().get(), ldC); + } + double cute_time = timer.seconds() / timing_iterations; + CUTE_CHECK_LAST(); + printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000); + return 0; +} diff --git a/examples/cute/tutorial/sgemm_2.cu b/examples/cute/tutorial/sgemm_2.cu new file mode 100644 index 00000000..b0d25bfe --- /dev/null +++ b/examples/cute/tutorial/sgemm_2.cu @@ -0,0 +1,523 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#include +#include +#include + +#include +#include + +#include + +#include "cutlass/util/print_error.hpp" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/helper_cuda.hpp" + +template +__global__ static +__launch_bounds__(decltype(size(TiledMma{}))::value) +void +gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler, + TA const* A, AStride dA, ASmemLayout sA_layout, TiledCopyA copy_a, + TB const* B, BStride dB, BSmemLayout sB_layout, TiledCopyB copy_b, + TC * C, CStride dC, CSmemLayout , TiledMma mma, + Alpha alpha, Beta beta) +{ + using namespace cute; + + // Preconditions + CUTE_STATIC_ASSERT_V(rank(shape_MNK) == Int<3>{}); // (M, N, K) + CUTE_STATIC_ASSERT_V(rank(cta_tiler) == Int<3>{}); // (BLK_M, BLK_N, BLK_K) + + CUTE_STATIC_ASSERT_V(size(copy_a) == size(mma)); // NumThreads + CUTE_STATIC_ASSERT_V(size(copy_b) == size(mma)); // NumThreads + + static_assert(is_static::value); + static_assert(is_static::value); + static_assert(is_static::value); + + CUTE_STATIC_ASSERT_V(size<0>(ASmemLayout{}) == size<0>(cta_tiler)); // BLK_M + CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<0>(cta_tiler)); // BLK_M + CUTE_STATIC_ASSERT_V(size<0>(BSmemLayout{}) == size<1>(cta_tiler)); // BLK_N + CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<1>(cta_tiler)); // BLK_N + CUTE_STATIC_ASSERT_V(size<1>(ASmemLayout{}) == size<2>(cta_tiler)); // BLK_K + CUTE_STATIC_ASSERT_V(size<1>(BSmemLayout{}) == size<2>(cta_tiler)); // BLK_K + + CUTE_STATIC_ASSERT_V(congruent(select<0,2>(shape_MNK), dA)); // dA strides for shape MK + CUTE_STATIC_ASSERT_V(congruent(select<1,2>(shape_MNK), dB)); // dB strides for shape NK + CUTE_STATIC_ASSERT_V(congruent(select<0,1>(shape_MNK), dC)); // dC strides for shape MN + + // + // Full and Tiled Tensors + // + + // Represent the full tensors + Tensor mA = make_tensor(make_gmem_ptr(A), select<0,2>(shape_MNK), dA); // (M,K) + Tensor mB = make_tensor(make_gmem_ptr(B), select<1,2>(shape_MNK), dB); // (N,K) + Tensor mC = make_tensor(make_gmem_ptr(C), select<0,1>(shape_MNK), dC); // (M,N) + + // Get the appropriate blocks for this thread block + auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k) + Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) + + // Shared memory buffers + __shared__ TA smemA[cosize_v]; + __shared__ TB smemB[cosize_v]; + Tensor sA = make_tensor(make_smem_ptr(smemA), sA_layout); // (BLK_M,BLK_K) + Tensor sB = make_tensor(make_smem_ptr(smemB), sB_layout); // (BLK_N,BLK_K) + + // + // Partition the copying of A and B tiles across the threads + // + + // TUTORIAL: Example of partitioning via a TiledCopy + + ThrCopy thr_copy_a = copy_a.get_slice(threadIdx.x); + Tensor tAgA = thr_copy_a.partition_S(gA); // (CPY,CPY_M,CPY_K,k) + Tensor tAsA = thr_copy_a.partition_D(sA); // (CPY,CPY_M,CPY_K) + // Allocate registers same shape/layout as partitioned data + Tensor tArA = make_fragment_like(tAsA); // (CPY,CPY_M,CPY_K) + + ThrCopy thr_copy_b = copy_b.get_slice(threadIdx.x); + Tensor tBgB = thr_copy_b.partition_S(gB); // (CPY,CPY_N,CPY_K,k) + Tensor tBsB = thr_copy_b.partition_D(sB); // (CPY,CPY_N,CPY_K) + // Allocate registers same shape/layout as partitioned data + Tensor tBrB = make_fragment_like(tBsB); // (CPY,CPY_N,CPY_K) + + CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tAsA)); // CPY_M + CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tArA)); // CPY_M + CUTE_STATIC_ASSERT_V(size<2>(tAgA) == size<2>(tAsA)); // CPY_K + CUTE_STATIC_ASSERT_V(size<2>(tAgA) == size<2>(tArA)); // CPY_K + CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBsB)); // CPY_N + CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBrB)); // CPY_N + CUTE_STATIC_ASSERT_V(size<2>(tBgB) == size<2>(tBsB)); // CPY_K + CUTE_STATIC_ASSERT_V(size<2>(tBgB) == size<2>(tBrB)); // CPY_K + + // Copy gmem to rmem for k_tile=0 + copy(copy_a, tAgA(_,_,_,0), tArA); + copy(copy_b, tBgB(_,_,_,0), tBrB); + // + // Define A/B partitioning and C accumulators + // + + // TUTORIAL: Example of partitioning via a TiledMMA + + ThrMMA thr_mma = mma.get_slice(threadIdx.x); + Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K) + Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K) + Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N) + + // Allocate the accumulators -- same size as the projected data + Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N) + + CUTE_STATIC_ASSERT_V( shape(tCrC) == shape(tCgC)); // (MMA,MMA_M,MMA_N) + CUTE_STATIC_ASSERT_V(size<1>(tCgC) == size<1>(tCsA)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(tCgC) == size<1>(tCsB)); // MMA_N + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // MMA_K + + // Clear the accumulators + clear(tCrC); + +#if 0 + if(thread0()) { + print(" mA : "); print( mA); print("\n"); + print(" gA : "); print( gA); print("\n"); + print(" sA : "); print( sA); print("\n"); + print("tAgA : "); print(tAgA); print("\n"); + print("tAsA : "); print(tAsA); print("\n"); + print("tArA : "); print(tArA); print("\n"); + } +#endif + +#if 0 + if(thread0()) { + print(" mB : "); print( mB); print("\n"); + print(" gB : "); print( gB); print("\n"); + print(" sB : "); print( sB); print("\n"); + print("tBgB : "); print(tBgB); print("\n"); + print("tBsB : "); print(tBsB); print("\n"); + print("tArA : "); print(tArA); print("\n"); + } +#endif + +#if 0 + if(thread0()) { + print(" mC : "); print( mC); print("\n"); + print(" gC : "); print( gC); print("\n"); + print("tCsA : "); print(tCsA); print("\n"); + print("tCsB : "); print(tCsB); print("\n"); + print("tCgC : "); print(tCgC); print("\n"); + print("tCrC : "); print(tCrC); print("\n"); + } +#endif + +#if 1 + + // TUTORIAL: Example of an inner loop that pipelines compute with reads + // from global memory by staging through register and shared memory. + // Data is read from global to registers, then to shared via the TiledCopy partitions + // gemm(.) operates on the shared memory directly via the TiledMMA partitions + + auto K_TILE_MAX = size<3>(tAgA); + + for (int k_tile = 0; k_tile < K_TILE_MAX; ++k_tile) + { + // Copy rmem to smem with tA|tB thread-partitioned tensors + __syncthreads(); // Wait for all threads to consume smem + copy(tArA, tAsA); + copy(tBrB, tBsB); + __syncthreads(); // Wait for all threads to consume smem + + // Copy gmem to rmem for k_tile+1 with tA|tB thread-partitioned tensors + int k_tile_next = (k_tile + 1 < K_TILE_MAX) ? k_tile + 1 : k_tile; + copy(copy_a, tAgA(_,_,_,k_tile_next), tArA); + copy(copy_b, tBgB(_,_,_,k_tile_next), tBrB); + // TUTORIAL: The above call to copy(copy_a, tAgA(_,_,_,k_tile_next), tArA) is equivalent to + // CUTE_UNROLL + // for (int k = 0; k < size<1>(tCsA); ++k) { + // CUTE_UNROLL + // for (int m = 0; m < size<0>(tCrC); ++m) { + // copy_a.call(tAgA(_,m,k), tArA(_,m,k); + // } + // } + + // Compute gemm on mma-partitioned smem + gemm(mma, tCsA, tCsB, tCrC); + // TUTORIAL: The above call to gemm(tCsA, tCsB, tCrC) is equivalent to + // CUTE_UNROLL + // for (int k = 0; k < size<1>(tCsA); ++k) { + // CUTE_UNROLL + // for (int m = 0; m < size<0>(tCrC); ++m) { + // CUTE_UNROLL + // for (int n = 0; n < size<1>(tCrC); ++n) { + // mma.call(tCsA(_,m,k), tCsB(_,n,k), tCrC(_,m,n); + // } + // } + // } + } + +#endif + + // + // Epilogue + // + + axpby(alpha, tCrC, beta, tCgC); +} + +// Setup params for a NT GEMM +template +void +gemm_nt(int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + using namespace cute; + + // Define shapes (dynamic) + auto M = int(m); + auto N = int(n); + auto K = int(k); + auto prob_shape = make_shape(M, N, K); // (M, N, K) + + // Define NT strides (mixed) + auto dA = make_stride(Int<1>{}, ldA); // (dM, dK) + auto dB = make_stride(Int<1>{}, ldB); // (dN, dK) + auto dC = make_stride(Int<1>{}, ldC); // (dM, dN) + + // Define CTA tile sizes (static) + auto bM = Int<128>{}; + auto bN = Int<128>{}; + auto bK = Int< 8>{}; + auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K) + + // Define the smem layouts (static) + auto sA = make_layout(make_shape(bM, bK)); // (m,k) -> smem_idx; m-major + auto sB = make_layout(make_shape(bN, bK)); // (n,k) -> smem_idx; n-major + auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx; m-major + + // Define the thread layouts (static) + + // TUTORIAL: Construct TiledCopy with a particular Copy_Atom to use and + // define the partitioning pattern to apply. + // Each thread will (try to) copy 4x1 elements of type TA using 128-bit copy. + // Use 32x8 of these threads. + + TiledCopy copyA = make_tiled_copy(Copy_Atom, TA>{}, + Layout>{}, // Thr layout 32x8 m-major + Layout>{}); // Val layout 4x1 m-major + TiledCopy copyB = make_tiled_copy(Copy_Atom, TB>{}, + Layout>{}, // Thr layout 32x8 n-major + Layout>{}); // Val layout 4x1 n-major + + // TUTORIAL: Construct TiledMMA with a particular MMA_Atom to use and + // define the partitioning pattern to apply. + // Use a 1x1x1 FMA on the types TC += TA * TB. Each atom requires a single thread. + // Reproduce that atom 16x16x1 times (m-major) across threads so that we use 256 threads. + + TiledMMA mmaC = make_tiled_mma(UniversalFMA{}, + Layout>{}); // 16x16x1 UniversalFMA + +#if 0 + print(copyA); + print(copyB); + print(mmaC); +#endif + +#if 0 + print_latex(copyA); + print_latex(copyB); + print_latex(mmaC); +#endif + + dim3 dimBlock(size(mmaC)); + dim3 dimGrid(size(ceil_div(M, bM)), + size(ceil_div(N, bN))); + gemm_device<<>> + (prob_shape, cta_tiler, + A, dA, sA, copyA, + B, dB, sB, copyB, + C, dC, sC, mmaC, + alpha, beta); +} + +// Setup params for a TN GEMM +template +void +gemm_tn(int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + using namespace cute; + + // Define shapes (dynamic) + auto M = int(m); + auto N = int(n); + auto K = int(k); + auto prob_shape = make_shape(M, N, K); // (M, N, K) + + // Define TN strides (mixed) + auto dA = make_stride(ldA, Int<1>{}); // (dM, dK) + auto dB = make_stride(ldB, Int<1>{}); // (dN, dK) + auto dC = make_stride(Int<1>{}, ldC); // (dM, dN) + + // Define CTA tile sizes (static) + auto bM = Int<128>{}; + auto bN = Int<128>{}; + auto bK = Int< 8>{}; + auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K) + + // Define the smem layouts (static) + auto sA = make_layout(make_shape ( bM, bK), + make_stride(Int<1>{}, bM+Int<1>{})); // (m,k) -> smem_idx; padded m-major + auto sB = make_layout(make_shape ( bN, bK), + make_stride(Int<1>{}, bN+Int<1>{})); // (n,k) -> smem_idx; padded n-major + auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx + + // TUTORIAL: Construct TiledCopy to define the Copy_Atom to use and the + // partitioning pattern to apply. + // Each thread will copy 1x1 elements of type TA. + // Use 32x8 of these threads arranged in k-major. + + TiledCopy copyA = make_tiled_copy(Copy_Atom, TA>{}, + Layout,Stride<_8,_1>>{}, // Thr layout 32x8 k-major + Layout>{}); // Val layout 1x1 + TiledCopy copyB = make_tiled_copy(Copy_Atom, TB>{}, + Layout,Stride<_8,_1>>{}, // Thr layout 32x8 k-major + Layout>{}); // Val layout 1x1 + + // TUTORIAL: Construct TiledMMA to define the MMA_Atom to use and the + // partitioning pattern to apply. + // Use a 1x1x1 FMA on the types TC += TA * TB. Each atom requires a single thread. + // Reproduce that atom 16x16x1 times (m-major) across threads so that we use 256 threads. + + TiledMMA mmaC = make_tiled_mma(UniversalFMA{}, + Layout>{}); // 16x16x1 TiledMMA + +#if 0 + print(copyA); + print(copyB); + print(mmaC); +#endif + +#if 0 + print_latex(copyA); + print_latex(copyB); + print_latex(mmaC); +#endif + + dim3 dimBlock(size(mmaC)); + dim3 dimGrid(size(ceil_div(M, bM)), + size(ceil_div(N, bN))); + gemm_device<<>> + (prob_shape, cta_tiler, + A, dA, sA, copyA, + B, dB, sB, copyB, + C, dC, sC, mmaC, + alpha, beta); +} + +template +void +gemm(char transA, char transB, int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + if (transA == 'N' && transB == 'T') { + return gemm_nt(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream); + } else + if (transA == 'T' && transB == 'N') { + return gemm_tn(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream); + } + assert(false && "Not implemented"); +} + + +int main(int argc, char** argv) +{ + int m = 5120; + if (argc >= 2) + sscanf(argv[1], "%d", &m); + + int n = 5120; + if (argc >= 3) + sscanf(argv[2], "%d", &n); + + int k = 4096; + if (argc >= 4) + sscanf(argv[3], "%d", &k); + + char transA = 'N'; + if (argc >= 5) + sscanf(argv[4], "%c", &transA); + + char transB = 'T'; + if (argc >= 6) + sscanf(argv[5], "%c", &transB); + + using TA = float; + using TB = float; + using TC = float; + using TI = float; + + TI alpha = 1.0; + TI beta = 0.0; + + std::cout << "M = " << m << std::endl; + std::cout << "N = " << n << std::endl; + std::cout << "K = " << k << std::endl; + std::cout << "C = A^" << transA << " B^" << transB << std::endl; + + cute::device_init(0); + + thrust::host_vector h_A(m*k); + thrust::host_vector h_B(n*k); + thrust::host_vector h_C(m*n); + + for (int j = 0; j < m*k; ++j) h_A[j] = static_cast( 2*(rand() / double(RAND_MAX)) - 1 ); + for (int j = 0; j < n*k; ++j) h_B[j] = static_cast( 2*(rand() / double(RAND_MAX)) - 1 ); + for (int j = 0; j < m*n; ++j) h_C[j] = static_cast(-1); + + thrust::device_vector d_A = h_A; + thrust::device_vector d_B = h_B; + thrust::device_vector d_C = h_C; + + double gflops = (2.0*m*n*k) * 1e-9; + + const int timing_iterations = 100; + GPU_Clock timer; + + int ldA = 0, ldB = 0, ldC = m; + + if (transA == 'N') { + ldA = m; + } else if (transA == 'T') { + ldA = k; + } else { + assert(false); + } + + if (transB == 'N') { + ldB = k; + } else if (transB == 'T') { + ldB = n; + } else { + assert(false); + } + + // Run once + d_C = h_C; + gemm(transA, transB, m, n, k, + alpha, + d_A.data().get(), ldA, + d_B.data().get(), ldB, + beta, + d_C.data().get(), ldC); + CUTE_CHECK_LAST(); + thrust::host_vector cute_result = d_C; + + // Timing iterations + timer.start(); + for (int i = 0; i < timing_iterations; ++i) { + gemm(transA, transB, m, n, k, + alpha, + d_A.data().get(), ldA, + d_B.data().get(), ldB, + beta, + d_C.data().get(), ldC); + } + double cute_time = timer.seconds() / timing_iterations; + CUTE_CHECK_LAST(); + printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000); + + return 0; +} diff --git a/examples/cute/tutorial/sgemm_nt_1.cu b/examples/cute/tutorial/sgemm_nt_1.cu deleted file mode 100644 index 65ce9930..00000000 --- a/examples/cute/tutorial/sgemm_nt_1.cu +++ /dev/null @@ -1,426 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ -#include -#include - -#include - -#include "cutlass/util/print_error.hpp" -#include "cutlass/util/GPU_Clock.hpp" -#if defined(CUTLASS_ENABLE_CUBLAS) && CUTLASS_ENABLE_CUBLAS != 0 -# include "cutlass/util/cublas_wrappers.hpp" -#endif -#include "cutlass/util/helper_cuda.hpp" - -template -__global__ static -__launch_bounds__(decltype(size(CThreadLayout{}))::value) -void -gemm_device(MShape M, NShape N, KShape K, - TA const* A, AStride dA, ABlockLayout blockA, AThreadLayout tA, - TB const* B, BStride dB, BBlockLayout blockB, BThreadLayout tB, - TC * C, CStride dC, CBlockLayout , CThreadLayout tC, - Alpha alpha, Beta beta) -{ - using namespace cute; - using X = Underscore; - - // Preconditions - CUTE_STATIC_ASSERT(is_static::value); - CUTE_STATIC_ASSERT(is_static::value); - CUTE_STATIC_ASSERT(is_static::value); - - CUTE_STATIC_ASSERT(is_static::value); - CUTE_STATIC_ASSERT(is_static::value); - CUTE_STATIC_ASSERT(is_static::value); - - CUTE_STATIC_ASSERT_V(size(tA) == size(tC)); - CUTE_STATIC_ASSERT_V(size(tB) == size(tC)); - - //CUTE_STATIC_ASSERT_V(shape<0>(blockA) == shape<0>(blockC)); // BLK_M - //CUTE_STATIC_ASSERT_V(shape<0>(blockB) == shape<1>(blockC)); // BLK_N - CUTE_STATIC_ASSERT_V(shape<1>(blockA) == shape<1>(blockB)); // BLK_K - - // Shared memory buffers - __shared__ TA smemA[cosize_v]; - __shared__ TB smemB[cosize_v]; - auto sA = make_tensor(make_smem_ptr(smemA), blockA); // (BLK_M,BLK_K) - auto sB = make_tensor(make_smem_ptr(smemB), blockB); // (BLK_N,BLK_K) - - // Represent the full tensors - auto mA = make_tensor(make_gmem_ptr(A), make_shape(M,K), dA); // (M,K) - auto mB = make_tensor(make_gmem_ptr(B), make_shape(N,K), dB); // (N,K) - auto mC = make_tensor(make_gmem_ptr(C), make_shape(M,N), dC); // (M,N) - - // Get the appropriate blocks for this thread block -- - // potential for thread block locality - auto blk_shape = make_shape(size<0>(sA), size<0>(sB), size<1>(sB));// (BLK_M,BLK_N,BLK_K) - auto blk_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k) - - auto gA = local_tile(mA, blk_shape, blk_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) - auto gB = local_tile(mB, blk_shape, blk_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - auto gC = local_tile(mC, blk_shape, blk_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) - - // - // Partition the copying of A and B tiles across the threads - // - - // TUTORIAL: Example of simple partitioning of A|B tiles over tA|tB - // Default is a raked partition, but can be changed with Step parameter - - auto tAgA = local_partition(gA, tA, threadIdx.x); // (THR_M,THR_K,k) - auto tAsA = local_partition(sA, tA, threadIdx.x); // (THR_M,THR_K) - - auto tBgB = local_partition(gB, tB, threadIdx.x); // (THR_N,THR_K,k) - auto tBsB = local_partition(sB, tB, threadIdx.x); // (THR_N,THR_K) - - // - // Define C accumulators and A/B partitioning - // - - // TUTORIAL: Example of partitioning via projections of tC - - // Partition sA (M,K) by the rows of tC - auto tCsA = local_partition(sA, tC, threadIdx.x, Step<_1, X>{}); // (THR_M,BLK_K) - // Partition sB (N,K) by the cols of tC - auto tCsB = local_partition(sB, tC, threadIdx.x, Step< X,_1>{}); // (THR_N,BLK_K) - // Partition gC (M,N) by the tile of tC - auto tCgC = local_partition(gC, tC, threadIdx.x, Step<_1,_1>{}); // (THR_M,THR_N) - - // Allocate the accumulators -- same size as the projected data - auto tCrC = make_fragment_like(tCgC); // (THR_M,THR_N) - - // Clear the accumulators - clear(tCrC); - -#if 0 - if(thread0()) { - print("mA\n"); - print(mA.shape()); print("\n"); print(mA.stride()); - print("\n\ngA\n"); - print(gA.shape()); print("\n"); print(gA.stride()); - print("\n\ntAgA\n"); - print(tAgA.shape()); print("\n"); print(tAgA.stride()); - print("\n\nsA\n"); - print(sA.shape()); print("\n"); print(sA.stride()); - print("\n\ntAsA\n"); - print(tAsA.shape()); print("\n"); print(tAsA.stride()); - print("\n\n"); - } -#endif - -#if 0 - if(thread0()) { - print("mB\n"); - print(mB.shape()); print("\n"); print(mB.stride()); - print("\n\ngB\n"); - print(gB.shape()); print("\n"); print(gB.stride()); - print("\n\ntBgB\n"); - print(tBgB.shape()); print("\n"); print(tBgB.stride()); - print("\n\nsB\n"); - print(sB.shape()); print("\n"); print(sB.stride()); - print("\n\ntBsB\n"); - print(tBsB.shape()); print("\n"); print(tBsB.stride()); - print("\n\n"); - } -#endif - -#if 0 - if(thread0()) { - print("mC\n"); - print(mC.shape()); print("\n"); print(mC.stride()); - print("\n\ngC\n"); - print(gC.shape()); print("\n"); print(gC.stride()); - print("\n\ntCsA\n"); - print(tCsA.shape()); print("\n"); print(tCsA.stride()); - print("\n\ntCsB\n"); - print(tCsB.shape()); print("\n"); print(tCsB.stride()); - print("\n\ntCgC\n"); - print(tCgC.shape()); print("\n"); print(tCgC.stride()); - print("\n\ntCrC\n"); - print(tCrC.shape()); print("\n"); print(tCrC.stride()); - print("\n\n"); - } -#endif - -#if 1 - - // TUTORIAL: Example of a very simple compute loop - // Data is read from global to shared memory via the tA|tB partitioning - // gemm(.) operates on the shared memory directly via the tC partitioning - - auto k_max = size<2>(tAgA); - - for (int k = 0; k < k_max; ++k) - { - // Copy gmem to smem - copy(tAgA(_,_,k), tAsA); - copy(tBgB(_,_,k), tBsB); - - // In case copy uses cp.async, make sure that the cp.async - // instructions are ordered with respect to other cp.async - // instructions (fence), then wait on all the outstanding copy - // operations (wait<0>()). __syncthreads() alone does not do - // this. - // - // NOTE: cp_async_wait<0>() currently issues cp.async.wait_all. - // This is equivalent to cp.async.commit_group followed by - // cp.async_wait_group 0. This should make the first - // cp_async_fence() (which also issues cp.async.commit_group) - // redundant. The tutorial works as-is, so we'll leave the - // redundant fence in for now and study its removal later. - cp_async_fence(); - cp_async_wait<0>(); - - __syncthreads(); - - // Compute gemm on smem - gemm(tCsA, tCsB, tCrC); - - __syncthreads(); - } - -#endif - - // - // Epilogue - // - - axpby(alpha, tCrC, beta, tCgC); -} - - -template -void -gemm(int m, int n, int k, - Alpha alpha, - TA const* A, int ldA, - TB const* B, int ldB, - Beta beta, - TC * C, int ldC, - cudaStream_t stream = 0) -{ - using namespace cute; - - // Define shapes (dynamic) - auto M = int(m); - auto N = int(n); - auto K = int(k); - - // Define strides (mixed) - auto dA = make_stride(Int<1>{}, ldA); - auto dB = make_stride(Int<1>{}, ldB); - auto dC = make_stride(Int<1>{}, ldC); - - // Define block sizes (static) - auto bM = Int<128>{}; - auto bN = Int<128>{}; - auto bK = Int< 8>{}; - - // Define the block layouts (static) - auto sA = make_layout(make_shape(bM,bK)); - auto sB = make_layout(make_shape(bN,bK)); - auto sC = make_layout(make_shape(bM,bN)); - - // Define the thread layouts (static) - auto tA = make_layout(make_shape(Int<32>{}, Int< 8>{})); - auto tB = make_layout(make_shape(Int<32>{}, Int< 8>{})); - auto tC = make_layout(make_shape(Int<16>{}, Int<16>{})); - - dim3 dimBlock(size(tC)); - dim3 dimGrid(ceil_div(size(M), size(bM)), - ceil_div(size(N), size(bN))); - gemm_device - <<< dimGrid, dimBlock, 0, stream >>> - (M, N, K, - A, dA, sA, tA, - B, dB, sB, tB, - C, dC, sC, tC, - alpha, beta); -} - -#include -#include -#include - -void test_gemm(int m, int n, int k) -{ - cute::device_init(0); - - std::cout << "M = " << m << std::endl; - std::cout << "N = " << n << std::endl; - std::cout << "K = " << k << std::endl; - - using TA = float; - using TB = float; - using TC = float; - using TI = float; - - thrust::host_vector h_A(m*k); - thrust::host_vector h_B(n*k); - thrust::host_vector h_C(m*n); - - for (int j = 0; j < m*k; ++j) h_A[j] = static_cast( 2*(rand() / double(RAND_MAX)) - 1 ); - for (int j = 0; j < n*k; ++j) h_B[j] = static_cast( 2*(rand() / double(RAND_MAX)) - 1 ); - for (int j = 0; j < m*n; ++j) h_C[j] = static_cast(-1); - - thrust::device_vector d_A = h_A; - thrust::device_vector d_B = h_B; - thrust::device_vector d_C = h_C; - - TI alpha = 1.0; - TI beta = 0.0; - - double gflops = (2.0*m*n*k) * 1e-9; - - const int timing_iterations = 100; - GPU_Clock timer; - -#if defined(CUTLASS_ENABLE_CUBLAS) && CUTLASS_ENABLE_CUBLAS != 0 - // - // cuBLas - // - - cublasHandle_t handle; - cublasCreate(&handle); - - // Run once - d_C = h_C; - blam::cublas::gemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, - m, n, k, - &alpha, - d_A.data().get(), m, - d_B.data().get(), n, - &beta, - d_C.data().get(), m); - CUTE_CHECK_LAST(); - - thrust::host_vector cublas_result = d_C; - - // Timing iterations - timer.start(); - for (int i = 0; i < timing_iterations; ++i) { - blam::cublas::gemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, - m, n, k, - &alpha, - d_A.data().get(), m, - d_B.data().get(), n, - &beta, - d_C.data().get(), m); - } - double cublas_time = timer.seconds() / timing_iterations; - CUTE_CHECK_LAST(); - printf("CUBLAS_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cublas_time, cublas_time*1000); - -#else - - std::cout << "Verification by comparison with cuBLAS is disabled, " - "either because the CMake option CUTLASS_ENABLE_CUBLAS " - "was explicitly set to OFF, or because CMake could not find cuBLAS. " - "If you would like to enable verification with cuBLAS, " - "please set the CMake option CUTLASS_ENABLE_CUBLAS to ON, " - "rerun CMake, and recompile this example.\n"; - -#endif // CUTLASS_ENABLE_CUBLAS - - // - // CuTe - // - - // Run once (and check) - d_C = h_C; - gemm(m, n, k, - alpha, - d_A.data().get(), m, - d_B.data().get(), n, - beta, - d_C.data().get(), m); - CUTE_CHECK_LAST(); - thrust::host_vector cute_result = d_C; - - // Timing iterations - timer.start(); - for (int i = 0; i < timing_iterations; ++i) { - gemm(m, n, k, - alpha, - d_A.data().get(), m, - d_B.data().get(), n, - beta, - d_C.data().get(), m); - } - double cute_time = timer.seconds() / timing_iterations; - CUTE_CHECK_LAST(); - printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000); - -#if defined(CUTLASS_ENABLE_CUBLAS) && CUTLASS_ENABLE_CUBLAS != 0 - printf("Empirical Perf: %.1f%%\n", (cublas_time / cute_time) * 100); - - auto host_matrix_to_const_column_major_cute_tensor = - [](const auto& X, int num_rows, int num_cols, int LDX) { - const auto shape = cute::Shape{num_rows, num_cols}; - const auto strides = cute::Stride{1, LDX}; - return cute::make_tensor(X.data(), cute::make_layout(shape, strides)); - }; - - const auto A_view = host_matrix_to_const_column_major_cute_tensor(h_A, m, k, m); - // B^T is k x n, so B is n x k. - const auto B_view = host_matrix_to_const_column_major_cute_tensor(h_B, n, k, n); - const auto C_computed_view = host_matrix_to_const_column_major_cute_tensor(cute_result, m, n, m); - const auto C_expected_view = host_matrix_to_const_column_major_cute_tensor(cublas_result, m, n, m); - print_matrix_multiply_mollified_relative_error("float", A_view, B_view, C_computed_view, C_expected_view); - -#endif // CUTLASS_ENABLE_CUBLAS -} - - -int main(int argc, char** argv) -{ - int m = 5120; - if (argc >= 2) - sscanf(argv[1], "%d", &m); - - int n = 5120; - if (argc >= 3) - sscanf(argv[2], "%d", &n); - - int k = 4096; - if (argc >= 4) - sscanf(argv[3], "%d", &k); - - test_gemm(m, n, k); - - return 0; -} diff --git a/examples/cute/tutorial/sgemm_sm70.cu b/examples/cute/tutorial/sgemm_sm70.cu new file mode 100644 index 00000000..8aba8132 --- /dev/null +++ b/examples/cute/tutorial/sgemm_sm70.cu @@ -0,0 +1,526 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#include +#include +#include + +#include +#include + +#include + +#include "cutlass/util/print_error.hpp" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/helper_cuda.hpp" + +template +__global__ static +__launch_bounds__(decltype(size(TiledMma{}))::value) +void +gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler, + TA const* A, AStride dA, ASmemLayout sA_layout, TiledCopyA copy_a, + TB const* B, BStride dB, BSmemLayout sB_layout, TiledCopyB copy_b, + TC * C, CStride dC, CSmemLayout , TiledMma mma, + Alpha alpha, Beta beta) +{ + using namespace cute; + + // Preconditions + CUTE_STATIC_ASSERT_V(rank(shape_MNK) == Int<3>{}); // (M, N, K) + CUTE_STATIC_ASSERT_V(rank(cta_tiler) == Int<3>{}); // (BLK_M, BLK_N, BLK_K) + + CUTE_STATIC_ASSERT_V(size(copy_a) == size(mma)); // NumThreads + CUTE_STATIC_ASSERT_V(size(copy_b) == size(mma)); // NumThreads + + static_assert(is_static::value); + static_assert(is_static::value); + static_assert(is_static::value); + + CUTE_STATIC_ASSERT_V(size<0>(ASmemLayout{}) == size<0>(cta_tiler)); // BLK_M + CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<0>(cta_tiler)); // BLK_M + CUTE_STATIC_ASSERT_V(size<0>(BSmemLayout{}) == size<1>(cta_tiler)); // BLK_N + CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<1>(cta_tiler)); // BLK_N + CUTE_STATIC_ASSERT_V(size<1>(ASmemLayout{}) == size<2>(cta_tiler)); // BLK_K + CUTE_STATIC_ASSERT_V(size<1>(BSmemLayout{}) == size<2>(cta_tiler)); // BLK_K + + CUTE_STATIC_ASSERT_V(congruent(select<0,2>(shape_MNK), dA)); // dA strides for shape MK + CUTE_STATIC_ASSERT_V(congruent(select<1,2>(shape_MNK), dB)); // dB strides for shape NK + CUTE_STATIC_ASSERT_V(congruent(select<0,1>(shape_MNK), dC)); // dC strides for shape MN + + // + // Full and Tiled Tensors + // + + // Represent the full tensors + Tensor mA = make_tensor(make_gmem_ptr(A), select<0,2>(shape_MNK), dA); // (M,K) + Tensor mB = make_tensor(make_gmem_ptr(B), select<1,2>(shape_MNK), dB); // (N,K) + Tensor mC = make_tensor(make_gmem_ptr(C), select<0,1>(shape_MNK), dC); // (M,N) + + // Get the appropriate blocks for this thread block + auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k) + Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) + + // Shared memory buffers + __shared__ TA smemA[cosize_v]; + __shared__ TB smemB[cosize_v]; + Tensor sA = make_tensor(make_smem_ptr(smemA), sA_layout); // (BLK_M,BLK_K) + Tensor sB = make_tensor(make_smem_ptr(smemB), sB_layout); // (BLK_N,BLK_K) + + // + // Partition the copying of A and B tiles across the threads + // + + // TUTORIAL: Example of partitioning via a TiledCopy + + ThrCopy thr_copy_a = copy_a.get_slice(threadIdx.x); + Tensor tAgA = thr_copy_a.partition_S(gA); // (CPY,CPY_M,CPY_K,k) + Tensor tAsA = thr_copy_a.partition_D(sA); // (CPY,CPY_M,CPY_K) + Tensor tArA = make_fragment_like(tAsA); // (CPY,CPY_M,CPY_K) + + ThrCopy thr_copy_b = copy_b.get_slice(threadIdx.x); + Tensor tBgB = thr_copy_b.partition_S(gB); // (CPY,CPY_N,CPY_K,k) + Tensor tBsB = thr_copy_b.partition_D(sB); // (CPY,CPY_N,CPY_K) + Tensor tBrB = make_fragment_like(tBsB); // (CPY,CPY_N,CPY_K) + + CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tAsA)); // CPY_M + CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tArA)); // CPY_M + CUTE_STATIC_ASSERT_V(size<2>(tAgA) == size<2>(tAsA)); // CPY_K + CUTE_STATIC_ASSERT_V(size<2>(tAgA) == size<2>(tArA)); // CPY_K + CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBsB)); // CPY_N + CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBrB)); // CPY_N + CUTE_STATIC_ASSERT_V(size<2>(tBgB) == size<2>(tBsB)); // CPY_K + CUTE_STATIC_ASSERT_V(size<2>(tBgB) == size<2>(tBrB)); // CPY_K + + // Copy gmem to rmem for k_tile=0 + copy(copy_a, tAgA(_,_,_,0), tArA); + copy(copy_b, tBgB(_,_,_,0), tBrB); + // + // Define A/B partitioning and C accumulators + // + + // TUTORIAL: Example of partitioning via a TiledMMA + + ThrMMA thr_mma = mma.get_slice(threadIdx.x); + Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K) + Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K) + Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N) + + // Allocate registers for pipelining + 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) + // Allocate the accumulators -- same size as the projected data + Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N) + + CUTE_STATIC_ASSERT_V( shape(tCrA) == shape(tCsA)); // (MMA,MMA_M,MMA_K) + CUTE_STATIC_ASSERT_V( shape(tCrB) == shape(tCsB)); // (MMA,MMA_N,MMA_K) + CUTE_STATIC_ASSERT_V( shape(tCrC) == shape(tCgC)); // (MMA,MMA_M,MMA_N) + CUTE_STATIC_ASSERT_V(size<1>(tCgC) == size<1>(tCsA)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(tCgC) == size<1>(tCsB)); // MMA_N + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // MMA_K + + // Clear the accumulators + clear(tCrC); + +#if 0 + if(thread0()) { + print(" mA : "); print( mA); print("\n"); + print(" gA : "); print( gA); print("\n"); + print(" sA : "); print( sA); print("\n"); + print("tAgA : "); print(tAgA); print("\n"); + print("tAsA : "); print(tAsA); print("\n"); + print("tArA : "); print(tArA); print("\n"); + } +#endif + +#if 0 + if(thread0()) { + print(" mB : "); print( mB); print("\n"); + print(" gB : "); print( gB); print("\n"); + print(" sB : "); print( sB); print("\n"); + print("tBgB : "); print(tBgB); print("\n"); + print("tBsB : "); print(tBsB); print("\n"); + print("tArA : "); print(tArA); print("\n"); + } +#endif + +#if 0 + if(thread0()) { + print(" mC : "); print( mC); print("\n"); + print(" gC : "); print( gC); print("\n"); + print("tCsA : "); print(tCsA); print("\n"); + print("tCsB : "); print(tCsB); print("\n"); + print("tCgC : "); print(tCgC); print("\n"); + print("tCrC : "); print(tCrC); print("\n"); + } +#endif + +#if 1 + + // Copy rmem to smem + copy(tArA, tAsA); + copy(tBrB, tBsB); + __syncthreads(); + + // + // PIPELINED MAIN LOOP + // TUTORIAL: Example of a gemm loop that pipelines shared memory AND register memory + // Data is read from global to registers, then to shared via the tA|tB partitions + // Data is then copied from shared to registers in multiple waves via the tC partitions + // and gemm(.) operates on the current register wave + // + + // Load A, B shmem->regs for k_block=0 + copy(tCsA(_,_,0), tCrA(_,_,0)); + copy(tCsB(_,_,0), tCrB(_,_,0)); + auto K_TILE_MAX = size<3>(tAgA); + auto K_BLOCK_MAX = size<2>(tCrA); + + CUTE_NO_UNROLL + for (int k_tile = 0; k_tile < K_TILE_MAX; ++k_tile) + { + // Pipeline the k-mode of the block registers + CUTE_UNROLL + for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) + { + if (k_block == K_BLOCK_MAX - 1) + { + // Copy rmem to smem + __syncthreads(); + copy(tArA, tAsA); + copy(tBrB, tBsB); + __syncthreads(); + } + + // Copy smem to rmem for k_block+1 + int k_block_next = (k_block + 1) % K_BLOCK_MAX; + copy(tCsA(_,_,k_block_next), tCrA(_,_,k_block_next)); + copy(tCsB(_,_,k_block_next), tCrB(_,_,k_block_next)); + if (k_block == 0) + { + // Copy gmem to rmem for k_tile+1 + int k_tile_next = (k_tile + 1 < K_TILE_MAX) ? k_tile + 1 : k_tile; + copy(copy_a, tAgA(_,_,_,k_tile_next), tArA); + copy(copy_b, tBgB(_,_,_,k_tile_next), tBrB); + } + // Thread-level register gemm for k_block + gemm(mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC); + } // k_block + } // k_tile + +#endif + + // + // Epilogue + // + + axpby(alpha, tCrC, beta, tCgC); +} + +// Setup params for a NT GEMM +template +void +gemm_nt(int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + using namespace cute; + + // Define shapes (dynamic) + auto M = int(m); + auto N = int(n); + auto K = int(k); + auto prob_shape = make_shape(M, N, K); // (M, N, K) + + // Define NT strides (mixed) + auto dA = make_stride(Int<1>{}, ldA); // (dM, dK) + auto dB = make_stride(Int<1>{}, ldB); // (dN, dK) + auto dC = make_stride(Int<1>{}, ldC); // (dM, dN) + + // Define CTA tile sizes (static) + auto bM = Int<128>{}; + auto bN = Int<128>{}; + auto bK = Int< 8>{}; + auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K) + + // Define the smem layouts (static) + auto sA = make_layout(make_shape(bM, bK)); // (m,k) -> smem_idx; m-major + auto sB = make_layout(make_shape(bN, bK)); // (n,k) -> smem_idx; n-major + auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx; m-major + + // Define the thread layouts (static) + TiledCopy copyA = make_tiled_copy(Copy_Atom, TA>{}, + Layout>{}, // Thr layout 32x8 m-major + Layout>{}); // Val layout 4x1 m-major + TiledCopy copyB = make_tiled_copy(Copy_Atom, TB>{}, + Layout>{}, // Thr layout 32x8 n-major + Layout>{}); // Val layout 4x1 n-major + + TiledMMA mmaC = make_tiled_mma(UniversalFMA{}, + Layout>{}); // 16x16x1 TiledMMA + +#if 0 + print(copyA); + print(copyB); + print(mmaC); +#endif + +#if 0 + print_latex(copyA); + print_latex(copyB); + print_latex(mmaC); +#endif + + dim3 dimBlock(size(mmaC)); + dim3 dimGrid(size(ceil_div(M, bM)), + size(ceil_div(N, bN))); + gemm_device<<>> + (prob_shape, cta_tiler, + A, dA, sA, copyA, + B, dB, sB, copyB, + C, dC, sC, mmaC, + alpha, beta); +} + +// Setup params for a TN GEMM +template +void +gemm_tn(int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + using namespace cute; + + // Define shapes (dynamic) + auto M = int(m); + auto N = int(n); + auto K = int(k); + auto prob_shape = make_shape(M, N, K); // (M, N, K) + + // Define TN strides (mixed) + auto dA = make_stride(ldA, Int<1>{}); // (dM, dK) + auto dB = make_stride(ldB, Int<1>{}); // (dN, dK) + auto dC = make_stride(Int<1>{}, ldC); // (dM, dN) + + // Define CTA tile sizes (static) + auto bM = Int<128>{}; + auto bN = Int<128>{}; + auto bK = Int< 8>{}; + auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K) + + // Define the smem layouts (static) + auto sA = make_layout(make_shape ( bM, bK), + make_stride(Int<1>{}, bM+Int<1>{})); // (m,k) -> smem_idx; padded m-major + auto sB = make_layout(make_shape ( bN, bK), + make_stride(Int<1>{}, bN+Int<1>{})); // (n,k) -> smem_idx; padded n-major + auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx + + // Define the thread layouts (static) + + TiledCopy copyA = make_tiled_copy(Copy_Atom, TA>{}, + Layout,Stride<_8,_1>>{}, // Thr layout 32x8 k-major + Layout>{}); // Val layout 1x1 + TiledCopy copyB = make_tiled_copy(Copy_Atom, TB>{}, + Layout,Stride<_8,_1>>{}, // Thr layout 32x8 k-major + Layout>{}); // Val layout 1x1 + + TiledMMA mmaC = make_tiled_mma(UniversalFMA{}, + Layout>{}); // 16x16x1 TiledMMA + +#if 0 + print(copyA); + print(copyB); + print(mmaC); +#endif + +#if 0 + print_latex(copyA); + print_latex(copyB); + print_latex(mmaC); +#endif + + dim3 dimBlock(size(mmaC)); + dim3 dimGrid(size(ceil_div(M, bM)), + size(ceil_div(N, bN))); + gemm_device<<>> + (prob_shape, cta_tiler, + A, dA, sA, copyA, + B, dB, sB, copyB, + C, dC, sC, mmaC, + alpha, beta); +} + +template +void +gemm(char transA, char transB, int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + if (transA == 'N' && transB == 'T') { + return gemm_nt(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream); + } else + if (transA == 'T' && transB == 'N') { + return gemm_tn(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream); + } + assert(false && "Not implemented"); +} + + +int main(int argc, char** argv) +{ + cudaDeviceProp props; + cudaError_t error = cudaGetDeviceProperties(&props, 0); + if (error != cudaSuccess) { + std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl; + return -1; + } + + if (props.major < 7) { + std::cout << "This example requires an Volta GPU or newer (CC >= 70)" << std::endl; + // Return 0 so tests pass if run on unsupported architectures or CUDA Toolkits. + return 0; + } + + int m = 5120; + if (argc >= 2) + sscanf(argv[1], "%d", &m); + + int n = 5120; + if (argc >= 3) + sscanf(argv[2], "%d", &n); + + int k = 4096; + if (argc >= 4) + sscanf(argv[3], "%d", &k); + + char transA = 'N'; + if (argc >= 5) + sscanf(argv[4], "%c", &transA); + + char transB = 'T'; + if (argc >= 6) + sscanf(argv[5], "%c", &transB); + + using TA = float; + using TB = float; + using TC = float; + using TI = float; + + TI alpha = 1.0; + TI beta = 0.0; + + std::cout << "M = " << m << std::endl; + std::cout << "N = " << n << std::endl; + std::cout << "K = " << k << std::endl; + std::cout << "C = A^" << transA << " B^" << transB << std::endl; + + thrust::host_vector h_A(m*k); + thrust::host_vector h_B(n*k); + thrust::host_vector h_C(m*n); + + for (int j = 0; j < m*k; ++j) h_A[j] = static_cast( 2*(rand() / double(RAND_MAX)) - 1 ); + for (int j = 0; j < n*k; ++j) h_B[j] = static_cast( 2*(rand() / double(RAND_MAX)) - 1 ); + for (int j = 0; j < m*n; ++j) h_C[j] = static_cast(-1); + + thrust::device_vector d_A = h_A; + thrust::device_vector d_B = h_B; + thrust::device_vector d_C = h_C; + + double gflops = (2.0*m*n*k) * 1e-9; + + const int timing_iterations = 100; + GPU_Clock timer; + + int ldA = 0, ldB = 0, ldC = m; + + if (transA == 'N') { + ldA = m; + } else if (transA == 'T') { + ldA = k; + } else { + assert(false); + } + + if (transB == 'N') { + ldB = k; + } else if (transB == 'T') { + ldB = n; + } else { + assert(false); + } + + // Run once + d_C = h_C; + gemm(transA, transB, m, n, k, + alpha, + d_A.data().get(), ldA, + d_B.data().get(), ldB, + beta, + d_C.data().get(), ldC); + CUTE_CHECK_LAST(); + thrust::host_vector cute_result = d_C; + + // Timing iterations + timer.start(); + for (int i = 0; i < timing_iterations; ++i) { + gemm(transA, transB, m, n, k, + alpha, + d_A.data().get(), ldA, + d_B.data().get(), ldB, + beta, + d_C.data().get(), ldC); + } + double cute_time = timer.seconds() / timing_iterations; + CUTE_CHECK_LAST(); + printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000); + + return 0; +} diff --git a/examples/cute/tutorial/sgemm_sm80.cu b/examples/cute/tutorial/sgemm_sm80.cu new file mode 100644 index 00000000..3adb042b --- /dev/null +++ b/examples/cute/tutorial/sgemm_sm80.cu @@ -0,0 +1,567 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#include +#include +#include + +#include +#include + +#include + +#include "cutlass/util/print_error.hpp" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/helper_cuda.hpp" + +template +__global__ static +__launch_bounds__(decltype(size(TiledMma{}))::value) +void +gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler, + TA const* A, AStride dA, ASmemLayout sA_layout, TiledCopyA copy_a, + TB const* B, BStride dB, BSmemLayout sB_layout, TiledCopyB copy_b, + TC * C, CStride dC, CSmemLayout , TiledMma mma, + Alpha alpha, Beta beta) +{ + using namespace cute; + + // Preconditions + CUTE_STATIC_ASSERT_V(rank(shape_MNK) == Int<3>{}); // (M, N, K) + CUTE_STATIC_ASSERT_V(rank(cta_tiler) == Int<3>{}); // (BLK_M, BLK_N, BLK_K) + + CUTE_STATIC_ASSERT_V(size(copy_a) == size(mma)); // NumThreads + CUTE_STATIC_ASSERT_V(size(copy_b) == size(mma)); // NumThreads + + static_assert(is_static::value); + static_assert(is_static::value); + static_assert(is_static::value); + + CUTE_STATIC_ASSERT_V(size<0>(ASmemLayout{}) == size<0>(cta_tiler)); // BLK_M + CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<0>(cta_tiler)); // BLK_M + CUTE_STATIC_ASSERT_V(size<0>(BSmemLayout{}) == size<1>(cta_tiler)); // BLK_N + CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<1>(cta_tiler)); // BLK_N + CUTE_STATIC_ASSERT_V(size<1>(ASmemLayout{}) == size<2>(cta_tiler)); // BLK_K + CUTE_STATIC_ASSERT_V(size<1>(BSmemLayout{}) == size<2>(cta_tiler)); // BLK_K + + CUTE_STATIC_ASSERT_V(congruent(select<0,2>(shape_MNK), dA)); // dA strides for shape MK + CUTE_STATIC_ASSERT_V(congruent(select<1,2>(shape_MNK), dB)); // dB strides for shape NK + CUTE_STATIC_ASSERT_V(congruent(select<0,1>(shape_MNK), dC)); // dC strides for shape MN + + // + // Full and Tiled Tensors + // + + // Represent the full tensors + Tensor mA = make_tensor(make_gmem_ptr(A), select<0,2>(shape_MNK), dA); // (M,K) + Tensor mB = make_tensor(make_gmem_ptr(B), select<1,2>(shape_MNK), dB); // (N,K) + Tensor mC = make_tensor(make_gmem_ptr(C), select<0,1>(shape_MNK), dC); // (M,N) + + // Get the appropriate blocks for this thread block + auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k) + Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) + + // Shared memory buffers + __shared__ TA smemA[cosize_v]; + __shared__ TB smemB[cosize_v]; + Tensor sA = make_tensor(make_smem_ptr(smemA), sA_layout); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(smemB), sB_layout); // (BLK_N,BLK_K,PIPE) + + // + // Partition the copying of A and B tiles across the threads + // + + ThrCopy thr_copy_a = copy_a.get_slice(threadIdx.x); + Tensor tAgA = thr_copy_a.partition_S(gA); // (CPY,CPY_M,CPY_K,k) + Tensor tAsA = thr_copy_a.partition_D(sA); // (CPY,CPY_M,CPY_K,PIPE) + + ThrCopy thr_copy_b = copy_b.get_slice(threadIdx.x); + Tensor tBgB = thr_copy_b.partition_S(gB); // (CPY,CPY_N,CPY_K,k) + Tensor tBsB = thr_copy_b.partition_D(sB); // (CPY,CPY_N,CPY_K,PIPE) + + CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tAsA)); // CPY_M + CUTE_STATIC_ASSERT_V(size<2>(tAgA) == size<2>(tAsA)); // CPY_K + CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBsB)); // CPY_N + CUTE_STATIC_ASSERT_V(size<2>(tBgB) == size<2>(tBsB)); // CPY_K + + // + // PREFETCH + // + + auto K_PIPE_MAX = size<3>(tAsA); + + // Total count of tiles + int k_tile_count = size<3>(tAgA); + // Current tile index in gmem to read from + int k_tile_next = 0; + + // Start async loads for all pipes but the last + CUTE_UNROLL + for (int k_pipe = 0; k_pipe < K_PIPE_MAX-1; ++k_pipe) { + copy(copy_a, tAgA(_,_,_,k_tile_next), tAsA(_,_,_,k_pipe)); + copy(copy_b, tBgB(_,_,_,k_tile_next), tBsB(_,_,_,k_pipe)); + cp_async_fence(); + --k_tile_count; + if (k_tile_count > 0) { ++k_tile_next; } + } + + // + // Define A/B partitioning and C accumulators + // + + ThrMMA thr_mma = mma.get_slice(threadIdx.x); + 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 tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N) + + // Allocate registers for pipelining + Tensor tCrA = thr_mma.make_fragment_A(tCsA(_,_,_,0)); // (MMA,MMA_M,MMA_K) + Tensor tCrB = thr_mma.make_fragment_B(tCsB(_,_,_,0)); // (MMA,MMA_N,MMA_K) + // Allocate the accumulators -- same size as the projected data + Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N) + + CUTE_STATIC_ASSERT_V( shape(tCrA) == shape(tCsA)); // (MMA,MMA_M,MMA_K) + CUTE_STATIC_ASSERT_V( shape(tCrB) == shape(tCsB)); // (MMA,MMA_N,MMA_K) + CUTE_STATIC_ASSERT_V( shape(tCrC) == shape(tCgC)); // (MMA,MMA_M,MMA_N) + CUTE_STATIC_ASSERT_V(size<1>(tCgC) == size<1>(tCsA)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(tCgC) == size<1>(tCsB)); // MMA_N + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // MMA_K + + // Clear the accumulators + clear(tCrC); + +#if 0 + if(thread0()) { + print(" mA : "); print( mA); print("\n"); + print(" gA : "); print( gA); print("\n"); + print(" sA : "); print( sA); print("\n"); + print("tAgA : "); print(tAgA); print("\n"); + print("tAsA : "); print(tAsA); print("\n"); + } +#endif + +#if 0 + if(thread0()) { + print(" mB : "); print( mB); print("\n"); + print(" gB : "); print( gB); print("\n"); + print(" sB : "); print( sB); print("\n"); + print("tBgB : "); print(tBgB); print("\n"); + print("tBsB : "); print(tBsB); print("\n"); + } +#endif + +#if 0 + if(thread0()) { + print(" mC : "); print( mC); print("\n"); + print(" gC : "); print( gC); print("\n"); + print("tCsA : "); print(tCsA); print("\n"); + print("tCsB : "); print(tCsB); print("\n"); + print("tCgC : "); print(tCgC); print("\n"); + print("tCrA : "); print(tCrA); print("\n"); + print("tCrB : "); print(tCrB); print("\n"); + print("tCrC : "); print(tCrC); print("\n"); + } +#endif + +#if 1 + + // 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 = K_PIPE_MAX-1; + + // Pipe slice + 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(tCsA_p(_,_,Int<0>{}), tCrA(_,_,Int<0>{})); + copy(tCsB_p(_,_,Int<0>{}), tCrB(_,_,Int<0>{})); + } + + // + // PIPELINED MAIN LOOP + // TUTORIAL: Example of a gemm loop that pipelines shared memory using SM80's cp.async instructions + // and explicit pipelines in shared memory. + // Data is read from global(k_tile_next) to shared(smem_pipe_write). + // Data is read from shared(smem_pipe_read) to registers(k_block_next). + // Data is computed on registers(b_block). + // + // This allows all copies and compute to overlap: + // Copy from gmem->smem can overlap with copies from smem->rmem and compute on rmem. + // Copy from smem->rmem can overlap with compute on rmem. + // + + CUTE_NO_UNROLL + while (k_tile_count > -(K_PIPE_MAX-1)) + { + CUTE_UNROLL + for (int k_block = 0; k_block < K_BLOCK_MAX; ++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(tCsA_p(_,_,k_block_next), tCrA(_,_,k_block_next)); + copy(tCsB_p(_,_,k_block_next), tCrB(_,_,k_block_next)); + // Copy gmem to smem before computing gemm on each k-pipe + if (k_block == 0) + { + copy(copy_a, tAgA(_,_,_,k_tile_next), tAsA(_,_,_,smem_pipe_write)); + copy(copy_b, tBgB(_,_,_,k_tile_next), tBsB(_,_,_,smem_pipe_write)); + cp_async_fence(); + + // Advance the gmem tile + --k_tile_count; + if (k_tile_count > 0) { ++k_tile_next; } + + // Advance the smem pipe + smem_pipe_write = smem_pipe_read; + ++smem_pipe_read; + smem_pipe_read = (smem_pipe_read == K_PIPE_MAX) ? 0 : smem_pipe_read; + } + // Thread-level register gemm for k_block + gemm(mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC); + } + + } + +#endif + + // + // Epilogue + // + + axpby(alpha, tCrC, beta, tCgC); +} + +// Setup params for a NT GEMM +template +void +gemm_nt(int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + using namespace cute; + + // Define shapes (dynamic) + auto M = int(m); + auto N = int(n); + auto K = int(k); + auto prob_shape = make_shape(M, N, K); // (M, N, K) + + // Define NT strides (mixed) + auto dA = make_stride(Int<1>{}, ldA); // (dM, dK) + auto dB = make_stride(Int<1>{}, ldB); // (dN, dK) + auto dC = make_stride(Int<1>{}, ldC); // (dM, dN) + + // Define CTA tile sizes (static) + auto bM = Int<128>{}; + auto bN = Int<128>{}; + auto bK = Int< 8>{}; + auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K) + auto bP = Int<3>{}; // Pipeline + + // Define the smem layouts (static) + auto sA = make_layout(make_shape(bM, bK, bP)); // (m,k,p) -> smem_idx; m-major + auto sB = make_layout(make_shape(bN, bK, bP)); // (n,k,p) -> smem_idx; n-major + auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx; m-major + + // Define the thread layouts (static) + + TiledCopy copyA = make_tiled_copy(Copy_Atom, TA>{}, + Layout>{}, // Thr layout 32x8 m-major + Layout>{});// Val layout 4x1 m-major + TiledCopy copyB = make_tiled_copy(Copy_Atom, TB>{}, + Layout>{}, // Thr layout 32x8 n-major + Layout>{});// Val layout 4x1 n-major + + TiledMMA mmaC = make_tiled_mma(UniversalFMA{}, + Layout>{}); // 16x16x1 TiledMMA + +#if 0 + print(copyA); + print(copyB); + print(mmaC); +#endif + +#if 0 + print_latex(copyA); + print_latex(copyB); + print_latex(mmaC); +#endif + + dim3 dimBlock(size(mmaC)); + dim3 dimGrid(size(ceil_div(M, bM)), + size(ceil_div(N, bN))); + gemm_device<<>> + (prob_shape, cta_tiler, + A, dA, sA, copyA, + B, dB, sB, copyB, + C, dC, sC, mmaC, + alpha, beta); +} + +// Setup params for a NT GEMM +template +void +gemm_tn(int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + using namespace cute; + + // Define shapes (dynamic) + auto M = int(m); + auto N = int(n); + auto K = int(k); + auto prob_shape = make_shape(M, N, K); // (M, N, K) + + // Define TN strides (mixed) + auto dA = make_stride(ldA, Int<1>{}); // (dM, dK) + auto dB = make_stride(ldB, Int<1>{}); // (dN, dK) + auto dC = make_stride(Int<1>{}, ldC); // (dM, dN) + + // Define CTA tile sizes (static) + auto bM = Int<128>{}; + auto bN = Int<128>{}; + auto bK = Int< 8>{}; + auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K) + auto bP = Int<3>{}; // Pipeline + + // Define the smem layouts (static) + auto sA_atom = make_layout(make_shape ( bM, bK), + make_stride(Int<1>{}, bM+Int<1>{})); // (m,k) -> smem_idx; padded m-major + auto sB_atom = make_layout(make_shape ( bN, bK), + make_stride(Int<1>{}, bN+Int<1>{})); // (n,k) -> smem_idx; padded n-major + auto sA = tile_to_shape(sA_atom, make_shape(bM, bK, bP)); + auto sB = tile_to_shape(sA_atom, make_shape(bN, bK, bP)); + auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx + + // Define the thread layouts (static) + + TiledCopy copyA = make_tiled_copy(Copy_Atom, TA>{}, + Layout,Stride<_8,_1>>{}, // Thr layout 32x8 k-major + Layout>{}); // Val layout 1x1 + TiledCopy copyB = make_tiled_copy(Copy_Atom, TB>{}, + Layout,Stride<_8,_1>>{}, // Thr layout 32x8 k-major + Layout>{}); // Val layout 1x1 + + TiledMMA mmaC = make_tiled_mma(UniversalFMA{}, + Layout>{}); // 16x16x1 TiledMMA + +#if 0 + print(copyA); + print(copyB); + print(mmaC); +#endif + +#if 0 + print_latex(copyA); + print_latex(copyB); + print_latex(mmaC); +#endif + + dim3 dimBlock(size(mmaC)); + dim3 dimGrid(size(ceil_div(M, bM)), + size(ceil_div(N, bN))); + gemm_device<<>> + (prob_shape, cta_tiler, + A, dA, sA, copyA, + B, dB, sB, copyB, + C, dC, sC, mmaC, + alpha, beta); +} + +template +void +gemm(char transA, char transB, int m, int n, int k, + Alpha alpha, + TA const* A, int ldA, + TB const* B, int ldB, + Beta beta, + TC * C, int ldC, + cudaStream_t stream = 0) +{ + if (transA == 'N' && transB == 'T') { + return gemm_nt(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream); + } else + if (transA == 'T' && transB == 'N') { + return gemm_tn(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream); + } + assert(false && "Not implemented"); +} + + +int main(int argc, char** argv) +{ + cudaDeviceProp props; + cudaError_t error = cudaGetDeviceProperties(&props, 0); + if (error != cudaSuccess) { + std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl; + return -1; + } + + if (props.major < 8) { + std::cout << "This example requires an Ampere GPU or newer (CC >= 80)" << std::endl; + // Return 0 so tests pass if run on unsupported architectures or CUDA Toolkits. + return 0; + } + + int m = 5120; + if (argc >= 2) + sscanf(argv[1], "%d", &m); + + int n = 5120; + if (argc >= 3) + sscanf(argv[2], "%d", &n); + + int k = 4096; + if (argc >= 4) + sscanf(argv[3], "%d", &k); + + char transA = 'N'; + if (argc >= 5) + sscanf(argv[4], "%c", &transA); + + char transB = 'T'; + if (argc >= 6) + sscanf(argv[5], "%c", &transB); + + using TA = float; + using TB = float; + using TC = float; + using TI = float; + + TI alpha = 1.0; + TI beta = 0.0; + + std::cout << "M = " << m << std::endl; + std::cout << "N = " << n << std::endl; + std::cout << "K = " << k << std::endl; + std::cout << "C = A^" << transA << " B^" << transB << std::endl; + + thrust::host_vector h_A(m*k); + thrust::host_vector h_B(n*k); + thrust::host_vector h_C(m*n); + + for (int j = 0; j < m*k; ++j) h_A[j] = static_cast( 2*(rand() / double(RAND_MAX)) - 1 ); + for (int j = 0; j < n*k; ++j) h_B[j] = static_cast( 2*(rand() / double(RAND_MAX)) - 1 ); + for (int j = 0; j < m*n; ++j) h_C[j] = static_cast(-1); + + thrust::device_vector d_A = h_A; + thrust::device_vector d_B = h_B; + thrust::device_vector d_C = h_C; + + double gflops = (2.0*m*n*k) * 1e-9; + + const int timing_iterations = 100; + GPU_Clock timer; + + int ldA = 0, ldB = 0, ldC = m; + + if (transA == 'N') { + ldA = m; + } else if (transA == 'T') { + ldA = k; + } else { + assert(false); + } + + if (transB == 'N') { + ldB = k; + } else if (transB == 'T') { + ldB = n; + } else { + assert(false); + } + + // Run once + d_C = h_C; + gemm(transA, transB, m, n, k, + alpha, + d_A.data().get(), ldA, + d_B.data().get(), ldB, + beta, + d_C.data().get(), ldC); + CUTE_CHECK_LAST(); + thrust::host_vector cute_result = d_C; + + // Timing iterations + timer.start(); + for (int i = 0; i < timing_iterations; ++i) { + gemm(transA, transB, m, n, k, + alpha, + d_A.data().get(), ldA, + d_B.data().get(), ldB, + beta, + d_C.data().get(), ldC); + } + double cute_time = timer.seconds() / timing_iterations; + CUTE_CHECK_LAST(); + printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000); + + return 0; +} diff --git a/examples/cute/tutorial/tiled_copy.cu b/examples/cute/tutorial/tiled_copy.cu index 97f8897a..d370320b 100644 --- a/examples/cute/tutorial/tiled_copy.cu +++ b/examples/cute/tutorial/tiled_copy.cu @@ -67,7 +67,7 @@ // // Uses local_partition() to partition a tile among threads arranged as (THR_M, THR_N). template -__global__ void copy_kernel(TensorS S, TensorD D, ThreadLayout) +__global__ void copy_kernel(TensorS S, TensorD D, ThreadLayout) { using namespace cute; @@ -77,12 +77,13 @@ __global__ void copy_kernel(TensorS S, TensorD D, ThreadLayout) // Construct a partitioning of the tile among threads with the given thread arrangement. - // Concept: Tensor Layout Index - Tensor thr_tile_S = local_partition(tile_S, ThreadLayout{}, threadIdx.x); - Tensor thr_tile_D = local_partition(tile_D, ThreadLayout{}, threadIdx.x); + // Concept: Tensor ThrLayout ThrIndex + Tensor thr_tile_S = local_partition(tile_S, ThreadLayout{}, threadIdx.x); // (ThrValM, ThrValN) + Tensor thr_tile_D = local_partition(tile_D, ThreadLayout{}, threadIdx.x); // (ThrValM, ThrValN) // Construct a register-backed Tensor with the same shape as each thread's partition - auto fragment = make_fragment_like(thr_tile_S); + // Use make_tensor to try to match the layout of thr_tile_S + Tensor fragment = make_tensor_like(thr_tile_S); // (ThrValM, ThrValN) // Copy from GMEM to RMEM and from RMEM to GMEM copy(thr_tile_S, fragment); @@ -95,17 +96,17 @@ __global__ void copy_kernel(TensorS S, TensorD D, ThreadLayout) /// has the precondition that pointers are aligned to the vector size. /// template -__global__ void copy_kernel_vectorized(TensorS S, TensorD D, ThreadLayout, VecLayout) +__global__ void copy_kernel_vectorized(TensorS S, TensorD D, ThreadLayout, VecLayout) { using namespace cute; using Element = typename TensorS::value_type; // Slice the tensors to obtain a view into each tile. - Tensor tile_S = S(make_coord(_, _), blockIdx.x, blockIdx.y); // (BlockShape_M, BlockShape_N) - Tensor tile_D = D(make_coord(_, _), blockIdx.x, blockIdx.y); // (BlockShape_M, BlockShape_N) + Tensor tile_S = S(make_coord(_, _), blockIdx.x, blockIdx.y); // (BlockShape_M, BlockShape_N) + Tensor tile_D = D(make_coord(_, _), blockIdx.x, blockIdx.y); // (BlockShape_M, BlockShape_N) // Define `AccessType` which controls the size of the actual memory access. - using AccessType = cutlass::AlignedArray; + using AccessType = cutlass::AlignedArray; // A copy atom corresponds to one hardware memory access. using Atom = Copy_Atom, Element>; @@ -125,29 +126,18 @@ __global__ void copy_kernel_vectorized(TensorS S, TensorD D, ThreadLayout, VecLa // Construct a Tensor corresponding to each thread's slice. auto thr_copy = tiled_copy.get_thread_slice(threadIdx.x); - Tensor thr_tile_S = thr_copy.partition_S(tile_S); - Tensor thr_tile_D = thr_copy.partition_D(tile_D); + Tensor thr_tile_S = thr_copy.partition_S(tile_S); // (CopyOp, CopyM, CopyN) + Tensor thr_tile_D = thr_copy.partition_D(tile_D); // (CopyOp, CopyM, CopyN) // Construct a register-backed Tensor with the same shape as each thread's partition - auto fragment = make_fragment_like(thr_tile_D); + // Use make_fragment because the first mode is the instruction-local mode + Tensor fragment = make_fragment_like(thr_tile_D); // (CopyOp, CopyM, CopyN) // Copy from GMEM to RMEM and from RMEM to GMEM copy(tiled_copy, thr_tile_S, fragment); copy(tiled_copy, fragment, thr_tile_D); } -/// Helper to convert a shape to a dim3 -template -dim3 shape_to_dim3(Shape shape) -{ - using namespace cute; - - CUTE_STATIC_ASSERT_V(rank(shape) <= Int<3>{}); - auto result = append<3>(product_each(shape), 1u); - - return dim3(get<0>(result), get<1>(result), get<2>(result)); -} - /// Main function int main(int argc, char** argv) { @@ -161,13 +151,13 @@ int main(int argc, char** argv) // Define a tensor shape with dynamic extents (m, n) auto tensor_shape = make_shape(256, 512); + // + // Allocate and initialize + // + thrust::host_vector h_S(size(tensor_shape)); thrust::host_vector h_D(size(tensor_shape)); - // - // Initialize - // - for (size_t i = 0; i < h_S.size(); ++i) { h_S[i] = static_cast(i); h_D[i] = Element{}; @@ -180,33 +170,36 @@ int main(int argc, char** argv) // Make tensors // - Tensor tensor_S = make_tensor(make_gmem_ptr(d_S.data().get()), make_layout(tensor_shape)); - Tensor tensor_D = make_tensor(make_gmem_ptr(d_D.data().get()), make_layout(tensor_shape)); + Tensor tensor_S = make_tensor(make_gmem_ptr(thrust::raw_pointer_cast(d_S.data())), make_layout(tensor_shape)); + Tensor tensor_D = make_tensor(make_gmem_ptr(thrust::raw_pointer_cast(d_D.data())), make_layout(tensor_shape)); // - // Partition + // Tile tensors // - // Define a statically sized block (M, N). - // // Note, by convention, capital letters are used to represent static modes. auto block_shape = make_shape(Int<128>{}, Int<64>{}); - if ((get<0>(tensor_shape) % get<0>(block_shape)) || (get<1>(tensor_shape) % get<1>(block_shape))) { + if ((size<0>(tensor_shape) % size<0>(block_shape)) || (size<1>(tensor_shape) % size<1>(block_shape))) { std::cerr << "The tensor shape must be divisible by the block shape." << std::endl; return -1; } + // Equivalent check to the above + if (not weakly_compatible(block_shape, tensor_shape)) { + std::cerr << "Expected the tensors to be weakly compatible with the block_shape." << std::endl; + return -1; + } - // Tile the tensor (m, m) ==> ((M, N), m', n') where (M, N) is the static tile + // Tile the tensor (m, n) ==> ((M, N), m', n') where (M, N) is the static tile // shape, and modes (m', n') correspond to the number of tiles. - // - // These will be used to determine the CUDA kernel grid dimensinos. - Tensor tiled_tensor_S = tiled_divide(tensor_S, block_shape); - Tensor tiled_tensor_D = tiled_divide(tensor_D, block_shape); + // + // These will be used to determine the CUDA kernel grid dimensions. + Tensor tiled_tensor_S = tiled_divide(tensor_S, block_shape); // ((M, N), m', n') + Tensor tiled_tensor_D = tiled_divide(tensor_D, block_shape); // ((M, N), m', n') // Thread arrangement - Layout thr_layout = make_layout(make_shape(Int<32>{}, Int< 8>{})); + Layout thr_layout = make_layout(make_shape(Int<32>{}, Int<8>{})); // Vector dimensions Layout vec_layout = make_layout(make_shape(Int<4>{}, Int<1>{})); @@ -215,16 +208,16 @@ int main(int argc, char** argv) // Determine grid and block dimensions // - dim3 gridDim = shape_to_dim3(select<1,2>(shape(tiled_tensor_D))); // Grid shape corresponds to modes m' and n' - dim3 blockDim(size(shape(thr_layout))); + dim3 gridDim (size<1>(tiled_tensor_D), size<2>(tiled_tensor_D)); // Grid shape corresponds to modes m' and n' + dim3 blockDim(size(thr_layout)); // // Launch the kernel // copy_kernel_vectorized<<< gridDim, blockDim >>>( - tiled_tensor_S, - tiled_tensor_D, - thr_layout, + tiled_tensor_S, + tiled_tensor_D, + thr_layout, vec_layout); cudaError result = cudaDeviceSynchronize(); diff --git a/include/cute/algorithm/axpby.hpp b/include/cute/algorithm/axpby.hpp index 11b4d0da..df9605b7 100644 --- a/include/cute/algorithm/axpby.hpp +++ b/include/cute/algorithm/axpby.hpp @@ -33,6 +33,7 @@ #include #include +#include namespace cute { @@ -43,15 +44,17 @@ namespace cute template + class YEngine, class YLayout, + class PrdTensor = TrivialPredTensor> CUTE_HOST_DEVICE void axpby(Alpha const& alpha, Tensor const& x, Beta const& beta, - Tensor && y) + Tensor && y, + PrdTensor const& p = {}) { - return axpby(alpha, x, beta, y); + return axpby(alpha, x, beta, y, p); } // @@ -60,13 +63,15 @@ axpby(Alpha const& alpha, template + class YEngine, class YLayout, + class PrdTensor = TrivialPredTensor> CUTE_HOST_DEVICE void axpby(Alpha const& alpha, Tensor const& x, Beta const& beta, - Tensor & y) + Tensor & y, + PrdTensor const& p = {}) { auto isBetaZero = [&] () { if constexpr (is_complex::value) { @@ -81,7 +86,9 @@ axpby(Alpha const& alpha, CUTE_UNROLL for (int i = 0; i < size(x); ++i) { - y(i) = (isBetaZero ? alpha * x(i) : alpha * x(i) + beta * y(i)); + if (p(i)) { + y(i) = (isBetaZero ? alpha * x(i) : alpha * x(i) + beta * y(i)); + } } } diff --git a/include/cute/algorithm/cooperative_copy.hpp b/include/cute/algorithm/cooperative_copy.hpp new file mode 100644 index 00000000..c0337aba --- /dev/null +++ b/include/cute/algorithm/cooperative_copy.hpp @@ -0,0 +1,196 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 +{ + +// cooperative_copy(thr_idx, src, dst) +// Use NumThreads to copy src to dst with element vectorization up to MaxVecBits. +// @pre 0 <= @a tid < NumThreads +// @pre Tensors @a src and @a dst are aligned up to MaxVecBits. +// +template +CUTE_HOST_DEVICE +void +cooperative_copy(uint32_t const& tid, + Tensor const& src, + Tensor & dst) +{ + // Assumes the shapes are static, can generalize + CUTE_STATIC_ASSERT_V(size(src) == size(dst)); + // Assumes the types are the same, can generalize + static_assert(sizeof_bits_v == sizeof_bits_v); + static_assert(MaxVecBits == sizeof_bits_v || + MaxVecBits == 8 || MaxVecBits == 16 || MaxVecBits == 32 || MaxVecBits == 64 || MaxVecBits == 128, + "Expected MaxVecBits to be value size or 8 or 16 or 32 or 64 or 128 for alignment and performance."); + // Check that the tensors are likely shared across threads: either gmem or smem + static_assert((is_gmem::value || is_smem::value), + "cooperative_copy expects shared gmem or smem source tensor."); + static_assert((is_gmem::value || is_smem::value), + "cooperative_copy expects shared gmem or smem destination tensor."); + + // Precondition on tid in DEBUG + assert(tid < NumThreads); + // Precondition on pointer alignment in DEBUG + assert(is_byte_aligned(raw_pointer_cast(src.data()))); + assert(is_byte_aligned(raw_pointer_cast(dst.data()))); + // + // Determine val+thr vectorization based on src/dst size and number of threads + // NOTE: This heuristic promotes parallelization over vectorization + // + + constexpr int elem_bits = sizeof_bits_v; + + // The number of elements that can be vectorized in values + constexpr int common_elem = decltype(max_common_vector(src, dst))::value; + constexpr int common_bits = common_elem * elem_bits; + constexpr int total_elem = decltype(size(src))::value; + constexpr int total_bits = total_elem * elem_bits; + static_assert(total_bits % NumThreads == 0); + constexpr int total_bits_per_thr = total_bits / NumThreads; + // If there are too many threads to allow a full elem copy, trunc the thrs and use elem_bits + constexpr int max_vec_bits_by_thr = cute::max(elem_bits, total_bits_per_thr); + + // Cap the vectorization to the common bits, the max_vec_bits_by_thr, and the MaxVecBits + constexpr int vec_bits = cute::min(common_bits, max_vec_bits_by_thr, static_cast(MaxVecBits)); + // Convert back to number of elements, safe_div + static_assert((vec_bits % elem_bits) == 0); + constexpr int vec_elem = vec_bits / elem_bits; + + // Use only part of threads if there's not enough work for all threads + constexpr int vec_thrs = (total_elem % (vec_elem * NumThreads) == 0) + ? NumThreads + : (total_elem / vec_elem); + + // The common layout of the two tensors that can be vectorized over threads + // vidx -> coord + auto common_layout = max_common_layout(get_nonswizzle_portion(src.layout()), + get_nonswizzle_portion(dst.layout())); + + // Scale up the common_layout to cover the entire tensors + // vidx -> coord + auto full_perm = tile_to_shape(make_layout(common_layout), size(src)); + + // Create the Tiler + // ((vid,tid),iter) + auto layout_vt = logical_divide(full_perm, Layout, Int>>{}); + + // Apply and slice + Tensor src_v = src.compose(layout_vt)(make_coord(_,tid),_); + Tensor dst_v = dst.compose(layout_vt)(make_coord(_,tid),_); + + // Should account for vec_bits < 8 and/or vec_elem <= 1 + // And also account for subbyte types, which could cause race conditions + // Want to ENFORCE sufficient vectorization in those cases + static_assert((vec_bits >= 8), "No support for subbyte copying"); + using VecType = uint_bit_t; + +#if 0 + if (thread0()) { + print(" "); print("NumThreads: "); print(NumThreads); print("\n"); + print(" "); print("src: "); print(src); print("\n"); + print(" "); print("dst: "); print(dst); print("\n"); + print(" "); print("common_layout: "); print(common_layout); print("\n"); + print(" "); print("full_perm: "); print(full_perm); print("\n"); + print(" "); print("Used vector: "); print(vec_elem); print("\n"); + print(" "); print("Used threads: "); print(vec_thrs); print("\n"); + print(" "); print("layout_vt: "); print(layout_vt); print("\n"); + print(" "); print("src.compose(layout_vt): "); print(src.compose(layout_vt)); print("\n"); + print(" "); print("dst.compose(layout_vt): "); print(dst.compose(layout_vt)); print("\n"); + print(" "); print("src_v: "); print(src_v); print("\n"); + print(" "); print("dst_v: "); print(dst_v); print("\n"); + print(" "); print("recast(src_v): "); print(recast(src_v)); print("\n"); + print(" "); print("recast(dst_v): "); print(recast(dst_v)); print("\n"); + } +#ifdef __CUDA_ARCH__ + __syncthreads(); +#endif +#endif + + // If we're using all threads (static) or the tid is in in-range (dynamic) + if (vec_thrs >= NumThreads or tid < vec_thrs) { + return copy_if(TrivialPredTensor{}, recast(src_v), recast(dst_v)); + } +} + +template +CUTE_HOST_DEVICE +void +cooperative_copy(uint32_t const& tid, + Tensor const& src, + Tensor & dst) +{ + constexpr uint32_t MaxVecBits = sizeof_bits_v; + return cooperative_copy(tid, src, dst); +} + +// Accept mutable temporaries +template +CUTE_HOST_DEVICE +void +cooperative_copy(uint32_t const& tid, + Tensor const& src, + Tensor && dst) +{ + return cooperative_copy(tid, src, dst); +} + +// Accept mutable temporaries +template +CUTE_HOST_DEVICE +void +cooperative_copy(uint32_t const& tid, + Tensor const& src, + Tensor && dst) +{ + return cooperative_copy(tid, src, dst); +} + +} // end namespace cute diff --git a/include/cute/algorithm/cooperative_gemm.hpp b/include/cute/algorithm/cooperative_gemm.hpp new file mode 100644 index 00000000..32cec54b --- /dev/null +++ b/include/cute/algorithm/cooperative_gemm.hpp @@ -0,0 +1,326 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 +{ + +// +// Collective Shared-Memory GEMMs +// + +template ::value && + BLayout::rank == 2 && is_smem::value && + CLayout::rank == 2 && is_smem::value)> +CUTE_HOST_DEVICE +void +cooperative_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: "); print(rounded_sA); print("\n"); + print("rounded_sB: "); print(rounded_sB); print("\n"); + print("rounded_sC: "); print(rounded_sC); 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: "); print(tCsA); print("\n"); + print("tCsB: "); print(tCsB); print("\n"); + print("tCsC: "); print(tCsC); print("\n"); + print("tCrA: "); print(tCrA); print("\n"); + print("tCrB: "); print(tCrB); print("\n"); + print("tCrC: "); print(tCrC); 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 * static_cast(tCrC(i,m,n)) : alpha * static_cast(tCrC(i,m,n)) + beta * static_cast(tCsC(i,m,n)); + } + } + } + } +} + +template ::value && + BLayout::rank == 2 && is_smem::value && + CLayout::rank == 2 && is_smem::value)> +CUTE_HOST_DEVICE +void +cooperative_gemm(ThrMMA const& thr_mma, + Alpha const& alpha, + Tensor sA, + Tensor sB, + Beta const& beta, + Tensor sC) +{ + cooperative_gemm(thr_mma, alpha, sA, sB, beta, sC, identity() /* sA_load_op */, identity() /* sB_load_op */); +} + +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 */) +{ + cooperative_gemm(thr_mma, alpha, sA, sB, beta, sC, sA_load_op, sB_load_op); +} + +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) +{ + cooperative_gemm(thr_mma, alpha, sA, sB, beta, sC, identity() /* sA_load_op */, identity() /* sB_load_op */); +} + +} // end namespace cute diff --git a/include/cute/algorithm/copy.hpp b/include/cute/algorithm/copy.hpp index 65c2caf6..50a092d0 100644 --- a/include/cute/algorithm/copy.hpp +++ b/include/cute/algorithm/copy.hpp @@ -145,10 +145,10 @@ copy_if(PrdTensor const& pred, namespace detail { // Trait that detects if atom's traits has a member function with(bool) -template +template constexpr bool has_with_bool = false; -template +template constexpr bool has_with_bool().with(declval()))>> = true; } // end namespace detail diff --git a/include/cute/algorithm/functional.hpp b/include/cute/algorithm/functional.hpp index c6713310..8e7a58a5 100644 --- a/include/cute/algorithm/functional.hpp +++ b/include/cute/algorithm/functional.hpp @@ -33,6 +33,7 @@ #include #include +#include /** C++14 extensions */ @@ -46,7 +47,7 @@ struct identity { template CUTE_HOST_DEVICE constexpr decltype(auto) operator()(T&& arg) const { - return std::forward(arg); + return static_cast(arg); } }; @@ -69,7 +70,7 @@ struct constant_fn { template \ CUTE_HOST_DEVICE constexpr \ decltype(auto) operator()(T&& arg) const { \ - return OP std::forward(arg); \ + return OP static_cast(arg); \ } \ } #define CUTE_RIGHT_UNARY_OP(NAME,OP) \ @@ -77,7 +78,7 @@ struct constant_fn { template \ CUTE_HOST_DEVICE constexpr \ decltype(auto) operator()(T&& arg) const { \ - return std::forward(arg) OP ; \ + return static_cast(arg) OP ; \ } \ } #define CUTE_NAMED_UNARY_OP(NAME,OP) \ @@ -85,7 +86,7 @@ struct constant_fn { template \ CUTE_HOST_DEVICE constexpr \ decltype(auto) operator()(T&& arg) const { \ - return OP (std::forward(arg)); \ + return OP (static_cast(arg)); \ } \ } @@ -115,7 +116,7 @@ struct shift_right_const { template CUTE_HOST_DEVICE constexpr decltype(auto) operator()(T&& arg) const { - return std::forward(arg) >> Shift; + return static_cast(arg) >> Shift; } }; @@ -126,7 +127,7 @@ struct shift_left_const { template CUTE_HOST_DEVICE constexpr decltype(auto) operator()(T&& arg) const { - return std::forward(arg) << Shift; + return static_cast(arg) << Shift; } }; @@ -139,7 +140,7 @@ struct shift_left_const { template \ CUTE_HOST_DEVICE constexpr \ decltype(auto) operator()(T&& lhs, U&& rhs) const { \ - return std::forward(lhs) OP std::forward(rhs); \ + return static_cast(lhs) OP static_cast(rhs); \ } \ } #define CUTE_NAMED_BINARY_OP(NAME,OP) \ @@ -147,7 +148,7 @@ struct shift_left_const { template \ CUTE_HOST_DEVICE constexpr \ decltype(auto) operator()(T&& lhs, U&& rhs) const { \ - return OP (std::forward(lhs), std::forward(rhs)); \ + return OP (static_cast(lhs), static_cast(rhs)); \ } \ } @@ -273,7 +274,7 @@ struct bound_fn { CUTE_HOST_DEVICE constexpr decltype(auto) operator()(T&& arg) { - return fn_(arg_, std::forward(arg)); + return fn_(arg_, static_cast(arg)); } Fn fn_; diff --git a/include/cute/algorithm/gemm.hpp b/include/cute/algorithm/gemm.hpp index b5ade49e..27c32216 100644 --- a/include/cute/algorithm/gemm.hpp +++ b/include/cute/algorithm/gemm.hpp @@ -252,7 +252,7 @@ gemm(MMA_Atom const& 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) @@ -451,6 +451,7 @@ gemm(MMA_Atom const& 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) @@ -496,245 +497,4 @@ gemm(MMA_Atom const& mma, } } -// -// 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/include/cute/algorithm/prefetch.hpp b/include/cute/algorithm/prefetch.hpp new file mode 100644 index 00000000..47aefa87 --- /dev/null +++ b/include/cute/algorithm/prefetch.hpp @@ -0,0 +1,153 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 +{ + +// +// Prefetch global tensors into L2 +// + +template +CUTE_HOST_DEVICE +void +cooperative_prefetch(uint32_t const& tid, + Tensor const& src) +{ + static_assert(is_gmem::value, "Expected global tensor for prefetch"); + + constexpr int V = decltype(max_common_vector(src, src))::value; + + if constexpr (V > 1) { + // L2 sector is 32B, default fetch granularity is 64B + using VecType = conditional_t<(V * sizeof_bits_v) < (FetchBytes * 8), + ArrayEngine, + uint8_t[FetchBytes] >; + + Tensor src_v = recast(src); + CUTE_UNROLL + for (int i = tid; i < size(src_v); i += NumThreads) { + prefetch(raw_pointer_cast(&src_v(i))); + } + } else { + CUTE_UNROLL + for (int i = tid; i < size(src); i += NumThreads) { + prefetch(raw_pointer_cast(&src(i))); + } + } +} + +template +CUTE_HOST_DEVICE +void +prefetch(Tensor const& src) +{ + return cooperative_prefetch<1>(0, src); +} + +// Prefetch with copy atom +namespace detail { + +template +constexpr bool has_prefetch = false; + +template +constexpr bool has_prefetch> = true; + +template +constexpr bool is_prefetch = false; + +template +constexpr bool is_prefetch> = is_same_v; + +} // end namespace detail + +template +CUTE_HOST_DEVICE +void +prefetch(Copy_Atom, CA_Args...> const& atom, + Tensor const& src) +{ + if constexpr (detail::has_prefetch) { + using Prefetch_Traits = Copy_Traits; + using Prefetch_Atom = Copy_Atom; + Prefetch_Atom prefetch_atom{atom}; + auto& dst = const_cast&>(src); // dst is ignored for prefetch atoms + return copy(prefetch_atom, src, dst); + } else { + return prefetch(src); + } +} + +#if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED) +template +CUTE_HOST_DEVICE +void +prefetch(Copy_Traits const& atom, + Tensor const& src) +{ + using SrcType = typename SrcEngine::value_type; + static_assert(is_gmem::value, "Expected global tensor for L2 prefetch"); + + auto tiler = max_common_layout(src, src); + constexpr int vec_elem = decltype(size(tiler))::value; + constexpr int vec_bits = vec_elem * sizeof_bits_v; + static_assert(vec_bits >= 128, "Expected at least 128-bits for BLKCP"); + + // Construct a new concrete Atom of the vector size + auto bulk_atom = Copy_Atom>, SrcType>{}; + + return prefetch(bulk_atom, logical_divide(src, tiler)); +} + +// Backwards-compat. Throw out any extra Copy_Atom args. +template +CUTE_HOST_DEVICE +void +prefetch(Copy_Atom, CA_Args...> const& atom, + Tensor const& src) +{ + return prefetch(static_cast const&>(atom), src); +} +#endif // #if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED) + +} // end namespace cute diff --git a/include/cute/algorithm/tensor_algorithms.hpp b/include/cute/algorithm/tensor_algorithms.hpp index f77d73f2..c37cdca0 100644 --- a/include/cute/algorithm/tensor_algorithms.hpp +++ b/include/cute/algorithm/tensor_algorithms.hpp @@ -50,7 +50,7 @@ for_each(Tensor const& tensor, UnaryOp&& op) { CUTE_UNROLL for (int i = 0; i < size(tensor); ++i) { - static_cast(op)(tensor(i)); + op(tensor(i)); } } @@ -61,7 +61,7 @@ for_each(Tensor& tensor, UnaryOp&& op) { CUTE_UNROLL for (int i = 0; i < size(tensor); ++i) { - static_cast(op)(tensor(i)); + op(tensor(i)); } } @@ -71,7 +71,7 @@ CUTE_HOST_DEVICE constexpr void for_each(Tensor&& tensor, UnaryOp&& op) { - return for_each(tensor, static_cast(op)); + return for_each(tensor, op); } // @@ -86,7 +86,7 @@ transform(Tensor& tensor, UnaryOp&& op) { CUTE_UNROLL for (int i = 0; i < size(tensor); ++i) { - tensor(i) = static_cast(op)(tensor(i)); + tensor(i) = op(tensor(i)); } } @@ -96,27 +96,34 @@ CUTE_HOST_DEVICE constexpr void transform(Tensor&& tensor, UnaryOp&& op) { - return transform(tensor, std::forward(op)); + return transform(tensor, op); } // Similar to std::transform transforms one tensors and assigns it to another -template +template CUTE_HOST_DEVICE constexpr void -transform(Tensor& tensor_in, Tensor& tensor_out, UnaryOp&& op) +transform(Tensor const& 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)); + tensor_out(i) = op(tensor_in(i)); } } // Accept mutable temporaries template + class EngineOut, class LayoutOut, + class UnaryOp> CUTE_HOST_DEVICE constexpr void -transform(Tensor&& tensor_in, Tensor&& tensor_out, UnaryOp&& op) +transform(Tensor const& tensor_in, + Tensor && tensor_out, + UnaryOp&& op) { return transform(tensor_in, tensor_out, op); } @@ -127,29 +134,31 @@ transform(Tensor&& tensor_in, Tensor&& t // assigns it to tensor_out template + class EngineOut, class LayoutOut, + class BinaryOp> CUTE_HOST_DEVICE constexpr void -transform(Tensor& tensor_in1, - Tensor& tensor_in2, - Tensor& tensor_out, +transform(Tensor const& tensor_in1, + Tensor const& tensor_in2, + Tensor & tensor_out, BinaryOp&& op) { CUTE_UNROLL for (int i = 0; i < size(tensor_in1); ++i) { - tensor_out(i) = static_cast(op)(tensor_in1(i), tensor_in2(i)); + tensor_out(i) = op(tensor_in1(i), tensor_in2(i)); } } // Accept mutable temporaries template + class EngineOut, class LayoutOut, + class BinaryOp> CUTE_HOST_DEVICE constexpr void -transform(Tensor&& tensor_in1, - Tensor&& tensor_in2, - Tensor&& tensor_out, +transform(Tensor const& tensor_in1, + Tensor const& tensor_in2, + Tensor && tensor_out, BinaryOp&& op) { return transform(tensor_in1, tensor_in2, tensor_out, op); diff --git a/include/cute/algorithm/tuple_algorithms.hpp b/include/cute/algorithm/tuple_algorithms.hpp index 2ada8b27..3157e897 100644 --- a/include/cute/algorithm/tuple_algorithms.hpp +++ b/include/cute/algorithm/tuple_algorithms.hpp @@ -204,36 +204,6 @@ for_each_leaf(T&& t, F&& f) CUTE_GCC_UNREACHABLE; } -// -// For Sequence -// (s, t, f) => (f(t[s_0]),f(t[s_1]),...,f(t[s_n])) -// - -namespace detail { - -template -CUTE_HOST_DEVICE constexpr -void -for_sequence(seq const&, F&& f) { - (f(Int{}), ...); -} - -}; // end namespace detail - -template -CUTE_HOST_DEVICE constexpr -void -for_sequence(seq const& s, T&& t, F&& f) { - detail::for_sequence(s, [&](auto&& i){ f(get::value>(static_cast(t))); }); -} - -template -CUTE_HOST_DEVICE constexpr -void -for_sequence(T&& t, F&& f) { - for_sequence(make_seq{}, static_cast(t), static_cast(f)); -} - // // Transform // (t, f) => (f(t_0),f(t_1),...,f(t_n)) @@ -551,15 +521,15 @@ take(T const& t) template CUTE_HOST_DEVICE constexpr auto -select(T const & t) +select(T const& t) { return cute::make_tuple(get(t)...); } -template +template CUTE_HOST_DEVICE constexpr auto -select(T const & t, Indices const & indices) +select(T const& t, Indices const& indices) { if constexpr (is_tuple::value) { return cute::transform(indices, [&t](auto i) { return select(t, i); }); @@ -655,7 +625,7 @@ flatten(T const& t) namespace detail { -template +template CUTE_HOST_DEVICE constexpr auto unflatten_impl(FlatTuple const& flat_tuple, TargetProfile const& target_profile) @@ -680,7 +650,7 @@ unflatten_impl(FlatTuple const& flat_tuple, TargetProfile const& target_profile) // @pre rank(flatten(@a target_profile)) == rank(@a flat_tuple) // @post congruent(@a result, @a target_profile) // @post flatten(@a result) == @a flat_tuple -template +template CUTE_HOST_DEVICE constexpr auto unflatten(FlatTuple const& flat_tuple, TargetProfile const& target_profile) @@ -865,6 +835,7 @@ append(T const& a, X const& x) CUTE_GCC_UNREACHABLE; } + template CUTE_HOST_DEVICE constexpr auto @@ -902,6 +873,7 @@ prepend(T const& a, X const& x) CUTE_GCC_UNREACHABLE; } + template CUTE_HOST_DEVICE constexpr auto @@ -1105,14 +1077,13 @@ zip2_by(T const& t, TG const& guide) /// @return A tuple of the elements of @c t in reverse order. template -CUTE_HOST_DEVICE constexpr auto -reverse(T const& t) { +CUTE_HOST_DEVICE constexpr +auto +reverse(T const& t) +{ if constexpr (is_tuple::value) { - return detail::apply(t, [] (auto const&... a) { - return cute::make_tuple(a...); - }, tuple_rseq{}); - } - else { + return detail::apply(t, [](auto const&... a){ return cute::make_tuple(a...); }, tuple_rseq{}); + } else { return t; } } diff --git a/include/cute/arch/cluster_sm90.hpp b/include/cute/arch/cluster_sm90.hpp index a720c374..27a34d77 100644 --- a/include/cute/arch/cluster_sm90.hpp +++ b/include/cute/arch/cluster_sm90.hpp @@ -49,7 +49,7 @@ 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"); + CUTE_INVALID_CONTROL_PATH("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined"); #endif } @@ -58,7 +58,7 @@ 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"); + CUTE_INVALID_CONTROL_PATH("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined"); #endif } @@ -67,7 +67,7 @@ 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"); + CUTE_INVALID_CONTROL_PATH("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined"); #endif } @@ -77,7 +77,7 @@ CUTE_DEVICE void cluster_sync() cluster_arrive(); cluster_wait(); #else - CUTE_RUNTIME_ASSERT("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined"); + CUTE_INVALID_CONTROL_PATH("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined"); #endif } @@ -94,7 +94,7 @@ CUTE_DEVICE dim3 cluster_grid_dims() // MSVC requires protecting use of gridDim with __CUDA_ARCH__. return gridDim; #elif defined(_MSC_VER) - CUTE_RUNTIME_ASSERT("cluster_grid_dims() can only be called on device"); + CUTE_INVALID_CONTROL_PATH("cluster_grid_dims() can only be called on device"); return {0, 0, 0}; #else return {0, 0, 0}; @@ -114,7 +114,7 @@ CUTE_DEVICE dim3 cluster_id_in_grid() // MSVC requires protecting use of blockIdx with __CUDA_ARCH__. return blockIdx; #elif defined(_MSC_VER) - CUTE_RUNTIME_ASSERT("cluster_id_in_grid() can only be called on device"); + CUTE_INVALID_CONTROL_PATH("cluster_id_in_grid() can only be called on device"); return {0, 0, 0}; #else return {0, 0, 0}; diff --git a/include/cute/arch/copy.hpp b/include/cute/arch/copy.hpp index 5266c6c1..b85e6a20 100644 --- a/include/cute/arch/copy.hpp +++ b/include/cute/arch/copy.hpp @@ -33,7 +33,7 @@ #include #include -#include +#include namespace cute { @@ -89,4 +89,17 @@ using AutoVectorizingCopy = AutoVectorizingCopyWithAssumedAlignment<8>; // Alias using DefaultCopy = AutoVectorizingCopy; + +// +// Global memory prefetch into L2 +// + +CUTE_HOST_DEVICE static void +prefetch(void const* gmem_ptr) +{ +#if defined(__CUDA_ARCH__) + asm volatile("prefetch.global.L2 [%0];\n" : : "l"(gmem_ptr) : "memory"); +#endif +} + } // end namespace cute diff --git a/include/cute/arch/copy_sm75.hpp b/include/cute/arch/copy_sm75.hpp index c3fe3b49..3d3d37ac 100644 --- a/include/cute/arch/copy_sm75.hpp +++ b/include/cute/arch/copy_sm75.hpp @@ -78,7 +78,7 @@ struct SM75_U32x1_LDSM_N : "=r"(dst) : "r"(smem_int_ptr)); #else - CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); + CUTE_INVALID_CONTROL_PATH("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); #endif } }; @@ -98,7 +98,7 @@ struct SM75_U32x2_LDSM_N : "=r"(dst0), "=r"(dst1) : "r"(smem_int_ptr)); #else - CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); + CUTE_INVALID_CONTROL_PATH("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); #endif } }; @@ -118,7 +118,7 @@ struct SM75_U32x4_LDSM_N : "=r"(dst0), "=r"(dst1), "=r"(dst2), "=r"(dst3) : "r"(smem_int_ptr)); #else - CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); + CUTE_INVALID_CONTROL_PATH("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); #endif } }; @@ -138,7 +138,7 @@ struct SM75_U16x2_LDSM_T : "=r"(dst) : "r"(smem_int_ptr)); #else - CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); + CUTE_INVALID_CONTROL_PATH("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); #endif } }; @@ -158,7 +158,7 @@ struct SM75_U16x4_LDSM_T : "=r"(dst0), "=r"(dst1) : "r"(smem_int_ptr)); #else - CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); + CUTE_INVALID_CONTROL_PATH("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); #endif } }; @@ -178,7 +178,7 @@ struct SM75_U16x8_LDSM_T : "=r"(dst0), "=r"(dst1), "=r"(dst2), "=r"(dst3) : "r"(smem_int_ptr)); #else - CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); + CUTE_INVALID_CONTROL_PATH("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); #endif } }; diff --git a/include/cute/arch/copy_sm80.hpp b/include/cute/arch/copy_sm80.hpp index 05011804..43e3d0d7 100644 --- a/include/cute/arch/copy_sm80.hpp +++ b/include/cute/arch/copy_sm80.hpp @@ -64,7 +64,7 @@ struct SM80_CP_ASYNC_CACHEALWAYS "l"(gmem_ptr), "n"(sizeof(TS))); #else - CUTE_RUNTIME_ASSERT("Support for cp.async instructions has not been enabled"); + CUTE_INVALID_CONTROL_PATH("Support for cp.async instructions has not been enabled"); #endif } }; @@ -91,7 +91,7 @@ struct SM80_CP_ASYNC_CACHEGLOBAL "l"(gmem_ptr), "n"(sizeof(TS))); #else - CUTE_RUNTIME_ASSERT("Support for cp.async instructions has not been enabled"); + CUTE_INVALID_CONTROL_PATH("Support for cp.async instructions has not been enabled"); #endif } }; @@ -121,7 +121,7 @@ struct SM80_CP_ASYNC_CACHEALWAYS_ZFILL "n"(sizeof(TS)), "r"(src_size)); #else - CUTE_RUNTIME_ASSERT("Support for cp.async instructions has not been enabled"); + CUTE_INVALID_CONTROL_PATH("Support for cp.async instructions has not been enabled"); #endif } }; @@ -151,7 +151,7 @@ struct SM80_CP_ASYNC_CACHEGLOBAL_ZFILL "n"(sizeof(TS)), "r"(src_size)); #else - CUTE_RUNTIME_ASSERT("Support for cp.async instructions has not been enabled"); + CUTE_INVALID_CONTROL_PATH("Support for cp.async instructions has not been enabled"); #endif } }; diff --git a/include/cute/arch/copy_sm90.hpp b/include/cute/arch/copy_sm90.hpp index f51b1fed..e5684ec4 100644 --- a/include/cute/arch/copy_sm90.hpp +++ b/include/cute/arch/copy_sm90.hpp @@ -63,7 +63,7 @@ struct SM90_U32x1_STSM_N :: "r"(smem_int_ptr), "r"(src)); #else - CUTE_RUNTIME_ASSERT("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); #endif } }; @@ -83,7 +83,7 @@ struct SM90_U32x2_STSM_N :: "r"(smem_int_ptr), "r"(src0), "r"(src1)); #else - CUTE_RUNTIME_ASSERT("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); #endif } }; @@ -103,7 +103,7 @@ struct SM90_U32x4_STSM_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."); + CUTE_INVALID_CONTROL_PATH("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); #endif } }; @@ -123,7 +123,7 @@ struct SM90_U16x2_STSM_T :: "r"(smem_int_ptr), "r"(src)); #else - CUTE_RUNTIME_ASSERT("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); #endif } }; @@ -143,7 +143,7 @@ struct SM90_U16x4_STSM_T :: "r"(smem_int_ptr), "r"(src0), "r"(src1)); #else - CUTE_RUNTIME_ASSERT("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); #endif } }; @@ -163,7 +163,7 @@ struct SM90_U16x8_STSM_T :: "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."); + CUTE_INVALID_CONTROL_PATH("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); #endif } }; diff --git a/include/cute/arch/copy_sm90_desc.hpp b/include/cute/arch/copy_sm90_desc.hpp index 9fa4f34b..856d4dd5 100644 --- a/include/cute/arch/copy_sm90_desc.hpp +++ b/include/cute/arch/copy_sm90_desc.hpp @@ -43,8 +43,7 @@ #include #include #include -#include // to_Format<[u]intX> -#include // to_Format +#include namespace cute { @@ -177,8 +176,10 @@ to_CUtensorMapSwizzle(SmemSwizzleBits const& t) { #if (__CUDACC_VER_MAJOR__ >= 12) && !defined(__CUDACC_RTC__) using TmaDescriptor = CUtensorMap; + using Im2ColTmaDescriptor = CUtensorMap; #else using TmaDescriptor = struct alignas(64) { char bytes[128]; }; + using Im2ColTmaDescriptor = struct alignas(64) { char bytes[128]; }; #endif //////////////////////////////////////////////////////////////////////////////////////////////////// /// Initiates a TensorMap Prefetch @@ -197,7 +198,7 @@ prefetch_tma_descriptor(TmaDescriptor const* desc_ptr) : "l"(gmem_int_desc) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use TMA Descriptor Prefetch without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use TMA Descriptor Prefetch without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } @@ -208,7 +209,7 @@ prefetch_tma_descriptor(TmaDescriptor const* desc_ptr) // Replace tensor pointer directly in GMEM CUTE_HOST_DEVICE void -tma_descriptor_replace_addr_in_global_mem(TmaDescriptor const* desc_ptr, +tma_descriptor_replace_addr_in_global_mem(TmaDescriptor const* desc_ptr, void const* const new_tensor_ptr) { #if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED) @@ -218,14 +219,14 @@ tma_descriptor_replace_addr_in_global_mem(TmaDescriptor const* desc_ptr, "tensormap.replace.tile.global_address.global.b1024.b64 [%0], %1;" :: "l"(gmem_int_desc), "l"(new_desc_addr)); #else - CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); + CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); #endif } // Replace tensor pointer by bringing the tensormap from GMEM into the shared memory CUTE_HOST_DEVICE void -tma_descriptor_replace_addr_in_shared_mem(TmaDescriptor& smem_desc, +tma_descriptor_replace_addr_in_shared_mem(TmaDescriptor& smem_desc, void const* const new_tensor_ptr) { #if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED) @@ -239,7 +240,7 @@ tma_descriptor_replace_addr_in_shared_mem(TmaDescriptor& smem_desc, "tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;" :: "l"(smem_int64_desc), "l"(new_desc_addr)); #else - CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); + CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); #endif } @@ -273,7 +274,7 @@ tma_descriptor_replace_dims_strides_in_shared_mem(TmaDescriptor "tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 1, %1;" :: "l"(smem_int64_desc), "l"(prob_stride[2] >> 4)); #else - CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); + CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); #endif } @@ -292,7 +293,7 @@ tma_descriptor_cp_fence_release(TmaDescriptor const* gmem_desc_ptr, TmaDescripto "tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned [%0], [%1], 128;" :: "l"(gmem_int_desc), "r"(smem_int_desc)); #else - CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); + CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); #endif } @@ -307,7 +308,7 @@ tma_descriptor_fence_release() #if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED) asm volatile ("fence.proxy.tensormap::generic.release.gpu;"); #else - CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); + CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); #endif } @@ -332,7 +333,7 @@ tma_descriptor_fence_acquire(TmaDescriptor const* desc_ptr) : "l"(gmem_int_desc), "l"(gmem_int_desc) : "memory"); #else - CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); + CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3"); #endif } diff --git a/include/cute/arch/copy_sm90_tma.hpp b/include/cute/arch/copy_sm90_tma.hpp index 23c6d951..1136c433 100644 --- a/include/cute/arch/copy_sm90_tma.hpp +++ b/include/cute/arch/copy_sm90_tma.hpp @@ -44,13 +44,13 @@ namespace cute 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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" @@ -60,21 +60,42 @@ struct SM90_TMA_LOAD_1D "r"(crd0) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } + + struct PREFETCH + { + CUTE_HOST_DEVICE static void + copy(void const* desc_ptr, + int32_t const& crd0) + { + #if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + asm volatile ( + "cp.async.bulk.prefetch.tensor.1d.L2.global" + " [%0, {%1}];" + : + : "l"(gmem_int_desc), + "r"(crd0) + : "memory"); + #else + CUTE_INVALID_CONTROL_PATH("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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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" @@ -84,21 +105,42 @@ struct SM90_TMA_LOAD_2D "r"(crd0), "r"(crd1) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } + + struct PREFETCH + { + CUTE_HOST_DEVICE static void + copy(void const* desc_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); + asm volatile ( + "cp.async.bulk.prefetch.tensor.2d.L2.global" + " [%0, {%1, %2}];" + : + : "l"(gmem_int_desc), + "r"(crd0), "r"(crd1) + : "memory"); + #else + CUTE_INVALID_CONTROL_PATH("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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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" @@ -108,21 +150,42 @@ struct SM90_TMA_LOAD_3D "r"(crd0), "r"(crd1), "r"(crd2) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } + + struct PREFETCH + { + CUTE_HOST_DEVICE static void + copy(void const* desc_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); + asm volatile ( + "cp.async.bulk.prefetch.tensor.3d.L2.global" + " [%0, {%1, %2, %3}];" + : + : "l"(gmem_int_desc), + "r"(crd0), "r"(crd1), "r"(crd2) + : "memory"); + #else + CUTE_INVALID_CONTROL_PATH("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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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" @@ -132,21 +195,42 @@ struct SM90_TMA_LOAD_4D "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } + + struct PREFETCH + { + CUTE_HOST_DEVICE static void + copy(void const* desc_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); + asm volatile ( + "cp.async.bulk.prefetch.tensor.4d.L2.global" + " [%0, {%1, %2, %3, %4}];" + : + : "l"(gmem_int_desc), + "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3) + : "memory"); + #else + CUTE_INVALID_CONTROL_PATH("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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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" @@ -156,48 +240,103 @@ struct SM90_TMA_LOAD_5D "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."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } + + struct PREFETCH + { + CUTE_HOST_DEVICE static void + copy(void const* desc_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); + asm volatile ( + "cp.async.bulk.prefetch.tensor.5d.L2.global" + " [%0, {%1, %2, %3, %4, %5}];" + : + : "l"(gmem_int_desc), + "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4) + : "memory"); + #else + CUTE_INVALID_CONTROL_PATH("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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * smem_ptr, int32_t const& crd0) { - return SM90_TMA_LOAD_1D::copy(desc_ptr, smem_mbar, smem_ptr, crd0); + return SM90_TMA_LOAD_1D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0); } CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, uint64_t& smem_mbar, - void const* const smem_ptr, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * smem_ptr, int32_t const& crd0, int32_t const& crd1) { - return SM90_TMA_LOAD_2D::copy(desc_ptr, smem_mbar, smem_ptr, crd0, crd1); + return SM90_TMA_LOAD_2D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0, crd1); } CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, uint64_t& smem_mbar, - void const* const smem_ptr, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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); + return SM90_TMA_LOAD_3D::copy(desc_ptr, mbar_ptr, 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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); + return SM90_TMA_LOAD_4D::copy(desc_ptr, mbar_ptr, 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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); + return SM90_TMA_LOAD_5D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0, crd1, crd2, crd3, crd4); } + + struct PREFETCH + { + CUTE_HOST_DEVICE static void + copy(void const* desc_ptr, + int32_t const& crd0) + { + return SM90_TMA_LOAD_1D::PREFETCH::copy(desc_ptr, crd0); + } + CUTE_HOST_DEVICE static void + copy(void const* desc_ptr, + int32_t const& crd0, int32_t const& crd1) + { + return SM90_TMA_LOAD_2D::PREFETCH::copy(desc_ptr, crd0, crd1); + } + CUTE_HOST_DEVICE static void + copy(void const* desc_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2) + { + return SM90_TMA_LOAD_3D::PREFETCH::copy(desc_ptr, crd0, crd1, crd2); + } + CUTE_HOST_DEVICE static void + copy(void const* desc_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3) + { + return SM90_TMA_LOAD_4D::PREFETCH::copy(desc_ptr, crd0, crd1, crd2, crd3); + } + CUTE_HOST_DEVICE static void + copy(void const* desc_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::PREFETCH::copy(desc_ptr, crd0, crd1, crd2, crd3, crd4); + } + }; }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -207,14 +346,14 @@ struct SM90_TMA_LOAD 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); // Copy from global to shared::cluster. asm volatile ( @@ -226,23 +365,45 @@ struct SM90_TMA_LOAD_IM2COL_3D "h"(offset_w) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } + + struct PREFETCH + { + CUTE_HOST_DEVICE static void + copy(void const* desc_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); + asm volatile ( + "cp.async.bulk.prefetch.tensor.3d.L2.global.im2col" + " [%0, {%1, %2, %3}], {%4};" + : + : "l"(gmem_int_desc), + "r"(coord_c), "r"(coord_w), "r"(coord_n), + "h"(offset_w) + : "memory"); + #else + CUTE_INVALID_CONTROL_PATH("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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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) + 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); // Copy from global to shared::cluster. asm volatile ( @@ -254,24 +415,45 @@ struct SM90_TMA_LOAD_IM2COL_4D "h"(offset_w), "h"(offset_h) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } + + struct PREFETCH + { + CUTE_HOST_DEVICE static void + copy(void const* desc_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); + asm volatile ( + "cp.async.bulk.prefetch.tensor.4d.L2.global.im2col" + " [%0, {%1, %2, %3, %4}], {%5, %6};" + : + : "l"(gmem_int_desc), + "r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n), + "h"(offset_w), "h"(offset_h) + : "memory"); + #else + CUTE_INVALID_CONTROL_PATH("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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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) + 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); // Copy from global to shared::cluster. asm volatile ( @@ -283,48 +465,97 @@ struct SM90_TMA_LOAD_IM2COL_5D "h"(offset_w), "h"(offset_h), "h"(offset_d) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } + + struct PREFETCH + { + CUTE_HOST_DEVICE static void + copy(void const* desc_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); + asm volatile ( + "cp.async.bulk.prefetch.tensor.5d.L2.global.im2col" + " [%0, {%1, %2, %3, %4, %5}], {%6, %7, %8};" + : + : "l"(gmem_int_desc), + "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_INVALID_CONTROL_PATH("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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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, + return SM90_TMA_LOAD_IM2COL_3D::copy(desc_ptr, mbar_ptr, 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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) + 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); + return SM90_TMA_LOAD_IM2COL_4D::copy(desc_ptr, mbar_ptr, smem_ptr, + coord_c, coord_w, coord_h, coord_n, + offset_w, offset_h); + } + CUTE_HOST_DEVICE static void + copy(void const* desc_ptr, uint64_t* mbar_ptr, + void * 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, mbar_ptr, smem_ptr, + coord_c, coord_w, coord_h, coord_d, coord_n, + offset_w, offset_h, offset_d); } - 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) + struct PREFETCH { - 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); - } + CUTE_HOST_DEVICE static void + copy(void const* desc_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::PREFETCH::copy(desc_ptr, + coord_c, coord_w, coord_n, + offset_w); + } + CUTE_HOST_DEVICE static void + copy(void const* desc_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::PREFETCH::copy(desc_ptr, + coord_c, coord_w, coord_h, coord_n, + offset_w, offset_h); + } + CUTE_HOST_DEVICE static void + copy(void const* desc_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::PREFETCH::copy(desc_ptr, + coord_c, coord_w, coord_h, coord_d, coord_n, + offset_w, offset_h, offset_d); + } + }; }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -334,13 +565,13 @@ struct SM90_TMA_LOAD_IM2COL 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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" @@ -351,7 +582,7 @@ struct SM90_TMA_LOAD_MULTICAST_1D "r"(crd0) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -359,13 +590,13 @@ struct SM90_TMA_LOAD_MULTICAST_1D 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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" @@ -376,7 +607,7 @@ struct SM90_TMA_LOAD_MULTICAST_2D "r"(crd0), "r"(crd1) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -384,13 +615,13 @@ struct SM90_TMA_LOAD_MULTICAST_2D 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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" @@ -401,7 +632,7 @@ struct SM90_TMA_LOAD_MULTICAST_3D "r"(crd0), "r"(crd1), "r"(crd2) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -409,13 +640,13 @@ struct SM90_TMA_LOAD_MULTICAST_3D 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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" @@ -426,7 +657,7 @@ struct SM90_TMA_LOAD_MULTICAST_4D "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -434,13 +665,13 @@ struct SM90_TMA_LOAD_MULTICAST_4D 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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" @@ -451,7 +682,7 @@ struct SM90_TMA_LOAD_MULTICAST_5D "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."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -459,40 +690,42 @@ struct SM90_TMA_LOAD_MULTICAST_5D 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * smem_ptr, int32_t const& crd0) { - return SM90_TMA_LOAD_MULTICAST_1D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0); + return SM90_TMA_LOAD_MULTICAST_1D::copy(desc_ptr, mbar_ptr, 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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); + return SM90_TMA_LOAD_MULTICAST_2D::copy(desc_ptr, mbar_ptr, 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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); + return SM90_TMA_LOAD_MULTICAST_3D::copy(desc_ptr, mbar_ptr, 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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); + return SM90_TMA_LOAD_MULTICAST_4D::copy(desc_ptr, mbar_ptr, 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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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); + return SM90_TMA_LOAD_MULTICAST_5D::copy(desc_ptr, mbar_ptr, multicast_mask, smem_ptr, crd0, crd1, crd2, crd3, crd4); } + + using PREFETCH = typename SM90_TMA_LOAD::PREFETCH; }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -502,14 +735,14 @@ struct SM90_TMA_LOAD_MULTICAST struct SM90_TMA_LOAD_IM2COL_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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); // Copy from global to shared::cluster. asm volatile ( @@ -519,10 +752,10 @@ struct SM90_TMA_LOAD_IM2COL_MULTICAST_3D : "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) + "h"(multicast_mask) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -530,14 +763,14 @@ struct SM90_TMA_LOAD_IM2COL_MULTICAST_3D struct SM90_TMA_LOAD_IM2COL_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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); // Copy from global to shared::cluster. asm volatile ( @@ -547,10 +780,10 @@ struct SM90_TMA_LOAD_IM2COL_MULTICAST_4D : "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) + "h"(multicast_mask) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -558,14 +791,14 @@ struct SM90_TMA_LOAD_IM2COL_MULTICAST_4D struct SM90_TMA_LOAD_IM2COL_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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); // Copy from global to shared::cluster. asm volatile ( @@ -575,10 +808,10 @@ struct SM90_TMA_LOAD_IM2COL_MULTICAST_5D : "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) + "h"(multicast_mask) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -586,40 +819,42 @@ struct SM90_TMA_LOAD_IM2COL_MULTICAST_5D struct SM90_TMA_LOAD_IM2COL_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, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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, + return SM90_TMA_LOAD_IM2COL_MULTICAST_3D::copy(desc_ptr, mbar_ptr, 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 multicast_mask, - void const* const smem_ptr, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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, + return SM90_TMA_LOAD_IM2COL_MULTICAST_4D::copy(desc_ptr, mbar_ptr, 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 multicast_mask, - void const* const smem_ptr, + copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, + void * 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, + return SM90_TMA_LOAD_IM2COL_MULTICAST_5D::copy(desc_ptr, mbar_ptr, multicast_mask, smem_ptr, coord_c, coord_w, coord_h, coord_d, coord_n, offset_w, offset_h, offset_d); } + + using PREFETCH = typename SM90_TMA_LOAD_IM2COL::PREFETCH; }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -629,8 +864,8 @@ struct SM90_TMA_LOAD_IM2COL_MULTICAST struct SM90_TMA_STORE_1D { CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void const* smem_ptr, int32_t const& crd0) { #if defined(CUTE_ARCH_TMA_SM90_ENABLED) @@ -643,7 +878,7 @@ struct SM90_TMA_STORE_1D "r"(crd0) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -651,8 +886,8 @@ struct SM90_TMA_STORE_1D struct SM90_TMA_STORE_2D { CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void const* smem_ptr, int32_t const& crd0, int32_t const& crd1) { #if defined(CUTE_ARCH_TMA_SM90_ENABLED) @@ -665,7 +900,7 @@ struct SM90_TMA_STORE_2D "r"(crd0), "r"(crd1) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -673,8 +908,8 @@ struct SM90_TMA_STORE_2D struct SM90_TMA_STORE_3D { CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void const* smem_ptr, int32_t const& crd0, int32_t const& crd1, int32_t const& crd2) { #if defined(CUTE_ARCH_TMA_SM90_ENABLED) @@ -687,7 +922,7 @@ struct SM90_TMA_STORE_3D "r"(crd0), "r"(crd1), "r"(crd2) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -695,8 +930,8 @@ struct SM90_TMA_STORE_3D struct SM90_TMA_STORE_4D { CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void 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) @@ -709,7 +944,7 @@ struct SM90_TMA_STORE_4D "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -717,8 +952,8 @@ struct SM90_TMA_STORE_4D struct SM90_TMA_STORE_5D { CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void 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) @@ -731,7 +966,7 @@ struct SM90_TMA_STORE_5D "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."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -739,36 +974,36 @@ struct SM90_TMA_STORE_5D struct SM90_TMA_STORE { CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void 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, + copy(void const* desc_ptr, + void 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, + copy(void const* desc_ptr, + void 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, + copy(void const* desc_ptr, + void 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, + copy(void const* desc_ptr, + void 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); @@ -782,8 +1017,8 @@ struct SM90_TMA_STORE struct SM90_TMA_STORE_IM2COL_3D { CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void const* smem_ptr, int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n) { #if defined(CUTE_ARCH_TMA_SM90_ENABLED) @@ -797,7 +1032,7 @@ struct SM90_TMA_STORE_IM2COL_3D "r"(coord_c), "r"(coord_w), "r"(coord_n) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -805,8 +1040,8 @@ struct SM90_TMA_STORE_IM2COL_3D struct SM90_TMA_STORE_IM2COL_4D { CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void const* smem_ptr, int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n) { #if defined(CUTE_ARCH_TMA_SM90_ENABLED) @@ -820,7 +1055,7 @@ struct SM90_TMA_STORE_IM2COL_4D "r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -828,8 +1063,8 @@ struct SM90_TMA_STORE_IM2COL_4D struct SM90_TMA_STORE_IM2COL_5D { CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void 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) { #if defined(CUTE_ARCH_TMA_SM90_ENABLED) @@ -843,7 +1078,7 @@ struct SM90_TMA_STORE_IM2COL_5D "r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_d), "r"(coord_n) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; @@ -851,35 +1086,45 @@ struct SM90_TMA_STORE_IM2COL_5D struct SM90_TMA_STORE_IM2COL { CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void const* smem_ptr, int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n) { return SM90_TMA_STORE_IM2COL_3D::copy(desc_ptr, smem_ptr, coord_c, coord_w, coord_n); } CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void const* smem_ptr, int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n) { return SM90_TMA_STORE_IM2COL_4D::copy(desc_ptr, smem_ptr, coord_c, coord_w, coord_h, coord_n); } CUTE_HOST_DEVICE static void - copy(void const* const desc_ptr, - void const* const smem_ptr, + copy(void const* desc_ptr, + void 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) { return SM90_TMA_STORE_IM2COL_5D::copy(desc_ptr, smem_ptr, coord_c, coord_w, coord_h, coord_d, coord_n); } }; +// Fence for smem stores for subsequent TMA_STORE +CUTE_HOST_DEVICE static void +tma_store_fence() { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + asm volatile ("fence.proxy.async.shared::cta;"); +#elif defined(__CUDA_ARCH__) + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif +} + // 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."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } @@ -894,10 +1139,163 @@ tma_store_wait() { : "n"(Count) : "memory"); #else - CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } +//////////////////////////////////////////////////////////////////////////////////////////////////// +/// TMA_REDUCE_ADD : Initiates a TMA reduce-add from shared memory to global memory +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct SM90_TMA_REDUCE_ADD_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.reduce.async.bulk.tensor.1d.global.shared::cta.add.bulk_group [%0, {%2}], [%1];" + : + : "l"(gmem_int_desc), "r"(smem_int_ptr), + "r"(crd0) + : "memory"); +#else + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_REDUCE_ADD_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.reduce.async.bulk.tensor.2d.global.shared::cta.add.bulk_group [%0, {%2, %3}], [%1];" + : + : "l"(gmem_int_desc), "r"(smem_int_ptr), + "r"(crd0), "r"(crd1) + : "memory"); +#else + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_REDUCE_ADD_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.reduce.async.bulk.tensor.3d.global.shared::cta.add.bulk_group [%0, {%2, %3, %4}], [%1];" + : + : "l"(gmem_int_desc), "r"(smem_int_ptr), + "r"(crd0), "r"(crd1), "r"(crd2) + : "memory"); +#else + CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_REDUCE_ADD_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.reduce.async.bulk.tensor.4d.global.shared::cta.add.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_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_REDUCE_ADD_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.reduce.async.bulk.tensor.5d.global.shared::cta.add.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_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_REDUCE_ADD +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0) + { + return SM90_TMA_REDUCE_ADD_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_REDUCE_ADD_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_REDUCE_ADD_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_REDUCE_ADD_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_REDUCE_ADD_5D::copy(desc_ptr, smem_ptr, crd0, crd1, crd2, crd3, crd4); + } +}; + //////////////////////////////////////////////////////////////////////////////////////////////////// /// BULK_COPY : Copy a bulk of memory between shared memory and global memory //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -905,27 +1303,43 @@ tma_store_wait() { 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) + copy(void const* gmem_ptr, uint64_t* mbar_ptr, + void * 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_mbar = cast_smem_ptr_to_uint(mbar_ptr); 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."); + CUTE_INVALID_CONTROL_PATH("Trying to use BULK_COPY without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } + + struct PREFETCH + { + CUTE_HOST_DEVICE static void + copy(void const* gmem_ptr, int32_t load_bytes) + { + #if defined(CUTE_ARCH_TMA_SM90_ENABLED) + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;\n" + : + : "l"(gmem_ptr), "r"(load_bytes) + : "memory"); + #else + CUTE_INVALID_CONTROL_PATH("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) + copy(void const* smem_ptr, + void * 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); @@ -934,7 +1348,7 @@ struct SM90_BULK_COPY_S2G : "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."); + CUTE_INVALID_CONTROL_PATH("Trying to use BULK_COPY without CUTE_ARCH_TMA_SM90_ENABLED."); #endif } }; diff --git a/include/cute/arch/mma_sm61.hpp b/include/cute/arch/mma_sm61.hpp index f3e8d745..f7bcb7d1 100644 --- a/include/cute/arch/mma_sm61.hpp +++ b/include/cute/arch/mma_sm61.hpp @@ -58,7 +58,7 @@ struct SM61_DP4A : "=r"(d) : "r"(a), "r"(b), "r"(c)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM61_DP4A without CUTE_ARCH_MMA_SM61_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM61_DP4A without CUTE_ARCH_MMA_SM61_ENABLED"); #endif } }; @@ -79,7 +79,7 @@ struct SM61_DP2A : "=r"(d) : "r"(a), "r"(b), "r"(c)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM61_DP2A without CUTE_ARCH_MMA_SM61_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM61_DP2A without CUTE_ARCH_MMA_SM61_ENABLED"); #endif } }; diff --git a/include/cute/arch/mma_sm70.hpp b/include/cute/arch/mma_sm70.hpp index 1fcd6ded..63d96cf5 100644 --- a/include/cute/arch/mma_sm70.hpp +++ b/include/cute/arch/mma_sm70.hpp @@ -74,7 +74,7 @@ struct SM70_8x8x4_F16F16F16F16_TN "r"(b0), "r"(b1), "r"(c0), "r"(c1), "r"(c2), "r"(c3)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM70_8x8x4_F16F16F16F16_TN without CUTE_ARCH_MMA_SM70_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM70_8x8x4_F16F16F16F16_TN without CUTE_ARCH_MMA_SM70_ENABLED"); #endif } }; @@ -106,7 +106,7 @@ struct SM70_8x8x4_F16F16F16F16_NT "r"(b0), "r"(b1), "r"(c0), "r"(c1), "r"(c2), "r"(c3)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM70_8x8x4_F16F16F16F16_NT without CUTE_ARCH_MMA_SM70_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM70_8x8x4_F16F16F16F16_NT without CUTE_ARCH_MMA_SM70_ENABLED"); #endif } }; @@ -138,7 +138,7 @@ struct SM70_8x8x4_F16F16F16F16_NN "r"(b0), "r"(b1), "r"(c0), "r"(c1), "r"(c2), "r"(c3)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM70_8x8x4_F16F16F16F16_NN without CUTE_ARCH_MMA_SM70_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM70_8x8x4_F16F16F16F16_NN without CUTE_ARCH_MMA_SM70_ENABLED"); #endif } }; @@ -170,7 +170,7 @@ struct SM70_8x8x4_F16F16F16F16_TT "r"(b0), "r"(b1), "r"(c0), "r"(c1), "r"(c2), "r"(c3)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM70_8x8x4_F16F16F16F16_TT without CUTE_ARCH_MMA_SM70_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM70_8x8x4_F16F16F16F16_TT without CUTE_ARCH_MMA_SM70_ENABLED"); #endif } }; @@ -210,7 +210,7 @@ struct SM70_8x8x4_F32F16F16F32_TN "f"(c0), "f"(c1), "f"(c2), "f"(c3), "f"(c4), "f"(c5), "f"(c6), "f"(c7)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM70_8x8x4_F32F16F16F32_TN without CUTE_ARCH_MMA_SM70_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM70_8x8x4_F32F16F16F32_TN without CUTE_ARCH_MMA_SM70_ENABLED"); #endif } }; @@ -246,7 +246,7 @@ struct SM70_8x8x4_F32F16F16F32_NT "f"(c0), "f"(c1), "f"(c2), "f"(c3), "f"(c4), "f"(c5), "f"(c6), "f"(c7)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM70_8x8x4_F32F16F16F32_NT without CUTE_ARCH_MMA_SM70_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM70_8x8x4_F32F16F16F32_NT without CUTE_ARCH_MMA_SM70_ENABLED"); #endif } }; @@ -282,7 +282,7 @@ struct SM70_8x8x4_F32F16F16F32_NN "f"(c0), "f"(c1), "f"(c2), "f"(c3), "f"(c4), "f"(c5), "f"(c6), "f"(c7)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM70_8x8x4_F32F16F16F32_NN without CUTE_ARCH_MMA_SM70_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM70_8x8x4_F32F16F16F32_NN without CUTE_ARCH_MMA_SM70_ENABLED"); #endif } }; @@ -318,7 +318,7 @@ struct SM70_8x8x4_F32F16F16F32_TT "f"(c0), "f"(c1), "f"(c2), "f"(c3), "f"(c4), "f"(c5), "f"(c6), "f"(c7)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM70_8x8x4_F32F16F16F32_TT without CUTE_ARCH_MMA_SM70_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM70_8x8x4_F32F16F16F32_TT without CUTE_ARCH_MMA_SM70_ENABLED"); #endif } diff --git a/include/cute/arch/mma_sm75.hpp b/include/cute/arch/mma_sm75.hpp index a3bb2742..c33f7b39 100644 --- a/include/cute/arch/mma_sm75.hpp +++ b/include/cute/arch/mma_sm75.hpp @@ -74,7 +74,7 @@ struct SM75_16x8x8_F32F16F16F32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM75_16x8x8_F32F16F16F32_TN without CUTE_ARCH_MMA_SM75_ENABLED"); #endif } }; @@ -110,7 +110,7 @@ struct SM75_8x8x16_S32S8S8S32_TN "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM75_8x8x16_S32S8S8S32_TN without CUTE_ARCH_MMA_SM75_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM75_8x8x16_S32S8S8S32_TN without CUTE_ARCH_MMA_SM75_ENABLED"); #endif } }; diff --git a/include/cute/arch/mma_sm80.hpp b/include/cute/arch/mma_sm80.hpp index 49ae6265..8c684b70 100644 --- a/include/cute/arch/mma_sm80.hpp +++ b/include/cute/arch/mma_sm80.hpp @@ -33,6 +33,7 @@ #include #include +#include // Config #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) @@ -80,7 +81,7 @@ struct SM80_16x8x8_F16F16F16F16_TN "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x8_F16F16F16F16_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x8_F16F16F16F16_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -113,7 +114,7 @@ struct SM80_16x8x16_F16F16F16F16_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_F16F16F16F16_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -146,7 +147,7 @@ struct SM80_16x8x8_F32F16F16F32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x8_F32F16F16F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -179,7 +180,7 @@ struct SM80_16x8x16_F32F16F16F32_TN "r"(b0), "r"(b1), "f"(c0), "f"(c1), "f"(c2), "f"(c3)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_F32F16F16F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_F32F16F16F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -212,7 +213,7 @@ struct SM80_16x8x8_F32BF16BF16F32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x8_F32BF16BF16F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -245,7 +246,7 @@ struct SM80_16x8x16_F32BF16BF16F32_TN "r"(b0), "r"(b1), "f"(c0), "f"(c1), "f"(c2), "f"(c3)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_F32BF16BF16F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_F32BF16BF16F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -278,7 +279,7 @@ struct SM80_16x8x4_F32TF32TF32F32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x4_F32TF32TF32F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -311,7 +312,7 @@ struct SM80_16x8x8_F32TF32TF32F32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x8_F32TF32TF32F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -344,7 +345,7 @@ struct SM80_8x8x4_F64F64F64F64_TN "d"(b0), "d"(c0), "d"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x4_F64F64F64F64_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x4_F64F64F64F64_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -385,14 +386,14 @@ struct SM80_8x8x4_C64C64C64C64_TN // d.real() = -a.imag() * b.imag() + d.real(); SM80_8x8x4_F64F64F64F64_TN::fma( - rd0, rd1, + 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, + id0, id1, a0.real(), b0.imag(), d0.imag(), d1.imag()); @@ -412,15 +413,15 @@ 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 @@ -481,7 +482,7 @@ struct SM80_8x8x16_S32S8S8S32_TN "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32S8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x16_S32S8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -514,7 +515,7 @@ struct SM80_8x8x16_S32S8S8S32_TN_SATURATE "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32S8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x16_S32S8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -547,7 +548,7 @@ struct SM80_16x8x16_S32S8S8S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_S32S8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -580,7 +581,7 @@ struct SM80_16x8x16_S32S8S8S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_S32S8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -613,7 +614,7 @@ struct SM80_16x8x32_S32S8S8S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32S8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -646,7 +647,7 @@ struct SM80_16x8x32_S32S8S8S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32S8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -679,7 +680,7 @@ struct SM80_8x8x16_S32S8U8S32_TN "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32S8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x16_S32S8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -712,7 +713,7 @@ struct SM80_8x8x16_S32S8U8S32_TN_SATURATE "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32S8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x16_S32S8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -745,7 +746,7 @@ struct SM80_16x8x16_S32S8U8S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_S32S8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -778,7 +779,7 @@ struct SM80_16x8x16_S32S8U8S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_S32S8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -811,7 +812,7 @@ struct SM80_16x8x32_S32S8U8S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32S8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -844,7 +845,7 @@ struct SM80_16x8x32_S32S8U8S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32S8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -877,7 +878,7 @@ struct SM80_8x8x16_S32U8S8S32_TN "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32U8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x16_S32U8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -910,7 +911,7 @@ struct SM80_8x8x16_S32U8S8S32_TN_SATURATE "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32U8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x16_S32U8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -943,7 +944,7 @@ struct SM80_16x8x16_S32U8S8S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_S32U8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -976,7 +977,7 @@ struct SM80_16x8x16_S32U8S8S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_S32U8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1009,7 +1010,7 @@ struct SM80_16x8x32_S32U8S8S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32U8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1042,7 +1043,7 @@ struct SM80_16x8x32_S32U8S8S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32U8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1075,7 +1076,7 @@ struct SM80_8x8x16_S32U8U8S32_TN "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32U8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x16_S32U8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1108,7 +1109,7 @@ struct SM80_8x8x16_S32U8U8S32_TN_SATURATE "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32U8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x16_S32U8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1141,7 +1142,7 @@ struct SM80_16x8x16_S32U8U8S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_S32U8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1174,7 +1175,7 @@ struct SM80_16x8x16_S32U8U8S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x16_S32U8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1207,7 +1208,7 @@ struct SM80_16x8x32_S32U8U8S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32U8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1240,7 +1241,7 @@ struct SM80_16x8x32_S32U8U8S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32U8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1273,7 +1274,7 @@ struct SM80_8x8x32_S32S4S4S32_TN "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32S4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x32_S32S4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1306,7 +1307,7 @@ struct SM80_8x8x32_S32S4S4S32_TN_SATURATE "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32S4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x32_S32S4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1339,7 +1340,7 @@ struct SM80_16x8x32_S32S4S4S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32S4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1372,7 +1373,7 @@ struct SM80_16x8x32_S32S4S4S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32S4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1405,7 +1406,7 @@ struct SM80_16x8x64_S32S4S4S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x64_S32S4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1438,7 +1439,7 @@ struct SM80_16x8x64_S32S4S4S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x64_S32S4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1471,7 +1472,7 @@ struct SM80_8x8x32_S32S4U4S32_TN "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32S4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x32_S32S4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1504,7 +1505,7 @@ struct SM80_8x8x32_S32S4U4S32_TN_SATURATE "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32S4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x32_S32S4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1537,7 +1538,7 @@ struct SM80_16x8x32_S32S4U4S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32S4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1570,7 +1571,7 @@ struct SM80_16x8x32_S32S4U4S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32S4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1603,7 +1604,7 @@ struct SM80_16x8x64_S32S4U4S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x64_S32S4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1636,7 +1637,7 @@ struct SM80_16x8x64_S32S4U4S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x64_S32S4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1669,7 +1670,7 @@ struct SM80_8x8x32_S32U4S4S32_TN "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32U4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x32_S32U4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1702,7 +1703,7 @@ struct SM80_8x8x32_S32U4S4S32_TN_SATURATE "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32U4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x32_S32U4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1735,7 +1736,7 @@ struct SM80_16x8x32_S32U4S4S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32U4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1768,7 +1769,7 @@ struct SM80_16x8x32_S32U4S4S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32U4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1801,7 +1802,7 @@ struct SM80_16x8x64_S32U4S4S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x64_S32U4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1834,7 +1835,7 @@ struct SM80_16x8x64_S32U4S4S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x64_S32U4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1867,7 +1868,7 @@ struct SM80_8x8x32_S32U4U4S32_TN "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32U4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x32_S32U4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1900,7 +1901,7 @@ struct SM80_8x8x32_S32U4U4S32_TN_SATURATE "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32U4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x32_S32U4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1933,7 +1934,7 @@ struct SM80_16x8x32_S32U4U4S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32U4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1966,7 +1967,7 @@ struct SM80_16x8x32_S32U4U4S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x32_S32U4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -1999,7 +2000,7 @@ struct SM80_16x8x64_S32U4U4S32_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x64_S32U4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -2032,7 +2033,7 @@ struct SM80_16x8x64_S32U4U4S32_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x64_S32U4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -2067,7 +2068,7 @@ struct SM80_8x8x128_S32U1U1S32_TN_XORPOPC "r"(b0), "r"(c0), "r"(c1)); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x128_S32U1U1S32_TN_XORPOPC without CUTE_ARCH_MMA_SM80_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_8x8x128_S32U1U1S32_TN_XORPOPC without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -2100,7 +2101,7 @@ struct SM80_16x8x128_S32U1U1S32_TN_XORPOPC "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x128_S32U1U1S32_TN_XORPOPC without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; @@ -2133,7 +2134,7 @@ struct SM80_16x8x256_S32U1U1S32_TN_XORPOPC "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM80_16x8x256_S32U1U1S32_TN_XORPOPC without CUTE_ARCH_MMA_SM80_ENABLED"); #endif } }; diff --git a/include/cute/arch/mma_sm90.hpp b/include/cute/arch/mma_sm90.hpp index 22e3c3f7..10bed48a 100644 --- a/include/cute/arch/mma_sm90.hpp +++ b/include/cute/arch/mma_sm90.hpp @@ -73,7 +73,7 @@ struct SM90_16x8x4_F64F64F64F64_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_16x8x4_F64F64F64F64_TN without CUTE_ARCH_MMA_SM90_ENABLED"); #endif } }; @@ -106,7 +106,7 @@ struct SM90_16x8x8_F64F64F64F64_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_16x8x8_F64F64F64F64_TN without CUTE_ARCH_MMA_SM90_ENABLED"); #endif } }; @@ -141,7 +141,7 @@ struct SM90_16x8x16_F64F64F64F64_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_16x8x16_F64F64F64F64_TN without CUTE_ARCH_MMA_SM90_ENABLED"); #endif } }; @@ -364,37 +364,185 @@ ss_op_selector() // 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."); + if constexpr (is_same_v && is_same_v) { + 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{}; + // 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."); + } } - 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 { + + // FP8 + // Input A: float_e4m3_t ; Input B: float_e4m3_t + else if constexpr (is_same_v && 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."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_F16E4M3E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_F16E4M3E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_F16E4M3E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_F16E4M3E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_F16E4M3E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_F16E4M3E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_F16E4M3E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_F16E4M3E4M3_SS_TN{}; + } + else { static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // FP8 + // Input A: float_e4m3_t ; Input B: float_e5m2_t + else if constexpr (is_same_v && 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."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_F16E4M3E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_F16E4M3E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_F16E4M3E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_F16E4M3E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_F16E4M3E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_F16E4M3E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_F16E4M3E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_F16E4M3E5M2_SS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // FP8 + // Input A: float_e5m2_t ; Input B: float_e5m2_t + else if constexpr (is_same_v && 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."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_F16E5M2E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_F16E5M2E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_F16E5M2E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_F16E5M2E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_F16E5M2E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_F16E5M2E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_F16E5M2E5M2_SS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_F16E5M2E5M2_SS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // FP8 + // Input A: float_e5m2_t ; Input B: float_e4m3_t + else if constexpr (is_same_v && 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."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_F16E5M2E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_F16E5M2E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_F16E5M2E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_F16E5M2E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_F16E5M2E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_F16E5M2E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_F16E5M2E4M3_SS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_F16E5M2E4M3_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."); } } diff --git a/include/cute/arch/mma_sm90_gmma.hpp b/include/cute/arch/mma_sm90_gmma.hpp index 213dce73..bdf0d70e 100644 --- a/include/cute/arch/mma_sm90_gmma.hpp +++ b/include/cute/arch/mma_sm90_gmma.hpp @@ -49,7 +49,7 @@ 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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use wgmma.fence without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } @@ -62,7 +62,7 @@ warpgroup_wait() #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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use wgmma.wait_group without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } @@ -74,7 +74,7 @@ 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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use wgmma.commit_group without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } @@ -156,7 +156,7 @@ struct SM90_64x8x16_F16F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -202,7 +202,7 @@ struct SM90_64x8x16_F16F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -245,7 +245,7 @@ struct SM90_64x16x16_F16F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -291,7 +291,7 @@ struct SM90_64x16x16_F16F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -336,7 +336,7 @@ struct SM90_64x32x16_F16F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -384,7 +384,7 @@ struct SM90_64x32x16_F16F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -434,7 +434,7 @@ struct SM90_64x64x16_F16F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -487,7 +487,7 @@ struct SM90_64x64x16_F16F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -542,7 +542,7 @@ struct SM90_64x96x16_F16F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -600,7 +600,7 @@ struct SM90_64x96x16_F16F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -660,7 +660,7 @@ struct SM90_64x128x16_F16F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -723,7 +723,7 @@ struct SM90_64x128x16_F16F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -793,7 +793,7 @@ struct SM90_64x192x16_F16F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -866,7 +866,7 @@ struct SM90_64x192x16_F16F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -946,7 +946,7 @@ struct SM90_64x256x16_F16F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1029,7 +1029,7 @@ struct SM90_64x256x16_F16F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1072,7 +1072,7 @@ struct SM90_64x8x16_F32F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1118,7 +1118,7 @@ struct SM90_64x8x16_F32F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1163,7 +1163,7 @@ struct SM90_64x16x16_F32F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1211,7 +1211,7 @@ struct SM90_64x16x16_F32F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1261,7 +1261,7 @@ struct SM90_64x32x16_F32F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1314,7 +1314,7 @@ struct SM90_64x32x16_F32F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1374,7 +1374,7 @@ struct SM90_64x64x16_F32F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1437,7 +1437,7 @@ struct SM90_64x64x16_F32F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1507,7 +1507,7 @@ struct SM90_64x96x16_F32F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1580,7 +1580,7 @@ struct SM90_64x96x16_F32F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1660,7 +1660,7 @@ struct SM90_64x128x16_F32F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1743,7 +1743,7 @@ struct SM90_64x128x16_F32F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1843,7 +1843,7 @@ struct SM90_64x192x16_F32F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -1946,7 +1946,7 @@ struct SM90_64x192x16_F32F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2066,7 +2066,7 @@ struct SM90_64x256x16_F32F16F16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2189,7 +2189,7 @@ struct SM90_64x256x16_F32F16F16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2232,7 +2232,7 @@ struct SM90_64x8x16_F32BF16BF16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2278,7 +2278,7 @@ struct SM90_64x8x16_F32BF16BF16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2323,7 +2323,7 @@ struct SM90_64x16x16_F32BF16BF16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2371,7 +2371,7 @@ struct SM90_64x16x16_F32BF16BF16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2421,7 +2421,7 @@ struct SM90_64x32x16_F32BF16BF16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2474,7 +2474,7 @@ struct SM90_64x32x16_F32BF16BF16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2534,7 +2534,7 @@ struct SM90_64x64x16_F32BF16BF16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2597,7 +2597,7 @@ struct SM90_64x64x16_F32BF16BF16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2667,7 +2667,7 @@ struct SM90_64x96x16_F32BF16BF16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2740,7 +2740,7 @@ struct SM90_64x96x16_F32BF16BF16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2820,7 +2820,7 @@ struct SM90_64x128x16_F32BF16BF16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -2903,7 +2903,7 @@ struct SM90_64x128x16_F32BF16BF16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3003,7 +3003,7 @@ struct SM90_64x192x16_F32BF16BF16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3106,7 +3106,7 @@ struct SM90_64x192x16_F32BF16BF16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3226,7 +3226,7 @@ struct SM90_64x256x16_F32BF16BF16_SS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3349,7 +3349,7 @@ struct SM90_64x256x16_F32BF16BF16_RS "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3390,7 +3390,7 @@ struct SM90_64x8x8_F32TF32TF32_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3431,7 +3431,7 @@ struct SM90_64x8x8_F32TF32TF32_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3474,7 +3474,7 @@ struct SM90_64x16x8_F32TF32TF32_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3517,7 +3517,7 @@ struct SM90_64x16x8_F32TF32TF32_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3565,7 +3565,7 @@ struct SM90_64x32x8_F32TF32TF32_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3613,7 +3613,7 @@ struct SM90_64x32x8_F32TF32TF32_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3671,7 +3671,7 @@ struct SM90_64x64x8_F32TF32TF32_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3729,7 +3729,7 @@ struct SM90_64x64x8_F32TF32TF32_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3797,7 +3797,7 @@ struct SM90_64x96x8_F32TF32TF32_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3865,7 +3865,7 @@ struct SM90_64x96x8_F32TF32TF32_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -3943,7 +3943,7 @@ struct SM90_64x128x8_F32TF32TF32_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4021,7 +4021,7 @@ struct SM90_64x128x8_F32TF32TF32_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4119,7 +4119,7 @@ struct SM90_64x192x8_F32TF32TF32_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4217,7 +4217,7 @@ struct SM90_64x192x8_F32TF32TF32_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4335,7 +4335,7 @@ struct SM90_64x256x8_F32TF32TF32_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4453,7 +4453,7 @@ struct SM90_64x256x8_F32TF32TF32_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4490,7 +4490,7 @@ struct SM90_64x8x32_S32S8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4527,7 +4527,7 @@ struct SM90_64x8x32_S32S8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4566,7 +4566,7 @@ struct SM90_64x16x32_S32S8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4605,7 +4605,7 @@ struct SM90_64x16x32_S32S8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4649,7 +4649,7 @@ struct SM90_64x32x32_S32S8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4693,7 +4693,7 @@ struct SM90_64x32x32_S32S8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4747,7 +4747,7 @@ struct SM90_64x64x32_S32S8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4801,7 +4801,7 @@ struct SM90_64x64x32_S32S8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4865,7 +4865,7 @@ struct SM90_64x96x32_S32S8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -4929,7 +4929,7 @@ struct SM90_64x96x32_S32S8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5003,7 +5003,7 @@ struct SM90_64x128x32_S32S8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5077,7 +5077,7 @@ struct SM90_64x128x32_S32S8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5171,7 +5171,7 @@ struct SM90_64x192x32_S32S8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5265,7 +5265,7 @@ struct SM90_64x192x32_S32S8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5379,7 +5379,7 @@ struct SM90_64x256x32_S32S8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5493,7 +5493,7 @@ struct SM90_64x256x32_S32S8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5530,7 +5530,7 @@ struct SM90_64x8x32_S32S8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5567,7 +5567,7 @@ struct SM90_64x8x32_S32S8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5606,7 +5606,7 @@ struct SM90_64x16x32_S32S8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5645,7 +5645,7 @@ struct SM90_64x16x32_S32S8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5689,7 +5689,7 @@ struct SM90_64x32x32_S32S8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5733,7 +5733,7 @@ struct SM90_64x32x32_S32S8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5787,7 +5787,7 @@ struct SM90_64x64x32_S32S8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5841,7 +5841,7 @@ struct SM90_64x64x32_S32S8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5905,7 +5905,7 @@ struct SM90_64x96x32_S32S8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -5969,7 +5969,7 @@ struct SM90_64x96x32_S32S8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6043,7 +6043,7 @@ struct SM90_64x128x32_S32S8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6117,7 +6117,7 @@ struct SM90_64x128x32_S32S8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6211,7 +6211,7 @@ struct SM90_64x192x32_S32S8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6305,7 +6305,7 @@ struct SM90_64x192x32_S32S8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6419,7 +6419,7 @@ struct SM90_64x256x32_S32S8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6533,7 +6533,7 @@ struct SM90_64x256x32_S32S8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6570,7 +6570,7 @@ struct SM90_64x8x32_S32S8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6607,7 +6607,7 @@ struct SM90_64x8x32_S32S8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6646,7 +6646,7 @@ struct SM90_64x16x32_S32S8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6685,7 +6685,7 @@ struct SM90_64x16x32_S32S8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6729,7 +6729,7 @@ struct SM90_64x32x32_S32S8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6773,7 +6773,7 @@ struct SM90_64x32x32_S32S8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6827,7 +6827,7 @@ struct SM90_64x64x32_S32S8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6881,7 +6881,7 @@ struct SM90_64x64x32_S32S8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -6945,7 +6945,7 @@ struct SM90_64x96x32_S32S8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7009,7 +7009,7 @@ struct SM90_64x96x32_S32S8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7083,7 +7083,7 @@ struct SM90_64x128x32_S32S8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7157,7 +7157,7 @@ struct SM90_64x128x32_S32S8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7251,7 +7251,7 @@ struct SM90_64x192x32_S32S8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7345,7 +7345,7 @@ struct SM90_64x192x32_S32S8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7459,7 +7459,7 @@ struct SM90_64x256x32_S32S8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7573,7 +7573,7 @@ struct SM90_64x256x32_S32S8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7610,7 +7610,7 @@ struct SM90_64x8x32_S32S8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7647,7 +7647,7 @@ struct SM90_64x8x32_S32S8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7686,7 +7686,7 @@ struct SM90_64x16x32_S32S8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7725,7 +7725,7 @@ struct SM90_64x16x32_S32S8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7769,7 +7769,7 @@ struct SM90_64x32x32_S32S8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7813,7 +7813,7 @@ struct SM90_64x32x32_S32S8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7867,7 +7867,7 @@ struct SM90_64x64x32_S32S8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7921,7 +7921,7 @@ struct SM90_64x64x32_S32S8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -7985,7 +7985,7 @@ struct SM90_64x96x32_S32S8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8049,7 +8049,7 @@ struct SM90_64x96x32_S32S8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8123,7 +8123,7 @@ struct SM90_64x128x32_S32S8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8197,7 +8197,7 @@ struct SM90_64x128x32_S32S8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8291,7 +8291,7 @@ struct SM90_64x192x32_S32S8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8385,7 +8385,7 @@ struct SM90_64x192x32_S32S8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8499,7 +8499,7 @@ struct SM90_64x256x32_S32S8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8613,7 +8613,7 @@ struct SM90_64x256x32_S32S8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8650,7 +8650,7 @@ struct SM90_64x8x32_S32U8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8687,7 +8687,7 @@ struct SM90_64x8x32_S32U8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8726,7 +8726,7 @@ struct SM90_64x16x32_S32U8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8765,7 +8765,7 @@ struct SM90_64x16x32_S32U8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8809,7 +8809,7 @@ struct SM90_64x32x32_S32U8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8853,7 +8853,7 @@ struct SM90_64x32x32_S32U8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8907,7 +8907,7 @@ struct SM90_64x64x32_S32U8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -8961,7 +8961,7 @@ struct SM90_64x64x32_S32U8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9025,7 +9025,7 @@ struct SM90_64x96x32_S32U8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9089,7 +9089,7 @@ struct SM90_64x96x32_S32U8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9163,7 +9163,7 @@ struct SM90_64x128x32_S32U8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9237,7 +9237,7 @@ struct SM90_64x128x32_S32U8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9331,7 +9331,7 @@ struct SM90_64x192x32_S32U8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9425,7 +9425,7 @@ struct SM90_64x192x32_S32U8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9539,7 +9539,7 @@ struct SM90_64x256x32_S32U8S8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9653,7 +9653,7 @@ struct SM90_64x256x32_S32U8S8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9690,7 +9690,7 @@ struct SM90_64x8x32_S32U8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9727,7 +9727,7 @@ struct SM90_64x8x32_S32U8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9766,7 +9766,7 @@ struct SM90_64x16x32_S32U8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9805,7 +9805,7 @@ struct SM90_64x16x32_S32U8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9849,7 +9849,7 @@ struct SM90_64x32x32_S32U8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9893,7 +9893,7 @@ struct SM90_64x32x32_S32U8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -9947,7 +9947,7 @@ struct SM90_64x64x32_S32U8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10001,7 +10001,7 @@ struct SM90_64x64x32_S32U8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10065,7 +10065,7 @@ struct SM90_64x96x32_S32U8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10129,7 +10129,7 @@ struct SM90_64x96x32_S32U8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10203,7 +10203,7 @@ struct SM90_64x128x32_S32U8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10277,7 +10277,7 @@ struct SM90_64x128x32_S32U8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10371,7 +10371,7 @@ struct SM90_64x192x32_S32U8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10465,7 +10465,7 @@ struct SM90_64x192x32_S32U8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10579,7 +10579,7 @@ struct SM90_64x256x32_S32U8S8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10693,7 +10693,7 @@ struct SM90_64x256x32_S32U8S8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10730,7 +10730,7 @@ struct SM90_64x8x32_S32U8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10767,7 +10767,7 @@ struct SM90_64x8x32_S32U8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10806,7 +10806,7 @@ struct SM90_64x16x32_S32U8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10845,7 +10845,7 @@ struct SM90_64x16x32_S32U8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10889,7 +10889,7 @@ struct SM90_64x32x32_S32U8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10933,7 +10933,7 @@ struct SM90_64x32x32_S32U8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -10987,7 +10987,7 @@ struct SM90_64x64x32_S32U8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11041,7 +11041,7 @@ struct SM90_64x64x32_S32U8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11105,7 +11105,7 @@ struct SM90_64x96x32_S32U8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11169,7 +11169,7 @@ struct SM90_64x96x32_S32U8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11243,7 +11243,7 @@ struct SM90_64x128x32_S32U8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11317,7 +11317,7 @@ struct SM90_64x128x32_S32U8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11411,7 +11411,7 @@ struct SM90_64x192x32_S32U8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11505,7 +11505,7 @@ struct SM90_64x192x32_S32U8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11619,7 +11619,7 @@ struct SM90_64x256x32_S32U8U8_SS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11733,7 +11733,7 @@ struct SM90_64x256x32_S32U8U8_SS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11770,7 +11770,7 @@ struct SM90_64x8x32_S32U8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11807,7 +11807,7 @@ struct SM90_64x8x32_S32U8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11846,7 +11846,7 @@ struct SM90_64x16x32_S32U8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11885,7 +11885,7 @@ struct SM90_64x16x32_S32U8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11929,7 +11929,7 @@ struct SM90_64x32x32_S32U8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -11973,7 +11973,7 @@ struct SM90_64x32x32_S32U8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12027,7 +12027,7 @@ struct SM90_64x64x32_S32U8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12081,7 +12081,7 @@ struct SM90_64x64x32_S32U8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12145,7 +12145,7 @@ struct SM90_64x96x32_S32U8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12209,7 +12209,7 @@ struct SM90_64x96x32_S32U8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12283,7 +12283,7 @@ struct SM90_64x128x32_S32U8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12357,7 +12357,7 @@ struct SM90_64x128x32_S32U8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12451,7 +12451,7 @@ struct SM90_64x192x32_S32U8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12545,7 +12545,7 @@ struct SM90_64x192x32_S32U8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12659,7 +12659,7 @@ struct SM90_64x256x32_S32U8U8_RS_TN "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12773,7 +12773,7 @@ struct SM90_64x256x32_S32U8U8_RS_TN_SATURATE "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"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12814,7 +12814,7 @@ struct SM90_64x8x32_F16E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12855,7 +12855,7 @@ struct SM90_64x8x32_F16E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12896,7 +12896,7 @@ struct SM90_64x8x32_F32E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12937,7 +12937,7 @@ struct SM90_64x8x32_F32E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -12978,7 +12978,7 @@ struct SM90_64x16x32_F16E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13019,7 +13019,7 @@ struct SM90_64x16x32_F16E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13062,7 +13062,7 @@ struct SM90_64x16x32_F32E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13105,7 +13105,7 @@ struct SM90_64x16x32_F32E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13148,7 +13148,7 @@ struct SM90_64x32x32_F16E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13191,7 +13191,7 @@ struct SM90_64x32x32_F16E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13239,7 +13239,7 @@ struct SM90_64x32x32_F32E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13287,7 +13287,7 @@ struct SM90_64x32x32_F32E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13335,7 +13335,7 @@ struct SM90_64x64x32_F16E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13383,7 +13383,7 @@ struct SM90_64x64x32_F16E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13441,7 +13441,7 @@ struct SM90_64x64x32_F32E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13499,7 +13499,7 @@ struct SM90_64x64x32_F32E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13552,7 +13552,7 @@ struct SM90_64x96x32_F16E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13605,7 +13605,7 @@ struct SM90_64x96x32_F16E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13673,7 +13673,7 @@ struct SM90_64x96x32_F32E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13741,7 +13741,7 @@ struct SM90_64x96x32_F32E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13799,7 +13799,7 @@ struct SM90_64x128x32_F16E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13857,7 +13857,7 @@ struct SM90_64x128x32_F16E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -13935,7 +13935,7 @@ struct SM90_64x128x32_F32E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14013,7 +14013,7 @@ struct SM90_64x128x32_F32E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14081,7 +14081,7 @@ struct SM90_64x192x32_F16E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14149,7 +14149,7 @@ struct SM90_64x192x32_F16E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14247,7 +14247,7 @@ struct SM90_64x192x32_F32E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14345,7 +14345,7 @@ struct SM90_64x192x32_F32E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14423,7 +14423,7 @@ struct SM90_64x256x32_F16E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F16E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14501,7 +14501,7 @@ struct SM90_64x256x32_F16E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F16E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14619,7 +14619,7 @@ struct SM90_64x256x32_F32E4M3E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F32E4M3E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14737,7 +14737,7 @@ struct SM90_64x256x32_F32E4M3E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F32E4M3E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14778,7 +14778,7 @@ struct SM90_64x8x32_F16E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14819,7 +14819,7 @@ struct SM90_64x8x32_F16E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14860,7 +14860,7 @@ struct SM90_64x8x32_F32E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14901,7 +14901,7 @@ struct SM90_64x8x32_F32E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14942,7 +14942,7 @@ struct SM90_64x16x32_F16E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -14983,7 +14983,7 @@ struct SM90_64x16x32_F16E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15026,7 +15026,7 @@ struct SM90_64x16x32_F32E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15069,7 +15069,7 @@ struct SM90_64x16x32_F32E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15112,7 +15112,7 @@ struct SM90_64x32x32_F16E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15155,7 +15155,7 @@ struct SM90_64x32x32_F16E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15203,7 +15203,7 @@ struct SM90_64x32x32_F32E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15251,7 +15251,7 @@ struct SM90_64x32x32_F32E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15299,7 +15299,7 @@ struct SM90_64x64x32_F16E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15347,7 +15347,7 @@ struct SM90_64x64x32_F16E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15405,7 +15405,7 @@ struct SM90_64x64x32_F32E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15463,7 +15463,7 @@ struct SM90_64x64x32_F32E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15516,7 +15516,7 @@ struct SM90_64x96x32_F16E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15569,7 +15569,7 @@ struct SM90_64x96x32_F16E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15637,7 +15637,7 @@ struct SM90_64x96x32_F32E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15705,7 +15705,7 @@ struct SM90_64x96x32_F32E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15763,7 +15763,7 @@ struct SM90_64x128x32_F16E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15821,7 +15821,7 @@ struct SM90_64x128x32_F16E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15899,7 +15899,7 @@ struct SM90_64x128x32_F32E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -15977,7 +15977,7 @@ struct SM90_64x128x32_F32E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16045,7 +16045,7 @@ struct SM90_64x192x32_F16E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16113,7 +16113,7 @@ struct SM90_64x192x32_F16E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16211,7 +16211,7 @@ struct SM90_64x192x32_F32E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16309,7 +16309,7 @@ struct SM90_64x192x32_F32E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16387,7 +16387,7 @@ struct SM90_64x256x32_F16E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F16E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16465,7 +16465,7 @@ struct SM90_64x256x32_F16E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F16E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16583,7 +16583,7 @@ struct SM90_64x256x32_F32E4M3E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F32E4M3E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16701,7 +16701,7 @@ struct SM90_64x256x32_F32E4M3E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F32E4M3E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16742,7 +16742,7 @@ struct SM90_64x8x32_F16E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16783,7 +16783,7 @@ struct SM90_64x8x32_F16E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16824,7 +16824,7 @@ struct SM90_64x8x32_F32E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16865,7 +16865,7 @@ struct SM90_64x8x32_F32E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16906,7 +16906,7 @@ struct SM90_64x16x32_F16E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16947,7 +16947,7 @@ struct SM90_64x16x32_F16E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -16990,7 +16990,7 @@ struct SM90_64x16x32_F32E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17033,7 +17033,7 @@ struct SM90_64x16x32_F32E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17076,7 +17076,7 @@ struct SM90_64x32x32_F16E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17119,7 +17119,7 @@ struct SM90_64x32x32_F16E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17167,7 +17167,7 @@ struct SM90_64x32x32_F32E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17215,7 +17215,7 @@ struct SM90_64x32x32_F32E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17263,7 +17263,7 @@ struct SM90_64x64x32_F16E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17311,7 +17311,7 @@ struct SM90_64x64x32_F16E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17369,7 +17369,7 @@ struct SM90_64x64x32_F32E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17427,7 +17427,7 @@ struct SM90_64x64x32_F32E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17480,7 +17480,7 @@ struct SM90_64x96x32_F16E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17533,7 +17533,7 @@ struct SM90_64x96x32_F16E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17601,7 +17601,7 @@ struct SM90_64x96x32_F32E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17669,7 +17669,7 @@ struct SM90_64x96x32_F32E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17727,7 +17727,7 @@ struct SM90_64x128x32_F16E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17785,7 +17785,7 @@ struct SM90_64x128x32_F16E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17863,7 +17863,7 @@ struct SM90_64x128x32_F32E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -17941,7 +17941,7 @@ struct SM90_64x128x32_F32E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18009,7 +18009,7 @@ struct SM90_64x192x32_F16E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18077,7 +18077,7 @@ struct SM90_64x192x32_F16E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18175,7 +18175,7 @@ struct SM90_64x192x32_F32E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18273,7 +18273,7 @@ struct SM90_64x192x32_F32E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18351,7 +18351,7 @@ struct SM90_64x256x32_F16E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F16E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18429,7 +18429,7 @@ struct SM90_64x256x32_F16E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F16E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18547,7 +18547,7 @@ struct SM90_64x256x32_F32E5M2E4M3_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F32E5M2E4M3_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18665,7 +18665,7 @@ struct SM90_64x256x32_F32E5M2E4M3_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F32E5M2E4M3_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18706,7 +18706,7 @@ struct SM90_64x8x32_F16E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18747,7 +18747,7 @@ struct SM90_64x8x32_F16E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18788,7 +18788,7 @@ struct SM90_64x8x32_F32E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18829,7 +18829,7 @@ struct SM90_64x8x32_F32E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x8x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18870,7 +18870,7 @@ struct SM90_64x16x32_F16E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18911,7 +18911,7 @@ struct SM90_64x16x32_F16E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18954,7 +18954,7 @@ struct SM90_64x16x32_F32E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -18997,7 +18997,7 @@ struct SM90_64x16x32_F32E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x16x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19040,7 +19040,7 @@ struct SM90_64x32x32_F16E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19083,7 +19083,7 @@ struct SM90_64x32x32_F16E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19131,7 +19131,7 @@ struct SM90_64x32x32_F32E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19179,7 +19179,7 @@ struct SM90_64x32x32_F32E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x32x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19227,7 +19227,7 @@ struct SM90_64x64x32_F16E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19275,7 +19275,7 @@ struct SM90_64x64x32_F16E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19333,7 +19333,7 @@ struct SM90_64x64x32_F32E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19391,7 +19391,7 @@ struct SM90_64x64x32_F32E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x64x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19444,7 +19444,7 @@ struct SM90_64x96x32_F16E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19497,7 +19497,7 @@ struct SM90_64x96x32_F16E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19565,7 +19565,7 @@ struct SM90_64x96x32_F32E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19633,7 +19633,7 @@ struct SM90_64x96x32_F32E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x96x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19691,7 +19691,7 @@ struct SM90_64x128x32_F16E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19749,7 +19749,7 @@ struct SM90_64x128x32_F16E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19827,7 +19827,7 @@ struct SM90_64x128x32_F32E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19905,7 +19905,7 @@ struct SM90_64x128x32_F32E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x128x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -19973,7 +19973,7 @@ struct SM90_64x192x32_F16E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -20041,7 +20041,7 @@ struct SM90_64x192x32_F16E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -20139,7 +20139,7 @@ struct SM90_64x192x32_F32E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -20237,7 +20237,7 @@ struct SM90_64x192x32_F32E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x192x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -20315,7 +20315,7 @@ struct SM90_64x256x32_F16E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F16E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -20393,7 +20393,7 @@ struct SM90_64x256x32_F16E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F16E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -20511,7 +20511,7 @@ struct SM90_64x256x32_F32E5M2E5M2_SS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F32E5M2E5M2_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; @@ -20629,7 +20629,7 @@ struct SM90_64x256x32_F32E5M2E5M2_RS_TN "l"(desc_b), "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); #else - CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); + CUTE_INVALID_CONTROL_PATH("Attempting to use SM90_64x256x32_F32E5M2E5M2_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); #endif } }; diff --git a/include/cute/atom/copy_atom.hpp b/include/cute/atom/copy_atom.hpp index 84efafcd..d1cd3d4b 100644 --- a/include/cute/atom/copy_atom.hpp +++ b/include/cute/atom/copy_atom.hpp @@ -81,7 +81,7 @@ struct Copy_Atom, CopyInternalType> CUTE_HOST_DEVICE auto with(TraitsArgs&&... args) const { - auto traits = Traits::with(std::forward(args)...); + auto traits = Traits::with(static_cast(args)...); return Copy_Atom{traits}; } @@ -351,7 +351,7 @@ struct ThrCopy partition_S(STensor&& stensor) const { //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())); + auto thr_tensor = make_tensor(static_cast(stensor).data(), TiledCopy::tidfrg_S(stensor.layout())); return thr_tensor(thr_idx_, _, repeat>(_)); } @@ -361,7 +361,7 @@ struct ThrCopy partition_D(DTensor&& dtensor) const { //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())); + auto thr_tensor = make_tensor(static_cast(dtensor).data(), TiledCopy::tidfrg_D(dtensor.layout())); return thr_tensor(thr_idx_, _, repeat>(_)); } @@ -371,7 +371,7 @@ struct ThrCopy 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())); + return make_tensor(static_cast(stensor).data(), TiledCopy::retile(stensor.layout())); } template @@ -380,7 +380,7 @@ struct ThrCopy 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())); + return make_tensor(static_cast(dtensor).data(), TiledCopy::retile(dtensor.layout())); } }; diff --git a/include/cute/atom/copy_traits.hpp b/include/cute/atom/copy_traits.hpp index 5a92d9b6..b6259b58 100644 --- a/include/cute/atom/copy_traits.hpp +++ b/include/cute/atom/copy_traits.hpp @@ -94,17 +94,55 @@ struct Copy_Traits> namespace detail { +// Utility for exploding pointers, arrays, or tensors into Operation::copy template + class PtrSrc, int... Is, + class PtrDst, int... Id> CUTE_HOST_DEVICE constexpr void -copy_explode(PtrS&& s, int_sequence, - PtrD&& d, int_sequence) +copy_explode_index(PtrSrc&& s, int_sequence, + PtrDst&& d, int_sequence) { return Operation::copy(s[Is]..., d[Id]...); } +// Utility for exploding tuples into ::copy +template +CUTE_HOST_DEVICE constexpr +void +copy_explode(TupleArg&& t, int_sequence) +{ + return Operation::copy(get(static_cast(t))...); +} + +template +CUTE_HOST_DEVICE constexpr +void +copy_explode(TupleSrc&& s, int_sequence, + TupleDst&& d, int_sequence) +{ + return Operation::copy(get(static_cast(s))..., + get(static_cast(d))...); +} + +template +CUTE_HOST_DEVICE constexpr +void +copy_explode(TupleAux&& a, int_sequence, + TupleSrc&& s, int_sequence, + TupleDst&& d, int_sequence) +{ + return Operation::copy(get(static_cast(a))..., + get(static_cast(s))..., + get(static_cast(d))...); +} + } // end namespace detail // @@ -139,8 +177,8 @@ copy_unpack(Copy_Traits const&, CUTE_STATIC_ASSERT_V(size(rD) == Int{}, "Copy_Traits: dst failed to vectorize into registers. Layout is incompatible with this CopyOp."); - detail::copy_explode(rS, make_int_sequence{}, - rD, make_int_sequence{}); + detail::copy_explode_index(rS, make_int_sequence{}, + rD, make_int_sequence{}); } // diff --git a/include/cute/atom/copy_traits_sm90_im2col.hpp b/include/cute/atom/copy_traits_sm90_im2col.hpp new file mode 100644 index 00000000..34e71ed6 --- /dev/null +++ b/include/cute/atom/copy_traits_sm90_im2col.hpp @@ -0,0 +1,879 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 im2col make_tma_copy +*/ + +#include "cute/arch/copy_sm90.hpp" +#include "cute/arch/copy_sm90_desc.hpp" +#include "cute/tensor.hpp" + +#include "cute/algorithm/prefetch.hpp" + +namespace cute +{ + +// Utility for unpacking TMA_LOAD_IM2COL arguments into a CopyOp +template +struct TMA_LOAD_IM2COL_Unpack +{ + /// Copy from src to dst. + /// + /// @param traits Copy traits created with a TMA descriptor that + /// correctly matches the input tensor and other convolution + /// parameters. + /// + /// @param src Tile of the im2col-transformed coordinate tensor + /// (result of get_tma_tensor), representing the global-memory + /// tensor from which to load. + /// + /// @param dst Shared memory tile, into which to load. + template + CUTE_HOST_DEVICE friend constexpr void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, // tile of the transformed global activation (A) tensor + Tensor & dst) // shared memory tile + { + auto src_coord_offset = src(Int<0>{}); + auto src_coord_cwhdn_offset_srt = flatten(src_coord_offset); + // Interpret the TMA IM2COL coordinate as (c, ([w,h,d]), n, ([s,r,t])) + CUTE_STATIC_ASSERT_V(rank(src_coord_offset) == _4{}); + CUTE_STATIC_ASSERT_V(rank<1>(src_coord_offset) == rank<3>(src_coord_offset)); + + if constexpr (detail::is_prefetch) { + return detail::copy_explode(traits.opargs_, tuple_seq{}, + src_coord_cwhdn_offset_srt, tuple_seq{}); + } else { + static_assert(is_smem::value, "SM90_TMA_LOAD_IM2COL requires the destination be shared memory."); + void* dst_ptr = cute::raw_pointer_cast(dst.data()); + return detail::copy_explode(traits.opargs_, tuple_seq{}, + make_tuple(dst_ptr), seq<0>{}, + src_coord_cwhdn_offset_srt, tuple_seq{}); + } + } +}; + +// Copy_Traits for SM90 im2col TMA load comes in two layers. +// +// 1. Copy_Traits +// 2. Copy_Traits +// +// Copy_Traits +// is the "outer" layer. It has a TMA descriptor, +// but no barrier ("tma_mbar"), so it's "nonexecutable." +// One calls its "with" member function with a barrier, +// to get an executable "inner"-layer +// Copy_Traits object. +// That object's "copy_unpack" member function +// actually invokes im2col TMA load. + +struct SM90_TMA_LOAD_IM2COL_OP : SM90_TMA_LOAD_IM2COL {}; + +/// @brief Non-executable specialization of Copy_Traits for SM90 +/// im2col TMA load, with TMA descriptor but no barrier. +/// +/// Use `.with(memory_barrier)` 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; + + Im2ColTmaDescriptor tma_desc_; + TMATensor tma_tensor_; + + CUTE_HOST_DEVICE constexpr + Im2ColTmaDescriptor const* + get_tma_descriptor() const + { + return &tma_desc_; + } + + template + CUTE_HOST_DEVICE constexpr + TMATensor const + get_tma_tensor(GShape const&) const + { + return tma_tensor_; + } + + /// @brief Get an executable specialization. + /// + /// Copy_Traits specializations with SM90_TMA_LOAD_IM2COL are not + /// directly executable. Instead, call this "with" member function + /// to get an executable specialization. "Executable" means that + /// @c copy_unpack works. + /// + /// @param tma_mbar Memory barrier for synchronization + /// + /// @param multicast_mask Multicast mask (unused; only exists + /// for interface compatibility with the actual multicast Copy_Traits) + /// + /// @return Executable specialization of @c Copy_Traits + CUTE_HOST_DEVICE constexpr + Copy_Traits + with(uint64_t& tma_mbar, [[maybe_unused]] uint16_t const& multicast_mask = 0) const + { + return {{}, {&tma_desc_, &tma_mbar}}; + } + + // Copy_Traits specializations with SM90_TMA_LOAD_IM2COL + // are not directly executable. Instead, call .with + // to get an executable specialization. + template + CUTE_HOST_DEVICE friend constexpr void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, + Tensor & dst) = delete; +}; + +/// @brief Executable specialization of Copy_Traits for SM90 im2col +/// TMA load, with TMA descriptor and barrier. +template +struct Copy_Traits + : TMA_LOAD_IM2COL_Unpack +{ + 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_IM2COL arguments + tuple< + Im2ColTmaDescriptor const*, + uint64_t* // smem mbarrier + > const opargs_; +}; + +template +struct Copy_Traits + : TMA_LOAD_IM2COL_Unpack +{ + 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_IM2COL::PREFETCH arguments + tuple const opargs_; + + CUTE_HOST_DEVICE + Copy_Traits(Copy_Traits const& traits) + : opargs_({&traits.tma_desc_}) {} +}; + +////////////////////////////////////////////////////////////////////////////// +///////////////////////////// TMA_LOAD_MULTICAST ///////////////////////////// +////////////////////////////////////////////////////////////////////////////// + +struct SM90_TMA_LOAD_IM2COL_MULTICAST_OP : SM90_TMA_LOAD_IM2COL_MULTICAST {}; + +/// @brief Non-executable specialization of Copy_Traits for SM90 +/// im2col TMA load, with TMA descriptor but no barrier or multicast +/// mask. +/// +/// Use `.with(memory_barrier)` 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; + + Im2ColTmaDescriptor tma_desc_; + TMATensor tma_tensor_; + + CUTE_HOST_DEVICE constexpr + Im2ColTmaDescriptor const* + get_tma_descriptor() const { + return &tma_desc_; + } + + template + CUTE_HOST_DEVICE constexpr + TMATensor const + get_tma_tensor(GShape const&) const + { + return tma_tensor_; + } + + /// @brief Get an executable specialization. + /// + /// Copy_Traits specializations with SM90_TMA_LOAD_IM2COL_MULTICAST + /// are not directly executable. Instead, call this "with" member + /// function to get an executable specialization. "Executable" + /// means that @c copy_unpack works. + /// + /// @param tma_mbar Memory barrier for synchronization + /// + /// @param multicast_mask Multicast mask (defaults to a single CTA) + /// + /// @return Executable specialization of @c Copy_Traits + CUTE_HOST_DEVICE constexpr + Copy_Traits + with(uint64_t& tma_mbar, uint16_t const& multicast_mask) const { + return {{}, {&tma_desc_, &tma_mbar, multicast_mask}}; + } + + // Copy_Traits specializations with SM90_TMA_LOAD_IM2COL_MULTICAST + // are not directly executable. Instead, call .with to get an + // executable specialization. + template + CUTE_HOST_DEVICE friend constexpr void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, + Tensor & dst) = delete; +}; + +/// @brief Executable specialization of Copy_Traits for SM90 multicast +/// im2col TMA load, with TMA descriptor, barrier, and multicast mask. +template +struct Copy_Traits + : TMA_LOAD_IM2COL_Unpack +{ + 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_IM2COL_MULTICAST arguments + tuple< + Im2ColTmaDescriptor const*, + uint64_t*, // smem mbarrier + uint16_t // multicast mask + > const opargs_; +}; + +////////////////////////////////////////////////////////////////////////////// +///////////////////////////// TMA_STORE IM2COL//////////////////////////////// +////////////////////////////////////////////////////////////////////////////// + +// The executable SM90_TMA_STORE_IM2COL 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_IM2COL arguments + Im2ColTmaDescriptor tma_desc_; + TMATensor tma_tensor_; + + // Return TmaDescriptor/TensorMap + CUTE_HOST_DEVICE constexpr + Im2ColTmaDescriptor const* + get_tma_descriptor() const { + return &tma_desc_; + } + + template + CUTE_HOST_DEVICE constexpr + TMATensor const + get_tma_tensor(GShape const&) const + { + return tma_tensor_; + } + + // 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_IM2COL"); + + void const* const desc_ptr = &(traits.tma_desc_); + void const* const src_ptr = cute::raw_pointer_cast(src.data()); + auto dst_coord = flatten(take<0,3>(dst(Int<0>{}))); + + return detail::copy_explode(make_tuple(desc_ptr, src_ptr), seq<0,1>{}, + dst_coord, tuple_seq{}); + } +}; + +namespace detail { + +/// @brief Creates a TMA descriptor for im2col TMA load. +/// +/// @param tensor_cwhdn Global activation tensor (A matrix of Fprop). +/// This is the original (not im2col-transformed) tensor in global +/// memory. +/// +/// @param slayout Rank 2 (M,K) shared memory layout of the activation +/// tensor. Here, K is "GEMM K," not the filter tensor's mode of +/// the same name. +////// +/// @param traversal_stride Traversal strides convolution parameter +////// +/// Each of padding_shape, traversal_stride, and dilation_shape is a +/// tuple whose size is the number of spatial modes (e.g., 3 for a 5-D +/// convolution). +/// +/// @return TMA descriptor for im2col TMA load +template +CUTE_HOST +auto +make_im2col_tma_copy_desc( + Tensor const& tensor_cwhdn, // (C,W,H,D,N) + uint32_t range_c, // TILE_C + uint32_t range_whdn, // TILE_WHDN + SmemSwizzle const& smem_swizzle, // Swizzle + TMALayout const& tma_layout_vt, // TMA layout + LowerCornerStride const& lower_corner_whd, // WHD offset of the "base pointer" + UpperCornerStride const& upper_corner_whd, // WHD upper corner + LowerPaddingStride const& lower_padding_whd, // WHD lower padding + UpperPaddingStride const& upper_padding_whd, // WHD upper padding + TraversalStride const& stride_whd, // WHD traversal stride + LowerSRTStride const& lower_srt, // SRT offset of the "base pointer" + DilationStride const& stride_srt) // SRT stride - dilation +{ + static_assert(is_gmem::value, "Tensor must point to GPU global memory."); + using value_type = typename EngineA::value_type; + + constexpr uint32_t num_total_modes = LayoutA::rank; + constexpr int num_spatial_modes = num_total_modes - 2; + + // Gmem starting address + void* gmem_address = (void*) raw_pointer_cast(tensor_cwhdn.data()); + + // Gmem extents are just the tensor shape + cute::array gmem_prob_shape = {1,1,1,1,1}; + for_each(make_seq{}, [&](auto i) { + gmem_prob_shape[i] = static_cast(shape(tensor_cwhdn)); + }); + + // Gmem strides are byte strides of the activation tensor in CWHDN order + cute::array gmem_prob_stride = {0,0,0,0,0}; + for_each(make_seq{}, [&](auto i) { + gmem_prob_stride[i] = sizeof(value_type) * stride(tensor_cwhdn); + }); + + // Traversal strides are a function of the dilation shape + // corresponding to spatial (WHD) modes. + cute::array tma_traversal_strides = {1,1,1,1,1}; + for_each(make_seq{}, [&](auto i) { + tma_traversal_strides[i+1] = static_cast(get(stride_whd)); + }); + + cute::array tma_lower_corner{}; + for_each(make_seq{}, [&](auto i) { + tma_lower_corner[i] = static_cast(get(lower_corner_whd)); + }); + + cute::array tma_upper_corner{}; + for_each(make_seq{}, [&](auto i) { + tma_upper_corner[i] = static_cast(get(upper_corner_whd)); + }); + + Im2ColTmaDescriptor tma_desc; + +#if (__CUDACC_VER_MAJOR__ >= 12) + + 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_oob_fill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE; + CUtensorMapSwizzle tma_swizzle = TMA::to_CUtensorMapSwizzle(detail::get_tma_swizzle_bits(smem_swizzle)); + + CUresult encode_result = cuTensorMapEncodeIm2col( + &tma_desc, + tma_format, + num_total_modes, + gmem_address, + gmem_prob_shape.data(), + gmem_prob_stride.data() + 1, // gmem_prob_stride[0] implicitly sizeof(value_type) + tma_lower_corner.data(), + tma_upper_corner.data(), + range_c, + range_whdn, + tma_traversal_strides.data(), + tma_interleave, + tma_swizzle, + tma_l2Promotion, + tma_oob_fill); + + // The extra asserts help indicate the error's cause. + assert(encode_result != CUDA_ERROR_DEINITIALIZED); + assert(encode_result != CUDA_ERROR_NOT_INITIALIZED); + assert(encode_result != CUDA_ERROR_INVALID_CONTEXT); + assert(encode_result != CUDA_ERROR_INVALID_VALUE); + assert(encode_result == CUDA_SUCCESS); + +#endif // (__CUDACC_VER_MAJOR__ >= 12) + // + // Calculate gemm shapes and linearized shapes based on tma layout tiling. + // + + // Compute [w, h, d, n] + // q/p/z = (w/h/d + (upper_corner_whd - lower_corner_whd - 1)) / stride_whd + 1 + auto gemm_mn_ = cute::transform(cute::make_seq{}, [&](auto i) { + return (shape(tensor_cwhdn) + get(upper_corner_whd) - get(lower_corner_whd) - Int<1>{}) / get(stride_whd) + Int<1>{}; + }); + auto gemm_mn = append(gemm_mn_, shape(tensor_cwhdn)); + + // Compute [c, s, r, t] + // fprop/wgrad, s/r/t = 1 + (upper_padding_whd - upper_corner_whd) / stride_srt + // wgrad, s/r/t = 1 + (lower_padding_whd - lower_corner_whd) / stride_srt + auto gemm_k_ = cute::transform(cute::make_seq{}, [&](auto i) { + auto padding_size = conditional_return(get(stride_srt) > Int<0>{}, + get(upper_padding_whd) - get(upper_corner_whd), + get(lower_corner_whd) - get(lower_padding_whd)); + return Int<1>{} + padding_size / get(stride_srt); + }); + auto gemm_k = prepend(gemm_k_, shape<0>(tensor_cwhdn)); + + // For fprop/dgrad kernel, gemm_shapes is ((q, p, z, n), (c, s, r, t)) + // For wgrad kernel, gemm_shapes is ((c, s, r, t), (q, p, z, n)) + auto gemm_shapes_common = make_shape(gemm_mn, gemm_k); + auto gemm_shapes = make_shape( + basis_get(stride<0,1>(tma_layout_vt), gemm_shapes_common), + basis_get(stride<0,0>(tma_layout_vt), gemm_shapes_common)); + + // For fprop/dgrad kernel, linearized shapes is (whdn, (c, s, r, t)) + // For wgrad kernel linearized shapes is ((c, s, r, t), whdn) + auto linear_shapes_common = make_shape(size(gemm_mn), gemm_k); + auto linear_shapes = make_shape( + basis_get(stride<0,1>(tma_layout_vt), linear_shapes_common), + basis_get(stride<0,0>(tma_layout_vt), linear_shapes_common)); + + // + // Calculate gmem basis stride based on tma layout tiling. + // + + auto tma_basis_scale = make_shape(Int<1>{}, stride_whd, Int<1>{}, stride_srt); + auto tma_basis = elem_scale(tma_basis_scale, make_basis_like(tma_basis_scale)); + + auto gbasis_strides_common = make_stride( + append(get<1>(tma_basis), get<2>(tma_basis)), + prepend(get<3>(tma_basis), get<0>(tma_basis))); // ((w,h,d,n),(c,s,r,t)) + auto gbasis_strides = make_stride( + basis_get(stride<0,1>(tma_layout_vt), gbasis_strides_common), + basis_get(stride<0,0>(tma_layout_vt), gbasis_strides_common)); + + // + // Create tma tensor + // + + auto lower_corner = make_arithmetic_tuple(Int<0>{}, lower_corner_whd, Int<0>{}, lower_srt); + + auto tensor_multimode = make_tensor(ArithmeticTupleIterator(lower_corner), gemm_shapes, gbasis_strides); + auto tensor_linear = make_identity_tensor(linear_shapes); + auto tma_tensor = make_tensor(tensor_multimode.data(), composition( + tensor_multimode.layout(), + tensor_linear(Int<0>{}), + tensor_linear.layout())); + + return cute::make_tuple(tma_desc, tma_tensor); +} + +/// Make a TiledCopy for im2col TMA load. +/// +/// @param copy_op The copy implementation: either +/// SM90_TMA_LOAD_IM2COL or SM90_TMA_LOAD_IM2COL_MULTICAST. +/// +/// @param tensor_cwhdn The global tensor to use for im2col TMA loads. +/// For Fprop convolutions, this is the activation tensor. This is +/// the "original tensor that points to global memory, not the +/// coordinate (im2col-transformed) tensor. +/// +/// @param slayout Layout of shared memory tile. +/// +/// @param stride_whd The traversal strides convolution +/// parameter. +/// +/// @return TiledCopy specialization for im2col TMA loads. +template +CUTE_HOST_RTC +auto +make_tma_copy_im2col(CopyOp const& copy_op, + Tensor const& gtensor, + SLayout const& slayout, + Layout const& cta_t_map, // CTA tid -> logical TMA tid + Layout const& cta_v_map, // CTA vid -> gmem coord + LowerCornerStride const& lower_corner_whd, + UpperCornerStride const& upper_corner_whd, + LowerPaddingStride const& lower_padding_whd, + UpperPaddingStride const& upper_padding_whd, + TraversalStride const& stride_whd, // traversal stride + LowerSRTStride const& lower_srt, + DilationStride const& stride_srt) // dilation +{ + // + // 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 + // + + // 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)); + +#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"); +#endif + + // + // TMA gtensor manipulation + // + + // Generate a TupleBasis for the gtensor + auto glayout_basis = make_identity_layout(product_each(shape(gtensor))); + + // Tile the modes of gtensor with the truncated cta_v_map o inv_smem_layout_trunc + auto tma_layout_full = flatten(composition(glayout_basis, sidx_to_gmode)); + + // Truncate any incompatibilities -- no starting in the middle of gmodes + auto smem_rank = find_if(stride(tma_layout_full), [](auto e) { + [[maybe_unused]] auto v = basis_value(e); + return not is_constant<1,decltype(v)>{}; + }); + static_assert(smem_rank >= 2, "IM2COL expects at least 2 modes of the smem to vectorize with gmem."); + // IM2COL uses a maximum of 2 modes + constexpr int smem_tma_rank = cute::min(int(smem_rank), 2); + + // Keep only the static-1 basis modes into gmem + auto tma_layout_trunc = take<0,smem_tma_rank>(tma_layout_full); + + // Split according to the portion each multicast CTA will be responsible for + auto tma_layout_vt = logical_divide(tma_layout_trunc, shape_div(size(tma_layout_trunc), cosize(cta_t_map))); + +#if 0 + print("glayout_basis : "); print(glayout_basis); print("\n"); + print("tma_layout_full : "); print(tma_layout_full); print("\n"); + + print("tma_layout_trunc: "); print(tma_layout_trunc); print("\n"); + print("tma_layout_vt : "); print(tma_layout_vt); print("\n"); +#endif + + auto range_c = size<0,0>(tma_layout_vt); + auto range_whdn = size<0,1>(tma_layout_vt); + + Tensor gtensor_cwhdn = make_tensor(gtensor.data(), + flatten(make_layout(basis_get(stride<0,0>(tma_layout_vt), gtensor.layout()), + basis_get(stride<0,1>(tma_layout_vt), gtensor.layout())))); + + auto [tma_desc, tma_tensor] = make_im2col_tma_copy_desc( + gtensor_cwhdn, + range_c, + range_whdn, + detail::get_swizzle_portion(slayout), + tma_layout_vt, + lower_corner_whd, + upper_corner_whd, + lower_padding_whd, + upper_padding_whd, + stride_whd, + lower_srt, + stride_srt); + + // + // Construct the Copy_Traits + // + + using T = typename GEngine::value_type; + constexpr int num_bits_per_tma = decltype(size<0>(tma_layout_vt))::value * sizeof(T) * 8; + + using Traits = Copy_Traits, decltype(tma_tensor)>; + +#if 0 + print("num_bits : "); print(NumBitsPerTMA{}); print("\n"); +#endif + + Traits tma_traits{tma_desc, tma_tensor}; + + // + // 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(tma_layout_vt))); + // Scale that up to cover all of the smem_coords + // + // The smem vector might not cover all of the tile, + // so multiply it up to cover the entire tile. + // "T" here (the parallel index) is a CTA index. + 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}; +} + +/// Make a TiledCopy for im2col TMA with no offsets. +/// E.g. im2col TMA load for C and im2col TMA store for D. +template +CUTE_HOST_RTC +auto +make_tma_copy_im2col(CopyOp const& copy_op, + Tensor const& gtensor, + SLayout const& slayout, + Layout const& cta_t_map, // CTA tid -> logical TMA tid + Layout const& cta_v_map) // CTA vid -> gmem coord +{ + constexpr int num_spatial_modes = rank<0>(GLayout{}) - 1; + return make_tma_copy_im2col(copy_op, gtensor, slayout, cta_t_map, cta_v_map, + append(Stride<_0>{}, Int<0>{}), // lower_corner_whd + append(Stride<_0>{}, Int<0>{}), // upper_corner_whd + append(Stride<_0>{}, Int<0>{}), // lower_padding_whd + append(Stride<_0>{}, Int<0>{}), // upper_padding_whd + append(Stride<_1>{}, Int<1>{}), // stride_whd + append(Stride<_0>{}, Int<0>{}), // lower_srt + append(Stride<_1>{}, Int<1>{})); // stride_srt +} + +} // namespace detail + + + +template +CUTE_HOST_RTC +auto +make_im2col_tma_copy(CopyOp const& copy_op, + Tensor const& tensor_cwhdn, + SLayout const& slayout, + CTATiler const& cta_tiler, + MulticastSize const& multicast_size, + LowerCornerStride const& lower_corner_whd, + UpperCornerStride const& upper_corner_whd, + LowerPaddingStride const& lower_padding_whd, + UpperPaddingStride const& upper_padding_whd, + TraversalStride const& stride_whd, + LowerSRTStride const& lower_srt, + DilationStride const& stride_srt) +{ + auto cta_v_tile = make_identity_layout(product_each(shape(tensor_cwhdn))).compose(cta_tiler); + auto cta_t_tile = make_layout(multicast_size); + + return detail::make_tma_copy_im2col(copy_op, tensor_cwhdn, + slayout, cta_t_tile, cta_v_tile, + lower_corner_whd, upper_corner_whd, lower_padding_whd, upper_padding_whd, stride_whd, lower_srt, stride_srt); +} + +// Explicit default for multicast_size +template +CUTE_HOST_RTC +auto +make_im2col_tma_copy(CopyOp const& copy_op, + Tensor const& tensor_cwhdn, + SLayout const& slayout, + CTATiler const& cta_tiler, + LowerCornerStride const& lower_corner_whd, + UpperCornerStride const& upper_corner_whd, + LowerPaddingStride const& lower_padding_whd, + UpperPaddingStride const& upper_padding_whd, + TraversalStride const& stride_whd, + LowerSRTStride const& lower_srt, + DilationStride const& stride_srt) +{ + return make_im2col_tma_copy(copy_op, tensor_cwhdn, slayout, cta_tiler, Int<1>{}, + lower_corner_whd, upper_corner_whd, lower_padding_whd, upper_padding_whd, stride_whd, lower_srt, stride_srt); +} + +// Explicit default for cta_tiler and multicast_size +template +CUTE_HOST_RTC +auto +make_im2col_tma_copy(CopyOp const& copy_op, + Tensor const& tensor_cwhdn, + SLayout const& slayout, + LowerCornerStride const& lower_corner_whd, + UpperCornerStride const& upper_corner_whd, + LowerPaddingStride const& lower_padding_whd, + UpperPaddingStride const& upper_padding_whd, + TraversalStride const& stride_whd, + LowerSRTStride const& lower_srt, + DilationStride const& stride_srt) +{ + return make_im2col_tma_copy(copy_op, tensor_cwhdn, slayout, product_each(shape(slayout)), Int<1>{}, + lower_corner_whd, upper_corner_whd, lower_padding_whd, upper_padding_whd, stride_whd, lower_srt, stride_srt); +} + +// No offsets copy. +template +CUTE_HOST_RTC +auto +make_im2col_tma_copy(CopyOp const& copy_op, + Tensor const& tensor_cwhdn, + SLayout const& slayout, + CTATiler const& cta_tiler, + MulticastSize const& multicast_size) +{ + auto cta_v_tile = make_identity_layout(product_each(shape(tensor_cwhdn))).compose(cta_tiler); + auto cta_t_tile = make_layout(multicast_size); + + return detail::make_tma_copy_im2col(copy_op, tensor_cwhdn, slayout, cta_t_tile, cta_v_tile); +} + +// Explicit default for multicast_size +template +CUTE_HOST_RTC +auto +make_im2col_tma_copy(CopyOp const& copy_op, + Tensor const& tensor_cwhdn, + SLayout const& slayout, + CTATiler const& cta_tiler) +{ + return make_im2col_tma_copy(copy_op, tensor_cwhdn, slayout, cta_tiler, Int<1>{}); +} + +// Explicit default for cta_tiler and multicast_size +template +CUTE_HOST_RTC +auto +make_im2col_tma_copy(CopyOp const& copy_op, + Tensor const& tensor_cwhdn, + SLayout const& slayout) +{ + return make_im2col_tma_copy(copy_op, tensor_cwhdn, slayout, product_each(shape(slayout)), Int<1>{}); +} + +} // namespace cute diff --git a/include/cute/atom/copy_traits_sm90_tma.hpp b/include/cute/atom/copy_traits_sm90_tma.hpp index 842641bc..16b2a648 100644 --- a/include/cute/atom/copy_traits_sm90_tma.hpp +++ b/include/cute/atom/copy_traits_sm90_tma.hpp @@ -38,6 +38,8 @@ #include #include +#include + #include namespace cute @@ -53,77 +55,55 @@ struct AuxTmaParams { static_assert(is_static::value); }; +// Utility for unpacking TMA_LOAD arguments into a CopyOp +template +struct TMA_LOAD_Unpack +{ + template + CUTE_HOST_DEVICE friend constexpr void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, + Tensor & dst) + { + auto src_coord = src.data().coord_; + if constexpr (detail::is_prefetch) { + return detail::copy_explode(traits.opargs_, tuple_seq{}, + src_coord, tuple_seq{}); + } else { + static_assert(is_smem::value, "SM90_TMA_LOAD requires the destination be shared memory."); + void* dst_ptr = cute::raw_pointer_cast(dst.data()); +#if 0 + auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0); + printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n", + threadIdx.x, threadIdx.y, threadIdx.z, + blockIdx.x, blockIdx.y, blockIdx.z, + int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), dst_ptr); +#endif + return detail::copy_explode(traits.opargs_, tuple_seq{}, + make_tuple(dst_ptr), seq<0>{}, + src_coord, tuple_seq{}); + } + } +}; + ////////////////////////////////////////////////////////////////////////////// ///////////////////////////// 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 - auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0); - printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n", - threadIdx.x, threadIdx.y, threadIdx.z, - blockIdx.x, blockIdx.y, blockIdx.z, - int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), dst_ptr); -#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_smem::value, "Expected smem dst for SM90_TMA_LOAD"); - - traits.copy_unpack_(cute::raw_pointer_cast(dst.data()), 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>; - + 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; @@ -144,7 +124,7 @@ struct Copy_Traits Copy_Traits with(uint64_t& tma_mbar, [[maybe_unused]] uint16_t const& multicast_mask = 0) const { // We accept multicast_mask here to keep the API for both atoms consistent - return {tma_desc_, tma_mbar}; + return {{}, {&tma_desc_, &tma_mbar}}; } // Construct an executable SM90_TMA_LOAD with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm) @@ -152,7 +132,7 @@ struct Copy_Traits Copy_Traits with(TmaDescriptor const* new_tma_desc, uint64_t& tma_mbar, [[maybe_unused]] uint16_t const& multicast_mask = 0) const { // We accept multicast_mask here to keep the API for both atoms consistent - return {*new_tma_desc, tma_mbar}; + return {{}, {new_tma_desc, &tma_mbar}}; } // Generate the TMA coord tensor @@ -173,72 +153,65 @@ struct Copy_Traits Tensor & dst) = delete; }; +// The executable SM90_TMA_LOAD with tma_desc and tma_mbar +template +struct Copy_Traits + : TMA_LOAD_Unpack +{ + 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 + tuple< + TmaDescriptor const*, + uint64_t* // smem mbarrier + > const opargs_; +}; + +// The prefetch for SM90_TMA_LOAD with tma_desc +template +struct Copy_Traits + : TMA_LOAD_Unpack +{ + 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::PREFETCH arguments + tuple const opargs_; + + // Construct with any other Traits' TMA Desc + template + CUTE_HOST_DEVICE + Copy_Traits(Copy_Traits const& traits) + : opargs_({&traits.tma_desc_}) {} +}; + ////////////////////////////////////////////////////////////////////////////// ///////////////////////////// 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 - auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0); - printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n", - threadIdx.x, threadIdx.y, threadIdx.z, - blockIdx.x, blockIdx.y, blockIdx.z, - int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), dst_ptr); -#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_smem::value, "Expected smem dst for SM90_TMA_LOAD_MULTICAST"); - - traits.copy_unpack_(cute::raw_pointer_cast(dst.data()), src.data().coord_, tuple_seq{}); - } -}; - +// The non-executable SM90_TMA_LOAD_MULTICAST with tma_desc and no tma_mbar +// Use .with(tma_mbar, multicast_mask) to construct an executable version template struct Copy_Traits { - using ThrID = Layout<_1>; - + 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; @@ -258,14 +231,14 @@ struct Copy_Traits 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}; + return {{}, {&tma_desc_, &tma_load_mbar, multicast_mask}}; } // Construct an executable SM90_TMA_LOAD_MULTICAST_OP with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm) CUTE_HOST_DEVICE constexpr Copy_Traits with(TmaDescriptor const* new_tma_desc, uint64_t& tma_load_mbar, uint16_t const& multicast_mask) const { - return {*new_tma_desc, tma_load_mbar, multicast_mask}; + return {{}, {new_tma_desc, &tma_load_mbar, multicast_mask}}; } // Generate the TMA coord tensor @@ -286,6 +259,27 @@ struct Copy_Traits Tensor & dst) = delete; }; +// The executable SM90_TMA_LOAD_MULTICAST with tma_desc and tma_mbar and multicast_mask +template +struct Copy_Traits + : TMA_LOAD_Unpack +{ + 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 + tuple< + TmaDescriptor const*, + uint64_t*, // smem mbarrier + uint16_t // multicast mask + > const opargs_; +}; + ////////////////////////////////////////////////////////////////////////////// ///////////////////////////// TMA_STORE ////////////////////////////////////// ////////////////////////////////////////////////////////////////////////////// @@ -293,6 +287,68 @@ struct Copy_Traits // 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_; + using AuxParams = AuxParams_; + AuxParams aux_params_; + + // 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); + return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_)); + } + + 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 + + void const* const desc_ptr = &(traits.tma_desc_); + void const* const src_ptr = cute::raw_pointer_cast(src.data()); + auto dst_coord = dst.data().coord_; +#if 0 + auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0); + printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n", + threadIdx.x, threadIdx.y, threadIdx.z, + blockIdx.x, blockIdx.y, blockIdx.z, + int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), src_ptr); +#endif + return detail::copy_explode(make_tuple(desc_ptr, src_ptr), seq<0,1>{}, + dst_coord, tuple_seq{}); + } +}; + +////////////////////////////////////////////////////////////////////////////// +///////////////////////////// TMA_REDUCE_ADD ////////////////////////////////////// +////////////////////////////////////////////////////////////////////////////// + +// The executable SM90_TMA_REDUCE_ADD with tma_desc +template +struct Copy_Traits { using ThrID = Layout<_1>; @@ -304,7 +360,7 @@ struct Copy_Traits // Reference map from (thr,val) to bit using RefLayout = SrcLayout; - // SM90_TMA_STORE arguments + // SM90_TMA_REDUCE_ADD arguments TmaDescriptor tma_desc_; using AuxParams = AuxParams_; AuxParams aux_params_; @@ -339,7 +395,7 @@ struct Copy_Traits int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), src_ptr); #endif - SM90_TMA_STORE::copy(&tma_desc_, + SM90_TMA_REDUCE_ADD::copy(&tma_desc_, src_ptr, get(dst_coord)...); } @@ -354,8 +410,8 @@ struct Copy_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 + static_assert(is_smem::value, "Expected smem src for SM90_TMA_REDUCE_ADD"); + //static_assert(is_gmem::value, "Expected gmem dst for SM90_TMA_REDUCE_ADD"); // TMA spoofed src tensor traits.copy_unpack_(cute::raw_pointer_cast(src.data()), dst.data().coord_, tuple_seq{}); } @@ -383,6 +439,13 @@ struct Copy_Traits // 0: uint64_t* bulk_load_memory_barrier cute::tuple bulk_load_mbar_; + // Record the memory barrier for the instruction + CUTE_HOST_DEVICE constexpr + Copy_Traits + with(uint64_t& bulk_mbar) const { + return {{&bulk_mbar}}; + } + template CUTE_HOST_DEVICE friend constexpr @@ -395,15 +458,29 @@ struct Copy_Traits "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(raw_pointer_cast(src.data()), *get<0>(traits.bulk_load_mbar_), + SM90_BULK_COPY_G2S::copy(raw_pointer_cast(src.data()), get<0>(traits.bulk_load_mbar_), raw_pointer_cast(dst.data()), int32_t(NumBitsPerTMA::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 + : Copy_Traits +{ + template + CUTE_HOST_DEVICE + Copy_Traits(Copy_Traits const& traits) {} + + 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_BULK_PREFETCH"); + SM90_BULK_COPY_G2S::PREFETCH::copy(raw_pointer_cast(src.data()), int32_t(NumBitsPerTMA::value / 8)); } }; @@ -653,7 +730,7 @@ template const& gtensor, // Gmem Shapes and Strides, in units of TmaInternalType - TmaGmemBasisStride const& tma_gbasis_stride, // Map Tma mode idx -> Gmem mode(s) + TmaGmemBasisStride const& tma_gbasis_stride, // Map Tma mode idx -> Gmem mode(s) cute::array & gmem_prob_shape, // Tma Shapes, uint32_t or uin64_t cute::array & gmem_prob_stride) // Tma Strides { @@ -663,7 +740,7 @@ fill_tma_gmem_shape_stride(Tensor const& gtensor, / using TmaInternalType = typename GEngine::value_type; constexpr int tma_rank = decltype(rank(tma_gbasis_stride))::value; static_assert(TmaRank >= tma_rank); - + auto gmem_shape = shape(gtensor); auto gmem_stride = stride(gtensor); // Use the indirections in tma_gbasis_stride into gtensor to construct the tma gmem shapes/strides @@ -703,12 +780,12 @@ template CUTE_HOST_DEVICE constexpr void -fill_tma_gmem_shape_stride(Copy_Traits const& tma_traits, +fill_tma_gmem_shape_stride(Copy_Traits const& tma_traits, Tensor const& gtensor, // Gmem Shapes and Strides, value_type = TmaInternalType cute::array & gmem_prob_shape, // Tma Shapes, uint32_t or uin64_t cute::array & gmem_prob_stride) // Tma Strides { - return fill_tma_gmem_shape_stride(gtensor, stride(typename Aux::TmaGmemBasis{}), + return fill_tma_gmem_shape_stride(gtensor, stride(typename Aux::TmaGmemBasis{}), gmem_prob_shape, gmem_prob_stride); } @@ -824,7 +901,7 @@ make_tma_copy_desc(Tensor const& gtensor, // The origin // Construct the descriptor // - TmaDescriptor tma_desc = {0}; + TmaDescriptor tma_desc{}; // // TMA general info @@ -897,7 +974,7 @@ make_tma_copy_desc(Tensor const& gtensor, // The origin if constexpr (decltype(rank(tma_gmem_basis_stride) == Int<1>{})::value) { return E{}; // Return TMA Coord basis -- known scale of Int<1>{} } else { - int32_t scale = ceil_div(int32_t(di * sizeof_bits_v / cute::max(gmem_prob_stride[j], 16)), 8); + int32_t scale = ceil_div(int32_t(di * sizeof_bits_v / cute::max(gmem_prob_stride[j], uint64_t{16})), 8); return E{} * scale; // Return TMA Coord basis -- with a dynamic scale factor } } @@ -948,7 +1025,7 @@ make_tma_copy_atom(CopyOp, // Construct the Copy_Traits // - constexpr int num_bits_per_tma = size(tma_gbasis) * sizeof_bits::value; + constexpr int num_bits_per_tma = size(tma_gbasis) * sizeof_bits_v; using Traits = Copy_Traits, decltype(aux_params)>; using Atom = Copy_Atom; @@ -1105,6 +1182,14 @@ make_tma_copy(CopyOp const& copy_op, CTA_Tiler const& cta_tiler, Cluster_Size const& cluster_size) { + if constexpr (cute::is_same_v || + cute::is_same_v) { + return make_im2col_tma_copy(copy_op, + gtensor, + slayout, + cta_tiler, + cluster_size); + } else { auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler); auto cta_t_tile = make_layout(cluster_size); // Prefer TmaInternalType if specified. Fallback to GEngine::value_type @@ -1112,6 +1197,7 @@ make_tma_copy(CopyOp const& copy_op, return detail::make_tma_copy_tiled(copy_op, gtensor, slayout, cta_t_tile, cta_v_tile); + } } // Explicit defaulting @@ -1179,9 +1265,11 @@ auto tma_partition(Copy_Atom const& copy_atom, CtaCoord const& cta_coord, Layout const& cta_layout, // T: CTA coord -> logical multicast id - Tensor const& stensor, // SMEM Tensor (TMATile, Iter) - Tensor const& gtensor) // GMEM Tensor (TMATile, Iter) + Tensor const& stensor, // SMEM Tensor (TMATile, Rest...) + Tensor const& gtensor) // GMEM Tensor (TMATile, Rest...) { + CUTE_STATIC_ASSERT_V(size<0>(stensor) == size<0>(gtensor)); + // Invert the smem to get the largest contiguous vector in the smem layout Layout inv_smem_layout = right_inverse(get_nonswizzle_portion(layout<0>(stensor))); // Scale that up to cover all of the smem_coords @@ -1189,14 +1277,19 @@ tma_partition(Copy_Atom const& copy_atom, // Factor out the single-instrucion portion Layout tma_layout_v = make_layout(Int::NumValSrc>{}); - Layout layout_V = logical_divide(layout_v, tma_layout_v); - + auto layout_V = make_tile(logical_divide(layout_v, tma_layout_v)); + + // Append with _ until we cover all Rest... modes + auto glayout_V = append>(layout_V, _); + auto slayout_V = append>(layout_V, _); // Transform tile mode and coalesce - Tensor gtensor_v = coalesce(gtensor.compose(layout_V, _), Shape,_1>{}); // ((TMA,TMA_Iter),Iter) - Tensor stensor_v = coalesce(stensor.compose(layout_V, _), Shape,_1>{}); // ((TMA,TMA_Iter),Iter) + Tensor gtensor_v = coalesce(gtensor.compose(glayout_V), Shape>{}); // ((TMA,TMA_Iter), Rest...) + Tensor stensor_v = coalesce(stensor.compose(slayout_V), Shape>{}); // ((TMA,TMA_Iter), Rest...) #if 0 if (thread0()) { + print("gtensor : "); print(gtensor); print("\n"); + print("stensor : "); print(stensor); print("\n"); print("layout_V : "); print(layout_V); print("\n"); print("gtensor_v : "); print(gtensor_v); print("\n"); print("stensor_v : "); print(stensor_v); print("\n"); @@ -1205,11 +1298,15 @@ tma_partition(Copy_Atom const& copy_atom, // Restride the cta-into-tma-instr layout Layout tma_layout_t = composition(make_layout(Int<1>{}, shape_div(size(tma_layout_v), cosize(cta_layout))), cta_layout); - Layout tma_layout_tv = make_layout(tma_layout_t, tma_layout_v); + auto tma_layout_tv = make_tile(make_tile(make_layout(tma_layout_t, tma_layout_v), _)); + + // Append with _ until we cover all Rest... modes + auto gtma_layout_tv = append>(tma_layout_tv, _); + auto stma_layout_tv = append>(tma_layout_tv, _); // Transform TMA mode - Tensor gtensor_tv = gtensor_v.compose(make_tile(tma_layout_tv, _), _); // (((Thr,Frg),TMA_Iter),Iter) - Tensor stensor_tv = stensor_v.compose(make_tile(tma_layout_tv, _), _); // (((Thr,Frg),TMA_Iter),Iter) + Tensor gtensor_tv = gtensor_v.compose(gtma_layout_tv); // (((Thr,Frg),TMA_Iter), Rest...) + Tensor stensor_tv = stensor_v.compose(stma_layout_tv); // (((Thr,Frg),TMA_Iter), Rest...) #if 0 if (thread0()) { @@ -1219,9 +1316,11 @@ tma_partition(Copy_Atom const& copy_atom, } #endif - // Slice and group Frg,TMA_Iter and return - auto c = make_coord(make_coord(make_coord(cta_coord, _), _), _); - return cute::make_tuple(group_modes<0,2>(gtensor_tv(c)), group_modes<0,2>(stensor_tv(c))); + auto c = make_coord(make_coord(make_coord(cta_coord, _), _)); + auto c_s = append>(c, _); + auto c_g = append>(c, _); + + return cute::make_tuple(group_modes<0,2>(gtensor_tv(c_g)), group_modes<0,2>(stensor_tv(c_s))); } } // end namespace cute diff --git a/include/cute/atom/mma_atom.hpp b/include/cute/atom/mma_atom.hpp index 098efd4a..674e3519 100644 --- a/include/cute/atom/mma_atom.hpp +++ b/include/cute/atom/mma_atom.hpp @@ -35,7 +35,6 @@ #include #include - #include #include @@ -78,7 +77,7 @@ struct MMA_Atom> CUTE_HOST_DEVICE auto with(TraitsArgs&&... args) const { - auto traits = Traits::with(std::forward(args)...); + auto traits = Traits::with(static_cast(args)...); return MMA_Atom{traits}; } @@ -157,7 +156,7 @@ struct MMA_Atom> // 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)); + return make_tensor(static_cast(atensor)); } else { // Else, the intended FrgTypeA is a value type, construct a new tensor with a fragment layout return make_fragment_like(atensor); @@ -179,7 +178,7 @@ struct MMA_Atom> // 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)); + return make_tensor(static_cast(btensor)); } else { // Else, the intended FrgTypeB is a value type, construct a new tensor with a fragment layout return make_fragment_like(btensor); @@ -213,7 +212,7 @@ struct TiledMMA : MMA_Atom static_assert( rank_v == 3, "TiledMMA requires rank-3 AtomLayoutMNK"); static_assert( rank_v == 3, "TiledMMA requires rank-3 PermutationMNK"); - static_assert( is_tile::value, "TiledMMA requires independent permutations of MNK."); + static_assert( is_tuple::value, "TiledMMA requires independent permutations of MNK."); static_assert(is_static::value, "TiledMMA requires static permutations of MNK."); using ThrLayoutVMNK = decltype(tiled_product(AtomThrID{}, AtomLayoutMNK{})); @@ -391,7 +390,7 @@ struct TiledMMA : MMA_Atom } else { return cute::max(core_size, perm_size); } - + CUTE_GCC_UNREACHABLE; } @@ -517,7 +516,7 @@ struct ThrMMA : TiledMMA auto partition_C(CTensor&& ctensor) const { - auto thr_tensor = make_tensor(std::forward(ctensor).data(), this->thrfrg_C(ctensor.layout())); + auto thr_tensor = make_tensor(static_cast(ctensor).data(), this->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)>(_))); @@ -528,7 +527,7 @@ struct ThrMMA : TiledMMA auto partition_A(ATensor&& atensor) const { - auto thr_tensor = make_tensor(std::forward(atensor).data(), this->thrfrg_A(atensor.layout())); + auto thr_tensor = make_tensor(static_cast(atensor).data(), this->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)>(_))); @@ -539,7 +538,7 @@ struct ThrMMA : TiledMMA auto partition_B(BTensor&& btensor) const { - auto thr_tensor = make_tensor(std::forward(btensor).data(), this->thrfrg_B(btensor.layout())); + auto thr_tensor = make_tensor(static_cast(btensor).data(), this->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)>(_))); @@ -744,7 +743,15 @@ print(ThrMMA const& thr_mma) template CUTE_HOST_DEVICE -auto +void +print_latex(MMA_Atom const& mma_atom) +{ + print_latex(make_tiled_mma(mma_atom)); +} + +template +CUTE_HOST_DEVICE +void print_latex(TiledMMA const& mma) { auto layout_and_thrid_C = mma.get_layoutC_MN(); @@ -764,7 +771,7 @@ print_latex(TiledMMA const& mma) layoutB_NK, thrID_B); } -// MNK MMA Layout to console printer -- 8-value color coded by thread +// MNK MMA Layout to console printer template diff --git a/include/cute/atom/mma_traits_sm80.hpp b/include/cute/atom/mma_traits_sm80.hpp index 4f97d520..ab402881 100644 --- a/include/cute/atom/mma_traits_sm80.hpp +++ b/include/cute/atom/mma_traits_sm80.hpp @@ -32,12 +32,8 @@ #include #include - #include - -#include - -#include +#include namespace cute { diff --git a/include/cute/config.hpp b/include/cute/config.hpp index 57c537b3..941f60d7 100644 --- a/include/cute/config.hpp +++ b/include/cute/config.hpp @@ -98,16 +98,16 @@ # endif #endif -#ifdef _MSC_VER +#if defined(_MSC_VER) // Provides support for alternative operators 'and', 'or', and 'not' -#include +# include #endif // _MSC_VER #if defined(__CUDACC_RTC__) -#define CUTE_STL_NAMESPACE cuda::std -#define CUTE_STL_NAMESPACE_IS_CUDA_STD +# define CUTE_STL_NAMESPACE cuda::std +# define CUTE_STL_NAMESPACE_IS_CUDA_STD #else -#define CUTE_STL_NAMESPACE std +# define CUTE_STL_NAMESPACE std #endif // @@ -115,9 +115,9 @@ // #if defined(__CUDACC_RTC__) -#include +# include #else -#include +# include #endif #define CUTE_STATIC_V(x) decltype(x)::value @@ -125,10 +125,11 @@ #define CUTE_STATIC_ASSERT static_assert #define CUTE_STATIC_ASSERT_V(x,...) static_assert(decltype(x)::value, ##__VA_ARGS__) +// Fail and print a message. Typically used for notification of a compiler misconfiguration. #if defined(__CUDA_ARCH__) -# define CUTE_RUNTIME_ASSERT(x) __brkpt() +# define CUTE_INVALID_CONTROL_PATH(x) assert(0 && x); printf(x); __brkpt() #else -# define CUTE_RUNTIME_ASSERT(x) assert(0 && x) +# define CUTE_INVALID_CONTROL_PATH(x) assert(0 && x); printf(x) #endif // @@ -136,9 +137,9 @@ // #if !defined(__CUDACC_RTC__) -#include -#include -#include +# include +# include +# include #endif // @@ -151,13 +152,8 @@ // Basic types // -#include -#include -#include -#include -#include -#include -#include +#include + // // Debugging utilities // diff --git a/include/cute/container/alignment.hpp b/include/cute/container/alignment.hpp index cf91e1fb..509579ee 100644 --- a/include/cute/container/alignment.hpp +++ b/include/cute/container/alignment.hpp @@ -32,7 +32,7 @@ #include -#include +#include #include namespace cute diff --git a/include/cute/container/array.hpp b/include/cute/container/array.hpp index 3466b481..9cdcf5f4 100644 --- a/include/cute/container/array.hpp +++ b/include/cute/container/array.hpp @@ -355,7 +355,7 @@ void clear(array& a) a.fill(T(0)); } -template +template CUTE_HOST_DEVICE constexpr void fill(array& a, T const& value) { @@ -370,14 +370,14 @@ void swap(array& a, array& b) } /// @return A cute::array of the elements of @c t in reverse order. -template -CUTE_HOST_DEVICE constexpr cute::array -reverse(cute::array const& t) { +template +CUTE_HOST_DEVICE constexpr +cute::array reverse(cute::array const& t) +{ if constexpr (N == 0u) { return t; - } - else { - cute::array t_r{}; + } else { + cute::array t_r{}; for (size_t k = 0; k < N; ++k) { t_r[k] = t[N - k - 1]; } @@ -422,7 +422,7 @@ CUTE_HOST_DEVICE constexpr T&& get(array&& a) { static_assert(I < N, "Index out of range"); - return std::move(a[I]); + return cute::move(a[I]); } } // end namespace cute @@ -442,12 +442,12 @@ struct tuple_element> }; template -struct tuple_size> +struct tuple_size const> : CUTE_STL_NAMESPACE::integral_constant {}; template -struct tuple_element> +struct tuple_element const> { using type = T; }; @@ -462,7 +462,7 @@ namespace std template struct tuple_size; -template +template struct tuple_element; #endif @@ -478,12 +478,12 @@ struct tuple_element> }; template -struct tuple_size> +struct tuple_size const> : CUTE_STL_NAMESPACE::integral_constant {}; template -struct tuple_element> +struct tuple_element const> { using type = T; }; diff --git a/include/cute/container/array_subbyte.hpp b/include/cute/container/array_subbyte.hpp index 347b3710..3ab3bc32 100644 --- a/include/cute/container/array_subbyte.hpp +++ b/include/cute/container/array_subbyte.hpp @@ -37,29 +37,20 @@ #include -#include // sizeof_bits +#include #include namespace cute { - -template -struct is_subbyte { - static constexpr bool value = sizeof_bits_v < 8; -}; - -template -constexpr bool is_subbyte_v = is_subbyte::value; - // // Underlying subbyte storage type // template -using subbyte_storage_type_t = conditional_t<(sizeof_bits_v <= 8), uint8_t, - conditional_t<(sizeof_bits_v <= 16), uint16_t, - conditional_t<(sizeof_bits_v <= 32), uint32_t, - conditional_t<(sizeof_bits_v <= 64), uint64_t, - conditional_t<(sizeof_bits_v <= 128), uint128_t, +using subbyte_storage_type_t = conditional_t<(cute::sizeof_bits_v <= 8), uint8_t, + conditional_t<(cute::sizeof_bits_v <= 16), uint16_t, + conditional_t<(cute::sizeof_bits_v <= 32), uint32_t, + conditional_t<(cute::sizeof_bits_v <= 64), uint64_t, + conditional_t<(cute::sizeof_bits_v <= 128), uint128_t, T>>>>>; template struct subbyte_iterator; @@ -183,6 +174,11 @@ public: operator element_type() const { return get(); } + + // Address + subbyte_iterator operator&() const { + return {ptr_, idx_}; + } }; // @@ -314,7 +310,7 @@ public: CUTE_HOST_DEVICE constexpr friend auto recast_ptr(subbyte_iterator const& x) { using NewT = conditional_t<(is_const_v), NewT_ const, NewT_>; - if constexpr (is_subbyte::value) { // Making subbyte_iter, preserve the subbyte idx + if constexpr (cute::is_subbyte_v) { // Making subbyte_iter, preserve the subbyte idx return subbyte_iterator(x.ptr_, x.idx_); } else { // Not subbyte, assume/assert subbyte idx 0 return reinterpret_cast(raw_pointer_cast(x)); @@ -323,7 +319,7 @@ public: } CUTE_HOST_DEVICE friend void print(subbyte_iterator x) { - printf("subptr[%db](%p.%u)", int(sizeof_bits::value), x.ptr_, x.idx_); + printf("subptr[%db](%p.%u)", int(sizeof_bits_v), x.ptr_, x.idx_); } }; @@ -369,8 +365,8 @@ private: public: - CUTE_HOST_DEVICE constexpr - array_subbyte() {} + constexpr + array_subbyte() = default; CUTE_HOST_DEVICE constexpr array_subbyte(array_subbyte const& x) { @@ -562,7 +558,7 @@ CUTE_HOST_DEVICE constexpr T&& get(array_subbyte&& a) { static_assert(I < N, "Index out of range"); - return std::move(a[I]); + return cute::move(a[I]); } } // end namespace cute @@ -608,7 +604,7 @@ namespace std template struct tuple_size; -template +template struct tuple_element; #endif diff --git a/include/cute/container/bit_field.hpp b/include/cute/container/bit_field.hpp index fc82039e..c5748d84 100644 --- a/include/cute/container/bit_field.hpp +++ b/include/cute/container/bit_field.hpp @@ -37,7 +37,7 @@ #include -#include // uint_bit_t +#include // uint_bit_t namespace cute { diff --git a/include/cute/container/cuda_types.hpp b/include/cute/container/cuda_types.hpp index a0c23ffe..8034cb27 100644 --- a/include/cute/container/cuda_types.hpp +++ b/include/cute/container/cuda_types.hpp @@ -96,11 +96,11 @@ uint32_t&& get(dim3&& a) { static_assert(I < 3, "Index out of range"); if constexpr (I == 0) { - return std::move(a.x); + return cute::move(a.x); } else if constexpr (I == 1) { - return std::move(a.y); + return cute::move(a.y); } else if constexpr (I == 2) { - return std::move(a.z); + return cute::move(a.z); } CUTE_GCC_UNREACHABLE; @@ -162,11 +162,11 @@ uint32_t&& get(uint3&& a) { static_assert(I < 3, "Index out of range"); if constexpr (I == 0) { - return std::move(a.x); + return cute::move(a.x); } else if constexpr (I == 1) { - return std::move(a.y); + return cute::move(a.y); } else if constexpr (I == 2) { - return std::move(a.z); + return cute::move(a.z); } CUTE_GCC_UNREACHABLE; diff --git a/include/cute/container/tuple.hpp b/include/cute/container/tuple.hpp index ad49f2c1..0af98f56 100644 --- a/include/cute/container/tuple.hpp +++ b/include/cute/container/tuple.hpp @@ -126,18 +126,14 @@ CUTE_HOST_DEVICE constexpr T& getv(EBO& x) template CUTE_HOST_DEVICE constexpr T&& getv(EBO&& x) -{ return static_cast(x.t_); } +{ return cute::move(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. +// Base class of cute::tuple binds each element to an index +// by inheriting from EBO for each (i, t) in (I..., T...). +// The storage (for nonempty t) lives in the base classes. template struct TupleBase, T...> : EBO... @@ -169,11 +165,6 @@ struct TupleBase, T...> // // 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...> @@ -365,10 +356,10 @@ tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3, T4 const& t4, return cute::make_tuple(get(t0)..., get(t1)..., get(t2)..., get(t3)..., get(t4)...); } -template +template struct tuple_cat_static; -template +template struct tuple_cat_static, tuple> { using type = tuple; }; @@ -630,11 +621,8 @@ 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)}; + ((void(print(Is == 0 ? s : ',')), void(print(get(t)))), ...); print(e); } #if !defined(__CUDACC_RTC__) @@ -642,11 +630,8 @@ 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; + (void(os << (Is == 0 ? s : ',') << get(t)), ...); + return os << e; } #endif // !defined(__CUDACC_RTC__) @@ -707,7 +692,7 @@ namespace std template struct tuple_size; -template +template struct tuple_element; #endif diff --git a/include/cute/container/type_list.hpp b/include/cute/container/type_list.hpp index 7e790822..41c499ec 100644 --- a/include/cute/container/type_list.hpp +++ b/include/cute/container/type_list.hpp @@ -108,7 +108,7 @@ namespace std template struct tuple_size; -template +template struct tuple_element; #endif diff --git a/include/cute/int_tuple.hpp b/include/cute/int_tuple.hpp index fc8258c1..110e233a 100644 --- a/include/cute/int_tuple.hpp +++ b/include/cute/int_tuple.hpp @@ -218,7 +218,7 @@ static constexpr int depth_v = depth_t::value; // product // -// Implementation of product (see below) as a function object +// Implementation of product as a function object struct Product { template @@ -232,7 +232,7 @@ struct Product } else { return cute::transform_apply(a, Product{}, multiplies_unary_lfold{}); } - } else { + } else if constexpr (cute::is_integral::value) { return a; } @@ -248,7 +248,7 @@ CUTE_HOST_DEVICE constexpr auto product_each(Tuple const& t) { - return transform(wrap(t), [](auto const& x) { return product(x); }); + return transform(wrap(t), product); } // Take the product of Tuple at the leaves of TupleG @@ -394,7 +394,7 @@ shape_div(IntTupleA const& a, IntTupleB const& b) static_assert(IntTupleA::value % IntTupleB::value == 0 || IntTupleB::value % IntTupleA::value == 0, "Static shape_div failure"); return C{}; } else { // int int - //assert(a % b == 0 || b % a == 0); // Wave dynamic assertion + //assert(a % b == 0 || b % a == 0); // Waive dynamic assertion return a / b != 0 ? a / b : signum(a) * signum(b); // Division with rounding away from zero } @@ -855,7 +855,10 @@ elem_geq(T const& t, U const& u) { return !elem_less(t, u); } +namespace detail { + /** Increment a (dynamic) coord lexicographically within a shape + * @pre is_congruent::value * \code * auto shape = make_shape(1,2,make_shape(2,3),3); * @@ -866,44 +869,26 @@ elem_geq(T const& t, U const& u) { * assert(i == size(shape)); * \endcode */ -template +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) +increment(Coord& coord, Shape const& shape) { - 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{}); + if constexpr (is_integral::value) { + ++coord; + } else { + increment(get(coord), get(shape)); + if constexpr (I+1 < tuple_size::value) { + if (back(get(coord)) == back(get(shape))) { + back(get(coord)) = 0; + increment(coord, shape); + } } } } } // 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 {}; @@ -918,7 +903,7 @@ struct ForwardCoordIterator Coord const& operator*() const { return coord; } CUTE_HOST_DEVICE constexpr - ForwardCoordIterator& operator++() { increment(coord, shape); return *this; } + ForwardCoordIterator& operator++() { detail::increment(coord, shape); return *this; } // Sentinel for the end of the implied range CUTE_HOST_DEVICE constexpr diff --git a/include/cute/layout.hpp b/include/cute/layout.hpp index 586ea037..71c4ce13 100644 --- a/include/cute/layout.hpp +++ b/include/cute/layout.hpp @@ -56,6 +56,9 @@ using Step = cute::tuple; template using Coord = cute::tuple; +template +using Tile = cute::tuple; + template CUTE_HOST_DEVICE constexpr Shape @@ -80,7 +83,17 @@ Coord make_coord(Ts const&... t) { return {t...}; } +template +CUTE_HOST_DEVICE constexpr +Tile +make_tile(Ts const&... t) +{ + return {t...}; +} +// +// Layout +// template > struct Layout @@ -366,59 +379,56 @@ 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 +// +// Advanced Layout constructions +// + +// Make a compact layout with shape @a shape and strides following the order induced by @a order. +// Dynamic values in @a order are ignored, considered large, and considered ordered from left to right. +// Example: +// make_ordered_layout(Shape<_2,_2,_2,_2>{}, Step<_0,_2,_3,_1>{}) +// -> (_2,_2,_2,_2):(_1,_4,_8,_2) +// make_ordered_layout(make_shape(2,3,4,5), make_step(Int<2>{}, 67, 42, Int<50>{})) +// -> (2,3,4,5):(_1,10,30,2) template CUTE_HOST_DEVICE constexpr auto make_ordered_layout(Shape const& shape, Order const& order) { - static_assert(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 +// Make a compact layout with the same shape as @a layout +// and strides following the order induced by @a layout.stride(). +// Static-0 strides in the input @a layout are preserved in the output. +// Example: +// make_layout_like(Layout, Stride<_0,_2,_4,_1>>{}) +// -> (_2,_2,_2,_2):(_0,_2,_4,_1) +// make_layout_like(make_layout(make_shape(2,3,4,5), make_stride(Int<0>{},42,Int<1>{},Int<0>{}))) +// -> (2,3,4,5):(_0,4,_1,_0) 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; + return make_layout(layout.shape(), + compact_order(filter_zeros(layout.stride(), layout.shape()), layout.stride())); } -// -// Make a layout of the same shape, -// with mode-0 being colmajor then following the mode order in layout -// +// Make a compact layout with the same shape as @a layout +// and strides following the order induced by @a layout.stride(), +// except mode-0 is always stride-1 and generated column-major. +// The 0th mode is commonly used for MMA_Atoms or Copy_Atoms +// so this generates the 0th mode with LayoutLeft regardless of the reference 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))); + if constexpr (R > 1 && is_static::value) { + return tiled_product(make_layout(shape<0>(layout)), + make_ordered_layout(take<1,R>(layout.shape()), take<1,R>(layout.stride()))); } else { return make_layout(layout.shape()); } @@ -458,11 +468,11 @@ CUTE_HOST_DEVICE constexpr auto get(Layout const& layout) { - return make_layout(get(layout.shape()), + return make_layout(get(layout.shape()), get(layout.stride())); } -// Return a new layout with only the modes in the range [B,E) +// Return a new layout with only the modes in the range [B,E) template CUTE_HOST_DEVICE constexpr auto @@ -470,7 +480,7 @@ take(Layout const& layout) { static_assert(B < E, "take: empty range error"); static_assert(0 <= B && E <= Layout::rank, "take: range out of bounds"); - return make_layout(take(layout.shape()), + return make_layout(take(layout.shape()), take(layout.stride())); } @@ -490,7 +500,7 @@ CUTE_HOST_DEVICE constexpr auto flatten(Layout const& layout) { - return make_layout(flatten(layout.shape()), + return make_layout(flatten(layout.shape()), flatten(layout.stride())); } @@ -1376,6 +1386,23 @@ logical_divide(Layout const& layout, CUTE_GCC_UNREACHABLE; } +// Generalization of ceil_div for Layout lhs +// is effectively the "rest mode" of logical_divide. +// Occurs in the calculation of gridDim, for example, for generalized tilers +// Example: +// dim3 gridDim(size(ceil_div(problem_shape_M, cta_tiler_M)), +// size(ceil_div(problem_shape_N, cta_tiler_N))); +// This does not consider compositional acceptance, so it may be the case that +// ceil_div produces a result while logical_divide (and friends) do not. +template +CUTE_HOST_DEVICE constexpr +auto +ceil_div(Target const& target, + Layout const& tiler) +{ + return complement(tiler, size(target)); +} + // // Convenience operator // that produces layouts like ((BLK_A,BLK_B,...),(a,b,...,x,y)) @@ -1425,7 +1452,6 @@ flat_divide(Layout const& layout, // Logical product // -// @post compatible() template CUTE_HOST_DEVICE constexpr @@ -1501,7 +1527,7 @@ flat_product(Layout const& block, // // Rank-sensitive products -// +// // blocked_product -- Reproduce a block over a tiler. // Think of every element of "tiler" as a "block" @@ -1517,7 +1543,7 @@ blocked_product(Layout const& block, constexpr int R = cute::max(rank_v, rank_v); auto result = logical_product(append(block), append(tiler)); - + return coalesce(zip(get<0>(result), get<1>(result)), tuple_repeat(Int<1>{})); } @@ -1545,7 +1571,7 @@ raked_product(Layout const& block, // @param block The layout to repeat // @param trg_shape The target shape of the result // @param ord_shape The order of the modes of @a trg_shape to tile @a layout with. -// Defaults to GenColMajor, so @a layout will repeat +// Defaults to GenColMajor, so @a layout will repeat // across the first mode first, the second mode second, etc // E.g. Step<_2,_1,_3> will cause @a layout to repeat // across the second mode first, the first mode second, and the third mode last. @@ -1659,7 +1685,7 @@ recast_layout(Layout const& layout) else if constexpr (scale::num == 1) { return downcast(layout); } - else if constexpr (scale::den == 1) { + else if constexpr (scale::den == 1) { return upcast(layout); } else { diff --git a/include/cute/layout_composed.hpp b/include/cute/layout_composed.hpp index e1c73aa5..93c60898 100644 --- a/include/cute/layout_composed.hpp +++ b/include/cute/layout_composed.hpp @@ -178,6 +178,16 @@ struct ComposedLayout : private cute::tuple // EBO fo tile(Layouts const&... layouts) const { return tiled_divide(*this, make_tile(layouts...)); } + + // Equality, return a static or dynamic boolean + template + CUTE_HOST_DEVICE constexpr + auto + operator==(ComposedLayout const& other) const { + return this->layout_a() == other.layout_a() && + this->layout_b() == other.layout_b() && + this->offset() == other.offset(); + } }; template diff --git a/include/cute/numeric/arithmetic_tuple.hpp b/include/cute/numeric/arithmetic_tuple.hpp index ac5af864..27d1cf8e 100644 --- a/include/cute/numeric/arithmetic_tuple.hpp +++ b/include/cute/numeric/arithmetic_tuple.hpp @@ -63,6 +63,9 @@ struct ArithmeticTuple : tuple template struct is_tuple> : true_type {}; +template +struct is_flat> : is_flat> {}; + template CUTE_HOST_DEVICE constexpr auto @@ -108,16 +111,45 @@ 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...); }); + return t + ArithmeticTuple(u); } template CUTE_HOST_DEVICE constexpr auto operator+(tuple const& t, ArithmeticTuple const& u) { + return ArithmeticTuple(t) + u; +} + +// Subtraction +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...); }); + return transform_apply(append(t,Int<0>{}), append(u,Int<0>{}), minus{}, [](auto const&... a){ return make_arithmetic_tuple(a...); }); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator-(ArithmeticTuple const& t, tuple const& u) { + return t - ArithmeticTuple(u); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator-(tuple const& t, ArithmeticTuple const& u) { + return ArithmeticTuple(t) - u; +} + +// Negation +template +CUTE_HOST_DEVICE constexpr +auto +operator-(ArithmeticTuple const& t) { + return transform_apply(t, negate{}, [](auto const&... a){ return make_arithmetic_tuple(a...); }); } // @@ -128,7 +160,7 @@ template CUTE_HOST_DEVICE constexpr ArithmeticTuple const& operator+(C, ArithmeticTuple const& u) { - static_assert(t == 0, "Artihmetic tuple op+ error!"); + static_assert(t == 0, "Arithmetic tuple op+ error!"); return u; } @@ -136,7 +168,23 @@ template CUTE_HOST_DEVICE constexpr ArithmeticTuple const& operator+(ArithmeticTuple const& t, C) { - static_assert(u == 0, "Artihmetic tuple op+ error!"); + static_assert(u == 0, "Arithmetic tuple op+ error!"); + return t; +} + +template +CUTE_HOST_DEVICE constexpr +ArithmeticTuple const& +operator-(C, ArithmeticTuple const& u) { + static_assert(t == 0, "Arithmetic tuple op- error!"); + return -u; +} + +template +CUTE_HOST_DEVICE constexpr +ArithmeticTuple const& +operator-(ArithmeticTuple const& t, C) { + static_assert(u == 0, "Arithmetic tuple op- error!"); return t; } @@ -531,7 +579,7 @@ namespace std template struct tuple_size; -template +template struct tuple_element; #endif diff --git a/include/cute/numeric/complex.hpp b/include/cute/numeric/complex.hpp index 5f1f41b5..8cc36253 100644 --- a/include/cute/numeric/complex.hpp +++ b/include/cute/numeric/complex.hpp @@ -30,9 +30,9 @@ **************************************************************************************************/ #pragma once -#include -#include #include +#include +#include namespace cute { @@ -44,6 +44,9 @@ using cutlass::real; using cutlass::imag; using cutlass::conj; +template +static constexpr auto is_complex_v = is_complex::value; + /// Fused multiply-add for complex numbers template CUTE_HOST_DEVICE constexpr diff --git a/include/cute/numeric/half.hpp b/include/cute/numeric/half.hpp deleted file mode 100644 index 8b904ec3..00000000 --- a/include/cute/numeric/half.hpp +++ /dev/null @@ -1,41 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (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 cutlass::half_t; - -} // end namespace cute diff --git a/include/cute/numeric/int.hpp b/include/cute/numeric/int.hpp index 00f62995..169e3a0e 100644 --- a/include/cute/numeric/int.hpp +++ b/include/cute/numeric/int.hpp @@ -36,8 +36,7 @@ #include #endif -#include -#include +#include namespace cute { @@ -46,16 +45,16 @@ 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; +using int2_t = cutlass::int2b_t; +using int4_t = cutlass::int4b_t; +using CUTE_STL_NAMESPACE::int8_t; +using CUTE_STL_NAMESPACE::int16_t; +using CUTE_STL_NAMESPACE::int32_t; +using 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< 2> { using type = cutlass::int2b_t; }; +template <> struct int_bit< 4> { using type = cutlass::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; }; @@ -74,24 +73,24 @@ 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; +using uint1_t = cutlass::uint1b_t; +using uint2_t = cutlass::uint2b_t; +using uint4_t = cutlass::uint4b_t; +using CUTE_STL_NAMESPACE::uint8_t; +using CUTE_STL_NAMESPACE::uint16_t; +using CUTE_STL_NAMESPACE::uint32_t; +using CUTE_STL_NAMESPACE::uint64_t; +using cutlass::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< 1> { using type = cutlass::uint1b_t; }; +template <> struct uint_bit< 2> { using type = cutlass::uint2b_t; }; +template <> struct uint_bit< 4> { using type = cutlass::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 <> struct uint_bit<128> { using type = cutlass::uint128_t; }; template using uint_bit_t = typename uint_bit::type; @@ -102,50 +101,4 @@ 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: sizeof_bits {}; - -template <> -struct sizeof_bits { - static constexpr size_t value = 0; -}; - -template <> -struct sizeof_bits { - static constexpr size_t value = 1; -}; - -template -struct sizeof_bits> { - static constexpr size_t value = Bits; -}; - -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/include/cute/numeric/integer_subbyte.hpp b/include/cute/numeric/integer_subbyte.hpp deleted file mode 100644 index 9f0d11a0..00000000 --- a/include/cute/numeric/integer_subbyte.hpp +++ /dev/null @@ -1,235 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (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(__CUDACC_RTC__) -#include -#else -#include -#endif - -#include - -#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(__CUDACC_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(__CUDACC_RTC__) diff --git a/include/cute/numeric/integral_constant.hpp b/include/cute/numeric/integral_constant.hpp index 977812b6..904a6726 100644 --- a/include/cute/numeric/integral_constant.hpp +++ b/include/cute/numeric/integral_constant.hpp @@ -443,4 +443,35 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, C const&) { } #endif + +namespace detail { + +// parse_int_digits takes a variadic number of digits and converts them into an int +template +constexpr uint64_t parse_int_digits(uint64_t result, int digit, Ts... digits) +{ + if constexpr (sizeof...(Ts) == 0) { + return 10 * result + digit; + } else { + return parse_int_digits(10 * result + digit, digits...); + } +} + +} // end namespace detail + + +// This user-defined literal operator allows cute::constant written as literals. For example, +// +// auto var = 32_c; +// +// var has type cute::constant. +// +template +constexpr cute::constant operator "" _c() +{ + static_assert((('0' <= digits && digits <= '9') && ...), + "Expected 0 <= digit <= 9 for each digit of the integer."); + return {}; +} + } // end namespace cute diff --git a/include/cute/numeric/integral_ratio.hpp b/include/cute/numeric/integral_ratio.hpp index e6cf502b..943b0049 100644 --- a/include/cute/numeric/integral_ratio.hpp +++ b/include/cute/numeric/integral_ratio.hpp @@ -130,6 +130,10 @@ nratio(R, R) { return {}; } +// +// Operators +// + template CUTE_HOST_DEVICE constexpr typename R::type @@ -227,14 +231,14 @@ abs(R) { template CUTE_HOST_DEVICE constexpr -auto +int32_t log_2(R) { static_assert(R::num > 0); static_assert(R::den > 0); return log_2(static_cast(R::num)) - log_2(static_cast(R::den)); } - +// @return A non-reduced ratio cute::R of the Trait0::value / Trait1::value template CUTE_HOST_DEVICE constexpr auto diff --git a/include/cute/numeric/math.hpp b/include/cute/numeric/math.hpp index 3f234e46..5be50339 100644 --- a/include/cute/numeric/math.hpp +++ b/include/cute/numeric/math.hpp @@ -316,11 +316,11 @@ safe_div(T const& t, U const& u) { template CUTE_HOST_DEVICE constexpr -auto +int32_t log_2(T x) { assert(x > 0); static_assert(is_unsigned::value, "Only to be used for unsigned integral types."); - return bit_width(x) - 1; + return static_cast(bit_width(x)) - 1; } } // namespace cute diff --git a/include/cute/numeric/float8.hpp b/include/cute/numeric/numeric_types.hpp similarity index 68% rename from include/cute/numeric/float8.hpp rename to include/cute/numeric/numeric_types.hpp index 1d3906fc..02c70025 100644 --- a/include/cute/numeric/float8.hpp +++ b/include/cute/numeric/numeric_types.hpp @@ -30,14 +30,46 @@ **************************************************************************************************/ #pragma once -#include - #include #include +#include + +#include +#include namespace cute { +template +struct sizeof_bits : public cutlass::sizeof_bits {}; + +// DO NOT change auto to int, sizeof_bits use integral_ratio instead of int +template +static constexpr auto sizeof_bits_v = sizeof_bits::value; + +using cutlass::bits_to_bytes; + +using cutlass::is_subbyte; + +template +static constexpr auto is_subbyte_v = is_subbyte::value; + +using cutlass::half_t; +using cutlass::bfloat16_t; + +using cutlass::tfloat32_t; + +// Umbrella floating-point 8-bit data type : type_erased_dynamic_float8_t +// This umbrella datatype can be enabled when a user provides a specific +// datatype in runtime argument list. +using cutlass::type_erased_dynamic_float8_t; using cutlass::float_e4m3_t; using cutlass::float_e5m2_t; +using cutlass::uint1b_t; +using cutlass::int2b_t; +using cutlass::uint2b_t; +using cutlass::int4b_t; +using cutlass::uint4b_t; +using cutlass::bin1_t; + } // end namespace cute diff --git a/include/cute/numeric/uint128.hpp b/include/cute/numeric/uint128.hpp deleted file mode 100644 index bd5ec07b..00000000 --- a/include/cute/numeric/uint128.hpp +++ /dev/null @@ -1,259 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are met: - * - * 1. Redistributions of source code must retain the above copyright notice, this - * list of conditions and the following disclaimer. - * - * 2. Redistributions in binary form must reproduce the above copyright notice, - * this list of conditions and the following disclaimer in the documentation - * and/or other materials provided with the distribution. - * - * 3. Neither the name of the copyright holder nor the names of its - * contributors may be used to endorse or promote products derived from - * this software without specific prior written permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" - * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE - * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE - * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL - * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR - * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER - * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, - * OR TORT (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(__CUDACC_RTC__) -#include -#else -#include -#include -#include -#include -#include -#endif - -#include - -/// Optionally enable GCC's built-in type -#if defined(__x86_64) && !defined(__CUDA_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/include/cute/pointer.hpp b/include/cute/pointer.hpp index 72812e45..5647f97c 100644 --- a/include/cute/pointer.hpp +++ b/include/cute/pointer.hpp @@ -33,7 +33,7 @@ #include #include -#include // sizeof_bits +#include // sizeof_bits #include #include @@ -58,7 +58,7 @@ CUTE_HOST_DEVICE constexpr auto recast_ptr(void* ptr) { - if constexpr (is_subbyte::value) { + if constexpr (cute::is_subbyte_v) { return subbyte_iterator(ptr); } else { return reinterpret_cast(ptr); @@ -71,7 +71,7 @@ CUTE_HOST_DEVICE constexpr auto recast_ptr(void const* ptr) { - if constexpr (is_subbyte::value) { + if constexpr (cute::is_subbyte_v) { return subbyte_iterator(ptr); } else { return reinterpret_cast(ptr); diff --git a/include/cute/pointer_base.hpp b/include/cute/pointer_base.hpp index 98109eda..db5d3dcf 100644 --- a/include/cute/pointer_base.hpp +++ b/include/cute/pointer_base.hpp @@ -33,7 +33,7 @@ #include #include -#include // sizeof_bits +#include // sizeof_bits namespace cute { diff --git a/include/cute/pointer_flagged.hpp b/include/cute/pointer_flagged.hpp index 1e32a581..aa917d9f 100644 --- a/include/cute/pointer_flagged.hpp +++ b/include/cute/pointer_flagged.hpp @@ -109,7 +109,7 @@ as_position_independent_swizzle_tensor(Tensor&& tensor) } else { #if !defined(NDEBUG) { - uint32_t address = cast_smem_ptr_to_uint(raw_pointer_cast(std::forward(tensor).data())); + uint32_t address = cast_smem_ptr_to_uint(raw_pointer_cast(static_cast(tensor).data())); uint32_t mask = ((uint32_t(1) << SwizzleFn::num_base) - 1) | SwizzleFn::swizzle_code; assert((address & mask) == 0); // Alignment to the Base, Z, and Y of Swizzle } @@ -118,7 +118,7 @@ as_position_independent_swizzle_tensor(Tensor&& tensor) // Recast swizzle from acting on byte-addressed pointers to elements of type-T auto new_swizzle = recast_layout(SwizzleFn{}); // Strip off everything and create a new smem_ptr for type-T - auto new_ptr = make_smem_ptr(raw_pointer_cast(std::forward(tensor).data())); + auto new_ptr = make_smem_ptr(raw_pointer_cast(static_cast(tensor).data())); return make_tensor(new_ptr, composition(new_swizzle, Int<0>{}, tensor.layout())); } CUTE_GCC_UNREACHABLE; diff --git a/include/cute/stride.hpp b/include/cute/stride.hpp index 3fb77b54..92020048 100644 --- a/include/cute/stride.hpp +++ b/include/cute/stride.hpp @@ -394,20 +394,26 @@ compact_row_major(Shape const& shape, namespace detail { -template +// @pre weakly_congruent(order, shape) +// @pre is_congruent +// @pre is_static +// @pre is_static +template CUTE_HOST_DEVICE constexpr auto compact_order(Shape const& shape, Order const& order, - OrigShape const& orig_shape, OrigOrder const& orig_order) + RefShape const& ref_shape, RefOrder const& ref_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); }); + static_assert(tuple_size::value == tuple_size::value, "Need equal rank of shape and order"); + return transform(shape, order, [&](auto const& s, auto const& o) { return compact_order(s, o, ref_shape, ref_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); + // Compute the starting stride for this shape by accumulating all shapes corresponding to lesser orders + auto stride_start = product(transform(ref_shape, ref_order, + [&](auto const& s, auto const& o) { + return conditional_return(o < order, s, Int<1>{}); + })); + return compact_col_major(shape, stride_start); } CUTE_GCC_UNREACHABLE; @@ -420,15 +426,30 @@ CUTE_HOST_DEVICE constexpr auto compact_order(Shape const& shape, Order const& order) { - if constexpr(is_congruent::value) { - return detail::compact_order(shape, order, flatten_to_tuple(shape), flatten_to_tuple(order)); - } - else - { - // Here we only want to apply order to top-level subshapes and default (col-major) order on other levels - static_assert(rank(Shape{}) == rank(Order{}), "Need equal rank of shape and order"); - return detail::compact_order(shape, order, shape, order); - } + auto ref_shape = flatten_to_tuple(product_like(shape, order)); + + auto flat_order = flatten_to_tuple(order); + // Find the largest static element of order + auto max_order = cute::fold(flat_order, Int<0>{}, [](auto v, auto order) { + if constexpr (is_constant::value) { + return order; + } else { + return v; + } + }); + // Replace any dynamic elements within order with large-static elements + auto max_seq = make_range{}; + auto ref_order = cute::transform(max_seq, flat_order, [](auto seq_v, auto order) { + if constexpr (is_static::value) { + return order; + } else { + return seq_v; + } + }); + + auto new_order = unflatten(ref_order, order); + + return detail::compact_order(shape, new_order, ref_shape, ref_order); } template diff --git a/include/cute/swizzle.hpp b/include/cute/swizzle.hpp index 90377445..57735ce1 100644 --- a/include/cute/swizzle.hpp +++ b/include/cute/swizzle.hpp @@ -87,6 +87,14 @@ struct Swizzle { return apply(offset); } + + template + CUTE_HOST_DEVICE constexpr + auto + operator==(Swizzle const&) const + { + return B == BBits && M == MBase && S == SShift; + } }; // diff --git a/include/cute/tensor.hpp b/include/cute/tensor.hpp index e5e2a8f2..28d3ee67 100644 --- a/include/cute/tensor.hpp +++ b/include/cute/tensor.hpp @@ -42,7 +42,6 @@ #include #include -#include namespace cute { @@ -481,7 +480,7 @@ CUTE_HOST_DEVICE constexpr decltype(auto) tensor(Tensor&& tensor) { - return std::forward(tensor); + return static_cast(tensor); } template (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 make_tensor(static_cast(tensor).data(), get(tensor.layout())); } // Return the layout of a mode @@ -567,7 +556,7 @@ CUTE_HOST_DEVICE constexpr auto flatten(Tensor&& tensor) { - return make_tensor(std::forward(tensor).data(), flatten(tensor.layout())); + return make_tensor(static_cast(tensor).data(), flatten(tensor.layout())); } template (tensor).data(), coalesce(tensor.layout())); + return make_tensor(static_cast(tensor).data(), coalesce(tensor.layout())); } template (tensor).data(), coalesce(tensor.layout(), profile)); + return make_tensor(static_cast(tensor).data(), coalesce(tensor.layout(), profile)); } template (tensor).data(), filter_zeros(tensor.layout())); + return make_tensor(static_cast(tensor).data(), filter_zeros(tensor.layout())); } template (tensor).data(), filter(tensor.layout())); + return make_tensor(static_cast(tensor).data(), filter(tensor.layout())); } template (tensor).data(), filter(tensor.layout(), profile)); + return make_tensor(static_cast(tensor).data(), filter(tensor.layout(), profile)); } // Return a tensor with the same shape as input but offset by a given coordinate @@ -623,7 +612,7 @@ auto domain_offset(Coord const& coord, Tensor&& tensor) { auto [layout, ptr_offset] = domain_offset(coord, tensor.layout()); - return make_tensor(std::forward(tensor).data() + ptr_offset, layout); + return make_tensor(static_cast(tensor).data() + ptr_offset, layout); } // Group the modes [B,E) into a single mode @@ -635,10 +624,20 @@ CUTE_HOST_DEVICE constexpr auto group_modes(Tensor&& tensor) { - return make_tensor(std::forward(tensor).data(), + return make_tensor(static_cast(tensor).data(), group(tensor.layout())); } +// Return the subtensor of a range of modes +template >::value)> +CUTE_HOST_DEVICE constexpr +decltype(auto) +take(Tensor&& tensor) +{ + return make_tensor(static_cast(tensor).data(), take(tensor.layout())); +} + // // Recast // @@ -662,9 +661,9 @@ recast(Tensor&& tensor) 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_ptr(std::forward(tensor).data() + offset), new_layout); + return make_tensor(recast_ptr(static_cast(tensor).data() + offset), new_layout); } else { - return make_tensor(recast_ptr(std::forward(tensor).data() ), new_layout); + return make_tensor(recast_ptr(static_cast(tensor).data() ), new_layout); } CUTE_GCC_UNREACHABLE; @@ -788,7 +787,7 @@ auto logical_divide(Tensor && tensor, Tiler const& tiler) // Layout or Tile or Shape { - return make_tensor(std::forward(tensor).data(), + return make_tensor(static_cast(tensor).data(), logical_divide(tensor.layout(), tiler)); } @@ -802,7 +801,7 @@ auto zipped_divide(Tensor && tensor, Tiler const& tiler) // Layout or Tile or Shape { - return make_tensor(std::forward(tensor).data(), + return make_tensor(static_cast(tensor).data(), zipped_divide(tensor.layout(), tiler)); } @@ -811,10 +810,10 @@ template >::value)> CUTE_HOST_DEVICE constexpr auto -tiled_divide(Tensor && tensor, +tiled_divide(Tensor && tensor, Tiler const& tiler) // Layout or Tile or Shape { - return make_tensor(std::forward(tensor).data(), + return make_tensor(static_cast(tensor).data(), tiled_divide(tensor.layout(), tiler)); } @@ -826,7 +825,7 @@ auto flat_divide(Tensor && tensor, Tiler const& tiler) // Layout or Tile or Shape { - return make_tensor(std::forward(tensor).data(), + return make_tensor(static_cast(tensor).data(), flat_divide(tensor.layout(), tiler)); } @@ -850,7 +849,7 @@ inner_partition(Tensor && tensor, Tiler const& tiler, Coord const& coord) { - auto tensor_tiled = zipped_divide(std::forward(tensor), tiler); + auto tensor_tiled = zipped_divide(static_cast(tensor), tiler); constexpr int R0 = decltype(rank<0>(tensor_tiled))::value; // The coord slices into the second mode (the "rest" mode), flatten the first @@ -877,7 +876,7 @@ outer_partition(Tensor && tensor, Tiler const& tiler, Coord const& coord) { - auto tensor_tiled = zipped_divide(std::forward(tensor), tiler); + auto tensor_tiled = zipped_divide(static_cast(tensor), tiler); constexpr int R1 = decltype(rank<1>(tensor_tiled))::value; // The coord slices into the first mode (the "tile" mode), flatten the second @@ -903,7 +902,7 @@ local_tile(Tensor && tensor, Tiler const& tiler, // tiler to apply Coord const& coord) // coord to slice into "remainder" { - return inner_partition(std::forward(tensor), + return inner_partition(static_cast(tensor), tiler, coord); } @@ -928,7 +927,7 @@ local_tile(Tensor && tensor, Coord const& coord, // coord to slice into "remainder" Proj const& proj) // projection to apply to tiler and coord { - return local_tile(std::forward(tensor), + return local_tile(static_cast(tensor), dice(proj, tiler), dice(proj, coord)); } @@ -946,7 +945,7 @@ local_partition(Tensor && tensor, Index const& index) // index to slice for { static_assert(is_integral::value); - return outer_partition(std::forward(tensor), + return outer_partition(static_cast(tensor), product_each(shape(tile)), tile.get_flat_coord(index)); } @@ -970,7 +969,7 @@ local_partition(Tensor && tensor, Index const& index, // index to slice for Projection const& proj) { - return local_partition(std::forward(tensor), + return local_partition(static_cast(tensor), dice(proj, tile), index); } @@ -986,9 +985,11 @@ CUTE_HOST_DEVICE void print(Tensor const& tensor) } template -CUTE_HOST_DEVICE void print_tensor(Tensor const& tensor) +CUTE_HOST_DEVICE void print_tensor(Tensor const& tensor, bool print_type = true) { - print(tensor); print(":\n"); + if (print_type) { + print(tensor); print(":\n"); + } if constexpr (Layout::rank == 1) { @@ -1008,18 +1009,18 @@ CUTE_HOST_DEVICE void print_tensor(Tensor const& tensor) } else if constexpr (Layout::rank == 3) { - print_tensor(tensor(_,_,0)); + print_tensor(tensor(_,_,0), false); for (int k = 1; k < size<2>(tensor); ++k) { for (int i = 0; i < 5*size<1>(tensor); ++i) { print("-"); } print("\n"); - print_tensor(tensor(_,_,k)); + print_tensor(tensor(_,_,k), false); } } else if constexpr (Layout::rank == 4) { - print_tensor(tensor(_,_,_,0)); + print_tensor(tensor(_,_,_,0), false); for (int p = 1; p < size<3>(tensor); ++p) { for (int i = 0; i < 5*size<1>(tensor); ++i) { print("="); } print("\n"); - print_tensor(tensor(_,_,_,p)); + print_tensor(tensor(_,_,_,p), false); } } } @@ -1090,5 +1091,9 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, Tensor const #include #include #include +#include #include #include + +#include +#include diff --git a/include/cute/underscore.hpp b/include/cute/underscore.hpp index 118effe8..212f42d7 100644 --- a/include/cute/underscore.hpp +++ b/include/cute/underscore.hpp @@ -45,6 +45,9 @@ struct Underscore : Int<0> {}; CUTE_INLINE_CONSTANT Underscore _; +// Convenient alias +using X = Underscore; + // Treat Underscore as an integral like integral_constant template <> struct is_integral : true_type {}; diff --git a/include/cute/util/type_traits.hpp b/include/cute/util/type_traits.hpp index 3107b811..a8cab903 100644 --- a/include/cute/util/type_traits.hpp +++ b/include/cute/util/type_traits.hpp @@ -131,22 +131,21 @@ using CUTE_STL_NAMESPACE::remove_pointer_t; // using CUTE_STL_NAMESPACE::declval; -template< class T > +template constexpr T&& forward(remove_reference_t& t) noexcept { return static_cast(t); } -template< class T > +template constexpr T&& forward(remove_reference_t&& t) noexcept { - static_assert(! is_lvalue_reference_v, - "T cannot be an lvalue reference (e.g., U&)."); + static_assert(! is_lvalue_reference_v, "T cannot be an lvalue reference (e.g., U&)."); return static_cast(t); } -template< class T > -constexpr remove_reference_t&& move( T&& t ) noexcept +template +constexpr remove_reference_t&& move(T&& t) noexcept { return static_cast&&>(t); } diff --git a/include/cutlass/arch/arch.h b/include/cutlass/arch/arch.h index 7f13ff3d..36d4676b 100644 --- a/include/cutlass/arch/arch.h +++ b/include/cutlass/arch/arch.h @@ -86,6 +86,9 @@ struct Sm80 { struct Sm86 { static int const kMinComputeCapability = 86; }; +struct Sm89 { + static int const kMinComputeCapability = 89; +}; struct Sm90 { static int const kMinComputeCapability = 90; }; diff --git a/include/cutlass/arch/barrier.h b/include/cutlass/arch/barrier.h index 4aa36f79..402ef50c 100644 --- a/include/cutlass/arch/barrier.h +++ b/include/cutlass/arch/barrier.h @@ -68,6 +68,7 @@ class NamedBarrier { uint32_t const num_threads_; // Range : [0, 15] + // Note that should be set to the final barrier ID, including ReserveNamedBarrierCount should be considered uint32_t const id_; public: @@ -88,12 +89,14 @@ class NamedBarrier { CUTLASS_DEVICE void arrive_and_wait() const { - NamedBarrier::arrive_and_wait(num_threads_, id_); + // Note: The value of id_ is already the final barrier id (set correctly in the constructor). + NamedBarrier::arrive_and_wait_internal(num_threads_, id_); } CUTLASS_DEVICE void arrive() const { - NamedBarrier::arrive(num_threads_, id_); + // Note: The value of id_ is already the final barrier id (set correctly in the constructor). + NamedBarrier::arrive_internal(num_threads_, id_); } CUTLASS_DEVICE @@ -384,8 +387,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier { // Performs an arrive operation + expected transaction bytes increment CUTLASS_DEVICE - void arrive_and_expect_tx(uint32_t transaction_bytes, uint32_t cta_id) const { - ClusterTransactionBarrier::arrive_and_expect_tx(&this->barrier_, transaction_bytes , cta_id, true); + void arrive_and_expect_tx(uint32_t transaction_bytes, uint32_t cta_id, uint32_t pred = 1u) const { + ClusterTransactionBarrier::arrive_and_expect_tx(&this->barrier_, transaction_bytes , cta_id, pred); } // Performs an expected transaction bytes increment without doing an arrive operation diff --git a/include/cutlass/arch/mma.h b/include/cutlass/arch/mma.h index 007fe655..007ba19b 100644 --- a/include/cutlass/arch/mma.h +++ b/include/cutlass/arch/mma.h @@ -86,12 +86,9 @@ 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 that staged accumulation is not to be used. This is valid only for SM89 +/// FP8 kernels. +struct OpMultiplyAddFastAccum; ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -250,5 +247,23 @@ struct SparseMma; #include "cutlass/arch/mma_sm75.h" #include "cutlass/arch/mma_sm80.h" #include "cutlass/arch/mma_sparse_sm80.h" +#include "cutlass/arch/mma_sm89.h" +#include "cutlass/arch/mma_sparse_sm89.h" #include "cutlass/arch/mma_sm90.h" ///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace arch { +namespace detail { +/// 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 || + is_sm89_staged_policy_v; +}; +} // namespace detail +} // namespace arch +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/arch/mma_sm89.h b/include/cutlass/arch/mma_sm89.h new file mode 100644 index 00000000..fe4b7eb7 --- /dev/null +++ b/include/cutlass/arch/mma_sm89.h @@ -0,0 +1,367 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE 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-accumulate specialzied for SM89 +*/ + +#pragma once + +#if defined(__CUDACC_RTC__) +#include +#else +#include +#endif + +#include "cutlass/cutlass.h" +#include "mma.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 4) + +# define CUTLASS_ARCH_MMA_SM89_SUPPORTED 1 +#endif + +#if defined(CUTLASS_ARCH_MMA_SM89_SUPPORTED) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 890) +# define CUTLASS_ARCH_MMA_SM89_ENABLED +#endif + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace arch { + +//////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +// Whether the Mma uses as SM89 staged accumulation policy +template +static constexpr bool is_sm89_staged_policy_v = + ( + // ElementA must be FP8 + platform::is_same::value || + platform::is_same::value + ) && + ( + // ElementB must be FP8 + platform::is_same::value || + platform::is_same::value + ) && + ( + // The instruction shape must be 16x8x32 + Operator::ArchMmaOperator::Shape::kM == 16 && + Operator::ArchMmaOperator::Shape::kN == 8 && + Operator::ArchMmaOperator::Shape::kK == 32 + ) && + ( + // The operator must be OpMultiplyAdd (default) + platform::is_same::value + ); +} // namespace detail + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 16832 - Float {E4M3, E5M2}, FP32 accumulation +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation - F32 = fe4m3 * fe4m3 + F32 +template +struct Mma< + gemm::GemmShape<16, 8, 32>, + 32, + cutlass::float_e4m3_t, + layout::RowMajor, + cutlass::float_e4m3_t, + layout::ColumnMajor, + float, + layout::RowMajor, + Operator_> { + static_assert(platform::is_same::value || + platform::is_same::value, + "Invalid operator for SM89 FP8 instruction"); + + using Shape = gemm::GemmShape<16, 8, 32>; + + using ElementA = cutlass::float_e4m3_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = cutlass::float_e4m3_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = Operator_; + using ArchTag = arch::Sm89; + + CUTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(CUTLASS_ARCH_MMA_SM89_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( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.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 + + CUTLASS_UNUSED(d); + CUTLASS_UNUSED(a); + CUTLASS_UNUSED(b); + CUTLASS_UNUSED(c); + CUTLASS_NOT_IMPLEMENTED(); + +#endif + } +}; + +/// Matrix multiply-add operation - F32 = fe4m3 * fe5m2 + F32 +template +struct Mma< + gemm::GemmShape<16, 8, 32>, + 32, + cutlass::float_e4m3_t, + layout::RowMajor, + cutlass::float_e5m2_t, + layout::ColumnMajor, + float, + layout::RowMajor, + Operator_> { + static_assert(platform::is_same::value || + platform::is_same::value, + "Invalid operator for SM89 FP8 instruction"); + + using Shape = gemm::GemmShape<16, 8, 32>; + + using ElementA = cutlass::float_e4m3_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = cutlass::float_e5m2_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = Operator_; + using ArchTag = arch::Sm89; + + CUTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(CUTLASS_ARCH_MMA_SM89_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( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e5m2.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 + + CUTLASS_UNUSED(d); + CUTLASS_UNUSED(a); + CUTLASS_UNUSED(b); + CUTLASS_UNUSED(c); + CUTLASS_NOT_IMPLEMENTED(); + +#endif + } +}; + +/// Matrix multiply-add operation - F32 = fe5m2 * fe4m3 + F32 +template +struct Mma< + gemm::GemmShape<16, 8, 32>, + 32, + cutlass::float_e5m2_t, + layout::RowMajor, + cutlass::float_e4m3_t, + layout::ColumnMajor, + float, + layout::RowMajor, + Operator_> { + static_assert(platform::is_same::value || + platform::is_same::value, + "Invalid operator for SM89 FP8 instruction"); + + using Shape = gemm::GemmShape<16, 8, 32>; + + using ElementA = cutlass::float_e5m2_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = cutlass::float_e4m3_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = Operator_; + using ArchTag = arch::Sm89; + + CUTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(CUTLASS_ARCH_MMA_SM89_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( + "mma.sync.aligned.m16n8k32.row.col.f32.e5m2.e4m3.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 + + CUTLASS_UNUSED(d); + CUTLASS_UNUSED(a); + CUTLASS_UNUSED(b); + CUTLASS_UNUSED(c); + CUTLASS_NOT_IMPLEMENTED(); + +#endif + } +}; + +/// Matrix multiply-add operation - F32 = fe5m2 * fe5m2 + F32 +template +struct Mma< + gemm::GemmShape<16, 8, 32>, + 32, + cutlass::float_e5m2_t, + layout::RowMajor, + cutlass::float_e5m2_t, + layout::ColumnMajor, + float, + layout::RowMajor, + Operator_> { + static_assert(platform::is_same::value || + platform::is_same::value, + "Invalid operator for SM89 FP8 instruction"); + + using Shape = gemm::GemmShape<16, 8, 32>; + + using ElementA = cutlass::float_e5m2_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = cutlass::float_e5m2_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = Operator_; + using ArchTag = arch::Sm89; + + CUTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(CUTLASS_ARCH_MMA_SM89_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( + "mma.sync.aligned.m16n8k32.row.col.f32.e5m2.e5m2.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 + + CUTLASS_UNUSED(d); + CUTLASS_UNUSED(a); + CUTLASS_UNUSED(b); + CUTLASS_UNUSED(c); + CUTLASS_NOT_IMPLEMENTED(); + +#endif + } +}; + +} // namespace arch +} // namespace cutlass diff --git a/include/cutlass/arch/mma_sparse_sm89.h b/include/cutlass/arch/mma_sparse_sm89.h new file mode 100644 index 00000000..c092df76 --- /dev/null +++ b/include/cutlass/arch/mma_sparse_sm89.h @@ -0,0 +1,409 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE 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 SM89 +*/ + +#pragma once + +#if defined(__CUDACC_RTC__) +#include +#else +#include +#endif + +#include "mma.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/numeric_types.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 4) + +# define CUTLASS_ARCH_SPARSE_MMA_SM89_SUPPORTED 1 +#endif + +#if defined(CUTLASS_ARCH_SPARSE_MMA_SM89_SUPPORTED) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 890) +# define CUTLASS_ARCH_SPARSE_MMA_SM89_ENABLED +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = fe4m3 * fe4m3 + F32 +template +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + cutlass::float_e4m3_t, + layout::RowMajor, + cutlass::float_e4m3_t, + layout::ColumnMajor, + float, + layout::RowMajor, + Operator_, + SPFormatType::Thread> { + + static_assert(platform::is_same::value || + platform::is_same::value, + "Invalid operator for SM89 FP8 instruction"); + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = cutlass::float_e4m3_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = cutlass::float_e4m3_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = Operator_; + using ArchTag = arch::Sm89; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + CUTLASS_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(CUTLASS_ARCH_SPARSE_MMA_SM89_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.m16n8k64.row.col.f32.e4m3.e4m3.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 { + assert(0); + } +#else + CUTLASS_UNUSED(a); + CUTLASS_UNUSED(b); + CUTLASS_UNUSED(c); + CUTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = fe4m3 * fe5m2 + F32 +template +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + cutlass::float_e4m3_t, + layout::RowMajor, + cutlass::float_e5m2_t, + layout::ColumnMajor, + float, + layout::RowMajor, + Operator_, + SPFormatType::Thread> { + + static_assert(platform::is_same::value || + platform::is_same::value, + "Invalid operator for SM89 FP8 instruction"); + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = cutlass::float_e4m3_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = cutlass::float_e5m2_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = Operator_; + using ArchTag = arch::Sm89; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + CUTLASS_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(CUTLASS_ARCH_SPARSE_MMA_SM89_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.m16n8k64.row.col.f32.e4m3.e5m2.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 { + assert(0); + } +#else + CUTLASS_UNUSED(a); + CUTLASS_UNUSED(b); + CUTLASS_UNUSED(c); + CUTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = fe5m2 * fe4m3 + F32 +template +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + cutlass::float_e5m2_t, + layout::RowMajor, + cutlass::float_e4m3_t, + layout::ColumnMajor, + float, + layout::RowMajor, + Operator_, + SPFormatType::Thread> { + + static_assert(platform::is_same::value || + platform::is_same::value, + "Invalid operator for SM89 FP8 instruction"); + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = cutlass::float_e5m2_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = cutlass::float_e4m3_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = Operator_; + using ArchTag = arch::Sm89; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + CUTLASS_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(CUTLASS_ARCH_SPARSE_MMA_SM89_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.m16n8k64.row.col.f32.e5m2.e4m3.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 { + assert(0); + } +#else + CUTLASS_UNUSED(a); + CUTLASS_UNUSED(b); + CUTLASS_UNUSED(c); + CUTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = fe5m2 * fe5m2 + F32 +template +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + cutlass::float_e5m2_t, + layout::RowMajor, + cutlass::float_e5m2_t, + layout::ColumnMajor, + float, + layout::RowMajor, + Operator_, + SPFormatType::Thread> { + + static_assert(platform::is_same::value || + platform::is_same::value, + "Invalid operator for SM89 FP8 instruction"); + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = cutlass::float_e5m2_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = cutlass::float_e5m2_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = Operator_; + using ArchTag = arch::Sm89; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + CUTLASS_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(CUTLASS_ARCH_SPARSE_MMA_SM89_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.m16n8k64.row.col.f32.e5m2.e5m2.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 { + assert(0); + } +#else + CUTLASS_UNUSED(a); + CUTLASS_UNUSED(b); + CUTLASS_UNUSED(c); + CUTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/array.h b/include/cutlass/array.h index 1c2ef4b2..48ad7cca 100644 --- a/include/cutlass/array.h +++ b/include/cutlass/array.h @@ -46,12 +46,7 @@ #pragma once #include "cutlass/cutlass.h" #include "cutlass/functional.h" -#include "cutlass/numeric_size.h" -#include "cutlass/half.h" -#include "cutlass/integer_subbyte.h" -#include "cutlass/tfloat32.h" -#include "cutlass/bfloat16.h" -#include "cutlass/half.h" +#include "cutlass/numeric_types.h" namespace cutlass { //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -69,8 +64,7 @@ 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); + static constexpr int value = sizeof(Array) * 8; }; //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -457,7 +451,7 @@ public: CUTLASS_HOST_DEVICE void fill(T const &value) { CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < kElements; ++i) { + for (int i = 0; i < int(kElements); ++i) { storage[i] = static_cast(value); } } @@ -1092,6 +1086,24 @@ struct multiply_add, Array, Array> { } }; +/// Fused square-and-plus +template +struct square_and_plus> { + + CUTLASS_HOST_DEVICE + Array operator()(Array const &lhs, Array const &rhs) const { + multiply_add, Array, Array> ma_op; + return ma_op(rhs, rhs, lhs); + } + + CUTLASS_HOST_DEVICE + Array operator()(Array const &lhs, T const &rhs) const { + plus> plus_op; + multiplies multiplies_op; + return plus_op(multiplies_op(rhs, rhs), lhs); + } +}; + /// Fused multiply-add-relu0 template struct multiply_add_relu0, Array, Array> { @@ -2609,7 +2621,7 @@ template < /// Number of elements in the array int N, /// Alignment requirement in bytes - int Alignment = sizeof_bits::value * N / 8 + int Alignment = ( sizeof_bits::value * N + 7 ) / 8 > class alignas(Alignment) AlignedArray: public Array { public: diff --git a/include/cutlass/array_subbyte.h b/include/cutlass/array_subbyte.h index b2b44649..25bbe355 100644 --- a/include/cutlass/array_subbyte.h +++ b/include/cutlass/array_subbyte.h @@ -402,7 +402,7 @@ public: private: /// Internal storage - Storage storage[kStorageElements]; + Storage storage[kStorageElements] = {Storage{0}}; public: diff --git a/include/cutlass/barrier.h b/include/cutlass/barrier.h index 3a9dfae8..94f300ad 100644 --- a/include/cutlass/barrier.h +++ b/include/cutlass/barrier.h @@ -276,10 +276,8 @@ struct NamedBarrierManager { private: CUTLASS_DEVICE static void - check_barrier_in_range(uint32_t idx) { - if (idx >= MaxNumNamedBarriers) { - CUTE_RUNTIME_ASSERT("Index exceeds barrier count"); - } + check_barrier_in_range([[maybe_unused]] uint32_t idx) { + assert((idx >= MaxNumNamedBarriers) && "Index exceeds barrier count"); } template diff --git a/include/cutlass/cluster_launch.hpp b/include/cutlass/cluster_launch.hpp index 6189bb39..0996fff8 100644 --- a/include/cutlass/cluster_launch.hpp +++ b/include/cutlass/cluster_launch.hpp @@ -38,7 +38,6 @@ #include #include "cutlass/cutlass.h" #include "cutlass/trace.h" - #if defined(__CUDACC_RTC__) #include #else @@ -98,6 +97,23 @@ struct ClusterLauncher { #endif { #if defined(CUTLASS_SM90_CLUSTER_LAUNCH_ENABLED) +#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1) + if (kernel_function == nullptr) { + CUTLASS_TRACE_HOST("kernel_function is null"); + return Status::kInvalid; + } + CUTLASS_TRACE_HOST("Checking previous error state before calling cudaFuncSetAttribute"); + cudaError_t prevStatus = cudaGetLastError(); + if (prevStatus != cudaSuccess) { + fprintf(stderr, + "[ ERROR: CUDA Runtime ] %s:%d: %s\n", + __FILE__, + __LINE__, + cudaGetErrorString(prevStatus)); + return Status::kInvalid; + } + CUTLASS_TRACE_HOST("Calling cudaFuncSetAttribute"); +#endif // This attribute was added in CUDA 11.8. cudaError_t status = cudaFuncSetAttribute( @@ -157,6 +173,7 @@ struct ClusterLauncher { return Status::kInvalid; #endif } + }; namespace detail { @@ -226,8 +243,11 @@ launch_kernel_on_cluster(const ClusterLaunchParams& params, // the parameters as an array of raw pointers. if constexpr (sizeof...(Args) == 0) { return cutlass::ClusterLauncher::launch( - params.grid_dims, params.cluster_dims, params.block_dims, - params.smem_size_in_bytes, params.cuda_stream, + params.grid_dims, + params.cluster_dims, + params.block_dims, + params.smem_size_in_bytes, + params.cuda_stream, kernel_ptr, nullptr); } else { @@ -235,9 +255,13 @@ launch_kernel_on_cluster(const ClusterLaunchParams& params, detail::checked_addressof(std::forward(args))... }; return cutlass::ClusterLauncher::launch( - params.grid_dims, params.cluster_dims, params.block_dims, - params.smem_size_in_bytes, params.cuda_stream, - kernel_ptr, kernel_params); + params.grid_dims, + params.cluster_dims, + params.block_dims, + params.smem_size_in_bytes, + params.cuda_stream, + kernel_ptr, + kernel_params); } } diff --git a/include/cutlass/complex.h b/include/cutlass/complex.h index 78fa9a30..32cfa5f7 100644 --- a/include/cutlass/complex.h +++ b/include/cutlass/complex.h @@ -64,9 +64,6 @@ namespace cutlass { - - - ///////////////////////////////////////////////////////////////////////////////////////////////// /// Enumeraed type describing a transformation on a complex value. enum class ComplexTransform { diff --git a/include/cutlass/conv/collective/builders/sm90_common.inl b/include/cutlass/conv/collective/builders/sm90_common.inl new file mode 100644 index 00000000..526db83e --- /dev/null +++ b/include/cutlass/conv/collective/builders/sm90_common.inl @@ -0,0 +1,96 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 "cutlass/layout/tensor.h" +#include "cutlass/arch/mma.h" +#include "cutlass/conv/convolution.h" +#include "cutlass/conv/dispatch_policy.hpp" +#include "cutlass/detail/layout.hpp" +#include "cutlass/gemm/collective/builders/sm90_common.inl" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::conv::collective::detail { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Maps a rank-1 cute::Shape<> representing the cluster shape on to the IM2COL TMA atom that should be used with it +template +constexpr auto +sm90_cluster_shape_to_im2col_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_IM2COL{}; + } + else { + return cute::SM90_TMA_LOAD_IM2COL_MULTICAST{}; + } +} + +// Collective tile traits struct that serves as a type list containing a tensor's mem layouts and atoms for the +template< + class GmemTiledCopy_, + class SmemLayout_, + class SmemCopyAtom_ = void +> +struct Sm90ImplicitGemmTileTraits { + using GmemTiledCopy = GmemTiledCopy_; + using SmemLayout = SmemLayout_; + using SmemCopyAtom = SmemCopyAtom_; +}; + +// Accepts a cutlass::layout::Tensor tag and computes the corresponding spatial dimension count +template +constexpr int +gmem_layout_tags_to_spatial_dims() { + static_assert(cute::is_same_v); + if constexpr (cute::is_same_v) { + return 1; + } + else if constexpr (cute::is_same_v) { + return 2; + } + else if constexpr (cute::is_same_v) { + return 3; + } + else { + static_assert(cutlass::detail::dependent_false); + } +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::conv::collective::detail + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/collective/builders/sm90_gmma_builder.inl b/include/cutlass/conv/collective/builders/sm90_gmma_builder.inl new file mode 100644 index 00000000..a08209ef --- /dev/null +++ b/include/cutlass/conv/collective/builders/sm90_gmma_builder.inl @@ -0,0 +1,257 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 "cutlass/conv/collective/builders/sm90_common.inl" + +// SM90 Collective Builders should be used only starting CUDA 12.0 +#if (__CUDACC_VER_MAJOR__ >= 12) +#define CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::conv::collective { +using namespace cute; + +namespace detail { + +// 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::Int 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) { + constexpr auto mainloop_pipeline_bytes = sizeof(typename cutlass::PipelineTmaAsync<1>::SharedStorage); + constexpr auto a_bits = cute::sizeof_bits_v; + constexpr auto b_bits = cute::sizeof_bits_v; + constexpr int stage_bytes = + cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + static_cast(mainloop_pipeline_bytes); + + return (CapacityBytes - carveout_bytes) / stage_bytes; +} + +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA_TMA_WS_SS_FPROP +template < + conv::Operator ConvOp, + 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, + ConvOp, + 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> +> { + static_assert(is_static::value); + static_assert(is_static::value); +#ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED + static_assert(cutlass::detail::dependent_false, "Unsupported Toolkit for SM90 Collective Builder\n"); +#endif + static_assert(cutlass::gemm::collective::detail::is_aligned(), + "Should meet TMA alignment requirement\n"); + + // For fp32 types, map to tf32 MMA value type + using ElementAMma = cute::conditional_t, tfloat32_t, ElementA>; + using ElementBMma = cute::conditional_t, tfloat32_t, ElementB>; + + // For fprop, majorA = K, major B = K; + // For wgrad, majorA = MN, major B = MN; + // For dgrad, majorA = K, major B = MN; + static constexpr cute::GMMA::Major GmmaMajorA = + (ConvOp == conv::Operator::kWgrad) ? cute::GMMA::Major::MN : cute::GMMA::Major::K; + static constexpr cute::GMMA::Major GmmaMajorB = + (ConvOp == conv::Operator::kFprop) ? cute::GMMA::Major::K : cute::GMMA::Major::MN; + + using AtomLayoutMNK = cute::conditional_t, + Layout>, Layout>>; + + using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::ss_op_selector< + ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>(), AtomLayoutMNK{})); + + // For wgrad kernel, tensor A uses tma tiled mode and tensor B uses tma im2col mode. + using GmemTiledCopyA = cute::conditional_t(ClusterShape_MNK{}))), + decltype(cutlass::conv::collective::detail::sm90_cluster_shape_to_im2col_tma_atom(cute::shape<1>(ClusterShape_MNK{})))>; + using GmemTiledCopyB = cute::conditional_t(ClusterShape_MNK{}))), + decltype(cutlass::gemm::collective::detail::sm90_cluster_shape_to_tma_atom(cute::shape<0>(ClusterShape_MNK{})))>; + + using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GmmaMajorA, ElementAMma, decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); + using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GmmaMajorB, ElementBMma, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); + + static constexpr int PipelineStages = detail::compute_stage_count_or_override(StageCountType{}); + + using SmemLayoutA = decltype(tile_to_shape( + SmemLayoutAtomA{}, + make_shape(shape<0>(TileShape_MNK{}), shape<2>(TileShape_MNK{}), Int{}), + Step<_2,_1,_3>{})); + using SmemLayoutB = decltype(tile_to_shape( + SmemLayoutAtomB{}, + make_shape(shape<1>(TileShape_MNK{}), shape<2>(TileShape_MNK{}), Int{}), + Step<_2,_1,_3>{})); + + constexpr static int NumSpatialDimensions = cutlass::conv::collective::detail::gmem_layout_tags_to_spatial_dims(); + + using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedImplicitGemm< + ConvOp, PipelineStages, NumSpatialDimensions, ClusterShape_MNK, KernelScheduleType>; + + using CollectiveOp = CollectiveConv< + DispatchPolicy, + TileShape_MNK, + ElementA, + ElementB, + TiledMma, + detail::Sm90ImplicitGemmTileTraits, + detail::Sm90ImplicitGemmTileTraits + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA auto kernel schedule +template < + conv::Operator ConvOp, + 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, + ConvOp, + 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 CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED + static_assert(cutlass::detail::dependent_false, "Unsupported Toolkit for SM90 Collective Builder\n"); +#endif + +/* +#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>{}, + KernelImplicitTmaWarpSpecializedSm90PingPong, KernelImplicitTmaWarpSpecializedSm90Cooperative>; +#else + using KernelWarpSpecializedSchedule = KernelImplicitTmaWarpSpecializedSm90; +#endif +*/ + using KernelWarpSpecializedSchedule = KernelImplicitTmaWarpSpecializedSm90; + + using CollectiveOp = typename CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + ConvOp, + ElementA, + GmemLayoutA, + AlignmentA, + ElementB, + GmemLayoutB, + AlignmentB, + ElementAccumulator, + TileShape_MNK, + ClusterShape_MNK, + StageCountType, + KernelWarpSpecializedSchedule + >::CollectiveOp; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::conv::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/collective/collective_builder.hpp b/include/cutlass/conv/collective/collective_builder.hpp new file mode 100644 index 00000000..9d6a16c0 --- /dev/null +++ b/include/cutlass/conv/collective/collective_builder.hpp @@ -0,0 +1,93 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 "cutlass/detail/dependent_false.hpp" +#include "cutlass/conv/collective/collective_conv.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::conv::collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Used to specify stage counts or dispatch to automatic computation of stage count +template +struct StageCount { + static constexpr int value = num_stages; + + StageCount() = default; + explicit StageCount(cute::Int) {} +}; + +template +struct StageCountAutoCarveout { + static constexpr int bytes = carveout_bytes; + + StageCountAutoCarveout() = default; + explicit StageCountAutoCarveout(cute::Int) {} +}; + +// Used to automatically let the builder pick the kernel schedule. +// Can be overridden with kernel schedule tags in cutlass/conv/dispatch_policy.hpp +struct KernelScheduleAuto {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + class ArchTag, + class OpClass, + conv::Operator, + 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(cutlass::detail::dependent_false, "Could not build a collective for given parameters."); +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::conv::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "builders/sm90_gmma_builder.inl" +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cute/tile.hpp b/include/cutlass/conv/collective/collective_conv.hpp similarity index 64% rename from include/cute/tile.hpp rename to include/cutlass/conv/collective/collective_conv.hpp index 59f33cd3..d187b5ec 100644 --- a/include/cute/tile.hpp +++ b/include/cutlass/conv/collective/collective_conv.hpp @@ -30,29 +30,33 @@ **************************************************************************************************/ #pragma once -#include +#include "cutlass/detail/dependent_false.hpp" +#include "cutlass/conv/collective/detail.hpp" -#include +///////////////////////////////////////////////////////////////////////////////////////////////// -namespace cute -{ +namespace cutlass::conv::collective { -// -// A Tile is not a Layout, it's a tuple of Layouts or Tiles or Underscores -// +///////////////////////////////////////////////////////////////////////////////////////////////// -template -using Tile = tuple; +template < + class DispatchPolicy, + class TileShape, + class ElementA, + class ElementB, + class TiledMma, + class TileTraitsA, + class TileTraitsB +> +struct CollectiveConv { + static_assert(cutlass::detail::dependent_false, "Could not find a mainloop specialization."); +}; -template -using is_tile = is_tuple; +///////////////////////////////////////////////////////////////////////////////////////////////// -template -CUTE_HOST_DEVICE constexpr -auto -make_tile(Layouts const&... layouts) -{ - return Tile(layouts...); -} +} // namespace cutlass::conv::collective -} // end namespace cute +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "sm90_implicit_gemm_gmma_ss_warpspecialized.hpp" +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/collective/detail.hpp b/include/cutlass/conv/collective/detail.hpp new file mode 100644 index 00000000..0f192209 --- /dev/null +++ b/include/cutlass/conv/collective/detail.hpp @@ -0,0 +1,251 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 "cutlass/conv/convnd_problem_shape.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::conv::collective::detail { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Construct the stride types for conv collectives based on the dispatch policy, strides 64b by default +template +constexpr auto +sm90_dispatch_policy_to_stride_A() { + if constexpr (DispatchPolicy::ConvOp == conv::Operator::kFprop) { + // Maps to modes ((w,n), C) + if constexpr (DispatchPolicy::NumSpatialDimensions == 1) { + return cute::Stride, + cute::Int<1>>{}; + } + // Maps to modes ((w,h,n), C) + else if constexpr (DispatchPolicy::NumSpatialDimensions == 2) { + return cute::Stride, + cute::Int<1>>{}; + } + // Maps to modes ((w,h,d,n), C) + else if constexpr (DispatchPolicy::NumSpatialDimensions == 3) { + return cute::Stride, + cute::Int<1>>{}; + } + // error dims assert + else { + static_assert(cutlass::detail::dependent_false, "Unsupported spatial dim count."); + } + } + else if constexpr (DispatchPolicy::ConvOp == conv::Operator::kWgrad) { + // Maps to modes (k, nq/npq/nzpq) + if constexpr (DispatchPolicy::NumSpatialDimensions == 1 || + DispatchPolicy::NumSpatialDimensions == 2 || + DispatchPolicy::NumSpatialDimensions == 3) { + return cute::Stride, int64_t>{}; + } + // error dims assert + else { + static_assert(cutlass::detail::dependent_false, "Unsupported spatial dim count."); + } + } + else if constexpr (DispatchPolicy::ConvOp == conv::Operator::kDgrad) { + // Maps to modes ((q,n), K) + if constexpr (DispatchPolicy::NumSpatialDimensions == 1) { + return cute::Stride, + cute::Int<1>>{}; + } + // Maps to modes ((q,p,n), K) + else if constexpr (DispatchPolicy::NumSpatialDimensions == 2) { + return cute::Stride, + cute::Int<1>>{}; + } + // Maps to modes ((q,p,z,n), K) + else if constexpr (DispatchPolicy::NumSpatialDimensions == 3) { + return cute::Stride, + cute::Int<1>>{}; + } + // error dims assert + else { + static_assert(cutlass::detail::dependent_false, "Unsupported spatial dim count."); + } + } + else { + static_assert(cutlass::detail::dependent_false, "Unsupported ConvOp."); + } +} + +// Construct the stirde types for conv collectives based on the dispatch policy, strides 64b by default +template +constexpr auto +sm90_dispatch_policy_to_stride_B() { + if constexpr (DispatchPolicy::ConvOp == conv::Operator::kFprop) { + // Maps to modes (k, (C,s)) + if constexpr (DispatchPolicy::NumSpatialDimensions == 1) { + return cute::Stride, int64_t>>{}; + } + // Maps to modes (k, (C,s,r)) + else if constexpr (DispatchPolicy::NumSpatialDimensions == 2) { + return cute::Stride, int64_t, int64_t>>{}; + } + // Maps to modes (k, (C,s,r,t)) + else if constexpr (DispatchPolicy::NumSpatialDimensions == 3) { + return cute::Stride, int64_t, int64_t, int64_t>>{}; + } + // error dims assert + else { + static_assert(cutlass::detail::dependent_false, "Unsupported spatial dim count."); + } + } + else if constexpr (DispatchPolicy::ConvOp == conv::Operator::kWgrad) { + // Maps to modes (C, (w,n)) + if constexpr (DispatchPolicy::NumSpatialDimensions == 1) { + return cute::Stride, + cute::Stride>{}; + } + // Maps to modes (C, (w,h,n)) + else if constexpr (DispatchPolicy::NumSpatialDimensions == 2) { + return cute::Stride, + cute::Stride>{}; + } + // Maps to modes (C, (w,h,d,n)) + else if constexpr (DispatchPolicy::NumSpatialDimensions == 3) { + return cute::Stride, + cute::Stride>{}; + } + // error dims assert + else { + static_assert(cutlass::detail::dependent_false, "Unsupported spatial dim count."); + } + } + else if constexpr (DispatchPolicy::ConvOp == conv::Operator::kDgrad) { + // Maps to modes (C, (k,s)) + if constexpr (DispatchPolicy::NumSpatialDimensions == 1) { + return cute::Stride, cute::Stride>{}; + } + // Maps to modes (C, (k,s,r)) + else if constexpr (DispatchPolicy::NumSpatialDimensions == 2) { + return cute::Stride, cute::Stride>{}; + } + // Maps to modes (C, (k,s,r,t)) + else if constexpr (DispatchPolicy::NumSpatialDimensions == 3) { + return cute::Stride, cute::Stride>{}; + } + // error dims assert + else { + static_assert(cutlass::detail::dependent_false, "Unsupported spatial dim count."); + } + } + else { + static_assert(cutlass::detail::dependent_false, "Unsupported ConvOp."); + } +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Compute the lower/near corner, returning it as a cute::array in [W,H,D] order +template +CUTLASS_HOST_DEVICE +constexpr auto +compute_lower_corner_whd(ConvProblemShape const& problem_shape) { + using cute::for_each; + using cute::make_seq; + + cute::array lower{}; + if constexpr (ConvOp == conv::Operator::kFprop || + ConvOp == conv::Operator::kWgrad) { + for_each(make_seq{}, [&](auto i) { + lower[NumSpatialDimensions-1-i] = -1 * problem_shape.lower_padding[i]; + }); + } + else if constexpr (ConvOp == conv::Operator::kDgrad) { + for_each(make_seq{}, [&](auto i) { + lower[NumSpatialDimensions-1-i] = problem_shape.lower_padding[i] - + (problem_shape.shape_B[i+1] - 1) * problem_shape.dilation[i]; + }); + } + return lower; +} + +// Computes the upper/far corner, returning it as a cute::array in [W,H,D] order +template +CUTLASS_HOST_DEVICE +constexpr auto +compute_upper_corner_whd(ConvProblemShape const& problem_shape) { + using cute::for_each; + using cute::make_seq; + + cute::array upper{}; + if constexpr (ConvOp == conv::Operator::kFprop) { + for_each(make_seq{}, [&](auto i) { + upper[NumSpatialDimensions-1-i] = problem_shape.upper_padding[i] - + (problem_shape.shape_B[i+1] - 1) * problem_shape.dilation[i]; + }); + } + else if constexpr (ConvOp == conv::Operator::kWgrad) { + for_each(make_seq{}, [&](auto i) { + upper[NumSpatialDimensions-1-i] = problem_shape.upper_padding[i] - + (problem_shape.shape_C[i+1] - 1) * problem_shape.dilation[i]; + }); + } + else if constexpr (ConvOp == conv::Operator::kDgrad) { + for_each(make_seq{}, [&](auto i) { + upper[NumSpatialDimensions-1-i] = problem_shape.lower_padding[i] - + (problem_shape.shape_B[i+1] - 1) * problem_shape.dilation[i] + problem_shape.shape_C[i+1] - problem_shape.shape_A[i+1]; + }); + } + return upper; +} + +// Compute the lower/near corner of (t,r,s), returning it as a cute::array in [S,R,T] order +template +CUTLASS_HOST_DEVICE +constexpr auto +compute_lower_srt(ConvProblemShape const& problem_shape) { + using cute::for_each; + using cute::make_seq; + + cute::array lower{}; + if constexpr (ConvOp == conv::Operator::kFprop || + ConvOp == conv::Operator::kWgrad) { + for_each(make_seq{}, [&](auto i) { + lower[NumSpatialDimensions-1-i] = 0; + }); + } + else if constexpr (ConvOp == conv::Operator::kDgrad) { + for_each(make_seq{}, [&](auto i) { + lower[NumSpatialDimensions-1-i] = (problem_shape.shape_B[i+1] - 1) * problem_shape.dilation[i]; + }); + } + return lower; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::conv::collective::detail diff --git a/include/cutlass/conv/collective/sm90_implicit_gemm_gmma_ss_warpspecialized.hpp b/include/cutlass/conv/collective/sm90_implicit_gemm_gmma_ss_warpspecialized.hpp new file mode 100644 index 00000000..124c781b --- /dev/null +++ b/include/cutlass/conv/collective/sm90_implicit_gemm_gmma_ss_warpspecialized.hpp @@ -0,0 +1,616 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 "cutlass/cutlass.h" + +#include "cute/tensor_predicate.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/arch/copy_sm90.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/atom/copy_traits_sm90_im2col.hpp" +#include "cute/numeric/arithmetic_tuple.hpp" +#include "cute/algorithm/functional.hpp" +#include "cute/algorithm/gemm.hpp" + +#include "cutlass/conv/convolution.h" +#include "cutlass/conv/convnd_problem_shape.hpp" +#include "cutlass/conv/dispatch_policy.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/util/packed_stride.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::conv::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + conv::Operator ConvOp, + int Stages, + int NumSpatialDims, + class ClusterShape, + class KernelSchedule, + int PipelineAsyncMmaStages, + class TileShape_, + class ElementA_, + class ElementB_, + class TiledMma_, + class TileTraitsA_, + class TileTraitsB_> +struct CollectiveConv< + MainloopSm90TmaGmmaWarpSpecializedImplicitGemm< + ConvOp, Stages, NumSpatialDims, ClusterShape, KernelSchedule, PipelineAsyncMmaStages>, + TileShape_, + ElementA_, + ElementB_, + TiledMma_, + TileTraitsA_, + TileTraitsB_> +{ + // + // Type Aliases + // + using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedImplicitGemm< + ConvOp, Stages, NumSpatialDims, ClusterShape, KernelSchedule, PipelineAsyncMmaStages>; + using TileShape = TileShape_; + using ElementA = ElementA_; + using ElementB = ElementB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; + using GmemTiledCopyA = typename TileTraitsA_::GmemTiledCopy; + using GmemTiledCopyB = typename TileTraitsB_::GmemTiledCopy; + using SmemLayoutA = typename TileTraitsA_::SmemLayout; + using SmemLayoutB = typename TileTraitsB_::SmemLayout; + using ArchTag = typename DispatchPolicy::ArchTag; + static constexpr int NumSpatialDimensions = DispatchPolicy::NumSpatialDimensions; + static constexpr int NumTensorDimensions = NumSpatialDimensions + 2; + // Deduce the kernel-facing stride tuple types based on the dispatch policy + // (which is a function of the number of spatial dimensions, the algorithm, etc.) + using StrideA = decltype(detail::sm90_dispatch_policy_to_stride_A()); + using StrideB = decltype(detail::sm90_dispatch_policy_to_stride_B()); + + using MainloopPipeline = cutlass::PipelineTmaAsync; + + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename cutlass::PipelineState; + + // TODO: move pipeline mode tiling into the collective setup phase instead + static_assert(rank(SmemLayoutA{}) == 3, "SmemLayout must be rank 3 (M/N, K, PIPE)"); + static_assert((size<0>(TileShape{}) == size<0>(SmemLayoutA{})), "SmemLayout must be compatible with the tile shape."); + static_assert((size<2>(TileShape{}) == size<1>(SmemLayoutA{})), "SmemLayout must be compatible with the tile shape."); + + static_assert(rank(SmemLayoutB{}) == 3, "SmemLayout must be rank 3 (M/N, K, PIPE)"); + static_assert((size<1>(TileShape{}) == size<0>(SmemLayoutB{})), "SmemLayout must be compatible with the tile shape."); + static_assert((size<2>(TileShape{}) == size<1>(SmemLayoutB{})), "SmemLayout must be compatible with the tile shape."); + + 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."); + + // The tma load mode of wgrad is tiled for tensor A and im2col for tensor B while the tma load mode of fprop and dgrad + // kernel is im2col for tensor A and tiled for tensor B. + static_assert((ConvOp == conv::Operator::kWgrad + && (cute::is_same_v || cute::is_same_v)) + || (ConvOp != conv::Operator::kWgrad + && (cute::is_same_v || cute::is_same_v)), + "GmemTiledCopyA - invalid SM90 TMA copy atom specified."); + static_assert((ConvOp == conv::Operator::kWgrad + && (cute::is_same_v || cute::is_same_v)) + || (ConvOp != conv::Operator::kWgrad + && (cute::is_same_v || cute::is_same_v)), + "GmemTiledCopyB - 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; + + static constexpr int K_PIPE_MAX = DispatchPolicy::Stages; + static constexpr int K_PIPE_MMAS = DispatchPolicy::PipelineAsyncMmaStages; + 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))); + + // Host side kernel arguments + struct Arguments { + using ProblemShape = ConvProblemShape; + ProblemShape problem_shape{}; + ElementA const* ptr_A{nullptr}; + ElementB const* ptr_B{nullptr}; + }; + +private: + // Note that for fprop and dgrad kernel, the tma load mode is im2col for tensor A and tiled for + // tensor B while for wgrad kernel, the tma load mode is tiled for tensor A and im2col for tensor + // B since operand A, B is swapped. + + // Get tma_load_a instantce. + template + static constexpr auto + get_tma_load_a_instance(TensorA const& tensor_a, typename Arguments::ProblemShape const& problem_shape) { + if constexpr (ConvOp == conv::Operator::kFprop || ConvOp == conv::Operator::kDgrad) { + // compute the upper and lower corners based on the conv padding + auto lower_corner_whd = detail::compute_lower_corner_whd(problem_shape); + auto upper_corner_whd = detail::compute_upper_corner_whd(problem_shape); + auto lower_srt = detail::compute_lower_srt(problem_shape); + + // The calculation of gbasis strides for dgrad kernel needs perform negate for dilation values. + cute::array stride_srt{}; + for (int i = 0; i < NumSpatialDimensions; ++i) { + stride_srt[i] = ConvOp == conv::Operator::kDgrad ? + -problem_shape.dilation[NumSpatialDimensions-1-i] : + problem_shape.dilation[NumSpatialDimensions-1-i]; + } + + return make_im2col_tma_copy( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,_0{}), + product_each(shape(SmemLayoutA{}(_,_,_0{}))), + size<1>(ClusterShape{}), + shape(lower_corner_whd), + shape(upper_corner_whd), + cute::reverse(shape(problem_shape.lower_padding)), + cute::reverse(shape(problem_shape.upper_padding)), + cute::reverse(shape(problem_shape.traversal_stride)), + shape(lower_srt), + shape(stride_srt)); + } + // TMA tiled mode for tensor A in wgrad kernel. + else if constexpr (ConvOp == conv::Operator::kWgrad) { + return make_tma_copy( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,_0{}), + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})), + size<1>(ClusterShape{})); + } + } + + // Get tma_load_b instantce. + template + static constexpr auto + get_tma_load_b_instance(TensorB const& tensor_b, typename Arguments::ProblemShape const& problem_shape) { + if constexpr (ConvOp == conv::Operator::kFprop || ConvOp == conv::Operator::kDgrad) { + return make_tma_copy( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_0{}), + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), + size<0>(ClusterShape{})); + } + // TMA im2col mode for tensor B in wgrad kernel. + else if constexpr (ConvOp == conv::Operator::kWgrad) { + // compute the upper and lower corners based on the conv padding + auto lower_corner_whd = detail::compute_lower_corner_whd(problem_shape); + auto upper_corner_whd = detail::compute_upper_corner_whd(problem_shape); + auto lower_srt = detail::compute_lower_srt(problem_shape); + + return make_im2col_tma_copy( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_0{}), + product_each(shape(SmemLayoutB{}(_,_,_0{}))), + size<0>(ClusterShape{}), + shape(lower_corner_whd), + shape(upper_corner_whd), + cute::reverse(shape(problem_shape.lower_padding)), + cute::reverse(shape(problem_shape.upper_padding)), + cute::reverse(shape(problem_shape.traversal_stride)), + shape(lower_srt), + cute::reverse(shape(problem_shape.dilation))); + } + } + +public: + + // Device side kernel params + struct Params { + using _Submode = decltype(take<0,NumTensorDimensions-1>(typename Arguments::ProblemShape::TensorExtent{})); + using ProblemShape = cute::conditional_t, + Shape<_Submode, int, _Submode>>; + + // Assumption: StrideA is congruent with Problem_MK + // Select TMA load type according to convolution operator. + using TensorShapeA = cute::conditional_t; + + using TensorShapeB = cute::conditional_t; + + using TMA_A = decltype(get_tma_load_a_instance( + make_tensor( + make_gmem_ptr(static_cast(nullptr)), + make_layout(TensorShapeA{}, StrideA{})), + ConvProblemShape{})); + + using TMA_B = decltype(get_tma_load_b_instance( + make_tensor( + make_gmem_ptr(static_cast(nullptr)), + make_layout(TensorShapeB{}, StrideB{})), + ConvProblemShape{})); + + // Members + TMA_A tma_load_a; + TMA_B tma_load_b; + ProblemShape problem_shape; + }; + + // + // Methods + // + + // Lowers the host side user facing arguments to the kernel facing lauch params + static constexpr Params + to_underlying_arguments(Arguments const& args, void* workspace) { + (void) workspace; + // from the flat problem shape arrays of ConvProblemShape, create a rank-3 MNK problem shape tuple + // tma desc creation depends on the original untransformed domain. + + // A extents. + auto shape_A_orig = args.problem_shape.get_shape_A(); + // B extents. + auto shape_B_orig = args.problem_shape.get_shape_B(); + + // Fill inferred cute strides from flat stride arrays + auto dA = make_cute_packed_stride(StrideA{}, args.problem_shape.stride_A, ConvOp); + auto dB = make_cute_packed_stride(StrideB{}, args.problem_shape.stride_B, ConvOp); + + auto ptr_A = reinterpret_cast(args.ptr_A); + auto ptr_B = reinterpret_cast(args.ptr_B); + + Tensor tensor_a = make_tensor(make_gmem_ptr(ptr_A), make_layout(shape_A_orig, dA)); + Tensor tensor_b = make_tensor(make_gmem_ptr(ptr_B), make_layout(shape_B_orig, dB)); + + auto tma_load_a = get_tma_load_a_instance(tensor_a, args.problem_shape); + auto tma_load_b = get_tma_load_b_instance(tensor_b, args.problem_shape); + + auto problem_shape_mnk = args.problem_shape.get_transformed_problem_shape_MNK(); + + return { + tma_load_a, + tma_load_b, + problem_shape_mnk + }; + } + + template + CUTLASS_HOST_DEVICE static bool + can_implement( + ProblemShape const& problem_shape, + Arguments const& args) { + // Activation and Filter channel mode extents much match + bool implementable = true; + // channel mode is major + implementable &= args.problem_shape.stride_A[NumTensorDimensions-1] == 1; + implementable &= args.problem_shape.stride_B[NumTensorDimensions-1] == 1; + + constexpr int tma_alignment_bits = 128; + // A extents. + auto shape_A_orig = args.problem_shape.get_shape_A(); + // B extents. + auto shape_B_orig = args.problem_shape.get_shape_B(); + constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(shape_A_orig, StrideA{}); + constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(shape_B_orig, StrideB{}); + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + return false; + } + + // Check valid padding values for TMA_LOAD_IM2COL + constexpr int padding_limit = (ProblemShape::RankS == 1) ? 65536 : (ProblemShape::RankS == 2 ? 256 : 16); + for (int i = 0; i < problem_shape.RankS; ++i) { + implementable = implementable && problem_shape.lower_padding[i] <= padding_limit && problem_shape.lower_padding[i] >= 0; + implementable = implementable && problem_shape.upper_padding[i] <= padding_limit && problem_shape.upper_padding[i] >= 0; + } + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Padding values don't meet requirements for TMA LOAD IM2COL.\n"); + return false; + } + + if (problem_shape.groups > 1) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: This kernel does not support conv groups > 1.\n"); + return false; + } + + return true; + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + CUTLASS_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 + > + CUTLASS_DEVICE void + load(MainloopPipeline pipeline, + PipelineState smem_pipe_producer_state, + 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 thad_idx, + TensorStorage& shared_tensors) { + int warp_idx = canonical_warp_idx_sync(); + 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 || + 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 || + 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 + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + // LOCK smem_pipe_producer_state for _writing_ + pipeline.producer_acquire(smem_pipe_producer_state); + + // + // Copy gmem to smem for *k_tile_iter + // + + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_producer_state); + + int write_stage = smem_pipe_producer_state.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_producer_state + ++smem_pipe_producer_state; + } + } + } + + /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster + CUTLASS_DEVICE void + load_tail(MainloopPipeline pipeline, PipelineState smem_pipe_producer_state) { + int warp_idx = canonical_warp_idx_sync(); + 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_producer_state); + } + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Consumer Perspective + template + CUTLASS_DEVICE void + mma(MainloopPipeline pipeline, + PipelineState smem_pipe_consumer_state, + FrgTensorC& accum, + int k_tile_count, + int thread_idx, + TensorStorage& shared_tensors, + Params const& mainloop_params) { + static_assert(is_rmem::value, "C tensor must be rmem resident."); + + 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_consumer_state; + + // Prologue GMMAs + int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count); + + tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + + warpgroup_fence_operand(accum); + CUTLASS_PRAGMA_UNROLL + for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue) { + // WAIT on smem_pipe_consumer_state until its data are available (phase bit flips from rdPhaseBit value) + pipeline.consumer_wait(smem_pipe_consumer_state); + + int read_stage = smem_pipe_consumer_state.index(); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 + CUTLASS_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_consumer_state; + } + + warpgroup_fence_operand(accum); + // Mainloop GMMAs + k_tile_count -= prologue_mma_count; + + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + // WAIT on smem_pipe_consumer_state until its data are available (phase bit flips from rdPhaseBit value) + pipeline.consumer_wait(smem_pipe_consumer_state); + + // + // Compute on k_tile + // + + int read_stage = smem_pipe_consumer_state.index(); + warpgroup_fence_operand(accum); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M) x (V,N) => (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_producer_state 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_consumer_state and smem_pipe_release + ++smem_pipe_consumer_state; + ++smem_pipe_release; + } + + warpgroup_fence_operand(accum); + } + + /// Perform a Consumer Epilogue to release all buffers + CUTLASS_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 cutlass::conv::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/conv2d_problem_size.h b/include/cutlass/conv/conv2d_problem_size.h index 5cfcb6de..ec864219 100644 --- a/include/cutlass/conv/conv2d_problem_size.h +++ b/include/cutlass/conv/conv2d_problem_size.h @@ -35,7 +35,7 @@ activation (NHWC), filter (KRSC), output (NPQK), - pading (pad_h, pad_w), + pading (pad_h, pad_w), stride (stride_h, stride_w), dilation (dilation_h, dilation_w). @@ -109,7 +109,7 @@ public: 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), + 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 @@ -133,9 +133,9 @@ public: Mode mode, int split_k_slices = 1, int groups = 1 - ): + ): N(N), H(H), W(W), C(C), P(P), Q(Q), K(K), R(R), S(S), - pad_h(pad_h), pad_w(pad_w), stride_h(stride_h), stride_w(stride_w), + 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) { } @@ -156,8 +156,8 @@ public: N(input_size.n()), H(input_size.h()), W(input_size.w()), C(input_size.c()), P(output_size.h()), Q(output_size.w()), 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()), + 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) {} @@ -167,7 +167,7 @@ public: Conv2dProblemSize( cutlass::Tensor4DCoord input_size, // NHWC cutlass::Tensor4DCoord filter_size, // KRSC - cutlass::Tensor4DCoord padding, // pad_h, _, pad_w, _ + cutlass::Tensor4DCoord padding, // pad_h, upper_pad_h, pad_w, upper_pad_w cutlass::MatrixCoord stride, // stride_h, stride_w cutlass::MatrixCoord dilation, // dilation_h, dilation_w cutlass::conv::Mode mode = cutlass::conv::Mode::kCrossCorrelation, @@ -177,12 +177,12 @@ public: 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()), + 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; + P = ((H + pad_h + padding[1] - R * dilation_h) / stride_h) + 1; + Q = ((W + pad_w + padding[3] - S * dilation_w) / stride_w) + 1; } /// Constructs convolution problem size from cutlass Tensor4DCoord and MatrixCoord @@ -199,7 +199,7 @@ public: N(input_size.n()), H(input_size.h()), W(input_size.w()), C(input_size.c()), P(output_size.h()), Q(output_size.w()), K(filter_size.n()), R(filter_size.h()), S(filter_size.w()), - pad_h(R / 2), pad_w(S / 2), stride_h(1), stride_w(1), + 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) {} diff --git a/include/cutlass/conv/conv3d_problem_size.h b/include/cutlass/conv/conv3d_problem_size.h index cd7ba5b2..56b16423 100644 --- a/include/cutlass/conv/conv3d_problem_size.h +++ b/include/cutlass/conv/conv3d_problem_size.h @@ -91,7 +91,7 @@ public: Conv3dProblemSize(): Conv2dProblemSize(), D(0), T(0), Z(0), - pad_d(0), + pad_d(0), stride_d(1), dilation_d(1) { } @@ -205,6 +205,34 @@ public: Z = ((D + pad_d * 2 - T * dilation_d) / stride_d) + 1; } + /// Constructs convolution problem size from cutlass Tensor5DCoord, Coord3D + // *computes* output size and sets Z, P and Q (include all data members in ctor) + CUTLASS_HOST_DEVICE + Conv3dProblemSize( + cutlass::Tensor5DCoord input_size, // NDHWC + cutlass::Tensor5DCoord filter_size, // KTRSC + CUTLASS_STL_NAMESPACE::tuple padding, // Coord3D {pad_d, pad_h, pad_w} & Coord3D {far pad_d, pad_h, pad_w} to calculate o/p/q + Coord3D stride, // stride_d, stride_h, stride_w + Coord3D dilation, // dilation_d, dilation_h, dilation_w + cutlass::conv::Mode mode = cutlass::conv::Mode::kCrossCorrelation, + int split_k_slices = 1, + int groups = 1 + ): + 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()}, + {CUTLASS_STL_NAMESPACE::get<0>(padding)[1], CUTLASS_STL_NAMESPACE::get<1>(padding)[1], + CUTLASS_STL_NAMESPACE::get<0>(padding)[2], CUTLASS_STL_NAMESPACE::get<1>(padding)[2]}, + {stride[1], stride[2]}, + {dilation[1], dilation[2]}, + mode, split_k_slices, groups), + D(input_size.d()), T(filter_size.d()), + pad_d(CUTLASS_STL_NAMESPACE::get<0>(padding)[0]), stride_d(stride[0]), dilation_d(dilation[0]) + { + // set output Z + Z = ((D + pad_d + CUTLASS_STL_NAMESPACE::get<1>(padding)[0] - T * dilation_d) / stride_d) + 1; + } + /// Equality operator (ignores mode and split_k_slice) CUTLASS_HOST_DEVICE bool operator==(Conv3dProblemSize const &conv) const { @@ -282,7 +310,7 @@ public: return (N * Z * P * Q * K); } - /// Returns output extent as Tensor5DCoord + /// Returns padding as Coord3D CUTLASS_HOST_DEVICE Coord3D padding() const { diff --git a/include/cutlass/conv/convnd_problem_shape.hpp b/include/cutlass/conv/convnd_problem_shape.hpp new file mode 100644 index 00000000..a32389f6 --- /dev/null +++ b/include/cutlass/conv/convnd_problem_shape.hpp @@ -0,0 +1,574 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE 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 shapes. +*/ +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/tensor_coord.h" +#include "cutlass/conv/convolution.h" + +#include "cute/container/array.hpp" + +#if ! defined(__CUDACC_RTC__) +#include +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::conv { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Implements the user facing argument for all CUTLASS 3.x convolutions in a rank agnostic fashion. +// All tensors are flat and by default treated as layout right (NDHWC, KTRSC, NZPQK) +// Supports asymmetric padding, traversal strides, dilations, and all conv algorithm types. +template < + conv::Operator ConvOp_, + int NumSpatialDimensions +> +struct ConvProblemShape { + // + // Alias types for members + // + static constexpr int RankS = NumSpatialDimensions; + static constexpr int RankT = NumSpatialDimensions + 2; + static constexpr conv::Operator ConvOp = ConvOp_; + using SpatialExtent = cute::array; + using TensorExtent = cute::array; + using TensorStride = cute::array; + using ShapePadding = SpatialExtent; + using TraversalStride = SpatialExtent; + using ShapeDilation = SpatialExtent; + using Corner = SpatialExtent; + + // + // Members + // + cutlass::conv::Mode mode{}; + TensorExtent shape_A{}; + TensorStride stride_A{}; + TensorExtent shape_B{}; + TensorStride stride_B{}; + TensorExtent shape_C{}; + TensorStride stride_C{}; + + // asymmetric padding, both upper and lower padding must be >= 0 + ShapePadding lower_padding{}; + ShapePadding upper_padding{}; + TraversalStride traversal_stride{}; + ShapeDilation dilation{}; + int groups = 1; + + // + // Methods + // + + ConvProblemShape() = default; + + // Constructor accepts user facing arguments and computes to stores the corners as its internal state + ConvProblemShape( + conv::Mode mode, // convolution/cross-correlation + TensorExtent shape_act, // [n,d,h,w,c] + TensorStride stride_act, // [n,d,h,w,c] + TensorExtent shape_flt, // [k,t,r,s,c] + TensorStride stride_flt, // [k,t,r,s,c] + ShapePadding lower_padding, // [pad_d, pad_h, pad_w] + ShapePadding upper_padding, // [pad_d, pad_h, pad_w] + TraversalStride tstride, // [stride_d, stride_h, stride_w] + ShapeDilation dilation, // [dilation_d, dilation_h, dilation_w] + int groups) + : mode(mode) + , lower_padding(lower_padding) + , upper_padding(upper_padding) + , traversal_stride(tstride) + , dilation(dilation) + , groups(groups) { + + auto [shape_xformed_act, stride_xformed_act] = calculate_xformed_act(shape_act, shape_flt); + set_shape_stride_ABC(shape_act, stride_act, shape_flt, stride_flt, shape_xformed_act, stride_xformed_act); + } + + // Allow user input of xformed activation stride to support non-packed strides. + ConvProblemShape( + conv::Mode mode, // convolution/cross-correlation + TensorExtent shape_act, // [n,d,h,w,c] + TensorStride stride_act, // [n,d,h,w,c] + TensorExtent shape_flt, // [k,t,r,s,c] + TensorStride stride_flt, // [k,t,r,s,c] + TensorStride stride_xformed_act, // [n,z,p,q,k] + ShapePadding lower_padding, // [pad_d, pad_h, pad_w] + ShapePadding upper_padding, // [pad_d, pad_h, pad_w] + TraversalStride tstride, // [stride_d, stride_h, stride_w] + ShapeDilation dilation, // [dilation_d, dilation_h, dilation_w] + int groups) + : mode(mode) + , lower_padding(lower_padding) + , upper_padding(upper_padding) + , traversal_stride(tstride) + , dilation(dilation) + , groups(groups) { + + CUTLASS_ASSERT(stride_act[RankT - 1] == 1); + CUTLASS_ASSERT(stride_flt[RankT - 1] == 1); + CUTLASS_ASSERT(stride_xformed_act[RankT - 1] == 1); + + auto stride_act_packed = packed_stride_right_major(shape_act); + auto stride_flt_packed = packed_stride_right_major(shape_flt); + auto [shape_xformed_act, stride_xformed_act_packed] = calculate_xformed_act(shape_act, shape_flt); + + CUTLASS_PRAGMA_UNROLL + for(int i = 0; i < RankT - 1; ++i) { + CUTLASS_ASSERT(stride_act[i] >= stride_act_packed[i]); + CUTLASS_ASSERT(stride_flt[i] >= stride_flt_packed[i]); + CUTLASS_ASSERT(stride_xformed_act[i] >= stride_xformed_act_packed[i]); + } + + set_shape_stride_ABC(shape_act, stride_act, shape_flt, stride_flt, shape_xformed_act, stride_xformed_act); + } + + // Constructor accepts user facing arguments and presume packed tensor strides in canonical (CWHDN) order. + ConvProblemShape( + conv::Mode mode, + TensorExtent shape_act, + TensorExtent shape_flt, + ShapePadding lower_padding, + ShapePadding upper_padding, + TraversalStride tstride, + ShapeDilation dilation, + int groups) + : ConvProblemShape( + mode, + shape_act, + packed_stride_right_major(shape_act), + shape_flt, + packed_stride_right_major(shape_flt), + lower_padding, + upper_padding, + tstride, + dilation, + groups) { + } + +#if ! defined(__CUDACC_RTC__) + // Constructor accepts user facing arguments and computes to stores the corners as its internal state + ConvProblemShape( + conv::Mode mode, + std::initializer_list shape_act_, + std::initializer_list stride_act_, + std::initializer_list shape_flt_, + std::initializer_list stride_flt_, + std::initializer_list lower_padding_, + std::initializer_list upper_padding_, + std::initializer_list traversal_stride_, + std::initializer_list dilation_, + int groups) + : mode(mode) + , groups(groups) { + + TensorExtent shape_act{}; + TensorStride stride_act{}; + TensorExtent shape_flt{}; + TensorStride stride_flt{}; + + assert(shape_act_.size() == shape_act.size()); + assert(stride_act_.size() == stride_act.size()); + assert(shape_flt_.size() == shape_flt.size()); + assert(stride_flt_.size() == stride_flt.size()); + assert(lower_padding_.size() == lower_padding.size()); + assert(upper_padding_.size() == upper_padding.size()); + assert(traversal_stride_.size() == traversal_stride.size()); + assert(dilation_.size() == dilation.size()); + + std::copy(shape_act_.begin(), shape_act_.end(), shape_act.begin()); + std::copy(stride_act_.begin(), stride_act_.end(), stride_act.begin()); + std::copy(shape_flt_.begin(), shape_flt_.end(), shape_flt.begin()); + std::copy(stride_flt_.begin(), stride_flt_.end(), stride_flt.begin()); + std::copy(lower_padding_.begin(), lower_padding_.end(), lower_padding.begin()); + std::copy(upper_padding_.begin(), upper_padding_.end(), upper_padding.begin()); + std::copy(traversal_stride_.begin(), traversal_stride_.end(), traversal_stride.begin()); + std::copy(dilation_.begin(), dilation_.end(), dilation.begin()); + + auto [shape_xformed_act, stride_xformed_act] = calculate_xformed_act(shape_act, shape_flt); + set_shape_stride_ABC(shape_act, stride_act, shape_flt, stride_flt, shape_xformed_act, stride_xformed_act); + } + + // Allow user input of xformed activation stride to support non-packed strides. + ConvProblemShape( + conv::Mode mode, + std::initializer_list shape_act_, + std::initializer_list stride_act_, + std::initializer_list shape_flt_, + std::initializer_list stride_flt_, + std::initializer_list stride_xformed_act_, + std::initializer_list lower_padding_, + std::initializer_list upper_padding_, + std::initializer_list traversal_stride_, + std::initializer_list dilation_, + int groups) + : mode(mode) + , groups(groups) { + TensorExtent shape_act{}; + TensorStride stride_act{}; + TensorExtent shape_flt{}; + TensorStride stride_flt{}; + TensorStride stride_xformed_act{}; + + std::copy(shape_act_.begin(), shape_act_.end(), shape_act.begin()); + std::copy(stride_act_.begin(), stride_act_.end(), stride_act.begin()); + std::copy(shape_flt_.begin(), shape_flt_.end(), shape_flt.begin()); + std::copy(stride_flt_.begin(), stride_flt_.end(), stride_flt.begin()); + std::copy(stride_xformed_act_.begin(), stride_xformed_act_.end(), stride_xformed_act.begin()); + std::copy(lower_padding_.begin(), lower_padding_.end(), lower_padding.begin()); + std::copy(upper_padding_.begin(), upper_padding_.end(), upper_padding.begin()); + std::copy(traversal_stride_.begin(), traversal_stride_.end(), traversal_stride.begin()); + std::copy(dilation_.begin(), dilation_.end(), dilation.begin()); + + CUTLASS_ASSERT(stride_act[RankT - 1] == 1); + CUTLASS_ASSERT(stride_flt[RankT - 1] == 1); + CUTLASS_ASSERT(stride_xformed_act[RankT - 1] == 1); + + auto stride_act_packed = packed_stride_right_major(shape_act); + auto stride_flt_packed = packed_stride_right_major(shape_flt); + auto [shape_xformed_act, stride_xformed_act_packed] = calculate_xformed_act(shape_act, shape_flt); + + CUTLASS_PRAGMA_UNROLL + for(int i = 0; i < RankT - 1; ++i) { + CUTLASS_ASSERT(stride_act[i] >= stride_act_packed[i]); + CUTLASS_ASSERT(stride_flt[i] >= stride_flt_packed[i]); + CUTLASS_ASSERT(stride_xformed_act[i] >= stride_xformed_act_packed[i]); + } + + set_shape_stride_ABC(shape_act, stride_act, shape_flt, stride_flt, shape_xformed_act, stride_xformed_act); + } + + // Constructor accepts user facing arguments and computes to stores the corners as its internal state + ConvProblemShape( + conv::Mode mode, + std::initializer_list shape_act_, + std::initializer_list shape_flt_, + std::initializer_list lower_padding_, + std::initializer_list upper_padding_, + std::initializer_list traversal_stride_, + std::initializer_list dilation_, + int groups) + : mode(mode) + , groups(groups) { + TensorExtent shape_act{}; + TensorStride stride_act{}; + TensorExtent shape_flt{}; + TensorStride stride_flt{}; + + assert(shape_act_.size() == shape_act.size()); + assert(shape_flt_.size() == shape_flt.size()); + assert(lower_padding_.size() == lower_padding.size()); + assert(upper_padding_.size() == upper_padding.size()); + assert(traversal_stride_.size() == traversal_stride.size()); + assert(dilation_.size() == dilation.size()); + + std::copy(shape_act_.begin(), shape_act_.end(), shape_act.begin()); + std::copy(shape_flt_.begin(), shape_flt_.end(), shape_flt.begin()); + std::copy(lower_padding_.begin(), lower_padding_.end(), lower_padding.begin()); + std::copy(upper_padding_.begin(), upper_padding_.end(), upper_padding.begin()); + std::copy(traversal_stride_.begin(), traversal_stride_.end(), traversal_stride.begin()); + std::copy(dilation_.begin(), dilation_.end(), dilation.begin()); + stride_act = packed_stride_right_major(shape_act); + stride_flt = packed_stride_right_major(shape_flt); + + auto [shape_xformed_act, stride_xformed_act] = calculate_xformed_act(shape_act, shape_flt); + set_shape_stride_ABC(shape_act, stride_act, shape_flt, stride_flt, shape_xformed_act, stride_xformed_act); + } +#endif // not defined(__CUDACC_RTC__) + + // Set shape and stride of tensor A/B/C according to following table: + // | | Fprop | Dgrad | Wgrad | + // | ------ | ------ | ------ | ------| + // | ShapeA | NDHWC | NZPQK | NZPQK | + // | ShapeB | KTRSC | KTRSC | NDHWC | + // | ShapeC | NZPQK | NDHWC | KTRSC | + // + CUTLASS_HOST_DEVICE + constexpr void + set_shape_stride_ABC( + TensorExtent shape_act, + TensorStride stride_act, + TensorExtent shape_flt, + TensorStride stride_flt, + TensorExtent shape_xformed_act, + TensorStride stride_xformed_act) { + + if constexpr (ConvOp == cutlass::conv::Operator::kFprop) { + shape_A = shape_act; + stride_A = stride_act; + shape_B = shape_flt; + stride_B = stride_flt; + shape_C = shape_xformed_act; + stride_C = stride_xformed_act; + } + else if constexpr (ConvOp == cutlass::conv::Operator::kDgrad) { + shape_A = shape_xformed_act; + stride_A = stride_xformed_act; + shape_B = shape_flt; + stride_B = stride_flt; + shape_C = shape_act; + stride_C = stride_act; + } + else if constexpr (ConvOp == cutlass::conv::Operator::kWgrad) { + shape_A = shape_xformed_act; + stride_A = stride_xformed_act; + shape_B = shape_act; + stride_B = stride_act; + shape_C = shape_flt; + stride_C = stride_flt; + } + } + + // Get problem shape MNK according to following table: + // | | Fprop | Dgrad | Wgrad | + // | ---- | --------- | -------- | -------- | + // | Shape_M | (Q,P,Z,N) | (W,H,D,N) | (K) | + // | Shape_N | (K) | (C) | (C,S,R,T) | + // | Shape_K | (C,S,R,T) | (K,S,R,T) | (Q,P,Z,N) | + CUTLASS_HOST_DEVICE + constexpr auto + get_transformed_problem_shape_MNK() const { + using cute::insert; + using cute::make_shape; + using cute::reverse; + using cute::take; + + if constexpr (ConvOp == conv::Operator::kWgrad) { + auto M_xformed = shape_C[0]; + auto N_xformed = reverse(take<1, RankT>(shape_C)); + auto K_xformed = reverse(take<0, RankT - 1>(shape_A)); + + return make_shape(M_xformed, N_xformed, K_xformed); + } + else if constexpr (ConvOp == conv::Operator::kFprop){ + auto M_xformed = reverse(take<0, RankT - 1>(shape_C)); + auto N_xformed = shape_C[RankT - 1]; + auto K_xformed = reverse(take<1, RankT>(shape_B)); + + return make_shape(M_xformed, N_xformed, K_xformed); + } + else if constexpr (ConvOp == conv::Operator::kDgrad) { + auto M_xformed = reverse(take<0,RankT - 1>(shape_C)); + auto N_xformed = shape_C[RankT - 1]; + // shape_B: [K,T,R,S,C], K_xformed: [K,S,R,T] + auto K_xformed = insert<0>( + (reverse(take<1,RankT - 1>(shape_B))), + shape_B[0]); + return make_shape(M_xformed, N_xformed, K_xformed); + } + } + + + // Get A extents. + // fprop: A extents array contains [N,D,H,W,C]. Turn that into ((W,H,D,N), (C)) + // wgrad: A extents array contains [N,Z,P,Q,K]. Turn that into ((K), (Q,P,Z,N)) + // dgrad: A extents array contains [N,Z,P,Q,K]. Turn that into ((Q,P,Z,N), (K)) + CUTLASS_HOST_DEVICE + constexpr auto + get_shape_A() const { + using cute::make_shape; + using cute::take; + + if constexpr (ConvOp == conv::Operator::kFprop || + ConvOp == conv::Operator::kDgrad) { + return make_shape( + cute::reverse(take<0, RankT - 1>(shape_A)), + shape_A[RankT - 1]); + } + // For wgrad kernel, we need to linearize NZPQ for tensor A + else if constexpr (ConvOp == conv::Operator::kWgrad) { + return make_shape( + shape_A[RankT - 1], + cute::product(take<0, RankT - 1>(shape_A))); + } + } + + // Get B extents. + // fprop: B extents array contains [K,T,R,S,C]. Turn that into ((K), (C,S,R,T)) + // wgrad: B extents array contains [N,D,H,W,C]. Turn that into ((C), (W,H,D,N)) + // dgrad: B extents array contains [K,T,R,S,C]. Turn that into ((C), (K,S,R,T)) + CUTLASS_HOST_DEVICE + constexpr auto + get_shape_B() const { + using cute::make_shape; + using cute::reverse; + using cute::take; + + if constexpr (ConvOp == conv::Operator::kFprop) { + return make_shape( + shape_B[0], + reverse(take<1, RankT>(shape_B))); + } + else if constexpr (ConvOp == conv::Operator::kWgrad) { + return make_shape( + shape_B[RankT - 1], + reverse(take<0, RankT - 1>(shape_B))); + } + else if constexpr (ConvOp == conv::Operator::kDgrad) { + // shape_B: [K,T,R,S,C], return: [(C),(K,S,R,T)] + return make_shape( + shape_B[RankT - 1], + cute::insert<0>( + reverse(take<1, RankT - 1>(shape_B)), + shape_B[0])); + } + } + + // Static method that returns the canonical strides of tensors (layouts are right major and compact) + CUTLASS_HOST_DEVICE + static constexpr TensorStride + packed_stride_right_major(TensorExtent const& extents) { + TensorStride strides{}; + strides[RankT-1] = 1; + cute::for_each(cute::make_rseq{}, [&](auto i) { + strides[i] = extents[i+1] * strides[i+1]; + }); + return strides; + } + + // Static method that returns the packed logical size of any TensorExtent + CUTLASS_HOST_DEVICE + static constexpr size_t + size(TensorExtent const& extents) { + size_t size = 1; + cute::for_each(cute::make_seq{}, [&](auto i) { + size *= extents[i]; + }); + return size; + } + + CUTLASS_HOST_DEVICE + constexpr size_t + size_A() const { + return shape_A[0] * stride_A[0]; + } + + CUTLASS_HOST_DEVICE + constexpr size_t + size_B() const { + return shape_B[0] * stride_B[0]; + } + + CUTLASS_HOST_DEVICE + constexpr size_t + size_C() const { + return shape_C[0] * stride_C[0]; + } + + // Equality operator + CUTLASS_HOST_DEVICE + bool operator==(ConvProblemShape const& rhs) const { + using cute::for_each; + using cute::make_seq; + + bool is_equal = true; + + // Compare all tensor extents + for_each(make_seq{}, [&](auto i) { + is_equal = is_equal + && (shape_A[i] == rhs.shape_A[i]) + && (shape_B[i] == rhs.shape_B[i]); + }); + + // Compare all spatial extents + for_each(make_seq{}, [&](auto i) { + is_equal = is_equal + && (lower_padding[i] == rhs.lower_padding[i]) + && (upper_padding[i] == rhs.upper_padding[i]) + && (traversal_stride[i] == rhs.traversal_stride[i]) + && (dilation[i] == rhs.dilation[i]); + }); + + return is_equal; + } + + /// Inequality operator + CUTLASS_HOST_DEVICE + bool operator!=(ConvProblemShape const &rhs) const { + return !(*this == rhs); + } + +private: + CUTLASS_HOST_DEVICE + constexpr auto + calculate_xformed_act(TensorExtent shape_act, TensorExtent shape_flt) { + TensorExtent shape_xformed_act{}; + // calculate n,z,p,q,k. + // a helper lambda to compute a single spatial extent of the nzpqk tensor + auto nzpqk_extent = [](int act_ext, int filter_ext, int pad_total, int dilation, int tstride) { + return 1 + (act_ext + pad_total - ((filter_ext -1) * dilation + 1)) / tstride; + }; + + shape_xformed_act[0] = shape_act[0]; // Activation N extent + cute::for_each(cute::make_seq{}, [&](auto i) { + shape_xformed_act[i+1] = nzpqk_extent( + shape_act[i+1], shape_flt[i+1], upper_padding[i] + lower_padding[i], dilation[i], traversal_stride[i]); + }); + shape_xformed_act[RankT-1] = shape_flt[0]; // Filter K extent + + TensorStride stride_xformed_act = packed_stride_right_major(shape_xformed_act); + + return cute::make_tuple(shape_xformed_act, stride_xformed_act); + } +}; + +template< + conv::Operator ConvOp, + int SpatialDim +> +void print(ConvProblemShape const& problem) { + printf("ConvProblemShape with %d spatial dimensions implementing cutlass::conv::Operator::%d\n", + SpatialDim, int(ConvOp)); + printf("\tTensorA: "); + cute::print(problem.shape_A); printf(":"); + cute::print(problem.stride_A); printf("\n"); + printf("\tTensorB: "); + cute::print(problem.shape_B); printf(":"); + cute::print(problem.stride_B); printf("\n"); + printf("\tTensorC: "); + cute::print(problem.shape_C); printf(":"); + cute::print(problem.stride_C); printf("\n"); + printf("\tLower padding: "); print(problem.lower_padding); printf("\n"); + printf("\tUpper padding: "); print(problem.upper_padding); printf("\n"); + printf("\tTraversal strides: "); print(problem.traversal_stride); printf("\n"); + printf("\tDilation: "); print(problem.dilation); printf("\n"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::conv + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/convolution.h b/include/cutlass/conv/convolution.h index dce38441..a61f573e 100644 --- a/include/cutlass/conv/convolution.h +++ b/include/cutlass/conv/convolution.h @@ -171,6 +171,30 @@ struct TensorNHWCShape { } }; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Shape of a conv2d stride, which controls how the filter convolves around the input volume +template < + /// Stride in horizontal direction + int u = 1, + /// Stride in vertical direction + int v = 1 +> +struct Stride2D { + static int const kU = u; + static int const kV = v; + + // + // Static member functions + // + + /// Returns a Coord object + CUTLASS_HOST_DEVICE + static Coord<2> toCoord() { + return make_Coord(kU, kV); + } +}; + //////////////////////////////////////////////////////////////////////////////////////////////////// } // namespace conv diff --git a/include/cutlass/conv/device/conv_universal_adapter.hpp b/include/cutlass/conv/device/conv_universal_adapter.hpp new file mode 100644 index 00000000..69cfbaba --- /dev/null +++ b/include/cutlass/conv/device/conv_universal_adapter.hpp @@ -0,0 +1,414 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 + +// common +#include "cutlass/arch/mma.h" +#include "cutlass/cutlass.h" +#include "cutlass/arch/mma.h" +#include "cutlass/trace.h" +#include "cutlass/cluster_launch.hpp" +#include "cutlass/device_kernel.h" + +#include "cutlass/conv/kernel/conv_universal.hpp" +#include "cutlass/gemm/gemm.h" +#include "cutlass/detail/layout.hpp" +#include "cutlass/cuda_host_adapter.hpp" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::conv::device { + +//////////////////////////////////////////////////////////////////////////////// + +/*! + ConvUniversalAdapter is a stateful, reusable handle built around a kernel + of type cutlass::conv::kernel::ConvUniversal. + + It manages the lifetime of the underlying `kernel::Params` struct, and exposes APIs + to create it from the host facing arguments. For power users, static methods + are exposed that bypass the stateful methods or args->params lowering. +*/ +template +class ConvUniversalAdapter +{ +public: + using ConvKernel = ConvKernel_; + using TileShape = typename ConvKernel::TileShape; + using ElementA = typename ConvKernel::ElementA; + using ElementB = typename ConvKernel::ElementB; + using ElementC = typename ConvKernel::ElementC; + using ElementD = typename ConvKernel::ElementD; + using ElementAccumulator = typename ConvKernel::TiledMma::ValTypeC; + using DispatchPolicy = typename ConvKernel::DispatchPolicy; + using CollectiveMainloop = typename ConvKernel::CollectiveMainloop; + using CollectiveEpilogue = typename ConvKernel::CollectiveEpilogue; + + static bool const kEnableCudaHostAdapter = CUTLASS_ENABLE_CUDA_HOST_ADAPTER; + + // Tease out meta-information about the conv algorithm + static constexpr conv::Operator kConvolutionalOperator = DispatchPolicy::ConvOp; + static constexpr int NumSpatialDimensions = ConvKernel::NumSpatialDimensions; + + // 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 ConvKernel::TiledMma::AtomThrID{}) > 1), + cutlass::arch::OpClassTensorOp, cutlass::arch::OpClassSimt>; + + using ArchTag = typename ConvKernel::ArchTag; + + // Assume TiledMma's ShapeMNK is the same as 2.x's ThreadblockShape + using ThreadblockShape = cutlass::gemm::GemmShape< + cute::size<0>(TileShape{}), + cute::size<1>(TileShape{}), + cute::size<2>(TileShape{})>; + + using ClusterShape = cutlass::gemm::GemmShape< + cute::size<0>(typename ConvKernel::DispatchPolicy::ClusterShape{}), + cute::size<1>(typename ConvKernel::DispatchPolicy::ClusterShape{}), + cute::size<2>(typename ConvKernel::DispatchPolicy::ClusterShape{})>; + + // Instruction shape is easy too, since we get that directly from our TiledMma's atom shape + using InstructionShape = cutlass::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 = ConvKernel::MaxThreadsPerBlock; + + // Warp shape is not a primary API type in 3.x + // But we can best approximate it by inspecting the TiledMma + // 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_STATIC_V(cute::size(typename ConvKernel::TiledMma{})) / 32); + static constexpr int WarpsInMmaM = 4; + static constexpr int WarpsInMmaN = cute::ceil_div(WarpsInMma, WarpsInMmaM); + using WarpCount = cutlass::gemm::GemmShape; + using WarpShape = cutlass::gemm::GemmShape< + CUTE_STATIC_V(cute::tile_size<0>(typename CollectiveMainloop::TiledMma{})) / WarpsInMmaM, + CUTE_STATIC_V(cute::tile_size<1>(typename CollectiveMainloop::TiledMma{})) / WarpsInMmaN, + CUTE_STATIC_V(cute::tile_size<2>(typename CollectiveMainloop::TiledMma{}))>; + + static int constexpr kStages = CollectiveMainloop::DispatchPolicy::Stages; + + // Inspect TiledCopy for A and B to compute the alignment size + static int constexpr kAlignmentA = detail::get_alignment_count_from_gmem_tiled_copy< + typename CollectiveMainloop::GmemTiledCopyA, ElementA>(); + static int constexpr kAlignmentB = detail::get_alignment_count_from_gmem_tiled_copy< + typename CollectiveMainloop::GmemTiledCopyB, ElementB>(); + static int constexpr kAlignmentC = detail::get_alignment_count_from_gmem_tiled_copy< + typename CollectiveEpilogue::GmemTiledCopyC, ElementC>(); + static int constexpr kAlignmentD = detail::get_alignment_count_from_gmem_tiled_copy< + typename CollectiveEpilogue::GmemTiledCopyD, ElementD>(); + + using EpilogueOutputOp = typename CollectiveEpilogue::ThreadEpilogueOp; + + /// Argument structure: User API + using Arguments = typename ConvKernel::Arguments; + /// Argument structure: Kernel API + using Params = typename ConvKernel::Params; + +private: + + /// Kernel API parameters object + Params params_; + +public: + + /// Determines whether the conv can execute the given problem. + static Status + can_implement(Arguments const& args) { + if (ConvKernel::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; + CUTLASS_TRACE_HOST(" workspace_bytes: " << workspace_bytes); + + workspace_bytes += ConvKernel::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 = ConvKernel::to_underlying_arguments(args, workspace); + return ConvKernel::get_grid_shape(tmp_params); + } + + /// Computes the grid shape + static dim3 + get_grid_shape(Params const& params) { + return ConvKernel::get_grid_shape(params); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int /* smem_capacity */ = -1) { + CUTLASS_TRACE_HOST("ConvUniversal::maximum_active_blocks()"); + int max_active_blocks = -1; + int smem_size = ConvKernel::SharedStorageSize; + + // first, account for dynamic smem capacity if needed + cudaError_t result; + if (smem_size >= (48 << 10)) { + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + result = cudaFuncSetAttribute( + device_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST( + " cudaFuncSetAttribute() returned error: " + << cudaGetErrorString(result)); + return -1; + } + } + + // query occupancy after setting smem size + result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, + device_kernel, + ConvKernel::MaxThreadsPerBlock, + smem_size); + + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST( + " cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: " + << cudaGetErrorString(result)); + return -1; + } + + CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); + return max_active_blocks; + } + + /// Initializes conv state from arguments. + Status + initialize( + Arguments const& args, + void* workspace = nullptr, + cudaStream_t stream = nullptr, + CudaHostAdapter *cuda_adapter = nullptr) { + CUTLASS_TRACE_HOST("ConvUniversal::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + size_t workspace_bytes = ConvKernel::get_workspace_size(args); + CUTLASS_TRACE_HOST(" workspace_bytes: " << workspace_bytes); + + if (workspace_bytes) { + if (!workspace) { + CUTLASS_TRACE_HOST(" error: device workspace must not be null"); + return Status::kErrorWorkspaceNull; + } + + CUTLASS_TRACE_HOST(" clearing device workspace"); + cudaError_t result = cudaMemsetAsync(workspace, 0, workspace_bytes, stream); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST(" cudaMemsetAsync() returned error " << cudaGetErrorString(result)); + return Status::kErrorInternal; + } + } + + // Initialize the Params structure + params_ = ConvKernel::to_underlying_arguments(args, workspace); + + // Don't set the function attributes - require the CudaHostAdapter to set it. + if constexpr (kEnableCudaHostAdapter) { + CUTLASS_ASSERT(cuda_adapter); + return Status::kSuccess; + } + else { + // account for dynamic smem capacity if needed + int smem_size = ConvKernel::SharedStorageSize; + if (smem_size >= (48 << 10)) { + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + cudaError_t result = cudaFuncSetAttribute( + device_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(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) { + CUTLASS_TRACE_HOST("ConvUniversal()::update() - workspace: " << workspace); + + size_t workspace_bytes = get_workspace_size(args); + if (workspace_bytes > 0 && nullptr == workspace) { + return Status::kErrorWorkspaceNull; + } + + params_ = ConvKernel::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 ConvKernel::to_underling_arguments() + static Status + run(Params& params, cudaStream_t stream = nullptr, CudaHostAdapter *cuda_adapter = nullptr) { + CUTLASS_TRACE_HOST("ConvUniversal::run()"); + dim3 const block = ConvKernel::get_block_shape(); + dim3 const grid = get_grid_shape(params); + + // configure smem size and carveout + int smem_size = ConvKernel::SharedStorageSize; + + Status launch_result; + // Use extended launch API only for mainloops that use it + if constexpr(ConvKernel::ArchTag::kMinComputeCapability >= 90) { + dim3 cluster(cute::size<0>(typename ConvKernel::DispatchPolicy::ClusterShape{}), + cute::size<1>(typename ConvKernel::DispatchPolicy::ClusterShape{}), + cute::size<2>(typename ConvKernel::DispatchPolicy::ClusterShape{})); + void* kernel_params[] = {¶ms}; + if constexpr (kEnableCudaHostAdapter) { + // + // Use the cuda host adapter + // + CUTLASS_ASSERT(cuda_adapter); + if (cuda_adapter) { + + launch_result = cuda_adapter->launch( + grid, cluster, block, smem_size, stream, kernel_params, 0 + ); + } + else { + return Status::kErrorInternal; + } + } + else { + + CUTLASS_ASSERT(cuda_adapter == nullptr); + void const* kernel = (void const*) device_kernel; + + launch_result = ClusterLauncher::launch( + grid, cluster, block, smem_size, stream, kernel, kernel_params); + + } + } + else { + launch_result = Status::kSuccess; + + if constexpr (kEnableCudaHostAdapter) { + CUTLASS_ASSERT(cuda_adapter); + if (cuda_adapter) { + void* kernel_params[] = {¶ms}; + + launch_result = cuda_adapter->launch( + grid, block, smem_size, stream, kernel_params, 0 + ); + + } + else { + return Status::kErrorInternal; + } + } + else { + CUTLASS_ASSERT(cuda_adapter == nullptr); + device_kernel<<>>(params); + } + } + + cudaError_t result = cudaGetLastError(); + if (cudaSuccess == result && Status::kSuccess == launch_result) { + return Status::kSuccess; + } + else { + CUTLASS_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, + cudaStream_t stream = nullptr, + CudaHostAdapter *cuda_adapter = nullptr + ) { + Status status = initialize(args, workspace, stream, cuda_adapter); + if (Status::kSuccess == status) { + status = run(params_, stream, cuda_adapter); + } + return status; + } + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status + operator()( + Arguments const& args, + void* workspace = nullptr, + cudaStream_t stream = nullptr, + CudaHostAdapter *cuda_adapter = nullptr) { + return run(args, workspace, stream, cuda_adapter); + } + + /// Overload that allows a user to re-launch the same kernel without updating internal params struct. + Status + run(cudaStream_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()(cudaStream_t stream = nullptr) { + return run(params_, stream); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::conv::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/device/implicit_gemm_convolution.h b/include/cutlass/conv/device/implicit_gemm_convolution.h index 9a0cf719..d2319f2c 100644 --- a/include/cutlass/conv/device/implicit_gemm_convolution.h +++ b/include/cutlass/conv/device/implicit_gemm_convolution.h @@ -39,6 +39,7 @@ #include "cutlass/cutlass.h" #include "cutlass/device_kernel.h" #include "cutlass/conv/convolution.h" +#include "cutlass/cuda_host_adapter.hpp" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -80,6 +81,8 @@ public: static cutlass::conv::StrideSupport const kStrideSupport = UnderlyingKernel::kStrideSupport; static cutlass::conv::GroupMode const kGroupMode = UnderlyingKernel::kGroupMode; + static bool const kEnableCudaHostAdapter = CUTLASS_ENABLE_CUDA_HOST_ADAPTER; + static int const kWarpCount = (ThreadblockShape::kM / WarpShape::kM) * (ThreadblockShape::kN / WarpShape::kN) * @@ -230,7 +233,8 @@ public: Status initialize( Arguments const &args, void *workspace = nullptr, - cudaStream_t stream = nullptr) { + cudaStream_t stream = nullptr, + CudaHostAdapter *cuda_adapter = nullptr) { if (args.problem_size.split_k_slices > 1) { @@ -250,16 +254,22 @@ public: args, static_cast(workspace) ); - - int smem_size = int(sizeof(typename UnderlyingKernel::SharedStorage)); - if (smem_size >= (48 << 10)) { - cudaError_t result = cudaFuncSetAttribute(cutlass::Kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - smem_size); - - if (result != cudaSuccess) { - return Status::kErrorInternal; + if constexpr (kEnableCudaHostAdapter) { + CUTLASS_ASSERT(cuda_adapter); + return Status::kSuccess; + } + else { + int smem_size = int(sizeof(typename UnderlyingKernel::SharedStorage)); + + if (smem_size >= (48 << 10)) { + cudaError_t result = cudaFuncSetAttribute(cutlass::Kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != cudaSuccess) { + return Status::kErrorInternal; + } } } @@ -281,7 +291,7 @@ public: } /// Runs the kernel using initialized state. - Status run(cudaStream_t stream = nullptr) { + Status run(cudaStream_t stream = nullptr, CudaHostAdapter *cuda_adapter = nullptr) { ThreadblockSwizzle threadblock_swizzle; @@ -290,29 +300,53 @@ public: dim3 block(32 * kWarpCount, 1, 1); int smem_size = int(sizeof(typename UnderlyingKernel::SharedStorage)); + cutlass::Status launch_result = cutlass::Status::kSuccess ; - cutlass::Kernel<<>>(params_); + if constexpr (kEnableCudaHostAdapter) { + // + // Use the cuda host adapter + // + CUTLASS_ASSERT(cuda_adapter); + if (cuda_adapter) { + + void* kernel_params[] = {¶ms_}; + launch_result = cuda_adapter->launch( + grid, dim3(1,1,1), block, smem_size, stream, kernel_params, 0 + ); + } + else { + launch_result = Status::kErrorInternal; + } + } + else { + cutlass::Kernel<<>>(params_); + } cudaError_t result = cudaGetLastError(); - - return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal; + if (cudaSuccess == result && Status::kSuccess == launch_result) { + return Status::kSuccess; + } + else { + CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result); + return Status::kErrorInternal; + } } /// Runs the kernel using initialized state. - Status operator()(cudaStream_t stream = nullptr) { - return run(stream); + Status operator()(cudaStream_t stream = nullptr, CudaHostAdapter *cuda_adapter = nullptr) { + return run(stream, cuda_adapter); } /// Runs the kernel using initialized state. Status operator()( Arguments const &args, void *workspace = nullptr, - cudaStream_t stream = nullptr) { + cudaStream_t stream = nullptr, CudaHostAdapter *cuda_adapter = nullptr) { - Status status = initialize(args, workspace, stream); + Status status = initialize(args, workspace, stream, cuda_adapter); if (status == Status::kSuccess) { - status = run(stream); + status = run(stream, cuda_adapter); } return status; diff --git a/include/cute/numeric/bfloat.hpp b/include/cutlass/conv/dispatch_policy.hpp similarity index 50% rename from include/cute/numeric/bfloat.hpp rename to include/cutlass/conv/dispatch_policy.hpp index 4dcd328b..6c7876dc 100644 --- a/include/cute/numeric/bfloat.hpp +++ b/include/cutlass/conv/dispatch_policy.hpp @@ -30,24 +30,54 @@ **************************************************************************************************/ #pragma once -#include +#include "cutlass/conv/convolution.h" +#include "cutlass/arch/arch.h" -#include -#include +#include "cute/layout.hpp" +#include "cute/numeric/integral_constant.hpp" -namespace cute { +////////////////////////////////////////////////////////////////////////////// -using cutlass::bfloat16_t; +namespace cutlass::conv { + +////////////////////////////////////////////////////////////////////////////// // -// Display utilities +// Policies for categorical dispatch of mainloop against kernel grid schedules +// +struct KernelImplicitTmaWarpSpecializedSm90 { }; +struct KernelImplicitTmaWarpSpecializedSm90Cooperative { }; +struct KernelImplicitTmaWarpSpecializedSm90Pingpong { }; + +// +// Collective Mainloop Policies // -#if !defined(__CUDACC_RTC__) -CUTE_HOST std::ostream& operator<<(std::ostream& os, bfloat16_t const& v) -{ - return os << float(v); -} -#endif +// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, static schedule between TMA and GMMA +// for fprop +template< + conv::Operator ConvOp_, + int Stages_, + int NumSpatialDimensions_, + class ClusterShape_ = cute::Shape,cute::C<1>,cute::C<1>>, + class KernelSchedule = KernelImplicitTmaWarpSpecializedSm90, + int PipelineAsyncMmaStages_ = 1 +> +struct MainloopSm90TmaGmmaWarpSpecializedImplicitGemm { + static constexpr int Stages = Stages_; + static constexpr int NumSpatialDimensions = NumSpatialDimensions_; + static constexpr Operator ConvOp = ConvOp_; + static constexpr int PipelineAsyncMmaStages = PipelineAsyncMmaStages_; + using ClusterShape = ClusterShape_; + using ArchTag = arch::Sm90; + using Schedule = KernelSchedule; -} // end namespace cute + static_assert(NumSpatialDimensions >= 1); + static_assert(! (cute::is_same_v || + cute::is_same_v), + "Persistent schedules not support for conv yet."); +}; + +////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::conv diff --git a/include/cute/numeric/tfloat.hpp b/include/cutlass/conv/kernel/conv_universal.hpp similarity index 63% rename from include/cute/numeric/tfloat.hpp rename to include/cutlass/conv/kernel/conv_universal.hpp index 0596364c..9d98dc9d 100644 --- a/include/cute/numeric/tfloat.hpp +++ b/include/cutlass/conv/kernel/conv_universal.hpp @@ -30,24 +30,34 @@ **************************************************************************************************/ #pragma once -#include +#include "cutlass/detail/dependent_false.hpp" -#include -#include +//////////////////////////////////////////////////////////////////////////////// -namespace cute { +namespace cutlass::conv::kernel { -using cutlass::tfloat32_t; +//////////////////////////////////////////////////////////////////////////////// -// -// Display utilities -// +/* + * Stateless universal device CONV kernel type that treats CONV as + * a composition of a collective mainloop and a collective epilogue. +**/ +template < + class CollectiveMainloop_, + class CollectiveEpilogue_, + class TileSchedulerTag_ = void, + class Enable = void +> +class ConvUniversal { + static_assert(cutlass::detail::dependent_false, + "Could not find a valid specialization at the kernel layer to dispatch against."); +}; -#if !defined(__CUDACC_RTC__) -CUTE_HOST std::ostream& operator<<(std::ostream& os, tfloat32_t const& v) -{ - return os << float(v); -} -#endif +//////////////////////////////////////////////////////////////////////////////// -} // end namespace cute +} // namespace cutlass::conv::kernel + +//////////////////////////////////////////////////////////////////////////////// + +#include "cutlass/conv/kernel/sm90_implicit_gemm_tma_warpspecialized.hpp" +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/kernel/default_conv2d.h b/include/cutlass/conv/kernel/default_conv2d.h index d331fed7..51310304 100644 --- a/include/cutlass/conv/kernel/default_conv2d.h +++ b/include/cutlass/conv/kernel/default_conv2d.h @@ -106,6 +106,27 @@ struct DefaultConvEpilogue< }; ///////////////////////////////////////////////////////////////////////////////////////////////// +template < + typename ArchTag, + typename Shape, + typename WarpMmaSimt, + typename ElementOutput, + typename ElementTensor, + typename ElementVector, + typename OutputOp, + int ElementsPerAccess +> +struct DefaultConvEpilogueWithBroadcastSimt { + using Epilogue = typename epilogue::threadblock::DefaultEpilogueWithBroadcastSimt< + Shape, + WarpMmaSimt, + ElementOutput, + ElementTensor, + ElementVector, + OutputOp, + ElementsPerAccess + >::Epilogue; +}; template < typename ArchTag, diff --git a/include/cutlass/conv/kernel/default_conv2d_fprop_with_absmax.h b/include/cutlass/conv/kernel/default_conv2d_fprop_with_absmax.h new file mode 100644 index 00000000..b0e0ae65 --- /dev/null +++ b/include/cutlass/conv/kernel/default_conv2d_fprop_with_absmax.h @@ -0,0 +1,127 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE 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 default configuration for convolution with absolute maximum calculation. +*/ + +#pragma once + +#include "cutlass/cutlass.h" + +#include "cutlass/conv/kernel/default_conv2d_fprop.h" +#include "cutlass/conv/kernel/implicit_gemm_convolution_with_absmax.h" + +#include "cutlass/epilogue/threadblock/default_epilogue_with_absmax.h" +#include "cutlass/epilogue/threadblock/epilogue_with_absmax.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +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 / cutlass::sizeof_bits::value, + /// Access granularity of B matrix in units of elements + int AlignmentB = 128 / cutlass::sizeof_bits::value +> +struct DefaultConv2dFpropWithAbsMax { + + 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 cutlass::epilogue::threadblock::DefaultEpilogueWithAbsMax< + typename ImplicitGemmBase::Epilogue::Shape, + typename ImplicitGemmBase::Epilogue::WarpMmaOperator, + ImplicitGemmBase::Epilogue::kPartitionsK, + ElementC, + typename EpilogueOutputOp::ElementAuxOutput, + ElementC, + EpilogueOutputOp, + ImplicitGemmBase::Epilogue::kElementsPerAccess + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::ImplicitGemmConvolutionWithAbsMax< + typename ImplicitGemmBase::Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h b/include/cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h index 714935c2..6e6127d7 100644 --- a/include/cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h +++ b/include/cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h @@ -121,6 +121,96 @@ struct DefaultConv2dFpropWithBroadcast { >; }; +///////////////////////////////////////////////////////////////////////////////////////////////// +// 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 DefaultConv2dFpropWithBroadcast < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + using ImplicitGemmBase = typename DefaultConv2dFprop< + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB + >::Kernel; + + // Define epilogue + using Epilogue = typename cutlass::conv::kernel::detail::DefaultConvEpilogueWithBroadcastSimt< + ArchTag, + typename ImplicitGemmBase::Epilogue::Shape, + typename ImplicitGemmBase::Epilogue::WarpMmaOperator, + ElementC, + typename EpilogueOutputOp::ElementT, + typename EpilogueOutputOp::ElementVector, + EpilogueOutputOp, + ImplicitGemmBase::Epilogue::kElementsPerAccess + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::ImplicitGemmConvolutionWithFusedEpilogue< + typename ImplicitGemmBase::Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + ///////////////////////////////////////////////////////////////////////////////////////////////// } // namespace kernel diff --git a/include/cutlass/conv/kernel/default_conv2d_group_fprop.h b/include/cutlass/conv/kernel/default_conv2d_group_fprop.h index 8dbff9d0..927f70ce 100644 --- a/include/cutlass/conv/kernel/default_conv2d_group_fprop.h +++ b/include/cutlass/conv/kernel/default_conv2d_group_fprop.h @@ -57,7 +57,7 @@ namespace conv { namespace kernel { ///////////////////////////////////////////////////////////////////////////////////////////////// -/// Defines a kernel for Conv2dGroupFpro +/// Defines a kernel for Conv2dGroupFprop template < typename ElementA, typename LayoutA, @@ -135,11 +135,11 @@ struct DefaultConv2dGroupFprop < AlignmentB > { - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); // Define the core components from GEMM @@ -269,11 +269,11 @@ struct DefaultConv2dGroupFprop < AlignmentB > { - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); // Define the core components from GEMM @@ -400,11 +400,11 @@ struct DefaultConv2dGroupFprop < AlignmentB > { - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); // Define the core components from GEMM @@ -530,11 +530,11 @@ struct DefaultConv2dGroupFprop < AlignmentB > { - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); - static_assert(std::is_same::value, + static_assert(platform::is_same::value, "Current group conv only support NHWC layout"); // Define the core components from GEMM diff --git a/include/cutlass/conv/kernel/default_conv3d_dgrad.h b/include/cutlass/conv/kernel/default_conv3d_dgrad.h index 69355154..cb50ba49 100644 --- a/include/cutlass/conv/kernel/default_conv3d_dgrad.h +++ b/include/cutlass/conv/kernel/default_conv3d_dgrad.h @@ -293,6 +293,439 @@ struct DefaultConv3dDgrad < }; +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassSimt convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// + +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 DefaultConv3dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + conv::StrideSupport::kStrided +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + cutlass::conv::threadblock::Conv3dDgradOutputGradientTileAccessIteratorAnalytic< + cutlass::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 = + cutlass::conv::threadblock::Conv3dDgradFilterTileAccessIteratorAnalytic< + cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad, + Conv3dProblemSize + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dDgrad 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 +> +struct DefaultConv3dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport::kUnity +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + cutlass::conv::threadblock::Conv3dDgradOutputGradientTileAccessIteratorOptimized< + cutlass::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 = + cutlass::conv::threadblock::Conv3dDgradFilterTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementB, + ThreadMapB + // 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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +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 DefaultConv3dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + conv::StrideSupport::kStrided +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + // cutlass::conv::threadblock::TileIteratorStridedDgrad< + cutlass::conv::threadblock::Conv3dDgradOutputGradientTileAccessIteratorAnalytic< + cutlass::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 = + // cutlass::conv::threadblock::TileIteratorStridedDgrad< + cutlass::conv::threadblock::Conv3dDgradFilterTileAccessIteratorAnalytic< + cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dDgrad 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 +> +struct DefaultConv3dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport::kUnity +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + // cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv3dDgradOutputGradientTileAccessIteratorOptimized< + cutlass::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 = + // cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv3dDgradFilterTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementB, + ThreadMapB + // 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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad, + Conv3dProblemSize + >; + +}; + ///////////////////////////////////////////////////////////////////////////////////////////////// } // namespace kernel diff --git a/include/cutlass/conv/kernel/default_conv3d_fprop.h b/include/cutlass/conv/kernel/default_conv3d_fprop.h index 68d75e6b..3ea1e11c 100644 --- a/include/cutlass/conv/kernel/default_conv3d_fprop.h +++ b/include/cutlass/conv/kernel/default_conv3d_fprop.h @@ -54,7 +54,7 @@ namespace conv { namespace kernel { ///////////////////////////////////////////////////////////////////////////////////////////////// -/// Defines a kernel for Conv2dFprop +/// Defines a kernel for Conv3dFprop template < typename ElementA, typename LayoutA, @@ -185,7 +185,7 @@ struct DefaultConv3dFprop < ///////////////////////////////////////////////////////////////////////////////////////////////// -/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm and multistage +/// Defines a kernel for Conv3dFprop specialization for Analytic IteratorAlgorithm and multistage // pipeline. template < typename ElementA, @@ -506,7 +506,437 @@ struct DefaultConv3dFprop < }; ///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassSimt convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv3dFprop 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 +> +struct DefaultConv3dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + cutlass::conv::threadblock::Conv3dFpropActivationTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + cutlass::conv::threadblock::Conv3dFpropFilterTileAccessIteratorAnalytic< + cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dFprop 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 +> +struct DefaultConv3dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + cutlass::conv::threadblock::Conv3dFpropActivationTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + cutlass::conv::threadblock::Conv3dFpropFilterTileAccessIteratorOptimized< + cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dFprop 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 +> +struct DefaultConv3dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv3dFpropActivationTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv3dFpropFilterTileAccessIteratorAnalytic< + cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dFprop 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 +> +struct DefaultConv3dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv3dFpropActivationTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv3dFpropFilterTileAccessIteratorOptimized< + cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// } // namespace kernel } // namespace conv } // namespace cutlass diff --git a/include/cutlass/conv/kernel/default_conv3d_fprop_with_broadcast.h b/include/cutlass/conv/kernel/default_conv3d_fprop_with_broadcast.h new file mode 100644 index 00000000..1d70a29e --- /dev/null +++ b/include/cutlass/conv/kernel/default_conv3d_fprop_with_broadcast.h @@ -0,0 +1,218 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE 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 "cutlass/cutlass.h" + +#include "cutlass/conv/kernel/default_conv3d_fprop.h" +#include "cutlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h" + +#include "cutlass/epilogue/threadblock/default_epilogue_with_broadcast.h" +#include "cutlass/epilogue/threadblock/epilogue_with_broadcast.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +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 / cutlass::sizeof_bits::value, + /// Access granularity of B matrix in units of elements + int AlignmentB = 128 / cutlass::sizeof_bits::value +> +struct DefaultConv3dFpropWithBroadcast { + + using ImplicitGemmBase = typename DefaultConv3dFprop< + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm, + StrideSupport + >::Kernel; + + // Define epilogue + using Epilogue = typename cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolutionWithFusedEpilogue< + typename ImplicitGemmBase::Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassSimt convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv3dFprop 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 DefaultConv3dFpropWithBroadcast < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + using ImplicitGemmBase = typename DefaultConv3dFprop< + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport + >::Kernel; + + // Define epilogue + using Epilogue = typename cutlass::conv::kernel::detail::DefaultConvEpilogueWithBroadcastSimt< + ArchTag, + typename ImplicitGemmBase::Epilogue::Shape, + typename ImplicitGemmBase::Epilogue::WarpMmaOperator, + ElementC, + typename EpilogueOutputOp::ElementT, + typename EpilogueOutputOp::ElementVector, + EpilogueOutputOp, + ImplicitGemmBase::Epilogue::kElementsPerAccess + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::ImplicitGemmConvolutionWithFusedEpilogue< + typename ImplicitGemmBase::Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/kernel/default_conv3d_wgrad.h b/include/cutlass/conv/kernel/default_conv3d_wgrad.h index d0c53d0a..4ed5e0c1 100644 --- a/include/cutlass/conv/kernel/default_conv3d_wgrad.h +++ b/include/cutlass/conv/kernel/default_conv3d_wgrad.h @@ -53,7 +53,7 @@ namespace kernel { ///////////////////////////////////////////////////////////////////////////////////////////////// -/// Defines a kernel for Conv2dWgrad +/// Defines a kernel for Conv3dWgrad template < typename ElementA, typename LayoutA, @@ -500,6 +500,433 @@ struct DefaultConv3dWgrad < Conv3dProblemSize >; }; + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassSimt convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv3dWgrad 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 +> +struct DefaultConv3dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + cutlass::conv::threadblock::Conv3dWgradOutputGradientTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + cutlass::conv::threadblock::Conv3dWgradActivationTileAccessIteratorAnalytic< + cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad, + Conv3dProblemSize + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dWgrad 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 +> +struct DefaultConv3dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + cutlass::conv::threadblock::Conv3dWgradOutputGradientTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + cutlass::conv::threadblock::Conv3dWgradActivationTileAccessIteratorOptimized< + cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad, + Conv3dProblemSize + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dWgrad 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 +> +struct DefaultConv3dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv3dWgradOutputGradientTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv3dWgradActivationTileAccessIteratorAnalytic< + cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dWgrad 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 +> +struct DefaultConv3dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::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 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv3dWgradOutputGradientTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv3dWgradActivationTileAccessIteratorOptimized< + cutlass::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 = cutlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad, + Conv3dProblemSize + >; + +}; ///////////////////////////////////////////////////////////////////////////////////////////////// } // namespace kernel diff --git a/include/cutlass/conv/kernel/implicit_gemm_convolution_fusion.h b/include/cutlass/conv/kernel/implicit_gemm_convolution_fusion.h index 57c85da1..74ecae40 100644 --- a/include/cutlass/conv/kernel/implicit_gemm_convolution_fusion.h +++ b/include/cutlass/conv/kernel/implicit_gemm_convolution_fusion.h @@ -202,32 +202,30 @@ struct ImplicitGemmConvolutionFusion { /// Parameters structure struct Params { - ConvProblemSize problem_size; - cutlass::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; + ConvProblemSize problem_size{}; + cutlass::gemm::GemmCoord grid_tiled_shape{}; + gemm::GemmCoord implicit_gemm_problem_size{}; + int swizzle_log_tile{0}; + int gemm_k_iterations{0}; + typename Mma::IteratorA::Params iterator_A{}; + typename Mma::IteratorA::Element const *ptr_A = nullptr; + typename Mma::IteratorB::Params iterator_B{}; + typename Mma::IteratorB::Element const *ptr_B = nullptr; + typename Mma::IteratorScaleBias::Params iterator_scale_bias{}; + typename Mma::IteratorScaleBias::Element const *ptr_scale = nullptr; + typename Mma::IteratorScaleBias::Element const *ptr_bias = nullptr; + typename Epilogue::OutputTileIterator::Params iterator_C {}; + typename Epilogue::OutputTileIterator::Element *ptr_C = nullptr; + typename Epilogue::OutputTileIterator::Params iterator_D {}; + typename Epilogue::OutputTileIterator::Element *ptr_D = nullptr; + typename EpilogueOutputOp::Params output_op {}; + int *semaphore = nullptr; + SplitKMode split_k_mode {}; // // Methods // - - CUTLASS_HOST_DEVICE - Params(): swizzle_log_tile(0), gemm_k_iterations(0) { } + Params() = default; /// CUTLASS_HOST_DEVICE diff --git a/include/cutlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h b/include/cutlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h index 06c6ad4d..bf00f90b 100644 --- a/include/cutlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h +++ b/include/cutlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h @@ -158,21 +158,20 @@ struct ImplicitGemmConvolutionStridedDgrad { // 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; + 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 - CUTLASS_HOST_DEVICE - Arguments() { } + Arguments() = default; CUTLASS_HOST_DEVICE Arguments( @@ -205,30 +204,28 @@ struct ImplicitGemmConvolutionStridedDgrad { /// Parameters structure struct Params { - ConvProblemSize problem_size; - cutlass::gemm::GemmCoord grid_tiled_shape; - int swizzle_log_tile; - 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; + ConvProblemSize problem_size{}; + cutlass::gemm::GemmCoord grid_tiled_shape{}; + int swizzle_log_tile{0}; + FastDivmod stride_h_divmod{}; + FastDivmod stride_w_divmod{}; + int gemm_k_iterations{0}; + typename Mma::IteratorA::Params iterator_A{}; + typename Mma::IteratorA::Element const *ptr_A = nullptr; + typename Mma::IteratorB::Params iterator_B{}; + typename Mma::IteratorB::Element const *ptr_B = nullptr; + typename Epilogue::OutputTileIterator::Params iterator_C{}; + typename Epilogue::OutputTileIterator::Element *ptr_C = nullptr; + typename Epilogue::OutputTileIterator::Params iterator_D{}; + typename Epilogue::OutputTileIterator::Element *ptr_D = nullptr; + typename EpilogueOutputOp::Params output_op {}; + int *semaphore = nullptr; + SplitKMode split_k_mode {}; // // Methods // - - CUTLASS_HOST_DEVICE - Params(): gemm_k_iterations(0) { } + Params() = default; /// CUTLASS_HOST_DEVICE diff --git a/include/cutlass/conv/kernel/implicit_gemm_convolution_with_absmax.h b/include/cutlass/conv/kernel/implicit_gemm_convolution_with_absmax.h new file mode 100644 index 00000000..b05fd2d3 --- /dev/null +++ b/include/cutlass/conv/kernel/implicit_gemm_convolution_with_absmax.h @@ -0,0 +1,494 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Convolution kernel with an epilogue that computes the absolute maximum value of the output + and a pre-activation-function auxiliary output. The auxiliary output is also (optionally) + stored to global memory. +*/ + +#pragma once + +#include "cutlass/cutlass.h" + +#include "cutlass/aligned_buffer.h" +#include "cutlass/array.h" +#include "cutlass/numeric_types.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/semaphore.h" +#include "cutlass/tensor_ref.h" +#include "cutlass/layout/tensor.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/conv/convolution.h" +#include "cutlass/conv/conv2d_problem_size.h" +#include "cutlass/conv/conv3d_problem_size.h" +#include "cutlass/epilogue/threadblock/output_iterator_parameter.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +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 ImplicitGemmConvolutionWithAbsMax { + + 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 = cutlass::TensorRef; + using TensorRefAux = cutlass::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; + TensorRefC ref_Aux; + + typename EpilogueOutputOp::Params output_op; + SplitKMode split_k_mode; + + void * ptr_Vector; + + typename LayoutC::Stride::Index ldr; + + // + // Methods + // + + /// Default ctor + CUTLASS_HOST_DEVICE + Arguments() { } + + CUTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size + ): + problem_size(problem_size) { } + + CUTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size, + TensorRefA const & ref_A, + TensorRefB const & ref_B, + TensorRefC const & ref_C, + TensorRefC const & ref_D, + TensorRefAux const & ref_Aux, + typename EpilogueOutputOp::Params const & output_op, + SplitKMode const & split_k_mode = SplitKMode::kSerial, + void * ptr_Vector = nullptr, + typename LayoutC::Stride::Index ldr = 0 + ): + problem_size(problem_size), + ref_A(ref_A), + ref_B(ref_B), + ref_C(ref_C), + ref_D(ref_D), + ref_Aux(ref_Aux), + output_op(output_op), + split_k_mode(split_k_mode), + ptr_Vector(ptr_Vector), + ldr(ldr) + { + + } + + }; + + /// Parameters structure + struct Params { + ConvProblemSize problem_size; + cutlass::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 Epilogue::AuxOutputTileIterator::Params iterator_Aux; + typename Epilogue::AuxOutputTileIterator::Element *ptr_Aux; + typename EpilogueOutputOp::Params output_op; + int *semaphore; + SplitKMode split_k_mode; + + void * ptr_Vector; + typename LayoutC::Stride::Index ldr; + + // + // Methods + // + + CUTLASS_HOST_DEVICE + Params(): + swizzle_log_tile(0), + gemm_k_iterations(0), + ptr_Vector(nullptr), + ldr(0) + { } + + /// + CUTLASS_HOST_DEVICE + Params( + Arguments const &args, + int *semaphore = nullptr + ): + problem_size(args.problem_size), + implicit_gemm_problem_size(cutlass::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()), + iterator_Aux(ConvOutputIteratorParameter::layout(args.ref_Aux)), + ptr_Aux(args.ref_Aux.data()), + output_op(args.output_op), + semaphore(semaphore), + split_k_mode(args.split_k_mode), + ptr_Vector(args.ptr_Vector), + ldr(args.ldr) + + { + 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 + // + + CUTLASS_HOST_DEVICE + ImplicitGemmConvolutionWithAbsMax() { } + + /// Executes one ImplicitGEMM + CUTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + cutlass::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 = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + int lane_idx = threadIdx.x % 32; + + // + // 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 writing to auxiliary tensor. + typename Epilogue::AuxOutputTileIterator iterator_Aux( + params.iterator_Aux, + params.ptr_Aux, + 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 + ); + + // Define the reduction output pointer and move to the appropriate place + typename Epilogue::ElementVector *ptr_Vector = + static_cast(params.ptr_Vector); + + + // 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() * + cutlass::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, + iterator_Aux, + 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 cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/kernel/sm90_implicit_gemm_tma_warpspecialized.hpp b/include/cutlass/conv/kernel/sm90_implicit_gemm_tma_warpspecialized.hpp new file mode 100644 index 00000000..c6996f15 --- /dev/null +++ b/include/cutlass/conv/kernel/sm90_implicit_gemm_tma_warpspecialized.hpp @@ -0,0 +1,391 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (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 "cutlass/cutlass.h" +#include "cutlass/fast_math.h" +#include "cutlass/kernel_hardware_info.hpp" + +#include "cute/tensor.hpp" +#include "cute/arch/cluster_sm90.hpp" + +#include "cutlass/conv/convolution.h" +#include "cutlass/conv/dispatch_policy.hpp" +#include "cutlass/pipeline/sm90_pipeline.hpp" +#include "cutlass/gemm/kernel/tile_scheduler.hpp" + +/////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::conv::kernel { + +/////////////////////////////////////////////////////////////////////////////// + +template < + class CollectiveMainloop_, + class CollectiveEpilogue_, + class TileSchedulerTag +> +class ConvUniversal< + CollectiveMainloop_, + CollectiveEpilogue_, + TileSchedulerTag, + cute::enable_if_t>> +{ +public: + // + // Type Aliases + // + + // 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 constexpr int NumSpatialDimensions = CollectiveMainloop::NumSpatialDimensions; + 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_void_v, + "TMA warp-specialized kernel does not support specializing the tile scheduler."); + using TileScheduler = typename cutlass::gemm::kernel::detail::TileSchedulerSelector< + TileSchedulerTag, ArchTag, TileShape, ClusterShape>::Scheduler; + using TileSchedulerArguments = typename TileScheduler::Arguments; + + // 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 = CUTE_STATIC_V(size(TiledMma{})) + (NumLoadWarpGroups * NumThreadsPerWarpGroup); + static constexpr uint32_t MinBlocksPerMultiprocessor = 1; + + // Host facing host arguments + struct Arguments { + MainloopArguments mainloop{}; + EpilogueArguments epilogue{}; + KernelHardwareInfo hw_info{}; + TileSchedulerArguments scheduler{}; + }; + + // Kernel device entry point API + struct Params { + MainloopParams mainloop; + EpilogueParams epilogue; + }; + + // + // Methods + // + + // Map user facing arguments to device facing params + static Params + to_underlying_arguments(Arguments const& args, void* workspace) { + (void) workspace; + auto mainloop_params = CollectiveMainloop::to_underlying_arguments(args.mainloop, workspace); + auto problem_shape_MNKL = append<4>(mainloop_params.problem_shape, Int<1>{}); + + return { + mainloop_params, + CollectiveEpilogue::to_underlying_arguments(problem_shape_MNKL, args.epilogue, workspace) + }; + } + + // Given arguemnts, returns true if the kernel can successfully compute upon them. False otherwise. + static bool + can_implement(Arguments const& args) { + bool implementable = true; + implementable &= CollectiveMainloop::can_implement(args.mainloop.problem_shape, args.mainloop); + implementable &= CollectiveEpilogue::can_implement(args.mainloop.problem_shape.get_transformed_problem_shape_MNK(), args.epilogue); + return implementable; + } + + static size_t + 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) { + // The CONV mainloop params problem shape will be the cute::Shape<> rank-3 MNK tuple we want for grid planning + // Although conv problems do not have an L mode, we add it here to comply with the scheduler API + auto linear_problem_shape_MNKL = make_shape( + size<0>(params.mainloop.problem_shape), // M mode is linearized. + shape<1>(params.mainloop.problem_shape), + shape<2>(params.mainloop.problem_shape), + Int<1>{}); + + return cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90::get_tiled_cta_shape_mnl( + linear_problem_shape_MNKL, TileShape{}, ClusterShape{}); + } + + static dim3 + get_block_shape() { + return dim3(MaxThreadsPerBlock, 1, 1); + } + + CUTLASS_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(__CUDA_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_sync(); + 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, ClusterShape{}); + + // 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 = cutlass::make_producer_start_state(); + PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state(); + PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state(); + + // Separate out problem shape for convenience + auto M = get<0>(params.mainloop.problem_shape); + auto N = get<1>(params.mainloop.problem_shape); + auto K = get<2>(params.mainloop.problem_shape); + // output strides are coalesced so we linearize the output shape to match the shape/stride profiles + auto linear_problem_shape_MNKL = make_shape(size(M), N, K, Int<1>{}); + + // TMA requires special handling of strides to deal with coord codomain mapping + // Represent the full tensors -- get these from TMA + Tensor mA_mk = params.mainloop.tma_load_a.get_tma_tensor(make_shape(M, size(K))); + Tensor mB_nk = params.mainloop.tma_load_b.get_tma_tensor(make_shape(N, K)); + + // Get the appropriate blocks for this thread block -- potential for thread block locality + auto cta_tile_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K) + TiledMma tiled_mma; + + // Make tiled views, defer the slice + Tensor gA_mk = local_tile(mA_mk, cta_tile_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k) + Tensor gB_nk = local_tile(mB_nk, cta_tile_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k) + + // Compute m_coord, n_coord, and l_coord with their post-tiled shapes + auto m_coord = idx2crd(int(blockIdx.x), shape<2>(gA_mk)); + auto n_coord = idx2crd(int(blockIdx.y), shape<2>(gB_nk)); + // The output shape M is linearized so the output coord M here should also be linearized. + auto output_tile_coord = make_coord(int(blockIdx.x), n_coord, _, Int<0>{}); + + // Slice with m_coord and n_coord + Tensor gA = gA_mk(_,_,m_coord,_); // (BLK_M,BLK_K,k) + Tensor gB = gB_nk(_,_,n_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(cta_tile_shape); + auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(cta_tile_shape); + + // Make sure pipeline init is visible to all producers and consumer CTAs in cluster + if constexpr (size(ClusterShape{}) > 1) { + cute::cluster_arrive_relaxed(); + cute::cluster_wait(); + } + else { + __syncthreads(); + } + + // In a warp specialized kernel, collectives expose data movement and compute operations separately + CollectiveMainloop collective_mainloop; + CollectiveEpilogue collective_epilogue{params.epilogue, shared_storage.tensors.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_producer_load_needed()) { + collective_epilogue.load( + epi_load_pipeline, + epi_load_pipe_producer_state, + linear_problem_shape_MNKL, + cta_tile_shape, + output_tile_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>(cta_tile_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, + linear_problem_shape_MNKL, + cta_tile_shape, + output_tile_coord, + accumulators, + tiled_mma, + warp_group_thread_idx, + shared_storage.tensors.epilogue + ); + } + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::conv::kernel diff --git a/include/cutlass/conv/threadblock/implicit_gemm_multistage.h b/include/cutlass/conv/threadblock/implicit_gemm_multistage.h index eec44996..eea7743a 100644 --- a/include/cutlass/conv/threadblock/implicit_gemm_multistage.h +++ b/include/cutlass/conv/threadblock/implicit_gemm_multistage.h @@ -139,7 +139,7 @@ public: // 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; + static bool const kStagedAccumulation = arch::detail::UseStagedAccumulation::value; }; private: diff --git a/include/cutlass/cuda_host_adapter.hpp b/include/cutlass/cuda_host_adapter.hpp index ca3531d9..a62ebee7 100644 --- a/include/cutlass/cuda_host_adapter.hpp +++ b/include/cutlass/cuda_host_adapter.hpp @@ -180,6 +180,38 @@ struct CudaHostAdapter { cudaStream_t cuda_stream, void** kernel_params, int32_t kernel_index) const = 0; + +protected: + + /** + * Fills a buffer in Global Memory with a byte sequence copied from host memory. + * This function can be overriden to dispatch to the appropriate cuMemsetD*Async API + */ + virtual Status memsetDeviceImpl( + void* destination, ///< Device memory pointer to be filled + void const* fill_value, ///< Value to be filled in the buffer + size_t fill_size, ///< Size of the data type to be used for filling the buffer + size_t count, ///< Number of elements of size fill_size + cudaStream_t stream) const = 0; + +public: + + /// Fills a buffer in Global Memory with a byte sequence copied from host memory + template + Status memsetDevice( + void* destination, + FillValueType fill_value, + size_t count, + cudaStream_t stream) const + { + return this->memsetDeviceImpl( + destination, + &fill_value, + sizeof(FillValueType), + count, + stream); + } + }; ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/detail/layout.hpp b/include/cutlass/detail/layout.hpp index f3b99638..49bad92d 100644 --- a/include/cutlass/detail/layout.hpp +++ b/include/cutlass/detail/layout.hpp @@ -32,8 +32,10 @@ #include "cutlass/layout/matrix.h" #include "cutlass/layout/tensor.h" +#include "cutlass/numeric_types.h" #include "cute/layout.hpp" +#include "cute/util/type_traits.hpp" #include "cute/arch/copy_sm90_tma.hpp" //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -118,6 +120,60 @@ struct TagToStrideB { template struct TagToStrideC : TagToStrideA { }; +// Conv: Maps to modes ((P,N), C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride +template <> +struct TagToStrideC { + using type = cute::Stride, cute::Int<1>, cute::Int<0>>; +}; + +// Conv: Maps to modes (PN, C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride +template <> +struct TagToStrideC { + using type = cute::Stride, cute::Int<0>>; +}; + +// Conv: Maps to modes ((P,Q,N), C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride +template <> +struct TagToStrideC { + using type = cute::Stride, cute::Int<1>, cute::Int<0>>; +}; + +// Conv: Maps to modes (PQN, C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride +template <> +struct TagToStrideC { + using type = cute::Stride, cute::Int<0>>; +}; + +// Conv: Maps to modes ((P,Q,Z,N), C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride +template <> +struct TagToStrideC { + using type = cute::Stride, cute::Int<1>, cute::Int<0>>; +}; + +// Conv: Maps to modes (PQZN, C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride +template <> +struct TagToStrideC { + using type = cute::Stride, cute::Int<0>>; +}; + +// Conv: Maps to modes (K, (C,S), _0) for compatiblity with GEMM epilogues expecting a batch mode stride +template <> +struct TagToStrideC { + using type = cute::Stride, int64_t>, cute::Int<0>>; +}; + +// Conv: Maps to modes (K, (C,S,R), _0) for compatiblity with GEMM epilogues expecting a batch mode stride +template <> +struct TagToStrideC { + using type = cute::Stride, int64_t, int64_t>, cute::Int<0>>; +}; + +// Conv: Maps to modes (K, (C,S,R,T), _0) for compatiblity with GEMM epilogues expecting a batch mode stride +template <> +struct TagToStrideC { + using type = cute::Stride, int64_t, int64_t, int64_t>, cute::Int<0>>; +}; + // Convenience aliases template using TagToStrideA_t = typename TagToStrideA::type; @@ -229,8 +285,15 @@ constexpr bool is_tma_copy_engine() { return false; } +template +struct RawDtype { using type = X; }; + +template +struct RawDtype> { using type = typename X::raw_type; }; + + // Inspects a TiledCopy and returns its alignment in terms of element count -template +template constexpr int get_alignment_count_from_gmem_tiled_copy() { diff --git a/include/cutlass/epilogue/collective/builders/sm90_builder.inl b/include/cutlass/epilogue/collective/builders/sm90_builder.inl index f1c47f44..b8952d15 100644 --- a/include/cutlass/epilogue/collective/builders/sm90_builder.inl +++ b/include/cutlass/epilogue/collective/builders/sm90_builder.inl @@ -67,11 +67,14 @@ sm90_get_tma_dispatch_policy() { constexpr int EpiTiles = size(shape_div(take<0,2>(TileShapeMNK{}), EpilogueTileMN{})); constexpr int FragmentSize = size(EpilogueTileMN{}) / (detail::sm90_is_cooperative_v ? 256 : 128); - constexpr int ReuseSmemC = (sizeof_bits_v == sizeof_bits_v) && (sizeof_bits_v > 8); - constexpr int StagesD = 2; - constexpr int StagesC = ReuseSmemC ? cute::max(EpiTiles, StagesD + 1) : EpiTiles; + // 8b residuals load fast and consume little smem, so the perf cost of waiting on stores to finish outweighs the cost of extra allocation + constexpr bool ReuseSmem = (sizeof_bits_v == sizeof_bits_v) && (sizeof_bits_v > 8); + constexpr bool DelayTmaStore = is_void_v; // TMA store delay performs worse with residual loads + constexpr int StagesD = cute::min(EpiTiles, 2); + constexpr int StagesC = ReuseSmem ? cute::max(cute::min(EpiTiles, 4), StagesD+1) + : cute::min(EpiTiles, 4); - return Sm90TmaWarpSpecialized{}; + return Sm90TmaWarpSpecialized{}; } // Returns the smem layout atom to be used for C or D matrix @@ -102,29 +105,42 @@ template ) { + auto epi_tile = [&] () { + if constexpr (detail::sm90_is_cooperative_v) { + auto tile_m = cute::min(_128{}, size<0>(TileShape_MNK{})); + auto tile_n = cute::min(_32{}, size<1>(TileShape_MNK{})); + return make_shape(tile_m, tile_n); + } + else if constexpr (detail::sm90_is_warp_specialized_v) { + constexpr int N_perf = sizeof_bits_v == 8 ? 64 : 32; + auto tile_m = cute::min(_64{}, size<0>(TileShape_MNK{})); + auto tile_n = cute::min(Int{}, size<1>(TileShape_MNK{})); + return make_shape(tile_m, tile_n); + } + else { + static_assert(cutlass::detail::dependent_false, "Unsupported schedule."); + } + }(); - if constexpr (detail::sm90_is_cooperative_v) { - using N_tile = decltype(cute::min(_32{}, get<1>(TileShape_MNK{}))); - if constexpr (size<0>(TileShape_MNK{}) >= 128) { - return Shape<_128, N_tile>{}; - } - else { - return Shape<_64, N_tile>{}; - } - } - else if constexpr (detail::sm90_is_warp_specialized_v) { - if constexpr (sizeof_bits_v == 8) { - using N_tile = decltype(cute::min(_64{}, get<1>(TileShape_MNK{}))); - return Shape<_64, N_tile>{}; - } - else { - using N_tile = decltype(cute::min(_32{}, get<1>(TileShape_MNK{}))); - return Shape<_64,N_tile>{}; - } - } - else { - static_assert(cutlass::detail::dependent_false, "Unsupported schedule."); - } + return cute::transform(epi_tile, seq<0,1>{}, + [] (auto epi_tiler, auto I) { + auto cta_tiler = make_layout(get(TileShape_MNK{})); + // This is a multimodal CTA tiler, transform before returning + if constexpr (depth(cta_tiler) > 0) { + // This is an implicit multimodal tiler, match profile and return + if constexpr (tuple_size_v == 1) { + return make_tile(epi_tiler); + } + // This is an explicit multimodal tiler, compose out epi tiler + else { + return composition(cta_tiler, epi_tiler); + } + } + // This is a flat CTA tiler, no need for transformation + else { + return epi_tiler; + } + }); } else if constexpr (cute::is_tuple::value) { EpilogueTileType epi_tile; @@ -178,7 +194,7 @@ sm90_get_smem_load_op_for_source() { } else { // auto-vectorizing load - return AutoVectorizingCopyWithAssumedAlignment{}; + return AutoVectorizingCopyWithAssumedAlignment<128>{}; } } @@ -188,19 +204,20 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, class FusionOp, class TileShape_MNK, class EpilogueTile_MN, class ElementAccumulator > struct CallbacksBuilder< - Sm90TmaWarpSpecialized, + Sm90TmaWarpSpecialized, FusionOp, TileShape_MNK, EpilogueTile_MN, ElementAccumulator, - enable_if_t<(FusionOp::IsAuxOutSupported ^ FusionOp::IsAuxInSupported) // only one aux tensor - && not is_subbyte_v> + cute::enable_if_t<(FusionOp::IsAuxOutSupported ^ FusionOp::IsAuxInSupported) // only one aux tensor + && not cute::is_subbyte_v> > { using GmemStrideTypeAux = gemm::TagToStrideC_t; using SmemLayoutAtomAux = decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom< @@ -209,10 +226,10 @@ struct CallbacksBuilder< GmemStrideTypeAux, typename FusionOp::ElementAux>()); using CopyOpS2R = decltype(detail::sm90_get_smem_load_op_for_source< GmemStrideTypeAux, typename FusionOp::ElementAux>()); - using SmemCopyOpAux = conditional_t; + using SmemCopyOpAux = cute::conditional_t; using Callbacks = fusion::FusionCallbacks< - Sm90TmaWarpSpecialized, + Sm90TmaWarpSpecialized, FusionOp, TileShape_MNK, EpilogueTile_MN, SmemLayoutAtomAux, SmemCopyOpAux >; @@ -223,22 +240,23 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, class FusionOp, class TileShape_MNK, class EpilogueTile_MN, class ElementAccumulator > struct CallbacksBuilder< - Sm90TmaWarpSpecialized, + Sm90TmaWarpSpecialized, FusionOp, TileShape_MNK, EpilogueTile_MN, ElementAccumulator, - enable_if_t<(FusionOp::IsAuxOutSupported ^ FusionOp::IsAuxInSupported) // only one aux tensor + cute::enable_if_t<(FusionOp::IsAuxOutSupported ^ FusionOp::IsAuxInSupported) // only one aux tensor && sizeof_bits_v == 1> > { using Callbacks = fusion::FusionCallbacks< - Sm90TmaWarpSpecialized, + Sm90TmaWarpSpecialized, FusionOp, TileShape_MNK, EpilogueTile_MN, Layout<_1,_0>, DefaultCopy // aux bit tensor doesn't use smem >; @@ -272,12 +290,14 @@ struct Sm90TmaBuilderImpl { using GmemStrideTypeC = cutlass::detail::TagToStrideC_t; using GmemStrideTypeD = cutlass::detail::TagToStrideC_t; - using CopyOpS2G = + using CopyOpS2G = cute::conditional_t, + SM90_TMA_STORE_IM2COL, SM90_TMA_STORE - ; - using CopyOpG2S = + >; + using CopyOpG2S = cute::conditional_t, + SM90_TMA_LOAD_IM2COL, SM90_TMA_LOAD - ; + >; // TMA builder allows for passing callbacks directly, which is either a fusion::FusionCallbacks // instance or a direct visitor implementation, e.g. fusion::Sm90LinearCombination @@ -697,7 +717,8 @@ private: /* StagesC = */ size(shape_div(take<0, 2>(TileShape_MNK{}), EpilogueTile_MN{})), /* StagesD = */ 2, /* FragmentSize = */ size(EpilogueTile_MN{}) / (detail::sm90_is_cooperative_v ? 256 : 128), - /* ReuseSmemC = */ sizeof_bits_v == sizeof_bits_v + /* ReuseSmemC = */ sizeof_bits_v == sizeof_bits_v, + false >; using GmemStrideTypeAux = gemm::TagToStrideC_t; diff --git a/include/cutlass/epilogue/collective/collective_builder.hpp b/include/cutlass/epilogue/collective/collective_builder.hpp index 5e76595a..10aad81d 100644 --- a/include/cutlass/epilogue/collective/collective_builder.hpp +++ b/include/cutlass/epilogue/collective/collective_builder.hpp @@ -100,7 +100,7 @@ struct CallbacksBuilder< TileShape_MNK, EpilogueTile_MN, ElementAccumulator, - enable_if_t> + cute::enable_if_t> > { using Callbacks = FusionCallbacks; }; diff --git a/include/cutlass/epilogue/collective/default_epilogue.hpp b/include/cutlass/epilogue/collective/default_epilogue.hpp index aad96de7..bbeeacac 100644 --- a/include/cutlass/epilogue/collective/default_epilogue.hpp +++ b/include/cutlass/epilogue/collective/default_epilogue.hpp @@ -39,7 +39,8 @@ #include "cutlass/epilogue/collective/detail.hpp" #include "cute/tensor.hpp" -#include "cute/numeric/int.hpp" +#include "cute/numeric/numeric_types.hpp" +#include "cutlass/cuda_host_adapter.hpp" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -64,7 +65,7 @@ public: // using EpilogueSchedule = EpilogueSchedule_; using DispatchPolicy = EpilogueSchedule_; - + // derived types of output thread level operator using ThreadEpilogueOp = ThreadEpilogueOp_; using ElementOutput = typename ThreadEpilogueOp::ElementOutput; @@ -87,6 +88,8 @@ public: struct SharedStorage { }; + using TensorStorage = SharedStorage; + // Host side epilogue arguments struct Arguments { typename ThreadEpilogueOp::Params thread{}; @@ -120,7 +123,8 @@ public: template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } diff --git a/include/cutlass/epilogue/collective/default_epilogue_array.hpp b/include/cutlass/epilogue/collective/default_epilogue_array.hpp index 5807f87e..c2f8423e 100644 --- a/include/cutlass/epilogue/collective/default_epilogue_array.hpp +++ b/include/cutlass/epilogue/collective/default_epilogue_array.hpp @@ -39,9 +39,10 @@ #include "cutlass/epilogue/collective/detail.hpp" #include "cute/tensor.hpp" -#include "cute/numeric/int.hpp" +#include "cute/numeric/numeric_types.hpp" #include "cutlass/trace.h" +#include "cutlass/cuda_host_adapter.hpp" ///////////////////////////////////////////////////////////////////////////////////////////////// namespace cutlass { @@ -123,7 +124,8 @@ public: template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } diff --git a/include/cutlass/epilogue/collective/detail.hpp b/include/cutlass/epilogue/collective/detail.hpp index 32e790fd..6b01a22e 100644 --- a/include/cutlass/epilogue/collective/detail.hpp +++ b/include/cutlass/epilogue/collective/detail.hpp @@ -38,7 +38,7 @@ #include "cutlass/epilogue/dispatch_policy.hpp" #include "cute/tensor.hpp" -#include "cute/numeric/int.hpp" +#include "cute/numeric/numeric_types.hpp" #include "cute/util/type_traits.hpp" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -63,6 +63,14 @@ is_n_major() { return cutlass::gemm::detail::is_major<1,Stride>(); } +template +constexpr bool +is_im2col() { + return cute::is_same_v> + || cute::is_same_v> + || cute::is_same_v>; +} + using cutlass::atomic_maximum; template @@ -76,6 +84,12 @@ template static constexpr bool sm90_is_warp_specialized_v = cute::is_base_of_v; +template +static constexpr bool is_im2col_mode = + cute::is_same_v || + cute::is_same_v || + cute::is_same_v; + template struct EmptyStorage { CUTLASS_HOST_DEVICE @@ -105,6 +119,28 @@ struct IsThreadEpilogueOpWithBias +struct IsThreadEpilogueOpWithPerChannelScaling { + static constexpr bool value = false; +}; + +template +struct IsThreadEpilogueOpWithPerChannelScaling > { + static constexpr bool value = true; +}; + +template +struct IsThreadEpilogueOpWithActivation { + static constexpr bool value = false; + using type = void; +}; + +template +struct IsThreadEpilogueOpWithActivation > { + static constexpr bool value = true; + using type = typename ThreadEpilogueOp::ActivationFn; +}; + // Wrapper class to use operator-style epilogues in sm90 TMA warp-specialized kernels template class Sm90TmaWarpSpecializedAdapter : public EpilogueOp { diff --git a/include/cutlass/epilogue/collective/epilogue_tensor_broadcast.hpp b/include/cutlass/epilogue/collective/epilogue_tensor_broadcast.hpp index 8c43156d..c870b706 100644 --- a/include/cutlass/epilogue/collective/epilogue_tensor_broadcast.hpp +++ b/include/cutlass/epilogue/collective/epilogue_tensor_broadcast.hpp @@ -55,6 +55,7 @@ #include "cutlass/epilogue/collective/detail.hpp" #include "cute/tensor.hpp" +#include "cutlass/cuda_host_adapter.hpp" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -142,7 +143,8 @@ public: template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } diff --git a/include/cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp b/include/cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp index 89993f50..be19944d 100644 --- a/include/cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp +++ b/include/cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp @@ -129,7 +129,8 @@ public: template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } diff --git a/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp b/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp index a463463b..f23a74e0 100644 --- a/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp +++ b/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp @@ -45,6 +45,7 @@ #include "cutlass/trace.h" #include "cute/tensor.hpp" +#include "cutlass/cuda_host_adapter.hpp" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -59,6 +60,7 @@ template < int StagesD_, int FragmentSize_, bool ReuseSmemC_, + bool DelayTmaStore_, class CtaTileMNK_, // (CTA_M,CTA_N,CTA_K) class EpilogueTile_, // (EPI_TILE_M,EPI_TILE_N) class ElementC_, @@ -74,7 +76,7 @@ template < class CopyOpR2S_ > class CollectiveEpilogue< - Sm90TmaWarpSpecialized, + Sm90TmaWarpSpecialized, CtaTileMNK_, EpilogueTile_, ElementC_, @@ -93,7 +95,7 @@ public: // // Type Aliases // - using DispatchPolicy = Sm90TmaWarpSpecialized; + using DispatchPolicy = Sm90TmaWarpSpecialized; using CtaTileMNK = CtaTileMNK_; using EpilogueTile = EpilogueTile_; using FusionCallbacks = FusionCallbacks_; @@ -129,10 +131,14 @@ private: constexpr static int StagesC = StagesC_; constexpr static int StagesD = StagesD_; constexpr static bool ReuseSmemC = ReuseSmemC_ and is_destination_supported; + constexpr static bool DelayTmaStore = DelayTmaStore_; constexpr static bool is_m_major_C = detail::is_m_major(); constexpr static bool is_m_major_D = detail::is_m_major(); + constexpr static bool is_im2col_C = cute::is_same_v; + constexpr static bool is_im2col_D = cute::is_same_v; + using SmemLayoutC = decltype(tile_to_shape( SmemLayoutAtomC{}, make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int{}), @@ -213,12 +219,16 @@ public: CopyOpG2S{}, make_tensor(make_gmem_ptr(static_cast(nullptr)), repeat_like(StrideC{}, int32_t(0)), StrideC{}), - SmemLayoutC{}(_,_,0))); + take<0,2>(SmemLayoutC{}), + EpilogueTile{}, + _1{})); using TMA_D = decltype(make_tma_copy( CopyOpS2G{}, make_tensor(make_gmem_ptr(static_cast(nullptr)), repeat_like(StrideD{}, int32_t(0)), StrideD{}), - SmemLayoutD{}(_,_,0))); + take<0,2>(SmemLayoutD{}), + EpilogueTile{}, + _1{})); typename FusionCallbacks::Params thread{}; TMA_C tma_load_c; @@ -238,23 +248,20 @@ public: // 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, 1); auto [M, N, K, L] = problem_shape_MNKL; - auto M_C = - size(M) - ; - auto M_D = - size(M) - ; + // For fprop/dgrad kernel, problem shape M is multimodal which should be linearized under tiled mode + auto M_C = conditional_return(M, size(M)); + auto M_D = conditional_return(M, size(M)); - typename Params::TMA_C tma_load_c; + typename Params::TMA_C tma_load_c = {}; if constexpr (is_source_supported) { Tensor tensor_c = make_tensor(make_gmem_ptr(args.ptr_C), make_layout(make_shape(M_C,N,L), args.dC)); - tma_load_c = make_tma_copy(CopyOpG2S{}, tensor_c, SmemLayoutC{}(_,_,0)); + tma_load_c = make_tma_copy(CopyOpG2S{}, tensor_c, take<0,2>(SmemLayoutC{}), EpilogueTile{}, _1{}); } typename Params::TMA_D tma_store_d; if constexpr (is_destination_supported) { Tensor tensor_d = make_tensor(make_gmem_ptr(args.ptr_D), make_layout(make_shape(M_D,N,L), args.dD)); - tma_store_d = make_tma_copy(CopyOpS2G{}, tensor_d, SmemLayoutD{}(_,_,0)); + tma_store_d = make_tma_copy(CopyOpS2G{}, tensor_d, take<0,2>(SmemLayoutD{}), EpilogueTile{}, _1{}); } return { @@ -272,8 +279,9 @@ public: template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { - return FusionCallbacks::initialize_workspace(problem_shape, args.thread, workspace, stream); + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { + return FusionCallbacks::initialize_workspace(problem_shape, args.thread, workspace, stream, cuda_adapter); } template @@ -308,10 +316,7 @@ public: static constexpr int get_load_pipe_increment(TileShapeMNK tile_shape_MNK) { // Compute number of epilogue subtiles - constexpr int epi_m = size<0>(tile_shape_MNK) / size<0>(EpilogueTile{}); - constexpr int epi_n = size<1>(tile_shape_MNK) / size<1>(EpilogueTile{}); - - return epi_m * epi_n; + return size<1>(zipped_divide(make_layout(take<0,2>(tile_shape_MNK)), EpilogueTile{})); } template @@ -366,18 +371,14 @@ public: auto [M, N, K, L] = problem_shape_mnkl; auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl; - auto coord_shape = - make_coord(m_coord, n_coord, l_coord) - ; + // The tma tensor C under im2col mode only has two modes (M, N) which + // should be local tiled with only (m_coord, n_coord). + auto coord_shape = conditional_return( + make_coord(m_coord, n_coord), + make_coord(m_coord, n_coord, l_coord)); // Tile residue - auto m_max_coord = unwrap(cute::transform(make_seq(tile_shape_MNK)>{}, [&](auto i) { - return get<0,i>(problem_shape_mnkl) - get<0,i>(tile_shape_MNK) * get<0,i>(tile_coord_mnkl); - })); - auto n_max_coord = unwrap(cute::transform(make_seq(tile_shape_MNK)>{}, [&](auto i) { - return get<1,i>(problem_shape_mnkl) - get<1,i>(tile_shape_MNK) * get<1,i>(tile_coord_mnkl); - })); - auto residue_mn = make_coord(m_max_coord, n_max_coord); + auto residue_mn = make_coord(M,N); // Represent the full source tensor, slice to get the tile this CTA is currently responsible for Tensor mC_mn = params.tma_load_c.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) @@ -511,9 +512,11 @@ public: auto epi_tile_m = size<0>(EpilogueTile{}); auto epi_tile_n = size<1>(EpilogueTile{}); - auto coord_shape = - make_coord(m_coord, n_coord, l_coord) - ; + // The tma tensor D under im2col mode only has two modes (M, N) which + // should be local tiled with only (m_coord, n_coord). + auto coord_shape = conditional_return( + make_coord(m_coord, n_coord), + make_coord(m_coord, n_coord, l_coord)); // Represent the full output tensor, slice to get the tile this CTA is responsible for Tensor mD_mn = params.tma_store_d.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) @@ -572,7 +575,7 @@ public: // Allocate C registers // If C smem load is a non-vectorized dst(i) = src(i) then we can allocate C registers directly in the compute type // to eliminate some redundant pack+unpack instruction sequences for sub-word types - constexpr bool IsDirectS2R = cute::is_same_v + constexpr bool IsDirectS2R = cute::is_same_v> && decltype(max_common_vector(tSR_rC_layout, tSR_sC.layout()))::value <= 1; using RegisterElementC = cute::conditional_t; Tensor tRS_rC = make_tensor(tRS_rD_layout); // (R2S,R2S_M,R2S_N) @@ -583,18 +586,11 @@ public: Tensor bSG_sD = thrblk_s2g.partition_S(sD_epi); // (S2G,S2G_M,S2G_N,PIPE_D) Tensor bSG_gD = thrblk_s2g.partition_D(gD_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) - // Coordinate tensors and residue for tile quantization - auto m_max_coord = unwrap(cute::transform(make_seq(CtaTileMNK{})>{}, [&](auto i) { - auto c_m = get<0,i>(problem_shape_mnkl) - get<0,i>(CtaTileMNK{}) * get<0,i>(tile_coord_mnkl); - return cute::max(0, c_m); - })); - auto n_max_coord = unwrap(cute::transform(make_seq(CtaTileMNK{})>{}, [&](auto i) { - auto c_n = get<1,i>(problem_shape_mnkl) - get<1,i>(CtaTileMNK{}) * get<1,i>(tile_coord_mnkl); - return cute::max(0, c_n); - })); - auto residue_mn = make_coord(m_max_coord, n_max_coord); - Tensor cD = make_identity_tensor(take<0,2>(CtaTileMNK{})); + // OOB predication for tile quantization "residue" + Tensor mD_crd = make_identity_tensor(make_shape(M,N)); + Tensor cD = local_tile(mD_crd, take<0,2>(CtaTileMNK{}), make_coord(m_coord, n_coord)); Tensor tRS_cD = thread_r2s.partition_S(flat_divide(cD, EpilogueTile{})); + auto residue_mn = make_coord(M,N); 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"); @@ -634,6 +630,52 @@ public: load_wait_state.phase_ ^= 1; } + // We can delay issue of TMA store by one iteration to achieve better interleaving of non-TMA instructions + // Sync requirements of smem reuse may preclude this optimization + // Delayed stores cause delayed stage releases which causes deadlock when StagesC == StagesD + int epi_m_prev = 0, epi_n_prev = 0; + static_assert(not (DelayTmaStore and ReuseSmemC and StagesC == StagesD), "This TMA epilogue configuration will deadlock"); + + // The TMA store sequence for one subtile iteration + auto tma_store_fn = [&] (int epi_m, int epi_n) { + // Write the tile from smem to gmem with TMA + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + if constexpr (is_destination_supported) { + if (issue_tma_store) { + copy(params.tma_store_d, bSG_sD(_,_,_,store_pipe_producer_state.index()), bSG_gD(_,_,_,epi_m,epi_n)); + } + } + + // Post async fence, pre TMA commit callback entry point + cst_callbacks.tma_store(epi_m, epi_n, store_pipe_producer_state.count(), issue_tma_store); + + // Commit the TMA stores for this stage + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + ++issued_stores; + + // Wait for the next smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + + if constexpr (ReuseSmemC) { + // producer_acquire returns when at most StagesD-1 committed stores are pending + bool store_finished = issued_stores > StorePipeline::UnacquiredStages; + // Let dma warp know earliest smem buffer is consumed and empty after StagesD producer commits + if (store_finished) { + if (is_producer_load_needed) { + load_pipeline.consumer_release(load_pipe_consumer_state); + } + ++load_pipe_consumer_state; + } + } + }; + // // BEGIN EPILOGUE // @@ -646,6 +688,7 @@ public: for (int epi_n = 0; epi_n < size<3>(gD_epi); ++epi_n) { CUTLASS_PRAGMA_UNROLL for (int epi_m = 0; epi_m < size<2>(gD_epi); ++epi_m) { + bool is_first_iteration = epi_m == 0 && epi_n == 0; bool is_last_iteration = epi_m == size<2>(gD_epi)-1 && epi_n == size<3>(gD_epi)-1; if (subtile_idx != -1 && (epi_n * static_cast(size<2>(gD_epi)) + epi_m) != subtile_idx) { @@ -686,72 +729,42 @@ public: tRS_rD_frg(epi_v) = cst_callbacks.visit(tRS_rAcc_frg_mn(r2s_v + epi_v), epi_v, epi_m, epi_n); } + // The latest we can delay the TMA store is right before the smem store of the next iteration + // since the current TMA store needs to be committed before we can acquire the next smem buffer + if constexpr (DelayTmaStore) { + // Issue TMA stores for the previous subtile + if (not is_first_iteration and subtile_idx == -1) { + tma_store_fn(epi_m_prev, epi_n_prev); + } + epi_m_prev = epi_m; + epi_n_prev = epi_n; + } + + // Smem reduction callback entry point using current store buffer for workspace + cst_callbacks.reduce(sD_epi(_,_,store_pipe_producer_state.index()), + synchronize, epi_m, epi_n, is_last_iteration); + // Copy tile from register to smem if constexpr (is_destination_supported) { copy(tiled_r2s, tRS_rD, tRS_sD(_,_,_,store_pipe_producer_state.index())); } - // Post visit, pre async fence callback entry point + // Post reduction, pre TMA store callback entry point constexpr bool issue_smem_store = true; // No smem store predication - cst_callbacks.postvisit(epi_m, epi_n, store_pipe_producer_state.count(), issue_smem_store); + cst_callbacks.postreduce(epi_m, epi_n, store_pipe_producer_state.count(), issue_smem_store); - // Write the tile from smem to gmem with TMA - cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA - synchronize(); // ensure all threads have issued their async fence - if constexpr (is_destination_supported) { - if (issue_tma_store) { - copy(params.tma_store_d, bSG_sD(_,_,_,store_pipe_producer_state.index()), bSG_gD(_,_,_,epi_m,epi_n)); - } - } - - // Post async fence, pre TMA commit callback entry point - cst_callbacks.step(epi_m, epi_n, store_pipe_producer_state.count(), issue_tma_store); - - // Commit the TMA stores for this stage - if (issue_tma_store) { - store_pipeline.producer_commit(store_pipe_producer_state); - } - ++store_pipe_producer_state; - ++issued_stores; - - // Wait for the next smem buffer to be available - if (issue_tma_store) { - store_pipeline.producer_acquire(store_pipe_producer_state); - } - synchronize(); - - if constexpr (ReuseSmemC) { - // producer_acquire returns when at most StagesD-1 committed stores are pending - bool store_finished = issued_stores > StorePipeline::UnacquiredStages; - - // Free an smem buffer for reduction if necessary - if (cst_callbacks.is_reduction_buffer_needed(epi_m, epi_n, is_last_iteration) && not store_finished) { - if (issue_tma_store) { - store_pipeline.producer_tail(store_pipe_producer_state); // wait for all TMA stores to finish - } - synchronize(); - } - - // Smem reduction callback entry point using least recently acquired load buffer for workspace - cst_callbacks.reduce(sC_epi(_,_,load_pipe_consumer_state.index()), - synchronize, epi_m, epi_n, is_last_iteration); - - // Let dma warp know earliest smem buffer is consumed and empty after StagesD producer commits - if (store_finished) { - if (is_producer_load_needed) { - load_pipeline.consumer_release(load_pipe_consumer_state); - } - ++load_pipe_consumer_state; - } - } - else { - // Smem reduction callback entry point using most recently acquired store buffer for workspace - cst_callbacks.reduce(sD_epi(_,_,store_pipe_producer_state.index()), - synchronize, epi_m, epi_n, is_last_iteration); + if constexpr (not DelayTmaStore) { + // Issue TMA stores for this subtile + tma_store_fn(epi_m, epi_n); } } // for epi_m } // for epi_n + if constexpr (DelayTmaStore) { + // Issue TMA stores for the last subtile + tma_store_fn(epi_m_prev, epi_n_prev); + } + // Post-loop fusion callback entry point cst_callbacks.end(); diff --git a/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp b/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp index 7eeae04c..8eeb43c2 100644 --- a/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp +++ b/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp @@ -65,7 +65,7 @@ template < > class Sm90EpilogueTmaWarpSpecializedBiasElementwise : public CollectiveEpilogue< - Sm90TmaWarpSpecialized, + Sm90TmaWarpSpecialized, BlockTileShape_, EpilogueTileShape_, ElementC_, @@ -83,7 +83,7 @@ class Sm90EpilogueTmaWarpSpecializedBiasElementwise private: using Impl = CollectiveEpilogue< - Sm90TmaWarpSpecialized, + Sm90TmaWarpSpecialized, BlockTileShape_, EpilogueTileShape_, ElementC_, @@ -111,17 +111,17 @@ public: struct [[deprecated("use Sm90TmaWarpSpecialized Arguments instead")]] Arguments { struct ThreadArgs { - ElementCompute alpha; - ElementCompute beta; - ElementCompute const *alpha_ptr; - ElementCompute const *beta_ptr; + ElementCompute alpha{1}; + ElementCompute beta{0}; + ElementCompute const *alpha_ptr{nullptr}; + ElementCompute const *beta_ptr{nullptr}; } thread; - ElementC_ const* ptr_C; - StrideC_ dC; - ElementD_* ptr_D; - StrideD_ dD; - ElementBias const* ptr_Bias = nullptr; - ElementT* ptr_T = nullptr; + ElementC_ const* ptr_C{nullptr}; + StrideC_ dC{}; + ElementD_* ptr_D{nullptr}; + StrideD_ dD{}; + ElementBias const* ptr_Bias{nullptr}; + ElementT* ptr_T{nullptr}; CUTLASS_HOST_DEVICE operator typename Impl::Arguments() const { diff --git a/include/cutlass/epilogue/dispatch_policy.hpp b/include/cutlass/epilogue/dispatch_policy.hpp index 2d36314e..409ff74d 100644 --- a/include/cutlass/epilogue/dispatch_policy.hpp +++ b/include/cutlass/epilogue/dispatch_policy.hpp @@ -132,13 +132,15 @@ template< int StagesC_, int StagesD_, int FragmentSize_, - bool ReuseSmemC_ + bool ReuseSmemC_, + bool DelayTmaStore_ > struct Sm90TmaWarpSpecialized { constexpr static int StagesC = StagesC_; constexpr static int StagesD = StagesD_; constexpr static int FragmentSize = FragmentSize_; constexpr static bool ReuseSmemC = ReuseSmemC_; + constexpr static bool DelayTmaStore = DelayTmaStore_; }; diff --git a/include/cutlass/epilogue/fusion/operations.hpp b/include/cutlass/epilogue/fusion/operations.hpp index 6417fbc4..a483b1ba 100644 --- a/include/cutlass/epilogue/fusion/operations.hpp +++ b/include/cutlass/epilogue/fusion/operations.hpp @@ -122,6 +122,7 @@ struct LinCombEltAct static constexpr bool IsEltActSupported = true; }; + // D = alpha * acc + beta * C + per-row bias template< class ElementOutput_, diff --git a/include/cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp b/include/cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp index fec0cd50..6c729c10 100644 --- a/include/cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp +++ b/include/cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp @@ -61,6 +61,7 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, class ElementOutput, class ElementCompute, class ElementScalar, @@ -69,7 +70,7 @@ template < class EpilogueTile > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::ScaledAcc, CtaTileShapeMNK, EpilogueTile @@ -132,6 +133,7 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, class ElementOutput, class ElementCompute, class ElementSource, @@ -141,7 +143,7 @@ template < class EpilogueTile > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::LinearCombination, CtaTileShapeMNK, EpilogueTile @@ -196,6 +198,7 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, template class ActivationFn, class ElementOutput, class ElementCompute, @@ -206,7 +209,7 @@ template < class EpilogueTile > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::LinCombEltAct, CtaTileShapeMNK, EpilogueTile @@ -275,6 +278,7 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, class ElementOutput, class ElementCompute, class ElementBias, @@ -286,7 +290,7 @@ template < class EpilogueTile > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::LinCombPerRowBias, CtaTileShapeMNK, EpilogueTile @@ -351,6 +355,7 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, template class ActivationFn, class ElementOutput, class ElementCompute, @@ -363,7 +368,7 @@ template < class EpilogueTile > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::LinCombPerRowBiasEltAct< ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle >, @@ -452,6 +457,7 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, class GmemLayoutTagAux, template class ActivationFn, class ElementOutput, @@ -469,7 +475,7 @@ template < class CopyOpR2S > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::LinCombPerRowBiasEltActAux< GmemLayoutTagAux, ActivationFn, ElementOutput, ElementCompute, ElementAux, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle @@ -553,10 +559,10 @@ template< > using Sm90PerRowLinCombPerRowBias = Sm90EVT, // beta * C + (alpha * acc + bias) - Sm90ColBroadcast<0, CtaTileShapeMNK, ElementScalar, Stride<_1,_0,_0>, AlignmentScalar>, // beta + Sm90ColBroadcast<0, CtaTileShapeMNK, ElementScalar, Stride<_1,_0,int>, AlignmentScalar>, // beta Sm90SrcFetch, // C Sm90EVT, // alpha * acc + bias - Sm90ColBroadcast<0, CtaTileShapeMNK, ElementScalar, Stride<_1,_0,_0>, AlignmentScalar>, // alpha + Sm90ColBroadcast<0, CtaTileShapeMNK, ElementScalar, Stride<_1,_0,int>, AlignmentScalar>, // alpha Sm90AccFetch, // acc Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,int>, AlignmentBias> // bias > @@ -586,6 +592,7 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, template class ActivationFn, class ElementOutput, class ElementCompute, @@ -599,7 +606,7 @@ template < class EpilogueTile > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::PerRowLinCombPerRowBiasEltAct< ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle >, @@ -619,10 +626,14 @@ struct FusionCallbacks< >; struct Arguments { + using StrideAlpha = Stride<_1,_0,int>; + using StrideBeta = Stride<_1,_0,int>; ElementScalar alpha = ElementScalar(1); ElementScalar beta = ElementScalar(0); ElementScalar const* alpha_ptr = nullptr; ElementScalar const* beta_ptr = nullptr; + StrideAlpha dAlpha = {}; + StrideBeta dBeta = {}; using StrideBias = Stride<_1,_0,int>; ElementBias const* bias_ptr = nullptr; @@ -635,14 +646,14 @@ struct FusionCallbacks< return { // unary op : activation(beta * C + (alpha * acc + bias)) { // ternary op : beta * C + (alpha * acc + bias) - {beta_ptr, beta}, // leaf args : beta - {}, // leaf args : C - { // ternary op : alpha * acc + bias - {alpha_ptr, alpha}, // leaf args : alpha - {}, // leaf args : acc + {beta_ptr, beta, dBeta}, // leaf args : beta + {}, // leaf args : C + { // ternary op : alpha * acc + bias + {alpha_ptr, alpha, dAlpha}, // leaf args : alpha + {}, // leaf args : acc {bias_ptr, ElementBias(0), dBias}, // leaf args : bias - {} // ternary args : multiply_add - }, // end ternary op + {} // ternary args : multiply_add + }, // end ternary op {} // ternary args : multiply_add }, // end ternary op activation // unary args : activation @@ -658,6 +669,9 @@ struct FusionCallbacks< namespace detail { +template +constexpr bool is_fp8_v = cute::is_same_v || cute::is_same_v; + // We only apply the scaling factor if output is fp8 template struct ScaleOutOp { template using Op = cutlass::first; }; @@ -723,6 +737,7 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, template class ActivationFn, class ElementOutput, class ElementCompute, @@ -735,7 +750,7 @@ template < class EpilogueTile > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::ScaledLinCombPerRowBiasEltAct< ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle >, @@ -822,6 +837,8 @@ struct FusionCallbacks< // Aux = scale_aux * Z // else // Aux = Z + +// fp8 aux specialization template< class CtaTileShapeMNK, class EpilogueTile, @@ -841,7 +858,7 @@ template< int AlignmentBias = 128 / sizeof_bits_v, FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest > -using Sm90ScaledLinCombPerRowBiasEltActAmaxAux = +using Sm90ScaledLinCombPerRowBiasEltActAmaxAuxFp8 = Sm90SplitTreeVisitor< // Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-row bias Sm90ScaledLinCombPerRowBias, @@ -856,7 +873,7 @@ using Sm90ScaledLinCombPerRowBiasEltActAmaxAux = >, // Aux = Z * scale_aux, amax_aux = max(abs(elements in Aux)) Sm90EVT, // store(Aux) - Sm90EVT::template Op, ElementCompute, ElementCompute, RoundStyle>, // Z * scale_aux + Sm90EVT, // Z * scale_aux Sm90EVT, // amax_aux Sm90SplitTreeFetch // Z >, @@ -865,11 +882,79 @@ using Sm90ScaledLinCombPerRowBiasEltActAmaxAux = > >; +// non-fp8 aux specialization +// lets us use some EVT specializations such as relu + uint1b_t aux +template< + class CtaTileShapeMNK, + class EpilogueTile, + int StagesD, + class StrideAux, + class SmemLayoutAtom, + class CopyOpR2S, + template class ActivationFn, + class ElementOutput, + class ElementCompute, + class ElementAux = ElementOutput, + class ElementAmax = ElementCompute, + class ElementBias = ElementOutput, + class ElementSource = ElementOutput, + class ElementScalar = ElementCompute, + int AlignmentAux = 128 / sizeof_bits_v, + int AlignmentBias = 128 / sizeof_bits_v, + FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest +> +using Sm90ScaledLinCombPerRowBiasEltActAmaxAuxNotFp8 = + // D = activation(Z) * scale_d, amax_d = max(abs(elements in D)) + Sm90EVT::template Op, ElementOutput, ElementCompute, RoundStyle>, // activation(Z) * scale_d + Sm90EVT, // amax_d + Sm90EVT, // activation(Z) + Sm90EVT, // Aux = Z + // Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-row bias + Sm90ScaledLinCombPerRowBias, + > + > + >, + Sm90ScalarBroadcast // scale_d + >; + +// dispatcher +template< + class CtaTileShapeMNK, + class EpilogueTile, + int StagesD, + class StrideAux, + class SmemLayoutAtom, + class CopyOpR2S, + template class ActivationFn, + class ElementOutput, + class ElementCompute, + class ElementAux = ElementOutput, + class ElementAmax = ElementCompute, + class ElementBias = ElementOutput, + class ElementSource = ElementOutput, + class ElementScalar = ElementCompute, + int AlignmentAux = 128 / sizeof_bits_v, + int AlignmentBias = 128 / sizeof_bits_v, + FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest +> +using Sm90ScaledLinCombPerRowBiasEltActAmaxAux = conditional_t, + Sm90ScaledLinCombPerRowBiasEltActAmaxAuxFp8< + CtaTileShapeMNK, EpilogueTile, StagesD, StrideAux, SmemLayoutAtom, CopyOpR2S, ActivationFn, + ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar,AlignmentAux, AlignmentBias, RoundStyle + >, + Sm90ScaledLinCombPerRowBiasEltActAmaxAuxNotFp8< + CtaTileShapeMNK, EpilogueTile, StagesD, StrideAux, SmemLayoutAtom, CopyOpR2S, ActivationFn, + ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle + > +>; + + template < int StagesC, int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, class GmemLayoutTagAux, template class ActivationFn, class ElementOutput, @@ -888,7 +973,7 @@ template < class CopyOpR2S > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::ScaledLinCombPerRowBiasEltActAmaxAux< GmemLayoutTagAux, ActivationFn, ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle @@ -948,70 +1033,101 @@ struct FusionCallbacks< StrideAux dAux = {}; operator typename Impl::Arguments() const { - typename Impl::Arguments args; - // always use structured binding to unpack DAG args since it may or may not be a tuple - auto& [Z_args, aux_args, D_args] = args; - - Z_args = - { // ternary op : (scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias) - {{scale_c, beta}, - {scale_c_ptr, beta_ptr} - }, // leaf args : (scale_c * beta) - {}, // leaf args : C - { // ternary op : (scale_a * scale_b * alpha) * acc + bias - {{scale_a, scale_b, alpha}, - {scale_a_ptr, scale_b_ptr, alpha_ptr} - }, // leaf args : (scale_a * scale_b * alpha) - {}, // leaf args : acc - {bias_ptr, ElementBias(0), dBias}, // leaf args : bias - {} // ternary args : multiply_add - }, // end ternary op - {} // ternary args : multiply_add - }; // end ternary op - // Only compute amax_d if D is fp8 ElementAmax* amax_D_ptr_ = nullptr; - if constexpr (cute::is_same_v || - cute::is_same_v) { + if constexpr (detail::is_fp8_v) { amax_D_ptr_ = amax_D_ptr; } - D_args = - { // binary op : activation(Z) * scale_d or activation(Z) - { // unary op : reduce(activation(Z)) - { // unary op : activation(Z) - {}, // leaf args : Z - activation // unary args : activation - }, // end unary op - {amax_D_ptr_} // unary args : reduce - }, // end unary op - {{scale_d}, - {scale_d_ptr} - }, // leaf args : scale_d - {} // binary args : multiplies or first - }; // end binary op - // Only compute amax_aux if aux is fp8 - ElementAmax* amax_aux_ptr_ = nullptr; - if constexpr (cute::is_same_v || - cute::is_same_v) { - amax_aux_ptr_ = amax_aux_ptr; - } - aux_args = - { // unary op : store(Aux) - { // binary op : Z * scale_d or Z - { // unary op : reduce(Z) - {}, // leaf args : Z - {amax_aux_ptr_} // unary args : reduce - }, // end unary op - {{scale_aux}, - {scale_aux_ptr} + // Aux is fp8 -> DAG arguments + if constexpr (detail::is_fp8_v) { + typename Impl::Arguments args; + // always use structured binding to unpack DAG args since it may or may not be a tuple + auto& [Z_args, aux_args, D_args] = args; + + Z_args = + { // ternary op : (scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias) + {{scale_c, beta}, + {scale_c_ptr, beta_ptr} + }, // leaf args : (scale_c * beta) + {}, // leaf args : C + { // ternary op : (scale_a * scale_b * alpha) * acc + bias + {{scale_a, scale_b, alpha}, + {scale_a_ptr, scale_b_ptr, alpha_ptr} + }, // leaf args : (scale_a * scale_b * alpha) + {}, // leaf args : acc + {bias_ptr, ElementBias(0), dBias}, // leaf args : bias + {} // ternary args : multiply_add + }, // end ternary op + {} // ternary args : multiply_add + }; // end ternary op + + D_args = + { // binary op : activation(Z) * scale_d or activation(Z) + { // unary op : reduce(activation(Z)) + { // unary op : activation(Z) + {}, // leaf args : Z + activation // unary args : activation + }, // end unary op + {amax_D_ptr_} // unary args : reduce + }, // end unary op + {{scale_d}, + {scale_d_ptr} }, // leaf args : scale_d {} // binary args : multiplies or first - }, // end binary op - {aux_ptr, dAux} // unary args : store - }; // end unary op + }; // end binary op - return args; + aux_args = + { // unary op : store(Aux) + { // binary op : Z * scale_d or Z + { // unary op : reduce(Z) + {}, // leaf args : Z + {amax_aux_ptr} // unary args : reduce + }, // end unary op + {{scale_aux}, + {scale_aux_ptr} + }, // leaf args : scale_d + {} // binary args : multiplies + }, // end binary op + {aux_ptr, dAux} // unary args : store + }; // end unary op + + return args; + } + + // Aux is not fp8 -> Tree arguments + else { + return + { // binary op : activation(Z) * scale_d or activation(Z) + { // unary op : reduce(activation(Z)) + { // unary op : activation(Z) + { // unary op : store(Z) + { // ternary op : (scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias) + {{scale_c, beta}, + {scale_c_ptr, beta_ptr} + }, // leaf args : (scale_c * beta) + {}, // leaf args : C + { // ternary op : (scale_a * scale_b * alpha) * acc + bias + {{scale_a, scale_b, alpha}, + {scale_a_ptr, scale_b_ptr, alpha_ptr} + }, // leaf args : (scale_a * scale_b * alpha) + {}, // leaf args : acc + {bias_ptr, ElementBias(0), dBias + }, // leaf args : bias + {} // ternary args : multiply_add + }, // end ternary op + {} // ternary args : multiply_add + }, // end ternary op + {aux_ptr, dAux} // unary args : store + }, // end unary op + activation // unary args : activation + }, // end unary op + {amax_D_ptr_} // unary args : reduce + }, // end unary op + {{scale_d},{scale_d_ptr}}, // leaf args : scale_d + {} // binary args : multiplies or first + }; // end binary op + } } }; @@ -1048,6 +1164,7 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, class GmemLayoutTagAux, template class ActivationFn, class ElementOutput, @@ -1063,7 +1180,7 @@ template < class CopyOpS2R > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::LinCombDeEltAct< GmemLayoutTagAux, ActivationFn, ElementOutput, ElementCompute, ElementAux, ElementSource, ElementScalar, AlignmentAux, RoundStyle @@ -1146,7 +1263,7 @@ template< > using Sm90LinCombDeEltActDePerRowBias = Sm90EVT, // Identity for final conversion - Sm90EVT, AlignmentBias>, Sm90LinCombDeEltAct @@ -1158,6 +1275,7 @@ template < int StagesD, int FragmentSize, bool ReuseSmemC, + bool DelayTmaStore, class GmemLayoutTagAux, template class ActivationFn, class ElementOutput, @@ -1175,7 +1293,7 @@ template < class CopyOpS2R > struct FusionCallbacks< - epilogue::Sm90TmaWarpSpecialized, + epilogue::Sm90TmaWarpSpecialized, fusion::LinCombDeEltActDePerRowBias< GmemLayoutTagAux, ActivationFn, ElementOutput, ElementCompute, ElementAux, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle diff --git a/include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp b/include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp index f4ebcaff..e3160fa1 100644 --- a/include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp +++ b/include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp @@ -111,19 +111,20 @@ public: template static constexpr Params - to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { + to_underlying_arguments(ProblemShape const&, Arguments const& args, void*) { return args; } template static size_t - get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + get_workspace_size(ProblemShape const&, Arguments const&) { return 0; } template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } @@ -177,7 +178,7 @@ public: ComputeOutput compute_output{}; ConvertOutput convert_output{}; - if constexpr (is_same_v) { + if constexpr (cute::is_same_v) { return convert_output(compute_output(cvt_frg_inputs...)); } else { @@ -211,20 +212,18 @@ template < class ElementOutput, class ElementCompute, FloatRoundStyle RoundStyle, - class ElementScalar, - class StrideScalar, - int ScalarCount, - template class ScalarReduceFn, - class ElementSource, - class InputAddOp // Z + class InputScaleOp, // beta + class ElementSource, // C + class InputAddOp // Z > struct Sm90TreeVisitor< - Sm90Compute, - Sm90ScalarBroadcast, + Sm90Compute().is_zero())>>, + InputScaleOp, Sm90SrcFetch, InputAddOp > : Sm90VisitorImpl< - Sm90ScalarBroadcast, + InputScaleOp, Sm90SrcFetch, InputAddOp, Sm90Compute @@ -232,7 +231,7 @@ struct Sm90TreeVisitor< { using Impl = Sm90VisitorImpl< - Sm90ScalarBroadcast, + InputScaleOp, Sm90SrcFetch, InputAddOp, Sm90Compute @@ -251,18 +250,16 @@ struct Sm90TreeVisitor< CUTLASS_DEVICE bool is_producer_load_needed() const { - auto const& bcast_op = get<0>(Impl::ops); auto const& added_op = get<2>(Impl::ops); - return not (bcast_op.params_ptr->dScalar == Stride<_0,_0,_0>{} && not is_C_load_needed()) || - added_op.is_producer_load_needed(); + return is_C_load_needed() || added_op.is_producer_load_needed(); } CUTLASS_DEVICE bool is_C_load_needed() const { - auto const& bcast_op = get<0>(Impl::ops); + auto const& scale_op = get<0>(Impl::ops); auto const& src_op = get<1>(Impl::ops); auto const& added_op = get<2>(Impl::ops); - return (bcast_op.scalar != 0 && src_op.is_C_load_needed()) || added_op.is_C_load_needed(); + return (not scale_op.is_zero() && src_op.is_C_load_needed()) || added_op.is_C_load_needed(); } template @@ -347,7 +344,8 @@ struct Sm90ReLUAuxStore : Sm90VisitorImpl<> { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } @@ -379,8 +377,8 @@ template < > struct Sm90TreeVisitor< Sm90Compute, cutlass::epilogue::thread::ReLu> || - is_same_v, cutlass::epilogue::thread::Clamp> >>, + cute::enable_if_t, cutlass::epilogue::thread::ReLu> || + cute::is_same_v, cutlass::epilogue::thread::Clamp> >>, Sm90TreeVisitor< Sm90AuxStore< Stages, @@ -474,7 +472,7 @@ struct Sm90TreeVisitor< CUTLASS_PRAGMA_UNROLL for (int i = 0; i < FragmentSize; ++i) { ElementCompute pre_relu = frg_compute[i]; - if constexpr (is_same_v, cutlass::epilogue::thread::Clamp>) { + if constexpr (cute::is_same_v, cutlass::epilogue::thread::Clamp>) { frg_compute[i] = relu(frg_compute[i], params_compute); } else { @@ -606,7 +604,8 @@ struct Sm90AuxLoad< template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } diff --git a/include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp b/include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp index 971d5145..1ea663f6 100644 --- a/include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp +++ b/include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp @@ -98,6 +98,11 @@ struct Sm90SrcFetch : Sm90VisitorImpl<> { return not is_void_v; } + CUTLASS_DEVICE bool + is_zero() const { + return is_void_v; + } + using Sm90VisitorImpl<>::Sm90VisitorImpl; template @@ -158,6 +163,9 @@ struct Sm90AuxLoad { SmemLayoutTma{}, make_shape(size<0>(shape(EpilogueTile{})), size<1>(shape(EpilogueTile{})), Int{}), cute::conditional_t, Step<_1,_2,_3>>{} )); + using CopyOpG2S = + SM90_TMA_LOAD + ; struct SharedStorage { alignas(cutlass::detail::alignment_for_swizzle(SmemLayout{})) @@ -172,9 +180,9 @@ struct Sm90AuxLoad { struct Params { using TMA_Aux = decltype(make_tma_copy( - SM90_TMA_LOAD{}, - make_tensor(static_cast(nullptr), repeat_like(StrideMNL{}, int32_t(0)), StrideMNL{}), - SmemLayoutTma{})); + CopyOpG2S{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), repeat_like(StrideMNL{}, int32_t(0)), append<3>(StrideMNL{}, _0{})), + take<0,2>(SmemLayoutTma{}))); TMA_Aux tma_load_aux; Element null_default = Element(0); bool use_default = false; @@ -186,9 +194,11 @@ struct Sm90AuxLoad { // 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, 1); auto [M, N, K, L] = problem_shape_mnkl; - - Tensor tensor_aux = make_tensor(args.ptr_aux, make_layout(make_shape(M,N,L), args.dAux)); - typename Params::TMA_Aux tma_load_aux = make_tma_copy(SM90_TMA_LOAD{}, tensor_aux, SmemLayoutTma{}); + auto M_AUX = + size(M) + ; + Tensor tensor_aux = make_tensor(make_gmem_ptr(args.ptr_aux), make_layout(make_shape(M_AUX,N,L), append<3>(args.dAux, _0{}))); + typename Params::TMA_Aux tma_load_aux = make_tma_copy(CopyOpG2S{}, tensor_aux, take<0,2>(SmemLayoutTma{})); bool use_default = false; if constexpr (EnableNullptr) { @@ -206,7 +216,8 @@ struct Sm90AuxLoad { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } @@ -231,6 +242,11 @@ struct Sm90AuxLoad { return false; } + CUTLASS_DEVICE bool + is_zero() const { + return (params_ptr->use_default && params_ptr->null_default == Element(0)); + } + template struct ProducerLoadCallbacks : EmptyProducerLoadCallbacks { CUTLASS_DEVICE @@ -253,7 +269,7 @@ struct Sm90AuxLoad { if (issue_tma_load) { // Increment the expected transaction bytes of the current stage's mbarrier by the subtile's byte-size - constexpr uint32_t copy_bytes = size(take<0,2>(SmemLayout{})) * sizeof_bytes_v; + constexpr uint32_t copy_bytes = size(take<0,2>(SmemLayout{})) * sizeof_bits_v / 8; cutlass::arch::ClusterTransactionBarrier::expect_transaction(full_mbarrier_ptr, copy_bytes); // Issue the TMA load constexpr uint16_t mcast_mask = 0; @@ -270,8 +286,12 @@ struct Sm90AuxLoad { auto [M, N, K, L] = args.problem_shape_mnkl; auto [m, n, k, l] = args.tile_coord_mnkl; - Tensor mAux = params_ptr->tma_load_aux.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) - Tensor gAux = local_tile(mAux, take<0,2>(args.tile_shape_mnk), make_coord(m,n,l)); // (CTA_M,CTA_N) + auto coord_shape = + make_coord(m, n, l) + ; + Tensor mAux_mn = params_ptr->tma_load_aux.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) + Tensor mAux = coalesce(mAux_mn, take<0,2>(args.tile_shape_mnk)); + Tensor gAux = local_tile(mAux, take<0,2>(args.tile_shape_mnk), coord_shape); // (CTA_M,CTA_N) Tensor gAux_epi = flat_divide(gAux, args.epi_tile); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) Tensor sAux_epi = make_tensor(make_smem_ptr(smem_aux), SmemLayout{}); // (EPI_TILE_M,EPI_TILE_N,PIPE) @@ -331,9 +351,11 @@ struct Sm90AuxLoad { get_consumer_store_callbacks(ConsumerStoreArgs const& args) { auto [M, N, K, L] = args.problem_shape_mnkl; - Tensor mAux = params_ptr->tma_load_aux.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) - Tensor tC_gAux = sm90_partition_for_epilogue( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) - mAux, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx); + + Tensor mAux_mn = params_ptr->tma_load_aux.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) + Tensor mAux = coalesce(mAux_mn, take<0,2>(args.tile_shape_mnk)); + Tensor tC_gAux = sm90_partition_for_epilogue(mAux, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx); Tensor tC_rAux = make_tensor(take<0,3>(shape(tC_gAux))); // (CPY,CPY_M,CPY_N) auto tiled_s2r = conditional_return( @@ -393,7 +415,8 @@ struct Sm90ScalarBroadcast { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter *cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } @@ -407,6 +430,12 @@ struct Sm90ScalarBroadcast { return false; } + // This must be called after update_scalar is called + CUTLASS_DEVICE bool + is_zero() const { + return scalar == Element(0); + } + CUTLASS_HOST_DEVICE Sm90ScalarBroadcast() { } @@ -414,7 +443,7 @@ struct Sm90ScalarBroadcast { Sm90ScalarBroadcast(Params const& params, SharedStorage const& shared_storage) : params_ptr(¶ms) { // Get the scalar for non-batched broadcast - if constexpr (cute::is_same_v>) { + if (get<2>(params_ptr->dScalar) == 0) { update_scalar(); } } @@ -426,9 +455,7 @@ struct Sm90ScalarBroadcast { CUTLASS_DEVICE auto get_producer_load_callbacks(ProducerLoadArgs const& args) { // Get the scalar for batched broadcast - if constexpr ( - cute::is_same_v> || - cute::is_same_v>) { + if (get<2>(params_ptr->dScalar) != 0) { auto [m_coord, n_coord, k_coord, l_coord] = args.tile_coord_mnkl; update_scalar(l_coord); } @@ -462,9 +489,7 @@ struct Sm90ScalarBroadcast { get_consumer_store_callbacks(ConsumerStoreArgs const& args) { // Get the scalar for batched broadcast - if constexpr ( - cute::is_same_v> || - cute::is_same_v>) { + if (get<2>(params_ptr->dScalar) != 0) { auto [m_coord, n_coord, k_coord, l_coord] = args.tile_coord_mnkl; update_scalar(l_coord); } @@ -496,10 +521,28 @@ private: } } } + + template + CUTLASS_DEVICE void + update_scalar(cute::tuple) { + // Only support multiple L-modes with fully-broadcast scalar + static_assert(cute::is_same_v>); + scalar = params_ptr->scalars[0]; + } }; ///////////////////////////////////////////////////////////////////////////////////////////////// +namespace detail { + +template +constexpr int +compute_row_broadcast_stages() { + return ceil_div(StagesC, size<1>(zipped_divide(make_layout(take<0,2>(CtaTileShapeMNK{})), EpilogueTile{}))) + 1; +} + +} + // Row vector broadcast template< // Row bcast reuses the mbarriers from the epilogue subtile load pipeline, so this must be at least @@ -544,7 +587,8 @@ struct Sm90RowBroadcast { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } @@ -569,6 +613,11 @@ struct Sm90RowBroadcast { return false; } + CUTLASS_DEVICE bool + is_zero() const { + return (params.ptr_row == nullptr && params.null_default == Element(0)); + } + template struct ProducerLoadCallbacks : EmptyProducerLoadCallbacks { CUTLASS_DEVICE @@ -591,7 +640,7 @@ struct Sm90RowBroadcast { if (issue_tma_load) { // Increment the expect-tx count of the first subtile's mbarrier by the row vector's byte-size - constexpr uint32_t copy_bytes = size<1>(CtaTileShapeMNK{}) * sizeof_bytes_v; + constexpr uint32_t copy_bytes = size<1>(CtaTileShapeMNK{}) * sizeof_bits_v / 8; cutlass::arch::ClusterTransactionBarrier::expect_transaction(full_mbarrier_ptr, copy_bytes); // Issue the TMA bulk copy auto bulk_copy = Copy_Atom{}.with(*full_mbarrier_ptr); @@ -614,7 +663,7 @@ struct Sm90RowBroadcast { make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{}), Stages), make_stride(_0{},_1{},size<1>(CtaTileShapeMNK{}))); - constexpr int EpiTiles = decltype(size(shape_div(take<0,2>(args.tile_shape_mnk), args.epi_tile)))::value; + constexpr int EpiTiles = decltype(size<1>(zipped_divide(make_layout(take<0,2>(args.tile_shape_mnk)), args.epi_tile)))::value; return ProducerLoadCallbacks( cute::move(gRow), cute::move(sRow), params); } @@ -676,7 +725,7 @@ struct Sm90RowBroadcast { sRow, args.epi_tile, args.tiled_copy, args.thread_idx); Tensor tCrRow = make_tensor_like(take<0,3>(tCsRow)); // (CPY,CPY_M,CPY_N) - constexpr int EpiTiles = decltype(size(shape_div(take<0,2>(args.tile_shape_mnk), args.epi_tile)))::value; + constexpr int EpiTiles = decltype(size<1>(zipped_divide(make_layout(take<0,2>(args.tile_shape_mnk)), args.epi_tile)))::value; return ConsumerStoreCallbacks( cute::move(tCrRow), cute::move(tCsRow), params); } @@ -725,7 +774,8 @@ struct Sm90ColBroadcast { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } @@ -739,6 +789,11 @@ struct Sm90ColBroadcast { return false; } + CUTLASS_DEVICE bool + is_zero() const { + return (params.ptr_col == nullptr && params.null_default == Element(0)); + } + CUTLASS_HOST_DEVICE Sm90ColBroadcast() { } diff --git a/include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp b/include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp index dbaaaa73..c8d941b6 100644 --- a/include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp +++ b/include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp @@ -132,7 +132,8 @@ struct Sm90AuxStore { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { return cutlass::Status::kSuccess; } @@ -207,7 +208,7 @@ struct Sm90AuxStore { } CUTLASS_DEVICE void - postvisit(int epi_m, int epi_n, int store_iteration, bool issue_smem_store) { + postreduce(int epi_m, int epi_n, int store_iteration, bool issue_smem_store) { if constexpr (EnableNullptr) { if (params_ptr->is_nullptr) { return; @@ -224,7 +225,7 @@ struct Sm90AuxStore { } CUTLASS_DEVICE void - step(int epi_m, int epi_n, int store_iteration, bool issue_tma_store) { + tma_store(int epi_m, int epi_n, int store_iteration, bool issue_tma_store) { if constexpr (EnableNullptr) { if (params_ptr->is_nullptr) { return; @@ -329,12 +330,13 @@ public: template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { if constexpr (IsAtomic) { auto [M, N, K, L] = problem_shape; Layout mScalar_layout = make_layout(make_shape(M,N,L), args.dScalar); if (args.ptr_scalar != nullptr) { - return fill_workspace(args.ptr_scalar, ElementOutput(args.reduction_identity), cosize(mScalar_layout), stream); + return fill_workspace(args.ptr_scalar, ElementOutput(args.reduction_identity), cosize(mScalar_layout), stream, cuda_adapter); } } @@ -452,6 +454,7 @@ public: // Row vector reduction template < template class RegReduceFn, + template class ShuffleReduceFn, template class GmemReduceFn, int Stages, class CtaTileShapeMNK, @@ -460,7 +463,12 @@ template < FloatRoundStyle RoundStyle, class StrideMNL = Stride<_0,_1,_0>, int Alignment = 128 / sizeof_bits_v, - bool EnableNullptr = true // Noop on nullptr params + bool EnableNullptr = true, // Noop on nullptr params + // If this is false, ptr_row is assumed to point to a compact n-major (ceil_div(M,CTA_M), round_nearest(N,CTA_N), L) + // tensor of ElementCompute. It is the user's responsibility to reduce this to a (N, L) tensor of ElementOutput + bool FinalReduction = true, + // False means skip OOB predication if OOB inputs are known to be the reduction identity + bool VisitCheckOOB = true > struct Sm90RowReduction { private: @@ -470,43 +478,92 @@ private: (cute::is_same_v>) || // row vector reduction, e.g. per-col sum over all batches (cute::is_same_v>)); // batched row vector reduction, e.g. per-col sum per batch static constexpr bool IsAtomic = is_atomic>::value; - static_assert(IsAtomic, "non-atomic row reduction not supported yet"); + static_assert(not (IsAtomic && not FinalReduction), "atomic reduction must be final"); public: struct SharedStorage { }; struct Arguments { - ElementOutput* ptr_row = nullptr; + void* ptr_row = nullptr; // ElementOutput* if FinalReduction, else ElementCompute* ElementCompute reduction_identity = 0; StrideMNL dRow = {}; }; - using Params = Arguments; + struct Params { + void* ptr_row = nullptr; + ElementCompute reduction_identity = 0; + StrideMNL dRow = {}; + ElementCompute* reduction_buffer = nullptr; + int* tile_counters = nullptr; + }; template static constexpr Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { - return args; + ElementCompute* reduction_buffer; + int* tile_counters = nullptr; + if constexpr (IsAtomic) { + reduction_buffer = nullptr; + } + else if constexpr (not FinalReduction) { + reduction_buffer = reinterpret_cast(args.ptr_row); + } + else { + auto [M, N, K, L] = problem_shape; + auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{}; + size_t tile_counters_offset = product(ceil_div(make_shape(size<>(M), size<>(N), L), make_shape(tile_M, tile_N))) * tile_N * sizeof(ElementCompute); + tile_counters_offset = round_nearest(tile_counters_offset, sizeof(int)); + + reduction_buffer = reinterpret_cast(workspace); + tile_counters = reinterpret_cast(reinterpret_cast(workspace) + tile_counters_offset); + } + + return { + args.ptr_row, + args.reduction_identity, + args.dRow, + reduction_buffer, + tile_counters + }; } template static size_t get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { - return 0; + if constexpr (IsAtomic || not FinalReduction) { + return 0; + } + + size_t workspace_size = 0; + auto [M, N, K, L] = problem_shape; + auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{}; + // Increment by size of reduction buffer + workspace_size += product(ceil_div(make_shape(size<>(M),size<>(N),L), make_shape(tile_M, tile_N))) * tile_N * sizeof(ElementCompute); + // Align and increment by size of tile counters + workspace_size = round_nearest(workspace_size, sizeof(int)); + workspace_size += cute::ceil_div(size<>(N), tile_N) * sizeof(int); + return workspace_size; } template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { if constexpr (IsAtomic) { auto [M, N, K, L] = problem_shape; Layout mRow_layout = make_layout(make_shape(M,N,L), args.dRow); if (args.ptr_row != nullptr) { - return fill_workspace(args.ptr_row, ElementOutput(args.reduction_identity), cosize(mRow_layout), stream); + return fill_workspace(args.ptr_row, ElementOutput(args.reduction_identity), cosize(mRow_layout), stream, cuda_adapter); } + return Status::kSuccess; } + auto [M, N, K, L] = problem_shape; + auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{}; + size_t tile_counters_offset = product(ceil_div(make_shape(size<>(M),size<>(N),L), make_shape(tile_M, tile_N))) * tile_N * sizeof(ElementCompute); - return cutlass::Status::kSuccess; + int* tile_counters = reinterpret_cast(reinterpret_cast(workspace) + tile_counters_offset); + size_t tile_counters_size = cute::ceil_div(size<>(N), tile_N) * sizeof(int); + return zero_workspace(tile_counters, tile_counters_size, stream); } CUTLASS_DEVICE bool @@ -534,53 +591,49 @@ public: return EmptyProducerLoadCallbacks{}; } - template + template struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks { CUTLASS_DEVICE - ConsumerStoreCallbacks( - RTensor&& tCrRow, - GTensor&& tCgRow, - CTensor tCcRow, - ResidueMN residue_mn, - Params const& params) - : tCrRow(cute::forward(tCrRow)), - tCgRow(cute::forward(tCgRow)), - tCcRow(tCcRow), - residue_mn(residue_mn), + ConsumerStoreCallbacks(ArgsTuple&& args_tuple, Params const& params) + : args_tuple(cute::forward(args_tuple)), params(params) {} - // gmem store after every column of subtiles, assuming M-major loop - // needed to reduce reg pressure, otherwise each thread stores up to a full row in RF - // since row-elements aren't evenly distributed amongst threads - RTensor tCrRow; // (CPY,CPY_M,CPY_N,EPI_M) - GTensor tCgRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) - CTensor tCcRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) - ResidueMN residue_mn; + ArgsTuple args_tuple; Params const& params; + bool do_final_reduction = false; + template CUTLASS_DEVICE auto visit(Array const& frg_acc, int epi_v, int epi_m, int epi_n, Array const& frg_input) { - if constexpr (EnableNullptr) { if (params.ptr_row == nullptr) { return frg_input; } } + auto& [ref_src, tCrRow, tCcRow, gRow_l, cRow, gBuf_ml, sBuf_layout, + lane_layout_MN, lane_mn, warp_layout_MN, warp_mn, + tile_coord_mnkl, residue_mn, epi_tile, tiled_copy, thread_idx] = args_tuple; + Tensor tCrRow_mn = tCrRow(_,_,_,epi_m,epi_n); + Tensor tCcRow_mn = tCcRow(_,_,_,epi_m,epi_n); + using ConvertInput = NumericArrayConverter; using ReduceInput = RegReduceFn; ConvertInput convert_input{}; ReduceInput reduce_input{}; Array frg_I = convert_input(frg_input); - Tensor tCrRow_mn = tCrRow(_,_,_,epi_m); - Tensor tCcRow_mn = tCcRow(_,_,_,epi_m,epi_n); - CUTLASS_PRAGMA_UNROLL for (int i = 0; i < FragmentSize; ++i) { - if (elem_less(tCcRow_mn(epi_v * FragmentSize + i), residue_mn)) { + if constexpr (VisitCheckOOB) { + if (elem_less(tCcRow_mn(epi_v * FragmentSize + i), residue_mn)) { + ElementCompute& tCrRow_vmn = tCrRow_mn(epi_v * FragmentSize + i); + tCrRow_vmn = reduce_input(tCrRow_vmn, frg_I[i]); + } + } + else { ElementCompute& tCrRow_vmn = tCrRow_mn(epi_v * FragmentSize + i); tCrRow_vmn = reduce_input(tCrRow_vmn, frg_I[i]); } @@ -589,45 +642,215 @@ public: return frg_input; } + template CUTLASS_DEVICE void - step(int epi_m, int epi_n, int store_iteration, bool issue_tma_store) { + reduce(STensor&& smem_buffer, SyncFn const& sync_fn, int epi_m, int epi_n, bool is_last_iteration) { + if (not is_last_iteration) { + return; + } + + auto& [ref_src, tCrRow, tCcRow, gRow_l, cRow, gBuf_ml, sBuf_layout, + lane_layout_MN, lane_mn, warp_layout_MN, warp_mn, + tile_coord_mnkl, residue_mn, epi_tile, tiled_copy, thread_idx] = args_tuple; + auto [m, n, k, l] = tile_coord_mnkl; + constexpr bool ReferenceSrc = decltype(ref_src)::value; if constexpr (EnableNullptr) { if (params.ptr_row == nullptr) { return; } } - if (epi_m == size<3>(tCrRow)-1) { // assumes M-major subtile loop - using ConvertI = NumericConverter; - using ReduceInput = GmemReduceFn; + // fully OOB CTA in partially OOB cluster + if (not elem_less(cRow(_0{},_0{}), residue_mn)) { + return; + } - ConvertI convert_I{}; - ReduceInput reduce_input{}; - - // Filter so we don't issue redunant copies over stride-0 modes - Tensor tCrRow_flt = filter_zeros(tCrRow(_,_,_,epi_m)); - Tensor tCgRow_flt = filter_zeros(tCgRow(_,_,_,epi_m,epi_n)); - Tensor tCcRow_mn = tCcRow(_,_,_,epi_m,epi_n); - Tensor tCcRow_flt = make_tensor(tCcRow_mn.data(), make_layout(tCgRow_flt.shape(), tCcRow_mn.stride())); - - - auto [residue_m, residue_n] = residue_mn; + // + // 1. Warp shuffle reduction + // + using FragmentShuffle = Array; + using ReduceShuffle = ShuffleReduceFn; + ReduceShuffle reduce_shuffle{}; + Tensor tCrRow_frg = recast(filter(tCrRow)); + CUTLASS_PRAGMA_UNROLL + for (int reduction_rows = size<0>(lane_layout_MN) / 2; reduction_rows > 0; reduction_rows /= 2) { CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(tCrRow_flt); ++i) { - // partially OOB in M must still issue gmem reduction, so only consider residue_n - // in case last epi tile in column is fully OOB in M and CTA tile is partially OOB in M - if (residue_n > get<1>(tCcRow_flt(i)) && - // fully OOB in M does not need to issue gmem reduction, skip - residue_m > 0) { - reduce_input(&tCgRow_flt(i), convert_I(tCrRow_flt(i))); + for (int frg_idx = 0; frg_idx < size(tCrRow_frg); ++frg_idx) { + uint64_t frg_shfl = reinterpret_cast(tCrRow_frg(frg_idx)); + frg_shfl = __shfl_down_sync(0xFFFFFFFF, frg_shfl, lane_layout_MN(reduction_rows, _0{})); + tCrRow_frg(frg_idx) = reduce_shuffle(tCrRow_frg(frg_idx), reinterpret_cast(frg_shfl)); + } + } + bool is_reduced_lane = get<0>(lane_mn) == 0; + + // + // 2. Atomic reduction + // + if constexpr (IsAtomic) { + // Filter so we don't issue redunant copies over stride-0 modes + Tensor tCrRow_flt = filter_zeros(tCrRow); + Tensor tCcRow_flt = make_tensor(tCcRow.data(), make_layout(tCrRow_flt.shape(), tCcRow.stride())); + + Tensor tCgRow = sm90_partition_for_epilogue(gRow_l(_,_,l), epi_tile, tiled_copy, thread_idx); + Tensor tCgRow_flt = filter_zeros(tCgRow); + // NOTE: atomic reduction is performed in the output type + using ConvertOutput = NumericConverter; + using ReduceOutput = GmemReduceFn; + ConvertOutput convert_output{}; + ReduceOutput reduce_output{}; + + if (is_reduced_lane) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tCrRow_flt); ++i) { + if (elem_less(tCcRow_flt(i), residue_mn)) { + reduce_output(&tCgRow_flt(i), convert_output(tCrRow_flt(i))); + } } } + sync_fn(); + } - // Reset the registers to the reduction identity - fill(tCrRow, params.reduction_identity); + // + // 2. One warp in M, skip threadblock smem reduction + // + else if constexpr (decltype(size<0>(warp_layout_MN))::value <= 1) { + // Dump warp reduction to gmem workspace + using ElementGmem = cute::conditional_t; + Tensor tCgBuf = sm90_partition_for_epilogue(gBuf_ml(_,_,m,l), epi_tile, tiled_copy, thread_idx); + if (is_reduced_lane) { + // Filter so we don't issue redundant copies over stride-0 modes + // (only works if 0-strides are in same location, which is by construction) + copy_aligned(filter(tCrRow), recast(filter(tCgBuf))); + } + sync_fn(); + } + + // + // 2. Multiple warps in M, do threadblock smem reduction + // + else { + Tensor sBuf = make_tensor(make_smem_ptr(raw_pointer_cast(smem_buffer.data())), sBuf_layout); + static_assert(decltype(cosize(sBuf.layout()))::value * sizeof(ElementCompute) <= + decltype(cosize(smem_buffer.layout()))::value * sizeof(typename remove_cvref_t::value_type), + "smem reduction buffer not large enough, use a larger epilogue tile"); + + // Dump warp reduction to smem workspace + Tensor tCsBuf = sm90_partition_for_epilogue(sBuf(_,_,get<0>(warp_mn)), epi_tile, tiled_copy, thread_idx); + if (is_reduced_lane) { + // Filter so we don't issue redunant copies over stride-0 modes + // (only works if 0-strides are in same location, which is by construction) + copy_aligned(filter(tCrRow), filter(tCsBuf)); + } + sync_fn(); + + constexpr int SmemFragSize = cute::max(size_t{1}, sizeof(uint32_t) / sizeof(ElementCompute)); + using FragmentSmem = Array; + using VectorSmem = uint_bit_t>; + using ReduceSmem = GmemReduceFn; + ReduceSmem reduce_smem{}; + + Tensor sBuf_frg = recast(filter_zeros(sBuf)); + Tensor sBuf_vec = recast(filter_zeros(sBuf)); + constexpr int FragsPerRow = decltype(size<1>(sBuf_frg))::value; + + // Do the threadblock smem reduction + CUTLASS_PRAGMA_UNROLL + for (int reduction_rows = size<0>(warp_layout_MN) / 2; reduction_rows > 1; reduction_rows /= 2) { + int FragsPerReduction = reduction_rows * FragsPerRow; + CUTLASS_PRAGMA_NO_UNROLL + for (int frg_idx = thread_idx; frg_idx < FragsPerReduction; frg_idx += size(tiled_copy)) { + FragmentSmem frg_smem = reduce_smem(sBuf_frg(frg_idx), sBuf_frg(frg_idx + FragsPerReduction)); + sBuf_vec(frg_idx) = reinterpret_cast(frg_smem); + } + sync_fn(); + } + + // Do final smem reduction and dump to gmem workspace + using VectorGmem = cute::conditional_t; + Tensor gBuf_vec = recast(filter(gBuf_ml(_,_,m,l))); + CUTLASS_PRAGMA_NO_UNROLL + for (int frg_idx = thread_idx; frg_idx < FragsPerRow; frg_idx += size(tiled_copy)) { + FragmentSmem frg_smem = reduce_smem(sBuf_frg(frg_idx), sBuf_frg(frg_idx + FragsPerRow)); + gBuf_vec(frg_idx) = reinterpret_cast(frg_smem); + } + sync_fn(); + } + + // + // 3. Increment atomic counters to signal final gmem reduction + // + if constexpr (not IsAtomic && FinalReduction) { + // Ensure gmem writes are visible to other threads before incrementing counter + __threadfence(); + sync_fn(); + // Collective thread 0 increments atomic tile counter and copies value to smem + int* prev_tile_count = reinterpret_cast(raw_pointer_cast(smem_buffer.data())); + if (thread_idx == 0) { + *prev_tile_count = atomicAdd(¶ms.tile_counters[n], 1); + } + sync_fn(); + // Broadcast tile count to other threads in CTA and determine final reduction status + do_final_reduction = *prev_tile_count == size<2>(gBuf_ml) * size<3>(gBuf_ml) - 1; + sync_fn(); } } + CUTLASS_DEVICE void + end() { + // + // 4. Do final gmem reduction if necessary + // + if constexpr (not IsAtomic && FinalReduction) { + if (not do_final_reduction) { + return; + } + + auto& [ref_src, tCrRow, tCcRow, gRow_l, cRow, gBuf_ml, sBuf_layout, + lane_layout_MN, lane_mn, warp_layout_MN, warp_mn, + tile_coord_mnkl, residue_mn, epi_tile, tiled_copy, thread_idx] = args_tuple; + + using ReduceOutput = GmemReduceFn; + using ConvertOutput = NumericConverter; + ReduceOutput reduce_output{}; + ConvertOutput convert_output{}; + + // Reduction over batches + if (size<2>(stride(gRow_l)) == 0) { + CUTLASS_PRAGMA_NO_UNROLL + for (int n = thread_idx; n < size<1>(gBuf_ml); n += size(tiled_copy)) { + Tensor tRgBuf_ml = gBuf_ml(_0{},n,_,_); + ElementCompute output = tRgBuf_ml(_0{}); + CUTLASS_PRAGMA_NO_UNROLL + for (int ml = 1; ml < size(tRgBuf_ml); ++ml) { + output = reduce_output(output, tRgBuf_ml(ml)); + } + if (elem_less(cRow(_0{},n), residue_mn)) { + gRow_l(_0{},n,_0{}) = convert_output(output); + } + } + } + // No reduction over batches + else { + CUTLASS_PRAGMA_NO_UNROLL + for (int n = thread_idx; n < size<1>(gBuf_ml); n += size(tiled_copy)) { + bool do_store = elem_less(cRow(_0{},n), residue_mn); + CUTLASS_PRAGMA_NO_UNROLL + for (int l = 0; l < size<3>(gBuf_ml); ++l) { + Tensor tRgBuf_m = gBuf_ml(_0{},n,_,l); + ElementCompute output = tRgBuf_m(_0{}); + CUTLASS_PRAGMA_NO_UNROLL + for (int m = 1; m < size(tRgBuf_m); ++m) { + output = reduce_output(output, tRgBuf_m(m)); + } + if (do_store) { + gRow_l(_0{},n,l) = convert_output(output); + } + } + } + } + + } + } }; template < @@ -636,16 +859,59 @@ public: > CUTLASS_DEVICE auto get_consumer_store_callbacks(ConsumerStoreArgs const& args) { + Layout ref_layout_MN = [&] () { + if constexpr (ReferenceSrc) { return get<0>(args.tiled_copy.get_layoutS_MN()); } + else { return get<0>(args.tiled_copy.get_layoutD_MN()); } + }(); // tile_mn -> tv_idx + // Get the MN layout + coord of lanes to determine shuffle reduction iterations + using _W = Int; + Layout tv2lane = Layout,_W,_1>,Stride<_1,_0,_0>>{}; // tv_idx -> lane_idx + Layout ref2lane = composition(tv2lane, ref_layout_MN); // tile_mn -> lane_idx + Layout lane_layout_MN = make_layout(filter(get<0>(ref2lane)), filter(get<1>(ref2lane))); // lane_mn -> lane_idx + Layout inv_lane_layout_MN = right_inverse(lane_layout_MN); // lane_idx -> lane_mn + int lane_idx = canonical_lane_idx(); + auto lane_mn = idx2crd(inv_lane_layout_MN(lane_idx), shape(lane_layout_MN)); + + // Get the MN layout + coord of warps to determine smem reduction iterations + Layout tv2warp = Layout,_W,_1>,Stride<_0,_1,_0>>{}; // tv_idx -> warp_idx + Layout ref2warp = composition(tv2warp, ref_layout_MN); // tile_mn -> warp_idx + Layout warp_layout_MN = make_layout(filter(get<0>(ref2warp)), filter(get<1>(ref2warp))); // warp_mn -> warp_idx + Layout inv_warp_layout_MN = right_inverse(warp_layout_MN); // warp_idx -> warp_mn + + int warp_idx = args.thread_idx / NumThreadsPerWarp; + auto warp_mn = idx2crd(inv_warp_layout_MN(warp_idx), shape(warp_layout_MN)); + + // Partition output gmem and register tensors + auto [tile_M, tile_N, tile_K] = args.tile_shape_mnk; auto [M, N, K, L] = args.problem_shape_mnkl; - Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow); + auto [m, n, k, l] = args.tile_coord_mnkl; + + Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow); // (M,N,L) + Tensor gRow_l = local_tile(mRow, take<0,2>(args.tile_shape_mnk), make_coord(m,n,_)); // (CTA_M,CTA_N,L) Tensor tCgRow = sm90_partition_for_epilogue( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) - mRow, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx); - Tensor tCrRow = make_tensor_like(tCgRow(_,_,_,_,_0{})); // (CPY,CPY_M,CPY_N,EPI_M) + gRow_l(_,_,l), args.epi_tile, args.tiled_copy, args.thread_idx); + Tensor tCrRow = make_tensor_like(tCgRow); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) + fill(tCrRow, params.reduction_identity); - return ConsumerStoreCallbacks( - cute::move(tCrRow), cute::move(tCgRow), args.tCcD, args.residue_mn, params); + // Partition gmem+smem reduction buffer tensors + Layout gBuf_layout = make_layout(take<0,2>(args.tile_shape_mnk), make_stride(_0{}, _1{})); + auto block_shape = ceil_div(make_shape(M,N,L), shape(gBuf_layout)); // (M_CNT, N_CNT, L_CNT) + + // Let the M_CNT (the num of partial reduction results) become the outer mode + Layout block_layout = make_layout(block_shape, make_stride(get<1>(block_shape), _1{}, get<0>(block_shape) * get<1>(block_shape))); + Layout mBuf_layout = blocked_product(gBuf_layout, block_layout); + Tensor mBuf = make_tensor(make_gmem_ptr(params.reduction_buffer), mBuf_layout); // (ceil_M,ceil_N,L) + Tensor gBuf_ml = local_tile(mBuf, take<0,2>(args.tile_shape_mnk), make_coord(_,n,_)); // (CTA_M,CTA_N,REST_M,L) + Layout sBuf_layout = blocked_product(gBuf_layout, // (CTA_M,CTA_N,WARPS_M) + make_layout(make_shape(_1{},_1{},size<0>(warp_layout_MN)))); + + auto args_tuple = make_tuple( + bool_constant{}, cute::move(tCrRow), args.tCcD, gRow_l, args.cD, gBuf_ml, sBuf_layout, + lane_layout_MN, lane_mn, warp_layout_MN, warp_mn, + args.tile_coord_mnkl, args.residue_mn, args.epi_tile, args.tiled_copy, args.thread_idx); + return ConsumerStoreCallbacks(cute::move(args_tuple), params); } }; @@ -654,6 +920,7 @@ public: // Col vector reduction template < template class RegReduceFn, + template class ShuffleReduceFn, template class GmemReduceFn, int Stages, class CtaTileShapeMNK, @@ -665,7 +932,9 @@ template < bool EnableNullptr = true, // Noop on nullptr params // If this is false, ptr_col is assumed to point to a compact m-major (round_nearest(M,CTA_M), ceil_div(N,CTA_N), L) // tensor of ElementCompute. It is the user's responsibility to reduce this to a (M, L) tensor of ElementOutput - bool FinalReduction = true + bool FinalReduction = true, + // False means skip OOB predication if OOB inputs are known to be the reduction identity + bool VisitCheckOOB = true > struct Sm90ColReduction { private: @@ -698,14 +967,12 @@ public: static constexpr Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { ElementCompute* reduction_buffer; - int* tile_counters; + int* tile_counters = nullptr; if constexpr (IsAtomic) { reduction_buffer = nullptr; - tile_counters = nullptr; } else if constexpr (not FinalReduction) { reduction_buffer = reinterpret_cast(args.ptr_col); - tile_counters = nullptr; } else { auto [M, N, K, L] = problem_shape; @@ -748,12 +1015,13 @@ public: template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { if constexpr (IsAtomic) { auto [M, N, K, L] = problem_shape; Layout mCol_layout = make_layout(make_shape(M,N,L), args.dCol); if (args.ptr_col != nullptr) { - return fill_workspace(args.ptr_col, ElementOutput(args.reduction_identity), cosize(mCol_layout), stream); + return fill_workspace(args.ptr_col, ElementOutput(args.reduction_identity), cosize(mCol_layout), stream, cuda_adapter); } return Status::kSuccess; } @@ -828,9 +1096,17 @@ public: Array frg_I = convert_input(frg_input); CUTLASS_PRAGMA_UNROLL for (int i = 0; i < FragmentSize; ++i) { - if (elem_less(tCcCol_mn(epi_v * FragmentSize + i), residue_mn)) { - ElementCompute& tCrCol_vmn = tCrCol_mn(epi_v * FragmentSize + i); - tCrCol_vmn = reduce_input(tCrCol_vmn, frg_I[i]); + if constexpr (VisitCheckOOB) { + if (elem_less(tCcCol_mn(epi_v * FragmentSize + i), residue_mn)) { + ElementCompute& tCrCol_vmn = tCrCol_mn(epi_v * FragmentSize + i); + tCrCol_vmn = reduce_input(tCrCol_vmn, frg_I[i]); + } + } + else { + if (elem_less(tCcCol_mn(epi_v * FragmentSize + i), residue_mn)) { + ElementCompute& tCrCol_vmn = tCrCol_mn(epi_v * FragmentSize + i); + tCrCol_vmn = reduce_input(tCrCol_vmn, frg_I[i]); + } } } @@ -866,7 +1142,7 @@ public: // 1. Warp shuffle reduction // using FragmentShuffle = Array; - using ReduceShuffle = RegReduceFn; + using ReduceShuffle = ShuffleReduceFn; ReduceShuffle reduce_shuffle{}; Tensor tCrCol_frg = recast(filter(tCrCol)); CUTLASS_PRAGMA_UNROLL @@ -913,7 +1189,7 @@ public: // else if constexpr (decltype(size<1>(warp_layout_MN))::value <= 1) { // Dump warp reduction to gmem workspace - using ElementGmem = conditional_t; + using ElementGmem = cute::conditional_t; Tensor tCgBuf = sm90_partition_for_epilogue(gBuf_nl(_,_,n,l), epi_tile, tiled_copy, thread_idx); if (is_reduced_lane) { // Filter so we don't issue redundant copies over stride-0 modes @@ -941,7 +1217,7 @@ public: } sync_fn(); - constexpr int SmemFragSize = cute::max(1, sizeof(uint32_t) / sizeof(ElementCompute)); + constexpr int SmemFragSize = cute::max(size_t{1}, sizeof(uint32_t) / sizeof(ElementCompute)); using FragmentSmem = Array; using VectorSmem = uint_bit_t>; using ReduceSmem = GmemReduceFn; @@ -964,7 +1240,7 @@ public: } // Do final smem reduction and dump to gmem workspace - using VectorGmem = conditional_t; + using VectorGmem = cute::conditional_t; Tensor gBuf_vec = recast(filter(gBuf_nl(_,_,n,l))); CUTLASS_PRAGMA_NO_UNROLL for (int frg_idx = thread_idx; frg_idx < FragsPerCol; frg_idx += size(tiled_copy)) { @@ -1050,18 +1326,6 @@ public: } } - CUTLASS_DEVICE bool - is_reduction_buffer_needed(int epi_m, int epi_n, bool is_last_iteration) const { - auto const& [ref_src, tCrCol, tCcCol, gCol_l, cCol, gBuf_nl, sBuf_layout, - lane_layout_MN, lane_mn, warp_layout_MN, warp_mn, - tile_coord_mnkl, residue_mn, epi_tile, tiled_copy, thread_idx] = args_tuple; - - return (not IsAtomic && // atomic reduction doesn't use smem - is_last_iteration && // smem reduction happens after epilogue loop - (decltype(size<1>(warp_layout_MN))::value > 1 || // smem reduction happens when multiple warps are in N - FinalReduction)); // smem is used to broadcast tile counters for final reduction - } - }; template < @@ -1096,6 +1360,7 @@ public: auto [tile_M, tile_N, tile_K] = args.tile_shape_mnk; auto [M, N, K, L] = args.problem_shape_mnkl; auto [m, n, k, l] = args.tile_coord_mnkl; + Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol); // (M,N,L) Tensor gCol_l = local_tile(mCol, take<0,2>(args.tile_shape_mnk), make_coord(m,n,_)); // (CTA_M,CTA_N,L) Tensor tCgCol = sm90_partition_for_epilogue( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) diff --git a/include/cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp b/include/cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp index c5700e03..1e07cc89 100644 --- a/include/cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp +++ b/include/cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp @@ -99,7 +99,10 @@ sm90_partition_for_epilogue( TiledCopy tiled_copy, int thread_idx) { auto [m, n, k, l] = tile_coord_mnkl; - Tensor cT = local_tile(mT, take<0,2>(tile_shape_mnk), make_coord(m,n,l)); // (CTA_M,CTA_N) + auto coord_shape = + make_coord(m, n, l) + ; + Tensor cT = local_tile(mT, take<0,2>(tile_shape_mnk), coord_shape); // (CTA_M,CTA_N) Tensor tCcT = sm90_partition_for_epilogue(cT, epi_tile, tiled_copy, thread_idx); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) @@ -134,7 +137,7 @@ struct ProducerLoadArgs { TileCoordMNKL tile_coord_mnkl, ResidueMN residue_mn, EpilogueTile epi_tile, - int thread_idx) + int thread_idx) : problem_shape_mnkl(problem_shape_mnkl), tile_shape_mnk(tile_shape_mnk), tile_coord_mnkl(tile_coord_mnkl), @@ -177,7 +180,7 @@ struct ConsumerStoreArgs { int thread_idx, CoordTensor cD, ThrCoordTensor tCcD, - ThrSrcTensor const& tCrC) + ThrSrcTensor const& tCrC) : problem_shape_mnkl(problem_shape_mnkl), tile_shape_mnk(tile_shape_mnk), tile_coord_mnkl(tile_coord_mnkl), @@ -202,10 +205,16 @@ struct Sm90VisitorImplBase { template static constexpr Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { + uint8_t* op_workspace = reinterpret_cast(workspace); return transform_apply(tuple{}, args, [&] (auto&& op, auto const& op_args) { using Op = cute::remove_cvref_t; - return Op::to_underlying_arguments(problem_shape, op_args, workspace); + auto ret = Op::to_underlying_arguments(problem_shape, op_args, op_workspace); + if (op_workspace != nullptr) { + size_t op_workspace_size = Op::get_workspace_size(problem_shape, op_args); + op_workspace += round_nearest(op_workspace_size, MinWorkspaceAlignment); + } + return ret; }, [] (auto&&... op_params) { return cute::make_tuple(op_params...); } ); @@ -228,7 +237,8 @@ struct Sm90VisitorImplBase { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { Status status = Status::kSuccess; uint8_t* op_workspace = reinterpret_cast(workspace); return transform_apply(tuple{}, args, @@ -239,7 +249,7 @@ struct Sm90VisitorImplBase { } using Op = cute::remove_cvref_t; - status = Op::initialize_workspace(problem_shape, op_args, op_workspace, stream); + status = Op::initialize_workspace(problem_shape, op_args, op_workspace, stream, cuda_adapter); if (op_workspace != nullptr) { size_t op_workspace_size = Op::get_workspace_size(problem_shape, op_args); op_workspace += round_nearest(op_workspace_size, MinWorkspaceAlignment); @@ -247,7 +257,7 @@ struct Sm90VisitorImplBase { return status; }, // Return the final status - [&] (auto const&...) { return status; } + [&] (auto const&...ops) { return status; } ); } @@ -296,7 +306,7 @@ struct Sm90VisitorImpl : Sm90VisitorImplBase { // e.g. for batched beta this must always be true regardless of current batch idx CUTLASS_DEVICE bool is_producer_load_needed() const { - return apply(ops, + return cute::apply(ops, [] (auto const&... op) { return (false || ... || op.is_producer_load_needed()); } @@ -310,7 +320,7 @@ struct Sm90VisitorImpl : Sm90VisitorImplBase { // e.g. for batched beta this can be false depending on current batch idx CUTLASS_DEVICE bool is_C_load_needed() const { - return apply(ops, + return cute::apply(ops, [] (auto const&... op) { return (false || ... || op.is_C_load_needed()); } @@ -417,29 +427,7 @@ struct Sm90VisitorImpl : Sm90VisitorImplBase { Array const&... frg_inputs) // depends on the N-naryness of the op = delete; // Must be implemented for each operation - // After visit call, before smem async fence. Smem stores usually performed here. - // Upon exit, all smem stores for TMA must have been issued - CUTLASS_DEVICE void - postvisit(int epi_m, int epi_n, int store_iteration, bool issue_smem_store) { - for_each(callbacks_tuple, - [&] (auto& callbacks) { - callbacks.postvisit(epi_m, epi_n, store_iteration, issue_smem_store); - } - ); - } - - // After async fence, before TMA store commit. Aux stores usually performed here - // Upon exit, all TMA stores for this subtile must have been issued - CUTLASS_DEVICE void - step(int epi_m, int epi_n, int store_iteration, bool issue_tma_store) { - for_each(callbacks_tuple, - [&] (auto& callbacks) { - callbacks.step(epi_m, epi_n, store_iteration, issue_tma_store); - } - ); - } - - // After TMA store commit. Smem reductions usually performed here + // After visit call. Smem reductions usually performed here // reduction_buffer is an arbitrary smem tensor that can be used for workspace // It is each nodes reponsibility to assert that this buffer is sufficiently sized // and to ensure that this buffer is no longer needed upon callback exit @@ -454,12 +442,26 @@ struct Sm90VisitorImpl : Sm90VisitorImplBase { ); } - // Collective can query this to determine whether a buffer needs to be freed for reduction - CUTLASS_DEVICE bool - is_reduction_buffer_needed(int epi_m, int epi_n, bool is_last_iteration) const { - return apply(callbacks_tuple, - [&] (auto const&... callbacks) { - return (false || ... || callbacks.is_reduction_buffer_needed(epi_m, epi_n, is_last_iteration)); + // After reduce call, before smem async fence. Smem stores usually performed here. + // Upon exit, all smem stores for TMA must have been issued + CUTLASS_DEVICE void + postreduce(int epi_m, int epi_n, int store_iteration, bool issue_smem_store) { + for_each(callbacks_tuple, + [&] (auto& callbacks) { + callbacks.postreduce(epi_m, epi_n, store_iteration, issue_smem_store); + } + ); + } + + // After smem async fence, before TMA store commit. Aux stores usually performed here + // Upon exit, all TMA stores for this subtile must have been issued + // Because of the TMA store delay optimization, this entry point must ONLY be used for TMA stores + // other gmem stores can be placed in the reduce or postreduce entry points + CUTLASS_DEVICE void + tma_store(int epi_m, int epi_n, int store_iteration, bool issue_tma_store) { + for_each(callbacks_tuple, + [&] (auto& callbacks) { + callbacks.tma_store(epi_m, epi_n, store_iteration, issue_tma_store); } ); } @@ -592,10 +594,9 @@ struct Sm90SplitTreeVisitor : Sm90VisitorImpl(callbacks_tuple).visit(frg_acc, epi_v, epi_m, epi_n); constexpr int Rm2 = sizeof...(AuxOutTrees); - cute::detail::for_sequence(make_seq{}, // restrict the sequence to aux out trees - [&] (auto&& _I) { - constexpr int i = remove_cvref_t::value; - get(callbacks_tuple).visit(frg_input, epi_v, epi_m, epi_n); + cute::for_each(make_seq{}, // restrict the sequence to aux out trees + [&] (auto I) { + get(callbacks_tuple).visit(frg_input, epi_v, epi_m, epi_n); } ); @@ -664,7 +665,7 @@ struct Sm90TopologicalVisitor : Sm90VisitorImpl { return frg_compute; // unused }, // Visit the last op - [&] (auto const&...) { + [&] (auto const&...ops) { return cute::detail::apply(frg_compute_tuple, // Compute the last op with children inputs [&] (auto const&... frg_inputs) { @@ -734,12 +735,13 @@ struct Sm90VisitorImplBase { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { Status status = Status::kSuccess; uint8_t* workspace_ptr = reinterpret_cast(workspace); size_t workspace_offset = 0; - status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream); + status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream, cuda_adapter); workspace_offset += Op0::get_workspace_size(problem_shape, args.op_0); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { @@ -782,9 +784,12 @@ struct Sm90VisitorImplBase { template static constexpr Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { + size_t op_0_workspace_size = Op0::get_workspace_size(problem_shape, args.op_0); + uint8_t* op_0_workspace = reinterpret_cast(workspace); + uint8_t* op_1_workspace = op_0_workspace + op_0_workspace_size; return Params{ - Op0::to_underlying_arguments(problem_shape, args.op_0, workspace), - Op1::to_underlying_arguments(problem_shape, args.op_1, workspace) + Op0::to_underlying_arguments(problem_shape, args.op_0, op_0_workspace), + Op1::to_underlying_arguments(problem_shape, args.op_1, op_1_workspace) }; } @@ -803,19 +808,20 @@ struct Sm90VisitorImplBase { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { Status status = Status::kSuccess; uint8_t* workspace_ptr = reinterpret_cast(workspace); size_t workspace_offset = 0; - status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream); + status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream, cuda_adapter); workspace_offset += Op0::get_workspace_size(problem_shape, args.op_0); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { return status; } - status = Op1::initialize_workspace(problem_shape, args.op_1, workspace_ptr + workspace_offset, stream); + status = Op1::initialize_workspace(problem_shape, args.op_1, workspace_ptr + workspace_offset, stream, cuda_adapter); workspace_offset += Op1::get_workspace_size(problem_shape, args.op_1); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { @@ -862,10 +868,15 @@ struct Sm90VisitorImplBase { template static constexpr Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { + size_t op_0_workspace_size = Op0::get_workspace_size(problem_shape, args.op_0); + size_t op_1_workspace_size = Op1::get_workspace_size(problem_shape, args.op_1); + uint8_t* op_0_workspace = reinterpret_cast(workspace); + uint8_t* op_1_workspace = op_0_workspace + op_0_workspace_size; + uint8_t* op_2_workspace = op_1_workspace + op_1_workspace_size; return Params{ - Op0::to_underlying_arguments(problem_shape, args.op_0, workspace), - Op1::to_underlying_arguments(problem_shape, args.op_1, workspace), - Op2::to_underlying_arguments(problem_shape, args.op_2, workspace) + Op0::to_underlying_arguments(problem_shape, args.op_0, op_0_workspace), + Op1::to_underlying_arguments(problem_shape, args.op_1, op_1_workspace), + Op2::to_underlying_arguments(problem_shape, args.op_2, op_2_workspace) }; } @@ -887,26 +898,27 @@ struct Sm90VisitorImplBase { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { Status status = Status::kSuccess; uint8_t* workspace_ptr = reinterpret_cast(workspace); size_t workspace_offset = 0; - status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream); + status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream, cuda_adapter); workspace_offset += Op0::get_workspace_size(problem_shape, args.op_0); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { return status; } - status = Op1::initialize_workspace(problem_shape, args.op_1, workspace_ptr + workspace_offset, stream); + status = Op1::initialize_workspace(problem_shape, args.op_1, workspace_ptr + workspace_offset, stream, cuda_adapter); workspace_offset += Op1::get_workspace_size(problem_shape, args.op_1); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { return status; } - status = Op2::initialize_workspace(problem_shape, args.op_2, workspace_ptr + workspace_offset, stream); + status = Op2::initialize_workspace(problem_shape, args.op_2, workspace_ptr + workspace_offset, stream, cuda_adapter); workspace_offset += Op2::get_workspace_size(problem_shape, args.op_2); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { @@ -957,11 +969,18 @@ struct Sm90VisitorImplBase { template static constexpr Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { + size_t op_0_workspace_size = Op0::get_workspace_size(problem_shape, args.op_0); + size_t op_1_workspace_size = Op1::get_workspace_size(problem_shape, args.op_1); + size_t op_2_workspace_size = Op2::get_workspace_size(problem_shape, args.op_2); + uint8_t* op_0_workspace = reinterpret_cast(workspace); + uint8_t* op_1_workspace = op_0_workspace + op_0_workspace_size; + uint8_t* op_2_workspace = op_1_workspace + op_1_workspace_size; + uint8_t* op_3_workspace = op_2_workspace + op_2_workspace_size; return Params{ - Op0::to_underlying_arguments(problem_shape, args.op_0, workspace), - Op1::to_underlying_arguments(problem_shape, args.op_1, workspace), - Op2::to_underlying_arguments(problem_shape, args.op_2, workspace), - Op3::to_underlying_arguments(problem_shape, args.op_3, workspace) + Op0::to_underlying_arguments(problem_shape, args.op_0, op_0_workspace), + Op1::to_underlying_arguments(problem_shape, args.op_1, op_1_workspace), + Op2::to_underlying_arguments(problem_shape, args.op_2, op_2_workspace), + Op3::to_underlying_arguments(problem_shape, args.op_3, op_3_workspace) }; } @@ -986,33 +1005,34 @@ struct Sm90VisitorImplBase { template static cutlass::Status - initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream) { + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { Status status = Status::kSuccess; uint8_t* workspace_ptr = reinterpret_cast(workspace); size_t workspace_offset = 0; - status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream); + status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream, cuda_adapter); workspace_offset += Op0::get_workspace_size(problem_shape, args.op_0); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { return status; } - status = Op1::initialize_workspace(problem_shape, args.op_1, workspace_ptr + workspace_offset, stream); + status = Op1::initialize_workspace(problem_shape, args.op_1, workspace_ptr + workspace_offset, stream, cuda_adapter); workspace_offset += Op1::get_workspace_size(problem_shape, args.op_1); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { return status; } - status = Op2::initialize_workspace(problem_shape, args.op_2, workspace_ptr + workspace_offset, stream); + status = Op2::initialize_workspace(problem_shape, args.op_2, workspace_ptr + workspace_offset, stream, cuda_adapter); workspace_offset += Op2::get_workspace_size(problem_shape, args.op_2); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { return status; } - status = Op3::initialize_workspace(problem_shape, args.op_3, workspace_ptr + workspace_offset, stream); + status = Op3::initialize_workspace(problem_shape, args.op_3, workspace_ptr + workspace_offset, stream, cuda_adapter); workspace_offset += Op3::get_workspace_size(problem_shape, args.op_3); workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment); if (status != Status::kSuccess) { diff --git a/include/cutlass/epilogue/thread/activation.h b/include/cutlass/epilogue/thread/activation.h index 6b8f1089..1a226a75 100644 --- a/include/cutlass/epilogue/thread/activation.h +++ b/include/cutlass/epilogue/thread/activation.h @@ -66,7 +66,7 @@ struct Identity { template struct Identity > { CUTLASS_HOST_DEVICE - Array operator()(Array const &value) const { + Array operator()(Array value) const { return value; } }; @@ -75,17 +75,18 @@ struct Identity > { template struct Scale { struct Arguments { + using scale_type = T; T scale = T(1); }; CUTLASS_HOST_DEVICE - T operator()(T const& value, T const& scale) const { + T operator()(T value, T scale) const { multiplies mul; return mul(scale, value); } CUTLASS_HOST_DEVICE - T operator()(T const& value, Arguments const& args = Arguments()) const { + T operator()(T value, Arguments args = Arguments()) const { return this->operator()(value, args.scale); } }; @@ -95,13 +96,13 @@ struct Scale> { using Arguments = typename Scale::Arguments; CUTLASS_HOST_DEVICE - Array operator()(Array const& values, T const& scale) const { + Array operator()(Array values, T scale) const { multiplies> mul; return mul(scale, values); } CUTLASS_HOST_DEVICE - Array operator()(Array const& values, Arguments const& args = Arguments()) const { + Array operator()(Array values, Arguments args = Arguments()) const { return this->operator()(values, args.scale); } }; @@ -113,14 +114,14 @@ struct Scale> { using Arguments = typename Scale::Arguments; CUTLASS_HOST_DEVICE - T operator()(T const &value, decltype(Arguments{}.scale) const& scale) const { + T operator()(T value, typename Arguments::scale_type scale) const { multiplies mul; Activation act; return mul(scale, act(value)); } CUTLASS_HOST_DEVICE - T operator()(T const& value, Arguments const& args = Arguments()) const { + T operator()(T value, Arguments args = Arguments()) const { return this->operator()(value, args.scale); } }; @@ -132,7 +133,7 @@ struct ReLu { static const bool kIsHeavy = false; CUTLASS_HOST_DEVICE - T operator()(T const & threshold, T value) const { + T operator()(T threshold, T value) const { maximum mx; return mx(value, threshold); @@ -666,6 +667,32 @@ struct dReLU> { } }; +/// Computes backwards pass for ReLU operator assuming d_t is the layer gradient and +/// z is computed from the forward pass. +template +struct dReLU_Z { + CUTLASS_HOST_DEVICE + T operator()(T d_t, T z) const { + return z < 0 ? T(0) : d_t; + } +}; + +template +struct dReLU_Z> { + CUTLASS_HOST_DEVICE + Array operator()(Array const& d_t, Array const& z) const { + Array y; + dReLU_Z relu_op; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + y[i] = relu_op(d_t[i], z[i]); + } + + return y; + } +}; + ///////////////////////////////////////////////////////////////////////////////////////////////// } // namespace thread diff --git a/include/cutlass/epilogue/thread/linear_combination_bias_elementwise.h b/include/cutlass/epilogue/thread/linear_combination_bias_elementwise.h index 61a5b7ba..192bc6d1 100644 --- a/include/cutlass/epilogue/thread/linear_combination_bias_elementwise.h +++ b/include/cutlass/epilogue/thread/linear_combination_bias_elementwise.h @@ -213,6 +213,74 @@ public: } } + /// Applies the operation when elementwise_op require arguments and is_source_needed() is true + template + CUTLASS_HOST_DEVICE + void operator()( + FragmentZ &frag_Z, + FragmentT &frag_T, + FragmentAccumulator const &AB, + FragmentC const &frag_C, + FragmentCompute const &V, + ElementwiseArgs const &elementwise_args) const { + + ElementwiseOp elementwise_op; + BinaryOp binary_op; + + FragmentCompute tmp_Accum = NumericArrayConverter()(AB); + FragmentCompute tmp_C = NumericArrayConverter()(frag_C); + FragmentCompute result_Z; + FragmentCompute result_T; + + CUTLASS_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, elementwise_args); + } + + NumericArrayConverter convert_z; + frag_Z = convert_z(result_Z); + + if constexpr (kStoreT) { + NumericArrayConverter convert_t; + frag_T = convert_t(result_T); + } + } + + /// Applies the operation when elementwise_op require arguments and is_source_needed() is false + template + CUTLASS_HOST_DEVICE + void operator()( + FragmentZ &frag_Z, + FragmentT &frag_T, + FragmentAccumulator const &AB, + FragmentCompute const &V, + ElementwiseArgs const &elementwise_args) const { + + ElementwiseOp elementwise_op; + BinaryOp binary_op; + + FragmentCompute tmp_Accum = NumericArrayConverter()(AB); + FragmentCompute result_Z; + FragmentCompute result_T; + + CUTLASS_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, elementwise_args); + } + + NumericArrayConverter convert_z; + frag_Z = convert_z(result_Z); + + if constexpr (kStoreT) { + NumericArrayConverter convert_t; + frag_T = convert_t(result_T); + } + } + /// Applies the operation when is_source_needed() is true CUTLASS_HOST_DEVICE void operator()( diff --git a/include/cutlass/epilogue/thread/linear_combination_generic_with_scaling.h b/include/cutlass/epilogue/thread/linear_combination_generic_with_scaling.h new file mode 100644 index 00000000..e1dde1a6 --- /dev/null +++ b/include/cutlass/epilogue/thread/linear_combination_generic_with_scaling.h @@ -0,0 +1,325 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE 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 with a generic element-wise activation + function. Scaling factors are applied to operands A, B, and C. The pre-activation auxiliary + output is also returned. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_types.h" +#include "cutlass/array.h" +#include "cutlass/functional.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/epilogue/thread/scale_type.h" +#include "cutlass/epilogue/thread/linear_combination_generic.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +/// +/// Aux = ((alpha * scale_a * scale_b) * accumulator) + ((beta * scale_c) * source) + bias +/// D = activation(Aux) +/// +template < + template class ActivationFunctor, + typename ElementOutput_, ///< Data type used to load and store tensors + typename ElementAuxOutput_, ///< Data type used to store auxiliary output + 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 LinearCombinationGenericWithScalingAndAbsMax { +public: + + using ElementOutput = ElementOutput_; + using ElementAuxOutput = ElementAuxOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementScalingFactor = ElementAccumulator_; + + /// Data type used for absolute maximum value + using ElementAbsmax = float; + + static bool const kIsScalingAndAmaxAuxOutputNeeded = (platform::is_same::value || + platform::is_same::value); + static bool const kIsScalingAndAmaxOutputNeeded = (platform::is_same::value || + platform::is_same::value); + + static bool const kIsHeavy = IsHeavy; + static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; + + using FragmentOutput = Array; + using FragmentAuxOutput = Array; + using FragmentAccumulator = Array; + using FragmentCompute = Array; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + struct Params { + struct ActivationParams + : LinearCombinationGenericParams, + GenericActivationTraits>::Arguments { + using LinearCombinationGenericParams::LinearCombinationGenericParams; + }; + + ActivationParams activation; + ElementScalingFactor const* scale_a_ptr = nullptr; ///< pointer to a scalar - if not null, loads it from memory + ElementScalingFactor const* scale_b_ptr = nullptr; ///< pointer to b scalar - if not null, loads it from memory + ElementScalingFactor const* scale_c_ptr = nullptr; ///< pointer to c scalar - if not null, loads it from memory + ElementScalingFactor const* scale_d_ptr = nullptr; ///< pointer to d scalar - if not null, loads it from memory + ElementScalingFactor const* scale_aux_ptr = nullptr; ///< pointer to aux scalar - if not null, loads it from memory + + ElementAbsmax * abs_max_aux_ptr = nullptr; ///< pointer to location to store amax of Aux + ElementAbsmax * abs_max_D_ptr = nullptr; ///< pointer to location to store amax of D + + CUTLASS_HOST_DEVICE + Params() : + scale_a_ptr(nullptr), + scale_b_ptr(nullptr), + scale_c_ptr(nullptr), + scale_d_ptr(nullptr), + scale_aux_ptr(nullptr), + abs_max_aux_ptr(nullptr), + abs_max_D_ptr(nullptr) {} + + CUTLASS_HOST_DEVICE + Params(ActivationParams activation_params, + ElementScalingFactor const* scale_a_ptr, + ElementScalingFactor const* scale_b_ptr, + ElementScalingFactor const* scale_c_ptr, + ElementScalingFactor const* scale_d_ptr, + ElementScalingFactor const* scale_aux_ptr, + ElementAbsmax * abs_max_aux_ptr, + ElementAbsmax * abs_max_D_ptr) : + activation(activation_params), + scale_a_ptr(scale_a_ptr), + scale_b_ptr(scale_b_ptr), + scale_c_ptr(scale_c_ptr), + scale_d_ptr(scale_d_ptr), + scale_aux_ptr(scale_aux_ptr), + abs_max_aux_ptr(abs_max_aux_ptr), + abs_max_D_ptr(abs_max_D_ptr) {} + }; + +private: + + // + // Data members + // + + Params params_; + bool skip_elementwise_; + + // Scaling factors for output and auxiliary output + ElementCompute scale_d_; + ElementCompute scale_aux_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + CUTLASS_HOST_DEVICE + LinearCombinationGenericWithScalingAndAbsMax(Params const ¶ms) : + params_(params), + skip_elementwise_(false), + scale_d_(ElementCompute(params.scale_d_ptr ? *(params.scale_d_ptr) : ElementScalingFactor(1))), + scale_aux_(ElementCompute(params.scale_aux_ptr ? *(params.scale_aux_ptr) : ElementScalingFactor(1))) + { + params_.activation.alpha = (params.activation.alpha_ptr ? *params.activation.alpha_ptr : params.activation.alpha); + params_.activation.beta = (params.activation.beta_ptr ? *params.activation.beta_ptr : params.activation.beta); + auto scale_a = + ElementCompute(params.scale_a_ptr ? *(params.scale_a_ptr) : ElementScalingFactor(1)); + auto scale_b = + ElementCompute(params.scale_b_ptr ? *(params.scale_b_ptr) : ElementScalingFactor(1)); + auto scale_c = + ElementCompute(params.scale_c_ptr ? *(params.scale_c_ptr) : ElementScalingFactor(1)); + + multiplies multiply; + params_.activation.alpha = multiply(params.activation.alpha, multiply(scale_a, scale_b)); + params_.activation.beta = multiply(params.activation.beta, scale_c); + } + + /// Returns true if source is needed + CUTLASS_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_.activation.beta != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + CUTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + params_.activation.beta = ElementCompute(1); + } + + // Only the final partition should perform the activation function + // and scale the output and auxiliary output values. + if (k_partition != k_partition_count - 1) { + skip_elementwise_ = true; + scale_d_ = ElementCompute(1.); + scale_aux_ = ElementCompute(1.); + } + } + + /// Computes linear scaling: + /// Aux = (alpha * scale_a * scale_b * accumulator) + (beta * scale_c * source) + bias + /// D = activation(Aux) + CUTLASS_HOST_DEVICE + void operator()( + FragmentCompute& output, + FragmentCompute& aux_output, + FragmentAccumulator const &accumulator, + FragmentCompute const& bias, + FragmentOutput const &source) { + + // 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 multiply; + plus add; + multiply_add mul_add_accumulator; + ActivationFunctor activation; + + if (Scale == ScaleType::NoBetaScaling) { + intermediate = converted_source; + intermediate = mul_add_accumulator(params_.activation.alpha, converted_accumulator, intermediate); + } else if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = multiply(params_.activation.beta, converted_source); + intermediate = mul_add_accumulator(params_.activation.alpha, converted_accumulator, intermediate); + } + + intermediate = add(intermediate, bias); + + aux_output = intermediate; + if constexpr (GenericActivationTraits>::IsArgumentsNeeded) { + output = skip_elementwise_ ? intermediate : activation(intermediate, params_.activation); + } else { + output = skip_elementwise_ ? intermediate : activation(intermediate); + } + } + + /// Computes linear scaling: + /// Aux = (alpha * scale_a * scale_b * accumulator) + bias + /// D = activation(Aux) + CUTLASS_DEVICE + void operator()( + FragmentCompute& output, + FragmentCompute& aux_output, + FragmentAccumulator const &accumulator, + FragmentCompute const& bias) { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + + FragmentCompute intermediate; + + multiplies multiply; + plus add; + ActivationFunctor activation; + + if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = multiply(params_.activation.alpha, converted_accumulator); + } + + intermediate = add(intermediate, bias); + + aux_output = intermediate; + if constexpr (GenericActivationTraits>::IsArgumentsNeeded) { + output = skip_elementwise_ ? intermediate : activation(intermediate, params_.activation); + } else { + output = skip_elementwise_ ? intermediate : activation(intermediate); + } + } + + CUTLASS_HOST_DEVICE + ElementAbsmax* get_ptr_output_abs_max() const { + return params_.abs_max_D_ptr; + } + + CUTLASS_HOST_DEVICE + ElementAbsmax* get_ptr_aux_output_abs_max() const { + return params_.abs_max_aux_ptr; + } + + CUTLASS_HOST_DEVICE + ElementCompute get_scale_d() const { + return scale_d_; + } + + CUTLASS_HOST_DEVICE + ElementCompute get_scale_aux() const { + return scale_aux_; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace cutlass diff --git a/include/cutlass/epilogue/thread/linear_combination_planar_complex.h b/include/cutlass/epilogue/thread/linear_combination_planar_complex.h index 4cfc4a3d..a6e2a3dc 100644 --- a/include/cutlass/epilogue/thread/linear_combination_planar_complex.h +++ b/include/cutlass/epilogue/thread/linear_combination_planar_complex.h @@ -40,6 +40,7 @@ #include "cutlass/array_planar_complex.h" #include "cutlass/functional.h" #include "cutlass/numeric_conversion.h" +#include "cutlass/epilogue/thread/scale_type.h" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -62,7 +63,8 @@ template < ///< 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 + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest, + ScaleType::Kind Scale = ScaleType::Default ///< Control Alpha and Beta scaling > class LinearCombinationPlanarComplex { public: @@ -70,8 +72,10 @@ public: using ElementOutput = ElementOutput_; using ElementAccumulator = ElementAccumulator_; using ElementCompute = ElementCompute_; + using ElementScalar = complex; static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; using FragmentOutput = ArrayPlanarComplex; using FragmentAccumulator = ArrayPlanarComplex; @@ -82,37 +86,30 @@ public: /// 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 + ElementScalar alpha{ElementCompute(1)}; ///< scales accumulators + ElementScalar beta{ElementCompute(0)}; ///< scales source tensor + ElementScalar const* alpha_ptr{nullptr}; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementScalar const* beta_ptr{nullptr}; ///< pointer to source scalar - if not null, loads it from memory // // Methods // - CUTLASS_HOST_DEVICE - Params(): - alpha(ElementCompute(1)), - beta(ElementCompute(0)), - alpha_ptr(nullptr), - beta_ptr(nullptr) { } + Params() = default; CUTLASS_HOST_DEVICE Params( - complex alpha, - complex beta - ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { - - } + ElementScalar alpha, + ElementScalar beta + ): alpha(alpha), beta(beta) + {} CUTLASS_HOST_DEVICE Params( - complex const *alpha_ptr, - complex const *beta_ptr - ): alpha(complex()), beta(complex()), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { - - } + ElementScalar const *alpha_ptr, + ElementScalar const *beta_ptr + ): alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) + {} }; private: @@ -121,15 +118,14 @@ private: // Data members // - complex alpha_; - complex beta_; + ElementScalar alpha_; + ElementScalar beta_; public: /// Constructs the function object, possibly loading from pointers in host memory CUTLASS_HOST_DEVICE LinearCombinationPlanarComplex(Params const ¶ms) { - alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); } @@ -137,6 +133,8 @@ public: /// Returns true if source is needed CUTLASS_HOST_DEVICE bool is_source_needed() const { + if (Scale == ScaleType::OnlyAlphaScaling) return false; + return beta_.real() != ElementCompute(0) || beta_.imag() != ElementCompute(0); } @@ -213,8 +211,8 @@ public: 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_op(alpha_.real(), converted_accumulator.real); + intermediate.imag = mul_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); diff --git a/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op.h b/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op.h index 7ce58115..f3b006a1 100644 --- a/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op.h +++ b/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op.h @@ -287,7 +287,7 @@ struct DefaultIteratorsTensorOp< platform::is_same::value, "ElementOutput needs to be 4 or 8 bit (unsigned) int."); - static_assert((ElementsPerAccess == 16 || ElementsPerAccess == 8), + static_assert((ElementsPerAccess == 16 || ElementsPerAccess == 8 || ElementsPerAccess == 4), "ElementsPerAccess needs to be 16 or 8."); using WarpTileIteratorMixed = cutlass::epilogue::warp::TileIteratorTensorOpMixed< @@ -308,7 +308,7 @@ struct DefaultIteratorsTensorOp< >; using WarpTileIterator = typename platform::conditional< - (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8) || (ElementsPerAccess == 4), WarpTileIteratorNotMixed, WarpTileIteratorMixed>::type; @@ -327,7 +327,7 @@ struct DefaultIteratorsTensorOp< >; using SharedLoadIterator = typename platform::conditional< - (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8) || (ElementsPerAccess == 4), SharedLoadIteratorNotMixed, SharedLoadIteratorMixed>::type; @@ -354,7 +354,7 @@ struct DefaultIteratorsTensorOp< using ElementOutput = cutlass::float_e4m3_t; - static_assert((ElementsPerAccess == 16 || ElementsPerAccess == 8), + static_assert((ElementsPerAccess == 16 || ElementsPerAccess == 8 || ElementsPerAccess == 4), "ElementsPerAccess needs to be 16 or 8."); using WarpTileIteratorMixed = cutlass::epilogue::warp::TileIteratorTensorOpMixed< @@ -375,7 +375,7 @@ struct DefaultIteratorsTensorOp< >; using WarpTileIterator = typename platform::conditional< - (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8) || (ElementsPerAccess == 4), WarpTileIteratorNotMixed, WarpTileIteratorMixed>::type; @@ -394,7 +394,7 @@ struct DefaultIteratorsTensorOp< >; using SharedLoadIterator = typename platform::conditional< - (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8) || (ElementsPerAccess == 4), SharedLoadIteratorNotMixed, SharedLoadIteratorMixed>::type; @@ -421,7 +421,7 @@ struct DefaultIteratorsTensorOp< using ElementOutput = cutlass::float_e5m2_t; - static_assert((ElementsPerAccess == 16 || ElementsPerAccess == 8), + static_assert((ElementsPerAccess == 16 || ElementsPerAccess == 8 || ElementsPerAccess == 4), "ElementsPerAccess needs to be 16 or 8."); using WarpTileIteratorMixed = cutlass::epilogue::warp::TileIteratorTensorOpMixed< @@ -442,7 +442,7 @@ struct DefaultIteratorsTensorOp< >; using WarpTileIterator = typename platform::conditional< - (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8) || (ElementsPerAccess == 4), WarpTileIteratorNotMixed, WarpTileIteratorMixed>::type; @@ -461,7 +461,7 @@ struct DefaultIteratorsTensorOp< >; using SharedLoadIterator = typename platform::conditional< - (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8) || (ElementsPerAccess == 4), SharedLoadIteratorNotMixed, SharedLoadIteratorMixed>::type; diff --git a/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h b/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h index bca7ee4e..e1ae5a24 100644 --- a/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h +++ b/include/cutlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h @@ -119,7 +119,7 @@ struct DefaultEpilogueTensorOpBlas3 { kBlasMode >; - using AccumulatorFragmentIterator = typename std::conditional::value, + using AccumulatorFragmentIterator = typename platform::conditional::value, cutlass::epilogue::warp::FragmentIteratorComplexTensorOp< typename WarpMmaTensorOp::Shape, typename WarpMmaTensorOp::Policy::Operator::Shape, diff --git a/include/cutlass/epilogue/threadblock/default_epilogue_with_absmax.h b/include/cutlass/epilogue/threadblock/default_epilogue_with_absmax.h new file mode 100644 index 00000000..b0e89a4e --- /dev/null +++ b/include/cutlass/epilogue/threadblock/default_epilogue_with_absmax.h @@ -0,0 +1,126 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE 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 configuration for epilogue computing absolute maximum of output and auxiliary outputs. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_types.h" +#include "cutlass/array.h" + +#include "cutlass/gemm/gemm.h" + +#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "cutlass/epilogue/threadblock/epilogue.h" +#include "cutlass/epilogue/threadblock/epilogue_with_absmax.h" + +#include "cutlass/layout/permute.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for absolute-maximum-computing epilogues with TensorOps +template < + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename ElementOutput, + typename ElementAuxOutput, + typename ElementVector, + typename OutputOp, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultEpilogueWithAbsMax { + + /// Use defaults related to the existing epilogue + using Base = DefaultEpilogueTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp, + ElementsPerAccess + >; + + // + // Stores the output + // + using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout + >; + + // + // Stores the auxiliary output + // + using AuxOutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementAuxOutput, + ScatterD, + PermuteDLayout + >; + + /// Define the epilogue + using Epilogue = EpilogueWithAbsMax< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputTileIterator, + AuxOutputTileIterator, + ElementVector, + typename Base::AccumulatorFragmentIterator, + typename Base::WarpTileIterator, + typename Base::SharedLoadIterator, + OutputOp, + typename Base::Padding, + Base::kFragmentsPerIteration + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/threadblock/default_epilogue_with_broadcast.h b/include/cutlass/epilogue/threadblock/default_epilogue_with_broadcast.h index 140e807b..550354b3 100644 --- a/include/cutlass/epilogue/threadblock/default_epilogue_with_broadcast.h +++ b/include/cutlass/epilogue/threadblock/default_epilogue_with_broadcast.h @@ -57,6 +57,63 @@ namespace cutlass { namespace epilogue { namespace threadblock { +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for SimtOps. +template < + typename Shape, + typename WarpMmaSimt, + typename ElementOutput, + typename ElementTensor, + typename ElementVector, + typename OutputOp, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultEpilogueWithBroadcastSimt { + + /// Use defaults related to the existing epilogue + using Base = DefaultEpilogueSimt< + Shape, + WarpMmaSimt, + OutputOp, + ElementsPerAccess + >; + + // + // Stores the result z = (y = GEMM(A, B, C), broadcast) + // + using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout + >; + + // + // Additional tensor tile iterator - stores t = Elementwise(z) + // + using TensorTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementTensor + >; + + /// Define the epilogue + using Epilogue = EpilogueWithBroadcast< + Shape, + WarpMmaSimt, + Base::kPartitionsK, + OutputTileIterator, + TensorTileIterator, + ElementVector, + typename Base::AccumulatorFragmentIterator, + typename Base::WarpTileIterator, + typename Base::SharedLoadIterator, + OutputOp, + typename Base::Padding + >; +}; //////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/threadblock/epilogue_with_absmax.h b/include/cutlass/epilogue/threadblock/epilogue_with_absmax.h new file mode 100644 index 00000000..9bae7a74 --- /dev/null +++ b/include/cutlass/epilogue/threadblock/epilogue_with_absmax.h @@ -0,0 +1,923 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + + \brief Threadblock-level epilogue computing: + Aux = ((alpha * scale_a * scale_b) * accumulator) + ((beta * scale_c) * source) + bias + D = activation(Aux) + + if Aux is fp8 type: + abs_max_output = max( abs(aux) | (for every aux in Aux)) + Aux = scale_aux * Aux + endif + + if D is fp8 type: + abs_max_output = max( abs(d) | (for every d in D)) + D = scale_d * D + endif + + Parameter Aux is optionally stored to global memory +*/ + +#pragma once + +#if defined(__CUDACC_RTC__) +#include +#include +#else +#include +#include +#endif + +#include "cutlass/cutlass.h" +#include "cutlass/array.h" +#include "cutlass/numeric_types.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/tensor_coord.h" +#include "cutlass/aligned_buffer.h" +#include "cutlass/functional.h" +#include "cutlass/fast_math.h" +#include "cutlass/layout/vector.h" +#include "cutlass/layout/tensor.h" + +#include "cutlass/gemm/gemm.h" + +#include "cutlass/transform/pitch_linear_thread_map.h" +#include "cutlass/transform/threadblock/regular_tile_iterator.h" + +#include "cutlass/epilogue/threadblock/epilogue_base.h" +#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h" + +#include "cutlass/numeric_types.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +/// Helper class for keeping track of absolute maximums and performing scaling +template < + typename Iterator, // Iterator type used for storing the data for which absolute maximum and scaling + // will be computed. This type is used for predicating absolute maximum calculations. + typename Fragment, // Type of input to be computed on + bool ScalingAndAmaxNeeded // Whether to perform absolute maximum and scaling operations +> +struct ScalingAndAmaxHelper; + +/// Partial specialization that does not perform scaling or calculate an absolute maximum +template +struct ScalingAndAmaxHelper { + using Element = typename Fragment::Element; + + CUTLASS_HOST_DEVICE + ScalingAndAmaxHelper(Element scale) { } + + CUTLASS_DEVICE + Fragment operator()(const Iterator& iterator, const Fragment& inp) { + return inp; + } + + CUTLASS_HOST_DEVICE + Element get_abs_max() const { + return Element(0.); + } + + CUTLASS_HOST_DEVICE + void set_scaling_factor(Element scale_) { } +}; + +/// Partial specialization that keeps track of an absolute maximum value of inputs seen +/// and scales inputs +template +struct ScalingAndAmaxHelper { + using Element = typename Fragment::Element; + using AccessType = typename Iterator::AccessType; + using ThreadMap = typename Iterator::ThreadMap; + + Element abs_max; + Element scale; + + // Operators + maximum_with_nan_propogation max_op; + absolute_value_op abs_op; + multiplies multiply; + + CUTLASS_HOST_DEVICE + ScalingAndAmaxHelper(Element scale_) : abs_max(0.), scale(scale_) { } + + // Compute the absolute maximum value between `abs_max` and the entries + // of `frag` for predicated-on entries of `iterator`. Return a scaled + // version of `inp`. + CUTLASS_DEVICE + Fragment operator()(const Iterator& iterator, const Fragment& frag) { + using PredicateGroup = Array; + PredicateGroup const *frag_ptr = reinterpret_cast(&frag); + + typename Iterator::Mask mask; + iterator.get_mask(mask); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + CUTLASS_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 + iterator.thread_start_row()) < iterator.extent_row()); + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + bool guard = row_guard && mask.predicates[column]; + + if (guard) { + int access_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < PredicateGroup::kElements; ++i) { + abs_max = max_op(abs_max, abs_op(frag_ptr[access_idx][i])); + } + } + } + } + } + } + + // Perform scaling + return multiply(scale, frag); + } + + CUTLASS_HOST_DEVICE + Element get_abs_max() const { + return abs_max; + } + + CUTLASS_HOST_DEVICE + void set_scaling_factor(Element scale_) { + scale = scale_; + } +}; + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +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 AuxOutputTileIterator_, ///< Tile iterator writing auxiliary output tensors + typename ElementVector_, ///< Data type of bias 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 Padding_, ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape) + int FragmentsPerPartition = 1, ///< Used to coarsen the epilogue granularity + int IterationsUnroll = ///< Used to reduce binary size when epilogue op is large + (!IsEpilogueFunctorHeavy::value) +> +class EpilogueWithAbsMax : + 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 AuxOutputTileIterator = AuxOutputTileIterator_; + 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; + + /// Data type used for absolute maximum value + using ElementAbsmax = typename OutputOp::ElementAbsmax; + + /// Compute data type produced by the output op + using ElementCompute = typename OutputOp::ElementCompute; + + /// Compute fragment + using FragmentCompute = Array; + + /// Helpers for (optionally) computing absolute maximums and scaling output and auxiliary output + using OutputScaler = detail::ScalingAndAmaxHelper; + + using AuxOutputScaler = detail::ScalingAndAmaxHelper; + + /// 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 auxiliary output + using ElementAuxOutput = typename AuxOutputTileIterator::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 cutlass::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; + + /// Auxiliary output access type + using AuxAccessType = 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 + CUTLASS_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; + }; + + CUTLASS_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 + CUTLASS_DEVICE + EpilogueWithAbsMax( + 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 + CUTLASS_DEVICE + void operator()( + OutputOp &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 + AuxOutputTileIterator aux_iterator, ///< Tile iterator for destination auxiliary output + 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); + + OutputScaler output_scaler(output_op.get_scale_d()); + + AuxOutputScaler aux_scaler(output_op.get_scale_aux()); + + if (!output_op.is_source_needed()) { + compute_source_not_needed_( + output_op, + broadcast_fragment, + destination_iterator, + accumulators, + aux_iterator, + output_scaler, + aux_scaler); + } + else { + compute_source_needed_( + output_op, + broadcast_fragment, + destination_iterator, + accumulators, + source_iterator, + aux_iterator, + output_scaler, + aux_scaler); + } + + // Store the absolute maximum values of the output and auxiliar tensors, if needed. + if (output_op.get_ptr_output_abs_max() != nullptr) { + ElementAbsmax local_abs_max = + NumericConverter{}(output_scaler.get_abs_max()); + atomic_maximum{}( + output_op.get_ptr_output_abs_max(), local_abs_max); + } + + if (output_op.get_ptr_aux_output_abs_max() != nullptr) { + ElementAbsmax local_abs_max = + NumericConverter{}(aux_scaler.get_abs_max()); + atomic_maximum{}( + output_op.get_ptr_aux_output_abs_max(), local_abs_max); + } + } + +private: + + CUTLASS_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); + + CUTLASS_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 + CUTLASS_DEVICE static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator &warp_tile_iterator) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + CUTLASS_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)); + } + } + + CUTLASS_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)...}; + + CUTLASS_UNUSED(dummy[0]); + } + }; + + /// Streams the result to global memory + CUTLASS_DEVICE + void compute_source_not_needed_( + OutputOp &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 + AuxOutputTileIterator aux_iterator, ///< Tile iterator for destination auxiliary output + OutputScaler& output_scaler, ///< Helper for (optionally) computing the absolute maximum and scaling output + AuxOutputScaler& aux_scaler ///< Helper for (optionally) computing the absolute maximum and scaling the auxiliary output + ) { + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + + // CUTLASS_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< + cutlass::make_index_sequence>::push(iter, + accum_fragment_iterator, + this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + CUTLASS_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; + + CUTLASS_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 + // + + FragmentCompute frag_Z_compute; + FragmentCompute frag_Aux_compute; + + apply_output_operator_source_not_needed_( + frag_Z_compute, + frag_Aux_compute, + output_op, + aligned_accum_fragment[0], + broadcast_fragment); + + // + // Conditionally store fragments + // + + // (Optionally) compute the absolute maximum of frag_Z and scale frag_Z + frag_Z_compute = output_scaler(destination_iterator, frag_Z_compute); + NumericArrayConverter cvt_to_dst; + typename OutputTileIterator::Fragment frag_Z = cvt_to_dst(frag_Z_compute); + + // Always store the output + destination_iterator.store(frag_Z); + ++destination_iterator; + + // Only store the auxiliary output if scaling and absolute-maximum calculation were needed + if (OutputOp::kIsScalingAndAmaxAuxOutputNeeded) { + frag_Aux_compute = aux_scaler(aux_iterator, frag_Aux_compute); + + NumericArrayConverter cvt_to_aux; + typename AuxOutputTileIterator::Fragment frag_Aux = cvt_to_aux(frag_Aux_compute); + aux_iterator.store(frag_Aux); + ++aux_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 + CUTLASS_DEVICE + static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator &warp_tile_iterator) { + CUTLASS_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); + } + + CUTLASS_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 + CUTLASS_DEVICE + void compute_source_needed_( + OutputOp &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 + AuxOutputTileIterator aux_iterator, ///< Tile iterator for destination auxiliary output + OutputScaler& output_scaler, ///< Helper for (optionally) computing the absolute maximum and scaling output + AuxOutputScaler& aux_scaler ///< Helper for (optionally) computing the absolute maximum and scaling the auxiliary output + ) { + + 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) + 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; + + CUTLASS_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 + // + + FragmentCompute frag_Z_compute; + FragmentCompute frag_Aux_compute; + + apply_output_operator_( + frag_Z_compute, + frag_Aux_compute, + output_op, + aligned_accum_fragment[0], + source_fragment, + broadcast_fragment); + + // + // Conditionally store fragments + // + + // (Optionally) compute the absolute maximum of frag_Z and scale frag_Z + frag_Z_compute = output_scaler(destination_iterator, frag_Z_compute); + NumericArrayConverter cvt_to_dst; + typename OutputTileIterator::Fragment frag_Z = cvt_to_dst(frag_Z_compute); + + // Always store the output + destination_iterator.store(frag_Z); + ++destination_iterator; + + // Only store the auxiliary output if scaling and absolute-maximum calculation were needed + if (OutputOp::kIsScalingAndAmaxAuxOutputNeeded) { + frag_Aux_compute = aux_scaler(aux_iterator, frag_Aux_compute); + + NumericArrayConverter cvt_to_aux; + typename AuxOutputTileIterator::Fragment frag_Aux = cvt_to_aux(frag_Aux_compute); + aux_iterator.store(frag_Aux); + ++aux_iterator; + } + } + } + + /// Helper to invoke the output functor over each vector of output + CUTLASS_DEVICE + void apply_output_operator_( + FragmentCompute &frag_Z, + FragmentCompute &frag_Aux, + OutputOp &output_op, + typename SharedLoadIterator::Fragment const &frag_AB, + typename OutputTileIterator::Fragment const &frag_C, + BroadcastFragment const &frag_Broadcast) { + + using AccessTypeZ = Array; + using AccessTypeAux = Array; + using AccessTypeBroadcast = Array; + + AccessTypeZ *frag_Z_ptr = reinterpret_cast(&frag_Z); + AccessTypeAux *frag_Aux_ptr = reinterpret_cast(&frag_Aux); + + 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; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + output_op( + frag_Z_ptr[i], + frag_Aux_ptr[i], + frag_AB_ptr[i], + frag_Broadcast_ptr[i % ThreadMap::Iterations::kColumn], + frag_C_ptr[i]); + } + } + + /// Helper to invoke the output functor over each vector of output + CUTLASS_DEVICE + void apply_output_operator_source_not_needed_( + FragmentCompute &frag_Z, + FragmentCompute &frag_Aux, + OutputOp &output_op, + typename SharedLoadIterator::Fragment const &frag_AB, + BroadcastFragment const &frag_Broadcast) { + + using AccessTypeZ = Array; + using AccessTypeAux = Array; + using AccessTypeBroadcast = Array; + + AccessTypeZ *frag_Z_ptr = reinterpret_cast(&frag_Z); + AccessTypeAux *frag_Aux_ptr = reinterpret_cast(&frag_Aux); + + 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; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + + output_op( + frag_Z_ptr[i], + frag_Aux_ptr[i], + frag_AB_ptr[i], + frag_Broadcast_ptr[i % ThreadMap::Iterations::kColumn]); + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/threadblock/fusion/visitor_2x.hpp b/include/cutlass/epilogue/threadblock/fusion/visitor_2x.hpp index dc84d1bf..8b1cd4fd 100644 --- a/include/cutlass/epilogue/threadblock/fusion/visitor_2x.hpp +++ b/include/cutlass/epilogue/threadblock/fusion/visitor_2x.hpp @@ -92,11 +92,11 @@ struct VisitorImpl2x: fusion::detail::Sm90VisitorImplBase { /// Called after accumulators have been exchanged for each accumulator vector template CUTLASS_DEVICE auto // returns an Array - visit(int iter_idx, int row_idx, int column_idx, int frg_idx, + visit(int iter_idx, int row_idx, int column_idx, int frg_idx, Array const& frg_acc, Array const&... frg_inputs) // depends on the N-naryness of the op = delete; // Must be implemented for each operation - + /// Called at the start of a row CUTLASS_DEVICE void end_row(int row_idx) { @@ -179,12 +179,12 @@ struct TreeVisitor2x : VisitorImpl2x { CUTLASS_DEVICE Callbacks(CallbacksImpl&& impl) : CallbacksImpl(cute::forward(impl)) {} - + using CallbacksImpl::callbacks_tuple; template CUTLASS_DEVICE auto - visit(int iter_idx, int row_idx, int column_idx, int frg_idx, + visit(int iter_idx, int row_idx, int column_idx, int frg_idx, Array const& frg_acc) { constexpr int Rm1 = sizeof...(ChildOps); return cute::detail::tapply(callbacks_tuple, @@ -242,16 +242,16 @@ struct TopologicalVisitor2x : VisitorImpl2x { CUTLASS_DEVICE Callbacks(CallbacksImpl&& impl) : CallbacksImpl(cute::forward(impl)) {} - + using CallbacksImpl::callbacks_tuple; template CUTLASS_DEVICE auto - visit(int iter_idx, int row_idx, int column_idx, int frg_idx, + visit(int iter_idx, int row_idx, int column_idx, int frg_idx, Array const& frg_acc) { constexpr int Rm1 = sizeof...(Ops) - 1; auto frg_compute_tuple = cute::repeat(Array{}); - + return cute::detail::tapply(EdgeTuple{}, callbacks_tuple, frg_compute_tuple, // Visit the first R-1 ops in topological order [&] (auto&& edge_seq, auto& callbacks, auto& frg_compute) { @@ -271,7 +271,7 @@ struct TopologicalVisitor2x : VisitorImpl2x { return frg_compute; }, // Visit the last op - [&] (auto const&...) { + [&] (auto const&...ops) { return cute::detail::apply(frg_compute_tuple, // Compute the last op with children inputs [&] (auto const&... frg_inputs) { @@ -343,7 +343,7 @@ struct OutputTileThreadLayout: DefaultThreadMapTensorOp< ThreadblockShape_::kK/WarpShape_::kK, Element_, ElementsPerAccess>::Type { - + using Base = typename DefaultThreadMapTensorOp< ThreadblockShape_, WarpShape_, diff --git a/include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp b/include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp index 01d3cb6f..7a332f11 100644 --- a/include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp +++ b/include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp @@ -113,6 +113,12 @@ struct VisitorScalarBroadcast { return args; } + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + return 0; + } + CUTLASS_HOST_DEVICE VisitorScalarBroadcast() { } @@ -216,6 +222,12 @@ struct VisitorAuxLoad{ return args; } + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + return 0; + } + // Software pipeline stages static const int Stages = ThreadMap::Stages; @@ -341,6 +353,12 @@ struct VisitorRowBroadcast { return args; } + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + return 0; + } + struct SharedStorage {}; // Global load type @@ -464,6 +482,12 @@ struct VisitorColBroadcast { return args; } + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + return 0; + } + struct SharedStorage { }; CUTLASS_HOST_DEVICE diff --git a/include/cutlass/epilogue/threadblock/fusion/visitor_store.hpp b/include/cutlass/epilogue/threadblock/fusion/visitor_store.hpp index 749b2a33..1c24e22d 100644 --- a/include/cutlass/epilogue/threadblock/fusion/visitor_store.hpp +++ b/include/cutlass/epilogue/threadblock/fusion/visitor_store.hpp @@ -74,6 +74,12 @@ struct VisitorAuxStore{ return args; } + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + return 0; + } + struct SharedStorage {}; static int constexpr vec_bits = ThreadMap::kElementsPerAccess * sizeof_bits::value; @@ -258,6 +264,12 @@ struct VisitorColReduction { return args; } + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + return 0; + } + struct SharedStorage { }; CUTLASS_HOST_DEVICE @@ -398,6 +410,12 @@ struct VisitorRowReduction { return args; } + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + return 0; + } + using SharedStorageShape = decltype(select<0,1,2,3,5,8,10>(typename ThreadMap::ThreadMapShape{})); struct SharedStorage { @@ -672,6 +690,12 @@ struct VisitorScalarReduction { return args; } + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + return 0; + } + struct SharedStorage { }; CUTLASS_HOST_DEVICE diff --git a/include/cutlass/epilogue/warp/tile_iterator_tensor_op_mixed.h b/include/cutlass/epilogue/warp/tile_iterator_tensor_op_mixed.h index 8b26116a..c512dd87 100644 --- a/include/cutlass/epilogue/warp/tile_iterator_tensor_op_mixed.h +++ b/include/cutlass/epilogue/warp/tile_iterator_tensor_op_mixed.h @@ -108,7 +108,10 @@ public: 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."); + // Currently support max 4 ptr + static constexpr int kMaxPointerCount{4}; + + static_assert(kPointerCount <= kMaxPointerCount, "Can only accommodate four pointers at present."); static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32)."); }; @@ -127,24 +130,18 @@ private: // /// Internal pointer to memory - AccessType *pointers_[Detail::kPointerCount]; + AccessType *pointers_[Detail::kPointerCount] = {nullptr}; /// Stride in units of AccessType - int stride_; + int stride_{0}; /// Logical column in which warp tile is aligned - int warp_column_; + int warp_column_{0}; public: /// Default constructor - CUTLASS_HOST_DEVICE - TileIteratorTensorOpMixed() { - CUTLASS_PRAGMA_UNROLL - for (int64_t i = 0; i < Detail::kPointerCount; ++i) { - pointers_[i] = nullptr; - } - } + TileIteratorTensorOpMixed() = default; /// Constructor from TensorRef CUTLASS_HOST_DEVICE @@ -165,18 +162,7 @@ public: 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; - } + pointers_[i % Detail::kPointerCount] = ptr; } } @@ -379,24 +365,18 @@ private: // /// Internal pointer to memory - AccessType *pointers_[Detail::kPointerCount]; + AccessType *pointers_[Detail::kPointerCount] = {nullptr}; /// Stride in units of AccessType - int stride_; + int stride_{0}; /// Uniform offset in bytes added to warp tile iterator - int uniform_offset_[Detail::kOffsetCount]; + int uniform_offset_[Detail::kOffsetCount] = {0}; public: /// Default constructor - CUTLASS_HOST_DEVICE - TileIteratorTensorOpMixed() { - CUTLASS_PRAGMA_UNROLL - for (int64_t i = 0; i < Detail::kPointerCount; ++i) { - pointers_[i] = nullptr; - } - } + TileIteratorTensorOpMixed() = default; /// Constructor from TensorRef CUTLASS_HOST_DEVICE @@ -583,21 +563,15 @@ private: // /// Internal pointer to memory - AccessType *pointers_[Detail::kPointerCount]; + AccessType *pointers_[Detail::kPointerCount] = {nullptr}; /// Stride in units of AccessType - int stride_; + int stride_{0}; public: /// Default constructor - CUTLASS_HOST_DEVICE - TileIteratorTensorOpMixed() { - CUTLASS_PRAGMA_UNROLL - for (int64_t i = 0; i < Detail::kPointerCount; ++i) { - pointers_[i] = nullptr; - } - } + TileIteratorTensorOpMixed() = default; /// Constructor from TensorRef CUTLASS_HOST_DEVICE @@ -716,6 +690,386 @@ public: }; ///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for float x 16 => float_e4m3_t/float_e5m2_t x 16 +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename OperatorShape_ ///< matrix multiply operation shape (concept: gemm::GemmShape), +> +class TileIteratorTensorOpMixed { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using Element = float; + 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] = {nullptr}; + + /// Stride in units of AccessType + int stride_{0}; + + /// Uniform offset in bytes added to warp tile iterator + int uniform_offset_[Detail::kOffsetCount] = {0}; + +public: + + /// Default constructor + TileIteratorTensorOpMixed() = default; + + /// Constructor from TensorRef + CUTLASS_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); + + CUTLASS_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 (i == 0) { + pointers_[0] = ptr; + } + else if (i == 1) { + pointers_[1] = ptr; + } + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Detail::kOffsetCount; ++i) { + uniform_offset_[i] = (i ^ 0) * 4 * sizeof(AccessType); + } + } + + /// Adds a pointer offset + CUTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) { + + CUTLASS_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 + CUTLASS_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; + + CUTLASS_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 + CUTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) { + return add_tile_offset(tile_offset); + } + + /// Store + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + CUTLASS_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 + // + AccessType *smem_ptr = pointers_[ptr_idx]; + smem_ptr[offset] = frag_ptr[n]; +#else + uint32_t smem_addr = arch::cutlass_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]) + ); +#endif + } + } + + /// Store + CUTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for float x 8 => float_e4m3_t/float_e5m2_t x 8 +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename OperatorShape_ ///< matrix multiply operation shape (concept: gemm::GemmShape) +> +class TileIteratorTensorOpMixed { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using Element = float; + 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] = {nullptr}; + + /// Stride in units of AccessType + int stride_{0}; + +public: + + /// Default constructor + TileIteratorTensorOpMixed() = default; + + /// Constructor from TensorRef + CUTLASS_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); + + CUTLASS_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 (i == 0) { + pointers_[0] = ptr; + } + else if (i == 1) { + pointers_[1] = ptr; + } + } + } + + /// Adds a pointer offset + CUTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) { + + CUTLASS_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 + CUTLASS_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 + CUTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) { + return add_tile_offset(tile_offset); + } + + /// Store + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + CUTLASS_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 + // + AccessType *smem_ptr = pointers_[ptr_idx]; + smem_ptr[offset] = frag_ptr[n]; +#else + uint32_t smem_addr = arch::cutlass_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]) + ); +#endif + } + } + + /// Store + CUTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + } // namespace warp } // namespace epilogue } // namespace cutlass diff --git a/include/cutlass/fast_math.h b/include/cutlass/fast_math.h index 86b7bf6b..3e842d6a 100644 --- a/include/cutlass/fast_math.h +++ b/include/cutlass/fast_math.h @@ -153,13 +153,15 @@ struct divide_assert { * Round dividend up to the nearest multiple of divisor */ template -CUTLASS_HOST_DEVICE dividend_t round_nearest(dividend_t dividend, divisor_t divisor) { +CUTLASS_HOST_DEVICE +CUTLASS_CONSTEXPR_IF_CXX17 +dividend_t round_nearest(dividend_t dividend, divisor_t divisor) { return ((dividend + divisor - 1) / divisor) * divisor; } template CUTLASS_HOST_DEVICE -constexpr +CUTLASS_CONSTEXPR_IF_CXX17 value_t abs_for_integer(value_t a) { return ((a > 0) ? a : -a); } @@ -185,9 +187,8 @@ template CUTLASS_HOST_DEVICE CUTLASS_CONSTEXPR_IF_CXX17 value_t lcm(value_t a, value_t b) { - value_t temp = gcd(a, b); - - return temp ? (cutlass::abs_for_integer(a) / temp * cutlass::abs_for_integer(b)) : value_t(); + value_t temp = cutlass::gcd(a, b); + return (temp != 0) ? value_t(cutlass::abs_for_integer(a) / temp * cutlass::abs_for_integer(b)) : value_t{}; } /** @@ -195,9 +196,9 @@ value_t lcm(value_t a, value_t b) { */ template CUTLASS_HOST_DEVICE -constexpr +CUTLASS_CONSTEXPR_IF_CXX17 value_t gcd_cxx11(value_t a, value_t b) { - return (a == 0 || b == 0) ? cutlass::abs_for_integer(a | b) : gcd_cxx11(b, a % b); + return (a == 0 || b == 0) ? cutlass::abs_for_integer(a | b) : cutlass::gcd_cxx11(b, a % b); } /** @@ -205,11 +206,11 @@ value_t gcd_cxx11(value_t a, value_t b) { */ template CUTLASS_HOST_DEVICE -constexpr +CUTLASS_CONSTEXPR_IF_CXX17 value_t lcm_cxx11(value_t a, value_t b) { - return gcd_cxx11(a, b) ? (cutlass::abs_for_integer(a) / gcd_cxx11(a, b) * - cutlass::abs_for_integer(b)) - : value_t(); + return cutlass::gcd_cxx11(a, b) ? (cutlass::abs_for_integer(a) / cutlass::gcd_cxx11(a, b) * + cutlass::abs_for_integer(b)) + : value_t{}; } /// Returns the smallest value in the half-open range [a, a+b) that is a multiple of b @@ -234,16 +235,20 @@ int ceil_div(int a, int b) { * log2_up/down codes? */ template -CUTLASS_HOST_DEVICE value_t clz(value_t x) { +CUTLASS_HOST_DEVICE +CUTLASS_CONSTEXPR_IF_CXX17 +value_t clz(value_t x) { for (int i = 31; i >= 0; --i) { if ((1 << i) & x) return value_t(31 - i); } - return 32; + return value_t(32); } template -CUTLASS_HOST_DEVICE value_t find_log2(value_t x) { +CUTLASS_HOST_DEVICE +CUTLASS_CONSTEXPR_IF_CXX17 +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; @@ -253,7 +258,8 @@ CUTLASS_HOST_DEVICE value_t find_log2(value_t x) { /** * Find divisor, using find_log2 */ -CUTLASS_HOST_DEVICE +CUTLASS_HOST_DEVICE +CUTLASS_CONSTEXPR_IF_CXX17 void find_divisor(unsigned int& mul, unsigned int& shr, unsigned int denom) { if (denom == 1) { mul = 0; @@ -270,7 +276,8 @@ void find_divisor(unsigned int& mul, unsigned int& shr, unsigned int denom) { /** * Find quotient and remainder using device-side intrinsics */ -CUTLASS_HOST_DEVICE +CUTLASS_HOST_DEVICE +CUTLASS_CONSTEXPR_IF_CXX17 void fast_divmod(int& quo, int& rem, int src, int div, unsigned int mul, unsigned int shr) { #if defined(__CUDA_ARCH__) @@ -286,6 +293,7 @@ void fast_divmod(int& quo, int& rem, int src, int div, unsigned int mul, unsigne // For long int input CUTLASS_HOST_DEVICE +CUTLASS_CONSTEXPR_IF_CXX17 void fast_divmod(int& quo, int64_t& rem, int64_t src, int div, unsigned int mul, unsigned int shr) { #if defined(__CUDA_ARCH__) diff --git a/include/cutlass/float8.h b/include/cutlass/float8.h index fa1f9259..0d925f26 100644 --- a/include/cutlass/float8.h +++ b/include/cutlass/float8.h @@ -55,6 +55,10 @@ # if (__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8)) # define CUDA_PTX_FP8_CVT_ENABLED 1 # endif // (__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8)) +# elif (__CUDA_ARCH__ == 890) +# if (__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 1)) +# define CUDA_PTX_FP8_CVT_ENABLED 1 +# endif // (__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 1)) # endif // (__CUDA_ARCH__ >= 900) #endif // defined(__CUDA_ARCH__) @@ -131,7 +135,7 @@ struct alignas(1) float8_base { 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_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; @@ -1039,6 +1043,34 @@ float_e5m2_t::float_e5m2_t(float_e4m3_t x) { storage = from_float(float_e4m3_t::to_float(x)).storage; } +/////////////////////////////////////////////////////////////// +/// +/// Umbrella floating-point 8-bit data type : type_erased_dynamic_float8_t +/// This umbrella datatype can be enabled when a user provides a specific +/// datatype in runtime argument list. +/// +/// Currently supported runtime datatypes compatible with type_erased_dynamic_float8_t: +/// QMMAFormat::E5M2 +/// QMMAFormat::E4M3 +/// +/////////////////////////////////////////////////////////////// + +union type_erased_dynamic_float8_t { + uint8_t data; + cutlass::float_e5m2_t e5m2; + cutlass::float_e4m3_t e4m3; + CUTLASS_HOST_DEVICE + explicit operator cutlass::float_e5m2_t() const { + return e5m2; + } + + CUTLASS_HOST_DEVICE + explicit operator cutlass::float_e4m3_t() const { + return e4m3; + } + +}; + /////////////////////////////////////////////////////////////////////////////////////////////////// } // namespace cutlass @@ -1073,24 +1105,31 @@ public: static int const digits = F8Type::FP8_NUM_MANTISSA_BITS; /// Least positive value + CUTLASS_HOST_DEVICE static F8Type min() { return F8Type::bitcast(0x01); } /// Maximum finite value + CUTLASS_HOST_DEVICE static F8Type max() { return F8Type::bitcast(F8Type::FP8_MAX_FLT); } /// Returns maximum rounding error + CUTLASS_HOST_DEVICE static F8Type round_error() { return F8Type(0.5f); } /// Returns positive infinity value + CUTLASS_HOST_DEVICE static F8Type infinity() { return F8Type::bitcast(F8Type::FP8_INFINITY_MASK); } /// Returns quiet NaN value + CUTLASS_HOST_DEVICE static F8Type quiet_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); } /// Returns signaling NaN value + CUTLASS_HOST_DEVICE static F8Type signaling_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); } /// Returns smallest positive subnormal value + CUTLASS_HOST_DEVICE static F8Type denorm_min() { return F8Type::bitcast(0x01); } }; @@ -1150,24 +1189,31 @@ public: static int const digits = F8Type::FP8_NUM_MANTISSA_BITS; /// Least positive value + CUTLASS_HOST_DEVICE static F8Type min() { return F8Type::bitcast(0x01); } /// Maximum finite value + CUTLASS_HOST_DEVICE static F8Type max() { return F8Type::bitcast(F8Type::FP8_MAX_FLT); } /// Returns maximum rounding error + CUTLASS_HOST_DEVICE static F8Type round_error() { return F8Type(0.5f); } /// Returns positive infinity value + CUTLASS_HOST_DEVICE static F8Type infinity() { return F8Type::bitcast(F8Type::FP8_INFINITY_MASK); } /// Returns quiet NaN value + CUTLASS_HOST_DEVICE static F8Type quiet_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); } /// Returns signaling NaN value + CUTLASS_HOST_DEVICE static F8Type signaling_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); } /// Returns smallest positive subnormal value + CUTLASS_HOST_DEVICE static F8Type denorm_min() { return F8Type::bitcast(0x01); } }; diff --git a/include/cutlass/functional.h b/include/cutlass/functional.h index 26e1249b..27c8de2d 100644 --- a/include/cutlass/functional.h +++ b/include/cutlass/functional.h @@ -45,9 +45,7 @@ #pragma once #include "cutlass/cutlass.h" -#include "cutlass/half.h" -#include "cutlass/tfloat32.h" -#include "cutlass/bfloat16.h" +#include "cutlass/numeric_types.h" #if defined(CUTLASS_ARCH_WMMA_ENABLED) #include @@ -106,7 +104,7 @@ struct multiplies { template struct scale { T const scaling_factor_; - + CUTLASS_HOST_DEVICE scale(float scaling_factor) : scaling_factor_(scaling_factor) { } @@ -228,19 +226,19 @@ struct divides { } }; -/// reciprocal_approximate +/// reciprocal_approximate template struct reciprocal_approximate { CUTLASS_HOST_DEVICE T operator()(T lhs) const { - return divide(T(1), lhs); + return divides{}(T(1), lhs); } }; template <> struct reciprocal_approximate { CUTLASS_HOST_DEVICE - float operator()(float lhs) const { + float operator()(float lhs) const { float ret; ret = 1.0f / lhs; return ret; @@ -256,7 +254,7 @@ struct negate { } }; -/// Greater equal +/// Greater equal template struct greater_equal { CUTLASS_HOST_DEVICE @@ -265,7 +263,7 @@ struct greater_equal { } }; -/// Greater +/// Greater template struct greater { CUTLASS_HOST_DEVICE @@ -274,7 +272,7 @@ struct greater { } }; -/// Less equal +/// Less equal template struct less_equal { CUTLASS_HOST_DEVICE @@ -283,7 +281,7 @@ struct less_equal { } }; -/// Less +/// Less template struct less { CUTLASS_HOST_DEVICE @@ -421,6 +419,15 @@ struct multiply_add { } }; +template +struct square_and_plus { + CUTLASS_HOST_DEVICE + T operator()(T lhs, T const &rhs) const { + multiply_add multiply_add_op; + return multiply_add_op(rhs, rhs, lhs); + } +}; + // Fused multiply-add that takes exactly one template parameter. // This is useful for working around a known Clang issue, // where a template template parameter with one template parameter @@ -473,6 +480,10 @@ struct first { T operator()(T const & first, T const &...) const { return first; } + CUTLASS_HOST_DEVICE + T operator()(T const & first) const { + return first; + } }; ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -481,7 +492,7 @@ template struct logical_and { CUTLASS_HOST_DEVICE T operator()(T const &a, T const &b) const { - return ((a && b) ? T(1) : T()); + return ((static_cast(a) && static_cast(b)) ? T(1) : T()); } }; @@ -489,7 +500,7 @@ template struct logical_or { CUTLASS_HOST_DEVICE T operator()(T const &a, T const &b) const { - return ((a || b) ? T(1) : T()); + return ((static_cast(a) || static_cast(b)) ? T(1) : T()); } }; diff --git a/include/cutlass/gemm/collective/builders/sm90_common.inl b/include/cutlass/gemm/collective/builders/sm90_common.inl index 945862bd..14ae7398 100644 --- a/include/cutlass/gemm/collective/builders/sm90_common.inl +++ b/include/cutlass/gemm/collective/builders/sm90_common.inl @@ -30,8 +30,12 @@ **************************************************************************************************/ #pragma once +#include "cutlass/arch/mma.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/dispatch_policy.hpp" #include "cutlass/detail/layout.hpp" #include "cutlass/detail/collective.hpp" +#include "cutlass/detail/dependent_false.hpp" #include "cute/atom/mma_traits_sm90_gmma.hpp" #include "cute/atom/copy_traits_sm90_tma.hpp" @@ -123,6 +127,8 @@ sm90_cluster_shape_to_tma_atom(UnimodalClusterShape) { template constexpr auto make_cp_async_gmem_tiled_copy() { + using namespace cute; + using AlignmentType = cute::uint_byte_t(sizeof(Element)) * Alignment>; constexpr int TileSizeMN = cute::size(TileMN{}); constexpr int TileSizeK = cute::size(TileK{}); @@ -166,9 +172,11 @@ make_cp_async_gmem_tiled_copy() { // or hierarchically // ((BLK_MN0,BLK_MN1,...),(BLK_K0,BLK_K1,...)) // and returns the optimal GMMA::Layout that fits BLK_MN0 and BLK_K0 -template +template constexpr auto rs_smem_selector() { + using namespace cute; + auto BLK_MN0 = size<0>(BLK_MN{}); auto BLK_K0 = size<0>(BLK_K{}); @@ -260,18 +268,19 @@ rs_smem_selector() { // or hierarchically // ((BLK_MN0,BLK_MN1,...),(BLK_K0,BLK_K1,...)) // and returns the largest GMMA::Layout that fits BLK_MN0 and BLK_K0 -template +template CUTE_HOST_DEVICE constexpr auto ss_smem_selector() { + using namespace cute; + 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{}; diff --git a/include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl b/include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl index 082ef915..25b1f848 100644 --- a/include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl +++ b/include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl @@ -30,10 +30,6 @@ **************************************************************************************************/ #pragma once -#include "cutlass/arch/mma.h" -#include "cutlass/gemm/gemm.h" -#include "cutlass/gemm/dispatch_policy.hpp" - #include "cutlass/gemm/collective/builders/sm90_common.inl" // SM90 Collective Builders should be used only starting CUDA 12.0 @@ -67,14 +63,13 @@ compute_stage_count_or_override(cute::Int stage_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_bits = static_cast(sizeof_bits::value); - constexpr int b_bits = static_cast(sizeof_bits::value); + constexpr auto mainloop_pipeline_bytes = sizeof(typename cutlass::PipelineTmaAsync<1>::SharedStorage); + constexpr auto a_bits = cute::sizeof_bits_v; + constexpr auto b_bits = cute::sizeof_bits_v; constexpr int stage_bytes = - (a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) / 8 + - (b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) / 8 + - stage_barrier_bytes; + cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + static_cast(mainloop_pipeline_bytes); return (CapacityBytes - carveout_bytes) / stage_bytes; } @@ -102,23 +97,23 @@ constexpr int compute_stage_count_or_override_single_affine_transformed_input(StageCountAutoCarveout stage_count) { // 32 bytes to account for barriers etc. - constexpr int stage_barrier_bytes = 32; + constexpr auto mainloop_pipeline_bytes = sizeof(typename cutlass::PipelineTmaAsync<1>::SharedStorage); constexpr int scale_zero_k_tile = 1; - constexpr int a_bits = static_cast(sizeof_bits::value); - constexpr int b_bits = static_cast(sizeof_bits::value); - constexpr int s_bits = get_bits_for_possibly_void_element(); - constexpr int z_bits = get_bits_for_possibly_void_element(); + constexpr auto a_bits = cute::sizeof_bits_v; + constexpr auto b_bits = cute::sizeof_bits_v; + constexpr auto s_bits = get_bits_for_possibly_void_element(); + constexpr auto z_bits = get_bits_for_possibly_void_element(); - constexpr int scale_bytes = (s_bits * size<0>(TileShapeMNK{}) * scale_zero_k_tile) / 8; - constexpr int zero_bytes = (z_bits * size<0>(TileShapeMNK{}) * scale_zero_k_tile) / 8; + constexpr auto scale_bytes = cutlass::bits_to_bytes(s_bits * size<0>(TileShapeMNK{}) * scale_zero_k_tile); + constexpr auto zero_bytes = cutlass::bits_to_bytes(z_bits * size<0>(TileShapeMNK{}) * scale_zero_k_tile); static_assert(scale_bytes % 128 == 0, "Scale bytes must be a multiple of 128"); static_assert(zero_bytes % 128 == 0, "Zero bytes must be a multiple of 128"); // When scales are void, s_bits will be 0 so no smem will be allocated for scales. constexpr int stage_bytes = - (a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) / 8 + - (b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) / 8 + - scale_bytes + zero_bytes + stage_barrier_bytes; + cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + static_cast(scale_bytes + zero_bytes + mainloop_pipeline_bytes); return (CapacityBytes - carveout_bytes) / stage_bytes; } @@ -156,10 +151,10 @@ is_warpspecialized_transpose_B(){ // GMMA_TMA_WS_SS template < class ElementA, - class GmemLayoutA, + class GmemLayoutATag, int AlignmentA, class ElementB, - class GmemLayoutB, + class GmemLayoutBTag, int AlignmentB, class ElementAccumulator, class TileShape_MNK, @@ -171,10 +166,10 @@ struct CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, ElementA, - GmemLayoutA, + GmemLayoutATag, AlignmentA, ElementB, - GmemLayoutB, + GmemLayoutBTag, AlignmentB, ElementAccumulator, TileShape_MNK, @@ -186,7 +181,7 @@ struct CollectiveBuilder< cute::is_same_v || cute::is_same_v || cute::is_same_v) && - not detail::is_use_rmem_A()> + not detail::is_use_rmem_A()> > { static_assert(is_static::value); static_assert(is_static::value); @@ -202,29 +197,29 @@ struct CollectiveBuilder< "Kernel[Array/Group]TmaWarpSpecializedCooperative is only compatible with FP8 FastAccum version right now\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>; + using ElementAMma = cute::conditional_t, tfloat32_t, ElementA>; + using ElementBMma = 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(); + 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< cute::is_same_v || IsArrayOfPointersGemm, Layout>, Layout>>; using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::ss_op_selector< - MmaElementA, MmaElementB, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>(), AtomLayoutMNK{})); + ElementAMma, ElementBMma, 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{}))>()); + GmmaMajorA, ElementAMma, 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{}))>()); + GmmaMajorB, ElementBMma, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); static constexpr int PipelineStages = detail::compute_stage_count_or_override(StageCountType{}); + ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{}); using DispatchPolicy = cute::conditional_t, /* For FP8 use a separate mainloop compared to other datatypes */ @@ -239,9 +234,9 @@ struct CollectiveBuilder< DispatchPolicy, TileShape_MNK, ElementA, - TagToStrideA_t, + TagToStrideA_t, ElementB, - TagToStrideB_t, + TagToStrideB_t, TiledMma, GmemTiledCopyA, SmemLayoutAtomA, @@ -259,10 +254,10 @@ struct CollectiveBuilder< // GMMA_TMA_WS_RS template < class ElementA, - class GmemLayoutA, + class GmemLayoutATag, int AlignmentA, class ElementB, - class GmemLayoutB, + class GmemLayoutBTag, int AlignmentB, class ElementAccumulator, class TileShape_MNK, @@ -274,10 +269,10 @@ struct CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, ElementA, - GmemLayoutA, + GmemLayoutATag, AlignmentA, ElementB, - GmemLayoutB, + GmemLayoutBTag, AlignmentB, ElementAccumulator, TileShape_MNK, @@ -288,7 +283,7 @@ struct CollectiveBuilder< (cute::is_same_v || cute::is_same_v || cute::is_same_v) && - detail::is_use_rmem_A()> + detail::is_use_rmem_A()> > { static_assert(is_static::value); static_assert(is_static::value); @@ -297,32 +292,32 @@ struct CollectiveBuilder< #ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED static_assert(cutlass::detail::dependent_false, "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 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>(); + ElementA, GmemLayoutATag, ElementB, GmemLayoutBTag, 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 ElementAMma = cute::conditional_t, tfloat32_t, ElementA>; + using ElementBMma = 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{})); + ElementAMma, ElementBMma, 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{}); + ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{}); using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecialized< PipelineStages, ClusterShape_MNK, KernelScheduleType>; @@ -334,9 +329,9 @@ struct CollectiveBuilder< DispatchPolicy, TileShape_MNK, ElementA, - TagToStrideA_t, + TagToStrideA_t, ElementB, - TagToStrideB_t, + TagToStrideB_t, TiledMma, GmemTiledCopyA, SmemLayoutAtomA, @@ -354,10 +349,10 @@ struct CollectiveBuilder< // GMMA_TMA_WS_RS Mixed Scaled GEMM template < class ElementPairA_, - class GmemLayoutA_, + class GmemLayoutATag_, int AlignmentA, class ElementPairB_, - class GmemLayoutB_, + class GmemLayoutBTag_, int AlignmentB, class ElementAccumulator, class TileShape_MNK, @@ -369,10 +364,10 @@ struct CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, ElementPairA_, - GmemLayoutA_, + GmemLayoutATag_, AlignmentA, ElementPairB_, - GmemLayoutB_, + GmemLayoutBTag_, AlignmentB, ElementAccumulator, TileShape_MNK, @@ -401,8 +396,8 @@ public: static constexpr bool IsANarrow = sizeof_bits::value < sizeof_bits::value; - using GmemLayoutA = GmemLayoutA_; - using GmemLayoutB = GmemLayoutB_; + using GmemLayoutATag = GmemLayoutATag_; + using GmemLayoutBTag = GmemLayoutBTag_; using ElementPairA = cute::conditional_t, ElementPairA_>; using ElementPairB = cute::conditional_t, ElementPairB_>; @@ -418,10 +413,10 @@ public: #ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED static_assert(cutlass::detail::dependent_false, "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 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 IsWarpSpecializedTransposeB = detail::is_warpspecialized_transpose_B< - ElementA, GmemLayoutA, ElementB, GmemLayoutB, KernelScheduleType>(); + ElementA, GmemLayoutATag, ElementB, GmemLayoutBTag, KernelScheduleType>(); static_assert(!IsWarpSpecializedTransposeB, "Mixed input GEMM does not support WS transpose B."); // If A is scaled, then we don't need to swap. Otherwise, we must ensure B goes to RF and we must swap the operands. @@ -456,8 +451,8 @@ public: using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput; // We pack the scale data with the operand that will be optionally scaled and converted before MMA. - using StrideA = TagToStrideA_t; - using StrideB = TagToStrideB_t; + using StrideA = TagToStrideA_t; + using StrideB = TagToStrideB_t; using CollectiveOp = CollectiveMma< DispatchPolicy, @@ -484,10 +479,10 @@ public: // GMMA_TMA_WS_FP8_FAST_ACCUM_SS template < class ElementA, - class GmemLayoutA, + class GmemLayoutATag, int AlignmentA, class ElementB, - class GmemLayoutB, + class GmemLayoutBTag, int AlignmentB, class ElementAccumulator, class TileShape_MNK, @@ -499,10 +494,10 @@ struct CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, ElementA, - GmemLayoutA, + GmemLayoutATag, AlignmentA, ElementB, - GmemLayoutB, + GmemLayoutBTag, AlignmentB, ElementAccumulator, TileShape_MNK, @@ -522,14 +517,14 @@ struct CollectiveBuilder< static_assert(detail::is_input_fp8(), "Only FP8 datatypes are compatible with these kernel schedules\n"); // Dispatch TN fp8 kernels only to TMA warp specialized FP8 builder - static_assert(!detail::is_use_rmem_A(), + static_assert(!detail::is_use_rmem_A(), "Not supported for fp8 non-TN warp specialized kernels yet\n"); #ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED static_assert(cutlass::detail::dependent_false, "Unsupported Toolkit for SM90 Collective Builder\n"); #endif - 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(); + 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(); static constexpr bool IsArrayOfPointersGemm = (cute::is_same_v); using AtomLayoutMNK = cute::conditional_t || @@ -560,9 +555,9 @@ struct CollectiveBuilder< DispatchPolicy, TileShape_MNK, ElementA, - TagToStrideA_t, + TagToStrideA_t, ElementB, - TagToStrideB_t, + TagToStrideB_t, TiledMma, GmemTiledCopyA, SmemLayoutAtomA, @@ -580,10 +575,10 @@ struct CollectiveBuilder< // GMMA_TMA_SS template < class ElementA, - class GmemLayoutA, + class GmemLayoutATag, int AlignmentA, class ElementB, - class GmemLayoutB, + class GmemLayoutBTag, int AlignmentB, class ElementAccumulator, class TileShape_MNK, @@ -595,10 +590,10 @@ struct CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, ElementA, - GmemLayoutA, + GmemLayoutATag, AlignmentA, ElementB, - GmemLayoutB, + GmemLayoutBTag, AlignmentB, ElementAccumulator, TileShape_MNK, @@ -606,7 +601,7 @@ struct CollectiveBuilder< StageCountType, KernelScheduleType, cute::enable_if_t && - not detail::is_use_rmem_A()> + not detail::is_use_rmem_A()> > { static_assert(is_static::value); static_assert(is_static::value); @@ -617,25 +612,25 @@ struct CollectiveBuilder< #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>; + using ElementAMma = cute::conditional_t, tfloat32_t, ElementA>; + using ElementBMma = 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(); + 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>())); + ElementAMma, ElementBMma, 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{}))>()); + GmmaMajorA, ElementAMma, 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{}))>()); + GmmaMajorB, ElementBMma, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); static constexpr int PipelineStages = detail::compute_stage_count_or_override(StageCountType{}); + ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{}); using DispatchPolicy = MainloopSm90TmaGmma; using SmemCopyAtomA = void; @@ -645,9 +640,9 @@ struct CollectiveBuilder< DispatchPolicy, TileShape_MNK, ElementA, - TagToStrideA_t, + TagToStrideA_t, ElementB, - TagToStrideB_t, + TagToStrideB_t, TiledMma, GmemTiledCopyA, SmemLayoutAtomA, @@ -667,10 +662,10 @@ struct CollectiveBuilder< // GMMA_CpAsync template < class ElementA, - class GmemLayoutA, + class GmemLayoutATag, int AlignmentA, class ElementB, - class GmemLayoutB, + class GmemLayoutBTag, int AlignmentB, class ElementAccumulator, class TileShape_MNK, @@ -683,10 +678,10 @@ CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, ElementA, - GmemLayoutA, + GmemLayoutATag, AlignmentA, ElementB, - GmemLayoutB, + GmemLayoutBTag, AlignmentB, ElementAccumulator, TileShape_MNK, @@ -701,10 +696,10 @@ CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, ElementA, - GmemLayoutA, + GmemLayoutATag, AlignmentA, ElementB, - GmemLayoutB, + GmemLayoutBTag, AlignmentB, ElementAccumulator, TileShape_MNK, @@ -719,10 +714,10 @@ CollectiveBuilder< // GMMA_CpAsync_WS_SS template < class ElementA, - class GmemLayoutA, + class GmemLayoutATag, int AlignmentA, class ElementB, - class GmemLayoutB, + class GmemLayoutBTag, int AlignmentB, class ElementAccumulator, class TileShape_MNK, @@ -734,10 +729,10 @@ struct CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, ElementA, - GmemLayoutA, + GmemLayoutATag, AlignmentA, ElementB, - GmemLayoutB, + GmemLayoutBTag, AlignmentB, ElementAccumulator, TileShape_MNK, @@ -748,7 +743,7 @@ struct CollectiveBuilder< (cute::is_same_v || cute::is_same_v || cute::is_same_v) && - not detail::is_use_rmem_A() + not detail::is_use_rmem_A() > > { static_assert(is_static::value); @@ -758,37 +753,37 @@ struct CollectiveBuilder< #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>; + using ElementAMma = cute::conditional_t, tfloat32_t, ElementA>; + using ElementBMma = 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(); + 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(TileShape_MNK{}) < 128) ? 1 : 2>,_1,_1>>, Layout>>; using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::ss_op_selector< - MmaElementA, MmaElementB, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>(), AtomLayoutMNK{})); + ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>(), AtomLayoutMNK{})); static constexpr int NumLoadWarpGroups = cute::is_same_v ? 2 : 1; using GmemTiledCopyA = decltype(detail::make_cp_async_gmem_tiled_copy< - NumThreadsPerWarpGroup * NumLoadWarpGroups, ElementA, AlignmentA, TagToStrideA_t, + NumThreadsPerWarpGroup * NumLoadWarpGroups, 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< - NumThreadsPerWarpGroup * NumLoadWarpGroups, ElementB, AlignmentB, TagToStrideB_t, + NumThreadsPerWarpGroup * NumLoadWarpGroups, 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{}))>()); + GmmaMajorA, ElementAMma, 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{}))>()); + GmmaMajorB, ElementBMma, 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{}); + detail::sm90_smem_capacity_bytes, ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{}); using DispatchPolicy = MainloopSm90CpAsyncGmmaWarpSpecialized< PipelineStages, ClusterShape_MNK, KernelScheduleType>; @@ -797,9 +792,9 @@ struct CollectiveBuilder< DispatchPolicy, TileShape_MNK, ElementA, - TagToStrideA_t, + TagToStrideA_t, ElementB, - TagToStrideB_t, + TagToStrideB_t, TiledMma, GmemTiledCopyA, SmemLayoutAtomA, @@ -817,10 +812,10 @@ struct CollectiveBuilder< // GMMA_CpAsync_WS_RS template < class ElementA, - class GmemLayoutA, + class GmemLayoutATag, int AlignmentA, class ElementB, - class GmemLayoutB, + class GmemLayoutBTag, int AlignmentB, class ElementAccumulator, class TileShape_MNK, @@ -832,10 +827,10 @@ struct CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, ElementA, - GmemLayoutA, + GmemLayoutATag, AlignmentA, ElementB, - GmemLayoutB, + GmemLayoutBTag, AlignmentB, ElementAccumulator, TileShape_MNK, @@ -846,7 +841,7 @@ struct CollectiveBuilder< (cute::is_same_v || cute::is_same_v || cute::is_same_v) && - detail::is_use_rmem_A() + detail::is_use_rmem_A() > > { static_assert(is_static::value); @@ -856,40 +851,40 @@ struct CollectiveBuilder< #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>; + using ElementAMma = cute::conditional_t, tfloat32_t, ElementA>; + using ElementBMma = 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_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 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>(); + ElementA, GmemLayoutATag, ElementB, GmemLayoutBTag, KernelScheduleType>(); using AtomLayoutMNK = cute::conditional_t, Layout(TileShape_MNK{}) < 128) ? 1 : 2>,_1,_1>>, 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{})); + ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK, GMMA::Major::K, GMMA::Major::K>(), AtomLayoutMNK{})); static constexpr int NumLoadWarpGroups = 1; using GmemTiledCopyA = decltype(detail::make_cp_async_gmem_tiled_copy< - NumThreadsPerWarpGroup * NumLoadWarpGroups, ElementA, AlignmentA, TagToStrideA_t, + NumThreadsPerWarpGroup * NumLoadWarpGroups, 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< - NumThreadsPerWarpGroup * NumLoadWarpGroups, ElementB, AlignmentB, TagToStrideB_t, + NumThreadsPerWarpGroup * NumLoadWarpGroups, ElementB, AlignmentB, TagToStrideB_t, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_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< - detail::sm90_smem_capacity_bytes, MmaElementA, MmaElementB, TileShape_MNK>(StageCountType{}); + detail::sm90_smem_capacity_bytes, ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{}); using DispatchPolicy = MainloopSm90CpAsyncGmmaRmemAWarpSpecialized< PipelineStages, ClusterShape_MNK, KernelScheduleType>; @@ -901,9 +896,9 @@ struct CollectiveBuilder< DispatchPolicy, TileShape_MNK, ElementA, - TagToStrideA_t, + TagToStrideA_t, ElementB, - TagToStrideB_t, + TagToStrideB_t, TiledMma, GmemTiledCopyA, SmemLayoutAtomA, @@ -921,10 +916,10 @@ struct CollectiveBuilder< // GMMA auto kernel schedule template < class ElementA, - class GmemLayoutA, + class GmemLayoutATag, int AlignmentA, class ElementB, - class GmemLayoutB, + class GmemLayoutBTag, int AlignmentB, class ElementAccumulator, class TileShape_MNK, @@ -936,10 +931,10 @@ struct CollectiveBuilder< arch::Sm90, arch::OpClassTensorOp, ElementA, - GmemLayoutA, + GmemLayoutATag, AlignmentA, ElementB, - GmemLayoutB, + GmemLayoutBTag, AlignmentB, ElementAccumulator, TileShape_MNK, @@ -988,10 +983,10 @@ static constexpr bool IsMixedWidthInput = IsDifferentWidth || (IsDifferentWidth arch::Sm90, arch::OpClassTensorOp, ElementA, - GmemLayoutA, + GmemLayoutATag, AlignmentA, ElementB, - GmemLayoutB, + GmemLayoutBTag, AlignmentB, ElementAccumulator, TileShape_MNK, diff --git a/include/cutlass/gemm/collective/collective_builder.hpp b/include/cutlass/gemm/collective/collective_builder.hpp index 8c3e44bf..d86fda53 100644 --- a/include/cutlass/gemm/collective/collective_builder.hpp +++ b/include/cutlass/gemm/collective/collective_builder.hpp @@ -39,10 +39,20 @@ namespace cutlass::gemm::collective { // Used to specify stage counts or dispatch to automatic computation of stage count template -struct StageCount { static constexpr int value = num_stages; }; +struct StageCount { + static constexpr int value = num_stages; + + StageCount() = default; + explicit StageCount(cute::Int) {} +}; template -struct StageCountAutoCarveout { static constexpr int bytes = carveout_bytes; }; +struct StageCountAutoCarveout { + static constexpr int bytes = carveout_bytes; + + StageCountAutoCarveout() = default; + explicit StageCountAutoCarveout(cute::Int) {} +}; using StageCountAuto = StageCountAutoCarveout<0>; diff --git a/include/cutlass/gemm/collective/sm80_mma_multistage.hpp b/include/cutlass/gemm/collective/sm80_mma_multistage.hpp index 304c57d4..a0038cdd 100644 --- a/include/cutlass/gemm/collective/sm80_mma_multistage.hpp +++ b/include/cutlass/gemm/collective/sm80_mma_multistage.hpp @@ -100,6 +100,9 @@ struct CollectiveMma< using TransformA = TransformA_; using TransformB = TransformB_; using ArchTag = typename DispatchPolicy::ArchTag; + // Follow the change in TestSmall: TileShape switch to CtaShape + // For sm80 arch, CtaShape should euqal to TileShape + using CtaShape_MNK = TileShape; static_assert(cute::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."); @@ -332,6 +335,8 @@ struct CollectiveMma< } + cp_async_wait<0>(); + __syncthreads(); } }; @@ -377,6 +382,9 @@ struct CollectiveMma< // using DispatchPolicy = MainloopSm80CpAsync; using TileShape = TileShape_; + // Follow the change in TestSmall: TileShape switch to CtaShape + // In legacy arch, it should be same + using CtaShape_MNK = TileShape; using ElementA = ElementA_; using StrideA = StrideA_; using ElementB = ElementB_; @@ -391,7 +399,6 @@ struct CollectiveMma< using TransformA = TransformA_; using TransformB = TransformB_; using ArchTag = typename DispatchPolicy::ArchTag; - static_assert(cute::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."); @@ -680,6 +687,8 @@ struct CollectiveMma< } + cp_async_wait<0>(); + __syncthreads(); } }; diff --git a/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp b/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp index 5ae84307..16255d70 100644 --- a/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp @@ -31,17 +31,18 @@ #pragma once #include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/trace.h" + #include "cute/arch/cluster_sm90.hpp" #include "cute/arch/copy_sm90.hpp" -#include "cutlass/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 "cutlass/pipeline/pipeline.hpp" -#include "cutlass/trace.h" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -113,7 +114,7 @@ struct CollectiveMma< using PipelineState = cutlass::PipelineState; using PipelineParams = typename MainloopPipeline::Params; - + using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{})); 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."); @@ -126,11 +127,11 @@ struct CollectiveMma< using SmemLayoutA = decltype(tile_to_shape( SmemLayoutAtomA{}, make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); using SmemLayoutB = decltype(tile_to_shape( SmemLayoutAtomB{}, make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more."); static_assert(cute::is_base_of::value && @@ -312,9 +313,8 @@ struct CollectiveMma< 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_bits::value)) / 8+ - (size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(sizeof_bits::value)) / 8; - + cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(sizeof_bits::value))+ + cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(sizeof_bits::value)); // Set up the data needed by this collective for load and mma. // Returns a tuple of tensors. The collective and the kernel layer have the contract that the diff --git a/include/cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp b/include/cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp index 4dc43022..bca824ae 100644 --- a/include/cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp @@ -31,18 +31,19 @@ #pragma once #include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/transform/collective/sm90_wgmma_transpose.hpp" +#include "cutlass/trace.h" + #include "cute/arch/cluster_sm90.hpp" #include "cute/arch/copy_sm90.hpp" -#include "cutlass/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 "cutlass/pipeline/pipeline.hpp" -#include "cutlass/transform/collective/sm90_wgmma_transpose.hpp" -#include "cutlass/trace.h" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -106,6 +107,7 @@ struct CollectiveMma< using SmemCopyAtomA = SmemCopyAtomA_; using SmemCopyAtomB = SmemCopyAtomB_; + using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{})); // 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, Fp8 WGMMA) static constexpr bool IsLayoutAkBmn = cute::is_same_v, layout::RowMajor> && @@ -178,7 +180,7 @@ struct CollectiveMma< static_assert(!SwapAB || !TransposeB, "Cannot SwapAB and TransposeB at the same time."); static_assert(TransposeB xor (cute::is_same_v), "Should be same layout if not TransposeB."); - static_assert(!TransposeB || ((size<1>(SmemLayoutB{}) * sizeof_bits::value) / 8) == 128, + static_assert(!TransposeB || (cutlass::bits_to_bytes(size<1>(SmemLayoutB{}) * sizeof_bits::value)) == 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."); diff --git a/include/cutlass/gemm/collective/sm90_mma_multistage_gmma_ss_warpspecialized.hpp b/include/cutlass/gemm/collective/sm90_mma_multistage_gmma_ss_warpspecialized.hpp index 6d2ec0f2..25a3ca4b 100644 --- a/include/cutlass/gemm/collective/sm90_mma_multistage_gmma_ss_warpspecialized.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_multistage_gmma_ss_warpspecialized.hpp @@ -108,6 +108,7 @@ struct CollectiveMma< using TransformB = TransformB_; using ArchTag = typename DispatchPolicy::ArchTag; + using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{})); using MainloopPipeline = cutlass::PipelineAsync; using PipelineState = typename MainloopPipeline::PipelineState; using PipelineParams = typename MainloopPipeline::Params; diff --git a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp index 83102382..dbac39a1 100644 --- a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp @@ -31,21 +31,22 @@ #pragma once #include "cutlass/cutlass.h" -#include "cute/arch/cluster_sm90.hpp" -#include "cute/arch/copy_sm90.hpp" #include "cutlass/gemm/gemm.h" #include "cutlass/detail/dependent_false.hpp" #include "cutlass/gemm/dispatch_policy.hpp" #include "cutlass/detail/layout.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/transform/collective/sm90_wgmma_transpose.hpp" +#include "cutlass/trace.h" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/arch/copy_sm90.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 "cutlass/pipeline/pipeline.hpp" -#include "cutlass/transform/collective/sm90_wgmma_transpose.hpp" -#include "cutlass/trace.h" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -108,6 +109,7 @@ struct CollectiveMma< using SmemCopyAtomA = SmemCopyAtomA_; using SmemCopyAtomB = SmemCopyAtomB_; + using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{})); // 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, Fp8 WGMMA) static constexpr bool IsLayoutAkBmn = cute::is_same_v, layout::RowMajor> && @@ -151,11 +153,11 @@ struct CollectiveMma< using SmemLayoutA = decltype(tile_to_shape( InternalSmemLayoutAtomA{}, make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); using SmemLayoutB = decltype(tile_to_shape( InternalSmemLayoutAtomB{}, make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideB>(), Step<_2,_1,_3>, Step<_1,_2,_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 = @@ -182,12 +184,12 @@ struct CollectiveMma< using GmmaSmemLayoutB = decltype(tile_to_shape( GmmaSmemLayoutAtomB{}, make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); static_assert(!SwapAB || !TransposeB, "Cannot SwapAB and TransposeB at the same time."); static_assert(TransposeB xor (cute::is_same_v), "Should be same layout if not TransposeB."); - static_assert(!TransposeB || (((size<1>(SmemLayoutB{}) * sizeof_bits::value)) / 8) == 128, + static_assert(!TransposeB || (cutlass::bits_to_bytes((size<1>(SmemLayoutB{}) * sizeof_bits::value))) == 128, "SmemLayoutB K must be 128bytes to be transposed."); static constexpr bool uses_universal_transposition() { @@ -329,8 +331,8 @@ struct CollectiveMma< static constexpr int K_PIPE_MAX = DispatchPolicy::Stages; static constexpr uint32_t TmaTransactionBytes = - (size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(sizeof_bits::value)) / 8 + - (size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(sizeof_bits::value)) / 8 ; + cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(sizeof_bits::value)) + + cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(sizeof_bits::value)) ; /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance CUTLASS_DEVICE @@ -547,6 +549,7 @@ struct CollectiveMma< 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 + CUTE_STATIC_ASSERT_V(size<2>(tCrA) > _2{}, "RS loops require more than 2 MMA k-iterations for correctness."); // // PIPELINED MAIN LOOP diff --git a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp index da80284d..29ee9605 100644 --- a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp @@ -32,13 +32,18 @@ #include "cutlass/cutlass.h" #include "cutlass/numeric_conversion.h" -#include "cute/arch/cluster_sm90.hpp" -#include "cute/arch/copy_sm90.hpp" #include "cutlass/gemm/gemm.h" #include "cutlass/detail/dependent_false.hpp" #include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/numeric_types.h" #include "cutlass/detail/layout.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/transform/collective/sm90_wgmma_transpose.hpp" +#include "cutlass/trace.h" +#include "cutlass/detail/collective.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/arch/copy_sm90.hpp" #include "cute/algorithm/functional.hpp" #include "cute/atom/mma_atom.hpp" #include "cute/atom/copy_traits_sm90_tma.hpp" @@ -46,9 +51,7 @@ #include "cute/tensor_predicate.hpp" #include "cute/numeric/arithmetic_tuple.hpp" #include "cutlass/pipeline/pipeline.hpp" -#include "cutlass/transform/collective/sm90_wgmma_transpose.hpp" #include "cutlass/trace.h" - #include "cutlass/detail/collective.hpp" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -156,6 +159,7 @@ public: static_assert(cutlass::gemm::detail::is_mn_major(), "Scale must be MN major [Col Major if A is scaled, Row Major if B is scaled]."); + using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{})); using TiledMma = TiledMma_; using ElementAccumulator = typename TiledMma::ValTypeC; @@ -172,13 +176,6 @@ public: using SmemCopyAtomB = SmemCopyAtomB_; using SmemCopyAtomScale = Copy_Atom; - // 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, Fp8 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; - // We must ensure the type to be scaled goes to RF static constexpr bool SwapAB = !IsATransformed; using InternalSmemLayoutAtomA = cute::conditional_t; @@ -233,26 +230,17 @@ public: using SmemLayoutA = decltype(tile_to_shape( InternalSmemLayoutAtomA{}, make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); using SmemLayoutB = decltype(tile_to_shape( InternalSmemLayoutAtomB{}, make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); // It is assumed that the scales and zero-points share the same smem layout using SmemLayoutScale = decltype(tile_to_shape( SmemLayoutAtomScale{}, make_shape(shape<0>(ScaleTileShape{}), shape<1>(ScaleTileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,NonVoidStrideScale>(), Step<_2,_1,_3>, Step<_1,_2,_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; - using TransposeOperandB = decltype(cutlass::transform::collective::detail::make_transpose_operand_b( - 0, 0, TiledMma{}, SmemLayoutB{}, InternalSmemLayoutAtomB{}, - InternalElementB{}, cute::bool_constant{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,NonVoidStrideScale>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more."); static_assert(not cute::is_base_of::value && @@ -263,27 +251,11 @@ public: static_assert(cute::is_same_v || cute::is_same_v, "GmemTiledCopy - invalid SM90 TMA copy atom specified."); - 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{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); - - // These two restrictions are related, so we place the assertions together. // To relax them, we need to handle loading more than 1 row of scales for every main loop iteration. // We must also handle updating the pipeline transaction bytes on the fly. // NOTE: Deleting this assertion without required changes will cause the code to hang. static_assert(size<1>(SmemLayoutAtomScale{}) == 1, "size<1>(SmemLayoutAtomScale) must be 1."); - static_assert(!SwapAB || !TransposeB, "Cannot SwapAB and TransposeB at the same time."); - static_assert(TransposeB xor (cute::is_same_v), - "Should be same layout if not TransposeB."); - static_assert(!TransposeB || size<1>(SmemLayoutB{}) * cute::sizeof_bits_v / 8 == 128, - "SmemLayoutB K must be 128bytes to be transposed."); - private: static constexpr ConversionMode get_conversion_mode() { @@ -332,8 +304,8 @@ private: // These methods use some the public members of the class. For that reason, we define them after the public section. static constexpr uint32_t compute_tma_transaction_bytes() { - constexpr uint32_t a_bytes = (size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(cute::sizeof_bits_v) / 8); - constexpr uint32_t b_bytes = (size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(cute::sizeof_bits_v) / 8); + constexpr uint32_t a_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(cute::sizeof_bits_v)); + constexpr uint32_t b_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(cute::sizeof_bits_v)); constexpr uint32_t baseline_bytes = a_bytes + b_bytes; @@ -341,14 +313,14 @@ private: return baseline_bytes; } else if constexpr (ModeHasScales) { - constexpr uint32_t scale_tx_bytes = (size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast(cute::sizeof_bits_v) / 8); + constexpr uint32_t scale_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast(cute::sizeof_bits_v)); static_assert(scale_tx_bytes % 128 == 0, "Each scale stage must be 128B aligned."); // required by TMA if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) { return baseline_bytes + scale_tx_bytes; } else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) { // Scale and zero share smem layout - constexpr uint32_t zero_tx_bytes = (size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast(cute::sizeof_bits_v) / 8); + constexpr uint32_t zero_tx_bytes = cutlass::bits_to_bytes(size<0>(SmemLayoutScale{}) * size<1>(SmemLayoutScale{}) * static_cast(cute::sizeof_bits_v)); static_assert(zero_tx_bytes % 128 == 0, "Each zero stage must be 128B aligned."); // required by TMA return baseline_bytes + scale_tx_bytes + zero_tx_bytes; } @@ -851,13 +823,7 @@ public: Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE) - Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) - Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_M,BLK_K,PIPE) - - // If TransposeB, GMMA will read from transposed B layout SMEM - Tensor gmma_sB_position_dependent = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), - GmmaSmemLayoutB{}); // (BLK_N,BLK_K,PIPE) - Tensor gmma_sB = as_position_independent_swizzle_tensor(gmma_sB_position_dependent); // (BLK_N,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) // // Define C accumulators and A/B partitioning @@ -871,7 +837,7 @@ public: Tensor tCrA_mma = thread_mma.partition_fragment_A(sA(_,_,Int<0>{})); // (MMA,MMA_M,MMA_K,PIPE) Tensor tCrA_load = make_fragment_like(tCrA_mma); - Tensor tCsB = thread_mma.partition_B(gmma_sB_position_dependent); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tCsB = thread_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE) Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE) // @@ -911,12 +877,9 @@ public: tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; - TransposeOperandB transpose = cutlass::transform::collective::detail::make_transpose_operand_b( - warp_idx, warp_group_thread_idx, tiled_mma, SmemLayoutB{}, - InternalSmemLayoutAtomB{}, InternalElementB{}, - cute::bool_constant{}); - warpgroup_fence_operand(accum); + + constexpr int K_BLOCK_MAX = size<2>(tCrA_load); ConsumerToken barrier_token = {BarrierStatus::WaitAgain}; // first k tile @@ -934,41 +897,31 @@ public: partitioned_extra_info, copy_partitions_extra_info, 0, read_stage); transform_A_kblock(tCrA_load, A_CPY_VEC{}, tCrA_mma, partitioned_extra_info, 0); - // transpose B operand in SMEM - transpose(sB, gmma_sB, read_stage, 0); // Unroll the K mode manually to set scale D to 1 CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA_load) - 1; ++k_block) { - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, - partitioned_extra_info, copy_partitions_extra_info, k_block + 1, read_stage); - transform_A_kblock(tCrA_load, A_CPY_VEC{}, tCrA_mma, partitioned_extra_info, k_block + 1); - transpose.synchronize(k_block); - transpose(sB, gmma_sB, read_stage, k_block + 1); + for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) { + if (k_block < K_BLOCK_MAX - 1) { + copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + partitioned_extra_info, copy_partitions_extra_info, k_block + 1, read_stage); + transform_A_kblock(tCrA_load, A_CPY_VEC{}, tCrA_mma, partitioned_extra_info, k_block + 1); + } warpgroup_arrive(); // (V,M) x (V,N) => (V,M,N) cute::gemm(tiled_mma, tCrA_mma(_,_,k_block), tCrB(_,_,k_block,read_stage), accum); tiled_mma.accumulate_ = GMMA::ScaleOut::One; warpgroup_commit_batch(); - } + } - warpgroup_wait<2>(); - --k_tile_count; if (k_tile_count > 0) { + // Wait for K_BLOCK_MAX - 1 to be in flight to ensure that it is safe to overwrite the A registers for the first mma. + warpgroup_wait(); pipeline.consumer_wait(smem_pipe_read, barrier_token); copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, partitioned_extra_info, copy_partitions_extra_info, 0, smem_pipe_read.index()); transform_A_kblock(tCrA_load, A_CPY_VEC{}, tCrA_mma, partitioned_extra_info, 0); - transpose(sB, gmma_sB, smem_pipe_read.index(), 0); } - warpgroup_arrive(); - // (V,M) x (V,N) => (V,M,N) - const int final_k = size<2>(tCrA_load) - 1; - cute::gemm(tiled_mma, tCrA_mma(_,_, final_k), tCrB(_,_,final_k,read_stage), accum); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - warpgroup_commit_batch(); - warpgroup_wait<2>(); } if (k_tile_count == 0) { @@ -990,38 +943,36 @@ public: warpgroup_fence_operand(accum); // Unroll the K mode manually to set scale D to 1 CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA_load); ++k_block) { - if (k_block == 0) { - barrier_token = pipeline.consumer_try_wait(smem_pipe_read); - } - if (k_block == size<2>(tCrA_load) - 1) { - pipeline.consumer_wait(smem_pipe_read, barrier_token); - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, - partitioned_extra_info, copy_partitions_extra_info, 0, smem_pipe_read.index()); - transform_A_kblock(tCrA_load, A_CPY_VEC{}, tCrA_mma, partitioned_extra_info, 0); - // transpose B operand in SMEM - transpose(sB, gmma_sB, smem_pipe_read.index(), 0); - } - else { - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, - partitioned_extra_info, copy_partitions_extra_info, k_block + 1, read_stage); - transform_A_kblock(tCrA_load, A_CPY_VEC{}, tCrA_mma, partitioned_extra_info, k_block + 1); - // transpose B operand in SMEM - transpose.synchronize(k_block); // make transpose of k_block available - transpose(sB, gmma_sB, read_stage, k_block + 1); - } + for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) { warpgroup_arrive(); // (V,M) x (V,N) => (V,M,N) cute::gemm(tiled_mma, tCrA_mma(_,_,k_block), tCrB(_,_,k_block,read_stage), accum); tiled_mma.accumulate_ = GMMA::ScaleOut::One; warpgroup_commit_batch(); - warpgroup_wait<2>(); - if (k_block == 1) { - // release prior barrier + + warpgroup_wait(); + if (k_block == K_BLOCK_MAX - 1) { + // We have K_BLOCK_MAX - 1 GMMA instructions pending for this stage, so we can release prior barrier pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it ++smem_pipe_release; } + + if (k_block == 0) { + barrier_token = pipeline.consumer_try_wait(smem_pipe_read); + } + + if (k_block == K_BLOCK_MAX - 1) { + pipeline.consumer_wait(smem_pipe_read, barrier_token); + copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + partitioned_extra_info, copy_partitions_extra_info, 0, smem_pipe_read.index()); + transform_A_kblock(tCrA_load, A_CPY_VEC{}, tCrA_mma, partitioned_extra_info, 0); + } + else { + copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + partitioned_extra_info, copy_partitions_extra_info, k_block + 1, read_stage); + transform_A_kblock(tCrA_load, A_CPY_VEC{}, tCrA_mma, partitioned_extra_info, k_block + 1); + } } warpgroup_fence_operand(accum); @@ -1040,32 +991,26 @@ public: // Unroll the K mode manually to set scale D to 1 CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA_load) - 1; ++k_block) { - - copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, - partitioned_extra_info, copy_partitions_extra_info, k_block + 1, read_stage); - transform_A_kblock(tCrA_load, A_CPY_VEC{}, tCrA_mma, partitioned_extra_info, k_block + 1); - transpose.synchronize(k_block); // make k_block transpose available - transpose(sB, gmma_sB, read_stage, k_block + 1); + for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) { + warpgroup_arrive(); // (V,M) x (V,N) => (V,M,N) cute::gemm(tiled_mma, tCrA_mma(_,_,k_block), tCrB(_,_,k_block,read_stage), accum); tiled_mma.accumulate_ = GMMA::ScaleOut::One; warpgroup_commit_batch(); - warpgroup_wait<2>(); - if (k_block == 1) { + warpgroup_wait(); + if (k_block == K_BLOCK_MAX - 1) { // release prior barrier pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it ++smem_pipe_release; } + + if (k_block < K_BLOCK_MAX - 1) { + copy_A_and_extra_info(smem_tiled_copy_A, tCsA, tCrA_copy_view, + partitioned_extra_info, copy_partitions_extra_info, k_block + 1, read_stage); + transform_A_kblock(tCrA_load, A_CPY_VEC{}, tCrA_mma, partitioned_extra_info, k_block + 1); + } } - - warpgroup_arrive(); - // (V,M) x (V,N) => (V,M,N) - const int final_k = size<2>(tCrA_load) - 1; - cute::gemm(tiled_mma, tCrA_mma(_,_,final_k), tCrB(_,_,final_k,read_stage), accum); - tiled_mma.accumulate_ = GMMA::ScaleOut::One; - warpgroup_commit_batch(); } warpgroup_fence_operand(accum); diff --git a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp index a1a0fae8..32589cb6 100644 --- a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp @@ -31,17 +31,18 @@ #pragma once #include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/trace.h" + #include "cute/arch/cluster_sm90.hpp" #include "cute/arch/copy_sm90.hpp" -#include "cutlass/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 "cutlass/pipeline/pipeline.hpp" -#include "cutlass/trace.h" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -106,6 +107,7 @@ struct CollectiveMma< using TransformB = TransformB_; using ArchTag = typename DispatchPolicy::ArchTag; + using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{})); using MainloopPipeline = cutlass::PipelineTmaAsync; using PipelineParams = typename MainloopPipeline::Params; @@ -125,11 +127,11 @@ struct CollectiveMma< using SmemLayoutA = decltype(tile_to_shape( SmemLayoutAtomA{}, make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); using SmemLayoutB = decltype(tile_to_shape( SmemLayoutAtomB{}, make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more."); static_assert(cute::is_base_of::value && @@ -320,9 +322,8 @@ struct CollectiveMma< // 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_bits::value) / 8 + - (size<0>(sB) * size<1>(sB) * sizeof_bits::value) / 8); - + cutlass::bits_to_bytes(size<0>(sA) * size<1>(sA) * sizeof_bits::value) + + cutlass::bits_to_bytes(size<0>(sB) * size<1>(sB) * sizeof_bits::value)); // Obtain warp index int warp_idx = canonical_warp_idx_sync(); diff --git a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp index 8bf8a23a..6ee9bf2b 100644 --- a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp @@ -31,17 +31,18 @@ #pragma once #include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/trace.h" + #include "cute/arch/cluster_sm90.hpp" #include "cute/arch/copy_sm90.hpp" -#include "cutlass/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 "cutlass/pipeline/pipeline.hpp" -#include "cutlass/trace.h" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -107,6 +108,7 @@ struct CollectiveMma< using TransformB = TransformB_; using ArchTag = typename DispatchPolicy::ArchTag; + using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{})); using MainloopPipeline = cutlass::PipelineTmaAsync; using PipelineState = cutlass::PipelineState; @@ -124,11 +126,11 @@ struct CollectiveMma< using SmemLayoutA = decltype(tile_to_shape( SmemLayoutAtomA{}, make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); using SmemLayoutB = decltype(tile_to_shape( SmemLayoutAtomB{}, make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more."); static_assert(cute::is_base_of::value && @@ -248,8 +250,8 @@ struct CollectiveMma< 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_bits::value)) / 8+ - (size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(sizeof_bits::value)) / 8; + cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(sizeof_bits::value))+ + cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(sizeof_bits::value)); /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance CUTLASS_DEVICE diff --git a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp index 8b47fec9..c0e2c907 100644 --- a/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp @@ -32,17 +32,18 @@ #pragma once #include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/collective/fp8_accumulation.hpp" +#include "cutlass/trace.h" +#include "cutlass/numeric_types.h" + #include "cute/arch/cluster_sm90.hpp" #include "cute/arch/copy_sm90.hpp" -#include "cutlass/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 "cutlass/gemm/collective/fp8_accumulation.hpp" -#include "cutlass/trace.h" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -108,6 +109,7 @@ struct CollectiveMma< using TransformB = TransformB_; using ArchTag = typename DispatchPolicy::ArchTag; + using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{})); using MainloopPipeline = cutlass::PipelineTmaAsync; using PipelineState = cutlass::PipelineState; @@ -125,11 +127,11 @@ struct CollectiveMma< using SmemLayoutA = decltype(tile_to_shape( SmemLayoutAtomA{}, make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); using SmemLayoutB = decltype(tile_to_shape( SmemLayoutAtomB{}, make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), - conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more."); static_assert(cute::is_base_of::value && @@ -246,8 +248,8 @@ struct CollectiveMma< 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_bits::value)) / 8+ - (size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(sizeof_bits::value)) / 8; + cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(sizeof_bits::value))+ + cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(sizeof_bits::value)); /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance CUTLASS_DEVICE diff --git a/include/cutlass/gemm/device/default_gemm_configuration.h b/include/cutlass/gemm/device/default_gemm_configuration.h index ad33913d..4197a6b0 100644 --- a/include/cutlass/gemm/device/default_gemm_configuration.h +++ b/include/cutlass/gemm/device/default_gemm_configuration.h @@ -763,6 +763,94 @@ struct DefaultGemmConfiguration< }; //////////////////////////////////////////////////////////////////////////////// + +/// Base configuration for all {fe4m3, fe5m2} x {fe4m3, fe5m2} combinations on SM89 +template < + typename ElementA, + typename ElementB, + typename ElementC, + typename ElementAccumulator> +struct DefaultGemmConfigurationSm89F8 { + static_assert((platform::is_same::value || + platform::is_same::value), + "ElementA must be of type float_e4m3_t or float_e5m2_t"); + static_assert((platform::is_same::value || + platform::is_same::value), + "ElementB must be of type float_e4m3_t or float_e5m2_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::LinearCombination< + ElementC, 128 / sizeof_bits::value, ElementAccumulator, + ElementAccumulator>; + + using Operator = arch::OpMultiplyAdd; +}; + +/// Partial specialization for SM89 fe4m3 x fe4m3 +template +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm89, + cutlass::float_e4m3_t, + cutlass::float_e4m3_t, + ElementC, + ElementAccumulator> : DefaultGemmConfigurationSm89F8< + cutlass::float_e4m3_t, + cutlass::float_e4m3_t, + ElementC, + ElementAccumulator> {}; + +/// Partial specialization for SM89 fe4m3 x fe5m2 +template +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm89, + cutlass::float_e4m3_t, + cutlass::float_e5m2_t, + ElementC, + ElementAccumulator> : DefaultGemmConfigurationSm89F8< + cutlass::float_e4m3_t, + cutlass::float_e5m2_t, + ElementC, + ElementAccumulator> {}; + +/// Partial specialization for SM89 fe5m2 x fe4m3 +template +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm89, + cutlass::float_e5m2_t, + cutlass::float_e4m3_t, + ElementC, + ElementAccumulator> : DefaultGemmConfigurationSm89F8< + cutlass::float_e5m2_t, + cutlass::float_e4m3_t, + ElementC, + ElementAccumulator> {}; + +/// Partial specialization for SM89 fe5m2 x fe5m2 +template +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm89, + cutlass::float_e5m2_t, + cutlass::float_e5m2_t, + ElementC, + ElementAccumulator> : DefaultGemmConfigurationSm89F8< + cutlass::float_e5m2_t, + cutlass::float_e5m2_t, + ElementC, + ElementAccumulator> {}; + +//////////////////////////////////////////////////////////////////////////////// + template struct DefaultGemmConfiguration::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 SparseGemmWithAbsmax { + 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::DefaultSparseGemmWithAbsmax< + 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; + + using Arguments = typename GemmKernel::Arguments; + +private: + + /// Kernel parameters object + typename GemmKernel::Params params_; + +public: + + /// Constructs the GEMM. + SparseGemmWithAbsmax() { } + + /// 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; + + cutlass::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, cudaStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + cutlass::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); + + cudaError_t result = cudaMemsetAsync(workspace, 0, bytes, stream); + + if (result != cudaSuccess) { + 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.ref_Aux, + args.ptr_Vector, + args.ldr, + args.epilogue, + static_cast(workspace) + }; + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + if (smem_size >= (48 << 10)) { + cudaError_t result = cudaFuncSetAttribute(Kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != cudaSuccess) { + 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(cudaStream_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)); + + cutlass::Kernel<<>>(params_); + + cudaError_t result = cudaGetLastError(); + + return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(cudaStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + cudaStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +} // namespace device +} // namespace gemm +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/device/gemm_universal.h b/include/cutlass/gemm/device/gemm_universal.h index 155c6ffb..55413b77 100644 --- a/include/cutlass/gemm/device/gemm_universal.h +++ b/include/cutlass/gemm/device/gemm_universal.h @@ -34,6 +34,7 @@ #pragma once +#include "cutlass/arch/mma.h" #include "cutlass/cutlass.h" #include "cutlass/numeric_types.h" #include "cutlass/arch/arch.h" diff --git a/include/cutlass/gemm/device/gemm_universal_adapter.h b/include/cutlass/gemm/device/gemm_universal_adapter.h index c6acdb54..750b16a6 100644 --- a/include/cutlass/gemm/device/gemm_universal_adapter.h +++ b/include/cutlass/gemm/device/gemm_universal_adapter.h @@ -111,7 +111,7 @@ public: static ComplexTransform const kTransformA = cute::is_same_v ? ComplexTransform::kConjugate : ComplexTransform::kNone; - static ComplexTransform const kTransformB = cute::is_same_v ? + static ComplexTransform const kTransformB = cute::is_same_v ? ComplexTransform::kConjugate : ComplexTransform::kNone; // Legacy: Assume MultiplyAdd only since we do not use this tag type in 3.0 @@ -161,9 +161,9 @@ public: // Inspect TiledCopy for A and B to compute the alignment size static int constexpr kAlignmentA = cutlass::detail::get_alignment_count_from_gmem_tiled_copy< - typename CollectiveMainloop::GmemTiledCopyA, ElementA>(); + typename CollectiveMainloop::GmemTiledCopyA, ElementA, typename CollectiveMainloop::TiledMma::ValTypeA>(); static int constexpr kAlignmentB = cutlass::detail::get_alignment_count_from_gmem_tiled_copy< - typename CollectiveMainloop::GmemTiledCopyB, ElementB>(); + typename CollectiveMainloop::GmemTiledCopyB, ElementB, typename CollectiveMainloop::TiledMma::ValTypeB>(); static int constexpr kAlignmentC = cutlass::detail::get_alignment_count_from_gmem_tiled_copy< typename CollectiveEpilogue::GmemTiledCopyC, ElementC>(); static int constexpr kAlignmentD = cutlass::detail::get_alignment_count_from_gmem_tiled_copy< @@ -278,20 +278,18 @@ public: Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr, - CudaHostAdapter *cuda_adapter = nullptr) { + CudaHostAdapter* cuda_adapter = nullptr) { CUTLASS_TRACE_HOST("GemmUniversal::initialize() - workspace " << workspace << ", stream: " << (stream ? "non-null" : "null")); // Initialize the workspace - Status status = GemmKernel::initialize_workspace(args, workspace, stream); + Status status = GemmKernel::initialize_workspace(args, workspace, stream, cuda_adapter); if (status != Status::kSuccess) { return status; } - // Initialize the Params structure params_ = GemmKernel::to_underlying_arguments(args, workspace); - // Don't set the function attributes - require the CudaHostAdapter to set it. if constexpr (kEnableCudaHostAdapter) { CUTLASS_ASSERT(cuda_adapter); @@ -318,7 +316,6 @@ public: } } } - return Status::kSuccess; } @@ -342,7 +339,6 @@ public: run(Params& params, cudaStream_t stream = nullptr, CudaHostAdapter *cuda_adapter = nullptr) { - CUTLASS_TRACE_HOST("GemmUniversal::run()"); dim3 const block = GemmKernel::get_block_shape(); dim3 const grid = get_grid_shape(params); @@ -350,13 +346,12 @@ public: // configure smem size and carveout int smem_size = GemmKernel::SharedStorageSize; - Status launch_result; + Status launch_result{ Status::kSuccess }; // Use extended launch API only for mainloops that use it - if constexpr(GemmKernel::ArchTag::kMinComputeCapability >= 90) { + 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* kernel_params[] = {¶ms}; if constexpr (kEnableCudaHostAdapter) { @@ -375,13 +370,12 @@ public: } } else { - CUTLASS_ASSERT(cuda_adapter == nullptr); void const* kernel = (void const*) device_kernel; - - launch_result = ClusterLauncher::launch( - grid, cluster, block, smem_size, stream, kernel, kernel_params); - + if constexpr (GemmKernel::ArchTag::kMinComputeCapability == 90) { + launch_result = ClusterLauncher::launch( + grid, cluster, block, smem_size, stream, kernel, kernel_params); + } } } else { @@ -428,7 +422,6 @@ public: cudaStream_t stream = nullptr, CudaHostAdapter *cuda_adapter = nullptr ) { - Status status = initialize(args, workspace, stream, cuda_adapter); if (Status::kSuccess == status) { @@ -556,15 +549,15 @@ public: } /// Determines whether the GEMM can execute the given problem. - static Status can_implement(Arguments const &args) { + static Status can_implement(Arguments const &args, CudaHostAdapter *cuda_adapter = nullptr) { - return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + return UnderlyingOperator::can_implement(to_underlying_arguments(args), cuda_adapter); } /// Gets the workspace size - static size_t get_workspace_size(Arguments const &args) { + static size_t get_workspace_size(Arguments const &args, CudaHostAdapter *cuda_adapter = nullptr) { - return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args), cuda_adapter); } /// Computes the grid shape @@ -604,7 +597,7 @@ public: /// Runs the kernel using initialized state. Status operator()( - cudaStream_t stream = nullptr, + cudaStream_t stream = nullptr, CudaHostAdapter *cuda_adapter = nullptr) { return run(stream); diff --git a/include/cutlass/gemm/device/gemm_universal_with_absmax.h b/include/cutlass/gemm/device/gemm_universal_with_absmax.h new file mode 100644 index 00000000..35f7b541 --- /dev/null +++ b/include/cutlass/gemm/device/gemm_universal_with_absmax.h @@ -0,0 +1,404 @@ +/*************************************************************************************************** + * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE 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 computes the absolute maximum of the output tensor + and applies additional scaling factors to operands. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/arch.h" +#include "cutlass/epilogue/thread/linear_combination_bias_elementwise.h" +#include "cutlass/device_kernel.h" + +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/threadblock/threadblock_swizzle.h" +#include "cutlass/gemm/kernel/gemm_universal.h" + +#include "cutlass/gemm/kernel/default_gemm_universal.h" +#include "cutlass/gemm/kernel/default_gemm_with_absmax.h" +#include "cutlass/gemm/device/default_gemm_configuration.h" +#include "cutlass/gemm/device/gemm_universal_base.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Universal GEMM with absolute-maximum calculation and scaling +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. 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::Sm89, + /// 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_ = cutlass::epilogue::thread::LinearCombinationBiasElementwise< + ElementC_, ElementAccumulator_, ElementAccumulator_, + ElementC_, ElementC_, 128 / cutlass::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 GemmUniversalWithAbsMax; + +// Partial specialization for SM89 +template < + typename ElementA_, + typename LayoutA_, + typename ElementB_, + typename LayoutB_, + typename ElementC_, + typename LayoutC_, + typename ElementAccumulator_, + typename ThreadblockShape_, + typename WarpShape_, + typename InstructionShape_, + typename EpilogueOutputOp_, + typename ThreadblockSwizzle_, + int Stages, + int AlignmentA, + int AlignmentB, + typename Operator_, + ComplexTransform TransformA, + ComplexTransform TransformB +> +class GemmUniversalWithAbsMax< + ElementA_, + LayoutA_, + ElementB_, + LayoutB_, + ElementC_, + LayoutC_, + ElementAccumulator_, + arch::OpClassTensorOp, + arch::Sm89, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + AlignmentA, + AlignmentB, + Operator_, + TransformA, + TransformB +> : + public GemmUniversalBase< + typename kernel::DefaultGemmWithAbsMax< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + ElementC_, + LayoutC_, + ElementAccumulator_, + arch::OpClassTensorOp, + arch::Sm89, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_ + >::GemmKernel + > { + + public: + + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = arch::OpClassTensorOp; + using ArchTag = arch::Sm89; + 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::DefaultGemmWithAbsMax< + 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 SM89 column-major output exchanges problem size and operand. +template < + typename ElementA_, + typename LayoutA_, + typename ElementB_, + typename LayoutB_, + typename ElementC_, + typename ElementAccumulator_, + typename ThreadblockShape_, + typename WarpShape_, + typename InstructionShape_, + typename EpilogueOutputOp_, + typename ThreadblockSwizzle_, + int Stages, + int AlignmentA, + int AlignmentB, + typename Operator_, + ComplexTransform TransformA, + ComplexTransform TransformB> +class GemmUniversalWithAbsMax { + 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 = arch::OpClassTensorOp; + using ArchTag = arch::Sm89; + 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 GemmUniversalWithAbsMax< + 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. + GemmUniversalWithAbsMax() { } + + /// 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, cudaStream_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(cudaStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(cudaStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + cudaStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/dispatch_policy.hpp b/include/cutlass/gemm/dispatch_policy.hpp index c8ba9ba1..3694d0a8 100644 --- a/include/cutlass/gemm/dispatch_policy.hpp +++ b/include/cutlass/gemm/dispatch_policy.hpp @@ -37,11 +37,37 @@ #include "cute/numeric/integral_constant.hpp" ////////////////////////////////////////////////////////////////////////////// +namespace cutlass::detail { + +template class U> +struct is_kernel_tag_of : cute::false_type {}; + +template