leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 29 Packages Projects Releases Wiki Activity

CUTLASS 2.10 (#615)

Browse Source 
Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2022-09-03 18:48:46 -04:00
committed by GitHub
co-authored by Aniket Shivam
parent ca23ff7924
commit b72cbf957d
289 changed files with 43708 additions and 2513 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/35_gemm_softmax/epilogue_with_visitor.h
-407
View File
@@ -1,407 +0,0 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2022 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 Generic epilogue for implementing certain kinds of fused epilogue behavior.
*/
#pragma once
/////////////////////////////////////////////////////////////////////////////////////////////////
#include "cutlass/cutlass.h"
#include "cutlass/fast_math.h"
#include "cutlass/matrix_coord.h"
#include "cutlass/semaphore.h"
#include "cutlass/epilogue/threadblock/epilogue_base.h"
////////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace epilogue {
namespace threadblock {
////////////////////////////////////////////////////////////////////////////////////////////////////
class EpilogueFusedVisitorConcept {
public:
static int const kIterations = 1;
static int const kElementsPerAccess = 4;
using ElementOutput = float;
using ElementAccumulator = float;
using AccumulatorFragment = Array<ElementAccumulator, kElementsPerAccess>;
/// Arguments structure
struct Arguments { };
/// Params structure
struct Params {
Params() { }
Params(Arguments const &args) { }
};
/// Shared storage
struct SharedStorage { };
public:
CUTLASS_DEVICE
EpilogueFusedVisitorConcept(
Params const &params, ///< Parameters routed to the epilogue
SharedStorage &shared_storage, ///< Shared storage needed by the functors here
MatrixCoord const &problem_size, ///< Problem size of the output
int thread_idx, ///< Thread index within the threadblock
int warp_idx, ///< Warp index within the threadblock
int lane_idx, ///< Lane index within the warp
MatrixCoord const &threadblock_offset = MatrixCoord(0, 0)) { ///< Coordinate
}
/// Helper to indicate split-K behavior
CUTLASS_DEVICE
void set_k_partition(
int split_k_index, ///< Index of this threadblock within split-K partitioned scheme
int split_k_slices) { ///< Total number of split-K slices
}
/// Called to set the batch index
CUTLASS_DEVICE
void set_batch_index(int batch_idx) {
}
/// Called at the start of the epilogue just before iterating over accumulator slices
CUTLASS_DEVICE
void begin_epilogue() {
}
/// Called at the start of one step before starting accumulator exchange
CUTLASS_DEVICE
void begin_step(int step_idx) {
}
/// Called at the start of a row
CUTLASS_DEVICE
void begin_row(int row_idx) {
}
/// Called after accumulators have been exchanged for each accumulator vector
CUTLASS_DEVICE
void visit(
int row_idx,
int column_idx,
int frag_idx,
AccumulatorFragment const &accum) {
}
/// Called at the start of a row
CUTLASS_DEVICE
void end_row(int row_idx) {
}
/// Called after all accumulator elements have been visited
CUTLASS_DEVICE
void end_step(int step_idx) {
}
/// Called after all steps have been completed
CUTLASS_DEVICE
void end_epilogue() {
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Epilogue operator
template <
typename Visitor_, ///< Functor containing fused operations (satisfies EpilogueFusedVisitorConcept)
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 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 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
(true || !IsEpilogueFunctorHeavy<Visitor_>::value)
>
class EpilogueWithVisitor :
public EpilogueBase<
Shape_,
typename WarpMmaOperator_::Shape,
PartitionsK,
AccumulatorFragmentIterator_,
WarpTileIterator_,
Padding_,
FragmentsPerPartition> {
public:
using Visitor = Visitor_;
using Base = EpilogueBase<
Shape_,
typename WarpMmaOperator_::Shape,
PartitionsK,
AccumulatorFragmentIterator_,
WarpTileIterator_,
Padding_,
FragmentsPerPartition>;
using Shape = Shape_;
using WarpMmaOperator = WarpMmaOperator_;
static int const kPartitionsK = PartitionsK;
using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
using WarpTileIterator = WarpTileIterator_;
using SharedLoadIterator = SharedLoadIterator_;
using Padding = Padding_;
using Layout = layout::RowMajor;
using LongIndex = typename Layout::LongIndex;
/// The complete warp-level accumulator tile
using AccumulatorTile = typename Base::AccumulatorTile;
/// Accumulator element
using ElementAccumulator = typename WarpTileIterator::Element;
/// Output access size
static int const kElementsPerAccess = Visitor::kElementsPerAccess;
/// Tensor reference to sync tensor
using SyncTensorRef = typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
/// Array type used by output functor
using AccumulatorAccessType = Array<
typename WarpTileIterator::Element, kElementsPerAccess>;
/// Number of warps
using WarpCount = typename Base::WarpCount;
static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1 ? Base::kFragmentsPerIteration : kPartitionsK;
static int constexpr kSmemPointerOffset = Base::SharedStorage::StorageShape::kCount / kSmemTiles;
using SharedStorage = typename Base::SharedStorage;
private:
/// Loads fragment from shared memory aligned with output tensor
SharedLoadIterator shared_load_iterator_;
public:
/// Constructor
CUTLASS_DEVICE
EpilogueWithVisitor(
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),
shared_load_iterator_(shared_storage.reference(), thread_idx)
{
}
/// Streams the result to global memory
CUTLASS_DEVICE
void operator()(
Visitor & visitor,
AccumulatorTile const &accumulators) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
visitor.begin_epilogue();
//
// Iterator over warp-level accumulator fragment
//
AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
//
// Iterate over accumulator tile
//
#pragma unroll(IterationsUnroll ? Visitor::kIterations : 1)
for (int iter_idx = 0; iter_idx < Visitor::kIterations; ++iter_idx) {
//
// Load the source
//
visitor.begin_step(iter_idx);
//
// Convert and store fragment
//
__syncthreads();
acc2smem_source_needed<cutlass::make_index_sequence<Visitor::kIterations>>::push(
iter_idx, accum_fragment_iterator, this->warp_tile_iterator_);
__syncthreads();
//
// Load fragments from shared memory
//
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
shared_load_iterator_.load(aligned_accum_fragment[0]);
// If the number of k-slices is > 1 - perform a reduction amongst the k-slices
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]);
}
shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset);
}
//
// Iterate over output fragments
//
AccumulatorAccessType const *accum_frag_ptr =
reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment[0]);
int const kAccumulatorFragmentCount = AccumulatorTile::kElements / (Visitor::kIterations * AccumulatorAccessType::kElements);
CUTLASS_PRAGMA_UNROLL
for (int idx = 0; idx < kAccumulatorFragmentCount; ++idx) {
int row_idx = idx / SharedLoadIterator::ThreadMap::Iterations::kColumn;
int col_idx = idx % SharedLoadIterator::ThreadMap::Iterations::kColumn;
// Start a new row of the output fragment
if (!col_idx) {
visitor.begin_row(row_idx);
}
visitor.visit(
row_idx,
col_idx,
idx,
accum_frag_ptr[idx]
);
// End the row of the output fragment
if (col_idx + 1 == SharedLoadIterator::ThreadMap::Iterations::kColumn) {
visitor.end_row(row_idx);
}
}
//
// Conclude the step
//
visitor.end_step(iter_idx);
}
visitor.end_epilogue();
}
private:
template<class Seq>
struct acc2smem_source_needed;
template <size_t... Seq>
struct acc2smem_source_needed<cutlass::index_sequence<Seq...>> {
template<int Advance>
CUTLASS_DEVICE
static void helper(AccumulatorFragmentIterator accum_fragment_iterator,
WarpTileIterator &warp_tile_iterator) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Advance; i++) {
++accum_fragment_iterator;
}
typename AccumulatorFragmentIterator::Fragment accum_fragment;
accum_fragment_iterator.load(accum_fragment);
warp_tile_iterator.store(accum_fragment);
}
CUTLASS_DEVICE
static void push(size_t pos,
AccumulatorFragmentIterator const &iterator_begin,
WarpTileIterator &warp_tile_iterator) {
int dummy[] = {(pos == Seq) && (helper<Seq>(iterator_begin, warp_tile_iterator), 0)...};
}
};
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Helper to create an EpilogueWithVisitor from an existing epilogue
template <typename Visitor_, typename Existing_, bool IterationsUnroll = true>
struct EpilogueWithVisitorFromExistingEpilogue {
using Epilogue = EpilogueWithVisitor<
Visitor_,
typename Existing_::Shape,
typename Existing_::WarpMmaOperator,
Existing_::kPartitionsK,
typename Existing_::AccumulatorFragmentIterator,
typename Existing_::WarpTileIterator,
typename Existing_::SharedLoadIterator,
typename Existing_::Padding,
Existing_::kFragmentsPerIteration,
IterationsUnroll
>;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
examples/35_gemm_softmax/gemm_softmax.cu
+209 -191
View File
@@ -47,14 +47,17 @@
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/reference/host/gemm_complex.h"
#include "cutlass/util/reference/device/gemm_complex.h"
#include "cutlass/util/reference/host/tensor_reduce.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "cutlass/util/reference/host/tensor_norm.h"
#include "cutlass/util/reference/host/tensor_copy.h"
#include "cutlass/util/reference/device/tensor_fill.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/error_metrics.h"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/epilogue/thread/linear_combination.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -85,18 +88,18 @@ struct Options {
float alpha;
float beta;
bool verification_enabled;
double tolerance;
float tolerance;
Options():
help(false),
problem_size({16, 24, 64}),
batch_count(1), // As a temporary limitation to the test bench, batch count must be 1. The kernels support arbitrary batching.
batch_count(16),
iterations(20),
seed(2022),
alpha(1),
beta(),
beta(0),
verification_enabled(true),
tolerance(0.01)
tolerance(1e-5f)
{ }
bool valid() {
@@ -116,6 +119,8 @@ struct Options {
cmd.get_cmd_line_argument("n", problem_size.n());
cmd.get_cmd_line_argument("k", problem_size.k());
cmd.get_cmd_line_argument("batch_count", batch_count);
cmd.get_cmd_line_argument("alpha", alpha);
cmd.get_cmd_line_argument("beta", beta);
@@ -135,6 +140,7 @@ struct Options {
<< " --m=<int> GEMM M dimension\n"
<< " --n=<int> GEMM N dimension\n"
<< " --k=<int> GEMM K dimension\n"
<< " --batch_count=<int> Batch number\n"
<< " --alpha=<f32> Epilogue scalar alpha\n"
<< " --beta=<f32> Epilogue scalar beta\n\n"
<< " --seed=<int> Random number seed (1*)\n\n"
@@ -198,13 +204,22 @@ struct Testbed {
using ElementA = cutlass::half_t;
using ElementB = cutlass::half_t;
using ElementC = cutlass::half_t;
using ElementD = cutlass::half_t;
using ElementCompute = float;
using ElementSoftmax = cutlass::half_t;
using ElementD = ElementC;
using ElementSoftmax = ElementC;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using ThreadblockShape = cutlass::gemm::GemmShape<128, 128, 32>;
using WarpShape = cutlass::gemm::GemmShape<64, 64, 32>;
using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>;
using OperatorClass = cutlass::arch::OpClassTensorOp;
using ArchTag = cutlass::arch::Sm80;
static int const kStages = 3;
/// Linear scaling operator
using EpilogueFunctorOp = cutlass::epilogue::thread::LinearCombination<
ElementC,
@@ -218,12 +233,21 @@ struct Testbed {
ElementB, LayoutB,
ElementC,
ElementCompute,
EpilogueFunctorOp
OperatorClass,
ArchTag,
ThreadblockShape,
WarpShape,
InstructionShape,
EpilogueFunctorOp,
kStages
>;
using ElementNorm = typename GemmSoftmax::ElementNorm;
using ElementSum = typename GemmSoftmax::ElementSum;
using LayoutC = typename GemmSoftmax::LayoutC;
using LayoutN = typename GemmSoftmax::LayoutN;
using LayoutS = typename GemmSoftmax::LayoutS;
using MatrixCoord = typename LayoutC::TensorCoord;
//
// Data members
@@ -231,20 +255,42 @@ struct Testbed {
Options const &options;
cutlass::HostTensor<ElementA, LayoutA> tensor_A;
cutlass::HostTensor<ElementB, LayoutB> tensor_B;
cutlass::HostTensor<ElementC, LayoutC> tensor_C;
cutlass::HostTensor<ElementD, LayoutC> tensor_D;
cutlass::HostTensor<ElementNorm, LayoutC> tensor_N;
cutlass::HostTensor<ElementSum, LayoutC> tensor_S;
cutlass::HostTensor<ElementSoftmax, LayoutC> tensor_Softmax;
cutlass::HostTensor<ElementD, LayoutC> reference_D;
cutlass::HostTensor<ElementNorm, LayoutC> reference_N;
cutlass::HostTensor<ElementSoftmax, LayoutC> reference_Softmax;
cutlass::DeviceAllocation<ElementA> block_A;
cutlass::DeviceAllocation<ElementB> block_B;
cutlass::DeviceAllocation<ElementC> block_C;
cutlass::DeviceAllocation<ElementD> block_D;
cutlass::DeviceAllocation<ElementD> block_Ref;
cutlass::DeviceAllocation<ElementSoftmax> block_Softmax;
cutlass::DeviceAllocation<ElementNorm> block_Norm;
cutlass::DeviceAllocation<ElementSum> block_Sum;
int block_num = (options.problem_size.n() + GemmSoftmax::ThreadblockShape::kN - 1) / GemmSoftmax::ThreadblockShape::kN;
cutlass::gemm::GemmCoord problem = options.problem_size;
int64_t lda = LayoutA::packed({problem.m(), problem.k()}).stride(0);
int64_t ldb = LayoutB::packed({problem.k(), problem.n()}).stride(0);
int64_t ldc = LayoutC::packed({problem.m(), problem.n()}).stride(0);
// fixed rowmajor for norm and sum
int64_t ldn = problem.m();
int64_t lds = ldn;
int64_t total_elements_A_per_batch = problem.m() * problem.k();
int64_t total_elements_B_per_batch = problem.k() * problem.n();
int64_t total_elements_C_per_batch = problem.m() * problem.n();
int64_t total_elements_D_per_batch = problem.m() * problem.n();
int64_t total_elements_partial_norm_per_batch = block_num * problem.m();
int64_t total_elements_A = total_elements_A_per_batch * options.batch_count;
int64_t total_elements_B = total_elements_B_per_batch * options.batch_count;
int64_t total_elements_C = total_elements_C_per_batch * options.batch_count;
int64_t total_elements_D = total_elements_D_per_batch * options.batch_count;
int64_t total_elements_partial_norm = total_elements_partial_norm_per_batch * options.batch_count;
//
// Methods
//
@@ -254,20 +300,7 @@ struct Testbed {
):
options(options_)
{
tensor_A.reset({options.problem_size.m(), options.problem_size.k()});
tensor_B.reset({options.problem_size.k(), options.problem_size.n()});
tensor_C.reset({options.problem_size.m(), options.problem_size.n()});
tensor_D.reset({options.problem_size.m(), options.problem_size.n()});
tensor_N.reset({block_num, options.problem_size.m()});
tensor_S.reset({block_num, options.problem_size.m()});
tensor_Softmax.reset({options.problem_size.m(), options.problem_size.n()});
reference_D.reset({options.problem_size.m(), options.problem_size.n()}, false);
reference_N.reset({options.problem_size.m(), 1}, false);
reference_Softmax.reset({options.problem_size.m(), options.problem_size.n()}, false);
}
/// Run
@@ -300,11 +333,6 @@ struct Testbed {
return disposition;
}
//
// Compute the reference
//
compute_reference();
//
// Verify
//
@@ -334,43 +362,38 @@ struct Testbed {
/// Random initialization
void initialize() {
cutlass::reference::host::TensorFillRandomUniform(
tensor_A.host_view(),
options.seed,
ElementD(5),
ElementD(-5),
0
);
block_A.reset(total_elements_A);
block_B.reset(total_elements_B);
block_C.reset(total_elements_C);
block_D.reset(total_elements_D);
block_Softmax.reset(total_elements_D);
block_Ref.reset(total_elements_D_per_batch);
block_Norm.reset(total_elements_partial_norm);
block_Sum.reset(total_elements_partial_norm);
cutlass::reference::host::TensorFillRandomUniform(
tensor_B.host_view(),
options.seed + 19,
ElementD(5),
ElementD(-5),
0
);
cutlass::reference::device::BlockFillRandomUniform(
block_A.get(), total_elements_A, options.seed, ElementA(5), ElementA(-5), 0);
cutlass::reference::host::TensorFill(
reference_D.host_view(),
ElementD()
);
cutlass::reference::device::BlockFillRandomUniform(
block_B.get(), total_elements_B, options.seed + 1, ElementB(5), ElementB(-5), 0);
cutlass::reference::device::BlockFillRandomUniform(
block_C.get(), total_elements_C, options.seed + 2, ElementC(5), ElementC(-5), 0);
cutlass::reference::device::BlockFillRandomUniform(
block_D.get(), total_elements_D, options.seed + 3, ElementD(5), ElementD(-5), 0);
cutlass::reference::device::BlockFillRandomUniform(
block_Ref.get(), total_elements_D_per_batch, options.seed + 3, ElementD(5), ElementD(-5), 0);
cutlass::reference::device::BlockFillRandomUniform(
block_Softmax.get(), total_elements_D, options.seed + 3, ElementSoftmax(5), ElementSoftmax(-5), 0);
cutlass::reference::host::TensorFill(
reference_N.host_view(),
ElementNorm()
);
cutlass::reference::host::TensorFill(
reference_Softmax.host_view(),
ElementSoftmax()
);
tensor_A.sync_device();
tensor_B.sync_device();
tensor_D.sync_device();
tensor_N.sync_device();
tensor_S.sync_device();
tensor_Softmax.sync_device();
}
cutlass::Status execute_device_kernel() {
@@ -384,17 +407,24 @@ struct Testbed {
GemmSoftmax::Arguments args(
options.problem_size,
options.batch_count,
tensor_A.device_ref(),
tensor_B.device_ref(),
tensor_C.device_ref(),
tensor_D.device_ref(),
{block_A.get(), lda},
{block_B.get(), ldb},
{block_C.get(), ldc},
{block_D.get(), ldc},
{
ElementCompute(options.alpha),
ElementCompute(options.beta)
},
tensor_N.device_ref(),
tensor_S.device_ref(),
tensor_Softmax.device_ref()
{block_Norm.get(), ldn},
{block_Sum.get(), lds},
{block_Softmax.get(), ldc},
total_elements_A_per_batch,
total_elements_B_per_batch,
total_elements_C_per_batch,
total_elements_D_per_batch,
total_elements_partial_norm_per_batch,
total_elements_partial_norm_per_batch,
total_elements_D_per_batch
);
//
@@ -415,68 +445,21 @@ struct Testbed {
return status;
}
/// Reference calculation
void compute_reference() {
template<typename Element>
bool verify_tensor(std::vector<Element> vector_Input, \
std::vector<Element> vector_Input_Ref) {
// Compute GEMM
cutlass::reference::host::GemmComplex(
options.problem_size,
options.alpha,
tensor_A.host_ref(),
cutlass::ComplexTransform::kNone,
tensor_B.host_ref(),
cutlass::ComplexTransform::kNone,
options.beta,
tensor_C.host_ref(),
reference_D.host_ref(),
double()
);
// Compute the norm
for (int m = 0; m < options.problem_size.m(); ++m) {
reference_N.at({m, 0}) = reference_D.at({m, 0});
for (int n = 1; n < options.problem_size.n(); ++n) {
reference_N.at({m, 0}) = std::max(reference_N.at({m, 0}), ElementNorm(reference_D.at({m, n})));
}
}
// Compute softmax
for (int m = 0; m < options.problem_size.m(); ++m) {
float sum = float();
for (int n = 0; n < options.problem_size.n(); ++n) {
sum += std::exp( float(reference_D.at({m, n})) - float(reference_N.at({m, 0})) );
}
float inv_sum = float(1.0f / sum);
for (int n = 0; n < options.problem_size.n(); ++n) {
reference_Softmax.at({m, n}) = ElementSoftmax(
std::exp( float(reference_D.at({m, n})) - float(reference_N.at({m, 0})) ) * inv_sum
);
}
}
}
/// Emits all tensor values
void emit_results() {
std::cout << "D = \n" << tensor_D.host_view() << "\n\n";
std::cout << "N = \n" << tensor_N.host_view() << "\n\n";
std::cout << "Softmax = \n" << tensor_Softmax.host_view() << "\n\n";
std::cout << "Reference N = \n" << reference_N.host_view() << "\n\n";
std::cout << "Reference D = \n" << reference_D.host_view() << "\n\n";
std::cout << "Reference Softmax = \n" << reference_Softmax.host_view() << "\n\n";
}
bool verify_tensor_N(cutlass::HostTensor<ElementNorm, LayoutC> tensor_N, \
cutlass::HostTensor<ElementNorm, LayoutC> reference_N) {
for (int m = 0; m < options.problem_size.m(); ++m) {
float diff = (float)(tensor_N.at({0, m}) - reference_N.at({m, 0}));
if (fabs(diff) > options.tolerance) {
int64_t size = (vector_Input.size() < vector_Input_Ref.size()) ? vector_Input.size() : vector_Input_Ref.size();
float abs_tol = options.tolerance;
float rel_tol = options.tolerance;
for (int64_t i = 0; i < size; ++i) {
float diff = (float)(vector_Input.at(i) - vector_Input_Ref.at(i));
float abs_diff = fabs(diff);
float abs_ref = fabs((float)vector_Input_Ref.at(i));
float relative_diff = abs_ref > abs_tol ? abs_diff / abs_ref : 0;
if ( (isnan(abs_diff) || isinf(abs_diff)) || (abs_diff > rel_tol && relative_diff > rel_tol)) {
printf("diff = %f, {%f, %f}.\n", abs_diff, (float)(vector_Input.at(i)), (float)(vector_Input_Ref.at(i)));
return false;
}
@@ -488,80 +471,112 @@ struct Testbed {
/// Verifies the reference matches
bool verify() {
tensor_D.sync_host();
tensor_N.sync_host();
tensor_Softmax.sync_host();
LayoutA layout_A(lda);
LayoutB layout_B(ldb);
LayoutC layout_C(ldc);
LayoutN Layout_N(ldn);
LayoutS Layout_S(lds);
double const kThreshold = options.tolerance;
MatrixCoord extent_A{problem.m(), problem.k()};
MatrixCoord extent_B{problem.k(), problem.n()};
MatrixCoord extent_C{problem.m(), problem.n()};
// Verification checks - set any of these to 'true' to override the verification checks.
bool verified_D = false;
bool verified_N = false;
bool verified_Softmax = false;
for (int batch_idx = 0; batch_idx < options.batch_count; batch_idx++) {
// Verify softmax output
if (!verified_D) {
cutlass::TensorView<ElementA, LayoutA> view_A(block_A.get() + total_elements_A_per_batch * batch_idx, layout_A, extent_A);
cutlass::TensorView<ElementB, LayoutB> view_B(block_B.get() + total_elements_B_per_batch * batch_idx, layout_B, extent_B);
cutlass::TensorView<ElementC, LayoutC> view_C(block_C.get() + total_elements_C_per_batch * batch_idx, layout_C, extent_C);
cutlass::TensorView<ElementC, LayoutC> view_Ref_device(block_Ref.get(), layout_C, extent_C);
double norm_diff = cutlass::reference::host::TensorNormDiff(
tensor_D.host_view(),
reference_D.host_view());
cutlass::reference::device::GemmComplex<
ElementA, LayoutA,
ElementB, LayoutB,
ElementC, LayoutC,
ElementCompute, ElementCompute
>(
problem,
options.alpha,
view_A,
cutlass::ComplexTransform::kNone,
view_B,
cutlass::ComplexTransform::kNone,
options.beta,
view_C,
view_Ref_device,
ElementCompute(0)
);
double norm_reference = cutlass::reference::host::TensorNorm(
reference_D.host_view());
// Copy reference results to host memory for verification
std::vector<ElementD> matrix_D_Ref(layout_C.capacity(extent_C));
cutlass::device_memory::copy_to_host(matrix_D_Ref.data(), block_Ref.get(), matrix_D_Ref.size());
cutlass::TensorView<ElementD, LayoutC> view_Ref(matrix_D_Ref.data(), layout_C, extent_C);
double rel_error = norm_diff / norm_reference;
std::vector<ElementSoftmax> matrix_Softmax_Ref(layout_C.capacity(extent_C));
cutlass::TensorView<ElementSoftmax, LayoutC> view_Softmax_Ref(matrix_Softmax_Ref.data(), layout_C, extent_C);
if (rel_error > kThreshold) {
std::cerr << "\n\nTensor D Relative error: " << rel_error << std::endl;
// Copy computed results to host memory
std::vector<ElementD> matrix_D(layout_C.capacity(extent_C));
cutlass::device_memory::copy_to_host(matrix_D.data(), block_D.get() + total_elements_D_per_batch * batch_idx, matrix_D.size());
std::vector<ElementD> matrix_Softmax(layout_C.capacity(extent_C));
cutlass::device_memory::copy_to_host(matrix_Softmax.data(), block_Softmax.get() + total_elements_D_per_batch * batch_idx, matrix_Softmax.size());
// Compute the norm
for (int m = 0; m < options.problem_size.m(); ++m) {
reference_N.at({m, 0}) = view_Ref.ref().at({m, 0});
for (int n = 1; n < options.problem_size.n(); ++n) {
reference_N.at({m, 0}) = std::max(reference_N.at({m, 0}), ElementNorm(view_Ref.ref().at({m, n})));
}
}
else {
verified_D = true;
// Compute softmax
for (int m = 0; m < options.problem_size.m(); ++m) {
float sum = float();
for (int n = 0; n < options.problem_size.n(); ++n) {
sum += std::exp( float(view_Ref.ref().at({m, n})) - float(reference_N.at({m, 0})) );
}
float inv_sum = float(1.0f / sum);
for (int n = 0; n < options.problem_size.n(); ++n) {
view_Softmax_Ref.ref().at({m, n}) = ElementSoftmax(
std::exp( float(view_Ref.ref().at({m, n})) - float(reference_N.at({m, 0})) ) * inv_sum
);
}
}
}
if (!verified_N) {
verified_N = verify_tensor_N(tensor_N, reference_N);
}
// Verification checks - set any of these to 'true' to override the verification checks.
bool verified_D = false;
bool verified_Softmax = false;
if (!verified_Softmax) {
double norm_diff = cutlass::reference::host::TensorNormDiff(
tensor_Softmax.host_view(),
reference_Softmax.host_view());
double norm_reference = cutlass::reference::host::TensorNorm(
reference_Softmax.host_view());
double rel_error = norm_diff / norm_reference;
if (rel_error > kThreshold) {
std::cerr << "\n\nSoftmax Relative error: " << rel_error << std::endl;
}
else {
verified_Softmax = true;
}
}
if (!verified_D || !verified_N || !verified_Softmax) {
std::cerr << "Verification check failed for tensor Softmax" << std::endl;
emit_results();
// Summarize which checks failed
// Verify softmax output
if (!verified_D) {
std::cerr << "Verification of D tensor failed\n";
}
if (!verified_N) {
std::cerr << "Verification of N tensor failed\n";
verified_D = verify_tensor<ElementC>(matrix_D, matrix_D_Ref);
}
if (!verified_Softmax) {
std::cerr << "Verification of Softmax tensor failed\n";
verified_Softmax = verify_tensor<ElementSoftmax>(matrix_Softmax, matrix_Softmax_Ref);
}
if (!verified_D || !verified_Softmax) {
std::cerr << "Verification check failed for tensor Softmax at batch " << batch_idx << "\n";
// Summarize which checks failed
if (!verified_D) {
std::cerr << "Verification of D tensor failed\n";
}
if (!verified_Softmax) {
std::cerr << "Verification of Softmax tensor failed\n";
}
return false;
}
return false;
}
return true;
@@ -637,14 +652,17 @@ struct Testbed {
int64_t flops = int64_t(options.problem_size.m()) * options.problem_size.n() * options.problem_size.k() * 2;
int64_t bytes = (sizeof(ElementD) * 2 + sizeof(ElementSoftmax)) * options.problem_size.m() * options.problem_size.n();
double gflops_per_second = double(flops) * kIterations / double(elapsed_ms / 1000.0f) / double(1.0e9);
double gbytes_per_second = double(bytes) * kIterations / double(elapsed_ms / 1000.0f) / double(1 << 30);
double gflops_per_second = double(flops) * kIterations * options.batch_count / double(elapsed_ms / 1000.0f) / double(1.0e9);
double gbytes_per_second = double(bytes) * kIterations * options.batch_count / double(elapsed_ms / 1000.0f) / double(1 << 30);
double elapsed_ms_per_iter = double(elapsed_ms) / kIterations;
std::cout << " Problem: "
<< options.problem_size.m() << "-by-" << options.problem_size.n() << "-by-" << options.problem_size.k()
<< ", batch size: " << options.batch_count
<< std::endl;
std::cout << " Runtime: " << elapsed_ms << " ms\n" << std::endl;
std::cout << " Runtime: " << elapsed_ms_per_iter << " ms\n" << std::endl;
std::cout << " GFLOPs: " << gflops_per_second << " GFLOPs" << std::endl;
std::cout << "Memory bandwidth: " << gbytes_per_second << " GiB/s" << std::endl;
examples/35_gemm_softmax/gemm_with_epilogue_visitor.h
+46 -2
View File
@@ -29,7 +29,8 @@
*
**************************************************************************************************/
/*! \file
\brief GEMM kernel to support the 'epilogue visitor' model for fusion.
\brief GEMM kernel to support the epilogue visitor model
for customized softmax partial reduction epilogue fusion.
This source file will likely be moved to `include/cutlass/gemm/kernel/` in the future once
its usage has been stabilized. For now, it is included in this example to demonstrate
@@ -78,6 +79,7 @@ public:
using ElementC = typename EpilogueVisitor::ElementOutput;
using LayoutC = typename Epilogue::Layout;
using TensorRefC = TensorRef<ElementC, LayoutC>;
static ComplexTransform const kTransformA = Mma::kTransformA;
static ComplexTransform const kTransformB = Mma::kTransformB;
@@ -89,6 +91,9 @@ public:
using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
using ArchTag = typename Mma::ArchTag;
using ElementNorm = typename EpilogueVisitor::ElementNorm;
using ElementSum = typename EpilogueVisitor::ElementSum;
static int const kStages = Mma::kStages;
static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
@@ -121,6 +126,11 @@ public:
TensorRefA ref_A;
TensorRefB ref_B;
TensorRefC ref_C;
TensorRefC ref_D;
ElementNorm *ptr_Max;
ElementSum *ptr_Sum;
int64_t batch_stride_A;
int64_t batch_stride_B;
@@ -144,6 +154,10 @@ public:
int batch_count_,
TensorRefA ref_A_,
TensorRefB ref_B_,
TensorRefC ref_C_,
TensorRefC ref_D_,
ElementNorm *ptr_Max_,
ElementSum *ptr_Sum_,
int64_t batch_stride_A_,
int64_t batch_stride_B_,
typename EpilogueVisitor::Arguments epilogue_visitor_
@@ -153,6 +167,10 @@ public:
batch_count(batch_count_),
ref_A(ref_A_),
ref_B(ref_B_),
ref_C(ref_C_),
ref_D(ref_D_),
ptr_Max(ptr_Max_),
ptr_Sum(ptr_Sum_),
batch_stride_A(batch_stride_A_),
batch_stride_B(batch_stride_B_),
epilogue_visitor(epilogue_visitor_)
@@ -174,6 +192,8 @@ public:
typename Mma::IteratorA::Params params_A;
typename Mma::IteratorB::Params params_B;
typename EpilogueVisitor::OutputTileIterator::Params params_C;
typename EpilogueVisitor::OutputTileIterator::Params params_D;
GemmUniversalMode mode;
int batch_count;
@@ -181,6 +201,11 @@ public:
void * ptr_A;
void * ptr_B;
ElementC * ptr_C;
ElementC * ptr_D;
ElementNorm * ptr_Max;
ElementSum * ptr_Sum;
int64_t batch_stride_A;
int64_t batch_stride_B;
@@ -196,11 +221,17 @@ public:
swizzle_log_tile(0),
params_A(0),
params_B(0),
params_C(0),
params_D(0),
batch_count(0),
gemm_k_size(0),
mode(cutlass::gemm::GemmUniversalMode::kGemm),
ptr_A(nullptr),
ptr_B(nullptr),
ptr_C(nullptr),
ptr_D(nullptr),
ptr_Max(nullptr),
ptr_Sum(nullptr),
batch_stride_A(0),
batch_stride_B(0)
{ }
@@ -213,11 +244,17 @@ public:
swizzle_log_tile(0),
params_A(args.ref_A.layout()),
params_B(args.ref_B.layout()),
params_C(args.ref_C.layout()),
params_D(args.ref_D.layout()),
mode(args.mode),
batch_count(args.batch_count),
gemm_k_size(args.problem_size.k()),
ptr_A(args.ref_A.data()),
ptr_B(args.ref_B.data()),
ptr_C(args.ref_C.data()),
ptr_D(args.ref_D.data()),
ptr_Max(args.ptr_Max),
ptr_Sum(args.ptr_Sum),
batch_stride_A(args.batch_stride_A),
batch_stride_B(args.batch_stride_B),
epilogue_visitor(args.epilogue_visitor)
@@ -467,7 +504,14 @@ public:
thread_idx,
warp_idx,
lane_idx,
threadblock_offset);
params.params_C,
params.params_D,
params.ptr_C,
params.ptr_D,
params.ptr_Max,
params.ptr_Sum,
threadblock_offset,
blockIdx.y *params.problem_size.m() );
if (params.mode == GemmUniversalMode::kGemm) {
// Indicate which position in a serial reduction the output operator is currently updating
examples/35_gemm_softmax/gemm_with_softmax.h
+54 -593
View File
@@ -49,10 +49,12 @@
#include "cutlass/gemm/kernel/default_gemm.h"
#include "cutlass/gemm/kernel/default_gemm_complex.h"
#include "cutlass/gemm/device/default_gemm_configuration.h"
#include "cutlass/epilogue/threadblock/epilogue_visitor_with_softmax.h"
#include "cutlass/epilogue/threadblock/epilogue_with_visitor.h"
#include "cutlass/reduction/kernel/reduce_softmax_final.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
#include "epilogue_with_visitor.h"
#include "gemm_with_epilogue_visitor.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -209,6 +211,9 @@ private:
int idx_m = block_m + thread_m;
int idx_n = block_n + thread_n;
int batch_offset_norm = block_batch * params.args.batch_stride_N;
int batch_offset_sum = block_batch * params.args.batch_stride_S;
// Kill off thread if it is outside the row boundary
if (params.args.extent.row() <= idx_m) {
return;
@@ -251,8 +256,8 @@ private:
params.args.batch_stride_Soft * block_batch +
params.args.ref_Soft.layout()({idx_m, idx_n}));
ElementSum inv_sum = (params.args.ref_S.data())[block_m];
ElementNorm norm = (params.args.ref_N.data())[block_m];
ElementSum inv_sum = (params.args.ref_S.data())[block_m + batch_offset_sum];
ElementNorm norm = (params.args.ref_N.data())[block_m + batch_offset_norm];
//
// Loop
@@ -281,556 +286,6 @@ private:
}
};
template <
typename ElementNorm_,
typename ElementSum_,
typename ElementSoftmaxCompute_,
typename ThreadblockShape_
>
class ApplyFinalReduction {
public:
using ElementNorm = ElementNorm_;
using ElementSum = ElementSum_;
using ElementSoftmaxCompute = ElementSoftmaxCompute_;
using ThreadblockShape = ThreadblockShape_;
using Layout = cutlass::layout::RowMajor;
using TensorRefN = TensorRef<ElementNorm, Layout>;
using TensorRefSum = TensorRef<ElementSum, Layout>;
//
// Arguments
//
struct Arguments {
MatrixCoord extent; ///< Extent of D and Softmax matrices
int batch_count; ///< Batch count
TensorRefN ref_N; ///< Norm tensor (input / output)
TensorRefSum ref_Sum; ///< Sum tensor (input / output)
int64_t batch_stride_N; ///< Batch stride for N tensor
int64_t batch_stride_Sum; ///< Batch stride for softmax tensor
//
// Methods
//
Arguments():
batch_count(1),
batch_stride_N(0),
batch_stride_Sum(0)
{ }
Arguments(
MatrixCoord extent_, ///< Extent of D and Softmax matrices
int batch_count_, ///< Batch count
TensorRefN ref_N_, ///< Output parameter for N
TensorRefSum ref_Sum_ , ///< Sum
int64_t batch_stride_N_ = 0,
int64_t batch_stride_Sum_ = 0
):
extent(extent_),
batch_count(batch_count_),
ref_N(ref_N_),
ref_Sum(ref_Sum_),
batch_stride_N(batch_stride_N_),
batch_stride_Sum(batch_stride_Sum_)
{
}
};
struct SharedStorage {
};
//
// Params struct
//
struct Params {
Arguments args;
//
// Methods
//
Params() { }
Params(Arguments const &args_): args(args_) { }
};
private:
public:
CUTLASS_DEVICE
ApplyFinalReduction() { }
CUTLASS_DEVICE
void operator()(Params const &params, SharedStorage &shared_storage) {
apply(params, shared_storage);
}
private:
/// Partial reduction
CUTLASS_DEVICE
void apply(Params const &params, SharedStorage &shared_storage) {
int threadblock_num = (params.args.extent.column() + ThreadblockShape::kN - 1) / ThreadblockShape::kN;
int block_batch = blockIdx.z;
int block_n = blockIdx.x * blockDim.x;
int thread_n = threadIdx.x;
int idx_n = block_n + thread_n;
if (idx_n >= params.args.extent.row()) {
return;
}
using ConvertSumOutput = cutlass::NumericConverter<ElementSum, ElementSoftmaxCompute>;
using ConvertNormOutput = cutlass::NumericConverter<ElementNorm, ElementSoftmaxCompute>;
using ConvertSum = cutlass::NumericConverter<ElementSoftmaxCompute, ElementSum>;
using ConvertNorm = cutlass::NumericConverter<ElementSoftmaxCompute, ElementNorm>;
ConvertSum convert_sum;
ConvertNorm convert_norm;
ConvertSumOutput convert_sum_output;
ConvertNormOutput convert_norm_output;
ElementNorm *access_n = params.args.ref_N.data() + params.args.batch_stride_N * block_batch + idx_n;
ElementSum *access_s = params.args.ref_Sum.data() + params.args.batch_stride_Sum * block_batch + idx_n;
ElementNorm *access_n_bak = access_n;
ElementSum *access_s_bak = access_s;
uint32_t float_max_bits = 0xff7fffff;
float min_float = reinterpret_cast<float const &>(float_max_bits);
ElementSoftmaxCompute max_val = ElementSoftmaxCompute(min_float);
ElementSoftmaxCompute sum_val = ElementSoftmaxCompute(0);
ElementNorm fetch_n;
ElementSum fetch_s;
CUTLASS_PRAGMA_UNROLL
for (int idx_m = 0; idx_m < threadblock_num; idx_m++) {
arch::global_load<ElementNorm, sizeof(ElementNorm)>(fetch_n, access_n, true);
max_val = fast_max(max_val, convert_norm(fetch_n));
access_n += params.args.extent.row();
}
access_n = access_n_bak;
CUTLASS_PRAGMA_UNROLL
for (int idx_m = 0; idx_m < threadblock_num; idx_m++) {
arch::global_load<ElementNorm, sizeof(ElementNorm)>(fetch_n, access_n, true);
arch::global_load<ElementSum, sizeof(ElementSum)>(fetch_s, access_s, true);
sum_val += convert_sum(fetch_s) * fast_exp(convert_norm(fetch_n) - max_val);
access_n += params.args.extent.row();
access_s += params.args.extent.row();
}
ElementSoftmaxCompute inv_sum = cutlass::constants::one<ElementSoftmaxCompute>() / sum_val;
access_n = access_n_bak;
access_s = access_s_bak;
access_n[0] = convert_norm_output(max_val);
access_s[0] = convert_sum_output(inv_sum);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
typename ThreadblockShape_,
int ThreadCount,
typename OutputTileIterator_,
typename ElementAccumulator_,
typename ElementNorm_,
typename ElementSum_,
typename ElementSoftmaxCompute_,
typename ElementwiseFunctor_
>
class EpilogueVisitorBiasMax {
public:
using ThreadblockShape = ThreadblockShape_;
static int const kThreadCount = ThreadCount;
using OutputTileIterator = OutputTileIterator_;
using ElementwiseFunctor = ElementwiseFunctor_;
static int const kIterations = OutputTileIterator::kIterations;
static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
using ElementOutput = typename OutputTileIterator::Element;
using LayoutOutput = cutlass::layout::RowMajor;
using ElementAccumulator = ElementAccumulator_;
using ElementNorm = ElementNorm_;
using ElementSum = ElementSum_;
using ElementSoftmaxCompute = ElementSoftmaxCompute_;
using AccumulatorFragment = Array<ElementAccumulator, kElementsPerAccess>;
using SoftmaxFragment = Array<ElementSoftmaxCompute, kElementsPerAccess>;
using OutputVector = Array<ElementOutput, kElementsPerAccess>;
using TensorRefD = TensorRef<ElementOutput, LayoutOutput>;
/// Argument structure
struct Arguments {
typename ElementwiseFunctor::Params elementwise;
TensorRefD ref_C;
TensorRefD ref_D;
ElementNorm *ptr_Max;
ElementSum *ptr_Sum;
int64_t batch_stride_C;
int64_t batch_stride_D;
int64_t batch_stride_Max;
int64_t batch_stride_Sum;
//
// Methods
//
Arguments():
ptr_Max(nullptr),
ptr_Sum(nullptr),
batch_stride_C(0),
batch_stride_D(0),
batch_stride_Max(0),
batch_stride_Sum(0)
{
}
Arguments(
typename ElementwiseFunctor::Params elementwise_,
TensorRefD ref_C_,
TensorRefD ref_D_,
ElementNorm *ptr_Max_,
ElementSum *ptr_Sum_,
int64_t batch_stride_C_,
int64_t batch_stride_D_,
int64_t batch_stride_Max_,
int64_t batch_stride_Sum_
):
elementwise(elementwise_),
ref_C(ref_C_),
ref_D(ref_D_),
ptr_Max(ptr_Max_),
ptr_Sum(ptr_Sum_),
batch_stride_C(batch_stride_C_),
batch_stride_D(batch_stride_D_),
batch_stride_Max(batch_stride_Max_),
batch_stride_Sum(batch_stride_Sum_)
{
}
};
struct Params {
typename ElementwiseFunctor::Params elementwise;
typename OutputTileIterator::Params params_C;
typename OutputTileIterator::Params params_D;
typename OutputTileIterator::Element *ptr_C;
typename OutputTileIterator::Element *ptr_D;
ElementNorm *ptr_Max;
ElementSum *ptr_Sum;
int64_t batch_stride_C;
int64_t batch_stride_D;
int64_t batch_stride_Max;
int64_t batch_stride_Sum;
//
// Methods
//
CUTLASS_HOST_DEVICE
Params():
ptr_D(nullptr),
ptr_Max(nullptr),
ptr_Sum(nullptr)
{
}
CUTLASS_HOST_DEVICE
Params(Arguments const &args):
elementwise(args.elementwise),
params_C(args.ref_C.layout()),
params_D(args.ref_D.layout()),
ptr_C(args.ref_C.data()),
ptr_D(args.ref_D.data()),
ptr_Max(args.ptr_Max),
ptr_Sum(args.ptr_Sum),
batch_stride_C(args.batch_stride_C),
batch_stride_D(args.batch_stride_D),
batch_stride_Max(args.batch_stride_Max),
batch_stride_Sum(args.batch_stride_Sum)
{
}
};
/// Shared storage
struct SharedStorage {
};
private:
Params const & params_;
SharedStorage & shared_storage_;
MatrixCoord extent_;
ElementwiseFunctor elementwise_;
OutputTileIterator iterator_C_;
OutputTileIterator iterator_D_;
typename OutputTileIterator::Fragment fragment_C_;
typename OutputTileIterator::Fragment fragment_D_;
ElementAccumulator alpha_;
ElementAccumulator beta_;
ElementSoftmaxCompute accum_max_;
int threadblock_row_;
public:
CUTLASS_DEVICE
EpilogueVisitorBiasMax(
Params const &params, ///< Parameters routed to the epilogue
SharedStorage &shared_storage, ///< Shared storage needed by the functors here
MatrixCoord const &problem_size, ///< Problem size of the output
int thread_idx, ///< Thread index within the threadblock
int warp_idx, ///< Warp index within the threadblock
int lane_idx, ///< Lane index within the warp
MatrixCoord const &threadblock_offset = MatrixCoord(0, 0)
):
params_(params),
shared_storage_(shared_storage),
extent_(problem_size),
elementwise_(params.elementwise),
iterator_C_(params.params_C, params.ptr_C, problem_size, thread_idx, threadblock_offset),
iterator_D_(params.params_D, params.ptr_D, problem_size, thread_idx, threadblock_offset),
threadblock_row_(threadblock_offset.row())
{
alpha_ = (params.elementwise.alpha_ptr ? *params.elementwise.alpha_ptr : params.elementwise.alpha);
beta_ = (params.elementwise.beta_ptr ? *params.elementwise.beta_ptr : params.elementwise.beta);
if (beta_ == ElementAccumulator()) {
iterator_C_.clear_mask();
}
}
/// Helper to indicate split-K behavior
CUTLASS_DEVICE
void set_k_partition(
int split_k_index, ///< Index of this threadblock within split-K partitioned scheme
int split_k_slices) { ///< Total number of split-K slices
}
/// Called to set the batch index
CUTLASS_DEVICE
void set_batch_index(int batch_idx) {
iterator_C_.add_pointer_offset(batch_idx * params_.batch_stride_C);
iterator_D_.add_pointer_offset(batch_idx * params_.batch_stride_D);
}
/// Called at the start of the epilogue just before iterating over accumulator slices
CUTLASS_DEVICE
void begin_epilogue() {
}
/// Called at the start of one step before starting accumulator exchange
CUTLASS_DEVICE
void begin_step(int step_idx) {
fragment_D_.clear();
fragment_C_.clear();
if (elementwise_.kScale != cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling) {
iterator_C_.load(fragment_C_);
++iterator_C_;
}
}
/// Called at the start of a row
CUTLASS_DEVICE
void begin_row(int row_idx) {
}
/// Called after accumulators have been exchanged for each accumulator vector
CUTLASS_DEVICE
void visit(
int row_idx,
int column_idx,
int frag_idx,
AccumulatorFragment const &accum) {
using Mul = cutlass::multiplies<SoftmaxFragment>;
using Minus = cutlass::minus<SoftmaxFragment>;
using Exp = cutlass::fast_exp_op<SoftmaxFragment>;
Minus minus;
Exp exponential;
SoftmaxFragment result;
using ConvertSumOutput = cutlass::NumericConverter<ElementSoftmaxCompute, ElementSum>;
using ConvertNormOutput = cutlass::NumericConverter<ElementSoftmaxCompute, ElementNorm>;
ConvertSumOutput convert_sum_output;
ConvertNormOutput convert_norm_output;
NumericArrayConverter<ElementSoftmaxCompute, ElementOutput, kElementsPerAccess> source_converter;
OutputVector &source_vector = reinterpret_cast<OutputVector *>(&fragment_C_)[frag_idx];
if (elementwise_.kScale == cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling) {
result = source_converter(elementwise_(accum));
}else{
result = source_converter(elementwise_(accum, source_vector));
}
MatrixCoord thread_offset =
iterator_D_.thread_start() +
OutputTileIterator::ThreadMap::iteration_offset(frag_idx);
int thread_in_row = OutputTileIterator::ThreadMap::Detail::RowArrangement::Detail::kShapeWidth;
int half_thread_in_row = (thread_in_row >> 1);
bool column_guard = (thread_offset.column() < extent_.column());
// Compute the maximum within one row
if (!column_idx) {
// This is the first fragment in a new row
if (column_guard) {
accum_max_ = maximum_accumulator_(result);
}
}
else {
// This is an additional fragment in the same row
if (column_guard) {
accum_max_ = maximum_accumulator_(result, accum_max_);
}
}
CUTLASS_PRAGMA_UNROLL
for (int i = half_thread_in_row; i > 0; i >>= 1) {
ElementSoftmaxCompute tmp = __shfl_xor_sync(0xFFFFFFFF, accum_max_, i);
accum_max_ = fast_max(accum_max_, tmp);
}
SoftmaxFragment sum_frag = exponential(minus(result, accum_max_));
ElementSoftmaxCompute reduction_sum = sum_accumulator_(sum_frag);
CUTLASS_PRAGMA_UNROLL
for (int i = half_thread_in_row; i > 0; i >>= 1) {
ElementSoftmaxCompute tmp = __shfl_xor_sync(0xFFFFFFFF, reduction_sum, i);
reduction_sum += tmp;
}
bool is_write_thread = (thread_offset.row() < extent_.row() && (threadIdx.x % thread_in_row) == 0);
ElementNorm *curr_ptr_max = params_.ptr_Max + thread_offset.row() + blockIdx.y * extent_.row();
ElementSum *curr_ptr_sum = params_.ptr_Sum + thread_offset.row() + blockIdx.y * extent_.row();
arch::global_store<ElementNorm, sizeof(ElementNorm)>(
convert_norm_output(accum_max_),
(void *)curr_ptr_max,
is_write_thread);
arch::global_store<ElementSum, sizeof(ElementSum)>(
convert_sum_output(reduction_sum),
(void *)curr_ptr_sum,
is_write_thread);
clear_accum_max_();
// Convert to the output
NumericArrayConverter<ElementOutput, ElementSoftmaxCompute, kElementsPerAccess> output_converter;
OutputVector &output = reinterpret_cast<OutputVector *>(&fragment_D_)[frag_idx];
output = output_converter(result);
}
/// Called at the start of a row
CUTLASS_DEVICE
void end_row(int row_idx) {
}
/// Called after all accumulator elements have been visited
CUTLASS_DEVICE
void end_step(int step_idx) {
iterator_D_.store(fragment_D_);
++iterator_D_;
}
/// Called after all steps have been completed
CUTLASS_DEVICE
void end_epilogue() {
}
private:
CUTLASS_DEVICE
void clear_accum_max_() {
uint32_t float_max_bits = 0xff7fffff; // -FLT_MAX
float min_float = reinterpret_cast<float const &>(float_max_bits);
accum_max_ = ElementSoftmaxCompute(min_float);
}
CUTLASS_DEVICE
ElementSoftmaxCompute sum_accumulator_(SoftmaxFragment const &accum) {
ElementSoftmaxCompute sum_ = ElementSoftmaxCompute(0);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < SoftmaxFragment::kElements; ++i) {
sum_ += ElementSoftmaxCompute(accum[i]);
}
return sum_;
}
CUTLASS_DEVICE
ElementSoftmaxCompute maximum_accumulator_(SoftmaxFragment const &accum) {
ElementSoftmaxCompute max_ = accum[0];
CUTLASS_PRAGMA_UNROLL
for (int i = 1; i < SoftmaxFragment::kElements; ++i) {
max_ = fast_max(max_, ElementSoftmaxCompute(accum[i]));
}
return max_;
}
CUTLASS_DEVICE
ElementSoftmaxCompute maximum_accumulator_(SoftmaxFragment const &accum, ElementSoftmaxCompute max_) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < SoftmaxFragment::kElements; ++i) {
max_ = fast_max(max_, ElementSoftmaxCompute(accum[i]));
}
return max_;
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -846,10 +301,18 @@ template <
typename LayoutB_,
typename ElementC_,
typename ElementCompute_,
typename OperatorClass_,
typename ArchTag_,
typename ThreadblockShape_,
typename WarpShape_,
typename InstructionShape_,
typename EpilogueFunctorOp_,
int kStages_,
int AlignmentA_ = 128 / cutlass::sizeof_bits<ElementA_>::value,
int AlignmentB_ = 128 / cutlass::sizeof_bits<ElementB_>::value,
int AlignmentSoftmax_ = 128 / cutlass::sizeof_bits<ElementC_>::value,
typename ElementNorm_ = float,
typename ElementSum_ = float,
int Alignment = 128 / cutlass::sizeof_bits<ElementA_>::value,
typename ElementSoftmax_ = ElementC_
>
class GemmSoftmax {
@@ -872,8 +335,6 @@ public:
using LayoutA = LayoutA_;
using LayoutB = LayoutB_;
static int const kAlignment = Alignment;
using EpilogueFunctorOp = EpilogueFunctorOp_;
using ElementNorm = ElementNorm_;
@@ -890,13 +351,17 @@ public:
using TensorRefSum = TensorRef<ElementSum, LayoutS>;
using TensorRefSoft = TensorRef<ElementSoft, LayoutSoft>;
using ThreadblockShape = cutlass::gemm::GemmShape<128, 128, 32>;
using WarpShape = cutlass::gemm::GemmShape<64, 64, 32>;
using InstructionShape = cutlass::gemm::GemmShape<16, 8, 16>;
using ThreadblockShape = ThreadblockShape_;
using WarpShape = WarpShape_;
using InstructionShape = InstructionShape_;
using OperatorClass = cutlass::arch::OpClassTensorOp;
using ArchTag = cutlass::arch::Sm80;
static int const kStages = 3;
using OperatorClass = OperatorClass_;
using ArchTag = ArchTag_;
static int const kStages = kStages_;
static int const AlignmentA = AlignmentA_;
static int const AlignmentB = AlignmentB_;
static int const AlignmentSoftmax = AlignmentSoftmax_;
using ThreadblockSwizzle = cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle;
@@ -906,10 +371,10 @@ public:
using DefaultGemmKernel = typename cutlass::gemm::kernel::DefaultGemm<
ElementA,
LayoutA,
kAlignment,
AlignmentA,
ElementB,
LayoutB,
kAlignment,
AlignmentB,
ElementC,
LayoutC,
ElementCompute,
@@ -930,7 +395,7 @@ public:
///////////////////////////////////////////////////////////////////////////////////////////////
// Epilogue visitor
using EpilogueVisitor = kernel::EpilogueVisitorBiasMax<
using EpilogueVisitor = typename cutlass::epilogue::threadblock::EpilogueVisitorSoftmax<
ThreadblockShape,
DefaultGemmKernel::kThreadCount,
typename DefaultGemmKernel::Epilogue::OutputTileIterator,
@@ -961,13 +426,13 @@ public:
ElementSum,
ElementSoft,
ElementSoftmaxCompute,
kAlignment,
AlignmentSoftmax,
MatrixShape<
1, 1024
>
>;
using ApplyFinalReductionKernel = kernel::ApplyFinalReduction<
using ApplyFinalReductionKernel = cutlass::reduction::kernel::ApplySoftmaxFinalReduction<
ElementNorm,
ElementSum,
ElementSoftmaxCompute,
@@ -983,6 +448,7 @@ public:
typename SoftmaxApplyKernel::Arguments softmax;
typename ApplyFinalReductionKernel::Arguments reduction;
cutlass::gemm::GemmCoord extend;
//
// Methods
//
@@ -1013,14 +479,14 @@ public:
batch_count_,
ref_A_,
ref_B_,
ref_C_,
ref_D_,
ref_N_.data(),
ref_S_.data(),
batch_stride_A_,
batch_stride_B_,
typename EpilogueVisitor::Arguments(
linear_scaling,
ref_C_,
ref_D_,
ref_N_.data(),
ref_S_.data(),
batch_stride_C_,
batch_stride_D_,
batch_stride_Max_,
@@ -1028,10 +494,9 @@ public:
)
),
reduction(
MatrixCoord(problem_size.m(), problem_size.n()),
batch_count_,
ref_N_,
ref_S_,
problem_size,
ref_N_.data(),
ref_S_.data(),
batch_stride_Max_,
batch_stride_Sum_
),
@@ -1127,28 +592,24 @@ public:
// Launch the ApplyFinalReductionKernel
//
int threadblock_num_in_column = (params_.extend.column() + ThreadblockShape::kN - 1) / ThreadblockShape::kN;
int thread_per_block = 128;
int block_per_row = (params_.extend.row() + thread_per_block - 1) / thread_per_block;
if (block_per_row < 4) {
thread_per_block = 32;
block_per_row = (params_.extend.row() + thread_per_block - 1) / thread_per_block;
}
if (threadblock_num_in_column > 1) {
int thread_per_block = 128;
int block_per_row = (params_.extend.row() + thread_per_block - 1) / thread_per_block;
if (block_per_row < 4) {
thread_per_block = 32;
block_per_row = (params_.extend.row() + thread_per_block - 1) / thread_per_block;
}
dim3 final_reduction_grid(block_per_row, 1, params_.softmax.args.batch_count);
dim3 final_reduction_block(thread_per_block);
dim3 final_reduction_grid(block_per_row);
dim3 final_reduction_block(thread_per_block);
Kernel<ApplyFinalReductionKernel><<<
final_reduction_grid, final_reduction_block, sizeof(typename ApplyFinalReductionKernel::SharedStorage), stream
>>>(params_.reduction);
Kernel<ApplyFinalReductionKernel><<<
final_reduction_grid, final_reduction_block, sizeof(typename ApplyFinalReductionKernel::SharedStorage), stream
>>>(params_.reduction);
result = cudaGetLastError();
result = cudaGetLastError();
if (result != cudaSuccess) {
return cutlass::Status::kErrorInternal;
}
if (result != cudaSuccess) {
return cutlass::Status::kErrorInternal;
}
//