From 808c25337a3ed4c97ac21895257b1addc72d6ca8 Mon Sep 17 00:00:00 2001 From: Manish Gupta Date: Fri, 19 Nov 2021 13:26:35 -0800 Subject: [PATCH] CUTLASS 2.8 (#363) CUTLASS 2.8 --- CHANGELOG.md | 24 + CUDA.cmake | 25 +- README.md | 102 +- ...6nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm75.h | 265 ++-- ...6nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.h | 153 -- ...hwx_s8cxrskx_s8ncxhwx_tensor_op_s32_sm75.h | 254 ++-- ...hwx_s8cxrskx_s8ncxhwx_tensor_op_s32_sm80.h | 156 -- .../13_two_tensor_op_fusion/b2b_conv2d_run.h | 96 +- .../b2b_interleaved_conv2d_run.h | 77 +- .../device/b2b_implicit_gemm_convolution.h | 24 + .../13_two_tensor_op_fusion/fused_conv2d.cu | 2 + .../kernel/b2b_implicit_gemm_convolution.h | 42 +- .../kernel/default_b2b_conv2d_fprop.h | 910 ++++++++++- .../b2b_implicit_gemm_multistage.h | 99 +- .../threadblock/b2b_implicit_gemm_pipelined.h | 84 +- ...implicit_gemm_pipelined_smem_accumulator.h | 532 +++++++ .../threadblock/b2b_mma_base.h | 12 +- .../b2b_mma_base_smem_accumulator.h | 174 +++ .../threadblock/b2b_mma_multistage.h | 6 +- .../threadblock/b2b_mma_pipelined.h | 6 +- .../threadblock/default_b2b_mma.h | 2 +- examples/24_gemm_grouped/CMakeLists.txt | 28 + examples/24_gemm_grouped/gemm_grouped.cu | 1326 +++++++++++++++++ .../CMakeLists.txt | 28 + .../ampere_fprop_mainloop_fusion.cu | 751 ++++++++++ .../CMakeLists.txt | 28 + .../ampere_wgrad_mainloop_fusion.cu | 749 ++++++++++ ...pere_3xtf32_fast_accurate_tensorop_gemm.cu | 744 +++++++++ .../CMakeLists.txt | 27 + .../CMakeLists.txt | 27 + ...ere_3xtf32_fast_accurate_tensorop_fprop.cu | 815 ++++++++++ ...f32_fast_accurate_tensorop_complex_gemm.cu | 686 +++++++++ .../CMakeLists.txt | 27 + examples/CMakeLists.txt | 6 + include/cutlass/arch/memory_sm75.h | 47 + include/cutlass/arch/memory_sm80.h | 54 + include/cutlass/arch/mma.h | 10 + include/cutlass/arch/mma_sm75.h | 1 - include/cutlass/arch/mma_sm80.h | 6 +- include/cutlass/array.h | 2 + include/cutlass/array_subbyte.h | 2 + .../device/implicit_gemm_convolution_fusion.h | 262 ++++ include/cutlass/conv/kernel/default_conv2d.h | 4 +- .../conv/kernel/default_conv2d_fprop_fusion.h | 351 +++++ .../conv/kernel/default_conv2d_wgrad_fusion.h | 319 ++++ .../conv/kernel/implicit_gemm_convolution.h | 1 - .../kernel/implicit_gemm_convolution_fusion.h | 455 ++++++ .../implicit_gemm_convolution_strided_dgrad.h | 1 - ...cit_gemm_convolution_with_fused_epilogue.h | 1 - .../implicit_gemm_fprop_fusion_multistage.h | 787 ++++++++++ .../threadblock/implicit_gemm_multistage.h | 50 +- .../implicit_gemm_wgrad_fusion_multistage.h | 718 +++++++++ ...icated_scale_bias_vector_access_iterator.h | 393 +++++ .../predicated_scale_bias_vector_iterator.h | 365 +++++ ...egular_scale_bias_vector_access_iterator.h | 247 +++ .../warp/conv2d_fprop_scale_bias_iterator.h | 568 +++++++ .../conv/warp/scale_bias_relu_transform.h | 217 +++ include/cutlass/core_io.h | 4 +- .../thread/linear_combination_bias_relu.h | 2 - .../epilogue/thread/linear_combination_relu.h | 142 +- include/cutlass/epilogue/thread/scale_type.h | 7 +- .../default_epilogue_direct_store.h | 68 + .../threadblock/direct_epilogue_tensor_op.h | 252 ---- .../direct_store_epilogue_iterator.h | 135 ++ .../threadblock/epilogue_direct_store.h | 342 +++++ .../threadblock/epilogue_smem_accumulator.h | 173 +++ .../warp/fragment_iterator_tensor_op.h | 4 + .../cutlass/epilogue/warp/tensor_op_policy.h | 11 + .../epilogue/warp/tile_iterator_tensor_op.h | 196 +++ include/cutlass/fast_math.h | 7 + include/cutlass/gemm/device/gemm_grouped.h | 313 ++++ .../gemm/device/gemm_universal_adapter.h | 67 +- .../gemm/kernel/default_gemm_complex.h | 1 - .../gemm/kernel/default_gemm_grouped.h | 356 +++++ include/cutlass/gemm/kernel/gemm.h | 2 - include/cutlass/gemm/kernel/gemm_grouped.h | 615 ++++++++ .../gemm/kernel/gemm_transpose_operands.h | 118 ++ include/cutlass/gemm/kernel/gemm_universal.h | 3 +- .../gemm/kernel/gemm_with_fused_epilogue.h | 9 +- .../gemm/kernel/gemm_with_k_reduction.h | 1 - .../cutlass/gemm/threadblock/default_mma.h | 12 +- .../gemm/threadblock/default_mma_core_sm80.h | 4 +- .../default_multistage_mma_complex.h | 15 +- ...default_multistage_mma_complex_core_sm80.h | 34 +- .../cutlass/gemm/threadblock/mma_multistage.h | 54 +- .../gemm/warp/default_mma_complex_tensor_op.h | 68 + .../gemm/warp/default_mma_tensor_op_sm80.h | 46 + .../cutlass/gemm/warp/mma_complex_tensor_op.h | 4 +- .../warp/mma_complex_tensor_op_fast_f32.h | 657 ++++++++ .../warp/mma_gaussian_complex_tensor_op.h | 2 +- .../gemm/warp/mma_simt_tile_iterator.h | 16 + include/cutlass/gemm/warp/mma_tensor_op.h | 3 + .../gemm/warp/mma_tensor_op_fast_f32.h | 465 ++++++ .../warp/mma_tensor_op_fragment_iterator.h | 102 +- .../gemm/warp/mma_with_reduction_tensor_op.h | 1 - include/cutlass/numeric_conversion.h | 43 + include/cutlass/tfloat32.h | 4 +- .../predicated_vector_access_iterator.h | 350 +++++ .../transform/threadblock/vector_iterator.h | 143 ++ .../transform/warp/vector_fragment_iterator.h | 277 ++++ media/docs/implicit_gemm_convolution.md | 8 +- test/unit/gemm/device/CMakeLists.txt | 10 + ...6n_f16n_direct_store_tensor_op_f32_sm80.cu | 108 ++ .../gemm_f16n_f16t_f16t_tensor_op_f32_sm80.cu | 2 - test/unit/gemm/device/gemm_grouped_sm80.cu | 598 ++++++++ test/unit/gemm/device/testbed_complex.h | 60 +- test/unit/gemm/device/testbed_grouped.h | 515 +++++++ test/unit/gemm/warp/gemm_complex_sm80.cu | 2 - test/unit/gemm/warp/gemm_sm80.cu | 4 +- .../library/include/cutlass/library/library.h | 2 + tools/library/scripts/conv2d_operation.py | 2 +- tools/library/scripts/gemm_operation.py | 2 +- tools/library/scripts/generator.py | 171 ++- tools/library/src/library_internal.h | 9 + tools/profiler/CMakeLists.txt | 2 +- .../{cublas_helpers.cpp => cublas_helpers.cu} | 8 +- tools/profiler/src/cublas_helpers.h | 2 +- tools/profiler/src/gemm_operation_profiler.cu | 4 + tools/profiler/src/operation_profiler.cu | 24 +- tools/profiler/src/options.cu | 7 +- tools/profiler/src/options.h | 3 + tools/profiler/src/performance_report.cpp | 27 +- tools/profiler/src/performance_report.h | 1 + .../util/include/cutlass/util/command_line.h | 2 +- .../util/include/cutlass/util/distribution.h | 2 +- .../util/reference/device/convolution.h | 41 +- .../util/reference/device/tensor_reduce.h | 1 - 127 files changed, 18555 insertions(+), 1338 deletions(-) create mode 100644 examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_pipelined_smem_accumulator.h create mode 100644 examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h create mode 100644 examples/24_gemm_grouped/CMakeLists.txt create mode 100644 examples/24_gemm_grouped/gemm_grouped.cu create mode 100644 examples/25_ampere_fprop_mainloop_fusion/CMakeLists.txt create mode 100644 examples/25_ampere_fprop_mainloop_fusion/ampere_fprop_mainloop_fusion.cu create mode 100644 examples/26_ampere_wgrad_mainloop_fusion/CMakeLists.txt create mode 100644 examples/26_ampere_wgrad_mainloop_fusion/ampere_wgrad_mainloop_fusion.cu create mode 100644 examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/27_ampere_3xtf32_fast_accurate_tensorop_gemm.cu create mode 100644 examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/CMakeLists.txt create mode 100644 examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/CMakeLists.txt create mode 100644 examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/ampere_3xtf32_fast_accurate_tensorop_fprop.cu create mode 100644 examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm.cu create mode 100644 examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/CMakeLists.txt create mode 100644 include/cutlass/conv/device/implicit_gemm_convolution_fusion.h create mode 100644 include/cutlass/conv/kernel/default_conv2d_fprop_fusion.h create mode 100644 include/cutlass/conv/kernel/default_conv2d_wgrad_fusion.h create mode 100644 include/cutlass/conv/kernel/implicit_gemm_convolution_fusion.h create mode 100644 include/cutlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h create mode 100644 include/cutlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h create mode 100644 include/cutlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h create mode 100644 include/cutlass/conv/threadblock/predicated_scale_bias_vector_iterator.h create mode 100644 include/cutlass/conv/threadblock/regular_scale_bias_vector_access_iterator.h create mode 100644 include/cutlass/conv/warp/conv2d_fprop_scale_bias_iterator.h create mode 100644 include/cutlass/conv/warp/scale_bias_relu_transform.h create mode 100644 include/cutlass/epilogue/threadblock/default_epilogue_direct_store.h delete mode 100644 include/cutlass/epilogue/threadblock/direct_epilogue_tensor_op.h create mode 100644 include/cutlass/epilogue/threadblock/direct_store_epilogue_iterator.h create mode 100644 include/cutlass/epilogue/threadblock/epilogue_direct_store.h create mode 100644 include/cutlass/epilogue/threadblock/epilogue_smem_accumulator.h create mode 100644 include/cutlass/gemm/device/gemm_grouped.h create mode 100644 include/cutlass/gemm/kernel/default_gemm_grouped.h create mode 100644 include/cutlass/gemm/kernel/gemm_grouped.h create mode 100644 include/cutlass/gemm/kernel/gemm_transpose_operands.h create mode 100644 include/cutlass/gemm/warp/mma_complex_tensor_op_fast_f32.h create mode 100644 include/cutlass/gemm/warp/mma_tensor_op_fast_f32.h create mode 100644 include/cutlass/transform/threadblock/predicated_vector_access_iterator.h create mode 100644 include/cutlass/transform/threadblock/vector_iterator.h create mode 100644 include/cutlass/transform/warp/vector_fragment_iterator.h create mode 100644 test/unit/gemm/device/gemm_f16n_f16n_f16n_direct_store_tensor_op_f32_sm80.cu create mode 100644 test/unit/gemm/device/gemm_grouped_sm80.cu create mode 100644 test/unit/gemm/device/testbed_grouped.h rename tools/profiler/src/{cublas_helpers.cpp => cublas_helpers.cu} (98%) diff --git a/CHANGELOG.md b/CHANGELOG.md index 90e6e6bd..d7a57e30 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,29 @@ # NVIDIA CUTLASS Changelog +## [2.8.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.8.0) (2021-11-19) + +* **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_ampere_3xtf32_fast_accurate_tensorop_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) +* **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. +* Optimal performance using [**CUDA 11.5**](https://developer.nvidia.com/cuda-downloads) +* Updates from the community (thanks!) + +* **Deprecation announcement:** CUTLASS plans to deprecate the following platforms in the future. Let us know if this affects your use case. + * Maxwell and Pascal GPU architectures + * Ubuntu 16.04 + * 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) diff --git a/CUDA.cmake b/CUDA.cmake index ca5b2a02..db1978a5 100644 --- a/CUDA.cmake +++ b/CUDA.cmake @@ -74,7 +74,7 @@ find_library( lib64 lib NO_DEFAULT_PATH - # We aren't going to search any system paths. We want to find the runtime + # We aren't going to search any system paths. We want to find the runtime # in the CUDA toolkit we're building against. ) @@ -89,10 +89,10 @@ if(NOT TARGET cudart AND CUDART_LIBRARY) # from the PATH search. else() add_library(cudart SHARED IMPORTED GLOBAL) - endif() + endif() add_library(nvidia::cudart ALIAS cudart) - + set_property( TARGET cudart PROPERTY IMPORTED_LOCATION @@ -120,7 +120,7 @@ find_library( lib64/stubs lib/stubs NO_DEFAULT_PATH - # We aren't going to search any system paths. We want to find the runtime + # We aren't going to search any system paths. We want to find the runtime # in the CUDA toolkit we're building against. ) @@ -135,10 +135,10 @@ if(NOT TARGET cuda_driver AND CUDA_DRIVER_LIBRARY) # from the PATH search. else() add_library(cuda_driver SHARED IMPORTED GLOBAL) - endif() + endif() add_library(nvidia::cuda_driver ALIAS cuda_driver) - + set_property( TARGET cuda_driver PROPERTY IMPORTED_LOCATION @@ -164,7 +164,7 @@ find_library( lib64 lib NO_DEFAULT_PATH - # We aren't going to search any system paths. We want to find the runtime + # We aren't going to search any system paths. We want to find the runtime # in the CUDA toolkit we're building against. ) @@ -179,10 +179,10 @@ if(NOT TARGET nvrtc AND NVRTC_LIBRARY) # from the PATH search. else() add_library(nvrtc SHARED IMPORTED GLOBAL) - endif() - + endif() + add_library(nvidia::nvrtc ALIAS nvrtc) - + set_property( TARGET nvrtc PROPERTY IMPORTED_LOCATION @@ -242,7 +242,7 @@ function(cutlass_unify_source_files TARGET_ARGS_VAR) set(CUDA_FILE_ARGS) set(TARGET_SOURCE_ARGS) - + foreach(ARG ${__UNPARSED_ARGUMENTS}) if(${ARG} MATCHES ".*\.cu$") list(APPEND CUDA_FILE_ARGS ${ARG}) @@ -250,7 +250,7 @@ function(cutlass_unify_source_files TARGET_ARGS_VAR) list(APPEND TARGET_SOURCE_ARGS ${ARG}) endif() endforeach() - + list(LENGTH CUDA_FILE_ARGS NUM_CUDA_FILE_ARGS) while(NUM_CUDA_FILE_ARGS GREATER 0) list(SUBLIST CUDA_FILE_ARGS 0 ${__BATCH_SIZE} CUDA_FILE_BATCH) @@ -280,7 +280,6 @@ function(cutlass_unify_source_files TARGET_ARGS_VAR) set(${TARGET_ARGS_VAR} ${TARGET_SOURCE_ARGS} PARENT_SCOPE) endfunction() - function(cutlass_add_library NAME) set(options) diff --git a/README.md b/README.md index a9fc680b..36eb9511 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") -# CUTLASS 2.7 +# CUTLASS 2.8 -_CUTLASS 2.7 - September 2021_ +_CUTLASS 2.8 - November 2021_ CUTLASS is a collection of CUDA C++ template abstractions for implementing high-performance matrix-multiplication (GEMM) and related computations at all levels @@ -34,77 +34,20 @@ 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 supported at each level of the execution model hierarchy. -See the [CHANGELOG](CHANGELOG.md) for descriptions of recent updates. +# What's New in CUTLASS 2.8 +CUTLASS 2.8 is an update to CUTLASS adding: +- [TF32x3:](/examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm) emulated single-precision using Tensor Cores; 45+ TFLOPs on NVIDIA A100 +- [Mainloop fusion for Convolution:](/examples/25_ampere_fprop_mainloop_fusion) convolution with fused per-channel bias-add +- [Grouped GEMM:](/examples/24_gemm_grouped) similar to batched GEMM with distinct problem size per group +- [Implicit GEMM Convolution fusion](/examples/13_two_tensor_op_fusion/) supports staging 1st convolution's output accumulator in the shared memory on Turing. +- Optimal performance using [CUDA 11.5](https://developer.nvidia.com/cuda-downloads) +- CUTLASS plans to **deprecate** the following platforms in the future. Let us know if this affects your use case. + - Maxwell and Pascal GPU architectures + - Ubuntu 16.04 + - CUDA 10.2 +- Updates and bugfixes from the community (thanks!) -# What's New in CUTLASS 2.7 -CUTLASS 2.7 is a minor update to CUTLASS adding: -- Mainloop fusion for GEMM: [summation over A or B](/examples/23_ampere_gemm_operand_reduction_fusion/ampere_gemm_operand_reduction_fusion.cu) -- [Optimizations for strided DGRAD](/include/cutlass/conv/kernel/default_conv2d_dgrad.h) -- [Half-precision GELU_taylor activation functions](/include/cutlass/epilogue/thread/activation.h#L196) -- 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) -- Numerous updates from the community (thanks!) - -# What's New in CUTLASS 2.6 -CUTLASS 2.6 is a minor update to CUTLASS adding: -- Fused [broadcast](test/unit/gemm/device/gemm_with_broadcast_f16n_f16n_f16n_tensorop_f32_sm75.cu) and [reductions](/test/unit/gemm/device/gemm_with_reduction_f16n_f16n_f16n_tensorop_f32_sm75.cu) in the epilogues of GEMM and Convolution -- [Quaternion-valued GEMM](/examples/21_quaternion_gemm/quaternion_gemm.cu) and [Convolution](/examples/22_quaternion_conv/quaternion_conv.cu) in single-precision -- [New strided Dgrad](test/unit/conv/device/conv2d_strided_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.cu) implementation offers up to 4x performance improvements over previous strided Dgrad -- 64-bit strides for large tensor allocations -- [General affine layouts](/examples/18_ampere_fp64_tensorop_affine2_gemm/ampere_fp64_tensorop_affine2_gemm.cu) fp64 tensor core and simt GEMM -- [Batched GEMV](/test/unit/gemm/device/gemv.cu) preview implementation -- Enhanced functionality, boosted performance, and bug fixes in the epilogue. -- Optimal performance when compiled with the [CUDA 11.4 Toolkit](https://developer.nvidia.com/cuda-toolkit) -- Adopt new L2 prefetch feature in [ptx instruction](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#ptx-isa-version-7-4). -- Enhanced Clang support and the combination of Clang 13 and CUDA 11.4 can build and run kernels from Pascal and Ampere. -- Numerous updates from the community (thanks!) - -# What's New in CUTLASS 2.5 -CUTLASS 2.5 is a minor update to CUTLASS adding: -- [Tensor reductions](/test/unit/reduction/device/tensor_reduce_contiguous.cu) -- [Optimizations for 3-D convolution](include/cutlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h) -- [Fused Convolution+Convolution example](/examples/13_two_tensor_op_fusion/README.md) - -# What's New in CUTLASS 2.4 -CUTLASS 2.4 is a significant update to CUTLASS adding: -- 1-D, 2-D, and 3-D convolution targeting Tensor and CUDA cores for NVIDIA Ampere, Turing, and Volta GPU architectures -- CUTLASS profiler support for convolution -- [Documentation](/media/docs/implicit_gemm_convolution.md) describing Implicit GEMM Convolution algorithm and implementation - -# What's New in CUTLASS 2.3 - -CUTLASS 2.3 is a minor update to CUTLASS adding: -- GEMMs targeting structured [Sparse Tensor Cores](test/unit/gemm/device/gemm_f16n_f16n_f32t_tensor_op_f32_sparse_sm80.cu) in NVIDIA Ampere Architecture GPUs -- Fast SGEMM kernels targeting GeForce RTX 30-series CUDA Cores -- Intended to be compiled with [CUDA 11.1 Toolkit](https://developer.nvidia.com/cuda-toolkit) or later - -# What's New in CUTLASS 2.2 - -CUTLASS 2.2 is a significant update to CUTLASS adding: - -- Coverage of [NVIDIA Ampere Architecture features](https://devblogs.nvidia.com/nvidia-ampere-architecture-in-depth/) -- Tensor Core-accelerated GEMMs targeting Tensor Float 32, BFloat16, and double-precision data types -- Deep software pipelines using asynchronous copy -- Described in [GTC 2020 Webinar (SR 21745)](https://developer.nvidia.com/gtc/2020/video/s21745) -- Intended to be compiled with [CUDA 11 Toolkit](https://developer.nvidia.com/cuda-toolkit) or later - -# What's New in CUTLASS 2.1 - -CUTLASS 2.1 is a minor update to CUTLASS adding: - -- [Planar complex GEMM kernels](/examples/10_planar_complex/planar_complex.cu) targeting Volta and Turing Tensor Cores -- BLAS-style API to launch kernels compiled into the [CUTLASS Library](/media/docs/quickstart.md#cutlass-library) - -# What's New in CUTLASS 2.0 - -CUTLASS 2.0 is a substantial refactoring from the previous version, intended to offer: - -- Better performance over 1.x, particularly for kernels targeting Turing Tensor Cores -- Robust and durable templates that reliably span the design space -- Encapsulated functionality that may be reusable in other contexts - -**See the [CHANGELOG](CHANGELOG.md) for more details.** +**See the [CHANGELOG](CHANGELOG.md) for a detailed listing of releases and updates.** # Performance @@ -120,8 +63,8 @@ using CUDA 11.0 Toolkit. Tensor Core operations are implemented using CUDA's # Compatibility CUTLASS requires a C++11 host compiler and -performs best when built with the [CUDA 11.4 Toolkit](https://developer.nvidia.com/cuda-toolkit). -It is also compatible with CUDA 10.2, CUDA 11.0, CUDA 11.1, CUDA 11.2, and CUDA 11.3. +performs best when built with the [CUDA 11.5 Toolkit](https://developer.nvidia.com/cuda-toolkit). +It is also compatible with CUDA 11.0, CUDA 11.1, CUDA 11.2, CUDA 11.3, and CUDA 11.4. We have tested the following environments. @@ -129,29 +72,26 @@ We have tested the following environments. |-----------------|----------| | Windows 10 | Microsoft Visual Studio 2015| | | Microsoft Visual Studio 2017| -| Ubuntu 16.04 | GCC 5.4.0 | | Ubuntu 18.04 | GCC 7.5.0 | -| Ubuntu 20.04 | GCC 10.2.0 | +| Ubuntu 20.04 | GCC 10.3.0 | Additionally, CUTLASS may be built with clang. See [these instructions](media/docs/quickstart.md#clang) for more details. CUTLASS runs successfully on the following NVIDIA GPUs, and it is expected to be efficient on -any Maxwell-, Pascal-, Volta-, Turing-, or NVIDIA Ampere- architecture NVIDIA GPU. +any Volta-, Turing-, or NVIDIA Ampere- architecture NVIDIA GPU. -For all GPUs, we recommend compiling with the [CUDA 11.4 Toolkit](https://developer.nvidia.com/cuda-toolkit) +For all GPUs, we recommend compiling with the [**CUDA 11.5 Toolkit**](https://developer.nvidia.com/cuda-toolkit) for best performance. |**GPU**|**CUDA Compute Capability**|**Minimum CUDA Toolkit**|**CUDA Toolkit Enabling Native Tensor Cores**| |---|---|---|---| -|NVIDIA Tesla P100|6.0|9.2| | -|NVIDIA GeForce 1080|6.1|9.2| | -|NVIDIA TitanXP|6.1|9.2| | |NVIDIA Tesla V100|7.0|9.2|10.1| |NVIDIA TitanV|7.0|9.2|10.1| |NVIDIA GeForce RTX 2080 TI, 2080, 2070|7.5|10.0|10.2| |NVIDIA Tesla T4|7.5|10.0|10.2| |NVIDIA A100|8.0|11.0|11.0| +|NVIDIA A10 |8.6|11.1|11.1| |NVIDIA GeForce 3090|8.6|11.1|11.1| # Documentation diff --git a/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm75.h b/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm75.h index 35ba83a6..42edd6aa 100644 --- a/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm75.h +++ b/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm75.h @@ -48,169 +48,13 @@ cutlass::conv::Conv2dProblemSize conv2d_f16_sm75_problem_size_0 ( ); cutlass::conv::Conv2dProblemSize conv2d_f16_sm75_problem_size_1 ( {128, 56, 56, 64}, // input size (NHWC) - {64, 1, 1, 64}, // filter size (KRSC) + {256, 1, 1, 64}, // filter size (KRSC) {0, 0, 0, 0}, // padding (pad_h, _, pad_w, _) {1, 1}, // stride (stride_h, stride_w) {1, 1}, // dilation (dilation_h, dilation_w) - {128, 56, 56, 64} // output size (NPQK) + {128, 56, 56, 256} // output size (NPQK) ); -bool run_nonfused_conv2d_fprop_f16_sm75() { - - using ElementA = cutlass::half_t; - using ElementB = cutlass::half_t; - using ElementC = cutlass::half_t; - using ElementAccumulator = cutlass::half_t; - using ElementCompute = cutlass::half_t; - - ElementCompute alpha0 = ElementCompute(1); - ElementCompute beta0 = ElementCompute(0); - ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); - - using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 32>; - using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 32>; - using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 64, 32>; - using WarpShape1 = cutlass::gemm::GemmShape<32, 64, 32>; - using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; - - using Conv2dFpropKernel0 = typename cutlass::conv::kernel::DefaultConv2dFprop< - ElementA, cutlass::layout::TensorNHWC, - ElementB, cutlass::layout::TensorNHWC, - ElementC, cutlass::layout::TensorNHWC, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm75, - ThreadblockShape0, - WarpShape0, - InstructionShape, - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 128 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute, - cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling - >, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 2, - cutlass::arch::OpMultiplyAdd, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using Conv2dFprop0 = cutlass::conv::device::ImplicitGemmConvolution; - - using Conv2dFpropKernel1 = typename cutlass::conv::kernel::DefaultConv2dFprop< - ElementA, cutlass::layout::TensorNHWC, - ElementB, cutlass::layout::TensorNHWC, - ElementC, cutlass::layout::TensorNHWC, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm75, - ThreadblockShape1, - WarpShape1, - InstructionShape, - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 128 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute - >, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 2, - cutlass::arch::OpMultiplyAdd, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using Conv2dFprop1 = cutlass::conv::device::ImplicitGemmConvolution; - - B2bNonFusedConv2dRun nonFusedConv2d; - - std::cout << "Running Non-fused back-to-back FP16 Analytic Convolution Fprops...\n"; - bool pass = nonFusedConv2d.run(conv2d_f16_sm75_problem_size_0, conv2d_f16_sm75_problem_size_1, cutlass::conv::SplitKMode::kSerial, - alpha0, beta0, alpha1, beta1); - - if(pass) - std::cout << "Pass\n"; - else - std::cout << "Fail\n"; - - return pass; -} - -bool run_fused_conv2d_fprop_f16_sm75() { - - using ElementA = cutlass::half_t; - using ElementB = cutlass::half_t; - using ElementC = cutlass::half_t; - using ElementAccumulator = cutlass::half_t; - using ElementCompute = cutlass::half_t; - - ElementCompute alpha0 = ElementCompute(1); - ElementCompute beta0 = ElementCompute(0); - ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); - - using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 32>; - using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 32>; - using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 64, 32>; - using WarpShape1 = cutlass::gemm::GemmShape<32, 64, 32>; - using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; - - using EpilogueOutputOp0 = - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - InstructionShape::kM * InstructionShape::kN / 32, - ElementAccumulator, - ElementCompute, - cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling - >; - - using EpilogueOutputOp1 = - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 128 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute - >; - - - - using B2bConv2dFpropKernel = typename cutlass::conv::kernel::DefaultB2bConv2dFprop< - ElementA, cutlass::layout::TensorNHWC, - ElementB, cutlass::layout::TensorNHWC, - ElementC, cutlass::layout::TensorNHWC, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm75, - ThreadblockShape0, - ThreadblockShape1, - WarpShape0, - WarpShape1, - InstructionShape, - EpilogueOutputOp0, - EpilogueOutputOp1, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 2, - cutlass::arch::OpMultiplyAdd, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using B2bConv2dFprop = cutlass::conv::device::B2bImplicitGemmConvolution; - - B2bFusedConv2dRun fusedConv2d; - - std::cout << "Running Fused back-to-back FP16 Analytic Convolution Fprops...\n"; - bool pass = fusedConv2d.run(conv2d_f16_sm75_problem_size_0, conv2d_f16_sm75_problem_size_1, cutlass::conv::SplitKMode::kSerial, - alpha0, beta0, alpha1, beta1); - - if(pass) - std::cout << "Pass\n"; - else - std::cout << "Fail\n"; - - return pass; -} - bool run_nonfused_conv2d_fprop_optimized_f16_sm75() { using ElementA = cutlass::half_t; @@ -220,9 +64,9 @@ bool run_nonfused_conv2d_fprop_optimized_f16_sm75() { using ElementCompute = cutlass::half_t; ElementCompute alpha0 = ElementCompute(1); - ElementCompute beta0 = ElementCompute(0); + ElementCompute beta0 = ElementCompute(1); //use beta for bias ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); + ElementCompute beta1 = ElementCompute(1); //use beta for bias using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 32>; using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 32>; @@ -245,7 +89,7 @@ bool run_nonfused_conv2d_fprop_optimized_f16_sm75() { 128 / cutlass::sizeof_bits::value, ElementAccumulator, ElementCompute, - cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling + cutlass::epilogue::thread::ScaleType::NoBetaScaling >, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, 2, @@ -269,7 +113,8 @@ bool run_nonfused_conv2d_fprop_optimized_f16_sm75() { ElementC, 128 / cutlass::sizeof_bits::value, ElementAccumulator, - ElementCompute + ElementCompute, + cutlass::epilogue::thread::ScaleType::NoBetaScaling >, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, 2, @@ -302,14 +147,14 @@ bool run_fused_conv2d_fprop_optimized_f16_sm75() { using ElementCompute = cutlass::half_t; ElementCompute alpha0 = ElementCompute(1); - ElementCompute beta0 = ElementCompute(0); + ElementCompute beta0 = ElementCompute(0); ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); + ElementCompute beta1 = ElementCompute(1); //use beta for bias using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 32>; - using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 32>; - using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 64, 32>; - using WarpShape1 = cutlass::gemm::GemmShape<32, 64, 32>; + using WarpShape0 = cutlass::gemm::GemmShape<32, 32, 32>; + using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 256, 32>; + using WarpShape1 = cutlass::gemm::GemmShape<64, 64, 32>; using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; using EpilogueOutputOp0 = @@ -326,10 +171,12 @@ bool run_fused_conv2d_fprop_optimized_f16_sm75() { ElementC, 128 / cutlass::sizeof_bits::value, ElementAccumulator, - ElementCompute + ElementCompute, + cutlass::epilogue::thread::ScaleType::NoBetaScaling >; + const bool SmemAccumulator = true; using B2bConv2dFpropKernel = typename cutlass::conv::kernel::DefaultB2bConv2dFprop< ElementA, cutlass::layout::TensorNHWC, @@ -348,14 +195,92 @@ bool run_fused_conv2d_fprop_optimized_f16_sm75() { cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, 2, cutlass::arch::OpMultiplyAdd, - cutlass::conv::IteratorAlgorithm::kOptimized + cutlass::conv::IteratorAlgorithm::kOptimized, + SmemAccumulator >::Kernel; using B2bConv2dFprop = cutlass::conv::device::B2bImplicitGemmConvolution; B2bFusedConv2dRun fusedConv2d; - std::cout << "Running Fused back-to-back FP16 Optimized Convolution Fprops...\n"; + std::cout << "Running Fused back-to-back FP16 Optimized Convolution Fprops with shared memory staging...\n"; + bool pass = fusedConv2d.run(conv2d_f16_sm75_problem_size_0, conv2d_f16_sm75_problem_size_1, cutlass::conv::SplitKMode::kSerial, + alpha0, beta0, alpha1, beta1); + + if(pass) + std::cout << "Pass\n"; + else + std::cout << "Fail\n"; + + return pass; +} + +bool run_fused_conv2d_fprop_optimized_f16_sm75_rf_res() { + + using ElementA = cutlass::half_t; + using ElementB = cutlass::half_t; + using ElementC = cutlass::half_t; + using ElementAccumulator = cutlass::half_t; + using ElementCompute = cutlass::half_t; + + ElementCompute alpha0 = ElementCompute(1); + ElementCompute beta0 = ElementCompute(0); + ElementCompute alpha1 = ElementCompute(1); + ElementCompute beta1 = ElementCompute(1); //use beta for bias + + using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 32>; + using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 32>; + using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 256, 32>; + using WarpShape1 = cutlass::gemm::GemmShape<32, 256, 32>; + using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; + + using EpilogueOutputOp0 = + cutlass::epilogue::thread::LinearCombinationRelu< + ElementC, + InstructionShape::kM * InstructionShape::kN / 32, + ElementAccumulator, + ElementCompute, + cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling + >; + + using EpilogueOutputOp1 = + cutlass::epilogue::thread::LinearCombinationRelu< + ElementC, + 128 / cutlass::sizeof_bits::value, + ElementAccumulator, + ElementCompute, + cutlass::epilogue::thread::ScaleType::NoBetaScaling + >; + + + const bool SmemAccumulator = false; + + using B2bConv2dFpropKernel = typename cutlass::conv::kernel::DefaultB2bConv2dFprop< + ElementA, cutlass::layout::TensorNHWC, + ElementB, cutlass::layout::TensorNHWC, + ElementC, cutlass::layout::TensorNHWC, + ElementAccumulator, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm75, + ThreadblockShape0, + ThreadblockShape1, + WarpShape0, + WarpShape1, + InstructionShape, + EpilogueOutputOp0, + EpilogueOutputOp1, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, + 2, + cutlass::arch::OpMultiplyAdd, + cutlass::conv::IteratorAlgorithm::kOptimized, + SmemAccumulator + >::Kernel; + + using B2bConv2dFprop = cutlass::conv::device::B2bImplicitGemmConvolution; + + B2bFusedConv2dRun fusedConv2d; + + std::cout << "Running Fused back-to-back FP16 Optimized Convolution Fprops with RF Residency...\n"; bool pass = fusedConv2d.run(conv2d_f16_sm75_problem_size_0, conv2d_f16_sm75_problem_size_1, cutlass::conv::SplitKMode::kSerial, alpha0, beta0, alpha1, beta1); diff --git a/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.h b/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.h index 3b2e720a..cdad8512 100644 --- a/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.h +++ b/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.h @@ -55,159 +55,6 @@ cutlass::conv::Conv2dProblemSize conv2d_f16_sm80_problem_size_1 ( {128, 56, 56, 64} // output size (NPQK) ); -bool run_nonfused_conv2d_fprop_f16_sm80() { - - using ElementA = cutlass::half_t; - using ElementB = cutlass::half_t; - using ElementC = cutlass::half_t; - using ElementAccumulator = cutlass::half_t; - using ElementCompute = cutlass::half_t; - - ElementCompute alpha0 = ElementCompute(1); - ElementCompute beta0 = ElementCompute(0); - ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); - - using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 64>; - using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape1 = cutlass::gemm::GemmShape<32, 64, 64>; - using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>; - - using Conv2dFpropKernel0 = typename cutlass::conv::kernel::DefaultConv2dFprop< - ElementA, cutlass::layout::TensorNHWC, - ElementB, cutlass::layout::TensorNHWC, - ElementC, cutlass::layout::TensorNHWC, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm80, - ThreadblockShape0, - WarpShape0, - InstructionShape, - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 128 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute, - cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling - >, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 3, - cutlass::arch::OpMultiplyAdd, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using Conv2dFprop0 = cutlass::conv::device::ImplicitGemmConvolution; - - using Conv2dFpropKernel1 = typename cutlass::conv::kernel::DefaultConv2dFprop< - ElementA, cutlass::layout::TensorNHWC, - ElementB, cutlass::layout::TensorNHWC, - ElementC, cutlass::layout::TensorNHWC, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm80, - ThreadblockShape1, - WarpShape1, - InstructionShape, - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 128 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute - >, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 3, - cutlass::arch::OpMultiplyAdd, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using Conv2dFprop1 = cutlass::conv::device::ImplicitGemmConvolution; - - B2bNonFusedConv2dRun nonFusedConv2d; - - std::cout << "Running Non-fused back-to-back FP16 Analytic Convolution Fprops...\n"; - bool pass = nonFusedConv2d.run(conv2d_f16_sm80_problem_size_0, conv2d_f16_sm80_problem_size_1, cutlass::conv::SplitKMode::kSerial, - alpha0, beta0, alpha1, beta1); - - if(pass) - std::cout << "Pass\n"; - else - std::cout << "Fail\n"; - - return pass; -} - -bool run_fused_conv2d_fprop_f16_sm80() { - - using ElementA = cutlass::half_t; - using ElementB = cutlass::half_t; - using ElementC = cutlass::half_t; - using ElementAccumulator = cutlass::half_t; - using ElementCompute = cutlass::half_t; - - ElementCompute alpha0 = ElementCompute(1); - ElementCompute beta0 = ElementCompute(0); - ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); - - using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 64>; - using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape1 = cutlass::gemm::GemmShape<32, 64, 64>; - using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>; - - using EpilogueOutputOp0 = - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - InstructionShape::kM * InstructionShape::kN / 32, - ElementAccumulator, - ElementCompute, - cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling - >; - - using EpilogueOutputOp1 = - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 128 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute - >; - - using B2bConv2dFpropKernel = typename cutlass::conv::kernel::DefaultB2bConv2dFprop< - ElementA, cutlass::layout::TensorNHWC, - ElementB, cutlass::layout::TensorNHWC, - ElementC, cutlass::layout::TensorNHWC, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm80, - ThreadblockShape0, - ThreadblockShape1, - WarpShape0, - WarpShape1, - InstructionShape, - EpilogueOutputOp0, - EpilogueOutputOp1, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 3, - cutlass::arch::OpMultiplyAdd, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using B2bConv2dFprop = cutlass::conv::device::B2bImplicitGemmConvolution; - - B2bFusedConv2dRun fusedConv2d; - - std::cout << "Running Fused back-to-back FP16 Analytic Convolution Fprops...\n"; - bool pass = fusedConv2d.run(conv2d_f16_sm80_problem_size_0, conv2d_f16_sm80_problem_size_1, cutlass::conv::SplitKMode::kSerial, - alpha0, beta0, alpha1, beta1); - - if(pass) - std::cout << "Pass\n"; - else - std::cout << "Fail\n"; - - return pass; -} bool run_nonfused_conv2d_fprop_optimized_f16_sm80() { diff --git a/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_s8ncxhwx_s8cxrskx_s8ncxhwx_tensor_op_s32_sm75.h b/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_s8ncxhwx_s8cxrskx_s8ncxhwx_tensor_op_s32_sm75.h index 11a68942..eb0ba622 100644 --- a/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_s8ncxhwx_s8cxrskx_s8ncxhwx_tensor_op_s32_sm75.h +++ b/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_s8ncxhwx_s8cxrskx_s8ncxhwx_tensor_op_s32_sm75.h @@ -48,169 +48,13 @@ cutlass::conv::Conv2dProblemSize conv2d_s8_sm75_problem_size_0 ( ); cutlass::conv::Conv2dProblemSize conv2d_s8_sm75_problem_size_1 ( {128, 56, 56, 64}, // input size (NHWC) - {64, 1, 1, 64}, // filter size (KRSC) + {256, 1, 1, 64}, // filter size (KRSC) {0, 0, 0, 0}, // padding (pad_h, _, pad_w, _) {1, 1}, // stride (stride_h, stride_w) {1, 1}, // dilation (dilation_h, dilation_w) - {128, 56, 56, 64} // output size (NPQK) + {128, 56, 56, 256} // output size (NPQK) ); -bool run_nonfused_conv2d_fprop_s8_sm75() { - - using ElementA = int8_t; - using ElementB = int8_t; - using ElementC = int8_t; - using ElementAccumulator = int32_t; - using ElementCompute = float; - - ElementCompute alpha0 = ElementCompute(1); - ElementCompute beta0 = ElementCompute(0); - ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); - - using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 64>; - using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape1 = cutlass::gemm::GemmShape<32, 64, 64>; - using InstructionShape = cutlass::gemm::GemmShape<8, 8, 16>; - - using Conv2dFpropKernel0 = typename cutlass::conv::kernel::DefaultConv2dFprop< - ElementA, cutlass::layout::TensorNCxHWx<32>, - ElementB, cutlass::layout::TensorCxRSKx<32>, - ElementC, cutlass::layout::TensorNCxHWx<32>, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm75, - ThreadblockShape0, - WarpShape0, - InstructionShape, - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 64 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute, - cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling - >, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 2, - cutlass::arch::OpMultiplyAddSaturate, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using Conv2dFprop0 = cutlass::conv::device::ImplicitGemmConvolution; - - using Conv2dFpropKernel1 = typename cutlass::conv::kernel::DefaultConv2dFprop< - ElementA, cutlass::layout::TensorNCxHWx<32>, - ElementB, cutlass::layout::TensorCxRSKx<32>, - ElementC, cutlass::layout::TensorNCxHWx<32>, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm75, - ThreadblockShape1, - WarpShape1, - InstructionShape, - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 64 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute - >, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 2, - cutlass::arch::OpMultiplyAddSaturate, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using Conv2dFprop1 = cutlass::conv::device::ImplicitGemmConvolution; - - B2bInterleavedNonFusedConv2dRun nonFusedConv2d; - - std::cout << "Running Non-fused back-to-back INT8 interleaved Analytic Convolution Fprops...\n"; - bool pass = nonFusedConv2d.run(conv2d_s8_sm75_problem_size_0, conv2d_s8_sm75_problem_size_1, cutlass::conv::SplitKMode::kSerial, - alpha0, beta0, alpha1, beta1); - - if(pass) - std::cout << "Pass\n"; - else - std::cout << "Fail\n"; - - return pass; -} - -bool run_fused_conv2d_fprop_s8_sm75() { - - using ElementA = int8_t; - using ElementB = int8_t; - using ElementC = int8_t; - using ElementAccumulator = int32_t; - using ElementCompute = float; - - ElementCompute alpha0 = ElementCompute(1); - ElementCompute beta0 = ElementCompute(0); - ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); - - using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 64>; - using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape1 = cutlass::gemm::GemmShape<32, 64, 64>; - using InstructionShape = cutlass::gemm::GemmShape<8, 8, 16>; - - using EpilogueOutputOp0 = - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - InstructionShape::kM * InstructionShape::kN / 32, - ElementAccumulator, - ElementCompute, - cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling - >; - - using EpilogueOutputOp1 = - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 64 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute - >; - - - - using B2bConv2dFpropKernel = typename cutlass::conv::kernel::DefaultB2bConv2dFprop< - ElementA, cutlass::layout::TensorNCxHWx<32>, - ElementB, cutlass::layout::TensorCxRSKx<32>, - ElementC, cutlass::layout::TensorNCxHWx<32>, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm75, - ThreadblockShape0, - ThreadblockShape1, - WarpShape0, - WarpShape1, - InstructionShape, - EpilogueOutputOp0, - EpilogueOutputOp1, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 2, - cutlass::arch::OpMultiplyAddSaturate, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using B2bConv2dFprop = cutlass::conv::device::B2bImplicitGemmConvolution; - - B2bInterleavedFusedConv2dRun fusedConv2d; - - std::cout << "Running Fused back-to-back INT8 interleaved Analytic Convolution Fprops...\n"; - bool pass = fusedConv2d.run(conv2d_s8_sm75_problem_size_0, conv2d_s8_sm75_problem_size_1, cutlass::conv::SplitKMode::kSerial, - alpha0, beta0, alpha1, beta1); - - if(pass) - std::cout << "Pass\n"; - else - std::cout << "Fail\n"; - - return pass; -} - bool run_nonfused_conv2d_fprop_optimized_s8_sm75() { using ElementA = int8_t; @@ -222,7 +66,7 @@ bool run_nonfused_conv2d_fprop_optimized_s8_sm75() { ElementCompute alpha0 = ElementCompute(1); ElementCompute beta0 = ElementCompute(0); ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); + ElementCompute beta1 = ElementCompute(1); using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 64>; using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 64>; @@ -304,12 +148,12 @@ bool run_fused_conv2d_fprop_optimized_s8_sm75() { ElementCompute alpha0 = ElementCompute(1); ElementCompute beta0 = ElementCompute(0); ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); + ElementCompute beta1 = ElementCompute(1); - using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 64>; - using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape1 = cutlass::gemm::GemmShape<32, 64, 64>; + using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 32>; + using WarpShape0 = cutlass::gemm::GemmShape<32, 32, 32>; + using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 256, 32>; + using WarpShape1 = cutlass::gemm::GemmShape<64, 64, 32>; using InstructionShape = cutlass::gemm::GemmShape<8, 8, 16>; using EpilogueOutputOp0 = @@ -330,6 +174,7 @@ bool run_fused_conv2d_fprop_optimized_s8_sm75() { >; + const bool SmemAccumulator = true; using B2bConv2dFpropKernel = typename cutlass::conv::kernel::DefaultB2bConv2dFprop< ElementA, cutlass::layout::TensorNCxHWx<32>, @@ -348,14 +193,91 @@ bool run_fused_conv2d_fprop_optimized_s8_sm75() { cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, 2, cutlass::arch::OpMultiplyAddSaturate, - cutlass::conv::IteratorAlgorithm::kOptimized + cutlass::conv::IteratorAlgorithm::kOptimized, + SmemAccumulator >::Kernel; using B2bConv2dFprop = cutlass::conv::device::B2bImplicitGemmConvolution; B2bInterleavedFusedConv2dRun fusedConv2d; - std::cout << "Running Fused back-to-back INT8 interleaved Optimized Convolution Fprops...\n"; + std::cout << "Running Fused back-to-back INT8 interleaved Optimized Convolution Fprops with shared memory staging...\n"; + bool pass = fusedConv2d.run(conv2d_s8_sm75_problem_size_0, conv2d_s8_sm75_problem_size_1, cutlass::conv::SplitKMode::kSerial, + alpha0, beta0, alpha1, beta1); + + if(pass) + std::cout << "Pass\n"; + else + std::cout << "Fail\n"; + + return pass; +} + +bool run_fused_conv2d_fprop_optimized_s8_sm75_rf_res() { + + using ElementA = int8_t; + using ElementB = int8_t; + using ElementC = int8_t; + using ElementAccumulator = int32_t; + using ElementCompute = float; + + ElementCompute alpha0 = ElementCompute(1); + ElementCompute beta0 = ElementCompute(0); + ElementCompute alpha1 = ElementCompute(1); + ElementCompute beta1 = ElementCompute(1); + + using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 32>; + using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 32>; + using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 256, 32>; + using WarpShape1 = cutlass::gemm::GemmShape<32, 256, 32>; + using InstructionShape = cutlass::gemm::GemmShape<8, 8, 16>; + + using EpilogueOutputOp0 = + cutlass::epilogue::thread::LinearCombinationRelu< + ElementC, + InstructionShape::kM * InstructionShape::kN / 32, + ElementAccumulator, + ElementCompute, + cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling + >; + + using EpilogueOutputOp1 = + cutlass::epilogue::thread::LinearCombinationRelu< + ElementC, + 64 / cutlass::sizeof_bits::value, + ElementAccumulator, + ElementCompute + >; + + + const bool SmemAccumulator = false; + + using B2bConv2dFpropKernel = typename cutlass::conv::kernel::DefaultB2bConv2dFprop< + ElementA, cutlass::layout::TensorNCxHWx<32>, + ElementB, cutlass::layout::TensorCxRSKx<32>, + ElementC, cutlass::layout::TensorNCxHWx<32>, + ElementAccumulator, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm75, + ThreadblockShape0, + ThreadblockShape1, + WarpShape0, + WarpShape1, + InstructionShape, + EpilogueOutputOp0, + EpilogueOutputOp1, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, + 2, + cutlass::arch::OpMultiplyAddSaturate, + cutlass::conv::IteratorAlgorithm::kOptimized, + SmemAccumulator + >::Kernel; + + using B2bConv2dFprop = cutlass::conv::device::B2bImplicitGemmConvolution; + + B2bInterleavedFusedConv2dRun fusedConv2d; + + std::cout << "Running Fused back-to-back INT8 interleaved Optimized Convolution Fprops with RF residency...\n"; bool pass = fusedConv2d.run(conv2d_s8_sm75_problem_size_0, conv2d_s8_sm75_problem_size_1, cutlass::conv::SplitKMode::kSerial, alpha0, beta0, alpha1, beta1); diff --git a/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_s8ncxhwx_s8cxrskx_s8ncxhwx_tensor_op_s32_sm80.h b/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_s8ncxhwx_s8cxrskx_s8ncxhwx_tensor_op_s32_sm80.h index 20cfe32b..740434f8 100644 --- a/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_s8ncxhwx_s8cxrskx_s8ncxhwx_tensor_op_s32_sm80.h +++ b/examples/13_two_tensor_op_fusion/b2b_conv2d_fprop_implicit_gemm_s8ncxhwx_s8cxrskx_s8ncxhwx_tensor_op_s32_sm80.h @@ -55,162 +55,6 @@ cutlass::conv::Conv2dProblemSize conv2d_s8_sm80_problem_size_1 ( {128, 56, 56, 64} // output size (NPQK) ); -bool run_nonfused_conv2d_fprop_s8_sm80() { - - using ElementA = int8_t; - using ElementB = int8_t; - using ElementC = int8_t; - using ElementAccumulator = int32_t; - using ElementCompute = float; - - ElementCompute alpha0 = ElementCompute(1); - ElementCompute beta0 = ElementCompute(0); - ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); - - using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 64>; - using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape1 = cutlass::gemm::GemmShape<32, 64, 64>; - using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32>; - - using Conv2dFpropKernel0 = typename cutlass::conv::kernel::DefaultConv2dFprop< - ElementA, cutlass::layout::TensorNCxHWx<32>, - ElementB, cutlass::layout::TensorCxRSKx<32>, - ElementC, cutlass::layout::TensorNCxHWx<32>, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm80, - ThreadblockShape0, - WarpShape0, - InstructionShape, - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 64 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute, - cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling - >, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 3, - cutlass::arch::OpMultiplyAddSaturate, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using Conv2dFprop0 = cutlass::conv::device::ImplicitGemmConvolution; - - using Conv2dFpropKernel1 = typename cutlass::conv::kernel::DefaultConv2dFprop< - ElementA, cutlass::layout::TensorNCxHWx<32>, - ElementB, cutlass::layout::TensorCxRSKx<32>, - ElementC, cutlass::layout::TensorNCxHWx<32>, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm80, - ThreadblockShape1, - WarpShape1, - InstructionShape, - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 64 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute - >, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 3, - cutlass::arch::OpMultiplyAddSaturate, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using Conv2dFprop1 = cutlass::conv::device::ImplicitGemmConvolution; - - B2bInterleavedNonFusedConv2dRun nonFusedConv2d; - - std::cout << "Running Non-fused back-to-back INT8 interleaved Analytic Convolution Fprops...\n"; - bool pass = nonFusedConv2d.run(conv2d_s8_sm80_problem_size_0, conv2d_s8_sm80_problem_size_1, cutlass::conv::SplitKMode::kSerial, - alpha0, beta0, alpha1, beta1); - - if(pass) - std::cout << "Pass\n"; - else - std::cout << "Fail\n"; - - return pass; -} - -bool run_fused_conv2d_fprop_s8_sm80() { - - using ElementA = int8_t; - using ElementB = int8_t; - using ElementC = int8_t; - using ElementAccumulator = int32_t; - using ElementCompute = float; - - ElementCompute alpha0 = ElementCompute(1); - ElementCompute beta0 = ElementCompute(0); - ElementCompute alpha1 = ElementCompute(1); - ElementCompute beta1 = ElementCompute(0); - - using ThreadblockShape0 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape0 = cutlass::gemm::GemmShape<32, 64, 64>; - using ThreadblockShape1 = cutlass::gemm::GemmShape<64, 64, 64>; - using WarpShape1 = cutlass::gemm::GemmShape<32, 64, 64>; - using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32>; - - using EpilogueOutputOp0 = - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 8 * InstructionShape::kN / 32, - ElementAccumulator, - ElementCompute, - cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling - >; - - using EpilogueOutputOp1 = - cutlass::epilogue::thread::LinearCombinationRelu< - ElementC, - 64 / cutlass::sizeof_bits::value, - ElementAccumulator, - ElementCompute - >; - - - - using B2bConv2dFpropKernel = typename cutlass::conv::kernel::DefaultB2bConv2dFprop< - ElementA, cutlass::layout::TensorNCxHWx<32>, - ElementB, cutlass::layout::TensorCxRSKx<32>, - ElementC, cutlass::layout::TensorNCxHWx<32>, - ElementAccumulator, - cutlass::arch::OpClassTensorOp, - cutlass::arch::Sm80, - ThreadblockShape0, - ThreadblockShape1, - WarpShape0, - WarpShape1, - InstructionShape, - EpilogueOutputOp0, - EpilogueOutputOp1, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>, - 3, - cutlass::arch::OpMultiplyAddSaturate, - cutlass::conv::IteratorAlgorithm::kAnalytic - >::Kernel; - - using B2bConv2dFprop = cutlass::conv::device::B2bImplicitGemmConvolution; - - B2bInterleavedFusedConv2dRun fusedConv2d; - - std::cout << "Running Fused back-to-back INT8 interleaved Analytic Convolution Fprops...\n"; - bool pass = fusedConv2d.run(conv2d_s8_sm80_problem_size_0, conv2d_s8_sm80_problem_size_1, cutlass::conv::SplitKMode::kSerial, - alpha0, beta0, alpha1, beta1); - - if(pass) - std::cout << "Pass\n"; - else - std::cout << "Fail\n"; - - return pass; -} - bool run_nonfused_conv2d_fprop_optimized_s8_sm80() { using ElementA = int8_t; diff --git a/examples/13_two_tensor_op_fusion/b2b_conv2d_run.h b/examples/13_two_tensor_op_fusion/b2b_conv2d_run.h index 3a8cdeb8..badcd25e 100644 --- a/examples/13_two_tensor_op_fusion/b2b_conv2d_run.h +++ b/examples/13_two_tensor_op_fusion/b2b_conv2d_run.h @@ -79,16 +79,19 @@ public: cutlass::Distribution::Kind init_A; cutlass::Distribution::Kind init_B; cutlass::Distribution::Kind init_C; + cutlass::Distribution::Kind init_Bias; uint64_t seed; cutlass::HostTensor tensor_A0; cutlass::HostTensor tensor_B0; cutlass::HostTensor tensor_C0; + cutlass::HostTensor tensor_Bias0; cutlass::HostTensor tensor_D0_computed; cutlass::HostTensor tensor_D0_reference; cutlass::HostTensor tensor_B1; cutlass::HostTensor tensor_C1; + cutlass::HostTensor tensor_Bias1; cutlass::HostTensor tensor_D1_computed; cutlass::HostTensor tensor_D1_reference; @@ -99,9 +102,10 @@ public: cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform, cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform, cutlass::Distribution::Kind init_C_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_Bias_ = cutlass::Distribution::Uniform, uint64_t seed_ = 2080 ): - init_A(init_A_), init_B(init_B_), init_C(init_C_), seed(seed_) { + init_A(init_A_), init_B(init_B_), init_C(init_C_), init_Bias(init_Bias_), seed(seed_) { } @@ -138,37 +142,50 @@ public: cutlass::reference::host::BlockFillSequential(view.data(), view.capacity()); } + else if (dist_kind == cutlass::Distribution::AllZeros) { + cutlass::reference::host::TensorFill(view, Element(0)); + } + else if (dist_kind == cutlass::Distribution::AllOnes) { + cutlass::reference::host::TensorFill(view, Element(1)); + } else { } } void initialize( cutlass::conv::Conv2dProblemSize const &problem_size_0, - cutlass::conv::Conv2dProblemSize const &problem_size_1, uint64_t seed = 2019) { + cutlass::conv::Conv2dProblemSize const &problem_size_1, + uint64_t seed = 2019) { tensor_A0.resize(implicit_gemm_tensor_a_extent(kConvolutionalOperator, problem_size_0)); tensor_B0.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_0)); tensor_C0.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_0)); + tensor_Bias0.resize({1, 1, 1, problem_size_0.K}); tensor_D0_computed.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_0)); tensor_D0_reference.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_0)); tensor_B1.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_1)); tensor_C1.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); + tensor_Bias1.resize({1, 1, 1, problem_size_1.K}); tensor_D1_computed.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); tensor_D1_reference.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); initialize_tensor(tensor_A0.host_view(), init_A, seed); initialize_tensor(tensor_B0.host_view(), init_B, seed * 17); initialize_tensor(tensor_C0.host_view(), init_C, seed * 39); + initialize_tensor(tensor_Bias0.host_view(), init_Bias, seed * 83); initialize_tensor(tensor_B1.host_view(), init_B, seed * 18); initialize_tensor(tensor_C1.host_view(), init_C, seed * 40); + initialize_tensor(tensor_Bias1.host_view(), init_Bias, seed * 84); tensor_A0.sync_device(); tensor_B0.sync_device(); tensor_C0.sync_device(); + tensor_Bias0.sync_device(); tensor_D0_computed.sync_device(); tensor_D0_reference.sync_device(); tensor_B1.sync_device(); tensor_C1.sync_device(); + tensor_Bias1.sync_device(); tensor_D1_computed.sync_device(); tensor_D1_reference.sync_device(); } @@ -196,7 +213,7 @@ public: problem_size_0, tensor_A0.device_ref(), tensor_B0.device_ref(), - tensor_C0.device_ref(), + {tensor_Bias0.device_data(), typename Conv2d0::LayoutC::Stride(0)}, tensor_D0_computed.device_ref(), {alpha0, beta0}, split_k_mode @@ -205,7 +222,7 @@ public: problem_size_1, tensor_D0_computed.device_ref(), tensor_B1.device_ref(), - tensor_C1.device_ref(), + {tensor_Bias1.device_data(), typename Conv2d1::LayoutC::Stride(0)}, tensor_D1_computed.device_ref(), {alpha1, beta1}, split_k_mode @@ -279,7 +296,7 @@ public: problem_size_0, tensor_A0.device_ref(), tensor_B0.device_ref(), - tensor_C0.device_ref(), + {tensor_Bias0.device_data(), typename Conv2d0::LayoutC::Stride(0)}, tensor_D0_reference.device_ref(), alpha0, beta0); @@ -302,7 +319,7 @@ public: problem_size_1, tensor_D0_reference.device_ref(), tensor_B1.device_ref(), - tensor_C1.device_ref(), + {tensor_Bias1.device_data(), typename Conv2d1::LayoutC::Stride(0)}, tensor_D1_reference.device_ref(), alpha1, beta1); @@ -344,10 +361,12 @@ public: << "\nA0:\n" << tensor_A0.host_view() << "\n" << "\nB0:\n" << tensor_B0.host_view() << "\n" << "\nC0:\n" << tensor_C0.host_view() << "\n" + << "\nBias0:\n" << tensor_Bias0.host_view() << "\n" << "\nD0 reference:\n" << tensor_D0_reference.host_view() << "\n" << "\nD0 computed:\n" << tensor_D0_computed.host_view() << "\n" << "\nB1:\n" << tensor_B1.host_view() << "\n" << "\nC1:\n" << tensor_C1.host_view() << "\n" + << "\nBias1:\n" << tensor_Bias1.host_view() << "\n" << "\nD1 reference:\n" << tensor_D1_reference.host_view() << "\n" << "\nD1 computed:\n" << tensor_D1_computed.host_view(); @@ -375,15 +394,20 @@ public: cutlass::Distribution::Kind init_A; cutlass::Distribution::Kind init_B; cutlass::Distribution::Kind init_C; + cutlass::Distribution::Kind init_Scale; + cutlass::Distribution::Kind init_Bias; uint64_t seed; cutlass::HostTensor tensor_A0; cutlass::HostTensor tensor_B0; cutlass::HostTensor tensor_C0; + cutlass::HostTensor tensor_Scale0; + cutlass::HostTensor tensor_Bias0; cutlass::HostTensor tensor_D0_reference; cutlass::HostTensor tensor_B1; cutlass::HostTensor tensor_C1; + cutlass::HostTensor tensor_Bias1; cutlass::HostTensor tensor_D1_computed; cutlass::HostTensor tensor_D1_reference; @@ -394,9 +418,12 @@ public: cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform, cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform, cutlass::Distribution::Kind init_C_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_Scale_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_Bias_ = cutlass::Distribution::Uniform, uint64_t seed_ = 2080 ): - init_A(init_A_), init_B(init_B_), init_C(init_C_), seed(seed_) { + init_A(init_A_), init_B(init_B_), init_C(init_C_), + init_Scale(init_Scale_), init_Bias(init_Bias_), seed(seed_) { } @@ -433,35 +460,56 @@ public: cutlass::reference::host::BlockFillSequential(view.data(), view.capacity()); } + else if (dist_kind == cutlass::Distribution::AllZeros) { + cutlass::reference::host::TensorFill(view, Element(0)); + } + else if (dist_kind == cutlass::Distribution::AllOnes) { + cutlass::reference::host::TensorFill(view, Element(1)); + } else { } } void initialize( cutlass::conv::Conv2dProblemSize const &problem_size_0, - cutlass::conv::Conv2dProblemSize const &problem_size_1, uint64_t seed = 2019) { + cutlass::conv::Conv2dProblemSize const &problem_size_1, + ElementCompute alpha0, + ElementCompute alpha1, + uint64_t seed = 2019) { tensor_A0.resize(implicit_gemm_tensor_a_extent(kConvolutionalOperator, problem_size_0)); tensor_B0.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_0)); tensor_C0.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_0)); + if(alpha0 == ElementCompute(0)) //per-channel scale + tensor_Scale0.resize({1, problem_size_0.K}); + tensor_Bias0.resize({1, problem_size_0.K}); tensor_D0_reference.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_0)); tensor_B1.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_1)); tensor_C1.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); + tensor_Bias1.resize({1, 1, 1, problem_size_1.K}); tensor_D1_computed.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); tensor_D1_reference.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); initialize_tensor(tensor_A0.host_view(), init_A, seed); initialize_tensor(tensor_B0.host_view(), init_B, seed * 17); initialize_tensor(tensor_C0.host_view(), init_C, seed * 39); + if(alpha0 == ElementCompute(0)) //per-channel scale + initialize_tensor(tensor_Scale0.host_view(), init_Scale, seed * 61); + initialize_tensor(tensor_Bias0.host_view(), init_Bias, seed * 83); initialize_tensor(tensor_B1.host_view(), init_B, seed * 18); initialize_tensor(tensor_C1.host_view(), init_C, seed * 40); + initialize_tensor(tensor_Bias1.host_view(), init_Bias, seed * 84); tensor_A0.sync_device(); tensor_B0.sync_device(); tensor_C0.sync_device(); + if(alpha0 == ElementCompute(0)) //per-channel scale + tensor_Scale0.sync_device(); + tensor_Bias0.sync_device(); tensor_D0_reference.sync_device(); tensor_B1.sync_device(); tensor_C1.sync_device(); + tensor_Bias1.sync_device(); tensor_D1_computed.sync_device(); tensor_D1_reference.sync_device(); } @@ -479,7 +527,7 @@ public: int warm_ups = 1, int runs = 100) { - initialize(problem_size_0, problem_size_1); + initialize(problem_size_0, problem_size_1, alpha0, alpha1); // configure the operator B2bConv2d b2b_conv2d_op; @@ -490,15 +538,31 @@ public: tensor_A0.device_ref(), tensor_B0.device_ref(), tensor_C0.device_ref(), + tensor_Scale0.device_ref(), + tensor_Bias0.device_ref(), tensor_B1.device_ref(), - tensor_C1.device_ref(), + {tensor_Bias1.device_data(), typename B2bConv2d::LayoutC::Stride(0)}, tensor_D1_computed.device_ref(), {alpha0, beta0}, {alpha1, beta1}, split_k_mode ); - cutlass::Status status = b2b_conv2d_op.initialize(b2b_conv2d_args); + cutlass::Status status = b2b_conv2d_op.can_implement(b2b_conv2d_args); + + if(status != cutlass::Status::kSuccess) { + std::cout << "Problem sizes not supported.\n" + << "Requirments:\n" + << " problem_size_0.N*P*Q = problem_size_1.N*P*Q\n" + << " problem_size_0.K = problem_size_1.C\n" + << " problem_size_1.R = problem_size_1.S = 1\n" + << " ThreadblockShape0::kN = problem_size_0.K\n" + << " ThreadblockShape1::kN = problem_size_1.K" << std::endl; + } + + CUTLASS_CHECK(status); + + status = b2b_conv2d_op.initialize(b2b_conv2d_args); CUTLASS_CHECK(status); @@ -551,7 +615,10 @@ public: tensor_C0.device_ref(), tensor_D0_reference.device_ref(), alpha0, - beta0); + beta0, + nullptr, // stream + tensor_Scale0.device_ref(), + tensor_Bias0.device_ref()); if(relu) { cutlass::reference::device::TensorReLu(tensor_D0_reference.device_view()); @@ -571,7 +638,7 @@ public: problem_size_1, tensor_D0_reference.device_ref(), tensor_B1.device_ref(), - tensor_C1.device_ref(), + {tensor_Bias1.device_data(), typename B2bConv2d::LayoutC::Stride(0)}, tensor_D1_reference.device_ref(), alpha1, beta1); @@ -612,8 +679,11 @@ public: << "\nA0:\n" << tensor_A0.host_view() << "\n" << "\nB0:\n" << tensor_B0.host_view() << "\n" << "\nC0:\n" << tensor_C0.host_view() << "\n" + << "\nScale0:\n" << tensor_Scale0.host_view() << "\n" + << "\nBias0:\n" << tensor_Bias0.host_view() << "\n" << "\nB1:\n" << tensor_B1.host_view() << "\n" << "\nC1:\n" << tensor_C1.host_view() << "\n" + << "\nBias1:\n" << tensor_Bias1.host_view() << "\n" << "\nD1 reference:\n" << tensor_D1_reference.host_view() << "\n" << "\nD1 computed:\n" << tensor_D1_computed.host_view(); diff --git a/examples/13_two_tensor_op_fusion/b2b_interleaved_conv2d_run.h b/examples/13_two_tensor_op_fusion/b2b_interleaved_conv2d_run.h index ff95f4e1..59521e49 100644 --- a/examples/13_two_tensor_op_fusion/b2b_interleaved_conv2d_run.h +++ b/examples/13_two_tensor_op_fusion/b2b_interleaved_conv2d_run.h @@ -80,18 +80,21 @@ public: cutlass::Distribution::Kind init_A; cutlass::Distribution::Kind init_B; cutlass::Distribution::Kind init_C; + cutlass::Distribution::Kind init_Bias; uint64_t seed; cutlass::HostTensor tensor_A0; cutlass::HostTensor tensor_B0; cutlass::HostTensor tensor_B0_reordered; cutlass::HostTensor tensor_C0; + cutlass::HostTensor tensor_Bias0; cutlass::HostTensor tensor_D0_computed; cutlass::HostTensor tensor_D0_reference; cutlass::HostTensor tensor_B1; cutlass::HostTensor tensor_B1_reordered; cutlass::HostTensor tensor_C1; + cutlass::HostTensor tensor_Bias1; cutlass::HostTensor tensor_D1_computed; cutlass::HostTensor tensor_D1_reference; @@ -102,9 +105,10 @@ public: cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform, cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform, cutlass::Distribution::Kind init_C_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_Bias_ = cutlass::Distribution::Uniform, uint64_t seed_ = 2080 ): - init_A(init_A_), init_B(init_B_), init_C(init_C_), seed(seed_) { + init_A(init_A_), init_B(init_B_), init_C(init_C_), init_Bias(init_Bias_), seed(seed_) { } @@ -141,6 +145,12 @@ public: cutlass::reference::host::BlockFillSequential(view.data(), view.capacity()); } + else if (dist_kind == cutlass::Distribution::AllZeros) { + cutlass::reference::host::TensorFill(view, Element(0)); + } + else if (dist_kind == cutlass::Distribution::AllOnes) { + cutlass::reference::host::TensorFill(view, Element(1)); + } else { } } @@ -153,17 +163,20 @@ public: tensor_B0.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_0)); tensor_B0_reordered.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_0)); tensor_C0.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_0)); + tensor_Bias0.resize({1, 1, 1, problem_size_0.K}); tensor_D0_computed.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_0)); tensor_D0_reference.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_0)); tensor_B1.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_1)); tensor_B1_reordered.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_1)); tensor_C1.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); + tensor_Bias1.resize({1, 1, 1, problem_size_1.K}); tensor_D1_computed.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); tensor_D1_reference.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); initialize_tensor(tensor_A0.host_view(), init_A, seed); initialize_tensor(tensor_B0.host_view(), init_B, seed * 17); initialize_tensor(tensor_C0.host_view(), init_C, seed * 39); + initialize_tensor(tensor_Bias0.host_view(), init_Bias, seed * 83); initialize_tensor(tensor_B1.host_view(), init_B, seed * 18); initialize_tensor(tensor_C1.host_view(), init_C, seed * 40); @@ -177,11 +190,13 @@ public: tensor_B0.sync_device(); tensor_B0_reordered.sync_device(); tensor_C0.sync_device(); + tensor_Bias0.sync_device(); tensor_D0_computed.sync_device(); tensor_D0_reference.sync_device(); tensor_B1.sync_device(); tensor_B1_reordered.sync_device(); tensor_C1.sync_device(); + tensor_Bias1.sync_device(); tensor_D1_computed.sync_device(); tensor_D1_reference.sync_device(); } @@ -392,17 +407,22 @@ public: cutlass::Distribution::Kind init_A; cutlass::Distribution::Kind init_B; cutlass::Distribution::Kind init_C; + cutlass::Distribution::Kind init_Scale; + cutlass::Distribution::Kind init_Bias; uint64_t seed; cutlass::HostTensor tensor_A0; cutlass::HostTensor tensor_B0; cutlass::HostTensor tensor_B0_reordered; cutlass::HostTensor tensor_C0; + cutlass::HostTensor tensor_Scale0; + cutlass::HostTensor tensor_Bias0; cutlass::HostTensor tensor_D0_reference; cutlass::HostTensor tensor_B1; cutlass::HostTensor tensor_B1_reordered; cutlass::HostTensor tensor_C1; + cutlass::HostTensor tensor_Bias1; cutlass::HostTensor tensor_D1_computed; cutlass::HostTensor tensor_D1_reference; @@ -413,9 +433,12 @@ public: cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform, cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform, cutlass::Distribution::Kind init_C_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_Scale_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_Bias_ = cutlass::Distribution::Uniform, uint64_t seed_ = 2080 ): - init_A(init_A_), init_B(init_B_), init_C(init_C_), seed(seed_) { + init_A(init_A_), init_B(init_B_), init_C(init_C_), + init_Scale(init_Scale_), init_Bias(init_Bias_), seed(seed_) { } @@ -452,30 +475,47 @@ public: cutlass::reference::host::BlockFillSequential(view.data(), view.capacity()); } + else if (dist_kind == cutlass::Distribution::AllZeros) { + cutlass::reference::host::TensorFill(view, Element(0)); + } + else if (dist_kind == cutlass::Distribution::AllOnes) { + cutlass::reference::host::TensorFill(view, Element(1)); + } else { } } void initialize( cutlass::conv::Conv2dProblemSize const &problem_size_0, - cutlass::conv::Conv2dProblemSize const &problem_size_1, uint64_t seed = 2019) { + cutlass::conv::Conv2dProblemSize const &problem_size_1, + ElementCompute alpha0, + ElementCompute alpha1, + uint64_t seed = 2019) { tensor_A0.resize(implicit_gemm_tensor_a_extent(kConvolutionalOperator, problem_size_0)); tensor_B0.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_0)); tensor_B0_reordered.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_0)); tensor_C0.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_0)); + if(alpha0 == ElementCompute(0)) //per-channel scale + tensor_Scale0.resize({1, problem_size_0.K}); + tensor_Bias0.resize({1, problem_size_0.K}); tensor_D0_reference.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_0)); tensor_B1.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_1)); tensor_B1_reordered.resize(implicit_gemm_tensor_b_extent(kConvolutionalOperator, problem_size_1)); tensor_C1.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); + tensor_Bias1.resize({1, 1, 1, problem_size_1.K}); tensor_D1_computed.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); tensor_D1_reference.resize(implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size_1)); initialize_tensor(tensor_A0.host_view(), init_A, seed); initialize_tensor(tensor_B0.host_view(), init_B, seed * 17); initialize_tensor(tensor_C0.host_view(), init_C, seed * 39); + if(alpha0 == ElementCompute(0)) //per-channel scale + initialize_tensor(tensor_Scale0.host_view(), init_Scale, seed * 61); + initialize_tensor(tensor_Bias0.host_view(), init_Bias, seed * 83); initialize_tensor(tensor_B1.host_view(), init_B, seed * 18); initialize_tensor(tensor_C1.host_view(), init_C, seed * 40); + initialize_tensor(tensor_Bias1.host_view(), init_Bias, seed * 84); //Reorder B0 and B1 cutlass::reorder_convK<16, InterleavedK>( @@ -487,10 +527,14 @@ public: tensor_B0.sync_device(); tensor_B0_reordered.sync_device(); tensor_C0.sync_device(); + if(alpha0 == ElementCompute(0)) //per-channel scale + tensor_Scale0.sync_device(); + tensor_Bias0.sync_device(); tensor_D0_reference.sync_device(); tensor_B1.sync_device(); tensor_B1_reordered.sync_device(); tensor_C1.sync_device(); + tensor_Bias1.sync_device(); tensor_D1_computed.sync_device(); tensor_D1_reference.sync_device(); } @@ -508,7 +552,7 @@ public: int warm_ups = 1, int runs = 100) { - initialize(problem_size_0, problem_size_1); + initialize(problem_size_0, problem_size_1, alpha0, alpha1); // configure the operator B2bConv2d b2b_conv2d_op; @@ -519,6 +563,8 @@ public: tensor_A0.device_ref(), tensor_B0_reordered.device_ref(), tensor_C0.device_ref(), + tensor_Scale0.device_ref(), + tensor_Bias0.device_ref(), tensor_B1_reordered.device_ref(), tensor_C1.device_ref(), tensor_D1_computed.device_ref(), @@ -527,7 +573,21 @@ public: split_k_mode ); - cutlass::Status status = b2b_conv2d_op.initialize(b2b_conv2d_args); + cutlass::Status status = b2b_conv2d_op.can_implement(b2b_conv2d_args); + + if(status != cutlass::Status::kSuccess) { + std::cout << "Problem sizes not supported.\n" + << "Requirments:\n" + << " problem_size_0.N*P*Q = problem_size_1.N*P*Q\n" + << " problem_size_0.K = problem_size_1.C\n" + << " problem_size_1.R = problem_size_1.S = 1\n" + << " ThreadblockShape0::kN = problem_size_0.K\n" + << " ThreadblockShape1::kN = problem_size_1.K" << std::endl; + } + + CUTLASS_CHECK(status); + + status = b2b_conv2d_op.initialize(b2b_conv2d_args); CUTLASS_CHECK(status); @@ -581,7 +641,10 @@ public: tensor_C0.device_ref(), tensor_D0_reference.device_ref(), alpha0, - beta0); + beta0, + nullptr, // stream + tensor_Scale0.device_ref(), + tensor_Bias0.device_ref()); if(relu) { cutlass::reference::device::TensorReLu(tensor_D0_reference.device_view()); @@ -644,6 +707,8 @@ public: << "\nB0:\n" << tensor_B0.host_view() << "\n" << "\nB0_reordered:\n" << tensor_B0_reordered.host_view() << "\n" << "\nC0:\n" << tensor_C0.host_view() << "\n" + << "\nScale0:\n" << tensor_Scale0.host_view() << "\n" + << "\nBias0:\n" << tensor_Bias0.host_view() << "\n" << "\nB1:\n" << tensor_B1.host_view() << "\n" << "\nB1_reordered:\n" << tensor_B1_reordered.host_view() << "\n" << "\nC1:\n" << tensor_C1.host_view() << "\n" diff --git a/examples/13_two_tensor_op_fusion/device/b2b_implicit_gemm_convolution.h b/examples/13_two_tensor_op_fusion/device/b2b_implicit_gemm_convolution.h index 1d539a0f..e078aab5 100644 --- a/examples/13_two_tensor_op_fusion/device/b2b_implicit_gemm_convolution.h +++ b/examples/13_two_tensor_op_fusion/device/b2b_implicit_gemm_convolution.h @@ -55,6 +55,8 @@ public: using LayoutC = typename B2bImplicitGemmKernel::LayoutC; using ElementAccumulator = typename B2bImplicitGemmKernel::ElementAccumulator; using ElementCompute = typename B2bImplicitGemmKernel::ElementCompute; + using ElementScaleBias = typename B2bImplicitGemmKernel::ElementScaleBias; + using LayoutScaleBias = typename B2bImplicitGemmKernel::LayoutScaleBias; using OperatorClass = typename B2bImplicitGemmKernel::OperatorClass; using ArchTag = typename B2bImplicitGemmKernel::ArchTag; using ThreadblockShape0 = typename B2bImplicitGemmKernel::ThreadblockShape0; @@ -126,6 +128,26 @@ public: return Status::kErrorInvalidProblem; } + // Determine if fusion sizes are valid + + cutlass::gemm::GemmCoord problem_size_0 = implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size_0); + cutlass::gemm::GemmCoord problem_size_1 = implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size_1); + + if(problem_size_0.m() != problem_size_1.m()) + return Status::kErrorInvalidProblem; + + if(problem_size_0.n() != problem_size_1.k()) + return Status::kErrorInvalidProblem; + + if(args.problem_size_1.R != 1 || args.problem_size_1.S != 1) + return Status::kErrorInvalidProblem; + + if(problem_size_0.n() > ThreadblockShape0::kN) + return Status::kErrorInvalidProblem; + + if(problem_size_1.n() > ThreadblockShape1::kN) + return Status::kErrorInvalidProblem; + return Status::kSuccess; } @@ -209,6 +231,8 @@ public: params_.ptr_A0 = args.ref_A0.data(); params_.ptr_B0 = args.ref_B0.data(); params_.ptr_C0 = args.ref_C0.data(); + params_.ptr_Scale0 = args.ref_Scale0.data(); + params_.ptr_Bias0 = args.ref_Bias0.data(); params_.ptr_B1 = args.ref_B1.data(); params_.ptr_C1 = args.ref_C1.data(); params_.ptr_D1 = args.ref_D1.data(); diff --git a/examples/13_two_tensor_op_fusion/fused_conv2d.cu b/examples/13_two_tensor_op_fusion/fused_conv2d.cu index 0f7f738f..05bee70c 100644 --- a/examples/13_two_tensor_op_fusion/fused_conv2d.cu +++ b/examples/13_two_tensor_op_fusion/fused_conv2d.cu @@ -60,8 +60,10 @@ int run_sm75() { std::cout << "Running on SM75" << std::endl; pass &= run_nonfused_conv2d_fprop_optimized_f16_sm75(); pass &= run_fused_conv2d_fprop_optimized_f16_sm75(); + pass &= run_fused_conv2d_fprop_optimized_f16_sm75_rf_res(); pass &= run_nonfused_conv2d_fprop_optimized_s8_sm75(); pass &= run_fused_conv2d_fprop_optimized_s8_sm75(); + pass &= run_fused_conv2d_fprop_optimized_s8_sm75_rf_res(); if(pass) return 1; diff --git a/examples/13_two_tensor_op_fusion/kernel/b2b_implicit_gemm_convolution.h b/examples/13_two_tensor_op_fusion/kernel/b2b_implicit_gemm_convolution.h index 926d2cee..b9246c20 100644 --- a/examples/13_two_tensor_op_fusion/kernel/b2b_implicit_gemm_convolution.h +++ b/examples/13_two_tensor_op_fusion/kernel/b2b_implicit_gemm_convolution.h @@ -79,6 +79,10 @@ struct B2bImplicitGemmConvolution { using ElementAccumulator = typename EpilogueOutputOp0::ElementAccumulator; using ElementCompute = typename EpilogueOutputOp0::ElementCompute; + /// Scale and Bias + using ElementScaleBias = typename B2bMma::IteratorAccumulatorScaleBias::Element; + using LayoutScaleBias = typename B2bMma::IteratorAccumulatorScaleBias::Layout; + using WarpMmaOperator0 = typename B2bMma::Policy0::Operator; using WarpMmaOperator1 = typename B2bMma::Policy1::Operator; @@ -103,13 +107,14 @@ struct B2bImplicitGemmConvolution { using TensorRefA0 = typename B2bMma::IteratorA0::TensorRef; using TensorRefB0 = typename B2bMma::IteratorB0::TensorRef; + using TensorRefScaleBias0 = typename B2bMma::IteratorAccumulatorScaleBias::TensorRef; using TensorRefB1 = typename B2bMma::IteratorB1::TensorRef; using TensorRefC = cutlass::TensorRef; /// Check iterator A and B convolution dimension are the same and // set device::B2bImplicitGemmConvolution::kConvDim static_assert(B2bMma::IteratorA0::kConvDim == B2bMma::IteratorB0::kConvDim, - "Convolution on different different dimensions is not supported"); + "Convolution on different dimensions is not supported"); static int const kConvDim = B2bMma::IteratorA0::kConvDim; /// Conv dimension and problem size structure (Conv2d or Conv3d) @@ -148,6 +153,8 @@ struct B2bImplicitGemmConvolution { TensorRefA0 ref_A0; TensorRefB0 ref_B0; TensorRefC ref_C0; + TensorRefScaleBias0 ref_Scale0; + TensorRefScaleBias0 ref_Bias0; TensorRefB1 ref_B1; TensorRefC ref_C1; TensorRefC ref_D1; @@ -178,6 +185,8 @@ struct B2bImplicitGemmConvolution { TensorRefA0 const & ref_A0, TensorRefB0 const & ref_B0, TensorRefC const & ref_C0, + TensorRefScaleBias0 const & ref_Scale0, + TensorRefScaleBias0 const & ref_Bias0, TensorRefB1 const & ref_B1, TensorRefC const & ref_C1, TensorRefC const & ref_D1, @@ -190,6 +199,8 @@ struct B2bImplicitGemmConvolution { ref_A0(ref_A0), ref_B0(ref_B0), ref_C0(ref_C0), + ref_Scale0(ref_Scale0), + ref_Bias0(ref_Bias0), ref_B1(ref_B1), ref_C1(ref_C1), ref_D1(ref_D1), @@ -218,6 +229,8 @@ struct B2bImplicitGemmConvolution { typename B2bMma::IteratorB0::Element const *ptr_B0; typename Epilogue::OutputTileIterator::Params iterator_C0; typename Epilogue::OutputTileIterator::Element *ptr_C0; + typename B2bMma::IteratorAccumulatorScaleBias::Element *ptr_Scale0; + typename B2bMma::IteratorAccumulatorScaleBias::Element *ptr_Bias0; typename B2bMma::IteratorB1::Params iterator_B1; typename B2bMma::IteratorB1::Element const *ptr_B1; typename Epilogue::OutputTileIterator::Params iterator_C1; @@ -252,6 +265,8 @@ struct B2bImplicitGemmConvolution { ptr_B0(args.ref_B0.data()), iterator_C0(ConvOutputIteratorParameter::layout(args.ref_C0)), ptr_C0(args.ref_C0.data()), + ptr_Scale0(args.ref_Scale0.data()), + ptr_Bias0(args.ref_Bias0.data()), iterator_B1(args.problem_size_1, args.ref_B1.layout()), ptr_B1(args.ref_B1.data()), iterator_C1(ConvOutputIteratorParameter::layout(args.ref_C1)), @@ -350,6 +365,28 @@ struct B2bImplicitGemmConvolution { int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); int lane_idx = threadIdx.x % 32; + // Construct iterators to accumulator scale/bias vector + typename B2bMma::IteratorAccumulatorScaleBias iterator_Scale0( + params.ptr_Scale0, + {1, params.problem_size_0.K}, + thread_idx, + warp_idx, + MatrixCoord( + 0, threadblock_tile_idx.n() * B2bMma::Shape0::kN + ) + ); + + typename B2bMma::IteratorAccumulatorScaleBias iterator_Bias0( + params.ptr_Bias0, + {1, params.problem_size_0.K}, + thread_idx, + warp_idx, + MatrixCoord( + 0, threadblock_tile_idx.n() * B2bMma::Shape0::kN + ) + ); + + // // Main loop // @@ -366,7 +403,8 @@ struct B2bImplicitGemmConvolution { accumulators.clear(); // Compute threadblock-scoped matrix multiply-add - b2bMma(params.gemm_k_iterations_0, accumulators, iterator_A0, iterator_B0, iterator_B1, src_accum, output_op_0); + b2bMma(params.gemm_k_iterations_0, accumulators, iterator_A0, iterator_B0, + iterator_Scale0, iterator_Bias0, iterator_B1, src_accum, output_op_0); // // Epilogue diff --git a/examples/13_two_tensor_op_fusion/kernel/default_b2b_conv2d_fprop.h b/examples/13_two_tensor_op_fusion/kernel/default_b2b_conv2d_fprop.h index b53898d1..1ae49ae1 100644 --- a/examples/13_two_tensor_op_fusion/kernel/default_b2b_conv2d_fprop.h +++ b/examples/13_two_tensor_op_fusion/kernel/default_b2b_conv2d_fprop.h @@ -38,11 +38,17 @@ #include "cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h" #include "cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h" #include "cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h" + +#include "cutlass/transform/threadblock/predicated_vector_access_iterator.h" +#include "cutlass/transform/threadblock/vector_iterator.h" +#include "cutlass/transform/warp/vector_fragment_iterator.h" + #include "cutlass/gemm/warp/mma_tensor_op_fragment_iterator.h" #include "kernel/b2b_implicit_gemm_convolution.h" #include "threadblock/b2b_implicit_gemm_pipelined.h" #include "threadblock/b2b_implicit_gemm_multistage.h" +#include "threadblock/b2b_implicit_gemm_pipelined_smem_accumulator.h" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -73,7 +79,7 @@ template < int Stages, typename MathOperatorTag, conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kAnalytic, - conv::StrideSupport StrideSupport = StrideSupport::kStrided + bool SmemAccumulator = false > struct DefaultB2bConv2dFprop; ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -166,6 +172,23 @@ struct DefaultB2bConv2dFprop < MmaCore1::Shape::kK, //kBlocksColumn ElementAccumulator, ElementA, AccumulatorLayout, InstructionShape, EpilogueOutputOp0>; + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 2; + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Warp-level iterators to load scale and bias vectors + using FragmentIteratorA1ScaleBias = cutlass::transform::warp::VectorFragmentIterator< + MatrixShape<1, IteratorAccumulatorScaleBias::Fragment::kElements>, ElementScaleBias, + LayoutScaleBias, InstructionShape, kElementsPerAccess>; + // Define iterators over tiles from the B operand using ThreadMapB1 = typename MmaCore1::IteratorThreadMapB; using IteratorB1 = @@ -193,6 +216,8 @@ struct DefaultB2bConv2dFprop < arch::CacheOperation::Global, ThreadblockShape1, FragmentIteratorA1, + IteratorAccumulatorScaleBias, + FragmentIteratorA1ScaleBias, IteratorB1, SmemIteratorB1, arch::CacheOperation::Global, @@ -319,10 +344,27 @@ struct DefaultB2bConv2dFprop < MmaCore1::Shape::kK, //kBlocksColumn ElementAccumulator, ElementA, AccumulatorLayout, InstructionShape, EpilogueOutputOp0>; + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 4; + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Warp-level iterators to load scale and bias vectors + using FragmentIteratorA1ScaleBias = cutlass::transform::warp::VectorFragmentIterator< + MatrixShape<1, IteratorAccumulatorScaleBias::Fragment::kElements>, ElementScaleBias, + LayoutScaleBias, InstructionShape, kElementsPerAccess>; + using ThreadMapB1 = typename MmaCore1::SmemThreadMapB; using IteratorB1 = cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< - cutlass::MatrixShape, + cutlass::MatrixShape, ElementB, layout::TensorCxRSKx, ThreadMapB1 >; @@ -346,6 +388,8 @@ struct DefaultB2bConv2dFprop < arch::CacheOperation::Global, ThreadblockShape1, FragmentIteratorA1, + IteratorAccumulatorScaleBias, + FragmentIteratorA1ScaleBias, IteratorB1, SmemIteratorB1, arch::CacheOperation::Global, @@ -465,6 +509,23 @@ struct DefaultB2bConv2dFprop < MmaCore1::Shape::kK, //kBlocksColumn ElementAccumulator, ElementA, AccumulatorLayout, InstructionShape, EpilogueOutputOp0>; + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 2; + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Warp-level iterators to load scale and bias vectors + using FragmentIteratorA1ScaleBias = cutlass::transform::warp::VectorFragmentIterator< + MatrixShape<1, IteratorAccumulatorScaleBias::Fragment::kElements>, ElementScaleBias, + LayoutScaleBias, InstructionShape, kElementsPerAccess>; + // Define iterators over tiles from the B operand using ThreadMapB1 = typename MmaCore1::IteratorThreadMapB; using IteratorB1 = @@ -492,6 +553,8 @@ struct DefaultB2bConv2dFprop < SmemIteratorB0, ThreadblockShape1, FragmentIteratorA1, + IteratorAccumulatorScaleBias, + FragmentIteratorA1ScaleBias, IteratorB1, SmemIteratorB1, ElementC, @@ -519,6 +582,187 @@ struct DefaultB2bConv2dFprop < >; }; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialzation for Analytic IteratorAlgorithm +/// and 2 stage pipeline. +/// Accumulator will be staged in shared memory. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape0, + typename ThreadblockShape1, + typename WarpShape0, + typename WarpShape1, + typename InstructionShape, + typename EpilogueOutputOp0, + typename EpilogueOutputOp1, + typename ThreadblockSwizzle, + typename MathOperatorTag +> +struct DefaultB2bConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape0, + ThreadblockShape1, + WarpShape0, + WarpShape1, + InstructionShape, + EpilogueOutputOp0, + EpilogueOutputOp1, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + true +> { + + // Define the core components from GEMM + using MmaCore0 = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape0, WarpShape0, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + using MmaCore1 = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape1, WarpShape1, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA0 = typename MmaCore0::IteratorThreadMapA; + using IteratorA0 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA0 + > + >; + + using SmemIteratorA0 = typename MmaCore0::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB0 = typename MmaCore0::IteratorThreadMapB; + using IteratorB0 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB0 + > + >; + + using SmemIteratorB0 = typename MmaCore0::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 2; + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Define iterators over tiles from the B operand + using ThreadMapB1 = typename MmaCore1::IteratorThreadMapB; + using IteratorB1 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB1 + > + >; + + using SmemIteratorB1 = typename MmaCore1::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp0 = typename MmaCore0::MmaTensorOp; + using WarpMmaTensorOp1 = typename MmaCore1::MmaTensorOp; + using MmaPolicy0 = typename MmaCore0::MmaPolicy; + using MmaPolicy1 = typename MmaCore1::MmaPolicy; + + // Use fragment iterator for the accumulator + using SmemAccumulatorLayout = cutlass::layout::RowMajor; + using FragmentIteratorAccumulator = cutlass::epilogue::warp::FragmentIteratorTensorOp< + WarpShape0, InstructionShape, + ElementAccumulator, + typename WarpMmaTensorOp0::Policy::Operator::FragmentC, + SmemAccumulatorLayout + >; + + // Store Accumulator tiles to Shared Memory + using SmemIteratorD0 = + cutlass::epilogue::warp::TileIteratorTensorOp< + WarpShape0, + InstructionShape, + ElementC, + SmemAccumulatorLayout + >; + + static int const kThreadCount = 32; + // load warp tile from Shared Memory accumulator + using WarpIteratorA1 = cutlass::gemm::warp::MmaTensorOpMultiplicandTileIterator< + MatrixShape, cutlass::gemm::Operand::kA, + ElementA, SmemAccumulatorLayout, + MatrixShape, + WarpMmaTensorOp1::Policy::OpDelta::kRow, kThreadCount>; + + // Define the Mma + using B2bMma = threadblock::B2bImplicitGemmPipelinedSmemAccumulator< + ThreadblockShape0, + IteratorA0, + SmemIteratorA0, + IteratorB0, + SmemIteratorB0, + IteratorAccumulatorScaleBias, + FragmentIteratorAccumulator, + SmemIteratorD0, + ThreadblockShape1, + WarpIteratorA1, + IteratorB1, + SmemIteratorB1, + ElementC, + LayoutC, + EpilogueOutputOp0, + MmaPolicy0, + MmaPolicy1 + >; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape1, + WarpMmaTensorOp1, + 1, + EpilogueOutputOp1 + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::B2bImplicitGemmConvolution< + B2bMma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + + ///////////////////////////////////////////////////////////////////////////////////////////////// /// Defines a kernel for Conv2dFprop specialzation for Analytic IteratorAlgorithm and 2 stage @@ -561,7 +805,8 @@ struct DefaultB2bConv2dFprop < ThreadblockSwizzle, 2, MathOperatorTag, - IteratorAlgorithm::kAnalytic + IteratorAlgorithm::kAnalytic, + false > { // Define the core components from GEMM @@ -621,6 +866,23 @@ struct DefaultB2bConv2dFprop < MmaCore1::Shape::kK, //kBlocksColumn ElementAccumulator, ElementA, AccumulatorLayout, InstructionShape, EpilogueOutputOp0>; + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 4; + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Warp-level iterators to load scale and bias vectors + using FragmentIteratorA1ScaleBias = cutlass::transform::warp::VectorFragmentIterator< + MatrixShape<1, IteratorAccumulatorScaleBias::Fragment::kElements>, ElementScaleBias, + LayoutScaleBias, InstructionShape, kElementsPerAccess>; + // Define iterators over tiles from the B operand using ThreadMapB1 = typename MmaCore1::SmemThreadMapB; using IteratorB1 = @@ -648,6 +910,199 @@ struct DefaultB2bConv2dFprop < SmemIteratorB0, ThreadblockShape1, FragmentIteratorA1, + IteratorAccumulatorScaleBias, + FragmentIteratorA1ScaleBias, + IteratorB1, + SmemIteratorB1, + ElementC, + LayoutC, + EpilogueOutputOp0, + MmaPolicy0, + MmaPolicy1 + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultInterleavedConvEpilogue< + ThreadblockShape1, + WarpMmaTensorOp1, + 1, + EpilogueOutputOp1, + EpilogueOutputOp1::kCount, + InterleavedK + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::B2bImplicitGemmConvolution< + B2bMma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +/// Defines a kernel for Conv2dFprop specialzation for Analytic IteratorAlgorithm and 2 stage +/// pipeline with interleaved layout. +/// Accumulator will be staged in shared memory. +template < + typename ElementA, + typename ElementB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape0, + typename ThreadblockShape1, + typename WarpShape0, + typename WarpShape1, + typename InstructionShape, + typename EpilogueOutputOp0, + typename EpilogueOutputOp1, + typename ThreadblockSwizzle, + typename MathOperatorTag, + int InterleavedK +> +struct DefaultB2bConv2dFprop < + ElementA, + layout::TensorNCxHWx, + ElementB, + layout::TensorCxRSKx, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape0, + ThreadblockShape1, + WarpShape0, + WarpShape1, + InstructionShape, + EpilogueOutputOp0, + EpilogueOutputOp1, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + true +> { + + // Define the core components from GEMM + using MmaCore0 = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape0, WarpShape0, InstructionShape, ElementA, layout::ColumnMajorInterleaved, + ElementB, layout::RowMajorInterleaved, + ElementAccumulator, LayoutC, arch::OpClassTensorOp, + 2, MathOperatorTag, true>; + using MmaCore1 = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape1, WarpShape1, InstructionShape, ElementA, layout::ColumnMajorInterleaved, + ElementB, layout::RowMajorInterleaved, + ElementAccumulator, LayoutC, arch::OpClassTensorOp, + 2, MathOperatorTag, true>; + + // Define iterators over tiles from the A operand + // Note GEMM shared memory threadmap is used here because conv global memory + // layout needs to be mapped to fprop which is similar to the crosswise + // layout which is used by the interleaved GEMM shared memory threadmap. + // The Interleaved GEMM global memory layout is similar to the congruous + // layout. + using ThreadMapA0 = typename MmaCore0::SmemThreadMapA; + using IteratorA0 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementA, layout::TensorNCxHWx, + ThreadMapA0 + > + >; + + using SmemIteratorA0 = typename MmaCore0::SmemIteratorA; + + // Define iterators over tiles from the B operand + // Note GEMM shared memory threadmap is used here because conv global memory + // layout needs to be mapped to fprop which is similar to the crosswise + // layout which is used by the interleaved GEMM shared memory threadmap. + // The Interleaved GEMM global memory layout is similar to the congruous + // layout. + using ThreadMapB0 = typename MmaCore0::SmemThreadMapB; + using IteratorB0 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementB, layout::TensorCxRSKx, + ThreadMapB0 + > + >; + + using SmemIteratorB0 = typename MmaCore0::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 4; //For interleaved layout + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Define iterators over tiles from the B operand + using ThreadMapB1 = typename MmaCore1::SmemThreadMapB; + using IteratorB1 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementB, layout::TensorCxRSKx, + ThreadMapB1 + > + >; + + using SmemIteratorB1 = typename MmaCore1::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp0 = typename MmaCore0::MmaTensorOp; + using WarpMmaTensorOp1 = typename MmaCore1::MmaTensorOp; + using MmaPolicy0 = typename MmaCore0::MmaPolicy; + using MmaPolicy1 = typename MmaCore1::MmaPolicy; + + // Use fragment iterator for the accumulator + using SmemAccumulatorLayout = cutlass::layout::ColumnMajorInterleaved<16>; + using FragmentIteratorAccumulator = cutlass::epilogue::warp::FragmentIteratorTensorOp< + WarpShape0, InstructionShape, + ElementAccumulator, + typename WarpMmaTensorOp0::Policy::Operator::FragmentC, + SmemAccumulatorLayout + >; + + + // Store Accumulator tiles to Shared Memory + using SmemIteratorD0 = + cutlass::epilogue::warp::TileIteratorTensorOp< + WarpShape0, + InstructionShape, + ElementC, + SmemAccumulatorLayout + >; + + static int const kThreadCount = 32; + // load warp tile from Shared Memory accumulator + using WarpIteratorA1 = cutlass::gemm::warp::MmaTensorOpMultiplicandTileIteratorCanonical< + MatrixShape, cutlass::gemm::Operand::kA, + ElementA, SmemAccumulatorLayout, + MatrixShape, + WarpMmaTensorOp1::Policy::OpDelta::kRow, kThreadCount>; + + // Define the Mma + using B2bMma = threadblock::B2bImplicitGemmPipelinedSmemAccumulator< + ThreadblockShape0, + IteratorA0, + SmemIteratorA0, + IteratorB0, + SmemIteratorB0, + IteratorAccumulatorScaleBias, + FragmentIteratorAccumulator, + SmemIteratorD0, + ThreadblockShape1, + WarpIteratorA1, IteratorB1, SmemIteratorB1, ElementC, @@ -764,6 +1219,23 @@ struct DefaultB2bConv2dFprop < MmaCore1::Shape::kK, //kBlocksColumn ElementAccumulator, ElementA, AccumulatorLayout, InstructionShape, EpilogueOutputOp0>; + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 2; + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Warp-level iterators to load scale and bias vectors + using FragmentIteratorA1ScaleBias = cutlass::transform::warp::VectorFragmentIterator< + MatrixShape<1, IteratorAccumulatorScaleBias::Fragment::kElements>, ElementScaleBias, + LayoutScaleBias, InstructionShape, kElementsPerAccess>; + // Define iterators over tiles from the B operand using ThreadMapB1 = typename MmaCore1::IteratorThreadMapB; using IteratorB1 = @@ -791,6 +1263,8 @@ struct DefaultB2bConv2dFprop < arch::CacheOperation::Global, ThreadblockShape1, FragmentIteratorA1, + IteratorAccumulatorScaleBias, + FragmentIteratorA1ScaleBias, IteratorB1, SmemIteratorB1, arch::CacheOperation::Global, @@ -917,10 +1391,27 @@ struct DefaultB2bConv2dFprop < MmaCore1::Shape::kK, //kBlocksColumn ElementAccumulator, ElementA, AccumulatorLayout, InstructionShape, EpilogueOutputOp0>; + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 4; + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Warp-level iterators to load scale and bias vectors + using FragmentIteratorA1ScaleBias = cutlass::transform::warp::VectorFragmentIterator< + MatrixShape<1, IteratorAccumulatorScaleBias::Fragment::kElements>, ElementScaleBias, + LayoutScaleBias, InstructionShape, kElementsPerAccess>; + using ThreadMapB1 = typename MmaCore1::SmemThreadMapB; using IteratorB1 = cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< - cutlass::MatrixShape, + cutlass::MatrixShape, ElementB, layout::TensorCxRSKx, ThreadMapB1 >; @@ -944,6 +1435,8 @@ struct DefaultB2bConv2dFprop < arch::CacheOperation::Global, ThreadblockShape1, FragmentIteratorA1, + IteratorAccumulatorScaleBias, + FragmentIteratorA1ScaleBias, IteratorB1, SmemIteratorB1, arch::CacheOperation::Global, @@ -1063,6 +1556,23 @@ struct DefaultB2bConv2dFprop < MmaCore1::Shape::kK, //kBlocksColumn ElementAccumulator, ElementA, AccumulatorLayout, InstructionShape, EpilogueOutputOp0>; + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 2; + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Warp-level iterators to load scale and bias vectors + using FragmentIteratorA1ScaleBias = cutlass::transform::warp::VectorFragmentIterator< + MatrixShape<1, IteratorAccumulatorScaleBias::Fragment::kElements>, ElementScaleBias, + LayoutScaleBias, InstructionShape, kElementsPerAccess>; + // Define iterators over tiles from the B operand using ThreadMapB1 = typename MmaCore1::IteratorThreadMapB; using IteratorB1 = @@ -1090,6 +1600,187 @@ struct DefaultB2bConv2dFprop < SmemIteratorB0, ThreadblockShape1, FragmentIteratorA1, + IteratorAccumulatorScaleBias, + FragmentIteratorA1ScaleBias, + IteratorB1, + SmemIteratorB1, + ElementC, + LayoutC, + EpilogueOutputOp0, + MmaPolicy0, + MmaPolicy1 + >; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape1, + WarpMmaTensorOp1, + 1, + EpilogueOutputOp1 + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::B2bImplicitGemmConvolution< + B2bMma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialzation for Optimized IteratorAlgorithm +/// and 2 stage pipeline. +/// Accumulator will be staged in shared memory. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape0, + typename ThreadblockShape1, + typename WarpShape0, + typename WarpShape1, + typename InstructionShape, + typename EpilogueOutputOp0, + typename EpilogueOutputOp1, + typename ThreadblockSwizzle, + typename MathOperatorTag +> +struct DefaultB2bConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape0, + ThreadblockShape1, + WarpShape0, + WarpShape1, + InstructionShape, + EpilogueOutputOp0, + EpilogueOutputOp1, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + true +> { + + // Define the core components from GEMM + using MmaCore0 = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape0, WarpShape0, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + using MmaCore1 = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape1, WarpShape1, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA0 = typename MmaCore0::IteratorThreadMapA; + using IteratorA0 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA0 + > + >; + + using SmemIteratorA0 = typename MmaCore0::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB0 = typename MmaCore0::IteratorThreadMapB; + using IteratorB0 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB0 + > + >; + + using SmemIteratorB0 = typename MmaCore0::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 2; + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Define iterators over tiles from the B operand + using ThreadMapB1 = typename MmaCore1::IteratorThreadMapB; + using IteratorB1 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB1 + > + >; + + using SmemIteratorB1 = typename MmaCore1::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp0 = typename MmaCore0::MmaTensorOp; + using WarpMmaTensorOp1 = typename MmaCore1::MmaTensorOp; + using MmaPolicy0 = typename MmaCore0::MmaPolicy; + using MmaPolicy1 = typename MmaCore1::MmaPolicy; + + // Use fragment iterator for the accumulator + using SmemAccumulatorLayout = cutlass::layout::RowMajor; + using FragmentIteratorAccumulator = cutlass::epilogue::warp::FragmentIteratorTensorOp< + WarpShape0, InstructionShape, + ElementAccumulator, + typename WarpMmaTensorOp0::Policy::Operator::FragmentC, + SmemAccumulatorLayout + >; + + // Store Accumulator tiles to Shared Memory + using SmemIteratorD0 = + cutlass::epilogue::warp::TileIteratorTensorOp< + WarpShape0, + InstructionShape, + ElementC, + SmemAccumulatorLayout + >; + + static int const kThreadCount = 32; + // load warp tile from Shared Memory accumulator + using WarpIteratorA1 = cutlass::gemm::warp::MmaTensorOpMultiplicandTileIterator< + MatrixShape, cutlass::gemm::Operand::kA, + ElementA, SmemAccumulatorLayout, + MatrixShape, + WarpMmaTensorOp1::Policy::OpDelta::kRow, kThreadCount>; + + // Define the Mma + using B2bMma = threadblock::B2bImplicitGemmPipelinedSmemAccumulator< + ThreadblockShape0, + IteratorA0, + SmemIteratorA0, + IteratorB0, + SmemIteratorB0, + IteratorAccumulatorScaleBias, + FragmentIteratorAccumulator, + SmemIteratorD0, + ThreadblockShape1, + WarpIteratorA1, IteratorB1, SmemIteratorB1, ElementC, @@ -1216,6 +1907,23 @@ struct DefaultB2bConv2dFprop < MmaCore1::Shape::kK, //kBlocksColumn ElementAccumulator, ElementA, AccumulatorLayout, InstructionShape, EpilogueOutputOp0>; + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 4; + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + // Warp-level iterators to load scale and bias vectors + using FragmentIteratorA1ScaleBias = cutlass::transform::warp::VectorFragmentIterator< + MatrixShape<1, IteratorAccumulatorScaleBias::Fragment::kElements>, ElementScaleBias, + LayoutScaleBias, InstructionShape, kElementsPerAccess>; + using ThreadMapB1 = typename MmaCore1::SmemThreadMapB; using IteratorB1 = cutlass::conv::threadblock::TileIterator< @@ -1228,7 +1936,6 @@ struct DefaultB2bConv2dFprop < using SmemIteratorB1 = typename MmaCore1::SmemIteratorB; - // Warp-level GEMM components using WarpMmaTensorOp1 = typename MmaCore1::MmaTensorOp; using MmaPolicy0 = typename MmaCore0::MmaPolicy; @@ -1243,6 +1950,8 @@ struct DefaultB2bConv2dFprop < SmemIteratorB0, ThreadblockShape1, FragmentIteratorA1, + IteratorAccumulatorScaleBias, + FragmentIteratorA1ScaleBias, IteratorB1, SmemIteratorB1, ElementC, @@ -1271,6 +1980,197 @@ struct DefaultB2bConv2dFprop < >; }; +/// Defines a kernel for Conv2dFprop specialzation for Optimized IteratorAlgorithm and 2 stage +/// pipeline with interleaved layout. +/// Accumulator will be staged in shared memory. +template < + typename ElementA, + typename ElementB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape0, + typename ThreadblockShape1, + typename WarpShape0, + typename WarpShape1, + typename InstructionShape, + typename EpilogueOutputOp0, + typename EpilogueOutputOp1, + typename ThreadblockSwizzle, + typename MathOperatorTag, + int InterleavedK +> +struct DefaultB2bConv2dFprop < + ElementA, + layout::TensorNCxHWx, + ElementB, + layout::TensorCxRSKx, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape0, + ThreadblockShape1, + WarpShape0, + WarpShape1, + InstructionShape, + EpilogueOutputOp0, + EpilogueOutputOp1, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + true +> { + + // Define the core components from GEMM + using MmaCore0 = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape0, WarpShape0, InstructionShape, ElementA, layout::ColumnMajorInterleaved, + ElementB, layout::RowMajorInterleaved, + ElementAccumulator, LayoutC, arch::OpClassTensorOp, + 2, MathOperatorTag, true>; + using MmaCore1 = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape1, WarpShape1, InstructionShape, ElementA, layout::ColumnMajorInterleaved, + ElementB, layout::RowMajorInterleaved, + ElementAccumulator, LayoutC, arch::OpClassTensorOp, + 2, MathOperatorTag, true>; + + // Define iterators over tiles from the A operand + // Note GEMM shared memory threadmap is used here because conv global memory + // layout needs to be mapped to fprop which is similar to the crosswise + // layout which is used by the interleaved GEMM shared memory threadmap. + // The Interleaved GEMM global memory layout is similar to the congruous + // layout. + + // Define iterators over tiles from the A operand + using ThreadMapA0 = typename MmaCore0::SmemThreadMapA; + using IteratorA0 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementA, layout::TensorNCxHWx, + ThreadMapA0 + > + >; + + using SmemIteratorA0 = typename MmaCore0::SmemIteratorA; + + // Define iterators over tiles from the B operand + // Note GEMM shared memory threadmap is used here because conv global memory + // layout needs to be mapped to fprop which is similar to the crosswise + // layout which is used by the interleaved GEMM shared memory threadmap. + // The Interleaved GEMM global memory layout is similar to the congruous + // layout. + using ThreadMapB0 = typename MmaCore0::SmemThreadMapB; + using IteratorB0 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementB, layout::TensorCxRSKx, + ThreadMapB0 + > + >; + + using SmemIteratorB0 = typename MmaCore0::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using ElementScaleBias = typename EpilogueOutputOp0::ElementCompute; + using LayoutScaleBias = layout::RowMajor; //vector layout doesn't really matter + static int const kElementsPerAccess = 4; //For interleaved layout + using IteratorAccumulatorScaleBias = + cutlass::transform::threadblock::VectorIterator< + cutlass::transform::threadblock::PredicatedVectorAccessIterator< + cutlass::MatrixShape, + cutlass::MatrixShape, + ElementScaleBias, LayoutScaleBias, kElementsPerAccess> + >; + + using ThreadMapB1 = typename MmaCore1::SmemThreadMapB; + using IteratorB1 = + cutlass::conv::threadblock::TileIterator< + cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementB, layout::TensorCxRSKx, + ThreadMapB1 + > + >; + + using SmemIteratorB1 = typename MmaCore1::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp0 = typename MmaCore0::MmaTensorOp; + using WarpMmaTensorOp1 = typename MmaCore1::MmaTensorOp; + using MmaPolicy0 = typename MmaCore0::MmaPolicy; + using MmaPolicy1 = typename MmaCore1::MmaPolicy; + + // Use fragment iterator for the accumulator + using SmemAccumulatorLayout = cutlass::layout::ColumnMajorInterleaved<16>; + using FragmentIteratorAccumulator = cutlass::epilogue::warp::FragmentIteratorTensorOp< + WarpShape0, InstructionShape, + ElementAccumulator, + typename WarpMmaTensorOp0::Policy::Operator::FragmentC, + SmemAccumulatorLayout + >; + + + // Store Accumulator tiles to Shared Memory + using SmemIteratorD0 = + cutlass::epilogue::warp::TileIteratorTensorOp< + WarpShape0, + InstructionShape, + ElementC, + SmemAccumulatorLayout + >; + + static int const kThreadCount = 32; + // load warp tile from Shared Memory accumulator + using WarpIteratorA1 = cutlass::gemm::warp::MmaTensorOpMultiplicandTileIteratorCanonical< + MatrixShape, cutlass::gemm::Operand::kA, + ElementA, SmemAccumulatorLayout, + MatrixShape, + WarpMmaTensorOp1::Policy::OpDelta::kRow, kThreadCount>; + + // Define the Mma + using B2bMma = threadblock::B2bImplicitGemmPipelinedSmemAccumulator< + ThreadblockShape0, + IteratorA0, + SmemIteratorA0, + IteratorB0, + SmemIteratorB0, + IteratorAccumulatorScaleBias, + FragmentIteratorAccumulator, + SmemIteratorD0, + ThreadblockShape1, + WarpIteratorA1, + IteratorB1, + SmemIteratorB1, + ElementC, + LayoutC, + EpilogueOutputOp0, + MmaPolicy0, + MmaPolicy1 + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultInterleavedConvEpilogue< + ThreadblockShape1, + WarpMmaTensorOp1, + 1, + EpilogueOutputOp1, + EpilogueOutputOp1::kCount, + InterleavedK + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::B2bImplicitGemmConvolution< + B2bMma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_multistage.h b/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_multistage.h index 9b0d7968..aebc87dd 100644 --- a/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_multistage.h +++ b/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_multistage.h @@ -77,6 +77,11 @@ template < /// Iterates over the intermediate accumulator tile // (concept::MmaTensorOpFragmentIterator) typename FragmentIteratorA1_, + /// Iterates over vectors of scale and bias vector in global memory + // (concept: VectorIterator) + typename IteratorAccumulatorScaleBias_, + /// WarpIterator to load Scale or Bias vector from threadblock fragment + typename FragmentIteratorA1ScaleBias_, /// Iterates over tiles of B operand in global memory // (concept: ReadableTileIterator | ForwardTileIterator | // MaskedTileIterator) @@ -117,6 +122,10 @@ public: using Shape1 = Shape1_; ///< Iterates over tiles of A operand in global memory using FragmentIteratorA1 = FragmentIteratorA1_; + ///< Iterates over tiles of the scale and bias vectors in global memory + using IteratorAccumulatorScaleBias = IteratorAccumulatorScaleBias_; + ///< WarpIterator to load Scale or Bias vector from threadblock fragment + using FragmentIteratorA1ScaleBias = FragmentIteratorA1ScaleBias_; ///< Iterates over tiles of B operand in global memory using IteratorB1 = IteratorB1_; ///< Policy describing tuning details @@ -126,6 +135,9 @@ public: ///< Epilogue after 1st Gemm using OutputOp = OutputOp_; + + static const bool PerChannelScale = (OutputOp::kScale == + epilogue::thread::ScaleType::OnlyAlphaPerChannelScaling); static cutlass::arch::CacheOperation::Kind const kCacheOpA0 = CacheOpA0; static cutlass::arch::CacheOperation::Kind const kCacheOpB0 = CacheOpB0; @@ -143,6 +155,9 @@ public: /// Warp-level Mma using Operator0 = typename Policy0::Operator; + /// Fragment of Scale and Bias loaded from global memory + using FragmentA1ScaleBias = typename IteratorAccumulatorScaleBias::Fragment; + /// Fragment of accumulator tile using FragmentC1 = typename Policy1::Operator::FragmentC; @@ -193,6 +208,8 @@ public: using WarpLoadedFragmentB0 = typename Operator0::FragmentB; /// Warp Fragment of operand A1 loaded from accmulator tile using WarpLoadedFragmentA1 = typename FragmentIteratorA1::Fragment; + using WarpLoadedFragmentA1ScaleBias = + typename FragmentIteratorA1ScaleBias::Fragment; using WarpLoadedFragmentB1 = typename Operator1::FragmentB; using WarpTransformedFragmentA0 = typename Operator0::TransformedFragmentA; using WarpTransformedFragmentB0 = typename Operator0::TransformedFragmentB; @@ -229,9 +246,9 @@ public: int lane_idx ): Base(shared_storage, thread_idx, warp_idx, lane_idx), - smem_iterator_A0_(shared_storage.sharedStorage0.operand_A_ref(), thread_idx), - smem_iterator_B0_(shared_storage.sharedStorage0.operand_B_ref(), thread_idx), - smem_iterator_B1_(shared_storage.sharedStorage1.operand_B_ref(), thread_idx) + smem_iterator_A0_(shared_storage.shared_storage0.operand_A_ref(), thread_idx), + smem_iterator_B0_(shared_storage.shared_storage0.operand_B_ref(), thread_idx), + smem_iterator_B1_(shared_storage.shared_storage1.operand_B_ref(), thread_idx) { // Compute warp location within threadblock tile by mapping the warp_id to // three coordinates: @@ -343,11 +360,15 @@ public: int gemm_k_iterations_0, ///< destination accumulator tile FragmentC1 &accum, - ///< iterator over A operand in global memory + ///< iterator over A0 operand in global memory IteratorA0 iterator_A0, - ///< iterator over B operand in global memory + ///< iterator over B0 operand in global memory IteratorB0 iterator_B0, - ///< iterator over B operand in global memory + ///< iterator over A1 operand scale vector in global memory + IteratorAccumulatorScaleBias iterator_A1_scale, + ///< iterator over A1 operand bias vector in global memory + IteratorAccumulatorScaleBias iterator_A1_bias, + ///< iterator over B1 operand in global memory IteratorB1 iterator_B1, ///< initial value of accumulator FragmentC0 const &src_accum, @@ -571,6 +592,20 @@ public: /// Iterator to load a warp-scoped tile of A1 operand from intermediate accumulator tile FragmentIteratorA1 warp_tile_iterator_A1_(accum0); + FragmentA1ScaleBias tb_frag_A1_scale; + FragmentA1ScaleBias tb_frag_A1_bias; + FragmentIteratorA1ScaleBias warp_tile_iterator_A1_scale_(tb_frag_A1_scale); + FragmentIteratorA1ScaleBias warp_tile_iterator_A1_bias_(tb_frag_A1_bias); + + if(PerChannelScale) { + tb_frag_A1_scale.clear(); + iterator_A1_scale.load(tb_frag_A1_scale); + ++iterator_A1_scale; + } + tb_frag_A1_bias.clear(); + iterator_A1_bias.load(tb_frag_A1_bias); + ++iterator_A1_bias; + // // Prologue @@ -619,18 +654,29 @@ public: // Pair of fragments used to overlap shared memory loads and math // instructions WarpLoadedFragmentA1 warp_loaded_frag_A1[2]; + WarpLoadedFragmentA1ScaleBias warp_loaded_frag_A1_scale[2]; + WarpLoadedFragmentA1ScaleBias warp_loaded_frag_A1_bias[2]; WarpLoadedFragmentB1 warp_loaded_frag_B1[2]; WarpTransformedFragmentA1 warp_transformed_frag_A1[2]; WarpTransformedFragmentB1 warp_transformed_frag_B1[2]; Operator1 warp_mma1; - this->warp_tile_iterator_B1_.set_kgroup_index(0); - - warp_tile_iterator_A1_.load(warp_loaded_frag_A1[0], output_op_0); - this->warp_tile_iterator_B1_.load(warp_loaded_frag_B1[0]); + if(PerChannelScale) { + warp_tile_iterator_A1_scale_.load(warp_loaded_frag_A1_scale[0]); + ++warp_tile_iterator_A1_scale_; + } + warp_tile_iterator_A1_bias_.load(warp_loaded_frag_A1_bias[0]); + ++warp_tile_iterator_A1_bias_; + warp_tile_iterator_A1_.load(warp_loaded_frag_A1[0], + warp_loaded_frag_A1_scale[0], + warp_loaded_frag_A1_bias[0], + output_op_0); ++warp_tile_iterator_A1_; + + this->warp_tile_iterator_B1_.set_kgroup_index(0); + this->warp_tile_iterator_B1_.load(warp_loaded_frag_B1[0]); ++this->warp_tile_iterator_B1_; // Start issuing the first group of the next stage outside of the mainloop @@ -660,17 +706,40 @@ public: for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations1; ++warp_mma_k) { + // Load threadblock-level scale/bias vector from global memory + if (warp_mma_k + 1 == Base::kWarpGemmIterations1) { + if(PerChannelScale) { + tb_frag_A1_scale.clear(); + iterator_A1_scale.load(tb_frag_A1_scale); + ++iterator_A1_scale; + } + tb_frag_A1_bias.clear(); + iterator_A1_bias.load(tb_frag_A1_bias); + ++iterator_A1_bias; + } + + // Load warp-level scale bias fragment from threadblock scale/bias vector + if(PerChannelScale) { + warp_tile_iterator_A1_scale_.load(warp_loaded_frag_A1_scale[(warp_mma_k + 1) % 2]); + ++warp_tile_iterator_A1_scale_; + } + warp_tile_iterator_A1_bias_.load(warp_loaded_frag_A1_bias[(warp_mma_k + 1) % 2]); + ++warp_tile_iterator_A1_bias_; + + // Load warp-level tile from accumulator fragment + warp_tile_iterator_A1_.load(warp_loaded_frag_A1[(warp_mma_k + 1) % 2], + warp_loaded_frag_A1_scale[(warp_mma_k + 1) % 2], + warp_loaded_frag_A1_bias[(warp_mma_k + 1) % 2], + output_op_0); + ++warp_tile_iterator_A1_; + // Load warp-level tiles from shared memory, wrapping to k offset if // this is the last group as the case may be. - this->warp_tile_iterator_B1_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations1); - - warp_tile_iterator_A1_.load(warp_loaded_frag_A1[(warp_mma_k + 1) % 2], output_op_0); this->warp_tile_iterator_B1_.load(warp_loaded_frag_B1[(warp_mma_k + 1) % 2]); - - ++warp_tile_iterator_A1_; ++this->warp_tile_iterator_B1_; + if (warp_mma_k > 0) warp_mma1.transform(warp_transformed_frag_A1[warp_mma_k % 2], warp_transformed_frag_B1[warp_mma_k % 2], diff --git a/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_pipelined.h b/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_pipelined.h index c26aa2fb..92b069bb 100644 --- a/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_pipelined.h +++ b/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_pipelined.h @@ -70,6 +70,12 @@ template < /// Iterates over the intermediate accumulator tile // (concept::MmaTensorOpFragmentIterator) typename FragmentIteratorA1_, + /// Iterates over vectors of scale and bias vector in global memory + // (concept: VectorIterator) + typename IteratorAccumulatorScaleBias_, + /// FragmentIterator to load Scale or Bias vector from threadblock fragment + typename FragmentIteratorA1ScaleBias_, + // (concept: VectorFragmentIterator) /// Iterates over tiles of B operand in global memory // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) typename IteratorB1_, @@ -92,13 +98,13 @@ template < typename IteratorA0_::Element, IteratorA0_::Fragment::kElements>, /// - /// Transformation applied to A operand + /// Transformation applied to B operand typename TransformB0_ = NumericArrayConverter< typename SmemIteratorB0_::Element, typename IteratorB0_::Element, IteratorB0_::Fragment::kElements>, /// - /// Transformation applied to A operand + /// Transformation applied to B operand typename TransformB1_ = NumericArrayConverter< typename SmemIteratorB1_::Element, typename IteratorB1_::Element, @@ -106,7 +112,8 @@ template < /// Used for partial specialization typename Enable = bool > -class B2bImplicitGemmPipelined : public gemm::threadblock::B2bMmaBase { +class B2bImplicitGemmPipelined : + public gemm::threadblock::B2bMmaBase { public: ///< Base class @@ -121,17 +128,24 @@ public: using SmemIteratorB0 = SmemIteratorB0_; using Shape1 = Shape1_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> - using FragmentIteratorA1 = FragmentIteratorA1_; ///< Iterates over tiles of A operand in global memory + using FragmentIteratorA1 = FragmentIteratorA1_; ///< Iterates over tiles of A1 operand from accumulator tile + using IteratorAccumulatorScaleBias = IteratorAccumulatorScaleBias_; ///< Iterates over tiles of the scale and bias vectors in global memory + using FragmentIteratorA1ScaleBias = + FragmentIteratorA1ScaleBias_; ///< WarpIterator to load Scale or Bias vector from the threadblock fragment using IteratorB1 = IteratorB1_; ///< Iterates over tiles of B operand in global memory using Policy1 = Policy1_; ///< Policy1 describing tuning details using SmemIteratorB1 = SmemIteratorB1_; + using ElementC = ElementC_; ///< Data type of accumulator matrix using LayoutC = LayoutC_; ///< Layout of accumulator matrix using OutputOp = OutputOp_; ///< Epilogue after 1st Gemm + static const bool PerChannelScale = (OutputOp::kScale == + epilogue::thread::ScaleType::OnlyAlphaPerChannelScaling); + using TransformA0 = TransformA0_; using TransformB0 = TransformB0_; using TransformB1 = TransformB1_; @@ -152,6 +166,9 @@ public: /// Warp-level Mma using Operator0 = typename Policy0::Operator; + /// Fragment of Scale and Bias loaded from global memory + using FragmentA1ScaleBias = typename IteratorAccumulatorScaleBias::Fragment; + /// Fragment of operand B loaded from global memory using FragmentB1 = typename IteratorB1::Fragment; @@ -182,6 +199,9 @@ private: using WarpFragmentB0 = typename Operator0::FragmentB; /// Warp Fragment of operand A1 loaded from accmulator tile using WarpFragmentA1 = typename FragmentIteratorA1::Fragment; + /// Warp Fragment of operand A1 scale and bias loaded from threadblock fragment + using WarpFragmentA1ScaleBias = + typename FragmentIteratorA1ScaleBias::Fragment; using WarpFragmentB1 = typename Operator1::FragmentB; protected: @@ -206,9 +226,9 @@ public: int lane_idx ///< ID of each thread within a warp ): Base(shared_storage, thread_idx, warp_idx, lane_idx), - smem_iterator_A_(shared_storage.sharedStorage0.operand_A_ref(), thread_idx), - smem_iterator_B0_(shared_storage.sharedStorage0.operand_B_ref(), thread_idx), - smem_iterator_B1_(shared_storage.sharedStorage1.operand_B_ref(), thread_idx) { + smem_iterator_A_(shared_storage.shared_storage0.operand_A_ref(), thread_idx), + smem_iterator_B0_(shared_storage.shared_storage0.operand_B_ref(), thread_idx), + smem_iterator_B1_(shared_storage.shared_storage1.operand_B_ref(), thread_idx) { // Compute warp location within threadblock tile by mapping the warp_id to // three coordinates: @@ -240,9 +260,11 @@ public: FragmentC1 &accum, ///< destination accumulator tile IteratorA0 iterator_A, ///< iterator over A operand in global memory IteratorB0 iterator_B0, ///< iterator over B0 operand in global memory + IteratorAccumulatorScaleBias iterator_A1_scale, ///< iterator over A1 operand scale vectors in global memory + IteratorAccumulatorScaleBias iterator_A1_bias, ///< iterator over A1 operand bias vectors in global memory IteratorB1 iterator_B1, ///< iterator over B1 operand in global memory FragmentC0 const &src_accum, ///< source accumulator tile - OutputOp output_op_0, ///< epilogue operation after 1st Gemm + OutputOp output_op_0, ///< epilogue operation after 1st Gemm TransformA0 transform_A0 = TransformA0(), ///< transformation applied to A0 fragment TransformB0 transform_B0 = TransformB0(), ///< transformation applied to B0 fragment TransformB1 transform_B1 = TransformB1()) { ///< transformation applied to B1 fragment @@ -370,18 +392,33 @@ public: /// Iterator to load a warp-scoped tile of A1 operand from intermediate accumulator tile FragmentIteratorA1 warp_tile_iterator_A1_(accum0); + + // // Prologue // + FragmentA1ScaleBias tb_frag_A1_scale; + FragmentA1ScaleBias tb_frag_A1_bias; + FragmentIteratorA1ScaleBias warp_tile_iterator_A1_scale_(tb_frag_A1_scale); + FragmentIteratorA1ScaleBias warp_tile_iterator_A1_bias_(tb_frag_A1_bias); FragmentB1 tb_frag_B1; + if(PerChannelScale) + tb_frag_A1_scale.clear(); + tb_frag_A1_bias.clear(); tb_frag_B1.clear(); // The last kblock is loaded in the prolog + if(PerChannelScale) + iterator_A1_scale.load(tb_frag_A1_scale); + iterator_A1_bias.load(tb_frag_A1_bias); iterator_B1.load(tb_frag_B1); + if(PerChannelScale) + ++iterator_A1_scale; + ++iterator_A1_bias; ++iterator_B1; this->smem_iterator_B1_.store(transform_B1(tb_frag_B1)); @@ -391,15 +428,24 @@ public: __syncthreads(); // Pair of fragments used to overlap shared memory loads and math instructions + WarpFragmentA1ScaleBias warp_frag_A1_scale[2]; + WarpFragmentA1ScaleBias warp_frag_A1_bias[2]; WarpFragmentA1 warp_frag_A1[2]; WarpFragmentB1 warp_frag_B1[2]; this->warp_tile_iterator_B1_.set_kgroup_index(0); - warp_tile_iterator_A1_.load(warp_frag_A1[0], output_op_0); + if(PerChannelScale) + warp_tile_iterator_A1_scale_.load(warp_frag_A1_scale[0]); + warp_tile_iterator_A1_bias_.load(warp_frag_A1_bias[0]); + warp_tile_iterator_A1_.load(warp_frag_A1[0], warp_frag_A1_scale[0], + warp_frag_A1_bias[0], output_op_0); this->warp_tile_iterator_B1_.load(warp_frag_B1[0]); ++warp_tile_iterator_A1_; + if(PerChannelScale) + ++warp_tile_iterator_A1_scale_; + ++warp_tile_iterator_A1_bias_; ++this->warp_tile_iterator_B1_; Operator1 warp_mma1; @@ -447,13 +493,31 @@ public: } smem_write_stage_idx ^= 1; + + if(PerChannelScale) { + tb_frag_A1_scale.clear(); + iterator_A1_scale.load(tb_frag_A1_scale); + ++iterator_A1_scale; + } + tb_frag_A1_bias.clear(); + iterator_A1_bias.load(tb_frag_A1_bias); + ++iterator_A1_bias; } this->warp_tile_iterator_B1_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations1); - warp_tile_iterator_A1_.load(warp_frag_A1[(warp_mma_k + 1) % 2], output_op_0); + if(PerChannelScale) + warp_tile_iterator_A1_scale_.load(warp_frag_A1_scale[(warp_mma_k + 1) % 2]); + warp_tile_iterator_A1_bias_.load(warp_frag_A1_bias[(warp_mma_k + 1) % 2]); + warp_tile_iterator_A1_.load(warp_frag_A1[(warp_mma_k + 1) % 2], + warp_frag_A1_scale[(warp_mma_k + 1) % 2], + warp_frag_A1_bias[(warp_mma_k + 1) % 2], + output_op_0); this->warp_tile_iterator_B1_.load(warp_frag_B1[(warp_mma_k + 1) % 2]); + if(PerChannelScale) + ++warp_tile_iterator_A1_scale_; + ++warp_tile_iterator_A1_bias_; ++warp_tile_iterator_A1_; ++this->warp_tile_iterator_B1_; diff --git a/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_pipelined_smem_accumulator.h b/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_pipelined_smem_accumulator.h new file mode 100644 index 00000000..ea462f12 --- /dev/null +++ b/examples/13_two_tensor_op_fusion/threadblock/b2b_implicit_gemm_pipelined_smem_accumulator.h @@ -0,0 +1,532 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/array.h" +#include "cutlass/aligned_buffer.h" +#include "cutlass/numeric_conversion.h" + +#include "cutlass/numeric_types.h" +#include "cutlass/matrix_shape.h" + +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/warp/mma_tensor_op_fragment_iterator.h" + +#include "threadblock/b2b_mma_base_smem_accumulator.h" +#include "cutlass/epilogue/threadblock/epilogue_smem_accumulator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape0_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorA0_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA0_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorB0_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB0_, + /// Iterates over vectors of scale and bias vector in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorAccumulatorScaleBias_, + /// Iterates over accumulator tile + typename FragmentIteratorAccumulator_, + /// Iterates over accumulator tile in shared memory + typename SmemIteratorD0_, + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape1_, + /// Iterates over the intermediate accumulator tile in shared memory + typename WarpIteratorA1_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorB1_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB1_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Output operator for 1st Gemm(concept: epilogue::thread::LinearCombinationClamp, etc...) + typename OutputOp_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy0_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy1_, + /// Transformation applied to A operand + typename TransformA0_ = NumericArrayConverter< + typename SmemIteratorA0_::Element, + typename IteratorA0_::Element, + IteratorA0_::Fragment::kElements>, + /// + /// Transformation applied to B operand + typename TransformB0_ = NumericArrayConverter< + typename SmemIteratorB0_::Element, + typename IteratorB0_::Element, + IteratorB0_::Fragment::kElements>, + /// + /// Transformation applied to B operand + typename TransformB1_ = NumericArrayConverter< + typename SmemIteratorB1_::Element, + typename IteratorB1_::Element, + IteratorB1_::Fragment::kElements>, + /// Used for partial specialization + typename Enable = bool +> +class B2bImplicitGemmPipelinedSmemAccumulator : + public gemm::threadblock::B2bMmaBaseSmemAccumulator { +public: + + ///< Base class + using Base = gemm::threadblock::B2bMmaBaseSmemAccumulator; + + using Shape0 = Shape0_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using IteratorA0 = IteratorA0_; ///< Iterates over tiles of A operand in global memory + using IteratorB0 = IteratorB0_; ///< Iterates over tiles of B operand in global memory + using IteratorAccumulatorScaleBias = IteratorAccumulatorScaleBias_; ///< Iterates over tiles of the scale and bias vectors in global memory + using Policy0 = Policy0_; ///< Policy0 describing tuning details + + using SmemIteratorA0 = SmemIteratorA0_; + using SmemIteratorB0 = SmemIteratorB0_; + using SmemIteratorD0 = SmemIteratorD0_; ///< Iterates over accumulator tile in shared memory + + using FragmentIteratorAccumulator = FragmentIteratorAccumulator_; ///< Iterates over accumulator tile + + using Shape1 = Shape1_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using IteratorB1 = IteratorB1_; ///< Iterates over tiles of B operand in global memory + using Policy1 = Policy1_; ///< Policy1 describing tuning details + + using SmemIteratorB1 = SmemIteratorB1_; + using WarpIteratorA1 = WarpIteratorA1_; ///< Iterates over the intermediate accumulator tile in shared memory + + + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + + using OutputOp = OutputOp_; ///< Epilogue after 1st Gemm + + using TransformA0 = TransformA0_; + using TransformB0 = TransformB0_; + using TransformB1 = TransformB1_; + + // + // Dependent types + // + + /// Fragment of operand A loaded from global memory + using FragmentA0 = typename IteratorA0::Fragment; + + /// Fragment of operand B loaded from global memory + using FragmentB0 = typename IteratorB0::Fragment; + + /// Fragment of accumulator tile + using FragmentC0 = typename Policy0::Operator::FragmentC; + + /// Warp-level Mma + using Operator0 = typename Policy0::Operator; + + /// Fragment of operand B loaded from global memory + using FragmentB1 = typename IteratorB1::Fragment; + + /// Fragment of accumulator tile + using FragmentC1 = typename Policy1::Operator::FragmentC; + + /// Warp-level Mma + using Operator1 = typename Policy1::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy0::Operator::ArchTag; + + /// Complex transform on A0 operand + static ComplexTransform const kTransformA0 = Operator0::kTransformA; + + /// Complex transform on B0 operand + static ComplexTransform const kTransformB0 = Operator0::kTransformB; + + /// Complex transform on B1 operand + static ComplexTransform const kTransformB1 = Operator1::kTransformB; + + /// staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) + static_assert((Base::kStages==2), "MmaPipelined requires kStages set to value 2"); + + /// Epilog in shared memory + using Epilogue0 = epilogue::threadblock::EpilogueSmemAccumulator< + SmemIteratorD0, ///< SmemTileIterator + FragmentIteratorAccumulator, ///< AccumulatorFragmentIterator + IteratorAccumulatorScaleBias, ///< ScaleBiasIterator + OutputOp>; ///< Output operator + + + +private: + + using WarpFragmentA0 = typename Operator0::FragmentA; + using WarpFragmentB0 = typename Operator0::FragmentB; + using WarpFragmentA1 = typename Operator1::FragmentA; + using WarpFragmentB1 = typename Operator1::FragmentB; + +protected: + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA0 smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B0 operand to shared memory + SmemIteratorB0 smem_iterator_B0_; + + /// Shared Memory Iterator to store accumulator tile + SmemIteratorD0 smem_iterator_D0_; + + /// Iterator to load a warp-scoped tile of A1 operand from intermediate accumulator tile + WarpIteratorA1 warp_tile_iterator_A1_; + + /// Iterator to write threadblock-scoped tile of B1 operand to shared memory + SmemIteratorB1 smem_iterator_B1_; + +public: + + /// Construct from tensor references + CUTLASS_DEVICE + B2bImplicitGemmPipelinedSmemAccumulator( + typename Base::B2bMmaSharedStorage &shared_storage, ///< Shared storage needed for internal use by threadblock-scoped GEMM + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx ///< ID of each thread within a warp + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.b2b_mma_shared_storage.shared_storage0.operand_A_ref(), thread_idx), + smem_iterator_B0_(shared_storage.b2b_mma_shared_storage.shared_storage0.operand_B_ref(), thread_idx), + smem_iterator_D0_(shared_storage.accumulator_shared_storage0.accum_ref(), lane_idx), + warp_tile_iterator_A1_(shared_storage.accumulator_shared_storage0.accum_ref(), lane_idx), + smem_iterator_B1_(shared_storage.b2b_mma_shared_storage.shared_storage1.operand_B_ref(), thread_idx) { + + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn_0 = warp_idx % (Base::WarpCount0::kM * Base::WarpCount0::kN); + int warp_idx_k_0 = warp_idx / (Base::WarpCount0::kM * Base::WarpCount0::kN); + + int warp_idx_m_0 = warp_idx_mn_0 % Base::WarpCount0::kM; + int warp_idx_n_0 = warp_idx_mn_0 / Base::WarpCount0::kM; + + int tile_offset_k_0 = Base::kWarpGemmIterations0 * warp_idx_k_0; + + int warp_idx_mn_1 = warp_idx % (Base::WarpCount1::kM * Base::WarpCount1::kN); + int warp_idx_k_1 = warp_idx / (Base::WarpCount1::kM * Base::WarpCount1::kN); + + int warp_idx_m_1 = warp_idx_mn_1 % Base::WarpCount1::kM; + int warp_idx_n_1 = warp_idx_mn_1 / Base::WarpCount1::kM; + + int tile_offset_k_1 = Base::kWarpGemmIterations1 * warp_idx_k_1; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A0_.add_tile_offset({warp_idx_m_0, tile_offset_k_0}); + this->warp_tile_iterator_B0_.add_tile_offset({tile_offset_k_0, warp_idx_n_0}); + warp_tile_iterator_A1_.add_tile_offset({warp_idx_m_1, tile_offset_k_1}); + this->warp_tile_iterator_B1_.add_tile_offset({tile_offset_k_1, warp_idx_n_1}); + + // Add smem accumulator iterator warp offset + smem_iterator_D0_.add_tile_offset({ warp_idx_m_0 * SmemIteratorD0::TileIterations::kRow, + warp_idx_n_0 * SmemIteratorD0::TileIterations::kColumn}); + + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + int gemm_k_iterations_0, ///< number of iterations of the mainloop + FragmentC1 &accum, ///< destination accumulator tile + IteratorA0 iterator_A, ///< iterator over A operand in global memory + IteratorB0 iterator_B0, ///< iterator over B0 operand in global memory + IteratorAccumulatorScaleBias iterator_accum0_scale, ///< iterator over D0 scale vector in global memory + IteratorAccumulatorScaleBias iterator_accum0_bias, ///< iterator over D0 bias vector in global memory + IteratorB1 iterator_B1, ///< iterator over B1 operand in global memory + FragmentC0 const &src_accum, ///< source accumulator tile + OutputOp output_op_0, ///< epilogue operation after 1st Gemm + TransformA0 transform_A0 = TransformA0(), ///< transformation applied to A0 fragment + TransformB0 transform_B0 = TransformB0(), ///< transformation applied to B0 fragment + TransformB1 transform_B1 = TransformB1()) { ///< transformation applied to B1 fragment + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + FragmentC0 accum0 = src_accum; + + FragmentA0 tb_frag_A; + FragmentB0 tb_frag_B0; + + tb_frag_A.clear(); + tb_frag_B0.clear(); + + // The last kblock is loaded in the prolog + iterator_A.load(tb_frag_A); + iterator_B0.load(tb_frag_B0); + + ++iterator_A; + ++iterator_B0; + + this->smem_iterator_A_.store(transform_A0(tb_frag_A)); + this->smem_iterator_B0_.store(transform_B0(tb_frag_B0)); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B0_; + + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math instructions + WarpFragmentA0 warp_frag_A0[2]; + WarpFragmentB0 warp_frag_B0[2]; + + this->warp_tile_iterator_A0_.set_kgroup_index(0); + this->warp_tile_iterator_B0_.set_kgroup_index(0); + + this->warp_tile_iterator_A0_.load(warp_frag_A0[0]); + this->warp_tile_iterator_B0_.load(warp_frag_B0[0]); + + ++this->warp_tile_iterator_A0_; + ++this->warp_tile_iterator_B0_; + + Operator0 warp_mma0; + + int smem_write_stage_idx = 1; + + // Issue loads during the first warp-level matrix multiply-add *AFTER* issuing + // shared memory loads (which have the tighest latency requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations_0 > 0; --gemm_k_iterations_0) { + // + // Loop over GEMM K dimension + // + + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations0; ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group + // as the case may be. + + if (warp_mma_k == Base::kWarpGemmIterations0 - 1) { + + // Write fragments to shared memory + this->smem_iterator_A_.store(transform_A0(tb_frag_A)); + + this->smem_iterator_B0_.store(transform_B0(tb_frag_B0)); + + __syncthreads(); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B0_; + + // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory + if (smem_write_stage_idx == 1) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B0_.add_tile_offset({-Base::kStages, 0}); + } + else { + this->warp_tile_iterator_A0_.add_tile_offset( + {0, -Base::kStages * Policy0::kPartitionsK * Base::kWarpGemmIterations0}); + this->warp_tile_iterator_B0_.add_tile_offset( + {-Base::kStages * Policy0::kPartitionsK * Base::kWarpGemmIterations0, + 0}); + } + + smem_write_stage_idx ^= 1; + } + + this->warp_tile_iterator_A0_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations0); + this->warp_tile_iterator_B0_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations0); + + this->warp_tile_iterator_A0_.load(warp_frag_A0[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B0_.load(warp_frag_B0[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A0_; + ++this->warp_tile_iterator_B0_; + + if (warp_mma_k == 0) { + + iterator_A.load(tb_frag_A); + iterator_B0.load(tb_frag_B0); + + ++iterator_A; + ++iterator_B0; + } + + warp_mma0(accum0, warp_frag_A0[warp_mma_k % 2], + warp_frag_B0[warp_mma_k % 2], accum0); + + } + } + + /// Epilogue for the first Implicit Gemm + Epilogue0 epilogue0; + + epilogue0(output_op_0, smem_iterator_D0_, accum0, iterator_accum0_scale, iterator_accum0_bias); + + __syncthreads(); + + /// 2nd Implicit Gemm + + + // + // Prologue + // + + FragmentB1 tb_frag_B1; + + tb_frag_B1.clear(); + + // The last kblock is loaded in the prolog + iterator_B1.load(tb_frag_B1); + + ++iterator_B1; + + this->smem_iterator_B1_.store(transform_B1(tb_frag_B1)); + + ++this->smem_iterator_B1_; + + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math instructions + WarpFragmentA1 warp_frag_A1[2]; + WarpFragmentB1 warp_frag_B1[2]; + + this->warp_tile_iterator_B1_.set_kgroup_index(0); + + warp_tile_iterator_A1_.load(warp_frag_A1[0]); + this->warp_tile_iterator_B1_.load(warp_frag_B1[0]); + + ++warp_tile_iterator_A1_; + ++this->warp_tile_iterator_B1_; + + Operator1 warp_mma1; + + smem_write_stage_idx = 1; + +// int gemm_k_iterations_1 = FragmentIteratorA1::Policy::kIterations / Base::kWarpGemmIterations1; + int gemm_k_iterations_1 = Shape0::kN / Shape1::kK; + + // Issue loads during the first warp-level matrix multiply-add *AFTER* issuing + // shared memory loads (which have the tighest latency requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + CUTLASS_PRAGMA_UNROLL + for (; gemm_k_iterations_1 > 0; --gemm_k_iterations_1) { + // + // Loop over GEMM K dimension + // + + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations1; ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group + // as the case may be. + + if (warp_mma_k == Base::kWarpGemmIterations1 - 1) { + + this->smem_iterator_B1_.store(transform_B1(tb_frag_B1)); + + __syncthreads(); + + ++this->smem_iterator_B1_; + + // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory + if (smem_write_stage_idx == 1) { + this->smem_iterator_B1_.add_tile_offset({-Base::kStages, 0}); + } + else { + this->warp_tile_iterator_B1_.add_tile_offset( + {-Base::kStages * Policy1::kPartitionsK * Base::kWarpGemmIterations1, + 0}); + } + + smem_write_stage_idx ^= 1; + + } + + this->warp_tile_iterator_B1_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations1); + + // skip warp tile loading for the last kgroup + if(gemm_k_iterations_1 > 1 || warp_mma_k < Base::kWarpGemmIterations1 - 1) + warp_tile_iterator_A1_.load(warp_frag_A1[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B1_.load(warp_frag_B1[(warp_mma_k + 1) % 2]); + + ++warp_tile_iterator_A1_; + ++this->warp_tile_iterator_B1_; + + if (warp_mma_k == 0) { + + iterator_B1.load(tb_frag_B1); + + ++iterator_B1; + } + + warp_mma1(accum, warp_frag_A1[warp_mma_k % 2], + warp_frag_B1[warp_mma_k % 2], accum); + + } + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base.h b/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base.h index 3e87ce73..fc2576f0 100644 --- a/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base.h +++ b/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base.h @@ -180,8 +180,8 @@ class B2bMmaBase { using SharedStorage0 = SharedStorage; using SharedStorage1 = SharedStorage; union B2bMmaSharedStorage { - SharedStorage0 sharedStorage0; - SharedStorage1 sharedStorage1; + SharedStorage0 shared_storage0; + SharedStorage1 shared_storage1; }; @@ -197,7 +197,7 @@ class B2bMmaBase { /// Iterator to load a warp-scoped tile of B0 operand from shared memory typename Operator0::IteratorB warp_tile_iterator_B0_; - /// Iterator to load a warp-scoped tile of B0 operand from shared memory + /// Iterator to load a warp-scoped tile of B1 operand from shared memory typename Operator1::IteratorB warp_tile_iterator_B1_; public: @@ -214,9 +214,9 @@ public: ///< ID of each thread within a warp int lane_idx ): - warp_tile_iterator_A0_(shared_storage.sharedStorage0.operand_A_ref(), lane_idx), - warp_tile_iterator_B0_(shared_storage.sharedStorage0.operand_B_ref(), lane_idx), - warp_tile_iterator_B1_(shared_storage.sharedStorage1.operand_B_ref(), lane_idx) { + warp_tile_iterator_A0_(shared_storage.shared_storage0.operand_A_ref(), lane_idx), + warp_tile_iterator_B0_(shared_storage.shared_storage0.operand_B_ref(), lane_idx), + warp_tile_iterator_B1_(shared_storage.shared_storage1.operand_B_ref(), lane_idx) { } }; diff --git a/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h b/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h new file mode 100644 index 00000000..67c2d22b --- /dev/null +++ b/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_base_smem_accumulator.h @@ -0,0 +1,174 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "cutlass/aligned_buffer.h" +#include "cutlass/arch/memory.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/numeric_types.h" +#include "threadblock/b2b_mma_base.h" +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape0_, + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape1_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy0_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy1_, + /// Shared Memory Accumulator Iterator + typename SmemAccumulatorIterator0_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class B2bMmaBaseSmemAccumulator : + public B2bMmaBase { + + public: + ///< Base class + using Base = B2bMmaBase; + + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape0 = Shape0_; + using Shape1 = Shape1_; + + ///< Policy describing tuning details + using Policy0 = Policy0_; + using Policy1 = Policy1_; + + + using SmemAccumulatorIterator0 = SmemAccumulatorIterator0_; + + // + // Nested structs + // + /// Shared storage object needed by accumulator + template< + typename Shape_, + typename Element_, + typename Layout_, + typename Padding_ + > + class AccumulatorSharedStorage { + public: + // + // Type definitions + // + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using Padding = Padding_; + + /// Tensor reference to the accumulator + using TensorRefAccum = TensorRef; + + /// Shape of the accumulator matrix in shared memory + using ShapeAccum = MatrixShape; + + public: + // + // Data members + // + + /// Buffer for accumulator + AlignedBuffer accum; + + public: + + // + // Methods + // + + /// Returns a layout object for the Accum matrix + CUTLASS_DEVICE + static Layout LayoutAccum() { + return Layout::packed({ShapeAccum::kRow, ShapeAccum::kColumn}); + } + + /// Returns a TensorRef to the Accumulator + CUTLASS_HOST_DEVICE + TensorRefAccum accum_ref() { + return TensorRefAccum{accum.data(), LayoutAccum()}; + } + + }; + + using AccumulatorSharedStorage0 = AccumulatorSharedStorage< + Shape0, typename SmemAccumulatorIterator0::Element, + typename SmemAccumulatorIterator0::TensorLayout, + typename SmemAccumulatorIterator0::Padding>; + + struct B2bMmaSharedStorage { + typename Base::B2bMmaSharedStorage b2b_mma_shared_storage; + AccumulatorSharedStorage0 accumulator_shared_storage0; + }; + + +public: + + /// Construct from tensor references + CUTLASS_DEVICE + B2bMmaBaseSmemAccumulator( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + B2bMmaSharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage.b2b_mma_shared_storage, thread_idx, warp_idx, lane_idx) { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_multistage.h b/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_multistage.h index edb7fbe8..9ef21275 100644 --- a/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_multistage.h +++ b/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_multistage.h @@ -247,9 +247,9 @@ public: int lane_idx ): Base(shared_storage, thread_idx, warp_idx, lane_idx), - smem_iterator_A0_(shared_storage.sharedStorage0.operand_A_ref(), thread_idx), - smem_iterator_B0_(shared_storage.sharedStorage0.operand_B_ref(), thread_idx), - smem_iterator_B1_(shared_storage.sharedStorage1.operand_B_ref(), thread_idx) + smem_iterator_A0_(shared_storage.shared_storage0.operand_A_ref(), thread_idx), + smem_iterator_B0_(shared_storage.shared_storage0.operand_B_ref(), thread_idx), + smem_iterator_B1_(shared_storage.shared_storage1.operand_B_ref(), thread_idx) { // Compute warp location within threadblock tile by mapping the warp_id to // three coordinates: diff --git a/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_pipelined.h b/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_pipelined.h index 3e190136..c8c35092 100644 --- a/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_pipelined.h +++ b/examples/13_two_tensor_op_fusion/threadblock/b2b_mma_pipelined.h @@ -207,9 +207,9 @@ public: int lane_idx ///< ID of each thread within a warp ): Base(shared_storage, thread_idx, warp_idx, lane_idx), - smem_iterator_A_(shared_storage.sharedStorage0.operand_A_ref(), thread_idx), - smem_iterator_B0_(shared_storage.sharedStorage0.operand_B_ref(), thread_idx), - smem_iterator_B1_(shared_storage.sharedStorage1.operand_B_ref(), thread_idx) { + smem_iterator_A_(shared_storage.shared_storage0.operand_A_ref(), thread_idx), + smem_iterator_B0_(shared_storage.shared_storage0.operand_B_ref(), thread_idx), + smem_iterator_B1_(shared_storage.shared_storage1.operand_B_ref(), thread_idx) { // Compute warp location within threadblock tile by mapping the warp_id to three coordinates: diff --git a/examples/13_two_tensor_op_fusion/threadblock/default_b2b_mma.h b/examples/13_two_tensor_op_fusion/threadblock/default_b2b_mma.h index 6f49182e..7cd95065 100644 --- a/examples/13_two_tensor_op_fusion/threadblock/default_b2b_mma.h +++ b/examples/13_two_tensor_op_fusion/threadblock/default_b2b_mma.h @@ -166,7 +166,7 @@ struct DefaultB2bMma, - ElementB, LayoutB, 0, typename MmaCore1::IteratorThreadMapB>; + ElementB, LayoutB, 0, typename MmaCore1::IteratorThreadMapB, kAlignmentB>; // Define the threadblock-scoped pipelined matrix multiply using ThreadblockB2bMma = cutlass::gemm::threadblock::B2bMmaPipelined< diff --git a/examples/24_gemm_grouped/CMakeLists.txt b/examples/24_gemm_grouped/CMakeLists.txt new file mode 100644 index 00000000..fd5f4073 --- /dev/null +++ b/examples/24_gemm_grouped/CMakeLists.txt @@ -0,0 +1,28 @@ +# Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without modification, are permitted +# provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright notice, this list of +# conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright notice, this list of +# conditions and the following disclaimer in the documentation and/or other materials +# provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used +# to endorse or promote products derived from this software without specific prior written +# permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR +# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; +# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +cutlass_example_add_executable( + 24_gemm_grouped + gemm_grouped.cu + ) + diff --git a/examples/24_gemm_grouped/gemm_grouped.cu b/examples/24_gemm_grouped/gemm_grouped.cu new file mode 100644 index 00000000..8a518b1b --- /dev/null +++ b/examples/24_gemm_grouped/gemm_grouped.cu @@ -0,0 +1,1326 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief GEMM Grouped Example. + + This workload computes a batch of GEMM operations with distinct problem sizes. Pointers to matrices + in Global Memory are passed to the kernel in array (also held in Global Memory). Similarly, + leading dimensions and problem sizes are stored in arrays in GMEM. + + This differs from "Batched Array" GEMM because the size of each GEMM problem in the Grouped GEMM + concept may be distinct. + + This benchmark program initializes a workspace with random problem sizes for a given number of + groups. Command line options enable overriding M, N, and/or K dimensions with uniform values to + model problems more similar to the traditional batched GEMM. + + Additionally, problem sizes are collected and binned to compute the same problem as a series of + conventional batched GEMMs (setup for this problem is not timed). This demonstrates the performance + enhancement achieved by implementing a specialized grouped GEMM kernel. + + Examples: + + # Runs a grouped GEMM with 100 random problem sizes + $ ./examples/24_gemm_grouped/24_gemm_grouped --groups 100 + + # Runs a grouped GEMM with 100 random problem sizes (with GEMM-K dimension equal to 1024) + $ ./examples/24_gemm_grouped/24_gemm_grouped --groups 100 --k=1024 --verbose=true + + # Runs a grouped GEMM that is equivalent to a batched GEMM + $ ./examples/24_gemm_grouped/24_gemm_grouped --groups 100 --m=2048 --n=1024 --k=1024 --verbose=true + + # Execute Grouped GEMM and profile with NSight + $ nv-nsight-cu-cli ./examples/24_gemm_grouped/24_gemm_grouped --m=256 --n=256 --k=256 --verbose=true \ + --iterations=1 --reference-check=false + +*/ + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include +#include +#include +#include +#include + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/kernel/gemm_grouped.h" +#include "cutlass/gemm/kernel/default_gemm_grouped.h" +#include "cutlass/gemm/device/gemm_grouped.h" +#include "cutlass/gemm/device/gemm_universal.h" + +#include "cutlass/util/command_line.h" +#include "cutlass/util/distribution.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/reference/host/gemm_complex.h" +#include "cutlass/util/reference/device/gemm_complex.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/device/tensor_fill.h" +#include "cutlass/util/reference/host/tensor_norm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Result structure +struct Result { + + double runtime_ms; + double gflops; + cutlass::Status status; + cudaError_t error; + bool passed; + + // + // Methods + // + + Result( + double runtime_ms = 0, + double gflops = 0, + cutlass::Status status = cutlass::Status::kSuccess, + cudaError_t error = cudaSuccess + ): + runtime_ms(runtime_ms), gflops(gflops), status(status), error(error), passed(true) { } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Hash function for cutlass::gemm::GemmCoord +struct HashGemmCoord { + size_t operator()(cutlass::gemm::GemmCoord const &problem) const { + std::hash hasher; + return (hasher(problem.m() * 3)) ^ (hasher(1 + problem.n() * 5)) ^ (hasher(2 + problem.k() * 7)); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Command line options parsing +struct Options { + + bool help; + bool reference_check; + + std::vector problem_sizes; + std::unordered_map, HashGemmCoord> problem_bins; + + int problem_count; + int iterations; + int cuda_streams; + bool verbose; + float alpha; + float beta; + + // + // Methods + // + + Options(): + help(false), + reference_check(true), + problem_count(15), + iterations(20), + cuda_streams(0), + verbose(false), + alpha(1), + beta() + { } + + // 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; + } + + cmd.get_cmd_line_argument("groups", problem_count, 15); + cmd.get_cmd_line_argument("alpha", alpha, 1.0f); + cmd.get_cmd_line_argument("beta", beta, 0.0f); + cmd.get_cmd_line_argument("iterations", iterations, 20); + cmd.get_cmd_line_argument("streams", cuda_streams, 0); + cmd.get_cmd_line_argument("verbose", verbose, false); + cmd.get_cmd_line_argument("reference-check", reference_check, true); + + int cmd_line_m = -1; + int cmd_line_n = -1; + int cmd_line_k = -1; + + cmd.get_cmd_line_argument("m", cmd_line_m); + cmd.get_cmd_line_argument("n", cmd_line_n); + cmd.get_cmd_line_argument("k", cmd_line_k); + + // + // For now, randomly choose the problem sizes. + // + + problem_sizes.reserve(problem_count); + + for (int i = 0; i < problem_count; ++i) { + + int m = cmd_line_m; + int n = cmd_line_n; + int k = cmd_line_k; + + if (m < 1) { + m = 8 * (rand() % 256) + 8; + } + + if (n < 1) { + n = 8 * (rand() % 256) + 8; + } + + if (k < 1) { + k = 8 * (rand() % 256) + 8; + } + + cutlass::gemm::GemmCoord problem(m, n, k); + + problem_sizes.push_back(problem); + } + + // + // Insert the problem sizes into a sorted container class. This is *NOT* necessary + // to run the CUTLASS kernel, but it enables the execution of cublas's batched GEMM. + // + for (int i = 0; i < problem_count; ++i) { + auto it = problem_bins.find(problem_sizes.at(i)); + if (it == problem_bins.end()) { + problem_bins.insert({problem_sizes.at(i), std::vector({i}) }); + } + else { + it->second.push_back(i); + } + } + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "24_gemm_grouped\n\n" + << " This example profiles the performance of a 'grouped' GEMM kernel. This is similar to batched GEMM\n" + << " in that multiple, independent GEMMs are computed by one grid launch. It differs in that each\n" + << " 'group' may compute a unique problem size. Problem sizes and pointers to matrices are both stored\n" + << " in device Global Memory and loaded by the kernel.\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement.\n\n" + << " --groups= Number of individual GEMM problems (default: --groups=15)\n" + << " --m= Sets the M dimension for all groups. Otherwise, it is selected randomly\n" + << " --n= Sets the N dimension for all groups. Otherwise, it is selected randomly\n" + << " --k= Sets the K dimension for all groups. Otherwise, it is selected randomly\n" + << " --alpha= Epilogue scalar alpha (real part)\n" + << " --beta= Epilogue scalar beta (real part)\n\n" + << " --iterations= Number of profiling iterations to perform.\n" + << " --reference-check= If true, performs reference check.\n" + << " --verbose= If true, prints problem sizes and batching structure.\n"; + + out << "\n\nExamples:\n\n" + + << "# Runs a grouped GEMM with 100 random problem sizes\n" + << "$ ./examples/24_gemm_grouped/24_gemm_grouped --groups 100\n\n" + + << "# Runs a grouped GEMM with 100 random problem sizes (with GEMM-K dimension equal to 1024)\n" + << "$ ./examples/24_gemm_grouped/24_gemm_grouped --groups 100 --k=1024 --verbose=true\n\n" + + << "# Runs a grouped GEMM that is equivalent to a batched GEMM\n" + << "$ ./examples/24_gemm_grouped/24_gemm_grouped --groups 100 --m=2048 --n=1024 --k=1024 --verbose=true\n\n" + + << "# Execute Grouped GEMM and profile with NSight\n" + << "$ nv-nsight-cu-cli ./examples/24_gemm_grouped/24_gemm_grouped --m=256 --n=256 --k=256 --verbose=true --iterations=1 --reference-check=false\n\n"; + + return out; + } + + /// Compute performance in GFLOP/s + double gflops(double runtime_s) const { + + // Number of real-valued multiply-adds + int64_t fmas = int64_t(); + + for (auto const & problem : problem_sizes) { + fmas += problem.product(); + } + + // Two flops per multiply-add + return 2.0 * double(fmas) / double(1.0e9) / runtime_s; + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +template +class TestbedGrouped { +public: + + // + // Type definitions + // + + using ElementA = typename Gemm::ElementA; + using ElementB = typename Gemm::ElementB; + using ElementC = typename Gemm::ElementC; + using ElementAccumulator = typename Gemm::ElementAccumulator; + + using EpilogueOutputOp = typename Gemm::GemmKernel::Epilogue::OutputOp; + using ElementCompute = typename EpilogueOutputOp::ElementCompute; + + using LayoutA = typename Gemm::LayoutA; + using LayoutB = typename Gemm::LayoutB; + using LayoutC = typename Gemm::LayoutC; + + using MatrixCoord = typename LayoutC::TensorCoord; + +private: + + // + // Data members + // + + Options options; + + /// Initialization + cutlass::Distribution::Kind init_A; + cutlass::Distribution::Kind init_B; + cutlass::Distribution::Kind init_C; + uint32_t seed; + + cutlass::DeviceAllocation problem_sizes_device; + + std::vector offset_A; + std::vector offset_B; + std::vector offset_C; + std::vector offset_D; + + std::vector lda_host; + std::vector ldb_host; + std::vector ldc_host; + std::vector ldd_host; + + cutlass::DeviceAllocation lda; + cutlass::DeviceAllocation ldb; + cutlass::DeviceAllocation ldc; + cutlass::DeviceAllocation ldd; + + cutlass::DeviceAllocation block_A; + cutlass::DeviceAllocation block_B; + cutlass::DeviceAllocation block_C; + cutlass::DeviceAllocation block_D; + + cutlass::DeviceAllocation ptr_A; + cutlass::DeviceAllocation ptr_B; + cutlass::DeviceAllocation ptr_C; + cutlass::DeviceAllocation ptr_D; + +public: + + // + // Methods + // + + TestbedGrouped( + Options const &options_, + cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_C_ = cutlass::Distribution::Uniform, + uint32_t seed_ = 3080 + ): + options(options_), init_A(init_A_), init_B(init_B_), init_C(init_C_), seed(seed_) { } + + int problem_count() const { + return options.problem_count; + } + +private: + + /// Helper to initialize a tensor view + template + void initialize_tensor_( + Element *ptr, + size_t capacity, + cutlass::Distribution::Kind dist_kind, + uint32_t seed) { + + if (dist_kind == cutlass::Distribution::Uniform) { + + Element 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) { + if (cutlass::sizeof_bits::value <= 16) { + scope_max = 5; + scope_min = -5; + } + else { + scope_max = 8; + scope_min = -8; + } + } else { + scope_max = 8; + scope_min = -8; + } + + cutlass::reference::device::BlockFillRandomUniform( + ptr, capacity, seed, scope_max, scope_min, 0); + } + else if (dist_kind == cutlass::Distribution::Gaussian) { + + cutlass::reference::device::BlockFillRandomGaussian( + ptr, capacity, seed, Element(), Element(0.5f)); + } + else if (dist_kind == cutlass::Distribution::Sequential) { + + // Fill with increasing elements + cutlass::reference::device::BlockFillSequential( + ptr, capacity, Element(1), Element()); + } + else { + + // Fill with all 1s + cutlass::reference::device::BlockFillSequential( + ptr, capacity, Element(), Element(1)); + } + } + + /// Verbose printing of problem sizes + void print_problem_sizes_() { + + // Print groups + std::cout << problem_count() << " groups:\n"; + + int32_t idx = 0; + int64_t total_tiles = 0; + + for (auto const & problem : options.problem_sizes) { + + int tiles = + ((problem.m() + Gemm::ThreadblockShape::kM - 1) / Gemm::ThreadblockShape::kM) * + ((problem.n() + Gemm::ThreadblockShape::kN - 1) / Gemm::ThreadblockShape::kN); + + total_tiles += tiles; + + std::cout << " [" << idx << "]: " + << problem.m() << "-by-" << problem.n() << "-by-" << problem.k() + << " (" << tiles << " threadblock tiles)" << "\n"; + + ++idx; + } + + // Print batched GEMM equivalent + size_t bin_idx = 0; + size_t problem_count_check = 0; + std::cout << "\nConventionally executed as " << options.problem_bins.size() << " batched GEMMs:\n"; + for (auto const & bin : options.problem_bins) { + + std::cout << " [" << bin_idx << "]: " + << bin.first.m() << "-by-" << bin.first.n() << "-by-" << bin.first.k() + << ", batch count: " << bin.second.size() << "\n"; + + ++bin_idx; + problem_count_check += bin.second.size(); + } + + if (problem_count_check != problem_count()) { + std::cout << "\n***\nERROR in BINNING LOGIC!\n***\n" << std::endl; + } + } + + /// Initializes data structures + void initialize_() { + + // + // Choose random problem sizes + // + + // construct a few problems of random sizes + srand(seed); + + int64_t total_elements_A = 0; + int64_t total_elements_B = 0; + int64_t total_elements_C = 0; + int64_t total_elements_D = 0; + + + lda_host.resize(problem_count()); + ldb_host.resize(problem_count()); + ldc_host.resize(problem_count()); + ldd_host.resize(problem_count()); + + for (int32_t i = 0; i < problem_count(); ++i) { + + auto problem = options.problem_sizes.at(i); + + lda_host.at(i) = LayoutA::packed({problem.m(), problem.k()}).stride(0); + ldb_host.at(i) = LayoutB::packed({problem.k(), problem.n()}).stride(0); + ldc_host.at(i) = LayoutC::packed({problem.m(), problem.n()}).stride(0); + ldd_host.at(i) = LayoutC::packed({problem.m(), problem.n()}).stride(0); + + offset_A.push_back(total_elements_A); + offset_B.push_back(total_elements_B); + offset_C.push_back(total_elements_C); + offset_D.push_back(total_elements_D); + + int64_t elements_A = problem.m() * problem.k(); + int64_t elements_B = problem.k() * problem.n(); + int64_t elements_C = problem.m() * problem.n(); + int64_t elements_D = problem.m() * problem.n(); + + total_elements_A += elements_A; + total_elements_B += elements_B; + total_elements_C += elements_C; + total_elements_D += elements_D; + } + + problem_sizes_device.reset(problem_count()); + problem_sizes_device.copy_from_host(options.problem_sizes.data()); + + lda.reset(problem_count()); + ldb.reset(problem_count()); + ldc.reset(problem_count()); + ldd.reset(problem_count()); + + lda.copy_from_host(lda_host.data()); + ldb.copy_from_host(ldb_host.data()); + ldc.copy_from_host(ldc_host.data()); + ldd.copy_from_host(ldd_host.data()); + + // + // Assign pointers + // + + block_A.reset(total_elements_A); + block_B.reset(total_elements_B); + block_C.reset(total_elements_C); + block_D.reset(total_elements_D); + + std::vector ptr_A_host(problem_count()); + std::vector ptr_B_host(problem_count()); + std::vector ptr_C_host(problem_count()); + std::vector ptr_D_host(problem_count()); + + for (int32_t i = 0; i < problem_count(); ++i) { + ptr_A_host.at(i) = block_A.get() + offset_A.at(i); + ptr_B_host.at(i) = block_B.get() + offset_B.at(i); + ptr_C_host.at(i) = block_C.get() + offset_C.at(i); + ptr_D_host.at(i) = block_D.get() + offset_D.at(i); + } + + ptr_A.reset(problem_count()); + ptr_A.copy_from_host(ptr_A_host.data()); + + ptr_B.reset(problem_count()); + ptr_B.copy_from_host(ptr_B_host.data()); + + ptr_C.reset(problem_count()); + ptr_C.copy_from_host(ptr_C_host.data()); + + ptr_D.reset(problem_count()); + ptr_D.copy_from_host(ptr_D_host.data()); + + // + // Initialize the problems of the workspace + // + + initialize_tensor_(block_A.get(), total_elements_A, init_A, seed * 2021); + initialize_tensor_(block_B.get(), total_elements_B, init_B, seed * 2022); + initialize_tensor_(block_C.get(), total_elements_C, init_C, seed * 2023); + + cutlass::reference::device::BlockFillSequential( + block_D.get(), total_elements_D, ElementC(), ElementC()); + } + + /// Verifies the result is a GEMM + bool verify_() { + + bool passed = true; + + for (int32_t i = 0; i < problem_count(); ++i) { + cutlass::gemm::GemmCoord problem = options.problem_sizes.at(i); + + LayoutA layout_A(lda_host.at(i)); + LayoutB layout_B(ldb_host.at(i)); + LayoutC layout_C(ldc_host.at(i)); + LayoutC layout_D(ldd_host.at(i)); + + MatrixCoord extent_A{problem.m(), problem.k()}; + MatrixCoord extent_B{problem.k(), problem.n()}; + MatrixCoord extent_C{problem.m(), problem.n()}; + + cutlass::TensorView view_A(block_A.get() + offset_A.at(i), layout_A, extent_A); + cutlass::TensorView view_B(block_B.get() + offset_B.at(i), layout_B, extent_B); + cutlass::TensorView view_C(block_C.get() + offset_C.at(i), layout_C, extent_C); + + cutlass::DeviceAllocation block_Ref(layout_D.capacity(extent_C)); + cutlass::TensorView view_Ref_device(block_Ref.get(), layout_D, extent_C); + + // Reference GEMM + cutlass::reference::device::GemmComplex< + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + ElementCompute, ElementAccumulator + >( + problem, + options.alpha, + view_A, + Gemm::kTransformA, + view_B, + Gemm::kTransformB, + options.beta, + view_C, + view_Ref_device, + ElementAccumulator(0) + ); + + // Copy to host memory + std::vector matrix_D(layout_D.capacity(extent_C)); + std::vector matrix_Ref(layout_D.capacity(extent_C)); + + cutlass::device_memory::copy_to_host(matrix_D.data(), block_D.get() + offset_D.at(i), matrix_D.size()); + cutlass::device_memory::copy_to_host(matrix_Ref.data(), block_Ref.get(), matrix_D.size()); + + cutlass::TensorView view_D( matrix_D.data(), layout_D, extent_C); + cutlass::TensorView view_Ref(matrix_Ref.data(), layout_D, extent_C); + + // Reference check + passed = cutlass::reference::host::TensorEquals(view_D, view_Ref); + + if (!passed) { + std::cerr << "\n***\nError - problem " << i << " failed the QA check\n***\n" << std::endl; + return passed; + } + } + + return passed; + } + +public: + + /// Returns the number of threadblocks to launch if the kernel can run on the target + /// device. Otherwise, returns zero. + int sufficient() const { + // + // Determine SMEM requirements and waive if not satisfied + // + + int smem_size = int(sizeof(typename Gemm::GemmKernel::SharedStorage)); + + cudaDeviceProp properties; + int device_idx; + cudaError_t result = cudaGetDevice(&device_idx); + + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDevice() API call failed."); + } + + result = cudaGetDeviceProperties(&properties, device_idx); + + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDeviceProperties() failed"); + } + + int occupancy = std::min(2, int(properties.sharedMemPerMultiprocessor / smem_size)); + + return properties.multiProcessorCount * occupancy; + } + + + /// Executes a Grouped GEMM kernel and measures runtime. + Result profile_grouped() { + + Result result; + + int threadblock_count = sufficient(); + + // Early exit + if (!threadblock_count) { + std::cout << "Active CUDA device lacks hardware resources to run CUTLASS Grouped GEMM kernel." << std::endl; + return result; + } + + if (options.verbose) { + print_problem_sizes_(); + } + + result.passed = false; + + // Initialize the problem + initialize_(); + + // Configure the GEMM arguments + typename EpilogueOutputOp::Params epilogue_op(options.alpha, options.beta); + + // Configure GEMM arguments + typename Gemm::Arguments args( + problem_sizes_device.get(), + problem_count(), + threadblock_count, + epilogue_op, + ptr_A.get(), + ptr_B.get(), + ptr_C.get(), + ptr_D.get(), + lda.get(), + ldb.get(), + ldc.get(), + ldd.get() + ); + + // Initialize the GEMM object + Gemm gemm; + + result.status = gemm.initialize(args); + + if (result.status != cutlass::Status::kSuccess) { + std::cerr << "Failed to initialize CUTLASS Grouped GEMM kernel." << std::endl; + return result; + } + + // Run the grouped GEMM object + result.status = gemm.run(); + + if (result.status != cutlass::Status::kSuccess) { + std::cerr << "Failed to run CUTLASS Grouped GEMM kernel." << std::endl; + return result; + } + + // Wait for completion + result.error = cudaDeviceSynchronize(); + + if (result.error != cudaSuccess) { + std::cerr << "Kernel execution error: " << cudaGetErrorString(result.error); + return result; + } + + // + // Verify correctness + // + result.passed = true; + + if (options.reference_check) { + result.passed = verify_(); + } + + // + // Warm-up run of the grouped GEMM object + // + result.status = gemm.run(); + + if (result.status != cutlass::Status::kSuccess) { + std::cerr << "Failed to run CUTLASS Grouped GEMM kernel." << std::endl; + return result; + } + + // + // Construct events + // + + cudaEvent_t events[2]; + + for (auto & event : events) { + result.error = cudaEventCreate(&event); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventCreate() failed: " << cudaGetErrorString(result.error) << std::endl; + return -1; + } + } + + // Record an event at the start of a series of GEMM operations + result.error = cudaEventRecord(events[0]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // + // Run profiling loop + // + + for (int iter = 0; iter < options.iterations; ++iter) { + gemm(); + } + + // + // Stop profiling loop + // + + // Record an event when the GEMM operations have been launched. + result.error = cudaEventRecord(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Wait for work on the device to complete. + result.error = cudaEventSynchronize(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventSynchronize() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Measure elapsed runtime + float runtime_ms = 0; + result.error = cudaEventElapsedTime(&runtime_ms, events[0], events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Compute average runtime and GFLOPs. + result.runtime_ms = double(runtime_ms) / double(options.iterations); + result.gflops = options.gflops(result.runtime_ms / 1000.0); + + // + // Cleanup + // + + for (auto event : events) { + (void)cudaEventDestroy(event); + } + + int32_t idx = 0; + int64_t total_tiles = 0; + + for (auto const & problem : options.problem_sizes) { + + int tiles = + ((problem.m() + Gemm::ThreadblockShape::kM - 1) / Gemm::ThreadblockShape::kM) * + ((problem.n() + Gemm::ThreadblockShape::kN - 1) / Gemm::ThreadblockShape::kN); + + total_tiles += tiles; + ++idx; + } + + std::cout << std::endl; + std::cout << "Grouped GEMM (CUTLASS):\n" + << "====================================================" << std::endl; + + std::cout << " " << total_tiles << " total threadblock tiles." << std::endl; + + std::cout << std::endl; + std::cout << " " << "Grouped Runtime: " << result.runtime_ms << " ms" << std::endl; + std::cout << " " << "Grouped GFLOPs: " << result.gflops << std::endl; + + return result; + } + + /// Executes a conventional batched GEMM kernel. + Result profile_batched() { + + Result result; + result.passed = false; + + // + // Prepare batched GEMM environment + // + + int32_t effective_streams = (options.cuda_streams ? options.cuda_streams : 1); + + // Array of leading dimensions used by batched GEMM calls + std::vector bin_problem_sizes; + std::vector bin_count; + std::vector bin_ldm_A; + std::vector bin_ldm_B; + std::vector bin_ldm_C; + std::vector bin_start; + + std::vector ptr_A_batched_host; + std::vector ptr_B_batched_host; + std::vector ptr_C_batched_host; + + for (auto const & bin : options.problem_bins) { + int first_idx = bin.second.front(); + + bin_problem_sizes.push_back(options.problem_sizes.at(first_idx)); + bin_count.push_back(int32_t(bin.second.size())); + + bin_ldm_A.push_back(static_cast(lda_host.at(first_idx))); + bin_ldm_B.push_back(static_cast(ldb_host.at(first_idx))); + bin_ldm_C.push_back(static_cast(ldc_host.at(first_idx))); + + if (ptr_A_batched_host.size() % 2) { + ptr_A_batched_host.push_back(nullptr); + ptr_B_batched_host.push_back(nullptr); + ptr_C_batched_host.push_back(nullptr); + } + + bin_start.push_back(int32_t(ptr_A_batched_host.size())); + + for (int idx : bin.second) { + + if (bin_problem_sizes.back() != options.problem_sizes.at(idx)) { + std::cerr << "Error - failed to group problems.\n"; + return result; + } + + if (bin_ldm_A.back() != lda_host.at(idx)) { + std::cerr << "Error - failed to group problems.\n"; + return result; + } + + if (bin_ldm_B.back() != ldb_host.at(idx)) { + std::cerr << "Error - failed to group problems.\n"; + return result; + } + + if (bin_ldm_C.back() != ldc_host.at(idx)) { + std::cerr << "Error - failed to group problems.\n"; + return result; + } + + ptr_A_batched_host.push_back(block_A.get() + offset_A.at(idx)); + ptr_B_batched_host.push_back(block_B.get() + offset_B.at(idx)); + ptr_C_batched_host.push_back(block_D.get() + offset_C.at(idx)); + } + } + + // Array of GMEM pointers used by batched array GEMM calls + cutlass::DeviceAllocation ptr_A_batched; + cutlass::DeviceAllocation ptr_B_batched; + cutlass::DeviceAllocation ptr_C_batched; + + ptr_A_batched.reset(ptr_A_batched_host.size()); + ptr_B_batched.reset(ptr_A_batched_host.size()); + ptr_C_batched.reset(ptr_A_batched_host.size()); + + ptr_A_batched.copy_from_host(ptr_A_batched_host.data()); + ptr_B_batched.copy_from_host(ptr_B_batched_host.data()); + ptr_C_batched.copy_from_host(ptr_C_batched_host.data()); + + // + // Create CUDA streams to maximize concurrency of batched-array GEMM kernels + // + std::vector cuda_streams; + + // + // Warmup run + // + + + if (options.cuda_streams) { + for (int i = 0; i < options.cuda_streams; ++i) { + cudaStream_t stream; + + result.error = cudaStreamCreate(&stream); + if (result.error != cudaSuccess) { + std::cerr << "Failed to create CUDA stream." << std::endl; + return result; + } + cuda_streams.push_back(stream); + + } + } + else { + cuda_streams.push_back(nullptr); + + } + + // Use 'D' for the in/out workspace + block_D.copy_from_device(block_C.get()); + + for (int bin_idx = 0; bin_idx < int32_t(bin_problem_sizes.size()); ++bin_idx) { + + cutlass::gemm::GemmCoord const & problem = bin_problem_sizes[bin_idx]; + int32_t batch_count = bin_count[bin_idx]; + int32_t bin_start_idx = bin_start[bin_idx]; + int32_t lda = bin_ldm_A[bin_idx]; + int32_t ldb = bin_ldm_B[bin_idx]; + int32_t ldc = bin_ldm_C[bin_idx]; + + void const ** ptr_A_array = ptr_A_batched.get() + bin_start[bin_idx]; + void const ** ptr_B_array = ptr_B_batched.get() + bin_start[bin_idx]; + void ** ptr_C_array = ptr_C_batched.get() + bin_start[bin_idx]; + + // + // Initialize the CUTLASS GEMM operator + // + + // Configure the GEMM arguments + typename EpilogueOutputOp::Params epilogue_op(options.alpha, options.beta); + + typename GemmBatched::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kArray, + problem, + batch_count, + epilogue_op, + (void const *)ptr_A_array, + (void const *)ptr_B_array, + (void const *)ptr_C_array, + (void *)ptr_C_array, + int64_t(), + int64_t(), + int64_t(), + int64_t(), + int64_t(lda), + int64_t(ldb), + int64_t(ldc), + int64_t(ldc) + }; + + GemmBatched gemm_op; + + cutlass::Status status = gemm_op.initialize(arguments); + + if (status != cutlass::Status::kSuccess) { + std::cerr << "CUTLASS error on line " << __LINE__ << std::endl; + return result; + } + + status = gemm_op(); + + if (status != cutlass::Status::kSuccess) { + std::cerr << "CUTLASS error on line " << __LINE__ << std::endl; + return result; + } + + } + + // + // Wait for completion + // + + result.error = cudaDeviceSynchronize(); + + if (result.error != cudaSuccess) { + std::cerr << "Kernel execution error: " << cudaGetErrorString(result.error); + return result; + } + + // + // Construct events + // + + cudaEvent_t events[2]; + + for (auto & event : events) { + result.error = cudaEventCreate(&event); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventCreate() failed: " << cudaGetErrorString(result.error) << std::endl; + return -1; + } + } + + // + // Wait for completion + // + + result.error = cudaDeviceSynchronize(); + + if (result.error != cudaSuccess) { + std::cerr << "Kernel execution error: " << cudaGetErrorString(result.error); + return result; + } + + // Record an event at the start of a series of GEMM operations + result.error = cudaEventRecord(events[0]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // + // Run profiling loop + // + + int last_stream_idx = 0; + + for (int iter = 0; iter < options.iterations; ++iter) { + + for (int bin_idx = 0; bin_idx < int32_t(bin_problem_sizes.size()); ++bin_idx) { + + cutlass::gemm::GemmCoord const & problem = bin_problem_sizes[bin_idx]; + int32_t batch_count = bin_count[bin_idx]; + int32_t bin_start_idx = bin_start[bin_idx]; + int32_t lda = bin_ldm_A[bin_idx]; + int32_t ldb = bin_ldm_B[bin_idx]; + int32_t ldc = bin_ldm_C[bin_idx]; + + void const ** ptr_A_array = ptr_A_batched.get() + bin_start[bin_idx]; + void const ** ptr_B_array = ptr_B_batched.get() + bin_start[bin_idx]; + void ** ptr_C_array = ptr_C_batched.get() + bin_start[bin_idx]; + + last_stream_idx = (bin_idx % effective_streams); + + // + // Initialize the CUTLASS GEMM operator + // + + // Configure the GEMM arguments + typename EpilogueOutputOp::Params epilogue_op(options.alpha, options.beta); + + typename GemmBatched::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kArray, + problem, + batch_count, + epilogue_op, + (void const *)ptr_A_array, + (void const *)ptr_B_array, + (void const *)ptr_C_array, + (void *)ptr_C_array, + int64_t(), + int64_t(), + int64_t(), + int64_t(), + int64_t(lda), + int64_t(ldb), + int64_t(ldc), + int64_t(ldc) + }; + + GemmBatched gemm_op; + + cutlass::Status status = gemm_op.initialize(arguments); + + if (status != cutlass::Status::kSuccess) { + std::cerr << "CUTLASS error on line " << __LINE__ << std::endl; + return result; + } + + status = gemm_op(cuda_streams[last_stream_idx]); + + if (status != cutlass::Status::kSuccess) { + std::cerr << "CUTLASS error on line " << __LINE__ << std::endl; + return result; + } + + } + } + + // + // Stop profiling loop + // + + // Record an event when the GEMM operations have been launched. + result.error = cudaEventRecord(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // + // Wait for work to be completed + // + + result.error = cudaDeviceSynchronize(); + + if (result.error != cudaSuccess) { + std::cerr << "Kernel execution error: " << cudaGetErrorString(result.error); + return result; + } + + // Wait for work on the device to complete. + result.error = cudaEventSynchronize(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventSynchronize() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Wait for work on the device to complete. + result.error = cudaEventSynchronize(events[0]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventSynchronize() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Measure elapsed runtime + float runtime_ms = 0; + result.error = cudaEventElapsedTime(&runtime_ms, events[0], events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Compute average runtime and GFLOPs. + result.runtime_ms = double(runtime_ms) / double(options.iterations); + result.gflops = options.gflops(result.runtime_ms / 1000.0); + + // + // Cleanup + // + + for (auto event : events) { + (void)cudaEventDestroy(event); + } + + for (auto stream : cuda_streams) { + if (stream) { + (void)cudaStreamDestroy(stream); + } + } + + std::cout << std::endl; + std::cout << "Batched GEMM:\n" + << "====================================================" << std::endl; + + std::cout << " " << bin_problem_sizes.size() << " batched GEMMs launched" << std::endl; + std::cout << std::endl; + std::cout << " " << "Batched Runtime: " << result.runtime_ms << " ms" << std::endl; + std::cout << " " << "Batched GFLOPs: " << result.gflops << std::endl; + + result.passed = true; + return result; + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +int main(int argc, char const **args) { + + // + // This example uses mma.sync to directly access Tensor Cores to achieve peak performance. + // + + 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__ < 11 || props.major < 8) { + + // + // This example requires an NVIDIA Ampere-architecture GPU. + // + + std::cout + << "CUTLASS's Grouped GEMM example requires a GPU of NVIDIA's Ampere Architecture or " + << "later (compute capability 80 or greater).\n"; + + return 0; + } + + // + // Parse options + // + + Options options; + + options.parse(argc, args); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + // + // Define the Grouped GEMM type + // + + using ElementOutput = cutlass::half_t; + using ElementAccumulator = float; + + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + cutlass::half_t, + cutlass::layout::ColumnMajor, + cutlass::ComplexTransform::kNone, + 8, + cutlass::half_t, + cutlass::layout::ColumnMajor, + cutlass::ComplexTransform::kNone, + 8, + ElementOutput, cutlass::layout::ColumnMajor, + ElementAccumulator, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<128, 128, 32>, + cutlass::gemm::GemmShape<64, 64, 32>, + cutlass::gemm::GemmShape<16, 8, 16>, + cutlass::epilogue::thread::LinearCombination< + ElementOutput, 128 / cutlass::sizeof_bits::value, + ElementAccumulator, ElementAccumulator>, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 4>::GemmKernel; + + using GemmGrouped = cutlass::gemm::device::GemmGrouped; + + // + // Define a conventional batched GEMM type + // + + // Gemm operator cutlass_tensorop_f16_s16816gemm_f16_128x128_32x4_nt_align8 + using GemmBatched = cutlass::gemm::device::GemmUniversal< + ElementOutput, cutlass::layout::ColumnMajor, + ElementOutput, cutlass::layout::ColumnMajor, + ElementOutput, cutlass::layout::ColumnMajor, + ElementAccumulator, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<128, 128, 32>, + cutlass::gemm::GemmShape<64, 64, 32>, + cutlass::gemm::GemmShape<16, 8, 16>, + cutlass::epilogue::thread::LinearCombination< + ElementOutput, + 128 / cutlass::sizeof_bits::value, + ElementAccumulator, + ElementAccumulator + >, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>, + 4 + >; + + // + // Profile it + // + + TestbedGrouped testbed(options); + + if (!testbed.sufficient()) { + std::cout << "The active CUDA device lacks sufficient hardware resources to execute this kernel.\n"; + return 0; + } + + Result result = testbed.profile_grouped(); + if (!result.passed) { + std::cout << "Profiling CUTLASS grouped GEMM has failed.\n"; + std::cout << "\nFailed\n"; + return -1; + } + + result = testbed.profile_batched(); + if (!result.passed) { + + std::cout << "Profiling batched GEMM has failed.\n"; + std::cout << "\nFailed\n"; + return -1; + } + + std::cout << "\nPassed\n"; + + return 0; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/examples/25_ampere_fprop_mainloop_fusion/CMakeLists.txt b/examples/25_ampere_fprop_mainloop_fusion/CMakeLists.txt new file mode 100644 index 00000000..6334a08b --- /dev/null +++ b/examples/25_ampere_fprop_mainloop_fusion/CMakeLists.txt @@ -0,0 +1,28 @@ +# Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without modification, are permitted +# provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright notice, this list of +# conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright notice, this list of +# conditions and the following disclaimer in the documentation and/or other materials +# provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used +# to endorse or promote products derived from this software without specific prior written +# permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR +# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; +# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +cutlass_example_add_executable( + 25_ampere_fprop_mainloop_fusion + ampere_fprop_mainloop_fusion.cu + ) + diff --git a/examples/25_ampere_fprop_mainloop_fusion/ampere_fprop_mainloop_fusion.cu b/examples/25_ampere_fprop_mainloop_fusion/ampere_fprop_mainloop_fusion.cu new file mode 100644 index 00000000..aba02625 --- /dev/null +++ b/examples/25_ampere_fprop_mainloop_fusion/ampere_fprop_mainloop_fusion.cu @@ -0,0 +1,751 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/** + +This example shows how to fuse per channel scale+bias+relu of the activations +into the fprop mainloop. + +Compared with original fprop kernel, this example has two more vectors, one for +the scale and one for the bias. The length of the vectors are the same as the +activation channel number. This kernels loads the vectors when the associated +activation channels are loaded in the mainloop. Between reading the +activations and scale/bias data from the shared memory and calling tensor core +instructions, scale+bias+relu is computed in the register file. + +This example is customized for Ampere 16816 fp16 tensor core instruction. +Changing to different data types or different tensor core instruction require +source code changing. See +include/cutlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h for more +technical details. + +This example is modified based on 16_ampere_tensorop_conv2dfprop. The command +line is the same. +*/ + +#include +#include + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/device/gemm.h" +#include "cutlass/conv/kernel/default_conv2d_fprop_fusion.h" +#include "cutlass/conv/device/implicit_gemm_convolution_fusion.h" + +#include "cutlass/util/command_line.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/reference/device/gemm.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/device/convolution.h" +#include "cutlass/util/tensor_view_io.h" + +#include "helper.h" + +// The code section below describes datatype for input, output tensors and computation between +// elements +using ElementAccumulator = float; // Data type of accumulator +using ElementComputeEpilogue = float; // Data type of epilogue computation (alpha, beta) +using ElementInputA = cutlass::half_t; // Data type of elements in input tensor +using ElementInputB = cutlass::half_t; // Data type of elements in input tensor +using ElementInputScaleBias = cutlass::half_t; // Data type of elements in input sclae and bias vectors +using ElementOutput = float; // Data type of elements in output tensor + +using LayoutInputA = cutlass::layout::TensorNHWC; +using LayoutInputB = cutlass::layout::TensorNHWC; +using LayoutInputScaleBias = cutlass::layout::RowMajor; +using LayoutOutput = cutlass::layout::TensorNHWC; + +// This code section describes whether you want to use tensor cores or regular SIMT cores on GPU SM +using MMAOp = cutlass::arch::OpClassTensorOp; + +// This code section describes CUDA SM architecture number +using SmArch = cutlass::arch::Sm80; + +// This code section describes the tile size a thread block will compute +using ThreadblockShape = cutlass::gemm::GemmShape<128, 128, 32>; // Threadblock tile shape + +// This code section describes tile size a warp will compute +using WarpShape = cutlass::gemm::GemmShape<64, 64, 32>; // Warp tile shape + +// This code section describes the size of MMA op +using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>; // TensorCore instruction shape + +// This code section describes how threadblocks are scheduled on GPU +using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; + +// Number of pipelines you want to use +constexpr int NumStages = 4; + +// This code section describe iterator algorithm selected is Analytic or Optimized +static cutlass::conv::IteratorAlgorithm const IteratorAlgorithm = cutlass::conv::IteratorAlgorithm::kOptimized; + +// This code section describes the epilogue part of the kernel, we use default value +using EpilogueOp = cutlass::epilogue::thread::LinearCombination< + ElementOutput, // Data type of output matrix. + 128 / cutlass::sizeof_bits::value, // The number of elements per vectorized. + // memory access. This becomes the vector width of + // math instructions in the epilogue too. + ElementAccumulator, // Data type of accumulator + ElementComputeEpilogue>; // Data type for alpha/beta in linear combination + +using Conv2dFpropFusionKernel = typename cutlass::conv::kernel::DefaultConv2dFpropFusion< + ElementInputA, LayoutInputA, + ElementInputB, LayoutInputB, + ElementInputScaleBias, LayoutInputScaleBias, + ElementOutput, LayoutOutput, + ElementAccumulator, + MMAOp, + SmArch, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOp, + SwizzleThreadBlock, + NumStages, + cutlass::arch::OpMultiplyAdd, + IteratorAlgorithm +>::Kernel; + +using ImplicitGemmFusion = cutlass::conv::device::ImplicitGemmConvolutionFusion; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Command line options parsing +struct Options { + + bool help; + cutlass::Tensor4DCoord input_size; + cutlass::Tensor4DCoord filter_size; + cutlass::Tensor4DCoord padding; + cutlass::MatrixCoord conv_stride; + cutlass::MatrixCoord dilation; + bool reference_check; + bool measure_performance; + int iterations; + bool save_workspace; + ElementComputeEpilogue alpha; + ElementComputeEpilogue beta; + bool benchmark; + std::string tag; + + Options(): + help(false), + input_size(1, 32, 32, 32), + filter_size(32, 3, 3, 32), + padding(1, 1, 1, 1), + conv_stride(1, 1), + dilation(1, 1), + reference_check(true), + measure_performance(false), + iterations(20), + save_workspace(false), + alpha(1), + beta(0), + benchmark(false) { } + + // Verify the problem size is compatible with the CUTLASS Convolution implementation. + bool valid() { + + // + // CUTLASS attempts to load 128b vectors of cutlass::half_t (F16) elements. Consequently, + // all pointers, strides, and tensor extents must be divisible by 8 elements. + // + int const kAlignment = 8; + + if ((input_size.c() % kAlignment) || + (filter_size.n() % kAlignment)) { + + // misaligned tensors + return false; + } + + // Invalid padding + if ((padding.h() != filter_size.h() / 2) || + (padding.w() != filter_size.w() / 2)) { + + return false; + } + + return true; + } + + /// Updates input and filter sizes + void update( + cutlass::Tensor4DCoord input_size, + cutlass::Tensor4DCoord filter_size, + cutlass::MatrixCoord stride) { + + this->input_size = input_size; + this->filter_size = filter_size; + conv_stride = stride; + + padding.n() = filter_size.h() / 2; + padding.h() = filter_size.h() / 2; + padding.w() = filter_size.w() / 2; + padding.c() = filter_size.w() / 2; + } + + // 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; + } + + if (cmd.check_cmd_line_flag("ref-check")) { + reference_check = true; + } + + if (cmd.check_cmd_line_flag("perf-check")) { + measure_performance = true; + } + + if (cmd.check_cmd_line_flag("save-workspace")) { + save_workspace = true; + } + + if (cmd.check_cmd_line_flag("benchmark")) { + benchmark = true; + } + + cmd.get_cmd_line_argument("n", input_size.n()); + cmd.get_cmd_line_argument("h", input_size.h()); + cmd.get_cmd_line_argument("w", input_size.w()); + cmd.get_cmd_line_argument("c", input_size.c()); + + cmd.get_cmd_line_argument("k", filter_size.n()); + cmd.get_cmd_line_argument("r", filter_size.h()); + cmd.get_cmd_line_argument("s", filter_size.w()); + filter_size.c() = input_size.c(); + + cmd.get_cmd_line_argument("alpha", alpha); + cmd.get_cmd_line_argument("beta", beta); + + cmd.get_cmd_line_argument("iterations", iterations); + cmd.get_cmd_line_argument("tag", tag); + + if (filter_size.h() == 3 && filter_size.w() == 3) { + padding = {1, 1, 1, 1}; + } + else { + filter_size.h() = 1; + filter_size.w() = 1; + padding = {0, 0, 0, 0}; + } + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "25_ampere_fprop_mainloop_fusion example\n\n" + << " This example fuses scale+bias+relu of the activations into Ampere's\n" + << " Tensor Core operators on F16 data types to compute\n" + << " forward convolution on tensors of layout NHWC.\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement.\n\n" + << " --n Input tensor extent N\n" + << " --h Input tensor extent H\n" + << " --w Input tensor extent W\n" + << " --c Input tensor extent C\n" + << " --k Filter extent K\n" + << " --r Filter extent R\n" + << " --s Filter extent S\n\n" + << " --alpha Epilogue scalar alpha\n" + << " --beta Epilogue scalar beta\n\n" + << " --ref-check If set (true), reference check on the host is computed\n" + << " --perf-check If set (true), performance is measured.\n" + << " --benchmark If set (true), performance benchmarking on several layers and batch-size.\n" + << " --iterations Number of profiling iterations to perform.\n" + << " --save-workspace If set, workspace is written to a text file.\n" + << " --tag String to replicate across the first column in the results table\n"; + + out << "\n\nExamples:\n\n" + << "$ ./examples/25_ampere_fprop_mainloop_fusion/25_ampere_fprop_mainloop_fusion --n=32 --h=224 --w=224 --c=128 --k=256 --r=1 --s=1\n\n" + << "$ ./examples/25_ampere_fprop_mainloop_fusion/25_ampere_fprop_mainloop_fusion --n=1 --h=224 --w=224 --c=32 --k=32 --r=3 --s=3 --ref-check\n\n"; + + return out; + } + + /// Computes the output tensor size (NPQK) + cutlass::Tensor4DCoord output_size() const { + return cutlass::Tensor4DCoord( + input_size.n(), + (input_size.h() + padding.n() + padding.h() - filter_size.h()) / conv_stride.row() + 1, + (input_size.w() + padding.w() + padding.c() - filter_size.w()) / conv_stride.column() + 1, + filter_size.n()); + } + + /// Compute performance in GFLOP/s + double gflops(double runtime_s) const { + + // Number of multiply-adds = NPQK * CRS + int64_t fmas = output_size().product() * int64_t(filter_size.h() * filter_size.w() * filter_size.c()); + + // Two flops per multiply-add + return 2.0 * double(fmas) / double(1.0e9) / runtime_s; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +struct Result { + double runtime_ms; + double gflops; + cutlass::Status status; + cutlass::Status reference_check; + cudaError_t error; + + Result(): + runtime_ms(0), + gflops(0), + status(cutlass::Status::kSuccess), + reference_check(cutlass::Status::kInvalid), + error(cudaSuccess) { } + + static std::ostream & print_header(std::ostream &out, Options const &options) { + + if (!options.tag.empty()) { + out << "Name,"; + } + + out << "Layer,N,H,W,C,K,R,S,Stride_H,Stride_W,Runtime,GFLOPs"; + + return out; + } + + std::ostream & print(std::ostream &out, int idx, Options const &options) { + + if (!options.tag.empty()) { + out << options.tag << ","; + } + + out + << "conv_" << idx << "," + << options.input_size.n() << "," + << options.input_size.h() << "," + << options.input_size.w() << "," + << options.input_size.c() << "," + << options.filter_size.n() << "," + << options.filter_size.h() << "," + << options.filter_size.w() << "," + << options.conv_stride.row() << "," + << options.conv_stride.column() << "," + << runtime_ms << "," + << gflops; + + return out; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Runs one benchmark +Result profile_convolution(Options const &options) { + + Result result; + + // + // Allocate host-device tensors using the CUTLASS Utilities. + // + + cutlass::HostTensor tensor_a(options.input_size); + cutlass::HostTensor tensor_transformed_a(options.input_size); + cutlass::HostTensor tensor_b(options.filter_size); + cutlass::HostTensor + tensor_a_scale({1, options.input_size.c()}); + cutlass::HostTensor + tensor_a_bias({1, options.input_size.c()}); + cutlass::HostTensor tensor_c(options.output_size()); + cutlass::HostTensor tensor_ref_c(options.output_size()); + + // + // Initialize tensors + // + + // Fill tensor A on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_a.host_view(), + 1, + ElementInputA(3), + ElementInputA(-4), + 0); + + // Fill scale vector for tensor A on host with uniform-distribution random + // data + cutlass::reference::host::TensorFillRandomUniform( + tensor_a_scale.host_view(), + 1, + ElementInputA(3), + ElementInputA(-4), + 0); + + // Fill bias vector for tensor A on host with uniform-distribution random + // data + cutlass::reference::host::TensorFillRandomUniform( + tensor_a_bias.host_view(), + 1, + ElementInputA(3), + ElementInputA(-4), + 0); + + // Fill tensor B on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_b.host_view(), + 1, + ElementInputB(7), + ElementInputB(-8), + 0); + + // Fill tensor C on host with zeros + cutlass::reference::host::TensorFill( + tensor_c.host_view()); + + // Fill tensor C for reference on host with zeros + cutlass::reference::host::TensorFill( + tensor_ref_c.host_view()); + + // Copy data from host to GPU + tensor_a.sync_device(); + tensor_a_scale.sync_device(); + tensor_a_bias.sync_device(); + tensor_b.sync_device(); + tensor_c.sync_device(); + tensor_ref_c.sync_device(); + + // + // Define arguments for CUTLASS Convolution + // + + cutlass::conv::Mode mode = cutlass::conv::Mode::kCrossCorrelation; + + // Split K dimension into 1 partitions + int split_k_slices = 1; + + // Construct Conv2dProblemSize with user defined output size + cutlass::conv::Conv2dProblemSize problem_size( + options.input_size, + options.filter_size, + options.padding, + options.conv_stride, + options.dilation, + options.output_size(), + mode, + split_k_slices + ); + + typename ImplicitGemmFusion::Arguments arguments{ + problem_size, + tensor_a.device_ref(), + tensor_b.device_ref(), + tensor_a_scale.device_ref(), + tensor_a_bias.device_ref(), + tensor_c.device_ref(), + tensor_c.device_ref(), + {options.alpha, options.beta}, + }; + + // + // Initialize CUTLASS Convolution + // + + ImplicitGemmFusion implicit_gemm_fusion_op; + + size_t workspace_size = implicit_gemm_fusion_op.get_workspace_size(arguments); + + // Allocate workspace memory + cutlass::device_memory::allocation workspace(workspace_size); + + result.status = implicit_gemm_fusion_op.can_implement(arguments); + CUTLASS_CHECK(result.status); + + result.status = implicit_gemm_fusion_op.initialize(arguments, workspace.get()); + CUTLASS_CHECK(result.status); + + // + // Launch initialized CUTLASS kernel + // + result.status = implicit_gemm_fusion_op(); + + CUTLASS_CHECK(result.status); + + // + // Optional reference check + // + + if (options.reference_check) { + std::cout << "Verification on device...\n"; + + // Compute scale + bias + relu in host code + for (int n = 0; n < options.input_size.n(); ++n) { + for (int h = 0; h < options.input_size.h(); ++h) { + for (int w = 0; w < options.input_size.w(); ++w) { + for (int c = 0; c < options.input_size.c(); ++c) { + tensor_transformed_a.at({n, h, w, c}) = std::max( + ElementOutput(0), ElementOutput(tensor_a.at({n, h, w, c}) * + tensor_a_scale.at({0, c}) + + tensor_a_bias.at({0, c}))); + } + } + } + } + + tensor_transformed_a.sync_device(); + + // Compute with reference implementation + cutlass::reference::device::Conv2dFprop< + ElementInputA, + LayoutInputA, + ElementInputB, + LayoutInputB, + ElementOutput, + LayoutOutput, + ElementComputeEpilogue, + ElementAccumulator, + cutlass::NumericConverter + >( + problem_size, + tensor_transformed_a.device_ref(), + tensor_b.device_ref(), + tensor_c.device_ref(), + tensor_ref_c.device_ref(), + options.alpha, + options.beta + ); + + // Check if output from CUTLASS kernel and reference kernel are equal or not + tensor_c.sync_host(); + tensor_ref_c.sync_host(); + + bool passed = cutlass::reference::host::TensorEquals( + tensor_c.host_view(), + tensor_ref_c.host_view()); + + if (!passed) { + result.reference_check = cutlass::Status::kErrorInternal; + std::cout << "ERROR - results miscompared.\n"; + } + else { + result.reference_check = cutlass::Status::kSuccess; + std::cout << "Passed.\n"; + } + } + else { + result.reference_check = cutlass::Status::kInvalid; + } + + if (options.save_workspace) { + + std::stringstream ss; + + ss << "18_ampere_fused_fprop_batch_normalization_" + << options.input_size.n() << "x" << options.input_size.h() << "x" << options.input_size.w() << "x" << options.input_size.c() + << "_" + << options.filter_size.n() << "x" << options.filter_size.h() << "x" << options.filter_size.w() << "x" << options.filter_size.c() + << ".dat"; + + std::ofstream output_workspace(ss.str()); + + output_workspace + << "Input = \n" << tensor_a.host_view() << "\n\n" + << "Filters = \n" << tensor_b.host_view() << "\n\n"; + + if (options.reference_check) { + output_workspace << "Reference = \n" << tensor_ref_c.host_view() << "\n\n"; + } + + output_workspace << "Computed = \n" << tensor_c.host_view() << std::endl; + + std::cout << "Results written to '" << ss.str() << "'." << std::endl; + } + + // + // Performance measurement + // + + if (options.measure_performance) { + + cudaEvent_t events[2]; + + for (auto & event : events) { + result.error = cudaEventCreate(&event); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventCreate() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + } + + // Record an event at the start of a series of convolution operations. + result.error = cudaEventRecord(events[0]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Launch a sequence of implicit GEMM operations on the device + for (int iteration = 0; iteration < options.iterations; ++iteration) { + result.status = implicit_gemm_fusion_op(); + CUTLASS_CHECK(result.status); + } + + // Record an event when the convolutions have been launched. + result.error = cudaEventRecord(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Wait for work on the device to complete. + result.error = cudaEventSynchronize(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventSynchronize() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Measure elapsed runtime + float runtime_ms = 0; + result.error = cudaEventElapsedTime(&runtime_ms, events[0], events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Print average runtime and GFLOPs. + result.runtime_ms = double(runtime_ms) / double(options.iterations); + result.gflops = options.gflops(result.runtime_ms / 1000.0); + + // Cleanup + for (auto event : events) { + (void)cudaEventDestroy(event); + } + } + + return result; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +int main(int argc, char const **args) { + + bool notSupported = false; + + // Ampere Tensor Core operations exposed with mma.sync are first available in CUDA 11.0. + // + // CUTLASS must be compiled with CUDA 11 Toolkit to run Conv2dFprop examples. + if (!(__CUDACC_VER_MAJOR__ > 11 || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 0))) { + std::cerr << "Ampere Tensor Core operations must be compiled with CUDA 11.0 Toolkit or later." << std::endl; + notSupported = true; + } + + cudaDeviceProp props; + CUDA_CHECK(cudaGetDeviceProperties(&props, 0)); + + if (!(props.major == 8 && props.minor == 0)) { + std::cerr << "This test must run on SM80 A100.\n"; + notSupported = true; + } + + if (notSupported) { + return 0; + } + + Options options; + + options.parse(argc, args); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + if (options.benchmark) { + // Benchmark several layers + + int batch_sizes[] = {34, 408}; + + struct Benchmark { + int h, w, c, k, r, s, stride_h, stride_w; + } layers[] = { + {56, 56, 64, 256, 1, 1, 1, 1}, + {56, 56, 64, 64, 1, 1, 1, 1}, + {56, 56, 64, 64, 3, 3, 1, 1}, + {56, 56, 256, 64, 1, 1, 1, 1}, + {56, 56, 256, 512, 1, 1, 2, 2}, + {56, 56, 256, 128, 1, 1, 1, 1}, + {56, 56, 128, 128, 3, 3, 2, 2}, + {28, 28, 128, 512, 1, 1, 1, 1}, + {28, 28, 512, 128, 1, 1, 1, 1}, + {28, 28, 128, 128, 3, 3, 1, 1}, + {28, 28, 512, 1024, 1, 1, 2, 2}, + {28, 28, 512, 256, 1, 1, 1, 1}, + {28, 28, 256, 256, 3, 3, 2, 2}, + {14, 14, 256, 1024, 1, 1, 1, 1}, + {14, 14, 1024, 256, 1, 1, 1, 1}, + {14, 14, 256, 256, 3, 3, 1, 1}, + {14, 14, 1024, 2048, 1, 1, 2, 2}, + {14, 14, 1024, 512, 1, 1, 1, 1}, + {14, 14, 512, 512, 3, 3, 2, 2}, + { 7, 7, 512, 2048, 1, 1, 1, 1}, + { 7, 7, 2048, 512, 1, 1, 1, 1}, + { 7, 7, 512, 512, 3, 3, 1, 1}, + }; + + Result::print_header(std::cout, options) << std::endl; + + int idx = 1; + + for (auto const &layer : layers) { + for (auto N : batch_sizes) { + options.update({N, layer.h, layer.w, layer.c}, + {layer.k, layer.r, layer.s, layer.c}, + {layer.stride_h, layer.stride_w}); + + Result result = profile_convolution(options); + result.print(std::cout, idx, options) << std::endl; + } + + ++idx; + } + } + else { + + // Execute one problem size + if (!options.valid()) { + std::cerr << "Invalid problem." << std::endl; + return -1; + } + + Result result = profile_convolution(options); + + Result::print_header(std::cout, options) << std::endl; + result.print(std::cout, 1, options) << std::endl; + } + + return 0; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/examples/26_ampere_wgrad_mainloop_fusion/CMakeLists.txt b/examples/26_ampere_wgrad_mainloop_fusion/CMakeLists.txt new file mode 100644 index 00000000..d35277d0 --- /dev/null +++ b/examples/26_ampere_wgrad_mainloop_fusion/CMakeLists.txt @@ -0,0 +1,28 @@ +# Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without modification, are permitted +# provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright notice, this list of +# conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright notice, this list of +# conditions and the following disclaimer in the documentation and/or other materials +# provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used +# to endorse or promote products derived from this software without specific prior written +# permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR +# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; +# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +cutlass_example_add_executable( + 26_ampere_wgrad_mainloop_fusion + ampere_wgrad_mainloop_fusion.cu + ) + diff --git a/examples/26_ampere_wgrad_mainloop_fusion/ampere_wgrad_mainloop_fusion.cu b/examples/26_ampere_wgrad_mainloop_fusion/ampere_wgrad_mainloop_fusion.cu new file mode 100644 index 00000000..001aa912 --- /dev/null +++ b/examples/26_ampere_wgrad_mainloop_fusion/ampere_wgrad_mainloop_fusion.cu @@ -0,0 +1,749 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/** + +This example shows how to fuse activation's per channel scale+bias+relu +into the wgrad mainloop. + +Compared with original fprop kernel, this example has two more vectors, one for +the scale and one for the bias. The length of the vectors are the same as the +activation channel number. This kernels loads the vectors when the associated +activation channels are loaded in the mainloop. Between reading the +activations and scale/bias data from the shared memory and calling tensor core +instructions, scale+bias+relu is computed in the register file. + +This example is customized for Ampere 16816 fp16 tensor core instruction. +Changing to different data types or different tensor core instruction require +source code changing. See +include/cutlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h for more +technical details. +*/ + +#include +#include + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/device/gemm.h" +#include "cutlass/conv/kernel/default_conv2d_wgrad_fusion.h" +#include "cutlass/conv/device/implicit_gemm_convolution_fusion.h" + +#include "cutlass/util/command_line.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/reference/device/gemm.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/device/convolution.h" +#include "cutlass/util/tensor_view_io.h" + +#include "helper.h" + +// The code section below describes datatype for input, output tensors and computation between +// elements +using ElementAccumulator = float; // Data type of accumulator +using ElementComputeEpilogue = float; // Data type of epilogue computation (alpha, beta) +using ElementInputA = cutlass::half_t; // Data type of elements in input tensor +using ElementInputB = cutlass::half_t; // Data type of elements in input tensor +using ElementInputScaleBias = cutlass::half_t; // Data type of elements in input sclae and bias vectors +using ElementOutput = float; // Data type of elements in output tensor + +using LayoutInputA = cutlass::layout::TensorNHWC; +using LayoutInputB = cutlass::layout::TensorNHWC; +using LayoutInputScaleBias = cutlass::layout::RowMajor; +using LayoutOutput = cutlass::layout::TensorNHWC; + +// This code section describes whether you want to use tensor cores or regular SIMT cores on GPU SM +using MMAOp = cutlass::arch::OpClassTensorOp; + +// This code section describes CUDA SM architecture number +using SmArch = cutlass::arch::Sm80; + +// This code section describes the tile size a thread block will compute +using ThreadblockShape = cutlass::gemm::GemmShape<128, 128, 32>; // Threadblock tile shape + +// This code section describes tile size a warp will compute +using WarpShape = cutlass::gemm::GemmShape<64, 64, 32>; // Warp tile shape + +// This code section describes the size of MMA op +using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>; // TensorCore instruction shape + +// This code section describes how threadblocks are scheduled on GPU +using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; + +// Number of pipelines you want to use +constexpr int NumStages = 5; + +// This code section describe iterator algorithm selected is Analytic or Optimized +static cutlass::conv::IteratorAlgorithm const IteratorAlgorithm = cutlass::conv::IteratorAlgorithm::kOptimized; + +// This code section describes the epilogue part of the kernel, we use default value +using EpilogueOp = cutlass::epilogue::thread::LinearCombination< + ElementOutput, // Data type of output matrix. + 128 / cutlass::sizeof_bits::value, // The number of elements per vectorized. + // memory access. This becomes the vector width of + // math instructions in the epilogue too. + ElementAccumulator, // Data type of accumulator + ElementComputeEpilogue>; // Data type for alpha/beta in linear combination + +using Conv2dWgradFusionKernel = typename cutlass::conv::kernel::DefaultConv2dWgradFusion< + ElementInputA, LayoutInputA, + ElementInputB, LayoutInputB, + ElementInputScaleBias, LayoutInputScaleBias, + ElementOutput, LayoutOutput, + ElementAccumulator, + MMAOp, + SmArch, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOp, + SwizzleThreadBlock, + NumStages, + cutlass::arch::OpMultiplyAdd, + IteratorAlgorithm +>::Kernel; + +using ImplicitGemmFusion = cutlass::conv::device::ImplicitGemmConvolutionFusion; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Command line options parsing +struct Options { + + bool help; + cutlass::Tensor4DCoord input_size; + cutlass::Tensor4DCoord filter_size; + cutlass::Tensor4DCoord padding; + cutlass::MatrixCoord conv_stride; + cutlass::MatrixCoord dilation; + bool reference_check; + bool measure_performance; + int iterations; + bool save_workspace; + ElementComputeEpilogue alpha; + ElementComputeEpilogue beta; + bool benchmark; + std::string tag; + + Options(): + help(false), + input_size(1, 32, 32, 32), + filter_size(32, 3, 3, 32), + padding(1, 1, 1, 1), + conv_stride(1, 1), + dilation(1, 1), + reference_check(true), + measure_performance(false), + iterations(20), + save_workspace(false), + alpha(1), + beta(0), + benchmark(false) { } + + // Verify the problem size is compatible with the CUTLASS Convolution implementation. + bool valid() { + + // + // CUTLASS attempts to load 128b vectors of cutlass::half_t (F16) elements. Consequently, + // all pointers, strides, and tensor extents must be divisible by 8 elements. + // + int const kAlignment = 8; + + if ((input_size.c() % kAlignment) || + (filter_size.n() % kAlignment)) { + + // misaligned tensors + return false; + } + + // Invalid padding + if ((padding.h() != filter_size.h() / 2) || + (padding.w() != filter_size.w() / 2)) { + + return false; + } + + return true; + } + + /// Updates input and filter sizes + void update( + cutlass::Tensor4DCoord input_size, + cutlass::Tensor4DCoord filter_size, + cutlass::MatrixCoord stride) { + + this->input_size = input_size; + this->filter_size = filter_size; + conv_stride = stride; + + padding.n() = filter_size.h() / 2; + padding.h() = filter_size.h() / 2; + padding.w() = filter_size.w() / 2; + padding.c() = filter_size.w() / 2; + } + + // 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; + } + + if (cmd.check_cmd_line_flag("ref-check")) { + reference_check = true; + } + + if (cmd.check_cmd_line_flag("perf-check")) { + measure_performance = true; + } + + if (cmd.check_cmd_line_flag("save-workspace")) { + save_workspace = true; + } + + if (cmd.check_cmd_line_flag("benchmark")) { + benchmark = true; + } + + cmd.get_cmd_line_argument("n", input_size.n()); + cmd.get_cmd_line_argument("h", input_size.h()); + cmd.get_cmd_line_argument("w", input_size.w()); + cmd.get_cmd_line_argument("c", input_size.c()); + + cmd.get_cmd_line_argument("k", filter_size.n()); + cmd.get_cmd_line_argument("r", filter_size.h()); + cmd.get_cmd_line_argument("s", filter_size.w()); + filter_size.c() = input_size.c(); + + cmd.get_cmd_line_argument("alpha", alpha); + cmd.get_cmd_line_argument("beta", beta); + + cmd.get_cmd_line_argument("iterations", iterations); + cmd.get_cmd_line_argument("tag", tag); + + if (filter_size.h() == 3 && filter_size.w() == 3) { + padding = {1, 1, 1, 1}; + } + else { + filter_size.h() = 1; + filter_size.w() = 1; + padding = {0, 0, 0, 0}; + } + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "26_ampere_fused_wgrad_batch_normalization example\n\n" + << " This example fuses scale+bias+relu from batch norm into Ampere's\n" + << " Tensor Core operators on F16 data types to compute\n" + << " backward convolution on tensors of layout NHWC.\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement.\n\n" + << " --n Input tensor extent N\n" + << " --h Input tensor extent H\n" + << " --w Input tensor extent W\n" + << " --c Input tensor extent C\n" + << " --k Filter extent K\n" + << " --r Filter extent R\n" + << " --s Filter extent S\n\n" + << " --alpha Epilogue scalar alpha\n" + << " --beta Epilogue scalar beta\n\n" + << " --ref-check If set (true), reference check on the host is computed\n" + << " --perf-check If set (true), performance is measured.\n" + << " --benchmark If set (true), performance benchmarking on several layers and batch-size.\n" + << " --iterations Number of profiling iterations to perform.\n" + << " --save-workspace If set, workspace is written to a text file.\n" + << " --tag String to replicate across the first column in the results table\n"; + + out << "\n\nExamples:\n\n" + << "$ ./examples/26_ampere_fused_fprop_batch_normalization/26_ampere_fused_wgrad_batch_normalization --n=32 --h=224 --w=224 --c=128 --k=256 --r=1 --s=1\n\n" + << "$ ./examples/26_ampere_fused_fprop_batch_normalization/26_ampere_fused_wgrad_batch_normalization --n=1 --h=224 --w=224 --c=32 --k=32 --r=3 --s=3 --ref-check\n\n"; + + return out; + } + + /// Computes the output tensor size (NPQK) + cutlass::Tensor4DCoord output_size() const { + return cutlass::Tensor4DCoord( + input_size.n(), + (input_size.h() + padding.n() + padding.h() - filter_size.h()) / conv_stride.row() + 1, + (input_size.w() + padding.w() + padding.c() - filter_size.w()) / conv_stride.column() + 1, + filter_size.n()); + } + + /// Compute performance in GFLOP/s + double gflops(double runtime_s) const { + + // Number of multiply-adds = NPQK * CRS + int64_t fmas = output_size().product() * int64_t(filter_size.h() * filter_size.w() * filter_size.c()); + + // Two flops per multiply-add + return 2.0 * double(fmas) / double(1.0e9) / runtime_s; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +struct Result { + double runtime_ms; + double gflops; + cutlass::Status status; + cutlass::Status reference_check; + cudaError_t error; + + Result(): + runtime_ms(0), + gflops(0), + status(cutlass::Status::kSuccess), + reference_check(cutlass::Status::kInvalid), + error(cudaSuccess) { } + + static std::ostream & print_header(std::ostream &out, Options const &options) { + + if (!options.tag.empty()) { + out << "Name,"; + } + + out << "Layer,N,H,W,C,K,R,S,Stride_H,Stride_W,Runtime,GFLOPs"; + + return out; + } + + std::ostream & print(std::ostream &out, int idx, Options const &options) { + + if (!options.tag.empty()) { + out << options.tag << ","; + } + + out + << "conv_" << idx << "," + << options.input_size.n() << "," + << options.input_size.h() << "," + << options.input_size.w() << "," + << options.input_size.c() << "," + << options.filter_size.n() << "," + << options.filter_size.h() << "," + << options.filter_size.w() << "," + << options.conv_stride.row() << "," + << options.conv_stride.column() << "," + << runtime_ms << "," + << gflops; + + return out; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Runs one benchmark +Result profile_convolution(Options const &options) { + + Result result; + + // + // Allocate host-device tensors using the CUTLASS Utilities. + // + + cutlass::HostTensor tensor_a(options.output_size()); + cutlass::HostTensor tensor_b(options.input_size); + cutlass::HostTensor tensor_transformed_b(options.input_size); + cutlass::HostTensor + tensor_b_scale({1, options.input_size.c()}); + cutlass::HostTensor + tensor_b_bias({1, options.input_size.c()}); + + cutlass::HostTensor tensor_c(options.filter_size); + cutlass::HostTensor tensor_ref_c(options.filter_size); + + // + // Initialize tensors + // + + // Fill tensor A on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_a.host_view(), + 1, + ElementInputA(3), + ElementInputA(-4), + 0); + + // Fill tensor B on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_b.host_view(), + 1, + ElementInputB(7), + ElementInputB(-8), + 0); + + // Fill scale vector for tensor B on host with uniform-distribution random + // data + cutlass::reference::host::TensorFillRandomUniform( + tensor_b_scale.host_view(), + 1, + ElementInputA(3), + ElementInputA(-4), + 0); + + // Fill bias vector for tensor B on host with uniform-distribution random + // data + cutlass::reference::host::TensorFillRandomUniform( + tensor_b_bias.host_view(), + 1, + ElementInputA(3), + ElementInputA(-4), + 0); + + // Fill tensor C on host with zeros + cutlass::reference::host::TensorFill( + tensor_c.host_view()); + + // Fill tensor C for reference on host with zeros + cutlass::reference::host::TensorFill( + tensor_ref_c.host_view()); + + // Copy data from host to GPU + tensor_a.sync_device(); + tensor_b.sync_device(); + tensor_b_scale.sync_device(); + tensor_b_bias.sync_device(); + tensor_c.sync_device(); + tensor_ref_c.sync_device(); + + // + // Define arguments for CUTLASS Convolution + // + + cutlass::conv::Mode mode = cutlass::conv::Mode::kCrossCorrelation; + + // Split K dimension into 1 partitions + int split_k_slices = 1; + + // Construct Conv2dProblemSize with user defined output size + cutlass::conv::Conv2dProblemSize problem_size( + options.input_size, + options.filter_size, + options.padding, + options.conv_stride, + options.dilation, + options.output_size(), + mode, + split_k_slices + ); + + typename ImplicitGemmFusion::Arguments arguments{ + problem_size, + tensor_a.device_ref(), + tensor_b.device_ref(), + tensor_b_scale.device_ref(), + tensor_b_bias.device_ref(), + tensor_c.device_ref(), + tensor_c.device_ref(), + {options.alpha, options.beta}, + }; + + // + // Initialize CUTLASS Convolution + // + + ImplicitGemmFusion implicit_gemm_fusion_op; + + size_t workspace_size = implicit_gemm_fusion_op.get_workspace_size(arguments); + + // Allocate workspace memory + cutlass::device_memory::allocation workspace(workspace_size); + + result.status = implicit_gemm_fusion_op.can_implement(arguments); + CUTLASS_CHECK(result.status); + + result.status = implicit_gemm_fusion_op.initialize(arguments, workspace.get()); + CUTLASS_CHECK(result.status); + + // + // Launch initialized CUTLASS kernel + // + result.status = implicit_gemm_fusion_op(); + + CUTLASS_CHECK(result.status); + + // + // Optional reference check + // + + if (options.reference_check) { + std::cout << "Verification on device...\n"; + + // Compute scale + bias + relu in host code + for (int n = 0; n < options.input_size.n(); ++n) { + for (int h = 0; h < options.input_size.h(); ++h) { + for (int w = 0; w < options.input_size.w(); ++w) { + for (int c = 0; c < options.input_size.c(); ++c) { + tensor_transformed_b.at({n, h, w, c}) = std::max( + ElementOutput(0), ElementOutput(tensor_b.at({n, h, w, c}) * + tensor_b_scale.at({0, c}) + + tensor_b_bias.at({0, c}))); + } + } + } + } + + tensor_transformed_b.sync_device(); + + // Compute with reference implementation + cutlass::reference::device::Conv2dWgrad< + ElementInputA, + LayoutInputA, + ElementInputB, + LayoutInputB, + ElementOutput, + LayoutOutput, + ElementComputeEpilogue, + ElementAccumulator, + cutlass::NumericConverter + >( + problem_size, + tensor_a.device_ref(), + tensor_transformed_b.device_ref(), + tensor_c.device_ref(), + tensor_ref_c.device_ref(), + options.alpha, + options.beta + ); + + // Check if output from CUTLASS kernel and reference kernel are equal or not + tensor_c.sync_host(); + tensor_ref_c.sync_host(); + + bool passed = cutlass::reference::host::TensorEquals( + tensor_c.host_view(), + tensor_ref_c.host_view()); + + if (!passed) { + result.reference_check = cutlass::Status::kErrorInternal; + std::cout << "ERROR - results miscompared.\n"; + } + else { + result.reference_check = cutlass::Status::kSuccess; + std::cout << "Passed.\n"; + } + } + else { + result.reference_check = cutlass::Status::kInvalid; + } + + if (options.save_workspace) { + + std::stringstream ss; + + ss << "26_ampere_fused_wgrad_batch_normalization_" + << options.input_size.n() << "x" << options.input_size.h() << "x" << options.input_size.w() << "x" << options.input_size.c() + << "_" + << options.filter_size.n() << "x" << options.filter_size.h() << "x" << options.filter_size.w() << "x" << options.filter_size.c() + << ".dat"; + + std::ofstream output_workspace(ss.str()); + + output_workspace + << "Input = \n" << tensor_a.host_view() << "\n\n" + << "Filters = \n" << tensor_b.host_view() << "\n\n"; + + if (options.reference_check) { + output_workspace << "Reference = \n" << tensor_ref_c.host_view() << "\n\n"; + } + + output_workspace << "Computed = \n" << tensor_c.host_view() << std::endl; + + std::cout << "Results written to '" << ss.str() << "'." << std::endl; + } + + // + // Performance measurement + // + + if (options.measure_performance) { + + cudaEvent_t events[2]; + + for (auto & event : events) { + result.error = cudaEventCreate(&event); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventCreate() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + } + + // Record an event at the start of a series of convolution operations. + result.error = cudaEventRecord(events[0]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Launch a sequence of implicit GEMM operations on the device + for (int iteration = 0; iteration < options.iterations; ++iteration) { + result.status = implicit_gemm_fusion_op(); + CUTLASS_CHECK(result.status); + } + + // Record an event when the convolutions have been launched. + result.error = cudaEventRecord(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Wait for work on the device to complete. + result.error = cudaEventSynchronize(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventSynchronize() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Measure elapsed runtime + float runtime_ms = 0; + result.error = cudaEventElapsedTime(&runtime_ms, events[0], events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Print average runtime and GFLOPs. + result.runtime_ms = double(runtime_ms) / double(options.iterations); + result.gflops = options.gflops(result.runtime_ms / 1000.0); + + // Cleanup + for (auto event : events) { + (void)cudaEventDestroy(event); + } + } + + return result; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +int main(int argc, char const **args) { + + bool notSupported = false; + + // Ampere Tensor Core operations exposed with mma.sync are first available in CUDA 11.0. + // + // CUTLASS must be compiled with CUDA 11 Toolkit to run Conv2dFprop examples. + if (!(__CUDACC_VER_MAJOR__ > 11 || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 0))) { + std::cerr << "Ampere Tensor Core operations must be compiled with CUDA 11.0 Toolkit or later." << std::endl; + notSupported = true; + } + + cudaDeviceProp props; + CUDA_CHECK(cudaGetDeviceProperties(&props, 0)); + + if (!(props.major == 8 && props.minor == 0)) { + std::cerr << "This test must run on SM80 A100.\n"; + notSupported = true; + } + + if (notSupported) { + return 0; + } + + Options options; + + options.parse(argc, args); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + if (options.benchmark) { + // Benchmark several layers + + int batch_sizes[] = {34, 408}; + + struct Benchmark { + int h, w, c, k, r, s, stride_h, stride_w; + } layers[] = { + {56, 56, 64, 256, 1, 1, 1, 1}, + {56, 56, 64, 64, 1, 1, 1, 1}, + {56, 56, 64, 64, 3, 3, 1, 1}, + {56, 56, 256, 64, 1, 1, 1, 1}, + {56, 56, 256, 512, 1, 1, 2, 2}, + {56, 56, 256, 128, 1, 1, 1, 1}, + {56, 56, 128, 128, 3, 3, 2, 2}, + {28, 28, 128, 512, 1, 1, 1, 1}, + {28, 28, 512, 128, 1, 1, 1, 1}, + {28, 28, 128, 128, 3, 3, 1, 1}, + {28, 28, 512, 1024, 1, 1, 2, 2}, + {28, 28, 512, 256, 1, 1, 1, 1}, + {28, 28, 256, 256, 3, 3, 2, 2}, + {14, 14, 256, 1024, 1, 1, 1, 1}, + {14, 14, 1024, 256, 1, 1, 1, 1}, + {14, 14, 256, 256, 3, 3, 1, 1}, + {14, 14, 1024, 2048, 1, 1, 2, 2}, + {14, 14, 1024, 512, 1, 1, 1, 1}, + {14, 14, 512, 512, 3, 3, 2, 2}, + { 7, 7, 512, 2048, 1, 1, 1, 1}, + { 7, 7, 2048, 512, 1, 1, 1, 1}, + { 7, 7, 512, 512, 3, 3, 1, 1}, + }; + + Result::print_header(std::cout, options) << std::endl; + + int idx = 1; + + for (auto const &layer : layers) { + for (auto N : batch_sizes) { + options.update({N, layer.h, layer.w, layer.c}, + {layer.k, layer.r, layer.s, layer.c}, + {layer.stride_h, layer.stride_w}); + + Result result = profile_convolution(options); + result.print(std::cout, idx, options) << std::endl; + } + + ++idx; + } + } + else { + + // Execute one problem size + if (!options.valid()) { + std::cerr << "Invalid problem." << std::endl; + return -1; + } + + Result result = profile_convolution(options); + + Result::print_header(std::cout, options) << std::endl; + result.print(std::cout, 1, options) << std::endl; + } + + return 0; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/27_ampere_3xtf32_fast_accurate_tensorop_gemm.cu b/examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/27_ampere_3xtf32_fast_accurate_tensorop_gemm.cu new file mode 100644 index 00000000..52260e44 --- /dev/null +++ b/examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/27_ampere_3xtf32_fast_accurate_tensorop_gemm.cu @@ -0,0 +1,744 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/** +NVIDIA Ampere architecture starts supporting tfloat32 (see include/cutlass/tfloat32.h) +data types in tensor cores. One big advantage is that we can load in fp32 data and convert them +implicitly to tf32 inside the GEMM kernel which means no change is needed to accelerate traditional +fp32 data by using NVIDIA Ampere architecture. + +We can use the tf32 mode of tensor core to emulate a fast accurate SGEMM kernel which is accelerated +using Ampere Tensor Cores (see include/cutlass/gemm/warp/mma_tensor_op_fast_f32.h). + +The trick is very simple + a x b = (a_big + a_small) x (b_big + b_small) = a_big x b_big + a_big x b_small + a_small x b_big + big = convert_to_tf32(fp32) + small = convert_to_tf32(fp32 - big) + +a_small x b_small is discarded because they are too small. + +This example demonstrates usage of this kernel, along with accuracy measurements w.r.t. actual FP32 +results (SGEMM using SIMT) and against FP64 results (DGEMM) + +To enable this feature, the only change needs to make is to change the default OpMultiplyAdd to +OpMultiplyAddFastF32. + +Now, we have several different flavors of sgemm now in the profiler for Ampere. Here are the difference + + sgemm // CUDA core SIMT kernel. FP32 in, accumulated in FP32, FP32 out. + s1688gemm // Use 3xTF32 to emulate FP32. FP32 in, converted in TF32-big and TF32-small internally, + // accumulated in FP32, FP32 out. + s1688tf32gemm // Use 1xTF32. FP32 in, converted to one TF32 internally, accumulated in FP32, FP32 out. + s1688gemm_tf32 // TF32 in, accumulated in FP32, FP32 out. +*/ + +#include +#include +#include + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/device/gemm.h" + +#include "cutlass/util/command_line.h" +#include "cutlass/util/host_tensor.h" + +#include "cutlass/util/reference/device/gemm.h" +#include "cutlass/util/reference/host/tensor_reduce.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_norm.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/error_metrics.h" +#include "cutlass/util/tensor_view_io.h" + +#include "helper.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Result structure +struct Result { + + double runtime_ms; + double gflops; + cutlass::Status status; + cudaError_t error; + + int m, n, k; + double l2_norm_3xtf32_vs_fp64; + double l2_norm_1xtf32_vs_fp64; + double l2_norm_fp32_vs_fp64; + + // ctor + Result( + int m, int n, int k, + double runtime_ms, double gflops, + double l2_norm_3xtf32_vs_fp64, + double l2_norm_1xtf32_vs_fp64, + double l2_norm_fp32_vs_fp64) : + m(m), n(n), k(k), + runtime_ms(runtime_ms), gflops(gflops), + l2_norm_3xtf32_vs_fp64(l2_norm_3xtf32_vs_fp64), + l2_norm_1xtf32_vs_fp64(l2_norm_1xtf32_vs_fp64), + l2_norm_fp32_vs_fp64(l2_norm_fp32_vs_fp64) {} + + Result() {} + + // + // Methods + // + static void print_csv_header() { + std::cout << "M,N,K,Runtime(ms),GFLOPS,3xTF32_vs_FP64,1xTF32_vs_FP64,FP32_vs_FP64" << std::endl; + } + + void print_csv_row() { + std::cout << m << "," + << n << "," + << k << "," + << runtime_ms << "," + << gflops << "," + << l2_norm_3xtf32_vs_fp64 << "," + << l2_norm_1xtf32_vs_fp64 << "," + << l2_norm_fp32_vs_fp64 << std::endl; + } +}; + +std::vector results; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// Command line options parsing +struct Options { + + bool help; + + cutlass::gemm::GemmCoord problem_size; + float alpha; + float beta; + std::string rand_mode; + + int iterations; + int seed; + bool benchmark; + + Options(): + help(false), + problem_size({3456, 4096, 4096}), + iterations(20), + seed(1), + alpha(1), + beta(), + rand_mode("uniform"), + benchmark(false) { } + + bool valid() { + // + // CUTLASS attempts to load 128b vectors of F32 elements. Consequently, + // all pointers, strides, and tensor extents must be divisible by 4 elements. + // + int const kAlignment = 4; + + if ((problem_size.m() % kAlignment) || + (problem_size.n() % kAlignment) || + (problem_size.k() % kAlignment)) { + + // misaligned tensors + return false; + } + + return true; + } + + // 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; + } + + cmd.get_cmd_line_argument("m", problem_size.m()); + cmd.get_cmd_line_argument("n", problem_size.n()); + cmd.get_cmd_line_argument("k", problem_size.k()); + + cmd.get_cmd_line_argument("alpha", alpha); + cmd.get_cmd_line_argument("beta", beta); + + cmd.get_cmd_line_argument("iterations", iterations); + cmd.get_cmd_line_argument("seed", seed); + cmd.get_cmd_line_argument("rand_mode", rand_mode); + + if (cmd.check_cmd_line_flag("benchmark")) { + benchmark = true; + } + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "27_ampere_3xtf32_fast_accurate_tensorop_gemm example\n\n" + << " This example uses the CUTLASS Library to emulate FP32 with TF32 tensorop GEMM computations.\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement.\n\n" + << " --m GEMM M dimension\n" + << " --n GEMM N dimension\n" + << " --k GEMM K dimension\n" + << " --alpha Epilogue scalar alpha\n" + << " --beta Epilogue scalar beta\n\n" + << " --rand_mode gauss / uniform*\n\n" + << " --seed Random number seed (1*)\n\n" + << " --iterations Number of profiling iterations to perform.\n\n" + << " --benchmark If set (true), performance benchmarking on several layers and batch-size.\n\n"; + + out << "\n\nExamples:\n\n" + << "$ ./examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/27_ampere_3xtf32_fast_accurate_tensorop_gemm --m=1024 --n=512 \\\n" + << " --alpha=2 --beta=0.707 \n\n"; + + return out; + } + + /// Compute performance in GFLOP/s + double gflops(double runtime_s) const { + + // Number of real-valued multiply-adds + int64_t fmas = problem_size.product(); + + // Two flops per multiply-add + return 2.0 * double(fmas) / double(1.0e9) / runtime_s; + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// The code section below describes matrix layout of input and output matrices. Column Major for +// Matrix A, Row Major for Matrix B and Row Major for Matrix C +using LayoutInputA = cutlass::layout::RowMajor; +using LayoutInputB = cutlass::layout::ColumnMajor; +using LayoutOutput = cutlass::layout::RowMajor; + +// This code section describes whether you want to use tensor cores or regular SIMT cores on GPU SM +using MMAOp = cutlass::arch::OpClassTensorOp; + +// This code section describes CUDA SM architecture number +using SmArch = cutlass::arch::Sm80; + +// This code section describes the tile size a thread block will compute +using ShapeMMAThreadBlock = + cutlass::gemm::GemmShape<128, 64, 16>; // <- threadblock tile M = 128, N = 128, K = 16 +// This code section describes tile size a warp will compute +using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 32, 16>; // <- warp tile M = 64, N = 64, K = 16 +// This code section describes the size of MMA op +using ShapeMMAOp = cutlass::gemm::GemmShape<16, 8, 8>; // <- MMA Op tile M = 16, N = 8, K = 8 + +// This code section describes how threadblocks are scheduled on GPU +using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; // <- ?? + +// This code section describes the epilogue part of the kernel +using EpilogueOp = cutlass::epilogue::thread::LinearCombination< + float, // <- data type of output matrix + 128 / cutlass::sizeof_bits::value, // <- the number of elements per vectorized + // memory access. For a byte, it's 16 + // elements. This becomes the vector width of + // math instructions in the epilogue too + float, // <- data type of accumulator + float>; // <- data type for alpha/beta in linear combination function + +// Number of pipelines you want to use +constexpr int NumStages = 3; +// Alignment +constexpr int Alignment = 4; + +// +// Gemm Operators (Gemm_3xTF32, Gemm_1xTF32, GEMM_F32, GEMM_F64) +// + +// Gemm_3xTF32 +using Gemm_3xTF32 = cutlass::gemm::device::Gemm< + float, + LayoutInputA, + float, + LayoutInputB, + float, + LayoutOutput, + float, + MMAOp, + SmArch, + ShapeMMAThreadBlock, + ShapeMMAWarp, + ShapeMMAOp, + EpilogueOp, + SwizzleThreadBlock, + NumStages, + Alignment, + Alignment, + false, + cutlass::arch::OpMultiplyAddFastF32>; + +// Gemm_1xTF32 +using Gemm_1xTF32 = cutlass::gemm::device::Gemm< + float, + LayoutInputA, + float, + LayoutInputB, + float, + LayoutOutput, + float, + MMAOp, + SmArch, + ShapeMMAThreadBlock, + ShapeMMAWarp, + ShapeMMAOp, + EpilogueOp, + SwizzleThreadBlock, + NumStages, + Alignment, + Alignment, + false, + cutlass::arch::OpMultiplyAdd>; + +// Gemm_F64 +using Gemm_F64 = cutlass::reference::device::Gemm< + double, + LayoutInputA, + double, + LayoutInputB, + double, + LayoutOutput, + double, + double>; + +// Gemm_F32 +using Gemm_F32 = cutlass::reference::device::Gemm< + float, + LayoutInputA, + float, + LayoutInputB, + float, + LayoutOutput, + float, + float>; + +bool run(Options &options) { + + // Create a tuple of problem size for matrix multiplication + cutlass::gemm::GemmCoord problem_size = options.problem_size; + + //////////////////////////////////////////////////////////////////////////////// + /// 1. Initialize F32 Precision input tensors using CUTLASS helper functions + //////////////////////////////////////////////////////////////////////////////// + cutlass::HostTensor tensor_a_F32(problem_size.mk()); // <- Create matrix A with dimensions M x K + cutlass::HostTensor tensor_b_F32(problem_size.kn()); // <- Create matrix B with dimensions K x N + cutlass::HostTensor tensor_c_F32(problem_size.mn()); // <- Create matrix C with dimensions M x N + cutlass::HostTensor tensor_d_F32(problem_size.mn()); // <- Create matrix D with dimensions M x N + + if (options.rand_mode == "uniform") { + const float min = -1; + const float max = 1; + // Fill input and output matrices on host using CUTLASS helper functions + cutlass::reference::host::TensorFillRandomUniform( + tensor_a_F32.host_view(), + options.seed, + double(max), + double(min)); // <- Fill matrix A on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_b_F32.host_view(), + options.seed, + double(max), + double(min)); // <- Fill matrix B on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_c_F32.host_view(), + options.seed, + double(max), + double(min)); // <- Fill matrix C on host with uniform-distribution random data + } else if (options.rand_mode == "gauss") { + // Fill input and output matrices on host using CUTLASS helper functions + cutlass::reference::host::TensorFillRandomGaussian( + tensor_a_F32.host_view(), + options.seed, + double(0), + double(5)); // <- Fill matrix A on host with gaussian-distribution random data + cutlass::reference::host::TensorFillRandomGaussian( + tensor_b_F32.host_view(), + options.seed, + double(0), + double(5)); // <- Fill matrix B on host with gaussian-distribution random data + cutlass::reference::host::TensorFillRandomGaussian( + tensor_c_F32.host_view(), + options.seed, + double(0), + double(5)); // <- Fill matrix C on host with gaussian-distribution random data + } + cutlass::reference::host::TensorFill( + tensor_d_F32.host_view()); // <- fill matrix D on host with zeros + + // Copy data from host to GPU + tensor_a_F32.sync_device(); + tensor_b_F32.sync_device(); + tensor_c_F32.sync_device(); + tensor_d_F32.sync_device(); + + //////////////////////////////////////////////////////////////////////////////// + /// 2. Initialize F64 tensors using the same values used for F32 + //////////////////////////////////////////////////////////////////////////////// + // Gemm input operands (A, B, C) + cutlass::HostTensor tensor_a_F64(problem_size.mk()); // <- Create matrix A with dimensions M x K + cutlass::HostTensor tensor_b_F64(problem_size.kn()); // <- Create matrix B with dimensions K x N + cutlass::HostTensor tensor_c_F64(problem_size.mn()); // <- Create matrix C with dimensions M x N + + // Gemm output (D) for GEMM_F64 + cutlass::HostTensor tensor_d_F64(problem_size.mn()); // <- Create matrix D with dimensions M x N + // Gemm output (D) for GEMM_3xTF32 + cutlass::HostTensor tensor_d_3xTF32(problem_size.mn()); // <- Create matrix D with dimensions M x N + // Gemm output (D) for GEMM_1xTF32 + cutlass::HostTensor tensor_d_1xTF32(problem_size.mn()); // <- Create matrix D with dimensions M x N + + // Copy values from the DP tensors + cutlass::reference::host::TensorCopy(tensor_a_F64.host_view(), tensor_a_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_b_F64.host_view(), tensor_b_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_c_F64.host_view(), tensor_c_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_d_F64.host_view(), tensor_d_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_d_3xTF32.host_view(), tensor_d_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_d_1xTF32.host_view(), tensor_d_F32.host_view()); + + // Copy data from host to GPU + tensor_a_F64.sync_device(); + tensor_b_F64.sync_device(); + tensor_c_F64.sync_device(); + tensor_d_F64.sync_device(); + tensor_d_3xTF32.sync_device(); + tensor_d_1xTF32.sync_device(); + + // Initialize alpha and beta for dot product computation + float alpha = float(options.alpha); + float beta = float(options.beta); + + // Split K dimension into 1 partitions + int split_k_slices = 1; + + //////////////////////////////////////////////////////////////////////////////// + /// 3. Run 3xTF32 kernel within a profiling loop + //////////////////////////////////////////////////////////////////////////////// + // Create a tuple of gemm kernel arguments. This is later passed as arguments to launch + // instantiated CUTLASS kernel + typename Gemm_3xTF32::Arguments arguments_3xtf32{problem_size, // <- problem size of matrix multiplication + tensor_a_F32.device_ref(), // <- reference to matrix A on device + tensor_b_F32.device_ref(), // <- reference to matrix B on device + tensor_c_F32.device_ref(), // <- reference to matrix C on device + tensor_d_3xTF32.device_ref(), // <- reference to matrix D on device + {alpha, beta}, // <- tuple of alpha and beta + split_k_slices}; // <- k-dimension split factor + + // Using the arguments, query for extra workspace required for matrix multiplication computation + size_t workspace_size_3xtf32 = Gemm_3xTF32::get_workspace_size(arguments_3xtf32); + + // Allocate workspace memory + cutlass::device_memory::allocation workspace_3xtf32(workspace_size_3xtf32); + + // Instantiate CUTLASS kernel depending on templates + Gemm_3xTF32 gemm_op_3xTF32; + + // Check the problem size is supported or not + cutlass::Status status_3xtf32 = gemm_op_3xTF32.can_implement(arguments_3xtf32); + CUTLASS_CHECK(status_3xtf32); + + // Initialize CUTLASS kernel with arguments and workspace pointer + status_3xtf32 = gemm_op_3xTF32.initialize(arguments_3xtf32, workspace_3xtf32.get()); + CUTLASS_CHECK(status_3xtf32); + + // Result structure + Result result; + + // + // Construct events + // + + cudaEvent_t events[2]; + + for (auto & event : events) { + result.error = cudaEventCreate(&event); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventCreate() failed: " << cudaGetErrorString(result.error) << std::endl; + return false; + } + } + + // Record an event at the start of a series of GEMMs + result.error = cudaEventRecord(events[0]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return false; + } + + // + // Run profiling loop + // + + for (int iter = 0; iter < options.iterations; ++iter) { + // Launch initialized CUTLASS kernel + status_3xtf32 = gemm_op_3xTF32(); + CUTLASS_CHECK(status_3xtf32); + } + + // + // Stop profiling loop + // + + // Record an event when the GEMMs are complete + result.error = cudaEventRecord(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return false; + } + + // Wait for work on the device to complete. + result.error = cudaEventSynchronize(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventSynchronize() failed: " << cudaGetErrorString(result.error) << std::endl; + return false; + } + + // Measure elapsed runtime + float runtime_ms = 0; + result.error = cudaEventElapsedTime(&runtime_ms, events[0], events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result.error) << std::endl; + return false; + } + + // Compute average runtime and GFLOPs. + result.m = problem_size.m(); + result.n = problem_size.n(); + result.k = problem_size.k(); + result.runtime_ms = double(runtime_ms) / double(options.iterations); + result.gflops = options.gflops(result.runtime_ms / 1000.0); + + // Cleanup + for (auto event : events) { + (void)cudaEventDestroy(event); + } + + tensor_d_3xTF32.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + /// 4. Run TF32 kernel without profiling loop + //////////////////////////////////////////////////////////////////////////////// + // Create a tuple of gemm kernel arguments. This is later passed as arguments to launch + // instantiated CUTLASS kernel + typename Gemm_1xTF32::Arguments arguments_1xtf32{problem_size, // <- problem size of matrix multiplication + tensor_a_F32.device_ref(), // <- reference to matrix A on device + tensor_b_F32.device_ref(), // <- reference to matrix B on device + tensor_c_F32.device_ref(), // <- reference to matrix C on device + tensor_d_1xTF32.device_ref(), // <- reference to matrix D on device + {alpha, beta}, // <- tuple of alpha and beta + split_k_slices}; // <- k-dimension split factor + + // Using the arguments, query for extra workspace required for matrix multiplication computation + size_t workspace_size_1xtf32 = Gemm_1xTF32::get_workspace_size(arguments_1xtf32); + + // Allocate workspace memory + cutlass::device_memory::allocation workspace_1xtf32(workspace_size_1xtf32); + + // Instantiate CUTLASS kernel depending on templates + Gemm_1xTF32 gemm_op_1xtf32; + + // Check the problem size is supported or not + cutlass::Status status_1xtf32 = gemm_op_1xtf32.can_implement(arguments_1xtf32); + CUTLASS_CHECK(status_1xtf32); + + // Initialize CUTLASS kernel with arguments and workspace pointer + status_1xtf32 = gemm_op_1xtf32.initialize(arguments_1xtf32, workspace_1xtf32.get()); + CUTLASS_CHECK(status_1xtf32); + + // Launch initialized CUTLASS kernel + status_1xtf32 = gemm_op_1xtf32(); + CUTLASS_CHECK(status_1xtf32); + + tensor_d_1xTF32.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + // Run reference kernel (F64) + //////////////////////////////////////////////////////////////////////////////// + + // Create instantiation for device reference gemm kernel + Gemm_F64 gemm_f64; + + // Launch device reference gemm kernel + gemm_f64(problem_size, + alpha, + tensor_a_F64.device_ref(), + tensor_b_F64.device_ref(), + beta, + tensor_c_F64.device_ref(), + tensor_d_F64.device_ref()); + + // Wait for kernels to finish + cudaDeviceSynchronize(); + + // Copy output data from CUTLASS and reference kernel to host for comparison + tensor_d_F64.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + // Run reference kernel (F32) + //////////////////////////////////////////////////////////////////////////////// + + // Create instantiation for device reference gemm kernel + Gemm_F32 gemm_f32; + + // Launch device reference gemm kernel + gemm_f32(problem_size, + alpha, + tensor_a_F32.device_ref(), + tensor_b_F32.device_ref(), + beta, + tensor_c_F32.device_ref(), + tensor_d_F32.device_ref()); + + // Wait for kernels to finish + cudaDeviceSynchronize(); + + // Copy output data from CUTLASS and reference kernel to host for comparison + tensor_d_F32.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + /////// Compute l2 norms + //////////////////////////////////////////////////////////////////////////////// + + // l2 norm 3xTF32 vs F64 + cutlass::HostTensor tensor_d_3xTF32_in_F64(problem_size.mn()); + cutlass::reference::host::TensorCopy(tensor_d_3xTF32_in_F64.host_view(), tensor_d_3xTF32.host_view()); + + result.l2_norm_3xtf32_vs_fp64 = cutlass::reference::host::TensorRelativeErrorMetric( + tensor_d_3xTF32_in_F64.host_view(), tensor_d_F64.host_view()); + + // l2 norm 1xTF32 vs F64 + cutlass::HostTensor tensor_d_1xTF32_in_F64(problem_size.mn()); + cutlass::reference::host::TensorCopy(tensor_d_1xTF32_in_F64.host_view(), tensor_d_1xTF32.host_view()); + + result.l2_norm_1xtf32_vs_fp64 = cutlass::reference::host::TensorRelativeErrorMetric( + tensor_d_1xTF32_in_F64.host_view(), tensor_d_F64.host_view()); + + // l2 norm F32 vs F64 + cutlass::HostTensor tensor_d_F32_in_F64(problem_size.mn()); + cutlass::reference::host::TensorCopy(tensor_d_F32_in_F64.host_view(), tensor_d_F32.host_view()); + + result.l2_norm_fp32_vs_fp64 = cutlass::reference::host::TensorRelativeErrorMetric( + tensor_d_F32_in_F64.host_view(), tensor_d_F64.host_view()); + + results.push_back(result); + + /////////////////////////////////////////////////////////////////////////////// + + // Check if output from CUTLASS kernel and reference kernel are equal or not + + std::cout << std::fixed; + std::cout.precision(4); + std::cout << "Runtime: " << result.runtime_ms << " ms" << std::endl; + std::cout.precision(2); + std::cout << "GFLOPs: " << result.gflops << std::endl; + std::cout << "Normalized L2 norm of" << std::endl; + std::cout.precision(8); + std::cout << std::scientific + << " - 3xTF32 error with FP64 reference : " << result.l2_norm_3xtf32_vs_fp64 << std::endl + << " - 1xTF32 error with FP64 reference : " << result.l2_norm_1xtf32_vs_fp64 << std::endl + << " - FP32 error with FP64 reference : " << result.l2_norm_fp32_vs_fp64 << std::endl; + + return true; +} + +int main(int argc, const char **argv) { + + bool notSupported = false; + + // Ampere Tensor Core operations exposed with mma.sync and ldmatrix are first available + // in CUDA 11.0. + // + // CUTLASS must be compiled with CUDA 11.0 Toolkit to run these examples. + if (!(__CUDACC_VER_MAJOR__ >= 11)) { + std::cerr << "Ampere Tensor Core operations must be compiled with CUDA 11.0 Toolkit or later." << std::endl; + notSupported = true; + } + + cudaDeviceProp props; + + cudaError_t error = cudaGetDeviceProperties(&props, 0); + if (error != cudaSuccess) { + std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl; + return false; + } + + if (!((props.major * 10 + props.minor) >= 80)) { + std::cerr << "Ampere Tensor Core operations must be run on a machine with compute capability at least 80." + << std::endl; + notSupported = true; + } + + if (notSupported) { + // Returning zero so this test passes on older Toolkits. Its actions are no-op. + return 0; + } + + Options options; + options.parse(argc, argv); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + bool result = true; + + if (options.benchmark) { + for (int k = 4; k <= 65536; k *= 2) { + + options.problem_size[2] = k; + + printf("Gemm problem size: %d x %d x %d\n", \ + options.problem_size.m(), options.problem_size.n(), options.problem_size.k()); + + if (!options.valid()) { + std::cerr << "Invalid problem." << std::endl; + return -1; + } + + result &= run(options); + } + } else { + // Execute one problem size + if (!options.valid()) { + std::cerr << "Invalid problem." << std::endl; + return -1; + } + + result = run(options); + } + + if (!result) return -1; + + std::cout << std::endl << "CSV results" << std::endl; + Result::print_csv_header(); + for(auto &r : results) + r.print_csv_row(); + + return 0; +} diff --git a/examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/CMakeLists.txt b/examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/CMakeLists.txt new file mode 100644 index 00000000..fb2b4f91 --- /dev/null +++ b/examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/CMakeLists.txt @@ -0,0 +1,27 @@ +# Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without modification, are permitted +# provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright notice, this list of +# conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright notice, this list of +# conditions and the following disclaimer in the documentation and/or other materials +# provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used +# to endorse or promote products derived from this software without specific prior written +# permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR +# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; +# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +cutlass_example_add_executable( + 27_ampere_3xtf32_fast_accurate_tensorop_gemm + 27_ampere_3xtf32_fast_accurate_tensorop_gemm.cu + ) + diff --git a/examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/CMakeLists.txt b/examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/CMakeLists.txt new file mode 100644 index 00000000..693e7c9b --- /dev/null +++ b/examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/CMakeLists.txt @@ -0,0 +1,27 @@ +# Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without modification, are permitted +# provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright notice, this list of +# conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright notice, this list of +# conditions and the following disclaimer in the documentation and/or other materials +# provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used +# to endorse or promote products derived from this software without specific prior written +# permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR +# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; +# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +cutlass_example_add_executable( + 28_ampere_3xtf32_fast_accurate_tensorop_fprop + ampere_3xtf32_fast_accurate_tensorop_fprop.cu + ) diff --git a/examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/ampere_3xtf32_fast_accurate_tensorop_fprop.cu b/examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/ampere_3xtf32_fast_accurate_tensorop_fprop.cu new file mode 100644 index 00000000..6208f1a7 --- /dev/null +++ b/examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/ampere_3xtf32_fast_accurate_tensorop_fprop.cu @@ -0,0 +1,815 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/** + +This example adopts example 16 to use 3xTF32 to bring FP32 accuracy with 2x performance +compared with CUDA Cores. See example 27 for the trick of 3xTF32. +*/ + +#include +#include + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/device/gemm.h" +#include "cutlass/conv/kernel/default_conv2d_fprop.h" +#include "cutlass/conv/device/implicit_gemm_convolution.h" + +#include "cutlass/util/command_line.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/reference/device/convolution.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/convolution.h" +#include "cutlass/util/reference/host/error_metrics.h" +#include "cutlass/util/tensor_view_io.h" + +#include "helper.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// The code section below describes datatype for input, output tensors and computation between +// elements +using ElementAccumulator = float; // Data type of accumulator +using ElementComputeEpilogue = float; // Data type of epilogue computation (alpha, beta) +using ElementInputA = float; // Data type of elements in input tensor +using ElementInputB = float; // Data type of elements in input tensor +using ElementOutput = float; // Data type of elements in output tensor + +using LayoutInputA = cutlass::layout::TensorNHWC; +using LayoutInputB = cutlass::layout::TensorNHWC; +using LayoutOutput = cutlass::layout::TensorNHWC; + +// This code section describes whether you want to use tensor cores or regular SIMT cores on GPU SM +using MMAOp = cutlass::arch::OpClassTensorOp; + +// This code section describes CUDA SM architecture number +using SmArch = cutlass::arch::Sm80; + +// This code section describes the tile size a thread block will compute +using ThreadblockShape = cutlass::gemm::GemmShape<128, 64, 16>; // Threadblock tile shape + +// This code section describes tile size a warp will compute +using WarpShape = cutlass::gemm::GemmShape<64, 32, 16>; // Warp tile shape + +// This code section describes the size of MMA op +using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; // TensorCore instruction shape + +// This code section describes how threadblocks are scheduled on GPU +using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; + +// Number of pipelines you want to use +constexpr int NumStages = 3; + +// This code section describe iterator algorithm selected is Analytic or Optimized +static cutlass::conv::IteratorAlgorithm const IteratorAlgorithm = cutlass::conv::IteratorAlgorithm::kOptimized; + +// This code section describes the epilogue part of the kernel, we use default value +using EpilogueOp = cutlass::epilogue::thread::LinearCombination< + ElementOutput, // Data type of output matrix. + 128 / cutlass::sizeof_bits::value, // The number of elements per vectorized. + // memory access. This becomes the vector width of + // math instructions in the epilogue too. + ElementAccumulator, // Data type of accumulator + ElementComputeEpilogue>; // Data type for alpha/beta in linear combination + +// 3xTF32 Fprop +using Conv2dFpropKernel_3xTF32 = typename cutlass::conv::kernel::DefaultConv2dFprop< + ElementInputA, LayoutInputA, + ElementInputB, LayoutInputB, + ElementOutput, LayoutOutput, + ElementAccumulator, + MMAOp, + SmArch, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOp, + SwizzleThreadBlock, + NumStages, + // Only thing needs to be changed from normal Fprop + cutlass::arch::OpMultiplyAddFastF32, + IteratorAlgorithm +>::Kernel; + +// 1xTF32 Fprop +using Conv2dFpropKernel_1xTF32 = typename cutlass::conv::kernel::DefaultConv2dFprop< + ElementInputA, LayoutInputA, + ElementInputB, LayoutInputB, + ElementOutput, LayoutOutput, + ElementAccumulator, + MMAOp, + SmArch, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOp, + SwizzleThreadBlock, + NumStages, + cutlass::arch::OpMultiplyAdd, + IteratorAlgorithm +>::Kernel; + +using ImplicitGemm_3xTF32 = cutlass::conv::device::ImplicitGemmConvolution; +using ImplicitGemm_1xTF32 = cutlass::conv::device::ImplicitGemmConvolution; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Command line options parsing +struct Options { + + bool help; + cutlass::Tensor4DCoord input_size; + cutlass::Tensor4DCoord filter_size; + cutlass::Tensor4DCoord padding; + cutlass::MatrixCoord conv_stride; + cutlass::MatrixCoord dilation; + int iterations; + bool save_workspace; + ElementComputeEpilogue alpha; + ElementComputeEpilogue beta; + bool benchmark; + std::string tag; + + Options(): + help(false), + input_size(1, 32, 32, 32), + filter_size(32, 3, 3, 32), + padding(1, 1, 1, 1), + conv_stride(1, 1), + dilation(1, 1), + iterations(20), + save_workspace(false), + alpha(1), + beta(0), + benchmark(false) { } + + // Verify the problem size is compatible with the CUTLASS Convolution implementation. + bool valid() { + + // + // CUTLASS attempts to load 128b vectors of cutlass::half_t (F16) elements. Consequently, + // all pointers, strides, and tensor extents must be divisible by 8 elements. + // + int const kAlignment = 4; + + if ((input_size.c() % kAlignment) || + (filter_size.n() % kAlignment)) { + + // misaligned tensors + return false; + } + + // Invalid padding + if ((padding.h() != filter_size.h() / 2) || + (padding.w() != filter_size.w() / 2)) { + + return false; + } + + return true; + } + + /// Updates input and filter sizes + void update( + cutlass::Tensor4DCoord input_size, + cutlass::Tensor4DCoord filter_size) { + + this->input_size = input_size; + this->filter_size = filter_size; + + padding.n() = filter_size.h() / 2; + padding.h() = filter_size.h() / 2; + padding.w() = filter_size.w() / 2; + padding.c() = filter_size.w() / 2; + } + + // 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; + } + + if (cmd.check_cmd_line_flag("save-workspace")) { + save_workspace = true; + } + + if (cmd.check_cmd_line_flag("benchmark")) { + benchmark = true; + } + + cmd.get_cmd_line_argument("n", input_size.n()); + cmd.get_cmd_line_argument("h", input_size.h()); + cmd.get_cmd_line_argument("w", input_size.w()); + cmd.get_cmd_line_argument("c", input_size.c()); + + cmd.get_cmd_line_argument("k", filter_size.n()); + cmd.get_cmd_line_argument("r", filter_size.h()); + cmd.get_cmd_line_argument("s", filter_size.w()); + filter_size.c() = input_size.c(); + + cmd.get_cmd_line_argument("alpha", alpha); + cmd.get_cmd_line_argument("beta", beta); + + cmd.get_cmd_line_argument("iterations", iterations); + cmd.get_cmd_line_argument("tag", tag); + + if (filter_size.h() == 3 && filter_size.w() == 3) { + padding = {1, 1, 1, 1}; + } + else { + filter_size.h() = 1; + filter_size.w() = 1; + padding = {0, 0, 0, 0}; + } + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "28_ampere_3xtf32_fast_accurate_tensorop_fprop example\n\n" + << " This example uses Ampere's Tensor Core operators on F16 data types to compute\n" + << " forward convolution on tensors of layout NHWC.\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement.\n\n" + << " --n Input tensor extent N\n" + << " --h Input tensor extent H\n" + << " --w Input tensor extent W\n" + << " --c Input tensor extent C\n" + << " --k Filter extent K\n" + << " --r Filter extent R\n" + << " --s Filter extent S\n\n" + << " --alpha Epilogue scalar alpha\n" + << " --beta Epilogue scalar beta\n\n" + << " --benchmark If set (true), performance benchmarking on several layers and batch-size.\n" + << " --iterations Number of profiling iterations to perform.\n" + << " --save-workspace If set, workspace is written to a text file.\n" + << " --tag String to replicate across the first column in the results table\n"; + + out << "\n\nExamples:\n\n" + << "$ ./examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/28_ampere_3xtf32_fast_accurate_tensorop_fprop --n=32 --h=224 --w=224 --c=128 --k=256 --r=1 --s=1\n\n" + << "$ ./examples/28_ampere_3xtf32_fast_accurate_tensorop_fprop/28_ampere_3xtf32_fast_accurate_tensorop_fprop --n=1 --h=224 --w=224 --c=32 --k=32 --r=3 --s=3 --ref-check\n\n"; + + return out; + } + + /// Computes the output tensor size (NPQK) + cutlass::Tensor4DCoord output_size() const { + return cutlass::Tensor4DCoord( + input_size.n(), + (input_size.h() + padding.n() + padding.h() - filter_size.h()) / conv_stride.row() + 1, + (input_size.w() + padding.w() + padding.c() - filter_size.w()) / conv_stride.column() + 1, + filter_size.n()); + } + + /// Compute performance in GFLOP/s + double gflops(double runtime_s) const { + + // Number of multiply-adds = NPQK * CRS + int64_t fmas = output_size().product() * int64_t(filter_size.h() * filter_size.w() * filter_size.c()); + + // Two flops per multiply-add + return 2.0 * double(fmas) / double(1.0e9) / runtime_s; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +struct Result { + double runtime_ms; + double gflops; + cutlass::Status status; + cudaError_t error; + + double l2_norm_3xtf32_vs_fp64; + double l2_norm_1xtf32_vs_fp64; + double l2_norm_fp32_vs_fp64; + + Result(): + runtime_ms(0), + gflops(0), + status(cutlass::Status::kSuccess), + error(cudaSuccess), + l2_norm_3xtf32_vs_fp64(0), + l2_norm_1xtf32_vs_fp64(0), + l2_norm_fp32_vs_fp64(0) { } + + static std::ostream & print_header(std::ostream &out, Options const &options) { + + if (!options.tag.empty()) { + out << "Name,"; + } + + out << "Layer,N,H,W,C,K,R,S,Runtime,GFLOPs,3xTF32_vs_FP64,1xTF32_vs_FP64,FP32_vs_FP64"; + + return out; + } + + std::ostream & print(std::ostream &out, int idx, Options const &options) { + + if (!options.tag.empty()) { + out << options.tag << ","; + } + + out + << "conv_" << idx << "," + << options.input_size.n() << "," + << options.input_size.h() << "," + << options.input_size.w() << "," + << options.input_size.c() << "," + << options.filter_size.n() << "," + << options.filter_size.h() << "," + << options.filter_size.w() << "," + << runtime_ms << "," + << gflops << "," + << l2_norm_3xtf32_vs_fp64 << "," + << l2_norm_1xtf32_vs_fp64 << "," + << l2_norm_fp32_vs_fp64; + + return out; + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Runs one benchmark +Result profile_convolution(Options const &options) { + + Result result; + + //////////////////////////////////////////////////////////////////////////////// + /// 1. Initialize F32 Precision input tensors using CUTLASS helper functions + //////////////////////////////////////////////////////////////////////////////// + + // + // Allocate host-device tensors using the CUTLASS Utilities. + // + + cutlass::HostTensor tensor_a_F32(options.input_size); + cutlass::HostTensor tensor_b_F32(options.filter_size); + cutlass::HostTensor tensor_c_F32(options.output_size()); + cutlass::HostTensor tensor_d_F32(options.output_size()); + + // + // Initialize tensors + // + + // Fill tensor A on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_a_F32.host_view(), + 1, + ElementInputA(7), + ElementInputA(-8)); + + // Fill tensor B on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_b_F32.host_view(), + 1, + ElementInputB(7), + ElementInputB(-8)); + + // Fill tensor C on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_c_F32.host_view(), + 1, + ElementInputB(7), + ElementInputB(-8)); + + // Fill tensor D on host with zeros + cutlass::reference::host::TensorFill( + tensor_d_F32.host_view()); + + // Copy data from host to GPU + tensor_a_F32.sync_device(); + tensor_b_F32.sync_device(); + tensor_c_F32.sync_device(); + tensor_d_F32.sync_device(); + + //////////////////////////////////////////////////////////////////////////////// + /// 2. Initialize F32 Precision input tensors using CUTLASS helper functions + //////////////////////////////////////////////////////////////////////////////// + + // + // Allocate host-device tensors using the CUTLASS Utilities. + // + + cutlass::HostTensor tensor_a_F64(options.input_size); + cutlass::HostTensor tensor_b_F64(options.filter_size); + cutlass::HostTensor tensor_c_F64(options.output_size()); + + cutlass::HostTensor tensor_d_F64(options.output_size()); + cutlass::HostTensor tensor_d_3xTF32(options.output_size()); + cutlass::HostTensor tensor_d_1xTF32(options.output_size()); + + // Copy values from the DP tensors + cutlass::reference::host::TensorCopy(tensor_a_F64.host_view(), tensor_a_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_b_F64.host_view(), tensor_b_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_c_F64.host_view(), tensor_c_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_d_F64.host_view(), tensor_d_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_d_3xTF32.host_view(), tensor_d_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_d_1xTF32.host_view(), tensor_d_F32.host_view()); + + // Copy data from host to GPU + tensor_a_F64.sync_device(); + tensor_b_F64.sync_device(); + tensor_c_F64.sync_device(); + tensor_d_F64.sync_device(); + tensor_d_3xTF32.sync_device(); + tensor_d_1xTF32.sync_device(); + + // + // Define arguments for CUTLASS Convolution + // + + cutlass::conv::Mode mode = cutlass::conv::Mode::kCrossCorrelation; + + // Split K dimension into 1 partitions + int split_k_slices = 1; + + // Construct Conv2dProblemSize with user defined output size + cutlass::conv::Conv2dProblemSize problem_size( + options.input_size, + options.filter_size, + options.padding, + options.conv_stride, + options.dilation, + options.output_size(), + mode, + split_k_slices + ); + + //////////////////////////////////////////////////////////////////////////////// + /// 3. Run 3xTF32 kernel within a profiling loop + //////////////////////////////////////////////////////////////////////////////// + + // Construct ImplicitGemm::Argument structure with conv2d + // problem size, data pointers, and epilogue values + typename ImplicitGemm_3xTF32::Arguments arguments_3xTF32{ + problem_size, + tensor_a_F32.device_ref(), + tensor_b_F32.device_ref(), + tensor_c_F32.device_ref(), + tensor_d_3xTF32.device_ref(), + {options.alpha, options.beta}, + }; + + // + // Initialize CUTLASS Convolution + // + + ImplicitGemm_3xTF32 implicit_gemm_op_3xTF32; + + size_t workspace_size_3xTF32 = implicit_gemm_op_3xTF32.get_workspace_size(arguments_3xTF32); + + // Allocate workspace memory + cutlass::device_memory::allocation workspace_3xTF32(workspace_size_3xTF32); + + result.status = implicit_gemm_op_3xTF32.can_implement(arguments_3xTF32); + CUTLASS_CHECK(result.status); + + result.status = implicit_gemm_op_3xTF32.initialize(arguments_3xTF32, workspace_3xTF32.get()); + CUTLASS_CHECK(result.status); + + // + // Launch initialized CUTLASS kernel + // + result.status = implicit_gemm_op_3xTF32(); + + CUTLASS_CHECK(result.status); + + // + // Performance measurement + // + + cudaEvent_t events[2]; + + for (auto & event : events) { + result.error = cudaEventCreate(&event); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventCreate() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + } + + // Record an event at the start of a series of convolution operations. + result.error = cudaEventRecord(events[0]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Launch a sequence of implicit GEMM operations on the device + for (int iteration = 0; iteration < options.iterations; ++iteration) { + result.status = implicit_gemm_op_3xTF32(); + CUTLASS_CHECK(result.status); + } + + // Record an event when the convolutions have been launched. + result.error = cudaEventRecord(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Wait for work on the device to complete. + result.error = cudaEventSynchronize(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventSynchronize() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Measure elapsed runtime + float runtime_ms = 0; + result.error = cudaEventElapsedTime(&runtime_ms, events[0], events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result.error) << std::endl; + return result; + } + + // Print average runtime and GFLOPs. + result.runtime_ms = double(runtime_ms) / double(options.iterations); + result.gflops = options.gflops(result.runtime_ms / 1000.0); + + // Cleanup + for (auto event : events) { + (void)cudaEventDestroy(event); + } + + tensor_d_3xTF32.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + /// 4. Run 1xTF32 kernel within a profiling loop + //////////////////////////////////////////////////////////////////////////////// + + // Construct ImplicitGemm::Argument structure with conv2d + // problem size, data pointers, and epilogue values + typename ImplicitGemm_1xTF32::Arguments arguments_1xTF32{ + problem_size, + tensor_a_F32.device_ref(), + tensor_b_F32.device_ref(), + tensor_c_F32.device_ref(), + tensor_d_1xTF32.device_ref(), + {options.alpha, options.beta}, + }; + + // + // Initialize CUTLASS Convolution + // + + ImplicitGemm_1xTF32 implicit_gemm_op_1xTF32; + + size_t workspace_size_1xTF32 = implicit_gemm_op_1xTF32.get_workspace_size(arguments_1xTF32); + + // Allocate workspace memory + cutlass::device_memory::allocation workspace_1xTF32(workspace_size_1xTF32); + + result.status = implicit_gemm_op_1xTF32.can_implement(arguments_1xTF32); + CUTLASS_CHECK(result.status); + + result.status = implicit_gemm_op_1xTF32.initialize(arguments_1xTF32, workspace_1xTF32.get()); + CUTLASS_CHECK(result.status); + + // + // Launch initialized CUTLASS kernel + // + result.status = implicit_gemm_op_1xTF32(); + + CUTLASS_CHECK(result.status); + + tensor_d_1xTF32.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + // Run reference kernel (F64) + //////////////////////////////////////////////////////////////////////////////// + + cutlass::reference::device::Conv2d< + double, + LayoutInputA, + double, + LayoutInputB, + double, + LayoutOutput, + double, + double + >( + cutlass::conv::Operator::kFprop, + problem_size, + tensor_a_F64.device_ref(), + tensor_b_F64.device_ref(), + tensor_c_F64.device_ref(), + tensor_d_F64.device_ref(), + options.alpha, + options.beta); + + // Wait for kernels to finish + cudaDeviceSynchronize(); + + // Copy output data from CUTLASS and reference kernel to host for comparison + tensor_d_F64.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + // Run reference kernel (F32) + //////////////////////////////////////////////////////////////////////////////// + + cutlass::reference::device::Conv2d< + float, + LayoutInputA, + float, + LayoutInputB, + float, + LayoutOutput, + float, + float + >( + cutlass::conv::Operator::kFprop, + problem_size, + tensor_a_F32.device_ref(), + tensor_b_F32.device_ref(), + tensor_c_F32.device_ref(), + tensor_d_F32.device_ref(), + options.alpha, + options.beta); + + // Wait for kernels to finish + cudaDeviceSynchronize(); + + // Copy output data from CUTLASS and reference kernel to host for comparison + tensor_d_F32.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + /////// Compute l2 norms + //////////////////////////////////////////////////////////////////////////////// + + // l2 norm 3xTF32 vs F64 + cutlass::HostTensor tensor_d_3xTF32_in_F64(options.output_size()); + cutlass::reference::host::TensorCopy(tensor_d_3xTF32_in_F64.host_view(), tensor_d_3xTF32.host_view()); + + result.l2_norm_3xtf32_vs_fp64 = cutlass::reference::host::TensorRelativeErrorMetric( + tensor_d_3xTF32_in_F64.host_view(), tensor_d_F64.host_view()); + + // l2 norm 1xTF32 vs F64 + cutlass::HostTensor tensor_d_1xTF32_in_F64(options.output_size()); + cutlass::reference::host::TensorCopy(tensor_d_1xTF32_in_F64.host_view(), tensor_d_1xTF32.host_view()); + + result.l2_norm_1xtf32_vs_fp64 = cutlass::reference::host::TensorRelativeErrorMetric( + tensor_d_1xTF32_in_F64.host_view(), tensor_d_F64.host_view()); + + // l2 norm F32 vs F64 + cutlass::HostTensor tensor_d_F32_in_F64(options.output_size()); + cutlass::reference::host::TensorCopy(tensor_d_F32_in_F64.host_view(), tensor_d_F32.host_view()); + + result.l2_norm_fp32_vs_fp64 = cutlass::reference::host::TensorRelativeErrorMetric( + tensor_d_F32_in_F64.host_view(), tensor_d_F64.host_view()); + + /////////////////////////////////////////////////////////////////////////////// + + if (options.save_workspace) { + + std::stringstream ss; + + ss << "28_ampere_3xtf32_fast_accurate_tensorop_fprop_" + << options.input_size.n() << "x" << options.input_size.h() << "x" << options.input_size.w() << "x" << options.input_size.c() + << "_" + << options.filter_size.n() << "x" << options.filter_size.h() << "x" << options.filter_size.w() << "x" << options.filter_size.c() + << ".dat"; + + std::ofstream output_workspace(ss.str()); + + output_workspace + << "Input = \n" << tensor_a_F32.host_view() << "\n\n" + << "Filters = \n" << tensor_b_F32.host_view() << "\n\n"; + + output_workspace << "TF32x3 = \n" << tensor_d_3xTF32.host_view() << std::endl; + output_workspace << "TF32x1 = \n" << tensor_d_1xTF32.host_view() << std::endl; + output_workspace << "FP32 = \n" << tensor_d_F32.host_view() << std::endl; + output_workspace << "FP64 = \n" << tensor_d_F64.host_view() << "\n\n"; + + std::cout << "Results written to '" << ss.str() << "'." << std::endl; + } + + return result; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +int main(int argc, char const **args) { + + bool notSupported = false; + + // Ampere Tensor Core operations exposed with mma.sync are first available in CUDA 11.0. + // + // CUTLASS must be compiled with CUDA 11 Toolkit to run Conv2dFprop examples. + if (!(__CUDACC_VER_MAJOR__ > 11 || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 0))) { + std::cerr << "Ampere Tensor Core operations must be compiled with CUDA 11.0 Toolkit or later." << std::endl; + notSupported = true; + } + + cudaDeviceProp props; + CUDA_CHECK(cudaGetDeviceProperties(&props, 0)); + + if (!(props.major > 8 || (props.major == 8 && props.minor >= 0))) { + std::cerr << "Ampere Tensor Ops must be run on a machine with compute capability at least 80." + << std::endl; + notSupported = true; + } + + if (notSupported) { + return 0; + } + + Options options; + + options.parse(argc, args); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + if (options.benchmark) { + // Benchmark several layers + + int batch_sizes[] = {1, 32, 64, 128, 256}; + + struct Benchmark { + int h, w, c, k, r, s; + } layers[] = { + {56, 56, 64, 256, 1, 1}, + {56, 56, 64, 64, 1, 1}, + {56, 56, 64, 64, 3, 3}, + {56, 56, 256, 64, 1, 1}, + {56, 56, 256, 512, 1, 1}, + {56, 56, 256, 128, 1, 1}, + {28, 28, 128, 128, 3, 3}, + {28, 28, 128, 512, 1, 1}, + {28, 28, 512, 128, 1, 1}, + {28, 28, 512, 1024, 1, 1}, + {28, 28, 512, 256, 1, 1}, + {14, 14, 256, 256, 3, 3}, + {14, 14, 256, 1024, 1, 1}, + {14, 14, 1024, 256, 1, 1}, + {14, 14, 1024, 2048, 1, 1}, + {14, 14, 1024, 512, 1, 1}, + {7, 7, 512, 512, 3, 3}, + }; + + Result::print_header(std::cout, options) << std::endl; + + int idx = 1; + + for (auto const &layer : layers) { + for (auto N : batch_sizes) { + + options.update({N, layer.h, layer.w, layer.c}, {layer.k, layer.r, layer.s, layer.c}); + + Result result = profile_convolution(options); + result.print(std::cout, idx, options) << std::endl; + } + + ++idx; + } + } + else { + + // Execute one problem size + if (!options.valid()) { + std::cerr << "Invalid problem." << std::endl; + return -1; + } + + Result result = profile_convolution(options); + + Result::print_header(std::cout, options) << std::endl; + result.print(std::cout, 1, options) << std::endl; + } + + return 0; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm.cu b/examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm.cu new file mode 100644 index 00000000..08e85c61 --- /dev/null +++ b/examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm.cu @@ -0,0 +1,686 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/** + This example is almost the same as example 27 which uses 3xTF32 to run GEMM. The only + difference is that this example uses 3xtf32 on complex gemm. + + To enable this feature, the only change needs to make is to change OpMultiplyAddComplex + to OpMultiplyAddComplexFastF32. +*/ + +#include +#include +#include + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/device/gemm_complex.h" + +#include "cutlass/util/command_line.h" +#include "cutlass/util/host_tensor.h" + +#include "cutlass/util/reference/device/gemm_complex.h" +#include "cutlass/util/reference/host/tensor_reduce.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_norm.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/error_metrics.h" +#include "cutlass/util/tensor_view_io.h" + +#include "helper.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Result structure +struct Result { + + double runtime_ms; + double gflops; + cutlass::Status status; + cudaError_t error; + + int m, n, k; + double l2_norm_3xtf32_vs_fp64; + double l2_norm_1xtf32_vs_fp64; + double l2_norm_fp32_vs_fp64; + + // ctor + Result( + int m, int n, int k, + double runtime_ms, double gflops, + double l2_norm_3xtf32_vs_fp64, + double l2_norm_1xtf32_vs_fp64, + double l2_norm_fp32_vs_fp64) : + m(m), n(n), k(k), + runtime_ms(runtime_ms), gflops(gflops), + l2_norm_3xtf32_vs_fp64(l2_norm_3xtf32_vs_fp64), + l2_norm_1xtf32_vs_fp64(l2_norm_1xtf32_vs_fp64), + l2_norm_fp32_vs_fp64(l2_norm_fp32_vs_fp64) {} + + Result() {} + + // + // Methods + // + static void print_csv_header() { + std::cout << "M,N,K,Runtime(ms),GFLOPS,3xTF32_vs_FP64,1xTF32_vs_FP64,FP32_vs_FP64" << std::endl; + } + + void print_csv_row() { + std::cout << m << "," + << n << "," + << k << "," + << runtime_ms << "," + << gflops << "," + << l2_norm_3xtf32_vs_fp64 << "," + << l2_norm_1xtf32_vs_fp64 << "," + << l2_norm_fp32_vs_fp64 << std::endl; + } +}; + +std::vector results; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// Command line options parsing +struct Options { + + bool help; + + cutlass::gemm::GemmCoord problem_size; + float alpha; + float beta; + std::string rand_mode; + + int iterations; + int seed; + bool benchmark; + + Options(): + help(false), + problem_size({3456, 4096, 4096}), + iterations(20), + seed(1), + alpha(1), + beta(), + rand_mode("uniform"), + benchmark(false) { } + + bool valid() { + return true; + } + + // 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; + } + + cmd.get_cmd_line_argument("m", problem_size.m()); + cmd.get_cmd_line_argument("n", problem_size.n()); + cmd.get_cmd_line_argument("k", problem_size.k()); + + cmd.get_cmd_line_argument("alpha", alpha); + cmd.get_cmd_line_argument("beta", beta); + + cmd.get_cmd_line_argument("iterations", iterations); + cmd.get_cmd_line_argument("seed", seed); + cmd.get_cmd_line_argument("rand_mode", rand_mode); + + if (cmd.check_cmd_line_flag("benchmark")) { + benchmark = true; + } + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm example\n\n" + << " This example uses the CUTLASS Library to emulate FP32 complex GEMM computations with TF32 tensor cores.\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement.\n\n" + << " --m GEMM M dimension\n" + << " --n GEMM N dimension\n" + << " --k GEMM K dimension\n" + << " --alpha Epilogue scalar alpha\n" + << " --beta Epilogue scalar beta\n\n" + << " --rand_mode gauss / uniform*\n\n" + << " --seed Random number seed (1*)\n\n" + << " --iterations Number of profiling iterations to perform.\n\n" + << " --benchmark If set (true), performance benchmarking on several layers and batch-size.\n\n"; + + out << "\n\nExamples:\n\n" + << "$ ./examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_ampere_3xtf32_fast_accurate_complex_gemm --m=1024 --n=512 \\\n" + << " --alpha=2 --beta=0.707 \n\n"; + + return out; + } + + /// Compute performance in GFLOP/s + double gflops(double runtime_s) const { + + // Number of real-valued multiply-adds + int64_t fmas = problem_size.product(); + + // Two flops per multiply-add + return 2.0 * double(fmas) / double(1.0e9) / runtime_s; + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// The code section below describes matrix layout of input and output matrices. Column Major for +// Matrix A, Row Major for Matrix B and Row Major for Matrix C +using LayoutInputA = cutlass::layout::ColumnMajor; +using LayoutInputB = cutlass::layout::RowMajor; +using LayoutOutput = cutlass::layout::RowMajor; + +// This code section describes whether you want to use tensor cores or regular SIMT cores on GPU SM +using MMAOp = cutlass::arch::OpClassTensorOp; + +// This code section describes CUDA SM architecture number +using SmArch = cutlass::arch::Sm80; + +// This code section describes the tile size a thread block will compute +using ShapeMMAThreadBlock = + cutlass::gemm::GemmShape<64, 64, 16>; // <- threadblock tile M = 128, N = 128, K = 16 +// This code section describes tile size a warp will compute +using ShapeMMAWarp = cutlass::gemm::GemmShape<32, 32, 16>; // <- warp tile M = 64, N = 64, K = 16 +// This code section describes the size of MMA op +using ShapeMMAOp = cutlass::gemm::GemmShape<16, 8, 8>; // <- MMA Op tile M = 16, N = 8, K = 8 + +// This code section describes how threadblocks are scheduled on GPU +using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; // <- ?? + +// This code section describes the epilogue part of the kernel +using EpilogueOp = cutlass::epilogue::thread::LinearCombination< + cutlass::complex, // <- data type of output matrix + 1, // <- the number of elements per vectorized + // memory access. For a byte, it's 16 + // elements. This becomes the vector width of + // math instructions in the epilogue too + cutlass::complex, // <- data type of accumulator + cutlass::complex>; // <- data type for alpha/beta in linear combination function + +// Number of pipelines you want to use +constexpr int NumStages = 3; +// Transform +constexpr cutlass::ComplexTransform TransformA = cutlass::ComplexTransform::kNone; +constexpr cutlass::ComplexTransform TransformB = cutlass::ComplexTransform::kNone; + +// +// Gemm Operators (Gemm_3xTF32, Gemm_1xTF32, GEMM_F32, GEMM_F64) +// + +// Gemm_3xTF32 +using Gemm_3xTF32 = cutlass::gemm::device::GemmComplex< + cutlass::complex, + LayoutInputA, + cutlass::complex, + LayoutInputB, + cutlass::complex, + LayoutOutput, + cutlass::complex, + MMAOp, + SmArch, + ShapeMMAThreadBlock, + ShapeMMAWarp, + ShapeMMAOp, + EpilogueOp, + SwizzleThreadBlock, + NumStages, + TransformA, + TransformB, + cutlass::arch::OpMultiplyAddComplexFastF32>; + +// Gemm_1xTF32 +using Gemm_1xTF32 = cutlass::gemm::device::GemmComplex< + cutlass::complex, + LayoutInputA, + cutlass::complex, + LayoutInputB, + cutlass::complex, + LayoutOutput, + cutlass::complex, + MMAOp, + SmArch, + ShapeMMAThreadBlock, + ShapeMMAWarp, + ShapeMMAOp, + EpilogueOp, + SwizzleThreadBlock, + NumStages, + TransformA, + TransformB, + cutlass::arch::OpMultiplyAddComplex>; + +bool run(Options &options) { + + // Create a tuple of problem size for matrix multiplication + cutlass::gemm::GemmCoord problem_size = options.problem_size; + + //////////////////////////////////////////////////////////////////////////////// + /// 1. Initialize F32 Precision input tensors using CUTLASS helper functions + //////////////////////////////////////////////////////////////////////////////// + cutlass::HostTensor, LayoutInputA> tensor_a_F32(problem_size.mk()); // <- Create matrix A with dimensions M x K + cutlass::HostTensor, LayoutInputB> tensor_b_F32(problem_size.kn()); // <- Create matrix B with dimensions K x N + cutlass::HostTensor, LayoutOutput> tensor_c_F32(problem_size.mn()); // <- Create matrix C with dimensions M x N + cutlass::HostTensor, LayoutOutput> tensor_d_F32(problem_size.mn()); // <- Create matrix D with dimensions M x N + + if (options.rand_mode == "uniform") { + const float min = -1; + const float max = 1; + // Fill input and output matrices on host using CUTLASS helper functions + cutlass::reference::host::TensorFillRandomUniform( + tensor_a_F32.host_view(), + options.seed, + double(max), + double(min)); // <- Fill matrix A on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_b_F32.host_view(), + options.seed, + double(max), + double(min)); // <- Fill matrix B on host with uniform-distribution random data + cutlass::reference::host::TensorFillRandomUniform( + tensor_c_F32.host_view(), + options.seed, + double(max), + double(min)); // <- Fill matrix C on host with uniform-distribution random data + } else if (options.rand_mode == "gauss") { + // Fill input and output matrices on host using CUTLASS helper functions + cutlass::reference::host::TensorFillRandomGaussian( + tensor_a_F32.host_view(), + options.seed, + double(0), + double(5)); // <- Fill matrix A on host with gaussian-distribution random data + cutlass::reference::host::TensorFillRandomGaussian( + tensor_b_F32.host_view(), + options.seed, + double(0), + double(5)); // <- Fill matrix B on host with gaussian-distribution random data + cutlass::reference::host::TensorFillRandomGaussian( + tensor_c_F32.host_view(), + options.seed, + double(0), + double(5)); // <- Fill matrix C on host with gaussian-distribution random data + } + cutlass::reference::host::TensorFill( + tensor_d_F32.host_view()); // <- fill matrix D on host with zeros + + // Copy data from host to GPU + tensor_a_F32.sync_device(); + tensor_b_F32.sync_device(); + tensor_c_F32.sync_device(); + tensor_d_F32.sync_device(); + + //////////////////////////////////////////////////////////////////////////////// + /// 2. Initialize F64 tensors using the same values used for F32 + //////////////////////////////////////////////////////////////////////////////// + // Gemm input operands (A, B, C) + cutlass::HostTensor, LayoutInputA> tensor_a_F64(problem_size.mk()); // <- Create matrix A with dimensions M x K + cutlass::HostTensor, LayoutInputB> tensor_b_F64(problem_size.kn()); // <- Create matrix B with dimensions K x N + cutlass::HostTensor, LayoutOutput> tensor_c_F64(problem_size.mn()); // <- Create matrix C with dimensions M x N + + // Gemm output (D) for GEMM_F64 + cutlass::HostTensor, LayoutOutput> tensor_d_F64(problem_size.mn()); // <- Create matrix D with dimensions M x N + // Gemm output (D) for GEMM_3xTF32 + cutlass::HostTensor, LayoutOutput> tensor_d_3xTF32(problem_size.mn()); // <- Create matrix D with dimensions M x N + // Gemm output (D) for GEMM_1xTF32 + cutlass::HostTensor, LayoutOutput> tensor_d_1xTF32(problem_size.mn()); // <- Create matrix D with dimensions M x N + + // Copy values from the DP tensors + cutlass::reference::host::TensorCopy(tensor_a_F64.host_view(), tensor_a_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_b_F64.host_view(), tensor_b_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_c_F64.host_view(), tensor_c_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_d_F64.host_view(), tensor_d_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_d_3xTF32.host_view(), tensor_d_F32.host_view()); + cutlass::reference::host::TensorCopy(tensor_d_1xTF32.host_view(), tensor_d_F32.host_view()); + + // Copy data from host to GPU + tensor_a_F64.sync_device(); + tensor_b_F64.sync_device(); + tensor_c_F64.sync_device(); + tensor_d_F64.sync_device(); + tensor_d_3xTF32.sync_device(); + tensor_d_1xTF32.sync_device(); + + // Initialize alpha and beta for dot product computation + cutlass::complex alpha = cutlass::complex(options.alpha); + cutlass::complex beta = cutlass::complex(options.beta); + + // Split K dimension into 1 partitions + int split_k_slices = 1; + + //////////////////////////////////////////////////////////////////////////////// + /// 3. Run 3xTF32 kernel within a profiling loop + //////////////////////////////////////////////////////////////////////////////// + // Create a tuple of gemm kernel arguments. This is later passed as arguments to launch + // instantiated CUTLASS kernel + typename Gemm_3xTF32::Arguments arguments_3xtf32{problem_size, // <- problem size of matrix multiplication + tensor_a_F32.device_ref(), // <- reference to matrix A on device + tensor_b_F32.device_ref(), // <- reference to matrix B on device + tensor_c_F32.device_ref(), // <- reference to matrix C on device + tensor_d_3xTF32.device_ref(), // <- reference to matrix D on device + {alpha, beta}, // <- tuple of alpha and beta + split_k_slices}; // <- k-dimension split factor + + // Using the arguments, query for extra workspace required for matrix multiplication computation + size_t workspace_size_3xtf32 = Gemm_3xTF32::get_workspace_size(arguments_3xtf32); + + // Allocate workspace memory + cutlass::device_memory::allocation workspace_3xtf32(workspace_size_3xtf32); + + // Instantiate CUTLASS kernel depending on templates + Gemm_3xTF32 gemm_op; + + // Check the problem size is supported or not + cutlass::Status status_3xtf32 = gemm_op.can_implement(arguments_3xtf32); + CUTLASS_CHECK(status_3xtf32); + + // Initialize CUTLASS kernel with arguments and workspace pointer + status_3xtf32 = gemm_op.initialize(arguments_3xtf32, workspace_3xtf32.get()); + CUTLASS_CHECK(status_3xtf32); + + // Result structure + Result result; + + // + // Construct events + // + + cudaEvent_t events[2]; + + for (auto & event : events) { + result.error = cudaEventCreate(&event); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventCreate() failed: " << cudaGetErrorString(result.error) << std::endl; + return false; + } + } + + // Record an event at the start of a series of GEMMs + result.error = cudaEventRecord(events[0]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return false; + } + + // + // Run profiling loop + // + + for (int iter = 0; iter < options.iterations; ++iter) { + // Launch initialized CUTLASS kernel + status_3xtf32 = gemm_op(); + CUTLASS_CHECK(status_3xtf32); + } + + // + // Stop profiling loop + // + + // Record an event when the GEMMs are complete + result.error = cudaEventRecord(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result.error) << std::endl; + return false; + } + + // Wait for work on the device to complete. + result.error = cudaEventSynchronize(events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventSynchronize() failed: " << cudaGetErrorString(result.error) << std::endl; + return false; + } + + // Measure elapsed runtime + float runtime_ms = 0; + result.error = cudaEventElapsedTime(&runtime_ms, events[0], events[1]); + if (result.error != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result.error) << std::endl; + return false; + } + + // Compute average runtime and GFLOPs. + result.m = problem_size.m(); + result.n = problem_size.n(); + result.k = problem_size.k(); + result.runtime_ms = double(runtime_ms) / double(options.iterations); + result.gflops = options.gflops(result.runtime_ms / 1000.0); + + // Cleanup + for (auto event : events) { + (void)cudaEventDestroy(event); + } + + tensor_d_3xTF32.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + /// 4. Run TF32 kernel without profiling loop + //////////////////////////////////////////////////////////////////////////////// + // Create a tuple of gemm kernel arguments. This is later passed as arguments to launch + // instantiated CUTLASS kernel + typename Gemm_1xTF32::Arguments arguments_1xtf32{problem_size, // <- problem size of matrix multiplication + tensor_a_F32.device_ref(), // <- reference to matrix A on device + tensor_b_F32.device_ref(), // <- reference to matrix B on device + tensor_c_F32.device_ref(), // <- reference to matrix C on device + tensor_d_1xTF32.device_ref(), // <- reference to matrix D on device + {alpha, beta}, // <- tuple of alpha and beta + split_k_slices}; // <- k-dimension split factor + + // Using the arguments, query for extra workspace required for matrix multiplication computation + size_t workspace_size_1xtf32 = Gemm_1xTF32::get_workspace_size(arguments_1xtf32); + + // Allocate workspace memory + cutlass::device_memory::allocation workspace_1xtf32(workspace_size_1xtf32); + + // Instantiate CUTLASS kernel depending on templates + Gemm_1xTF32 gemm_op_1xtf32; + + // Check the problem size is supported or not + cutlass::Status status_1xtf32 = gemm_op_1xtf32.can_implement(arguments_1xtf32); + CUTLASS_CHECK(status_1xtf32); + + // Initialize CUTLASS kernel with arguments and workspace pointer + status_1xtf32 = gemm_op_1xtf32.initialize(arguments_1xtf32, workspace_1xtf32.get()); + CUTLASS_CHECK(status_1xtf32); + + // Launch initialized CUTLASS kernel + status_1xtf32 = gemm_op_1xtf32(); + CUTLASS_CHECK(status_1xtf32); + + tensor_d_1xTF32.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + // Run reference kernel (F64) + //////////////////////////////////////////////////////////////////////////////// + + // Launch device reference gemm kernel + cutlass::reference::device::GemmComplex( + problem_size, + alpha, + tensor_a_F64.device_ref(), + TransformA, + tensor_b_F64.device_ref(), + TransformB, + beta, + tensor_c_F64.device_ref(), + tensor_d_F64.device_ref(), + cutlass::complex(0.f)); + + // Wait for kernels to finish + cudaDeviceSynchronize(); + + // Copy output data from CUTLASS and reference kernel to host for comparison + tensor_d_F64.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + // Run reference kernel (F32) + //////////////////////////////////////////////////////////////////////////////// + + // Launch device reference gemm kernel + cutlass::reference::device::GemmComplex( + problem_size, + alpha, + tensor_a_F32.device_ref(), + TransformA, + tensor_b_F32.device_ref(), + TransformB, + beta, + tensor_c_F32.device_ref(), + tensor_d_F32.device_ref(), + cutlass::complex(0.f)); + + // Wait for kernels to finish + cudaDeviceSynchronize(); + + // Copy output data from CUTLASS and reference kernel to host for comparison + tensor_d_F32.sync_host(); + + //////////////////////////////////////////////////////////////////////////////// + /////// Compute l2 norms + //////////////////////////////////////////////////////////////////////////////// + + // l2 norm 3xTF32 vs F64 + cutlass::HostTensor, LayoutOutput> tensor_d_3xTF32_in_F64(problem_size.mn()); + cutlass::reference::host::TensorCopy(tensor_d_3xTF32_in_F64.host_view(), tensor_d_3xTF32.host_view()); + + result.l2_norm_3xtf32_vs_fp64 = cutlass::reference::host::TensorRelativeErrorMetric( + tensor_d_3xTF32_in_F64.host_view(), tensor_d_F64.host_view()); + + // l2 norm 1xTF32 vs F64 + cutlass::HostTensor, LayoutOutput> tensor_d_1xTF32_in_F64(problem_size.mn()); + cutlass::reference::host::TensorCopy(tensor_d_1xTF32_in_F64.host_view(), tensor_d_1xTF32.host_view()); + + result.l2_norm_1xtf32_vs_fp64 = cutlass::reference::host::TensorRelativeErrorMetric( + tensor_d_1xTF32_in_F64.host_view(), tensor_d_F64.host_view()); + + // l2 norm F32 vs F64 + cutlass::HostTensor, LayoutOutput> tensor_d_F32_in_F64(problem_size.mn()); + cutlass::reference::host::TensorCopy(tensor_d_F32_in_F64.host_view(), tensor_d_F32.host_view()); + + result.l2_norm_fp32_vs_fp64 = cutlass::reference::host::TensorRelativeErrorMetric( + tensor_d_F32_in_F64.host_view(), tensor_d_F64.host_view()); + + results.push_back(result); + + /////////////////////////////////////////////////////////////////////////////// + + // Check if output from CUTLASS kernel and reference kernel are equal or not + + std::cout << std::fixed; + std::cout.precision(4); + std::cout << "Runtime: " << result.runtime_ms << " ms" << std::endl; + std::cout.precision(2); + std::cout << "GFLOPs: " << result.gflops << std::endl; + std::cout << "Normalized L2 norm of" << std::endl; + std::cout.precision(8); + std::cout << std::scientific + << " - 3xTF32 error with FP64 reference : " << result.l2_norm_3xtf32_vs_fp64 << std::endl + << " - 1xTF32 error with FP64 reference : " << result.l2_norm_1xtf32_vs_fp64 << std::endl + << " - FP32 error with FP64 reference : " << result.l2_norm_fp32_vs_fp64 << std::endl; + + return true; +} + +int main(int argc, const char **argv) { + + bool notSupported = false; + + // Ampere Tensor Core operations exposed with mma.sync and ldmatrix are first available + // in CUDA 11.0. + // + // CUTLASS must be compiled with CUDA 11.0 Toolkit to run these examples. + if (!(__CUDACC_VER_MAJOR__ >= 11)) { + std::cerr << "Ampere Tensor Core operations must be compiled with CUDA 11.0 Toolkit or later." << std::endl; + notSupported = true; + } + + cudaDeviceProp props; + + cudaError_t error = cudaGetDeviceProperties(&props, 0); + if (error != cudaSuccess) { + std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl; + return false; + } + + if (!((props.major * 10 + props.minor) >= 80)) { + std::cerr << "Ampere Tensor Core operations must be run on a machine with compute capability at least 80." + << std::endl; + notSupported = true; + } + + if (notSupported) { + // Returning zero so this test passes on older Toolkits. Its actions are no-op. + return 0; + } + + Options options; + options.parse(argc, argv); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + bool result = true; + + if (options.benchmark) { + for (int k = 4; k <= 65536; k *= 2) { + + options.problem_size[2] = k; + + printf("Gemm problem size: %d x %d x %d\n", \ + options.problem_size.m(), options.problem_size.n(), options.problem_size.k()); + + if (!options.valid()) { + std::cerr << "Invalid problem." << std::endl; + return -1; + } + + result &= run(options); + } + } else { + // Execute one problem size + if (!options.valid()) { + std::cerr << "Invalid problem." << std::endl; + return -1; + } + + result = run(options); + } + + if (!result) return -1; + + std::cout << std::endl << "CSV results" << std::endl; + Result::print_csv_header(); + for(auto &r : results) + r.print_csv_row(); + + return 0; +} diff --git a/examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/CMakeLists.txt b/examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/CMakeLists.txt new file mode 100644 index 00000000..d7b59a39 --- /dev/null +++ b/examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/CMakeLists.txt @@ -0,0 +1,27 @@ +# Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without modification, are permitted +# provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright notice, this list of +# conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright notice, this list of +# conditions and the following disclaimer in the documentation and/or other materials +# provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used +# to endorse or promote products derived from this software without specific prior written +# permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR +# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; +# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +cutlass_example_add_executable( + 29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm + 29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm.cu + ) + diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 2dec82d5..6663d37a 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -94,6 +94,12 @@ foreach(EXAMPLE 21_quaternion_gemm 22_quaternion_conv 23_ampere_gemm_operand_reduction_fusion + 24_gemm_grouped + 25_ampere_fprop_mainloop_fusion + 26_ampere_wgrad_mainloop_fusion + 27_ampere_3xtf32_fast_accurate_tensorop_gemm + 28_ampere_3xtf32_fast_accurate_tensorop_fprop + 29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm ) add_subdirectory(${EXAMPLE}) diff --git a/include/cutlass/arch/memory_sm75.h b/include/cutlass/arch/memory_sm75.h index 35fdd393..96b208f3 100644 --- a/include/cutlass/arch/memory_sm75.h +++ b/include/cutlass/arch/memory_sm75.h @@ -282,5 +282,52 @@ inline __device__ void ldsm( ///////////////////////////////////////////////////////////////////////////////////////////////// +template +struct shared_load_op { + CUTLASS_DEVICE + shared_load_op(AccessType &D, void const *ptr) { + D = *reinterpret_cast(ptr); + } +}; + +template +CUTLASS_DEVICE void shared_load(AccessType &D, void const *ptr) { + shared_load_op(D, ptr); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct shared_load_op { + CUTLASS_DEVICE + shared_load_op(AccessType &D, void const *ptr) { + unsigned addr = cutlass_get_smem_pointer(ptr); + + uint4 v; + asm volatile ("ld.shared.v4.b32 {%0, %1, %2, %3}, [%4];" : + "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "r"(addr)); + + D = reinterpret_cast(v); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct shared_load_op { + CUTLASS_DEVICE + shared_load_op(AccessType &D, void const *ptr) { + unsigned addr = cutlass_get_smem_pointer(ptr); + + uint2 v; + asm volatile ("ld.shared.v2.b32 {%0, %1}, [%2];" : + "=r"(v.x), "=r"(v.y) : "r"(addr)); + + D = reinterpret_cast(v); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + } // namespace arch } // namespace cutlass diff --git a/include/cutlass/arch/memory_sm80.h b/include/cutlass/arch/memory_sm80.h index 6f9c3c2c..f26e05e3 100644 --- a/include/cutlass/arch/memory_sm80.h +++ b/include/cutlass/arch/memory_sm80.h @@ -68,6 +68,18 @@ template < CacheOperation::Kind cache_op = CacheOperation::Always> struct cp_async_zfill; +/// Initiates an asynchronous copy from global memory to shared memory. Rather than predicate +/// the entire transfer, nans (0x7eff) are written to SMEM if the guard predicate is false. +/// +/// LDGSTS +/// +template < + /// Size of the access in bytes + int SizeInBytes, + /// Cache operation + CacheOperation::Kind cache_op = CacheOperation::Always> +struct cp_async_nan; + //////////////////////////////////////////////////////////////////////////////////////////////////// /// Partial specialization @@ -150,6 +162,48 @@ struct cp_async_zfill { } }; +__device__ __constant__ uint4 OOB_NAN_F16x8 = {0x7eff7eff, 0x7eff7eff, + 0x7eff7eff, 0x7eff7eff}; + +/// Partial specialization +template <> +struct cp_async_nan<16, CacheOperation::Always> { + static int const kSizeInBytes = 16; + + /// Copy with nan fill + CUTLASS_DEVICE + cp_async_nan(void *smem_ptr, void const *global_ptr, bool pred_guard) { + #if CUDA_CP_ASYNC_ACTIVATED + + unsigned smem_int_ptr = cutlass_get_smem_pointer(smem_ptr); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" +#if CUTLASS_ENABLE_L2_PREFETCH + " @p cp.async.ca.shared.global.L2::128B [%1], [%2], %3;\n" +#else + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" +#endif + " @!p st.shared.v4.u32 [%1], {%4, %5, %6, %7};\n" + "}\n" + : + : "r"((int)pred_guard), "r"(smem_int_ptr), "l"(global_ptr), + "n"(kSizeInBytes), "r"(OOB_NAN_F16x8.x), "r"(OOB_NAN_F16x8.y), "r"(OOB_NAN_F16x8.z), + "r"(OOB_NAN_F16x8.w)); + + #else + + CUTLASS_UNUSED(smem_ptr); + CUTLASS_UNUSED(global_ptr); + CUTLASS_UNUSED(pred_guard); + CUTLASS_NOT_IMPLEMENTED(); + + #endif + } +}; + //////////////////////////////////////////////////////////////////////////////////////////////////// /// Partial specialization diff --git a/include/cutlass/arch/mma.h b/include/cutlass/arch/mma.h index 33dad7f2..eefa3596 100644 --- a/include/cutlass/arch/mma.h +++ b/include/cutlass/arch/mma.h @@ -62,6 +62,16 @@ struct OpMultiplyAddFastF16; ///////////////////////////////////////////////////////////////////////////////////////////////// +/// Tag indicating the input is converted to 2 (big and small) TF32 components +// Perform 3xTF32 or 4xTF32 for every F32 output element +struct OpMultiplyAddFastF32; + +/// Tag indicating the input is converted to 2 (big and small) TF32 components +// Perform 3xTF32 or 4xTF32 for every complex output element +struct OpMultiplyAddComplexFastF32; + +///////////////////////////////////////////////////////////////////////////////////////////////// + /// Tag indicating the complex multiply-add operation struct OpMultiplyAddComplex; diff --git a/include/cutlass/arch/mma_sm75.h b/include/cutlass/arch/mma_sm75.h index 817b7d2a..d2a39574 100644 --- a/include/cutlass/arch/mma_sm75.h +++ b/include/cutlass/arch/mma_sm75.h @@ -1170,7 +1170,6 @@ struct Mma< ) const { #if defined(CUTLASS_ARCH_MMA_SM75_ENABLED) - #if defined(CUTLASS_ARCH_WMMA_ENABLED) using WmmaFragmentA = nvcuda::wmma::fragment< nvcuda::wmma::matrix_a, diff --git a/include/cutlass/arch/mma_sm80.h b/include/cutlass/arch/mma_sm80.h index 31b2662c..a50ec049 100644 --- a/include/cutlass/arch/mma_sm80.h +++ b/include/cutlass/arch/mma_sm80.h @@ -34,6 +34,7 @@ #include #endif +#include "cutlass/cutlass.h" #include "mma.h" #include "cutlass/layout/matrix.h" #include "cutlass/numeric_types.h" @@ -2109,7 +2110,6 @@ struct Mma< int const *C = reinterpret_cast(&c); int *D = reinterpret_cast(&d); - asm volatile( "mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.xor.popc {%0,%1,%2,%3}, " "{%4,%5,%6,%7}, " @@ -2119,8 +2119,10 @@ struct Mma< "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); #else + assert(0); -#endif + +#endif // defined(CUTLASS_ARCH_MMA_SM80_ENABLED) } }; diff --git a/include/cutlass/array.h b/include/cutlass/array.h index bc17e527..ecca8913 100644 --- a/include/cutlass/array.h +++ b/include/cutlass/array.h @@ -338,6 +338,7 @@ private: public: + #if 0 CUTLASS_HOST_DEVICE Array() { } @@ -348,6 +349,7 @@ public: storage[i] = x.storage[i]; } } + #endif /// Efficient clear method CUTLASS_HOST_DEVICE diff --git a/include/cutlass/array_subbyte.h b/include/cutlass/array_subbyte.h index 1b01e533..a0c59526 100644 --- a/include/cutlass/array_subbyte.h +++ b/include/cutlass/array_subbyte.h @@ -395,6 +395,7 @@ private: public: + #if 0 CUTLASS_HOST_DEVICE Array() { } @@ -405,6 +406,7 @@ public: storage[i] = x.storage[i]; } } + #endif /// Efficient clear method CUTLASS_HOST_DEVICE diff --git a/include/cutlass/conv/device/implicit_gemm_convolution_fusion.h b/include/cutlass/conv/device/implicit_gemm_convolution_fusion.h new file mode 100644 index 00000000..1b481055 --- /dev/null +++ b/include/cutlass/conv/device/implicit_gemm_convolution_fusion.h @@ -0,0 +1,262 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/* \file + \brief Template for device-level fused activation's scale+bias+relu and Implicit GEMM Convolution +*/ + +#pragma once + +#include + +#include "cutlass/cutlass.h" +#include "cutlass/device_kernel.h" +#include "cutlass/conv/convolution.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +class ImplicitGemmConvolutionFusion { +public: + + using ImplicitGemmFusionKernel = ImplicitGemmFusionKernel_; + + using ElementA = typename ImplicitGemmFusionKernel::ElementA; + using LayoutA = typename ImplicitGemmFusionKernel::LayoutA; + using ElementB = typename ImplicitGemmFusionKernel::ElementB; + using LayoutB = typename ImplicitGemmFusionKernel::LayoutB; + +// using ElementScaleBias = typename ImplicitGemmFusionKernel::ElementScaleBias; +// using LayoutScaleBias = typename ImplicitGemmFusionKernel::LayoutScaleBias; + + using ElementC = typename ImplicitGemmFusionKernel::ElementC; + using LayoutC = typename ImplicitGemmFusionKernel::LayoutC; + using ElementAccumulator = typename ImplicitGemmFusionKernel::ElementAccumulator; + using ElementCompute = typename ImplicitGemmFusionKernel::ElementCompute; + using OperatorClass = typename ImplicitGemmFusionKernel::OperatorClass; + using ArchTag = typename ImplicitGemmFusionKernel::ArchTag; + using ThreadblockShape = typename ImplicitGemmFusionKernel::ThreadblockShape; + using WarpShape = typename ImplicitGemmFusionKernel::WarpShape; + using InstructionShape = typename ImplicitGemmFusionKernel::InstructionShape; + using ThreadblockSwizzle = typename ImplicitGemmFusionKernel::ThreadblockSwizzle; + using EpilogueOutputOp = typename ImplicitGemmFusionKernel::EpilogueOutputOp; + static int const kStages = ImplicitGemmFusionKernel::kStages; + static int const kConvDim = ImplicitGemmFusionKernel::kConvDim; + using WarpMmaOperator = typename ImplicitGemmFusionKernel::WarpMmaOperator; + using ArchMmaOperator = typename ImplicitGemmFusionKernel::ArchMmaOperator; + using MathOperator = typename ImplicitGemmFusionKernel::MathOperator; + + static cutlass::conv::Operator const kConvolutionalOperator = ImplicitGemmFusionKernel::kConvolutionalOperator; + static cutlass::conv::IteratorAlgorithm const kIteratorAlgorithm = ImplicitGemmFusionKernel::kIteratorAlgorithm; + + static int const kWarpCount = + (ThreadblockShape::kM / WarpShape::kM) * + (ThreadblockShape::kN / WarpShape::kN) * + (ThreadblockShape::kK / WarpShape::kK); + + /// Argument structure + using Arguments = typename ImplicitGemmFusionKernel::Arguments; + +private: + + /// Kernel parameters object + typename ImplicitGemmFusionKernel::Params params_; + +public: + + /// Constructs Implicit GEMM + ImplicitGemmConvolutionFusion() { } + + /// Determines whether the Implicit GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + // dispatch to iterators + Status status = ImplicitGemmFusionKernel::Mma::IteratorA::can_implement(args.problem_size); + if (Status::kSuccess != status) { + return status; + } + + status = ImplicitGemmFusionKernel::Mma::IteratorB::can_implement(args.problem_size); + if (Status::kSuccess != status) { + return status; + } + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape( + threadblock_swizzle.get_tiled_shape( + cutlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size), + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices)); + + if (!(grid.y <= std::numeric_limits::max() && + grid.z <= std::numeric_limits::max())) { + + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t workspace_bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + cutlass::gemm::GemmCoord grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + cutlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size), + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices); + + if(args.split_k_mode == SplitKMode::kParallel) { + + // Split-K parallel: CTAs in k-dimension write the partial results in a temporary workspace. + // The user needs to call a reduction operator to optain the final output tensor + workspace_bytes = + sizeof(ElementAccumulator) * + size_t(cutlass::conv::implicit_gemm_tensor_c_size(kConvolutionalOperator, args.problem_size)) * + size_t(grid_tiled_shape.k()); + } + + else if(args.split_k_mode == SplitKMode::kSerial && args.problem_size.split_k_slices > 1) { + + // Split-K serial: The user workspace is used to store semaphore and serialize writing the + // final reduced output to user's output tensor + workspace_bytes = sizeof(int) * size_t(grid_tiled_shape.m()) * size_t(grid_tiled_shape.n()); + } + + return workspace_bytes; + } + + /// Initializes GEMM state from arguments. + Status initialize( + Arguments const &args, + void *workspace = nullptr, + cudaStream_t stream = nullptr) { + + if (args.problem_size.split_k_slices > 1) { + + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + cudaError_t status = cudaMemsetAsync(workspace, 0, get_workspace_size(args), stream); + + if (status != cudaSuccess) { + return Status::kErrorInternal; + } + } + + // initialize the params structure from the arguments + params_ = typename ImplicitGemmFusionKernel::Params( + args, + static_cast(workspace) + ); + + int smem_size = int(sizeof(typename ImplicitGemmFusionKernel::SharedStorage)); + + if (smem_size >= (48 << 10)) { + cudaError_t result = cudaFuncSetAttribute(cutlass::Kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != cudaSuccess) { + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + /// Initializes Impicit GEMM state from arguments. + Status update(Arguments const &args, void *workspace = nullptr) { + + // update the params structure from the arguments + params_.ptr_A = args.ref_A.data(); + params_.ptr_B = args.ref_B.data(); + params_.ptr_scale = args.ref_A_scale.data(); + params_.ptr_bias = args.ref_A_bias.data(); + params_.ptr_C = args.ref_C.data(); + params_.ptr_D = args.ref_D.data(); + params_.output_op = args.output_op; + params_.semaphore = static_cast(workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(cudaStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(32 * kWarpCount, 1, 1); + + int smem_size = int(sizeof(typename ImplicitGemmFusionKernel::SharedStorage)); + + 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; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} +} +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/kernel/default_conv2d.h b/include/cutlass/conv/kernel/default_conv2d.h index 8679f7aa..5605162c 100644 --- a/include/cutlass/conv/kernel/default_conv2d.h +++ b/include/cutlass/conv/kernel/default_conv2d.h @@ -41,10 +41,12 @@ #include "cutlass/conv/threadblock/conv2d_tile_iterator.h" #include "cutlass/conv/threadblock/implicit_gemm_pipelined.h" #include "cutlass/conv/threadblock/implicit_gemm_multistage.h" +#include "cutlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h" +#include "cutlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h" #include "cutlass/conv/kernel/implicit_gemm_convolution.h" +#include "cutlass/conv/kernel/implicit_gemm_convolution_fusion.h" #include "cutlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h" - ///////////////////////////////////////////////////////////////////////////////////////////////// namespace cutlass { diff --git a/include/cutlass/conv/kernel/default_conv2d_fprop_fusion.h b/include/cutlass/conv/kernel/default_conv2d_fprop_fusion.h new file mode 100644 index 00000000..142b64c3 --- /dev/null +++ b/include/cutlass/conv/kernel/default_conv2d_fprop_fusion.h @@ -0,0 +1,351 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + Default kernel-level fused activation's scale+bias+relu and implicit GEMM convolution + definitions that combine threadblock-scoped matrix multiply-add with the + appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/conv/kernel/default_conv2d.h" + +#include "cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h" +#include "cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h" +#include "cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h" +#include "cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h" +#include "cutlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h" +#include "cutlass/conv/threadblock/regular_scale_bias_vector_access_iterator.h" +#include "cutlass/conv/warp/conv2d_fprop_scale_bias_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for fused batch norm and Conv2dFprop +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided +> struct DefaultConv2dFpropFusion; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassTensorOp convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialzation for Analytic IteratorAlgorithm and multistage +/// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv2dFpropFusion < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementScaleBias, + LayoutScaleBias, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + cutlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + 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::Conv2dFpropFilterTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using IteratorScaleBias = + cutlass::conv::threadblock::PredicatedScaleBiasVectorAccessIterator< + cutlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + using SmemIteratorScaleBias = + cutlass::conv::threadblock::RegularScaleBiasVectorAccessIterator< + cutlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static int const kThreadCount = 32; + + // Warp-level iterators to load scale and bias vectors + using WarpIteratorScaleBias = cutlass::conv::warp::WarpIteratorScaleBias< + MatrixShape, ElementScaleBias, + LayoutScaleBias, MatrixShape, + typename WarpMmaTensorOp::IteratorA::Base::Policy, kThreadCount, + MmaCore::WarpCount::kK>; + + // Define the Mma + using Mma = threadblock::ImplicitGemmFpropFusionMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + IteratorScaleBias, + SmemIteratorScaleBias, + arch::CacheOperation::Always, + MmaPolicy, + WarpIteratorScaleBias, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::ImplicitGemmConvolutionFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialzation for Optimzed IteratorAlgorithm and +/// multistage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv2dFpropFusion < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementScaleBias, + LayoutScaleBias, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag + >; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + cutlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + 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::Conv2dFpropFilterTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementB, + LayoutB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using IteratorScaleBias = + cutlass::conv::threadblock::PredicatedScaleBiasVectorAccessIterator< + cutlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + using SmemIteratorScaleBias = + cutlass::conv::threadblock::RegularScaleBiasVectorAccessIterator< + cutlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static int const kThreadCount = 32; + + // Warp-level iterators to load scale and bias vectors + using WarpIteratorScaleBias = cutlass::conv::warp::WarpIteratorScaleBias< + MatrixShape, ElementScaleBias, + LayoutScaleBias, MatrixShape, + typename WarpMmaTensorOp::IteratorA::Base::Policy, kThreadCount, + MmaCore::WarpCount::kK>; + + // Define the Mma + using Mma = threadblock::ImplicitGemmFpropFusionMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + IteratorScaleBias, + SmemIteratorScaleBias, + arch::CacheOperation::Always, + MmaPolicy, + WarpIteratorScaleBias, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::ImplicitGemmConvolutionFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/kernel/default_conv2d_wgrad_fusion.h b/include/cutlass/conv/kernel/default_conv2d_wgrad_fusion.h new file mode 100644 index 00000000..5aed546d --- /dev/null +++ b/include/cutlass/conv/kernel/default_conv2d_wgrad_fusion.h @@ -0,0 +1,319 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + Default kernel-level implicit GEMM convolution definitions combine threadblock-scoped + matrix multiply-add with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/conv/kernel/default_conv2d.h" + +#include "cutlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_analytic.h" +#include "cutlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_analytic.h" +#include "cutlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_optimized.h" +#include "cutlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_optimized.h" +#include "cutlass/conv/threadblock/conv2d_tile_iterator.h" +#include "cutlass/conv/threadblock/predicated_scale_bias_vector_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided +> struct DefaultConv2dWgradFusion; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassTensorOp convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialzation for Analytic IteratorAlgorithm and multistage +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv2dWgradFusion < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementScaleBias, + LayoutScaleBias, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + cutlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorAnalytic< + 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::Conv2dWgradActivationTileAccessIteratorAnalytic< + cutlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using IteratorScaleBias = + cutlass::conv::threadblock::PredicatedScaleBiasVectorIterator< + cutlass::MatrixShape<1, WarpShape::kN>, + ElementScaleBias, + LayoutScaleBias>; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmWgradFusionMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + IteratorScaleBias, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::ImplicitGemmConvolutionFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialzation for Optimized IteratorAlgorithm and multistage +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv2dWgradFusion < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementScaleBias, + LayoutScaleBias, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + cutlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorOptimized< + 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::Conv2dWgradActivationTileAccessIteratorOptimized< + cutlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using IteratorScaleBias = + cutlass::conv::threadblock::PredicatedScaleBiasVectorIterator< + cutlass::MatrixShape<1, WarpShape::kN>, + ElementScaleBias, + LayoutScaleBias>; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmWgradFusionMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + IteratorScaleBias, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = cutlass::conv::kernel::ImplicitGemmConvolutionFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/kernel/implicit_gemm_convolution.h b/include/cutlass/conv/kernel/implicit_gemm_convolution.h index fd6cc7e2..28658541 100644 --- a/include/cutlass/conv/kernel/implicit_gemm_convolution.h +++ b/include/cutlass/conv/kernel/implicit_gemm_convolution.h @@ -385,7 +385,6 @@ struct ImplicitGemmConvolution { semaphore.wait(threadblock_tile_idx.k()); - __threadfence(); } // Each split-k-slice writes to a unique tensor location else if (params.split_k_mode == SplitKMode::kParallel) { diff --git a/include/cutlass/conv/kernel/implicit_gemm_convolution_fusion.h b/include/cutlass/conv/kernel/implicit_gemm_convolution_fusion.h new file mode 100644 index 00000000..8d961048 --- /dev/null +++ b/include/cutlass/conv/kernel/implicit_gemm_convolution_fusion.h @@ -0,0 +1,455 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined fused activation's scale+bias+relu and Implicit GEMM kernel. +*/ + +#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 ImplicitGemmConvolutionFusion { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static Operator const kConvolutionalOperator = ConvOperator; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + + using ElementScaleBias = typename Mma::IteratorScaleBias::Element; + using LayoutScaleBias = typename Mma::IteratorScaleBias::Layout; + + using ElementC = typename EpilogueOutputOp::ElementOutput; + using LayoutC = LayoutA; + + using ElementAccumulator = typename EpilogueOutputOp::ElementAccumulator; + using ElementCompute = typename EpilogueOutputOp::ElementCompute; + + using WarpMmaOperator = typename Mma::Policy::Operator; + + using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator; + using MathOperator = typename ArchMmaOperator::Operator; + + using OperatorClass = typename WarpMmaOperator::OperatorClass; + using ArchTag = typename WarpMmaOperator::ArchTag; + + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename WarpMmaOperator::Shape; + using InstructionShape = typename ArchMmaOperator::Shape; + + static int const kStages = Mma::kStages; + static IteratorAlgorithm const kIteratorAlgorithm = Mma::IteratorA::kIteratorAlgorithm; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + using TensorRefA = typename Mma::IteratorA::TensorRef; + using TensorRefB = typename Mma::IteratorB::TensorRef; + using TensorRefScaleBias = typename Mma::IteratorScaleBias::TensorRef; + using TensorRefC = 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_; + + /// Wgrad C stride idx for implicit gemm algorithm + // Conv2d row-major matrix C (KxRSC) + // Conv3d row-major matrix C (KxTRSC) + static int const kWgradCStrideIdx = + cutlass::platform::is_same::value ? 2 : 3; + + /// This chooses the appropriate stride element of the C tensor. + static int const kTensorCStrideIdx = + (kConvolutionalOperator == conv::Operator::kWgrad ? kWgradCStrideIdx : 0); + + // + // + // + using ConvOutputIteratorParameter = epilogue::threadblock::ConvOutputIteratorParameter< + LayoutC, + typename Epilogue::OutputTileIterator::Layout, + TensorRefC, + ConvOperator, + ConvProblemSize + >; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + ConvProblemSize problem_size; + TensorRefA ref_A; + TensorRefB ref_B; + TensorRefScaleBias ref_scale; + TensorRefScaleBias ref_bias; + TensorRefC ref_C; + TensorRefC ref_D; + typename EpilogueOutputOp::Params output_op; + SplitKMode split_k_mode; + + // + // Methods + // + + /// Default ctor + 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, + TensorRefScaleBias const & ref_scale, + TensorRefScaleBias const & ref_bias, + TensorRefC const & ref_C, + TensorRefC const & ref_D, + typename EpilogueOutputOp::Params const & output_op, + SplitKMode const & split_k_mode = SplitKMode::kSerial + ): + problem_size(problem_size), + ref_A(ref_A), + ref_B(ref_B), + ref_scale(ref_scale), + ref_bias(ref_bias), + ref_C(ref_C), + ref_D(ref_D), + output_op(output_op), + split_k_mode(split_k_mode) + { + + } + + }; + + /// Parameters structure + struct Params { + ConvProblemSize problem_size; + 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; + + // + // Methods + // + + CUTLASS_HOST_DEVICE + Params(): swizzle_log_tile(0), gemm_k_iterations(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_scale_bias(args.problem_size, args.ref_scale.layout()), + ptr_scale(args.ref_scale.data()), + ptr_bias(args.ref_bias.data()), + iterator_C(ConvOutputIteratorParameter::layout(args.ref_C)), + ptr_C(args.ref_C.data()), + iterator_D(ConvOutputIteratorParameter::layout(args.ref_D)), + ptr_D(args.ref_D.data()), + output_op(args.output_op), + semaphore(semaphore), + split_k_mode(args.split_k_mode) + { + gemm_k_iterations = implicit_gemm_k_iterations(kConvolutionalOperator, ThreadblockShape::kK, args.problem_size); + + ThreadblockSwizzle threadblock_swizzle; + + grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + implicit_gemm_problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices); + + swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape); + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + CUTLASS_HOST_DEVICE + ImplicitGemmConvolutionFusion() { } + + /// 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 operand + typename Mma::IteratorA iterator_A( + params.iterator_A, + params.problem_size, + params.ptr_A, + thread_idx, + MatrixCoord( + threadblock_tile_idx.m() * Mma::Shape::kM, + threadblock_tile_idx.k() * Mma::Shape::kK + ) + ); + + // Construct iterators to B operand + typename Mma::IteratorB iterator_B( + params.iterator_B, + params.problem_size, + params.ptr_B, + thread_idx, + MatrixCoord( + threadblock_tile_idx.k() * Mma::Shape::kK, + threadblock_tile_idx.n() * Mma::Shape::kN + ) + ); + + // Construct iterators to A scale/bias vector + typename Mma::IteratorScaleBias iterator_scale_bias( + params.iterator_scale_bias, + params.problem_size, + params.ptr_scale, + params.ptr_bias, + thread_idx, + MatrixCoord( + 0, (kConvolutionalOperator == conv::Operator::kFprop) ? + (threadblock_tile_idx.k() * Mma::Shape::kK) : + // Wgrad + (threadblock_tile_idx.n() * Mma::Shape::kN) + ) + ); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = __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, iterator_scale_bias, accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // Construct the semaphore. + int block_idx = threadblock_tile_idx.m() + threadblock_tile_idx.n() * params.grid_tiled_shape.m(); + + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // Compute logical position within grid + threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_idx.k(), params.grid_tiled_shape.k()); + } + + MatrixCoord threadblock_offset( + threadblock_tile_idx.m() * Mma::Shape::kM, + threadblock_tile_idx.n() * Mma::Shape::kN + ); + + // Tile iterator writing to destination tensor + typename Epilogue::OutputTileIterator iterator_D( + params.iterator_D, + params.ptr_D, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset + ); + + // Tile iterator reading from source accumulator tensor + typename Epilogue::OutputTileIterator iterator_C( + params.iterator_C, + params.ptr_C, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset + ); + + // Construct the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_idx.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_idx.k()); + + } + // Each split-k-slice writes to a unique tensor location + else if (params.split_k_mode == SplitKMode::kParallel) { + iterator_D.add_pointer_offset(threadblock_tile_idx.k() * + cutlass::conv::implicit_gemm_tensor_c_size(ConvOperator, params.problem_size)); + } + + // Run efficient epilogue + epilogue(output_op, iterator_D, accumulators, iterator_C); + + // + // Release the semaphore + // + + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_idx.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_idx.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// 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 5ad1304f..d9529311 100644 --- a/include/cutlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h +++ b/include/cutlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h @@ -423,7 +423,6 @@ struct ImplicitGemmConvolutionStridedDgrad { semaphore.wait(threadblock_tile_idx.k()); - __threadfence(); } // Each split-k-slice writes to a unique tensor location else if (params.split_k_mode == SplitKMode::kParallel) { diff --git a/include/cutlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h b/include/cutlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h index 2db42453..1b5716e8 100644 --- a/include/cutlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h +++ b/include/cutlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h @@ -439,7 +439,6 @@ struct ImplicitGemmConvolutionWithFusedEpilogue { semaphore.wait(threadblock_tile_idx.k()); - __threadfence(); } // Each split-k-slice writes to a unique tensor location else if (params.split_k_mode == SplitKMode::kParallel) { diff --git a/include/cutlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h b/include/cutlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h new file mode 100644 index 00000000..5a32f7d9 --- /dev/null +++ b/include/cutlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h @@ -0,0 +1,787 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a multistage threadblock-scoped fused activation's + scale+bias+relu and Implicit GEMM Convolution kernel. + + The original implicit gemm will store out-of-bound data as zeroes in the + shared memory because zeros into the tensor core, zeroes out of the tensor + cores. The result is remained the same. When fusing scale+bias+relu + into the mainloop, it is no longer true because + + 0 x scale + bias = bias + + which is no longer always 0. So, instead of storing zeroes, this fused + kernel stores the out-of-bound data as a special NaN (0x7eff), when applying + scale+bias+relu, the code is like + + if (data == 0x7eff) + data = 0; + else + data = scale+bias+relu(data, scale, bias); + + See include/cutlass/conv/warp/scale_bias_relu_transformation.h for the + elementwise computation. See include/cutlass/arch/memory_sm80.h for nan fill. +*/ + +#pragma once + +#include "cutlass/aligned_buffer.h" +#include "cutlass/arch/memory.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/cache_operation.h" +#include "cutlass/gemm/gemm.h" + +#include "cutlass/conv/warp/conv2d_fprop_scale_bias_iterator.h" +#include "cutlass/conv/warp/scale_bias_relu_transform.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace threadblock { + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Element type of scale and bias vectors + typename ElementScaleBias_, + /// Layout of scale and bias vectors + typename LayoutScaleBias_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// WarpIterator to load Scale or Bias vector from the shared memory + typename WarpIteratorScaleBias_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class MmaFpropFusionBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Element type of scale and bias vectors + using ElementScaleBias = ElementScaleBias_; + + /// Layout of scale and bias vectors + using LayoutScaleBias = LayoutScaleBias_; + + ///< Policy describing tuning details + using Policy = Policy_; + + ///< WarpIterator to load Scale or Bias vector from the shared memory + using WarpIteratorScaleBias = WarpIteratorScaleBias_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = cutlass::gemm::GemmShape; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the scale and bias vectors + using TensorRefScaleBias = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape; + + /// Shape of the A scale and bias vectors in shared memory + using ShapeScaleBias = + MatrixShape<1 + Policy::SmemPaddingA::kRow, + 2 * Shape::kK * kStages + Policy::SmemPaddingA::kColumn>; + + /// Shape of the B matrix operand in shared memory + using ShapeB = + MatrixShape; + + public: + // + // Data members + // + + /// Buffer for A operand + AlignedBuffer operand_A; + + /// Buffer for B operand + AlignedBuffer operand_B; + + /// Buffer for A operand Scale and Bias + AlignedBuffer operand_A_scale_bias; + + public: + + // + // Methods + // + + /// Returns a layout object for the A matrix + CUTLASS_DEVICE + static typename Operator::LayoutA LayoutA() { + return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn}); + } + + /// Returns a layout object for the B matrix + CUTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a layout object for the A scale and bias vectors + CUTLASS_DEVICE + static LayoutScaleBias LayoutScaleBias() { + return LayoutScaleBias::packed( + {ShapeScaleBias::kRow, ShapeScaleBias::kColumn}); + } + + /// Returns a TensorRef to the A operand + CUTLASS_HOST_DEVICE + TensorRefA operand_A_ref() { + return TensorRefA{operand_A.data(), LayoutA()}; + } + + /// Returns a TensorRef to the B operand + CUTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + + /// Returns a TensorRef to the A operand Scale vector + CUTLASS_HOST_DEVICE + TensorRefScaleBias operand_A_scale_bias_ref() { + return TensorRefScaleBias{operand_A_scale_bias.data(), LayoutScaleBias()}; + } + }; + + protected: + + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of A operand scale and bias vector + /// from shared memory + WarpIteratorScaleBias warp_tile_iterator_A_scale_bias_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + +public: + + /// Construct from tensor references + CUTLASS_DEVICE + MmaFpropFusionBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx), + warp_tile_iterator_A_scale_bias_( + shared_storage.operand_A_scale_bias_ref(), lane_idx), + warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + cutlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB, + /// Iterates over vectors of scale and bias vector in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorScaleBias_, + /// Iterates over vectors of scale and bias vector in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorScaleBias_, + /// Cache operation for scale/bias operand + cutlass::arch::CacheOperation::Kind CacheOpScaleBias, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// WarpIterator to load Scale or Bias vector from the shared memory + typename WarpIteratorScaleBias_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class ImplicitGemmFpropFusionMultistage + : public MmaFpropFusionBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Iterates over tiles of the scale and bias vectors in global memory + using IteratorScaleBias = IteratorScaleBias_; + ///< WarpIterator to load Scale or Bias vector from the shared memory + using WarpIteratorScaleBias = WarpIteratorScaleBias_; + ///< Policy describing tuning details + using Policy = Policy_; + ///< Base class + using Base = MmaFpropFusionBase; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + using SmemIteratorScaleBias = SmemIteratorScaleBias_; + + static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + static cutlass::arch::CacheOperation::Kind const kCacheOpScaleBias = + CacheOpScaleBias; + + // + // Dependent types + // + + /// Fragment of accumulator tile + + using ElementC = typename Policy::Operator::ElementC; + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Internal structure exposed for introspection. + struct Detail { + + static_assert(Base::kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpLoadedFragmentScaleBias = + typename WarpIteratorScaleBias::Fragment; + + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of A operand scale vector to shared memory + SmemIteratorScaleBias smem_iterator_A_scale_bias_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + CUTLASS_DEVICE + ImplicitGemmFpropFusionMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_A_scale_bias_(shared_storage.operand_A_scale_bias_ref(), + thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_A_scale_bias_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + CUTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, + IteratorScaleBias &iterator_A_scale_bias, + IteratorB &iterator_B, int group_start_A = 0, + int group_start_B = 0) { + iterator_A.set_iteration_index(group_start_A); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / 8; + + // Uses nan fill for out of bound data + cutlass::arch::cp_async_nan( + dst_ptr, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + + ++this->smem_iterator_A_; + } + } + + // Async Copy for operand A scale and bias vector. Scale and bias vectors + // are small. One iteration is enough. + if (group_start_A == 0) { + typename IteratorScaleBias::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_scale_bias_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorScaleBias::kElementsPerAccess / 8; + + cutlass::arch::cp_async( + dst_ptr, iterator_A_scale_bias.get(), iterator_A_scale_bias.valid()); + } + + iterator_B.set_iteration_index(group_start_B); + + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / 8; + + cutlass::arch::cp_async_zfill( + dst_ptr, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + ++this->smem_iterator_B_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< iterator over scale and bias vectors in global memory + IteratorScaleBias iterator_A_scale_bias, + ///< initial value of accumulator + FragmentC const &src_accum, + ///< Imaginary strides used for planar-complex only - ignored here + int64_t imag_stride_A = 0, + int64_t imag_stride_B = 0) { + + // + // Prologue + // + + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / 8; + + // Uses Nan fill for out of bound data + cutlass::arch::cp_async_nan( + dst_ptr, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + ++this->smem_iterator_A_; + } + + // Async Copy for operand A scale and bias vectors. Scale and bias + // vectors are small. One iteration is enough. + { + typename IteratorScaleBias::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_scale_bias_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorScaleBias::kElementsPerAccess / 8; + + cutlass::arch::cp_async( + dst_ptr, iterator_A_scale_bias.get(), iterator_A_scale_bias.valid()); + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / 8; + + cutlass::arch::cp_async_zfill( + dst_ptr, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A.advance(); + iterator_A_scale_bias.advance(); + iterator_B.advance(); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_A_scale_bias_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Inserts a fence to group cp.async instructions into stages. + cutlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Waits until kStages-2 stages have committed. + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpLoadedFragmentScaleBias warp_loaded_frag_A_scale_bias[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + cutlass::conv::warp::FpropScaleBiasReluTransform + elementwise_transform; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_A_scale_bias_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_A_scale_bias_.load( + warp_loaded_frag_A_scale_bias[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_scale_bias_; + ++this->warp_tile_iterator_B_; + + // Start issuing the first group of the next stage outside of the mainloop + copy_tiles_and_advance(iterator_A, iterator_A_scale_bias, iterator_B); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + elementwise_transform(warp_transformed_frag_A[0], + warp_loaded_frag_A_scale_bias[0]); + + // + // Mainloop + // + + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_A_scale_bias_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_A_scale_bias_.load( + warp_loaded_frag_A_scale_bias[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_scale_bias_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) { + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + elementwise_transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_loaded_frag_A_scale_bias[warp_mma_k % 2]); + } + + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + + // Issue global->shared copies for the next stage + int group_start_iteration_A, group_start_iteration_B; + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + group_start_iteration_A = 0; + group_start_iteration_B = 0; + } else { + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + } + + copy_tiles_and_advance(iterator_A, iterator_A_scale_bias, iterator_B, + group_start_iteration_A, + group_start_iteration_B); + + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + elementwise_transform( + warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_A_scale_bias[(warp_mma_k + 1) % 2]); + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + // Inserts a fence to group cp.async instructions into stages. + cutlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages of cp.async have committed + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.advance(); + iterator_A_scale_bias.advance(); + iterator_B.advance(); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_A_scale_bias_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_A_scale_bias_.add_tile_offset( + {0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_A_scale_bias_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + } + } + + } + + // Insert fence and wait for all outstanding cp.async operations to commit. + cutlass::arch::cp_async_fence(); + cutlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/threadblock/implicit_gemm_multistage.h b/include/cutlass/conv/threadblock/implicit_gemm_multistage.h index 03b4c1ab..6033772b 100644 --- a/include/cutlass/conv/threadblock/implicit_gemm_multistage.h +++ b/include/cutlass/conv/threadblock/implicit_gemm_multistage.h @@ -376,6 +376,20 @@ public: warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + // tf32x3 kernels use staging accumulation. warp_mma uses a temporary + // accumulator and this temporary accumulator is added to the final + // accumulator once in every mainloop iteration. + plus plus_accum; + + FragmentC tmp_accum; + + if (platform::is_same::value + || platform::is_same::value) { + tmp_accum.clear(); + } + // // Mainloop // @@ -426,12 +440,29 @@ public: copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B); - warp_mma( - accum, - warp_transformed_frag_A[warp_mma_k % 2], - warp_transformed_frag_B[warp_mma_k % 2], - accum - ); + if (platform::is_same::value + || platform::is_same::value) { + warp_mma( + tmp_accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + tmp_accum + ); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + } if (warp_mma_k + 1 == Base::kWarpGemmIterations) warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], @@ -483,6 +514,13 @@ public: } + if (platform::is_same::value + || platform::is_same::value) { + accum = plus_accum(accum, tmp_accum); + } + // Insert fence and wait for all outstanding cp.async operations to commit. cutlass::arch::cp_async_fence(); cutlass::arch::cp_async_wait<0>(); diff --git a/include/cutlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h b/include/cutlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h new file mode 100644 index 00000000..5cd87b1b --- /dev/null +++ b/include/cutlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h @@ -0,0 +1,718 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a multistage threadblock-scoped fused activation's scale+bias+relu and + Implicit GEMM Convolution kernel. + + The original implicit gemm will store out-of-bound data as zeroes in the + shared memory because zeros into the tensor core, zeroes out of the tensor + cores. The result is remained the same. When fusing scale+bias+relu + into the mainloop, it is no longer true because + + 0 x scale + bias = bias + + which is no longer always 0. So, instead of storing zeroes, this fused + kernel stores the out-of-bound data as a special NaN (0x7eff), when applying + scale+bias+relu, the code is like + + if (data == 0x7eff) + data = 0; + else + data = scale+bias+relu(data, scale, bias); + + The biggest difference compared with the fused Fprop and scale+bias+relu is + that scale and bias are loop invariant in Wgrad so that they only needs to + be loaded once before the mainloop. + + See include/cutlass/conv/warp/scale_bias_relu_transformation.h for the + elementwise computation. See include/cutlass/arch/memory_sm80.h for nan fill. + + +*/ + +#pragma once + +#include "cutlass/aligned_buffer.h" +#include "cutlass/arch/memory.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/cache_operation.h" +#include "cutlass/gemm/gemm.h" + +#include "cutlass/conv/warp/conv2d_fprop_scale_bias_iterator.h" +#include "cutlass/conv/warp/scale_bias_relu_transform.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace threadblock { + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Element type of scale and bias vectors + typename ElementScaleBias_, + /// Layout of scale and bias vectors + typename LayoutScaleBias_, + /// Element type of scale and bias vectors + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class MmaWgradFusionBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Element type of scale and bias vectors + using ElementScaleBias = ElementScaleBias_; + + /// Layout of scale and bias vectors + using LayoutScaleBias = LayoutScaleBias_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = cutlass::gemm::GemmShape; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape; + + /// Shape of the B matrix operand in shared memory + using ShapeB = + MatrixShape; + + public: + // + // Data members + // + + /// Buffer for A operand + AlignedBuffer operand_A; + + /// Buffer for B operand + AlignedBuffer operand_B; + + public: + + // + // Methods + // + + /// Returns a layout object for the A matrix + CUTLASS_DEVICE + static typename Operator::LayoutA LayoutA() { + return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn}); + } + + /// Returns a layout object for the B matrix + CUTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a TensorRef to the A operand + CUTLASS_HOST_DEVICE + TensorRefA operand_A_ref() { + return TensorRefA{operand_A.data(), LayoutA()}; + } + + /// Returns a TensorRef to the B operand + CUTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + }; + + protected: + + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + +public: + + /// Construct from tensor references + CUTLASS_DEVICE + MmaWgradFusionBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx), + warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + cutlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB, + /// Iterates over vectors of scale and bias vector in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorScaleBias_, + /// Iterates over vectors of scale and bias vector i + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class ImplicitGemmWgradFusionMultistage + : public MmaWgradFusionBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Iterates over tiles of the scale and bias vectors in global memory + using IteratorScaleBias = IteratorScaleBias_; + ///< Policy describing tuning details + using Policy = Policy_; + ///< Base class + using Base = MmaWgradFusionBase; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + + using ElementC = typename Policy::Operator::ElementC; + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Internal structure exposed for introspection. + struct Detail { + + static_assert(Base::kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + static int const kBBufferSize = + ((sizeof(typename Operator::ElementC) == 4) && + ((platform::is_same::value && + platform::is_same::value)) && + (Operator::Shape::kM >= 64 && Operator::Shape::kN >= 64)) + ? 1 + : 2; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpLoadedFragmentScaleBias = typename IteratorScaleBias::Fragment; + + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + int warp_idx_m_; + + int warp_idx_n_; + +public: + + /// Construct from tensor references + CUTLASS_DEVICE + ImplicitGemmWgradFusionMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) { + + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + warp_idx_m_ = warp_idx_mn % Base::WarpCount::kM; + warp_idx_n_ = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m_, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n_}); + } + + CUTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, + IteratorB &iterator_B, + int group_start_A = 0, int group_start_B = 0) { + + iterator_A.set_iteration_index(group_start_A); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / 8; + + cutlass::arch::cp_async_zfill( + dst_ptr, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index(group_start_B); + + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / 8; + + // Uses nan fill for out of bound data + cutlass::arch::cp_async_nan( + dst_ptr, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + ++this->smem_iterator_B_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< iterator over scale and bias vectors in global memory + IteratorScaleBias iterator_B_scale_bias, + ///< initial value of accumulator + FragmentC const &src_accum, + ///< Imaginary strides used for planar-complex only - ignored here + int64_t imag_stride_A = 0, + int64_t imag_stride_B = 0) { + + // + // Prologue + // + + WarpLoadedFragmentScaleBias warp_loaded_frag_B_scale_bias; + iterator_B_scale_bias.add_tile_offset({0, warp_idx_n_}); + iterator_B_scale_bias.load(warp_loaded_frag_B_scale_bias); + + // Issue several complete stages + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / 8; + + cutlass::arch::cp_async_zfill( + dst_ptr, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + ++this->smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / 8; + + // Uses Nan fill for out of bound data + cutlass::arch::cp_async_nan( + dst_ptr, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A.advance(); + iterator_B.advance(); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Inserts a fence to group cp.async instructions into stages. + cutlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Waits until kStages-2 stages have committed. + cutlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[Detail::kBBufferSize]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[Detail::kBBufferSize]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + cutlass::conv::warp::WgradScaleBiasReluTransform + elementwise_transform; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + // Start issuing the first group of the next stage outside of the mainloop + copy_tiles_and_advance(iterator_A, iterator_B); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + elementwise_transform(warp_transformed_frag_B[0], + warp_loaded_frag_B_scale_bias); + + // + // Mainloop + // + + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + if (Detail::kBBufferSize == 2) { + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % Detail::kBBufferSize]); + ++this->warp_tile_iterator_A_; + } + + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) { + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % Detail::kBBufferSize], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % Detail::kBBufferSize], + warp_loaded_frag_B[warp_mma_k % 2]); + + elementwise_transform(warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_B_scale_bias); + } + + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % Detail::kBBufferSize], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + + if (Detail::kBBufferSize == 1) { + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + ++this->warp_tile_iterator_A_; + + } + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % Detail::kBBufferSize], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % Detail::kBBufferSize], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + elementwise_transform( + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_B_scale_bias); + } + + // Issue global->shared copies for the next stage + int group_start_iteration_A, group_start_iteration_B; + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + group_start_iteration_A = 0; + group_start_iteration_B = 0; + } else { + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + } + + copy_tiles_and_advance(iterator_A, iterator_B, + group_start_iteration_A, + group_start_iteration_B); + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + // Inserts a fence to group cp.async instructions into stages. + cutlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages of cp.async have committed + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.advance(); + iterator_B.advance(); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + } + } + + } + + // Insert fence and wait for all outstanding cp.async operations to commit. + cutlass::arch::cp_async_fence(); + cutlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h b/include/cutlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h new file mode 100644 index 00000000..db2c2cd6 --- /dev/null +++ b/include/cutlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h @@ -0,0 +1,393 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT,INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates calculating the address and predicates to the load of scale and bias vectors. + + This iterator uses masks to guard out-of-bounds accesses. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "cutlass/array.h" +#include "cutlass/coord.h" +#include "cutlass/cutlass.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/pitch_linear.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/predicate_vector.h" +#include "cutlass/tensor_ref.h" +#include "cutlass/tensor_view.h" +#include "cutlass/conv/threadblock/conv2d_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedScaleBiasVectorAccessIterator +/// +template +class PredicatedScaleBiasVectorAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for fprop pitch-linear data. +/// +template +class PredicatedScaleBiasVectorAccessIterator { + public: + + using ThreadblockShape = ThreadblockShape_; + using Element = Element_; + using Layout = layout::PitchLinear; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + static int const kElementsPerAccess = 128 / sizeof_bits::value; + static int const kThreads = ThreadblockShape::kContiguous / kElementsPerAccess; + + using AccessType = AlignedArray; + + using Params = PredicatedScaleBiasVectorAccessIteratorParams; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Parameters object with precomputed internal state + Params const ¶ms_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + /// Size of tensor + Conv2dProblemSize problem_size_; + + int filter_c_; + int filter_r_; + int filter_s_; + + TensorCoord thread_offset_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Extent of tensor + Conv2dProblemSize const &problem_size, + /// Pointer to the start of the scale vector + ConstPointer scale_pointer, + /// Pointer to the start of the bias vector + ConstPointer bias_pointer, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : params_(params), + problem_size_(problem_size), + filter_c_(0), + filter_r_(0), + filter_s_(0) { + pointer_ = (thread_id < kThreads) + ? reinterpret_cast( + const_cast(scale_pointer)) + : reinterpret_cast( + const_cast(bias_pointer)); + + // Per-thread offset in logical coordinates of tensor + int thread_base = (thread_id < kThreads) ? 0 : kThreads; + + thread_offset_ = + threadblock_offset + + TensorCoord((thread_id - thread_base) * kElementsPerAccess, 0); + + set_iteration_index(0); + } + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Extent of tensor + Conv2dProblemSize const &problem_size, + /// Pointer to start of scale vector + ConstPointer scale_pointer, + /// Pointer to start of scale vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id) + : PredicatedScaleBiasVectorAccessIterator(params, problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) {} + + /// Advances an iterator along logical dimensions of matrix in units of whole threadblock tiles + CUTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + thread_offset_ = + thread_offset_ + + TensorCoord(ThreadblockShape::kContiguous * tile_offset.contiguous(), 0); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType *get() const { + + return reinterpret_cast( + pointer_ + + (thread_offset_.contiguous() * sizeof_bits::value / 8)); + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator &operator++() { + return *this; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + void advance() { + // moves to the next tile + ++filter_s_; + if (filter_s_ == problem_size_.S) { + filter_s_ = 0; + ++filter_r_; + + if (filter_r_ < problem_size_.R) { + } else { + filter_r_ = 0; + add_tile_offset(TensorCoord(1, 0)); + } + } + } + + /// Increment and return an instance to self. + CUTLASS_DEVICE + PredicatedScaleBiasVectorAccessIterator operator++(int) { + PredicatedScaleBiasVectorAccessIterator self(*this); + operator++(); + return self; + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + uint32_t enabled = 0; + +#if defined(_MSC_VER) || (__CUDACC_VER_MAJOR__ < 11) + enabled = threadIdx.x < kThreads * 2; +#else + asm volatile( + "{\n" + " .reg .u32 tid_reg;\n" + " .reg .pred p;\n" + " mov.u32 tid_reg, %%tid.x;\n" + " setp.lt.u32 p, tid_reg, %1;\n" + " selp.u32 %0, 1, 0, p;\n" + "}\n" : "+r"(enabled) :"n"(kThreads * 2)); +#endif + + return ((thread_offset_.contiguous() < problem_size_.C) && enabled); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedScaleBiasVectorAccessIterator { + public: + + using ThreadblockShape = ThreadblockShape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedScaleBiasVectorAccessIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear>; + + using AccessType = typename UnderlyingIterator::AccessType; + static int const kElementsPerAccess = UnderlyingIterator::kElementsPerAccess; + + using Params = PredicatedScaleBiasVectorAccessIteratorParams; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Extent of tensor + Conv2dProblemSize const &problem_size, + ///< Pointer to the start of the scale vector + ConstPointer scale_pointer, + ///< Pointer to the start of the bias vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params, problem_size, scale_pointer, bias_pointer, + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Conv2dProblemSize const &problem_size, ///< Extent of tensor + ConstPointer scale_pointer, ///< Pointer to the start of the scale vector + ConstPointer bias_pointer, ///< Pointer to the start of the bias vector + int thread_id ///< ID of each participating thread + ) + : PredicatedScaleBiasVectorAccessIterator(params, problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// threadblock tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator operator++(int) { + PredicatedScaleBiasVectorAccessIterator self(*this); + operator++(); + return self; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + void advance() { + iterator_.advance(); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/threadblock/predicated_scale_bias_vector_iterator.h b/include/cutlass/conv/threadblock/predicated_scale_bias_vector_iterator.h new file mode 100644 index 00000000..1719afee --- /dev/null +++ b/include/cutlass/conv/threadblock/predicated_scale_bias_vector_iterator.h @@ -0,0 +1,365 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT,INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates calculating the address and predicates to the load of scale and bias vectors. + + This iterator uses masks to guard out-of-bounds accesses. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "cutlass/array.h" +#include "cutlass/coord.h" +#include "cutlass/cutlass.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/pitch_linear.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/predicate_vector.h" +#include "cutlass/tensor_ref.h" +#include "cutlass/tensor_view.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedScaleBiasVectorIterator +/// +template +class PredicatedScaleBiasVectorIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for wgrad pitch-linear data. +/// +template +class PredicatedScaleBiasVectorIterator { + public: + + using WarpShape = WarpShape_; + using Element = Element_; + using Layout = layout::PitchLinear; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + static int const kElementsPerAccess = 1; + + using AccessType = AlignedArray; + + static int const kIterations = WarpShape::kContiguous / 8; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array<__half2, 2 * kIterations * kElementsPerAccess>; + + /// Parameters object is precomputed state and is host-constructible + using Params = Conv2dWgradActivationIteratorOptimizedParams; + + private: + // + // Data members + // + + /// Parameters object with precomputed internal state + Params const ¶ms_; + + /// Internal pointer to first access of tile + ConstPointer scale_pointer_; + ConstPointer bias_pointer_; + + /// Size of tensor + Conv2dProblemSize problem_size_; + + int32_t thread_offset_; + + // Channel dimension in contiguous dimension stays constant for each gemm_iteration_k + int32_t filter_c_[kIterations]; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Extent of tensor + Conv2dProblemSize const &problem_size, + /// Pointer to the start of the scale vector + ConstPointer scale_pointer, + /// Pointer to the start of the bias vector + ConstPointer bias_pointer, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : params_(params), + problem_size_(problem_size), + scale_pointer_(scale_pointer), + bias_pointer_(bias_pointer) { + + thread_offset_ = threadblock_offset.contiguous() + (thread_id % 32) / 4; + } + + /// Construct a PredicatedTileIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Extent of tensor + Conv2dProblemSize const &problem_size, + /// Pointer to start of scale vector + ConstPointer scale_pointer, + /// Pointer to start of scale vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id) + : PredicatedScaleBiasVectorIterator(params, problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Advances an iterator along logical dimensions of matrix in units of whole warp tiles + CUTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + + thread_offset_ += (WarpShape::kContiguous * tile_offset.contiguous()); + + CUTLASS_PRAGMA_UNROLL + for(int c = 0; c < kIterations; ++c) { + int rsc_offset = thread_offset_ + c * 8; + + int residual, tmp; + params_.sc_divmod(tmp, residual, rsc_offset); + params_.c_divmod(tmp, filter_c_[c], residual); + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + frag.fill(__float2half2_rn(0.0f)); + __half2 *frag_ptr = reinterpret_cast<__half2 *>(&frag); + + // load scale + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + + cutlass::arch::global_load< + __half, + sizeof(AccessType) + >( + frag_ptr[c * 2].x, + scale_pointer_ + filter_c_[c], + true + ); + } + + // load bias + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + + cutlass::arch::global_load< + __half, + sizeof(AccessType) + >( + frag_ptr[c * 2 + 1].x, + bias_pointer_ + filter_c_[c], + true + ); + } + + // duplicate scale + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + frag_ptr[c * 2].y = frag_ptr[c * 2].x; + } + + // duplicate bias + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + frag_ptr[c * 2 + 1].y = frag_ptr[c * 2 + 1].x; + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedScaleBiasVectorIterator { + public: + + using WarpShape = WarpShape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedScaleBiasVectorIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear>; + + using AccessType = typename UnderlyingIterator::AccessType; + static int const kElementsPerAccess = UnderlyingIterator::kElementsPerAccess; + using Fragment = typename UnderlyingIterator::Fragment; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedScaleBiasVectorIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default ctor + CUTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Conv2dProblemSize const &problem_size, Layout const &layout) + : params_(problem_size, layout::TensorNHWC(0, 0, 0)){}; + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Extent of tensor + Conv2dProblemSize const &problem_size, + ///< Pointer to the start of the scale vector + ConstPointer scale_pointer, + ///< Pointer to the start of the bias vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, problem_size, scale_pointer, bias_pointer, + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + Params const ¶ms, ///< Precomputed parameters object + Conv2dProblemSize const &problem_size, ///< Extent of tensor + ConstPointer scale_pointer, ///< Pointer to the start of the scale vector + ConstPointer bias_pointer, ///< Pointer to the start of the bias vector + int thread_id ///< ID of each participating thread + ) + : PredicatedScaleBiasVectorIterator(params, problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// threadblock tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment &frag) { + iterator_.load(frag); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/threadblock/regular_scale_bias_vector_access_iterator.h b/include/cutlass/conv/threadblock/regular_scale_bias_vector_access_iterator.h new file mode 100644 index 00000000..5f3116a9 --- /dev/null +++ b/include/cutlass/conv/threadblock/regular_scale_bias_vector_access_iterator.h @@ -0,0 +1,247 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates implementing computing the addresses of storing of small + scale and bias vectors in the shared memory. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/array.h" +#include "cutlass/layout/pitch_linear.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/matrix_coord.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/tensor_ref.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// RegularScaleBiasVectorAccessIterator +/// +template +class RegularScaleBiasVectorAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularScaleBiasVectorAccessIterator { + public: + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + /// Element type per access + static int const kElementsPerAccess = 128 / sizeof_bits::value; + static int const kThreads = Shape::kContiguous / kElementsPerAccess; + using AccessType = Array; + + private: + // + // Data members + // + + /// Internal pointer + AccessType *pointer_; + + /// Internal byte offset + Index byte_offset_; + + public: + /// Construct a TileIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator( + TensorRef scale_bias_ref, ///< Pointer to the start of the scale and bias + ///< vector + int thread_id ///< ID of each participating thread + ) + : byte_offset_(0) { + // Per-thread offset in logical coordinates of tensor + int thread_offset = thread_id * kElementsPerAccess; + + // initialize pointer + pointer_ = + reinterpret_cast(scale_bias_ref.data() + thread_offset); + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) {} + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Returns a pointer + CUTLASS_DEVICE + AccessType *get() const { + + char *access_byte_ptr = + reinterpret_cast(pointer_); + + return reinterpret_cast(access_byte_ptr + byte_offset_); + } + + /// Advances to the next tile in memory. + CUTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator &operator++() { return *this; } + + /// Advances to the next tile in memory. + CUTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator operator++(int) { + RegularScaleBiasVectorAccessIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset in the unit of tile. + CUTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + // Multiply by 2 because we store sclae and bias belong to the same stage + // next to each other. + add_pointer_offset(coord.contiguous() * Shape::kContiguous * 2); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for row major layouts +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularScaleBiasVectorAccessIterator< + Shape_, Element_, + layout::RowMajor> { + public: + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + /// Underlying iterator type + using UnderlyingIterator = RegularScaleBiasVectorAccessIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator( + TensorRef scale_bias_ref, ///< Pointer to the start of the scale and bias + ///< vector + int thread_id ///< ID of each participating thread + ) + : iterator_({scale_bias_ref.data(), scale_bias_ref.stride()}, thread_id) { + } + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + CUTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + CUTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + CUTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator operator++(int) { + RegularScaleBiasVectorAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/warp/conv2d_fprop_scale_bias_iterator.h b/include/cutlass/conv/warp/conv2d_fprop_scale_bias_iterator.h new file mode 100644 index 00000000..5fa48afd --- /dev/null +++ b/include/cutlass/conv/warp/conv2d_fprop_scale_bias_iterator.h @@ -0,0 +1,568 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Defines iterators used by warp-level loading scale and bias vectors. + Every scale/bias data only needs to be loaded once for every channel. +*/ + +#pragma once + +#include "cutlass/cutlass.h" + +#include "cutlass/array.h" +#include "cutlass/numeric_types.h" +#include "cutlass/tensor_ref.h" +#include "cutlass/matrix_shape.h" + +#include "cutlass/arch/memory_sm75.h" +#include "cutlass/gemm/gemm.h" + +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/tensor.h" +#include "cutlass/layout/pitch_linear.h" +#include "cutlass/layout/tensor_op_multiplicand_sm75.h" + +#include "cutlass/platform/platform.h" +#include "cutlass/fast_math.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Policy of the details of LDSM shape and iterations + typename Policy_, + /// Number of threads participating in one matrix operation + int Threads, + /// Number of partitions along K dimension + int PartitionsK_ = 1> +class WarpIteratorScaleBias; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Policy of the details of LDSM shape and iterations + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK_> +class WarpIteratorScaleBias { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = cutlass::layout::PitchLinear; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// Number of partitions along K dimension + static int const kElementsPerAccess = 128 / sizeof_bits::value; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + using Policy = Policy_; + + private: + + /// Pointer type used for accesses + using AccessType = Array; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + private: + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Internal counter used to determine when to increment byte offset and when + /// to XOR it + int k_group_idx_; + + public: + /// Default ctor constructs null iterator + CUTLASS_HOST_DEVICE + WarpIteratorScaleBias() + : pointer_(nullptr), + byte_offset_(0), + k_group_idx_(0) {} + + /// Constructor from TensorRef + CUTLASS_DEVICE + WarpIteratorScaleBias(TensorRef const &ref_scale_bias, + int lane_id) + : byte_offset_(0), k_group_idx_(0) { + /// 16816 only + pointer_ = reinterpret_cast(ref_scale_bias.data()) + + ((lane_id >> 3) & 1) * Shape::kContiguous / kElementsPerAccess + + (lane_id >> 4); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + CUTLASS_DEVICE + WarpIteratorScaleBias &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof_bits::value / 8; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_DEVICE + WarpIteratorScaleBias &add_tile_offset( + TensorCoord const &tile_offset) { + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + + byte_offset_ += k_groups_delta * sizeof_bits::value * + kElementsPerAccess * Policy::LdsmShape::kContiguous / 8; + + // Multiply by 2 because scale and bias belonging to the same stage are next + // to each other in the shared memory. + pointer_ += (2 * whole_tiles * Shape::kContiguous / kElementsPerAccess); + + return *this; + } + + /// Advances the iterator along the advance dimension + CUTLASS_DEVICE + WarpIteratorScaleBias &operator++() { + byte_offset_ += Policy::LdsmShape::kContiguous * + sizeof_bits::value * kElementsPerAccess / 8; + + k_group_idx_++; + + if (k_group_idx_ == (Policy::kGroupsPerTile / kPartitionsK)) { + k_group_idx_ = 0; + byte_offset_ -= (Policy::kGroupsPerTile / kPartitionsK) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * kElementsPerAccess / 8; + add_tile_offset({Policy::kGroupsPerTile, 0}); + } + + return *this; + } + + /// Advances the iterator along the advance dimension + CUTLASS_HOST_DEVICE + WarpIteratorScaleBias &operator--() { assert(0); } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + CUTLASS_DEVICE + WarpIteratorScaleBias &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + CUTLASS_DEVICE + WarpIteratorScaleBias &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + CUTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + CUTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + Array *fetch_ptr = + reinterpret_cast *>(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < 1; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + + AccessType const *source_ptr = + pointer_ + Policy::LdsmShape::kContiguous * c; + + char const *source_byte_ptr = + reinterpret_cast(source_ptr) + byte_offset + + byte_offset_; + + cutlass::arch::ldsm( + fetch_ptr[access_idx], source_byte_ptr); + } + } + } + + /// Loads a fragment from memory with additional logical offset + CUTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + CUTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + CUTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + CUTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + kElementsPerAccess; + + byte_offset += sizeof_bits::value * pointer_offset / 8; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + CUTLASS_DEVICE + void set_kgroup_index(int k_group) { + k_group_idx_ = k_group % (Policy::kGroupsPerTile / kPartitionsK); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Policy of the details of LDSM shape and iterations + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK_> +class WarpIteratorScaleBias { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = cutlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + using Policy = Policy_; + + /// Underlying tile iterator implementation + using Base = WarpIteratorScaleBias< + layout::PitchLinearShape, Element, + layout::PitchLinear, + layout::PitchLinearShape, + Policy, kThreads, PartitionsK_>; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + + private: + /// Underlying tile iterator + Base iterator_; + + public: + /// Default ctor constructs null iterator + CUTLASS_HOST_DEVICE + WarpIteratorScaleBias() {} + + /// Constructor from TensorRef + CUTLASS_HOST_DEVICE + WarpIteratorScaleBias(TensorRef const &ref_scale_bias, int lane_id) + : iterator_({ref_scale_bias.data(), ref_scale_bias.stride()}, lane_id) {} + + /// Adds a pointer offset to internal pointer(s) to advance through memory + CUTLASS_HOST_DEVICE + WarpIteratorScaleBias &add_pointer_offset(LongIndex offset) { + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + WarpIteratorScaleBias &add_tile_offset( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_DEVICE + WarpIteratorScaleBias &add_tile_offset_negative( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset_negative({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + CUTLASS_HOST_DEVICE + WarpIteratorScaleBias &operator++() { + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + CUTLASS_HOST_DEVICE + WarpIteratorScaleBias &operator--() { + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + CUTLASS_DEVICE + WarpIteratorScaleBias &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + CUTLASS_DEVICE + WarpIteratorScaleBias &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + CUTLASS_HOST_DEVICE + void load(Fragment &frag) const { iterator_.load(frag); } + + /// Loads a fragment from memory with additional logical offset + CUTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + CUTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + CUTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + // TODO + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + CUTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + // TODO + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + CUTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, {tile_offset.strided(), tile_offset.contiguous()}, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + CUTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace conv +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/conv/warp/scale_bias_relu_transform.h b/include/cutlass/conv/warp/scale_bias_relu_transform.h new file mode 100644 index 00000000..0f7e04ce --- /dev/null +++ b/include/cutlass/conv/warp/scale_bias_relu_transform.h @@ -0,0 +1,217 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level per channel scale+bias+relu before + matrix multiply-accumulate operations targeting Tensor Cores. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/array.h" +#include "cutlass/platform/platform.h" + +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" +#include "cutlass/matrix_shape.h" + +#include "cutlass/arch/memory_sm75.h" +#include "cutlass/arch/mma_sm75.h" +#include "cutlass/arch/mma_sm80.h" + +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/warp/mma.h" + +#include "cutlass/gemm/warp/mma_tensor_op_policy.h" + +#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace conv { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct FpropScaleBiasReluTransform { + + using T = typename FragmentActivations::Element; + + static int const NumActivations = FragmentActivations::kElements; + static int const NumScaleBias = FragmentScaleBias::kElements; + static int const MmaElements = 2; + // One element has one scale and one bias + static int const MmaScaleBiasPair = 2; + // 16816 has 2 columns + static int const MmaCols = 2; + + using MmaOperand = Array; + using ScaleBiasOperand = Array; + + CUTLASS_DEVICE + void transform(MmaOperand &activations, ScaleBiasOperand const &scale_bias) { + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) + uint32_t *ptr_activations = reinterpret_cast(&activations); + uint32_t const *ptr_scale_bias = reinterpret_cast(&scale_bias); + + // Apply per channel scale+bias+relu if the data is not a special NaN + // (0x7eff). If it is a special NaN (0x7eff), hard code the output to 0. + + // We assumes the pair of FP16 are either both inbound or both out-of-bound. + // It requires C to be an even number. + asm volatile( + "{\n\t" + " .reg .pred %%p;\n\t" + " .reg .b32 t1;\n\t" + " setp.eq.u32 %%p, %2, %4;\n\t" + " fma.rn.f16x2.relu t1, %1, %2, %3;\n" + " selp.u32 %0, 0, t1, %%p;\n\t" + "}\n" + : "=r"(ptr_activations[0]) + : "r"(ptr_scale_bias[0]), "r"(ptr_activations[0]), + "r"(ptr_scale_bias[1]), "n"(0x7eff7eff)); +#else + // TODO: write emulation code + assert(0); +#endif + } + + CUTLASS_DEVICE + void operator()(FragmentActivations &activations, + FragmentScaleBias const &scale_bias) { + MmaOperand *ptr_activations = reinterpret_cast(&activations); + ScaleBiasOperand const *ptr_scale_bias = + reinterpret_cast(&scale_bias); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < (NumActivations / MmaElements); ++i) { + transform(ptr_activations[i], ptr_scale_bias[(i / MmaScaleBiasPair) % MmaCols]); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct WgradScaleBiasReluTransform { + + using T = typename FragmentActivations::Element; + + static int const NumActivations = FragmentActivations::kElements; + static int const NumScaleBias = FragmentScaleBias::kElements; + static int const MmaElements = 2; + // One element has one scale and one bias + static int const MmaScaleBiasPair = 2; + // 16816 has 2 rows + static int const MmaRows = 2; + + using MmaOperand = Array; + using ScaleBiasOperand = Array<__half2, MmaScaleBiasPair>; + + CUTLASS_DEVICE + void transform(MmaOperand &activations, ScaleBiasOperand const &scale_bias) { + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) + + __half2 *ptr_activations = reinterpret_cast<__half2 *>(&activations); + uint32_t const *ptr_scale_bias = reinterpret_cast(&scale_bias); + +#if 1 + // CUDA + PTX version + + bool h1_oob = (reinterpret_cast(ptr_activations[0].x) == 0x7eff); + bool h2_oob = (reinterpret_cast(ptr_activations[0].y) == 0x7eff); + + // Apply per channel scale+bias+relu if the data is not a special NaN + // (0x7eff). If it is a special NaN (0x7eff), hard code the output to 0. + + // We cannot gurantee that the pair of F16 are both in bound or both + // out-of-bound because C x R x S can be an odd number. + asm volatile( + "{\n\t" + " fma.rn.f16x2.relu %0 , %1, %2, %3;\n" + "}" + : "=r"(reinterpret_cast(ptr_activations[0])) + : "r"(ptr_scale_bias[0]), "r"(reinterpret_cast(ptr_activations[0])), + "r"(ptr_scale_bias[1])); + + reinterpret_cast(ptr_activations[0]) = h1_oob ? + (reinterpret_cast(ptr_activations[0]) & 0xffff0000) : + reinterpret_cast(ptr_activations[0]); + + reinterpret_cast(ptr_activations[0]) = h2_oob ? + (reinterpret_cast(ptr_activations[0]) & 0xffff) : + reinterpret_cast(ptr_activations[0]); +#else + // pure PTX version + + // Apply per channel scale+bias+relu if the data is not a special NaN + // (0x7eff). If it is a special NaN (0x7eff), hard code the output to 0. + asm volatile( + "{\n" + " .reg .b16 t1, t2;\n" + " .reg .b32 t3, t4, t5, t6;\n" + " .reg .pred p1, p2;\n" + " mov.b32 {t1, t2}, %2;\n" + " setp.eq.s16 p1, t1, %4;\n" + " setp.eq.s16 p2, t2, %4;\n" + " fma.rn.f16x2.relu t3, %1, %2, %3;\n" + " and.b32 t4, t3, %5;\n" + " selp.b32 t5, t4, t3, p1;\n" + " and.b32 t6, t5, %6;\n" + " selp.b32 %0, t6, t5, p2;\n" + "}\n" + : "=r"(reinterpret_cast(ptr_activations[0])) + : "r"(ptr_scale_bias[0]), "r"(reinterpret_cast(ptr_activations[0])), + "r"(ptr_scale_bias[1]), "n"(0x7eff), "n"(0xffff0000), "n"(0x0000ffff)); +#endif +#else + // TODO: write emulation code + assert(0); +#endif + } + + CUTLASS_DEVICE + void operator()(FragmentActivations &activations, + FragmentScaleBias const &scale_bias) { + MmaOperand *ptr_activations = reinterpret_cast(&activations); + ScaleBiasOperand const *ptr_scale_bias = + reinterpret_cast(&scale_bias); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < (NumActivations / MmaElements); ++i) { + transform(ptr_activations[i], ptr_scale_bias[(i / MmaRows)]); + } + } +}; +} // namespace warp +} // namespace conv +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/core_io.h b/include/cutlass/core_io.h index 22585dbc..4ad4f3f4 100644 --- a/include/cutlass/core_io.h +++ b/include/cutlass/core_io.h @@ -210,7 +210,7 @@ std::ostream & operator<<(std::ostream &out, GemmShape const &gemm_shape) /// Default printing to ostream for GemmCoord inline std::ostream & operator<<(std::ostream &out, GemmCoord const &gemm_coord) { - out << "cutlass::gemm::GemmCoord:: {" + out << "cutlass::gemm::GemmCoord {" << gemm_coord.m() <<"," << gemm_coord.n() <<"," << gemm_coord.k() << "}"; @@ -230,7 +230,7 @@ namespace layout { template < int Contiguous, int Strided> inline std::ostream & operator<<(std::ostream &out, PitchLinearShape const &pitch_linear_shape) { - out << "cutlass::layout::PitchLinearShape::(kContiguous, kStrided) {" + out << "cutlass::layout::PitchLinearShape:(kContiguous, kStrided) {" << cutlass::layout::PitchLinearShape::kContiguous <<"," << cutlass::layout::PitchLinearShape::kStrided <<"}"; return out; diff --git a/include/cutlass/epilogue/thread/linear_combination_bias_relu.h b/include/cutlass/epilogue/thread/linear_combination_bias_relu.h index 0c6d789f..cb3a4faa 100644 --- a/include/cutlass/epilogue/thread/linear_combination_bias_relu.h +++ b/include/cutlass/epilogue/thread/linear_combination_bias_relu.h @@ -352,7 +352,6 @@ public: FragmentCompute tmp_Accum = NumericArrayConverter()(AB); FragmentCompute tmp_C = NumericArrayConverter()(frag_C); FragmentCompute result_Z; - FragmentCompute result_T; bool conditions[kElementsPerAccess]; @@ -397,7 +396,6 @@ public: FragmentCompute tmp_Accum = NumericArrayConverter()(AB); FragmentCompute result_Z; - FragmentCompute result_T; bool conditions[kElementsPerAccess]; diff --git a/include/cutlass/epilogue/thread/linear_combination_relu.h b/include/cutlass/epilogue/thread/linear_combination_relu.h index 6805c1d9..6479c31e 100644 --- a/include/cutlass/epilogue/thread/linear_combination_relu.h +++ b/include/cutlass/epilogue/thread/linear_combination_relu.h @@ -78,10 +78,12 @@ public: using ElementCompute = ElementCompute_; static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; using FragmentOutput = Array; using FragmentAccumulator = Array; - using ComputeFragment = Array; + using FragmentCompute = Array; + using FragmentScaleBias = Array; static FloatRoundStyle const kRound = Round; @@ -181,15 +183,15 @@ public: NumericArrayConverter source_converter; NumericArrayConverter accumulator_converter; - ComputeFragment converted_source = source_converter(source); - ComputeFragment converted_accumulator = accumulator_converter(accumulator); + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); // Perform binary operations - ComputeFragment intermediate; + FragmentCompute intermediate; - multiplies mul_add_source; - multiply_add mul_add_accumulator; - ReLu relu; + multiplies mul_add_source; + multiply_add mul_add_accumulator; + ReLu relu; if (Scale == ScaleType::NoBetaScaling) intermediate = converted_source; @@ -215,13 +217,13 @@ public: // Convert source to interal compute numeric type NumericArrayConverter accumulator_converter; - ComputeFragment converted_accumulator = accumulator_converter(accumulator); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); // Perform binary operations - ComputeFragment intermediate; + FragmentCompute intermediate; - multiplies mul_accumulator; - ReLu relu; + multiplies mul_accumulator; + ReLu relu; intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum @@ -233,6 +235,42 @@ public: return destination_converter(intermediate); } + + /// Computes per-channel linear scaling and bias : D = scale * accumulator + bias + /// Scale and Bias are from input Fragment + CUTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentScaleBias const &scale, + FragmentScaleBias const &bias) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform per-channel scale and bias + FragmentCompute intermediate; + + multiply_add mul_add_accumulator; + + if(Scale == ScaleType::OnlyAlphaPerChannelScaling) + intermediate = mul_add_accumulator(scale, converted_accumulator, bias); // D = scale * Accum + bias + else + intermediate = mul_add_accumulator(alpha_, converted_accumulator, bias); // D = alpha * Accum + bias + + ReLu relu; + + // Compute threshold optionally + intermediate = relu(threshold_, intermediate); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } + + }; ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -262,10 +300,12 @@ public: static bool const kIsHeavy = detail::LinearCombinationReluIsHeavy(); static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; using FragmentOutput = Array; using FragmentAccumulator = Array; - using ComputeFragment = Array; + using FragmentCompute = Array; + using FragmentScaleBias = Array; static FloatRoundStyle const kRound = Round; @@ -363,15 +403,15 @@ public: NumericArrayConverter source_converter; NumericArrayConverter accumulator_converter; - ComputeFragment converted_source = source_converter(source); - ComputeFragment converted_accumulator = accumulator_converter(accumulator); + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); // Perform binary operations - ComputeFragment intermediate; + FragmentCompute intermediate; - multiplies mul_add_source; - multiply_add mul_add_accumulator; - ReLu relu; + multiplies mul_add_source; + multiply_add mul_add_accumulator; + ReLu relu; if (Scale == ScaleType::NoBetaScaling) intermediate = converted_source; @@ -411,13 +451,13 @@ public: // Convert source to interal compute numeric type NumericArrayConverter accumulator_converter; - ComputeFragment converted_accumulator = accumulator_converter(accumulator); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); // Perform binary operations - ComputeFragment intermediate; + FragmentCompute intermediate; - multiplies mul_accumulator; - ReLu relu; + multiplies mul_accumulator; + ReLu relu; intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum @@ -443,6 +483,64 @@ public: return destination_converter(intermediate); } } + + /// Computes per-channel linear scaling and bias : D = scale * accumulator + bias + /// Scale and Bias are from input Fragment + CUTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentScaleBias const &scale, + FragmentScaleBias const &bias) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform per-channel scale and bias + FragmentCompute intermediate; + + multiply_add mul_add_accumulator; + + if(Scale == ScaleType::OnlyAlphaPerChannelScaling) + intermediate = mul_add_accumulator(scale, converted_accumulator, bias); // D = scale * Accum + bias + else + intermediate = mul_add_accumulator(alpha_, converted_accumulator, bias); // D = alpha * Accum + bias + + ReLu relu; + + // Compute threshold optionally + intermediate = relu(threshold_, intermediate); + + if (platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value) { + // Convert floats back to INT + FragmentAccumulator scaled_accumulator; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + scaled_accumulator[i] = __float2int_rn(intermediate[i]); + } + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + return destination_converter(scaled_accumulator); + } else { + NumericArrayConverter + destination_converter; + return destination_converter(intermediate); + } + } + }; #endif // Conditional guards to enable partial specialization for packed integers diff --git a/include/cutlass/epilogue/thread/scale_type.h b/include/cutlass/epilogue/thread/scale_type.h index ca66464c..576b4194 100644 --- a/include/cutlass/epilogue/thread/scale_type.h +++ b/include/cutlass/epilogue/thread/scale_type.h @@ -41,9 +41,10 @@ namespace thread { /// Specifies internal data type for computation struct ScaleType { enum Kind { - Default, // alpha x C + beta x D - NoBetaScaling, // alpha x C + D - OnlyAlphaScaling, // alpha x C + Default, // alpha x C + beta x D + NoBetaScaling, // alpha x C + D + OnlyAlphaScaling, // alpha x C + OnlyAlphaPerChannelScaling, // alpha_vec x C Nothing // C }; }; diff --git a/include/cutlass/epilogue/threadblock/default_epilogue_direct_store.h b/include/cutlass/epilogue/threadblock/default_epilogue_direct_store.h new file mode 100644 index 00000000..62744d51 --- /dev/null +++ b/include/cutlass/epilogue/threadblock/default_epilogue_direct_store.h @@ -0,0 +1,68 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Direct store epilogue +*/ + +#pragma once + +//////////////////////////////////////////////////////////////////////////////// + +#include "cutlass/epilogue/threadblock/epilogue_direct_store.h" +#include "cutlass/epilogue/threadblock/direct_store_epilogue_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Given a properly constructed epilogue, returns a direct store epilogue +template +struct DefaultEpilogueDirectStore { + + using OutputTileIterator = DirectStoreEpilogueIterator; + + using Epilogue = EpilogueDirectStore< + typename EpilogueTensorOp::Shape, + typename EpilogueTensorOp::WarpMmaOperator, + EpilogueTensorOp::kPartitionsK, + OutputTileIterator, + typename EpilogueTensorOp::AccumulatorFragmentIterator, + typename EpilogueTensorOp::WarpTileIterator, + typename EpilogueTensorOp::SharedLoadIterator, + typename EpilogueTensorOp::OutputOp + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/threadblock/direct_epilogue_tensor_op.h b/include/cutlass/epilogue/threadblock/direct_epilogue_tensor_op.h deleted file mode 100644 index 7636e559..00000000 --- a/include/cutlass/epilogue/threadblock/direct_epilogue_tensor_op.h +++ /dev/null @@ -1,252 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. - * - * Redistribution and use in source and binary forms, with or without modification, are permitted - * provided that the following conditions are met: - * * Redistributions of source code must retain the above copyright notice, this list of - * conditions and the following disclaimer. - * * Redistributions in binary form must reproduce the above copyright notice, this list of - * conditions and the following disclaimer in the documentation and/or other materials - * provided with the distribution. - * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used - * to endorse or promote products derived from this software without specific prior written - * permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR - * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND - * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, - * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; - * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, - * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ -/*! \file - \brief Epilogue for tensor operations - -*/ - -#pragma once - -#include "cutlass/cutlass.h" -#include "cutlass/numeric_types.h" -#include "cutlass/array.h" - -#include "cutlass/gemm/gemm.h" - -///////////////////////////////////////////////////////////////////////////////////////////////// - -namespace cutlass { -namespace epilogue { -namespace threadblock { - -///////////////////////////////////////////////////////////////////////////////////////////////// - -/// Epilogue operator -template < - typename Shape_, /// Threadblock-scoped tile size (concept: GemmShape) - typename Operator_, /// Warp-scoped epilogue components (concept: gemm::warp::Mma) - int PartitionsK, /// Number of partitions of the threadblock's K dimension - typename Element_, /// Data type of source and destination tensor - typename OutputOp_, /// Function object computing final output - typename ConvertOp_ /// Conversion operator between accumulators and shared memory -> -class DirectEpilogueTensorOp { -public: - - using Shape = Shape_; - using Operator = Operator_; - - /// Number of warps spanning threadblock-scoped tile - using WarpCount = gemm::GemmShape< - Shape::kM / Operator::Shape::kM, - Shape::kN / Operator::Shape::kN, - PartitionsK, - >; - - static_assert(PartitionsK == 1, - "Direct epilogue cannot be used with when the threadblock tile is partitioned along the K dimension."); - - /// Accumulator tile is really the warp-scoped tile - using FragmentC = typename Operator::FragmentC; - - /// Data type of output tensor - using Element = Element_; - - /// Output layout is always row-major - using Layout = layout::RowMajor; - - /// Function operator computing final output - using OutputOp = OutputOp_; - - /// Conversion operator to shared memory - using ConvertOp = ConvertOp_; - - /// Reference to source and destination tensors - using TensorRef = TensorRef; - -public: - - /// Parameters structure for host-constructible state - struct Params { - - // - // Data members - // - - TensorRef destination_ref; - TensorRef source_ref; - - typename OutputOp::Params output_op; - typename ConvertOp::Params convert_op; - - // - // Methods - // - - /// Constructs a Params object - CUTLASS_HOST_DEVICE - Params( - TensorRef destination_ref_, - TensorRef source_ref_, - typename OutputOp::Params output_op_, - typename ConvertOp::Params convert_op_ - ): - destination_ref(destination_ref_), - source_ref(source_ref_), - output_op(output_op_), - convert_op(convert_op_) { - - } - - /// Constructs a Params object - CUTLASS_HOST_DEVICE - Params( - TensorRef destination_ref_, - TensorRef source_ref_, - typename OutputOp::Params output_op_ - ): - Params( - destination_ref, - source_ref, - output_op, - ConvertOp::Params() - ) { } - }; - - /// Shared storage allocation needed by the epilogue - struct SharedStorage { }; - -private: - - OutputOp output_op; ///< output operator - ConvertOp convert_op; ///< conversion operator - - TensorRef destination_ref_; ///< destination tensor reference - TensorRef source_ref_; ///< source tensor reference - - MatrixCoord warp_origin_; ///< position of the warp within the threadblock - -public: - - /// Constructor - CUTLASS_DEVICE - DirectEpilogueTensorOp( - Params const ¶ms, ///< Host-constructable params object - 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 - ): - output_op(params.output_op), - convert_op(params.convert_op), - destination_ref_(params.destination_ref), - source_ref_(params.source_ref) { - - - // Compute warp location within threadblock tile by mapping the warp_id to three coordinates: - // - // _m: the warp's position within the threadblock along the M dimension - // _n: the warp's position within the threadblock along the N dimension - // _k: the warp's position within the threadblock along the K dimension - - int warp_mn = warp_idx % (WarpCount::kM * WarpCount::kN); - int warp_m = warp_mn % WarpCount::kM; - int warp_n = warp_mn / WarpCount::kM; - - warp_origin_ = MatrixCoord{ - warp_m * Operator::Shape::kM, - warp_n * Operator::Shape::kN - }; - - destination_ref_.add_coord_offset(warp_origin_); - source_ref_.add_coord_offset(warp_origin_); - } - - /// Streams the result to global memory - CUTLASS_DEVICE - void operator()( - gemm::GemmCoord problem_size, ///< Problem size of GEMM (units of ElementC) - gemm::GemmCoord tb_tile_coord, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) - FragmentC const &accumulators) { ///< Accumulator tile - - MatrixCoord thread_origin = - MatrixCoord{tb_tile_coord.m() * Shape::kM, tb_tile_coord.n() * Shape::kN} + warp_origin_; - - /// Number of mma operations performed - using MmaIterations = MatrixShape< - Operator::Shape::kM / Operator::Policy::Operator::Shape::kM, - Operator::Shape::kN / Operator::Policy::Operator::Shape::kN - >; - - // Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire - // shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements - // of that row. The accumulators within one row are assumed to be consecutive. - int const kElementsPerAccess = Operator::Policy::Operator::Shape::kN / 4; - int const kRowsPerTile = 8; - int const kAccumulatorRows = Operator::Policy::Operator::Shape::kM / kRowsPerTile; - - CUTLASS_PRAGMA_UNROLL - for (int mma_n = 0; mma_n < MmaIterations::kN; ++mma_n) { - CUTLASS_PRAGMA_UNROLL - for (int mma_m = 0; mma_m < MmaIterations::kM; ++mma_m) { - - int mma_accum_start = kAccumulatorRows * kElementsPerAccess * - (mma_m * MmaIterations::kN + mma_n); - - CUTLASS_PRAGMA_UNROLL - for (int row = 0; row < kAccumulatorRows; ++row) { - CUTLASS_PRAGMA_UNROLL - for (int col = 0; col < kElementsPerAccess; ++col) { - - int accum_m = mma_m * Operator::Policy::Operator::Shape::kM + row * kRowsPerTile; - int accum_n = mma_n * Operator::Policy::Operator::Shape::kN + col; - int idx = mma_accum_start + row * kElementsPerAccess + col; - - MatrixCoord accum_coord = MatrixCoord{accum_m, accum_n}; - - MatrixCoord thread_coord = thread_origin + accum_coord; - - if (thread_coord < MatrixCoord{problem_size.m(), problem_size.n()}) { - - typename ConvertOp::result_type converted_accum = convert_op(accumulators[idx]); - - typename OutputOp::result_type output = output_op(converted_accum, source_ref_.at(accum_coord)); - - destination_ref_.at(accum_coord) = output; - } - } - } - } - } - } -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace threadblock -} // namespace epilogue -} // namespace cutlass - -///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/threadblock/direct_store_epilogue_iterator.h b/include/cutlass/epilogue/threadblock/direct_store_epilogue_iterator.h new file mode 100644 index 00000000..6e3a418b --- /dev/null +++ b/include/cutlass/epilogue/threadblock/direct_store_epilogue_iterator.h @@ -0,0 +1,135 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_types.h" +#include "cutlass/array.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/tensor.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/tensor_ref.h" +#include "cutlass/transform/pitch_linear_thread_map.h" +#include "cutlass/epilogue/threadblock/output_tile_thread_map.h" +#include "cutlass/arch/arch.h" +#include "cutlass/arch/memory.h" +#include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template +class DirectStoreEpilogueIterator { +public: + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = 1; + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + using Base = PredicatedTileIteratorParams; + + CUTLASS_HOST_DEVICE + Params() { } + + CUTLASS_HOST_DEVICE + Params(Layout const &layout) { + stride = layout.stride(0) * sizeof(Element); + } + + CUTLASS_HOST_DEVICE + Params(Base const &base) : + Base(base) { } + }; + +public: + + // + // Data members + // + + Element *pointer; // pointer to the output matrix + + LongIndex stride; // stride in elements between rows + + TensorCoord extent; // extent of output matrix + + int thread_idx; // thread index + + TensorCoord threadblock_offset; + +public: + + /// Constructor + CUTLASS_DEVICE + DirectStoreEpilogueIterator( + PredicatedTileIteratorParams const & params, + Element *pointer_, + TensorCoord extent_, + int thread_idx_, + TensorCoord threadblock_offset_ = TensorCoord() + ): + pointer(pointer_), + stride(params.stride / sizeof(Element)), + extent(extent_), + thread_idx(thread_idx_), + threadblock_offset(threadblock_offset_) + { + + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/threadblock/epilogue_direct_store.h b/include/cutlass/epilogue/threadblock/epilogue_direct_store.h new file mode 100644 index 00000000..028aee42 --- /dev/null +++ b/include/cutlass/epilogue/threadblock/epilogue_direct_store.h @@ -0,0 +1,342 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs and convolution using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_types.h" +#include "cutlass/array.h" + +#include "cutlass/gemm/gemm.h" + +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/epilogue/thread/conversion_op.h" +#include "cutlass/epilogue/thread/reduction_op.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template < + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename OutputTileIterator_, ///< Tile iterator reading and writing output tensors + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM + typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading from SMEM + typename OutputOp_ ///< Output operator +> +class EpilogueDirectStore { +public: + + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + using WarpShape = typename WarpMmaOperator_::Shape; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + using Padding = MatrixShape<0, 0>; + + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = typename 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; + + /// Number of warps + using WarpCount = gemm::GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + kPartitionsK + >; + + /// Use this to control the granularity of one epilogue 'iteration' + static int const kFragmentsPerIteration = 1; + + static int constexpr kSmemTiles = 1; + static int constexpr kSmemPointerOffset = 0; + + /// Shared storage allocation needed by the epilogue + struct SharedStorage { } ; + +private: + + // Assume accumulator tile is multipile interleaved 32x32 tile. + static int const kElementsPerPartial = 4; + using EleShapePerPatial = typename platform::conditional< + platform::is_same::value, + MatrixShape<2, 2>, + MatrixShape<1, 4> >::type; + static int const kElementsPerMma = 8; + static int const kAccumulatorPatials = 2; + using QuadShapePerPatialMma = MatrixShape<4, 4>; + + static_assert(OutputOp::kCount >= 2, + "The direct store epilogue for Tensor Ops requires the output functor have kCount >= 2."); + +private: + + LongIndex warp_offset; + int thread_idx; + int warp_idx; + int lane_idx; + int warp_m, warp_n; // warp coordinates within a cta + int tid_m, tid_n; // thread coordinates within a warp + +public: + + /// Constructor + CUTLASS_DEVICE + EpilogueDirectStore( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx_, ///< ID of a thread within the threadblock + int warp_idx_, ///< ID of warp within threadblock + int lane_idx_ ///< Id of thread within warp + ): + thread_idx(thread_idx_), + warp_idx(warp_idx_), + lane_idx(lane_idx_) + { + + // warp offsetting calculations + warp_offset = warp_idx * WarpShape::kM * WarpShape::kN; + int warp_id_mn = warp_idx % (WarpCount::kM * WarpShape::kN); + warp_m = warp_id_mn % WarpCount::kM; + warp_n = warp_id_mn / WarpCount::kM; + MatrixCoord warp_offset_coord(warp_m*WarpShape::kM, warp_n*WarpShape::kN); + + // thread offsetting calculations + int quad = (lane_idx >> 2); + int lane_in_quad = (lane_idx & 3); + + // this seems to be te correct layout + tid_m = quad; + tid_n = 2 * lane_in_quad; + } + + /// Streams the result to global memory + CUTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + if (!output_op.is_source_needed()) { + compute_source_not_needed_(output_op, destination_iterator, accumulators); + } + else { + compute_source_needed_(output_op, destination_iterator, accumulators, source_iterator); + } + } + +private: + + /// Streams the result to global memory + CUTLASS_DEVICE + void compute_source_needed_( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + const int kAccumBlockN = 2; + const int kThreadsM = 8; + const int kThreadsN = 4; + const int kBlockM = WarpShape::kM / kThreadsM; + + /// Array type used to output + using OutputAccessType = AlignedArray; + + /// Array type passed to the output operator - unused elements are optimized away + using OutputFragmentType = Array; + + /// Array type used by output functor + using AccumulatorAccessType = Array; + + /// Array type used by output functor + using AccumulatorFragmentType = Array; + + AccumulatorAccessType const *accumulator_pair = reinterpret_cast(&accumulators); + + CUTLASS_PRAGMA_UNROLL + for (int accum_m_idx = 0; accum_m_idx < WarpShape::kM / kThreadsM; accum_m_idx++) { + + int accum_m = kThreadsM * accum_m_idx; + int mL = destination_iterator.threadblock_offset.row() + WarpShape::kM * warp_m + tid_m + accum_m; + int nL_base = destination_iterator.threadblock_offset.column() + WarpShape::kN * warp_n + tid_n; + + ElementOutput *output_ptr = destination_iterator.pointer + mL * destination_iterator.stride; + ElementOutput *source_ptr = source_iterator.pointer + mL * source_iterator.stride; + + int const kIterationsN = WarpShape::kN / kThreadsN / kAccumBlockN; + + CUTLASS_PRAGMA_UNROLL + for (int accum_n_idx = 0; accum_n_idx < kIterationsN; accum_n_idx++) { + + int accum_idx = accum_m_idx + kBlockM * accum_n_idx; + int accum_n = kThreadsM * accum_n_idx; + + // mL and nL are logical coordinate in 2D mapping of epilogue's 4D output + int nL = nL_base + accum_n; + + bool guard = (mL < destination_iterator.extent.row()) && (nL < destination_iterator.extent.column()); + + AccumulatorFragmentType accum_fragment; + reinterpret_cast(accum_fragment) = accumulator_pair[accum_idx]; + + OutputFragmentType output_fragment; + + if(guard) { + reinterpret_cast(output_fragment) = + *reinterpret_cast(source_ptr + nL); + } + + // Perform output operator + output_fragment = output_op(accum_fragment, output_fragment); + + if(guard) { + // Store + *reinterpret_cast(output_ptr + nL) = reinterpret_cast(output_fragment); + } + } + } + } + + /// Streams the result to global memory + CUTLASS_DEVICE + void compute_source_not_needed_( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + const int kAccumBlockN = 2; + const int kThreadsM = 8; + const int kThreadsN = 4; + const int kBlockM = WarpShape::kM / kThreadsM; + + /// Array type used to output + using OutputAccessType = AlignedArray; + + /// Array type passed to the output operator - unused elements are optimized away + using OutputFragmentType = Array; + + /// Array type used by output functor + using AccumulatorAccessType = Array; + + /// Array type used by output functor + using AccumulatorFragmentType = Array; + + AccumulatorAccessType const *accumulator_pair = reinterpret_cast(&accumulators); + + CUTLASS_PRAGMA_UNROLL + for (int accum_m_idx = 0; accum_m_idx < WarpShape::kM / kThreadsM; accum_m_idx++) { + + int accum_m = kThreadsM * accum_m_idx; + int mL = destination_iterator.threadblock_offset.row() + WarpShape::kM * warp_m + tid_m + accum_m; + int nL_base = destination_iterator.threadblock_offset.column() + WarpShape::kN * warp_n + tid_n; + + ElementOutput *output_ptr = destination_iterator.pointer + mL * destination_iterator.stride; + + int const kIterationsN = WarpShape::kN / kThreadsN / kAccumBlockN; + + CUTLASS_PRAGMA_UNROLL + for (int accum_n_idx = 0; accum_n_idx < kIterationsN; accum_n_idx++) { + + int accum_idx = accum_m_idx + kBlockM * accum_n_idx; + int accum_n = kThreadsM * accum_n_idx; + + // mL and nL are logical coordinate in 2D mapping of epilogue's 4D output + int nL = nL_base + accum_n; + + bool guard = (mL < destination_iterator.extent.row()) && (nL < destination_iterator.extent.column()); + + AccumulatorFragmentType accum_fragment; + reinterpret_cast(accum_fragment) = accumulator_pair[accum_idx]; + + OutputFragmentType output_fragment; + + // Perform output operator + output_fragment = output_op(accum_fragment); + + if(guard) { + + // Store + *reinterpret_cast(output_ptr + nL) = + reinterpret_cast(output_fragment); + } + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/threadblock/epilogue_smem_accumulator.h b/include/cutlass/epilogue/threadblock/epilogue_smem_accumulator.h new file mode 100644 index 00000000..77d21f05 --- /dev/null +++ b/include/cutlass/epilogue/threadblock/epilogue_smem_accumulator.h @@ -0,0 +1,173 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMM/CONV to store accumulator in shared memory after + applying scale, bias loaded from global memory and element-wise operations. + + This Epilogue is typically used in fused GEMM/CONV to stage the intermediate accumulator. + +*/ + +#pragma once + +#if defined(__CUDACC_RTC__) +#include +#else +#include +#endif + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_types.h" +#include "cutlass/array.h" +#include "cutlass/layout/vector.h" +#include "cutlass/layout/tensor.h" +#include "cutlass/tensor_coord.h" +#include "cutlass/aligned_buffer.h" +#include "cutlass/functional.h" + +#include "cutlass/epilogue/warp/fragment_iterator_tensor_op.h" +#include "cutlass/epilogue/warp/tile_iterator_tensor_op.h" +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template < + typename SmemTileIterator_, ///< Shared memory Tile iterator to output to shared memory + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename ScaleBiasIterator_, ///< Iterator to load scale and bias from global memory + typename OutputOp_ ///< Output operator +> +class EpilogueSmemAccumulator { + +public: + + using SmemTileIterator = SmemTileIterator_; + + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + + using ScaleBiasIterator = ScaleBiasIterator_; + + using OutputOp = OutputOp_; + + /// Fragment of accumulator tile + using FragmentAccumulator = typename AccumulatorFragmentIterator::Fragment; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile; + + /// Fragment of Scale and Bias loaded from global memory + using FragmentScaleBias = typename ScaleBiasIterator::Fragment; + + static const bool PerChannelScale = (OutputOp::kScale == + epilogue::thread::ScaleType::OnlyAlphaPerChannelScaling); + + /// Constructor + CUTLASS_DEVICE + EpilogueSmemAccumulator() {} + + /// Streams the result to shared memory + CUTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + SmemTileIterator smem_iterator, ///< Tile iterator for destination in shared memory + AccumulatorTile const &accumulator, ///< Complete warp-level accumulator tile + ScaleBiasIterator scale_iterator, ///< iterator for scale vector in global memory + ScaleBiasIterator bias_iterator) { ///< iterator for bias vector in global memory + + + // Fragment to load scale bias from global memory + FragmentScaleBias tb_frag_scale; + FragmentScaleBias tb_frag_bias; + + /// Fragment Iterator to load slice of accumulator tile + AccumulatorFragmentIterator frag_iterator_accum(accumulator); + FragmentAccumulator tb_frag_accum; + + /// Epilogue output fragment + typename SmemTileIterator::Fragment tb_frag_smem; + + /// Load scale and bias from global memory + + if(PerChannelScale) + scale_iterator.load(tb_frag_scale); + + bias_iterator.load(tb_frag_bias); + + /// Iterate over the accumulator tile and store to shared memory + CUTLASS_PRAGMA_UNROLL + for (int rid = 0; rid < AccumulatorFragmentIterator::TileIterations::kRow; ++rid) { + + CUTLASS_PRAGMA_UNROLL + for (int cid = 0; cid < AccumulatorFragmentIterator::TileIterations::kColumn; ++cid) { + + using AccumulatorAccessType = typename OutputOp::FragmentAccumulator; + using ScaleBiasAccessType = typename OutputOp::FragmentScaleBias; + using FragmentSmemAccessType = typename OutputOp::FragmentOutput; + + + ScaleBiasAccessType const * scale_frag_ptr = + reinterpret_cast(&tb_frag_scale); + ScaleBiasAccessType const * bias_frag_ptr = + reinterpret_cast(&tb_frag_bias); + + FragmentSmemAccessType * smem_frag_ptr = + reinterpret_cast(&tb_frag_smem); + + for (int idx = 0; idx < AccumulatorFragmentIterator::kIterationsPerTile; ++idx) { + frag_iterator_accum.load(tb_frag_accum); + ++frag_iterator_accum; + + AccumulatorAccessType const * accumulator_frag_ptr = + reinterpret_cast(&tb_frag_accum); + const int kOutputIterations = FragmentAccumulator::kElements / OutputOp::kCount; + + CUTLASS_PRAGMA_UNROLL + for (int it = 0; it < kOutputIterations; it++) { + smem_frag_ptr[idx * kOutputIterations + it] = output_op(accumulator_frag_ptr[it], + scale_frag_ptr[cid * kOutputIterations + it], bias_frag_ptr[cid * kOutputIterations + it]); + } + } + + smem_iterator.store(tb_frag_smem); + ++smem_iterator; + + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// + diff --git a/include/cutlass/epilogue/warp/fragment_iterator_tensor_op.h b/include/cutlass/epilogue/warp/fragment_iterator_tensor_op.h index 835041fb..f234806c 100644 --- a/include/cutlass/epilogue/warp/fragment_iterator_tensor_op.h +++ b/include/cutlass/epilogue/warp/fragment_iterator_tensor_op.h @@ -94,6 +94,8 @@ public: /// Number of times this iterator can be incremented static int const kIterations = Policy::kIterations; + using TileIterations = typename Policy::TileIterations; + static int const kIterationsPerTile = kIterations / TileIterations::kCount; private: @@ -193,6 +195,8 @@ class FragmentIteratorTensorOp { // Number of externally visible iterations static int const kIterations = OperatorCount::kRow * kIterationsPerInstruction; + using TileIterations = MatrixShape; + static int const kAccumulatorRowStride = kElementsPerAccess; static int const kAccumulatorColumnStride = kElementsPerAccess * OperatorCount::kRow * kIterationsPerInstruction; + }; //////////////////////////////////////////////////////////////////////////////// @@ -120,6 +123,14 @@ struct TensorOpPolicy; }; //////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/warp/tile_iterator_tensor_op.h b/include/cutlass/epilogue/warp/tile_iterator_tensor_op.h index 4d34287d..5ec15575 100644 --- a/include/cutlass/epilogue/warp/tile_iterator_tensor_op.h +++ b/include/cutlass/epilogue/warp/tile_iterator_tensor_op.h @@ -68,6 +68,7 @@ public: using Element = Element_; using Layout = layout::RowMajor; + using TensorLayout = Layout; using TensorRef = TensorRef; ///< Tensor Reference object using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor using Index = typename TensorRef::Index; @@ -92,6 +93,9 @@ public: /// Number of times this iterator can be incremented static int const kIterations = Policy::kIterations; + /// Number of times this iterator can be incremented + using TileIterations = typename Policy::TileIterations; + // Internal constants struct Detail { static int const kLanesInQuad = 4; @@ -220,6 +224,198 @@ public: } }; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename Element_, ///< data type of element to be written + int InterleavedK ///< number of interleaved k +> +class TileIteratorTensorOp > { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using Element = Element_; + using Layout = layout::ColumnMajorInterleaved; + using TensorLayout = Layout; ///< shared memory tensor ref layout + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = TensorOpPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< +// Policy::kRowsPerIteration, + WarpShape::kM, + InterleavedK + >; + + /// This is the fragment size produced by one tile + using Fragment = Array< + Element, + Policy::OperatorCount::kRow * Policy::kElementsPerIteration>; + + /// This is the fragment size produced by one iteration +// using Fragment = Array< +// Element, Policy::kElementsPerIteration >; + + /// This is the complete warp-level accumulator tile. + //using AccumulatorTile = typename Operator::FragmentC; + + /// Number of times this iterator can be incremented + using TileIterations = typename Policy::TileIterations; + + // Internal constants + struct Detail { + static int const kLanesInQuad = 4; + }; + + /// Padding quantity + using Padding = MatrixShape< + 0, + Detail::kLanesInQuad * Policy::kElementsPerIteration>; + +private: + + /// Storage type for accessing memory + using AccessType = AlignedArray; + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointer_; + + /// Internal layout object + TensorLayout layout_; + + /// Thread offset + MatrixCoord thread_offset_; + +public: + + /// Default constructor + CUTLASS_HOST_DEVICE + TileIteratorTensorOp(): pointer_(nullptr) { } + + /// Constructor from TensorRef + CUTLASS_HOST_DEVICE + TileIteratorTensorOp( + TensorRef const &ref, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0]) { + + int quad_id = (lane_id / Detail::kLanesInQuad); + int lane_in_quad = (lane_id % Detail::kLanesInQuad); + + thread_offset_ = { + quad_id, lane_in_quad * Policy::kElementsPerIteration + }; + + pointer_ += (layout_({thread_offset_.row(), thread_offset_.column()}) / Policy::kElementsPerAccess); + } + + /// Adds a pointer offset + CUTLASS_HOST_DEVICE + TileIteratorTensorOp & add_pointer_offset(Index pointer_offset) { + pointer_ += pointer_offset / Policy::kElementsPerAccess; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + CUTLASS_HOST_DEVICE + TileIteratorTensorOp & add_tile_offset(TensorCoord const &tile_offset) { + + MatrixCoord coord_offset( + tile_offset.row() * Shape::kRow, + tile_offset.column() * Shape::kColumn + ); + + thread_offset_ += coord_offset; + + pointer_ += (layout_({ + coord_offset.row(), + coord_offset.column() + }) / Policy::kElementsPerAccess); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + CUTLASS_HOST_DEVICE + TileIteratorTensorOp & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + /// Store + CUTLASS_HOST_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::kRow * Policy::kIterationsPerInstruction; n++ ) { + + AccessType *ptr = pointer_ + layout_({n * Policy::kRowsPerIteration, 0}) / Policy::kElementsPerAccess; + + CUTLASS_PRAGMA_UNROLL + for (int a = 0; a < Policy::kAccessPerIteration; ++a) { + ptr[a + pointer_offset / Policy::kElementsPerAccess] = frag_ptr[n * Policy::kAccessPerIteration + a]; + +// printf("store thread %d, address %p, bank %ld\n", threadIdx.x, pointer_+a+n*Detail::kLanesInQuad, +// ((long long)(pointer_+a+n*Detail::kLanesInQuad)>>2)&0x1f); + } + } + } + + /// Store + CUTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Load + CUTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kRow * Policy::kIterationsPerInstruction; n++ ) { + + AccessType *ptr = pointer_ + layout_({n * Policy::kRowsPerIteration, 0}) / Policy::kElementsPerAccess; + + CUTLASS_PRAGMA_UNROLL + for (int a = 0; a < Policy::kAccessPerIteration; ++a) { + frag_ptr[n * Policy::kAccessPerIteration + a] = ptr[a + pointer_offset / Policy::kElementsPerAccess]; + } + } + } + + /// Load + CUTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + CUTLASS_HOST_DEVICE + TileIteratorTensorOp & operator++() { + return add_tile_offset({0, 1}); + } +}; + + ///////////////////////////////////////////////////////////////////////////////////////////////// /// Template for reading and writing tiles of accumulators to shared memory diff --git a/include/cutlass/fast_math.h b/include/cutlass/fast_math.h index d4ccf7de..b1650b94 100644 --- a/include/cutlass/fast_math.h +++ b/include/cutlass/fast_math.h @@ -48,6 +48,13 @@ namespace cutlass { ///////////////////////////////////////////////////////////////////////////////////////////////// +template +CUTLASS_HOST_DEVICE void swap(T &lhs, T &rhs) { + T tmp = lhs; + lhs = rhs; + rhs = tmp; +} + /****************************************************************************** * Static math utilities ******************************************************************************/ diff --git a/include/cutlass/gemm/device/gemm_grouped.h b/include/cutlass/gemm/device/gemm_grouped.h new file mode 100644 index 00000000..11290237 --- /dev/null +++ b/include/cutlass/gemm/device/gemm_grouped.h @@ -0,0 +1,313 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and + batched array variants. +*/ + +#pragma once + +#include + +#include "cutlass/cutlass.h" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/arch.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/device/default_gemm_configuration.h" + +#include "cutlass/trace.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// GEMM Grouped +template +class GemmGrouped { +public: + + using GemmKernel = GemmKernel_; + using ThreadblockShape = typename GemmKernel::Mma::Shape; + + using ElementA = typename GemmKernel::ElementA; + using LayoutA = typename GemmKernel::LayoutA; + using TensorRefA = TensorRef; + static ComplexTransform const kTransformA = GemmKernel::kTransformA; + + using ElementB = typename GemmKernel::ElementB; + using LayoutB = typename GemmKernel::LayoutB; + using TensorRefB = TensorRef; + static ComplexTransform const kTransformB = GemmKernel::kTransformB; + + using ElementC = typename GemmKernel::ElementC; + using LayoutC = typename GemmKernel::LayoutC; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + + using ElementAccumulator = typename GemmKernel::Mma::Policy::Operator::ElementC; + + using EpilogueOutputOp = typename GemmKernel::EpilogueOutputOp; + using ThreadblockSwizzle = typename GemmKernel::ThreadblockSwizzle; + using Operator = typename GemmKernel::Operator; + + /// Argument structure + using Arguments = typename GemmKernel::Arguments; + +protected: + + /// Kernel parameters object + typename GemmKernel::Params params_; + +public: + + /// Constructs the GEMM. + GemmGrouped() { } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + // Determine grid shape + cutlass::gemm::GemmCoord grid_tiled_shape; + int gemm_k_size = 0; + + get_grid_shape_(grid_tiled_shape, gemm_k_size, args); + + ThreadblockSwizzle threadblock_swizzle; + dim3 grid = threadblock_swizzle.get_grid_shape(grid_tiled_shape); + + if (!(grid.y <= std::numeric_limits::max() && + grid.z <= std::numeric_limits::max())) { + + return Status::kErrorInvalidProblem; + } + + return GemmKernel::can_implement(args); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + // This kerenl does not utilize a workspace + return size_t(); + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + + return dim3(args.threadblock_count, 1, 1); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + + CUTLASS_TRACE_HOST("GemmUniversalBase::maximum_active_blocks()"); + + int max_active_blocks = -1; + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + CUTLASS_TRACE_HOST(" smem_size: " << smem_size << " bytes"); + + if (smem_size <= (48 << 10)) { + + cudaError_t result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, + Kernel, + GemmKernel::kThreadCount, + smem_size); + + if (result == cudaSuccess) { + CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); + return max_active_blocks; + } + } + else { + + // Query assuming zero shared memory then compute occupancy limit based on SMEM + cudaError_t result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, + Kernel, + GemmKernel::kThreadCount, + 0); + + if (result != cudaSuccess) { + + CUTLASS_TRACE_HOST( + " cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error " + << cudaGetErrorString(result)); + + return -1; + } + + if (smem_capacity < 0) { + int device_idx = 0; + result = cudaGetDevice(&device_idx); + + if (result != cudaSuccess) { + return -1; + } + + cudaDeviceProp properties; + result = cudaGetDeviceProperties(&properties, device_idx); + + if (result != cudaSuccess) { + return -1; + } + + smem_capacity = static_cast(properties.sharedMemPerMultiprocessor); + } + + int occupancy = std::min(max_active_blocks, smem_capacity / smem_size); + + CUTLASS_TRACE_HOST(" occupancy: " << occupancy); + + return occupancy; + } + + CUTLASS_TRACE_HOST(" returning internal error"); + + return -1; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) { + + CUTLASS_TRACE_HOST("GemmUniversalBase::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + // Workspace + size_t workspace_bytes = get_workspace_size(args); + + if (workspace_bytes && !workspace) { + return Status::kErrorWorkspaceNull; + } + + // Initialize the Params structure + params_ = typename GemmKernel::Params(args, workspace); + + // Specify shared memory capacity for kernel. + 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) { + + size_t workspace_bytes = get_workspace_size(args); + + if (workspace_bytes && !workspace) { + return Status::kErrorWorkspaceNull; + } + + params_.update(args, workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(cudaStream_t stream = nullptr) { + + // + // Configure grid and block dimensions + // + + if (!params_.problem_visitor.problem_count) { + return Status::kSuccess; + } + + dim3 grid(params_.threadblock_count, 1, 1); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + // + // Launch kernel + // + + // Launch + cutlass::Kernel<<>>(params_); + + // + // Query for errors + // + cudaError_t result = cudaGetLastError(); + + if (result != cudaSuccess) { + CUTLASS_TRACE_HOST(" grid launch failed with error " << cudaGetErrorString(result)); + return Status::kErrorInternal; + } + + return Status::kSuccess; + } + + /// 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_adapter.h b/include/cutlass/gemm/device/gemm_universal_adapter.h index 3f0e3756..28375af2 100644 --- a/include/cutlass/gemm/device/gemm_universal_adapter.h +++ b/include/cutlass/gemm/device/gemm_universal_adapter.h @@ -32,6 +32,7 @@ #include "cutlass/cutlass.h" #include "cutlass/gemm/device/gemm_universal_base.h" +#include "cutlass/gemm/kernel/gemm_transpose_operands.h" //////////////////////////////////////////////////////////////////////////////// @@ -41,70 +42,6 @@ namespace device { ///////////////////////////////////////////////////////////////////////////////////////////////// -namespace detail { - - template < - typename ElementA_, - typename LayoutA_, - ComplexTransform TransformA, - int AlignmentA, - typename ElementB_, - typename LayoutB_, - ComplexTransform TransformB, - int AlignmentB, - typename LayoutC_, - bool Transpose - > - struct MapArguments { - using ElementA = ElementA_; - using LayoutA = LayoutA_; - static ComplexTransform const kTransformA = TransformA; - static int const kAlignmentA = AlignmentA; - using ElementB = ElementB_; - using LayoutB = LayoutB_; - static ComplexTransform const kTransformB = TransformB; - static int const kAlignmentB = AlignmentB; - using LayoutC = LayoutC_; - }; - - template < - typename ElementA_, - typename LayoutA_, - ComplexTransform TransformA, - int AlignmentA, - typename ElementB_, - typename LayoutB_, - ComplexTransform TransformB, - int AlignmentB, - typename LayoutC_ - > - struct MapArguments< - ElementA_, - LayoutA_, - TransformA, - AlignmentA, - ElementB_, - LayoutB_, - TransformB, - AlignmentB, - LayoutC_, - true - > { - using ElementA = ElementB_; - using LayoutA = typename layout::LayoutTranspose::type; - static ComplexTransform const kTransformA = TransformB; - static int const kAlignmentA = AlignmentB; - using ElementB = ElementA_; - using LayoutB = typename layout::LayoutTranspose::type; - static ComplexTransform const kTransformB = TransformA; - static int const kAlignmentB = AlignmentA; - using LayoutC = typename layout::LayoutTranspose::type; - }; - -}; - -///////////////////////////////////////////////////////////////////////////////////////////////// - template class GemmUniversalAdapter { public: @@ -129,7 +66,7 @@ public: using ArchTag = typename WarpMmaOperator::ArchTag; // Type, layout, and complex transform deliberately exchanged with B - using MapArguments = detail::MapArguments< + using MapArguments = kernel::detail::MapArguments< typename GemmKernel::ElementA, typename GemmKernel::LayoutA, GemmKernel::kTransformA, diff --git a/include/cutlass/gemm/kernel/default_gemm_complex.h b/include/cutlass/gemm/kernel/default_gemm_complex.h index f164c554..04cc2579 100644 --- a/include/cutlass/gemm/kernel/default_gemm_complex.h +++ b/include/cutlass/gemm/kernel/default_gemm_complex.h @@ -266,7 +266,6 @@ struct DefaultGemmComplex< using GemmKernel = kernel::Gemm; }; - //////////////////////////////////////////////////////////////////////////////// /// Partial specialization for Ampere Architecture diff --git a/include/cutlass/gemm/kernel/default_gemm_grouped.h b/include/cutlass/gemm/kernel/default_gemm_grouped.h new file mode 100644 index 00000000..664583f9 --- /dev/null +++ b/include/cutlass/gemm/kernel/default_gemm_grouped.h @@ -0,0 +1,356 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, CUTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. + +*/ + +#pragma once + +#include "cutlass/cutlass.h" + +#include "cutlass/complex.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/kernel/gemm_grouped.h" +#include "cutlass/gemm/kernel/gemm_transpose_operands.h" +#include "cutlass/gemm/kernel/default_gemm.h" +#include "cutlass/gemm/kernel/default_gemm_complex.h" +#include "cutlass/gemm/device/default_gemm_configuration.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator = typename device::DefaultGemmConfiguration< + OperatorClass, ArchTag, ElementA_, ElementB_, ElementC_, + ElementAccumulator>::Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// + typename Enable = void + > +struct DefaultGemmGrouped; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Real-valued GEMM kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear +> +struct DefaultGemmGrouped< + ElementA, + LayoutA, + ComplexTransform::kNone, // transform A + kAlignmentA, + ElementB, + LayoutB, + ComplexTransform::kNone, // transform B + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator, + SharedMemoryClear, + typename std::enable_if< ! cutlass::is_complex::value>::type +> { + + // If true, we must construct a 'transposed-and-exchanged' Mma operator. + static bool const kInternalTranspose = std::is_same::value; + + using MapArguments = kernel::detail::MapArguments< + ElementA, + LayoutA, + ComplexTransform::kNone, + kAlignmentA, + ElementB, + LayoutB, + ComplexTransform::kNone, + kAlignmentB, + LayoutC, + kInternalTranspose + >; + + // Define the default GEMM kernel + using DefaultGemmKernel = typename kernel::DefaultGemm< + typename MapArguments::ElementA, + typename MapArguments::LayoutA, + MapArguments::kAlignmentA, + typename MapArguments::ElementB, + typename MapArguments::LayoutB, + MapArguments::kAlignmentB, + ElementC, + typename MapArguments::LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + true, + Operator, + SharedMemoryClear + >::GemmKernel; + + /// Define the kernel in terms of the default kernel + using GemmKernel = kernel::GemmGrouped< + typename DefaultGemmKernel::Mma, + typename DefaultGemmKernel::Epilogue, + ThreadblockSwizzle, + kInternalTranspose + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Complex-valued GEMM kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear + > +struct DefaultGemmGrouped< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator, + SharedMemoryClear, + typename std::enable_if::value>::type +> { + + // If true, we must construct a 'transposed-and-exchanged' Mma operator. + static bool const kInternalTranspose = std::is_same::value; + + using MapArguments = kernel::detail::MapArguments< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + LayoutC, + kInternalTranspose + >; + + using DefaultGemmKernel = typename kernel::DefaultGemmComplex< + typename MapArguments::ElementA, + typename MapArguments::LayoutA, + typename MapArguments::ElementB, + typename MapArguments::LayoutB, + ElementC, + typename MapArguments::LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MapArguments::kTransformA, + MapArguments::kTransformB, + Operator, + false + >::GemmKernel; + + /// Define the kernel in terms of the default kernel + using GemmKernel = kernel::GemmGrouped< + typename DefaultGemmKernel::Mma, + typename DefaultGemmKernel::Epilogue, + ThreadblockSwizzle, + kInternalTranspose + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/kernel/gemm.h b/include/cutlass/gemm/kernel/gemm.h index b781dbcc..63c56db1 100644 --- a/include/cutlass/gemm/kernel/gemm.h +++ b/include/cutlass/gemm/kernel/gemm.h @@ -326,7 +326,6 @@ struct Gemm { semaphore.wait(threadblock_tile_offset.k()); - __threadfence(); } // Execute the epilogue operator to update the destination tensor. @@ -349,7 +348,6 @@ struct Gemm { lock = threadblock_tile_offset.k() + 1; } - __threadfence(); semaphore.release(lock); } } diff --git a/include/cutlass/gemm/kernel/gemm_grouped.h b/include/cutlass/gemm/kernel/gemm_grouped.h new file mode 100644 index 00000000..8b68c9d3 --- /dev/null +++ b/include/cutlass/gemm/kernel/gemm_grouped.h @@ -0,0 +1,615 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/fast_math.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/matrix_coord.h" +#include "cutlass/complex.h" +#include "cutlass/semaphore.h" + +#include "cutlass/layout/matrix.h" +#include "cutlass/trace.h" +#include "cutlass/gemm/kernel/gemm_transpose_operands.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Visitor class to abstract away the algorithm for iterating over tiles +template +struct GemmGroupedProblemVisitor { + + static bool const kTransposed = Transposed; + + struct Params { + cutlass::gemm::GemmCoord const *problem_sizes; + int32_t problem_count; + + // + // Methods + // + + /// Ctor + CUTLASS_HOST_DEVICE + Params(): problem_sizes(nullptr), problem_count(0) { } + + /// Ctor + CUTLASS_HOST_DEVICE + Params( + cutlass::gemm::GemmCoord const *problem_sizes, + int32_t problem_count + ): + problem_sizes(problem_sizes), + problem_count(problem_count) + {} + + }; + + struct SharedStorage { + // + // Nothing for now. As an optimization step, we could consider parallel + // argmin or prefix sums across the block. + // + }; + + // + // Data members + // + + Params const ¶ms; + SharedStorage &shared_storage; + cutlass::MatrixCoord threadblock_shape; + + int64_t tile_idx; + int64_t tile_count_sum; + int64_t problem_tile_start; + int32_t problem_idx; + + // + // Methods + // + CUTLASS_DEVICE + GemmGroupedProblemVisitor( + Params const ¶ms_, + SharedStorage &shared_storage_, + cutlass::MatrixCoord threadblock_shape_, + int32_t block_idx + ): + shared_storage(shared_storage_), + params(params_), + threadblock_shape(threadblock_shape_), + tile_idx(block_idx), + tile_count_sum(0), + problem_idx(0) + { + + cutlass::gemm::GemmCoord problem = problem_size(); + cutlass::gemm::GemmCoord grid = grid_shape(problem); + + problem_tile_start = 0; + tile_count_sum = grid.m() * grid.n(); + } + + /// Get the grid shape + CUTLASS_HOST_DEVICE + static cutlass::gemm::GemmCoord grid_shape( + cutlass::gemm::GemmCoord problem, + cutlass::MatrixCoord const & block_shape) { + + return cutlass::gemm::GemmCoord( + ((problem.m() - 1 + block_shape.row()) / block_shape.row()), + ((problem.n() - 1 + block_shape.column()) / block_shape.column()), + 1); + } + + /// Get the grid shape + CUTLASS_DEVICE + cutlass::gemm::GemmCoord grid_shape(cutlass::gemm::GemmCoord const &problem) const { + return grid_shape(problem, threadblock_shape); + } + + /// Returns true if there is a tile to compute + CUTLASS_DEVICE + bool next_tile() { + + if (tile_idx < tile_count_sum) { + return true; + } + + do { + ++problem_idx; + + if (problem_idx >= params.problem_count) { + return false; + } + + cutlass::gemm::GemmCoord problem = problem_size(); + cutlass::gemm::GemmCoord grid = grid_shape(problem); + + int64_t tile_count = grid.m() * grid.n(); + + problem_tile_start = tile_count_sum; + tile_count_sum += tile_count; + + } while (tile_count_sum <= tile_idx); + + return true; + } + + /// Gets the global tile index + CUTLASS_HOST_DEVICE + int64_t tile_index() const { + return tile_idx; + } + + /// Gets the index of the problem + CUTLASS_HOST_DEVICE + int32_t problem_index() const { + return problem_idx; + } + + /// Returns the problem size for the current problem + CUTLASS_HOST_DEVICE + cutlass::gemm::GemmCoord problem_size() const { + GemmCoord problem = params.problem_sizes[problem_idx]; + + if (kTransposed) { + swap(problem.m(), problem.n()); + } + + return problem; + } + + CUTLASS_HOST_DEVICE + int64_t threadblock_index() const { + return tile_idx - problem_tile_start; + } + + CUTLASS_DEVICE + void advance(int32_t grid_size) { + tile_idx += grid_size; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + bool Transposed = false +> +struct GemmGrouped { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static bool const kTransposed = Transposed; + + // Optional transpose + using MapArguments = kernel::detail::MapArguments< + typename Mma::IteratorA::Element, + typename Mma::IteratorA::Layout, + Mma::kTransformA, + Mma::IteratorA::AccessType::kElements, + typename Mma::IteratorB::Element, + typename Mma::IteratorB::Layout, + Mma::kTransformB, + Mma::IteratorB::AccessType::kElements, + typename Mma::LayoutC, + kTransposed + >; + + // Public-facing type definitions related to operand element type, layout, and complex conjugate + // operation. Must interact with the 'kTransposed' notion. + using ElementA = typename MapArguments::ElementA; + using LayoutA = typename MapArguments::LayoutA; + using ElementB = typename MapArguments::ElementB; + using LayoutB = typename MapArguments::LayoutB; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename MapArguments::LayoutC; + + static ComplexTransform const kTransformA = MapArguments::kTransformA; + static ComplexTransform const kTransformB = MapArguments::kTransformB; + + // Type definitions about the mainloop. + using Operator = typename Mma::Operator; + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = MapArguments::kAlignmentA; + static int const kAlignmentB = MapArguments::kAlignmentB; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + // + // Structures + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord *problem_sizes; + int problem_count; + int threadblock_count; + + typename EpilogueOutputOp::Params output_op; + + ElementA ** ptr_A; + ElementB ** ptr_B; + ElementC ** ptr_C; + ElementC ** ptr_D; + + typename LayoutA::Stride::LongIndex *lda; + typename LayoutB::Stride::LongIndex *ldb; + typename LayoutC::Stride::LongIndex *ldc; + typename LayoutC::Stride::LongIndex *ldd; + + // + // Methods + // + + /// Default ctor + CUTLASS_HOST_DEVICE + Arguments(): + problem_count(0), + threadblock_count(0), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + lda(nullptr), + ldb(nullptr), + ldc(nullptr), + ldd(nullptr) + { + + } + + /// Ctor + CUTLASS_HOST_DEVICE + Arguments( + GemmCoord *problem_sizes, + int problem_count, + int threadblock_count, + typename EpilogueOutputOp::Params output_op, + ElementA ** ptr_A, + ElementB ** ptr_B, + ElementC ** ptr_C, + ElementC ** ptr_D, + typename LayoutA::Stride::LongIndex *lda, + typename LayoutB::Stride::LongIndex *ldb, + typename LayoutC::Stride::LongIndex *ldc, + typename LayoutC::Stride::LongIndex *ldd + ): + problem_sizes(problem_sizes), + problem_count(problem_count), + threadblock_count(threadblock_count), + output_op(output_op), + ptr_A(ptr_A), + ptr_B(ptr_B), + ptr_C(ptr_C), + ptr_D(ptr_D), + lda(lda), + ldb(ldb), + ldc(ldc), + ldd(ldd) + { + + } + }; + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params { + + typename GemmGroupedProblemVisitor::Params problem_visitor; + int threadblock_count; + + typename EpilogueOutputOp::Params output_op; + + ElementA ** ptr_A; + ElementB ** ptr_B; + ElementC ** ptr_C; + ElementC ** ptr_D; + + typename LayoutA::Stride::LongIndex *lda; + typename LayoutB::Stride::LongIndex *ldb; + typename LayoutC::Stride::LongIndex *ldc; + typename LayoutC::Stride::LongIndex *ldd; + + + // + // Methods + // + + CUTLASS_HOST_DEVICE + Params(): + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + lda(nullptr), + ldb(nullptr), + ldc(nullptr), + ldd(nullptr) + { } + + CUTLASS_HOST_DEVICE + Params(Arguments const &args, void *workspace = nullptr): + problem_visitor(args.problem_sizes, args.problem_count), + threadblock_count(args.threadblock_count), + output_op(args.output_op), + ptr_A(args.ptr_A), + ptr_B(args.ptr_B), + ptr_C(args.ptr_C), + ptr_D(args.ptr_D), + lda(args.lda), + ldb(args.ldb), + ldc(args.ldc), + ldd(args.ldd) + { + + } + + CUTLASS_HOST_DEVICE + void update( + Arguments const &args, + void *workspace = nullptr) { + + problem_visitor = typename GemmGroupedProblemVisitor::Params(args.problem_sizes, args.problem_count); + threadblock_count = args.threadblock_count; + output_op = args.output_op; + ptr_A = args.ptr_A; + ptr_B = args.ptr_B; + ptr_C = args.ptr_C; + ptr_D = args.ptr_D; + lda = args.lda; + ldb = args.ldb; + ldc = args.ldc; + ldd = args.ldd; + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename GemmGroupedProblemVisitor::SharedStorage problem_visitor; + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + +public: + + // + // Methods + // + + CUTLASS_DEVICE + GemmGrouped() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement(cutlass::gemm::GemmCoord const & problem_size) { + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return Status::kSuccess; + } + + static size_t get_extra_workspace_size( + Arguments const &args, + cutlass::gemm::GemmCoord const &grid_tiled_shape) { + + return 0; + } + + /// Executes one GEMM + CUTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // + // These types shadow the type-level definitions and support the ability to implement + // a 'transposed' GEMM that computes the transposed problems. + // + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + + // + // Problem visitor. + // + GemmGroupedProblemVisitor problem_visitor( + params.problem_visitor, + shared_storage.problem_visitor, + {Mma::Shape::kM, Mma::Shape::kN}, + blockIdx.x); + + // Outer 'persistent' loop to iterate over tiles + while (problem_visitor.next_tile()) { + + GemmCoord problem_size = problem_visitor.problem_size(); + int32_t problem_idx = problem_visitor.problem_index(); + int32_t cta_idx = int32_t(problem_visitor.threadblock_index()); + + GemmCoord grid_shape = problem_visitor.grid_shape(problem_size); + + cutlass::gemm::GemmCoord threadblock_offset( + int(cta_idx / grid_shape.n()) * Mma::Shape::kM, + int(cta_idx % grid_shape.n()) * Mma::Shape::kN, + 0); + + // Load element pointers. Exchange pointers and strides if working on the transpose + ElementA *ptr_A = reinterpret_cast((kTransposed ? params.ptr_B[problem_idx] : params.ptr_A[problem_idx])); + typename LayoutA::LongIndex ldm_A = (kTransposed ? params.ldb[problem_idx] : params.lda[problem_idx]); + + ElementB *ptr_B = reinterpret_cast((kTransposed ? params.ptr_A[problem_idx] : params.ptr_B[problem_idx])); + typename LayoutB::LongIndex ldm_B = (kTransposed ? params.lda[problem_idx] : params.ldb[problem_idx]); + + // Compute initial location in logical coordinates + cutlass::MatrixCoord tb_offset_A{ + threadblock_offset.m(), + 0, + }; + + cutlass::MatrixCoord tb_offset_B{ + 0, + threadblock_offset.n() + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + LayoutA(ldm_A), + ptr_A, + {problem_size.m(), problem_size.k()}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + LayoutB(ldm_B), + ptr_B, + {problem_size.k(), problem_size.n()}, + thread_idx, + tb_offset_B); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); + + int lane_idx = threadIdx.x % 32; + + // + // Matrix multiply phase + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_B, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + ElementC *ptr_C = params.ptr_C[problem_idx]; + ElementC *ptr_D = params.ptr_D[problem_idx]; + + LayoutC layout_C(params.ldc[problem_idx]); + LayoutC layout_D(params.ldd[problem_idx]); + + typename Epilogue::OutputTileIterator::Params params_C(layout_C); + typename Epilogue::OutputTileIterator::Params params_D(layout_D); + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params_C, + ptr_C, + problem_size.mn(), + thread_idx, + threadblock_offset.mn() + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params_D, + ptr_D, + problem_size.mn(), + thread_idx, + threadblock_offset.mn() + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + // Next tile + problem_visitor.advance(gridDim.x); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/kernel/gemm_transpose_operands.h b/include/cutlass/gemm/kernel/gemm_transpose_operands.h new file mode 100644 index 00000000..69e25eee --- /dev/null +++ b/include/cutlass/gemm/kernel/gemm_transpose_operands.h @@ -0,0 +1,118 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and + batched array variants. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA_, + typename LayoutA_, + ComplexTransform TransformA, + int AlignmentA, + typename ElementB_, + typename LayoutB_, + ComplexTransform TransformB, + int AlignmentB, + typename LayoutC_, + bool Transpose +> +struct MapArguments { + using ElementA = ElementA_; + using LayoutA = LayoutA_; + static ComplexTransform const kTransformA = TransformA; + static int const kAlignmentA = AlignmentA; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + static ComplexTransform const kTransformB = TransformB; + static int const kAlignmentB = AlignmentB; + using LayoutC = LayoutC_; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA_, + typename LayoutA_, + ComplexTransform TransformA, + int AlignmentA, + typename ElementB_, + typename LayoutB_, + ComplexTransform TransformB, + int AlignmentB, + typename LayoutC_ +> +struct MapArguments< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + LayoutC_, + true +> { + using ElementA = ElementB_; + using LayoutA = typename layout::LayoutTranspose::type; + static ComplexTransform const kTransformA = TransformB; + static int const kAlignmentA = AlignmentB; + using ElementB = ElementA_; + using LayoutB = typename layout::LayoutTranspose::type; + static ComplexTransform const kTransformB = TransformA; + static int const kAlignmentB = AlignmentA; + using LayoutC = typename layout::LayoutTranspose::type; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} +} +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/kernel/gemm_universal.h b/include/cutlass/gemm/kernel/gemm_universal.h index 2da93857..eeeeb05c 100644 --- a/include/cutlass/gemm/kernel/gemm_universal.h +++ b/include/cutlass/gemm/kernel/gemm_universal.h @@ -36,6 +36,8 @@ #include "cutlass/complex.h" #include "cutlass/semaphore.h" +#include "cutlass/layout/matrix.h" + #include "cutlass/trace.h" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -574,7 +576,6 @@ public: semaphore.wait(threadblock_tile_offset.k()); - __threadfence(); } diff --git a/include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h b/include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h index d7963c0b..77e81fcb 100644 --- a/include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h +++ b/include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h @@ -626,8 +626,12 @@ public: else if (params.mode == GemmUniversalMode::kArray) { ptr_C = static_cast(params.ptr_C)[threadblock_tile_offset.k()]; ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; - ptr_Tensor = static_cast(params.ptr_Tensor)[threadblock_tile_offset.k()]; - ptr_Vector = static_cast(params.ptr_Vector)[threadblock_tile_offset.k()]; + if (ptr_Tensor) { + ptr_Tensor = static_cast(params.ptr_Tensor)[threadblock_tile_offset.k()]; + } + if (ptr_Vector) { + ptr_Vector = static_cast(params.ptr_Vector)[threadblock_tile_offset.k()]; + } } #endif @@ -679,7 +683,6 @@ public: semaphore.wait(threadblock_tile_offset.k()); - __threadfence(); } #endif diff --git a/include/cutlass/gemm/kernel/gemm_with_k_reduction.h b/include/cutlass/gemm/kernel/gemm_with_k_reduction.h index 1d376d2c..8f1600c2 100644 --- a/include/cutlass/gemm/kernel/gemm_with_k_reduction.h +++ b/include/cutlass/gemm/kernel/gemm_with_k_reduction.h @@ -579,7 +579,6 @@ public: semaphore.wait(threadblock_tile_offset.k()); - __threadfence(); } // Execute the epilogue operator to update the destination tensor. diff --git a/include/cutlass/gemm/threadblock/default_mma.h b/include/cutlass/gemm/threadblock/default_mma.h index d5af3aca..40b9b34b 100644 --- a/include/cutlass/gemm/threadblock/default_mma.h +++ b/include/cutlass/gemm/threadblock/default_mma.h @@ -378,21 +378,11 @@ struct DefaultMma>::value, "simt epilogue must be row major"); - static cutlass::arch::CacheOperation::Kind const CacheOpA = - ((sizeof_bits::value * kAlignmentA) == 128) - ? cutlass::arch::CacheOperation::Global - : cutlass::arch::CacheOperation::Always; - - static cutlass::arch::CacheOperation::Kind const CacheOpB = - ((sizeof_bits::value * kAlignmentB) == 128) - ? cutlass::arch::CacheOperation::Global - : cutlass::arch::CacheOperation::Always; - // Define the MmaCore components using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore< ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, ElementB, LayoutB, ElementAccumulator, LayoutC, arch::OpClassSimt, - Stages, Operator, false, CacheOpA, CacheOpB>; + Stages, Operator>; // Define iterators over tiles from the A operand using ThreadMapA = typename MmaCore::IteratorThreadMapA; diff --git a/include/cutlass/gemm/threadblock/default_mma_core_sm80.h b/include/cutlass/gemm/threadblock/default_mma_core_sm80.h index 70b1764a..221f5a8e 100644 --- a/include/cutlass/gemm/threadblock/default_mma_core_sm80.h +++ b/include/cutlass/gemm/threadblock/default_mma_core_sm80.h @@ -1111,8 +1111,8 @@ struct DefaultMmaCore< using ElementC = complex; using LayoutC = LayoutC_; static int const kStages = Stages; - static cutlass::arch::CacheOperation::Kind const kCacheOpA = cutlass::arch::CacheOperation::Global; - static cutlass::arch::CacheOperation::Kind const kCacheOpB = cutlass::arch::CacheOperation::Global; + static cutlass::arch::CacheOperation::Kind const kCacheOpA = cutlass::arch::CacheOperation::Always; + static cutlass::arch::CacheOperation::Kind const kCacheOpB = cutlass::arch::CacheOperation::Always; static const ComplexTransform TransformA = TransformA_; static const ComplexTransform TransformB = TransformB_; diff --git a/include/cutlass/gemm/threadblock/default_multistage_mma_complex.h b/include/cutlass/gemm/threadblock/default_multistage_mma_complex.h index 870123b9..f8a3e544 100644 --- a/include/cutlass/gemm/threadblock/default_multistage_mma_complex.h +++ b/include/cutlass/gemm/threadblock/default_multistage_mma_complex.h @@ -116,22 +116,11 @@ struct DefaultMultistageMmaComplex { - - static cutlass::arch::CacheOperation::Kind const CacheOpA = - (sizeof_bits::value == 128) - ? cutlass::arch::CacheOperation::Global - : cutlass::arch::CacheOperation::Always; - - static cutlass::arch::CacheOperation::Kind const CacheOpB = - (sizeof_bits::value == 128) - ? cutlass::arch::CacheOperation::Global - : cutlass::arch::CacheOperation::Always; - // Define the MmaCore components using MmaCore = typename cutlass::gemm::threadblock::DefaultMultistageMmaComplexCore< ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, ElementB, LayoutB, ElementAccumulator, layout::RowMajor, OperatorClass, - Stages, TransformA, TransformB, Operator, CacheOpA, CacheOpB>; + Stages, TransformA, TransformB, Operator>; // Define iterators over tiles from the A operand using ThreadMapA = typename MmaCore::IteratorThreadMapA; @@ -157,8 +146,6 @@ struct DefaultMultistageMmaComplex; }; -//////////////////////////////////////////////////////////////////////////////// - } // namespace threadblock } // namespace gemm } // namespace cutlass diff --git a/include/cutlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h b/include/cutlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h index 266b7824..45a068cd 100644 --- a/include/cutlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h +++ b/include/cutlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h @@ -113,8 +113,8 @@ struct DefaultMultistageMmaComplexCore< static ComplexTransform const kTransformA = TransformA; static ComplexTransform const kTransformB = TransformB; using Operator = Operator_; - static cutlass::arch::CacheOperation::Kind const kCacheOpA = cutlass::arch::CacheOperation::Global; - static cutlass::arch::CacheOperation::Kind const kCacheOpB = cutlass::arch::CacheOperation::Global; + static cutlass::arch::CacheOperation::Kind const kCacheOpA = cutlass::arch::CacheOperation::Always; + static cutlass::arch::CacheOperation::Kind const kCacheOpB = cutlass::arch::CacheOperation::Always; /// Number of warps present using WarpCount = GemmShape plus_accum; + + FragmentC tmp_accum; + + if (platform::is_same::value + || platform::is_same::value) { + + tmp_accum.clear(); + } + // // Mainloop // @@ -488,12 +503,30 @@ public: warp_loaded_frag_A[warp_mma_k % 2], warp_loaded_frag_B[warp_mma_k % 2]); - warp_mma( - accum, - warp_transformed_frag_A[warp_mma_k % 2], - warp_transformed_frag_B[warp_mma_k % 2], - accum - ); + if (platform::is_same::value + || platform::is_same::value) { + + warp_mma( + tmp_accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + tmp_accum + ); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + } // Issue global->shared copies for the this stage if (warp_mma_k < Base::kWarpGemmIterations - 1) { @@ -568,7 +601,14 @@ public: } } - + + if (platform::is_same::value + || platform::is_same::value) { + accum = plus_accum(accum, tmp_accum); + } + if (SharedMemoryClear == SharedMemoryClearOption::kZfill) { // commit and drain all pending and predicated LDGSTS pnz from the GEMM mainloop cutlass::arch::cp_async_fence(); diff --git a/include/cutlass/gemm/warp/default_mma_complex_tensor_op.h b/include/cutlass/gemm/warp/default_mma_complex_tensor_op.h index f90278c3..d9f2010e 100644 --- a/include/cutlass/gemm/warp/default_mma_complex_tensor_op.h +++ b/include/cutlass/gemm/warp/default_mma_complex_tensor_op.h @@ -30,6 +30,7 @@ #include "cutlass/cutlass.h" #include "cutlass/gemm/warp/mma_complex_tensor_op.h" +#include "cutlass/gemm/warp/mma_complex_tensor_op_fast_f32.h" #include "cutlass/gemm/warp/mma_gaussian_complex_tensor_op.h" #include "cutlass/layout/tensor_op_multiplicand_sm80.h" @@ -396,6 +397,73 @@ struct DefaultMmaComplexTensorOp< TransformB>; }; +///////////////////////////////////////////////////////////////////////////////////////////////// +/// 3xTF32 or 4xTF32 (fast and accurate complex operation) +/// Partial specialization - input and output types are complex * complex +// Use 3xTF32 or 4xTF32 tensor operation internally +// 4 real-valued MMA.1688.F32.TF32 operations on TF32 +// A = (ar + j ai), B (br +j bi), D = AB +// D = dr + j di = 3x[(ar*br - ai*bi) + j (ar*bi + ai*br)] +///////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB> +struct DefaultMmaComplexTensorOp< + WarpShape_, + InstructionShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + TransformA, + TransformB, + arch::OpMultiplyAddComplexFastF32> { + + // Complex floating point tensor operation use MMA.1688.F32.TF32 mma instruction + using Policy = cutlass::gemm::warp::MmaTensorOpPolicy< + cutlass::arch::Mma< + InstructionShape_, + 32, + tfloat32_t, + cutlass::layout::RowMajor, + tfloat32_t, + cutlass::layout::ColumnMajor, + float, + cutlass::layout::RowMajor, + arch::OpMultiplyAdd>, + cutlass::MatrixShape<1, 1> + >; + + // Define the warp-level tensor op + using Type = cutlass::gemm::warp::MmaComplexTensorOpFastF32< + WarpShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + Policy, + TransformA, + TransformB>; +}; + } // namespace warp } // namespace gemm } // namespace cutlass diff --git a/include/cutlass/gemm/warp/default_mma_tensor_op_sm80.h b/include/cutlass/gemm/warp/default_mma_tensor_op_sm80.h index 888f0b6e..479a0046 100644 --- a/include/cutlass/gemm/warp/default_mma_tensor_op_sm80.h +++ b/include/cutlass/gemm/warp/default_mma_tensor_op_sm80.h @@ -32,6 +32,7 @@ #include "cutlass/numeric_types.h" #include "cutlass/arch/mma.h" #include "cutlass/gemm/warp/mma_tensor_op.h" +#include "cutlass/gemm/warp/mma_tensor_op_fast_f32.h" #include "cutlass/gemm/warp/default_mma_tensor_op.h" ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -175,6 +176,51 @@ struct DefaultMmaTensorOp< ///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial Specialization - inputs and output types are float - uses TF32 for Fast Accurate FP32 +template < + /// Shape of one matrix production operation (concept: GemmShape) + typename WarpShape_, + /// Shape of target matrix multiply instruction (concept: GemmShape) + typename InstructionShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Number of partitions along K dimension + int PartitionsK, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor> +struct DefaultMmaTensorOp< + WarpShape_, + InstructionShape_, + float, LayoutA, + float, LayoutB, + float, LayoutC, + arch::OpMultiplyAddFastF32, PartitionsK, AccumulatorsInRowMajor> { + + // Uses TF32 internally + using Policy = cutlass::gemm::warp::MmaTensorOpPolicy< + cutlass::arch::Mma< + InstructionShape_, + 32, + cutlass::tfloat32_t, cutlass::layout::RowMajor, + cutlass::tfloat32_t, cutlass::layout::ColumnMajor, + float, cutlass::layout::RowMajor, + arch::OpMultiplyAdd + >, + cutlass::MatrixShape<1, 1> >; + + // Define the warp-level tensor op + using Type = cutlass::gemm::warp::MmaTensorOpFastF32< + WarpShape_, float, LayoutA, float, LayoutB, float, LayoutC, + Policy, PartitionsK, AccumulatorsInRowMajor>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + } // namespace warp } // namespace gemm } // namespace cutlass diff --git a/include/cutlass/gemm/warp/mma_complex_tensor_op.h b/include/cutlass/gemm/warp/mma_complex_tensor_op.h index 2f66eae8..2467e8f3 100644 --- a/include/cutlass/gemm/warp/mma_complex_tensor_op.h +++ b/include/cutlass/gemm/warp/mma_complex_tensor_op.h @@ -327,7 +327,7 @@ public: using InstructionShape = typename ArchMmaOperator::Shape; /// Indicates math operator - using MathOperator = typename ArchMmaOperator::Operator; + using MathOperator = arch::OpMultiplyAddComplex; /// Complex transform on A operand static ComplexTransform const kTransformA = TransformA; @@ -622,7 +622,7 @@ public: using OperatorClass = arch::OpClassTensorOp; /// Indicates math operator - using MathOperator = typename ArchMmaOperator::Operator; + using MathOperator = typename arch::OpMultiplyAddComplex; /// Complex transform on A operand static ComplexTransform const kTransformA = TransformA; diff --git a/include/cutlass/gemm/warp/mma_complex_tensor_op_fast_f32.h b/include/cutlass/gemm/warp/mma_complex_tensor_op_fast_f32.h new file mode 100644 index 00000000..666b00d5 --- /dev/null +++ b/include/cutlass/gemm/warp/mma_complex_tensor_op_fast_f32.h @@ -0,0 +1,657 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations targeting + Tensor Cores. +*/ + +#pragma once + +#include "cutlass/cutlass.h" + +#include "cutlass/array.h" +#include "cutlass/complex.h" +#include "cutlass/numeric_types.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/functional.h" + +#include "cutlass/arch/memory_sm75.h" +#include "cutlass/arch/mma_sm75.h" +#include "cutlass/arch/mma_sm80.h" + +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/warp/mma.h" + +#include "cutlass/gemm/warp/mma_tensor_op_policy.h" +#include "cutlass/gemm/warp/mma_tensor_op.h" + +#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" +#include "cutlass/gemm/warp/mma_complex_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace warp { + +namespace detail { + +template < + /// Data type of real & imag members of complex numbers in the SourceFragment + typename RealElement, + /// Destination fragment required by the mma operation + typename DestinationFragment, + /// Source fragment holding complex elements + typename SourceFragment, + /// Number of mma operations performed + typename MmaIterations, + /// Shape of operand elements + typename MmaOperandShape, + /// Complex transform on A operand + ComplexTransform Transform_, + /// Operand A or Operand B + Operand Operand_, + /// Floating-point rounding style for big part + FloatRoundStyle RoundBig_, + /// Floating-point rounding style for small part + FloatRoundStyle RoundSmall_> +struct UnpackComplexConvertAndPackForMmaFastF32; + +// Partial specialization for OperandA and Congruous smem layout +template < + typename RealElement, + typename DestinationFragment, + typename SourceFragment, + typename MmaIterations, + typename MmaOperandShape, + ComplexTransform Transform_, + FloatRoundStyle RoundBig_, + FloatRoundStyle RoundSmall_> +struct UnpackComplexConvertAndPackForMmaFastF32 < + RealElement, + DestinationFragment, + SourceFragment, + MmaIterations, + MmaOperandShape, + Transform_, + Operand::kA, + RoundBig_, + RoundSmall_> { + + // + // Type definitions + // + static Operand const kOperand = Operand::kA; + static ComplexTransform const kTransform = Transform_; + static FloatRoundStyle const kRoundBig = RoundBig_; + static FloatRoundStyle const kRoundSmall = RoundSmall_; + + // Data type of elements in the destination fragment + using MmaElement = typename DestinationFragment::Element; + + // Numeric convertor MmaElementBig, MmaElementSmall <= RealElement + using Converter = NumericConverterFastF32; + + // Operand layout parameters + using SourceFragmentLayout = layout::ColumnMajor; + static int const kLdm = MmaIterations::kRow * MmaOperandShape::kRow; + + // BigSmall Fragment holding two TF32 elements (big, small) for every float + using BigSmallFragment = Array; + + /// Index in fargments for the big and small part + static int const kBigIndex = 0; + static int const kSmallIndex = 1; + + /// Ctor + CUTLASS_DEVICE + UnpackComplexConvertAndPackForMmaFastF32() {} + + CUTLASS_DEVICE + void operator()(DestinationFragment *dest, SourceFragment const &source) { + + Converter convert_op; + SourceFragmentLayout layout(kLdm); + + DestinationFragment *dest_big_ = reinterpret_cast(dest); + DestinationFragment *dest_small_ = reinterpret_cast(&dest[MmaIterations::kRow * 2]); + + CUTLASS_PRAGMA_UNROLL + for(int i=0; i and apply rounding on real and imag parts + BigSmallFragment a = convert_op(source[layout(MatrixCoord{row,col})].real()); + BigSmallFragment b = convert_op(source[layout(MatrixCoord{row,col})].imag()); + + // Unpack rounded complex and pack into DestinationFragment for mma operation + dest_big_[i][pos] = a[kBigIndex]; + dest_big_[i+MmaIterations::kRow][pos] = (kTransform == ComplexTransform::kConjugate ? -b[kBigIndex] : b[kBigIndex]); + + // Unpack rounded complex and pack into DestinationFragment for mma operation + dest_small_[i][pos] = a[kSmallIndex]; + dest_small_[i+MmaIterations::kRow][pos] = (kTransform == ComplexTransform::kConjugate ? -b[kSmallIndex] : b[kSmallIndex]); + + // Next position + pos++; + } + } + } + } +}; + +// Partial specialization for OperandB and Congruous smem layout +template < + typename RealElement, + typename DestinationFragment, + typename SourceFragment, + typename MmaIterations, + typename MmaOperandShape, + ComplexTransform Transform_, + FloatRoundStyle RoundBig_, + FloatRoundStyle RoundSmall_> +struct UnpackComplexConvertAndPackForMmaFastF32 < + RealElement, + DestinationFragment, + SourceFragment, + MmaIterations, + MmaOperandShape, + Transform_, + Operand::kB, + RoundBig_, + RoundSmall_> { + + // + // Type definitions + // + static Operand const kOperand = Operand::kB; + static ComplexTransform const kTransform = Transform_; + static FloatRoundStyle const kRoundBig = RoundBig_; + static FloatRoundStyle const kRoundSmall = RoundSmall_; + + // Data type of elements in the destination fragment + using MmaElement = typename DestinationFragment::Element; + + // Numeric convertor MmaElementBig, MmaElementSmall <= RealElement + using Converter = NumericConverterFastF32; + + // Operand layout parameters + using SourceFragmentLayout = layout::RowMajor; + static int const kLdm = MmaIterations::kColumn * MmaOperandShape::kColumn; + + // BigSmall Fragment holding two TF32 elements (big, small) for every float + using BigSmallFragment = Array; + + /// Index in fargments for the big and small part + static int const kBigIndex = 0; + static int const kSmallIndex = 1; + + /// Ctor + CUTLASS_DEVICE + UnpackComplexConvertAndPackForMmaFastF32() {} + + CUTLASS_HOST_DEVICE + void operator()(DestinationFragment *dest, SourceFragment const &source) { + + Converter convert_op; + SourceFragmentLayout layout(kLdm); + + DestinationFragment *dest_big_ = reinterpret_cast(dest); + DestinationFragment *dest_small_ = reinterpret_cast(&dest[MmaIterations::kColumn * 2]); + + CUTLASS_PRAGMA_UNROLL + for(int i=0; i apply rounding on real and imag parts + BigSmallFragment a = convert_op(source[layout(MatrixCoord{row,col})].real()); + BigSmallFragment b = convert_op(source[layout(MatrixCoord{row,col})].imag()); + + // Unpack rounded complex and pack into DestinationFragment for mma operation + dest_big_[i][pos] = a[kBigIndex]; + dest_big_[i+MmaIterations::kColumn][pos] = (kTransform == ComplexTransform::kConjugate ? -b[kBigIndex] : b[kBigIndex]); + + // Unpack rounded complex and pack into DestinationFragment for mma operation + dest_small_[i][pos] = a[kSmallIndex]; + dest_small_[i+MmaIterations::kColumn][pos] = (kTransform == ComplexTransform::kConjugate ? -b[kSmallIndex] : b[kSmallIndex]); + + // next position + pos++; + } + } + } + } +}; +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transform on B operand + ComplexTransform TransformB = ComplexTransform::kNone, + /// Used for partial specialization + typename Enable = bool +> +class MmaComplexTensorOpFastF32; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex*complex+complex => complex: +// Operands data type: complex +// Rounding: float -> tfloat32_t (round half_ulp_truncate nearest) +// Math instruction: MMA.1688.F32.TF32 +// Output data type: complex +// +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB, + /// Used for partial specialization + typename Enable +> +class MmaComplexTensorOpFastF32< + Shape_, + complex, + LayoutA_, + complex, + LayoutB_, + complex, + LayoutC_, + Policy_, + TransformA, + TransformB, + Enable> { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of members of complex multiplicand A + using RealElementA = float; + + /// Data type of multiplicand A + using ElementA = complex; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of members of complex multiplicand B + using RealElementB = float; + + /// Data type of multiplicand B + using ElementB = complex; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of members of complex accumulator matrix C + using RealElementC = float; + + /// Data type of accumulator matrix C + using ElementC = complex; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Underlying arch tag + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Indicates math operator + using MathOperator = arch::OpMultiplyAddComplexFastF32; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = TransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = TransformB; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + + + /// Tune F32 to TF32 big small conversion for complex operation + /// Different combination of big small conversin can cause different tradeoff + /// between speed and accuracy. Generally, use round_half_ulp_truncate can + /// improve the performance but hur the accuracy. + using ComplexFastF32 = FastF32 < + FloatRoundStyle::round_toward_zero, // kRoundBigA + FloatRoundStyle::round_half_ulp_truncate, // kRoundSmallA + FloatRoundStyle::round_toward_zero, // kRoundBigB + FloatRoundStyle::round_half_ulp_truncate, // kRoundSmallB + TensorFloat32Op::k3xTF32 // Number of TF32 operations + >; + + /// Index in fargments for the big and small part + static int const kBigIndex = 0; + static int const kSmallIndex = 1; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + MatrixShape, + Policy::OpDelta::kRow, + 32, + 1 + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + // (4 times the original FragmentA::kElements) + // (real_big), (imag_big), (real_small), (imag_small) + using TransformedFragmentA = Array; + + // Fragment bisecting big and small sections + // (real_big, imag_big), (real_small, imag_small) + using AccessTypeFragmentA = Array; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + MatrixShape, + Policy::OpDelta::kColumn, + 32, + 1 + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + // (4 times the original FragmentB::kElements) + // (real_big), (imag_big), (real_small), (imag_small) + using TransformedFragmentB = Array; + + // Fragment bisecting big and small sections + // (real_big, imag_big), (real_small, imag_small) + using AccessTypeFragmentB = Array; + + static_assert( + !(Shape::kM % ArchMmaOperator::Shape::kM) && + !(Shape::kN % ArchMmaOperator::Shape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + /// Number of complex products operations performed (one complex product needs four mma instructions) + using MmaIterations = MatrixShape< + Shape::kM / ArchMmaOperator::Shape::kM, + Shape::kN / ArchMmaOperator::Shape::kN + >; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpAccumulatorTileIterator< + MatrixShape, + ElementC, + LayoutC, + typename ArchMmaOperator::Shape, + typename Policy::OpDelta>; + + /// Storage for C tile, the accumulator. Note, regardless of multiplicand type, this + /// storage arrangement is to be considered 'planar complex' in the sense that all real-valued + /// parts are stored consecutively followed by all imaginary parts. This matches the structure + /// of Tensor Cores which are always real-valued matrix multiplies. + using FragmentC = typename IteratorC::Fragment; + + // + // Alias types for underlying real-valued matrix multiply operator + // + using InstMmaOperandA = typename ArchMmaOperator::FragmentA; + using InstMmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + static_assert(platform::is_same, typename ArchMmaOperator::Shape>::value, + "This implementation only supports MMA.1688 math instructions."); + + static_assert(InstMmaOperandA::kElements == 4, + "This implementation only supports math instructions in which exactly four element is needed for the A operand." + "We can geneneralize later."); + + static_assert(InstMmaOperandB::kElements == 2, + "This implementation only supports math instructions in which exactly two element is needed for the B operand." + "We can geneneralize later."); + +private: + + // + // Data members + // + + /// Underlying real-valued matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + CUTLASS_DEVICE + MmaComplexTensorOpFastF32() {} + + /// Performs a warp-level matrix multiply-accumulate operation + CUTLASS_DEVICE + void operator()( + FragmentC &D, + TransformedFragmentA const &A, + TransformedFragmentB const &B, + FragmentC const &C + ) const { + + AccessTypeFragmentA const *complex_A = reinterpret_cast(&A); + AccessTypeFragmentB const *complex_B = reinterpret_cast(&B); + + // + // Accumulate in place + // + D = C; + + + complex_mma_operator(D, complex_A[kSmallIndex], complex_B[kBigIndex], D); + + complex_mma_operator(D, complex_A[kBigIndex], complex_B[kSmallIndex], D); + + complex_mma_operator(D, complex_A[kBigIndex], complex_B[kBigIndex], D); + + if (ComplexFastF32::kPrecision == TensorFloat32Op::k4xTF32) + complex_mma_operator(D, complex_A[kSmallIndex], complex_B[kSmallIndex], D); + } + + /// Performs a warp-level matrix multiply-accumulate operation + CUTLASS_DEVICE + void complex_mma_operator( + FragmentC &D, + AccessTypeFragmentA const &complex_A, + AccessTypeFragmentB const &complex_B, + FragmentC const &C + ) const { + + // Instruction Operands A & B holding real part followed by imaginary part for mma operations + InstMmaOperandA const *operand_A = reinterpret_cast(&complex_A); + InstMmaOperandB const *operand_B = reinterpret_cast(&complex_B); + + + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + // mma(accum.real(), a.real(), b.real(), accum.real()); + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Real-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_A[m], operand_B[n], *accum); + } + + // mma(accum.imag(), a.real(), b.imag(), accum.imag()); + CUTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Complex-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_A[m], operand_B[n+MmaIterations::kColumn], *accum); + } + + // mma(accum.real(), a.imag(), -b.imag(), accum.real()) + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // negate OperandB to accumulate -(a.imag()*b.imag()) + // negating OperandB emits less instrucitons than negating OperandA as OperandB has less elements + negate negate_op; + + // Real-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_A[m+MmaIterations::kRow], negate_op(operand_B[n+MmaIterations::kColumn]), *accum); + } + + // mma(accum.imag(), a.imag(), b.real(), accum.imag()) + CUTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Complex-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_A[m+MmaIterations::kRow], operand_B[n], *accum); + } + } + } + + /// Transform the mma operands to the required types + CUTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + + detail::UnpackComplexConvertAndPackForMmaFastF32 < + RealElementA, + InstMmaOperandA, + FragmentA, + MmaIterations, + MatrixShape<2, 2>, + kTransformA, + Operand::kA, + ComplexFastF32::kRoundBigA, + ComplexFastF32::kRoundSmallA> convert_A; + + detail::UnpackComplexConvertAndPackForMmaFastF32 < + RealElementB, + InstMmaOperandB, + FragmentB, + MmaIterations, + MatrixShape<2, 1>, + kTransformB, + Operand::kB, + ComplexFastF32::kRoundBigB, + ComplexFastF32::kRoundSmallB> convert_B; + + // Convert Fragment[A|B] holding complex to InstMmaOperand[A|B] holding InstMmaOperand[A|B]::Element + convert_A(reinterpret_cast(&dst_A), A); + convert_B(reinterpret_cast(&dst_B), B); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/warp/mma_gaussian_complex_tensor_op.h b/include/cutlass/gemm/warp/mma_gaussian_complex_tensor_op.h index 3de0113b..e1c94f7b 100644 --- a/include/cutlass/gemm/warp/mma_gaussian_complex_tensor_op.h +++ b/include/cutlass/gemm/warp/mma_gaussian_complex_tensor_op.h @@ -160,7 +160,7 @@ public: using OperatorClass = arch::OpClassTensorOp; /// Indicates math operator - using MathOperator = typename ArchMmaOperator::Operator; + using MathOperator = arch::OpMultiplyAddGaussianComplex; /// Complex transform on A operand static ComplexTransform const kTransformA = TransformA; diff --git a/include/cutlass/gemm/warp/mma_simt_tile_iterator.h b/include/cutlass/gemm/warp/mma_simt_tile_iterator.h index 466cc7d4..727b4949 100644 --- a/include/cutlass/gemm/warp/mma_simt_tile_iterator.h +++ b/include/cutlass/gemm/warp/mma_simt_tile_iterator.h @@ -33,6 +33,9 @@ #include "cutlass/array.h" #include "cutlass/tensor_ref.h" #include "cutlass/matrix_shape.h" + +#include "cutlass/arch/memory_sm75.h" + #include "cutlass/layout/matrix.h" #include "cutlass/gemm/gemm.h" @@ -224,8 +227,15 @@ public: for (int k = 0; k < Iterations::kColumn; ++k) { CUTLASS_PRAGMA_UNROLL for (int m = 0; m < Iterations::kRow; ++m) { + + // This logic has been replaced with calls to inline PTX to guarantee vectorization. + #if 0 dst_ptr[m + k * Iterations::kRow] = *(ref_.data() + ref_.offset({m * Policy::WarpShape::kRow, k}) + pointer_offset / Policy::LaneMmaShape::kM); + #endif + + auto ptr = ref_.data() + ref_.offset({m * Policy::WarpShape::kRow, k}) + pointer_offset / Policy::LaneMmaShape::kM; + arch::shared_load(dst_ptr[m + k * Iterations::kRow], ptr); } } } @@ -680,8 +690,14 @@ public: for (int k = 0; k < Iterations::kRow; ++k) { CUTLASS_PRAGMA_UNROLL for (int n = 0; n < Iterations::kColumn; ++n) { + + #if 0 dst_ptr[n + k * Iterations::kColumn] = *(ref_.data() + ref_.offset({k, n * Policy::WarpShape::kColumn}) + pointer_offset / Policy::LaneMmaShape::kN); + #endif + + void const *ptr = ref_.data() + ref_.offset({k, n * Policy::WarpShape::kColumn}) + pointer_offset / Policy::LaneMmaShape::kN; + arch::shared_load(dst_ptr[n + k * Iterations::kColumn], ptr); } } } diff --git a/include/cutlass/gemm/warp/mma_tensor_op.h b/include/cutlass/gemm/warp/mma_tensor_op.h index 9a4d402c..1eb61fef 100644 --- a/include/cutlass/gemm/warp/mma_tensor_op.h +++ b/include/cutlass/gemm/warp/mma_tensor_op.h @@ -404,3 +404,6 @@ public: ///////////////////////////////////////////////////////////////////////////////////////////////// +#include "cutlass/gemm/warp/mma_tensor_op_fast_f32.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/warp/mma_tensor_op_fast_f32.h b/include/cutlass/gemm/warp/mma_tensor_op_fast_f32.h new file mode 100644 index 00000000..bbaf39a7 --- /dev/null +++ b/include/cutlass/gemm/warp/mma_tensor_op_fast_f32.h @@ -0,0 +1,465 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations targeting + Tensor Cores. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/array.h" +#include "cutlass/platform/platform.h" + +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" +#include "cutlass/matrix_shape.h" + +#include "cutlass/arch/mma_sm80.h" + +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/warp/mma.h" + +#include "cutlass/gemm/warp/mma_tensor_op_policy.h" +#include "cutlass/gemm/warp/mma_tensor_op.h" + +#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +enum class TensorFloat32Op { + k3xTF32, + k4xTF32 +}; + +template < + /// Floating-point rounding style + FloatRoundStyle RoundBigA_, + /// Floating-point rounding style + FloatRoundStyle RoundSmallA_, + /// Floating-point rounding style + FloatRoundStyle RoundBigB_ = RoundBigA_, + /// Floating-point rounding style + FloatRoundStyle RoundSmallB_ = RoundSmallA_, + /// Precision for TensorFloat32Op + // (k3xTF32: BigxBig, BigxSmall, SmallxBig) + // (k4xTF32: BigxBig, BigxSmall, SmallxBig, SmallxSmall) + TensorFloat32Op Precision_ = TensorFloat32Op::k3xTF32 + > +struct FastF32 { + + static FloatRoundStyle const kRoundBigA = RoundBigA_; + static FloatRoundStyle const kRoundSmallA = RoundSmallA_; + static FloatRoundStyle const kRoundBigB = RoundBigB_; + static FloatRoundStyle const kRoundSmallB = RoundSmallB_; + static TensorFloat32Op const kPrecision = Precision_; +}; + + +namespace detail { + + template< + int N, + FloatRoundStyle RoundBig = FloatRoundStyle::round_toward_zero, + FloatRoundStyle RoundSmall = FloatRoundStyle::round_half_ulp_truncate + > + struct ConvertAndPackAccurateF32 { + + /// Rounding styles for big and small part + static FloatRoundStyle const kRoundBig = RoundBig; + static FloatRoundStyle const kRoundSmall = RoundSmall; + + /// Converter type + using Converter = NumericConverterFastF32; + + /// Source fragement + using SourceFragment = Array; + + /// Destination fragment + using DestinationFragment = Array; + + /// Converter Fragment holding two tfloat32_t elements for every float + using ConverterFragment = Array; + + /// Index in fargments for the big and small part + static int const kBigIndex = 0; + static int const kSmallIndex = 1; + + CUTLASS_HOST_DEVICE + void operator()(SourceFragment const &source, + DestinationFragment &dst_big, + DestinationFragment &dst_small) { + + Converter convert_; + ConverterFragment result_; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + // convert source to result fragment + result_ = convert_(source[i]); + + // store converted result fragments to destination fragment + dst_big[i] = result_[kBigIndex]; + dst_small[i] = result_[kSmallIndex]; + } + } + }; +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK_ = 1, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Used for partial specialization + typename Enable = bool +> +class MmaTensorOpFastF32; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for float*float+float => float using TF32 TensorOps +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK_, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor, + /// Used for partial specialization + typename Enable +> +class MmaTensorOpFastF32< + Shape_, + float, LayoutA_, + float, LayoutB_, + float, LayoutC_, + Policy_, PartitionsK_, + AccumulatorsInRowMajor, Enable> { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = float; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = float; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = float; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Indicates math operator + using MathOperator = arch::OpMultiplyAddFastF32; + + /// Architecture tag from underlying instruction + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = ComplexTransform::kNone; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// Tune F32 to TF32 big small conversion for float operation + /// Different combination of big small conversin can cause different tradeoff + /// between speed and accuracy. Generally, use round_half_ulp_truncate can + /// improve the performance but hur the accuracy. + using MmaFastF32 = FastF32 < + FloatRoundStyle::round_toward_zero, // kRoundBigA + FloatRoundStyle::round_half_ulp_truncate, // kRoundSmallA + FloatRoundStyle::round_toward_zero, // kRoundBigB + FloatRoundStyle::round_half_ulp_truncate, // kRoundSmallB + TensorFloat32Op::k3xTF32 // Number of TF32 operations + >; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + MatrixShape, + Policy::OpDelta::kRow, + kThreadCount, + kPartitionsK + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = + Array; + + /// Fragment bisecting big and small sections + using AccessTypeFragmentA = + Array; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + MatrixShape, + Policy::OpDelta::kRow, + kThreadCount, + kPartitionsK + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + using TransformedFragmentB = + Array; + + /// Fragment bisecting big and small sections + using AccessTypeFragmentB = + Array; + + /// Index in fargments for the big and small part + static int const kBigIndex = 0; + static int const kSmallIndex = 1; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpAccumulatorTileIterator< + MatrixShape, ElementC, LayoutC, + typename ArchMmaOperator::Shape, typename Policy::OpDelta>; + + /// Storage for C tile + using FragmentC = typename IteratorC::Fragment; + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM, + (Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN + >; + +public: + + /// Underlying matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + CUTLASS_DEVICE + MmaTensorOpFastF32() {} + + /// Performs a warp-level matrix multiply-accumulate operation + CUTLASS_DEVICE + void operator()( + FragmentC &D, + TransformedFragmentA const &A, + TransformedFragmentB const &B, + FragmentC const &C + ) const { + + AccessTypeFragmentA const *ptr_A = reinterpret_cast(&A); + AccessTypeFragmentB const *ptr_B = reinterpret_cast(&B); + + // + // Accumulate in place + // + D = C; + + mma_operator(D, ptr_A[kSmallIndex], ptr_B[kBigIndex], D); + + mma_operator(D, ptr_A[kBigIndex], ptr_B[kSmallIndex], D); + + mma_operator(D, ptr_A[kBigIndex], ptr_B[kBigIndex], D); + + if (MmaFastF32::kPrecision == TensorFloat32Op::k4xTF32) + mma_operator(D, ptr_A[kSmallIndex], ptr_B[kSmallIndex], D); + } + + /// Performs a warp-level matrix multiply-accumulate operation + CUTLASS_DEVICE + void mma_operator( + FragmentC &D, + AccessTypeFragmentA const &A, + AccessTypeFragmentB const &B, + FragmentC const &C + ) const { + + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + + using MmaOperandA = typename ArchMmaOperator::FragmentA; + using MmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + MmaOperandA const *ptr_A = reinterpret_cast(&A); + MmaOperandB const *ptr_B = reinterpret_cast(&B); + MmaOperandC *ptr_D = reinterpret_cast(&D); + + // Serpentine visitation order maximizing reuse of Ra + CUTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // This allows to reuse of Rb when at serpentine turns + int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n); + + if (AccumulatorsInRowMajor) { // matrix B is reordered + mma( + ptr_D[n_serpentine + m * MmaIterations::kColumn], + ptr_A[m], + ptr_B[n_serpentine], + ptr_D[n_serpentine + m * MmaIterations::kColumn]); + } else { + mma( + ptr_D[m + n_serpentine * MmaIterations::kRow], + ptr_A[m], + ptr_B[n_serpentine], + ptr_D[m + n_serpentine * MmaIterations::kRow]); + } + } // end n loop + } // end m loop + #else + assert(0); + #endif + } + + /// Transform the mma operands to the required types + CUTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + + // + // Define conversions from source type to instruction type + // + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + + detail::ConvertAndPackAccurateF32< + FragmentA::kElements / 2, + MmaFastF32::kRoundBigA, + MmaFastF32::kRoundSmallA> convert_A; + + detail::ConvertAndPackAccurateF32< + FragmentB::kElements, + MmaFastF32::kRoundBigB, + MmaFastF32::kRoundSmallB> convert_B; + + Array *ptr_dst_B = + reinterpret_cast *>(&dst_B); + + convert_B(B, ptr_dst_B[0], ptr_dst_B[1]); + + Array *ptr_dst_A = + reinterpret_cast *>(&dst_A); + + Array const *ptr_A = + reinterpret_cast const *>(&A); + + convert_A(ptr_A[0], ptr_dst_A[0], ptr_dst_A[2]); + + convert_A(ptr_A[1], ptr_dst_A[1], ptr_dst_A[3]); + #else + assert(0); + #endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/warp/mma_tensor_op_fragment_iterator.h b/include/cutlass/gemm/warp/mma_tensor_op_fragment_iterator.h index aa0439f0..aa2806db 100644 --- a/include/cutlass/gemm/warp/mma_tensor_op_fragment_iterator.h +++ b/include/cutlass/gemm/warp/mma_tensor_op_fragment_iterator.h @@ -149,6 +149,12 @@ public: /// Accumulator Fragment object using AccumulatorFragment = Array; + /// Scale Bias Element Type + using ElementScaleBias = typename OutputOp::ElementCompute; + + /// Scale Bias Fragment object + using ScaleBiasFragment = Array; + private: @@ -156,6 +162,8 @@ private: using AccessType = Array; using FragmentAccessType = Array; + using ScaleBiasAccessType = Array; + private: // // Data members @@ -208,9 +216,6 @@ public: if (output_op.is_source_needed()) //beta must be zero assert(0); - FragmentAccessType src_fragment; - src_fragment.clear(); - FragmentAccessType *frag_ptr = reinterpret_cast(&frag); int index = index_ * MmaIterations::kCount; @@ -223,11 +228,43 @@ public: frag_ptr[m * MmaIterations::kColumn + n].clear(); if(!(is_residual_tile_ && index_ >= kResidualIndex)) - frag_ptr[m * MmaIterations::kColumn + n] = output_op(accumulators_[accumulator_access_offset], src_fragment); + frag_ptr[m * MmaIterations::kColumn + n] = output_op(accumulators_[accumulator_access_offset]); } } } + /// Loads a fragment from the referenced part of the accumulator tile + /// Then apply per-channel scale and bias + CUTLASS_HOST_DEVICE + void load(Fragment &frag, ScaleBiasFragment &scale, + ScaleBiasFragment &bias, OutputOp output_op) const { + + if (output_op.is_source_needed()) //beta must be zero + assert(0); + + FragmentAccessType *frag_ptr = reinterpret_cast(&frag); + ScaleBiasAccessType * scale_ptr = reinterpret_cast(&scale); + ScaleBiasAccessType * bias_ptr = reinterpret_cast(&bias); + + int index = index_ * MmaIterations::kCount; + + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; n++) { + for (int m = 0; m < MmaIterations::kRow; m++) { + int accumulator_access_offset = + n * AccumulatorIterations::kRow + m + index; + + frag_ptr[m * MmaIterations::kColumn + n].clear(); + if(!(is_residual_tile_ && index_ >= kResidualIndex)) + frag_ptr[m * MmaIterations::kColumn + n] = + output_op(accumulators_[accumulator_access_offset], + scale_ptr[n] /*scale*/, bias_ptr[n] /*bias*/); + } + } + } + + + }; // Partial specialization for row-major accumulator tile @@ -266,7 +303,7 @@ class MmaTensorOpFragmentIterator; + /// Scale Bias Element Type + using ElementScaleBias = typename OutputOp::ElementCompute; + + /// Scale Bias Fragment object + using ScaleBiasFragment = Array; + private: /// Internal access type using AccessType = Array; using FragmentAccessType = Array; + using ScaleBiasAccessType = Array; private: // @@ -406,9 +450,6 @@ public: if (output_op.is_source_needed()) //beta must be zero assert(0); - FragmentAccessType src_fragment; - src_fragment.clear(); - FragmentAccessType *frag_ptr = reinterpret_cast(&frag); int index = index_ * AccessIterations::kCount; @@ -427,7 +468,50 @@ public: frag_ptr[i*kIterationsPerAccess + j].clear(); if(!(is_residual_tile_ && index_ >= kResidualIndex)) - frag_ptr[i*kIterationsPerAccess + j] = output_op(accumulators_[accumulator_access_offset + j * kAccessStride], src_fragment); + frag_ptr[i*kIterationsPerAccess + j] = output_op(accumulators_[accumulator_access_offset + j * kAccessStride]); + } + index++; + } + } + + /// Loads a fragment from the referenced part of the accumulator tile + /// Then apply per-channel scale and bias + CUTLASS_HOST_DEVICE + void load(Fragment &frag, ScaleBiasFragment &scale, + ScaleBiasFragment & bias, OutputOp output_op) const { + + if (output_op.is_source_needed()) //beta must be zero + assert(0); + + FragmentAccessType *frag_ptr = reinterpret_cast(&frag); + ScaleBiasAccessType * scale_ptr = reinterpret_cast(&scale); + ScaleBiasAccessType * bias_ptr = reinterpret_cast(&bias); + + int index = index_ * AccessIterations::kCount; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < AccessIterations::kCount; i++) { + + int accumulator_access_offset = index / AccessIterations::kCount * (MmaIterations::kColumn * kIterationsPerInstruction) + + (index % AccessIterations::kCount) / (AccessIterations::kColumn * kIterationsPerInstruction) * + AccumulatorIterations::kColumn * kIterationsPerInstruction + + (index % (AccessIterations::kColumn * kIterationsPerInstruction)) / kIterationsPerInstruction * + (kIterationsPerInstruction * kIterationsPerAccess) + + (index % kIterationsPerInstruction); + + int scale_bias_offset = (index + % (kIterationsPerInstruction * AccessIterations::kColumn)) + * kIterationsPerAccess; + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < kIterationsPerAccess; j++) { + + + frag_ptr[i*kIterationsPerAccess + j].clear(); + if(!(is_residual_tile_ && index_ >= kResidualIndex)) + frag_ptr[i*kIterationsPerAccess + j] = output_op( + accumulators_[accumulator_access_offset + j * kAccessStride], + scale_ptr[scale_bias_offset + j], bias_ptr[scale_bias_offset + j]); } index++; } diff --git a/include/cutlass/gemm/warp/mma_with_reduction_tensor_op.h b/include/cutlass/gemm/warp/mma_with_reduction_tensor_op.h index 31eff8cd..b9bcf632 100644 --- a/include/cutlass/gemm/warp/mma_with_reduction_tensor_op.h +++ b/include/cutlass/gemm/warp/mma_with_reduction_tensor_op.h @@ -402,4 +402,3 @@ public: } // namespace cutlass ///////////////////////////////////////////////////////////////////////////////////////////////// - diff --git a/include/cutlass/numeric_conversion.h b/include/cutlass/numeric_conversion.h index 3110aae9..66722c2d 100644 --- a/include/cutlass/numeric_conversion.h +++ b/include/cutlass/numeric_conversion.h @@ -617,6 +617,49 @@ struct NumericConverter { } }; +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Conversion operator for float to tfloat32_t big and small values +// +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + FloatRoundStyle RoundBig = FloatRoundStyle::round_toward_zero, + FloatRoundStyle RoundSmall = FloatRoundStyle::round_half_ulp_truncate +> +struct NumericConverterFastF32 { + + // result_type holds big tfloat32_t at idx(0) and small tfloat32_t at idx(1) + using result_type = Array; + + // source data type + using source_type = float; + + // rounding styles for big and small part + static FloatRoundStyle const kRoundBig = RoundBig; + static FloatRoundStyle const kRoundSmall = RoundSmall; + + CUTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + result_type result; + NumericConverter convert_big_; + NumericConverter convert_small_; + + // convert and fill tfloat32_t big at idx 0 + result[0] = convert_big_(source); + + // convert and fill tfloat32_t small at idx 1 + result[1] = convert_small_(source - static_cast(result[0])); + + return result; + } + + CUTLASS_HOST_DEVICE + result_type operator()(source_type const &s) { + return convert(s); + } +}; + ///////////////////////////////////////////////////////////////////////////////////////////////// // // Conversion and Clamp operator for Integers diff --git a/include/cutlass/tfloat32.h b/include/cutlass/tfloat32.h index 97918c01..18afc72d 100644 --- a/include/cutlass/tfloat32.h +++ b/include/cutlass/tfloat32.h @@ -363,8 +363,8 @@ tfloat32_t operator+(tfloat32_t const& lhs, tfloat32_t const& rhs) { CUTLASS_HOST_DEVICE tfloat32_t operator-(tfloat32_t const& lhs) { - float x = -static_cast(lhs); - return static_cast(x); + float x = -reinterpret_cast(lhs); + return *reinterpret_cast(&x); } CUTLASS_HOST_DEVICE diff --git a/include/cutlass/transform/threadblock/predicated_vector_access_iterator.h b/include/cutlass/transform/threadblock/predicated_vector_access_iterator.h new file mode 100644 index 00000000..6eeb0a10 --- /dev/null +++ b/include/cutlass/transform/threadblock/predicated_vector_access_iterator.h @@ -0,0 +1,350 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates implementing computing the addresses of loading small + vectors from the global memory. +*/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/array.h" +#include "cutlass/coord.h" +#include "cutlass/layout/pitch_linear.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/tensor.h" +#include "cutlass/matrix_coord.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/tensor_ref.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedVectorAccessIterator +/// +template +class PredicatedVectorAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Vector access iterator specialized for vectors, e.g. scale and bias +/// Thread arrangements are for TensorOps +/// +template +class PredicatedVectorAccessIterator { + public: + + using Shape = Shape_; + using WarpShape = WarpShape_; + using Element = Element_; + using Layout = layout::PitchLinear; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + +// static int const kElementsPerAccess = 128 / sizeof_bits::value; + static int const kElementsPerAccess = ElementsPerAccess; + static int const kThreads = 32; + static int const kRowsPerIteration = 8; + static int const kThreadsPerRow = kThreads / kRowsPerIteration; + static int const kThreadsPerRowMask = 0x3; + static int const kIterations = WarpShape::kContiguous / (kThreadsPerRow * kElementsPerAccess); + static int const kWarpCountStrided = Shape::kStrided / WarpShape::kStrided; + + using AccessType = AlignedArray; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Internal pointer to first access of tile + BytePointer pointer_; + + /// Extent of tensor + TensorCoord extent_; + + /// pointer offset of each thread + TensorCoord thread_offset_; + + /// iteration index + LongIndex iteration_; + + public: + /// Constructs a vector access iterator + CUTLASS_HOST_DEVICE + PredicatedVectorAccessIterator( + /// Pointer to the start of the vector + ConstPointer pointer, + /// Extent of vector + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// ID of each participating warp + int warp_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : pointer_(reinterpret_cast( + const_cast(pointer))), + extent_(extent) { + + + int warp_offset = (warp_id / kWarpCountStrided) * WarpShape::kContiguous; + + // Per-thread offset in logical coordinates of tensor + + thread_offset_ = threadblock_offset + TensorCoord(warp_offset, 0) + + TensorCoord((thread_id & kThreadsPerRowMask) * kElementsPerAccess, 0); + + set_iteration_index(0); + } + + /// Construct a PredicatedVectorAccessIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + PredicatedVectorAccessIterator( + /// Pointer to start of vector + ConstPointer pointer, + /// Extent of vector + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + /// ID of each participating warp + int warp_id) + : PredicatedVectorAccessIterator(pointer, extent, thread_id, warp_id, + make_Coord(0, 0)) {} + + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iteration_ = index; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + CUTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + thread_offset_ = + thread_offset_ + + TensorCoord(WarpShape::kContiguous * tile_offset.contiguous(), 0); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType *get() const { + + return reinterpret_cast( + pointer_ + + ((thread_offset_.contiguous() + iteration_ * kThreadsPerRow * kElementsPerAccess) + * sizeof_bits::value / 8)); + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedVectorAccessIterator &operator++() { + ++iteration_; + if(iteration_ >= kIterations) + iteration_ = 0; + + return *this; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + void advance() { + add_tile_offset(TensorCoord(1, 0)); + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + PredicatedVectorAccessIterator operator++(int) { + PredicatedVectorAccessIterator self(*this); + operator++(); + return self; + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return ((thread_offset_.contiguous() + + iteration_ * kThreadsPerRow * kElementsPerAccess) < extent_.contiguous()); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedVectorAccessIterator for row-major data. +/// +template +class PredicatedVectorAccessIterator { + public: + + using Shape = Shape_; + using WarpShape = WarpShape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedVectorAccessIterator< + layout::PitchLinearShape, + layout::PitchLinearShape, + Element, + layout::PitchLinear, + ElementsPerAccess>; + + using AccessType = typename UnderlyingIterator::AccessType; + static int const kElementsPerAccess = UnderlyingIterator::kElementsPerAccess; + static int const kRowsPerIteration = UnderlyingIterator::kRowsPerIteration; + static int const kThreads = UnderlyingIterator::kThreads; + static int const kIterations = UnderlyingIterator::kIterations; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedVectorAccessIterator( + ///< Pointer to the start of the vector + ConstPointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< ID of each participating warp + int warp_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(pointer, layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, warp_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedVectorAccessIterator with zero threadblock offset + CUTLASS_HOST_DEVICE + PredicatedVectorAccessIterator( + ConstPointer pointer, ///< Pointer to the start of the vector + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + int warp_id ///< ID of each participating warp + ) + : PredicatedVectorAccessIterator(pointer, extent, thread_id, warp_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + CUTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + CUTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedVectorAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedVectorAccessIterator operator++(int) { + PredicatedVectorAccessIterator self(*this); + operator++(); + return self; + } + + /// Increment and return an instance to self. + CUTLASS_HOST_DEVICE + void advance() { + iterator_.advance(); + } + + /// Returns whether access is valid or not + CUTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + diff --git a/include/cutlass/transform/threadblock/vector_iterator.h b/include/cutlass/transform/threadblock/vector_iterator.h new file mode 100644 index 00000000..06c2d88e --- /dev/null +++ b/include/cutlass/transform/threadblock/vector_iterator.h @@ -0,0 +1,143 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template wraps the vector access iterator concept to load whole vector from tensors in + memory. This is typically used for per-channel scale and bias in convolution kernels. +*/ + +#pragma once + +#include "cutlass/transform/threadblock/predicated_vector_access_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +class VectorIterator { +public: + using VectorAccessIterator = VectorAccessIterator_; + + using Shape = typename VectorAccessIterator::Shape; + using Element = typename VectorAccessIterator::Element; + using Layout = typename VectorAccessIterator::Layout; + using TensorCoord = typename Layout::TensorCoord; + using AccessType = typename VectorAccessIterator::AccessType; + using TensorRef = typename VectorAccessIterator::TensorRef; + using Index = typename VectorAccessIterator::Index; + using LongIndex = typename VectorAccessIterator::LongIndex; + + static int const kElementsPerAccess = VectorAccessIterator::kElementsPerAccess; + static int const kRowsPerIteration = VectorAccessIterator::kRowsPerIteration; + static int const kThreads = VectorAccessIterator::kThreads; + static int const kIterations = VectorAccessIterator::kIterations; + + /// Fragment object to be loaded or stored + using Fragment = cutlass::Array< + Element, kElementsPerAccess * kIterations>; + +private: + + /// Internal state + VectorAccessIterator vector_access_iterator_; + +public: + + /// Constructor + CUTLASS_HOST_DEVICE + VectorIterator( + Element const *ptr, + TensorCoord extent, + int thread_idx, + int warp_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + vector_access_iterator_(ptr, extent, thread_idx, warp_idx, threadblock_offset) { } + + /// Advances to the next tile in memory. + CUTLASS_HOST_DEVICE + VectorIterator &operator++() { + vector_access_iterator_.advance(); + return *this; + } + + /// Advances to the next tile in memory. + CUTLASS_HOST_DEVICE + VectorIterator operator++(int) { + VectorIterator self(*this); + operator++(); + return self; + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + frag.clear(); + AccessType *frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + + cutlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[c], + vector_access_iterator_.get() + pointer_offset, + vector_access_iterator_.valid() + ); + + ++vector_access_iterator_; + } +// } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment &frag) { + vector_access_iterator_.set_iteration_index(0); + load_with_pointer_offset(frag, 0); + } + + CUTLASS_DEVICE + void advance() { + vector_access_iterator_.advance(); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/include/cutlass/transform/warp/vector_fragment_iterator.h b/include/cutlass/transform/warp/vector_fragment_iterator.h new file mode 100644 index 00000000..62d9ef82 --- /dev/null +++ b/include/cutlass/transform/warp/vector_fragment_iterator.h @@ -0,0 +1,277 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + + +/*! \file + \brief This defines a "fragment" iterator for visiting the fragments of a warp vector + that participate in one warp-level mma operation. + + Typically, this is used to access the scale/bias fragement of a warp-level mma operation. + The scale/bias vector is then partitioned into smaller fragments that can be fed into + next warp-level mma operation. + + This iterator is necessary to accomplish warp-level mma fusion where the scale/bias vector is + applied to the multiplicand for the next mma. + +*/ + +#pragma once + +#include "cutlass/cutlass.h" + +#include "cutlass/array.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/tensor.h" +#include "cutlass/numeric_conversion.h" + +namespace cutlass { +namespace transform { +namespace warp { + + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the input fragment tile shape (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + //// Number of elements per access when loading fragment + int ElementsPerAccess> +class VectorFragmentIterator; + + +// Partial specialization for PitchLinear layout tile + +template < + /// Size of the input fragment vector shape (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + //// Number of elements per access when loading fragment + int ElementsPerAccess> +class VectorFragmentIterator { + public: + + /// Size of the input threadblock tile shape (concept: MatrixShape) + using Shape = Shape_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = cutlass::layout::PitchLinear; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Number of participating threads + static int const kThreads = 32; + + static int const kElementsPerAccess = ElementsPerAccess; + static int const kRowsPerIteration = 8; + static int const kColumnsPerAccess = 8; + static int const kElementsPerIteration = kRowsPerIteration * InstructionShape::kK / kThreads; + static int const kAccessPerIteration = kElementsPerIteration / kElementsPerAccess; + + /// Number of iterations + using Iterations = MatrixShape; + +public: + + // + // Derived quantities + // + // All fragments have kElementsPerAccess scale followed by bias + + /// Fragment object holding a thread's part of a tile + /// This is the fragment size produced by one iteration of the iterator. + using Fragment = Array; + + /// Input threadblock fragment tile + using ThreadblockFragment = Array; + +private: + + /// Internal access type + using AccessType = Array; + +private: + // + // Data members + // + + /// Input threadblock fragment tile + AccessType const *iterator_; + + /// Internal index + int index_; + +public: + /// Constructs an iterator + CUTLASS_HOST_DEVICE + VectorFragmentIterator(ThreadblockFragment const &threadblock_frag) + : iterator_(reinterpret_cast(&threadblock_frag)), + index_(0) {} + + /// Add offset + CUTLASS_HOST_DEVICE + void add_offset(int index_offset) { + index_ += index_offset; + + if(index_ >= Iterations::kColumn) + index_ = 0; + } + + /// Increments + CUTLASS_HOST_DEVICE + VectorFragmentIterator &operator++() { + add_offset(1); + return *this; + } + + CUTLASS_HOST_DEVICE + void set_index(int idx) { + index_ = idx; + } + + /// Loads a fragment from the referenced part of the accumulator tile + CUTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int r = 0; r < Iterations::kRow; r++) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kAccessPerIteration; i++) { + + frag_ptr[i * Iterations::kRow + r].clear(); + frag_ptr[i * Iterations::kRow + r] = iterator_[index_ * kAccessPerIteration + i]; + } + } + } + +}; + +// Partial specialization for Row-Major layout tile + +template < + /// Size of the input fragment tile shape (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + //// Number of elements per access when loading fragment + int ElementsPerAccess> +class VectorFragmentIterator { + public: + + /// Size of the input threadblock tile shape (concept: MatrixShape) + using Shape = Shape_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = cutlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Underlying iterator + using Base = VectorFragmentIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear, InstructionShape, ElementsPerAccess>; + + + public: + + // + // Derived quantities + // + /// Fragment object holding a thread's part of a tile + /// This is the fragment size produced by one iteration of the iterator. + using Fragment = typename Base::Fragment; + + /// Input threadblock fragment tile + using ThreadblockFragment = typename Base::ThreadblockFragment; + + private: + /// Underlying iterator + Base iterator_; + +public: + /// Constructs an iterator + CUTLASS_HOST_DEVICE + VectorFragmentIterator(ThreadblockFragment const &threadblock_frag) + : iterator_(threadblock_frag) {} + + /// Add offset + CUTLASS_HOST_DEVICE + void add_offset(int index_offset) { + iterator_.add_offset(index_offset); + } + + /// Increments + CUTLASS_HOST_DEVICE + VectorFragmentIterator &operator++() { + add_offset(1); + return *this; + } + + CUTLASS_HOST_DEVICE + void set_index(int idx) { + iterator_.set_index(idx); + } + + /// Loads a fragment from the referenced part of the accumulator tile + CUTLASS_HOST_DEVICE + void load(Fragment &frag) const { + iterator_.load(frag); + } + +}; + + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace conv +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/media/docs/implicit_gemm_convolution.md b/media/docs/implicit_gemm_convolution.md index 8d8272ef..69862b3f 100644 --- a/media/docs/implicit_gemm_convolution.md +++ b/media/docs/implicit_gemm_convolution.md @@ -152,13 +152,13 @@ a hierarchical GEMM kernel targeting Tensor Cores. # CUTLASS Convolution Implementation -The CUTLASS Implicit GEMM implementation makes several assumptions. +To get the best performance, the following parameters are recommended. - All tensors are 128-bit aligned NHWC tensors - Channel count (C) is a multiple of 32 elements - Filter count (K) is a multiple of 32 elements -This enables 128-bit vector memory acceses which lead to efficient CUDA kernels. +This enables 128-bit vector memory acceses which lead to efficient CUDA kernels. Smaller alignment is supported even on tensor cores by setting AlignmentA and AlignmentB in conv::kernel::DefaultConv2dFprop, but the performance is lower than 128-bit aligned tesnors. # CUTLASS Device-level Convolution Operator @@ -187,7 +187,7 @@ using Conv2dFpropKernel = typename cutlass::conv::kernel::DefaultConv2dFprop< SwizzleThreadBlock, // optional function to reorder threadblocks for locality NumStages, // number of pipeline stages in threadblock-scoped GEMM cutlass::arch::OpMultiplyAddSaturate, // math operation on data of element a and b - cutlass::conv::IteratorAlgorithm::kAnalytic // globabl memory iterator algorithm + cutlass::conv::IteratorAlgorithm::kOptimized // globabl memory iterator algorithm >::Kernel ``` @@ -219,7 +219,7 @@ using Conv2dFpropKernel = typename cutlass::conv::kernel::DefaultConv2dFprop< SwizzleThreadBlock, // optional function to reorder threadblocks for locality 2, // number of pipeline stages in threadblock-scoped GEMM cutlass::arch::OpMultiplyAddSaturate, // math operation on data of element a and b - cutlass::conv::IteratorAlgorithm::kAnalytic // globabl memory iterator algorithm + cutlass::conv::IteratorAlgorithm::kOptimized // globabl memory iterator algorithm >::Kernel ``` diff --git a/test/unit/gemm/device/CMakeLists.txt b/test/unit/gemm/device/CMakeLists.txt index 9505fdb1..51f28f39 100644 --- a/test/unit/gemm/device/CMakeLists.txt +++ b/test/unit/gemm/device/CMakeLists.txt @@ -192,6 +192,7 @@ cutlass_test_unit_add_executable( gemm_f16t_f16t_f32t_tensor_op_f32_sm80.cu gemm_bf16n_bf16n_f32t_tensor_op_f32_sm80.cu gemm_bf16t_bf16t_bf16t_tensor_op_f32_sm80.cu + gemm_f16n_f16n_f16n_direct_store_tensor_op_f32_sm80.cu ) cutlass_test_unit_add_executable( @@ -319,6 +320,15 @@ cutlass_test_unit_add_executable( gemm_planar_complex_f16_f16_f32_tensor_op_sm80.cu ) +cutlass_test_unit_add_executable( + cutlass_test_unit_gemm_device_grouped + + BATCH_SOURCES ON + BATCH_SIZE 4 + + gemm_grouped_sm80.cu +) + cutlass_test_unit_add_executable( cutlass_test_unit_gemm_device_sparse_tensorop_sm80 diff --git a/test/unit/gemm/device/gemm_f16n_f16n_f16n_direct_store_tensor_op_f32_sm80.cu b/test/unit/gemm/device/gemm_f16n_f16n_f16n_direct_store_tensor_op_f32_sm80.cu new file mode 100644 index 00000000..98ec1baf --- /dev/null +++ b/test/unit/gemm/device/gemm_f16n_f16n_f16n_direct_store_tensor_op_f32_sm80.cu @@ -0,0 +1,108 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Tests for device-wide GEMM interface +*/ + +#include + +#include "../../common/cutlass_unit_test.h" +#include "cutlass/cutlass.h" + +#include "cutlass/gemm/kernel/gemm_universal.h" +#include "cutlass/gemm/device/gemm_universal.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" + +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/reference/host/gemm.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/tensor_view_io.h" + +#include "testbed_universal.h" + +//////////////////////////////////////////////////////////////////////////////// + +#include "cutlass/epilogue/threadblock/epilogue_direct_store.h" +#include "cutlass/epilogue/threadblock/default_epilogue_direct_store.h" + +//////////////////////////////////////////////////////////////////////////////// + +#if defined(CUTLASS_ARCH_MMA_SM80_SUPPORTED) + +//////////////////////////////////////////////////////////////////////////////// + +TEST(SM80_Device_GemmUniversal_DirectStore_f16n_f16t_f32n_tensor_op_f32, 128x128x32_64x64x32) { + + using ElementOutput = float; + using ElementAccumulator = float; + + // Define the GEMM kernel + using GemmBase = cutlass::gemm::device::GemmUniversal< + cutlass::half_t, + cutlass::layout::ColumnMajor, + cutlass::half_t, + cutlass::layout::RowMajor, + ElementOutput, cutlass::layout::ColumnMajor, + ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, + cutlass::gemm::GemmShape<128, 128, 32>, + cutlass::gemm::GemmShape<64, 64, 32>, + cutlass::gemm::GemmShape<16, 8, 16>, + cutlass::epilogue::thread::LinearCombination< + ElementOutput, + 4, // This is the vector size of the epilogue. + ElementAccumulator, + ElementAccumulator>, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 3, + 8, + 8 + >; + + // Define the direct store epilogue + using EpilogueDirectStore = typename cutlass::epilogue::threadblock::DefaultEpilogueDirectStore< + typename GemmBase::GemmKernel::Epilogue + >::Epilogue; + + // Define a new kernel + using Kernel = cutlass::gemm::kernel::GemmUniversal< + typename GemmBase::GemmKernel::Mma, + EpilogueDirectStore, + typename GemmBase::GemmKernel::ThreadblockSwizzle + >; + + // Define the adaptor + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + EXPECT_TRUE(test::gemm::device::TestAllGemmUniversal()); +} + +//////////////////////////////////////////////////////////////////////////////// + +#endif // #if defined(CUTLASS_ARCH_MMA_SM80_SUPPORTED) + +//////////////////////////////////////////////////////////////////////////////// + diff --git a/test/unit/gemm/device/gemm_f16n_f16t_f16t_tensor_op_f32_sm80.cu b/test/unit/gemm/device/gemm_f16n_f16t_f16t_tensor_op_f32_sm80.cu index 99b74189..e7f56a81 100644 --- a/test/unit/gemm/device/gemm_f16n_f16t_f16t_tensor_op_f32_sm80.cu +++ b/test/unit/gemm/device/gemm_f16n_f16t_f16t_tensor_op_f32_sm80.cu @@ -50,7 +50,6 @@ TEST(SM80_Device_GemmUniversal_f16n_f16t_f32t_tensor_op_f32, 64x64x32_32x32x32) { - /* using ElementOutput = cutlass::half_t; using ElementAccumulator = cutlass::half_t; @@ -66,7 +65,6 @@ TEST(SM80_Device_GemmUniversal_f16n_f16t_f32t_tensor_op_f32, 64x64x32_32x32x32) cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle, 10>; EXPECT_TRUE(test::gemm::device::TestAllGemm()); - */ } //////////////////////////////////////////////////////////////////////////////// diff --git a/test/unit/gemm/device/gemm_grouped_sm80.cu b/test/unit/gemm/device/gemm_grouped_sm80.cu new file mode 100644 index 00000000..609732c4 --- /dev/null +++ b/test/unit/gemm/device/gemm_grouped_sm80.cu @@ -0,0 +1,598 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Tests for device-wide GEMM interface + +*/ + +#include + +#include "../../common/cutlass_unit_test.h" +#include "cutlass/cutlass.h" + +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/kernel/gemm_grouped.h" +#include "cutlass/gemm/kernel/default_gemm_grouped.h" +#include "cutlass/gemm/device/gemm_grouped.h" + +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/reference/host/gemm.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/tensor_view_io.h" + +#include "testbed_grouped.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(CUTLASS_ARCH_MMA_SM80_SUPPORTED) + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Visitor class to abstract away the algorithm for iterating over tiles. +// +// This is the prototype. We will delete this when the efficient kernel is +// available. +struct GemmGroupedProblemVisitor { + + struct Params { + cutlass::gemm::GemmCoord const *problem_sizes; + int32_t problem_count; + int64_t const *tile_count; + }; + + struct SharedStorage { + // + // Nothing for now. As an optimization step, we could consider parallel + // argmin or prefix sums across the block. + // + }; + + // + // Data members + // + + SharedStorage &shared_storage; + Params const ¶ms; + cutlass::MatrixCoord threadblock_shape; + + int64_t tile_idx; + int64_t tile_count_sum; + int64_t problem_tile_start; + int32_t problem_idx; + + // + // Methods + // + CUTLASS_DEVICE + GemmGroupedProblemVisitor( + SharedStorage &shared_storage_, + Params const ¶ms_, + cutlass::MatrixCoord threadblock_shape_, + int32_t block_idx + ): + shared_storage(shared_storage_), + params(params_), + threadblock_shape(threadblock_shape_), + tile_idx(block_idx), + tile_count_sum(0), + problem_idx(0) + { + + cutlass::gemm::GemmCoord problem = params.problem_sizes[problem_idx]; + + cutlass::gemm::GemmCoord grid = grid_shape(problem); + + problem_tile_start = 0; + tile_count_sum = grid.m() * grid.n(); + } + + /// Get the grid shape + CUTLASS_HOST_DEVICE + static cutlass::gemm::GemmCoord grid_shape( + cutlass::gemm::GemmCoord const &problem, + cutlass::MatrixCoord const & block_shape) { + + return cutlass::gemm::GemmCoord( + ((problem.m() - 1 + block_shape.row()) / block_shape.row()), + ((problem.n() - 1 + block_shape.column()) / block_shape.column()), + 1); + } + + /// Get the grid shape + CUTLASS_DEVICE + cutlass::gemm::GemmCoord grid_shape(cutlass::gemm::GemmCoord const &problem) const { + return grid_shape(problem, threadblock_shape); + } + + /// Returns true if there is a tile to compute + CUTLASS_DEVICE + bool next_tile() { + + if (tile_idx < tile_count_sum) { + return true; + } + + do { + ++problem_idx; + + if (problem_idx >= params.problem_count) { + return false; + } + + cutlass::gemm::GemmCoord problem = params.problem_sizes[problem_idx]; + cutlass::gemm::GemmCoord grid = grid_shape(problem); + + int64_t tile_count = grid.m() * grid.n(); + + problem_tile_start = tile_count_sum; + tile_count_sum += tile_count; + + } while (tile_count_sum <= tile_idx); + + return true; + } + + /// Gets the global tile index + CUTLASS_HOST_DEVICE + int64_t tile_index() const { + return tile_idx; + } + + /// Gets the index of the problem + CUTLASS_HOST_DEVICE + int32_t problem_index() const { + return problem_idx; + } + + /// Returns the problem size for the current problem + CUTLASS_HOST_DEVICE + cutlass::gemm::GemmCoord problem_size() const { + return params.problem_sizes[problem_idx]; + } + + CUTLASS_HOST_DEVICE + int64_t threadblock_index() const { + return tile_idx - problem_tile_start; + } + + CUTLASS_DEVICE + void advance(int32_t grid_size) { + tile_idx += grid_size; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +__global__ void GroupedBatchedKernel(GemmGroupedProblemVisitor::Params params) { + + __shared__ GemmGroupedProblemVisitor::SharedStorage shared_storage; + + GemmGroupedProblemVisitor problem_visitor( + shared_storage, + params, + {CtaShapeM, CtaShapeN}, + blockIdx.x); + + while (problem_visitor.next_tile()) { + + cutlass::gemm::GemmCoord problem_size = problem_visitor.problem_size(); + int64_t cta_idx = problem_visitor.threadblock_index(); + + cutlass::gemm::GemmCoord grid_shape = problem_visitor.grid_shape(problem_size); + + int cta_tile_m_idx = int(cta_idx / grid_shape.n()); + int cta_tile_n_idx = int(cta_idx % grid_shape.n()); + + // + // Do the MMA + // + + if (threadIdx.x == 0) { + #if 0 + printf("Block %d - tile: %lld, problem %d, cta_idx: %lld, cta(m: %d, n: %d)\n", + blockIdx.x, + problem_visitor.tile_index(), + problem_visitor.problem_index(), + cta_idx, + cta_tile_m_idx, + cta_tile_n_idx); + #endif + } + + // Next tile + problem_visitor.advance(gridDim.x); + } +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM80_Device_GemmGrouped_scheduler, 64x64x32_32x32x32) { + + int32_t problem_count = 16; + + int const kCtaShapeM = 64; + int const kCtaShapeN = 64; + + std::vector problem_sizes(problem_count); + std::vector tile_counts(problem_count); + + // construct a few problems of random sizes + srand(1921); + for (int32_t i = 0; i < problem_count; ++i) { + problem_sizes.at(i) = cutlass::gemm::GemmCoord( + 8 * (rand() % 48) + 64, + 8 * (rand() % 48) + 64, + 8 * (rand() % 48) + 64); + } + + // compute prefix sum + int64_t tile_count = 0; + + for (int32_t i = 0; i < problem_count; ++i) { + + cutlass::gemm::GemmCoord grid_shape = GemmGroupedProblemVisitor::grid_shape( + problem_sizes.at(i), {kCtaShapeM, kCtaShapeN}); + + int32_t problem_tile_count = (grid_shape.m() * grid_shape.n()); + + int64_t tile_start = tile_count; + + tile_count += problem_tile_count; + tile_counts.at(i) = tile_count; + + if (false) { + std::cout << "Problem " << i << " size(" + << problem_sizes.at(i).m() << "-by-" << problem_sizes.at(i).n() + << ") - tiles: " << problem_tile_count << ", grid(" << grid_shape.m() << ", " << grid_shape.n() + << "), tiles[" << tile_start << ", " << tile_count << ")" << std::endl; + } + } + + // Copy to device memory + cutlass::DeviceAllocation problem_sizes_device(problem_count); + cutlass::DeviceAllocation tile_counts_device(problem_count); + + problem_sizes_device.copy_from_host(problem_sizes.data()); + tile_counts_device.copy_from_host(tile_counts.data()); + + GemmGroupedProblemVisitor::Params params; + params.problem_sizes = problem_sizes_device.get(); + params.problem_count = problem_count; + params.tile_count = tile_counts_device.get(); + + // Launch the kernel + dim3 grid(108, 1, 1); + dim3 block(128, 1, 1); + + GroupedBatchedKernel<<< grid, block >>>(params); + + // wait + cudaDeviceSynchronize(); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM80_Device_GemmGrouped_f16n_f16t_f32n_tensor_op_f32, 128x128x32_64x64x32) { + + using ElementOutput = float; + using ElementAccumulator = float; + + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + cutlass::half_t, + cutlass::layout::ColumnMajor, + cutlass::ComplexTransform::kNone, + 8, + cutlass::half_t, + cutlass::layout::ColumnMajor, + cutlass::ComplexTransform::kNone, + 8, + ElementOutput, cutlass::layout::ColumnMajor, + ElementAccumulator, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<128, 128, 32>, + cutlass::gemm::GemmShape<64, 64, 32>, + cutlass::gemm::GemmShape<16, 8, 16>, + cutlass::epilogue::thread::LinearCombination< + ElementOutput, 128 / cutlass::sizeof_bits::value, + ElementAccumulator, ElementAccumulator>, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 3>::GemmKernel; + + using Gemm = cutlass::gemm::device::GemmGrouped; + + // + // Test + // + + test::gemm::device::TestbedGrouped testbed; + + bool passed = testbed.run(24); + EXPECT_TRUE(passed); + +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM80_Device_GemmGrouped_f16t_f16n_f32n_tensor_op_f32, 128x64x32_64x32x32) { + + using ElementOutput = cutlass::half_t; + using ElementAccumulator = float; + + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + cutlass::half_t, + cutlass::layout::RowMajor, + cutlass::ComplexTransform::kNone, + 8, + cutlass::half_t, + cutlass::layout::ColumnMajor, + cutlass::ComplexTransform::kNone, + 8, + ElementOutput, cutlass::layout::RowMajor, + ElementAccumulator, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<128, 64, 32>, + cutlass::gemm::GemmShape<64, 32, 32>, + cutlass::gemm::GemmShape<16, 8, 16>, + cutlass::epilogue::thread::LinearCombination< + ElementOutput, 128 / cutlass::sizeof_bits::value, + ElementAccumulator, ElementAccumulator>, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 4>::GemmKernel; + + using Gemm = cutlass::gemm::device::GemmGrouped; + + // + // Test + // + + test::gemm::device::TestbedGrouped testbed; + + bool passed = testbed.run(27); + EXPECT_TRUE(passed); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM80_Device_GemmGrouped_f64t_f64t_f64n_tensor_op_f64, 64x64x16_32x32x16) { + + using ElementInput = double; + using ElementOutput = double; + using ElementAccumulator = double; + + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + ElementInput, + cutlass::layout::RowMajor, + cutlass::ComplexTransform::kNone, + 1, + ElementInput, + cutlass::layout::RowMajor, + cutlass::ComplexTransform::kNone, + 1, + ElementOutput, cutlass::layout::ColumnMajor, + ElementAccumulator, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<64, 64, 16>, + cutlass::gemm::GemmShape<32, 32, 16>, + cutlass::gemm::GemmShape<8, 8, 4>, + cutlass::epilogue::thread::LinearCombination< + ElementOutput, 1, + ElementAccumulator, ElementAccumulator>, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 4>::GemmKernel; + + using Gemm = cutlass::gemm::device::GemmGrouped; + + // + // Test + // + + test::gemm::device::TestbedGrouped testbed; + + bool passed = testbed.run(27); + EXPECT_TRUE(passed); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM80_Device_GemmGrouped_f32t_f32t_f32n_simt_f32, 128x128x8_64x32x1) { + + using ElementInput = float; + using ElementOutput = float; + using ElementAccumulator = float; + + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + ElementInput, + cutlass::layout::RowMajor, + cutlass::ComplexTransform::kNone, + 1, + ElementInput, + cutlass::layout::RowMajor, + cutlass::ComplexTransform::kNone, + 1, + ElementOutput, cutlass::layout::ColumnMajor, + ElementAccumulator, + cutlass::arch::OpClassSimt, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<128, 128, 8>, + cutlass::gemm::GemmShape<64, 32, 8>, + cutlass::gemm::GemmShape<1, 1, 1>, + cutlass::epilogue::thread::LinearCombination< + ElementOutput, 1, + ElementAccumulator, ElementAccumulator>, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 3>::GemmKernel; + + using Gemm = cutlass::gemm::device::GemmGrouped; + + // + // Test + // + + test::gemm::device::TestbedGrouped testbed; + + bool passed = testbed.run(27); + EXPECT_TRUE(passed); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM80_Device_GemmGrouped_cf32n_cf32n_cf32n_tensorop_f32, 64x64x16_32x32x16) { + + using ElementInput = cutlass::complex; + using ElementOutput = cutlass::complex; + using ElementAccumulator = cutlass::complex; + + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + ElementInput, + cutlass::layout::ColumnMajor, + cutlass::ComplexTransform::kNone, + 1, + ElementInput, + cutlass::layout::ColumnMajor, + cutlass::ComplexTransform::kNone, + 1, + ElementOutput, cutlass::layout::ColumnMajor, + ElementAccumulator, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<64, 64, 16>, + cutlass::gemm::GemmShape<32, 32, 16>, + cutlass::gemm::GemmShape<16, 8, 8>, + cutlass::epilogue::thread::LinearCombination< + ElementOutput, 1, + ElementAccumulator, ElementAccumulator>, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 3, + cutlass::arch::OpMultiplyAddComplex>::GemmKernel; + + using Gemm = cutlass::gemm::device::GemmGrouped; + + // + // Test + // + + test::gemm::device::TestbedGrouped testbed; + + bool passed = testbed.run(27); + EXPECT_TRUE(passed); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM80_Device_GemmGrouped_cf32c_cf32t_cf32c_tensorop_f32, 64x64x16_32x32x16) { + + using ElementInput = cutlass::complex; + using ElementOutput = cutlass::complex; + using ElementAccumulator = cutlass::complex; + + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + ElementInput, + cutlass::layout::ColumnMajor, + cutlass::ComplexTransform::kConjugate, + 1, + ElementInput, + cutlass::layout::ColumnMajor, + cutlass::ComplexTransform::kConjugate, + 1, + ElementOutput, cutlass::layout::ColumnMajor, + ElementAccumulator, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<64, 64, 16>, + cutlass::gemm::GemmShape<32, 32, 16>, + cutlass::gemm::GemmShape<16, 8, 8>, + cutlass::epilogue::thread::LinearCombination< + ElementOutput, 1, + ElementAccumulator, ElementAccumulator>, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 3, + cutlass::arch::OpMultiplyAddComplex>::GemmKernel; + + using Gemm = cutlass::gemm::device::GemmGrouped; + + // + // Test + // + + test::gemm::device::TestbedGrouped testbed; + + bool passed = testbed.run(27); + EXPECT_TRUE(passed); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(SM80_Device_GemmGrouped_cf32t_cf32h_cf32n_tensorop_f32, 64x64x16_16x16x16) { + + using ElementInput = cutlass::complex; + using ElementOutput = cutlass::complex; + using ElementAccumulator = cutlass::complex; + + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + ElementInput, + cutlass::layout::RowMajor, + cutlass::ComplexTransform::kNone, + 1, + ElementInput, + cutlass::layout::RowMajor, + cutlass::ComplexTransform::kConjugate, + 1, + ElementOutput, cutlass::layout::ColumnMajor, + ElementAccumulator, + cutlass::arch::OpClassTensorOp, + cutlass::arch::Sm80, + cutlass::gemm::GemmShape<32, 32, 16>, + cutlass::gemm::GemmShape<16, 16, 16>, + cutlass::gemm::GemmShape<8, 8, 4>, + cutlass::epilogue::thread::LinearCombination< + ElementOutput, 1, + ElementAccumulator, ElementAccumulator>, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 3, + cutlass::arch::OpMultiplyAddComplex>::GemmKernel; + + using Gemm = cutlass::gemm::device::GemmGrouped; + + // + // Test + // + + test::gemm::device::TestbedGrouped testbed; + + bool passed = testbed.run(27); + EXPECT_TRUE(passed); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#endif // #if defined(CUTLASS_ARCH_MMA_SM80_SUPPORTED) + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/test/unit/gemm/device/testbed_complex.h b/test/unit/gemm/device/testbed_complex.h index 6e4ae7dd..4834a6af 100644 --- a/test/unit/gemm/device/testbed_complex.h +++ b/test/unit/gemm/device/testbed_complex.h @@ -102,33 +102,33 @@ struct TestbedComplex : public Testbed { } /// Returns true if the CUDA device is sufficient to execute the kernel. - bool sufficient() const { - // - // Determine SMEM requirements and waive if not satisfied - // - - int smem_size = int(sizeof(typename Gemm::GemmKernel::SharedStorage)); - - cudaDeviceProp properties; - int device_idx; - cudaError_t result = cudaGetDevice(&device_idx); - - if (result != cudaSuccess) { - throw std::runtime_error("cudaGetDevice() API call failed."); - } - - result = cudaGetDeviceProperties(&properties, device_idx); - - if (result != cudaSuccess) { - throw std::runtime_error("cudaGetDeviceProperties() failed"); - } - - if (properties.sharedMemPerMultiprocessor < smem_size) { - return false; - } - - return true; - } + bool sufficient() const { + // + // Determine SMEM requirements and waive if not satisfied + // + + int smem_size = int(sizeof(typename Gemm::GemmKernel::SharedStorage)); + + cudaDeviceProp properties; + int device_idx; + cudaError_t result = cudaGetDevice(&device_idx); + + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDevice() API call failed."); + } + + result = cudaGetDeviceProperties(&properties, device_idx); + + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDeviceProperties() failed"); + } + + if (properties.sharedMemPerMultiprocessor < smem_size) { + return false; + } + + return true; + } /// Executes one test bool run( @@ -145,9 +145,9 @@ struct TestbedComplex : public Testbed { return true; } - // - // Initialize workspace - // + // + // Initialize workspace + // this->initialize(problem_size); diff --git a/test/unit/gemm/device/testbed_grouped.h b/test/unit/gemm/device/testbed_grouped.h new file mode 100644 index 00000000..a5064c07 --- /dev/null +++ b/test/unit/gemm/device/testbed_grouped.h @@ -0,0 +1,515 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Tests for device-wide GEMM interface + +*/ + +#pragma once + +#include + +#include "../../common/cutlass_unit_test.h" +#include "cutlass/cutlass.h" + +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/kernel/gemm_grouped.h" +#include "cutlass/gemm/kernel/default_gemm_grouped.h" +#include "cutlass/gemm/device/gemm_grouped.h" + +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/reference/host/gemm_complex.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/tensor_norm.h" +#include "cutlass/util/tensor_view_io.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace test { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct TestbedGrouped { + + // + // Type definitions + // + + using ElementA = typename Gemm::ElementA; + using ElementB = typename Gemm::ElementB; + using ElementC = typename Gemm::ElementC; + using ElementAccumulator = typename Gemm::ElementAccumulator; + + using EpilogueOutputOp = typename Gemm::GemmKernel::Epilogue::OutputOp; + using ElementCompute = typename EpilogueOutputOp::ElementCompute; + + using LayoutA = typename Gemm::LayoutA; + using LayoutB = typename Gemm::LayoutB; + using LayoutC = typename Gemm::LayoutC; + + using MatrixCoord = typename LayoutC::TensorCoord; + + // + // Data members + // + + /// Initialization + cutlass::Distribution::Kind init_A; + cutlass::Distribution::Kind init_B; + cutlass::Distribution::Kind init_C; + uint32_t seed; + + int problem_count; + + std::vector problem_sizes_host; + cutlass::DeviceAllocation problem_sizes_device; + + std::vector offset_A; + std::vector offset_B; + std::vector offset_C; + std::vector offset_D; + + std::vector lda_host; + std::vector ldb_host; + std::vector ldc_host; + std::vector ldd_host; + + cutlass::DeviceAllocation lda; + cutlass::DeviceAllocation ldb; + cutlass::DeviceAllocation ldc; + cutlass::DeviceAllocation ldd; + + cutlass::DeviceAllocation block_A; + cutlass::DeviceAllocation block_B; + cutlass::DeviceAllocation block_C; + cutlass::DeviceAllocation block_D; + + cutlass::DeviceAllocation ptr_A; + cutlass::DeviceAllocation ptr_B; + cutlass::DeviceAllocation ptr_C; + cutlass::DeviceAllocation ptr_D; + + // + // Methods + // + + TestbedGrouped( + cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_C_ = cutlass::Distribution::Uniform, + uint32_t seed_ = 3080 + ): + init_A(init_A_), init_B(init_B_), init_C(init_C_), seed(seed_) { } + + /// Helper to initialize a tensor view + template + bool initialize_tensor( + cutlass::TensorView view, + cutlass::Distribution::Kind dist_kind, + uint32_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) { + if (cutlass::sizeof_bits::value <= 16) { + scope_max = 5; + scope_min = -5; + } + else { + scope_max = 8; + scope_min = -8; + } + } 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 { + // no fill - remain zero + } + + return true; + } + + /// Initializes data structures + void initialize() { + + // + // Choose random problem sizes + // + + // construct a few problems of random sizes + srand(seed); + + int64_t total_elements_A = 0; + int64_t total_elements_B = 0; + int64_t total_elements_C = 0; + int64_t total_elements_D = 0; + + + lda_host.resize(problem_count); + ldb_host.resize(problem_count); + ldc_host.resize(problem_count); + ldd_host.resize(problem_count); + + problem_sizes_host.clear(); + problem_sizes_host.resize(problem_count); + + for (int32_t i = 0; i < problem_count; ++i) { + + cutlass::gemm::GemmCoord problem( + 8 * (rand() % 64) + 24, + 8 * (rand() % 64) + 24, + 8 * (rand() % 64) + 24); + + if (!i) { + problem = cutlass::gemm::GemmCoord(48, 16, 8); + } + + problem_sizes_host.at(i) = problem; + + // std::cout << "Problem[" << i << "]: " << problem << std::endl; + + lda_host.at(i) = LayoutA::packed({problem.m(), problem.k()}).stride(0); + ldb_host.at(i) = LayoutB::packed({problem.k(), problem.n()}).stride(0); + ldc_host.at(i) = LayoutC::packed({problem.m(), problem.n()}).stride(0); + ldd_host.at(i) = LayoutC::packed({problem.m(), problem.n()}).stride(0); + + offset_A.push_back(total_elements_A); + offset_B.push_back(total_elements_B); + offset_C.push_back(total_elements_C); + offset_D.push_back(total_elements_D); + + int64_t elements_A = problem.m() * problem.k(); + int64_t elements_B = problem.k() * problem.n(); + int64_t elements_C = problem.m() * problem.n(); + int64_t elements_D = problem.m() * problem.n(); + + total_elements_A += elements_A; + total_elements_B += elements_B; + total_elements_C += elements_C; + total_elements_D += elements_D; + + // Random strides between problems? + } + + problem_sizes_device.reset(problem_count); + problem_sizes_device.copy_from_host(problem_sizes_host.data()); + + lda.reset(problem_count); + ldb.reset(problem_count); + ldc.reset(problem_count); + ldd.reset(problem_count); + + lda.copy_from_host(lda_host.data()); + ldb.copy_from_host(ldb_host.data()); + ldc.copy_from_host(ldc_host.data()); + ldd.copy_from_host(ldd_host.data()); + + // + // Assign pointers + // + + block_A.reset(total_elements_A); + block_B.reset(total_elements_B); + block_C.reset(total_elements_C); + block_D.reset(total_elements_D); + + std::vector ptr_A_host(problem_count); + std::vector ptr_B_host(problem_count); + std::vector ptr_C_host(problem_count); + std::vector ptr_D_host(problem_count); + + for (int32_t i = 0; i < problem_count; ++i) { + ptr_A_host.at(i) = block_A.get() + offset_A.at(i); + ptr_B_host.at(i) = block_B.get() + offset_B.at(i); + ptr_C_host.at(i) = block_C.get() + offset_C.at(i); + ptr_D_host.at(i) = block_D.get() + offset_D.at(i); + } + + ptr_A.reset(problem_count); + ptr_A.copy_from_host(ptr_A_host.data()); + + ptr_B.reset(problem_count); + ptr_B.copy_from_host(ptr_B_host.data()); + + ptr_C.reset(problem_count); + ptr_C.copy_from_host(ptr_C_host.data()); + + ptr_D.reset(problem_count); + ptr_D.copy_from_host(ptr_D_host.data()); + + // + // Initialize the problems of the workspace + // + + for (int32_t i = 0; i < problem_count; ++i) { + cutlass::gemm::GemmCoord problem = problem_sizes_host.at(i); + + LayoutA layout_A(lda_host.at(i)); + LayoutB layout_B(ldb_host.at(i)); + LayoutC layout_C(ldc_host.at(i)); + LayoutC layout_D(ldd_host.at(i)); + + MatrixCoord extent_A{problem.m(), problem.k()}; + MatrixCoord extent_B{problem.k(), problem.n()}; + MatrixCoord extent_C{problem.m(), problem.n()}; + + std::vector matrix_A(layout_A.capacity(extent_A)); + std::vector matrix_B(layout_B.capacity(extent_B)); + std::vector matrix_C(layout_C.capacity(extent_C)); + std::vector matrix_D(layout_D.capacity(extent_C)); + + initialize_tensor(cutlass::TensorView(matrix_A.data(), layout_A, extent_A), init_A, seed * 2021); + initialize_tensor(cutlass::TensorView(matrix_B.data(), layout_B, extent_B), init_B, seed * 2022); + initialize_tensor(cutlass::TensorView(matrix_C.data(), layout_C, extent_C), init_C, seed * 2023); + + cutlass::device_memory::copy_to_device(ptr_A_host.at(i), matrix_A.data(), matrix_A.size()); + cutlass::device_memory::copy_to_device(ptr_B_host.at(i), matrix_B.data(), matrix_B.size()); + cutlass::device_memory::copy_to_device(ptr_C_host.at(i), matrix_C.data(), matrix_C.size()); + cutlass::device_memory::copy_to_device(ptr_D_host.at(i), matrix_D.data(), matrix_D.size()); + } + } + + /// Verifies the result is a GEMM + bool verify( + ElementCompute alpha, + ElementCompute beta) { + + bool passed = true; + + for (int32_t i = 0; i < problem_count; ++i) { + cutlass::gemm::GemmCoord problem = problem_sizes_host.at(i); + + LayoutA layout_A(lda_host.at(i)); + LayoutB layout_B(ldb_host.at(i)); + LayoutC layout_C(ldc_host.at(i)); + LayoutC layout_D(ldd_host.at(i)); + + MatrixCoord extent_A{problem.m(), problem.k()}; + MatrixCoord extent_B{problem.k(), problem.n()}; + MatrixCoord extent_C{problem.m(), problem.n()}; + + std::vector matrix_A(layout_A.capacity(extent_A)); + std::vector matrix_B(layout_B.capacity(extent_B)); + std::vector matrix_C(layout_C.capacity(extent_C)); + std::vector matrix_D(layout_D.capacity(extent_C)); + std::vector matrix_Ref(layout_D.capacity(extent_C)); + + cutlass::device_memory::copy_to_host(matrix_A.data(), block_A.get() + offset_A.at(i), matrix_A.size()); + cutlass::device_memory::copy_to_host(matrix_B.data(), block_B.get() + offset_B.at(i), matrix_B.size()); + cutlass::device_memory::copy_to_host(matrix_C.data(), block_C.get() + offset_C.at(i), matrix_C.size()); + cutlass::device_memory::copy_to_host(matrix_D.data(), block_D.get() + offset_D.at(i), matrix_D.size()); + + cutlass::TensorView view_A(matrix_A.data(), layout_A, extent_A); + cutlass::TensorView view_B(matrix_B.data(), layout_B, extent_B); + cutlass::TensorView view_C(matrix_C.data(), layout_C, extent_C); + cutlass::TensorView view_D(matrix_D.data(), layout_D, extent_C); + cutlass::TensorView view_Ref(matrix_Ref.data(), layout_D, extent_C); + + // Reference GEMM + cutlass::reference::host::GemmComplex< + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + ElementCompute, ElementAccumulator + >( + problem, + alpha, + view_A, + Gemm::kTransformA, + view_B, + Gemm::kTransformB, + beta, + view_C, + view_Ref, + ElementAccumulator(0) + ); + + // Ensure that no input or output is entirely zero + EXPECT_GT(cutlass::reference::host::TensorNorm(view_A), 0); + EXPECT_GT(cutlass::reference::host::TensorNorm(view_B), 0); + EXPECT_GT(cutlass::reference::host::TensorNorm(view_C), 0); + EXPECT_GT(cutlass::reference::host::TensorNorm(view_D), 0); + EXPECT_GT(cutlass::reference::host::TensorNorm(view_Ref), 0); + + // Compare against reference + passed = cutlass::reference::host::TensorEquals(view_D, view_Ref); + + if (!passed) { + std::ofstream file("testbed_grouped_errors.txt"); + + file + << "problem: " << problem << " [group: " << i << "]\n" + << ", alpha: " << alpha << ", beta: " << beta << "\n\n"; + + file + << "A =\n" << view_A + << "\nB =\n" << view_B + << "\nC =\n" << view_C + << "\n\nReference =\n" << view_Ref + << "\nComputed =\n" << view_D; + + return passed; + } + } + + return passed; + } + + /// Returns the number of threadblocks to launch if the kernel can run on the target + /// device. Otherwise, returns zero. + int sufficient() const { + // + // Determine SMEM requirements and waive if not satisfied + // + + int smem_size = int(sizeof(typename Gemm::GemmKernel::SharedStorage)); + + cudaDeviceProp properties; + int device_idx; + cudaError_t result = cudaGetDevice(&device_idx); + + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDevice() API call failed."); + } + + result = cudaGetDeviceProperties(&properties, device_idx); + + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDeviceProperties() failed"); + } + + int occupancy = std::min(2, int(properties.sharedMemPerMultiprocessor / smem_size)); + + return properties.multiProcessorCount * occupancy; + } + + /// Executes one test + bool run( + int problem_count, + ElementCompute alpha = ElementCompute(1), + ElementCompute beta = ElementCompute(0)) { + + int threadblock_count = sufficient(); + + // Early exit + if (!threadblock_count) { + return false; + } + + this->problem_count = problem_count; + + // Initialize the problem + initialize(); + + // Configure the GEMM arguments + typename EpilogueOutputOp::Params epilogue_op(alpha, beta); + + // Configure GEMM arguments + typename Gemm::Arguments args( + problem_sizes_device.get(), + problem_count, + threadblock_count, + epilogue_op, + ptr_A.get(), + ptr_B.get(), + ptr_C.get(), + ptr_D.get(), + lda.get(), + ldb.get(), + ldc.get(), + ldd.get() + ); + + // Initialize the GEMM object + Gemm gemm; + + cutlass::Status status = gemm.initialize(args); + + if (status != cutlass::Status::kSuccess) { + return false; + } + + // Run the GEMM object + status = gemm.run(); + + if (status != cutlass::Status::kSuccess) { + return false; + } + + // Wait for completion + cudaError_t result = cudaDeviceSynchronize(); + + EXPECT_EQ(result, cudaSuccess) + << "Kernel execution error: " << cudaGetErrorString(result); + + if (result != cudaSuccess) { + return false; + } + + // Verify correctness + return verify(alpha, beta); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // device +} // gemm +} // test + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/test/unit/gemm/warp/gemm_complex_sm80.cu b/test/unit/gemm/warp/gemm_complex_sm80.cu index c6fd7440..06f1c019 100644 --- a/test/unit/gemm/warp/gemm_complex_sm80.cu +++ b/test/unit/gemm/warp/gemm_complex_sm80.cu @@ -689,6 +689,4 @@ TEST(SM80_warp_gemm_complex_tensor_op_f64, 32x32x8_8x8x4_nt) { //////////////////////////////////////////////////////////////////////////////////////////////// - #endif // #if defined(CUTLASS_ARCH_MMA_SM80_SUPPORTED) - diff --git a/test/unit/gemm/warp/gemm_sm80.cu b/test/unit/gemm/warp/gemm_sm80.cu index bb8a098b..956a1d70 100644 --- a/test/unit/gemm/warp/gemm_sm80.cu +++ b/test/unit/gemm/warp/gemm_sm80.cu @@ -635,7 +635,7 @@ TEST(SM80_warp_gemm_tensor_op_congruous_tf32, 128x128x32_32x32x32_16x8x8) { } //////////////////////////////////////////////////////////////////////////////// -TEST(SM80_warp_gemm_tensor_op_tn, tf32_round_128x128x32_16x16x32_16x8x8) { +TEST(SM80_warp_gemm_tensor_op_tn, tf32_round_128x128x32_64x64x32_16x8x8) { using Shape = cutlass::gemm::GemmShape<64, 64, 32>; using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; @@ -657,7 +657,7 @@ TEST(SM80_warp_gemm_tensor_op_tn, tf32_round_128x128x32_16x16x32_16x8x8) { .run(); } -TEST(SM80_warp_gemm_tensor_op_nt, tf32_round_128x128x32_16x16x32_16x8x8) { +TEST(SM80_warp_gemm_tensor_op_nt, tf32_round_128x128x32_64x64x32_16x8x8) { using Shape = cutlass::gemm::GemmShape<64, 64, 32>; using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>; diff --git a/tools/library/include/cutlass/library/library.h b/tools/library/include/cutlass/library/library.h index 3c7da39f..24f91256 100644 --- a/tools/library/include/cutlass/library/library.h +++ b/tools/library/include/cutlass/library/library.h @@ -193,7 +193,9 @@ enum class MathOperationID { kMultiplyAddSaturate, kMultiplyAddFastBF16, kMultiplyAddFastF16, + kMultiplyAddFastF32, kMultiplyAddComplex, + kMultiplyAddComplexFastF32, kMultiplyAddGaussianComplex, kXorPopc, kInvalid diff --git a/tools/library/scripts/conv2d_operation.py b/tools/library/scripts/conv2d_operation.py index 11d7f384..0ba4307f 100644 --- a/tools/library/scripts/conv2d_operation.py +++ b/tools/library/scripts/conv2d_operation.py @@ -103,7 +103,7 @@ class Conv2dOperation: ) if self.stride_support == StrideSupport.Unity: - configuration_name = "cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}_align${alignment}_unity_stride" + configuration_name = "cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}_unity_stride_align${alignment}" else: configuration_name = "cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}_align${alignment}" diff --git a/tools/library/scripts/gemm_operation.py b/tools/library/scripts/gemm_operation.py index 0a76f36b..fe7462a3 100644 --- a/tools/library/scripts/gemm_operation.py +++ b/tools/library/scripts/gemm_operation.py @@ -40,7 +40,7 @@ class GemmOperation: def is_complex(self): complex_operators = [ MathOperation.multiply_add_complex, - MathOperation.multiply_add_complex_gaussian + MathOperation.multiply_add_complex_gaussian, ] return self.tile_description.math_instruction.math_operation in complex_operators diff --git a/tools/library/scripts/generator.py b/tools/library/scripts/generator.py index daf793fe..333dae7b 100644 --- a/tools/library/scripts/generator.py +++ b/tools/library/scripts/generator.py @@ -117,9 +117,10 @@ def CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, data_t gemm_kinds = [GemmKind.PlanarComplex, GemmKind.PlanarComplexArray] - # by default, planar complex gemm kernels are not generated + # by default, only generate the largest tile and largest alignment if manifest.args.kernels == '': - return + tile_descriptions = [tile_descriptions[0],] + alignment_constraints = [alignment_constraints[0],] for gemm_kind in gemm_kinds: for layout in layouts: @@ -924,6 +925,8 @@ def GenerateSM75_TensorOp_8816_TN(manifest, args): TileDescription([256, 128, 64], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([ 64, 256, 64], 2, [1, 4, 1], math_inst, min_cc, max_cc), + TileDescription([256, 64, 64], 2, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), @@ -1188,7 +1191,6 @@ def GenerateSM75_TensorOp_88128(manifest, args): min_cc = 75 max_cc = 1024 - alignment_constraints = [128,] for math_inst in math_instructions: @@ -1953,28 +1955,30 @@ def GenerateSM80_TensorOp_168256(manifest, args): ] min_cc = 80 - max_cc = 1024 + max_cc = { + MathOperation.xor_popc: 1024 + } alignment_constraints = [128,] for math_inst in math_instructions: tile_descriptions = [ - TileDescription([256, 128, 512], 3, [4, 2, 1], math_inst, min_cc, max_cc), - TileDescription([128, 256, 512], 3, [2, 4, 1], math_inst, min_cc, max_cc), - TileDescription([256, 64, 512], 4, [4, 1, 1], math_inst, min_cc, max_cc), - TileDescription([ 64, 256, 512], 4, [1, 4, 1], math_inst, min_cc, max_cc), - TileDescription([128, 128, 512], 5, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([128, 64, 512], 6, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([ 64, 128, 512], 6, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([ 64, 64, 512], 10, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([256, 128, 1024], 3, [4, 2, 1], math_inst, min_cc, max_cc), - TileDescription([128, 256, 1024], 3, [2, 4, 1], math_inst, min_cc, max_cc), - TileDescription([256, 64, 1024], 4, [4, 1, 1], math_inst, min_cc, max_cc), - TileDescription([ 64, 256, 1024], 4, [1, 4, 1], math_inst, min_cc, max_cc), - TileDescription([128, 128, 1024], 4, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([128, 64, 1024], 3, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([ 64, 128, 1024], 3, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([ 64, 64, 1024], 5, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([256, 128, 512], 3, [4, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([128, 256, 512], 3, [2, 4, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([256, 64, 512], 4, [4, 1, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([ 64, 256, 512], 4, [1, 4, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([128, 128, 512], 5, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([128, 64, 512], 6, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([ 64, 128, 512], 6, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([ 64, 64, 512], 10, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([256, 128, 1024], 3, [4, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([128, 256, 1024], 3, [2, 4, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([256, 64, 1024], 4, [4, 1, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([ 64, 256, 1024], 4, [1, 4, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([128, 128, 1024], 4, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([128, 64, 1024], 3, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([ 64, 128, 1024], 3, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), + TileDescription([ 64, 64, 1024], 5, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), ] data_type = [DataType.b1, DataType.b1, DataType.s32, DataType.s32] @@ -2130,6 +2134,105 @@ def GenerateSM80_TensorOp_1688_fast_math(manifest, args): # # +def GenerateSM80_TensorOp_1688_fast_fp32_math(manifest, args): + + if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0): + return + + layouts = [ + (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), + (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), + (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), + (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), + ] + + math_instructions = [ + MathInstruction( \ + [16, 8, 8], \ + DataType.f32, DataType.f32, DataType.f32, \ + OpcodeClass.TensorOp, \ + MathOperation.multiply_add_fast_f32), + ] + + min_cc = 80 + max_cc = 1024 + max_cc_smem_limited = 80 + + alignment_constraints = [4, 2, 1] + + for math_inst in math_instructions: + tile_descriptions = [ + TileDescription([128, 128, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc), + TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), + TileDescription([256, 64, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), + TileDescription([ 64, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), + TileDescription([128, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([ 64, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([ 64, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([128, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), + TileDescription([256, 64, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc_smem_limited), + TileDescription([ 64, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc_smem_limited), + TileDescription([128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([ 64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([ 64, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), + ] + + data_type = [DataType.f32, DataType.f32, DataType.f32, DataType.f32] + + CreateGemmOperator(manifest, layouts, tile_descriptions, \ + data_type, alignment_constraints) + + conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) + CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) +# + +def GenerateSM80_TensorOp_1688_fast_fp32_math_complex(manifest, args): + + if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0): + return + + layouts = [ + (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), + (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), + (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), + (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), + ] + + math_inst = MathInstruction( \ + [16, 8, 8], \ + DataType.f32, DataType.f32, DataType.f32, \ + OpcodeClass.TensorOp, \ + MathOperation.multiply_add_complex_fast_f32) + + min_cc = 80 + max_cc = 1024 + + tile_descriptions = [ + TileDescription([128, 64, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), + TileDescription([64, 128, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), + TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([32, 32, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), + ] + + data_type = [ + DataType.cf32, DataType.cf32, DataType.cf32, DataType.cf32 + ] + + alignment_constraints = [1,] + + complex_transforms = [ + (ComplexTransform.none, ComplexTransform.none), + (ComplexTransform.conj, ComplexTransform.none), + (ComplexTransform.none, ComplexTransform.conj), + (ComplexTransform.conj, ComplexTransform.conj) + ] + + CreateGemmOperator(manifest, layouts, tile_descriptions, \ + data_type, alignment_constraints, complex_transforms) + + # def GenerateSM80_SparseTensorOp_16816_fast_math(manifest, args): @@ -2195,7 +2298,7 @@ def GenerateSM80_TensorOp_1688_complex(manifest, args): math_inst = MathInstruction( \ [16, 8, 8], \ - DataType.f32, DataType.f32, DataType.f32, \ + DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex) @@ -2203,10 +2306,12 @@ def GenerateSM80_TensorOp_1688_complex(manifest, args): max_cc = 1024 tile_descriptions = [ + TileDescription([128, 128, 16], 4, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 4, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 1, 1], math_inst, min_cc, max_cc), + TileDescription([32, 64, 16], 4, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] @@ -2297,14 +2402,22 @@ def GenerateSM80_TensorOp_884_complex(manifest, args): alignment_constraints = [1,] tile_descriptions = [ - TileDescription([128, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), - TileDescription([64, 128, 8], 3, [2, 4, 1], math_inst, min_cc, max_cc), - TileDescription([64, 64, 8], 3, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), - TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), - TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), + TileDescription([128, 64, 8 ], 3, [4, 2, 1], math_inst, min_cc, max_cc), + TileDescription([64, 128, 8 ], 3, [2, 4, 1], math_inst, min_cc, max_cc), + TileDescription([64, 64, 8 ], 3, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([64, 32, 8 ], 4, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([32, 64, 8 ], 4, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([32, 32, 8 ], 4, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([16, 32, 8 ], 4, [1, 2, 1], math_inst, min_cc, max_cc), + TileDescription([32, 16, 8 ], 4, [2, 1, 1], math_inst, min_cc, max_cc), + TileDescription([128, 64, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), + TileDescription([64, 128, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), + TileDescription([64, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([64, 32, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([32, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([32, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), + TileDescription([16, 32, 16], 4, [1, 2, 1], math_inst, min_cc, max_cc), + TileDescription([32, 16, 16], 3, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] diff --git a/tools/library/src/library_internal.h b/tools/library/src/library_internal.h index 11617cb2..8050e9ac 100644 --- a/tools/library/src/library_internal.h +++ b/tools/library/src/library_internal.h @@ -168,6 +168,15 @@ template <> struct MathOperationMap { static MathOperationID const kId = MathOperationID::kXorPopc; }; + +template <> struct MathOperationMap { + static MathOperationID const kId = MathOperationID::kMultiplyAddFastF32; +}; + +template <> struct MathOperationMap { + static MathOperationID const kId = MathOperationID::kMultiplyAddComplexFastF32; +}; + ///////////////////////////////////////////////////////////////////////////////////////////////// template struct LayoutMap; diff --git a/tools/profiler/CMakeLists.txt b/tools/profiler/CMakeLists.txt index de5c9924..ae922168 100644 --- a/tools/profiler/CMakeLists.txt +++ b/tools/profiler/CMakeLists.txt @@ -33,7 +33,7 @@ set(CUTLASS_TOOLS_PROFILER_SOURCES src/gpu_timer.cpp src/device_allocation.cu src/device_context.cu - src/cublas_helpers.cpp + src/cublas_helpers.cu src/cudnn_helpers.cpp src/problem_space.cpp src/operation_profiler.cu diff --git a/tools/profiler/src/cublas_helpers.cpp b/tools/profiler/src/cublas_helpers.cu similarity index 98% rename from tools/profiler/src/cublas_helpers.cpp rename to tools/profiler/src/cublas_helpers.cu index afab873b..94261e18 100644 --- a/tools/profiler/src/cublas_helpers.cpp +++ b/tools/profiler/src/cublas_helpers.cu @@ -220,7 +220,7 @@ cublasGemmExDispatcher::cublasGemmExDispatcher( // cuBLAS introduces a separate cublasComputeType enumerant to more precisely describe // internal numerical data types used in the computation. -#if (__CUDA_VER_MAJOR__ >= 11) +#if (__CUDACC_VER_MAJOR__ >= 11) library::OpcodeClassID const & opcode_class = op_desc.tile_description.math_instruction.opcode_class; @@ -253,7 +253,7 @@ cublasGemmExDispatcher::cublasGemmExDispatcher( break; } } -#endif // __CUDA_VER_MAJOR__ >= 11 +#endif // __CUDACC_VER_MAJOR__ >= 11 if (!good) { status = Status::kErrorNotSupported; @@ -286,7 +286,7 @@ cublasStatus_t cublasGemmExDispatcher::operator()(cublasHandle_t handle) { int(configuration.ldc), arguments.batch_stride_C, configuration.batch_count, - #if (__CUDA_VER_MAJOR__ >= 11) + #if (__CUDACC_VER_MAJOR__ >= 11) compute_type, #else compute_data_type, @@ -313,7 +313,7 @@ cublasStatus_t cublasGemmExDispatcher::operator()(cublasHandle_t handle) { arguments.D, data_type_C, int(configuration.ldc), - #if (__CUDA_VER_MAJOR__ >= 11) + #if (__CUDACC_VER_MAJOR__ >= 11) compute_type, #else compute_data_type, diff --git a/tools/profiler/src/cublas_helpers.h b/tools/profiler/src/cublas_helpers.h index 200d9130..a899f9ed 100644 --- a/tools/profiler/src/cublas_helpers.h +++ b/tools/profiler/src/cublas_helpers.h @@ -187,7 +187,7 @@ struct cublasGemmExDispatcher { cudaDataType_t data_type_C; cudaDataType_t compute_data_type; -#if (__CUDA_VER_MAJOR__ >= 11) +#if (__CUDACC_VER_MAJOR__ >= 11) cublasComputeType_t compute_type; #endif diff --git a/tools/profiler/src/gemm_operation_profiler.cu b/tools/profiler/src/gemm_operation_profiler.cu index fab38346..7afea996 100644 --- a/tools/profiler/src/gemm_operation_profiler.cu +++ b/tools/profiler/src/gemm_operation_profiler.cu @@ -254,6 +254,10 @@ int64_t GemmOperationProfiler::GemmProblem::flops(library::GemmDescription const flops_ *= 4; break; + case library::MathOperationID::kMultiplyAddComplexFastF32: + flops_ *= 4; + break; + case library::MathOperationID::kMultiplyAddGaussianComplex: flops_ *= 3; break; diff --git a/tools/profiler/src/operation_profiler.cu b/tools/profiler/src/operation_profiler.cu index 4bbb2774..84c3aeb1 100644 --- a/tools/profiler/src/operation_profiler.cu +++ b/tools/profiler/src/operation_profiler.cu @@ -260,6 +260,9 @@ int OperationProfiler::profile_all( auto min_cc = operation->description().tile_description.minimum_compute_capability; auto max_cc = operation->description().tile_description.maximum_compute_capability; + // Clear named allocations + device_context.free(); + // Execute compatible cutlass operations if they satisfy the current device's compute capability if (operation->description().kind == kind_ && operation->description().provider == library::Provider::kCUTLASS && @@ -301,9 +304,12 @@ int OperationProfiler::profile_all( problem); if (status == Status::kErrorInternal) { - // Stop profiling if there was an internal error - internal_error = true; - break; + + // If there was an internal error, consume the CUDA error and move to the next operation. + (void)cudaGetLastError(); + + report.append_results(results_); + continue; } else if (status != Status::kSuccess) { // If the workspace could not be initialized for any other reason, continue to @@ -322,9 +328,12 @@ int OperationProfiler::profile_all( problem); if (status == Status::kErrorInternal) { - // Stop profiling if there was an internal error - internal_error = true; - break; + + // If there was an internal error, consume the CUDA error and move to the next operation. + (void)cudaGetLastError(); + + report.append_results(results_); + continue; } else if (status != Status::kSuccess) { // If the workspace could not be initialized for any other reason, continue to @@ -383,9 +392,6 @@ int OperationProfiler::profile_all( problem); } - // Clear named allocations - device_context.free(); - report.append_results(results_); results_.clear(); } diff --git a/tools/profiler/src/options.cu b/tools/profiler/src/options.cu index a21ce50e..404c165f 100644 --- a/tools/profiler/src/options.cu +++ b/tools/profiler/src/options.cu @@ -577,6 +577,8 @@ Options::Report::Report(cutlass::CommandLine const &cmdline) { cmdline.get_cmd_line_argument("report-not-run", report_not_run, false); cmdline.get_cmd_line_argument("verbose", verbose, true); + + cmdline.get_cmd_line_argument("sort-results", sort_results, false); } void Options::Report::print_usage(std::ostream &out) const { @@ -602,7 +604,10 @@ void Options::Report::print_usage(std::ostream &out) const { << " column. Useful for generating pivot tables.\n\n" << " --verbose= " - << " Prints human-readable text to stdout. If false, nothing is written to stdout.\n\n"; + << " Prints human-readable text to stdout. If false, nothing is written to stdout.\n\n" + + << " --sort-results= " + << " Sorts results (by flops-per-byte).\n\n"; } void Options::Report::print_options(std::ostream &out, int indent) const { diff --git a/tools/profiler/src/options.h b/tools/profiler/src/options.h index f0f6b746..d8b69fc6 100644 --- a/tools/profiler/src/options.h +++ b/tools/profiler/src/options.h @@ -231,6 +231,9 @@ public: /// Prints human-readable text to stdout. If false, nothing is written to stdout bool verbose; + /// Sort results by (currently by flops-per-byte) + bool sort_results; + // // Methods // diff --git a/tools/profiler/src/performance_report.cpp b/tools/profiler/src/performance_report.cpp index 06bdecb1..0724db50 100644 --- a/tools/profiler/src/performance_report.cpp +++ b/tools/profiler/src/performance_report.cpp @@ -154,6 +154,22 @@ void PerformanceReport::append_result(PerformanceResult result) { } } +void PerformanceReport::sort_results(PerformanceResultVector &results) { + + struct FlopsPerByteCompare + { + bool operator()(const PerformanceResult &a, const PerformanceResult &b) + { + double a_flops_per_byte = double(a.flops) / double(a.bytes); + double b_flops_per_byte = double(b.flops) / double(b.bytes); + + return (a_flops_per_byte < b_flops_per_byte); + } + }; + + std::stable_sort(results.begin(), results.end(), FlopsPerByteCompare()); +} + void PerformanceReport::append_results(PerformanceResultVector const &results) { if (options_.report.verbose) { @@ -173,12 +189,16 @@ PerformanceReport::~PerformanceReport() { // if (options_.report.verbose && !concatenated_results_.empty()) { + if (options_.report.sort_results) { + sort_results(concatenated_results_); + } + std::cout << "\n\n"; std::cout << "=============================\n\n"; std::cout << "CSV Results:\n\n"; print_csv_header_(std::cout) << std::endl; - + for (auto const &result : concatenated_results_) { print_result_csv_(std::cout, result) << "\n"; } @@ -273,7 +293,8 @@ std::ostream & PerformanceReport::print_result_pretty_( out << " Bytes: " << result.bytes << " bytes\n" - << " FLOPs: " << result.flops << " flops\n\n"; + << " FLOPs: " << result.flops << " flops\n" + << " FLOPs/Byte: " << (result.flops / result.bytes) << "\n\n"; if (result.good()) { @@ -309,6 +330,7 @@ std::ostream & PerformanceReport::print_csv_header_( out << ",Bytes" << ",Flops" + << ",Flops/Byte" << ",Runtime" << ",GB/s" << ",GFLOPs" @@ -345,6 +367,7 @@ std::ostream & PerformanceReport::print_result_csv_( out << "," << result.bytes << "," << result.flops + << "," << result.flops / result.bytes << "," << result.runtime; if (result.good()) { diff --git a/tools/profiler/src/performance_report.h b/tools/profiler/src/performance_report.h index fad89e14..c2c7ec3b 100644 --- a/tools/profiler/src/performance_report.h +++ b/tools/profiler/src/performance_report.h @@ -86,6 +86,7 @@ public: void next_problem(); void append_result(PerformanceResult result); + void sort_results(PerformanceResultVector &results); void append_results(PerformanceResultVector const &results); public: diff --git a/tools/util/include/cutlass/util/command_line.h b/tools/util/include/cutlass/util/command_line.h index cf7b66f9..83270f8a 100644 --- a/tools/util/include/cutlass/util/command_line.h +++ b/tools/util/include/cutlass/util/command_line.h @@ -120,7 +120,7 @@ struct CommandLine { /** * Obtains the boolean value specified for a given commandline parameter --= */ - void get_cmd_line_argument(const char* arg_name, bool& val, bool _default = true) const { + void get_cmd_line_argument(const char* arg_name, bool& val, bool _default) const { val = _default; if (check_cmd_line_flag(arg_name)) { std::string value; diff --git a/tools/util/include/cutlass/util/distribution.h b/tools/util/include/cutlass/util/distribution.h index 3f676a01..b135ee05 100644 --- a/tools/util/include/cutlass/util/distribution.h +++ b/tools/util/include/cutlass/util/distribution.h @@ -37,7 +37,7 @@ namespace cutlass { /// Distribution type struct Distribution { /// Variant types - enum Kind { Invalid, Uniform, Gaussian, Identity, Sequential }; + enum Kind { Invalid, Uniform, Gaussian, Identity, Sequential, AllZeros, AllOnes }; /// Distribution state union { diff --git a/tools/util/include/cutlass/util/reference/device/convolution.h b/tools/util/include/cutlass/util/reference/device/convolution.h index 212fa08d..f6e04c1c 100644 --- a/tools/util/include/cutlass/util/reference/device/convolution.h +++ b/tools/util/include/cutlass/util/reference/device/convolution.h @@ -76,7 +76,9 @@ __global__ void Conv2dFprop( TensorRef tensor_y_in, TensorRef tensor_y_out, ElementCompute alpha, - ElementCompute beta + ElementCompute beta, + TensorRef tensor_scale, + TensorRef tensor_bias ) { ConvertOp convert_op; @@ -178,13 +180,26 @@ __global__ void Conv2dFprop( int thread_k = k_start + n; if (thread_k < problem_size.K) { - ElementCompute c_ref = ElementCompute(); - if (beta != ElementCompute()) { - c_ref = ElementCompute(tensor_y_in.at({thread_n[m], thread_p[m], thread_q[m], thread_k})); + if(alpha == ElementCompute()) { // use per-channel scale and bias + ElementCompute scale = tensor_scale.at({0, thread_k}); + ElementCompute bias = tensor_bias.at({0, thread_k}); + tensor_y_out.at({thread_n[m], thread_p[m], thread_q[m], thread_k}) = convert_op( + scale * ElementCompute(accum[m][n]) + bias); } + else if(tensor_bias.good()) { // use per-channel bias + ElementCompute bias = tensor_bias.at({0, thread_k}); + tensor_y_out.at({thread_n[m], thread_p[m], thread_q[m], thread_k}) = convert_op( + alpha * ElementCompute(accum[m][n]) + bias); + } + else { + ElementCompute c_ref = ElementCompute(); + if (beta != ElementCompute()) { + c_ref = ElementCompute(tensor_y_in.at({thread_n[m], thread_p[m], thread_q[m], thread_k})); + } - tensor_y_out.at({thread_n[m], thread_p[m], thread_q[m], thread_k}) = convert_op( - alpha * ElementCompute(accum[m][n]) + beta * c_ref); + tensor_y_out.at({thread_n[m], thread_p[m], thread_q[m], thread_k}) = convert_op( + alpha * ElementCompute(accum[m][n]) + beta * c_ref); + } } } } @@ -994,7 +1009,9 @@ Status Conv2dFprop( TensorRef tensor_y_out, ElementCompute alpha, ElementCompute beta, - cudaStream_t stream = nullptr) { + cudaStream_t stream = nullptr, + TensorRef tensor_scale = TensorRef(), + TensorRef tensor_bias = TensorRef() ) { // // Blocking factors improve performance of reference implementation @@ -1033,7 +1050,9 @@ Status Conv2dFprop( tensor_y_in, tensor_y_out, alpha, - beta + beta, + tensor_scale, + tensor_bias ); cudaError_t result = cudaPeekAtLastError(); @@ -1423,7 +1442,9 @@ Status Conv2d( TensorRef tensor_D, ElementCompute alpha, ElementCompute beta, - cudaStream_t stream = nullptr) { + cudaStream_t stream = nullptr, + TensorRef tensor_scale = TensorRef(), + TensorRef tensor_bias = TensorRef() ) { switch (convolutional_operator) { case conv::Operator::kFprop: @@ -1434,7 +1455,7 @@ Status Conv2d( ElementCompute, ElementAccumulator, ConvertOp, InnerProductOp - >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, stream); + >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, stream, tensor_scale, tensor_bias); break; case conv::Operator::kDgrad: diff --git a/tools/util/include/cutlass/util/reference/device/tensor_reduce.h b/tools/util/include/cutlass/util/reference/device/tensor_reduce.h index 80aab7c6..79422e34 100644 --- a/tools/util/include/cutlass/util/reference/device/tensor_reduce.h +++ b/tools/util/include/cutlass/util/reference/device/tensor_reduce.h @@ -502,4 +502,3 @@ ComputeType TensorNormDiff( } // namespace cutlass ///////////////////////////////////////////////////////////////////////////////////////////////// -