leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 25 Packages Projects Releases Wiki Activity

CUTLASS 3.5.1 (#1623)

Browse Source 
* CUTLASS 3.5.1

* updates, optimizations, fixes
This commit is contained in:
Vijay Thakkar
2024-07-29 08:46:24 -04:00
committed by GitHub
parent 56b46e2d13
commit be60a0b272
312 changed files with 19793 additions and 6775 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+17
View File
@@ -1,5 +1,22 @@
# NVIDIA CUTLASS Changelog
## [3.5.1](https://github.com/NVIDIA/cutlass/releases/tag/v3.5.1) (2024-07-25)
- [Minimal SM90 WGMMA + TMA GEMM example in 100 lines of code](./examples/cute/tutorial/wgmma_sm90.cu)
- [Exposure of L2 `cache_hint`s in TMA copy atoms](./include/cute/arch/copy_sm90_tma.hpp#L48)
- Exposure of raster order and tile swizzle extent in [CUTLASS library profiler](./media/docs/profiler.md#GEMM), and
[example 48](./examples/48_hopper_warp_specialized_gemm/48_hopper_warp_specialized_gemm.cu).
- [TMA store based and EVT supported epilogues](./include/cutlass/epilogue/collective/sm90_epilogue_array_tma_warpspecialized.hpp) for [Hopper pointer array batched kernels](./test/unit/gemm/device/sm90_gemm_f16_f16_f16_tensor_op_f32_ptr_array.cu).
- A new [`GemmSparseUniversal` API for CUTLASS 2.x Ampere kernels](./include/cutlass/gemm/device/gemm_sparse_universal.h) leveraging 2:4 structured sparsity and [support for LLM friendly tile sizes](./test/unit/gemm/device/gemm_f16n_f16t_f32t_tensor_op_f32_sparse_sm80.cu).
- [CUDA host adapter](./include/cutlass/cuda_host_adapter.hpp) extensions to support TMA descriptor construction driver APIs.
- Inclusion of more [Hopper fprop, dgrad, and wgrad convolution kernels in CUTLASS library and profiler](./python/cutlass_library/generator.py).
- Support for residual add (beta != 0) in convolution kernels.
- A refactor of [include files throughout CUTLASS core directories](./include/cutlass/gemm/collective/collective_mma_decl.hpp) to reduce circular dependencies and [tests to guard against them](./test/self_contained_includes/CMakeLists.txt).
- [A guide for setting up VSCode to work well with CUTLASS](./media/docs/ide_setup.md) and [expanded code style guide](./media/docs/programming_guidelines.md).
- Better support for MSVC as a host compiler.
- Many performance optimizations, improvements, and bug fixes including fixes for FlashAttention-2.
- Optimal code generation with CUDA toolkit versions 12.4 and 12.5u1.
## [3.5.0](https://github.com/NVIDIA/cutlass/releases/tag/v3.5.0) (2024-04-09)
- Implicit GEMM Convolutions targeting Hopper SM90A via WGMMA + [TMA im2col](./include/cute/atom/copy_traits_sm90_im2col.hpp)
CMakeLists.txt
+65 -41
View File
@@ -92,7 +92,7 @@ if(CUTLASS_NATIVE_CUDA)
else()
list(APPEND CUTLASS_CUDA_NVCC_FLAGS --std=c++17)
endif()
if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
set(CMAKE_INSTALL_PREFIX install CACHE PATH "Default installation location." FORCE)
endif()
@@ -134,6 +134,16 @@ set(CUTLASS_ENABLE_PERFORMANCE ${CUTLASS_ENABLE_PROFILER} CACHE BOOL "Enable CUT
set(CUTLASS_ENABLE_TESTS ${CUTLASS_ENABLE_TESTS_INIT} CACHE BOOL "Enable CUTLASS Tests")
set(CUTLASS_ENABLE_GTEST_UNIT_TESTS ${CUTLASS_ENABLE_TESTS} CACHE BOOL "Enable CUTLASS GTest-based Unit Tests")
set(CUTLASS_USE_SYSTEM_GOOGLETEST OFF CACHE BOOL "Use system/external installation of GTest")
set(CUTLASS_USE_PACKED_TUPLE ON CACHE BOOL "If ON, make cute::tuple be new standard-layout tuple type; if OFF, use the original cute::tuple implementation that is _not_ standard-layout.")
if (CUTLASS_USE_PACKED_TUPLE)
list(APPEND CUTLASS_CUDA_NVCC_FLAGS -DCUTE_USE_PACKED_TUPLE=1)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DCUTLASS_USE_PACKED_TUPLE=1")
message(STATUS "Make cute::tuple be the new standard-layout tuple type")
elseif()
message(STATUS "Use the original cute::tuple implementation that is _not_ standard-layout")
endif()
################################################################################
set(CUTLASS_NVCC_ARCHS_SUPPORTED "")
@@ -216,7 +226,7 @@ if (${CUTLASS_NVCC_VERBOSE})
endif()
#
# CUTLASS NAMESPACE
# CUTLASS NAMESPACE
#
set(CUTLASS_NAMESPACE "cutlass" CACHE STRING "Top level namespace of CUTLASS")
@@ -234,15 +244,15 @@ set(CUTLASS_ENABLE_F16C OFF CACHE BOOL "Enable F16C x86 extensions in host code.
set(KERNEL_FILTER_FILE "" CACHE STRING "KERNEL FILTER FILE FULL PATH")
if (KERNEL_FILTER_FILE AND NOT CUTLASS_LIBRARY_KERNELS)
if (KERNEL_FILTER_FILE AND NOT CUTLASS_LIBRARY_KERNELS)
# If a kernel filter file is specified, we want to generate and then
# filter on the entire kernel set, not the default kernel
# (sub)set. The user may overried CUTLASS_LIBRRARY_KERNELS, in which
# (sub)set. The user may have overridden CUTLASS_LIBRRARY_KERNELS, in which
# case the resulting kernel set will be the intersection of the two
# options differenced against CUTLASS_LIBRARY_IGNORE_KERNELS.
set(CUTLASS_LIBRARY_KERNELS_INIT "*")
else()
set(CUTLASS_LIBRARY_KERNELS_INIT "")
else()
set(CUTLASS_LIBRARY_KERNELS_INIT "")
endif()
if (KERNEL_FILTER_FILE)
@@ -256,9 +266,10 @@ if(KERNEL_FILTER_FILE)
message(STATUS "Full path of filter file: ${KERNEL_FILTER_FILE}")
endif()
set(CUTLASS_LIBRARY_OPERATIONS "all" CACHE STRING "Comma delimited list of operation name filters. Default '' means all operations are enabled.")
set(CUTLASS_LIBRARY_KERNELS ${CUTLASS_LIBRARY_KERNELS_INIT} CACHE STRING "Comma delimited list of kernel name filters. If unspecified, only the largest tile size is enabled. If 'all' is specified, all kernels are enabled.")
set(CUTLASS_LIBRARY_IGNORE_KERNELS "" CACHE STRING "Comma delimited list of kernel names to exclude from build.")
set(CUTLASS_LIBRARY_OPERATIONS "all" CACHE STRING "Comma-delimited list of operation name filters. Default '' means all operations are enabled.")
set(CUTLASS_LIBRARY_KERNELS ${CUTLASS_LIBRARY_KERNELS_INIT} CACHE STRING "Comma-delimited list of kernel name filters. If unspecified, only the largest tile size is enabled. If the string 'all' is specified, all kernels are enabled.")
set(CUTLASS_LIBRARY_IGNORE_KERNELS "" CACHE STRING "Comma-delimited list of kernels to exclude from build. This option ONLY takes effect if CUTLASS_LIBRARY_KERNELS is set.")
set(CUTLASS_LIBRARY_EXCLUDE_KERNELS "" CACHE STRING "Comma-delimited list of kernels to exclude from build. This option always takes effect, whether or not CUTLASS_LIBRARY_KERNELS is set. It also can exclude kernels from the filter file (see KERNEL_FILTER_FILE).")
################################################################################
@@ -330,6 +341,11 @@ if (CUTLASS_ENABLE_TENSOR_CORE_MMA)
list(APPEND CUTLASS_CUDA_FLAGS -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1)
endif()
set(CUTLASS_PROFILER_DISABLE_REFERENCE OFF CACHE BOOL "Disable compilation of reference kernels in the CUTLASS profiler.")
if (CUTLASS_PROFILER_DISABLE_REFERENCE)
list(APPEND CUTLASS_CUDA_NVCC_FLAGS -DCUTLASS_PROFILER_DISABLE_REFERENCE=1)
endif()
@@ -398,8 +414,8 @@ if(CUDA_COMPILER MATCHES "[Cc]lang")
message(FATAL_ERROR "Clang CUDA compilation requires Clang CXX compilation. Currently CMAKE_CXX_COMPILER is ${CMAKE_CXX_COMPILER_ID}" )
endif()
# There are numerous Clang versions that can work with each CUDA toolkit and the
# the checks are not very useful so we are turning them off and using testing to
# There are numerous Clang versions that can work with each CUDA toolkit and the
# the checks are not very useful so we are turning them off and using testing to
# ensure the various combinations work properly.
list(APPEND CUTLASS_CUDA_CLANG_FLAGS --cuda-path=${CUDA_TOOLKIT_ROOT_DIR})
@@ -425,7 +441,14 @@ if(CUDA_COMPILER MATCHES "[Cc]lang")
link_libraries(nvidia::cuda_driver)
endif()
# Support for 128-bit integers if using NVIDIA C++ compiler
# Known gcc 8.1-8.3 SFINAE issue (fixed in gcc 8.4), check https://gcc.gnu.org/bugzilla/show_bug.cgi?id=87748
# Also see https://github.com/NVIDIA/nccl/issues/835 for nvtx3.hpp
if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 8.1 AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS_EQUAL 8.3)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DNVTX3_USE_CHECKED_OVERLOADS_FOR_GET=0")
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -DNVTX3_USE_CHECKED_OVERLOADS_FOR_GET=0")
endif()
# Support for 128-bit integers if using NVIDIA C++ compiler
if (${CMAKE_CXX_COMPILER_ID} MATCHES "PGI" OR ${CMAKE_CXX_COMPILER_ID} MATCHES "NVHPC")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Mint128 ")
endif()
@@ -433,24 +456,24 @@ endif()
if (CMAKE_VERSION VERSION_GREATER_EQUAL 3.18)
# CMake 3.18 added support for CUDA_ARCHITECTURES target property. We will use this
# property for CMake 3.18+, so we request the NEW behavior for correct compatibility.
# https://cmake.org/cmake/help/v3.18/policy/CMP0104.html#policy:CMP0104
# https://cmake.org/cmake/help/v3.18/policy/CMP0104.html#policy:CMP0104
cmake_policy(SET CMP0104 NEW)
endif()
if (MSVC)
# MSVC by default does not apply the correct __cplusplus version as specified by the C++ standard
# because MSVC is not a completely compliant implementation. This option forces MSVC to use the
# because MSVC is not a completely compliant implementation. This option forces MSVC to use the
# appropriate value given the requested --std option. This fixes a compilation issue mismatch
# between GCC/Clang and MSVC.
#
# error : a constexpr function cannot have a nonliteral return type "dim3"
#
#
# See https://developercommunity.visualstudio.com/t/msvc-incorrectly-defines-cplusplus/139261
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /Zc:__cplusplus")
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler /Zc:__cplusplus")
endif()
# Some tests require this build option in order to link.
@@ -488,7 +511,7 @@ function(cutlass_apply_cuda_gencode_flags TARGET)
list(JOIN CODES "," CODES_STR)
list(APPEND NVCC_FLAGS -gencode=arch=compute_${ARCH},code=[${CODES_STR}])
endforeach()
if (NOT __SM_ARCHS)
if (CUDA_COMPILER MATCHES "[Cc]lang")
target_compile_options(
@@ -523,7 +546,7 @@ function(cutlass_apply_cuda_gencode_flags TARGET)
endfunction()
# Cache the flags so they are available when the function below is called anywhere globally.
# Cache the flags so they are available when the function below is called anywhere globally.
set(__CUTLASS_CUDA_FLAGS ${CUTLASS_CUDA_FLAGS} CACHE INTERNAL "")
set(__CUTLASS_CUDA_FLAGS_RELEASE ${CUTLASS_CUDA_FLAGS_RELEASE} CACHE INTERNAL "")
@@ -694,6 +717,7 @@ if(NOT WIN32)
"-Wl,-rpath,'$ORIGIN/../lib'"
"-Wl,-rpath,'${CUDA_TOOLKIT_ROOT_DIR}/lib64'"
"-Wl,-rpath,'${CUDA_TOOLKIT_ROOT_DIR}/lib'"
${CMAKE_DL_LIBS}
)
endif()
@@ -757,24 +781,24 @@ set(CUTLASS_CTEST_TEMPLATE_FILE ${CMAKE_CURRENT_LIST_DIR}/cmake/CTestTestfile.co
set(CUTLASS_CTEST_GENERATED_FILES "" CACHE INTERNAL "")
function(cutlass_add_executable_tests NAME TARGET)
#
# Generates test rules for `make test`, `make test_all`, and `ctest` invoked from either the
#
# Generates test rules for `make test`, `make test_all`, and `ctest` invoked from either the
# <CMAKE_BINARY_DIR> or the <CMAKE_INSTALL_PREFIX>/<CUTLASS_TEST_INSTALL_PREFIX> after installation.
#
#
# NAME: The base name for the test. Can be run with `make <NAME>` or `ctest -R 'c<NAME>'`.
# TARGET: The target corresponding to the executable under test.
# DISABLE_EXECUTABLE_INSTALL_RULE: An option, if given, that disables creating an install rule for TARGET.
# DEPENDS: A list of targets or files on which this test is dependent.
# DEPENDEES: A list of targets which should depend on this test.
# TEST_COMMAND_OPTIONS: A list of variables (i.e. by reference params) which contain command line arguments
# to pass to the test executable. A unique test is generated for each set of
# to pass to the test executable. A unique test is generated for each set of
# options given. If this option is not used, a single test with no arguments is generated.
# TEST_COMMAND_OPTIONS_PREFIX: If provided, is added as a prefix to each TEST_COMMAND_OPTIONS value for
# TEST_COMMAND_OPTIONS_PREFIX: If provided, is added as a prefix to each TEST_COMMAND_OPTIONS value for
# generating the full variable name to be referenced.
# RESULT_CACHE_FILE: A file to be installed alongside the test executable with pre-computed
# test results to speed up test runtime.
# TEST_SETS_SUPPORTED: A list of test set names these tests support.
#
# TEST_SETS_SUPPORTED: A list of test set names these tests support.
#
set(options DISABLE_EXECUTABLE_INSTALL_RULE)
set(oneValueArgs DISABLE_TESTS RESULT_CACHE_FILE TEST_COMMAND_OPTIONS_PREFIX)
@@ -806,9 +830,9 @@ function(cutlass_add_executable_tests NAME TARGET)
endif()
if (NOT __DISABLE_EXECUTABLE_INSTALL_RULE AND CUTLASS_INSTALL_TESTS)
# file(RELATIVE_PATH CMAKE_CURRENT_BINARY_RELATIVE_DIR ${CMAKE_BINARY_DIR} ${CMAKE_CURRENT_BINARY_DIR})
install(
TARGETS ${TARGET}
RUNTIME DESTINATION ${CUTLASS_TEST_INSTALL_BINDIR}
@@ -822,7 +846,7 @@ function(cutlass_add_executable_tests NAME TARGET)
)
endif()
endif()
if (NOT __TEST_COMMAND_OPTIONS)
@@ -856,10 +880,10 @@ function(cutlass_add_executable_tests NAME TARGET)
string(TOLOWER "${NAME}" TEST_NAME)
endif()
# The following rigmarole is needed to deal with spaces and possible quotes in
# The following rigmarole is needed to deal with spaces and possible quotes in
# command line arguments. The options are passed "by reference" as the actual
# variable names holding the real options. We then expand these in a way that
# preserves any quotes. Note, they have to be in this order for it to work for
# preserves any quotes. Note, they have to be in this order for it to work for
# all the use cases below.
set(TEST_COMMAND_OPTIONS ${${__TEST_COMMAND_OPTIONS_PREFIX}${CMD_OPTIONS_VAR}})
@@ -889,7 +913,7 @@ function(cutlass_add_executable_tests NAME TARGET)
endforeach()
# To run the tests from an install package with tests enabled, we need to generate test files
# that don't rely on the current directory structure in build.
# that don't rely on the current directory structure in build.
set(TEST_NAME c${NAME})
set(TEST_GEN_DIR ${CMAKE_CURRENT_BINARY_DIR}/ctest/${TEST_NAME})
@@ -906,14 +930,14 @@ function(cutlass_add_executable_tests NAME TARGET)
# The following line imports the tests for immediate run via `make test`.
include(${TEST_GEN_DIR}/CTestTestfile.${TEST_NAME}.cmake)
set(CUTLASS_CTEST_GENERATED_FILES ${CUTLASS_CTEST_GENERATED_FILES};ctest/${TEST_NAME}/CTestTestfile.${TEST_NAME}.cmake CACHE INTERNAL "")
if (CUTLASS_INSTALL_TESTS)
file(GENERATE
OUTPUT "${TEST_GEN_DIR}/CTestTestfile.${TEST_NAME}.install.cmake"
INPUT "${TEST_GEN_DIR}/CTestTestfile.${TEST_NAME}.install.cmake.in"
file(GENERATE
OUTPUT "${TEST_GEN_DIR}/CTestTestfile.${TEST_NAME}.install.cmake"
INPUT "${TEST_GEN_DIR}/CTestTestfile.${TEST_NAME}.install.cmake.in"
)
install(
@@ -971,19 +995,19 @@ endif()
include(CMakePackageConfigHelpers)
write_basic_package_version_file(
${CMAKE_CURRENT_BINARY_DIR}/NvidiaCutlassConfigVersion.cmake
${CMAKE_CURRENT_BINARY_DIR}/NvidiaCutlassConfigVersion.cmake
COMPATIBILITY AnyNewerVersion)
configure_file(
${CMAKE_CURRENT_SOURCE_DIR}/cmake/NvidiaCutlassConfig.cmake.in
${CMAKE_CURRENT_BINARY_DIR}/NvidiaCutlassConfig.cmake
${CMAKE_CURRENT_BINARY_DIR}/NvidiaCutlassConfig.cmake
@ONLY
)
install(
FILES
${CMAKE_CURRENT_BINARY_DIR}/NvidiaCutlassConfig.cmake
${CMAKE_CURRENT_BINARY_DIR}/NvidiaCutlassConfigVersion.cmake
FILES
${CMAKE_CURRENT_BINARY_DIR}/NvidiaCutlassConfig.cmake
${CMAKE_CURRENT_BINARY_DIR}/NvidiaCutlassConfigVersion.cmake
DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/NvidiaCutlass/
)
README.md
+25 -3
View File
@@ -1,8 +1,8 @@
![ALT](./media/images/gemm-hierarchy-with-epilogue-no-labels.png "Complete CUDA GEMM decomposition")
# CUTLASS 3.5
# CUTLASS 3.5.1
_CUTLASS 3.5 - April 2024_
_CUTLASS 3.5.1 - July 2024_
CUTLASS is a collection of CUDA C++ template abstractions for implementing
high-performance matrix-matrix multiplication (GEMM) and related computations at all levels
@@ -41,9 +41,30 @@ and improves code composability and readability. More documentation specific to
In addition to GEMMs, CUTLASS implements high-performance convolution via the implicit GEMM algorithm. Implicit GEMM is the formulation of a convolution operation as a GEMM thereby taking advantage of CUTLASS's modular GEMM pipeline. This allows CUTLASS to build convolutions by reusing highly-optimized GEMM components.
# What's New in CUTLASS 3.5
CUTLASS 3.5 is an update to CUTLASS adding:
CUTLASS 3.5.1 is an update to CUTLASS adding:
- [Minimal SM90 WGMMA + TMA GEMM example in 100 lines of code](./examples/cute/tutorial/wgmma_sm90.cu).
- [Exposure of L2 `cache_hint`s in TMA copy atoms](./include/cute/arch/copy_sm90_tma.hpp#L48)
- Exposure of raster order and tile swizzle extent in [CUTLASS library profiler](./media/docs/profiler.md#GEMM), and
[example 48](./examples/48_hopper_warp_specialized_gemm/48_hopper_warp_specialized_gemm.cu).
- [TMA store based and EVT supported epilogues](./include/cutlass/epilogue/collective/sm90_epilogue_array_tma_warpspecialized.hpp) for [Hopper pointer array batched kernels](./test/unit/gemm/device/sm90_gemm_f16_f16_f16_tensor_op_f32_ptr_array.cu).
- A new [`GemmSparseUniversal` API for CUTLASS 2.x Ampere kernels](./include/cutlass/gemm/device/gemm_sparse_universal.h) leveraging 2:4 structured sparsity and [support for LLM friendly tile sizes](./test/unit/gemm/device/gemm_f16n_f16t_f32t_tensor_op_f32_sparse_sm80.cu).
- [CUDA host adapter](./include/cutlass/cuda_host_adapter.hpp) extensions to support TMA descriptor construction driver APIs.
- Inclusion of more [Hopper fprop, dgrad, and wgrad convolution kernels in CUTLASS library and profiler](./python/cutlass_library/generator.py).
- Support for residual add (beta != 0) in convolution kernels.
- A refactor of [include files throughout CUTLASS core directories](./include/cutlass/gemm/collective/collective_mma_decl.hpp) to reduce circular dependencies and [tests to guard against them](./test/self_contained_includes/CMakeLists.txt).
- [A guide for setting up VSCode to work well with CUTLASS](./media/docs/ide_setup.md) and [expanded code style guide](./media/docs/programming_guidelines.md).
- Better support for MSVC as a host compiler.
- Many performance optimizations, improvements, and bug fixes including fixes for FlashAttention-2.
- Optimal code generation with CUDA toolkit versions 12.4 and 12.5u1.
- NOTICE:
+ Upcoming CUTLASS 3.6 release will include a breaking refactor to the CUTLASS 3.x convolution `kernel::ConvUniversal` API to bring it in line with `gemm::GemmUniversal`. After this, the 3.x convolution API will no longer be considered as a beta API.
+ Upcoming CUTLASS 3.6 release will include a breaking refactor to the Hopper TMA pointer array batched epilogue in order to support grouped GEMMs.
CUTLASS 3.5.0 is an update to CUTLASS adding:
- Implicit GEMM Convolutions targeting Hopper SM90A via WGMMA + [TMA im2col](./include/cute/atom/copy_traits_sm90_im2col.hpp).
+ Native implementation in CUTLASS 3.x using CuTe, mirroring the [same design hierarchy as that of GEMMs](./media/docs/gemm_api_3x.md).
@@ -61,6 +82,7 @@ CUTLASS 3.5 is an update to CUTLASS adding:
- Remove C++11 requirement on a few CUTLASS 2.x API header files. All CUTLASS files now require C++17.
- Fixes to greatly reduce build warnings.
- Updates and bugfixes from the community (thanks!)
- CUTLASS 3.5.1 is a minor update to CUTLASS containing small bug fixes and improvements, including fixes for FlashAttention-2 builds.
Minimum requirements:
examples/07_volta_tensorop_gemm/volta_tensorop_gemm.cu
+1 -1
View File
@@ -162,7 +162,7 @@ using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 64, 32>; // <- warp tile M =
using ShapeMMAOp = cutlass::gemm::GemmShape<8, 8, 4>; // <- MMA Op tile M = 8, N = 8, K = 4
// This code section describes how threadblocks are scheduled on GPU
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; // <- ??
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
// This code section describes ?
using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
examples/08_turing_tensorop_gemm/turing_tensorop_gemm.cu
+1 -1
View File
@@ -161,7 +161,7 @@ using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 64, 64>; // <- warp tile M =
using ShapeMMAOp = cutlass::gemm::GemmShape<8, 8, 16>; // <- MMA Op tile M = 8, N = 8, K = 16
// This code section describes how threadblocks are scheduled on GPU
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; // <- ??
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
// This code section describes the epilogue part of the kernel
using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
examples/12_gemm_bias_relu/gemm_bias_relu.cu
+1 -1
View File
@@ -84,7 +84,7 @@ using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 64, 32>; // <- warp tile M =
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<>; // <- ??
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
// Define the epilogue operation as LinearCombinationRelu. This is approximately equal to
//
examples/13_two_tensor_op_fusion/b2b_grouped_gemm_run.h
+2 -2
View File
@@ -231,7 +231,7 @@ struct B2bFusedGroupedGemmRun
host_tensor_ref_D1.at(i).sync_device();
ref_A0.at(i) = (host_tensor_A0.at(i).device_ref());
ref_B0.at(i) = (host_tensor_B0.at(i).device_ref());;
ref_B0.at(i) = (host_tensor_B0.at(i).device_ref());
ref_C0.at(i) = (host_tensor_C0.at(i).device_ref());
if (alpha0 == ElementCompute(0)) //per-channel scale
ref_Scale0.at(i) = (host_tensor_Scale0.at(i).device_ref());
@@ -340,7 +340,7 @@ struct B2bFusedGroupedGemmRun
std::cout << "Fusion time " << gemmTime / (float)runs << " ms\n";
for (int i = 0; i < problem_count; ++i) {
host_tensor_D1.at(i).sync_host();;
host_tensor_D1.at(i).sync_host();
//
// Verify
examples/14_ampere_tf32_tensorop_gemm/ampere_tf32_tensorop_gemm.cu
+1 -1
View File
@@ -194,7 +194,7 @@ using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 64, 16>; // <- warp tile M =
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<>; // <- ??
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
// This code section describes the epilogue part of the kernel
using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
examples/15_ampere_sparse_tensorop_gemm/CMakeLists.txt
+5
View File
@@ -33,6 +33,11 @@ cutlass_example_add_executable(
ampere_sparse_tensorop_gemm.cu
)
cutlass_example_add_executable(
15_ampere_sparse_tensorop_gemm_universal
ampere_sparse_tensorop_gemm_universal.cu
)
cutlass_example_add_executable(
15_ampere_sparse_tensorop_gemm_with_visitor
ampere_sparse_tensorop_gemm_with_visitor.cu
examples/15_ampere_sparse_tensorop_gemm/ampere_sparse_tensorop_gemm.cu
+1 -1
View File
@@ -84,7 +84,7 @@ using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 64, 256>; // <- warp tile M =
using ShapeMMAOp = cutlass::gemm::GemmShape<16, 8, 128>; // <- MMA Op tile M = 16, N = 8, K = 128
// This code section describes how threadblocks are scheduled on GPU
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; // <- ??
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
// This code section describes the epilogue part of the kernel
using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
examples/15_ampere_sparse_tensorop_gemm/ampere_sparse_tensorop_gemm_universal.cu
+329
View File
@@ -0,0 +1,329 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/**
Please check example 07, 08 and 17 for the basics of dense tensor op gemm kernels. NVIDIA Ampere
architecture also supports structured sparse tensor op for tf32, fp16, int8 and int4.
Sparse GEMM kernels needs to takes an additional E matrix which stores the meta data. The format of
meta data is different for every data types. CUTLASS templates can automatically infer it based on
input A and B. Check code below.
Moreover, matrix E needs to be preprocessed so that it can use ldmatrix to load into the registers
efficiently.
*/
#include <iostream>
#include "cutlass/cutlass.h"
#include "cutlass/gemm/device/gemm_sparse_universal.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/reference/host/gemm.h"
#include "cutlass/util/host_reorder.h"
#include "cutlass/util/host_uncompress.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 "helper.h"
// The code section below describes datatype for input, output matrices and computation between
// elements in input matrices.
using ElementAccumulator = int32_t; // <- data type of accumulator
using ElementComputeEpilogue = ElementAccumulator; // <- data type of epilogue operations
using ElementInputA = cutlass::int4b_t; // <- data type of elements in input matrix A
using ElementInputB = cutlass::int4b_t; // <- data type of elements in input matrix B
using ElementOutput = int32_t; // <- data type of elements in output matrix D
// The code section below describes matrix layout of input and output matrices. Row Major for
// Matrix A, Column 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, 128, 256>; // <- threadblock tile M = 128, N = 128, K = 256
// This code section describes tile size a warp will compute
using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 64, 256>; // <- warp tile M = 64, N = 64, K = 256
// This code section describes the size of MMA op
using ShapeMMAOp = cutlass::gemm::GemmShape<16, 8, 128>; // <- MMA Op tile M = 16, N = 8, K = 128
// 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<
ElementOutput, // <- data type of output matrix
128 / cutlass::sizeof_bits<ElementOutput>::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
ElementAccumulator, // <- data type of accumulator
ElementComputeEpilogue>; // <- data type for alpha/beta in linear combination function
// Number of pipelines you want to use
constexpr int NumStages = 3;
using Gemm = cutlass::gemm::device::GemmSparseUniversal<ElementInputA,
LayoutInputA,
ElementInputB,
LayoutInputB,
ElementOutput,
LayoutOutput,
ElementAccumulator,
MMAOp,
SmArch,
ShapeMMAThreadBlock,
ShapeMMAWarp,
ShapeMMAOp,
EpilogueOp,
SwizzleThreadBlock,
NumStages>;
// Data type and layout of meta data matrix E can be inferred from template Gemm.
using ElementInputE = typename Gemm::ElementE;
using LayoutInputE = cutlass::layout::RowMajor;
using ReorderedLayoutInputE = typename Gemm::LayoutE;
// Blow property is defined in include/cutlass/arch/sp_mma_sm80.h
// 50% Sparsity on Ampere
constexpr int kSparse = Gemm::kSparse;
// How many elements of A are covered per ElementE
constexpr int kElementsPerElementE = Gemm::kElementsPerElementE;
// The size of individual meta data
constexpr int kMetaSizeInBits = Gemm::kMetaSizeInBits;
int run() {
const int length_m = 512;
const int length_n = 512;
const int length_k = 1024;
// Create a tuple of problem size for matrix multiplication
cutlass::gemm::GemmCoord problem_size(length_m, length_n, length_k);
// Initialize tensors using CUTLASS helper functions
cutlass::HostTensor<ElementInputA, LayoutInputA> tensor_a(
cutlass::make_Coord(problem_size.m(), problem_size.k() / kSparse)); // <- Create matrix A with dimensions M x (K / 2)
cutlass::HostTensor<ElementInputA, LayoutInputA> tensor_a_uncompressed(
problem_size.mk()); // <- Create uncompressed matrix A with dimensions M x K for reference computing
cutlass::HostTensor<ElementInputB, LayoutInputB> tensor_b(
problem_size.kn()); // <- Create matrix B with dimensions K x N
cutlass::HostTensor<ElementOutput, LayoutOutput> tensor_c(
problem_size.mn()); // <- Create matrix C with dimensions M x N
cutlass::HostTensor<ElementOutput, LayoutOutput> tensor_d(
problem_size.mn()); // <- Create matrix D with dimensions M x N used to store output from
// CUTLASS kernel
cutlass::HostTensor<ElementOutput, LayoutOutput> tensor_ref_d(
problem_size.mn()); // <- Create matrix D with dimensions M x N used to store output from
// reference kernel
// Create matrix E with dimensions M x (K / 2 / kElementsPerElementE). This one is used by reference computing.
cutlass::HostTensor<ElementInputE, LayoutInputE> tensor_e(
cutlass::make_Coord(problem_size.m(), problem_size.k() / kSparse / kElementsPerElementE));
// Same size as the above. The above one needs to be reordered and stored in this one.
cutlass::HostTensor<ElementInputE, ReorderedLayoutInputE> tensor_e_reordered(
cutlass::make_Coord(problem_size.m(), problem_size.k() / kSparse / kElementsPerElementE));
// Fill input and output matrices on host using CUTLASS helper functions
cutlass::reference::host::TensorFillRandomUniform(
tensor_a.host_view(),
1,
ElementInputA(2),
ElementInputA(-2),
0); // <- Fill matrix A on host with uniform-distribution random data
cutlass::reference::host::TensorFillRandomUniform(
tensor_b.host_view(),
1,
ElementInputB(2),
ElementInputB(-2),
0); // <- Fill matrix B on host with uniform-distribution random data
cutlass::reference::host::TensorFillRandomUniform(
tensor_c.host_view(),
1,
ElementOutput(2),
ElementOutput(-2),
0); // <- Fill matrix C on host with uniform-distribution random data
cutlass::reference::host::TensorFillRandomSparseMeta(
tensor_e.host_view(),
1,
kMetaSizeInBits); // <- Fill matrix E on host with uniform-distribution random meta data
cutlass::reference::host::TensorFill(
tensor_d.host_view()); // <- fill matrix D on host with zeros
cutlass::reference::host::TensorFill(
tensor_ref_d.host_view()); // <- fill matrix D for reference on host with zeros
// Reorder the meta data matrix so that we can use ldmatrix to load them to tensor core
// instructions.
cutlass::reorder_meta(tensor_e_reordered.host_ref(), tensor_e.host_ref(),
{problem_size.m(), problem_size.n(),
problem_size.k() / kSparse / kElementsPerElementE});
// Copy data from host to GPU
tensor_a.sync_device();
tensor_b.sync_device();
tensor_c.sync_device();
tensor_d.sync_device();
tensor_e_reordered.sync_device();
tensor_ref_d.sync_device();
// Initialize alpha and beta for dot product computation
ElementComputeEpilogue alpha = ElementComputeEpilogue(1);
ElementComputeEpilogue beta = ElementComputeEpilogue(0);
// Split K dimension into 1 partitions
int split_k_slices = 2;
// Create a tuple of gemm kernel arguments. This is later passed as arguments to launch
// instantiated CUTLASS kernel
typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGemm,
problem_size, // <- problem size of matrix multiplication
split_k_slices,// <- k-dimension split factor
{alpha, beta}, // <- tuple of alpha and beta
tensor_a.device_data(), // <- reference to matrix A on device
tensor_b.device_data(), // <- reference to matrix B on device
tensor_c.device_data(), // <- reference to matrix C on device
tensor_d.device_data(), // <- reference to matrix D on device
tensor_e_reordered.device_data(), // <- reference to matrix E on device
int64_t(),
int64_t(),
int64_t(),
int64_t(),
int64_t(),
tensor_a.layout().stride(0),
tensor_b.layout().stride(0),
tensor_c.layout().stride(0),
tensor_d.layout().stride(0),
tensor_e_reordered.layout().stride(0)
};
// Using the arguments, query for extra workspace required for matrix multiplication computation
size_t workspace_size = Gemm::get_workspace_size(arguments);
// Allocate workspace memory
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
// Instantiate CUTLASS kernel depending on templates
Gemm gemm_op;
// Check the problem size is supported or not
cutlass::Status status = gemm_op.can_implement(arguments);
CUTLASS_CHECK(status);
// Initialize CUTLASS kernel with arguments and workspace pointer
status = gemm_op.initialize(arguments, workspace.get());
CUTLASS_CHECK(status);
// Launch initialized CUTLASS kernel
status = gemm_op();
CUTLASS_CHECK(status);
// uncompress tensor_a based on meta data tensor_e. We need it for reference computing.
cutlass::uncompress(tensor_a_uncompressed.host_ref(), tensor_a.host_ref(),
tensor_e.host_ref(), problem_size.m(), problem_size.k());
// Create instantiation for host reference gemm kernel
cutlass::reference::host::Gemm<ElementInputA,
LayoutInputA,
ElementInputB,
LayoutInputB,
ElementOutput,
LayoutOutput,
ElementComputeEpilogue,
ElementComputeEpilogue,
typename Gemm::Operator>
gemm_host;
// Launch host reference gemm kernel
gemm_host(problem_size,
alpha,
tensor_a_uncompressed.host_ref(),
tensor_b.host_ref(),
beta,
tensor_c.host_ref(),
tensor_ref_d.host_ref());
// Copy output data from CUTLASS host for comparison
tensor_d.sync_host();
// Check if output from CUTLASS kernel and reference kernel are equal or not
bool passed = cutlass::reference::host::TensorEquals(
tensor_d.host_view(),
tensor_ref_d.host_view());
std::cout << (passed ? "Passed" : "Failed") << std::endl;
return (passed ? 0 : -1);
}
int main() {
bool notSupported = false;
// Ampere Sparse Tensor Core operations exposed with mma.sync and ldmatrix are first available
// in CUDA 11.1.
//
// CUTLASS must be compiled with CUDA 11.1 Toolkit to run these examples.
if (!(__CUDACC_VER_MAJOR__ > 11 || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 1))) {
std::cerr << "Ampere Tensor Core operations must be compiled with CUDA 11.1 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 -1;
}
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;
}
return run();
}
examples/15_ampere_sparse_tensorop_gemm/ampere_sparse_tensorop_gemm_with_visitor.cu
+1 -1
View File
@@ -94,7 +94,7 @@ using ShapeMMAOp = cutlass::gemm::GemmShape<16, 8, 64>; // <- MMA Op tile M = 1
// This code section describes how threadblocks are scheduled on GPU
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
using Operator = cutlass::arch::OpMultiplyAdd;
using Operator = cutlass::arch::OpMultiplyAddSaturate;
// Number of pipelines you want to use
constexpr int NumStages = 3;
examples/23_ampere_gemm_operand_reduction_fusion/ampere_gemm_operand_reduction_fusion.cu
+2 -2
View File
@@ -138,7 +138,7 @@ using Gemm = typename cutlass::gemm::device::GemmWithKReduction<
>;
// Below is the reduction kernel used in the case of parallel split-k
using ReduceGemmSplitKShape = cutlass::MatrixShape<4, 64>;;
using ReduceGemmSplitKShape = cutlass::MatrixShape<4, 64>;
using ReduceOp = cutlass::reduction::thread::ReduceAdd<
ElementAccumulator,
@@ -154,7 +154,7 @@ using ReduceGemmSplitKKernel = cutlass::reduction::kernel::ReduceSplitK<
using ReduceGemmSplitK = cutlass::reduction::device::ReduceSplitK<ReduceGemmSplitKKernel>;
using ReduceVectorSplitKShape = cutlass::MatrixShape<1, 256>;;
using ReduceVectorSplitKShape = cutlass::MatrixShape<1, 256>;
// This code section describes the epilogue part of the kernel, we use default value
using DummyEpilogueOp = cutlass::epilogue::thread::LinearCombination<
examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/27_ampere_3xtf32_fast_accurate_tensorop_gemm.cu
+1 -1
View File
@@ -258,7 +258,7 @@ using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 32, 16>; // <- warp tile M =
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<>; // <- ??
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
// This code section describes the epilogue part of the kernel
using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
examples/29_ampere_3xtf32_fast_accurate_tensorop_complex_gemm/29_3xtf32_complex_gemm.cu
+1 -1
View File
@@ -221,7 +221,7 @@ using ShapeMMAWarp = cutlass::gemm::GemmShape<32, 32, 16>; // <- warp tile M =
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<>; // <- ??
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
// This code section describes the epilogue part of the kernel
using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
examples/33_ampere_3xtf32_tensorop_symm/ampere_3xtf32_tensorop_symm.cu
+1 -1
View File
@@ -193,7 +193,7 @@ using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 32, 16>; // <- warp tile M =
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<>; // <- ??
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
// This code section describes the epilogue part of the kernel
using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
examples/36_gather_scatter_fusion/gather_scatter_fusion.cu
+1 -1
View File
@@ -215,7 +215,7 @@ using ShapeMMAOp = cutlass::gemm::GemmShape<16, 8, 16>; // <- MMA Op tile M = 8
// 16, 8, 16 -> Ampere
// This code section describes how threadblocks are scheduled on GPU
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; // <- ??
using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
// Define the epilogue operation as LinearCombination. This is approximately equal to
//
examples/40_cutlass_py/conv2d.py
+1 -1
View File
@@ -118,7 +118,7 @@ operation = Conv2dOperation(
conv_kind=cutlass_bindings.conv.Operator.fprop,
iterator_algorithm=cutlass_bindings.conv.IteratorAlgorithm.optimized,
arch=cc, tile_description=tile_description,
A=A, B=B, C=C, stride_support=StrideSupport.Strided,
A=A, B=B, C=C, stride_support=StrideSupport.Unity,
epilogue_functor=epilogue_functor
)
examples/41_fused_multi_head_attention/gemm/mma_from_smem.h
+2 -2
View File
@@ -1799,7 +1799,7 @@ struct B2bGemm<
if (rowIdx == 1) {
lse_prefetched[colIdx] = accum_n < lse_extent
? lse[accum_n]
: platform::numeric_limits<accum_t>::infinity();
: cutlass::platform::numeric_limits<accum_t>::infinity();
}
accum[idx] = expf(accum[idx] - lse_prefetched[colIdx]);
++colIdx;
@@ -1938,7 +1938,7 @@ struct B2bGemm<
if (rowIdx == 1) {
lse_prefetched[colIdx] = accum_n < lse_extent
? lse[accum_n]
: platform::numeric_limits<accum_t>::infinity();
: cutlass::platform::numeric_limits<accum_t>::infinity();
}
accum[idx] = expf(accum[idx] - lse_prefetched[colIdx]);
++colIdx;
examples/48_hopper_warp_specialized_gemm/48_hopper_warp_specialized_gemm.cu
+51 -10
View File
@@ -35,19 +35,23 @@
This example demonstrate a simple way to instantiate and run a TF32 GEMM using the new CUTLASS 3.0
APIs on NVIDIA Hopper architecture. New features that will be showcased in this example are as follows:
1. NVIDIA Hopper architecture introduces a new series of tensor core instructions (GMMA)
1. NVIDIA Hopper architecture introduces a new series of tensor core instructions (GMMA)
which are more efficient than the Ampere tensor core instructions.
2. NVIDIA Hopper architecture includes new Tensor Memory Accelerator (TMA) unit to transfer large
2. NVIDIA Hopper architecture includes new Tensor Memory Accelerator (TMA) unit to transfer large
blocks of data efficiently between global memory and shared memory. TMA also supports asynchronous
copies between thread blocks in a cluster. Another advantage is that TMA can load in FP32 data and
convert them implicitly to TF32.
3. This example uses the Warp Specialized kernel design (see /media/docs/efficient_gemm.md for details).
4. A simple way to tune the CTA rasterization direction and swizzle pattern of Hopper kernels. Both the
CTA rasterization direction and swizzle pattern impact cross-CTA locality of accesses. By tuning we can
improve performance.
Examples:
$ ./examples/48_hopper_warp_specialized_gemm/48_hopper_warp_specialized_gemm --m=2048 --n=2048 --k=2048
$ ./examples/48_hopper_warp_specialized_gemm/48_hopper_warp_specialized_gemm --m=2048 --n=2048 --k=2048 --rasterization=N --swizzle=2
*/
#include <iostream>
@@ -63,6 +67,7 @@
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/kernel/tile_scheduler_params.h"
#include "cutlass/util/command_line.h"
#include "cutlass/util/distribution.h"
@@ -105,7 +110,7 @@ using OperatorClass = cutlass::arch::OpClassTensorOp; // O
using TileShape = Shape<_128,_128,_32>; // Threadblock-level tile size
using ClusterShape = Shape<_1,_2,_1>; // Shape of the threadblocks in a cluster
using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based on the tile size
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; // Kernel to launch based on the default setting in the Collective Builder
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; // Kernel to launch based on the default setting in the Collective Builder
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
@@ -175,6 +180,8 @@ cutlass::DeviceAllocation<typename Gemm::EpilogueOutputOp::ElementOutput> block_
/// Testbed utility types
/////////////////////////////////////////////////////////////////////////////////////////////////
using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90Params::RasterOrderOptions;
// Command line options parsing
struct Options {
@@ -183,12 +190,16 @@ struct Options {
float alpha, beta;
int iterations;
int m, n, k;
RasterOrderOptions raster;
int swizzle;
Options():
help(false),
m(5120), n(4096), k(4096),
alpha(1.f), beta(0.f),
iterations(1000)
iterations(1000),
raster(RasterOrderOptions::Heuristic),
swizzle(1)
{ }
// Parses the command line
@@ -206,6 +217,21 @@ struct Options {
cmd.get_cmd_line_argument("alpha", alpha, 1.f);
cmd.get_cmd_line_argument("beta", beta, 0.f);
cmd.get_cmd_line_argument("iterations", iterations);
char raster_char;
cmd.get_cmd_line_argument("raster", raster_char);
if (raster_char == 'N' || raster_char == 'n') {
raster = RasterOrderOptions::AlongN;
}
else if (raster_char == 'M' || raster_char == 'm') {
raster = RasterOrderOptions::AlongM;
}
else if (raster_char == 'H' || raster_char == 'h') {
raster = RasterOrderOptions::Heuristic;
}
cmd.get_cmd_line_argument("swizzle", swizzle, 1);
}
/// Prints the usage statement.
@@ -220,6 +246,8 @@ struct Options {
<< " --k=<int> Sets the K extent of the GEMM\n"
<< " --alpha=<f32> Epilogue scalar alpha\n"
<< " --beta=<f32> Epilogue scalar beta\n\n"
<< " --raster=<char> CTA Rasterization direction (N for along N, M for along M, and H for heuristic)\n\n"
<< " --swizzle=<int> CTA Rasterization swizzle\n\n"
<< " --iterations=<int> Number of profiling iterations to perform.\n\n";
out
@@ -294,10 +322,10 @@ bool initialize_block(
/// Initialize operands to be used in the GEMM and reference GEMM
void initialize(const Options &options) {
stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(options.m, options.k, Int<1>{}));
stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(options.n, options.k, Int<1>{}));
stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(options.m, options.n, Int<1>{}));
stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(options.m, options.n, Int<1>{}));
stride_A = cutlass::make_cute_packed_stride(StrideA{}, {options.m, options.k, 1});
stride_B = cutlass::make_cute_packed_stride(StrideB{}, {options.n, options.k, 1});
stride_C = cutlass::make_cute_packed_stride(StrideC{}, {options.m, options.n, 1});
stride_D = cutlass::make_cute_packed_stride(StrideD{}, {options.m, options.n, 1});
block_A.reset(options.m * options.k);
block_B.reset(options.k * options.n);
@@ -320,6 +348,10 @@ typename Gemm::Arguments args_from_options(const Options &options)
{{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D}
};
arguments.scheduler.raster_order = options.raster;
// The tile scheduler will swizzle up to 8 and with the nearest multiple of 2 (i.e., 1, 2, 4, and 8)
arguments.scheduler.max_swizzle_size = options.swizzle;
return arguments;
}
@@ -408,7 +440,17 @@ int run(Options &options)
result.avg_runtime_ms = double(elapsed_ms) / double(options.iterations);
result.gflops = options.gflops(result.avg_runtime_ms / 1000.0);
std::string raster = "Heuristic";
if (options.raster == RasterOrderOptions::AlongN) {
raster = "Along N";
}
else if (options.raster == RasterOrderOptions::AlongM) {
raster = "Along M";
}
std::cout << " Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << std::endl;
std::cout << " Rasterization: " << raster << " with a maximum CTA swizzle of " << options.swizzle << std::endl;
std::cout << " Avg runtime: " << result.avg_runtime_ms << " ms" << std::endl;
std::cout << " GFLOPS: " << result.gflops << std::endl;
}
@@ -441,7 +483,6 @@ int main(int argc, char const **args) {
<< "later (compute capability 90 or greater).\n";
return 0;
}
//
// Parse options
//
examples/49_hopper_gemm_with_collective_builder/49_collective_builder.cu
+1 -2
View File
@@ -538,9 +538,8 @@ int main(int argc, char const **args) {
std::cout
<< "This example requires a GPU of NVIDIA's Hopper Architecture or "
<< "later (compute capability 90 or greater) and CUDA 12.0 or greater.\n";
return 0;
return 0;
}
//
// Parse options
//
examples/50_hopper_gemm_with_epilogue_swizzle/50_hopper_gemm_with_epilogue_swizzle.cu
+1 -2
View File
@@ -354,9 +354,8 @@ int main(int argc, char const **args) {
std::cout
<< "This example requires a GPU of NVIDIA's Hopper Architecture or "
<< "later (compute capability 90 or greater) and CUDA 12.0 or greater.\n";
return 0;
return 0;
}
//
// Parse options
//
examples/52_hopper_gather_scatter_fusion/52_hopper_gather_scatter_fusion.cu
-1
View File
@@ -627,7 +627,6 @@ int main(int argc, const char ** argv) {
std::cerr << "This example requires a device with compute capability 90 or higher.\n";
notSupported = true;
}
if (notSupported) {
return EXIT_SUCCESS; // Do not fail CI checks on unsupported systems
}
examples/52_hopper_gather_scatter_fusion/gather_gemm.hpp
+2 -3
View File
@@ -166,7 +166,7 @@ public:
to_underlying_arguments(Arguments const& args, void* workspace) {
(void) workspace;
auto problem_shape = args.problem_shape;
if constexpr (detail::IF_SWAP_AB<CollectiveMainloop>::value) {
if constexpr (detail::Has_SwapAB_v<CollectiveMainloop>) {
// swap M/N
get<0>(problem_shape) = get<1>(args.problem_shape);
get<1>(problem_shape) = get<0>(args.problem_shape);
@@ -181,8 +181,7 @@ public:
};
}
CUTLASS_HOST_DEVICE static
bool
static bool
can_implement(Arguments const& args) {
bool implementable = (args.mode == GemmUniversalMode::kGemm) or
(args.mode == GemmUniversalMode::kBatched && cute::rank(ProblemShape{}) == 4);
examples/52_hopper_gather_scatter_fusion/scatter_epilogue.hpp
+1 -1
View File
@@ -119,7 +119,7 @@ public:
}
template<class ProblemShape>
CUTLASS_HOST_DEVICE static bool
static bool
can_implement(
[[maybe_unused]] ProblemShape const& problem_shape,
[[maybe_unused]] Arguments const& args) {
examples/53_hopper_gemm_permute/53_hopper_gemm_permute.cu
-1
View File
@@ -750,7 +750,6 @@ int main(int argc, char const **argv)
std::cerr << "This example requires a device with compute capability 90 or higher.\n";
notSupported = true;
}
if (notSupported) {
return EXIT_SUCCESS; // Do not fail CI checks on unsupported systems
}
examples/54_hopper_fp8_warp_specialized_gemm/54_hopper_fp8_warp_specialized_gemm.cu
+27 -7
View File
@@ -47,9 +47,13 @@
4. This example shows all important fusions used by FP8 gemm kernels,
i.e., scale factor for A, B, C, D tensor, the abs_max value of D tensor.
5. A simple way to tune the CTA rasterization direction and swizzle pattern of Hopper kernels. Both the
CTA rasterization direction and swizzle pattern impact cross-CTA locality of accesses. By tuning we can
improve performance.
Examples:
$ ./examples/54_hopper_fp8_warp_specialized_gemm/54_hopper_fp8_warp_specialized_gemm --m=2048 --n=2048 --k=2048
$ ./examples/54_hopper_fp8_warp_specialized_gemm/54_hopper_fp8_warp_specialized_gemm --m=2048 --n=2048 --k=2048 --rasterization=N --swizzle=2
*/
#include <iostream>
@@ -63,6 +67,7 @@
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/kernel/tile_scheduler_params.h"
#include "cutlass/epilogue/dispatch_policy.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
@@ -214,6 +219,8 @@ cutlass::HostTensor<ElementAmax , LayoutScalar> reference_abs_max_aux;
/// Testbed utility types
/////////////////////////////////////////////////////////////////////////////////////////////////
using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90Params::RasterOrderOptions;
/// Result structure
struct Result
{
@@ -273,7 +280,7 @@ bool initialize_tensor(
}
/// Initialize operands to be used in the GEMM and reference GEMM
void initialize(const Options &options) {
void initialize(const Options<RasterOrderOptions> &options) {
stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(options.m, options.k, options.l));
stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(options.n, options.k, options.l));
@@ -346,7 +353,7 @@ void initialize(const Options &options) {
}
/// Populates a Gemm::Arguments structure from the given commandline options
typename Gemm::Arguments args_from_options(const Options &options)
typename Gemm::Arguments args_from_options(const Options<RasterOrderOptions> &options)
{
typename Gemm::Arguments arguments{
cutlass::gemm::GemmUniversalMode::kGemm,
@@ -392,10 +399,14 @@ typename Gemm::Arguments args_from_options(const Options &options)
fusion_args.amax_D_ptr = abs_max_D.device_data();
}
arguments.scheduler.raster_order = options.raster;
// The tile scheduler will swizzle up to 8 and with the nearest multiple of 2 (i.e., 1, 2, 4, and 8)
arguments.scheduler.max_swizzle_size = options.swizzle;
return arguments;
}
bool verify(const Options &options) {
bool verify(const Options<RasterOrderOptions> &options) {
//
// Compute reference output
//
@@ -468,7 +479,7 @@ bool verify(const Options &options) {
/// Execute a given example GEMM computation
template <typename Gemm>
int run(Options &options)
int run(Options<RasterOrderOptions> &options)
{
initialize(options);
@@ -518,7 +529,17 @@ int run(Options &options)
result.avg_runtime_ms = double(elapsed_ms) / double(options.iterations);
result.gflops = options.gflops(result.avg_runtime_ms / 1000.0);
std::string raster = "Heuristic";
if (options.raster == RasterOrderOptions::AlongN) {
raster = "Along N";
}
else if (options.raster == RasterOrderOptions::AlongM) {
raster = "Along M";
}
std::cout << " Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << 'x' << options.l << std::endl;
std::cout << " Rasterization: " << raster << " with a maximum CTA swizzle of " << options.swizzle << std::endl;
std::cout << " Avg runtime: " << result.avg_runtime_ms << " ms" << std::endl;
std::cout << " GFLOPS: " << result.gflops << std::endl;
}
@@ -551,12 +572,11 @@ int main(int argc, char const **args) {
<< "later (compute capability 90 or greater).\n";
return 0;
}
//
// Parse options
//
Options options;
Options<RasterOrderOptions> options;
options.parse(argc, args);
examples/54_hopper_fp8_warp_specialized_gemm/hopper_fp8_commandline.hpp
+20
View File
@@ -30,6 +30,7 @@
**************************************************************************************************/
// Command line options parsing
template<typename RasterOrderOptions>
struct Options {
bool help = false;
@@ -41,6 +42,8 @@ struct Options {
bool save_amax = true;
int iterations = 1000;
int m = 1024, n = 512, k = 1024, l = 1;
RasterOrderOptions raster;
int swizzle;
// Parses the command line
void parse(int argc, char const **args) {
@@ -66,6 +69,21 @@ struct Options {
cmd.get_cmd_line_argument("save_aux", save_aux, true);
cmd.get_cmd_line_argument("save_amax", save_amax, true);
cmd.get_cmd_line_argument("iterations", iterations);
char raster_char;
cmd.get_cmd_line_argument("raster", raster_char);
if (raster_char == 'N' || raster_char == 'n') {
raster = RasterOrderOptions::AlongN;
}
else if (raster_char == 'M' || raster_char == 'm') {
raster = RasterOrderOptions::AlongM;
}
else if (raster_char == 'H' || raster_char == 'h') {
raster = RasterOrderOptions::Heuristic;
}
cmd.get_cmd_line_argument("swizzle", swizzle, 1);
}
/// Prints the usage statement.
@@ -89,6 +107,8 @@ struct Options {
<< " --device_scale=<bool> Copy scalars to device memory before kernel launch (default: false)\n"
<< " --save_aux=<bool> Save the pre-activation as an auxiliary tensor (default: true)\n"
<< " --save_amax=<bool> Save the pre-scaled max absolute value of any fp8 outputs (aux and/or D) (default: true)\n"
<< " --raster=<char> CTA Rasterization direction (N for along N, M for along M, and H for heuristic)\n\n"
<< " --swizzle=<int> CTA Rasterization swizzle\n\n"
<< " --iterations=<int> Number of profiling iterations to perform.\n\n";
out
examples/55_hopper_mixed_dtype_gemm/55_hopper_mixed_dtype_gemm.cu
-1
View File
@@ -687,7 +687,6 @@ int main(int argc, char const **args) {
<< "later (compute capability 90 or greater).\n";
return 0;
}
//
// Parse options
//
examples/56_hopper_ptr_array_batched_gemm/56_hopper_ptr_array_batched_gemm.cu
+1 -2
View File
@@ -99,7 +99,7 @@ using TileShape = Shape<_256,_128,_64>; // T
using ClusterShape = Shape<_1,_2,_1>; // Shape of the threadblocks in a cluster
using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based on the tile size
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative; // Kernel to launch
using EpilogueSchedule = cutlass::epilogue::PtrArrayNoSmemWarpSpecialized; // Epilogue to launch
using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative; // Epilogue to launch
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
@@ -492,7 +492,6 @@ int main(int argc, char const **args) {
<< "later (compute capability 90 or greater).\n";
return 0;
}
//
// Parse options
//
examples/56_hopper_ptr_array_batched_gemm/CMakeLists.txt
+2 -2
View File
@@ -30,10 +30,10 @@
set(TEST_SQUARE --m=2048 --n=2048 --k=2048 -l=10 --iterations=1) # Square problem sizes
set(TEST_SQUARE_LARGE_BATCH --m=2048 --n=2048 --k=2048 -l=500 --iterations=1) # Square problem sizes
set(TEST_EPILOGUE --alpha=0.5 --beta=0.7 --iterations=1) # Default problem sizes
set(TEST_EPILOGUE --alpha=0.5 --beta=0.5 --iterations=1) # Default problem sizes
set(TEST_EPILOGUE_LARGE_BATCH --alpha=1.5 --beta=2.0 -l=500 --iterations=1) # Default problem sizes
set(TEST_EPILOGUE_OP --beta=0.7 --iterations=1) # Default problem sizes w/ Epilogue Op test
set(TEST_EPILOGUE_OP --beta=0.5 --iterations=1) # Default problem sizes w/ Epilogue Op test
set(TEST_EPILOGUE_OP_LARGE_BATCH --alpha=1.5 -l=500 --iterations=1) # Default problem sizes w/ Epilogue Op test
set(TEST_SMALLK --m=2048 --n=5120 --k=128 --l=5 --iterations=1) # Small-k problem sizes
examples/57_hopper_grouped_gemm/57_hopper_grouped_gemm.cu
+13 -14
View File
@@ -32,7 +32,7 @@
/*! \file
\brief Hopper Grouped GEMM example using CUTLASS 3 APIs for NVIDIA Hopper architecture.
This example demonstrates an implementation of Grouped GEMM using a TMA + GMMA
This example demonstrates an implementation of Grouped GEMM using a TMA + GMMA
warp-specialized cooperative kernel.
For this example all scheduling work is performed on the device.
The new feature showcased in this example is on-the-fly modification of TMA descriptors
@@ -42,7 +42,7 @@
$ ./examples/57_hopper_grouped_gemm/57_hopper_grouped_gemm --m=2048 --n=2048 --k=2048 --groups=10
The above example command makes all 10 groups to be sized at the given m, n, k sizes.
The above example command makes all 10 groups to be sized at the given m, n, k sizes.
Skipping any of the problem dimensions randomizes it across the different groups.
Same applies for alpha and beta values that are randomized across the different groups.
@@ -117,7 +117,7 @@ constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value; // A
using ElementAccumulator = float; // Element type for internal accumulation
using ArchTag = cutlass::arch::Sm90; // Tag indicating the minimum SM that supports the intended feature
using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag
using TileShape = Shape<_256,_128,_64>; // Threadblock-level tile size
using TileShape = Shape<_256,_128,_128>; // Threadblock-level tile size
using ClusterShape = Shape<_2,_2,_1>; // Shape of the threadblocks in a cluster
using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based on the tile size
using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum; // Kernel to launch
@@ -163,10 +163,10 @@ using DeviceGemmReference = cutlass::reference::device::Gemm<
ElementAccumulator,
ElementAccumulator>;
using StrideA = typename Gemm::GemmKernel::UnderlyingStrideA;
using StrideB = typename Gemm::GemmKernel::UnderlyingStrideB;
using StrideC = typename Gemm::GemmKernel::UnderlyingStrideC;
using StrideD = typename Gemm::GemmKernel::UnderlyingStrideD;
using StrideA = typename Gemm::GemmKernel::InternalStrideA;
using StrideB = typename Gemm::GemmKernel::InternalStrideB;
using StrideC = typename Gemm::GemmKernel::InternalStrideC;
using StrideD = typename Gemm::GemmKernel::InternalStrideD;
// Host-side allocations
std::vector<int64_t> offset_A;
@@ -226,7 +226,7 @@ struct Options {
std::string benchmark_path;
std::vector<typename ProblemShape::UnderlyingProblemShape> problem_sizes_host;
int const tma_alignment_bits = 128;
int const alignment = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
int const alignment = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
// Parses the command line
void parse(int argc, char const **args) {
@@ -438,10 +438,10 @@ void allocate(const Options &options) {
total_elements_C += elements_C;
total_elements_D += elements_D;
stride_A_host.push_back(cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, Int<1>{})));
stride_B_host.push_back(cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, Int<1>{})));
stride_C_host.push_back(cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, Int<1>{})));
stride_D_host.push_back(cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, Int<1>{})));
stride_A_host.push_back(cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1}));
stride_B_host.push_back(cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}));
stride_C_host.push_back(cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1}));
stride_D_host.push_back(cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1}));
}
@@ -456,7 +456,7 @@ void allocate(const Options &options) {
/// Initialize operands to be used in the GEMM and reference GEMM
void initialize(const Options &options) {
uint64_t seed = 2020;
problem_sizes.reset(options.groups);
@@ -695,7 +695,6 @@ int main(int argc, char const **args) {
<< "later (compute capability 90 or greater).\n";
return 0;
}
//
// Parse options
//
examples/58_ada_fp8_gemm/ada_fp8_gemm.cu
+2 -2
View File
@@ -97,7 +97,7 @@ using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWith
cutlass::epilogue::thread::ReLu,
ElementOutput,
ElementAuxOutput,
128 / cutlass::sizeof_bits<ElementOutput>::value,
8,
ElementAccumulator,
ElementAccumulator
>;
@@ -106,7 +106,7 @@ template <typename MathOperator>
using Gemm_ = cutlass::gemm::device::GemmUniversalWithAbsMax<
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
cutlass::gemm::GemmShape<128, 64, 128>, cutlass::gemm::GemmShape<64, 32, 128>, cutlass::gemm::GemmShape<16, 8, 32>,
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages,
kAlignmentA, kAlignmentB, MathOperator
>;
examples/cute/tutorial/CMakeLists.txt
+5
View File
@@ -53,3 +53,8 @@ cutlass_example_add_executable(
tiled_copy.cu
)
cutlass_example_add_executable(
wgmma_sm90
wgmma_sm90.cu
)
examples/cute/tutorial/sgemm_sm80.cu
+11 -11
View File
@@ -153,12 +153,12 @@ gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler,
// Allocate the accumulators -- same size as the projected data
Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N)
CUTE_STATIC_ASSERT_V( shape(tCrA) == shape(tCsA)); // (MMA,MMA_M,MMA_K)
CUTE_STATIC_ASSERT_V( shape(tCrB) == shape(tCsB)); // (MMA,MMA_N,MMA_K)
CUTE_STATIC_ASSERT_V( shape(tCrC) == shape(tCgC)); // (MMA,MMA_M,MMA_N)
CUTE_STATIC_ASSERT_V(size<1>(tCgC) == size<1>(tCsA)); // MMA_M
CUTE_STATIC_ASSERT_V(size<2>(tCgC) == size<1>(tCsB)); // MMA_N
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // MMA_K
CUTE_STATIC_ASSERT_V(( shape(tCrA) == take<0,3>(shape(tCsA)))); // (MMA,MMA_M,MMA_K)
CUTE_STATIC_ASSERT_V(( shape(tCrB) == take<0,3>(shape(tCsB)))); // (MMA,MMA_N,MMA_K)
CUTE_STATIC_ASSERT_V(( shape(tCrC) == take<0,3>(shape(tCgC)))); // (MMA,MMA_M,MMA_N)
CUTE_STATIC_ASSERT_V((size<1>(tCgC) == size<1>(tCsA))); // MMA_M
CUTE_STATIC_ASSERT_V((size<2>(tCgC) == size<1>(tCsB))); // MMA_N
CUTE_STATIC_ASSERT_V((size<2>(tCsA) == size<2>(tCsB))); // MMA_K
// Clear the accumulators
clear(tCrC);
@@ -358,7 +358,7 @@ gemm_nt(int m, int n, int k,
alpha, beta);
}
// Setup params for a NT GEMM
// Setup params for a TN GEMM
template <class TA, class TB, class TC,
class Alpha, class Beta>
void
@@ -391,10 +391,10 @@ gemm_tn(int m, int n, int k,
auto bP = Int<3>{}; // Pipeline
// Define the smem layouts (static)
auto sA_atom = make_layout(make_shape ( bM, bK),
make_stride(Int<1>{}, bM+Int<1>{})); // (m,k) -> smem_idx; padded m-major
auto sB_atom = make_layout(make_shape ( bN, bK),
make_stride(Int<1>{}, bN+Int<1>{})); // (n,k) -> smem_idx; padded n-major
auto sA_atom = make_layout(make_shape ( bM, bK),
make_stride(Int<1>{}, bM+Int<1>{})); // (m,k) -> smem_idx; padded m-major
[[maybe_unused]] auto sB_atom = make_layout(make_shape ( bN, bK),
make_stride(Int<1>{}, bN+Int<1>{})); // (n,k) -> smem_idx; padded n-major
auto sA = tile_to_shape(sA_atom, make_shape(bM, bK, bP));
auto sB = tile_to_shape(sA_atom, make_shape(bN, bK, bP));
auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx
examples/cute/tutorial/wgmma_sm90.cu
+562
View File
@@ -0,0 +1,562 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#include <cstdlib>
#include <cstdio>
#include <cassert>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include "cutlass/cluster_launch.hpp"
#include "cutlass/arch/barrier.h"
#include "cutlass/pipeline/sm90_pipeline.hpp"
#include "cutlass/util/print_error.hpp"
#include "cutlass/util/GPU_Clock.hpp"
#include "cutlass/util/helper_cuda.hpp"
#include "cutlass/arch/mma_sm90.h"
#include "cutlass/device_kernel.h"
using namespace cute;
template <class ElementA,
class ElementB,
class SmemLayoutA, // (M,K,P)
class SmemLayoutB> // (N,K,P)
struct SharedStorage
{
array_aligned<ElementA, cosize_v<SmemLayoutA>> smem_A;
array_aligned<ElementB, cosize_v<SmemLayoutB>> smem_B;
uint64_t tma_barrier[size<2>(SmemLayoutA{})];
uint64_t mma_barrier[size<2>(SmemLayoutA{})];
};
template <class ProblemShape, class CtaTiler,
class TA, class SmemLayoutA, class TmaA,
class TB, class SmemLayoutB, class TmaB,
class TC, class CStride, class TiledMma,
class Alpha, class Beta>
__global__ static
__launch_bounds__(decltype(size(TiledMma{}))::value)
void
gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler,
TA const* A, CUTLASS_GRID_CONSTANT TmaA const tma_a,
TB const* B, CUTLASS_GRID_CONSTANT TmaB const tma_b,
TC * C, CStride dC, TiledMma mma,
Alpha alpha, Beta beta)
{
// Preconditions
CUTE_STATIC_ASSERT_V(rank(shape_MNK) == Int<3>{}); // (M, N, K)
CUTE_STATIC_ASSERT_V(rank(cta_tiler) == Int<3>{}); // (BLK_M, BLK_N, BLK_K)
static_assert(is_static<SmemLayoutA>::value);
static_assert(is_static<SmemLayoutB>::value);
CUTE_STATIC_ASSERT_V(size<0>(SmemLayoutA{}) == size<0>(cta_tiler)); // BLK_M
CUTE_STATIC_ASSERT_V(size<0>(SmemLayoutB{}) == size<1>(cta_tiler)); // BLK_N
CUTE_STATIC_ASSERT_V(size<1>(SmemLayoutA{}) == size<2>(cta_tiler)); // BLK_K
CUTE_STATIC_ASSERT_V(size<1>(SmemLayoutB{}) == size<2>(cta_tiler)); // BLK_K
CUTE_STATIC_ASSERT_V(congruent(select<0,1>(shape_MNK), dC)); // dC strides for shape MN
//
// Full and Tiled Tensors
//
// Represent the full tensors
auto [M, N, K] = shape_MNK;
Tensor mA = tma_a.get_tma_tensor(make_shape(M,K)); // (M,K) TMA Tensor
Tensor mB = tma_b.get_tma_tensor(make_shape(N,K)); // (N,K) TMA Tensor
Tensor mC = make_tensor(make_gmem_ptr(C), make_shape(M,N), dC); // (M,N)
// Get the appropriate blocks for this thread block
auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k)
Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k)
Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N)
// Shared memory tensors
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<TA, TB, SmemLayoutA, SmemLayoutB>;
SharedStorage& smem = *reinterpret_cast<SharedStorage*>(shared_memory);
Tensor sA = make_tensor(make_smem_ptr(smem.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(smem.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
//
// Partition the copying of A and B tiles
//
// TUTORIAL:
// These are TMA partitionings, which have a dedicated custom partitioner.
// The Int<0>, Layout<_1> indicates that the TMAs are not multicasted.
// Any multicasting must be in conformance with tma_x constructed with make_tma_atom on host.
// The group_modes<0,2> transforms the (X,Y,Z)-shaped tensors into ((X,Y),Z)-shaped tensors
// with the understanding that the TMA is responsible for everything in mode-0.
// The tma_partition reorders and offsets mode-0 according to the tma_x atom and the multicast info.
//
auto [tAgA, tAsA] = tma_partition(tma_a, Int<0>{}, Layout<_1>{},
group_modes<0,2>(sA), group_modes<0,2>(gA)); // (TMA,k) and (TMA,PIPE)
auto [tBgB, tBsB] = tma_partition(tma_b, Int<0>{}, Layout<_1>{},
group_modes<0,2>(sB), group_modes<0,2>(gB)); // (TMA,k) and (TMA,PIPE)
// The TMA is responsible for copying everything in mode-0 of tAsA and tBsB
constexpr int kTmaTransactionBytes = CUTE_STATIC_V(size<0>(tAsA)) * sizeof(TA) +
CUTE_STATIC_V(size<0>(tBsB)) * sizeof(TB);
//
// PREFETCH
//
auto K_PIPE_MAX = size<1>(tAsA);
// Total count of tiles
int k_tile_count = size<1>(tAgA);
// Current tile index in gmem to read from
int k_tile = 0;
// Initialize Barriers
int warp_idx = cutlass::canonical_warp_idx_sync();
int lane_predicate = cute::elect_one_sync();
uint64_t* producer_mbar = smem.tma_barrier;
uint64_t* consumer_mbar = smem.mma_barrier;
using ProducerBarType = cutlass::arch::ClusterTransactionBarrier; // TMA
using ConsumerBarType = cutlass::arch::ClusterBarrier; // MMA
CUTE_UNROLL
for (int pipe = 0; pipe < K_PIPE_MAX; ++pipe) {
if ((warp_idx == 0) && lane_predicate) {
ProducerBarType::init(&producer_mbar[pipe], 1);
ConsumerBarType::init(&consumer_mbar[pipe], 128);
}
}
// Ensure barrier init is complete on all CTAs
cluster_sync();
// Start async loads for all pipes
CUTE_UNROLL
for (int pipe = 0; pipe < K_PIPE_MAX; ++pipe)
{
if ((warp_idx == 0) && lane_predicate)
{
// Set expected Tx Bytes after each reset / init
ProducerBarType::arrive_and_expect_tx(&producer_mbar[pipe], kTmaTransactionBytes);
copy(tma_a.with(producer_mbar[pipe]), tAgA(_,k_tile), tAsA(_,pipe));
copy(tma_b.with(producer_mbar[pipe]), tBgB(_,k_tile), tBsB(_,pipe));
}
--k_tile_count;
++k_tile;
}
//
// Define A/B partitioning and C accumulators
//
// TUTORIAL:
// The tCrA and tCrB are actually Tensors of MMA Descriptors constructed as views of SMEM.
// The MMA Descriptor generation is automatic via inspection and validation of the SMEM Layouts.
// Because the MMA reads directly from SMEM and the fragments are descriptors rather than registers,
// there is no need for copy(tCsA, tCrA) in the mainloop.
//
ThrMMA thr_mma = mma.get_thread_slice(threadIdx.x);
Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N)
// Allocate accumulators and clear them
Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N)
clear(tCrC);
// Allocate "fragments"
Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE)
//
// PIPELINED MAIN LOOP
//
// TUTORIAL:
// Rather than interleaving the stages and instructions like in SM70 and SM80,
// the SM90 mainloops rely on explicit producer-consumer synchronization
// on the purely async instructions TMA and MMA.
// More advanced pipeline and warp-specialization strategies are available in CUTLASS mainloops.
//
// A PipelineState is a circular pipe index [.index()] and a pipe phase [.phase()]
// that flips each cycle through K_PIPE_MAX.
auto write_state = cutlass::PipelineState<K_PIPE_MAX>(); // TMA writes
auto read_state = cutlass::PipelineState<K_PIPE_MAX>(); // MMA reads
CUTE_NO_UNROLL
while (k_tile_count > -K_PIPE_MAX)
{
// Wait for Producer to complete
int read_pipe = read_state.index();
ProducerBarType::wait(&producer_mbar[read_pipe], read_state.phase());
// MMAs to cover 1 K_TILE
warpgroup_arrive();
gemm(mma, tCrA(_,_,_,read_pipe), tCrB(_,_,_,read_pipe), tCrC); // (V,M) x (V,N) => (V,M,N)
warpgroup_commit_batch();
// Wait for all MMAs in a K_TILE to complete
warpgroup_wait<0>();
// Notify that consumption is done
ConsumerBarType::arrive(&consumer_mbar[read_pipe]);
++read_state;
if ((warp_idx == 0) && lane_predicate)
{
int pipe = write_state.index();
// Wait for Consumer to complete consumption
ConsumerBarType::wait(&consumer_mbar[pipe], write_state.phase());
// Set expected Tx Bytes after each reset / init
ProducerBarType::arrive_and_expect_tx(&producer_mbar[pipe], kTmaTransactionBytes);
copy(tma_a.with(producer_mbar[pipe]), tAgA(_,k_tile), tAsA(_,pipe));
copy(tma_b.with(producer_mbar[pipe]), tBgB(_,k_tile), tBsB(_,pipe));
++write_state;
}
--k_tile_count;
++k_tile;
}
//
// Epilogue (unpredicated)
//
axpby(alpha, tCrC, beta, tCgC);
}
// Setup params for an NT GEMM
template <class TA, class TB, class TC,
class Alpha, class Beta>
void
gemm_nt(int m, int n, int k,
Alpha alpha,
TA const* A, int ldA,
TB const* B, int ldB,
Beta beta,
TC * C, int ldC,
cudaStream_t stream = 0)
{
// Define shapes (dynamic)
auto M = int(m);
auto N = int(n);
auto K = int(k);
auto prob_shape = make_shape(M, N, K); // (M, N, K)
// Define TN strides (mixed)
auto dA = make_stride(Int<1>{}, ldA); // (dM, dK)
auto dB = make_stride(Int<1>{}, ldB); // (dN, dK)
auto dC = make_stride(Int<1>{}, ldC); // (dM, dN)
// Define CTA tile sizes (static)
auto bM = Int<128>{};
auto bN = Int<128>{};
auto bK = Int< 64>{};
auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K)
auto bP = Int< 3>{}; // Pipeline
// Define the smem layouts (static)
auto sA = tile_to_shape(GMMA::Layout_MN_SW128_Atom<TA>{}, make_shape(bM,bK,bP));
auto sB = tile_to_shape(GMMA::Layout_MN_SW128_Atom<TB>{}, make_shape(bN,bK,bP));
// Define the MMA
TiledMMA tiled_mma = make_tiled_mma(SM90_64x64x16_F16F16F16_SS<GMMA::Major::MN,GMMA::Major::MN>{});
// Define the TMAs
// Create Global memory tensors for TMA inspection
Tensor mA = make_tensor(A, make_shape(M,K), dA);
Tensor mB = make_tensor(B, make_shape(N,K), dB);
// Create TMA Atoms with the desired copy operation on the source and destination
Copy_Atom tmaA = make_tma_atom(SM90_TMA_LOAD{}, mA, sA(_,_,0), make_shape(bM,bK));
Copy_Atom tmaB = make_tma_atom(SM90_TMA_LOAD{}, mB, sB(_,_,0), make_shape(bN,bK));
//
// Setup and Launch
//
// Launch parameter setup
int smem_size = int(sizeof(SharedStorage<TA, TB, decltype(sA), decltype(sB)>));
dim3 dimBlock(size(tiled_mma));
dim3 dimCluster(2, 1, 1);
dim3 dimGrid(round_up(size(ceil_div(m, bM)), dimCluster.x),
round_up(size(ceil_div(n, bN)), dimCluster.y));
cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size};
void const* kernel_ptr = reinterpret_cast<void const*>(
&gemm_device<decltype(prob_shape), decltype(cta_tiler),
TA, decltype(sA), decltype(tmaA),
TB, decltype(sB), decltype(tmaB),
TC, decltype(dC), decltype(tiled_mma),
decltype(alpha), decltype(beta)>);
CUTE_CHECK_ERROR(cudaFuncSetAttribute(
kernel_ptr,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size));
// Kernel Launch
cutlass::Status status = cutlass::launch_kernel_on_cluster(params, kernel_ptr,
prob_shape, cta_tiler,
A, tmaA,
B, tmaB,
C, dC, tiled_mma,
alpha, beta);
CUTE_CHECK_LAST();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Error: Failed at kernel Launch" << std::endl;
}
}
// Setup params for a TN GEMM
template <class TA, class TB, class TC,
class Alpha, class Beta>
void
gemm_tn(int m, int n, int k,
Alpha alpha,
TA const* A, int ldA,
TB const* B, int ldB,
Beta beta,
TC * C, int ldC,
cudaStream_t stream = 0)
{
// Define shapes (dynamic)
auto M = int(m);
auto N = int(n);
auto K = int(k);
auto prob_shape = make_shape(M, N, K); // (M, N, K)
// Define TN strides (mixed)
auto dA = make_stride(ldA, Int<1>{}); // (dM, dK)
auto dB = make_stride(ldB, Int<1>{}); // (dN, dK)
auto dC = make_stride(Int<1>{}, ldC); // (dM, dN)
// Define CTA tile sizes (static)
auto bM = Int<128>{};
auto bN = Int<128>{};
auto bK = Int< 64>{};
auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K)
auto bP = Int<3>{}; // Pipeline
// Define the smem layouts (static)
auto sA = tile_to_shape(GMMA::Layout_K_SW128_Atom<TA>{}, make_shape(bM,bK,bP));
auto sB = tile_to_shape(GMMA::Layout_K_SW128_Atom<TB>{}, make_shape(bN,bK,bP));
// Define the MMA
TiledMMA tiled_mma = make_tiled_mma(SM90_64x64x16_F16F16F16_SS<GMMA::Major::K,GMMA::Major::K>{});
// Define the TMAs
// Create Global memory tensors for TMA inspection
Tensor mA = make_tensor(A, make_shape(M,K), dA);
Tensor mB = make_tensor(B, make_shape(N,K), dB);
// Create TMA Atoms with the desired copy operation on the source and destination
Copy_Atom tmaA = make_tma_atom(SM90_TMA_LOAD{}, mA, sA(_,_,0), make_shape(bM,bK));
Copy_Atom tmaB = make_tma_atom(SM90_TMA_LOAD{}, mB, sB(_,_,0), make_shape(bN,bK));
//
// Setup and Launch
//
// Launch parameter setup
int smem_size = int(sizeof(SharedStorage<TA, TB, decltype(sA), decltype(sB)>));
dim3 dimBlock(size(tiled_mma));
dim3 dimCluster(2, 1, 1);
dim3 dimGrid(round_up(size(ceil_div(m, bM)), dimCluster.x),
round_up(size(ceil_div(n, bN)), dimCluster.y));
cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size};
void const* kernel_ptr = reinterpret_cast<void const*>(
&gemm_device<decltype(prob_shape), decltype(cta_tiler),
TA, decltype(sA), decltype(tmaA),
TB, decltype(sB), decltype(tmaB),
TC, decltype(dC), decltype(tiled_mma),
decltype(alpha), decltype(beta)>);
CUTE_CHECK_ERROR(cudaFuncSetAttribute(
kernel_ptr,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size));
// Kernel Launch
cutlass::Status status = cutlass::launch_kernel_on_cluster(params, kernel_ptr,
prob_shape, cta_tiler,
A, tmaA,
B, tmaB,
C, dC, tiled_mma,
alpha, beta);
CUTE_CHECK_LAST();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Error: Failed at kernel Launch" << std::endl;
}
}
template <class TA, class TB, class TC,
class Alpha, class Beta>
void
gemm(char transA, char transB, int m, int n, int k,
Alpha alpha,
TA const* A, int ldA,
TB const* B, int ldB,
Beta beta,
TC * C, int ldC,
cudaStream_t stream = 0)
{
if (transA == 'N' && transB == 'T') {
return gemm_nt(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream);
} else
if (transA == 'T' && transB == 'N') {
return gemm_tn(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream);
}
assert(false && "Not implemented");
}
int main(int argc, char** argv)
{
cudaDeviceProp props;
cudaError_t error = cudaGetDeviceProperties(&props, 0);
if (error != cudaSuccess) {
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
return -1;
}
if (props.major != 9) {
std::cout << "This example requires NVIDIA's Hopper Architecture GPU with compute capability 90a\n" << std::endl;
return 0;
}
#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
int m = 512;
if (argc >= 2)
sscanf(argv[1], "%d", &m);
int n = 256;
if (argc >= 3)
sscanf(argv[2], "%d", &n);
int k = 1024;
if (argc >= 4)
sscanf(argv[3], "%d", &k);
char transA = 'N';
if (argc >= 5)
sscanf(argv[4], "%c", &transA);
char transB = 'T';
if (argc >= 6)
sscanf(argv[5], "%c", &transB);
using TA = cute::half_t;
using TB = cute::half_t;
using TC = cute::half_t;
using TI = cute::half_t;
TI alpha = TI(1.0f);
TI beta = TI(0.0f);
thrust::host_vector<TA> h_A(m*k);
thrust::host_vector<TB> h_B(n*k);
thrust::host_vector<TC> h_C(m*n);
// Initialize the tensors
for (int j = 0; j < m*k; ++j) h_A[j] = TA(int((rand() % 2) ? 1 : -1));
for (int j = 0; j < n*k; ++j) h_B[j] = TB(int((rand() % 2) ? 1 : -1));
for (int j = 0; j < m*n; ++j) h_C[j] = TC(0);
thrust::device_vector<TA> d_A = h_A;
thrust::device_vector<TB> d_B = h_B;
thrust::device_vector<TC> d_C = h_C;
double gflops = (2.0*m*n*k) * 1e-9;
const int timing_iterations = 100;
GPU_Clock timer;
int ldA = 0, ldB = 0, ldC = m;
if (transA == 'N') {
ldA = m;
} else if (transA == 'T') {
ldA = k;
} else {
assert(false);
}
if (transB == 'N') {
ldB = k;
} else if (transB == 'T') {
ldB = n;
} else {
assert(false);
}
// Run once
d_C = h_C;
gemm(transA, transB, m, n, k,
alpha,
d_A.data().get(), ldA,
d_B.data().get(), ldB,
beta,
d_C.data().get(), ldC);
CUTE_CHECK_LAST();
thrust::host_vector<TC> cute_result = d_C;
// Timing iterations
timer.start();
for (int i = 0; i < timing_iterations; ++i) {
gemm(transA, transB, m, n, k,
alpha,
d_A.data().get(), ldA,
d_B.data().get(), ldB,
beta,
d_C.data().get(), ldC);
}
double cute_time = timer.seconds() / timing_iterations;
CUTE_CHECK_LAST();
printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000);
#else
std::cout << "CUTLASS_ARCH_MMA_SM90_SUPPORTED must be enabled, but it is not. Test is waived \n" << std::endl;
#endif
return 0;
}
include/cute/algorithm/axpby.hpp
+1 -1
View File
@@ -32,7 +32,7 @@
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/tensor_predicate.hpp>
namespace cute
include/cute/algorithm/clear.hpp
+1 -3
View File
@@ -31,9 +31,7 @@
#pragma once
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/algorithm/fill.hpp>
namespace cute
include/cute/algorithm/cooperative_copy.hpp
+236 -119
View File
@@ -1,51 +1,117 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
* Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/atom/copy_atom.hpp>
#include <cute/algorithm/copy.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/tensor_predicate.hpp>
namespace cute
{
template <uint32_t NumThreads,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE void
naive_cooperative_copy(uint32_t const& tid,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
auto N = size(src);
if (tid < N) {
uint32_t upper_bound = (N / NumThreads) * NumThreads;
CUTE_UNROLL
for (uint32_t i = 0; i < upper_bound; i += NumThreads) { // All in-bounds
dst[tid + i] = src[tid + i];
}
if (N % NumThreads != 0) { // Likely static condition
uint32_t final_idx = tid + upper_bound;
if (final_idx < N) { // Final in-bounds
dst[final_idx] = src[final_idx];
}
}
}
}
// Accept mutable temporaries
template <uint32_t NumThreads,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE void
naive_cooperative_copy(uint32_t const& tid,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> && dst)
{
return naive_cooperative_copy(tid, src, dst);
}
// A heuristic to determine a "good" permutation of two tensors for later vectorization and thr-assignment
template <class AEngine, class ALayout,
class BEngine, class BLayout>
CUTE_HOST_DEVICE constexpr
auto
heuristic_permutation(Tensor<AEngine, ALayout> const& a,
Tensor<BEngine, BLayout> const& b)
{
constexpr bool swizzleA = get_swizzle_t<AEngine>::num_bits != 0 or
get_swizzle_t<ALayout>::num_bits != 0;
constexpr bool swizzleB = get_swizzle_t<BEngine>::num_bits != 0 or
get_swizzle_t<BLayout>::num_bits != 0;
auto a_inv = right_inverse(get_nonswizzle_portion(a.layout()));
auto b_inv = right_inverse(get_nonswizzle_portion(b.layout()));
constexpr uint8_t scoreA = (uint8_t(swizzleA) << 2) |
(uint8_t(is_smem<AEngine>::value) << 1) |
(uint8_t(size(a_inv) > size(b_inv)) << 0);
constexpr uint8_t scoreB = (uint8_t(swizzleB) << 2) |
(uint8_t(is_smem<BEngine>::value) << 1) |
(uint8_t(size(b_inv) > size(a_inv)) << 0);
if constexpr (scoreA >= scoreB) {
return a_inv;
} else {
return b_inv;
}
}
// cooperative_copy<NumThreads, MaxVecBits>(thr_idx, src, dst)
// Use NumThreads to copy src to dst with element vectorization up to MaxVecBits.
// Use NumThreads to copy Tensor src to Tensor dst with element-wise vectorization up to MaxVecBits.
// @pre 0 <= @a tid < NumThreads
// @pre Tensors @a src and @a dst are aligned up to MaxVecBits.
// That is, pointers and dynamic strides are assumed to be aligned up to MaxVecBits.
//
template <uint32_t NumThreads, uint32_t MaxVecBits,
class SrcEngine, class SrcLayout,
@@ -56,121 +122,171 @@ cooperative_copy(uint32_t const& tid,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
// Assumes the shapes are static, can generalize
// Assumes the shapes are static, can generalize/fallback
CUTE_STATIC_ASSERT_V(is_static<decltype(shape(src))>{} && is_static<decltype(shape(dst))>{});
CUTE_STATIC_ASSERT_V(size(src) == size(dst));
// Assumes the types are the same, can generalize
static_assert(sizeof_bits_v<typename SrcEngine::value_type> == sizeof_bits_v<typename DstEngine::value_type>);
// Assumes the types are the same, can generalize/fallback
static_assert(cute::is_same<typename SrcEngine::value_type, typename DstEngine::value_type>::value);
static_assert(MaxVecBits == sizeof_bits_v<typename SrcEngine::value_type> ||
MaxVecBits == 8 || MaxVecBits == 16 || MaxVecBits == 32 || MaxVecBits == 64 || MaxVecBits == 128,
"Expected MaxVecBits to be value size or 8 or 16 or 32 or 64 or 128 for alignment and performance.");
// Check that the tensors are likely shared across threads: either gmem or smem
static_assert((is_gmem<SrcEngine>::value || is_smem<SrcEngine>::value),
"cooperative_copy expects shared gmem or smem source tensor.");
"cooperative_copy expects shared gmem or smem source tensor.");
static_assert((is_gmem<DstEngine>::value || is_smem<DstEngine>::value),
"cooperative_copy expects shared gmem or smem destination tensor.");
"cooperative_copy expects shared gmem or smem destination tensor.");
// Precondition on tid in DEBUG
assert(tid < NumThreads);
// Precondition on pointer alignment in DEBUG
assert(is_byte_aligned<ceil_div(MaxVecBits,8u)>(raw_pointer_cast(src.data())));
assert(is_byte_aligned<ceil_div(MaxVecBits,8u)>(raw_pointer_cast(dst.data())));
// Fallback - slow path, naive copy, vectorization disabled
if constexpr(size(SrcLayout{}) % NumThreads != 0) {
int index = static_cast<int>(tid);
CUTE_UNROLL
for(int i = 0; i < ceil_div(size(SrcLayout{}), NumThreads); i++) {
if(index < size(SrcLayout{})) {
dst[index] = src[index];
#if 0
if (thread0()) {
print(" "); print("cooperative_copy\n");
print(" "); print("NumThreads: "); print(NumThreads); print("\n");
print(" "); print("MaxVecBits: "); print(MaxVecBits); print("\n");
print(" "); print("src: "); print(src); print("\n");
print(" "); print("dst: "); print(dst); print("\n");
}
index += NumThreads;
#ifdef __CUDA_ARCH__
__syncthreads();
#endif
#endif
// The common layout of the two tensors that can be vectorized over elements and threads
// vidx -> coord
auto common_layout = heuristic_permutation(src, dst);
// Apply
// (V, rest)
Tensor src_a = coalesce(logical_divide(src, common_layout), Shape<_1,_1>{});
Tensor dst_a = coalesce(logical_divide(dst, common_layout), Shape<_1,_1>{});
//
// Determine vectorization of elems and thrs based on src/dst size and number of threads
// NOTE: This heuristic promotes parallelization over vectorization
//
// The number of elements and number of bits
constexpr int elem_bits = sizeof_bits_v<typename SrcEngine::value_type>;
constexpr int total_elem = size(SrcLayout{});
// The number of elements that can be vectorized in values
constexpr int common_elem = decltype(max_common_vector(src_a, dst_a))::value;
#if 0
if (thread0()) {
print(" "); print("common_layout: "); print(common_layout); print("\n");
print(" "); print("src_a: "); print(src_a); print("\n");
print(" "); print("dst_a: "); print(dst_a); print("\n");
}
#ifdef __CUDA_ARCH__
__syncthreads();
#endif
#endif
//
if constexpr (total_elem % NumThreads != 0) {
// Not attempting to find a partitioning pattern, fallback to dynamically indexed slowpath
if constexpr (common_elem > 1 && MaxVecBits > elem_bits) {
// If the vectorization is non-trivial and divides the maximum vectorizations, then vectorize
constexpr auto max_align_src = elem_bits * decltype(max_alignment(src_a.layout()))::value;
constexpr auto max_align_dst = elem_bits * decltype(max_alignment(dst_a.layout()))::value;
constexpr auto vec_bits = gcd(max_align_src, max_align_dst, MaxVecBits);
using VecType = uint_bit_t<vec_bits>;
static_assert(vec_bits % elem_bits == 0, "Expected divisibility");
static_assert((vec_bits >= 8), "No support for subbyte copying");
Tensor src_v = recast<VecType const>(src_a);
Tensor dst_v = recast<VecType >(dst_a);
#if 0
if (thread0()) {
print(" "); print("cooperative_copy -- naive\n");
print(" "); print("src_v: "); print(src_v); print("\n");
print(" "); print("dst_v: "); print(dst_v); print("\n");
}
#ifdef __CUDA_ARCH__
__syncthreads();
#endif
#endif
naive_cooperative_copy<NumThreads>(tid, src_v, dst_v);
} else {
naive_cooperative_copy<NumThreads>(tid, src_a, dst_a);
}
} else {
// Fast path with vectorization
// If the tensors can be equally partitioned by the threads,
// compute vectorization widths in elements and threads.
// Precondition on pointer alignment in DEBUG
assert(is_byte_aligned<max(MaxVecBits/8, 1u)>(raw_pointer_cast(src.data())));
assert(is_byte_aligned<max(MaxVecBits/8, 1u)>(raw_pointer_cast(dst.data())));
constexpr int elem_bits = sizeof_bits_v<typename SrcEngine::value_type>;
//
// Determine val+thr vectorization based on src/dst size and number of threads
// NOTE: This heuristic promotes parallelization over vectorization
//
// The number of elements that can be vectorized in values
constexpr int common_elem = decltype(max_common_vector(src, dst))::value;
constexpr int common_bits = common_elem * elem_bits;
constexpr int total_elem = decltype(size(src))::value;
// If there are too many threads to allow a full vectorized copy, trunc the vectorization
constexpr int total_bits = total_elem * elem_bits;
static_assert(total_bits % NumThreads == 0);
constexpr int total_bits_per_thr = total_bits / NumThreads;
// If there are too many threads to allow a full elem copy, trunc the thrs and use elem_bits
constexpr int max_vec_bits_by_thr = cute::max(elem_bits, total_bits_per_thr);
// Cap the vectorization to the common bits, the max_vec_bits_by_thr, and the MaxVecBits
constexpr int vec_bits = cute::min(common_bits, max_vec_bits_by_thr, static_cast<int>(MaxVecBits));
// Convert back to number of elements, safe_div
static_assert((vec_bits % elem_bits) == 0);
constexpr int vec_elem = vec_bits / elem_bits;
// Use only part of threads if there's not enough work for all threads
constexpr int vec_thrs = (total_elem % (vec_elem * NumThreads) == 0)
? NumThreads
: (total_elem / vec_elem);
static_assert(vec_thrs <= NumThreads);
// The common layout of the two tensors that can be vectorized over threads
// vidx -> coord
auto common_layout = max_common_layout(get_nonswizzle_portion(src.layout()),
get_nonswizzle_portion(dst.layout()));
// Scale up the common_layout to cover the entire tensors
// vidx -> coord
auto full_perm = tile_to_shape(make_layout(common_layout), size(src));
// Create the Tiler
// ((vid,tid),iter)
auto layout_vt = logical_divide(full_perm, Layout<Shape<Int<vec_elem>, Int<vec_thrs>>>{});
// Apply and slice
Tensor src_v = src.compose(layout_vt)(make_coord(_,tid),_);
Tensor dst_v = dst.compose(layout_vt)(make_coord(_,tid),_);
constexpr int max_bits_per_thr = total_bits / NumThreads;
// At least elem_bits, at most common_bits
constexpr int common_bits = common_elem * elem_bits;
constexpr int vec_bits = cute::max(elem_bits, cute::gcd(common_bits, int(MaxVecBits), max_bits_per_thr));
// Should account for vec_bits < 8 and/or vec_elem <= 1
// And also account for subbyte types, which could cause race conditions
// Want to ENFORCE sufficient vectorization in those cases
static_assert((vec_bits >= 8), "No support for subbyte copying");
static_assert(vec_bits % elem_bits == 0, "Expected divisibility");
static_assert(vec_bits >= 8, "No support for subbyte copying");
using VecType = uint_bit_t<vec_bits>;
constexpr int vec_elem = vec_bits / elem_bits;
constexpr int vec_thrs = cute::min(int(NumThreads), total_elem / vec_elem);
//
// Determine the partitioning patterns for the vec_elems and vec_thrs
//
// Distribute the rest of the V*T to some consistent portion outside of the common_layout, if needed
auto common_domain_src = domain_distribute(shape(src_a), Int<vec_elem*vec_thrs>{});
auto common_domain_dst = domain_distribute(shape(dst_a), Int<vec_elem*vec_thrs>{});
// Make sure for now, could fall back here instead
CUTE_STATIC_ASSERT_V(size(common_domain_src) == Int<vec_elem*vec_thrs>{});
CUTE_STATIC_ASSERT_V(compatible(common_domain_src, common_domain_dst) ||
compatible(common_domain_dst, common_domain_src));
// Use the "more specific" domain for the extra elements of V*T
auto common_domain = conditional_return(compatible(common_domain_src, common_domain_dst),
common_domain_dst, common_domain_src);
// Construct the tiler
auto tiler_vt = common_domain.with_shape(Int<vec_elem>{}, Int<vec_thrs>{});
// Apply and slice
Tensor src_v = logical_divide(src_a, tiler_vt)(make_coord(_,tid),_);
Tensor dst_v = logical_divide(dst_a, tiler_vt)(make_coord(_,tid),_);
#if 0
if (thread0()) {
print(" "); print("cooperative_copy -- vec\n");
print(" "); print("NumThreads: "); print(NumThreads); print("\n");
print(" "); print("MaxVecBits: "); print(MaxVecBits); print("\n");
print(" "); print("src: "); print(src); print("\n");
print(" "); print("dst: "); print(dst); print("\n");
print(" "); print("common_layout: "); print(common_layout); print("\n");
print(" "); print("full_perm: "); print(full_perm); print("\n");
print(" "); print("Used vector: "); print(vec_elem); print("\n");
print(" "); print("Used threads: "); print(vec_thrs); print("\n");
print(" "); print("layout_vt: "); print(layout_vt); print("\n");
print(" "); print("src.compose(layout_vt): "); print(src.compose(layout_vt)); print("\n");
print(" "); print("dst.compose(layout_vt): "); print(dst.compose(layout_vt)); print("\n");
print(" "); print("src_v: "); print(src_v); print("\n");
print(" "); print("dst_v: "); print(dst_v); print("\n");
print(" "); print("recast<VecType const>(src_v): "); print(recast<VecType const>(src_v)); print("\n");
print(" "); print("recast<VecType const>(dst_v): "); print(recast<VecType const>(dst_v)); print("\n");
}
if (thread0()) {
print(" "); print("cooperative_copy -- vec\n");
print(" "); print("Used vector: "); print(vec_elem); print("\n");
print(" "); print("Used threads: "); print(vec_thrs); print("\n");
print(" "); print("tiler_vt: "); print(tiler_vt); print("\n");
print(" "); print("src_v: "); print(src_v); print("\n");
print(" "); print("dst_v: "); print(dst_v); print("\n");
print(" "); print("recast<VecType const>(src_v): "); print(recast<VecType const>(src_v)); print("\n");
print(" "); print("recast<VecType >(dst_v): "); print(recast<VecType >(dst_v)); print("\n");
}
#ifdef __CUDA_ARCH__
__syncthreads();
__syncthreads();
#endif
#endif
// If we're using all threads (static) or the tid is in in-range (dynamic)
if (vec_thrs >= NumThreads or tid < vec_thrs) {
// If we're using all threads (static) or the tid is in-range (dynamic)
if (vec_thrs == NumThreads or tid < vec_thrs) {
return copy_if(TrivialPredTensor{}, recast<VecType const>(src_v), recast<VecType>(dst_v));
}
}
}
// Default max-vectorization size to value_type size
template <uint32_t NumThreads,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
@@ -184,7 +300,10 @@ cooperative_copy(uint32_t const& tid,
return cooperative_copy<NumThreads, MaxVecBits>(tid, src, dst);
}
//
// Accept mutable temporaries
//
template <uint32_t NumThreads,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
@@ -197,9 +316,7 @@ cooperative_copy(uint32_t const& tid,
return cooperative_copy<NumThreads>(tid, src, dst);
}
// Accept mutable temporaries
template <uint32_t NumThreads,
uint32_t MaxVecBits,
template <uint32_t NumThreads, uint32_t MaxVecBits,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
include/cute/algorithm/cooperative_gemm.hpp
+8 -25
View File
@@ -39,7 +39,7 @@
#include <cute/algorithm/functional.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
namespace cute
{
@@ -76,29 +76,15 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
using TypeB = typename TB::value_type;
using TypeC = typename TC::value_type;
// Original, static size of the problem
auto M = size<0>(sC);
auto N = size<1>(sC);
auto K = size<1>(sA);
// Block size of the compute tile
auto BLK_M = tile_size<0>(thr_mma);
auto BLK_N = tile_size<1>(thr_mma);
auto BLK_K = tile_size<2>(thr_mma);
//
// MMA Partitioning
//
// Round the layout extents up to BLK_X to satisfy MMA partitioning safety
Tensor rounded_sA = sA.compose(make_shape(round_up(M, BLK_M), round_up(K, BLK_K)));
Tensor rounded_sB = sB.compose(make_shape(round_up(N, BLK_N), round_up(K, BLK_K)));
Tensor rounded_sC = sC.compose(make_shape(round_up(M, BLK_M), round_up(N, BLK_N)));
// Partition the sA, sB, and sC tiles across the threads for the MMA
Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K)
Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K)
Tensor tCsC = thr_mma.partition_C(sC); // (MMA,MMA_M,MMA_N)
// Partition the sA and sB tiles across the threads for the MMA
Tensor tCsA = thr_mma.partition_A(rounded_sA); // (MMA,MMA_M,MMA_K)
Tensor tCsB = thr_mma.partition_B(rounded_sB); // (MMA,MMA_N,MMA_K)
Tensor tCsC = thr_mma.partition_C(rounded_sC); // (MMA,MMA_M,MMA_N)
// Create register tensors for the MMA to operate on
Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K)
Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K)
@@ -109,9 +95,6 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
print(" sA: "); print( sA); print("\n");
print(" sB: "); print( sB); print("\n");
print(" sC: "); print( sC); print("\n");
print("r_sA: "); print(rounded_sA); print("\n");
print("r_sB: "); print(rounded_sB); print("\n");
print("r_sC: "); print(rounded_sC); print("\n");
print(thr_mma);
print("tCsA: "); print(tCsA); print("\n");
print("tCsB: "); print(tCsB); print("\n");
@@ -127,8 +110,8 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
//
// Create coordinate tensors for the problem
Tensor cA = make_identity_tensor(shape(rounded_sA)); // (M,K) -> (m,k)
Tensor cB = make_identity_tensor(shape(rounded_sB)); // (N,K) -> (n,k)
Tensor cA = make_identity_tensor(shape(sA)); // (M,K) -> (m,k)
Tensor cB = make_identity_tensor(shape(sB)); // (N,K) -> (n,k)
// Repeat partitioning with thr_mma
Tensor tCcA = thr_mma.partition_A(cA); // (MMA,MMA_M,MMA_K) -> (m,k)
@@ -222,7 +205,7 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
//
// Create coordinate tensors for the problem
Tensor cC = make_identity_tensor(shape(rounded_sC)); // (M,N) -> (m,n)
Tensor cC = make_identity_tensor(shape(sC)); // (M,N) -> (m,n)
// Repeat partitioning with thr_mma
Tensor tCcC = thr_mma.partition_C(cC); // (MMA,MMA_M,MMA_N) -> (m,n)
include/cute/algorithm/copy.hpp
+28 -35
View File
@@ -34,7 +34,7 @@
#include <cute/container/alignment.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/tensor_predicate.hpp>
#include <cute/atom/copy_atom.hpp>
@@ -199,14 +199,14 @@ copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
{
static_assert(sizeof_bits_v<VecType> >= 8 && sizeof_bits_v<VecType> % 8 == 0,
"Expected a vectorization type of at least a byte.");
using SrcType = typename SrcEngine::element_type;
using DstType = typename DstEngine::element_type;
if constexpr (sizeof_bits_v<SrcType> == sizeof_bits_v<DstType> &&
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
if constexpr (cute::is_same<SrcType, DstType>::value &&
sizeof_bits_v<VecType> > sizeof_bits_v<DstType>)
{
// Preserve volatility of Src/Dst types.
using SrcVecType = conditional_t<is_volatile_v<SrcType>, VecType const volatile, VecType const>;
using DstVecType = conditional_t<is_volatile_v<DstType>, VecType volatile, VecType >;
using SrcVecType = conditional_t<is_volatile_v<typename SrcEngine::element_type>, VecType const volatile, VecType const>;
using DstVecType = conditional_t<is_volatile_v<typename DstEngine::element_type>, VecType volatile, VecType >;
Tensor src_v = recast<SrcVecType>(src);
Tensor dst_v = recast<DstVecType>(dst);
@@ -264,22 +264,22 @@ copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits> const&,
{
constexpr int vec_elem = decltype(max_common_vector(src, dst))::value;
constexpr int src_bits = sizeof_bits<typename SrcEngine::value_type>::value;
// When layouts are static, accept vec_bits up to 128
// When layouts are dynamic, accept vec_bits up to MaxVecBits
constexpr int vec_bits = (is_static<SrcLayout>::value && is_static<DstLayout>::value) ?
cute::min(vec_elem * src_bits, 128) :
cute::min(vec_elem * src_bits, MaxVecBits);
constexpr int max_align_src = decltype(max_alignment(src.layout()))::value;
constexpr int max_align_dst = decltype(max_alignment(dst.layout()))::value;
constexpr int max_align = gcd(vec_elem, max_align_src, max_align_dst);
#if 0
if (thread0()) {
print("copy -- found max_common_vector of %d elems and vectorization to %d bits\n", vec_elem, vec_bits);
print(" "); print(src); print("\n");
print(" "); print(dst); print("\n");
}
#endif
constexpr int src_bits = sizeof_bits<typename SrcEngine::value_type>::value;
constexpr int vec_bits = gcd(src_bits * max_align, MaxVecBits);
if constexpr (vec_elem > 1 && vec_bits >= 8) {
// If more than one element vectorizes to 8bits or more, then copy_vec
#if 0
if (thread0()) {
print("copy -- found max_common_vector of %d elems and vectorization to %d bits\n", vec_elem, vec_bits);
print(" "); print(src); print("\n");
print(" "); print(dst); print("\n");
}
#endif
return copy_vec<uint_bit_t<vec_bits>>(src, dst);
} else {
return copy_if(TrivialPredTensor{}, src, dst);
@@ -294,10 +294,16 @@ void
copy(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(AutoVectorizingCopy{}, src, dst);
if constexpr (is_static<SrcLayout>::value && is_static<DstLayout>::value) {
// Assume Tensors with static layouts (e.g. registers) have pointers that are 128b aligned
return copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, src, dst);
} else {
// Do not assume that dynamic layouts are aligned.
return copy(AutoVectorizingCopyWithAssumedAlignment<8>{}, src, dst);
}
}
// Auto-vectorizing copy with assumed alignment of dynamic layout strides up to 128bit.
// Auto-vectorizing copy with assumed alignment up to 128bit.
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
@@ -308,19 +314,6 @@ copy_aligned(Tensor<SrcEngine, SrcLayout> const& src,
return copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, src, dst);
}
// Specializaton for Atom AutoVectorizingCopy
template <class... Args,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Atom<AutoVectorizingCopy, Args...> const&,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(AutoVectorizingCopy{}, src, dst);
}
// Specializaton for Atom AutoVectorizingCopyAssumedAlignment
template <int MaxVecBits, class... Args,
class SrcEngine, class SrcLayout,
@@ -346,7 +339,7 @@ copy(Copy_Traits<SM90_BULK_COPY_AUTO, CT_Args...> const& atom, // Copy_Traits m
{
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
static_assert(sizeof_bits<SrcType>::value == sizeof_bits<DstType>::value);
static_assert(cute::is_same<SrcType, DstType>::value);
static_assert((is_gmem<SrcEngine>::value && is_smem<DstEngine>::value) ||
(is_smem<SrcEngine>::value && is_gmem<DstEngine>::value),
"Bulk Copy only supports gmem -> smem or smem -> gmem movement.");
include/cute/algorithm/fill.hpp
+1 -1
View File
@@ -32,7 +32,7 @@
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/algorithm/prefer.hpp>
namespace cute
include/cute/algorithm/gemm.hpp
+1 -1
View File
@@ -35,7 +35,7 @@
#include <cute/util/type_traits.hpp>
#include <cute/algorithm/functional.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/atom/mma_atom.hpp>
include/cute/algorithm/prefetch.hpp
+1 -7
View File
@@ -32,7 +32,7 @@
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/atom/copy_atom.hpp>
@@ -90,12 +90,6 @@ constexpr bool has_prefetch = false;
template <class CopyOp>
constexpr bool has_prefetch<CopyOp, void_t<typename CopyOp::PREFETCH>> = true;
template <class CopyOp, class = void>
constexpr bool is_prefetch = false;
template <class CopyOp>
constexpr bool is_prefetch<CopyOp, void_t<typename CopyOp::PREFETCH>> = is_same_v<CopyOp, typename CopyOp::PREFETCH>;
} // end namespace detail
template <class CopyOp, class... CT_Args, class... CA_Args,
include/cute/algorithm/tensor_algorithms.hpp
+13 -14
View File
@@ -33,8 +33,7 @@
#pragma once
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
namespace cute
{
@@ -100,13 +99,13 @@ transform(Tensor<Engine,Layout>&& tensor, UnaryOp&& op)
}
// Similar to std::transform transforms one tensors and assigns it to another
template <class EngineIn, class LayoutIn,
class EngineOut, class LayoutOut,
template <class EngineIn, class LayoutIn,
class EngineOut, class LayoutOut,
class UnaryOp>
CUTE_HOST_DEVICE constexpr
void
transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
Tensor<EngineOut,LayoutOut> & tensor_out,
transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
Tensor<EngineOut,LayoutOut> & tensor_out,
UnaryOp&& op)
{
CUTE_UNROLL
@@ -117,30 +116,30 @@ transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
// Accept mutable temporaries
template <class EngineIn, class LayoutIn,
class EngineOut, class LayoutOut,
class EngineOut, class LayoutOut,
class UnaryOp>
CUTE_HOST_DEVICE constexpr
void
transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
Tensor<EngineOut,LayoutOut> && tensor_out,
transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
Tensor<EngineOut,LayoutOut> && tensor_out,
UnaryOp&& op)
{
return transform(tensor_in, tensor_out, op);
}
// Similar to std::transform with a binary operation
// Takes two tensors as input and one tensor as output.
// Takes two tensors as input and one tensor as output.
// Applies the binary_op to tensor_in1 and tensor_in2 and
// assigns it to tensor_out
template <class EngineIn1, class LayoutIn1,
class EngineIn2, class LayoutIn2,
class EngineOut, class LayoutOut,
class EngineOut, class LayoutOut,
class BinaryOp>
CUTE_HOST_DEVICE constexpr
void
transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
Tensor<EngineIn2,LayoutIn2> const& tensor_in2,
Tensor<EngineOut,LayoutOut> & tensor_out,
Tensor<EngineOut,LayoutOut> & tensor_out,
BinaryOp&& op)
{
CUTE_UNROLL
@@ -152,11 +151,11 @@ transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
// Accept mutable temporaries
template <class EngineIn1, class LayoutIn1,
class EngineIn2, class LayoutIn2,
class EngineOut, class LayoutOut,
class EngineOut, class LayoutOut,
class BinaryOp>
CUTE_HOST_DEVICE constexpr
void
transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
Tensor<EngineIn2,LayoutIn2> const& tensor_in2,
Tensor<EngineOut,LayoutOut> && tensor_out,
BinaryOp&& op)
include/cute/algorithm/tuple_algorithms.hpp
+53 -21
View File
@@ -404,29 +404,54 @@ namespace detail {
// This impl compiles much faster than cute::apply and variadic args
template <class T, class V, class F>
CUTE_HOST_DEVICE constexpr
decltype(auto)
fold(T&& t, V&& v, F&& f, seq<>)
auto
fold(T&&, V&& v, F&&, seq<>)
{
return static_cast<V&&>(v);
return v;
}
template <class T, class V, class F, int I, int... Is>
template <class T, class V, class F, int I0>
CUTE_HOST_DEVICE constexpr
decltype(auto)
fold(T&& t, V&& v, F&& f, seq<I,Is...>)
auto
fold(T&& t, V&& v, F&& f, seq<I0>)
{
if constexpr (sizeof...(Is) == 0) {
return f(static_cast<V&&>(v), get<I>(static_cast<T&&>(t)));
} else {
return fold(static_cast<T&&>(t),
f(static_cast<V&&>(v), get<I>(static_cast<T&&>(t))),
f,
seq<Is...>{});
}
CUTE_GCC_UNREACHABLE;
return f(static_cast<V&&>(v), get<I0>(static_cast<T&&>(t)));
}
template <class T, class V, class F, int I0, int I1>
CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f, seq<I0,I1>)
{
return f(f(static_cast<V&&>(v), get<I0>(static_cast<T&&>(t))), get<I1>(static_cast<T&&>(t)));
}
template <class T, class V, class F, int I0, int I1, int I2>
CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f, seq<I0,I1,I2>)
{
return f(f(f(static_cast<V&&>(v), get<I0>(static_cast<T&&>(t))), get<I1>(static_cast<T&&>(t))), get<I2>(static_cast<T&&>(t)));
}
template <class T, class V, class F, int I0, int I1, int I2, int I3>
CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f, seq<I0,I1,I2,I3>)
{
return f(f(f(f(static_cast<V&&>(v), get<I0>(static_cast<T&&>(t))), get<I1>(static_cast<T&&>(t))), get<I2>(static_cast<T&&>(t))), get<I3>(static_cast<T&&>(t)));
}
template <class T, class V, class F, int I0, int I1, int I2, int I3, int... Is>
CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f, seq<I0,I1,I2,I3,Is...>)
{
return fold(static_cast<T&&>(t),
f(f(f(f(static_cast<V&&>(v), get<I0>(static_cast<T&&>(t))), get<I1>(static_cast<T&&>(t))), get<I2>(static_cast<T&&>(t))), get<I3>(static_cast<T&&>(t))),
f,
seq<Is...>{});
}
} // end namespace detail
template <class T, class V, class F>
@@ -448,7 +473,7 @@ fold(T&& t, V&& v, F&& f)
template <class T, class F>
CUTE_HOST_DEVICE constexpr
decltype(auto)
auto
fold_first(T&& t, F&& f)
{
if constexpr (is_tuple<remove_cvref_t<T>>::value) {
@@ -457,7 +482,7 @@ fold_first(T&& t, F&& f)
f,
make_range<1,tuple_size<remove_cvref_t<T>>::value>{});
} else {
return static_cast<T&&>(t);
return t;
}
CUTE_GCC_UNREACHABLE;
@@ -701,7 +726,14 @@ CUTE_HOST_DEVICE constexpr
auto
replace(T const& t, X const& x)
{
return detail::construct(t, x, make_seq<N>{}, seq<0>{}, make_range<N+1,tuple_size<T>::value>{});
if constexpr (is_tuple<T>::value) {
return detail::construct(t, x, make_seq<N>{}, seq<0>{}, make_range<N+1,tuple_size<T>::value>{});
} else {
static_assert(N == 0);
return x;
}
CUTE_GCC_UNREACHABLE;
}
// Replace the first element of the tuple with x
@@ -1077,9 +1109,9 @@ zip2_by(T const& t, TG const& guide)
/// @return A tuple of the elements of @c t in reverse order.
template <class T>
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
auto
reverse(T const& t)
reverse(T const& t)
{
if constexpr (is_tuple<T>::value) {
return detail::apply(t, [](auto const&... a){ return cute::make_tuple(a...); }, tuple_rseq<T>{});
include/cute/arch/copy.hpp
+8 -6
View File
@@ -68,7 +68,7 @@ struct UniversalCopy
//
// Placeholder for the copy algorithm's stronger auto-vectorizing behavior
// that assumes alignment of dynamic layouts up to MaxVecBits
// that assumes alignment of pointers and dynamic layouts up to MaxVecBits
//
template <int MaxVecBits = 128>
@@ -80,15 +80,17 @@ struct AutoVectorizingCopyWithAssumedAlignment
};
//
// Placeholder for the copy algorithm's default auto-vectorizing behavior
// that does not assume alignment of dynamic layouts
// AutoVectorizingCopy alias assumes maximal alignment of pointers and dynamic strides.
// If this is not the case then AutoVectorizingCopyWithAssumedAlignment should be used instead
//
using AutoVectorizingCopy = AutoVectorizingCopyWithAssumedAlignment<8>;
using AutoVectorizingCopy = AutoVectorizingCopyWithAssumedAlignment<128>;
// Alias
using DefaultCopy = AutoVectorizingCopy;
//
// DefaultCopy alias does not assume alignment of pointers or dynamic strides.
//
using DefaultCopy = AutoVectorizingCopyWithAssumedAlignment<8>;
//
// Global memory prefetch into L2
include/cute/arch/copy_sm90_desc.hpp
+91 -18
View File
@@ -95,8 +95,8 @@ wait_barrier(uint64_t& smem_barrier, // 64 bits user-mange
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"bra.uni LAB_WAIT;\n"
"@P1 bra DONE;\n"
"bra LAB_WAIT;\n"
"DONE:\n"
"}\n"
:: "r"(smem_int_ptr),
@@ -134,6 +134,48 @@ enum class SmemSwizzleBits : uint8_t {
B128 = 3,
};
enum class OOBFill : uint8_t {
ZERO = 0,
CONSTANT = 1,
};
CUTE_HOST_DEVICE char const* to_string(OOBFill const& t) {
switch (t) {
case OOBFill::ZERO: return "ZERO";
case OOBFill::CONSTANT: return "CONSTANT";
}
return nullptr;
}
enum class L2Promotion : uint8_t {
DISABLE = 0,
B64 = 1,
B128 = 2,
B256 = 3,
};
CUTE_HOST_DEVICE char const* to_string(L2Promotion const& t) {
switch (t) {
case L2Promotion::DISABLE: return "DISABLE";
case L2Promotion::B64: return "B64";
case L2Promotion::B128: return "B128";
case L2Promotion::B256: return "B256";
}
return nullptr;
}
// Aux parameters which are independent with the problem size
struct DescriptorAuxParams {
OOBFill oobfill_ = OOBFill::ZERO;
L2Promotion l2promo_ = L2Promotion::DISABLE;
};
enum class CacheHintSm90 : uint64_t {
EVICT_NORMAL = 0x1000000000000000,
EVICT_FIRST = 0x12F0000000000000,
EVICT_LAST = 0x14F0000000000000,
};
#if (__CUDACC_VER_MAJOR__ >= 12)
#if !defined(__CUDACC_RTC__)
@@ -168,6 +210,27 @@ to_CUtensorMapSwizzle(SmemSwizzleBits const& t) {
case SmemSwizzleBits::B128: return CU_TENSOR_MAP_SWIZZLE_128B;
}
}
inline CUtensorMapFloatOOBfill
to_CUtensorMapFloatOOBfill(OOBFill const& t) {
switch(t) {
default: assert(false && "Unknown OOBFill!");
case OOBFill::ZERO: return CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE;
case OOBFill::CONSTANT: return CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA;
}
}
inline CUtensorMapL2promotion
to_CUtensorMapL2promotion(L2Promotion const& t) {
switch(t) {
default: assert(false && "Unknown L2Promotion!");
case L2Promotion::DISABLE: return CU_TENSOR_MAP_L2_PROMOTION_NONE;
case L2Promotion::B64: return CU_TENSOR_MAP_L2_PROMOTION_L2_64B;
case L2Promotion::B128: return CU_TENSOR_MAP_L2_PROMOTION_L2_128B;
case L2Promotion::B256: return CU_TENSOR_MAP_L2_PROMOTION_L2_256B;
}
}
#endif // !defined(__CUDACC_RTC__)
#endif // (__CUDACC_VER_MAJOR__ >= 12)
@@ -257,22 +320,32 @@ tma_descriptor_replace_dims_strides_in_shared_mem(TmaDescriptor
asm volatile (
"cvt.u64.u32 %0, %1;"
:: "l"(smem_int64_desc), "r"(smem_int_desc));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[0]));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[1]));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 2, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[2]));
// Strides must be a multiple of 16. Also, stride for the intermost dimension is implicitly 1
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[1] >> 4));
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 1, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[2] >> 4));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[0]));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[1]));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 2, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[2]));
// Strides must be a multiple of 16. Also, stride for the intermost dimension is implicitly 1
#if ((__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 5)))
// 4 LSBs are not included
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[1]));
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 1, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[2]));
#else
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[1] >> 4));
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 1, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[2] >> 4));
#endif
#else
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
#endif
include/cute/arch/copy_sm90_tma.hpp
+30 -30
View File
@@ -44,7 +44,7 @@ namespace cute
struct SM90_TMA_LOAD_1D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0)
{
@@ -53,11 +53,11 @@ struct SM90_TMA_LOAD_1D
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3}], [%2];"
"cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3}], [%2], %4;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0)
"r"(crd0), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -89,7 +89,7 @@ struct SM90_TMA_LOAD_1D
struct SM90_TMA_LOAD_2D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1)
{
@@ -98,11 +98,11 @@ struct SM90_TMA_LOAD_2D
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4}], [%2];"
"cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4}], [%2], %5;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1)
"r"(crd0), "r"(crd1), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -134,7 +134,7 @@ struct SM90_TMA_LOAD_2D
struct SM90_TMA_LOAD_3D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2)
{
@@ -143,11 +143,11 @@ struct SM90_TMA_LOAD_3D
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5}], [%2];"
"cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5}], [%2], %6;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2)
"r"(crd0), "r"(crd1), "r"(crd2), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -179,7 +179,7 @@ struct SM90_TMA_LOAD_3D
struct SM90_TMA_LOAD_4D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3)
{
@@ -188,11 +188,11 @@ struct SM90_TMA_LOAD_4D
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6}], [%2];"
"cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6}], [%2], %7;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3)
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -224,7 +224,7 @@ struct SM90_TMA_LOAD_4D
struct SM90_TMA_LOAD_5D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4)
{
@@ -233,11 +233,11 @@ struct SM90_TMA_LOAD_5D
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2];"
"cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4)
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -269,39 +269,39 @@ struct SM90_TMA_LOAD_5D
struct SM90_TMA_LOAD
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0)
{
return SM90_TMA_LOAD_1D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0);
return SM90_TMA_LOAD_1D::copy(desc_ptr, mbar_ptr, cache_hint, smem_ptr, crd0);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1)
{
return SM90_TMA_LOAD_2D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0, crd1);
return SM90_TMA_LOAD_2D::copy(desc_ptr, mbar_ptr, cache_hint, smem_ptr, crd0, crd1);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2)
{
return SM90_TMA_LOAD_3D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0, crd1, crd2);
return SM90_TMA_LOAD_3D::copy(desc_ptr, mbar_ptr, cache_hint, smem_ptr, crd0, crd1, crd2);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3)
{
return SM90_TMA_LOAD_4D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0, crd1, crd2, crd3);
return SM90_TMA_LOAD_4D::copy(desc_ptr, mbar_ptr, cache_hint, smem_ptr, crd0, crd1, crd2, crd3);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4)
{
return SM90_TMA_LOAD_5D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0, crd1, crd2, crd3, crd4);
return SM90_TMA_LOAD_5D::copy(desc_ptr, mbar_ptr, cache_hint, smem_ptr, crd0, crd1, crd2, crd3, crd4);
}
struct PREFETCH
include/cute/arch/mma_sm90_desc.hpp
+14 -14
View File
@@ -85,7 +85,6 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, LayoutType const& t) {
union GmmaDescriptor
{
CUTE_HOST_DEVICE constexpr
GmmaDescriptor() noexcept : desc_(0) {}
CUTE_HOST_DEVICE constexpr
@@ -135,21 +134,22 @@ union GmmaDescriptor
// Decay to a uint64_t
CUTE_HOST_DEVICE constexpr
operator uint64_t() const noexcept { return desc_; }
// Printer
CUTE_HOST_DEVICE friend void print(GmmaDescriptor const& t)
{
#if !defined(__CUDACC_RTC__)
printf("GmmaDescriptor: 0x%016llx\n", static_cast<unsigned long long>(t.desc_));
printf(" start_addr : 0x%04x\n", t.bitfield.start_address_);
printf(" leading_off: 0x%04x (%d)\n", t.bitfield.leading_byte_offset_, t.bitfield.leading_byte_offset_);
printf(" stride_off : 0x%04x (%d)\n", t.bitfield.stride_byte_offset_, t.bitfield.stride_byte_offset_);
printf(" base_offset: 0x%01x\n", t.bitfield.base_offset_);
printf(" layout_type: 0x%01x (%s)\n", t.bitfield.layout_type_, to_string(static_cast<GMMA::LayoutType>(t.bitfield.layout_type_)));
#endif
}
};
// Printer
CUTE_HOST_DEVICE void
print(GmmaDescriptor const& t)
{
#if !defined(__CUDACC_RTC__)
printf("GmmaDescriptor: 0x%016llx\n", static_cast<unsigned long long>(t.desc_));
printf(" start_addr : 0x%04x\n", t.bitfield.start_address_);
printf(" leading_off: 0x%04x (%d)\n", t.bitfield.leading_byte_offset_, t.bitfield.leading_byte_offset_);
printf(" stride_off : 0x%04x (%d)\n", t.bitfield.stride_byte_offset_, t.bitfield.stride_byte_offset_);
printf(" base_offset: 0x%01x\n", t.bitfield.base_offset_);
printf(" layout_type: 0x%01x (%s)\n", t.bitfield.layout_type_, to_string(static_cast<GMMA::LayoutType>(t.bitfield.layout_type_)));
#endif // !defined(__CUDACC_RTC__)
}
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cute
include/cute/arch/util.hpp
+14 -13
View File
@@ -235,24 +235,25 @@ explode(Fn fn,
}
template <class Fn,
class PtrD, int... Id,
class PtrA, int... Ia,
class PtrB, int... Ib,
class PtrC, int... Ic,
class PtrSFA, int... Isfa,
class PtrSFB, int... Isfb>
class PtrD, int... Id,
class PtrA, int... Ia,
class PtrB, int... Ib,
class PtrC, int... Ic,
class PtrE, int... Ie,
class PtrF, int... If>
CUTE_HOST_DEVICE constexpr
void
explode(Fn fn,
PtrD&& d, int_sequence<Id...>,
PtrA&& a, int_sequence<Ia...>,
PtrB&& b, int_sequence<Ib...>,
PtrC&& c, int_sequence<Ic...>,
PtrSFA&& sfa, int_sequence<Isfa...>,
PtrSFB&& sfb, int_sequence<Isfb...>)
PtrD&& d, int_sequence<Id...>,
PtrA&& a, int_sequence<Ia...>,
PtrB&& b, int_sequence<Ib...>,
PtrC&& c, int_sequence<Ic...>,
PtrE&& e, int_sequence<Ie...>,
PtrF&& f, int_sequence<If...>)
{
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., sfa[Isfa]..., sfb[Isfb]...);
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., e[Ie]..., f[If]...);
}
//
// Utility for exploding tuples into functions
//
include/cute/atom/copy_atom.hpp
+1 -1
View File
@@ -39,7 +39,7 @@
#include <cute/util/type_traits.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
namespace cute
{
include/cute/atom/copy_traits.hpp
+12 -1
View File
@@ -32,7 +32,7 @@
#include <cute/arch/copy.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
namespace cute
{
@@ -145,4 +145,15 @@ copy_unpack(Copy_Traits<CopyOp,Args...> const& traits,
copy_unpack(traits, src, dst);
}
namespace detail {
template <class CopyOp, class = void>
constexpr bool is_prefetch = false;
template <class CopyOp>
constexpr bool is_prefetch<CopyOp, void_t<typename CopyOp::PREFETCH>> = is_same_v<CopyOp, typename CopyOp::PREFETCH>;
} // end namespace detail
} // end namespace cute
include/cute/atom/copy_traits_sm90_im2col.hpp
+49 -41
View File
@@ -39,7 +39,7 @@
#include "cute/tensor.hpp"
#include "cute/algorithm/prefetch.hpp"
#include "cutlass/fast_math.h"
namespace cute
{
@@ -388,18 +388,19 @@ template <class EngineA, class LayoutA,
CUTE_HOST
auto
make_im2col_tma_copy_desc(
Tensor<EngineA, LayoutA> const& tensor_cwhdn, // (C,W,H,D,N)
uint32_t range_c, // TILE_C
uint32_t range_whdn, // TILE_WHDN
SmemSwizzle const& smem_swizzle, // Swizzle
TMALayout const& tma_layout_vt, // TMA layout
LowerCornerStride const& lower_corner_whd, // WHD offset of the "base pointer"
UpperCornerStride const& upper_corner_whd, // WHD upper corner
LowerPaddingStride const& lower_padding_whd, // WHD lower padding
UpperPaddingStride const& upper_padding_whd, // WHD upper padding
TraversalStride const& stride_whd, // WHD traversal stride
LowerSRTStride const& lower_srt, // SRT offset of the "base pointer"
DilationStride const& stride_srt) // SRT stride - dilation
Tensor<EngineA, LayoutA> const& tensor_cwhdn, // (C,W,H,D,N)
uint32_t range_c, // TILE_C
uint32_t range_whdn, // TILE_WHDN
SmemSwizzle const& smem_swizzle, // Swizzle
TMALayout const& tma_layout_vt, // TMA layout
LowerCornerStride const& lower_corner_whd, // WHD offset of the "base pointer"
UpperCornerStride const& upper_corner_whd, // WHD upper corner
LowerPaddingStride const& lower_padding_whd, // WHD lower padding
UpperPaddingStride const& upper_padding_whd, // WHD upper padding
TraversalStride const& stride_whd, // WHD traversal stride
LowerSRTStride const& lower_srt, // SRT offset of the "base pointer"
DilationStride const& stride_srt, // SRT stride - dilation
TMA::DescriptorAuxParams const& aux_params = {})
{
static_assert(is_gmem<EngineA>::value, "Tensor must point to GPU global memory.");
using value_type = typename EngineA::value_type;
@@ -445,8 +446,8 @@ make_im2col_tma_copy_desc(
CUtensorMapDataType tma_format = TMA::to_CUtensorMapDataType<value_type>();
CUtensorMapInterleave tma_interleave = CU_TENSOR_MAP_INTERLEAVE_NONE;
CUtensorMapL2promotion tma_l2Promotion = CU_TENSOR_MAP_L2_PROMOTION_NONE;
CUtensorMapFloatOOBfill tma_oob_fill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE;
CUtensorMapL2promotion tma_l2Promotion = to_CUtensorMapL2promotion(aux_params.l2promo_);
CUtensorMapFloatOOBfill tma_oob_fill = to_CUtensorMapFloatOOBfill(aux_params.oobfill_);
CUtensorMapSwizzle tma_swizzle = TMA::to_CUtensorMapSwizzle(detail::get_tma_swizzle_bits(smem_swizzle));
CUresult encode_result = cuTensorMapEncodeIm2col(
@@ -498,7 +499,11 @@ make_im2col_tma_copy_desc(
// For fprop/dgrad kernel, gemm_shapes is ((q, p, z, n), (c, s, r, t))
// For wgrad kernel, gemm_shapes is ((c, s, r, t), (q, p, z, n))
auto gemm_shapes_common = make_shape(gemm_mn, gemm_k);
auto gemm_shapes_common = make_shape(
transform_leaf(gemm_mn, [](auto s) {
return conditional_return(cute::is_static<decltype(s)>{}, s, cutlass::FastDivmod(s));
}),
gemm_k);
auto gemm_shapes = make_shape(
basis_get(stride<0,1>(tma_layout_vt), gemm_shapes_common),
basis_get(stride<0,0>(tma_layout_vt), gemm_shapes_common));
@@ -554,17 +559,18 @@ template <class CopyOp,
CUTE_HOST_RTC
auto
make_tma_atom_im2col(CopyOp,
Tensor<GEngine,GLayout> const& gtensor, // Full GMEM Tensor: ((w, h, d, n), c)
SLayout const& slayout, // CTA Tile of SMEM, potentially swizzled
int32_t const& num_multicast, // The number of CTAs involved in multicasting
Layout<VShape,VStride> const& cta_v_map, // V: CTA val idx -> gmem mode
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd, // traversal stride
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt) // dilation
Tensor<GEngine,GLayout> const& gtensor, // Full GMEM Tensor: ((w, h, d, n), c)
SLayout const& slayout, // CTA Tile of SMEM, potentially swizzled
int32_t const& num_multicast, // The number of CTAs involved in multicasting
Layout<VShape,VStride> const& cta_v_map, // V: CTA val idx -> gmem mode
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd, // traversal stride
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt, // dilation
TMA::DescriptorAuxParams const& aux_params = {})
{
//
// TMA parameter checking
@@ -645,7 +651,8 @@ make_tma_atom_im2col(CopyOp,
upper_padding_whd,
stride_whd,
lower_srt,
stride_srt);
stride_srt,
aux_params);
//
// Construct the Copy_Traits
@@ -697,18 +704,19 @@ template <class CopyOp,
class DilationStride>
CUTE_HOST_RTC
auto
make_tma_copy_im2col(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
Layout<TShape,TStride> const& cta_t_map, // CTA tid -> logical TMA tid
Layout<VShape,VStride> const& cta_v_map, // CTA vid -> gmem coord
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd, // traversal stride
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt) // dilation
make_tma_copy_im2col(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
Layout<TShape,TStride> const& cta_t_map, // CTA tid -> logical TMA tid
Layout<VShape,VStride> const& cta_v_map, // CTA vid -> gmem coord
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd, // traversal stride
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt, // dilation
TMA::DescriptorAuxParams const& aux_params = {})
{
//
// TMA parameter checking
@@ -719,7 +727,7 @@ make_tma_copy_im2col(CopyOp const& copy_op,
Copy_Atom atom = make_tma_atom_im2col(copy_op, gtensor, slayout, cosize(cta_t_map), cta_v_map,
lower_corner_whd, upper_corner_whd, lower_padding_whd,
upper_padding_whd, stride_whd, lower_srt, stride_srt);
upper_padding_whd, stride_whd, lower_srt, stride_srt, aux_params);
//
// Construct the TiledCopy
include/cute/atom/copy_traits_sm90_tma.hpp
+186 -10
View File
@@ -124,17 +124,24 @@ struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
// Construct an executable SM90_TMA_LOAD with tma_mbar
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
with(uint64_t& tma_mbar, [[maybe_unused]] uint16_t const& multicast_mask = 0) const {
with(
uint64_t& tma_mbar,
[[maybe_unused]] uint16_t const& multicast_mask = 0,
TMA::CacheHintSm90 const& cache_hint = TMA::CacheHintSm90::EVICT_NORMAL) const {
// We accept multicast_mask here to keep the API for both atoms consistent
return {{}, {&tma_desc_, &tma_mbar}};
return {{}, {&tma_desc_, &tma_mbar, static_cast<uint64_t>(cache_hint)}};
}
// Construct an executable SM90_TMA_LOAD with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
with(TmaDescriptor const* new_tma_desc, uint64_t& tma_mbar, [[maybe_unused]] uint16_t const& multicast_mask = 0) const {
with(
TmaDescriptor const* new_tma_desc,
uint64_t& tma_mbar,
[[maybe_unused]] uint16_t const& multicast_mask = 0,
TMA::CacheHintSm90 const& cache_hint = TMA::CacheHintSm90::EVICT_NORMAL) const {
// We accept multicast_mask here to keep the API for both atoms consistent
return {{}, {new_tma_desc, &tma_mbar}};
return {{}, {new_tma_desc, &tma_mbar, static_cast<uint64_t>(cache_hint)}};
}
// Generate the TMA coord tensor
@@ -171,7 +178,8 @@ struct Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
// SM90_TMA_LOAD arguments
tuple<
TmaDescriptor const*,
uint64_t* // smem mbarrier
uint64_t*, // smem mbarrier
uint64_t // cache hint
> const opargs_;
};
@@ -286,6 +294,38 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
///////////////////////////// TMA_STORE //////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////
// Utility for unpacking TMA_STORE arguments into a CopyOp
template <class CopyOp>
struct TMA_STORE_Unpack
{
template <class... Args,
class TS, class SLayout,
class TD, class DLayout>
CUTE_HOST_DEVICE friend constexpr void
copy_unpack(Copy_Traits<CopyOp, Args...> const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst)
{
static_assert(is_smem<TS>::value, "Expected smem src for SM90_TMA_STORE");
void const* const desc_ptr = traits.tma_desc_;
void const* const src_ptr = cute::raw_pointer_cast(src.data());
auto dst_coord = dst.data().coord_;
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
threadIdx.x, threadIdx.y, threadIdx.z,
blockIdx.x, blockIdx.y, blockIdx.z,
int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), src_ptr);
#endif
return detail::explode_tuple(detail::CallCOPY<SM90_TMA_STORE>{},
make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
}
};
struct SM90_TMA_STORE_OP : SM90_TMA_STORE {};
// The executable SM90_TMA_STORE with tma_desc
template <class NumBitsPerTMA, class AuxParams_>
struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
@@ -343,6 +383,30 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
}
// Construct Copy_Traits executable (w/ swapped out TMA descriptor) for SM90_TMA_STORE (for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_STORE_OP, NumBitsPerTMA>
with(TmaDescriptor const* new_tma_desc) const {
return {{}, new_tma_desc};
}
};
// The executable SM90_TMA_STORE with tma_desc
template <class NumBitsPerTMA>
struct Copy_Traits<SM90_TMA_STORE_OP, NumBitsPerTMA>
: TMA_STORE_Unpack<SM90_TMA_STORE_OP>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// SM90_TMA_STORE arguments
TmaDescriptor const* tma_desc_;
};
//////////////////////////////////////////////////////////////////////////////
@@ -1240,14 +1304,14 @@ template <class TmaInternalType = void,
class GEngine, class GLayout,
class SLayout,
class CTA_Tiler,
class Cluster_Size>
class Cluster_Size = Int<1>>
CUTE_HOST_RTC
auto
make_tma_atom(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
CTA_Tiler const& cta_tiler,
Cluster_Size const& cluster_size)
Cluster_Size const& cluster_size = {})
{
auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler);
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
@@ -1283,8 +1347,8 @@ tma_partition(Copy_Atom<Args...> const& copy_atom,
auto layout_V = make_tile(logical_divide(layout_v, tma_layout_v));
// Append with _ until we cover all Rest... modes
auto glayout_V = append<rank_v<decltype(gtensor)>>(layout_V, _);
auto slayout_V = append<rank_v<decltype(stensor)>>(layout_V, _);
auto glayout_V = append<GLayout::rank>(layout_V, _);
auto slayout_V = append<SLayout::rank>(layout_V, _);
// Transform tile mode and coalesce
Tensor gtensor_v = coalesce(gtensor.compose(glayout_V), Shape<Shape<_1,_1>>{}); // ((TMA,TMA_Iter), Rest...)
Tensor stensor_v = coalesce(stensor.compose(slayout_V), Shape<Shape<_1,_1>>{}); // ((TMA,TMA_Iter), Rest...)
@@ -1304,8 +1368,8 @@ tma_partition(Copy_Atom<Args...> const& copy_atom,
// Offset inside the TMA-mode for the multicast
auto multicast_offset = cta_layout(cta_coord) * (size(tma_layout_v) / cosize(cta_layout));
auto multicast_coord = make_coord(make_coord(multicast_offset, Int<0>{}));
auto scoord = append<SLayout::rank>(multicast_coord, Int<0>{});
auto gcoord = append<GLayout::rank>(multicast_coord, Int<0>{});
auto scoord = append<SLayout::rank>(multicast_coord, Int<0>{});
Tensor gresult = domain_offset(gcoord, gtensor_v);
Tensor sresult = domain_offset(scoord, stensor_v);
@@ -1332,4 +1396,116 @@ create_tma_multicast_mask(CtaLayout const& cta_layout_vmnk,
return mcast_mask;
}
////////////////////////////////////
// Make TMA copy A/B/C
///////////////////////////////////
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class CTA_Tiler,
class Cluster_Size>
CUTE_HOST_RTC
auto
make_tma_copy_A_sm90(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
CTA_Tiler const& cta_tiler,
Cluster_Size const& cluster_size)
{
// Keep only MK modes from MNK
auto cta_tiler_mk = remove<1>(cta_tiler);
// mcast along N mode for this M load, if any
auto cluster_size_n = size<1>(cluster_size);
if constexpr (cute::is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL>) {
return make_im2col_tma_copy(copy_op,
gtensor,
slayout,
cta_tiler_mk,
cluster_size_n);
} else {
auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler_mk);
auto cta_t_tile = make_layout(cluster_size_n);
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
auto tma_copy = detail::make_tma_copy_tiled<TmaType>(copy_op, gtensor, slayout, cta_t_tile, cta_v_tile);
return tma_copy;
}
}
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class CTA_Tiler,
class Cluster_Size>
CUTE_HOST_RTC
auto
make_tma_copy_B_sm90(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
CTA_Tiler const& cta_tiler,
Cluster_Size const& cluster_size)
{
// Keep only NK modes from MNK
auto cta_tiler_nk = remove<0>(cta_tiler);
// mcast along M mode for this N load, if any
auto cluster_size_m = size<0>(cluster_size);
if constexpr (cute::is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL>) {
return make_im2col_tma_copy(copy_op,
gtensor,
slayout,
cta_tiler_nk,
cluster_size_m);
} else {
auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler_nk);
auto cta_t_tile = make_layout(cluster_size_m);
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
auto tma_copy = detail::make_tma_copy_tiled<TmaType>(copy_op, gtensor, slayout, cta_t_tile, cta_v_tile);
return tma_copy;
}
}
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class CTA_Tiler>
CUTE_HOST_RTC
auto
make_tma_copy_C_sm90(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
CTA_Tiler const& cta_tiler)
{
// Keep only MN modes from MNK
auto cta_tiler_mn = remove<2>(cta_tiler);
if constexpr (cute::is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL> ||
cute::is_same_v<CopyOp, SM90_TMA_STORE_IM2COL>) {
return make_im2col_tma_copy(copy_op,
gtensor,
slayout,
cta_tiler_mn,
_1{});
} else {
auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler_mn);
// No multicast, so only 1 CTA involved
auto cta_t_map = Layout<_1,_0>{};
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
auto tma_copy = detail::make_tma_copy_tiled<TmaType>(copy_op, gtensor, slayout, cta_t_map, cta_v_tile);
return tma_copy;
}
}
} // end namespace cute
include/cute/atom/mma_atom.hpp
+19 -28
View File
@@ -31,11 +31,9 @@
#pragma once
#include <cute/config.hpp>
#include <cute/arch/mma.hpp>
#include <cute/atom/mma_traits.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/util/type_traits.hpp>
namespace cute {
@@ -102,7 +100,7 @@ struct MMA_Atom<MMA_Traits<Args...>>
static_assert(BLayout::rank == 1, "Expected rank-1 B tensor");
static_assert(CLayout::rank == 1, "Expected rank-1 C tensor");
return mma_unpack(*this, D, A, B, C);
return mma_unpack(static_cast<Traits const&>(*this), D, A, B, C);
}
// Three arguments reproduces C
@@ -245,12 +243,9 @@ struct TiledMMA : MMA_Atom
thrfrg_C(CTensor&& ctensor) const
{
CUTE_STATIC_ASSERT_V(rank(ctensor) >= Int<2>{});
//CUTE_STATIC_ASSERT_V(size<0>(ctensor) % size<0>(TiledShape_MNK{}) == Int<0>{});
//CUTE_STATIC_ASSERT_V(size<1>(ctensor) % size<1>(TiledShape_MNK{}) == Int<0>{});
// Reorder the tensor for the TiledAtom
auto t_tile = make_tile(get<0>(PermutationMNK{}),
get<1>(PermutationMNK{}));
auto t_tile = make_tile(permutation_mnk<0>(),
permutation_mnk<1>());
auto t_tensor = logical_divide(ctensor, t_tile); // (PermM,PermN)
// Tile the tensor for the Atom
@@ -287,12 +282,9 @@ struct TiledMMA : MMA_Atom
thrfrg_A(ATensor&& atensor) const
{
CUTE_STATIC_ASSERT_V(rank(atensor) >= Int<2>{});
//CUTE_STATIC_ASSERT_V(size<0>(atensor) % size<0>(TiledShape_MNK{}) == Int<0>{});
//CUTE_STATIC_ASSERT_V(size<1>(atensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
// Reorder the tensor for the TiledAtom
auto t_tile = make_tile(get<0>(PermutationMNK{}),
get<2>(PermutationMNK{}));
auto t_tile = make_tile(permutation_mnk<0>(),
permutation_mnk<2>());
auto t_tensor = logical_divide(atensor, t_tile); // (PermM,PermK)
// Tile the tensor for the Atom
@@ -329,12 +321,9 @@ struct TiledMMA : MMA_Atom
thrfrg_B(BTensor&& btensor) const
{
CUTE_STATIC_ASSERT_V(rank(btensor) >= Int<2>{});
//CUTE_STATIC_ASSERT_V(size<0>(btensor) % size<1>(TiledShape_MNK{}) == Int<0>{});
//CUTE_STATIC_ASSERT_V(size<1>(btensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
// Reorder the tensor for the TiledAtom
auto t_tile = make_tile(get<1>(PermutationMNK{}),
get<2>(PermutationMNK{}));
auto t_tile = make_tile(permutation_mnk<1>(),
permutation_mnk<2>());
auto t_tensor = logical_divide(btensor, t_tile); // (PermN,PermK)
// Tile the tensor for the Atom
@@ -377,21 +366,23 @@ struct TiledMMA : MMA_Atom
// Utility for printing and visualization
//
// The permutation applied to the MNK-mode data
template <int I>
CUTE_HOST_DEVICE constexpr
auto
permutation_mnk() const {
static_assert(0 <= I && I < 3);
auto perm = get<I>(PermutationMNK{});
return conditional_return(is_underscore<decltype(perm)>{}, size<I>(AtomShape_MNK{}) * size<I+1>(get_thr_layout_vmnk()), perm);
}
// The size of the MNK-mode
template <int I>
CUTE_HOST_DEVICE constexpr
auto
tile_size_mnk() const {
static_assert(0 <= I && I < 3);
auto core_size = size<I>(AtomShape_MNK{}) * size<I+1>(get_thr_layout_vmnk());
[[maybe_unused]] auto perm_size = size<I>(PermutationMNK{});
if constexpr (is_underscore<decltype(perm_size)>::value) {
return core_size;
} else {
return cute::max(core_size, perm_size);
}
CUTE_GCC_UNREACHABLE;
return size(permutation_mnk<I>());
}
CUTE_HOST_DEVICE constexpr
include/cute/atom/mma_traits.hpp
+1 -1
View File
@@ -32,7 +32,7 @@
#include <cute/arch/mma.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
namespace cute
{
include/cute/atom/mma_traits_sm90_gmma.hpp
+5 -3
View File
@@ -332,9 +332,6 @@ struct DescriptorIterator
{
return { GmmaDescriptor{desc_ + uint64_t(offset)} };
}
CUTE_HOST_DEVICE friend void
print(DescriptorIterator) { printf("GMMA::DescriptorIterator"); }
};
template <class T>
@@ -353,6 +350,11 @@ recast_ptr(DescriptorIterator const& iter) {
return iter; // Do nothing, it will still dereference to GmmaDescriptor and decay to uint64_t
}
CUTE_HOST_DEVICE void
print(DescriptorIterator) {
printf("GMMA::DescriptorIterator");
}
// The GMMA Traits below have custom fragment type flags for their smem desc tensors.
// These flags specialize a MakeTensor customization point to correctly make the fragment that is desired.
template <GMMA::Major>
include/cute/container/alignment.hpp
+1 -1
View File
@@ -44,7 +44,7 @@ CUTE_HOST_DEVICE constexpr
bool
is_byte_aligned(void const* const ptr)
{
static_assert(N > 0 && (N & (N - 1)) == 0, "N must be a power of 2 in alignment check");
static_assert(has_single_bit(N), "N must be a power of 2 in alignment check");
return (reinterpret_cast<uintptr_t>(ptr) & (N-1)) == 0;
}
include/cute/container/array_subbyte.hpp
+46 -57
View File
@@ -205,18 +205,22 @@ struct subbyte_iterator
private:
template <class, class> friend struct swizzle_ptr;
template <class U> friend CUTE_HOST_DEVICE constexpr U* raw_pointer_cast(subbyte_iterator<U> const&);
template <class N, class U> friend CUTE_HOST_DEVICE constexpr auto recast_ptr(subbyte_iterator<U> const&);
template <class U> friend CUTE_HOST_DEVICE void print(subbyte_iterator<U> const&);
// Pointer to storage element
storage_type* ptr_ = nullptr;
storage_type* ptr_;
// Bit index of value_type starting position within storage_type element.
// RI: 0 <= idx_ < sizeof_bit<storage_type>
uint8_t idx_ = 0;
uint8_t idx_;
public:
// Ctor
subbyte_iterator() = default;
// Default Ctor
CUTE_HOST_DEVICE constexpr
subbyte_iterator() : ptr_{nullptr}, idx_{0} {};
// Ctor
template <class PointerType>
@@ -286,43 +290,48 @@ public:
return x.ptr_ == y.ptr_ && x.idx_ == y.idx_;
}
CUTE_HOST_DEVICE constexpr friend
bool operator!=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x == y); }
CUTE_HOST_DEVICE constexpr friend
bool operator< (subbyte_iterator const& x, subbyte_iterator const& y) {
return x.ptr_ < y.ptr_ || (x.ptr_ == y.ptr_ && x.idx_ < y.idx_);
}
CUTE_HOST_DEVICE constexpr friend
bool operator!=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x == y); }
CUTE_HOST_DEVICE constexpr friend
bool operator<=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(y < x); }
CUTE_HOST_DEVICE constexpr friend
bool operator> (subbyte_iterator const& x, subbyte_iterator const& y) { return (y < x); }
CUTE_HOST_DEVICE constexpr friend
bool operator>=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x < y); }
// Conversion to raw pointer with loss of subbyte index
CUTE_HOST_DEVICE constexpr friend
T* raw_pointer_cast(subbyte_iterator const& x) {
assert(x.idx_ == 0);
return reinterpret_cast<T*>(x.ptr_);
}
// Conversion to NewT_ with possible loss of subbyte index
template <class NewT_>
CUTE_HOST_DEVICE constexpr friend
auto recast_ptr(subbyte_iterator const& x) {
using NewT = conditional_t<(is_const_v<T>), NewT_ const, NewT_>;
if constexpr (cute::is_subbyte_v<NewT>) { // Making subbyte_iter, preserve the subbyte idx
return subbyte_iterator<NewT>(x.ptr_, x.idx_);
} else { // Not subbyte, assume/assert subbyte idx 0
return reinterpret_cast<NewT*>(raw_pointer_cast(x));
}
CUTE_GCC_UNREACHABLE;
}
CUTE_HOST_DEVICE friend void print(subbyte_iterator x) {
printf("subptr[%db](%p.%u)", int(sizeof_bits_v<T>), x.ptr_, x.idx_);
}
};
// Conversion to raw pointer with loss of subbyte index
template <class T>
CUTE_HOST_DEVICE constexpr
T*
raw_pointer_cast(subbyte_iterator<T> const& x) {
assert(x.idx_ == 0);
return reinterpret_cast<T*>(x.ptr_);
}
// Conversion to NewT_ with possible loss of subbyte index
template <class NewT_, class T>
CUTE_HOST_DEVICE constexpr
auto
recast_ptr(subbyte_iterator<T> const& x) {
using NewT = conditional_t<(is_const_v<T>), NewT_ const, NewT_>;
if constexpr (cute::is_subbyte_v<NewT>) { // Making subbyte_iter, preserve the subbyte idx
return subbyte_iterator<NewT>(x.ptr_, x.idx_);
} else { // Not subbyte, assume/assert subbyte idx 0
return reinterpret_cast<NewT*>(raw_pointer_cast(x));
}
CUTE_GCC_UNREACHABLE;
}
template <class T>
CUTE_HOST_DEVICE void
print(subbyte_iterator<T> const& x) {
printf("subptr[%db](%p.%u)", int(sizeof_bits_v<T>), x.ptr_, x.idx_);
}
//
// array_subbyte
// Statically sized array for non-byte-aligned data types
@@ -365,17 +374,6 @@ private:
public:
constexpr
array_subbyte() = default;
CUTE_HOST_DEVICE constexpr
array_subbyte(array_subbyte const& x) {
CUTE_UNROLL
for (size_type i = 0; i < StorageElements; ++i) {
storage[i] = x.storage[i];
}
}
CUTE_HOST_DEVICE constexpr
size_type size() const {
return N;
@@ -448,25 +446,16 @@ public:
return at(N-1);
}
// In analogy to std::vector<bool>::data(), these functions are deleted to prevent bugs.
// Instead, prefer
// auto* data = raw_pointer_cast(my_subbyte_array.begin());
// where the type of auto* is implementation-defined and
// with the knowledge that [data, data + my_subbyte_array.size()) may not be a valid range.
CUTE_HOST_DEVICE constexpr
pointer data() {
return reinterpret_cast<pointer>(storage);
}
pointer data() = delete;
CUTE_HOST_DEVICE constexpr
const_pointer data() const {
return reinterpret_cast<const_pointer>(storage);
}
CUTE_HOST_DEVICE constexpr
storage_type* raw_data() {
return storage;
}
CUTE_HOST_DEVICE constexpr
storage_type const* raw_data() const {
return storage;
}
const_pointer data() const = delete;
CUTE_HOST_DEVICE constexpr
iterator begin() {
include/cute/container/packed_tuple.hpp
+254
View File
@@ -0,0 +1,254 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/container/type_list.hpp>
namespace cute {
namespace detail {
// Empty Structure Optimization
template <bool IsFirstEmpty, bool IsRestEmpty, class... T>
struct ESO;
template <class First, class... Rest>
static constexpr bool is_first_empty_v = cute::is_empty<First>::value;
template <class First, class... Rest>
static constexpr bool is_rest_empty_v = (cute::is_empty<Rest>::value && ...);
template <class... T>
using ESO_t = ESO<is_first_empty_v<T...>, is_rest_empty_v<T...>, T...>;
// Empty First and Empty Rest...
template <class First, class... Rest>
struct ESO<true, true, First, Rest...> {
CUTE_HOST_DEVICE constexpr
ESO() {}
CUTE_HOST_DEVICE constexpr
ESO(First const&, Rest const&...) {}
};
// NonEmpty First and Empty Rest...
template <class First, class... Rest>
struct ESO<false, true, First, Rest...> {
CUTE_HOST_DEVICE constexpr
ESO() : first_{} {}
CUTE_HOST_DEVICE constexpr
ESO(First const& first, Rest const&...) : first_{first} {}
First first_;
};
// Empty First and NonEmpty Rest...
template <class First, class... Rest>
struct ESO<true, false, First, Rest...> {
CUTE_HOST_DEVICE constexpr
ESO() : rest_{} {}
CUTE_HOST_DEVICE constexpr
ESO(First const&, Rest const&... rest) : rest_{rest...} {}
ESO_t<Rest...> rest_;
};
// NonEmpty T and NonEmpty Rest...
template <class First, class... Rest>
struct ESO<false, false, First, Rest...> {
CUTE_HOST_DEVICE constexpr
ESO() : first_{}, rest_{} {}
CUTE_HOST_DEVICE constexpr
ESO(First const& first, Rest const&... rest) : first_{first}, rest_{rest...} {}
First first_;
ESO_t<Rest...> rest_;
};
// Get Nth value from ESO
template <size_t N, class T, class... Rest, bool F, bool R>
CUTE_HOST_DEVICE constexpr decltype(auto) getv(ESO<F, R, T, Rest...> const& s) {
if constexpr (N == 0) {
if constexpr (F) { return T{}; }
else { return static_cast<T const&>(s.first_); }
} else {
if constexpr (R) { return cute::tuple_element_t<N-1, cute::type_list<Rest...>>{}; }
else { return getv<N-1>(s.rest_); }
}
}
template <size_t N, class T, class... Rest, bool F, bool R>
CUTE_HOST_DEVICE constexpr decltype(auto) getv(ESO<F, R, T, Rest...>& s) {
if constexpr (N == 0) {
if constexpr (F) { return T{}; }
else { return static_cast<T&>(s.first_); }
} else {
if constexpr (R) { return cute::tuple_element_t<N-1, cute::type_list<Rest...>>{}; }
else { return getv<N-1>(s.rest_); }
}
}
template <size_t N, class T, class... Rest, bool F, bool R>
CUTE_HOST_DEVICE constexpr decltype(auto) getv(ESO<F, R, T, Rest...>&& s) {
if constexpr (N == 0) {
if constexpr (F) { return T{}; }
else { return static_cast<T&&>(s.first_); }
} else {
if constexpr (R) { return cute::tuple_element_t<N-1, cute::type_list<Rest...>>{}; }
else { return getv<N-1>(static_cast<ESO_t<Rest...>&&>(s.rest_)); }
}
}
// findt: Implementation detail of cute::find.
// If X is the first template argument of the tuple, findt returns C<N>.
template <class X, size_t N,
bool IsFirstEmpty, bool IsRestEmpty, class First, class... Rest>
CUTE_HOST_DEVICE constexpr
auto
findt(ESO<IsFirstEmpty, IsRestEmpty, First, Rest...> const& t) noexcept
{
if constexpr (cute::is_same_v<X, First>) {
return C<N>{};
}
else {
static_assert(sizeof...(Rest) != 0,
"The type does not appear in the argument list of the tuple.");
if constexpr (IsRestEmpty) {
// The rest is empty, so creating an instance of it is cheap.
return cute::detail::findt<X, N+1>(ESO_t<Rest...>{});
}
else {
return cute::detail::findt<X, N+1>(t.rest_);
}
}
}
} // end namespace detail
// packed_tuple<T...> is a tuple type that is a standard-layout type
// whenever all of its template arguments are standard layout types:
// (cute::is_standard_layout_v<T> && ...) implies (cute::is_standard_layout_v<packed_tuple<T...>>)
template <class... T>
struct packed_tuple : detail::ESO_t<T...>
{
CUTE_HOST_DEVICE constexpr
packed_tuple() {}
CUTE_HOST_DEVICE constexpr
packed_tuple(T const&... ts)
: detail::ESO_t<T...>(ts...)
{}
};
template <>
struct packed_tuple<> {};
template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr
decltype(auto)
get(packed_tuple<T...> const& t) {
static_assert(I < sizeof...(T), "Index out of range");
return detail::getv<I>(t);
}
template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr
decltype(auto)
get(packed_tuple<T...>& t) {
static_assert(I < sizeof...(T), "Index out of range");
return detail::getv<I>(t);
}
template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr
decltype(auto)
get(packed_tuple<T...>&& t) {
static_assert(I < sizeof...(T), "Index out of range");
return detail::getv<I>(static_cast<detail::ESO_t<T...>&&>(t));
}
template <class... T>
CUTE_HOST_DEVICE constexpr
packed_tuple<T...>
make_packed_tuple(T const&... t)
{
return {t...};
}
// Returns the position of type X (as a static integer) in the tuple
// type's argument list. X must be unique in the argument list.
template <class X, class... T>
CUTE_HOST_DEVICE constexpr
auto
find(packed_tuple<T...> const& t) noexcept
{
return detail::findt<X, 0>(t);
}
} // end namespace cute
namespace CUTE_STL_NAMESPACE
{
template <class... T>
struct tuple_size<cute::packed_tuple<T...>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
{};
template <size_t I, class... T>
struct tuple_element<I, cute::packed_tuple<T...>>
: CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>
{};
} // end namespace CUTE_STL_NAMESPACE
#ifdef CUTE_STL_NAMESPACE_IS_CUDA_STD
namespace std {
template <class ... T>
struct tuple_size<cute::packed_tuple<T...>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
{};
template <size_t I, class ... T>
struct tuple_element<I, cute::packed_tuple<T...>>
: CUTE_STL_NAMESPACE::tuple_element<I, cute::packed_tuple<T...>>
{};
} // end namespace std
#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD
include/cute/container/tuple.hpp
+58 -43
View File
@@ -36,10 +36,13 @@
#include <cute/numeric/integer_sequence.hpp>
#include <cute/container/cuda_types.hpp>
#include <cute/container/type_list.hpp>
#if defined(CUTLASS_USE_PACKED_TUPLE)
# include <cute/container/packed_tuple.hpp>
#endif
//#include <cute/container/array.hpp> // Advanced optimizations
//
// cute::tuple is like std::tuple, with two differences.
//
// 1. It works on both host and device.
@@ -50,19 +53,30 @@
// but do _not_ include references like int& or float&.
// (See std::tie for an example of a tuple of references.)
//
// This is simplified over the implementations in std::, cuda::std::, and thrust:: by ignoring much of
// the conversion SFINAE, special overloading, and avoiding cvref template types.
// Furthermore, the empty base optimization (EBO) is MORE aggressive by avoiding
// construction calls, and ignoring any need for unique element addresses.
//
// Over standard-conforming tuple implementations, this appears to accelerate compilation times by over 3x.
// If the template arguments of cute::tuple are all empty types (in
// the sense of std::is_empty_v), then the cute::tuple is also an
// empty type. Furthermore, if CUTLASS_USE_PACKED_TUPLE is defined,
// cute::tuple is always a standard-layout type if all of its template
// arguments are standard-layout types.
namespace cute
{
#if defined(CUTLASS_USE_PACKED_TUPLE)
template<class... T>
using tuple = packed_tuple<T...>;
#else
namespace detail
{
// This is simplified over the implementations in std::, cuda::std::, and thrust:: by ignoring much of
// the conversion SFINAE, special overloading, and avoiding cvref template types.
//
// Over standard-conforming tuple implementations, this appears to accelerate compilation times by over 3x.
// EBO stands for "empty base optimization."
// We use this technique to ensure that cute::tuple
// doesn't need to waste space storing any template arguments
@@ -70,6 +84,12 @@ namespace detail
// Otherwise, cute::tuple would need to spend at least 1 byte
// for each of its template arguments.
//
// This is one way in which cute::tuple differs from std::tuple.
// Empty types in the template argument list are not even constructed,
// and do not have unique element addresses. In fact, they are not
// even members of the tuple or stored in any way. Calling `get`
// constructs and returns an instance of an empty type on demand.
//
// EBO always "holds" a single value of type T.
// N is like an array index that TupleBase uses
// to access the desired tuple element.
@@ -109,9 +129,8 @@ struct EBO<N, T, false>
CUTE_HOST_DEVICE constexpr
EBO() : t_{} {}
template <class U>
CUTE_HOST_DEVICE constexpr
EBO(U const& u) : t_{u} {}
EBO(T const& t) : t_{t} {}
T t_;
};
@@ -141,15 +160,8 @@ struct TupleBase<index_sequence<I...>, T...>
CUTE_HOST_DEVICE constexpr
TupleBase() {}
template <class... U>
CUTE_HOST_DEVICE constexpr explicit
TupleBase(U const&... u)
: EBO<I,T>(u)... {}
template <class... U>
CUTE_HOST_DEVICE constexpr
TupleBase(TupleBase<index_sequence<I...>, U...> const& u)
: EBO<I,T>(getv(static_cast<EBO<I,U> const&>(u)))... {}
TupleBase(T const&... t) : EBO<I,T>(t)... {}
};
} // end namespace detail
@@ -172,16 +184,14 @@ struct tuple : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>
CUTE_HOST_DEVICE constexpr
tuple() {}
template <class... U>
CUTE_HOST_DEVICE constexpr
tuple(U const&... u) : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>(u...) {}
template <class... U>
CUTE_HOST_DEVICE constexpr
tuple(tuple<U...> const& u)
: detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>(static_cast<detail::TupleBase<make_index_sequence<sizeof...(U)>, U...> const&>(u)) {}
tuple(T const&... t) : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>(t...) {}
};
template <>
struct tuple<>
{};
//
// get for cute::tuple (just like std::get for std::tuple)
//
@@ -227,6 +237,8 @@ find(tuple<T...> const& t) noexcept
return detail::findt<X>(t);
}
#endif // CUTLASS_USE_PACKED_TUPLE
//
// Custom is_tuple trait simply checks the existence of tuple_size
// and assumes std::get<I>(.), std::tuple_element<I,.>
@@ -242,6 +254,9 @@ auto has_tuple_size(...) -> false_type;
template <class T>
struct is_tuple : decltype(detail::has_tuple_size((T*)0)) {};
template<typename T>
constexpr bool is_tuple_v = cute::is_tuple<T>::value;
//
// make_tuple (value-based implementation)
//
@@ -540,20 +555,12 @@ tuple_cat(Tuples const&... ts)
namespace detail {
template <size_t I, class TupleA, class TupleB>
template <class TupleA, class TupleB, size_t... I>
CUTE_HOST_DEVICE constexpr
auto
equal_impl(TupleA const& a, TupleB const& b)
equal_impl(TupleA const& a, TupleB const& b, index_sequence<I...>)
{
if constexpr (I == tuple_size<TupleA>::value) {
return cute::true_type{}; // Terminal: TupleA is exhausted
} else if constexpr (I == tuple_size<TupleB>::value) {
return cute::false_type{}; // Terminal: TupleA is not exhausted, TupleB is exhausted
} else {
return (get<I>(a) == get<I>(b)) && equal_impl<I+1>(a,b);
}
CUTE_GCC_UNREACHABLE;
return (cute::true_type{} && ... && (get<I>(a) == get<I>(b)));
}
} // end namespace detail
@@ -564,7 +571,13 @@ CUTE_HOST_DEVICE constexpr
auto
operator==(TupleT const& t, TupleU const& u)
{
return detail::equal_impl<0>(t, u);
if constexpr (tuple_size<TupleT>::value == tuple_size<TupleU>::value) {
return detail::equal_impl(t, u, make_index_sequence<tuple_size<TupleT>::value>{});
} else {
return cute::false_type{};
}
CUTE_GCC_UNREACHABLE;
}
template <class TupleT, class TupleU,
@@ -618,19 +631,17 @@ operator!=(TupleT const& t, TupleU const& u)
namespace detail {
template <class Tuple, size_t... Is>
CUTE_HOST_DEVICE void print_tuple(Tuple const& t,
index_sequence<Is...>, char s = '(', char e = ')')
CUTE_HOST_DEVICE void print_tuple(Tuple const& t, index_sequence<Is...>, char s = '(', char e = ')')
{
using cute::print;
((void(print(Is == 0 ? s : ',')), void(print(get<Is>(t)))), ...); print(e);
print(s); ((void(print(Is == 0 ? '\0' : ',')), void(print(get<Is>(t)))), ...); print(e);
}
#if !defined(__CUDACC_RTC__)
template <class Tuple, std::size_t... Is>
CUTE_HOST std::ostream& print_tuple_os(std::ostream& os, Tuple const& t,
index_sequence<Is...>, char s = '(', char e = ')')
CUTE_HOST std::ostream& print_tuple_os(std::ostream& os, Tuple const& t, index_sequence<Is...>, char s = '(', char e = ')')
{
(void(os << (Is == 0 ? s : ',') << get<Is>(t)), ...);
os << s; (void(os << (Is == 0 ? '\0' : ',') << get<Is>(t)), ...);
return os << e;
}
#endif // !defined(__CUDACC_RTC__)
@@ -655,6 +666,8 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, Tuple const& t)
} // end namespace cute
#if ! defined(CUTLASS_USE_PACKED_TUPLE)
namespace CUTE_STL_NAMESPACE
{
@@ -716,5 +729,7 @@ struct tuple_element<I, const cute::tuple<T...>>
: CUTE_STL_NAMESPACE::tuple_element<I, const CUTE_STL_NAMESPACE::tuple<T...>>
{};
} // end namepsace std
} // end namespace std
#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD
#endif // CUTLASS_USE_PACKED_TUPLE
include/cute/container/type_list.hpp
+23 -34
View File
@@ -30,19 +30,24 @@
**************************************************************************************************/
#pragma once
#include <cute/numeric/integral_constant.hpp>
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
namespace cute
{
template <class T>
struct type_c {
using type = T;
};
template <class... T>
struct type_list {};
// get<I> for type_list<T...>
// requires tuple_element_t<I,type_list<T...>> to have std::is_default_constructible
template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr
CUTE_STL_NAMESPACE::tuple_element_t<I, type_list<T...>>
get(type_list<T...> const& t) noexcept {
return {};
}
} // end namespace cute
//
@@ -55,26 +60,6 @@ struct type_list {};
#include <tuple>
#endif
#include <cute/container/tuple.hpp>
namespace cute
{
template <int I, class... T>
CUTE_HOST_DEVICE constexpr
CUTE_STL_NAMESPACE::tuple_element_t<I, type_list<T...>>
get(type_list<T...>&) noexcept {
return {};
}
template <int I, class... T>
CUTE_HOST_DEVICE constexpr
CUTE_STL_NAMESPACE::tuple_element_t<I, type_list<T...>>
get(type_list<T...> const& t) noexcept {
return {};
}
} // end namespace cute
namespace CUTE_STL_NAMESPACE
{
@@ -85,8 +70,9 @@ struct tuple_size<cute::type_list<T...>>
template <size_t I, class... T>
struct tuple_element<I, cute::type_list<T...>>
: cute::type_c<typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type>
{};
{
using type = typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type;
};
template <class... T>
struct tuple_size<const cute::type_list<T...>>
@@ -95,8 +81,9 @@ struct tuple_size<const cute::type_list<T...>>
template <size_t I, class... T>
struct tuple_element<I, const cute::type_list<T...>>
: cute::type_c<typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type>
{};
{
using type = typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type;
};
} // end namespace std
@@ -119,8 +106,9 @@ struct tuple_size<cute::type_list<T...>>
template <size_t I, class... T>
struct tuple_element<I, cute::type_list<T...>>
: cute::type_c<typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type>
{};
{
using type = typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type;
};
template <class... T>
struct tuple_size<const cute::type_list<T...>>
@@ -129,8 +117,9 @@ struct tuple_size<const cute::type_list<T...>>
template <size_t I, class... T>
struct tuple_element<I, const cute::type_list<T...>>
: cute::type_c<typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type>
{};
{
using type = typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type;
};
} // end namespace std
#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD
include/cute/int_tuple.hpp
+33
View File
@@ -493,6 +493,7 @@ using is_weakly_congruent = decltype(weakly_congruent(declval<A>(), declval<B>()
/** Test if Shape A is compatible with Shape B:
* the size of A and B are the same, and
* any coordinate into A can also be used as a coordinate into B
* Equivalently, the size of Shape B is the same as Shape A at each terminal of Shape A.
* compatible is a partial order on A and B: A <= B
*/
template <class IntTupleA, class IntTupleB>
@@ -523,6 +524,7 @@ using is_compatible = decltype(compatible(declval<A>(), declval<B>()));
/** Test if Shape A is weakly compatible with Shape B:
* there exists a Shape C congruent to A such that compatible(elem_scale(A,C), B)
* Equivalently, the size of Shape B is a multiple of Shape A at each terminal of Shape A.
* weakly_compatible is a partial order on A and B: A <= B
*/
template <class IntTupleA, class IntTupleB>
@@ -551,6 +553,37 @@ weakly_compatible(IntTupleA const& a, IntTupleB const& b)
template <class A, class B>
using is_weakly_compatible = decltype(weakly_compatible(declval<A>(), declval<B>()));
/** Test if Shape A is softly compatible with Shape B:
* there exists a Shape C congruent to A such that compatible(shape_div(A,C), B)
* Equivalently, the size of Shape B divides Shape A at each terminal of Shape A.
* softly_compatible is a partial order on A and B: A <= B
*/
template <class IntTupleA, class IntTupleB>
CUTE_HOST_DEVICE constexpr
auto
softly_compatible(IntTupleA const& a, IntTupleB const& b)
{
if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
if constexpr (tuple_size<IntTupleA>::value != tuple_size<IntTupleB>::value) {
return false_type{};
} else {
return transform_apply(a, b, [](auto const& x, auto const& y) { return softly_compatible(x,y); },
[](auto const&... z) { return (true_type{} && ... && z); });
}
} else if constexpr (is_integral<IntTupleA>::value) {
return a % size(b) == Int<0>{};
} else if constexpr (is_integral<IntTupleB>::value) {
return false_type{};
} else {
return softly_compatible(shape(a), shape(b));
}
CUTE_GCC_UNREACHABLE;
}
template <class A, class B>
using is_softly_compatible = decltype(softly_compatible(declval<A>(), declval<B>()));
/** Replace the elements of Tuple B that are paired with an Int<0> with an Int<1>
*/
template <class IntTupleA, class IntTupleB>
include/cute/layout.hpp
+119 -26
View File
@@ -329,34 +329,23 @@ struct is_layout<Layout<Shape,Stride>> : true_type {};
// Layout construction
//
template <class Shape, class Stride,
__CUTE_REQUIRES((is_tuple<Shape >::value || is_integral<Shape >::value) &&
(is_tuple<Stride>::value || is_integral<Stride>::value))>
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Shape const& shape, Stride const& stride)
{
static_assert(is_tuple<Shape >::value || is_integral<Shape >::value);
static_assert(is_tuple<Stride>::value || is_integral<Stride>::value);
return Layout<Shape,Stride>(shape, stride);
}
template <class Shape,
__CUTE_REQUIRES(is_tuple<Shape>::value || is_integral<Shape>::value)>
template <class Shape>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Shape const& shape)
{
return make_layout(shape, compact_col_major(shape));
}
// Construct a layout from multiple layouts by
// concatenating each layout as an independent mode
template <class... Shapes, class... Strides>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Layout<Shapes,Strides> const&... layouts)
{
return make_layout(make_shape (layouts.shape()...),
make_stride(layouts.stride()...));
static_assert(is_tuple<Shape >::value || is_integral<Shape >::value);
return make_layout(shape, compact_major<LayoutLeft>(shape));
}
//
@@ -366,17 +355,57 @@ make_layout(Layout<Shapes,Strides> const&... layouts)
template <class Shape>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Shape const& shape, GenColMajor)
make_layout(Shape const& shape, LayoutLeft)
{
return make_layout(shape, compact_col_major(shape));
return make_layout(shape, compact_major<LayoutLeft>(shape));
}
template <class Shape>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Shape const& shape, GenRowMajor)
make_layout(Shape const& shape, LayoutRight)
{
return make_layout(shape, compact_row_major(shape));
return make_layout(shape, compact_major<LayoutRight>(shape));
}
//
// Construct a layout from multiple layouts by concatenation
//
// One argument overload
template <class Shape0, class Stride0>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Layout<Shape0,Stride0> const& layout0)
{
return make_layout(make_shape (layout0.shape() ),
make_stride(layout0.stride()));
}
// Two argument overload
template <class Shape0, class Stride0,
class Shape1, class Stride1>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Layout<Shape0,Stride0> const& layout0,
Layout<Shape1,Stride1> const& layout1)
{
return make_layout(make_shape (layout0.shape() , layout1.shape() ),
make_stride(layout0.stride(), layout1.stride()));
}
// Var argument overload
template <class Shape0, class Stride0,
class Shape1, class Stride1,
class... Shapes, class... Strides>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Layout<Shape0,Stride0> const& layout0,
Layout<Shape1,Stride1> const& layout1,
Layout<Shapes,Strides> const&... layouts)
{
return make_layout(make_shape (layout0.shape() , layout1.shape() , layouts.shape()... ),
make_stride(layout0.stride(), layout1.stride(), layouts.stride()...));
}
//
@@ -428,7 +457,7 @@ make_fragment_like(Layout<Shape,Stride> const& layout)
constexpr int R = Layout<Shape,Stride>::rank;
if constexpr (R > 1 && is_static<Shape>::value) {
return tiled_product(make_layout(get<0>(layout.shape()),
compact_col_major(filter_zeros(get<0>(layout.stride()), get<0>(layout.shape())))),
compact_major<LayoutLeft>(filter_zeros(get<0>(layout.stride()), get<0>(layout.shape())))),
make_ordered_layout(take<1,R>(layout.shape()), take<1,R>(layout.stride())));
} else {
return make_layout(layout.shape());
@@ -1131,7 +1160,7 @@ complement(Shape const& shape, Stride const& stride, CoTarget const& cotarget)
// Compute the rest_shape and rest_stride
auto new_stride = get<0>(stride_) * get<0>(shape_); // new stride = min_stride * curr_shape
auto rest_shape = coalesce(ceil_div(cotarget, new_stride));
auto rest_stride = compact_col_major(rest_shape, new_stride);
auto rest_stride = compact_major<LayoutLeft>(rest_shape, new_stride);
// Coalesce and append (rest_shape, rest_stride)
return coalesce(make_layout(make_shape (result_shape , rest_shape ),
@@ -1220,7 +1249,7 @@ right_inverse(Layout<Shape,Stride> const& layout)
return Layout<_1,_0>{}; // Empty case, nothing found
} else {
// Generate the corresponding new strides and construct
auto rstride = compact_col_major(flat_layout.shape());
auto rstride = compact_major<LayoutLeft>(flat_layout.shape());
return make_layout(unwrap(transform(iseq, [&](auto i) { return shape<i>(flat_layout); })),
unwrap(transform(iseq, [&](auto i) { return signum(stride<i>(flat_layout)) * get<i>(rstride); })));
}
@@ -1318,6 +1347,50 @@ max_common_vector(Layout<ShapeA,StrideA> const& a,
CUTE_GCC_UNREACHABLE;
}
/* Return a layout that distributes ShapeB over ShapeA.
*
* @returns Layout result
* @post softly_compatible(@a b, @a result)
* @post For all i,j in [0,size(@a result)) with i < j, @a result(i) < @a result(j). Surjective and Ordered.
* @post composition(make_layout(shape(@a a)), @a result) is admissible
* \code
* // Note that 6 does not divide this shape
* Layout layoutA = Layout<Shape<Int<15>,Int<14>>>{};
*
* // Want to tile any 6 elements and don't care where they come from
* Layout dist = domain_distribute(layoutA, Int<6>{}); // (_3,_2):(_1,_15)
*
* // Not guaranteed to find all 6 though...
* CUTE_STATIC_ASSERT_V(Int<6>{} == size(dist));
*
* Layout result = zipped_divide(layoutA, dist); // (_6,Rest)
* \endcode
*/
template <class ShapeA, class ShapeB>
CUTE_HOST_DEVICE constexpr
auto
domain_distribute(ShapeA const& a, ShapeB const& b)
{
static_assert(is_integral<ShapeB>::value);
static_assert(is_static<ShapeB>::value);
auto flat_shape_a = flatten(shape(a));
static_assert(is_static<decltype(flat_shape_a)>::value);
// Compute the shape of the result
auto [result_shape, b_rest] = cute::fold(flat_shape_a, cute::make_tuple(cute::tuple<>{}, size(b)), [](auto init, auto a_) {
auto [result, b_] = init;
auto gcd_ = gcd(a_, b_);
return cute::make_tuple(append(result, gcd_), b_ / gcd_);
});
// Compute the stride of the result
auto result_stride = compact_major<LayoutLeft>(flat_shape_a);
return coalesce(make_layout(result_shape, result_stride));
}
//
// Kernel (Nullspace) of a Layout
//
@@ -1363,7 +1436,7 @@ nullspace(Layout<Shape,Stride> const& layout)
return Layout<_1,_0>{}; // Empty case, nothing found
} else {
// Generate the corresponding new strides and construct
auto rstride = compact_col_major(flat_layout.shape());
auto rstride = compact_major<LayoutLeft>(flat_layout.shape());
return make_layout(unwrap(transform(iseq, [&](auto i) { return shape<i>(flat_layout); })),
unwrap(transform(iseq, [&](auto i) { return get<i>(rstride); })));
}
@@ -1458,7 +1531,7 @@ auto
ceil_div(Target const& target,
Layout<TShape,TStride> const& tiler)
{
return complement(tiler, size(target));
return shape(complement(tiler, shape(target)));
}
//
@@ -1753,6 +1826,26 @@ recast_layout(Layout<Shape,Stride> const& layout)
CUTE_GCC_UNREACHABLE;
}
// Determine the maximum alignment of a Layout.
// The maximum alignment is the largest N for which upcast<N>(layout) will compile.
// upcast<N>(layout) compiles when the static shapes and strides pass divisibility checks.
// Therefore, upcast<M>(layout) will also compile for all divisors M of N.
// Note that this only considers the static shapes and strides of the Layout
// in symmetry with upcast<N> only checking against static shapes and strides and assuming all
// dynamic shapes and strides are large and multiples of N.
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(Layout<Shape,Stride> const& layout)
{
auto flat_layout = coalesce(layout);
auto static_shape = transform( shape(flat_layout), [](auto s){ return conditional_return<is_static<decltype(s)>::value>(s, Int<1>{}); });
auto static_stride = transform(stride(flat_layout), [](auto d){ return conditional_return<is_static<decltype(d)>::value>(d, Int<0>{}); });
auto filter_layout = make_layout(static_shape, static_stride);
auto permuted = logical_divide(filter_layout, right_inverse(filter_layout));
return gcd(size<0>(permuted), stride<1>(permuted));
}
//
// Display utilities
//
include/cute/layout_composed.hpp
+12
View File
@@ -577,6 +577,7 @@ coalesce(ComposedLayout<A,O,B> const& layout, Shape const& trg_profile)
return composition(layout.layout_a(), layout.offset(), coalesce(layout.layout_b(), trg_profile));
}
//
// Upcast and Downcast
//
@@ -597,6 +598,7 @@ downcast(ComposedLayout<A,O,B> const& layout)
return composition(downcast<N>(layout.layout_a()), downcast<N>(layout.offset()), downcast<N>(layout.layout_b()));
}
template <class OldType, class NewType,
class A, class O, class B>
CUTE_HOST_DEVICE constexpr
@@ -619,6 +621,16 @@ recast_layout(ComposedLayout<A,O,B> const& layout)
CUTE_GCC_UNREACHABLE;
}
template <class A, class O, class B>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(ComposedLayout<A,O,B> const& layout)
{
// Do not attempt for general ComposedLayouts
//return gcd(max_alignment(layout.layout_a()), max_alignment(layout.offset()), max_alignment(layout.layout_b()));
return Int<1>{};
}
//
// Display utilities
//
include/cute/numeric/complex.hpp
+9 -9
View File
@@ -48,13 +48,13 @@ template <class T>
static constexpr auto is_complex_v = is_complex<T>::value;
/// Fused multiply-add for complex numbers
template <class T>
template <class D, class A, class B, class C>
CUTE_HOST_DEVICE constexpr
void
fma(complex<T> & d,
complex<T> const& a,
complex<T> const& b,
complex<T> const& c)
fma(complex<D> & d,
complex<A> const& a,
complex<B> const& b,
complex<C> const& c)
{
fma(d.real(), a.real(), b.real(), c.real());
fma(d.imag(), a.real(), b.imag(), c.imag());
@@ -63,12 +63,12 @@ fma(complex<T> & d,
}
/// Fused multiply-add for triplets
template <class T>
template <class A, class B, class C>
CUTE_HOST_DEVICE constexpr
void
fma(complex<T> const& a,
complex<T> const& b,
complex<T> & c)
fma(complex<A> const& a,
complex<B> const& b,
complex<C> & c)
{
return fma(c, a, b, c);
}
include/cute/numeric/integral_constant.hpp
+5 -1
View File
@@ -33,6 +33,7 @@
#include "cute/util/print.hpp"
#include "cute/util/type_traits.hpp"
#include "cute/numeric/math.hpp"
#include "cutlass/fast_math.h"
namespace cute
{
@@ -82,8 +83,11 @@ struct is_integral<C<v> > : true_type {};
template <class T, T v>
struct is_integral<integral_constant<T,v>> : true_type {};
// is_static detects if an (abstract) value is defined completely by it's type (no members)
// Register FastDivmod as the integral type
template<>
struct is_integral<cutlass::FastDivmod> : true_type {};
// is_static detects if an (abstract) value is defined completely by its type (no members)
template <class T>
struct is_static : bool_constant<is_empty<remove_cvref_t<T>>::value> {};
include/cute/numeric/math.hpp
+30
View File
@@ -33,6 +33,7 @@
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cutlass/fast_math.h>
namespace cute
{
@@ -323,4 +324,33 @@ log_2(T x) {
return static_cast<int32_t>(bit_width(x)) - 1;
}
template <class IntDiv, class IntMod>
struct DivModReturnType {
IntDiv div_;
IntMod mod_;
CUTE_HOST_DEVICE constexpr
DivModReturnType(IntDiv const& div, IntMod const& mod) : div_(div), mod_(mod) {}
};
// General divmod
template <class CInt0, class CInt1>
CUTE_HOST_DEVICE constexpr
auto
divmod(CInt0 const& a, CInt1 const& b) {
return DivModReturnType{a / b, a % b};
}
// Specialized function with fastDivmod input
template <class CInt>
CUTE_HOST_DEVICE constexpr
auto
divmod(CInt const& a, cutlass::FastDivmod const& b) {
using val_div_type = typename cutlass::FastDivmod::value_div_type;
using val_mod_type = typename cutlass::FastDivmod::value_mod_type;
val_div_type div = 0;
val_mod_type mod = 0;
b(div, mod, a);
return DivModReturnType{div, mod};
}
} // namespace cute
include/cute/stride.hpp
+6 -4
View File
@@ -31,8 +31,9 @@
#pragma once
#include <cute/config.hpp>
#include <cute/int_tuple.hpp>
#include <cute/numeric/int.hpp>
#include <cute/numeric/math.hpp>
namespace cute
{
@@ -79,8 +80,9 @@ crd2idx_itt(CInt const& coord,
return crd2idx(_0{}, get<I0>(shape), get<I0>(stride))
+ (_0{} + ... + crd2idx(_0{}, get<Is>(shape), get<Is>(stride)));
} else { // General case
return crd2idx(coord % product(get<I0>(shape)), get<I0>(shape), get<I0>(stride))
+ crd2idx_itt(coord / product(get<I0>(shape)), shape, stride, seq<Is...>{});
auto [div, mod] = divmod(coord, product(get<I0>(shape)));
return crd2idx(mod, get<I0>(shape), get<I0>(stride))
+ crd2idx_itt(div, shape, stride, seq<Is...>{});
}
CUTE_GCC_UNREACHABLE;
@@ -229,7 +231,7 @@ idx2crd(Index const& idx,
}
} else {
if constexpr (is_tuple<Shape>::value) { // "int" tuple
return idx2crd(idx, shape, compact_col_major(shape));
return transform_leaf(as_arithmetic_tuple(crd2idx(idx, shape, make_basis_like(shape))), identity{});
} else { // "int" "int"
return idx;
}
include/cute/swizzle.hpp
+17 -1
View File
@@ -360,7 +360,7 @@ shiftr(MixedBits<S0,F0> const& m, C<S1> s)
}
//
// upcast and downcast
// Upcast and Downcast
//
template <uint32_t S0, uint32_t F0, auto S1>
@@ -410,6 +410,22 @@ downcast(T const& m)
return m * C<N>{};
}
template <uint32_t S0, uint32_t F0>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(MixedBits<S0,F0> const&)
{
return C<uint32_t(1) << countr_zero(S0 | F0)>{};
}
template <auto v>
CUTE_HOST_DEVICE constexpr
C<v>
max_alignment(C<v> const& c)
{
return c;
}
//
// Convert a Pow2Layout+Coord to a MixedBits
//
include/cute/swizzle_layout.hpp
+32 -7
View File
@@ -36,6 +36,7 @@
#include <cute/layout_composed.hpp>
#include <cute/swizzle.hpp>
#include <cute/pointer_swizzle.hpp> // get_swizzle
/* Specialized functionality for a ComposedLayout of the form
* InvolutionFn o Offset o LayoutB
@@ -56,6 +57,9 @@
namespace cute
{
template <int B, int M, int S, class Offset, class LayoutB>
struct get_swizzle<ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB>> { using type = Swizzle<B,M,S>; };
//
// Constructors
//
@@ -117,7 +121,7 @@ CUTE_HOST_DEVICE constexpr
auto
make_fragment_like(ComposedLayout<Swizzle<B,M,S>,Offset,Layout> const& layout)
{
return detail::transfer_swizzle<B,M,S>(layout.layout_b(), make_fragment_like(layout.layout_b()));
return make_fragment_like(layout.layout_b());
}
//
@@ -441,7 +445,7 @@ recast_layout(Swizzle<B,M,S> const& swizzle)
else if constexpr (scale::num == 1) {
return downcast<scale::den>(swizzle);
}
else if constexpr (scale::den == 1) {
else if constexpr (scale::den == 1) {
return upcast<scale::num>(swizzle);
}
else {
@@ -450,6 +454,24 @@ recast_layout(Swizzle<B,M,S> const& swizzle)
CUTE_GCC_UNREACHABLE;
}
template <int B, int M, int S>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(Swizzle<B,M,S> const&)
{
return Int<M>{};
}
template <int B, int M, int S, class Offset, class LayoutB>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& layout)
{
return gcd(max_alignment(layout.layout_a()),
max_alignment(layout.offset()),
max_alignment(layout.layout_b()));
}
//
// Other operations
//
@@ -485,7 +507,7 @@ max_common_vector(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& a,
Layout<Shape,Stride> const& b)
{
// This assumes that Offset is in the YZ domain of the Swizzle...
return cute::min(Int<(1 << M)>{}, max_common_vector(a.layout_b(), b));
return cute::min(max_common_vector(a.layout_b(), b), Int<(1 << M)>{});
}
template <class Shape, class Stride, int B, int M, int S, class Offset, class LayoutB>
@@ -504,12 +526,15 @@ auto
max_common_vector(ComposedLayout<Swizzle<B0,M0,S0>,Offset0,LayoutB0> const& a,
ComposedLayout<Swizzle<B1,M1,S1>,Offset1,LayoutB1> const& b)
{
auto result = coalesce(composition(a, right_inverse(b)));
// Typical impl is composition(a, right_inverse(b))
// so this is Sw0 o B0 o rinv(Sw1 o B1) = Sw0 o B0 o rinv(B1) o Sw1
auto vec = max_common_vector(a.layout_b(), b.layout_b());
if constexpr (is_constant<1, decltype(stride<0>(result.layout_b()))>::value) {
return shape<0>(result);
// This assumes that Offset is in the YZ domain of the Swizzle...
if constexpr (Swizzle<B0,M0,S0>{} == Swizzle<B1,M1,S1>{}) {
return vec;
} else {
return Int<1>{};
return cute::min(vec, Int<(1 << M0)>{}, Int<(1 << M1)>{});
}
CUTE_GCC_UNREACHABLE;
include/cute/tensor.hpp
+2 -1047
View File
File diff suppressed because it is too large Load Diff
include/cute/tensor_impl.hpp
+1153
View File
File diff suppressed because it is too large Load Diff
include/cute/util/debug.hpp
+2 -2
View File
@@ -120,7 +120,7 @@ print_type(T&&...) {
CUTE_HOST_DEVICE
bool
block(int bid)
block([[maybe_unused]] int bid)
{
#if defined(__CUDA_ARCH__)
return blockIdx.x + blockIdx.y*gridDim.x + blockIdx.z*gridDim.x*gridDim.y == bid;
@@ -131,7 +131,7 @@ block(int bid)
CUTE_HOST_DEVICE
bool
thread(int tid, int bid)
thread([[maybe_unused]] int tid, [[maybe_unused]] int bid)
{
#if defined(__CUDA_ARCH__)
return (threadIdx.x + threadIdx.y*blockDim.x + threadIdx.z*blockDim.x*blockDim.y == tid) && block(bid);
include/cute/util/type_traits.hpp
+13
View File
@@ -85,6 +85,8 @@ using CUTE_STL_NAMESPACE::is_volatile_v;
using CUTE_STL_NAMESPACE::conditional;
using CUTE_STL_NAMESPACE::conditional_t;
using CUTE_STL_NAMESPACE::add_const_t;
using CUTE_STL_NAMESPACE::remove_const_t;
using CUTE_STL_NAMESPACE::remove_cv_t;
using CUTE_STL_NAMESPACE::remove_reference_t;
@@ -107,6 +109,13 @@ using CUTE_STL_NAMESPACE::is_convertible_v;
using CUTE_STL_NAMESPACE::is_same;
using CUTE_STL_NAMESPACE::is_same_v;
using CUTE_STL_NAMESPACE::is_constructible;
using CUTE_STL_NAMESPACE::is_constructible_v;
using CUTE_STL_NAMESPACE::is_default_constructible;
using CUTE_STL_NAMESPACE::is_default_constructible_v;
using CUTE_STL_NAMESPACE::is_standard_layout;
using CUTE_STL_NAMESPACE::is_standard_layout_v;
using CUTE_STL_NAMESPACE::is_arithmetic;
using CUTE_STL_NAMESPACE::is_unsigned;
using CUTE_STL_NAMESPACE::is_unsigned_v;
@@ -131,6 +140,9 @@ using CUTE_STL_NAMESPACE::common_type_t;
using CUTE_STL_NAMESPACE::remove_pointer;
using CUTE_STL_NAMESPACE::remove_pointer_t;
using CUTE_STL_NAMESPACE::alignment_of;
using CUTE_STL_NAMESPACE::alignment_of_v;
// <utility>
using CUTE_STL_NAMESPACE::declval;
@@ -261,4 +273,5 @@ struct conditional_template<false, True, False> {
template <class... U>
using type = False<U...>;
};
} // end namespace cute
include/cutlass/arch/barrier.h
+24 -23
View File
@@ -50,11 +50,11 @@ namespace arch {
// Enumerates the reserved named barriers to avoid potential conflicts
// This enum class specifies the NamedBarriers reserved by CUTLASS.
enum class ReservedNamedBarriers {
EpilogueBarrier = 0,
TransposeBarrier = 1,
TransformBarrier = 2,
StreamkBarrier0 = 3,
StreamkBarrier1 = 4
EpilogueBarrier = 1,
TransposeBarrier = 2,
TransformBarrier = 3,
StreamkBarrier0 = 4,
StreamkBarrier1 = 5
, FirstUserBarrier = StreamkBarrier1 + 1
};
@@ -204,12 +204,12 @@ public:
}
CUTLASS_DEVICE
uint32_t test_wait(uint32_t phase, uint32_t pred=true) const {
bool test_wait(uint32_t phase, uint32_t pred=true) const {
return ClusterBarrier::test_wait(&this->barrier_, phase, pred);
}
CUTLASS_DEVICE
uint32_t try_wait(uint32_t phase) const {
bool try_wait(uint32_t phase) const {
return ClusterBarrier::try_wait(&this->barrier_, phase);
}
@@ -260,8 +260,8 @@ public:
".reg .pred P1; \n\t"
"LAB_WAIT: \n\t"
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2; \n\t"
"@P1 bra.uni DONE; \n\t"
"bra.uni LAB_WAIT; \n\t"
"@P1 bra DONE; \n\t"
"bra LAB_WAIT; \n\t"
"DONE: \n\t"
"}"
:
@@ -273,7 +273,7 @@ public:
}
CUTLASS_DEVICE
static uint32_t test_wait(ValueType const* smem_ptr, uint32_t phase, uint32_t pred) {
static bool test_wait(ValueType const* smem_ptr, uint32_t phase, uint32_t pred) {
#if CUDA_BARRIER_ENABLED
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
uint32_t waitComplete;
@@ -289,7 +289,7 @@ public:
: "=r"(waitComplete)
: "r"(smem_addr), "r"(phase), "r"(pred));
return waitComplete;
return static_cast<bool>(waitComplete);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#endif
@@ -297,7 +297,7 @@ public:
}
CUTLASS_DEVICE
static uint32_t try_wait(ValueType const* smem_ptr, uint32_t phase) {
static bool try_wait(ValueType const* smem_ptr, uint32_t phase) {
#if CUDA_BARRIER_ENABLED
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
uint32_t waitComplete;
@@ -311,7 +311,7 @@ public:
: "=r"(waitComplete)
: "r"(smem_addr), "r"(phase));
return waitComplete;
return static_cast<bool>(waitComplete);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#endif
@@ -323,16 +323,17 @@ public:
static void arrive(ValueType const* smem_ptr, uint32_t cta_id, uint32_t pred) {
#if CUDA_BARRIER_ENABLED
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
".reg .b32 remAddr32;\n\t"
"setp.eq.u32 p, %2, 1;\n\t"
"@p mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t"
"@p mbarrier.arrive.shared::cluster.b64 _, [remAddr32];\n\t"
"}"
:
: "r"(smem_addr), "r"(cta_id), "r"(pred));
if (pred) {
asm volatile(
"{\n\t"
".reg .b32 remAddr32;\n\t"
"mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t"
"mbarrier.arrive.shared::cluster.b64 _, [remAddr32];\n\t"
"}"
:
: "r"(smem_addr), "r"(cta_id));
}
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#endif
include/cutlass/arch/mma_sm75.h
-492
View File
@@ -201,252 +201,6 @@ struct Mma<
}
};
////////////////////////////////////////////////////////////////////////////////
//
// Integer matrix multiply .8816 (8b)
//
////////////////////////////////////////////////////////////////////////////////
/// Matrix multiply-add operation: S32 = S8 * S8 + S32
template <>
struct Mma<
gemm::GemmShape<8, 8, 16>,
32,
int8_t,
layout::RowMajor,
int8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8, 8, 16>;
using ElementA = int8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<int8_t, 4>;
using ElementB = int8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<int8_t, 4>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 2>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm75;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM75_ENABLED)
unsigned const & A = reinterpret_cast<unsigned const &>(a);
unsigned const & B = reinterpret_cast<unsigned const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
/// Matrix multiply-add operation: S32 = U8 * S8 + S32
template <>
struct Mma<
gemm::GemmShape<8, 8, 16>,
32,
uint8_t,
layout::RowMajor,
int8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8, 8, 16>;
using ElementA = uint8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<uint8_t, 4>;
using ElementB = int8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<int8_t, 4>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 2>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm75;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM75_ENABLED)
unsigned const & A = reinterpret_cast<unsigned const &>(a);
unsigned const & B = reinterpret_cast<unsigned const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.u8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
/// Matrix multiply-add operation: S32 = S8 * U8 + S32
template <>
struct Mma<
gemm::GemmShape<8, 8, 16>,
32,
int8_t,
layout::RowMajor,
uint8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8, 8, 16>;
using ElementA = int8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<int8_t, 4>;
using ElementB = uint8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<uint8_t, 4>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 2>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm75;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM75_ENABLED)
unsigned const & A = reinterpret_cast<unsigned const &>(a);
unsigned const & B = reinterpret_cast<unsigned const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile("mma.sync.aligned.m8n8k16.row.col.s8.u8 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
/// Matrix multiply-add operation: S32 = U8 * U8 + S32
template <>
struct Mma<
gemm::GemmShape<8, 8, 16>,
32,
uint8_t,
layout::RowMajor,
uint8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8, 8, 16>;
using ElementA = uint8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<uint8_t, 4>;
using ElementB = uint8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<uint8_t, 4>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 2>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm75;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM75_ENABLED)
unsigned const & A = reinterpret_cast<unsigned const &>(a);
unsigned const & B = reinterpret_cast<unsigned const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.u8.u8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
////////////////////////////////////////////////////////////////////////////////
//
// Integer matrix multiply (8b) with SATURATE
@@ -693,252 +447,6 @@ struct Mma<
}
};
////////////////////////////////////////////////////////////////////////////////
//
// Integer matrix multiply (4b)
//
////////////////////////////////////////////////////////////////////////////////
/// Matrix multiply-add operation: S32 = S4 * S4 + S32
template <>
struct Mma<
gemm::GemmShape<8, 8, 32>,
32,
int4b_t,
layout::RowMajor,
int4b_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = int4b_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<int4b_t, 8>;
using ElementB = int4b_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<int4b_t, 8>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 2>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm75;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM75_ENABLED)
unsigned const & A = reinterpret_cast<unsigned const &>(a);
unsigned const & B = reinterpret_cast<unsigned const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.s4.s4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
/// Matrix multiply-add operation: S32 = U4 * S4 + S32
template <>
struct Mma<
gemm::GemmShape<8, 8, 32>,
32,
uint4b_t,
layout::RowMajor,
int4b_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = uint4b_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<uint4b_t, 8>;
using ElementB = int4b_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<int4b_t, 8>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 2>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm75;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM75_ENABLED)
unsigned const & A = reinterpret_cast<unsigned const &>(a);
unsigned const & B = reinterpret_cast<unsigned const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.u4.s4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
/// Matrix multiply-add operation: S32 = S4 * U4 + S32
template <>
struct Mma<
gemm::GemmShape<8, 8, 32>,
32,
int4b_t,
layout::RowMajor,
uint4b_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = int4b_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<int4b_t, 8>;
using ElementB = uint4b_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<uint4b_t, 8>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 2>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm75;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM75_ENABLED)
unsigned const & A = reinterpret_cast<unsigned const &>(a);
unsigned const & B = reinterpret_cast<unsigned const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.s4.u4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
/// Matrix multiply-add operation: S32 = U4 * U4 + S32
template <>
struct Mma<
gemm::GemmShape<8, 8, 32>,
32,
uint4b_t,
layout::RowMajor,
uint4b_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<8, 8, 32>;
using ElementA = uint4b_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<uint4b_t, 8>;
using ElementB = uint4b_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<uint4b_t, 8>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 2>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm75;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM75_ENABLED)
unsigned const & A = reinterpret_cast<unsigned const &>(a);
unsigned const & B = reinterpret_cast<unsigned const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.u4.u4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
: "=r"(D[0]), "=r"(D[1])
: "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
////////////////////////////////////////////////////////////////////////////////
//
// Integer matrix multiply (4b) - SATURATE
include/cutlass/arch/mma_sm80.h
+4 -764
View File
@@ -535,256 +535,6 @@ struct Mma<
}
};
////////////////////////////////////////////////////////////////////////////////
//
// Matrix Multiply 16816 - S8 input, S32 accumulation
//
////////////////////////////////////////////////////////////////////////////////
/// Matrix multiply-add operation: S32 = S8 * S8 + S32
template <>
struct Mma<
gemm::GemmShape<16,8,16>,
32,
int8_t,
layout::RowMajor,
int8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16,8,16>;
using ElementA = int8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<int8_t, 8>;
using ElementB = int8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<int8_t, 4>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const &B = reinterpret_cast<uint32_t const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, "
"{%7,%8,%9,%10};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]),
"r"(C[3]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
#endif
}
};
/// Matrix multiply-add operation: S32 = U8 * S8 + S32
template <>
struct Mma<
gemm::GemmShape<16,8,16>,
32,
uint8_t,
layout::RowMajor,
int8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16,8,16>;
using ElementA = uint8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<uint8_t, 8>;
using ElementB = int8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<int8_t, 4>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const &B = reinterpret_cast<uint32_t const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.s32.u8.s8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, "
"{%7,%8,%9,%10};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]),
"r"(C[3]));
#else
assert(0);
#endif
}
};
/// Matrix multiply-add operation: S32 = S8 * U8 + S32
template <>
struct Mma<
gemm::GemmShape<16,8,16>,
32,
int8_t,
layout::RowMajor,
uint8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16,8,16>;
using ElementA = int8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<int8_t, 8>;
using ElementB = uint8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<uint8_t, 4>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const &B = reinterpret_cast<uint32_t const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.s32.s8.u8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, "
"{%7,%8,%9,%10};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]),
"r"(C[3]));
#else
assert(0);
#endif
}
};
/// Matrix multiply-add operation: S32 = U8 * U8 + S32
template <>
struct Mma<
gemm::GemmShape<16,8,16>,
32,
uint8_t,
layout::RowMajor,
uint8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16,8,16>;
using ElementA = uint8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<uint8_t, 8>;
using ElementB = uint8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<uint8_t, 4>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const &B = reinterpret_cast<uint32_t const &>(b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.s32.u8.u8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, "
"{%7,%8,%9,%10};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]),
"r"(C[3]));
#else
assert(0);
#endif
}
};
////////////////////////////////////////////////////////////////////////////////
//
// Matrix Multiply 16816 - S8 input, S32 accumulation - SATURATE
@@ -1031,252 +781,6 @@ struct Mma<
}
};
////////////////////////////////////////////////////////////////////////////////
//
// Matrix Multiply 16832 - S8 input, S32 accumulation
//
////////////////////////////////////////////////////////////////////////////////
/// Matrix multiply-add operation: S32 = S8 * S8 + S32
template <>
struct Mma<
gemm::GemmShape<16,8,32>,
32,
int8_t,
layout::RowMajor,
int8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16,8,32>;
using ElementA = int8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<int8_t, 16>;
using ElementB = int8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<int8_t, 8>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, "
"{%8,%9}, {%10,%11,%12,%13};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
"r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]));
#else
assert(0);
#endif
}
};
/// Matrix multiply-add operation: S32 = U8 * S8 + S32
template <>
struct Mma<
gemm::GemmShape<16,8,32>,
32,
uint8_t,
layout::RowMajor,
int8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16,8,32>;
using ElementA = uint8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<uint8_t, 16>;
using ElementB = int8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<int8_t, 8>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.u8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, "
"{%8,%9}, {%10,%11,%12,%13};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
"r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]));
#else
assert(0);
#endif
}
};
/// Matrix multiply-add operation: S32 = S8 * U8 + S32
template <>
struct Mma<
gemm::GemmShape<16,8,32>,
32,
int8_t,
layout::RowMajor,
uint8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16,8,32>;
using ElementA = int8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<int8_t, 16>;
using ElementB = uint8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<uint8_t, 8>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.u8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, "
"{%8,%9}, {%10,%11,%12,%13};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
"r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]));
#else
assert(0);
#endif
}
};
/// Matrix multiply-add operation: S32 = U8 * U8 + S32
template <>
struct Mma<
gemm::GemmShape<16,8,32>,
32,
uint8_t,
layout::RowMajor,
uint8_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16,8,32>;
using ElementA = uint8_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<uint8_t, 16>;
using ElementB = uint8_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<uint8_t, 8>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.u8.u8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, "
"{%8,%9}, {%10,%11,%12,%13};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
"r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]));
#else
assert(0);
#endif
}
};
////////////////////////////////////////////////////////////////////////////////
//
// Matrix Multiply 16832 - S8 input, S32 accumulation - SATURATE
@@ -1310,7 +814,7 @@ struct Mma<
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using Operator = OpMultiplyAddSaturate;
using ArchTag = arch::Sm80;
/// Computes multiply-add
@@ -1430,7 +934,7 @@ struct Mma<
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using Operator = OpMultiplyAddSaturate;
using ArchTag = arch::Sm80;
/// Computes multiply-add
@@ -1523,269 +1027,6 @@ struct Mma<
}
};
////////////////////////////////////////////////////////////////////////////////
//
// Matrix Multiply 16864 - S4 input, S32 accumulation
//
////////////////////////////////////////////////////////////////////////////////
/// Matrix multiply-add operation: S32 = S4 * S4 + S32
template <>
struct Mma<
gemm::GemmShape<16, 8, 64>,
32,
cutlass::int4b_t,
layout::RowMajor,
cutlass::int4b_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16, 8, 64>;
using ElementA = cutlass::int4b_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<cutlass::int4b_t, 32>;
using ElementB = cutlass::int4b_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<cutlass::int4b_t, 16>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, "
"{%8,%9}, {%10,%11,%12,%13};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
"r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
#endif
}
};
/// Matrix multiply-add operation: S32 = U4 * S4 + S32
template <>
struct Mma<
gemm::GemmShape<16, 8, 64>,
32,
cutlass::uint4b_t,
layout::RowMajor,
cutlass::int4b_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16, 8, 64>;
using ElementA = cutlass::uint4b_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<cutlass::uint4b_t, 32>;
using ElementB = cutlass::int4b_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<cutlass::int4b_t, 16>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.u4.s4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, "
"{%8,%9}, {%10,%11,%12,%13};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
"r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
#endif
}
};
/// Matrix multiply-add operation: S32 = S4 * U4 + S32
template <>
struct Mma<
gemm::GemmShape<16, 8, 64>,
32,
cutlass::int4b_t,
layout::RowMajor,
cutlass::uint4b_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16, 8, 64>;
using ElementA = cutlass::int4b_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<cutlass::int4b_t, 32>;
using ElementB = cutlass::uint4b_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<cutlass::uint4b_t, 16>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.s4.u4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, "
"{%8,%9}, {%10,%11,%12,%13};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
"r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
#endif
}
};
/// Matrix multiply-add operation: S32 = U4 * U4 + S32
template <>
struct Mma<
gemm::GemmShape<16, 8, 64>,
32,
cutlass::uint4b_t,
layout::RowMajor,
cutlass::uint4b_t,
layout::ColumnMajor,
int,
layout::RowMajor,
OpMultiplyAdd> {
using Shape = gemm::GemmShape<16, 8, 64>;
using ElementA = cutlass::uint4b_t;
using LayoutA = layout::RowMajor;
using FragmentA = Array<cutlass::uint4b_t, 32>;
using ElementB = cutlass::uint4b_t;
using LayoutB = layout::ColumnMajor;
using FragmentB = Array<cutlass::uint4b_t, 16>;
using ElementC = int;
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using ArchTag = arch::Sm80;
/// Computes multiply-add
CUTLASS_HOST_DEVICE
void operator()(
FragmentC &d,
FragmentA const &a,
FragmentB const &b,
FragmentC const &c
) const {
#if defined(CUTLASS_ARCH_MMA_SM80_ENABLED)
uint32_t const *A = reinterpret_cast<uint32_t const *>(&a);
uint32_t const *B = reinterpret_cast<uint32_t const *>(&b);
int const *C = reinterpret_cast<int const *>(&c);
int *D = reinterpret_cast<int *>(&d);
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.u4.u4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, "
"{%8,%9}, {%10,%11,%12,%13};\n"
: "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
"r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]));
#else
CUTLASS_UNUSED(a);
CUTLASS_UNUSED(b);
CUTLASS_UNUSED(c);
CUTLASS_UNUSED(d);
assert(0);
#endif
}
};
////////////////////////////////////////////////////////////////////////////////
//
// Matrix Multiply 16864 - S4 input, S32 accumulation - SATURATE
@@ -1819,7 +1060,7 @@ struct Mma<
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using Operator = OpMultiplyAddSaturate;
using ArchTag = arch::Sm80;
/// Computes multiply-add
@@ -1947,7 +1188,7 @@ struct Mma<
using LayoutC = layout::RowMajor;
using FragmentC = Array<int, 4>;
using Operator = OpMultiplyAdd;
using Operator = OpMultiplyAddSaturate;
using ArchTag = arch::Sm80;
/// Computes multiply-add
@@ -2261,5 +1502,4 @@ struct Mma<
} // namespace arch
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/arch/mma_sparse_sm80.h
+303 -750
View File
File diff suppressed because it is too large Load Diff
include/cutlass/arch/reg_reconfig.h
+3 -5
View File
@@ -37,12 +37,9 @@
#include "cutlass/cutlass.h"
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDACC_VER_MAJOR__ >= 12))
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL))
#if (defined(__CUDA_ARCH__) &&\
(__CUDA_ARCH__ >= 900) && (__CUDACC_VER_MAJOR__ >= 12) && defined(__CUDA_ARCH_FEAT_SM90_ALL))
#define CUDA_CTA_RECONFIG_ACTIVATED 1
#endif
#else
#define CUDA_CTA_RECONFIG_ACTIVATED 0
#endif
namespace cutlass {
@@ -55,6 +52,7 @@ void warpgroup_reg_alloc(){
asm volatile( "setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount) );
#endif
}
template<uint32_t RegCount>
CUTLASS_DEVICE
void warpgroup_reg_dealloc(){
include/cutlass/array.h
+22
View File
@@ -785,6 +785,24 @@ struct reciprocal_approximate<Array<T, N>> {
}
};
template <typename T, int N>
struct reciprocal_approximate_ftz<Array<T, N>> {
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &lhs) const {
Array<T, N> result;
reciprocal_approximate_ftz<T> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = scalar_op(lhs[i]);
}
return result;
}
};
template <typename T, int N>
struct maximum<Array<T, N>, false> {
@@ -979,6 +997,10 @@ struct minimum<Array<T, N>, true> {
}
};
template <typename T, int N>
struct minimum_with_nan_propagation<Array<T, N>> : minimum<Array<T, N>, true>
{};
template <typename T, int N>
struct negate<Array<T, N>> {
include/cutlass/bfloat16.h
+11
View File
@@ -100,6 +100,17 @@ public:
/// Default constructor
bfloat16_t() = default;
/// Reinterpret cast from CUDA's __nv_bfloat16 type
CUTLASS_HOST_DEVICE
explicit bfloat16_t(__nv_bfloat16 const & x) {
#if defined(__CUDA_ARCH__)
storage = reinterpret_cast<uint16_t const &>(x);
#else
__nv_bfloat16_raw raw(x);
std::memcpy(&storage, &raw.x, sizeof(storage));
#endif
}
/// Floating-point conversion - round toward nearest
CUTLASS_HOST_DEVICE
explicit bfloat16_t(float x) {
include/cutlass/cluster_launch.hpp
+9 -3
View File
@@ -133,7 +133,8 @@ struct ClusterLauncher {
size_t const smem_size,
cudaStream_t cuda_stream,
void const* kernel,
void** kernel_params) {
void** kernel_params,
bool launch_with_pdl = false) {
#if defined(CUTLASS_SM90_CLUSTER_LAUNCH_ENABLED)
if (check_cluster_dims(grid_dims, cluster_dims) != Status::kSuccess) {
CUTLASS_TRACE_HOST("ClusterLauncher: check_cluster_dims() failed. Aborting.");
@@ -152,14 +153,19 @@ struct ClusterLauncher {
launch_config.dynamicSmemBytes = smem_size;
launch_config.stream = cuda_stream;
cudaLaunchAttribute launch_attribute[1];
cudaLaunchAttribute launch_attribute[2];
launch_attribute[0].id = cudaLaunchAttributeClusterDimension;
launch_attribute[0].val.clusterDim.x = cluster_dims.x;
launch_attribute[0].val.clusterDim.y = cluster_dims.y;
launch_attribute[0].val.clusterDim.z = cluster_dims.z;
launch_attribute[1].id = cudaLaunchAttributeProgrammaticStreamSerialization;
launch_attribute[1].val.programmaticStreamSerializationAllowed = 1;
launch_config.numAttrs = launch_with_pdl ? 2 : 1;
launch_config.attrs = launch_attribute;
launch_config.numAttrs = 1;
CUTLASS_TRACE_HOST("ClusterLauncher: Launching GPC_CLUSTER_GRID GridDims = "
"(" << grid_dims.x << ", " << grid_dims.y << ", " << grid_dims.z << "), "
include/cutlass/complex.h
+183 -95
View File
@@ -43,6 +43,7 @@
#include "cutlass/cutlass.h"
#include "cutlass/functional.h"
#include "cutlass/platform/platform.h"
#include "cutlass/real.h"
#include "cutlass/numeric_types.h"
@@ -117,6 +118,18 @@ double const &imag(cuDoubleComplex const &z) { return z.y; }
/// Returns the imaginary part of the complex number
CUTLASS_HOST_DEVICE
double &imag(cuDoubleComplex &z) { return z.y; }
// Returns the conjugate of the complex number
CUTLASS_HOST_DEVICE cuFloatComplex
conj(cuFloatComplex const& z) {
return make_cuFloatComplex(z.x, -z.y);
}
// Returns the conjugate of the complex number
CUTLASS_HOST_DEVICE cuDoubleComplex
conj(cuDoubleComplex const& z) {
return make_cuDoubleComplex(z.x, -z.y);
}
#endif
///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -315,60 +328,90 @@ class complex
#endif
};
// Complex conjugate
template<class T>
CUTLASS_HOST_DEVICE complex<T> conj(complex<T> const& z) {
return {z.real(), -z.imag()};
}
///////////////////////////////////////////////////////////////////////////////////////////////////
//
// Accessors for complex template
//
/// Returns the real part of the complex number
template <typename T>
CUTLASS_HOST_DEVICE T const &real(complex<T> const &z) {
return z.real();
}
// Nonmember real and imag need to work for non-complex numbers too.
// That means cutlass::complex, std::complex, cuda::std::complex, and
// any user-defined complex number type that looks like std::complex.
// It's reasonable to assume that a "complex number type" has
// zero-argument real() and imag() member functions returning
// non-void. While cuFloatComplex and cuDoubleComplex lack those
// member functions, one-argument nonmember real and imag overloads
// for those types are defined above.
/// Returns the real part of the complex number
template <typename T>
CUTLASS_HOST_DEVICE T &real(complex<T> &z) {
return z.real();
}
namespace detail {
/// Returns the imaginary part of the complex number
template <typename T>
CUTLASS_HOST_DEVICE T const &imag(complex<T> const &z) {
return z.imag();
}
template <typename T, typename Enable = void>
struct has_zero_argument_real_member_function :
cutlass::platform::false_type
{};
/// Returns the imaginary part of the complex number
template <typename T>
CUTLASS_HOST_DEVICE T &imag(complex<T> &z) {
return z.imag();
}
struct has_zero_argument_real_member_function<T,
cutlass::platform::enable_if_t<
! cutlass::platform::is_void_v<
decltype(cutlass::platform::declval<T>().real())
>
>
> : cutlass::platform::true_type
{};
/// Returns the real part of the real number
template <typename T>
CUTLASS_HOST_DEVICE T const &real(T const &r) {
return r;
}
constexpr bool has_zero_argument_real_member_function_v =
has_zero_argument_real_member_function<T>::value;
template <typename T, typename Enable = void>
struct has_zero_argument_imag_member_function :
cutlass::platform::false_type
{};
/// Returns the real part of the real number
template <typename T>
CUTLASS_HOST_DEVICE T &real(T &r) {
return r;
}
struct has_zero_argument_imag_member_function<T,
cutlass::platform::enable_if_t<
! cutlass::platform::is_void_v<
decltype(cutlass::platform::declval<T>().imag())
>
>
> : cutlass::platform::true_type
{};
/// Returns the imaginary part of the real number
template <typename T>
CUTLASS_HOST_DEVICE T const &imag(T const &r) {
return T();
}
constexpr bool has_zero_argument_imag_member_function_v =
has_zero_argument_imag_member_function<T>::value;
/// Returns the imaginary part of the complex number
template <typename T>
CUTLASS_HOST_DEVICE T &imag(T &r) {
return T();
}
} // namespace detail
template<typename T>
CUTLASS_HOST_DEVICE auto real(T z) {
if constexpr (detail::has_zero_argument_real_member_function_v<T>) {
return z.real();
} else {
return z;
}
}
template<typename T>
CUTLASS_HOST_DEVICE auto imag(T z) {
if constexpr (detail::has_zero_argument_imag_member_function_v<T>) {
return z.imag();
} else {
// Imaginary part of a non-complex input has the same type as the
// input, and its value is zero. CUTLASS assumes in this case
// that value-initializing T is well-formed and results in zero.
return T{};
}
}
//
// Output operators
//
@@ -395,10 +438,36 @@ std::ostream &operator<<(std::ostream &out, complex<T> const &z) {
// Non-member functions defined for complex numbers
//
/// Returns the magnitude of the complex number
// abs returns the magnitude of the complex number.
CUTLASS_HOST_DEVICE float abs(complex<float> const &z) {
return ::hypot(z.real(), z.imag());
}
CUTLASS_HOST_DEVICE double abs(complex<double> const &z) {
return ::hypot(z.real(), z.imag());
}
// In theory, it would make sense to add a complex<long double>
// specialization of abs here, since hypot works for long double too.
// In practice, long double doesn't have a portable number of bits or
// behavior, so users who care about higher-precision floating-point
// computation should probably insist on an actual FP128 type.
template <typename T>
CUTLASS_HOST_DEVICE T abs(complex<T> const &z) {
return sqrt(norm(z));
// cutlass::complex permits all kinds of T, including types that
// don't have NaN. For a generic floating-point type with Inf
// and/or NaN, LAPACK's DLAPY2 algorithm would make sense, as it
// would handle issues like avoiding unwarranted overflow if
// z.real() or z.imag() is slightly bigger than the square root of
// the max finite number. That could be a future improvement; for
// now, the code just uses the naive algorithm.
//
// Use the "swap two-step" idiom so that argument-dependent lookup
// can find any CUTLASS-specific overloads.
using cutlass::sqrt;
return sqrt(z.real() * z.real() + z.imag() * z.imag());
}
/// Returns the magnitude of the complex number
@@ -438,67 +507,66 @@ CUTLASS_HOST_DEVICE R norm_accumulate(complex<T> const &z, R const &accumulator)
static_cast<R>(imag(z)) * static_cast<R>(imag(z));
}
CUTLASS_HOST_DEVICE float conj(float const &z) {
namespace detail {
template<class T>
CUTLASS_HOST_DEVICE T conj_impl(T const& z, cutlass::platform::true_type) {
return conj(z);
}
template<class T>
CUTLASS_HOST_DEVICE T conj_impl(T const& z, cutlass::platform::false_type) {
return z;
}
CUTLASS_HOST_DEVICE double conj(double const &z) {
return z;
template<class T>
CUTLASS_HOST_DEVICE T conj_impl(T const& z) {
constexpr bool use_unqualified_conj =
! cutlass::platform::is_arithmetic_v<T> &&
! detail::has_cutlass_conj_v<T> &&
detail::has_unqualified_conj_v<T>;
return conj_impl(z, cutlass::platform::bool_constant<use_unqualified_conj>{});
}
} // namespace detail
CUTLASS_HOST_DEVICE half_t conj(half_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE int32_t conj(int32_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE uint32_t conj(uint32_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE int64_t conj(int64_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE uint64_t conj(uint64_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE int4b_t conj(int4b_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE uint4b_t conj(uint4b_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE bfloat16_t conj(bfloat16_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE uint1b_t conj(uint1b_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE tfloat32_t conj(tfloat32_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE float_e4m3_t conj(float_e4m3_t const& z) {
return z;
}
CUTLASS_HOST_DEVICE float_e5m2_t conj(float_e5m2_t const& z) {
return z;
}
/// Returns the complex conjugate
template <typename T>
CUTLASS_HOST_DEVICE complex<T> conj(complex<T> const &z) {
return complex<T>(real(z), -imag(z));
// Return the complex conjugate of the input.
//
// This MUST be a function and not a function object, because it may
// be common practice for downstream types to define specifically
// cutlass::conj overloads, instead of overloads in their namespace.
//
// As a result of this being a function and not a function object,
// CUTLASS code needs to declare "using cutlass::conj;" in scope and
// then call this function unqualified, just like std::swap.
//
// If an overload already exists for cutlass::conj(T), that overload
// will be called instead of this one. Otherwise:
//
// 1. for arithmetic types, return z;
//
// 2. for types where (namespace-unqualified) conj(z) is well formed
// and cutlass::conj(z) is NOT well formed, return conj(z); and,
//
// 3. for everything else, return z.
//
// Regarding (1), the C++ Standard Library makes std::conj always
// return std::complex, even for (noncomplex) arithmetic types.
// cutlass::conj(T t) needs to return type T. This follows the
// convention of linear algebra software like the BLAS, where
// "conjugate transpose" means the same thing as "transpose" for a
// matrix of noncomplex numbers.
//
// Case (2) covers std::complex, cuda::std::complex, and non-Standard
// (including user-defined) complex number types (for which "conj(z)"
// is findable via argument-dependent lookup, but does not live in the
// cutlass namespace). It excludes cutlass::conj(z) in order to
// prevent infinite recursion.
//
// Case (3) covers non-Standard non-complex number types.
template<class T>
CUTLASS_HOST_DEVICE T conj(T const& z) {
return detail::conj_impl(z);
}
/// Projects the complex number z onto the Riemann sphere
@@ -699,10 +767,30 @@ template <typename T>
struct conjugate<complex<T>> {
CUTLASS_HOST_DEVICE
complex<T> operator()(complex<T> const &a) const {
return conj(a);
// Invoke the complex<T> overload specifically, rather than
// wasting the compiler's effort on overload resolution.
return cutlass::conj(a);
}
};
#if ! defined(__CUDACC_RTC__)
template <>
struct conjugate<cuFloatComplex> {
CUTLASS_HOST_DEVICE
cuFloatComplex operator()(cuFloatComplex const& z) const {
return make_cuFloatComplex(z.x, -z.y);
}
};
template <>
struct conjugate<cuDoubleComplex> {
CUTLASS_HOST_DEVICE
cuDoubleComplex operator()(cuDoubleComplex const& z) const {
return make_cuDoubleComplex(z.x, -z.y);
}
};
#endif
/// Computes the square of a difference with optional conversion
template <typename T, typename Output>
struct magnitude_squared_difference<complex<T>, Output> {
include/cutlass/conv/collective/detail.hpp
+3
View File
@@ -246,6 +246,9 @@ compute_lower_srt(ConvProblemShape<ConvOp, NumSpatialDimensions> const& problem_
return lower;
}
template <class CopyOp> struct is_im2col_load { static constexpr bool value = false; };
template <> struct is_im2col_load<SM90_TMA_LOAD_IM2COL > { static constexpr bool value = true; };
template <> struct is_im2col_load<SM90_TMA_LOAD_IM2COL_MULTICAST> { static constexpr bool value = true; };
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::conv::collective::detail
include/cutlass/conv/collective/sm90_implicit_gemm_gmma_ss_warpspecialized.hpp
+38 -26
View File
@@ -131,6 +131,9 @@ struct CollectiveConv<
&& (cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>)),
"GmemTiledCopyB - invalid SM90 TMA copy atom specified.");
static constexpr bool is_im2col_A = detail::is_im2col_load<GmemTiledCopyA>::value;
static constexpr bool is_im2col_B = detail::is_im2col_load<GmemTiledCopyB>::value;
// TMA converts f32 input to tf32 when copying from GMEM to SMEM
// For all other types, cast to size equivalent uint type to avoid any rounding by TMA.
static constexpr bool ConvertF32toTF32A = cute::is_same_v<float, ElementA>;
@@ -169,12 +172,11 @@ private:
// Note that for fprop and dgrad kernel, the tma load mode is im2col for tensor A and tiled for
// tensor B while for wgrad kernel, the tma load mode is tiled for tensor A and im2col for tensor
// B since operand A, B is swapped.
// Get tma_load_a instantce.
template <class TensorA>
static constexpr auto
get_tma_load_a_instance(TensorA const& tensor_a, typename Arguments::ProblemShape const& problem_shape) {
if constexpr (ConvOp == conv::Operator::kFprop || ConvOp == conv::Operator::kDgrad) {
if constexpr (is_im2col_A) {
// compute the upper and lower corners based on the conv padding
auto lower_corner_whd = detail::compute_lower_corner_whd(problem_shape);
auto upper_corner_whd = detail::compute_upper_corner_whd(problem_shape);
@@ -203,7 +205,7 @@ private:
shape(stride_srt));
}
// TMA tiled mode for tensor A in wgrad kernel.
else if constexpr (ConvOp == conv::Operator::kWgrad) {
else {
return make_tma_copy(
GmemTiledCopyA{},
tensor_a,
@@ -217,16 +219,8 @@ private:
template <class TensorB>
static constexpr auto
get_tma_load_b_instance(TensorB const& tensor_b, typename Arguments::ProblemShape const& problem_shape) {
if constexpr (ConvOp == conv::Operator::kFprop || ConvOp == conv::Operator::kDgrad) {
return make_tma_copy(
GmemTiledCopyB{},
tensor_b,
SmemLayoutB{}(_,_,_0{}),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
size<0>(ClusterShape{}));
}
// TMA im2col mode for tensor B in wgrad kernel.
else if constexpr (ConvOp == conv::Operator::kWgrad) {
if constexpr (is_im2col_B) {
// compute the upper and lower corners based on the conv padding
auto lower_corner_whd = detail::compute_lower_corner_whd(problem_shape);
auto upper_corner_whd = detail::compute_upper_corner_whd(problem_shape);
@@ -246,6 +240,26 @@ private:
shape(lower_srt),
cute::reverse(shape(problem_shape.dilation)));
}
else {
return make_tma_copy(
GmemTiledCopyB{},
tensor_b,
SmemLayoutB{}(_,_,_0{}),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
size<0>(ClusterShape{}));
}
}
static constexpr auto
get_problem_shape_MNKL(typename Arguments::ProblemShape const& problem_shape) {
if constexpr (is_im2col_A || is_im2col_B) {
// transformation + im2col linearization
return problem_shape.get_linearized_problem_shape_MNKL();
}
else {
// transformation
return problem_shape.get_transformed_problem_shape_MNKL();
}
}
public:
@@ -253,9 +267,7 @@ public:
// Device side kernel params
struct Params {
using _Submode = decltype(take<0,NumTensorDimensions-1>(typename Arguments::ProblemShape::TensorExtent{}));
using ProblemShape = cute::conditional_t<DispatchPolicy::ConvOp == conv::Operator::kWgrad,
Shape<int, _Submode, _Submode>,
Shape<_Submode, int, _Submode>>;
using ProblemShape = decltype(get_problem_shape_MNKL(typename Arguments::ProblemShape{}));
// Assumption: StrideA is congruent with Problem_MK
// Select TMA load type according to convolution operator.
@@ -283,6 +295,7 @@ public:
TMA_A tma_load_a;
TMA_B tma_load_b;
ProblemShape problem_shape;
uint32_t tma_transaction_bytes = TmaTransactionBytes;
};
//
@@ -314,17 +327,18 @@ public:
auto tma_load_a = get_tma_load_a_instance(tensor_a, args.problem_shape);
auto tma_load_b = get_tma_load_b_instance(tensor_b, args.problem_shape);
auto problem_shape_mnk = args.problem_shape.get_transformed_problem_shape_MNK();
auto problem_shape_mnkl = get_problem_shape_MNKL(args.problem_shape);
return {
tma_load_a,
tma_load_b,
problem_shape_mnk
problem_shape_mnkl,
TmaTransactionBytes
};
}
template<class ProblemShape>
CUTLASS_HOST_DEVICE static bool
static bool
can_implement(
ProblemShape const& problem_shape,
Arguments const& args) {
@@ -389,13 +403,12 @@ public:
TensorA const& gA, TMA_LOAD_A& tma_load_a,
TensorB const& gB, TMA_LOAD_B& tma_load_b,
KTileIterator k_tile_iter, int k_tile_count,
int thad_idx,
int thread_idx,
uint32_t block_rank_in_cluster,
TensorStorage& shared_tensors) {
int warp_idx = canonical_warp_idx_sync();
int warp_idx_in_warp_group = warp_idx % 4;
int lane_predicate = cute::elect_one_sync();
if (warp_idx_in_warp_group == 0 and lane_predicate) {
if (lane_predicate) {
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
@@ -403,7 +416,8 @@ public:
// Prepare the TMA loads for A and B
//
dim3 cluster_local_block_id = cute::block_id_in_cluster();
constexpr uint32_t cluster_shape_x = get<0>(ClusterShape());
uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
auto block_tma_a = tma_load_a.get_slice(cluster_local_block_id.y);
auto block_tma_b = tma_load_b.get_slice(cluster_local_block_id.x);
@@ -462,12 +476,10 @@ public:
/// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster
CUTLASS_DEVICE void
load_tail(MainloopPipeline pipeline, PipelineState smem_pipe_producer_state) {
int warp_idx = canonical_warp_idx_sync();
int warp_idx_in_warp_group = warp_idx % 4;
int lane_predicate = cute::elect_one_sync();
// Issue the epilogue waits
if (warp_idx_in_warp_group == 0 and lane_predicate) {
if (lane_predicate) {
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* Consumer UNLOCKs), or if the stage was never used
include/cutlass/conv/convnd_problem_shape.hpp
+65 -14
View File
@@ -352,15 +352,16 @@ struct ConvProblemShape {
}
}
// Get problem shape MNK according to following table:
// | | Fprop | Dgrad | Wgrad |
// | ---- | --------- | -------- | -------- |
// | Shape_M | (Q,P,Z,N) | (W,H,D,N) | (K) |
// | Shape_N | (K) | (C) | (C,S,R,T) |
// | Shape_K | (C,S,R,T) | (K,S,R,T) | (Q,P,Z,N) |
// Get problem shape MNKL according to following table:
// | | Fprop | Dgrad | Wgrad |
// | ---- | --------- | -------- | -------- |
// | Shape_M | (Q,P,Z,N) | (W/V,H/U,D/O,N) | (K) |
// | Shape_N | (K) | (C) | (C,S,R,T) |
// | Shape_K | (C,S,R,T) | (K,S,R,T) | (Q,P,Z,N) |
// | Shape_L | _1 | (V,U,O) | _1 |
CUTLASS_HOST_DEVICE
constexpr auto
get_transformed_problem_shape_MNK() const {
get_transformed_problem_shape_MNKL() const {
using cute::insert;
using cute::make_shape;
using cute::reverse;
@@ -370,32 +371,56 @@ struct ConvProblemShape {
auto M_xformed = shape_C[0];
auto N_xformed = reverse(take<1, RankT>(shape_C));
auto K_xformed = reverse(take<0, RankT - 1>(shape_A));
auto L_xformed = cute::Int<1>{};
return make_shape(M_xformed, N_xformed, K_xformed);
return make_shape(M_xformed, N_xformed, K_xformed, L_xformed);
}
else if constexpr (ConvOp == conv::Operator::kFprop){
auto M_xformed = reverse(take<0, RankT - 1>(shape_C));
auto N_xformed = shape_C[RankT - 1];
auto K_xformed = reverse(take<1, RankT>(shape_B));
auto L_xformed = cute::Int<1>{};
return make_shape(M_xformed, N_xformed, K_xformed);
return make_shape(M_xformed, N_xformed, K_xformed, L_xformed);
}
else if constexpr (ConvOp == conv::Operator::kDgrad) {
auto M_xformed = reverse(take<0,RankT - 1>(shape_C));
auto L_xformed = reverse(traversal_stride); // (V,U,O)
auto M_xformed = ceil_div(reverse(take<0,RankT - 1>(shape_C)), L_xformed);
auto N_xformed = shape_C[RankT - 1];
// shape_B: [K,T,R,S,C], K_xformed: [K,S,R,T]
auto K_xformed = insert<0>(
(reverse(take<1,RankT - 1>(shape_B))),
shape_B[0]);
return make_shape(M_xformed, N_xformed, K_xformed);
return make_shape(M_xformed, N_xformed, K_xformed, L_xformed);
}
}
// Assuming im2col linearization
// Get problem shape MNKL according to following table:
// | | Fprop | Dgrad | Wgrad |
// | ---- | --------- | -------- | -------- |
// | Shape_M | (Q*P*Z*N) | ([W/V]*[H/U]*[D/O]*N) | (K) |
// | Shape_N | (K) | (C) | (C,S,R,T) |
// | Shape_K | (C,S,R,T) | (K,S,R,T) | (Q*P*Z*N) |
// | Shape_L | _1 | (V*U*O) | _1 |
CUTLASS_HOST_DEVICE
constexpr auto
get_linearized_problem_shape_MNKL() const {
auto [M, N, K, L] = get_transformed_problem_shape_MNKL();
if constexpr (ConvOp == conv::Operator::kFprop || ConvOp == conv::Operator::kDgrad) {
return cute::make_shape(cute::product(M), N, K, cute::product(L));
}
else if constexpr (ConvOp == conv::Operator::kWgrad) {
return cute::make_shape(M, N, cute::product(K), L);
}
}
// Get A extents.
// fprop: A extents array contains [N,D,H,W,C]. Turn that into ((W,H,D,N), (C))
// wgrad: A extents array contains [N,Z,P,Q,K]. Turn that into ((K), (Q,P,Z,N))
// dgrad: A extents array contains [N,Z,P,Q,K]. Turn that into ((Q,P,Z,N), (K))
// wgrad: A extents array contains [N,Z,P,Q,K]. Turn that into ((K), (Q,P,Z,N))
CUTLASS_HOST_DEVICE
constexpr auto
get_shape_A() const {
@@ -418,8 +443,8 @@ struct ConvProblemShape {
// Get B extents.
// fprop: B extents array contains [K,T,R,S,C]. Turn that into ((K), (C,S,R,T))
// wgrad: B extents array contains [N,D,H,W,C]. Turn that into ((C), (W,H,D,N))
// dgrad: B extents array contains [K,T,R,S,C]. Turn that into ((C), (K,S,R,T))
// wgrad: B extents array contains [N,D,H,W,C]. Turn that into ((C), (W,H,D,N))
CUTLASS_HOST_DEVICE
constexpr auto
get_shape_B() const {
@@ -447,6 +472,30 @@ struct ConvProblemShape {
}
}
// Get C extents.
// fprop: C extents array contains [N,Z,P,Q,K]. Turn that into ((Q,P,Z,N), (K))
// dgrad: C extents array contains [N,D,H,W,C]. Turn that into ((W,H,D,N), (C))
// wgrad: C extents array contains [K,T,R,S,C]. Turn that into ((K), (C,S,R,T))
CUTLASS_HOST_DEVICE
constexpr auto
get_shape_C() const {
using cute::make_shape;
using cute::reverse;
using cute::take;
if constexpr (ConvOp == conv::Operator::kFprop ||
ConvOp == conv::Operator::kDgrad) {
return make_shape(
reverse(take<0, RankT - 1>(shape_C)),
shape_C[RankT - 1]);
}
else if constexpr (ConvOp == conv::Operator::kWgrad) {
return make_shape(
shape_C[0],
reverse(take<1, RankT>(shape_C)));
}
}
// Static method that returns the canonical strides of tensors (layouts are right major and compact)
CUTLASS_HOST_DEVICE
static constexpr TensorStride
@@ -529,7 +578,9 @@ private:
// calculate n,z,p,q,k.
// a helper lambda to compute a single spatial extent of the nzpqk tensor
auto nzpqk_extent = [](int act_ext, int filter_ext, int pad_total, int dilation, int tstride) {
return 1 + (act_ext + pad_total - ((filter_ext -1) * dilation + 1)) / tstride;
auto tmp = act_ext + pad_total - ((filter_ext -1) * dilation + 1);
CUTLASS_ASSERT(tmp % tstride == 0);
return 1 + tmp / tstride;
};
shape_xformed_act[0] = shape_act[0]; // Activation N extent
include/cutlass/conv/device/conv_universal_adapter.hpp
+20 -26
View File
@@ -228,29 +228,18 @@ public:
/// Initializes conv state from arguments.
Status
initialize(
Arguments const& args,
void* workspace = nullptr,
cudaStream_t stream = nullptr,
CudaHostAdapter *cuda_adapter = nullptr) {
Arguments const& args,
void* workspace = nullptr,
cudaStream_t stream = nullptr,
CudaHostAdapter *cuda_adapter = nullptr) {
CUTLASS_TRACE_HOST("ConvUniversal::initialize() - workspace "
<< workspace << ", stream: " << (stream ? "non-null" : "null"));
size_t workspace_bytes = ConvKernel::get_workspace_size(args);
CUTLASS_TRACE_HOST(" workspace_bytes: " << workspace_bytes);
if (workspace_bytes) {
if (!workspace) {
CUTLASS_TRACE_HOST(" error: device workspace must not be null");
return Status::kErrorWorkspaceNull;
}
CUTLASS_TRACE_HOST(" clearing device workspace");
cudaError_t result = cudaMemsetAsync(workspace, 0, workspace_bytes, stream);
if (cudaSuccess != result) {
result = cudaGetLastError(); // to clear the error bit
CUTLASS_TRACE_HOST(" cudaMemsetAsync() returned error " << cudaGetErrorString(result));
return Status::kErrorInternal;
}
// Initialize the workspace
Status status = ConvKernel::initialize_workspace(args, workspace, stream, cuda_adapter);
if (status != Status::kSuccess) {
return status;
}
// Initialize the Params structure
@@ -297,7 +286,7 @@ public:
/// Primary run() entry point API that is static allowing users to create and manage their own params.
/// Supplied params struct must be construct by calling ConvKernel::to_underling_arguments()
static Status
run(Params& params, cudaStream_t stream = nullptr, CudaHostAdapter *cuda_adapter = nullptr) {
run(Params& params, cudaStream_t stream = nullptr, CudaHostAdapter *cuda_adapter = nullptr, int32_t kernel_index = 0) {
CUTLASS_TRACE_HOST("ConvUniversal::run()");
dim3 const block = ConvKernel::get_block_shape();
dim3 const grid = get_grid_shape(params);
@@ -319,9 +308,13 @@ public:
CUTLASS_ASSERT(cuda_adapter);
if (cuda_adapter) {
launch_result = cuda_adapter->launch(
grid, cluster, block, smem_size, stream, kernel_params, 0
);
launch_result = cuda_adapter->launch(grid,
cluster,
block,
smem_size,
stream,
kernel_params,
kernel_index);
}
else {
return Status::kErrorInternal;
@@ -379,11 +372,12 @@ public:
Arguments const& args,
void* workspace = nullptr,
cudaStream_t stream = nullptr,
CudaHostAdapter *cuda_adapter = nullptr
CudaHostAdapter *cuda_adapter = nullptr,
int32_t kernel_index = 0
) {
Status status = initialize(args, workspace, stream, cuda_adapter);
if (Status::kSuccess == status) {
status = run(params_, stream, cuda_adapter);
status = run(params_, stream, cuda_adapter, kernel_index);
}
return status;
}
include/cutlass/conv/device/direct_convolution.h
+1 -1
View File
@@ -197,7 +197,7 @@ public:
params_.ptr_C = args.ref_C.data();
params_.ptr_D = args.ref_D.data();
params_.output_op = args.output_op;
params_.ptr_reordered_B = args.ref_reordered_B.data();;
params_.ptr_reordered_B = args.ref_reordered_B.data();
params_.semaphore = static_cast<int *>(workspace);
return Status::kSuccess;
include/cutlass/conv/dispatch_policy.hpp
+3
View File
@@ -31,6 +31,7 @@
#pragma once
#include "cutlass/conv/convolution.h"
#include "cutlass/epilogue/thread/activation.h"
#include "cutlass/arch/arch.h"
#include "cute/layout.hpp"
@@ -38,6 +39,8 @@
//////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////
namespace cutlass::conv {
//////////////////////////////////////////////////////////////////////////////
include/cutlass/conv/kernel/default_conv2d.h
+9 -2
View File
@@ -114,7 +114,10 @@ template <
typename ElementTensor,
typename ElementVector,
typename OutputOp,
int ElementsPerAccess
int ElementsPerAccess,
typename PermuteDLayout = layout::NoPermute,
conv::StrideSupport StrideSupport = conv::StrideSupport::kUnity,
int Rank = 4
>
struct DefaultConvEpilogueWithBroadcastSimt {
using Epilogue = typename epilogue::threadblock::DefaultEpilogueWithBroadcastSimt<
@@ -124,7 +127,11 @@ struct DefaultConvEpilogueWithBroadcastSimt {
ElementTensor,
ElementVector,
OutputOp,
ElementsPerAccess
ElementsPerAccess,
false,
PermuteDLayout,
StrideSupport,
Rank
>::Epilogue;
};
include/cutlass/conv/kernel/default_conv3d_fprop_with_broadcast.h
+4 -1
View File
@@ -197,7 +197,10 @@ struct DefaultConv3dFpropWithBroadcast <
typename EpilogueOutputOp::ElementT,
typename EpilogueOutputOp::ElementVector,
EpilogueOutputOp,
ImplicitGemmBase::Epilogue::kElementsPerAccess
ImplicitGemmBase::Epilogue::kElementsPerAccess,
layout::NoPermute,
StrideSupport,
5
>::Epilogue;
// Define the kernel
include/cutlass/conv/kernel/default_deconv2d.h
+20 -4
View File
@@ -181,7 +181,11 @@ struct DefaultDeconv2d <
ThreadblockShape,
WarpMmaSimtOp,
EpilogueOutputOp,
EpilogueOutputOp::kCount
EpilogueOutputOp::kCount,
false,
layout::NoPermute,
StrideSupport::kStrided,
4
>::Epilogue;
// Define the kernel
@@ -405,7 +409,11 @@ struct DefaultDeconv2d <
ThreadblockShape,
WarpMmaSimtOp,
EpilogueOutputOp,
EpilogueOutputOp::kCount
EpilogueOutputOp::kCount,
false,
layout::NoPermute,
StrideSupport::kStrided,
4
>::Epilogue;
// Define the kernel
@@ -627,7 +635,11 @@ struct DefaultDeconv2d <
ThreadblockShape,
WarpMmaSimtOp,
EpilogueOutputOp,
EpilogueOutputOp::kCount
EpilogueOutputOp::kCount,
false,
layout::NoPermute,
StrideSupport::kStrided,
4
>::Epilogue;
// Define the kernel
@@ -852,7 +864,11 @@ struct DefaultDeconv2d <
ThreadblockShape,
WarpMmaSimtOp,
EpilogueOutputOp,
EpilogueOutputOp::kCount
EpilogueOutputOp::kCount,
false,
layout::NoPermute,
StrideSupport::kStrided,
4
>::Epilogue;
// Define the kernel
include/cutlass/conv/kernel/default_deconv3d.h
+20 -4
View File
@@ -170,7 +170,11 @@ struct DefaultDeconv3d <
ThreadblockShape,
WarpMmaSimtOp,
EpilogueOutputOp,
EpilogueOutputOp::kCount
EpilogueOutputOp::kCount,
false,
layout::NoPermute,
StrideSupport::kStrided,
5
>::Epilogue;
// Define the kernel
@@ -282,7 +286,11 @@ struct DefaultDeconv3d <
ThreadblockShape,
WarpMmaSimtOp,
EpilogueOutputOp,
EpilogueOutputOp::kCount
EpilogueOutputOp::kCount,
false,
layout::NoPermute,
StrideSupport::kStrided,
5
>::Epilogue;
// Define the kernel
@@ -389,7 +397,11 @@ struct DefaultDeconv3d <
ThreadblockShape,
WarpMmaSimtOp,
EpilogueOutputOp,
EpilogueOutputOp::kCount
EpilogueOutputOp::kCount,
false,
layout::NoPermute,
StrideSupport::kStrided,
5
>::Epilogue;
// Define the kernel
@@ -501,7 +513,11 @@ struct DefaultDeconv3d <
ThreadblockShape,
WarpMmaSimtOp,
EpilogueOutputOp,
EpilogueOutputOp::kCount
EpilogueOutputOp::kCount,
false,
layout::NoPermute,
StrideSupport::kStrided,
5
>::Epilogue;
// Define the kernel
include/cutlass/conv/kernel/default_deconv3d_with_broadcast.h
+9 -3
View File
@@ -196,7 +196,10 @@ struct DefaultDeconv3dWithBroadcast <
typename EpilogueOutputOp::ElementT,
typename EpilogueOutputOp::ElementVector,
EpilogueOutputOp,
ImplicitGemmBase::Epilogue::kElementsPerAccess
ImplicitGemmBase::Epilogue::kElementsPerAccess,
layout::NoPermute,
StrideSupport::kStrided,
5
>::Epilogue;
// Define the kernel
@@ -273,7 +276,7 @@ struct DefaultDeconv3dWithBroadcast <
>::Kernel;
// Define epilogue
using Epilogue = typename cutlass::conv::kernel::detail::DefaultConvEpilogueWithBroadcastSimtStridedDgrad<
using Epilogue = typename cutlass::conv::kernel::detail::DefaultConvEpilogueWithBroadcastSimt<
ArchTag,
typename ImplicitGemmBase::Epilogue::Shape,
typename ImplicitGemmBase::Epilogue::WarpMmaOperator,
@@ -281,7 +284,10 @@ struct DefaultDeconv3dWithBroadcast <
typename EpilogueOutputOp::ElementT,
typename EpilogueOutputOp::ElementVector,
EpilogueOutputOp,
ImplicitGemmBase::Epilogue::kElementsPerAccess
ImplicitGemmBase::Epilogue::kElementsPerAccess,
layout::NoPermute,
StrideSupport::kStrided,
5
>::Epilogue;
// Define the kernel
include/cutlass/conv/kernel/implicit_gemm_convolution.h
+2 -2
View File
@@ -233,9 +233,9 @@ struct ImplicitGemmConvolution {
ptr_A(args.ref_A.data()),
iterator_B(args.problem_size, args.ref_B.layout()),
ptr_B(args.ref_B.data()),
iterator_C(ConvOutputIteratorParameter::layout(args.ref_C), args.problem_size),
iterator_C(ConvOutputIteratorParameter::layout(args.ref_C), implicit_gemm_tensor_c_extent(kConvolutionalOperator, args.problem_size)),
ptr_C(args.ref_C.data()),
iterator_D(ConvOutputIteratorParameter::layout(args.ref_D), args.problem_size),
iterator_D(ConvOutputIteratorParameter::layout(args.ref_D), implicit_gemm_tensor_c_extent(kConvolutionalOperator, args.problem_size)),
ptr_D(args.ref_D.data()),
output_op(args.output_op),
semaphore(semaphore),
include/cutlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h
+2 -2
View File
@@ -257,9 +257,9 @@ struct ImplicitGemmConvolutionWithFusedEpilogue {
ptr_A(args.ref_A.data()),
iterator_B(args.problem_size, args.ref_B.layout()),
ptr_B(args.ref_B.data()),
iterator_C(ConvOutputIteratorParameter::layout(args.ref_C)),
iterator_C(ConvOutputIteratorParameter::layout(args.ref_C), implicit_gemm_tensor_c_extent(kConvolutionalOperator, args.problem_size)),
ptr_C(args.ref_C.data()),
iterator_D(ConvOutputIteratorParameter::layout(args.ref_D)),
iterator_D(ConvOutputIteratorParameter::layout(args.ref_D), implicit_gemm_tensor_c_extent(kConvolutionalOperator, args.problem_size)),
ptr_D(args.ref_D.data()),
output_op(args.output_op),
semaphore(semaphore),
include/cutlass/conv/kernel/sm90_implicit_gemm_tma_warpspecialized.hpp
+88 -65
View File
@@ -51,12 +51,12 @@ namespace cutlass::conv::kernel {
template <
class CollectiveMainloop_,
class CollectiveEpilogue_,
class TileSchedulerTag
class TileSchedulerTag_
>
class ConvUniversal<
CollectiveMainloop_,
CollectiveEpilogue_,
TileSchedulerTag,
TileSchedulerTag_,
cute::enable_if_t<cute::is_base_of_v<cutlass::conv::KernelImplicitTmaWarpSpecializedSm90,
typename CollectiveMainloop_::DispatchPolicy::Schedule>>>
{
@@ -90,6 +90,7 @@ public:
using EpilogueArguments = typename CollectiveEpilogue::Arguments;
using EpilogueParams = typename CollectiveEpilogue::Params;
using TileSchedulerTag = TileSchedulerTag_;
static_assert(cute::is_void_v<TileSchedulerTag>,
"TMA warp-specialized kernel does not support specializing the tile scheduler.");
using TileScheduler = typename cutlass::gemm::kernel::detail::TileSchedulerSelector<
@@ -144,7 +145,7 @@ public:
to_underlying_arguments(Arguments const& args, void* workspace) {
(void) workspace;
auto mainloop_params = CollectiveMainloop::to_underlying_arguments(args.mainloop, workspace);
auto problem_shape_MNKL = append<4>(mainloop_params.problem_shape, Int<1>{});
auto problem_shape_MNKL = args.mainloop.problem_shape.get_transformed_problem_shape_MNKL();
return {
mainloop_params,
@@ -157,7 +158,7 @@ public:
can_implement(Arguments const& args) {
bool implementable = true;
implementable &= CollectiveMainloop::can_implement(args.mainloop.problem_shape, args.mainloop);
implementable &= CollectiveEpilogue::can_implement(args.mainloop.problem_shape.get_transformed_problem_shape_MNK(), args.epilogue);
implementable &= CollectiveEpilogue::can_implement(args.mainloop.problem_shape.get_transformed_problem_shape_MNKL(), args.epilogue);
return implementable;
}
@@ -166,19 +167,17 @@ public:
return 0;
}
static cutlass::Status
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
CudaHostAdapter* cuda_adapter = nullptr) {
return Status::kSuccess;
}
// Computes the kernel launch grid shape based on runtime parameters
static dim3
get_grid_shape(Params const& params) {
// The CONV mainloop params problem shape will be the cute::Shape<> rank-3 MNK tuple we want for grid planning
// Although conv problems do not have an L mode, we add it here to comply with the scheduler API
auto linear_problem_shape_MNKL = make_shape(
size<0>(params.mainloop.problem_shape), // M mode is linearized.
shape<1>(params.mainloop.problem_shape),
shape<2>(params.mainloop.problem_shape),
Int<1>{});
return cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90::get_tiled_cta_shape_mnl(
linear_problem_shape_MNKL, TileShape{}, ClusterShape{});
params.mainloop.problem_shape, TileShape{}, ClusterShape{});
}
static dim3
@@ -205,14 +204,25 @@ public:
Consumer = 1,
};
enum class ProducerWarpRole {
MainloopEpilogue = 0,
Warp1 = 1,
Warp2 = 2,
Warp3 = 3
};
// Kernel level shared memory storage
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
int thread_idx = int(threadIdx.x);
int lane_idx = canonical_lane_idx();
int warp_idx = canonical_warp_idx_sync();
int warp_idx_in_warp_group = warp_idx % NumWarpsPerWarpGroup;
int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup;
auto warp_group_role = WarpGroupRole(canonical_warp_group_idx());
auto producer_warp_role = ProducerWarpRole(warp_idx_in_warp_group);
int lane_predicate = cute::elect_one_sync();
uint32_t block_rank_in_cluster = cute::block_rank_in_cluster();
// Issue Tma Descriptor Prefetch from a single thread
if ((warp_idx == 0) && lane_predicate) {
@@ -223,7 +233,7 @@ public:
// Mainloop Load pipeline
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
typename MainloopPipeline::Params mainloop_pipeline_params;
if (warp_group_role == WarpGroupRole::Producer) {
if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::MainloopEpilogue) {
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
}
if (warp_group_role == WarpGroupRole::Consumer) {
@@ -231,22 +241,24 @@ public:
}
mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0;
mainloop_pipeline_params.num_consumers = NumThreadsPerWarpGroup;
mainloop_pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytes;
mainloop_pipeline_params.transaction_bytes = params.mainloop.tma_transaction_bytes;
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, ClusterShape{});
// Epilogue Load pipeline
using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
typename EpiLoadPipeline::Params epi_load_pipeline_params;
if (warp_group_role == WarpGroupRole::Producer) {
if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::MainloopEpilogue) {
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer;
}
if (warp_group_role == WarpGroupRole::Consumer) {
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer;
}
epi_load_pipeline_params.dst_blockid = cute::block_rank_in_cluster();
epi_load_pipeline_params.producer_arv_count = 1; // 1 thread issues TMA load
epi_load_pipeline_params.producer_arv_count = NumThreadsPerWarp;
epi_load_pipeline_params.consumer_arv_count = NumThreadsPerWarpGroup;
epi_load_pipeline_params.transaction_bytes = CollectiveEpilogue::TmaTransactionBytes;
if constexpr (CollectiveEpilogue::RequiresTransactionBytes) {
epi_load_pipeline_params.transaction_bytes = params.epilogue.tma_transaction_bytes;
}
EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params);
// Epilogue Store pipeline
@@ -266,16 +278,26 @@ public:
PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
auto cluster_wait_fn = [&] () {
// We need this to guarantee that the Pipeline init is visible
// To all producers and consumer thread blocks in the Cluster
if constexpr (size(ClusterShape{}) > 1) {
cute::cluster_arrive_relaxed();
return [] () { cute::cluster_wait(); };
}
else {
__syncthreads();
return [] () {}; // do nothing
}
} ();
// Separate out problem shape for convenience
auto M = get<0>(params.mainloop.problem_shape);
auto N = get<1>(params.mainloop.problem_shape);
auto K = get<2>(params.mainloop.problem_shape);
// output strides are coalesced so we linearize the output shape to match the shape/stride profiles
auto linear_problem_shape_MNKL = make_shape(size(M), N, K, Int<1>{});
auto problem_shape_MNKL = append<4>(params.mainloop.problem_shape, _1{});
auto [M, N, K, L] = problem_shape_MNKL;
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA_mk = params.mainloop.tma_load_a.get_tma_tensor(make_shape(M, size(K)));
Tensor mA_mk = params.mainloop.tma_load_a.get_tma_tensor(make_shape(M, K));
Tensor mB_nk = params.mainloop.tma_load_b.get_tma_tensor(make_shape(N, K));
// Get the appropriate blocks for this thread block -- potential for thread block locality
@@ -288,7 +310,8 @@ public:
// Compute m_coord, n_coord, and l_coord with their post-tiled shapes
auto m_coord = idx2crd(int(blockIdx.x), shape<2>(gA_mk));
auto n_coord = idx2crd(int(blockIdx.y), shape<2>(gB_nk));
auto n_coord = idx2crd(int(blockIdx.y), shape<2>(gB_nk), compact_col_major(shape<2>(gB_nk)));
// The output shape M is linearized so the output coord M here should also be linearized.
auto output_tile_coord = make_coord(int(blockIdx.x), n_coord, _, Int<0>{});
@@ -300,51 +323,43 @@ public:
auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA));
auto k_tile_count = size<2>(gA);
auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(cta_tile_shape);
auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(cta_tile_shape);
// Make sure pipeline init is visible to all producers and consumer CTAs in cluster
if constexpr (size(ClusterShape{}) > 1) {
cute::cluster_arrive_relaxed();
cute::cluster_wait();
}
else {
__syncthreads();
}
// In a warp specialized kernel, collectives expose data movement and compute operations separately
CollectiveMainloop collective_mainloop;
CollectiveEpilogue collective_epilogue{params.epilogue, shared_storage.tensors.epilogue};
if (warp_group_role == WarpGroupRole::Producer) {
collective_mainloop.load(
mainloop_pipeline,
mainloop_pipe_producer_state,
gA, params.mainloop.tma_load_a,
gB, params.mainloop.tma_load_b,
k_tile_iter, k_tile_count,
thread_idx,
shared_storage.tensors.mainloop
);
// Update starting mainloop pipeline state for the pipeline drain
mainloop_pipe_producer_state.advance(k_tile_count);
// Make sure mainloop consumer has been waited upon before issuing epilogue load
collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
// Wait for all thread blocks in Cluster
cluster_wait_fn();
if (collective_epilogue.is_producer_load_needed()) {
collective_epilogue.load(
epi_load_pipeline,
epi_load_pipe_producer_state,
linear_problem_shape_MNKL,
cta_tile_shape,
output_tile_coord,
tiled_mma,
warp_group_thread_idx,
shared_storage.tensors.epilogue
if (warp_group_role == WarpGroupRole::Producer) {
if (producer_warp_role == ProducerWarpRole::MainloopEpilogue) {
collective_mainloop.load(
mainloop_pipeline,
mainloop_pipe_producer_state,
gA, params.mainloop.tma_load_a,
gB, params.mainloop.tma_load_b,
k_tile_iter, k_tile_count,
lane_idx,
block_rank_in_cluster,
shared_storage.tensors.mainloop
);
// Update starting load pipeline state for the pipeline drain
epi_load_pipe_producer_state.advance(c_tile_count);
collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state);
// Update starting mainloop pipeline state for the pipeline drain
mainloop_pipe_producer_state.advance(k_tile_count);
// Make sure mainloop consumer has been waited upon before issuing epilogue load
collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
if (collective_epilogue.is_producer_load_needed()) {
epi_load_pipe_producer_state = collective_epilogue.load(
epi_load_pipeline,
epi_load_pipe_producer_state,
problem_shape_MNKL,
cta_tile_shape,
output_tile_coord,
tiled_mma,
lane_idx,
shared_storage.tensors.epilogue
);
collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state);
}
}
}
else if (warp_group_role == WarpGroupRole::Consumer) {
@@ -368,12 +383,13 @@ public:
);
// Epilogue and write to gD
auto [epi_load_pipe_consumer_state_next, epi_store_pipe_producer_state_next] =
collective_epilogue.store(
epi_load_pipeline,
epi_load_pipe_consumer_state,
epi_store_pipeline,
epi_store_pipe_producer_state,
linear_problem_shape_MNKL,
problem_shape_MNKL,
cta_tile_shape,
output_tile_coord,
accumulators,
@@ -381,6 +397,13 @@ public:
warp_group_thread_idx,
shared_storage.tensors.epilogue
);
collective_epilogue.store_tail(
epi_load_pipeline,
epi_load_pipe_consumer_state_next,
epi_store_pipeline,
epi_store_pipe_producer_state_next
);
}
}
};
include/cutlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_analytic.h
+1 -1
View File
@@ -251,7 +251,7 @@ public:
static Status can_implement(Conv3dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.C % (128/sizeof_bits<Element>::value)) {
if (problem_size.C % AccessType::kElements) {
return Status::kErrorInvalidProblem;
}

Some files were not shown because too many files have changed in this diff Show More