v3.8.0 update (#2082)

* 3.8 update

* fix Markus' name

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
Yujia Zhai
2025-02-06 18:33:40 -08:00
committed by GitHub
parent affd1b693d
commit 833f6990e0
168 changed files with 24945 additions and 3436 deletions

View File

@@ -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

View File

@@ -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)

View File

@@ -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()

View File

@@ -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

View File

@@ -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

View File

@@ -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

View 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

View File

@@ -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

View File

@@ -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

View 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

View File

@@ -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

View File

@@ -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

View File

@@ -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),

View File

@@ -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);

View File

@@ -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

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View 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)

View File

@@ -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

View 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

View 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

View 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

View 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_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
)

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View 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

View File

@@ -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

View File

@@ -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));