Stream-K with broadcast (#892)
* [WIP] GEMM StreamK w/ Fused Epilogue * Adds Gemm Streamk with Fused Epilogue kernel level struct. * Mostly based on Gemm with Fused Epilogue, * Requires a new epilogue * Work in progress * [WIP] StreamK support for GemmUniversalWithBroadcast * Just based off of how StreamK is allowed in GemmUniversal * Untested and a work in progress * Minor fixes * [WIP] It compiles! It is almost certainly incorrect, but we're past getting the templates to match, so checkpointing. * Correction to reference kernel * Fix typo * Added MSE measurement * Switch back to reference kernel + host for loop Still WIP. Now we're getting even a larger MSE, but it's both on basic Split-K and Stream-K. * Fix typos * Fix broadcast vector + requested changes * Comment typo * Small int option and more * Fix incorrect condition on source needed * Requested changes * I think I got it? * Bias vector should be stride 0 * Two source added! * Typos * Merge examples * Bring back vector row offset Just to ensure consistency with universal gemm with fused epilogue * Base arguments and params structs for StreamK * StreamK epilogue with broadcast now inherits the original * undo params_streamk_base.h --------- Co-authored-by: Ali Hassani <ahassanijr@gmail.com> Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Ali Hassani
@@ -48,6 +48,7 @@
|
||||
#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/epilogue.h"
|
||||
#include "cutlass/epilogue/threadblock/epilogue_with_broadcast.h"
|
||||
#include "cutlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h"
|
||||
|
||||
#include "cutlass/layout/permute.h"
|
||||
|
||||
@@ -120,6 +121,67 @@ struct DefaultEpilogueWithBroadcastTensorOp {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for streamk epilogues for TensorOps.
|
||||
template <
|
||||
typename Shape,
|
||||
typename WarpMmaTensorOp,
|
||||
int PartitionsK,
|
||||
typename ElementOutput,
|
||||
typename ElementTensor,
|
||||
typename ElementVector,
|
||||
typename OutputOp,
|
||||
int ElementsPerAccess,
|
||||
bool ScatterD = false,
|
||||
typename PermuteDLayout = layout::NoPermute
|
||||
>
|
||||
struct DefaultStreamkEpilogueWithBroadcastTensorOp {
|
||||
|
||||
/// Use defaults related to the existing epilogue
|
||||
using Base = DefaultEpilogueTensorOp<
|
||||
Shape,
|
||||
WarpMmaTensorOp,
|
||||
PartitionsK,
|
||||
OutputOp,
|
||||
ElementsPerAccess
|
||||
>;
|
||||
|
||||
//
|
||||
// Stores the result z = (y = GEMM(A, B, C), broadcast)
|
||||
//
|
||||
using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<
|
||||
typename Base::OutputTileThreadMap,
|
||||
ElementOutput,
|
||||
ScatterD,
|
||||
PermuteDLayout
|
||||
>;
|
||||
|
||||
//
|
||||
// Additional tensor tile iterator - stores t = Elementwise(z)
|
||||
//
|
||||
using TensorTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<
|
||||
typename Base::OutputTileThreadMap,
|
||||
ElementTensor
|
||||
>;
|
||||
|
||||
/// Define the epilogue
|
||||
using Epilogue = EpilogueStreamkWithBroadcast<
|
||||
Shape,
|
||||
WarpMmaTensorOp,
|
||||
PartitionsK,
|
||||
OutputTileIterator,
|
||||
TensorTileIterator,
|
||||
ElementVector,
|
||||
typename Base::AccumulatorFragmentIterator,
|
||||
typename Base::WarpTileIterator,
|
||||
typename Base::SharedLoadIterator,
|
||||
OutputOp,
|
||||
typename Base::Padding,
|
||||
Base::kFragmentsPerIteration
|
||||
>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues for VoltaTensorOps.
|
||||
template <
|
||||
typename Shape,
|
||||
|
||||
@@ -0,0 +1,443 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2023 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 Epilogue for threadblock scoped GEMMs using Tensor Ops.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#if defined(__CUDACC_RTC__)
|
||||
#include <cuda/std/cassert>
|
||||
#include <cuda/std/utility>
|
||||
#else
|
||||
#include <assert.h>
|
||||
#include <utility>
|
||||
#endif
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
#include "cutlass/tensor_coord.h"
|
||||
#include "cutlass/aligned_buffer.h"
|
||||
#include "cutlass/functional.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cutlass/layout/vector.h"
|
||||
#include "cutlass/layout/tensor.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/transform/pitch_linear_thread_map.h"
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/epilogue_base.h"
|
||||
#include "cutlass/epilogue/threadblock/epilogue_base_streamk.h"
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// This base class is meant to define the concept required of the
|
||||
/// EpilogueStreamkWithBroadcast::OutputOp
|
||||
template <
|
||||
typename ElementC_,
|
||||
typename ElementAccumulator_,
|
||||
typename ElementCompute_,
|
||||
typename ElementZ_,
|
||||
typename ElementT_,
|
||||
int ElementsPerAccess,
|
||||
bool StoreZ = true,
|
||||
bool StoreT = true
|
||||
>
|
||||
struct EpilogueStreamkWithBroadcastOpBase : EpilogueWithBroadcastOpBase<
|
||||
ElementC_,
|
||||
ElementAccumulator_,
|
||||
ElementCompute_,
|
||||
ElementZ_,
|
||||
ElementT_,
|
||||
ElementsPerAccess,
|
||||
StoreZ,
|
||||
StoreT
|
||||
>
|
||||
{
|
||||
|
||||
/// Parameters structure - required
|
||||
struct Params { };
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor from Params
|
||||
EpilogueStreamkWithBroadcastOpBase(Params const ¶ms_) { }
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Epilogue operator with bias vector broadcast over columns.
|
||||
///
|
||||
/// Computes the following:
|
||||
///
|
||||
///
|
||||
/// Z, T = OutputOp(AB, C, Broadcast)
|
||||
///
|
||||
/// if (ElementwiseOp::kStoreZ) {
|
||||
/// store(converted_u);
|
||||
/// }
|
||||
///
|
||||
/// if (ElementwiseOp::kStoreT) {
|
||||
/// store(v);
|
||||
/// }
|
||||
///
|
||||
template <
|
||||
typename Shape_, ///< Shape of threadblock tile (concept: GemmShape)
|
||||
typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp)
|
||||
int PartitionsK, ///< Number of partitions of the K dimension
|
||||
typename OutputTileIterator_, ///< Tile iterator reading and writing output tensors (z)
|
||||
typename TensorTileIterator_, ///< Additional tile iterator for tensor-valued operands (t)
|
||||
typename ElementVector_, ///< Pointer to broadcast vector
|
||||
typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators
|
||||
typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM
|
||||
typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading from SMEM
|
||||
typename OutputOp_, ///< Output operator - concept is EpilogueWithBroadcastOp
|
||||
typename Padding_, ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape)
|
||||
int FragmentsPerPartition = 1, ///< Used to coarsten the epilogue granularity
|
||||
int IterationsUnroll = ///< Used to reduce binary size when epilogue op is large
|
||||
(!IsEpilogueFunctorHeavy<OutputOp_>::value),
|
||||
bool IsSingleSource = OutputOp_::kIsSingleSource
|
||||
>
|
||||
class EpilogueStreamkWithBroadcast;
|
||||
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// EpilogueStreamkWithBroadcast: Two sources
|
||||
|
||||
template <
|
||||
typename Shape_,
|
||||
typename WarpMmaOperator_,
|
||||
int PartitionsK,
|
||||
typename OutputTileIterator_,
|
||||
typename TensorTileIterator_,
|
||||
typename ElementVector_,
|
||||
typename AccumulatorFragmentIterator_,
|
||||
typename WarpTileIterator_,
|
||||
typename SharedLoadIterator_,
|
||||
typename OutputOp_,
|
||||
typename Padding_,
|
||||
int FragmentsPerPartition,
|
||||
int IterationsUnroll
|
||||
>
|
||||
class EpilogueStreamkWithBroadcast<
|
||||
Shape_,
|
||||
WarpMmaOperator_,
|
||||
PartitionsK,
|
||||
OutputTileIterator_,
|
||||
TensorTileIterator_,
|
||||
ElementVector_,
|
||||
AccumulatorFragmentIterator_,
|
||||
WarpTileIterator_,
|
||||
SharedLoadIterator_,
|
||||
OutputOp_,
|
||||
Padding_,
|
||||
FragmentsPerPartition,
|
||||
IterationsUnroll,
|
||||
false
|
||||
> :
|
||||
public EpilogueWithBroadcast<
|
||||
Shape_,
|
||||
WarpMmaOperator_,
|
||||
PartitionsK,
|
||||
OutputTileIterator_,
|
||||
TensorTileIterator_,
|
||||
ElementVector_,
|
||||
AccumulatorFragmentIterator_,
|
||||
WarpTileIterator_,
|
||||
SharedLoadIterator_,
|
||||
OutputOp_,
|
||||
Padding_,
|
||||
FragmentsPerPartition,
|
||||
IterationsUnroll,
|
||||
false>,
|
||||
public EpilogueBaseStreamK<
|
||||
Shape_,
|
||||
PartitionsK,
|
||||
WarpMmaOperator_,
|
||||
AccumulatorFragmentIterator_>
|
||||
{
|
||||
|
||||
public:
|
||||
|
||||
using Base = EpilogueWithBroadcast<
|
||||
Shape_,
|
||||
WarpMmaOperator_,
|
||||
PartitionsK,
|
||||
OutputTileIterator_,
|
||||
TensorTileIterator_,
|
||||
ElementVector_,
|
||||
AccumulatorFragmentIterator_,
|
||||
WarpTileIterator_,
|
||||
SharedLoadIterator_,
|
||||
OutputOp_,
|
||||
Padding_,
|
||||
FragmentsPerPartition,
|
||||
IterationsUnroll,
|
||||
false>;
|
||||
|
||||
using BaseStreamK = EpilogueBaseStreamK<
|
||||
Shape_,
|
||||
PartitionsK,
|
||||
WarpMmaOperator_,
|
||||
AccumulatorFragmentIterator_>;
|
||||
|
||||
using Shape = Shape_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputTileIterator = OutputTileIterator_;
|
||||
using TensorTileIterator = TensorTileIterator_;
|
||||
using ElementVector = ElementVector_;
|
||||
using SharedLoadIterator = SharedLoadIterator_;
|
||||
using OutputOp = OutputOp_;
|
||||
|
||||
/// Fragment type used by the accumulator tile's fragment iterator
|
||||
using AccumulatorFragment = typename Base::AccumulatorFragmentIterator::Fragment;
|
||||
|
||||
/// Shared storage structure (shadows base) with additional SMEM buffer for reduction
|
||||
using SharedStorage = typename Base::SharedStorage;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
EpilogueStreamkWithBroadcast(
|
||||
SharedStorage &shared_storage, ///< Shared storage object
|
||||
int thread_idx, ///< ID of a thread within the threadblock
|
||||
int warp_idx, ///< ID of warp within threadblock
|
||||
int lane_idx ///< Id of thread within warp
|
||||
):
|
||||
Base(shared_storage, thread_idx, warp_idx, lane_idx),
|
||||
BaseStreamK(thread_idx)
|
||||
{ }
|
||||
|
||||
|
||||
/// Aggregates the accumulator sets shared by peer blocks in the global workspace,
|
||||
/// performing epilogue computations, writing to output
|
||||
CUTLASS_DEVICE
|
||||
void reduce(
|
||||
int peer_idx_begin,
|
||||
int peer_idx_end,
|
||||
int reduce_fragment_idx,
|
||||
void *element_workspace,
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
ElementVector const * broadcast_ptr, ///< Broadcast vector
|
||||
OutputTileIterator destination_iterator, ///< Tile iterator for destination
|
||||
OutputTileIterator source_iterator1, ///< Tile iterator for first source accumulator matrix
|
||||
OutputTileIterator source_iterator2, ///< Tile iterator for second source accumulator matrix
|
||||
TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand
|
||||
MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses
|
||||
MatrixCoord(Shape::kM, Shape::kN),
|
||||
MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space
|
||||
MatrixCoord())
|
||||
{
|
||||
// Reduce peer accumulator fragments into one fragment
|
||||
AccumulatorFragment accum_fragment;
|
||||
BaseStreamK::reduce(accum_fragment, peer_idx_begin, peer_idx_end, reduce_fragment_idx, element_workspace);
|
||||
|
||||
// Store fragment to shared memory
|
||||
this->warp_tile_iterator_.store(accum_fragment);
|
||||
|
||||
__syncthreads();
|
||||
|
||||
Base::reduce(reduce_fragment_idx, output_op, broadcast_ptr, destination_iterator, source_iterator1, source_iterator2, tensor_iterator, problem_size, threadblock_offset);
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// EpilogueStreamkWithBroadcast: Single source
|
||||
|
||||
template <
|
||||
typename Shape_,
|
||||
typename WarpMmaOperator_,
|
||||
int PartitionsK,
|
||||
typename OutputTileIterator_,
|
||||
typename TensorTileIterator_,
|
||||
typename ElementVector_,
|
||||
typename AccumulatorFragmentIterator_,
|
||||
typename WarpTileIterator_,
|
||||
typename SharedLoadIterator_,
|
||||
typename OutputOp_,
|
||||
typename Padding_,
|
||||
int FragmentsPerPartition,
|
||||
int IterationsUnroll
|
||||
>
|
||||
class EpilogueStreamkWithBroadcast<
|
||||
Shape_,
|
||||
WarpMmaOperator_,
|
||||
PartitionsK,
|
||||
OutputTileIterator_,
|
||||
TensorTileIterator_,
|
||||
ElementVector_,
|
||||
AccumulatorFragmentIterator_,
|
||||
WarpTileIterator_,
|
||||
SharedLoadIterator_,
|
||||
OutputOp_,
|
||||
Padding_,
|
||||
FragmentsPerPartition,
|
||||
IterationsUnroll,
|
||||
true
|
||||
> :
|
||||
public EpilogueWithBroadcast<
|
||||
Shape_,
|
||||
WarpMmaOperator_,
|
||||
PartitionsK,
|
||||
OutputTileIterator_,
|
||||
TensorTileIterator_,
|
||||
ElementVector_,
|
||||
AccumulatorFragmentIterator_,
|
||||
WarpTileIterator_,
|
||||
SharedLoadIterator_,
|
||||
OutputOp_,
|
||||
Padding_,
|
||||
FragmentsPerPartition,
|
||||
IterationsUnroll,
|
||||
true>,
|
||||
public EpilogueBaseStreamK<
|
||||
Shape_,
|
||||
PartitionsK,
|
||||
WarpMmaOperator_,
|
||||
AccumulatorFragmentIterator_>
|
||||
{
|
||||
|
||||
public:
|
||||
|
||||
using Base = EpilogueWithBroadcast<
|
||||
Shape_,
|
||||
WarpMmaOperator_,
|
||||
PartitionsK,
|
||||
OutputTileIterator_,
|
||||
TensorTileIterator_,
|
||||
ElementVector_,
|
||||
AccumulatorFragmentIterator_,
|
||||
WarpTileIterator_,
|
||||
SharedLoadIterator_,
|
||||
OutputOp_,
|
||||
Padding_,
|
||||
FragmentsPerPartition,
|
||||
IterationsUnroll,
|
||||
true>;
|
||||
|
||||
using BaseStreamK = EpilogueBaseStreamK<
|
||||
Shape_,
|
||||
PartitionsK,
|
||||
WarpMmaOperator_,
|
||||
AccumulatorFragmentIterator_>;
|
||||
|
||||
using Shape = Shape_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputTileIterator = OutputTileIterator_;
|
||||
using TensorTileIterator = TensorTileIterator_;
|
||||
using ElementVector = ElementVector_;
|
||||
using SharedLoadIterator = SharedLoadIterator_;
|
||||
using OutputOp = OutputOp_;
|
||||
|
||||
/// Fragment type used by the accumulator tile's fragment iterator
|
||||
using AccumulatorFragment = typename Base::AccumulatorFragmentIterator::Fragment;
|
||||
|
||||
/// Shared storage structure (shadows base) with additional SMEM buffer for reduction
|
||||
using SharedStorage = typename Base::SharedStorage;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
EpilogueStreamkWithBroadcast(
|
||||
SharedStorage &shared_storage, ///< Shared storage object
|
||||
int thread_idx, ///< ID of a thread within the threadblock
|
||||
int warp_idx, ///< ID of warp within threadblock
|
||||
int lane_idx ///< Id of thread within warp
|
||||
):
|
||||
Base(shared_storage, thread_idx, warp_idx, lane_idx),
|
||||
BaseStreamK(thread_idx)
|
||||
{ }
|
||||
|
||||
|
||||
/// Aggregates the accumulator sets shared by peer blocks in the global workspace,
|
||||
/// performing epilogue computations, writing to output
|
||||
CUTLASS_DEVICE
|
||||
void reduce(
|
||||
int peer_idx_begin,
|
||||
int peer_idx_end,
|
||||
int reduce_fragment_idx,
|
||||
void *element_workspace,
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
ElementVector const * broadcast_ptr, ///< Broadcast vector
|
||||
OutputTileIterator destination_iterator, ///< Tile iterator for destination
|
||||
OutputTileIterator source_iterator, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand
|
||||
MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses
|
||||
MatrixCoord(Shape::kM, Shape::kN),
|
||||
MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space
|
||||
MatrixCoord())
|
||||
{
|
||||
// Reduce peer accumulator fragments into one fragment
|
||||
AccumulatorFragment accum_fragment;
|
||||
BaseStreamK::reduce(accum_fragment, peer_idx_begin, peer_idx_end, reduce_fragment_idx, element_workspace);
|
||||
|
||||
// Store fragment to shared memory
|
||||
this->warp_tile_iterator_.store(accum_fragment);
|
||||
|
||||
__syncthreads();
|
||||
|
||||
Base::reduce(reduce_fragment_idx, output_op, broadcast_ptr, destination_iterator, source_iterator, tensor_iterator, problem_size, threadblock_offset);
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -863,6 +863,98 @@ private:
|
||||
frag_Broadcast_ptr[i % ThreadMap::Iterations::kColumn]);
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
/// Stream-K reduce helper
|
||||
CUTLASS_DEVICE
|
||||
void reduce(
|
||||
int reduce_fragment_idx, ///< Reduce fragment index
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
ElementVector const * broadcast_ptr, ///< Broadcast vector
|
||||
OutputTileIterator destination_iterator, ///< Tile iterator for destination
|
||||
OutputTileIterator source_iterator1, ///< Tile iterator for first source accumulator matrix
|
||||
OutputTileIterator source_iterator2, ///< Tile iterator for second source accumulator matrix
|
||||
TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand
|
||||
MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses
|
||||
MatrixCoord(Shape::kM, Shape::kN),
|
||||
MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space
|
||||
MatrixCoord())
|
||||
{
|
||||
|
||||
BroadcastFragment broadcast_fragment;
|
||||
load_broadcast_fragment_(broadcast_fragment, broadcast_ptr, problem_size, threadblock_offset);
|
||||
|
||||
// Initialize/load source-fragment data
|
||||
typename OutputTileIterator::Fragment source_fragment1;
|
||||
source_fragment1.clear();
|
||||
typename OutputTileIterator::Fragment source_fragment2;
|
||||
source_fragment2.clear();
|
||||
|
||||
if (output_op.is_source_needed())
|
||||
{
|
||||
source_iterator1 += reduce_fragment_idx;
|
||||
source_iterator1.load(source_fragment1);
|
||||
|
||||
source_iterator2 += reduce_fragment_idx;
|
||||
source_iterator2.load(source_fragment2);
|
||||
}
|
||||
|
||||
// Load fragment from shared memory
|
||||
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
|
||||
shared_load_iterator_.load(aligned_accum_fragment[0]);
|
||||
|
||||
// Add fragments shared by other k partitions
|
||||
if (kPartitionsK > 1)
|
||||
{
|
||||
plus <typename SharedLoadIterator::Fragment> add_fragments;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for ( int i = 1; i < kPartitionsK; ++i) {
|
||||
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
|
||||
shared_load_iterator_.load(aligned_accum_fragment[i]);
|
||||
aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]);
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Apply output operation
|
||||
//
|
||||
|
||||
typename OutputTileIterator::Fragment frag_Z;
|
||||
typename TensorTileIterator::Fragment frag_T;
|
||||
|
||||
if (!output_op.is_source_needed()) {
|
||||
apply_output_operator_source_not_needed_(
|
||||
frag_Z,
|
||||
frag_T,
|
||||
output_op,
|
||||
aligned_accum_fragment[0],
|
||||
broadcast_fragment);
|
||||
} else {
|
||||
apply_output_operator_(
|
||||
frag_Z,
|
||||
frag_T,
|
||||
output_op,
|
||||
aligned_accum_fragment[0],
|
||||
source_fragment1,
|
||||
source_fragment2,
|
||||
broadcast_fragment);
|
||||
}
|
||||
|
||||
//
|
||||
// Conditionally store fragments
|
||||
//
|
||||
|
||||
if (OutputOp::kStoreZ) {
|
||||
destination_iterator.store(frag_Z);
|
||||
++destination_iterator;
|
||||
}
|
||||
|
||||
if (OutputOp::kStoreT) {
|
||||
tensor_iterator.store(frag_T);
|
||||
++tensor_iterator;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -1529,6 +1621,92 @@ private:
|
||||
frag_Broadcast_ptr[i % ThreadMap::Iterations::kColumn]);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
public:
|
||||
/// Stream-K reduce helper
|
||||
CUTLASS_DEVICE
|
||||
void reduce(
|
||||
int reduce_fragment_idx, ///< Reduce fragment index
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
ElementVector const * broadcast_ptr, ///< Broadcast vector
|
||||
OutputTileIterator destination_iterator, ///< Tile iterator for destination
|
||||
OutputTileIterator source_iterator, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand
|
||||
MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses
|
||||
MatrixCoord(Shape::kM, Shape::kN),
|
||||
MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space
|
||||
MatrixCoord())
|
||||
{
|
||||
|
||||
BroadcastFragment broadcast_fragment;
|
||||
load_broadcast_fragment_(broadcast_fragment, broadcast_ptr, problem_size, threadblock_offset);
|
||||
|
||||
// Initialize/load source-fragment data
|
||||
typename OutputTileIterator::Fragment source_fragment;
|
||||
source_fragment.clear();
|
||||
|
||||
if (output_op.is_source_needed())
|
||||
{
|
||||
source_iterator += reduce_fragment_idx;
|
||||
source_iterator.load(source_fragment);
|
||||
}
|
||||
|
||||
// Load fragment from shared memory
|
||||
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
|
||||
shared_load_iterator_.load(aligned_accum_fragment[0]);
|
||||
|
||||
// Add fragments shared by other k partitions
|
||||
if (kPartitionsK > 1)
|
||||
{
|
||||
plus <typename SharedLoadIterator::Fragment> add_fragments;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for ( int i = 1; i < kPartitionsK; ++i) {
|
||||
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
|
||||
shared_load_iterator_.load(aligned_accum_fragment[i]);
|
||||
aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]);
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Apply output operation
|
||||
//
|
||||
|
||||
typename OutputTileIterator::Fragment frag_Z;
|
||||
typename TensorTileIterator::Fragment frag_T;
|
||||
|
||||
if (!output_op.is_source_needed()) {
|
||||
apply_output_operator_source_not_needed_(
|
||||
frag_Z,
|
||||
frag_T,
|
||||
output_op,
|
||||
aligned_accum_fragment[0],
|
||||
broadcast_fragment);
|
||||
} else {
|
||||
apply_output_operator_(
|
||||
frag_Z,
|
||||
frag_T,
|
||||
output_op,
|
||||
aligned_accum_fragment[0],
|
||||
source_fragment,
|
||||
broadcast_fragment);
|
||||
}
|
||||
|
||||
//
|
||||
// Conditionally store fragments
|
||||
//
|
||||
|
||||
if (OutputOp::kStoreZ) {
|
||||
destination_iterator.store(frag_Z);
|
||||
++destination_iterator;
|
||||
}
|
||||
|
||||
if (OutputOp::kStoreT) {
|
||||
tensor_iterator.store(frag_T);
|
||||
++tensor_iterator;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -0,0 +1,386 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2023 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 Template for a Stream-K GEMM kernel that can broadcast bias vector in the
|
||||
epilogue.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/arch/arch.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination_bias_elementwise.h"
|
||||
#include "cutlass/device_kernel.h"
|
||||
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
|
||||
#include "cutlass/gemm/kernel/gemm_universal.h"
|
||||
|
||||
#include "cutlass/gemm/kernel/default_gemm_universal.h"
|
||||
#include "cutlass/gemm/kernel/default_gemm_streamk_with_broadcast.h"
|
||||
#include "cutlass/gemm/device/default_gemm_configuration.h"
|
||||
#include "cutlass/gemm/device/gemm_universal_base.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace device {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/*!
|
||||
The universal GEMM with a broadcast epilogue.
|
||||
Supports
|
||||
*/
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA_,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA_,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB_,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB_,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC_,
|
||||
/// Layout type for C and D matrix operands
|
||||
typename LayoutC_,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator_ = ElementC_,
|
||||
/// Operator class tag
|
||||
typename OperatorClass_ = arch::OpClassSimt,
|
||||
/// Tag indicating architecture to tune for. This is the minimum SM that
|
||||
/// supports the intended feature. The device kernel can be built
|
||||
/// targeting any SM larger than this number.
|
||||
typename ArchTag_ = arch::Sm70,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::WarpShape,
|
||||
/// Instruction-level tile size (concept: GemmShape)
|
||||
typename InstructionShape_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::InstructionShape,
|
||||
/// Epilogue output operator - must satisfy concept of 'EpilogueWithBroadcastOp'
|
||||
typename EpilogueOutputOp_ = cutlass::epilogue::thread::LinearCombinationBiasElementwise<
|
||||
ElementC_, ElementAccumulator_, ElementAccumulator_,
|
||||
ElementC_, ElementC_, 128 / cutlass::sizeof_bits<ElementC_>::value>,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
ElementC_, ElementAccumulator_>::kStages,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int AlignmentA =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
ElementC_, ElementAccumulator_>::kAlignmentA,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int AlignmentB =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
ElementC_, ElementAccumulator_>::kAlignmentB,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator_ = typename DefaultGemmConfiguration<
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::Operator,
|
||||
/// Complex elementwise transformation on A operand
|
||||
ComplexTransform TransformA = ComplexTransform::kNone,
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB = ComplexTransform::kNone
|
||||
>
|
||||
class GemmUniversalStreamkWithBroadcast :
|
||||
public GemmUniversalBase<
|
||||
typename kernel::DefaultGemmStreamkWithBroadcast<
|
||||
ElementA_,
|
||||
LayoutA_,
|
||||
TransformA,
|
||||
AlignmentA,
|
||||
ElementB_,
|
||||
LayoutB_,
|
||||
TransformB,
|
||||
AlignmentB,
|
||||
ElementC_,
|
||||
LayoutC_,
|
||||
ElementAccumulator_,
|
||||
OperatorClass_,
|
||||
ArchTag_,
|
||||
ThreadblockShape_,
|
||||
WarpShape_,
|
||||
InstructionShape_,
|
||||
EpilogueOutputOp_,
|
||||
ThreadblockSwizzle_,
|
||||
Stages,
|
||||
Operator_
|
||||
>::GemmKernel
|
||||
> {
|
||||
|
||||
public:
|
||||
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using OperatorClass = OperatorClass_;
|
||||
using ArchTag = ArchTag_;
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = InstructionShape_;
|
||||
using EpilogueOutputOp = EpilogueOutputOp_;
|
||||
using ThreadblockSwizzle = ThreadblockSwizzle_;
|
||||
using Operator = Operator_;
|
||||
static int const kStages = Stages;
|
||||
static int const kAlignmentA = AlignmentA;
|
||||
static int const kAlignmentB = AlignmentB;
|
||||
static int const kAlignmentC = EpilogueOutputOp::kCount;
|
||||
static ComplexTransform const kTransformA = TransformA;
|
||||
static ComplexTransform const kTransformB = TransformB;
|
||||
|
||||
using Base = GemmUniversalBase<
|
||||
typename kernel::DefaultGemmStreamkWithBroadcast<
|
||||
ElementA_,
|
||||
LayoutA_,
|
||||
TransformA,
|
||||
AlignmentA,
|
||||
ElementB_,
|
||||
LayoutB_,
|
||||
TransformB,
|
||||
AlignmentB,
|
||||
ElementC_,
|
||||
LayoutC_,
|
||||
ElementAccumulator_,
|
||||
OperatorClass_,
|
||||
ArchTag_,
|
||||
ThreadblockShape_,
|
||||
WarpShape_,
|
||||
InstructionShape_,
|
||||
EpilogueOutputOp_,
|
||||
ThreadblockSwizzle_,
|
||||
Stages,
|
||||
Operator_
|
||||
>::GemmKernel
|
||||
>;
|
||||
|
||||
using Arguments = typename Base::Arguments;
|
||||
using GemmKernel = typename Base::GemmKernel;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for column-major output exchanges problem size and operand.
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA_,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA_,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB_,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB_,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC_,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator_,
|
||||
/// Operator class tag
|
||||
typename OperatorClass_,
|
||||
/// Tag indicating architecture to tune for. This is the minimum SM that
|
||||
/// supports the intended feature. The device kernel can be built
|
||||
/// targeting any SM larger than this number.
|
||||
typename ArchTag_,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape_,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape_,
|
||||
/// Instruction-level tile size (concept: GemmShape)
|
||||
typename InstructionShape_,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp_,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle_,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int AlignmentA,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int AlignmentB,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator_,
|
||||
/// Complex elementwise transformation on A operand
|
||||
ComplexTransform TransformA,
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB>
|
||||
class GemmUniversalStreamkWithBroadcast<ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_,
|
||||
layout::ColumnMajor, // partially specialized on LayoutC
|
||||
ElementAccumulator_, OperatorClass_, ArchTag_, ThreadblockShape_,
|
||||
WarpShape_, InstructionShape_, EpilogueOutputOp_,
|
||||
ThreadblockSwizzle_, Stages, AlignmentA, AlignmentB,
|
||||
Operator_, TransformA, TransformB> {
|
||||
public:
|
||||
|
||||
using ElementA = ElementA_;
|
||||
using LayoutA = LayoutA_;
|
||||
using TensorRefA = TensorRef<ElementA const, LayoutA>;
|
||||
using ElementB = ElementB_;
|
||||
using LayoutB = LayoutB_;
|
||||
using TensorRefB = TensorRef<ElementB const, LayoutB>;
|
||||
using ElementC = ElementC_;
|
||||
using LayoutC = layout::ColumnMajor;
|
||||
using TensorRefC = TensorRef<ElementC const, LayoutC>;
|
||||
using TensorRefD = TensorRef<ElementC, LayoutC>;
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using OperatorClass = OperatorClass_;
|
||||
using ArchTag = ArchTag_;
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
using InstructionShape = InstructionShape_;
|
||||
using EpilogueOutputOp = EpilogueOutputOp_;
|
||||
using ThreadblockSwizzle = ThreadblockSwizzle_;
|
||||
using Operator = Operator_;
|
||||
static int const kStages = Stages;
|
||||
static int const kAlignmentA = AlignmentA;
|
||||
static int const kAlignmentB = AlignmentB;
|
||||
static ComplexTransform const kTransformA = TransformA;
|
||||
static ComplexTransform const kTransformB = TransformB;
|
||||
|
||||
using UnderlyingOperator = typename GemmUniversalStreamkWithBroadcast<
|
||||
ElementB,
|
||||
typename layout::LayoutTranspose<LayoutB>::type,
|
||||
ElementA,
|
||||
typename layout::LayoutTranspose<LayoutA>::type,
|
||||
ElementC,
|
||||
layout::RowMajor,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
kAlignmentB,
|
||||
kAlignmentA,
|
||||
Operator,
|
||||
kTransformB,
|
||||
kTransformA
|
||||
>::Base;
|
||||
|
||||
using GemmKernel = typename UnderlyingOperator::GemmKernel;
|
||||
static int const kAlignmentC = EpilogueOutputOp::kCount;
|
||||
|
||||
/// Argument structure
|
||||
using Arguments = typename UnderlyingOperator::Arguments;
|
||||
|
||||
private:
|
||||
|
||||
UnderlyingOperator underlying_operator_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs the GEMM.
|
||||
GemmUniversalStreamkWithBroadcast() { }
|
||||
|
||||
/// Helper to construct a transposed equivalent for the underying GEMM operator
|
||||
static Arguments to_underlying_arguments(Arguments const &args) {
|
||||
return args.transposed_problem();
|
||||
}
|
||||
|
||||
/// Determines whether the GEMM can execute the given problem.
|
||||
static Status can_implement(Arguments const &args) {
|
||||
|
||||
return UnderlyingOperator::can_implement(to_underlying_arguments(args));
|
||||
}
|
||||
|
||||
/// Gets the workspace size
|
||||
static size_t get_workspace_size(Arguments const &args) {
|
||||
|
||||
return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args));
|
||||
}
|
||||
|
||||
/// Computes the grid shape
|
||||
static dim3 get_grid_shape(Arguments const &args) {
|
||||
return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args));
|
||||
}
|
||||
|
||||
/// Computes the maximum number of active blocks per multiprocessor
|
||||
static int maximum_active_blocks(int smem_capacity = -1) {
|
||||
return UnderlyingOperator::maximum_active_blocks(smem_capacity);
|
||||
}
|
||||
|
||||
/// Initializes GEMM state from arguments.
|
||||
Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
|
||||
return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream);
|
||||
}
|
||||
|
||||
/// Lightweight update given a subset of arguments
|
||||
Status update(Arguments const &args, void *workspace = nullptr) {
|
||||
|
||||
return underlying_operator_.update(to_underlying_arguments(args), workspace);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status run(cudaStream_t stream = nullptr) {
|
||||
|
||||
return underlying_operator_.run(stream);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status operator()(cudaStream_t stream = nullptr) {
|
||||
return run(stream);
|
||||
}
|
||||
|
||||
/// Runs the kernel using initialized state.
|
||||
Status operator()(
|
||||
Arguments const &args,
|
||||
void *workspace = nullptr,
|
||||
cudaStream_t stream = nullptr) {
|
||||
|
||||
Status status = initialize(args, workspace, stream);
|
||||
|
||||
if (status == Status::kSuccess) {
|
||||
status = run(stream);
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace device
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,146 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2023 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 a Stream-K GEMM that can broadcast a bias vector in the epilogue.
|
||||
Similar structure to DefaultGemmWithBroadcast, but uses its own epilogue
|
||||
(DefaultStreamkEpilogueWithBroadcastTensorOp) and its own GEMM kernel
|
||||
(GemmStreamkWithFusedEpilogue).
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
#include "cutlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h"
|
||||
#include "cutlass/gemm/kernel/default_gemm_universal.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/default_epilogue_with_broadcast.h"
|
||||
#include "cutlass/epilogue/threadblock/epilogue_with_broadcast.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace kernel {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA_,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA_,
|
||||
/// Complex elementwise transformation on A operand
|
||||
ComplexTransform TransformA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB_,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB_,
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC_,
|
||||
/// Layout type for C and D matrix operands
|
||||
typename LayoutC_,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Operator class tag
|
||||
typename OperatorClass,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
/// Epilogue output operator - must satisfy concept of 'EpilogueWithBroadcastOp'
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator,
|
||||
///
|
||||
typename Enable = void
|
||||
>
|
||||
struct DefaultGemmStreamkWithBroadcast {
|
||||
|
||||
using GemmBase = typename DefaultGemmUniversal<
|
||||
ElementA_, LayoutA_, TransformA, kAlignmentA,
|
||||
ElementB_, LayoutB_, TransformB, kAlignmentB,
|
||||
ElementC_, LayoutC_, ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
Operator
|
||||
>::GemmKernel;
|
||||
|
||||
// Replace epilogue
|
||||
using Epilogue = typename cutlass::epilogue::threadblock::DefaultStreamkEpilogueWithBroadcastTensorOp<
|
||||
typename GemmBase::Epilogue::Shape,
|
||||
typename GemmBase::Epilogue::WarpMmaOperator,
|
||||
GemmBase::Epilogue::kPartitionsK,
|
||||
ElementC_,
|
||||
typename EpilogueOutputOp::ElementT,
|
||||
typename EpilogueOutputOp::ElementVector,
|
||||
EpilogueOutputOp,
|
||||
GemmBase::Epilogue::kElementsPerAccess
|
||||
>::Epilogue;
|
||||
|
||||
// Compose the GEMM kernel
|
||||
using GemmKernel = GemmStreamkWithFusedEpilogue<
|
||||
typename GemmBase::Mma,
|
||||
Epilogue,
|
||||
ThreadblockSwizzle
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace kernel
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
2405
include/cutlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h
Normal file
2405
include/cutlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user