v3.8.0 update (#2082)
* 3.8 update * fix Markus' name --------- Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
@@ -198,7 +198,6 @@ set(header_files_to_check
|
||||
cutlass/version.h
|
||||
cutlass/wmma_array.h
|
||||
cutlass/workspace.h
|
||||
|
||||
cutlass/exmy_base.h
|
||||
cutlass/float_subbyte.h
|
||||
|
||||
|
||||
@@ -125,6 +125,50 @@ function(cutlass_test_unit_add_executable NAME)
|
||||
|
||||
endfunction()
|
||||
|
||||
|
||||
function(cutlass_test_unit_add_executable_split_file NAME)
|
||||
# Given the input arguments to cutlass_test_unit_add_executable, creates
|
||||
# a new set of arguments in which each file has at most one TEST definition,
|
||||
# and calls cutlass_test_unit_add_executable with the newly-formed arguments.
|
||||
# The goal of this is to reduce the memory consumed while building CUTLASS
|
||||
# tests with a high degree of parallelism while not requiring developers
|
||||
# to split unit tests across multiple files artificially.
|
||||
|
||||
# Get all arguments other than the NAME of the target
|
||||
list(SUBLIST ARGV 1 ${ARGC} SUBARGV)
|
||||
|
||||
if (CUTLASS_UNIT_TEST_SPLIT_FILES)
|
||||
execute_process(
|
||||
WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
|
||||
COMMAND ${Python3_EXECUTABLE} ${CUTLASS_SOURCE_DIR}/tools/scripts/split_test_cmake.py
|
||||
${NAME}
|
||||
${CMAKE_CURRENT_SOURCE_DIR}
|
||||
--src_files ${SUBARGV}
|
||||
--dst_dir ${CMAKE_CURRENT_BINARY_DIR}
|
||||
RESULT_VARIABLE cutlass_test_SPLIT_RESULT
|
||||
OUTPUT_VARIABLE cutlass_test_SPLIT_OUTPUT
|
||||
OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/test_split_files.txt
|
||||
ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/test_split_error.log
|
||||
)
|
||||
|
||||
if(NOT cutlass_test_SPLIT_RESULT EQUAL 0)
|
||||
message(FATAL_ERROR "Error splitting unit test. See ${CMAKE_CURRENT_BINARY_DIR}/test_split_error.log")
|
||||
endif()
|
||||
|
||||
# Forward the values printed by split_test_cmake.py as arguments to cutlass_test_unit_add_executable.
|
||||
# We additionally specify to add -I${CMAKE_CURRENT_SOURCE_DIR} to the target. This is necessary because
|
||||
# the splitting process writes new files to ${CMAKE_CURRENT_BINARY_DIR}, but many CUTLASS unit tests
|
||||
# use relative imports for including testbeds (e.g., '#include "../testbed.hpp"'). These headers are
|
||||
# not written to ${CMAKE_CURRENT_BINARY_DIR} during the splitting process, so we must indicate that
|
||||
# headers can also be searched for from ${CMAKE_CURRENT_SOURCE_DIR}.
|
||||
file(STRINGS ${CMAKE_CURRENT_BINARY_DIR}/test_split_files.txt NEW_OPTIONS)
|
||||
cutlass_test_unit_add_executable(${NAME} ${NEW_OPTIONS} EXTRA_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR})
|
||||
else()
|
||||
# Simply pass arguments through
|
||||
cutlass_test_unit_add_executable(${ARGV})
|
||||
endif()
|
||||
endfunction()
|
||||
|
||||
add_custom_target(cutlass_test_unit)
|
||||
add_custom_target(test_unit)
|
||||
|
||||
|
||||
@@ -276,11 +276,12 @@ endif()
|
||||
|
||||
if (CUTLASS_NVCC_MAX_ARCH GREATER_EQUAL 89)
|
||||
|
||||
# Conv - F8 input, F8 output, F32 accumulation
|
||||
# Conv - F8 input, F8 output
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_conv_device_tensorop_f8_sm89
|
||||
|
||||
conv2d_fprop_implicit_gemm_f8nhwc_f8nhwc_f8nhwc_tensor_op_f32_sm89.cu
|
||||
conv2d_fprop_implicit_gemm_f8nhwc_f8nhwc_f8nhwc_tensor_op_f16_sm89.cu
|
||||
)
|
||||
|
||||
endif()
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Tests for device-wide Conv2d fprop interface with:
|
||||
A: NHWC, of type FE4M4 or FE5M2
|
||||
B: NHWC, of type FE4M3 or FE5M2
|
||||
C: NHWC, of FE4M3 or FE5M2
|
||||
Accum: F16
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination_generic_with_scaling.h"
|
||||
#include "cutlass/conv/kernel/default_conv2d_fprop_with_absmax.h"
|
||||
#include "cutlass/conv/device/implicit_gemm_convolution.h"
|
||||
#include "cutlass/util/tensor_view_io.h"
|
||||
|
||||
#include "conv2d_with_absmax_testbed.h"
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_F16_SM89_SUPPORTED)
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Conv2d_Fprop_Analytic_ImplicitGemm_fe4m3nhwc_fe4mnhwc_fe4mnhwc_tensor_op_f16,
|
||||
identity_128x256x64_64x3_64x64x64) {
|
||||
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Conv2dFpropKernel = typename cutlass::conv::kernel::DefaultConv2dFpropWithAbsMax<
|
||||
ElementA, cutlass::layout::TensorNHWC,
|
||||
ElementB, cutlass::layout::TensorNHWC,
|
||||
ElementOutput, cutlass::layout::TensorNHWC,
|
||||
ElementAccumulator,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>,
|
||||
cutlass::gemm::GemmShape<64, 64, 64>,
|
||||
cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
kStages,
|
||||
cutlass::arch::OpMultiplyAdd,
|
||||
cutlass::conv::IteratorAlgorithm::kAnalytic
|
||||
>::Kernel;
|
||||
|
||||
using Conv2dFprop = cutlass::conv::device::ImplicitGemmConvolution<Conv2dFpropKernel>;
|
||||
|
||||
bool passed = test::conv::device::TestAllConv2dWithAbsmax<Conv2dFprop, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Conv2d_Fprop_Optimized_ImplicitGemm_fe4m3nhwc_fe4mnhwc_fe4mnhwc_tensor_op_f16,
|
||||
relu_128x256x64_64x3_64x64x64) {
|
||||
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::ReLu,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Conv2dFpropKernel = typename cutlass::conv::kernel::DefaultConv2dFpropWithAbsMax<
|
||||
ElementA, cutlass::layout::TensorNHWC,
|
||||
ElementB, cutlass::layout::TensorNHWC,
|
||||
ElementOutput, cutlass::layout::TensorNHWC,
|
||||
ElementAccumulator,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>,
|
||||
cutlass::gemm::GemmShape<64, 64, 64>,
|
||||
cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
kStages,
|
||||
cutlass::arch::OpMultiplyAdd,
|
||||
cutlass::conv::IteratorAlgorithm::kOptimized
|
||||
>::Kernel;
|
||||
|
||||
using Conv2dFprop = cutlass::conv::device::ImplicitGemmConvolution<Conv2dFpropKernel>;
|
||||
|
||||
bool passed = test::conv::device::TestAllConv2dWithAbsmax<Conv2dFprop, cutlass::epilogue::thread::ReLu>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Conv2d_Fprop_Optimized_ImplicitGemm_fe4m3nhwc_fe4mnhwc_fe4mnhwc_tensor_op_f16,
|
||||
identity_fastacc_128x256x64_64x3_64x64x64) {
|
||||
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Conv2dFpropKernel = typename cutlass::conv::kernel::DefaultConv2dFpropWithAbsMax<
|
||||
ElementA, cutlass::layout::TensorNHWC,
|
||||
ElementB, cutlass::layout::TensorNHWC,
|
||||
ElementOutput, cutlass::layout::TensorNHWC,
|
||||
ElementAccumulator,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>,
|
||||
cutlass::gemm::GemmShape<64, 64, 64>,
|
||||
cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
kStages,
|
||||
cutlass::arch::OpMultiplyAddFastAccum,
|
||||
cutlass::conv::IteratorAlgorithm::kOptimized
|
||||
>::Kernel;
|
||||
|
||||
using Conv2dFprop = cutlass::conv::device::ImplicitGemmConvolution<Conv2dFpropKernel>;
|
||||
|
||||
bool passed = test::conv::device::TestAllConv2dWithAbsmax<Conv2dFprop, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Conv2d_Fprop_Optimized_ImplicitGemm_fe4m3nhwc_fe4mnhwc_fe4mnhwc_tensor_op_f16,
|
||||
identity_noScale_128x256x64_64x3_64x64x64) {
|
||||
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Conv2dFpropKernel = typename cutlass::conv::kernel::DefaultConv2dFpropWithAbsMax<
|
||||
ElementA, cutlass::layout::TensorNHWC,
|
||||
ElementB, cutlass::layout::TensorNHWC,
|
||||
ElementOutput, cutlass::layout::TensorNHWC,
|
||||
ElementAccumulator,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>,
|
||||
cutlass::gemm::GemmShape<64, 64, 64>,
|
||||
cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
kStages,
|
||||
cutlass::arch::OpMultiplyAdd,
|
||||
cutlass::conv::IteratorAlgorithm::kOptimized
|
||||
>::Kernel;
|
||||
|
||||
using Conv2dFprop = cutlass::conv::device::ImplicitGemmConvolution<Conv2dFpropKernel>;
|
||||
|
||||
bool passed = test::conv::device::TestAllConv2dWithAbsmax<Conv2dFprop, cutlass::epilogue::thread::Identity>(
|
||||
/* scaleA = */false,
|
||||
/* scaleB = */false,
|
||||
/* scaleC = */false
|
||||
);
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_F16_SM89_SUPPORTED
|
||||
@@ -49,7 +49,7 @@
|
||||
|
||||
#include "conv2d_with_absmax_testbed.h"
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM89_SUPPORTED)
|
||||
#if defined(CUTLASS_ARCH_MMA_F32_SM89_SUPPORTED)
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -365,4 +365,4 @@ TEST(SM89_Device_Conv2d_Fprop_Optimized_ImplicitGemm_fe4m3nhwc_fe4mnhwc_fe4mnhwc
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_SM89_SUPPORTED
|
||||
#endif // CUTLASS_ARCH_MMA_F32_SM89_SUPPORTED
|
||||
|
||||
@@ -30,8 +30,6 @@ add_custom_target(cutlass_test_unit_gemm_device)
|
||||
add_custom_target(test_unit_gemm_device)
|
||||
|
||||
|
||||
add_subdirectory(sm100_blockscaled_tensorop_gemm)
|
||||
|
||||
|
||||
################################################################################
|
||||
|
||||
@@ -53,6 +51,12 @@ endfunction()
|
||||
|
||||
################################################################################
|
||||
|
||||
|
||||
add_subdirectory(sm100_blockscaled_tensorop_gemm)
|
||||
add_subdirectory(sm100_tensorop_gemm)
|
||||
|
||||
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable(
|
||||
cutlass_test_unit_gemm_device_simt
|
||||
|
||||
@@ -548,8 +552,10 @@ cutlass_test_unit_gemm_device_add_executable(
|
||||
|
||||
gemm_f8t_f8n_f32t_tensor_op_f32_sm89.cu
|
||||
gemm_f8t_f8n_f32t_tensor_op_f32_sparse_sm89.cu
|
||||
gemm_f8t_f8n_f16t_tensor_op_f16_sm89.cu
|
||||
gemm_f8t_f8n_f8t_tensor_op_f32_sm89.cu
|
||||
# gemm_f8t_f8n_f8t_tensor_op_f32_sparse_sm89.cu
|
||||
gemm_f8t_f8n_f8t_tensor_op_f16_sm89.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable(
|
||||
@@ -829,56 +835,6 @@ endif()
|
||||
|
||||
if(NOT CUTLASS_NVCC_ARCHS STREQUAL "100")
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_sm100_fp16_gemm
|
||||
|
||||
# No batching of source to control compiler memory usage
|
||||
BATCH_SOURCES ON
|
||||
BATCH_SIZE 1
|
||||
|
||||
sm100_gemm_f16_f16_f32_tensor_op_f32.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable(
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_stream_k
|
||||
|
||||
sm100_gemm_f16_f16_f16_tensor_op_f32_stream_k.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable(
|
||||
cutlass_test_unit_gemm_device_sm100_bf16_gemm
|
||||
|
||||
# No batching of source to control compiler memory usage
|
||||
BATCH_SOURCES ON
|
||||
BATCH_SIZE 1
|
||||
|
||||
sm100_gemm_bf16_bf16_f32_tensor_op_f32.cu
|
||||
)
|
||||
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable(
|
||||
cutlass_test_unit_gemm_device_tensorop_stride_batch_alpha_beta_sm100
|
||||
|
||||
# No batching of source to control compiler memory usage
|
||||
BATCH_SOURCES ON
|
||||
BATCH_SIZE 1
|
||||
|
||||
sm100_gemm_f8_f8_f8_tensor_op_s32_batch_alpha_beta.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable(
|
||||
cutlass_test_unit_gemm_device_tensorop_runtime_datatype_sm100
|
||||
|
||||
# No batching of source to control compiler memory usage
|
||||
BATCH_SOURCES ON
|
||||
BATCH_SIZE 1
|
||||
|
||||
sm100_gemm_f8_f8_f8_tensor_op_f32_runtime_datatype.cu
|
||||
sm100_gemm_f6_f6_f32_tensor_op_f32_runtime_datatype.cu
|
||||
sm100_gemm_f4_f4_f32_tensor_op_f32_runtime_datatype.cu
|
||||
sm100_gemm_f8_f4_f32_tensor_op_f32_runtime_datatype.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable(
|
||||
cutlass_test_unit_gemm_device_16b_tensorop_sm100_ptr_array
|
||||
|
||||
|
||||
154
test/unit/gemm/device/gemm_f8t_f8n_f16t_tensor_op_f16_sm89.cu
Normal file
154
test/unit/gemm/device/gemm_f8t_f8n_f16t_tensor_op_f16_sm89.cu
Normal file
@@ -0,0 +1,154 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Tests for device-wide GEMM interface with:
|
||||
A: row major, of type FE4M4 or FE5M2
|
||||
B: column major, of type FE4M3 or FE5M2
|
||||
C: row major, of type F16
|
||||
Accum: F16
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/gemm/device/gemm.h"
|
||||
#include "cutlass/util/host_tensor.h"
|
||||
#include "cutlass/util/reference/host/gemm.h"
|
||||
#include "cutlass/util/reference/host/tensor_compare.h"
|
||||
#include "cutlass/util/reference/host/tensor_copy.h"
|
||||
#include "cutlass/util/reference/host/tensor_fill.h"
|
||||
#include "cutlass/util/tensor_view_io.h"
|
||||
|
||||
#include "testbed.h"
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_F16_SM89_SUPPORTED)
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe4m3t_fe4m3n_f16t_tensor_op_f16, 128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::half_t;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using Gemm = cutlass::gemm::device::Gemm<
|
||||
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::epilogue::thread::LinearCombination<
|
||||
ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator, ElementAccumulator>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages>;
|
||||
|
||||
EXPECT_TRUE(test::gemm::device::TestAllGemm<Gemm>());
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe4m3t_fe5m2n_f16t_tensor_op_f16, 128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e5m2_t;
|
||||
using ElementOutput = cutlass::half_t;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using Gemm = cutlass::gemm::device::Gemm<
|
||||
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::epilogue::thread::LinearCombination<
|
||||
ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator, ElementAccumulator>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages>;
|
||||
|
||||
EXPECT_TRUE(test::gemm::device::TestAllGemm<Gemm>());
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe5m2t_fe4m3n_f16t_tensor_op_f16, 128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::half_t;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using Gemm = cutlass::gemm::device::Gemm<
|
||||
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::epilogue::thread::LinearCombination<
|
||||
ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator, ElementAccumulator>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages>;
|
||||
|
||||
EXPECT_TRUE(test::gemm::device::TestAllGemm<Gemm>());
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe5m2t_fe5m2n_f16t_tensor_op_f16, 128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
using ElementB = cutlass::float_e5m2_t;
|
||||
using ElementOutput = cutlass::half_t;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using Gemm = cutlass::gemm::device::Gemm<
|
||||
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::epilogue::thread::LinearCombination<
|
||||
ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator, ElementAccumulator>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages>;
|
||||
|
||||
EXPECT_TRUE(test::gemm::device::TestAllGemm<Gemm>());
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_F16_SM89_SUPPORTED
|
||||
@@ -51,7 +51,7 @@
|
||||
|
||||
#include "testbed.h"
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM89_SUPPORTED)
|
||||
#if defined(CUTLASS_ARCH_MMA_F32_SM89_SUPPORTED)
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -151,4 +151,4 @@ TEST(SM89_Device_Gemm_fe5m2t_fe5m2n_f32t_tensor_op_f32, 128x256x64_64x64x64) {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_SM89_SUPPORTED
|
||||
#endif // CUTLASS_ARCH_MMA_F32_SM89_SUPPORTED
|
||||
|
||||
@@ -51,7 +51,7 @@
|
||||
|
||||
#include "testbed_sparse.h"
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM89_SUPPORTED)
|
||||
#if defined(CUTLASS_ARCH_SPARSE_MMA_F32_SM89_SUPPORTED)
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -151,4 +151,4 @@ TEST(SM89_Device_Sparse_Gemm_fe5m2t_fe5m2n_f32t_tensor_op_f32, 128x128x128_64x64
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_SM89_SUPPORTED
|
||||
#endif // CUTLASS_ARCH_MMA_F32_SM89_SUPPORTED
|
||||
|
||||
430
test/unit/gemm/device/gemm_f8t_f8n_f8t_tensor_op_f16_sm89.cu
Normal file
430
test/unit/gemm/device/gemm_f8t_f8n_f8t_tensor_op_f16_sm89.cu
Normal file
@@ -0,0 +1,430 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Tests for device-wide GEMM interface with:
|
||||
A: row major, of type FE4M4 or FE5M2
|
||||
B: column major, of type FE4M3 or FE5M2
|
||||
C: row major, of FE4M3 or FE5M2
|
||||
Accum: F16
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination_generic_with_scaling.h"
|
||||
#include "cutlass/gemm/device/gemm_universal_with_absmax.h"
|
||||
#include "cutlass/util/host_tensor.h"
|
||||
#include "cutlass/util/reference/host/gemm.h"
|
||||
#include "cutlass/util/reference/host/tensor_compare.h"
|
||||
#include "cutlass/util/reference/host/tensor_copy.h"
|
||||
#include "cutlass/util/reference/host/tensor_fill.h"
|
||||
#include "cutlass/util/tensor_view_io.h"
|
||||
|
||||
#include "testbed.h"
|
||||
#include "testbed_with_absmax.h"
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_F16_SM89_SUPPORTED)
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe4m3t_fe4m3n_fe4m3t_tensor_op_f16, identity_128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe4m3t_fe4m3n_fe4m3t_tensor_op_f16, identity_fastacc_128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
static int const kAlignment = 16;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages,
|
||||
kAlignment, kAlignment, cutlass::arch::OpMultiplyAddFastAccum
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe4m3t_fe4m3n_fe4m3t_tensor_op_f16, relu_128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::ReLu,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::ReLu>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe4m3t_fe5m2n_fe4m3t_tensor_op_f16, identity_128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e5m2_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe5m2t_fe4m3n_fe4m3t_tensor_op_f16, identity_128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe5m2t_fe5m2n_fe4m3t_tensor_op_f16, identity_128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
using ElementB = cutlass::float_e5m2_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe4m3t_fe4m3n_fe5m2t_tensor_op_f16, identity_128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e5m2_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe5m2t_fe5m2n_fe5m2t_tensor_op_f16, identity_diff_aux_output_types_128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
using ElementB = cutlass::float_e5m2_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = cutlass::float_e5m2_t;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe4m3t_fe4m3n_fe4m3t_tensor_op_f16, identity_128x128x64_32x64x64) {
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 128, 64>, cutlass::gemm::GemmShape<32, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe4m3t_fe4m3n_fe4m3t_tensor_op_f16, identity_noScale_128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::Identity>(
|
||||
/* scaleA = */false,
|
||||
/* scaleB = */false,
|
||||
/* scaleC = */false
|
||||
);
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM89_Device_Gemm_fe4m3t_fe4m3n_fe4m3t_tensor_op_f16, identity_noAux_128x256x64_64x64x64) {
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = float;
|
||||
using ElementAccumulator = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::Identity,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages
|
||||
>;
|
||||
|
||||
bool passed = test::gemm::device::TestAllGemmWithAbsmax<Gemm, test::gemm::device::Testbed<Gemm>, cutlass::epilogue::thread::Identity>();
|
||||
EXPECT_TRUE(passed);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_F16_SM89_SUPPORTED
|
||||
@@ -54,7 +54,7 @@
|
||||
#include "testbed.h"
|
||||
#include "testbed_with_absmax.h"
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM89_SUPPORTED)
|
||||
#if defined(CUTLASS_ARCH_MMA_F32_SM89_SUPPORTED)
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -427,4 +427,4 @@ TEST(SM89_Device_Gemm_fe4m3t_fe4m3n_fe4m3t_tensor_op_f32, identity_noAux_128x256
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_SM89_SUPPORTED
|
||||
#endif // CUTLASS_ARCH_MMA_F32_SM89_SUPPORTED
|
||||
|
||||
@@ -54,7 +54,7 @@
|
||||
#include "testbed_sparse.h"
|
||||
#include "testbed_with_absmax.h"
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM89_SUPPORTED)
|
||||
#if defined(CUTLASS_ARCH_SPARSE_MMA_F32_SM89_SUPPORTED)
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -461,4 +461,4 @@ TEST(SM89_Device_Sparse_Gemm_fe4m3t_fe4m3n_fe4m3t_tensor_op_f32, identity_noAux_
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_SM89_SUPPORTED
|
||||
#endif // CUTLASS_ARCH_MMA_F32_SM89_SUPPORTED
|
||||
|
||||
@@ -893,7 +893,6 @@ struct HostCollectiveMainloop<ScheduleType_, Gemm, ElementA_, ElementB_,
|
||||
};
|
||||
|
||||
|
||||
|
||||
//
|
||||
// Block Scaled Gemm Input Operands : A , B, scalefactorA, scalefactorB
|
||||
//
|
||||
@@ -1412,8 +1411,12 @@ struct HostCollectiveEpilogue {
|
||||
cutlass::epilogue::thread::Identity<ElementCompute>>;
|
||||
|
||||
static constexpr bool IsRowBiasEnabled = FusionOp::IsPerRowBiasSupported;
|
||||
static constexpr bool IsColBiasEnabled = FusionOp::IsPerColBiasSupported;
|
||||
static_assert(not (IsColBiasEnabled && IsRowBiasEnabled));
|
||||
|
||||
static constexpr bool IsDeBiasEnabled = FusionOp::IsDePerRowBiasSupported;
|
||||
static constexpr bool IsPerRowScaleEnabled = FusionOp::IsPerRowScaleSupported;
|
||||
static constexpr bool IsPerColScaleEnabled = FusionOp::IsPerColScaleSupported;
|
||||
static constexpr bool IsScaleFactorEnabled = FusionOp::IsScaleFactorSupported;
|
||||
static constexpr bool IsAuxInEnabled = FusionOp::IsAuxInSupported;
|
||||
static constexpr bool IsAuxOutEnabled = FusionOp::IsAuxOutSupported;
|
||||
@@ -1462,7 +1465,7 @@ struct HostCollectiveEpilogue {
|
||||
CheckEquality check_relative_equality = CheckEquality::EXACT;
|
||||
// Are scalars copied to device memory before kernel launch
|
||||
ScalarLoc use_device_scalars = ScalarLoc::ON_HOST;
|
||||
// If per-row scale is enabled and this is disabled, alpha/beta are passed as a host or device scalar instead of device vector
|
||||
// If vector scale is supported and this is disabled, alpha/beta are passed as a host or device scalar instead of device vector
|
||||
VectorScale vector_scale_mode = VectorScale::DISABLED;
|
||||
|
||||
// Random distribution with which to initialize the A/B/C/D/Aux scaling factors
|
||||
@@ -1555,8 +1558,7 @@ struct HostCollectiveEpilogue {
|
||||
auto col_vector_coord = cutlass::make_Coord(M);
|
||||
auto row_vector_coord = cutlass::make_Coord(N);
|
||||
auto batch_vector_coord = cutlass::make_Coord(L);
|
||||
auto ML_coord = cutlass::make_Coord(M * L);
|
||||
if constexpr (IsPerRowScaleEnabled) {
|
||||
if constexpr (IsPerRowScaleEnabled or IsPerColScaleEnabled) {
|
||||
// scalars
|
||||
if (vector_scale_mode == VectorScale::DISABLED) {
|
||||
// batched scalars
|
||||
@@ -1581,8 +1583,9 @@ struct HostCollectiveEpilogue {
|
||||
}
|
||||
// batched vectors
|
||||
else {
|
||||
alpha.resize(ML_coord, true);
|
||||
beta.resize(ML_coord, true);
|
||||
auto batched_vector_coord = cutlass::make_Coord((IsPerRowScaleEnabled ? M : N) * L);
|
||||
alpha.resize(batched_vector_coord, true);
|
||||
beta.resize(batched_vector_coord, true);
|
||||
EXPECT_TRUE(initialize_tensor(alpha.host_view(), init_scale, seed + 2023));
|
||||
if (beta_ != ElementScalar(0)) {
|
||||
EXPECT_TRUE(initialize_tensor(beta.host_view(), init_scale, seed + 2024));
|
||||
@@ -1627,9 +1630,7 @@ struct HostCollectiveEpilogue {
|
||||
scale_D.sync_device();
|
||||
}
|
||||
|
||||
if constexpr (
|
||||
IsRowBiasEnabled
|
||||
) {
|
||||
if constexpr (IsRowBiasEnabled or IsColBiasEnabled) {
|
||||
bias.resize(IsRowBiasEnabled ? col_vector_coord : row_vector_coord);
|
||||
EXPECT_TRUE(initialize_tensor(bias.host_view(), init_bias, seed + 2023));
|
||||
bias.sync_device();
|
||||
@@ -1810,7 +1811,6 @@ struct HostCollectiveEpilogue {
|
||||
}
|
||||
passed &= passed_sf;
|
||||
}
|
||||
|
||||
|
||||
return passed;
|
||||
}
|
||||
@@ -1823,7 +1823,7 @@ struct HostCollectiveEpilogue {
|
||||
<< ", scale_b: " << scale_B.at(coord_0)
|
||||
<< ", scale_c: " << scale_C.at(coord_0);
|
||||
}
|
||||
if constexpr (IsPerRowScaleEnabled) {
|
||||
if constexpr (IsPerRowScaleEnabled or IsPerColScaleEnabled) {
|
||||
file << "\n\nvalpha = \n" << alpha.host_view();
|
||||
file << "\n\nvbeta = \n" << beta.host_view();
|
||||
} else {
|
||||
@@ -1853,9 +1853,10 @@ struct HostCollectiveEpilogue {
|
||||
file << "\n\n";
|
||||
}
|
||||
|
||||
if constexpr (IsRowBiasEnabled) {
|
||||
if constexpr (IsRowBiasEnabled or IsColBiasEnabled) {
|
||||
file << "\n\nBias = \n" << bias.host_view();
|
||||
}
|
||||
|
||||
if constexpr (IsAuxInEnabled) {
|
||||
file << "\n\nAux Input = \n" << tensor_Aux.host_view();
|
||||
}
|
||||
@@ -1876,7 +1877,6 @@ struct HostCollectiveEpilogue {
|
||||
<< "\n\nSFD Reference =\n" << reference_SFD.host_view()
|
||||
<< "\n\nSFD Computed =\n" << tensor_SFD.host_view();
|
||||
}
|
||||
|
||||
|
||||
file
|
||||
<< "\nC =\n" << tensor_C.host_view()
|
||||
@@ -1921,6 +1921,12 @@ struct HostCollectiveEpilogue {
|
||||
fusion_args.dAlpha = cute::make_stride(bool(m_stride),cute::_0{}, l_stride);
|
||||
fusion_args.dBeta = cute::make_stride(bool(m_stride),cute::_0{}, l_stride);
|
||||
}
|
||||
else if constexpr (IsPerColScaleEnabled) {
|
||||
int32_t n_stride = vector_scale_mode == VectorScale::ENABLED ? 1 : 0;
|
||||
int64_t l_stride = vector_scale_mode == VectorScale::ENABLED ? N : (use_device_scalars == ScalarLoc::ON_DEVICE ? 1 : 0);
|
||||
fusion_args.dAlpha = cute::make_stride(cute::_0{}, bool(n_stride), l_stride);
|
||||
fusion_args.dBeta = cute::make_stride(cute::_0{}, bool(n_stride), l_stride);
|
||||
}
|
||||
else {
|
||||
if constexpr (not IsFfma2Kernel) {
|
||||
if (use_device_scalars == ScalarLoc::ON_DEVICE) {
|
||||
@@ -1943,9 +1949,7 @@ struct HostCollectiveEpilogue {
|
||||
fusion_args.scale_d_ptr = scale_D.device_data();
|
||||
}
|
||||
|
||||
if constexpr (
|
||||
IsRowBiasEnabled
|
||||
) {
|
||||
if constexpr (IsRowBiasEnabled or IsColBiasEnabled) {
|
||||
fusion_args.bias_ptr = bias.device_data();
|
||||
}
|
||||
|
||||
@@ -1993,7 +1997,6 @@ struct HostCollectiveEpilogue {
|
||||
arguments.thread.block_scale_factor_ptr = tensor_SFD.device_data();
|
||||
arguments.thread.norm_constant_ptr = norm_constant.device_data();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
return arguments;
|
||||
@@ -2025,6 +2028,12 @@ struct HostCollectiveEpilogue {
|
||||
return cute::make_tensor(detail::make_iterator(alpha.host_data()),
|
||||
cute::make_layout(cute::make_shape(M, N, L), make_stride(m_stride, cute::_0{}, l_stride)));
|
||||
}
|
||||
else if constexpr (IsPerColScaleEnabled) {
|
||||
int n_stride = vector_scale_mode == VectorScale::ENABLED ? 1 : 0;
|
||||
int l_stride = vector_scale_mode == VectorScale::ENABLED ? N : (use_device_scalars == ScalarLoc::ON_DEVICE ? 1 : 0);
|
||||
return cute::make_tensor(detail::make_iterator(alpha.host_data()),
|
||||
cute::make_layout(cute::make_shape(M, N, L), make_stride(cute::_0{}, n_stride, l_stride)));
|
||||
}
|
||||
else {
|
||||
return cute::make_tensor(detail::make_iterator(alpha.host_data()),
|
||||
cute::make_layout(cute::make_shape(M, N, L), make_stride(cute::_0{}, cute::_0{}, cute::_1{})));
|
||||
@@ -2038,6 +2047,12 @@ struct HostCollectiveEpilogue {
|
||||
return cute::make_tensor(detail::make_iterator(beta.host_data()),
|
||||
cute::make_layout(cute::make_shape(M, N, L), make_stride(m_stride, cute::_0{}, l_stride)));
|
||||
}
|
||||
else if constexpr (IsPerColScaleEnabled) {
|
||||
int n_stride = vector_scale_mode == VectorScale::ENABLED ? 1 : 0;
|
||||
int l_stride = vector_scale_mode == VectorScale::ENABLED ? N : (use_device_scalars == ScalarLoc::ON_DEVICE ? 1 : 0);
|
||||
return cute::make_tensor(detail::make_iterator(beta.host_data()),
|
||||
cute::make_layout(cute::make_shape(M, N, L), make_stride(cute::_0{}, n_stride, l_stride)));
|
||||
}
|
||||
else {
|
||||
return cute::make_tensor(detail::make_iterator(beta.host_data()),
|
||||
cute::make_layout(cute::make_shape(M, N, L), make_stride(cute::_0{}, cute::_0{}, cute::_1{})));
|
||||
@@ -2069,8 +2084,8 @@ struct HostCollectiveEpilogue {
|
||||
ActivationFunctor,
|
||||
decltype(SfD),
|
||||
Int<SFD_VectorSize>,
|
||||
cutlass::plus<ElementCompute>
|
||||
, false /*PerColumnBias_*/
|
||||
cutlass::plus<ElementCompute>,
|
||||
IsColBiasEnabled
|
||||
, SfGenStrategy
|
||||
> epilogue_params{};
|
||||
|
||||
@@ -2086,8 +2101,7 @@ struct HostCollectiveEpilogue {
|
||||
epilogue_params.scale_d = scale_D.at(coord_0);
|
||||
}
|
||||
|
||||
if constexpr (IsRowBiasEnabled
|
||||
or IsDeBiasEnabled)
|
||||
if constexpr (IsRowBiasEnabled or IsColBiasEnabled or IsDeBiasEnabled)
|
||||
{
|
||||
epilogue_params.Bias = Bias;
|
||||
}
|
||||
@@ -2110,7 +2124,7 @@ struct HostCollectiveEpilogue {
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (IsPerRowScaleEnabled) {
|
||||
if constexpr (IsPerRowScaleEnabled or IsPerColScaleEnabled) {
|
||||
epilogue_params.Valpha = Valpha;
|
||||
if (vector_scale_mode == VectorScale::ENABLED) {
|
||||
epilogue_params.Vbeta = Vbeta;
|
||||
@@ -3294,9 +3308,14 @@ bool TestAll(double alpha = 1.0, double beta = cute::is_same_v<typename Gemm::Ge
|
||||
|
||||
Testbed3x<Gemm, ActivationFunctor> testbed(check_relative_equality, ScalarLoc::ON_HOST, VectorScale::DISABLED);
|
||||
|
||||
int max_alignment = std::max(Gemm::kAlignmentA, Gemm::kAlignmentB);
|
||||
std::vector<int> problem_size_m = {max_alignment, 512 - 3 * max_alignment};
|
||||
std::vector<int> problem_size_n = {max_alignment, 512 - 2 * max_alignment};
|
||||
int max_alignment_m = std::max({Gemm::kAlignmentA, Gemm::kAlignmentC, Gemm::kAlignmentD});
|
||||
int max_alignment_n = std::max({Gemm::kAlignmentB, Gemm::kAlignmentC, Gemm::kAlignmentD});
|
||||
if constexpr (std::is_base_of_v<cutlass::epilogue::fusion::FusionOperation, typename Gemm::EpilogueOutputOp>) {
|
||||
max_alignment_m = std::max(max_alignment_m, Gemm::EpilogueOutputOp::AlignmentAux);
|
||||
max_alignment_n = std::max(max_alignment_n, Gemm::EpilogueOutputOp::AlignmentAux);
|
||||
}
|
||||
std::vector<int> problem_size_m = {max_alignment_m, 512 - 3 * max_alignment_m};
|
||||
std::vector<int> problem_size_n = {max_alignment_n, 512 - 2 * max_alignment_n};
|
||||
|
||||
if constexpr (cute::is_same_v<typename Gemm::GemmKernel::DispatchPolicy::Schedule,
|
||||
cutlass::gemm::KernelTmaWarpSpecializedPingpong>) {
|
||||
@@ -3307,7 +3326,8 @@ bool TestAll(double alpha = 1.0, double beta = cute::is_same_v<typename Gemm::Ge
|
||||
constexpr int Stages = Gemm::GemmKernel::DispatchPolicy::Stages;
|
||||
constexpr int TileShapeK = cute::size<2>(typename Gemm::GemmKernel::TileShape{});
|
||||
|
||||
std::vector<int> problem_size_k = {max_alignment, TileShapeK * (Stages + 1) - max_alignment};
|
||||
int max_alignment_k = std::max(Gemm::kAlignmentA, Gemm::kAlignmentB);
|
||||
std::vector<int> problem_size_k = {max_alignment_k, TileShapeK * (Stages + 1) - max_alignment_k};
|
||||
|
||||
using DecompositionMode = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::DecompositionMode;
|
||||
std::vector<DecompositionMode> decomposition_modes = {DecompositionMode::Heuristic};
|
||||
@@ -3323,7 +3343,7 @@ bool TestAll(double alpha = 1.0, double beta = cute::is_same_v<typename Gemm::Ge
|
||||
|
||||
// Use larger K sizes for stream-K tests
|
||||
static constexpr int min_tiles_per_sk_unit = cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::min_iters_per_sk_unit_;
|
||||
problem_size_k = {TileShapeK * min_tiles_per_sk_unit, TileShapeK * 3 * min_tiles_per_sk_unit - max_alignment};
|
||||
problem_size_k = {TileShapeK * min_tiles_per_sk_unit, TileShapeK * 3 * min_tiles_per_sk_unit - max_alignment_k};
|
||||
}
|
||||
|
||||
using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90::RasterOrderOptions;
|
||||
@@ -3421,7 +3441,7 @@ bool TestAll(double alpha = 1.0, double beta = cute::is_same_v<typename Gemm::Ge
|
||||
|
||||
// if we do support batched GEMM, just run one test on it to save on test time
|
||||
if constexpr (cute::rank(ProblemShapeType{}) == 4) {
|
||||
auto problem_size = ProblemShapeType{256 + max_alignment, 256 + max_alignment, 160 + max_alignment, /* l */ 3};
|
||||
auto problem_size = ProblemShapeType{256 + max_alignment_m, 256 + max_alignment_n, 160 + max_alignment_k, /* l */ 3};
|
||||
passed = testbed.run(
|
||||
problem_size,
|
||||
cutlass::from_real<ElementScalar>(alpha),
|
||||
|
||||
@@ -1119,7 +1119,6 @@ struct HostCollectiveEpilogue {
|
||||
std::vector<cutlass::HostTensor<ElementSFD, LayoutTagD>> tensors_SFD;
|
||||
std::vector<cutlass::HostTensor<ElementSFD, LayoutTagD>> references_SFD;
|
||||
cutlass::DeviceAllocation<ElementSFD *> device_tensors_SFD;
|
||||
|
||||
|
||||
using ElementCompute = typename FusionOp::ElementCompute;
|
||||
using ElementScalar = typename FusionOp::ElementScalar;
|
||||
@@ -2205,8 +2204,7 @@ bool TestSmall(double alpha = 1.0, double beta = 1.0,
|
||||
|
||||
static constexpr bool IsF8F6F4 = cutlass::gemm::collective::detail::is_sm100_mma_f8f6f4<TiledMma, ElementA, ElementB>();
|
||||
alignment_bits = cutlass::detail::get_input_alignment_bits<ElementA, IsF8F6F4>();
|
||||
// For fp4 and fp6 mx kernels, the min alignment_input is 128 elements, so we don't need to add alignment_input in test problem sizes.
|
||||
|
||||
// For fp4 and fp6 QMMA kernels, the min alignment_input is 128 elements, so we don't need to add alignment_input in test problem sizes.
|
||||
int alignment_input = (alignment_bits / cute::sizeof_bits<ElementA>::value == 128) ? 0 : (alignment_bits / cute::sizeof_bits<ElementA>::value);
|
||||
|
||||
|
||||
|
||||
@@ -46,7 +46,7 @@ add_custom_target(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_mxf4xmxf6
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_add_executable(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_nvf4xnvf4
|
||||
|
||||
BATCH_SOURCES ON
|
||||
@@ -57,7 +57,7 @@ cutlass_test_unit_add_executable(
|
||||
nvf4_nvf4_f16_nvfp4_epilogue.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_mxf4xmxf4
|
||||
|
||||
BATCH_SOURCES ON
|
||||
@@ -67,7 +67,7 @@ cutlass_test_unit_add_executable(
|
||||
mxf4_mxf4_void_f16_nt_layout.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_mxf6xmxf6
|
||||
|
||||
BATCH_SOURCES ON
|
||||
@@ -77,7 +77,7 @@ cutlass_test_unit_add_executable(
|
||||
mxf6_mxf6_void_bf16_nt_layout.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_mxf8xmxf8
|
||||
|
||||
BATCH_SOURCES ON
|
||||
@@ -87,7 +87,7 @@ cutlass_test_unit_add_executable(
|
||||
mxf8_mxf8_void_f8_nt_layout.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_mxf6xmxf8
|
||||
|
||||
BATCH_SOURCES ON
|
||||
@@ -97,7 +97,7 @@ cutlass_test_unit_add_executable(
|
||||
mxf6_mxf8_void_f32_nt_layout.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_mxf8xmxf6
|
||||
|
||||
BATCH_SOURCES ON
|
||||
@@ -107,7 +107,7 @@ cutlass_test_unit_add_executable(
|
||||
mxf8_mxf6_f16_f8_nt_layout.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_mxf4xmxf8
|
||||
|
||||
BATCH_SOURCES ON
|
||||
@@ -117,7 +117,7 @@ cutlass_test_unit_add_executable(
|
||||
mxf4_mxf8_bf16_bf16_nt_layout.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_mxf8xmxf4
|
||||
|
||||
BATCH_SOURCES ON
|
||||
@@ -127,7 +127,7 @@ cutlass_test_unit_add_executable(
|
||||
mxf8_mxf4_f16_bf16_nt_layout.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_mxf6xmxf4
|
||||
|
||||
BATCH_SOURCES ON
|
||||
@@ -137,7 +137,7 @@ cutlass_test_unit_add_executable(
|
||||
mxf6_mxf4_f16_f16_nt_layout.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_bstensorop_sm100_mxf4xmxf6
|
||||
|
||||
BATCH_SOURCES ON
|
||||
|
||||
@@ -1,323 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Tests for device-wide GEMM interface
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
/// A Row B Col
|
||||
TEST(SM100_Device_Gemm_f16t_f16n_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) {
|
||||
using ElementA = cutlass::bfloat16_t;
|
||||
using ElementB = cutlass::bfloat16_t;
|
||||
using ElementC = void;
|
||||
using ElementD = float;
|
||||
using ElementCompute = float;
|
||||
using ElementAccumulator = float;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ClusterTileShape_MNK = Shape<_512,_512,_128>;
|
||||
using ClusterShape_MNK = Shape<_4,_4,_1>;
|
||||
using MmaTileShape_MNK = Shape<_256,_128,_128>;
|
||||
using OutputCtaShape = decltype(shape_div(ClusterTileShape_MNK{}, ClusterShape_MNK{}));
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape_MNK,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator, ElementCompute,
|
||||
ElementC, GmemLayoutC, 16,
|
||||
ElementD, GmemLayoutC, 16,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
ElementA, GmemLayoutA, 8,
|
||||
ElementB, GmemLayoutB, 8,
|
||||
ElementAccumulator,
|
||||
MmaTileShape_MNK, ClusterShape_MNK,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
auto pass = test::gemm::device::TestSmallFusion<Gemm>(1.0, 0);
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
/// A Col B Row
|
||||
TEST(SM100_Device_Gemm_f16n_f16t_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) {
|
||||
using ElementA = cutlass::bfloat16_t;
|
||||
using ElementB = cutlass::bfloat16_t;
|
||||
using ElementC = void;
|
||||
using ElementD = float;
|
||||
using ElementCompute = float;
|
||||
using ElementAccumulator = float;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ClusterTileShape_MNK = Shape<_512,_512,_128>;
|
||||
using ClusterShape_MNK = Shape<_4,_4,_1>;
|
||||
using MmaTileShape_MNK = Shape<_256,_128,_128>;
|
||||
using OutputCtaShape = decltype(shape_div(ClusterTileShape_MNK{}, ClusterShape_MNK{}));
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape_MNK,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator, ElementCompute,
|
||||
ElementC, GmemLayoutC, 16,
|
||||
ElementD, GmemLayoutC, 16,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
ElementA, GmemLayoutA, 8,
|
||||
ElementB, GmemLayoutB, 8,
|
||||
ElementAccumulator,
|
||||
MmaTileShape_MNK, ClusterShape_MNK,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
auto pass = test::gemm::device::TestSmallFusion<Gemm>(1.0, 0);
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
/// A Row B Row
|
||||
TEST(SM100_Device_Gemm_f16t_f16t_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) {
|
||||
using ElementA = cutlass::bfloat16_t;
|
||||
using ElementB = cutlass::bfloat16_t;
|
||||
using ElementC = void;
|
||||
using ElementD = float;
|
||||
using ElementCompute = float;
|
||||
using ElementAccumulator = float;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ClusterTileShape_MNK = Shape<_512,_512,_128>;
|
||||
using ClusterShape_MNK = Shape<_4,_4,_1>;
|
||||
using MmaTileShape_MNK = Shape<_256,_128,_128>;
|
||||
using OutputCtaShape = decltype(shape_div(ClusterTileShape_MNK{}, ClusterShape_MNK{}));
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape_MNK,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator, ElementCompute,
|
||||
ElementC, GmemLayoutC, 16,
|
||||
ElementD, GmemLayoutC, 16,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
ElementA, GmemLayoutA, 8,
|
||||
ElementB, GmemLayoutB, 8,
|
||||
ElementAccumulator,
|
||||
MmaTileShape_MNK, ClusterShape_MNK,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
auto pass = test::gemm::device::TestSmallFusion<Gemm>(1.0, 0);
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
/// A Col B Col
|
||||
TEST(SM100_Device_Gemm_f16n_f16n_f32t_tensorop_2sm_f32, 512x512x128_4x4x1) {
|
||||
using ElementA = cutlass::bfloat16_t;
|
||||
using ElementB = cutlass::bfloat16_t;
|
||||
using ElementC = void;
|
||||
using ElementD = float;
|
||||
using ElementCompute = float;
|
||||
using ElementAccumulator = float;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ClusterTileShape_MNK = Shape<_512,_512,_128>;
|
||||
using ClusterShape_MNK = Shape<_4,_4,_1>;
|
||||
using MmaTileShape_MNK = Shape<_256,_128,_128>;
|
||||
using OutputCtaShape = decltype(shape_div(ClusterTileShape_MNK{}, ClusterShape_MNK{}));
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape_MNK,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator, ElementCompute,
|
||||
ElementC, GmemLayoutC, 16,
|
||||
ElementD, GmemLayoutC, 16,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
ElementA, GmemLayoutA, 8,
|
||||
ElementB, GmemLayoutB, 8,
|
||||
ElementAccumulator,
|
||||
MmaTileShape_MNK, ClusterShape_MNK,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
auto pass = test::gemm::device::TestSmallFusion<Gemm>(1.0, 0);
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100_Device_Gemm_bf16t_bf16t_bf32_void_f32n_tensor_op, 128x256x64_1x2x1) {
|
||||
using ElementA = cutlass::bfloat16_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using ElementB = cutlass::bfloat16_t;
|
||||
using LayoutB = cutlass::layout::RowMajor;
|
||||
using ElementAccumulator = float;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
using MmaTileShape = Shape<_128,_128,_64>;
|
||||
using TileShape_MNK = Shape<_128,_256,_64>;
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
using OutputCtaShape = decltype(shape_div(TileShape_MNK{}, ClusterShape_MNK{}));
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape_MNK,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
void, LayoutC, 8,
|
||||
float, LayoutC, 8,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::half_t, LayoutA, 8,
|
||||
cutlass::half_t, LayoutB, 8,
|
||||
float,
|
||||
MmaTileShape, ClusterShape_MNK,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
auto pass = test::gemm::device::TestSmall<Gemm>(1.0, 0.0);
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
@@ -1,250 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief Tests for device-wide GEMM interface with stream-K scheduling
|
||||
*/
|
||||
|
||||
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/kernel/tile_scheduler.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
#include "cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp"
|
||||
#include "cutlass/epilogue/collective/default_epilogue.hpp"
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "gemm_testbed_3x.hpp"
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
using namespace cute;
|
||||
|
||||
TEST(SM100_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_stream_k, 128x256x64_1x2x1) {
|
||||
using ElementA = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using ElementB = cutlass::half_t;
|
||||
using LayoutB = cutlass::layout::RowMajor;
|
||||
using ElementAccumulator = float;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
using TileShape_MNK = Shape<_128,_256,_64>;
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
using AtomThrShape = decltype(shape_div(ClusterShape_MNK{}, Shape<_1,_1,_1>{}));
|
||||
using OutputCtaShape = decltype(shape_div(TileShape_MNK{}, ClusterShape_MNK{}));
|
||||
using MmaTileShape = decltype(shape_div(TileShape_MNK{}, AtomThrShape{}));
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape_MNK,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
cutlass::half_t, LayoutC, 8,
|
||||
cutlass::half_t, LayoutC, 8,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::half_t, LayoutA, 8,
|
||||
cutlass::half_t, LayoutB, 8,
|
||||
float,
|
||||
MmaTileShape, ClusterShape_MNK,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
cutlass::gemm::StreamKScheduler
|
||||
>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
using Testbed = Testbed3x<Gemm, cutlass::epilogue::thread::Identity>;
|
||||
bool result = TestSmall<Gemm, false /*force_legacy_epilogue*/, false /*apply_alignment_offset*/>(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048});
|
||||
EXPECT_TRUE(result);
|
||||
}
|
||||
|
||||
TEST(SM100_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_stream_k, 256x128x64_2x1x1) {
|
||||
using ElementA = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using ElementB = cutlass::half_t;
|
||||
using LayoutB = cutlass::layout::RowMajor;
|
||||
using ElementAccumulator = float;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
using TileShape_MNK = Shape<_256,_128,_64>;
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
using AtomThrShape = decltype(shape_div(ClusterShape_MNK{}, Shape<_1,_1,_1>{}));
|
||||
using OutputCtaShape = decltype(shape_div(TileShape_MNK{}, ClusterShape_MNK{}));
|
||||
using MmaTileShape = decltype(shape_div(TileShape_MNK{}, AtomThrShape{}));
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape_MNK,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
cutlass::half_t, LayoutC, 8,
|
||||
cutlass::half_t, LayoutC, 8,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::half_t, LayoutA, 8,
|
||||
cutlass::half_t, LayoutB, 8,
|
||||
float,
|
||||
MmaTileShape, ClusterShape_MNK,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
cutlass::gemm::StreamKScheduler
|
||||
>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
using Testbed = Testbed3x<Gemm, cutlass::epilogue::thread::Identity>;
|
||||
bool result = TestSmall<Gemm, false /*force_legacy_epilogue*/, false /*apply_alignment_offset*/>(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048});
|
||||
EXPECT_TRUE(result);
|
||||
}
|
||||
|
||||
TEST(SM100_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_stream_k, 256x256x64_2x2x1) {
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::RowMajor;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
using TileShape_MNK = Shape<_256,_256,_64>;
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
using AtomThrShape = decltype(shape_div(ClusterShape_MNK{}, Shape<_1,_1,_1>{}));
|
||||
using OutputCtaShape = decltype(shape_div(TileShape_MNK{}, ClusterShape_MNK{}));
|
||||
using MmaTileShape = decltype(shape_div(TileShape_MNK{}, AtomThrShape{}));
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape_MNK,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
cutlass::half_t, LayoutC, 8,
|
||||
cutlass::half_t, LayoutC, 8,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::half_t, LayoutA, 8,
|
||||
cutlass::half_t, LayoutB, 8,
|
||||
float,
|
||||
MmaTileShape, ClusterShape_MNK,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
cutlass::gemm::StreamKScheduler
|
||||
>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
using Testbed = Testbed3x<Gemm, cutlass::epilogue::thread::Identity>;
|
||||
bool result = TestSmall<Gemm, false /*force_legacy_epilogue*/, false /*apply_alignment_offset*/>(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048});
|
||||
EXPECT_TRUE(result);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM100_Device_Gemm_f16t_f16n_f32n_tensor_op_gmma_f32_stream_k, 256x128x64_2x4x1) {
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
using TileShape_MNK = Shape<_256,_256,_64>;
|
||||
using ClusterShape_MNK = Shape<_2,_4,_1>;
|
||||
using AtomThrShape = decltype(shape_div(ClusterShape_MNK{}, Shape<_1,_1,_1>{}));
|
||||
using OutputCtaShape = decltype(shape_div(TileShape_MNK{}, ClusterShape_MNK{}));
|
||||
using MmaTileShape = decltype(shape_div(TileShape_MNK{}, AtomThrShape{}));
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape_MNK,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
cutlass::half_t, LayoutC, 8,
|
||||
cutlass::half_t, LayoutC, 8,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::half_t, LayoutA, 8,
|
||||
cutlass::half_t, LayoutB, 8,
|
||||
float,
|
||||
MmaTileShape, ClusterShape_MNK,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
cutlass::gemm::StreamKScheduler
|
||||
>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
using Testbed = Testbed3x<Gemm, cutlass::epilogue::thread::Identity>;
|
||||
bool result = TestSmall<Gemm, false /*force_legacy_epilogue*/, false /*apply_alignment_offset*/>(1.0, 0.0, CheckEquality::EXACT, ScalarLoc::ON_DEVICE, VectorScale::ENABLED, {64, 1024, 2048});
|
||||
EXPECT_TRUE(result);
|
||||
}
|
||||
|
||||
#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
@@ -1,104 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (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 <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100_Device_Gemm_f16t_f16t_f32_void_f16n_tensor_op, 128x256x64_1x2x1) {
|
||||
using ElementA = cutlass::half_t;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using ElementB = cutlass::half_t;
|
||||
using LayoutB = cutlass::layout::RowMajor;
|
||||
using ElementAccumulator = float;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
using TileShape_MNK = Shape<_128,_256,_64>;
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
using AtomThrShape = decltype(shape_div(ClusterShape_MNK{}, Shape<_1,_1,_1>{}));
|
||||
using OutputCtaShape = decltype(shape_div(TileShape_MNK{}, ClusterShape_MNK{}));
|
||||
using MmaTileShape = decltype(shape_div(TileShape_MNK{}, AtomThrShape{}));
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape_MNK,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
void, LayoutC, 8,
|
||||
cutlass::half_t, LayoutC, 8,
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::half_t, LayoutA, 8,
|
||||
cutlass::half_t, LayoutB, 8,
|
||||
float,
|
||||
MmaTileShape, ClusterShape_MNK,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
auto pass = test::gemm::device::TestSmall<Gemm>(1.0, 0.0);
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
@@ -1,156 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Tests for device-wide GEMM interface
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100_Device_Gemm_e2m1t_e2m1n_f32t_tensorop_2sm_f32_runtime_datatype, 512x512x128_4x4x1) {
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_4,cute::_4,cute::_1>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
float, cutlass::layout::RowMajor, 4,
|
||||
float, cutlass::layout::RowMajor, 4,
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm,
|
||||
|
||||
cutlass::epilogue::fusion::LinearCombination<
|
||||
float,
|
||||
float,
|
||||
float,
|
||||
float
|
||||
>
|
||||
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::type_erased_dynamic_float4_t, cutlass::layout::RowMajor, 128,
|
||||
cutlass::type_erased_dynamic_float4_t, cutlass::layout::ColumnMajor, 128,
|
||||
float,
|
||||
cute::Shape<cute::_256, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_4,cute::_4,cute::_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
auto pass = TestRuntimeDataTypeSmall<Gemm>(cute::UMMA::MXF8F6F4Format::E2M1, cute::UMMA::MXF8F6F4Format::E2M1);
|
||||
EXPECT_TRUE(pass);
|
||||
|
||||
}
|
||||
|
||||
|
||||
TEST(SM100_Device_Gemm_e2m1t_e2m1n_f32t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) {
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
float, cutlass::layout::RowMajor, 4,
|
||||
float, cutlass::layout::RowMajor, 4,
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm,
|
||||
|
||||
cutlass::epilogue::fusion::LinearCombination<
|
||||
float,
|
||||
float,
|
||||
float,
|
||||
float
|
||||
>
|
||||
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::type_erased_dynamic_float4_t, cutlass::layout::RowMajor, 128,
|
||||
cutlass::type_erased_dynamic_float4_t, cutlass::layout::ColumnMajor, 128,
|
||||
float,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
auto pass = TestRuntimeDataTypeSmall<Gemm>(cute::UMMA::MXF8F6F4Format::E2M1, cute::UMMA::MXF8F6F4Format::E2M1);
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
@@ -1,156 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Tests for device-wide GEMM interface
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100_Device_Gemm_e3m2t_e2m3n_f32t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) {
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
float, cutlass::layout::RowMajor, 4,
|
||||
float, cutlass::layout::RowMajor, 4,
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm,
|
||||
|
||||
cutlass::epilogue::fusion::LinearCombination<
|
||||
float,
|
||||
float,
|
||||
float,
|
||||
float
|
||||
>
|
||||
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::type_erased_dynamic_float6_t, cutlass::layout::RowMajor, 128,
|
||||
cutlass::type_erased_dynamic_float6_t, cutlass::layout::ColumnMajor, 128,
|
||||
float,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
auto pass = TestRuntimeDataTypeSmall<Gemm>(cute::UMMA::MXF8F6F4Format::E3M2, cute::UMMA::MXF8F6F4Format::E2M3);
|
||||
EXPECT_TRUE(pass);
|
||||
|
||||
}
|
||||
|
||||
TEST(SM100_Device_Gemm_e3m2t_e2m3n_f32t_tensorop_1sm_f32_runtime_datatype, 512x512x128_4x4x1) {
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_4,cute::_4,cute::_1>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
float, cutlass::layout::RowMajor, 4,
|
||||
float, cutlass::layout::RowMajor, 4,
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm,
|
||||
|
||||
cutlass::epilogue::fusion::LinearCombination<
|
||||
float,
|
||||
float,
|
||||
float,
|
||||
float
|
||||
>
|
||||
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::type_erased_dynamic_float6_t, cutlass::layout::RowMajor, 128,
|
||||
cutlass::type_erased_dynamic_float6_t, cutlass::layout::ColumnMajor, 128,
|
||||
float,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_4,cute::_4,cute::_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
auto pass = TestRuntimeDataTypeSmall<Gemm>(cute::UMMA::MXF8F6F4Format::E3M2, cute::UMMA::MXF8F6F4Format::E2M3);
|
||||
EXPECT_TRUE(pass);
|
||||
|
||||
}
|
||||
|
||||
#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
@@ -1,109 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Tests for device-wide GEMM interface
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100_Device_Gemm_e4m3t_e2m1n_f32t_tensorop_2sm_f32_runtime_datatype, 256x128x128_2x2x1) {
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cute::Shape<cute::_256, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_1,cute::_1>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
float, cutlass::layout::RowMajor, 4,
|
||||
float, cutlass::layout::RowMajor, 4,
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm,
|
||||
|
||||
cutlass::epilogue::fusion::LinearCombination<
|
||||
float,
|
||||
float,
|
||||
float,
|
||||
float
|
||||
>
|
||||
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::type_erased_dynamic_float4_t, cutlass::layout::ColumnMajor, 128,
|
||||
float,
|
||||
cute::Shape<cute::_256, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_1,cute::_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
auto pass = TestRuntimeDataTypeSmall<Gemm>(cute::UMMA::MXF8F6F4Format::E4M3, cute::UMMA::MXF8F6F4Format::E2M1);
|
||||
EXPECT_TRUE(pass);
|
||||
|
||||
}
|
||||
|
||||
#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
@@ -1,297 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Tests for device-wide GEMM interface
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100_Device_Gemm_e5m2t_e4m3n_e4m3t_tensorop_2sm_f32_runtime_datatype, 256x128x128_2x2x1) {
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cute::Shape<cute::_256, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_1,cute::_1>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm,
|
||||
|
||||
cutlass::epilogue::fusion::LinearCombination<
|
||||
cutlass::float_e4m3_t,
|
||||
float,
|
||||
cutlass::float_e4m3_t,
|
||||
float
|
||||
>
|
||||
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16,
|
||||
float,
|
||||
cute::Shape<cute::_256, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_1,cute::_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
auto pass = TestRuntimeDataTypeSmall<Gemm>(cute::UMMA::MXF8F6F4Format::E5M2, cute::UMMA::MXF8F6F4Format::E4M3);
|
||||
EXPECT_TRUE(pass);
|
||||
|
||||
}
|
||||
|
||||
TEST(SM100_Device_Gemm_e5m2t_e4m3n_e4m3t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) {
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm,
|
||||
|
||||
cutlass::epilogue::fusion::LinearCombination<
|
||||
cutlass::float_e4m3_t,
|
||||
float,
|
||||
cutlass::float_e4m3_t,
|
||||
float
|
||||
>
|
||||
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16,
|
||||
float,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
auto pass = TestRuntimeDataTypeSmall<Gemm>(cute::UMMA::MXF8F6F4Format::E5M2, cute::UMMA::MXF8F6F4Format::E4M3);
|
||||
EXPECT_TRUE(pass);
|
||||
|
||||
}
|
||||
|
||||
TEST(SM100_Device_Gemm_e4m3t_e5m2n_e4m3t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) {
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm,
|
||||
|
||||
cutlass::epilogue::fusion::LinearCombination<
|
||||
cutlass::float_e4m3_t,
|
||||
float,
|
||||
cutlass::float_e4m3_t,
|
||||
float
|
||||
>
|
||||
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16,
|
||||
float,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
auto pass = TestRuntimeDataTypeSmall<Gemm>(cute::UMMA::MXF8F6F4Format::E4M3, cute::UMMA::MXF8F6F4Format::E5M2);
|
||||
EXPECT_TRUE(pass);
|
||||
|
||||
}
|
||||
|
||||
TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3t_tensorop_1sm_f32_runtime_datatype, 256x256x128_2x2x1) {
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::float_e4m3_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm,
|
||||
|
||||
cutlass::epilogue::fusion::LinearCombination<
|
||||
cutlass::float_e4m3_t,
|
||||
float,
|
||||
cutlass::float_e4m3_t,
|
||||
float
|
||||
>
|
||||
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16,
|
||||
float,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
auto pass = TestRuntimeDataTypeSmall<Gemm>(cute::UMMA::MXF8F6F4Format::E4M3, cute::UMMA::MXF8F6F4Format::E4M3);
|
||||
EXPECT_TRUE(pass);
|
||||
|
||||
}
|
||||
|
||||
TEST(SM100_Device_Gemm_e5m2t_e5m2n_e5m2t_tensorop_2sm_f32_runtime_datatype, 256x256x128_2x2x1) {
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cute::Shape<cute::_128, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
float, float,
|
||||
cutlass::float_e5m2_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::float_e5m2_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm,
|
||||
|
||||
cutlass::epilogue::fusion::LinearCombination<
|
||||
cutlass::float_e5m2_t,
|
||||
float,
|
||||
cutlass::float_e5m2_t,
|
||||
float
|
||||
>
|
||||
|
||||
>::CollectiveOp;
|
||||
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::RowMajor, 16,
|
||||
cutlass::type_erased_dynamic_float8_t, cutlass::layout::ColumnMajor, 16,
|
||||
float,
|
||||
cute::Shape<cute::_256, cute::_128, cute::_128>,
|
||||
cute::Shape<cute::_2,cute::_2,cute::_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using namespace test::gemm::device;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
auto pass = TestRuntimeDataTypeSmall<Gemm>(cute::UMMA::MXF8F6F4Format::E5M2, cute::UMMA::MXF8F6F4Format::E5M2);
|
||||
EXPECT_TRUE(pass);
|
||||
|
||||
}
|
||||
|
||||
#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
@@ -1,230 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Tests for device-wide GEMM interface
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
|
||||
#include "../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
////////////////////////////////////////// Test Batch alpha and beta //////////////////////////////////////////
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1cta_s32_batch_alpha_beta, 128x64x128_1x1x1) {
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementC = cutlass::float_e4m3_t;
|
||||
using ElementD = cutlass::float_e4m3_t;
|
||||
using ElementAccumulator = float;
|
||||
using ElementCompute = float;
|
||||
using ElementBias = cutlass::half_t;
|
||||
using ClusterTileShape = cute::Shape<_128,_64,Int<128 / sizeof(ElementA)>>;
|
||||
using ClusterShape = Shape<_1,_1,_1>;
|
||||
using AtomThrShape = decltype(shape_div(ClusterShape{}, Shape<_1,_1,_1>{}));
|
||||
using OutputCtaShape = decltype(shape_div(ClusterTileShape{}, ClusterShape{}));
|
||||
using MmaTileShape = decltype(shape_div(ClusterTileShape{}, AtomThrShape{}));
|
||||
|
||||
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm;
|
||||
|
||||
using FusionOperation = cutlass::epilogue::fusion::LinearCombination<
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementC,
|
||||
ElementBias
|
||||
>;
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator, ElementCompute,
|
||||
ElementC, LayoutC, 16 / sizeof(ElementC),
|
||||
ElementD, LayoutC, 16 / sizeof(ElementD),
|
||||
EpilogueSchedule,
|
||||
FusionOperation
|
||||
>::CollectiveOp;
|
||||
|
||||
using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100;
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
ElementA, LayoutA, 16 / sizeof(ElementA),
|
||||
ElementB, LayoutB, 16 / sizeof(ElementB),
|
||||
ElementAccumulator,
|
||||
MmaTileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
MainloopSchedule
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
auto pass = test::gemm::device::TestSmallFusion<Gemm, false, true, true>(1.0, 1.0); // beta is [1.0, 2.0]
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu_batch_alpha_beta, 128x128x128_1x1x1) {
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementC = cutlass::float_e4m3_t;
|
||||
using ElementD = cutlass::float_e4m3_t;
|
||||
using ElementAccumulator = float;
|
||||
using ElementCompute = float;
|
||||
using ElementBias = cutlass::half_t;
|
||||
using ClusterTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>;
|
||||
using ClusterShape = Shape<_1,_1,_1>;
|
||||
using AtomThrShape = decltype(shape_div(ClusterShape{}, Shape<_1,_1,_1>{}));
|
||||
using OutputCtaShape = decltype(shape_div(ClusterTileShape{}, ClusterShape{}));
|
||||
using MmaTileShape = decltype(shape_div(ClusterTileShape{}, AtomThrShape{}));
|
||||
|
||||
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm;
|
||||
using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct<
|
||||
cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>;
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator, ElementCompute,
|
||||
ElementC, LayoutC, 16 / sizeof(ElementC),
|
||||
ElementD, LayoutC, 16 / sizeof(ElementD),
|
||||
EpilogueSchedule,
|
||||
FusionOperation
|
||||
>::CollectiveOp;
|
||||
|
||||
using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100;
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
ElementA, LayoutA, 16 / sizeof(ElementA),
|
||||
ElementB, LayoutB, 16 / sizeof(ElementB),
|
||||
ElementAccumulator,
|
||||
MmaTileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
MainloopSchedule
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
auto pass = test::gemm::device::TestSmallFusion<Gemm, false, false, true>(1.0, 0.5); // beta is [0.5, 1.5]
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100_Device_Gemm_e4m3t_e4m3n_e4m3n_tensorop_1sm_f32_bias_relu__batch_alpha_beta0, 128x128x128_1x1x1) {
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementC = cutlass::float_e4m3_t;
|
||||
using ElementD = cutlass::float_e4m3_t;
|
||||
using ElementAccumulator = float;
|
||||
using ElementCompute = float;
|
||||
using ElementBias = cutlass::half_t;
|
||||
using ClusterTileShape = cute::Shape<_128,_128,Int<128 / sizeof(ElementA)>>;
|
||||
using ClusterShape = Shape<_1,_1,_1>;
|
||||
using AtomThrShape = decltype(shape_div(ClusterShape{}, Shape<_1,_1,_1>{}));
|
||||
using OutputCtaShape = decltype(shape_div(ClusterTileShape{}, ClusterShape{}));
|
||||
using MmaTileShape = decltype(shape_div(ClusterTileShape{}, AtomThrShape{}));
|
||||
|
||||
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm;
|
||||
using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct<
|
||||
cutlass::epilogue::thread::ReLU, ElementD, ElementCompute, ElementBias>;
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
OutputCtaShape, ClusterShape,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAccumulator, ElementCompute,
|
||||
ElementC, LayoutC, 16 / sizeof(ElementC),
|
||||
ElementD, LayoutC, 16 / sizeof(ElementD),
|
||||
EpilogueSchedule,
|
||||
FusionOperation
|
||||
>::CollectiveOp;
|
||||
|
||||
using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100;
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
ElementA, LayoutA, 16 / sizeof(ElementA),
|
||||
ElementB, LayoutB, 16 / sizeof(ElementB),
|
||||
ElementAccumulator,
|
||||
MmaTileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
MainloopSchedule
|
||||
>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
auto pass = test::gemm::device::TestSmallFusion<Gemm, false, false, true>(1.0, -1.0); // beta is [-1.0, 0.0]
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
71
test/unit/gemm/device/sm100_tensorop_gemm/CMakeLists.txt
Normal file
71
test/unit/gemm/device/sm100_tensorop_gemm/CMakeLists.txt
Normal file
@@ -0,0 +1,71 @@
|
||||
# Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
#
|
||||
|
||||
#
|
||||
|
||||
add_custom_target(
|
||||
cutlass_test_unit_gemm_device_sm100_tensorop
|
||||
DEPENDS
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_f16xf16
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_f8xf8
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_s8xs8
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_f16xf16
|
||||
|
||||
BATCH_SOURCES ON
|
||||
BATCH_SIZE 1
|
||||
|
||||
f16_f16_void_f32.cu
|
||||
f16_f16_f16_f16_fusion.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_f8xf8
|
||||
|
||||
BATCH_SOURCES ON
|
||||
BATCH_SIZE 1
|
||||
|
||||
f8_f8_void_f32.cu
|
||||
f8_f8_f16_f8_fusion.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_s8xs8
|
||||
|
||||
BATCH_SOURCES ON
|
||||
BATCH_SIZE 1
|
||||
|
||||
s8_s8_void_s32.cu
|
||||
s8_s8_s32_s32_fusion.cu
|
||||
)
|
||||
|
||||
add_subdirectory(narrow_precision)
|
||||
@@ -0,0 +1,607 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (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 <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Inference fprop fusions
|
||||
//
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16t_f16n_f16t_f16t_tensor_op_f32, 128x128x64_1x2x1_1sm_bias_relu) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = cutlass::half_t;
|
||||
constexpr int AlignC = 8;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = cutlass::half_t;
|
||||
constexpr int AlignD = 8;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// Z = alpha * acc + beta * C + per-row bias
|
||||
// D = ReLU(Z)
|
||||
using ElementBias = cutlass::half_t;
|
||||
using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltAct<
|
||||
cutlass::epilogue::thread::ReLU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16t_f16n_f16t_f16t_tensor_op_f32, 128x128x64_1x2x1_1sm_bias_gelu) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = cutlass::half_t;
|
||||
constexpr int AlignC = 8;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = cutlass::half_t;
|
||||
constexpr int AlignD = 8;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// Z = alpha * acc + beta * C + per-row bias
|
||||
// D = GELU(Z)
|
||||
using ElementBias = cutlass::half_t;
|
||||
using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltAct<
|
||||
cutlass::epilogue::thread::GELU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Training fprop fusions
|
||||
//
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16t_f16n_f16t_f16t_tensor_op_f32, 128x128x64_1x2x1_1sm_bias_relu_aux) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = cutlass::half_t;
|
||||
constexpr int AlignC = 8;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = cutlass::half_t;
|
||||
constexpr int AlignD = 8;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// Z = alpha * acc + beta * C + per-row bias
|
||||
// D = ReLU(Z)
|
||||
// For ReLU with uint1b_t aux, aux computes the dReLU/dZ gradient, i.e.
|
||||
// Aux(i) = Z(i) >= 0 ? 1 : 0
|
||||
using ElementBias = cutlass::half_t;
|
||||
using ElementAux = cutlass::uint1b_t;
|
||||
using GmemLayoutAux = GmemLayoutC;
|
||||
using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltActAux<
|
||||
GmemLayoutAux,
|
||||
cutlass::epilogue::thread::ReLU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementAux,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16t_f16n_f16t_f16t_tensor_op_f32, 128x128x64_1x2x1_1sm_bias_gelu_aux) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = cutlass::half_t;
|
||||
constexpr int AlignC = 8;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = cutlass::half_t;
|
||||
constexpr int AlignD = 8;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// Z = alpha * acc + beta * C + per-row bias
|
||||
// D = GELU(Z)
|
||||
// Aux = Z
|
||||
using ElementBias = cutlass::half_t;
|
||||
using ElementAux = cutlass::half_t;
|
||||
using GmemLayoutAux = GmemLayoutC;
|
||||
using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltActAux<
|
||||
GmemLayoutAux,
|
||||
cutlass::epilogue::thread::GELU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementAux,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Backprop fusions
|
||||
//
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16t_f16n_f16t_f16t_tensor_op_f32, 128x128x64_1x2x1_1sm_dbias_drelu) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = cutlass::half_t;
|
||||
constexpr int AlignC = 8;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = cutlass::half_t;
|
||||
constexpr int AlignD = 8;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// dY = alpha * acc + beta * C
|
||||
// D = dReLU(dY, Aux)
|
||||
// dBias = sum of columns of D
|
||||
using ElementBias = cutlass::half_t;
|
||||
using ElementAux = cutlass::uint1b_t;
|
||||
using GmemLayoutAux = GmemLayoutC;
|
||||
using FusionOperation = cutlass::epilogue::fusion::LinCombDeEltActDePerRowBias<
|
||||
GmemLayoutAux,
|
||||
cutlass::epilogue::thread::dReLU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementAux,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16t_f16n_f16t_f16t_tensor_op_f32, 128x128x64_1x2x1_1sm_dbias_dgelu) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = cutlass::half_t;
|
||||
constexpr int AlignC = 8;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = cutlass::half_t;
|
||||
constexpr int AlignD = 8;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// dY = alpha * acc + beta * C
|
||||
// D = dGELU(dY, Aux)
|
||||
// dBias = sum of columns of D
|
||||
using ElementBias = cutlass::half_t;
|
||||
using ElementAux = cutlass::half_t;
|
||||
using GmemLayoutAux = GmemLayoutC;
|
||||
using FusionOperation = cutlass::epilogue::fusion::LinCombDeEltActDePerRowBias<
|
||||
GmemLayoutAux,
|
||||
cutlass::epilogue::thread::dGELU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementAux,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
#endif
|
||||
655
test/unit/gemm/device/sm100_tensorop_gemm/f16_f16_void_f32.cu
Normal file
655
test/unit/gemm/device/sm100_tensorop_gemm/f16_f16_void_f32.cu
Normal file
@@ -0,0 +1,655 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (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 <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16n_f16t_void_f32n_tensor_op_f32, 64x64x64_4x1x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_64,_64,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_4,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_64,_64>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy <=== NEEDS TO BE 1SM otherwise ambigous.
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16t_f16n_void_f32t_tensor_op_f32, 64x128x64_1x4x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_64,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_4,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_128,_64>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16n_f16n_void_f32t_tensor_op_f32, 128x64x64_1x8x1_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_64,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_8,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_64,_64>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16t_f16t_void_f32n_tensor_op_f32, 128x128x64_2x8x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_8,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16n_f16t_void_f32n_tensor_op_f32, 128x64x64_2x4x1_2sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_64,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_4,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_64,_64>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16t_f16n_void_f32n_tensor_op_f32, 128x128x64_16x1x1_2sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_16,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_128,_64>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16n_f16n_void_f32n_tensor_op_f32, 256x64x64_4x1x1) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_64,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_4,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_64,_64>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_f16t_f16t_void_f32n_tensor_op_f32, 256x256x64_2x1x1) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::half_t;
|
||||
constexpr int AlignA = 8;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 8;
|
||||
using ElementB = cutlass::half_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_256,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_64>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
#endif
|
||||
430
test/unit/gemm/device/sm100_tensorop_gemm/f8_f8_f16_f8_fusion.cu
Normal file
430
test/unit/gemm/device/sm100_tensorop_gemm/f8_f8_f16_f8_fusion.cu
Normal file
@@ -0,0 +1,430 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (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 <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Inference fprop fusions
|
||||
//
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3t_e4m3n_f16t_e4m3t_tensor_op_f32, 128x128x128_1x2x1_1sm_bias_relu) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = cutlass::half_t;
|
||||
constexpr int AlignC = 8;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = cutlass::float_e4m3_t;
|
||||
constexpr int AlignD = 16;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// Z = alpha * scale_a * scale_b * acc + beta * scale_c * C + per-row bias
|
||||
// D = scale_d * ReLU(Z)
|
||||
using ElementBias = cutlass::half_t;
|
||||
using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct<
|
||||
cutlass::epilogue::thread::ReLU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3t_e4m3n_f16t_f32t_tensor_op_f32, 128x128x128_1x2x1_1sm_bias_relu) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = cutlass::half_t;
|
||||
constexpr int AlignC = 8;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// Z = alpha * scale_a * scale_b * acc + beta * scale_c * C + per-row bias
|
||||
// D = ReLU(Z)
|
||||
// scale_d is only applied if D is an fp8 type
|
||||
using ElementBias = float;
|
||||
using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltAct<
|
||||
cutlass::epilogue::thread::ReLU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Training fprop fusions
|
||||
//
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3t_e4m3n_f16t_e4m3t_tensor_op_f32, 128x128x128_1x2x1_1sm_bias_relu_amax_aux) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = cutlass::half_t;
|
||||
constexpr int AlignC = 8;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = cutlass::float_e4m3_t;
|
||||
constexpr int AlignD = 16;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// Z = alpha * scale_a * scale_b * acc + beta * scale_c * C + per-row bias
|
||||
// D = scale_d * ReLU(Z)
|
||||
// Amax_D = max absolute value of ReLU(Z)
|
||||
// Aux = Z
|
||||
// scale_d and Amax_D are only computed if D is fp8
|
||||
using ElementBias = cutlass::half_t;
|
||||
using ElementAmax = float;
|
||||
using ElementAux = float;
|
||||
using GmemLayoutAux = GmemLayoutC;
|
||||
using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux<
|
||||
GmemLayoutAux,
|
||||
cutlass::epilogue::thread::ReLU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementAux,
|
||||
ElementAmax,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3t_e4m3n_f16t_f32t_tensor_op_f32, 128x128x128_1x2x1_1sm_bias_relu_amax_aux) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = cutlass::half_t;
|
||||
constexpr int AlignC = 8;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_64>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_64>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// Z = alpha * scale_a * scale_b * acc + beta * scale_c * C + per-row bias
|
||||
// D = ReLU(Z)
|
||||
// Aux = scale_aux * Z
|
||||
// Amax_Aux = max absolute value of Z
|
||||
// scale_aux and Amax_Aux are only computed if Aux is fp8
|
||||
using ElementBias = float;
|
||||
using ElementAmax = float;
|
||||
using ElementAux = cutlass::float_e4m3_t;
|
||||
using GmemLayoutAux = GmemLayoutC;
|
||||
using FusionOperation = cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux<
|
||||
GmemLayoutAux,
|
||||
cutlass::epilogue::thread::ReLU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementAux,
|
||||
ElementAmax,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
|
||||
#endif
|
||||
659
test/unit/gemm/device/sm100_tensorop_gemm/f8_f8_void_f32.cu
Normal file
659
test/unit/gemm/device/sm100_tensorop_gemm/f8_f8_void_f32.cu
Normal file
@@ -0,0 +1,659 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (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 <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3n_e4m3t_void_f32n_tensor_op_f32, 64x64x128_4x1x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_64,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_4,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_64,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3t_e5m2n_void_f32t_tensor_op_f32, 64x128x128_1x4x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e5m2_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_64,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_4,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_128,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e5m2n_e4m3n_void_f32t_tensor_op_f32, 128x64x128_1x8x1) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_8,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_64,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e5m2t_e5m2t_void_f32n_tensor_op_f32, 128x128x128_2x8x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e5m2_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_8,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e5m2n_e4m3t_void_f32n_tensor_op_f32, 128x64x128_2x4x1_2sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_4,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_64,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3t_e4m3n_void_f32n_tensor_op_f32, 128x128x128_16x1x1_2sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_16,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_128,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3n_e4m3n_void_f32n_tensor_op_f32, 256x64x128_4x1x1) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_4,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_64,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3t_e4m3t_void_f32n_tensor_op_f32, 256x256x128_2x1x1_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,71 @@
|
||||
# Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
#
|
||||
|
||||
#
|
||||
|
||||
add_custom_target(
|
||||
cutlass_test_unit_gemm_device_sm100_tensorop_narrow_precision
|
||||
DEPENDS
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_f6f4xf6f4
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_f6f4xf8
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_f8xf6f4
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_f6f4xf6f4
|
||||
|
||||
BATCH_SOURCES ON
|
||||
BATCH_SIZE 1
|
||||
|
||||
f6f4_f6f4_void_f32_tn_layout.cu
|
||||
f6f4_f6f4_void_f32_nn_layout.cu
|
||||
f6f4_f6f4_void_f32_nt_layout.cu
|
||||
f6f4_f6f4_void_f32_tt_layout.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_f6f4xf8
|
||||
|
||||
BATCH_SOURCES ON
|
||||
BATCH_SIZE 1
|
||||
|
||||
f6f4_f8_void_f32_tn_layout.cu
|
||||
f6f4_f8_void_f32_nt_layout.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_gemm_device_add_executable_split_file(
|
||||
cutlass_test_unit_gemm_device_tensorop_sm100_f8xf6f4
|
||||
|
||||
BATCH_SOURCES ON
|
||||
BATCH_SIZE 1
|
||||
|
||||
f8_f6f4_void_f32_tn_layout.cu
|
||||
f8_f6f4_void_f32_nt_layout.cu
|
||||
)
|
||||
@@ -0,0 +1,687 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Unit tests for {f6f4}x{f6f4} Gemm
|
||||
|
||||
* A tensor:
|
||||
* Types: {e2m1,e2m3,e3m2}
|
||||
* Alignment: 128 elements
|
||||
* B tensor:
|
||||
* Types: {e2m1,e2m3,e3m2}
|
||||
* Alignment: 128 elements
|
||||
* Mma Tile Shapes supported:
|
||||
Support Matrix (Y: Yes, N: No)
|
||||
| 1/2 SM | Mma Tile Shape | TN | TT | NT | NN | Dispatch Policy |
|
||||
|--------|----------------|----|----|----|----|------------------------------------|
|
||||
| 1SM | 64x64x128 | Y | N | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x128x128 | Y | Y | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x192x128 | Y | N | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x256x128 | Y | Y | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x64x128 | Y | N | N | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x128x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x192x128 | Y | N | N | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 2SM | 128x64x128 | Y | N | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x128x128 | Y | N | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x192x128 | Y | N | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x256x128 | Y | Y | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x64x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x128x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x192x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../../common/cutlass_unit_test.h"
|
||||
#include "../../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e2m3n_void_f32n_tensor_op_f32, 128x64x128_4x1x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m3_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_4,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_64,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e3m2n_e2m1n_void_f32n_tensor_op_f32, 128x128x128_2x1x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e3m2_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e2m1n_void_f32n_tensor_op_f32, 128x192x128_2x4x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_192,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_4,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_192,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m3n_e3m2n_void_f32n_tensor_op_f32, 128x256x128_2x2x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m3_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e3m2_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e3m2n_e3m2n_void_f32n_tensor_op_f32, 256x64x128_4x1x1_2sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e3m2_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e3m2_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_4,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_64,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e2m1n_void_f32n_tensor_op_f32, 256x128x128_2x1x1_2sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e2m3n_void_f32n_tensor_op_f32, 256x192x128_2x4x1_2sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m3_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_192,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_4,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_192,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e2m1n_void_f32n_tensor_op_f32, 256x256x128_2x2x1_2sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,310 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Unit tests for {f6f4}x{f6f4} Gemm
|
||||
|
||||
* A tensor:
|
||||
* Types: {e2m1,e2m3,e3m2}
|
||||
* Alignment: 128 elements
|
||||
* B tensor:
|
||||
* Types: {e2m1,e2m3,e3m2}
|
||||
* Alignment: 128 elements
|
||||
* Mma Tile Shapes supported:
|
||||
Support Matrix (Y: Yes, N: No)
|
||||
| 1/2 SM | Mma Tile Shape | TN | TT | NT | NN | Dispatch Policy |
|
||||
|--------|----------------|----|----|----|----|------------------------------------|
|
||||
| 1SM | 64x64x128 | Y | N | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x128x128 | Y | Y | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x192x128 | Y | N | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x256x128 | Y | Y | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x64x128 | Y | N | N | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x128x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x192x128 | Y | N | N | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 2SM | 128x64x128 | Y | N | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x128x128 | Y | N | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x192x128 | Y | N | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x256x128 | Y | Y | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x64x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x128x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x192x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../../common/cutlass_unit_test.h"
|
||||
#include "../../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e2m3t_void_f32n_tensor_op_f32, 128x128x128_2x1x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e2m1t_void_f32n_tensor_op_f32, 128x256x128_2x2x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e3m2n_e2m1t_void_f32n_tensor_op_f32, 256x256x128_2x2x1_2sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e3m2_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,536 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Unit tests for {f6f4}x{f6f4} Gemm
|
||||
|
||||
* A tensor:
|
||||
* Types: {e2m1,e2m3,e3m2}
|
||||
* Alignment: 128 elements
|
||||
* B tensor:
|
||||
* Types: {e2m1,e2m3,e3m2}
|
||||
* Alignment: 128 elements
|
||||
* Mma Tile Shapes supported:
|
||||
Support Matrix (Y: Yes, N: No)
|
||||
| 1/2 SM | Mma Tile Shape | TN | TT | NT | NN | Dispatch Policy |
|
||||
|--------|----------------|----|----|----|----|------------------------------------|
|
||||
| 1SM | 64x64x128 | Y | N | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x128x128 | Y | Y | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x192x128 | Y | N | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x256x128 | Y | Y | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x64x128 | Y | N | N | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x128x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x192x128 | Y | N | N | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 2SM | 128x64x128 | Y | N | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x128x128 | Y | N | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x192x128 | Y | N | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x64x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x128x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x192x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../../common/cutlass_unit_test.h"
|
||||
#include "../../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
TEST(SM100Only_Device_Gemm_e2m1t_e2m1t_void_f32n_tensor_op_f32, 64x128x128_2x1x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_64,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_128,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m3t_e2m3t_void_f32n_tensor_op_f32, 64x256x128_2x2x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m3_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_64,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1t_e2m1t_void_f32n_tensor_op_f32, 128x128x128_2x1x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e3m2t_e3m2t_void_f32n_tensor_op_f32, 128x256x128_2x2x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e3m2_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e3m2_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1t_e2m1t_void_f32n_tensor_op_f32, 128x256x128_2x2x1_2sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_256,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1t_e2m3t_void_f32n_tensor_op_f32, 256x256x128_2x2x1_2sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,686 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Unit tests for {f6f4}xf8 Gemm
|
||||
|
||||
* A tensor:
|
||||
* Types: {e2m1,e2m3,e3m2}
|
||||
* Alignment: 128 elements
|
||||
* B tensor:
|
||||
* Types: {e5m2,e4m3}
|
||||
* Alignment: 16 elements
|
||||
* Mma Tile Shapes supported:
|
||||
Support Matrix (Y: Yes, N: No)
|
||||
| 1/2 SM | Mma Tile Shape | TN | TT | NT | NN | Dispatch Policy |
|
||||
|--------|----------------|----|----|----|----|------------------------------------|
|
||||
| 1SM | 64x64x128 | Y | Y | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x128x128 | Y | Y | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x192x128 | Y | Y | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x256x128 | Y | Y | N | N | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x64x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x128x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x192x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 2SM | 128x64x128 | Y | Y | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x128x128 | Y | Y | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x192x128 | Y | Y | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x256x128 | Y | Y | N | N | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x64x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x128x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x192x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../../common/cutlass_unit_test.h"
|
||||
#include "../../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e4m3t_void_f32n_tensor_op_f32, 128x64x128_4x1x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_4,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_64,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m3n_e5m2t_void_f32n_tensor_op_f32, 128x128x128_2x1x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m3_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e5m2_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m3n_e4m3t_void_f32n_tensor_op_f32, 128x192x128_2x4x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m3_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_192,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_4,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_192,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e3m2n_e5m2t_void_f32n_tensor_op_f32, 128x256x128_2x2x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e3m2_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e5m2_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e4m3t_void_f32n_tensor_op_f32, 256x64x128_4x1x1_2sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_4,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_64,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e5m2t_void_f32n_tensor_op_f32, 256x128x128_2x1x1_2sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e5m2_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m1n_e4m3t_void_f32n_tensor_op_f32, 256x192x128_2x4x1_2sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_192,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_4,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_192,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e2m3n_e4m3t_void_f32n_tensor_op_f32, 256x256x128_2x2x1_2sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e2m3_t;
|
||||
constexpr int AlignA = 128;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,538 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
|
||||
|
||||
/*! \file
|
||||
\brief Unit tests for {f8}x{f6f4} Gemm
|
||||
|
||||
* A tensor:
|
||||
* Types: {e5m2,e4m3}
|
||||
* Alignment: 16 elements
|
||||
* B tensor:
|
||||
* Types: {e2m1,e2m3,e3m2}
|
||||
* Alignment: 128 elements
|
||||
* Mma Tile Shapes supported:
|
||||
Support Matrix (Y: Yes, N: No)
|
||||
| 1/2 SM | Mma Tile Shape | TN | TT | NT | NN | Dispatch Policy |
|
||||
|--------|----------------|----|----|----|----|------------------------------------|
|
||||
| 1SM | 64x64x128 | Y | N | N | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x128x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x192x128 | Y | N | N | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 64x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x64x128 | Y | N | N | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x128x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x192x128 | Y | N | N | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 1SM | 128x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized1SmSm100` |
|
||||
| 2SM | 128x64x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x128x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x192x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 128x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x64x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x128x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x192x128 | Y | N | N | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
| 2SM | 256x256x128 | Y | Y | Y | Y | `KernelTmaWarpSpecialized2SmSm100` |
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../../common/cutlass_unit_test.h"
|
||||
#include "../../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
TEST(SM100Only_Device_Gemm_e4m3n_e2m3t_void_f32n_tensor_op_f32, 64x128x128_2x1x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_64,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_128,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e5m2n_e3m2t_void_f32n_tensor_op_f32, 64x256x128_2x2x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e3m2_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_64,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3n_e2m1t_void_f32n_tensor_op_f32, 128x128x128_2x1x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e5m2n_e2m3t_void_f32n_tensor_op_f32, 128x256x128_2x2x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m3_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e4m3n_e3m2t_void_f32n_tensor_op_f32, 128x256x128_2x2x1_2sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e3m2_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_256,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_e5m2n_e2m1t_void_f32n_tensor_op_f32, 256x256x128_2x2x1_2sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = cutlass::float_e5m2_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 128;
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = float;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = float;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,226 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (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 <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100Only_Device_Gemm_s8t_s8n_s32t_s32t_tensor_op_f32, 128x128x128_1x2x1_1sm_rowscale_bias_relu) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = int8_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = int8_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = int32_t;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = int32_t;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = int32_t;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// Z = per-row alpha * acc + per-row beta * C + per-row bias
|
||||
// D = ReLU(Z)
|
||||
using ElementBias = int32_t;
|
||||
using FusionOperation = cutlass::epilogue::fusion::PerRowLinCombPerRowBiasEltAct<
|
||||
cutlass::epilogue::thread::ReLU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_s8t_s8n_s32t_s32t_tensor_op_f32, 128x128x128_1x2x1_1sm_colscale_bias_relu) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = int8_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = int8_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = int32_t;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = int32_t;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = int32_t;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = float;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_2,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
// Epilogue fusion operation
|
||||
// Z = per-col alpha * acc + per-col beta * C + per-col bias
|
||||
// D = ReLU(Z)
|
||||
using ElementBias = int32_t;
|
||||
using FusionOperation = cutlass::epilogue::fusion::PerColLinCombPerColBiasEltAct<
|
||||
cutlass::epilogue::thread::ReLU,
|
||||
ElementD,
|
||||
ElementCompute,
|
||||
ElementBias,
|
||||
ElementC>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto, // Epilogue schedule policy
|
||||
FusionOperation // Epilogue fusion operation
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
#endif
|
||||
659
test/unit/gemm/device/sm100_tensorop_gemm/s8_s8_void_s32.cu
Normal file
659
test/unit/gemm/device/sm100_tensorop_gemm/s8_s8_void_s32.cu
Normal file
@@ -0,0 +1,659 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (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 <iostream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/dispatch_policy.hpp"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "../../../common/cutlass_unit_test.h"
|
||||
|
||||
#include "../gemm_testbed_3x.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
TEST(SM100Only_Device_Gemm_s8n_s8t_void_s32n_tensor_op_f32, 64x64x128_4x1x1_1sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = int8_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = int8_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = int32_t;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = int32_t;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = int32_t;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_64,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_4,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_64,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_s8t_s8n_void_s32t_tensor_op_f32, 64x128x128_1x4x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = int8_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = int8_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = int32_t;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = int32_t;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = int32_t;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_64,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_4,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_128,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_s8n_s8n_void_s32t_tensor_op_f32, 128x64x128_1x8x1_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = int8_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = int8_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::RowMajor;
|
||||
using ElementD = int32_t;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::RowMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = int32_t;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = int32_t;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_1,_8,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_64,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_s8t_s8t_void_s32n_tensor_op_f32, 128x128x128_2x8x1_1sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = int8_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = int8_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = int32_t;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = int32_t;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = int32_t;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_8,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_128,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized1Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized1SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
|
||||
TEST(SM100Only_Device_Gemm_s8n_s8t_void_s32n_tensor_op_f32, 128x64x128_2x4x1_2sm_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = int8_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = int8_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = int32_t;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = int32_t;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = int32_t;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_4,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_64,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
|
||||
TEST(SM100Only_Device_Gemm_s8t_s8n_void_s32n_tensor_op_f32, 128x128x128_16x1x1_2sm) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = int8_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = int8_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = int32_t;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = int32_t;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = int32_t;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_128,_128,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_16,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_64,_128,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::TmaWarpSpecialized2Sm // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::KernelTmaWarpSpecialized2SmSm100 // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_s8n_s8n_void_s32n_tensor_op_f32, 256x64x128_4x1x1_streamK) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = int8_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::ColumnMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = int8_t;
|
||||
using GmemLayoutB = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = int32_t;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = int32_t;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = int32_t;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_64,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_4,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_64,_128>;
|
||||
|
||||
// Tile Scheduler
|
||||
using TileScheduler = cutlass::gemm::StreamKScheduler;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
TileScheduler
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
|
||||
TEST(SM100Only_Device_Gemm_s8t_s8t_void_s32n_tensor_op_f32, 256x256x128_2x1x1) {
|
||||
// Describe A and B tensors
|
||||
using ElementA = int8_t;
|
||||
constexpr int AlignA = 16;
|
||||
using GmemLayoutA = cutlass::layout::RowMajor;
|
||||
constexpr int AlignB = 16;
|
||||
using ElementB = int8_t;
|
||||
using GmemLayoutB = cutlass::layout::RowMajor;
|
||||
|
||||
// Describe C and D tensors
|
||||
using ElementC = void;
|
||||
constexpr int AlignC = 4;
|
||||
using GmemLayoutC = cutlass::layout::ColumnMajor;
|
||||
using ElementD = int32_t;
|
||||
constexpr int AlignD = 4;
|
||||
using GmemLayoutD = cutlass::layout::ColumnMajor;
|
||||
|
||||
// Mma's accumulator type
|
||||
using ElementAccumulator = int32_t;
|
||||
// Epilogue computation's precision type
|
||||
using ElementCompute = int32_t;
|
||||
|
||||
// Tile and cluster shapes
|
||||
// Collective MMA takes tile shape of the MMA operation as input
|
||||
using MmaTileShape_MNK = Shape<_256,_256,_128>;
|
||||
// Cluster size for multicast
|
||||
using ClusterShape_MNK = Shape<_2,_1,_1>;
|
||||
// Collective Epilogue takes the output tile shape for 1 CTA
|
||||
using PerSmTileShape_MNK = Shape<_128,_256,_128>;
|
||||
|
||||
//
|
||||
// Construct CollectiveEpilogue
|
||||
//
|
||||
|
||||
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
PerSmTileShape_MNK, ClusterShape_MNK, // Epilogue tile shape, and cluster shape
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, // Epilogue subtile shape. Auto will find a suitable tile shape
|
||||
ElementAccumulator, ElementCompute, // Mma instr's accumulator type and compute precision for epilogue
|
||||
ElementC, GmemLayoutC, AlignC, // C tensor description
|
||||
ElementD, GmemLayoutD, AlignD, // D tensor description
|
||||
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
|
||||
>::CollectiveOp;
|
||||
|
||||
//
|
||||
// Construct CollectiveMainloop
|
||||
//
|
||||
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, // Arch and Tensorop spec
|
||||
ElementA, GmemLayoutA, AlignA, // A tensor elem type, layout and alignment requirement
|
||||
ElementB, GmemLayoutB, AlignB, // B tensor elem type, layout and alignment requirement
|
||||
ElementAccumulator, // Mma instruction accumulator type
|
||||
MmaTileShape_MNK, ClusterShape_MNK, // Mma instruction tile shape, cluster shape
|
||||
// Epilogue's SMEM usage that needs to be subtracted from overall SMEM capacity
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto // Kernel schedule policy. Auto or using targeted scheduling policy
|
||||
>::CollectiveOp;
|
||||
|
||||
// Create Gemm Kernel using CollectiveEpilogue and CollectiveMainloop created by the builders
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
// Run tests
|
||||
auto pass = test::gemm::device::TestAll<Gemm>();
|
||||
// Check results
|
||||
EXPECT_TRUE(pass);
|
||||
}
|
||||
#endif
|
||||
@@ -282,100 +282,100 @@ struct PipelineTest {
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
//Cluster1x2 Stage4
|
||||
TEST(SM100_Verify_PipelineClusterLaunchControlAsync_WS, Cluster1x2_Stage4) {
|
||||
Options options;
|
||||
OptionsClusterLaunch options;
|
||||
options.grid_dim = {32,32,1};
|
||||
using ClusterShape = cutlass::gemm::GemmShape<1, 2, 1>;
|
||||
static constexpr uint32_t Stages = 4;
|
||||
using Test = PipelineTest<Stages, ClusterShape>;
|
||||
Testbed<Test> testbed(options);
|
||||
TestbedClusterLaunch<Test> testbed(options);
|
||||
EXPECT_TRUE(testbed.verification());
|
||||
}
|
||||
|
||||
//Cluster2x1 Stage4
|
||||
TEST(SM100_Verify_PipelineClusterLaunchControlAsync_WS, Cluster2x1_Stage4) {
|
||||
Options options;
|
||||
OptionsClusterLaunch options;
|
||||
options.grid_dim = {32,32,1};
|
||||
using ClusterShape = cutlass::gemm::GemmShape<2, 1, 1>;
|
||||
static constexpr uint32_t Stages = 4;
|
||||
using Test = PipelineTest<Stages, ClusterShape>;
|
||||
Testbed<Test> testbed(options);
|
||||
TestbedClusterLaunch<Test> testbed(options);
|
||||
EXPECT_TRUE(testbed.verification());
|
||||
}
|
||||
|
||||
//Cluster2x2 Stage4
|
||||
TEST(SM100_Verify_PipelineClusterLaunchControlAsync_WS, Cluster2x2_Stage4) {
|
||||
Options options;
|
||||
OptionsClusterLaunch options;
|
||||
options.grid_dim = {32,32,1};
|
||||
using ClusterShape = cutlass::gemm::GemmShape<2, 2, 1>;
|
||||
static constexpr uint32_t Stages = 4;
|
||||
using Test = PipelineTest<Stages, ClusterShape>;
|
||||
Testbed<Test> testbed(options);
|
||||
TestbedClusterLaunch<Test> testbed(options);
|
||||
EXPECT_TRUE(testbed.verification());
|
||||
}
|
||||
|
||||
//Cluster1x1 Stage3
|
||||
TEST(SM100_Verify_PipelineClusterLaunchControlAsync_WS, Cluster1x1_Stage3) {
|
||||
Options options;
|
||||
OptionsClusterLaunch options;
|
||||
options.grid_dim = {32,32,1};
|
||||
using ClusterShape = cutlass::gemm::GemmShape<1, 1, 1>;
|
||||
static constexpr uint32_t Stages = 3;
|
||||
using Test = PipelineTest<Stages, ClusterShape>;
|
||||
Testbed<Test> testbed(options);
|
||||
TestbedClusterLaunch<Test> testbed(options);
|
||||
EXPECT_TRUE(testbed.verification());
|
||||
}
|
||||
|
||||
//Cluster1x4 Stage4
|
||||
TEST(SM100_Verify_PipelineClusterLaunchControlAsync_WS, Cluster1x4_Stage4) {
|
||||
Options options;
|
||||
OptionsClusterLaunch options;
|
||||
options.grid_dim = {32,32,1};
|
||||
using ClusterShape = cutlass::gemm::GemmShape<1, 4, 1>;
|
||||
static constexpr uint32_t Stages = 4;
|
||||
using Test = PipelineTest<Stages, ClusterShape>;
|
||||
Testbed<Test> testbed(options);
|
||||
TestbedClusterLaunch<Test> testbed(options);
|
||||
EXPECT_TRUE(testbed.verification());
|
||||
}
|
||||
|
||||
//Cluster4x1 Stage4
|
||||
TEST(SM100_Verify_PipelineClusterLaunchControlAsync_WS, Cluster4x1_Stage4) {
|
||||
Options options;
|
||||
OptionsClusterLaunch options;
|
||||
options.grid_dim = {32,32,1};
|
||||
using ClusterShape = cutlass::gemm::GemmShape<4, 1, 1>;
|
||||
static constexpr uint32_t Stages = 4;
|
||||
using Test = PipelineTest<Stages, ClusterShape>;
|
||||
Testbed<Test> testbed(options);
|
||||
TestbedClusterLaunch<Test> testbed(options);
|
||||
EXPECT_TRUE(testbed.verification());
|
||||
}
|
||||
|
||||
//Cluster2x4 Stage4
|
||||
TEST(SM100_Verify_PipelineClusterLaunchControlAsync_WS, Cluster2x4_Stage4) {
|
||||
Options options;
|
||||
OptionsClusterLaunch options;
|
||||
options.grid_dim = {32,32,1};
|
||||
using ClusterShape = cutlass::gemm::GemmShape<2, 4, 1>;
|
||||
static constexpr uint32_t Stages = 4;
|
||||
using Test = PipelineTest<Stages, ClusterShape>;
|
||||
Testbed<Test> testbed(options);
|
||||
TestbedClusterLaunch<Test> testbed(options);
|
||||
EXPECT_TRUE(testbed.verification());
|
||||
}
|
||||
|
||||
//Cluster4x2 Stage4
|
||||
TEST(SM100_Verify_PipelineClusterLaunchControlAsync_WS, Cluster4x2_Stage4) {
|
||||
Options options;
|
||||
OptionsClusterLaunch options;
|
||||
options.grid_dim = {32,32,1};
|
||||
using ClusterShape = cutlass::gemm::GemmShape<4, 2, 1>;
|
||||
static constexpr uint32_t Stages = 4;
|
||||
using Test = PipelineTest<Stages, ClusterShape>;
|
||||
Testbed<Test> testbed(options);
|
||||
TestbedClusterLaunch<Test> testbed(options);
|
||||
EXPECT_TRUE(testbed.verification());
|
||||
}
|
||||
|
||||
//Cluster4x4 Stage4
|
||||
TEST(SM100_Verify_PipelineClusterLaunchControlAsync_WS, Cluster4x4_Stage4) {
|
||||
Options options;
|
||||
OptionsClusterLaunch options;
|
||||
options.grid_dim = {32,32,1};
|
||||
using ClusterShape = cutlass::gemm::GemmShape<4, 4, 1>;
|
||||
static constexpr uint32_t Stages = 4;
|
||||
using Test = PipelineTest<Stages, ClusterShape>;
|
||||
Testbed<Test> testbed(options);
|
||||
TestbedClusterLaunch<Test> testbed(options);
|
||||
EXPECT_TRUE(testbed.verification());
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -51,7 +51,7 @@
|
||||
#include "cutlass/util/command_line.h"
|
||||
|
||||
// Command line test options
|
||||
struct Options {
|
||||
struct OptionsClusterLaunch {
|
||||
//
|
||||
// Data Members
|
||||
//
|
||||
@@ -95,10 +95,10 @@ struct Options {
|
||||
//
|
||||
|
||||
template<typename Pipeline>
|
||||
class Testbed {
|
||||
class TestbedClusterLaunch {
|
||||
private:
|
||||
// Commandline options
|
||||
Options options;
|
||||
OptionsClusterLaunch options;
|
||||
|
||||
bool run_test() {
|
||||
|
||||
@@ -114,7 +114,7 @@ private:
|
||||
|
||||
|
||||
public:
|
||||
Testbed(Options const &options_) : options(options_) {
|
||||
TestbedClusterLaunch(OptionsClusterLaunch const &options_) : options(options_) {
|
||||
int device_id = 0;
|
||||
cudaDeviceProp device_prop;
|
||||
CUTE_CHECK_ERROR(cudaSetDevice(device_id));
|
||||
|
||||
Reference in New Issue
Block a user