From 04a68a7ec85cff6233127c7fb24ea4f422875b2b Mon Sep 17 00:00:00 2001 From: Wang Zupeng Date: Sat, 12 Sep 2026 15:07:03 +0800 Subject: [PATCH] Expose resolved GEMM schedules in library descriptions --- media/docs/cpp/code_organization.md | 42 +++ test/unit/CMakeLists.txt | 1 + test/unit/library/CMakeLists.txt | 40 +++ test/unit/library/gemm_description.cu | 327 ++++++++++++++++++ .../include/cutlass/library/descriptions.h | 18 + tools/library/include/cutlass/library/types.h | 39 +++ .../src/block_scaled_gemm_operation_3x.hpp | 2 + .../src/blockwise_gemm_operation_3x.hpp | 2 + tools/library/src/gemm_operation_3x.hpp | 3 + .../library/src/gemm_schedule_description.hpp | 243 +++++++++++++ 10 files changed, 717 insertions(+) create mode 100644 test/unit/library/CMakeLists.txt create mode 100644 test/unit/library/gemm_description.cu create mode 100644 tools/library/src/gemm_schedule_description.hpp diff --git a/media/docs/cpp/code_organization.md b/media/docs/cpp/code_organization.md index 03a56875d9..3d1b614037 100644 --- a/media/docs/cpp/code_organization.md +++ b/media/docs/cpp/code_organization.md @@ -121,6 +121,48 @@ python/ When CMake is executed, the CUTLASS Instance Library generator scripts are executed to construct a set of instantiations in `build/tools/library/generated/`. +### Runtime GEMM schedule metadata + +The instance library exposes the selected CUTLASS 3.x mainloop and epilogue strategies +through `GemmDescription`. These fields describe the resolved collective policies, +including when the kernel was constructed with `KernelScheduleAuto` or +`EpilogueScheduleAuto`: + +| Field | Meaning | +| --- | --- | +| `mainloop_schedule` | Multistage, TMA, warp-specialized, ping-pong, cooperative, or Blackwell 1-SM / 2-SM execution | +| `mainloop_load` | `cp.async`, TMA, or mixed TMA / `cp.async` operand loads | +| `epilogue_schedule` | Direct stores without shared memory, or a TMA epilogue | +| `accumulation_kind` | Default accumulation or the opt-in Hopper FP8 fast-accumulation policy | + +For example, a client can select cooperative Hopper kernels without parsing their names: + +```cpp +using namespace cutlass::library; +auto const& description = static_cast(operation.description()); +if (description.mainloop_schedule == MainloopScheduleKind::kWarpSpecializedCooperative && + description.accumulation_kind == AccumulationKind::kDefault) { + // Consider this kernel for the application's dispatch table. +} +``` + +Use the description type matching the operation kind. `BlockScaledGemmDescription` +and `BlockwiseGemmDescription` expose the same fields; grouped operations expose them +through `GroupedGemmDescription::gemm`. The 1-SM / 2-SM distinction comes from the +selected MMA atom's CTA group, rather than the cluster size. + +Unrecognized custom policies and descriptions populated by older operation wrappers +retain `kUnknown` fields. Clients should handle `kUnknown` explicitly. In particular, +the pre-Hopper `KernelMultistage` tag alone does not distinguish synchronous loads +from `cp.async`, so its load kind is unknown. Existing constructor arguments are +unchanged; rebuild the instance library and its clients together when updating the +public description structures. + +The `cutlass_test_unit_library` target checks this metadata using real collective +builders and operation descriptions. It requires a CUDA toolkit for compilation but +runs on the host without launching kernels or requiring a GPU. Architecture-specific +cases are enabled according to the toolkit's support for the respective builders. + ### CUTLASS Profiler The CUTLASS Profiler is designed to load the CUTLASS Instance Library and execute all operations contained therein. diff --git a/test/unit/CMakeLists.txt b/test/unit/CMakeLists.txt index 627d52a6f6..b1633e5d7c 100644 --- a/test/unit/CMakeLists.txt +++ b/test/unit/CMakeLists.txt @@ -173,6 +173,7 @@ add_custom_target(cutlass_test_unit) add_custom_target(test_unit) set(SUBDIRS + library core cute gemm diff --git a/test/unit/library/CMakeLists.txt b/test/unit/library/CMakeLists.txt new file mode 100644 index 0000000000..095e632dfb --- /dev/null +++ b/test/unit/library/CMakeLists.txt @@ -0,0 +1,40 @@ +# Copyright (c) 2017 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +cutlass_test_unit_add_executable( + cutlass_test_unit_library + WITHOUT_CUDA + gemm_description.cu + EXTRA_INCLUDE_DIRS + ${CUTLASS_SOURCE_DIR}/tools/library/include + ${CUTLASS_SOURCE_DIR}/tools/library/src +) + +target_link_libraries(cutlass_test_unit_library PRIVATE + CUTLASS cutlass_tools_util_includes GTest::gtest_main cuda_driver) diff --git a/test/unit/library/gemm_description.cu b/test/unit/library/gemm_description.cu new file mode 100644 index 0000000000..80d03c4b68 --- /dev/null +++ b/test/unit/library/gemm_description.cu @@ -0,0 +1,327 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 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 Statically sized array of elements that accommodates all CUTLASS-supported numeric types + and is safe to use in a union. +*/ + + +#include +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "gemm_operation_3x.hpp" +#include "grouped_gemm_operation_3x.hpp" +#include "block_scaled_gemm_operation_3x.hpp" +#include "blockwise_gemm_operation_3x.hpp" + +using namespace cute; +using namespace cutlass::library; +namespace cg = cutlass::gemm; +namespace ce = cutlass::epilogue; + +namespace { + +// Construct real library descriptions without instantiating a launch path or +// querying a GPU. The operators and collectives below are the actual builders. +template +struct DescriptionOnlyCommon : Base { + using Base::Base; + cutlass::Status can_implement(void const*, void const*) const override { + return cutlass::Status::kErrorNotSupported; + } + uint64_t get_device_workspace_size(void const*, void const*) const override { return 0; } + cutlass::Status initialize(void const*, void*, void*, cudaStream_t) const override { + return cutlass::Status::kErrorNotSupported; + } + cutlass::Status run(void const*, void*, void*, cudaStream_t) const override { + return cutlass::Status::kErrorNotSupported; + } +}; + +template +struct DescriptionOnly : DescriptionOnlyCommon { + using DescriptionOnlyCommon::DescriptionOnlyCommon; + uint64_t get_host_workspace_size(void const*) const override { return 0; } +}; + +template , class Cluster = Shape<_1, _1, _1>, + bool Grouped = false, class OpClass = cutlass::arch::OpClassTensorOp> +struct GemmConfig { + using LayoutA = conditional_t; + using LayoutB = conditional_t; + using ProblemShape = conditional_t>, Shape>; + static constexpr int Alignment = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentA = is_same_v ? 2 * Alignment : Alignment; + using Epilogue = typename ce::collective::CollectiveBuilder< + Arch, cutlass::arch::OpClassTensorOp, Tile, Cluster, + ce::collective::EpilogueTileAuto, float, float, + cutlass::half_t, LayoutA, 8, + cutlass::half_t, LayoutA, 8, EpilogueSchedule>::CollectiveOp; + using Mainloop = typename cg::collective::CollectiveBuilder< + Arch, OpClass, + Element, LayoutA, AlignmentA, + Element, LayoutB, Alignment, + float, Tile, Cluster, cg::collective::StageCount<2>, Schedule>::CollectiveOp; + using Kernel = cg::kernel::GemmUniversal; + using Gemm = cg::device::GemmUniversalAdapter; +}; + +template +void check_description(MainloopScheduleKind schedule, MainloopLoadKind load, + EpilogueScheduleKind epilogue, AccumulationKind accumulation = AccumulationKind::kDefault) { + DescriptionOnly> operation("schedule_test"); + auto const& desc = operation.get_gemm_description(); + EXPECT_STREQ(desc.name, "schedule_test"); + EXPECT_EQ(desc.mainloop_schedule, schedule); + EXPECT_EQ(desc.mainloop_load, load); + EXPECT_EQ(desc.epilogue_schedule, epilogue); + EXPECT_EQ(desc.accumulation_kind, accumulation); + EXPECT_EQ(desc.A.element, NumericTypeID(NumericTypeMap::kId)); + EXPECT_EQ(desc.A.layout, LayoutTypeID::kRowMajor); + EXPECT_EQ(desc.B.layout, LayoutTypeID::kColumnMajor); + EXPECT_EQ(desc.tile_description.math_instruction.instruction_shape, + cutlass::make_Coord(Config::Gemm::InstructionShape::kM, + Config::Gemm::InstructionShape::kN, + Config::Gemm::InstructionShape::kK)); +} + +TEST(GemmDescription, LegacyConstructorsDefaultToUnknown) { + GemmDescription descriptions[] = { + GemmDescription(), + GemmDescription(GemmKind::kUniversal), + GemmDescription(OperationDescription(), GemmKind::kUniversal, + TensorDescription(), TensorDescription(), TensorDescription(), TensorDescription(), + NumericTypeID::kF32, SplitKMode::kNone, ComplexTransform::kNone, ComplexTransform::kNone)}; + for (auto const& d : descriptions) { + EXPECT_EQ(d.mainloop_schedule, MainloopScheduleKind::kUnknown); + EXPECT_EQ(d.mainloop_load, MainloopLoadKind::kUnknown); + EXPECT_EQ(d.epilogue_schedule, EpilogueScheduleKind::kUnknown); + EXPECT_EQ(d.accumulation_kind, AccumulationKind::kUnknown); + } +} + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) +TEST(GemmDescription, HopperCooperative) { + using Config = GemmConfig; + check_description(MainloopScheduleKind::kWarpSpecializedCooperative, + MainloopLoadKind::kTma, EpilogueScheduleKind::kTma); +} + +TEST(GemmDescription, HopperPingpongDirectStore) { + using Config = GemmConfig; + check_description(MainloopScheduleKind::kWarpSpecializedPingpong, + MainloopLoadKind::kTma, EpilogueScheduleKind::kNoSmem); +} + +TEST(GemmDescription, HopperFp8DefaultAndFastAccumulation) { + using Default = GemmConfig; + using Fast = GemmConfig; + check_description(MainloopScheduleKind::kWarpSpecializedCooperative, + MainloopLoadKind::kTma, EpilogueScheduleKind::kTma); + check_description(MainloopScheduleKind::kWarpSpecializedCooperative, + MainloopLoadKind::kTma, EpilogueScheduleKind::kTma, AccumulationKind::kFastAccum); +} + +#endif + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) +TEST(GemmDescription, BlackwellOneSm) { + using Config = GemmConfig; + check_description(MainloopScheduleKind::kWarpSpecialized1Sm, + MainloopLoadKind::kTma, EpilogueScheduleKind::kTma); +} + +TEST(GemmDescription, BlackwellTwoSm) { + using Config = GemmConfig, Shape<_2, _1, _1>>; + check_description(MainloopScheduleKind::kWarpSpecialized2Sm, + MainloopLoadKind::kTma, EpilogueScheduleKind::kTma); +} + +TEST(GemmDescription, BlackwellAutoResolvesSchedule) { + using Config = GemmConfig, Shape<_2, _1, _1>>; + check_description(MainloopScheduleKind::kWarpSpecialized2Sm, + MainloopLoadKind::kTma, EpilogueScheduleKind::kTma); +} + +#endif + + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) +TEST(GemmDescription, HopperCpAsync) { + using Config = GemmConfig; + check_description(MainloopScheduleKind::kWarpSpecialized, + MainloopLoadKind::kCpAsync, EpilogueScheduleKind::kNoSmem); +} + +TEST(GemmDescription, HopperGroupedFp8FastAccumulation) { + using Config = GemmConfig, Shape<_1, _1, _1>, true>; + DescriptionOnlyCommon> operation; + auto const& desc = static_cast(operation.description()); + EXPECT_EQ(desc.kind, OperationKind::kGroupedGemm); + EXPECT_EQ(desc.gemm.mainloop_schedule, MainloopScheduleKind::kWarpSpecializedCooperative); + EXPECT_EQ(desc.gemm.mainloop_load, MainloopLoadKind::kTma); + EXPECT_EQ(desc.gemm.epilogue_schedule, EpilogueScheduleKind::kTma); + EXPECT_EQ(desc.gemm.accumulation_kind, AccumulationKind::kFastAccum); +} + +TEST(GemmDescription, HopperBlockwiseDescription) { + using Tile = Shape<_128, _128, _128>; + using Cluster = Shape<_1, _1, _1>; + using ScaleConfig = cutlass::detail::Sm90BlockwiseScaleConfig<1, 128, 128>; + using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); + using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); + using Epilogue = typename GemmConfig::Epilogue; + using Mainloop = typename cg::collective::CollectiveBuilder< + cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, + cutlass::float_e4m3_t, tuple, 16, + cutlass::float_e4m3_t, tuple, 16, + float, Tile, Cluster, cg::collective::StageCountAutoCarveout, + cg::KernelTmaWarpSpecializedCooperativeFP8Blockwise>::CollectiveOp; + using Kernel = cg::kernel::GemmUniversal, Mainloop, Epilogue>; + using Gemm = cg::device::GemmUniversalAdapter; + DescriptionOnly> operation; + auto const& desc = static_cast(operation.description()); + EXPECT_EQ(desc.kind, OperationKind::kBlockwiseGemm); + EXPECT_EQ(desc.mainloop_schedule, MainloopScheduleKind::kWarpSpecializedCooperative); + EXPECT_EQ(desc.mainloop_load, MainloopLoadKind::kTma); + EXPECT_EQ(desc.epilogue_schedule, EpilogueScheduleKind::kTma); + EXPECT_EQ(desc.accumulation_kind, AccumulationKind::kDefault); + EXPECT_EQ(desc.SFKVecSize, 128); +} +#endif + +#if defined(CUTLASS_ARCH_MMA_SPARSE_SM90_SUPPORTED) +TEST(GemmDescription, HopperSparse) { + using Config = GemmConfig, Shape<_1, _1, _1>, false, cutlass::arch::OpClassSparseTensorOp>; + check_description(MainloopScheduleKind::kWarpSpecializedCooperative, + MainloopLoadKind::kTma, EpilogueScheduleKind::kTma); +} +#endif + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) +TEST(GemmDescription, BlackwellDirectStore) { + using Config = GemmConfig; + check_description(MainloopScheduleKind::kWarpSpecialized1Sm, + MainloopLoadKind::kTma, EpilogueScheduleKind::kNoSmem); +} + +TEST(GemmDescription, BlackwellOneSmWithTwoCtaCluster) { + using Config = GemmConfig, Shape<_2, _1, _1>>; + check_description(MainloopScheduleKind::kWarpSpecialized1Sm, + MainloopLoadKind::kTma, EpilogueScheduleKind::kTma); +} + +TEST(GemmDescription, BlackwellCpAsync) { + using Config = GemmConfig; + check_description(MainloopScheduleKind::kWarpSpecialized1Sm, + MainloopLoadKind::kCpAsync, EpilogueScheduleKind::kNoSmem); +} + +TEST(GemmDescription, BlackwellMixedTmaCpAsync) { + using Config = GemmConfig; + check_description(MainloopScheduleKind::kWarpSpecialized1Sm, + MainloopLoadKind::kMixedTmaCpAsync, EpilogueScheduleKind::kTma); +} + +TEST(GemmDescription, BlackwellGroupedTwoSm) { + using Config = GemmConfig, Shape<_2, _1, _1>, true>; + DescriptionOnlyCommon> operation; + auto const& desc = static_cast(operation.description()); + EXPECT_EQ(desc.kind, OperationKind::kGroupedGemm); + EXPECT_EQ(desc.gemm.mainloop_schedule, MainloopScheduleKind::kWarpSpecialized2Sm); + EXPECT_EQ(desc.gemm.mainloop_load, MainloopLoadKind::kTma); + EXPECT_EQ(desc.gemm.epilogue_schedule, EpilogueScheduleKind::kTma); + EXPECT_EQ(desc.gemm.accumulation_kind, AccumulationKind::kDefault); +} + +TEST(GemmDescription, BlackwellBlockScaledDescription) { + using Config = GemmConfig, + cg::collective::KernelScheduleAuto, ce::collective::EpilogueScheduleAuto, + Shape<_256, _128, _128>, Shape<_2, _1, _1>, false, cutlass::arch::OpClassBlockScaledTensorOp>; + DescriptionOnly> operation; + auto const& desc = operation.get_gemm_description(); + EXPECT_EQ(desc.kind, OperationKind::kBlockScaledGemm); + EXPECT_EQ(desc.mainloop_schedule, MainloopScheduleKind::kWarpSpecialized2Sm); + EXPECT_EQ(desc.mainloop_load, MainloopLoadKind::kTma); + EXPECT_EQ(desc.epilogue_schedule, EpilogueScheduleKind::kTma); + EXPECT_EQ(desc.accumulation_kind, AccumulationKind::kDefault); + EXPECT_EQ(desc.SFVecSize, 32); +} +#endif + +TEST(GemmDescription, QuantizedConstructorsDefaultToUnknown) { + BlockScaledGemmDescription block_scaled; + EXPECT_EQ(block_scaled.mainloop_schedule, MainloopScheduleKind::kUnknown); + EXPECT_EQ(block_scaled.mainloop_load, MainloopLoadKind::kUnknown); + EXPECT_EQ(block_scaled.epilogue_schedule, EpilogueScheduleKind::kUnknown); + EXPECT_EQ(block_scaled.accumulation_kind, AccumulationKind::kUnknown); + BlockwiseGemmDescription blockwise; + EXPECT_EQ(blockwise.mainloop_schedule, MainloopScheduleKind::kUnknown); + EXPECT_EQ(blockwise.mainloop_load, MainloopLoadKind::kUnknown); + EXPECT_EQ(blockwise.epilogue_schedule, EpilogueScheduleKind::kUnknown); + EXPECT_EQ(blockwise.accumulation_kind, AccumulationKind::kUnknown); +} + +struct CustomSchedule {}; +struct CustomEpilogue {}; +TEST(GemmDescription, CustomPoliciesStayUnknown) { + EXPECT_EQ(KernelScheduleMap::kSchedule, MainloopScheduleKind::kUnknown); + EXPECT_EQ(KernelScheduleMap::kLoad, MainloopLoadKind::kUnknown); + EXPECT_EQ(KernelScheduleMap::kAccumulation, AccumulationKind::kUnknown); + EXPECT_EQ(CollectiveEpilogueScheduleMap::kId, EpilogueScheduleKind::kUnknown); +} + +} // namespace diff --git a/tools/library/include/cutlass/library/descriptions.h b/tools/library/include/cutlass/library/descriptions.h index 89ee41f3c0..9611a222b2 100644 --- a/tools/library/include/cutlass/library/descriptions.h +++ b/tools/library/include/cutlass/library/descriptions.h @@ -253,6 +253,12 @@ struct GemmDescription : public OperationDescription { /// Transformation on B operand ComplexTransform transform_B; + /// Schedule metadata for CUTLASS 3.x kernels. Older/custom kernels default to unknown. + MainloopScheduleKind mainloop_schedule = MainloopScheduleKind::kUnknown; + MainloopLoadKind mainloop_load = MainloopLoadKind::kUnknown; + EpilogueScheduleKind epilogue_schedule = EpilogueScheduleKind::kUnknown; + AccumulationKind accumulation_kind = AccumulationKind::kUnknown; + // // Methods // @@ -379,6 +385,12 @@ struct BlockScaledGemmDescription : public OperationDescription { /// Describes the Output ScaleFactor VectorSize int EpilogueSFVecSize; + /// Schedule metadata for CUTLASS 3.x kernels. Older/custom kernels default to unknown. + MainloopScheduleKind mainloop_schedule = MainloopScheduleKind::kUnknown; + MainloopLoadKind mainloop_load = MainloopLoadKind::kUnknown; + EpilogueScheduleKind epilogue_schedule = EpilogueScheduleKind::kUnknown; + AccumulationKind accumulation_kind = AccumulationKind::kUnknown; + // // Methods // @@ -469,6 +481,12 @@ struct BlockwiseGemmDescription : public OperationDescription { int SFNVecSize; int SFKVecSize; + /// Schedule metadata for CUTLASS 3.x kernels. Older/custom kernels default to unknown. + MainloopScheduleKind mainloop_schedule = MainloopScheduleKind::kUnknown; + MainloopLoadKind mainloop_load = MainloopLoadKind::kUnknown; + EpilogueScheduleKind epilogue_schedule = EpilogueScheduleKind::kUnknown; + AccumulationKind accumulation_kind = AccumulationKind::kUnknown; + // // Methods // diff --git a/tools/library/include/cutlass/library/types.h b/tools/library/include/cutlass/library/types.h index 0c6c9e1ca3..d42a5af42f 100644 --- a/tools/library/include/cutlass/library/types.h +++ b/tools/library/include/cutlass/library/types.h @@ -288,6 +288,45 @@ enum class RasterOrder { kInvalid }; +/// Mainloop execution strategy. Unknown is used for unrecognized/custom schedules. +enum class MainloopScheduleKind { + kUnknown, + kMultistage, + kTma, + kWarpSpecialized, + kWarpSpecializedPingpong, + kWarpSpecializedCooperative, + kWarpSpecialized1Sm, + kWarpSpecialized2Sm, + kInvalid +}; + +/// Instructions used to move the mainloop operands from global memory. +enum class MainloopLoadKind { + kUnknown, + kCpAsync, + kTma, + kMixedTmaCpAsync, + kInvalid +}; + +/// Epilogue data movement strategy (including stores). +enum class EpilogueScheduleKind { + kUnknown, + kNoSmem, + kTma, + kInvalid +}; + +/// Distinguishes the opt-in Hopper FP8 fast-accumulation mainloops. +enum class AccumulationKind { + kUnknown, + kDefault, + kFastAccum, + kInvalid +}; + + ///////////////////////////////////////////////////////////////////////////////////////////////// } // namespace library diff --git a/tools/library/src/block_scaled_gemm_operation_3x.hpp b/tools/library/src/block_scaled_gemm_operation_3x.hpp index a1550848b2..541fa14556 100644 --- a/tools/library/src/block_scaled_gemm_operation_3x.hpp +++ b/tools/library/src/block_scaled_gemm_operation_3x.hpp @@ -167,6 +167,8 @@ class BlockScaledGemmUniversal3xOperationBase : public GemmOperation3xBase::kId; description_.split_k_mode = SplitKMode::kNone; + + initialize_gemm_schedule_description(description_); } /// Returns the description of the GEMM operation diff --git a/tools/library/src/blockwise_gemm_operation_3x.hpp b/tools/library/src/blockwise_gemm_operation_3x.hpp index cdc9ca03be..2d9a448b84 100644 --- a/tools/library/src/blockwise_gemm_operation_3x.hpp +++ b/tools/library/src/blockwise_gemm_operation_3x.hpp @@ -158,6 +158,8 @@ class BlockwiseGemmUniversal3xOperation : public GemmOperation3xBase description_.element_epilogue = NumericTypeMap::kId; description_.split_k_mode = SplitKMode::kNone; + + initialize_gemm_schedule_description(description_); } /// Returns the description of the GEMM operation diff --git a/tools/library/src/gemm_operation_3x.hpp b/tools/library/src/gemm_operation_3x.hpp index ff03975675..a89dd3604c 100644 --- a/tools/library/src/gemm_operation_3x.hpp +++ b/tools/library/src/gemm_operation_3x.hpp @@ -40,6 +40,7 @@ #include "cutlass/array_subbyte.h" #include "cutlass/library/library.h" #include "library_internal.h" +#include "gemm_schedule_description.hpp" #include "cutlass/gemm/dispatch_policy.hpp" #include "cutlass/util/packed_stride.hpp" #include "cutlass/util/mixed_dtype_utils.hpp" @@ -134,6 +135,8 @@ class GemmOperation3xBase : public Operation { description_.split_k_mode = SplitKMode::kNone; description_.transform_A = ComplexTransformMap::kId; description_.transform_B = ComplexTransformMap::kId; + + initialize_gemm_schedule_description(description_); } /// Returns the description of the GEMM operation diff --git a/tools/library/src/gemm_schedule_description.hpp b/tools/library/src/gemm_schedule_description.hpp new file mode 100644 index 0000000000..f91f33be64 --- /dev/null +++ b/tools/library/src/gemm_schedule_description.hpp @@ -0,0 +1,243 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 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. + * + **************************************************************************************************/ +/// Maps the selected collective policies to runtime GEMM schedule metadata. +#pragma once + +#include +#include "cutlass/library/descriptions.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" + +namespace cutlass::library { + +// Inspect the resolved schedule, not the builder's input tag (which may be Auto). +template +struct KernelScheduleMap { + template + static constexpr bool derives = cute::is_base_of_v; + + template