ANIKET SHIVAM
co-authored by
Aniket Shivam
Aniket Shivam
parent
ca23ff7924
commit
b72cbf957d
@@ -0,0 +1,36 @@
|
||||
|
||||
# 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.
|
||||
|
||||
|
||||
|
||||
cutlass_example_add_executable(
|
||||
41_multi_head_attention
|
||||
fused_multihead_attention.cu
|
||||
)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,626 @@
|
||||
/***************************************************************************************************
|
||||
* 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 holdvr 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 the FusedMultiHeadAttention Class
|
||||
|
||||
The class contains the following:
|
||||
1) GEMM0 with epilogue fusion,
|
||||
2) GEMM1 with mainloop fusion, and
|
||||
3) A lightweight full softmax reduction kernel.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <limits>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/arch/memory.h"
|
||||
#include "cutlass/arch/memory_sm75.h"
|
||||
#include "cutlass/epilogue/threadblock/epilogue_visitor_with_softmax.h"
|
||||
#include "cutlass/epilogue/thread/scale_type.h"
|
||||
#include "cutlass/gemm/kernel/default_gemm_grouped_softmax_mainloop_fusion.h"
|
||||
#include "cutlass/reduction/kernel/reduce_softmax_final.h"
|
||||
#include "gemm_grouped_with_softmax_visitor.h"
|
||||
|
||||
namespace cutlass {
|
||||
|
||||
template <
|
||||
typename ElementQ_,
|
||||
typename LayoutQ_,
|
||||
typename ElementK_,
|
||||
typename LayoutK_,
|
||||
typename ElementP_,
|
||||
typename LayoutP_,
|
||||
typename ElementCompute_,
|
||||
typename OperatorClass_,
|
||||
typename ArchTag_,
|
||||
typename ThreadblockShape0_,
|
||||
typename ThreadblockShape1_,
|
||||
typename WarpShape0_,
|
||||
typename WarpShape1_,
|
||||
typename InstructionShape_,
|
||||
int kStages0_,
|
||||
int kStages1_,
|
||||
bool UseMasking_ = false,
|
||||
cutlass::gemm::kernel::GroupScheduleMode GroupScheduleMode0_ = cutlass::gemm::kernel::GroupScheduleMode::kHostPrecompute,
|
||||
cutlass::gemm::kernel::GroupScheduleMode GroupScheduleMode1_ = cutlass::gemm::kernel::GroupScheduleMode::kHostPrecompute,
|
||||
int Alignment = 128 / cutlass::sizeof_bits<ElementQ_>::value,
|
||||
typename ElementSoftmax_ = ElementP_
|
||||
>
|
||||
class FusedMultiHeadAttention {
|
||||
public:
|
||||
|
||||
using ElementQ = ElementQ_;
|
||||
using ElementK = ElementK_;
|
||||
using ElementP = ElementP_;
|
||||
using ElementV = ElementK;
|
||||
using ElementOutput = ElementP;
|
||||
using ElementAccumulator = ElementCompute_;
|
||||
|
||||
using LayoutQ = LayoutQ_;
|
||||
using LayoutK = LayoutK_;
|
||||
using LayoutP = LayoutP_;
|
||||
using LayoutV = LayoutK;
|
||||
using LayoutO = LayoutP;
|
||||
|
||||
using ElementNorm = cutlass::half_t;
|
||||
using ElementSum = cutlass::half_t;
|
||||
using ElementSoftmaxCompute = float;
|
||||
using LayoutNorm = cutlass::layout::RowMajor;
|
||||
|
||||
using ThreadblockSwizzle = cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle;
|
||||
|
||||
using OperatorClass = OperatorClass_;
|
||||
using ArchTag = ArchTag_;
|
||||
|
||||
using ThreadblockShape0 = ThreadblockShape0_;
|
||||
using WarpShape0 = WarpShape0_;
|
||||
|
||||
using ThreadblockShape1 = ThreadblockShape1_;
|
||||
using WarpShape1 = WarpShape1_;
|
||||
|
||||
static int const Stages0 = kStages0_;
|
||||
static int const Stages1 = kStages1_;
|
||||
|
||||
using InstructionShape = InstructionShape_;
|
||||
|
||||
using EpilogueOutputOp0 = cutlass::epilogue::thread::LinearCombination<
|
||||
ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator, ElementAccumulator, cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling>;
|
||||
|
||||
using EpilogueOutputOp1 = cutlass::epilogue::thread::LinearCombination<
|
||||
ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator, ElementAccumulator, cutlass::epilogue::thread::ScaleType::Nothing>;
|
||||
|
||||
using Operator = typename cutlass::gemm::device::DefaultGemmConfiguration<
|
||||
OperatorClass, ArchTag, ElementQ, ElementK, ElementP,
|
||||
ElementAccumulator>::Operator;
|
||||
static bool const kInternalTranspose = cutlass::platform::is_same<LayoutP, cutlass::layout::ColumnMajor>::value;
|
||||
|
||||
static bool const kUseMasking = UseMasking_;
|
||||
|
||||
static cutlass::gemm::kernel::GroupScheduleMode const kGroupScheduleMode0 = GroupScheduleMode0_;
|
||||
static cutlass::gemm::kernel::GroupScheduleMode const kGroupScheduleMode1 = GroupScheduleMode1_;
|
||||
|
||||
using MapArguments = cutlass::gemm::kernel::detail::MapArguments<
|
||||
ElementQ,
|
||||
LayoutQ,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
8,
|
||||
ElementK,
|
||||
LayoutK,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
8,
|
||||
LayoutP,
|
||||
kInternalTranspose
|
||||
>;
|
||||
|
||||
using DefaultGemmKernel = typename cutlass::gemm::kernel::DefaultGemm<
|
||||
typename MapArguments::ElementA,
|
||||
typename MapArguments::LayoutA,
|
||||
MapArguments::kAlignmentA,
|
||||
typename MapArguments::ElementB,
|
||||
typename MapArguments::LayoutB,
|
||||
MapArguments::kAlignmentB,
|
||||
ElementP,
|
||||
typename MapArguments::LayoutC,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape0,
|
||||
WarpShape0,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp0,
|
||||
ThreadblockSwizzle,
|
||||
Stages0,
|
||||
true,
|
||||
Operator,
|
||||
cutlass::gemm::SharedMemoryClearOption::kNone
|
||||
>::GemmKernel;
|
||||
|
||||
using EpilogueVisitor = typename cutlass::epilogue::threadblock::EpilogueVisitorSoftmax<
|
||||
ThreadblockShape0,
|
||||
DefaultGemmKernel::kThreadCount,
|
||||
typename DefaultGemmKernel::Epilogue::OutputTileIterator,
|
||||
typename EpilogueOutputOp0::ElementCompute,
|
||||
ElementNorm,
|
||||
ElementSum,
|
||||
ElementSoftmaxCompute,
|
||||
EpilogueOutputOp0,
|
||||
kUseMasking
|
||||
>;
|
||||
|
||||
using Epilogue = typename cutlass::epilogue::threadblock::EpilogueWithVisitorFromExistingEpilogue<
|
||||
EpilogueVisitor,
|
||||
typename DefaultGemmKernel::Epilogue
|
||||
>::Epilogue;
|
||||
|
||||
using GemmKernel0 = cutlass::gemm::kernel::GemmGroupedWithEpilogueVistor<
|
||||
typename DefaultGemmKernel::Mma,
|
||||
Epilogue,
|
||||
ThreadblockSwizzle,
|
||||
kGroupScheduleMode0,
|
||||
kInternalTranspose,
|
||||
kUseMasking
|
||||
>;
|
||||
|
||||
using GemmGrouped0 = cutlass::gemm::device::GemmGrouped<GemmKernel0>;
|
||||
|
||||
using ApplyFinalReductionDevice = cutlass::reduction::kernel::ApplySoftmaxFinalReduction<
|
||||
ElementNorm,
|
||||
ElementSum,
|
||||
typename GemmGrouped0::GemmKernel::EpilogueVisitor::ElementSoftmaxCompute,
|
||||
typename GemmGrouped0::GemmKernel::EpilogueVisitor::ThreadblockShape,
|
||||
true
|
||||
>;
|
||||
|
||||
using GemmKernel1 = typename cutlass::gemm::kernel::DefaultGemmGroupedSoftmaxMainloopFusion<
|
||||
ElementP,
|
||||
LayoutP,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
128 / cutlass::sizeof_bits<ElementQ>::value,
|
||||
ElementV,
|
||||
LayoutV,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
128 / cutlass::sizeof_bits<ElementK>::value,
|
||||
ElementNorm,
|
||||
LayoutNorm,
|
||||
ElementOutput,
|
||||
LayoutO,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape1,
|
||||
WarpShape1,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp1,
|
||||
ThreadblockSwizzle,
|
||||
Stages1,
|
||||
kGroupScheduleMode1
|
||||
>::GemmKernel;
|
||||
|
||||
using GemmGrouped1 = cutlass::gemm::device::GemmGrouped<GemmKernel1>;
|
||||
|
||||
public:
|
||||
|
||||
/// Arguments class
|
||||
struct Arguments {
|
||||
cutlass::gemm::GemmCoord *problem_sizes0;
|
||||
cutlass::gemm::GemmCoord *problem_sizes0_real;
|
||||
cutlass::gemm::GemmCoord *problem_sizes1;
|
||||
int problem_count;
|
||||
int threadblock_count;
|
||||
|
||||
ElementQ ** ptr_Q;
|
||||
ElementK ** ptr_K;
|
||||
ElementP ** ptr_P;
|
||||
ElementP ** ptr_V;
|
||||
ElementP ** ptr_O;
|
||||
|
||||
ElementNorm **ptr_Max;
|
||||
ElementSum **ptr_Sum;
|
||||
|
||||
ElementP *block_P;
|
||||
ElementNorm *block_Norm;
|
||||
ElementSum *block_Sum;
|
||||
int64_t *offset_P;
|
||||
int64_t *offset_Norm_Device;
|
||||
int64_t *offset_Sum_Device;
|
||||
|
||||
typename LayoutQ::Stride::LongIndex *ldq;
|
||||
typename LayoutK::Stride::LongIndex *ldk;
|
||||
typename LayoutP::Stride::LongIndex *ldp;
|
||||
typename LayoutP::Stride::LongIndex *ldv;
|
||||
typename LayoutP::Stride::LongIndex *ldo;
|
||||
|
||||
cutlass::gemm::GemmCoord *problem_sizes0_host;
|
||||
cutlass::gemm::GemmCoord *problem_sizes1_host;
|
||||
|
||||
ElementAccumulator alpha0;
|
||||
ElementAccumulator alpha1;
|
||||
ElementAccumulator beta;
|
||||
|
||||
int head_number;
|
||||
int batch_size;
|
||||
int seq_length;
|
||||
|
||||
typename ApplyFinalReductionDevice::Arguments reduction;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
Arguments():
|
||||
problem_count(0),
|
||||
threadblock_count(0),
|
||||
ptr_Q(nullptr),
|
||||
ptr_K(nullptr),
|
||||
ptr_P(nullptr),
|
||||
ptr_V(nullptr),
|
||||
ptr_O(nullptr),
|
||||
ptr_Max(nullptr),
|
||||
ptr_Sum(nullptr),
|
||||
block_P(nullptr),
|
||||
block_Norm(nullptr),
|
||||
block_Sum(nullptr),
|
||||
offset_P(nullptr),
|
||||
offset_Norm_Device(nullptr),
|
||||
offset_Sum_Device(nullptr),
|
||||
ldq(nullptr),
|
||||
ldk(nullptr),
|
||||
ldp(nullptr),
|
||||
ldv(nullptr),
|
||||
ldo(nullptr),
|
||||
head_number(0),
|
||||
batch_size(0),
|
||||
seq_length(0)
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
Arguments(
|
||||
cutlass::gemm::GemmCoord *problem_sizes0,
|
||||
cutlass::gemm::GemmCoord *problem_sizes1,
|
||||
int problem_count,
|
||||
int threadblock_count,
|
||||
ElementQ ** ptr_Q,
|
||||
ElementK ** ptr_K,
|
||||
ElementP ** ptr_P,
|
||||
ElementP ** ptr_V,
|
||||
ElementP ** ptr_O,
|
||||
ElementNorm **ptr_Max,
|
||||
ElementSum **ptr_Sum,
|
||||
ElementP *block_P,
|
||||
ElementNorm *block_Norm,
|
||||
ElementSum *block_Sum,
|
||||
int64_t *offset_P,
|
||||
int64_t *offset_Norm_Device,
|
||||
int64_t *offset_Sum_Device,
|
||||
typename LayoutQ::Stride::LongIndex *ldq,
|
||||
typename LayoutK::Stride::LongIndex *ldk,
|
||||
typename LayoutP::Stride::LongIndex *ldp,
|
||||
typename LayoutP::Stride::LongIndex *ldv,
|
||||
typename LayoutP::Stride::LongIndex *ldo,
|
||||
ElementAccumulator alpha0,
|
||||
ElementAccumulator alpha1,
|
||||
ElementAccumulator beta,
|
||||
int head_number,
|
||||
int batch_size,
|
||||
int seq_length,
|
||||
cutlass::gemm::GemmCoord *problem_sizes0_host = nullptr,
|
||||
cutlass::gemm::GemmCoord *problem_sizes1_host = nullptr,
|
||||
cutlass::gemm::GemmCoord *problem_sizes0_real = nullptr
|
||||
):
|
||||
problem_sizes0(problem_sizes0),
|
||||
problem_sizes1(problem_sizes1),
|
||||
problem_count(problem_count),
|
||||
threadblock_count(threadblock_count),
|
||||
ptr_Q(ptr_Q),
|
||||
ptr_K(ptr_K),
|
||||
ptr_P(ptr_P),
|
||||
ptr_V(ptr_V),
|
||||
ptr_O(ptr_O),
|
||||
ptr_Max(ptr_Max),
|
||||
ptr_Sum(ptr_Sum),
|
||||
block_P(block_P),
|
||||
block_Norm(block_Norm),
|
||||
block_Sum(block_Sum),
|
||||
offset_P(offset_P),
|
||||
offset_Norm_Device(offset_Norm_Device),
|
||||
offset_Sum_Device(offset_Sum_Device),
|
||||
ldq(ldq),
|
||||
ldk(ldk),
|
||||
ldp(ldp),
|
||||
ldv(ldv),
|
||||
ldo(ldo),
|
||||
alpha0(alpha0),
|
||||
alpha1(alpha1),
|
||||
beta(beta),
|
||||
head_number(head_number),
|
||||
batch_size(batch_size),
|
||||
seq_length(seq_length),
|
||||
problem_sizes0_host(problem_sizes0_host),
|
||||
problem_sizes1_host(problem_sizes1_host),
|
||||
problem_sizes0_real(problem_sizes0_real),
|
||||
reduction(
|
||||
problem_sizes0,
|
||||
block_Norm,
|
||||
block_Sum,
|
||||
offset_Norm_Device,
|
||||
offset_Sum_Device
|
||||
)
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
|
||||
};
|
||||
|
||||
struct Params {
|
||||
cutlass::gemm::GemmCoord *problem_sizes0;
|
||||
cutlass::gemm::GemmCoord *problem_sizes0_real;
|
||||
cutlass::gemm::GemmCoord *problem_sizes1;
|
||||
int problem_count;
|
||||
int threadblock_count;
|
||||
|
||||
ElementQ ** ptr_Q;
|
||||
ElementK ** ptr_K;
|
||||
ElementP ** ptr_P;
|
||||
ElementP ** ptr_V;
|
||||
ElementP ** ptr_O;
|
||||
|
||||
ElementNorm **ptr_Max;
|
||||
ElementSum **ptr_Sum;
|
||||
|
||||
ElementP *block_P;
|
||||
ElementNorm *block_Norm;
|
||||
ElementSum *block_Sum;
|
||||
int64_t *offset_P;
|
||||
int64_t *offset_Norm_Device;
|
||||
int64_t *offset_Sum_Device;
|
||||
|
||||
typename LayoutQ::Stride::LongIndex *ldq;
|
||||
typename LayoutK::Stride::LongIndex *ldk;
|
||||
typename LayoutP::Stride::LongIndex *ldp;
|
||||
typename LayoutP::Stride::LongIndex *ldv;
|
||||
typename LayoutP::Stride::LongIndex *ldo;
|
||||
|
||||
cutlass::gemm::GemmCoord *problem_sizes0_host;
|
||||
cutlass::gemm::GemmCoord *problem_sizes1_host;
|
||||
|
||||
ElementAccumulator alpha0;
|
||||
ElementAccumulator alpha1;
|
||||
ElementAccumulator beta;
|
||||
|
||||
int head_number;
|
||||
int batch_size;
|
||||
int seq_length;
|
||||
|
||||
typename ApplyFinalReductionDevice::Params reduction;
|
||||
|
||||
Params():
|
||||
problem_count(0),
|
||||
threadblock_count(0),
|
||||
ptr_Q(nullptr),
|
||||
ptr_K(nullptr),
|
||||
ptr_P(nullptr),
|
||||
ptr_V(nullptr),
|
||||
ptr_O(nullptr),
|
||||
ptr_Max(nullptr),
|
||||
ptr_Sum(nullptr),
|
||||
block_P(nullptr),
|
||||
block_Norm(nullptr),
|
||||
block_Sum(nullptr),
|
||||
offset_P(nullptr),
|
||||
offset_Norm_Device(nullptr),
|
||||
offset_Sum_Device(nullptr),
|
||||
ldq(nullptr),
|
||||
ldk(nullptr),
|
||||
ldp(nullptr),
|
||||
ldv(nullptr),
|
||||
ldo(nullptr),
|
||||
problem_sizes0(nullptr),
|
||||
problem_sizes1(nullptr),
|
||||
problem_sizes0_real(nullptr),
|
||||
head_number(0),
|
||||
batch_size(0),
|
||||
seq_length(0)
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
Params(Arguments const &args, void *workspace = nullptr):
|
||||
problem_sizes0(args.problem_sizes0),
|
||||
problem_sizes1(args.problem_sizes1),
|
||||
problem_count(args.problem_count),
|
||||
threadblock_count(args.threadblock_count),
|
||||
ptr_Q(args.ptr_Q),
|
||||
ptr_K(args.ptr_K),
|
||||
ptr_P(args.ptr_P),
|
||||
ptr_V(args.ptr_V),
|
||||
ptr_O(args.ptr_O),
|
||||
ptr_Max(args.ptr_Max),
|
||||
ptr_Sum(args.ptr_Sum),
|
||||
block_P(args.block_P),
|
||||
block_Norm(args.block_Norm),
|
||||
block_Sum(args.block_Sum),
|
||||
offset_P(args.offset_P),
|
||||
offset_Norm_Device(args.offset_Norm_Device),
|
||||
offset_Sum_Device(args.offset_Sum_Device),
|
||||
ldq(args.ldq),
|
||||
ldk(args.ldk),
|
||||
ldp(args.ldp),
|
||||
ldv(args.ldv),
|
||||
ldo(args.ldo),
|
||||
problem_sizes0_host(args.problem_sizes0_host),
|
||||
problem_sizes1_host(args.problem_sizes1_host),
|
||||
problem_sizes0_real(args.problem_sizes0_real),
|
||||
alpha0(args.alpha0),
|
||||
alpha1(args.alpha1),
|
||||
beta(args.beta),
|
||||
head_number(args.head_number),
|
||||
batch_size(args.batch_size),
|
||||
seq_length(args.seq_length),
|
||||
reduction(args.reduction)
|
||||
{
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
private:
|
||||
|
||||
Params params_;
|
||||
GemmGrouped0 gemm_grouped0;
|
||||
GemmGrouped1 gemm_grouped1;
|
||||
|
||||
|
||||
public:
|
||||
|
||||
/// Ctor
|
||||
FusedMultiHeadAttention() {
|
||||
|
||||
}
|
||||
|
||||
/// Initialize
|
||||
Status initialize(Arguments const &args,
|
||||
void *workspace0 = nullptr,
|
||||
void *workspace1 = nullptr) {
|
||||
|
||||
params_ = Params(args);
|
||||
|
||||
typename GemmGrouped0::Arguments args_gemm0(
|
||||
params_.problem_sizes0,
|
||||
params_.problem_count,
|
||||
params_.threadblock_count,
|
||||
params_.ptr_Q,
|
||||
params_.ptr_K,
|
||||
params_.ptr_P,
|
||||
params_.ptr_P,
|
||||
params_.ptr_Max,
|
||||
params_.ptr_Sum,
|
||||
params_.ldq,
|
||||
params_.ldk,
|
||||
params_.ldp,
|
||||
params_.ldp,
|
||||
typename GemmGrouped0::GemmKernel::EpilogueVisitor::Arguments(
|
||||
{
|
||||
params_.alpha0,
|
||||
params_.beta
|
||||
}
|
||||
),
|
||||
params_.problem_sizes0_host,
|
||||
params_.problem_sizes0_real
|
||||
);
|
||||
|
||||
|
||||
Status result0 = gemm_grouped0.initialize(args_gemm0, workspace0);
|
||||
|
||||
typename EpilogueOutputOp1::Params epilogue_op1(params_.alpha1, params_.beta);
|
||||
|
||||
typename GemmGrouped1::Arguments args_gemm1(
|
||||
params_.problem_sizes1,
|
||||
params_.problem_count,
|
||||
params_.threadblock_count,
|
||||
epilogue_op1,
|
||||
params_.ptr_P,
|
||||
params_.ptr_V,
|
||||
params_.ptr_O,
|
||||
params_.ptr_O,
|
||||
(void**)params_.ptr_Max,
|
||||
(void**)params_.ptr_Sum,
|
||||
params_.ldp,
|
||||
params_.ldv,
|
||||
params_.ldo,
|
||||
params_.ldo,
|
||||
params_.problem_sizes1_host
|
||||
);
|
||||
|
||||
Status result1 = gemm_grouped1.initialize(args_gemm1, workspace1);
|
||||
|
||||
if ((result0 == cutlass::Status::kSuccess) && (result1 == cutlass::Status::kSuccess) ) {
|
||||
return cutlass::Status::kSuccess;
|
||||
}else{
|
||||
if (result0 != cutlass::Status::kSuccess) {
|
||||
return result0;
|
||||
}else{
|
||||
return result1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Run
|
||||
Status run(cudaStream_t stream = nullptr) {
|
||||
|
||||
Status result = gemm_grouped0.run();
|
||||
cudaError_t error_info;
|
||||
|
||||
if (result != cutlass::Status::kSuccess) {
|
||||
return cutlass::Status::kErrorInternal;
|
||||
}
|
||||
|
||||
int thread_per_block = 1024;
|
||||
|
||||
dim3 final_reduction_grid(params_.head_number * params_.batch_size);
|
||||
dim3 final_reduction_block(thread_per_block);
|
||||
|
||||
cutlass::Kernel<ApplyFinalReductionDevice><<<
|
||||
final_reduction_grid, final_reduction_block, sizeof(typename ApplyFinalReductionDevice::SharedStorage), stream
|
||||
>>>(params_.reduction);
|
||||
|
||||
error_info = cudaGetLastError();
|
||||
|
||||
if (error_info != cudaSuccess) {
|
||||
return cutlass::Status::kErrorInternal;
|
||||
}
|
||||
|
||||
result = gemm_grouped1.run();
|
||||
|
||||
if (result != cutlass::Status::kSuccess) {
|
||||
return cutlass::Status::kErrorInternal;
|
||||
}
|
||||
|
||||
return cutlass::Status::kSuccess;
|
||||
}
|
||||
|
||||
/// Function call operator
|
||||
Status operator()(cudaStream_t stream = nullptr) {
|
||||
return run(stream);
|
||||
}
|
||||
};
|
||||
|
||||
}
|
||||
@@ -0,0 +1,522 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Grouped GEMM kernel with epilogue visitor customized for softmax
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/matrix_coord.h"
|
||||
#include "cutlass/complex.h"
|
||||
#include "cutlass/semaphore.h"
|
||||
#include "cutlass/util/device_memory.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/trace.h"
|
||||
#include "cutlass/gemm/kernel/gemm_transpose_operands.h"
|
||||
#include "cutlass/gemm/kernel/gemm_grouped_problem_visitor.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace gemm {
|
||||
namespace kernel {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate
|
||||
typename Epilogue_, ///! Epilogue
|
||||
typename ThreadblockSwizzle_, ///! Threadblock swizzling function
|
||||
GroupScheduleMode GroupScheduleMode_, ///! Type of scheduling to perform
|
||||
bool Transposed_ = false,
|
||||
bool UseMask_ = false
|
||||
>
|
||||
struct GemmGroupedWithEpilogueVistor {
|
||||
public:
|
||||
|
||||
using Mma = Mma_;
|
||||
using Epilogue = Epilogue_;
|
||||
using ThreadblockSwizzle = ThreadblockSwizzle_;
|
||||
|
||||
static GroupScheduleMode const kGroupScheduleMode = GroupScheduleMode_;
|
||||
|
||||
using EpilogueVisitor = typename Epilogue::Visitor;
|
||||
using EpilogueOutputOp = typename EpilogueVisitor::ElementwiseFunctor;
|
||||
static bool const kTransposed = Transposed_;
|
||||
|
||||
// Optional transpose
|
||||
using MapArguments = kernel::detail::MapArguments<
|
||||
typename Mma::IteratorA::Element,
|
||||
typename Mma::IteratorA::Layout,
|
||||
Mma::kTransformA,
|
||||
Mma::IteratorA::AccessType::kElements,
|
||||
typename Mma::IteratorB::Element,
|
||||
typename Mma::IteratorB::Layout,
|
||||
Mma::kTransformB,
|
||||
Mma::IteratorB::AccessType::kElements,
|
||||
typename Mma::LayoutC,
|
||||
kTransposed
|
||||
>;
|
||||
|
||||
// Public-facing type definitions related to operand element type, layout, and complex conjugate
|
||||
// operation. Must interact with the 'kTransposed' notion.
|
||||
using ElementA = typename MapArguments::ElementA;
|
||||
using LayoutA = typename MapArguments::LayoutA;
|
||||
using ElementB = typename MapArguments::ElementB;
|
||||
using LayoutB = typename MapArguments::LayoutB;
|
||||
using ElementC = typename EpilogueVisitor::ElementOutput;
|
||||
using LayoutC = typename MapArguments::LayoutC;
|
||||
|
||||
using ElementNorm = typename EpilogueVisitor::ElementNorm;
|
||||
using ElementSum = typename EpilogueVisitor::ElementSum;
|
||||
|
||||
static ComplexTransform const kTransformA = MapArguments::kTransformA;
|
||||
static ComplexTransform const kTransformB = MapArguments::kTransformB;
|
||||
|
||||
// Type definitions about the mainloop.
|
||||
using Operator = typename Mma::Operator;
|
||||
using OperatorClass = typename Mma::Operator::OperatorClass;
|
||||
using ThreadblockShape = typename Mma::Shape;
|
||||
using WarpShape = typename Mma::Operator::Shape;
|
||||
using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
|
||||
using ArchTag = typename Mma::ArchTag;
|
||||
|
||||
static int const kStages = Mma::kStages;
|
||||
static int const kAlignmentA = MapArguments::kAlignmentA;
|
||||
static int const kAlignmentB = MapArguments::kAlignmentB;
|
||||
static int const kAlignmentC = EpilogueVisitor::kElementsPerAccess;
|
||||
|
||||
/// Warp count (concept: GemmShape)
|
||||
using WarpCount = typename Mma::WarpCount;
|
||||
static int const kThreadCount = 32 * WarpCount::kCount;
|
||||
|
||||
using ProblemVisitor = GemmGroupedProblemVisitor<
|
||||
ThreadblockShape,
|
||||
kGroupScheduleMode,
|
||||
kThreadCount,
|
||||
kThreadCount,
|
||||
kTransposed>;
|
||||
|
||||
//
|
||||
// Structures
|
||||
//
|
||||
|
||||
/// Argument structure
|
||||
struct Arguments {
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
GemmCoord *problem_sizes;
|
||||
// when using mask, real problem sizes may not be aligned
|
||||
// then we need to mask out unpadded elements in softmax
|
||||
GemmCoord *problem_sizes_real;
|
||||
int problem_count;
|
||||
int threadblock_count;
|
||||
|
||||
ElementA ** ptr_A;
|
||||
ElementB ** ptr_B;
|
||||
ElementC ** ptr_C;
|
||||
ElementC ** ptr_D;
|
||||
|
||||
ElementNorm **ptr_Max;
|
||||
ElementSum **ptr_Sum;
|
||||
|
||||
typename LayoutA::Stride::LongIndex *lda;
|
||||
typename LayoutB::Stride::LongIndex *ldb;
|
||||
typename LayoutC::Stride::LongIndex *ldc;
|
||||
typename LayoutC::Stride::LongIndex *ldd;
|
||||
|
||||
typename EpilogueVisitor::Arguments epilogue_visitor;
|
||||
|
||||
// Only used by device-level operator
|
||||
GemmCoord *host_problem_sizes;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Default ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Arguments():
|
||||
problem_count(0),
|
||||
threadblock_count(0),
|
||||
ptr_A(nullptr),
|
||||
ptr_B(nullptr),
|
||||
ptr_C(nullptr),
|
||||
ptr_D(nullptr),
|
||||
ptr_Max(nullptr),
|
||||
ptr_Sum(nullptr),
|
||||
lda(nullptr),
|
||||
ldb(nullptr),
|
||||
ldc(nullptr),
|
||||
ldd(nullptr),
|
||||
host_problem_sizes(nullptr)
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
/// Ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Arguments(
|
||||
GemmCoord *problem_sizes,
|
||||
int problem_count,
|
||||
int threadblock_count,
|
||||
ElementA ** ptr_A,
|
||||
ElementB ** ptr_B,
|
||||
ElementC ** ptr_C,
|
||||
ElementC ** ptr_D,
|
||||
ElementNorm **ptr_Max,
|
||||
ElementSum **ptr_Sum,
|
||||
typename LayoutA::Stride::LongIndex *lda,
|
||||
typename LayoutB::Stride::LongIndex *ldb,
|
||||
typename LayoutC::Stride::LongIndex *ldc,
|
||||
typename LayoutC::Stride::LongIndex *ldd,
|
||||
typename EpilogueVisitor::Arguments epilogue_visitor_,
|
||||
GemmCoord *host_problem_sizes=nullptr,
|
||||
GemmCoord *problem_sizes_real=nullptr
|
||||
):
|
||||
problem_sizes(problem_sizes),
|
||||
problem_count(problem_count),
|
||||
threadblock_count(threadblock_count),
|
||||
ptr_A(ptr_A),
|
||||
ptr_B(ptr_B),
|
||||
ptr_C(ptr_C),
|
||||
ptr_D(ptr_D),
|
||||
ptr_Max(ptr_Max),
|
||||
ptr_Sum(ptr_Sum),
|
||||
lda(lda),
|
||||
ldb(ldb),
|
||||
ldc(ldc),
|
||||
ldd(ldd),
|
||||
epilogue_visitor(epilogue_visitor_),
|
||||
host_problem_sizes(host_problem_sizes),
|
||||
problem_sizes_real(problem_sizes_real)
|
||||
{
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
//
|
||||
// Structure for precomputing values in host memory and passing to kernels
|
||||
//
|
||||
|
||||
/// Parameters structure
|
||||
struct Params {
|
||||
|
||||
typename ProblemVisitor::Params problem_visitor;
|
||||
GemmCoord *problem_sizes_real;
|
||||
int threadblock_count;
|
||||
|
||||
ElementA ** ptr_A;
|
||||
ElementB ** ptr_B;
|
||||
ElementC ** ptr_C;
|
||||
ElementC ** ptr_D;
|
||||
|
||||
ElementNorm **ptr_Max;
|
||||
ElementSum **ptr_Sum;
|
||||
|
||||
typename LayoutA::Stride::LongIndex *lda;
|
||||
typename LayoutB::Stride::LongIndex *ldb;
|
||||
typename LayoutC::Stride::LongIndex *ldc;
|
||||
typename LayoutC::Stride::LongIndex *ldd;
|
||||
|
||||
typename EpilogueVisitor::Params epilogue_visitor;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params():
|
||||
ptr_A(nullptr),
|
||||
ptr_B(nullptr),
|
||||
ptr_C(nullptr),
|
||||
ptr_D(nullptr),
|
||||
ptr_Max(nullptr),
|
||||
ptr_Sum(nullptr),
|
||||
lda(nullptr),
|
||||
ldb(nullptr),
|
||||
ldc(nullptr),
|
||||
ldd(nullptr),
|
||||
problem_sizes_real(problem_sizes_real)
|
||||
{ }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Arguments const &args, void *workspace = nullptr, int32_t tile_count = 0):
|
||||
problem_visitor(args.problem_sizes, args.problem_count, workspace, tile_count),
|
||||
threadblock_count(args.threadblock_count),
|
||||
ptr_A(args.ptr_A),
|
||||
ptr_B(args.ptr_B),
|
||||
ptr_C(args.ptr_C),
|
||||
ptr_D(args.ptr_D),
|
||||
ptr_Max(args.ptr_Max),
|
||||
ptr_Sum(args.ptr_Sum),
|
||||
lda(args.lda),
|
||||
ldb(args.ldb),
|
||||
ldc(args.ldc),
|
||||
ldd(args.ldd),
|
||||
epilogue_visitor(args.epilogue_visitor),
|
||||
problem_sizes_real(args.problem_sizes_real)
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void update(
|
||||
Arguments const &args,
|
||||
void *workspace = nullptr,
|
||||
int32_t tile_count = -1) {
|
||||
|
||||
problem_visitor = typename ProblemVisitor::Params(args.problem_sizes, args.problem_count, workspace, tile_count);
|
||||
threadblock_count = args.threadblock_count;
|
||||
ptr_A = args.ptr_A;
|
||||
ptr_B = args.ptr_B;
|
||||
ptr_C = args.ptr_C;
|
||||
ptr_D = args.ptr_D;
|
||||
ptr_Max = args.ptr_Max;
|
||||
ptr_Sum = args.ptr_Sum;
|
||||
lda = args.lda;
|
||||
ldb = args.ldb;
|
||||
ldc = args.ldc;
|
||||
ldd = args.ldd;
|
||||
problem_sizes_real = args.problem_sizes_real;
|
||||
}
|
||||
};
|
||||
|
||||
/// Shared memory storage structure
|
||||
struct SharedStorage {
|
||||
union {
|
||||
typename Mma::SharedStorage main_loop;
|
||||
struct {
|
||||
typename Epilogue::SharedStorage epilogue;
|
||||
typename EpilogueVisitor::SharedStorage visitor;
|
||||
} epilogue;
|
||||
} kernel;
|
||||
|
||||
// ProblemVisitor shared storage can't be overlapped with others
|
||||
typename ProblemVisitor::SharedStorage problem_visitor;
|
||||
};
|
||||
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_DEVICE
|
||||
GemmGroupedWithEpilogueVistor() { }
|
||||
|
||||
/// Determines whether kernel satisfies alignment
|
||||
static Status can_implement(cutlass::gemm::GemmCoord const & problem_size) {
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
static Status can_implement(Arguments const &args) {
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
static size_t get_extra_workspace_size(
|
||||
Arguments const &args,
|
||||
cutlass::gemm::GemmCoord const &grid_tiled_shape) {
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
/// Executes one GEMM
|
||||
CUTLASS_DEVICE
|
||||
void operator()(Params const ¶ms, SharedStorage &shared_storage) {
|
||||
|
||||
//
|
||||
// These types shadow the type-level definitions and support the ability to implement
|
||||
// a 'transposed' GEMM that computes the transposed problems.
|
||||
//
|
||||
using ElementA = typename Mma::IteratorA::Element;
|
||||
using LayoutA = typename Mma::IteratorA::Layout;
|
||||
using ElementB = typename Mma::IteratorB::Element;
|
||||
using LayoutB = typename Mma::IteratorB::Layout;
|
||||
using ElementC = typename EpilogueVisitor::ElementOutput;
|
||||
using LayoutC = typename Mma::LayoutC;
|
||||
|
||||
//
|
||||
// Problem visitor.
|
||||
//
|
||||
ProblemVisitor problem_visitor(
|
||||
params.problem_visitor,
|
||||
shared_storage.problem_visitor,
|
||||
blockIdx.x);
|
||||
|
||||
// Outer 'persistent' loop to iterate over tiles
|
||||
while (problem_visitor.next_tile()) {
|
||||
|
||||
GemmCoord problem_size = problem_visitor.problem_size();
|
||||
int32_t problem_idx = problem_visitor.problem_index();
|
||||
int32_t threadblock_idx = int32_t(problem_visitor.threadblock_idx());
|
||||
|
||||
GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
|
||||
|
||||
cutlass::gemm::GemmCoord threadblock_offset(
|
||||
int(threadblock_idx / grid_shape.n()) * Mma::Shape::kM,
|
||||
int(threadblock_idx % grid_shape.n()) * Mma::Shape::kN,
|
||||
0);
|
||||
|
||||
// Load element pointers. Exchange pointers and strides if working on the transpose
|
||||
ElementA *ptr_A = reinterpret_cast<ElementA *>((kTransposed ? params.ptr_B[problem_idx] : params.ptr_A[problem_idx]));
|
||||
typename LayoutA::LongIndex ldm_A = (kTransposed ? params.ldb[problem_idx] : params.lda[problem_idx]);
|
||||
|
||||
ElementB *ptr_B = reinterpret_cast<ElementB *>((kTransposed ? params.ptr_A[problem_idx] : params.ptr_B[problem_idx]));
|
||||
typename LayoutB::LongIndex ldm_B = (kTransposed ? params.lda[problem_idx] : params.ldb[problem_idx]);
|
||||
|
||||
// Compute initial location in logical coordinates
|
||||
cutlass::MatrixCoord tb_offset_A{
|
||||
threadblock_offset.m(),
|
||||
0,
|
||||
};
|
||||
|
||||
cutlass::MatrixCoord tb_offset_B{
|
||||
0,
|
||||
threadblock_offset.n()
|
||||
};
|
||||
|
||||
// Compute position within threadblock
|
||||
int thread_idx = threadIdx.x;
|
||||
|
||||
// Construct iterators to A and B operands
|
||||
typename Mma::IteratorA iterator_A(
|
||||
LayoutA(ldm_A),
|
||||
ptr_A,
|
||||
{problem_size.m(), problem_size.k()},
|
||||
thread_idx,
|
||||
tb_offset_A);
|
||||
|
||||
typename Mma::IteratorB iterator_B(
|
||||
LayoutB(ldm_B),
|
||||
ptr_B,
|
||||
{problem_size.k(), problem_size.n()},
|
||||
thread_idx,
|
||||
tb_offset_B);
|
||||
|
||||
typename Mma::FragmentC accumulators;
|
||||
|
||||
accumulators.clear();
|
||||
|
||||
// Broadcast the warp_id computed by lane 0 to ensure dependent code
|
||||
// is compiled as warp-uniform.
|
||||
int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
|
||||
|
||||
int lane_idx = threadIdx.x % 32;
|
||||
|
||||
//
|
||||
// Matrix multiply phase
|
||||
//
|
||||
|
||||
// Construct thread-scoped matrix multiply
|
||||
Mma mma(shared_storage.kernel.main_loop, thread_idx, warp_idx, lane_idx);
|
||||
|
||||
// Compute threadblock-scoped matrix multiply-add
|
||||
int gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK;
|
||||
|
||||
// Wait for all threads to finish their epilogue phases from the previous tile.
|
||||
__syncthreads();
|
||||
|
||||
// Compute threadblock-scoped matrix multiply-add
|
||||
mma(
|
||||
gemm_k_iterations,
|
||||
accumulators,
|
||||
iterator_A,
|
||||
iterator_B,
|
||||
accumulators);
|
||||
|
||||
ElementC *ptr_C = params.ptr_C[problem_idx];
|
||||
ElementC *ptr_D = params.ptr_D[problem_idx];
|
||||
|
||||
ElementNorm *ptr_Max = params.ptr_Max[problem_idx];
|
||||
ElementSum *ptr_Sum = params.ptr_Sum[problem_idx];
|
||||
|
||||
LayoutC layout_C(params.ldc[problem_idx]);
|
||||
LayoutC layout_D(params.ldd[problem_idx]);
|
||||
|
||||
int column_offset = (threadblock_offset.n() / ThreadblockShape::kN) * problem_size.m();
|
||||
|
||||
typename EpilogueVisitor::OutputTileIterator::Params params_C(layout_C);
|
||||
typename EpilogueVisitor::OutputTileIterator::Params params_D(layout_D);
|
||||
|
||||
//
|
||||
// Construct the epilogue visitor
|
||||
//
|
||||
|
||||
EpilogueVisitor epilogue_visitor(
|
||||
params.epilogue_visitor,
|
||||
shared_storage.kernel.epilogue.visitor,
|
||||
problem_size.mn(),
|
||||
thread_idx,
|
||||
warp_idx,
|
||||
lane_idx,
|
||||
params_C,
|
||||
params_D,
|
||||
ptr_C,
|
||||
ptr_D,
|
||||
ptr_Max,
|
||||
ptr_Sum,
|
||||
threadblock_offset.mn(),
|
||||
column_offset,
|
||||
params.problem_sizes_real[problem_idx].mn()
|
||||
);
|
||||
|
||||
// Construct the epilogue
|
||||
Epilogue epilogue(
|
||||
shared_storage.kernel.epilogue.epilogue,
|
||||
thread_idx,
|
||||
warp_idx,
|
||||
lane_idx);
|
||||
|
||||
// Execute the epilogue operator to update the destination tensor
|
||||
epilogue(epilogue_visitor, accumulators);
|
||||
|
||||
// Next tile
|
||||
problem_visitor.advance(gridDim.x);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace kernel
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
Reference in New Issue
Block a user