b78588d163
* CUTLASS 3.7 * clean up changelog --------- Co-authored-by: yuzhai <yuzhai@nvidia.com> Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
981 lines
40 KiB
C++
981 lines
40 KiB
C++
/***************************************************************************************************
|
|
* 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 Defines operations for all CONV operation kinds in CUTLASS Library.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include "cutlass/cutlass.h"
|
|
#include "cutlass/library/library.h"
|
|
#include "library_internal.h"
|
|
#include "cutlass/conv/convnd_problem_shape.hpp"
|
|
#include "cutlass/util/packed_stride.hpp"
|
|
#include "cutlass/detail/dependent_false.hpp"
|
|
#include "cutlass/trace.h"
|
|
#include <utility>
|
|
#include <variant>
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL)
|
|
#include <sstream>
|
|
#endif
|
|
|
|
namespace cutlass::library {
|
|
|
|
namespace detail {
|
|
|
|
template<class ValueType, size_t ... Indices>
|
|
constexpr cute::array<ValueType, 1u + sizeof...(Indices)>
|
|
vector_to_array_strides_helper(const std::vector<ValueType>& v,
|
|
std::index_sequence<Indices...>)
|
|
{
|
|
return {v[(sizeof...(Indices) - 1u) - Indices]..., ValueType(1)};
|
|
}
|
|
|
|
template<class ValueType, size_t Size>
|
|
cute::array<ValueType, Size>
|
|
vector_to_array_strides(const std::vector<ValueType>& v, std::integral_constant<size_t, Size>)
|
|
{
|
|
static_assert(Size != 0);
|
|
CUTLASS_ASSERT(v.size() + 1u == Size);
|
|
return vector_to_array_strides_helper(v, std::make_index_sequence<Size - 1u>{});
|
|
}
|
|
|
|
template<class Index, class LongIndex, size_t ... Indices>
|
|
constexpr cute::array<int64_t, 1u + sizeof...(Indices)>
|
|
coord_to_array_strides_helper(
|
|
const ::cutlass::Coord<int(sizeof...(Indices)), Index, LongIndex> coord,
|
|
std::index_sequence<Indices...>)
|
|
{
|
|
return {int64_t(coord[(sizeof...(Indices) - 1u) - Indices])..., int64_t(1)};
|
|
}
|
|
|
|
template<int Rank, class Index, class LongIndex>
|
|
cute::array<int64_t, 1u + size_t(Rank)>
|
|
coord_to_array_strides(const ::cutlass::Coord<Rank, Index, LongIndex>& coord)
|
|
{
|
|
static_assert(Rank >= 0);
|
|
return coord_to_array_strides_helper(coord, std::make_index_sequence<Rank>{});
|
|
}
|
|
|
|
} // namespace detail
|
|
|
|
// Tells the profiler about CUTLASS 3's 2-D and 3-D convolutions.
|
|
// For CUTLASS 2's 2-D convolutions, see Conv2dOperation.
|
|
// For CUTLASS 2's 3-D convolutions, see Conv3dOperation.
|
|
template<class Operator_>
|
|
class ConvOperation3x : public Operation {
|
|
public:
|
|
using Operator = Operator_;
|
|
|
|
static_assert(Operator::NumSpatialDimensions == 2 ||
|
|
Operator::NumSpatialDimensions == 3,
|
|
"The profiler currently only supports convolutions with 2 or 3 spatial dimensions.");
|
|
using LayoutA = cute::conditional_t<Operator::NumSpatialDimensions == 3,
|
|
cutlass::layout::TensorNDHWC,
|
|
cute::conditional_t<Operator::NumSpatialDimensions == 2,
|
|
cutlass::layout::TensorNHWC,
|
|
cutlass::layout::TensorNWC>
|
|
>;
|
|
using LayoutB = LayoutA;
|
|
using LayoutC = LayoutA;
|
|
|
|
using ElementA = typename Operator::ElementA;
|
|
using ElementB = typename Operator::ElementB;
|
|
using ElementC = typename Operator::ElementC;
|
|
using ElementD = typename Operator::ElementD;
|
|
using ElementAccumulator = typename Operator::ElementAccumulator;
|
|
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
|
|
static cutlass::conv::Operator const kConvolutionalOperator = Operator::kConvolutionalOperator;
|
|
|
|
ConvOperation3x(const char* name = "unknown_cutlass_3_conv") {
|
|
// Initialize OperationDescription (the base class)
|
|
description_.name = name;
|
|
description_.provider = Provider::kCUTLASS;
|
|
|
|
if constexpr (Operator::NumSpatialDimensions == 2) {
|
|
description_.kind = OperationKind::kConv2d;
|
|
}
|
|
else if constexpr (Operator::NumSpatialDimensions == 3) {
|
|
description_.kind = OperationKind::kConv3d;
|
|
}
|
|
else {
|
|
static_assert(::cutlass::detail::dependent_false<Operator>,
|
|
"This class currently only supports 2-D and 3-D convolutions.");
|
|
}
|
|
|
|
description_.tile_description.threadblock_shape = make_Coord(
|
|
Operator::ThreadblockShape::kM,
|
|
Operator::ThreadblockShape::kN,
|
|
Operator::ThreadblockShape::kK);
|
|
|
|
description_.tile_description.threadblock_stages = Operator::kStages;
|
|
|
|
description_.tile_description.warp_count = make_Coord(
|
|
Operator::WarpCount::kM,
|
|
Operator::WarpCount::kN,
|
|
Operator::WarpCount::kK);
|
|
|
|
description_.tile_description.math_instruction.instruction_shape = make_Coord(
|
|
Operator::InstructionShape::kM,
|
|
Operator::InstructionShape::kN,
|
|
Operator::InstructionShape::kK);
|
|
|
|
description_.tile_description.math_instruction.element_accumulator =
|
|
NumericTypeMap<ElementAccumulator>::kId;
|
|
|
|
description_.tile_description.math_instruction.opcode_class =
|
|
OpcodeClassMap<typename Operator::OperatorClass>::kId;
|
|
|
|
description_.tile_description.math_instruction.math_operation =
|
|
MathOperationID::kMultiplyAdd;
|
|
|
|
description_.tile_description.minimum_compute_capability =
|
|
ArchMap<typename Operator::ArchTag, typename Operator::OperatorClass>::kMin;
|
|
|
|
description_.tile_description.maximum_compute_capability =
|
|
ArchMap<typename Operator::ArchTag, typename Operator::OperatorClass>::kMax;
|
|
|
|
// Initialize ConvDescription (the subclass)
|
|
|
|
// kConvDim does not exist in Operator for CUTLASS 3 convolutions.
|
|
// For CUTLASS 2 convolutions, it is the number of spatial dimensions.
|
|
description_.conv_dim = Operator::NumSpatialDimensions;
|
|
description_.conv_kind = ConvKindMap<kConvolutionalOperator>::kId;
|
|
|
|
description_.iterator_algorithm = {};
|
|
|
|
description_.A = make_TensorDescription<ElementA, LayoutA>();
|
|
description_.B = make_TensorDescription<ElementB, LayoutB>();
|
|
description_.C = make_TensorDescription<ElementC, LayoutC>();
|
|
description_.element_epilogue = NumericTypeMap<ElementCompute>::kId;
|
|
}
|
|
|
|
~ConvOperation3x() override = default;
|
|
|
|
OperationDescription const& description() const override {
|
|
return static_cast<OperationDescription const&>(description_);
|
|
}
|
|
|
|
private:
|
|
Status update_operator_arguments_from_configuration_2d_or_3d(
|
|
typename Operator::Arguments& out_args,
|
|
void const* configuration) const {
|
|
Status status = Status::kInvalid;
|
|
|
|
CUTLASS_ASSERT(configuration != nullptr);
|
|
|
|
if constexpr (Operator::NumSpatialDimensions == 2) {
|
|
CUTLASS_ASSERT(description_.kind == OperationKind::kConv2d);
|
|
// tools/library/include/cutlass/library/library.h
|
|
// defines Conv2dConfiguration.
|
|
// tools/profiler/include/cutlass/profiler/conv2d_operation_profiler.h
|
|
// uses Conv2dConfiguration.
|
|
auto* conf_ptr = reinterpret_cast<Conv2dConfiguration const*>(configuration);
|
|
status = update_operator_arguments_from_configuration(out_args, *conf_ptr);
|
|
}
|
|
else if constexpr (Operator::NumSpatialDimensions == 3) {
|
|
CUTLASS_ASSERT(description_.kind == OperationKind::kConv3d);
|
|
auto* conf_ptr = reinterpret_cast<Conv3dConfiguration const*>(configuration);
|
|
status = update_operator_arguments_from_configuration(out_args, *conf_ptr);
|
|
}
|
|
else {
|
|
static_assert(::cutlass::detail::dependent_false<Operator>,
|
|
"This class currently only supports 2-D and 3-D convolutions.");
|
|
}
|
|
|
|
return status;
|
|
}
|
|
|
|
public:
|
|
Status can_implement(
|
|
void const* configuration,
|
|
void const* arguments) const override {
|
|
Status status = Status::kInvalid;
|
|
|
|
// gemm_operation_3x.hpp accesses "configuration" as
|
|
// GemmUniversalConfiguration (which lives in
|
|
// tools/library/include/cutlass/library/library.h) and
|
|
// "arguments" as GemmUniversalArguments (which lives in
|
|
// tools/library/include/cutlass/library/library.h).
|
|
// Those things don't apply to convolutions.
|
|
// Despite the existence of ConvUniversal, there's no
|
|
// corresponding "ConvUniversalConfiguration" or
|
|
// "ConvUniversalArguments."
|
|
|
|
CUTLASS_ASSERT(configuration != nullptr);
|
|
CUTLASS_ASSERT(arguments != nullptr);
|
|
|
|
typename Operator::Arguments out_args{};
|
|
status = update_operator_arguments_from_configuration_2d_or_3d(out_args, configuration);
|
|
if (status != Status::kSuccess) {
|
|
CUTLASS_TRACE_HOST("*** can_implement: update_operator_arguments_from_configuration_2d_or_3d failed");
|
|
return status;
|
|
}
|
|
|
|
auto* in_args_ptr = reinterpret_cast<ConvArguments const*>(arguments);
|
|
status = update_operator_arguments_from_arguments(out_args, *in_args_ptr);
|
|
if (status != Status::kSuccess) {
|
|
CUTLASS_TRACE_HOST("*** can_implement: update_operator_arguments_from_arguments failed");
|
|
return status;
|
|
}
|
|
|
|
return Operator::can_implement(out_args);
|
|
}
|
|
|
|
uint64_t get_host_workspace_size(void const* /* configuration */) const override {
|
|
return sizeof(Operator);
|
|
}
|
|
|
|
uint64_t get_device_workspace_size(
|
|
void const* configuration,
|
|
void const* arguments = nullptr) const override
|
|
{
|
|
// This presumes that at least one of configuration or arguments is nonnull.
|
|
Status status = Status::kInvalid;
|
|
|
|
// gemm_operation_3x.hpp has get_device_workspace_size return 0 on
|
|
// error. It's not clear that this is what we want -- perhaps we
|
|
// should return something like expected<uint64_t, Status>? -- but
|
|
// it's the only option that preserves the current interface.
|
|
constexpr uint64_t error_indication = 0;
|
|
|
|
typename Operator::Arguments out_args{};
|
|
if (configuration != nullptr) {
|
|
status = update_operator_arguments_from_configuration_2d_or_3d(out_args, configuration);
|
|
if (status != Status::kSuccess) {
|
|
return error_indication;
|
|
}
|
|
}
|
|
if (arguments != nullptr) {
|
|
auto* in_args_ptr = reinterpret_cast<ConvArguments const*>(arguments);
|
|
status = update_operator_arguments_from_arguments(out_args, *in_args_ptr);
|
|
if (status != Status::kSuccess) {
|
|
return error_indication;
|
|
}
|
|
}
|
|
|
|
if (status == Status::kSuccess) {
|
|
return static_cast<uint64_t>(Operator::get_workspace_size(out_args));
|
|
}
|
|
else {
|
|
return error_indication;
|
|
}
|
|
}
|
|
|
|
Status initialize(
|
|
void const* configuration,
|
|
void* host_workspace,
|
|
void* /* device_workspace */ = nullptr,
|
|
cudaStream_t stream = nullptr) const override
|
|
{
|
|
Status status = Status::kInvalid;
|
|
|
|
if (configuration == nullptr) {
|
|
CUTLASS_TRACE_HOST("Input configuration is null.");
|
|
return Status::kInvalid;
|
|
}
|
|
|
|
typename Operator::Arguments out_args{};
|
|
status = update_operator_arguments_from_configuration_2d_or_3d(out_args, configuration);
|
|
if (status != Status::kSuccess) {
|
|
// Any kind of failure invalidates the last successful configuration.
|
|
clear_last_successful_config();
|
|
return status;
|
|
}
|
|
else {
|
|
set_last_successful_config(configuration);
|
|
}
|
|
|
|
if (host_workspace == nullptr) {
|
|
CUTLASS_TRACE_HOST("host_workspace is null.");
|
|
return Status::kInvalid;
|
|
}
|
|
(void) new (host_workspace) Operator;
|
|
return status;
|
|
|
|
// CUTLASS 2 convolutions call the Operator's initialize function
|
|
// here, like this.
|
|
//
|
|
//return op->initialize(args, device_workspace, stream);
|
|
//
|
|
// CUTLASS 3 convolutions (ConvUniversal), like CUTLASS 3 Gemms
|
|
// (GemmUniversal), lack an "initialize" member function.
|
|
}
|
|
|
|
Status run(
|
|
void const* arguments,
|
|
void* host_workspace,
|
|
void* device_workspace = nullptr,
|
|
cudaStream_t stream = nullptr) const override
|
|
{
|
|
auto status = Status::kInvalid;
|
|
|
|
// The Operator doesn't appear to save the last configuration (it
|
|
// doesn't have a way to do that, since it lacks an initialize()
|
|
// member function), so we have to use the stored configuration
|
|
// from the last successful initialize() call (if any).
|
|
typename Operator::Arguments out_args{};
|
|
status = update_operator_arguments_from_stored_configuration(out_args);
|
|
if (status != Status::kSuccess) {
|
|
CUTLASS_TRACE_HOST("Updating from previous successful configuration failed.");
|
|
return status;
|
|
}
|
|
|
|
if (arguments == nullptr) {
|
|
CUTLASS_TRACE_HOST("Input argument 'arguments' is null.");
|
|
return Status::kInvalid;
|
|
}
|
|
auto* in_args_ptr = reinterpret_cast<ConvArguments const*>(arguments);
|
|
status = update_operator_arguments_from_arguments(out_args, *in_args_ptr);
|
|
if (status != Status::kSuccess) {
|
|
return status;
|
|
}
|
|
|
|
auto* op = reinterpret_cast<Operator*>(host_workspace);
|
|
return op->run(out_args, device_workspace, stream, nullptr, in_args_ptr->use_pdl);
|
|
}
|
|
|
|
private:
|
|
ConvDescription description_;
|
|
// Result of initialize() calling
|
|
// update_operator_arguments_from_configuration() successfully.
|
|
// This is needed because run() doesn't take a configuration, just
|
|
// arguments, and the kernel doesn't appear to save the
|
|
// configuration from the last initialize() call.
|
|
//
|
|
// Unfortunately, this must be declared mutable, because it must be
|
|
// set in initialize(), and initialize() is inherited as const.
|
|
mutable std::variant<
|
|
std::monostate,
|
|
Conv2dConfiguration,
|
|
Conv3dConfiguration> last_successful_config_{std::monostate{}};
|
|
|
|
// Clear the last configuration resulting from a successful initialize() call.
|
|
//
|
|
// Unfortunately, this must be declared const, because initialize() is.
|
|
void clear_last_successful_config() const {
|
|
last_successful_config_ = std::monostate{};
|
|
}
|
|
|
|
// Set the last configuration resulting from a successful initialize() call.
|
|
//
|
|
// Unfortunately, this must be declared const, because initialize() is.
|
|
void set_last_successful_config(void const* configuration) const {
|
|
CUTLASS_ASSERT(configuration != nullptr);
|
|
|
|
if constexpr (Operator::NumSpatialDimensions == 2) {
|
|
CUTLASS_ASSERT(description_.kind == OperationKind::kConv2d);
|
|
auto* conf_ptr = reinterpret_cast<Conv2dConfiguration const*>(configuration);
|
|
last_successful_config_ = *conf_ptr;
|
|
} else if constexpr (Operator::NumSpatialDimensions == 3) {
|
|
CUTLASS_ASSERT(description_.kind == OperationKind::kConv3d);
|
|
auto* conf_ptr = reinterpret_cast<Conv3dConfiguration const*>(configuration);
|
|
last_successful_config_ = *conf_ptr;
|
|
}
|
|
else {
|
|
static_assert(::cutlass::detail::dependent_false<Operator>,
|
|
"This class currently only supports 2-D and 3-D convolutions.");
|
|
}
|
|
}
|
|
|
|
// Whether a configuration from a successful initialize() call exists.
|
|
bool last_successful_config_exists() const {
|
|
return not std::holds_alternative<std::monostate>(last_successful_config_);
|
|
}
|
|
|
|
// Visitor for update_operator_arguments_from_stored_configuration.
|
|
struct ConfigurationVisitor {
|
|
typename Operator::Arguments& out_args;
|
|
|
|
Status operator() (std::monostate const&) const {
|
|
CUTLASS_TRACE_HOST("No successful previous configuration exists. "
|
|
"One cause is calling run() before a successful initialize() call.");
|
|
return Status::kInvalid;
|
|
}
|
|
Status operator() (Conv2dConfiguration const& conf2d) const {
|
|
return update_operator_arguments_from_configuration(out_args, conf2d);
|
|
}
|
|
Status operator() (Conv3dConfiguration const& conf3d) const {
|
|
return update_operator_arguments_from_configuration(out_args, conf3d);
|
|
}
|
|
};
|
|
|
|
// Like update_operator_arguments_from_configuration, but on the
|
|
// stored configuration from the last successful initialize() call,
|
|
// if any. If there was no last successful initialize() call,
|
|
// then return Status::kInvalid.
|
|
//
|
|
// Unfortunately, this must be declared const, because run() is.
|
|
Status update_operator_arguments_from_stored_configuration(
|
|
typename Operator::Arguments& out_args) const
|
|
{
|
|
return std::visit(ConfigurationVisitor{out_args}, last_successful_config_);
|
|
}
|
|
|
|
template<class FusionArgs, class = void>
|
|
struct UpdateFusionArgs {
|
|
static Status update_(
|
|
FusionArgs const&,
|
|
ConvArguments const&)
|
|
{
|
|
// For custom EVT, it is the user's responsibility to ensure
|
|
// that alpha and beta are updated appropriately.
|
|
return Status::kSuccess;
|
|
}
|
|
};
|
|
|
|
template<class FusionArgs>
|
|
struct UpdateFusionArgs<FusionArgs, cute::void_t<decltype(FusionArgs{}.alpha)>> {
|
|
static Status update_(
|
|
FusionArgs& fusion_args,
|
|
ConvArguments const& arguments)
|
|
{
|
|
if (arguments.pointer_mode == ScalarPointerMode::kHost) {
|
|
fusion_args.alpha = *static_cast<ElementCompute const *>(arguments.alpha);
|
|
fusion_args.beta = *static_cast<ElementCompute const *>(arguments.beta);
|
|
fusion_args.alpha_ptr = nullptr;
|
|
fusion_args.beta_ptr = nullptr;
|
|
|
|
return Status::kSuccess;
|
|
}
|
|
else if (arguments.pointer_mode == ScalarPointerMode::kDevice) {
|
|
fusion_args.alpha = 0;
|
|
fusion_args.beta = 0;
|
|
fusion_args.alpha_ptr = static_cast<ElementCompute const *>(arguments.alpha);
|
|
fusion_args.beta_ptr = static_cast<ElementCompute const *>(arguments.beta);
|
|
|
|
return Status::kSuccess;
|
|
}
|
|
else {
|
|
return Status::kErrorInvalidProblem;
|
|
}
|
|
}
|
|
};
|
|
|
|
static Status update_operator_arguments_from_configuration(
|
|
typename Operator::Arguments& out_args,
|
|
Conv2dConfiguration const& config)
|
|
{
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
CUTLASS_TRACE_HOST("ConvOperator3x::"
|
|
"update_operator_arguments_from_configuration"
|
|
"(Conv2dConfiguration)\n");
|
|
#endif
|
|
using detail::vector_to_array_strides;
|
|
|
|
constexpr int num_spatial_dims = Operator::NumSpatialDimensions;
|
|
if constexpr (num_spatial_dims != 2) {
|
|
CUTLASS_TRACE_HOST("You can only use Conv2dConfiguration "
|
|
"with an Operator whose NumSpatialDimensions is exactly 2.");
|
|
return Status::kInvalid;
|
|
}
|
|
else {
|
|
// Convolutions split the metadata (in Conv2dConfiguration) from
|
|
// the data (ConvArguments, which only has pointers and a single
|
|
// enum value). Thus, this class will need both the
|
|
// configuration and the (user's input) arguments to set up the
|
|
// kernel's arguments. This function can fill in what the
|
|
// configuration has now, but the class will need the user's
|
|
// input arguments later.
|
|
if (config.split_k_mode != conv::SplitKMode::kSerial) {
|
|
CUTLASS_TRACE_HOST("CUTLASS 3 convolutions currently only support split_k_mode = kSerial.");
|
|
return Status::kInvalid;
|
|
}
|
|
// config.problem_size.split_k_slices is only meaningful if
|
|
// split_k_mode != kSerial. If this code later supports other
|
|
// split_k_mode values, then it will also need to read
|
|
// split_k_slices.
|
|
|
|
const int N = config.problem_size.N;
|
|
const int H = config.problem_size.H;
|
|
const int W = config.problem_size.W;
|
|
const int C = config.problem_size.C;
|
|
const int K = config.problem_size.K;
|
|
const int R = config.problem_size.R;
|
|
const int S = config.problem_size.S;
|
|
const int pad_h = config.problem_size.pad_h;
|
|
const int pad_w = config.problem_size.pad_w;
|
|
const int traversal_stride_h = config.problem_size.stride_h;
|
|
const int traversal_stride_w = config.problem_size.stride_w;
|
|
const int dilation_h = config.problem_size.dilation_h;
|
|
const int dilation_w = config.problem_size.dilation_w;
|
|
|
|
// CUTLASS 3's implicit GEMM convolution kernels currently only
|
|
// support cross correlation (passing over the activation and
|
|
// filter tensors in the same order). The convolution mode is
|
|
// future work.
|
|
const auto mode = config.problem_size.mode;
|
|
if (mode != cutlass::conv::Mode::kCrossCorrelation) {
|
|
CUTLASS_TRACE_HOST("Convolution modes other than kCrossCorrelation "
|
|
"are not currently supported.");
|
|
return Status::kInvalid;
|
|
}
|
|
|
|
constexpr int num_spatial_dims = Operator::NumSpatialDimensions;
|
|
constexpr size_t stride_size = size_t(num_spatial_dims) + 2u;
|
|
constexpr auto the_stride_size = std::integral_constant<size_t, stride_size>{};
|
|
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
std::cerr << " num_spatial_dims = " << num_spatial_dims << "\n"
|
|
<< " stride_size = " << stride_size << "\n";
|
|
auto print_stride = [] (auto const& stride, char const variable_name[]) {
|
|
std::cerr << " " << variable_name << ": [";
|
|
for (size_t k = 0; k < stride.size(); ++k) {
|
|
std::cerr << stride[k];
|
|
if (k + 1u < stride.size()) {
|
|
std::cerr << ", ";
|
|
}
|
|
}
|
|
std::cerr << "]\n";
|
|
};
|
|
print_stride(config.stride_a, "config.stride_a");
|
|
print_stride(config.stride_b, "config.stride_b");
|
|
print_stride(config.stride_c, "config.stride_c");
|
|
#endif
|
|
|
|
// Conv2dConfiguration stores the strides as std::vector,
|
|
// so the code needs to check the run-time vector lengths.
|
|
if (config.stride_a.size() + 1u != stride_size) {
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL)
|
|
std::ostringstream os;
|
|
os << "config.stride_a.size() + 1u = "
|
|
<< (config.stride_a.size() + 1u)
|
|
<< " != num_spatial_dims + 2u = " << stride_size;
|
|
CUTLASS_TRACE_HOST( os.str() );
|
|
#endif
|
|
return Status::kInvalid;
|
|
}
|
|
if (config.stride_b.size() + 1u != stride_size) {
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL)
|
|
std::ostringstream os;
|
|
os << "config.stride_b.size() + 1u = "
|
|
<< (config.stride_b.size() + 1u)
|
|
<< " != num_spatial_dims + 2u = " << stride_size;
|
|
CUTLASS_TRACE_HOST( os.str() );
|
|
#endif
|
|
return Status::kInvalid;
|
|
}
|
|
if (config.stride_c.size() + 1u != stride_size) {
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL)
|
|
std::ostringstream os;
|
|
os << "config.stride_c.size() + 1u = "
|
|
<< (config.stride_c.size() + 1u)
|
|
<< " != num_spatial_dims + 2u = " << stride_size;
|
|
CUTLASS_TRACE_HOST( os.str() );
|
|
#endif
|
|
return Status::kInvalid;
|
|
}
|
|
|
|
constexpr cutlass::conv::Operator conv_op = Operator::DispatchPolicy::ConvOp;
|
|
using problem_shape_type =
|
|
cutlass::conv::ConvProblemShape<conv_op, num_spatial_dims>;
|
|
// cute::array<int64_t, RankT>; must convert to the kernel's native strides
|
|
using TensorStride = typename problem_shape_type::TensorStride;
|
|
|
|
const TensorStride stride_A = vector_to_array_strides(config.stride_a, the_stride_size);
|
|
const TensorStride stride_B = vector_to_array_strides(config.stride_b, the_stride_size);
|
|
const TensorStride stride_C = vector_to_array_strides(config.stride_c, the_stride_size);
|
|
|
|
// cutlass::library::Conv2dConfiguration has no member stride_d.
|
|
// The code below imitates the testbed,
|
|
// which just sets D's strides to C's strides.
|
|
|
|
const int num_groups = config.problem_size.groups;
|
|
if (num_groups != 1) {
|
|
CUTLASS_TRACE_HOST("CUTLASS 3 kernels currently only support groups = 1.");
|
|
return Status::kInvalid;
|
|
}
|
|
// ConvProblemShape is how CUTLASS 3 kernels represent
|
|
// convolution problems. ConvProblemShape's constructors take
|
|
// shape_act, stride_act, shape_flt, and stride_flt, and set
|
|
// shape_A, stride_A, shape_B, stride_B, shape_C, and stride_C
|
|
// according to Fprop / Dgrad / Wgrad.
|
|
//
|
|
// This means that stride_act isn't always config.stride_A,
|
|
// depending on Fprop / Dgrad / Wgrad. The code here "undoes"
|
|
// the logic in Conv2dWorkspace::set_stride_vector so that we
|
|
// can recover the strides of the activation and filter tensors.
|
|
// It doesn't need to worry about the so-called "output" tensor
|
|
// (which might not be C), as ConvProblemShape's constructor
|
|
// figures out its shapes and strides.
|
|
using TensorExtent = typename problem_shape_type::TensorExtent;
|
|
TensorExtent shape_act{N, H, W, C};
|
|
auto stride_act = [&] () {
|
|
// Some compilers consider conv_op (defined above), as
|
|
// captured by this lambda, as "not a constant expression."
|
|
constexpr auto conv_kind = Operator::DispatchPolicy::ConvOp;
|
|
if constexpr (conv_kind == cutlass::conv::Operator::kFprop) {
|
|
return stride_A;
|
|
}
|
|
else if constexpr (conv_kind == cutlass::conv::Operator::kDgrad) {
|
|
return stride_C;
|
|
}
|
|
else { // conv_kind == cutlass::conv::Operator::kWgrad
|
|
return stride_B;
|
|
}
|
|
} ();
|
|
TensorExtent shape_flt{K, R, S, C};
|
|
auto stride_flt = [&] () {
|
|
// Some compilers consider conv_op (defined above), as
|
|
// captured by this lambda, as "not a constant expression."
|
|
constexpr auto conv_kind = Operator::DispatchPolicy::ConvOp;
|
|
if constexpr (conv_kind == cutlass::conv::Operator::kFprop) {
|
|
return stride_B;
|
|
}
|
|
else if constexpr (conv_kind == cutlass::conv::Operator::kDgrad) {
|
|
return stride_B;
|
|
}
|
|
else { // conv_kind == cutlass::conv::Operator::kWgrad
|
|
return stride_C;
|
|
}
|
|
} ();
|
|
|
|
problem_shape_type problem_shape(
|
|
/* mode = */ mode,
|
|
/* shape_act = */ shape_act,
|
|
/* stride_act = */ stride_act,
|
|
/* shape_flt = */ shape_flt,
|
|
/* stride_flt = */ stride_flt,
|
|
/* lower_padding = */ {pad_h, pad_w},
|
|
/* upper_padding = */ {pad_h, pad_w},
|
|
/* traversal_stride = */ {traversal_stride_h, traversal_stride_w},
|
|
/* dilation = */ {dilation_h, dilation_w},
|
|
num_groups);
|
|
out_args.problem_shape = problem_shape;
|
|
|
|
// ConvProblemShape's constructor sets its shape_C member.
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
printf("\n problem_shape.shape_C: ");
|
|
print(problem_shape.shape_C);
|
|
printf("\n problem_shape.stride_C: ");
|
|
print(problem_shape.stride_C);
|
|
printf("\n");
|
|
#endif
|
|
// Initialization of C's and D's strides follows the CUTLASS 3
|
|
// convolutions testbed (test/unit/conv/device_3x/testbed_conv.hpp).
|
|
{
|
|
using StrideC = typename Operator::ConvKernel::StrideC;
|
|
using StrideD = typename Operator::ConvKernel::StrideD;
|
|
auto stride_C = StrideC{};
|
|
auto stride_D = StrideD{};
|
|
|
|
if constexpr (conv_op == cutlass::conv::Operator::kWgrad) {
|
|
stride_C = cutlass::make_cute_packed_stride(
|
|
StrideC{}, problem_shape.shape_C, problem_shape.stride_C, conv_op);
|
|
stride_D = cutlass::make_cute_packed_stride(
|
|
StrideD{}, problem_shape.shape_C, problem_shape.stride_C, conv_op);
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
std::cerr << " Wgrad: stride_C: " << stride_C << "\n";
|
|
#endif
|
|
}
|
|
else {
|
|
cute::for_each(cute::make_seq<cute::rank<0>(StrideC{})>{}, [&](auto i) {
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
const auto stride_C_i = problem_shape.stride_C[problem_shape_type::RankT-2-i];
|
|
std::cerr << " Fprop or Dgrad: get<0, " << i << ">(stride_C): "
|
|
<< stride_C_i << "\n";
|
|
#endif
|
|
cute::get<0, i>(stride_C) = problem_shape.stride_C[problem_shape_type::RankT-2-i];
|
|
});
|
|
cute::for_each(cute::make_seq<cute::rank<0>(StrideD{})>{}, [&](auto i) {
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
const auto stride_D_i = problem_shape.stride_C[problem_shape_type::RankT-2-i];
|
|
std::cerr << " Fprop or Dgrad: get<0, " << i << ">(stride_D): "
|
|
<< stride_D_i << "\n";
|
|
#endif
|
|
cute::get<0, i>(stride_D) = problem_shape.stride_C[problem_shape_type::RankT-2-i];
|
|
});
|
|
}
|
|
out_args.epilogue.dC = stride_C;
|
|
out_args.epilogue.dD = stride_D;
|
|
}
|
|
return Status::kSuccess;
|
|
}
|
|
}
|
|
|
|
static Status update_operator_arguments_from_configuration(
|
|
typename Operator::Arguments& out_args,
|
|
Conv3dConfiguration const& config)
|
|
{
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
CUTLASS_TRACE_HOST("ConvOperator3x::"
|
|
"update_operator_arguments_from_configuration"
|
|
"(Conv3dConfiguration)\n");
|
|
#endif
|
|
using detail::coord_to_array_strides;
|
|
|
|
constexpr int num_spatial_dims = Operator::NumSpatialDimensions;
|
|
if constexpr (num_spatial_dims != 3) {
|
|
CUTLASS_TRACE_HOST("You can only use Conv3dConfiguration "
|
|
"with an Operator whose NumSpatialDimensions is exactly 3.");
|
|
return Status::kInvalid;
|
|
}
|
|
else {
|
|
// Convolutions split the metadata (in Conv3dConfiguration) from
|
|
// the data (ConvArguments, which only has pointers and a single
|
|
// enum value). Thus, this class will need both the
|
|
// configuration and the (user's input) arguments to set up the
|
|
// kernel's arguments. This function can fill in what the
|
|
// configuration has now, but the class will need the user's
|
|
// input arguments later.
|
|
if (config.split_k_mode != conv::SplitKMode::kSerial) {
|
|
CUTLASS_TRACE_HOST("CUTLASS 3 convolutions currently only support split_k_mode = kSerial.");
|
|
return Status::kInvalid;
|
|
}
|
|
// config.problem_size.split_k_slices is only meaningful if
|
|
// split_k_mode != kSerial. If this code later supports other
|
|
// split_k_mode values, then it will also need to read
|
|
// split_k_slices.
|
|
|
|
const int N = config.problem_size.N;
|
|
const int D = config.problem_size.D;
|
|
const int H = config.problem_size.H;
|
|
const int W = config.problem_size.W;
|
|
const int C = config.problem_size.C;
|
|
const int K = config.problem_size.K;
|
|
const int T = config.problem_size.T;
|
|
const int R = config.problem_size.R;
|
|
const int S = config.problem_size.S;
|
|
const int pad_d = config.problem_size.pad_d;
|
|
const int pad_h = config.problem_size.pad_h;
|
|
const int pad_w = config.problem_size.pad_w;
|
|
const int traversal_stride_d = config.problem_size.stride_d;
|
|
const int traversal_stride_h = config.problem_size.stride_h;
|
|
const int traversal_stride_w = config.problem_size.stride_w;
|
|
const int dilation_d = config.problem_size.dilation_d;
|
|
const int dilation_h = config.problem_size.dilation_h;
|
|
const int dilation_w = config.problem_size.dilation_w;
|
|
|
|
// CUTLASS 3's implicit GEMM convolution kernels currently only
|
|
// support cross correlation (passing over the activation and
|
|
// filter tensors in the same order). The convolution mode is
|
|
// future work.
|
|
const auto mode = config.problem_size.mode;
|
|
if (mode != cutlass::conv::Mode::kCrossCorrelation) {
|
|
CUTLASS_TRACE_HOST("Convolution modes other than kCrossCorrelation "
|
|
"are not currently supported.");
|
|
return Status::kInvalid;
|
|
}
|
|
|
|
using Stride = cutlass::layout::TensorNDHWC::Stride;
|
|
static_assert(std::is_same_v<Stride, cutlass::Coord<4>>);
|
|
|
|
const cutlass::library::ConvKind conv_kind = [] () {
|
|
constexpr cutlass::conv::Operator op = Operator::DispatchPolicy::ConvOp;
|
|
if constexpr (op == cutlass::conv::Operator::kFprop) {
|
|
return library::ConvKind::kFprop;
|
|
}
|
|
else if constexpr (op == cutlass::conv::Operator::kDgrad) {
|
|
return library::ConvKind::kDgrad;
|
|
}
|
|
else /* if constexpr (op == cutlass::conv::Operator::kWgrad) */ {
|
|
return library::ConvKind::kWgrad;
|
|
}
|
|
} ();
|
|
const Stride input_stride_a = config.layout_a(conv_kind).stride();
|
|
const Stride input_stride_b = config.layout_b(conv_kind).stride();
|
|
const Stride input_stride_c = config.layout_c(conv_kind).stride();
|
|
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
constexpr size_t stride_size = size_t(num_spatial_dims) + 2u;
|
|
std::cerr << " num_spatial_dims = " << num_spatial_dims << "\n"
|
|
<< " stride_size = " << stride_size << "\n";
|
|
auto print_stride = [] (Stride const& stride, char const variable_name[]) {
|
|
std::cerr << " " << variable_name << ": [";
|
|
for (size_t k = 0; k < Stride::kRank; ++k) {
|
|
std::cerr << stride[static_cast<int>(k)];
|
|
if (k + 1u < Stride::kRank) {
|
|
std::cerr << ", ";
|
|
}
|
|
}
|
|
std::cerr << "]\n";
|
|
};
|
|
print_stride(input_stride_a, "input_stride_a");
|
|
print_stride(input_stride_b, "input_stride_b");
|
|
print_stride(input_stride_c, "input_stride_c");
|
|
#endif
|
|
// Conv3dConfiguration stores the strides as Coord (with
|
|
// compile-time size), so there's no need to check sizes here
|
|
// (unlike Conv2dConfiguration, which stores strides as
|
|
// std::vector).
|
|
|
|
constexpr cutlass::conv::Operator conv_op = Operator::DispatchPolicy::ConvOp;
|
|
using problem_shape_type =
|
|
cutlass::conv::ConvProblemShape<conv_op, num_spatial_dims>;
|
|
// cute::array<int64_t, RankT>; must convert to the kernel's native strides
|
|
using TensorStride = typename problem_shape_type::TensorStride;
|
|
|
|
const TensorStride stride_A = coord_to_array_strides(input_stride_a);
|
|
const TensorStride stride_B = coord_to_array_strides(input_stride_b);
|
|
const TensorStride stride_C = coord_to_array_strides(input_stride_c);
|
|
|
|
const int num_groups = config.problem_size.groups;
|
|
if (num_groups != 1) {
|
|
CUTLASS_TRACE_HOST("CUTLASS 3 kernels currently only support groups = 1.");
|
|
return Status::kInvalid;
|
|
}
|
|
// ConvProblemShape is how CUTLASS 3 kernels represent
|
|
// convolution problems. ConvProblemShape's constructors take
|
|
// shape_act, stride_act, shape_flt, and stride_flt, and set
|
|
// shape_A, stride_A, shape_B, stride_B, shape_C, and stride_C
|
|
// according to Fprop / Dgrad / Wgrad.
|
|
//
|
|
// Conv3dConfiguration differs a bit from Conv2dConfiguration,
|
|
// but the idea is the same: the "input_stride_a" from config
|
|
// depends on conv_kind (Fprop, Dgrad, or Wgrad), so stride_act
|
|
// isn't always input_stride_a. Analogously, stride_flt isn't
|
|
// always input_stride_b. The code here "undoes" the logic in
|
|
// config.layout_a(conv_kind) and config.layout_b(conv_kind)
|
|
// (analogous to Conv2dWorkspace::set_stride_vector) so that we
|
|
// can recover the strides of the activation and filter tensors.
|
|
// It doesn't need to worry about the so-called "output" tensor
|
|
// (which might not be C), as ConvProblemShape's constructor
|
|
// figures out its shapes and strides.
|
|
using TensorExtent = typename problem_shape_type::TensorExtent;
|
|
TensorExtent shape_act{N, D, H, W, C};
|
|
auto stride_act = [&] () {
|
|
// Some compilers consider conv_op (defined above), as
|
|
// captured by this lambda, as "not a constant expression."
|
|
constexpr auto conv_kind = Operator::DispatchPolicy::ConvOp;
|
|
if constexpr (conv_kind == cutlass::conv::Operator::kFprop) {
|
|
return stride_A;
|
|
}
|
|
else if constexpr (conv_kind == cutlass::conv::Operator::kDgrad) {
|
|
return stride_C;
|
|
}
|
|
else { // conv_kind == cutlass::conv::Operator::kWgrad
|
|
return stride_B;
|
|
}
|
|
} ();
|
|
TensorExtent shape_flt{K, T, R, S, C};
|
|
auto stride_flt = [&] () {
|
|
// Some compilers consider conv_op (defined above), as
|
|
// captured by this lambda, as "not a constant expression."
|
|
constexpr auto conv_kind = Operator::DispatchPolicy::ConvOp;
|
|
if constexpr (conv_kind == cutlass::conv::Operator::kFprop) {
|
|
return stride_B;
|
|
}
|
|
else if constexpr (conv_kind == cutlass::conv::Operator::kDgrad) {
|
|
return stride_B;
|
|
}
|
|
else { // conv_kind == cutlass::conv::Operator::kWgrad
|
|
return stride_C;
|
|
}
|
|
} ();
|
|
|
|
problem_shape_type problem_shape(
|
|
/* mode = */ mode,
|
|
/* shape_act = */ shape_act,
|
|
/* stride_act = */ stride_act,
|
|
/* shape_flt = */ shape_flt,
|
|
/* stride_flt = */ stride_flt,
|
|
/* lower_padding = */ {pad_d, pad_h, pad_w},
|
|
/* upper_padding = */ {pad_d, pad_h, pad_w},
|
|
/* traversal_stride = */ {traversal_stride_d, traversal_stride_h, traversal_stride_w},
|
|
/* dilation = */ {dilation_d, dilation_h, dilation_w},
|
|
num_groups);
|
|
out_args.problem_shape = problem_shape;
|
|
|
|
// ConvProblemShape's constructor sets its shape_C member.
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
printf("\n problem_shape.shape_C: ");
|
|
print(problem_shape.shape_C);
|
|
printf("\n problem_shape.stride_C: ");
|
|
print(problem_shape.stride_C);
|
|
printf("\n");
|
|
#endif
|
|
// Initialization of C's and D's strides follows the CUTLASS 3
|
|
// convolutions testbed (test/unit/conv/device_3x/testbed_conv.hpp).
|
|
{
|
|
using StrideC = typename Operator::ConvKernel::StrideC;
|
|
using StrideD = typename Operator::ConvKernel::StrideD;
|
|
auto stride_C = StrideC{};
|
|
auto stride_D = StrideD{};
|
|
|
|
if constexpr (conv_op == cutlass::conv::Operator::kWgrad) {
|
|
stride_C = cutlass::make_cute_packed_stride(
|
|
StrideC{}, problem_shape.shape_C, problem_shape.stride_C, conv_op);
|
|
stride_D = cutlass::make_cute_packed_stride(
|
|
StrideD{}, problem_shape.shape_C, problem_shape.stride_C, conv_op);
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
std::cerr << " Wgrad: stride_C: " << stride_C << "\n";
|
|
#endif
|
|
}
|
|
else {
|
|
cute::for_each(cute::make_seq<cute::rank<0>(StrideC{})>{}, [&](auto i) {
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
const auto stride_C_i = problem_shape.stride_C[problem_shape_type::RankT-2-i];
|
|
std::cerr << " Fprop or Dgrad: get<0, " << i << ">(stride_C): "
|
|
<< stride_C_i << "\n";
|
|
#endif
|
|
cute::get<0, i>(stride_C) = problem_shape.stride_C[problem_shape_type::RankT-2-i];
|
|
});
|
|
cute::for_each(cute::make_seq<cute::rank<0>(StrideD{})>{}, [&](auto i) {
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
const auto stride_D_i = problem_shape.stride_C[problem_shape_type::RankT-2-i];
|
|
std::cerr << " Fprop or Dgrad: get<0, " << i << ">(stride_D): "
|
|
<< stride_D_i << "\n";
|
|
#endif
|
|
cute::get<0, i>(stride_D) = problem_shape.stride_C[problem_shape_type::RankT-2-i];
|
|
});
|
|
}
|
|
out_args.epilogue.dC = stride_C;
|
|
out_args.epilogue.dD = stride_D;
|
|
}
|
|
return Status::kSuccess;
|
|
}
|
|
}
|
|
|
|
Status update_operator_arguments_from_arguments(
|
|
typename Operator::Arguments& out_args,
|
|
ConvArguments const& in_args) const
|
|
{
|
|
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
|
|
CUTLASS_TRACE_HOST("ConvOperation3x::update_operator_arguments_from_arguments\n");
|
|
#endif
|
|
auto status = UpdateFusionArgs<decltype(out_args.epilogue.thread)>::update_(
|
|
out_args.epilogue.thread, in_args);
|
|
if (status != Status::kSuccess) {
|
|
return status;
|
|
}
|
|
|
|
out_args.mainloop.ptr_A = reinterpret_cast<ElementA const*>(in_args.A);
|
|
out_args.mainloop.ptr_B = reinterpret_cast<ElementB const*>(in_args.B);
|
|
|
|
out_args.epilogue.ptr_C = reinterpret_cast<ElementC const*>(in_args.C);
|
|
out_args.epilogue.ptr_D = reinterpret_cast<ElementD*>(in_args.D);
|
|
|
|
return Status::kSuccess;
|
|
}
|
|
};
|
|
|
|
} // namespace cutlass::library
|