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

More updates for 3.1 (#958)

Browse Source 
* Updates for 3.1

* Minor change

* doc link fix

* Minor updates
This commit is contained in:
ANIKET SHIVAM
2023-05-24 10:17:16 -04:00
committed by GitHub
parent 13f413493a
commit f079619f5e
48 changed files with 1611 additions and 1858 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/gemm/device/gemm_universal_adapter.h
+21 -21
View File
@@ -28,9 +28,9 @@
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*!
/*!
\file
\brief The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and
\brief The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and
batched array variants.
*/
@@ -57,7 +57,7 @@ namespace cutlass::gemm::device {
////////////////////////////////////////////////////////////////////////////////
/*!
/*!
GemmUniversalAdapter is a stateful, reusable GEMM handle built around a kernel
of type cutlass::gemm::kernel::Gemm or cutlass::gemm::kernel::GemmUniversal.
@@ -159,10 +159,10 @@ public:
typename CollectiveMainloop::GmemTiledCopyA, ElementA>();
static int constexpr kAlignmentB = gemm::detail::get_alignment_count_from_gmem_tiled_copy<
typename CollectiveMainloop::GmemTiledCopyB, ElementB>();
// NOTE: 3.0 DefaultEpilogues don't support vectorized stores (yet)
static int constexpr kAlignmentC = 1;
static int constexpr kAlignmentD = 1;
static int constexpr kAlignmentC = gemm::detail::get_alignment_count_from_gmem_tiled_copy<
typename CollectiveEpilogue::GmemTiledCopyC, ElementC>();
static int constexpr kAlignmentD = gemm::detail::get_alignment_count_from_gmem_tiled_copy<
typename CollectiveEpilogue::GmemTiledCopyD, ElementD>();
using EpilogueOutputOp = typename CollectiveEpilogue::ThreadEpilogueOp;
@@ -327,7 +327,7 @@ public:
static Status
run(Params& params, cudaStream_t stream = nullptr) {
CUTLASS_TRACE_HOST("GemmUniversal::run()");
dim3 constexpr block = GemmKernel::get_block_shape();
dim3 const block = GemmKernel::get_block_shape();
dim3 const grid = get_grid_shape(params);
// configure smem size and carveout
@@ -404,19 +404,19 @@ public:
using GemmKernel = GemmKernel_;
static bool const kInternalTranspose =
static bool const kInternalTranspose =
cute::is_same<typename GemmKernel::LayoutC, cutlass::layout::RowMajor>::value;
using ThreadblockShape = typename GemmKernel::Mma::Shape;
using WarpShape = typename GemmKernel::WarpShape;
using InstructionShape = typename GemmKernel::InstructionShape;
// warp-level, arch-level (instruction), math operator
// warp-level, arch-level (instruction), math operator
using WarpMmaOperator = typename GemmKernel::Mma::Policy::Operator;
using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator;
using MathOperator = typename WarpMmaOperator::MathOperator;
// Operator class and arch tag extract bottom-up
// Operator class and arch tag extract bottom-up
// set it for top-level gemm device-level template
using OperatorClass = typename WarpMmaOperator::OperatorClass;
using ArchTag = typename WarpMmaOperator::ArchTag;
@@ -444,15 +444,15 @@ public:
using LayoutB = typename MapArguments::LayoutB;
static ComplexTransform const kTransformB = MapArguments::kTransformB;
static int const kAlignmentB = MapArguments::kAlignmentB;
using ElementC = typename GemmKernel::ElementC;
using LayoutC = typename MapArguments::LayoutC;
static int const kAlignmentC = GemmKernel::kAlignmentC;
// C and D same type for 2.x kernel
using ElementD = ElementC;
using LayoutD = LayoutC;
using TensorRefA = TensorRef<ElementA const, LayoutA>;
using TensorRefB = TensorRef<ElementB const, LayoutB>;
using TensorRefC = TensorRef<ElementC const, LayoutC>;
@@ -493,12 +493,12 @@ public:
/// Gets the workspace size
static size_t get_workspace_size(Arguments const &args) {
return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args));
}
/// Computes the grid shape
static dim3 get_grid_shape(Arguments const &args) {
static dim3 get_grid_shape(Arguments const &args) {
return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args));
}
@@ -532,12 +532,12 @@ public:
/// Runs the kernel using initialized state.
Status operator()(
Arguments const &args,
void *workspace = nullptr,
Arguments const &args,
void *workspace = nullptr,
cudaStream_t stream = nullptr) {
Status status = initialize(args, workspace, stream);
if (status == Status::kSuccess) {
status = run(stream);
}
include/cutlass/gemm/device/gemv.h
+22 -9
View File
@@ -75,6 +75,8 @@ public:
static ComplexTransform const kTransformB = GemvKernel::kTransformB;
static int const kThreadCount = GemvKernel::kThreadCount;
static int const kThreadsPerRow = GemvKernel::kThreadsPerRow;
static int const kStages = GemvKernel::kStages;
static int const kAlignmentA = GemvKernel::kAlignmentA;
@@ -106,8 +108,23 @@ public:
}
/// Computes the grid shape
static dim3 get_grid_shape(Arguments const &args) {
return dim3((args.problem_size.row() + (kThreadCount - 1)) / kThreadCount, 1, args.batch_count % 65565);
static dim3 get_grid_shape(Arguments const &args, dim3 const &block) {
if(platform::is_same<LayoutA, layout::ColumnMajor>::value) {
return dim3((args.problem_size.row() + (block.x - 1)) / block.x, 1, args.batch_count % 65536);
}
else {
return dim3((args.problem_size.row() + (block.y - 1)) / block.y, 1, args.batch_count % 65536);
}
}
/// Computes the block shape
static dim3 get_block_shape() {
if(platform::is_same<LayoutA, layout::ColumnMajor>::value) {
return dim3(kThreadCount, 1, 1);
}
else {
return dim3(kThreadsPerRow, kThreadCount / kThreadsPerRow, 1);
}
}
/// Initializes Gemv state from arguments.
@@ -124,8 +141,8 @@ public:
/// Runs the kernel using initialized state.
Status run(cudaStream_t stream = nullptr) {
dim3 grid = get_grid_shape(params_);
dim3 block(GemvKernel::kThreadCount, 1, 1);
dim3 block = get_block_shape();
dim3 grid = get_grid_shape(params_, block);
int smem_size = int(sizeof(typename GemvKernel::SharedStorage));
@@ -137,11 +154,7 @@ public:
//
cudaError_t result = cudaGetLastError();
if (result != cudaSuccess) {
return Status::kErrorInternal;
}
return Status::kSuccess;
return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal;
}
/// Runs the kernel using initialized state.
include/cutlass/gemm/device/gemv_strided_batched.h
-167
View File
@@ -1,167 +0,0 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/arch/arch.h"
#include "cutlass/device_kernel.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
#include "cutlass/gemm/kernel/gemm_universal.h"
#include "cutlass/gemm/kernel/default_gemm_universal.h"
#include "cutlass/gemm/device/default_gemm_configuration.h"
#include "cutlass/gemm/device/gemm_universal_base.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace device {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename GemvKernel_>
class GemvStridedBatched {
public:
using GemvKernel = GemvKernel_;
using ElementA = typename GemvKernel::ElementA;
using LayoutA = typename GemvKernel::LayoutA;
using ElementB = typename GemvKernel::ElementB;
using ElementC = typename GemvKernel::ElementC;
using ElementAccumulator = typename GemvKernel::ElementAccumulator;
using EpilogueOutputOp = typename GemvKernel::EpilogueOutputOp;
static ComplexTransform const kTransformA = GemvKernel::kTransformA;
static ComplexTransform const kTransformB = GemvKernel::kTransformB;
static int const kThreadCount = GemvKernel::kThreadCount;
static int const mThreadCount = GemvKernel::mThreadCount;
static int const kStages = GemvKernel::kStages;
static int const kAlignmentA = GemvKernel::kAlignmentA;
static int const kAlignmentB = GemvKernel::kAlignmentB;
static int const kAlignmentC = GemvKernel::kAlignmentC;
using Arguments = typename GemvKernel::Arguments;
using Params = typename GemvKernel::Params;
private:
Params params_;
public:
/// Constructs the Gemv.
GemvStridedBatched() {}
/// Determines whether the Gemv can execute the given problem.
static Status can_implement(Arguments const& args) {
return GemvKernel::can_implement(args);
}
/// Gets the workspace size
static size_t get_workspace_size(Arguments const& args) { return 0; }
/// Initializes Gemv state from arguments.
Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
params_ = Params(args);
if (args.problem_size.column() % GemvKernel::kElementsPerAccess) {
return Status::kErrorMisalignedOperand;
}
return Status::kSuccess;
}
/// Lightweight update given a subset of arguments
Status update(Arguments const &args, void *workspace = nullptr) {
return params_.update(args);
}
/// Runs the kernel using initialized state.
Status run(cudaStream_t stream = nullptr) {
dim3 grid(1, 1, params_.batch_count % 65536);
dim3 block(kThreadCount, mThreadCount, 1);
int smem_size = 0;
// Launch
cutlass::Kernel<GemvKernel><<<grid, block, smem_size, stream>>>(params_);
//
// Query for errors
//
cudaError_t result = cudaGetLastError();
return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal;
}
/// Runs the kernel using initialized state.
Status operator()(cudaStream_t stream = nullptr) { return run(stream); }
/// Runs the kernel using initialized state.
Status operator()(
Arguments const &args,
void *workspace = nullptr,
cudaStream_t stream = nullptr) {
Status status = initialize(args, workspace, stream);
if (status == Status::kSuccess) {
status = run(stream);
}
return status;
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace device
} // namespace gemm
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/gemm.h
+19 -7
View File
@@ -522,15 +522,27 @@ stride_to_layout_tag_B() {
template <class GmemTiledCopy, class Element>
constexpr int
get_alignment_count_from_gmem_tiled_copy() {
// For TMA tiled copies, we know the alignment has to be 128 bits
if constexpr ( cute::is_base_of_v<cute::SM90_TMA_LOAD, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM90_TMA_LOAD_MULTICAST, GmemTiledCopy>
) {
return 128 / sizeof_bits<Element>::value;
if constexpr (cute::is_void_v<GmemTiledCopy>) {
return 1;
}
// Account for ElementC = void kernels
else if constexpr (cute::is_void_v<Element>) {
return 0;
}
else {
// For non-TMA tiled copies, TiledCopy holds the alignment count directly in its TiledShape_MN
return GmemTiledCopy::NumValSrc;
// For TMA tiled copies, we know the alignment has to be 128 bits
if constexpr ( cute::is_base_of_v<cute::SM90_TMA_LOAD, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM90_TMA_LOAD_MULTICAST, GmemTiledCopy>
|| cute::is_base_of_v<cute::SM90_TMA_STORE, GmemTiledCopy>
) {
return 128 / sizeof_bits<Element>::value;
}
else {
// For non-TMA tiled copies, TiledCopy holds the alignment count directly in its TiledShape_MN
return GmemTiledCopy::NumValSrc;
}
}
}
include/cutlass/gemm/kernel/gemv.h
+375 -26
View File
@@ -41,9 +41,13 @@
#include "cutlass/complex.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/arch/memory.h"
#include "cutlass/arch/cache_operation.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/numeric_conversion.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
@@ -58,18 +62,49 @@ template <
typename ElementB_,
typename ElementC_,
typename ElementAccumulator_,
typename EpilogueOutputOp_
typename EpilogueOutputOp_,
int kElementsPerAccess_ = 1, ///< Number of elements involved in a global access.
int kThreadCount_ = 0, ///< Number of threads in the thread block.
/// It will be calculated automatically if set to 0.
int kThreadsPerRow_ = 0 ///< Number of threads in the k dimension.
/// It will be calculated automatically if set to 0.
>
struct Gemv {
struct Gemv;
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Specializations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
// GEMV for column-major A matrix
template <
typename ElementA_,
typename ElementB_,
typename ElementC_,
typename ElementAccumulator_,
typename EpilogueOutputOp_,
int kElementsPerAccess_,
int kThreadCount_,
int kThreadsPerRow_
>
struct Gemv <
ElementA_,
layout::ColumnMajor,
ElementB_,
ElementC_,
ElementAccumulator_,
EpilogueOutputOp_,
kElementsPerAccess_,
kThreadCount_,
kThreadsPerRow_
>{
public:
using ElementA = ElementA_;
using LayoutA = layout::ColumnMajor;
using TensorRefA = TensorRef<ElementA, LayoutA>;
static_assert(platform::is_same<LayoutA, LayoutA_>::value,
"Only supported for column-major A matrix");
using ElementB = ElementB_;
using ElementC = ElementC_;
@@ -79,7 +114,10 @@ public:
static ComplexTransform const kTransformA = ComplexTransform::kNone;
static ComplexTransform const kTransformB = ComplexTransform::kNone;
static int const kThreadCount = 32;
// thread block shape (kThreadCount, 1, 1)
static int const kThreadCount = (kThreadCount_ == 0) ? 32 : kThreadCount_;
static int const kThreadsPerRow = kThreadsPerRow_;
static int const kStages = 1;
static int const kAlignmentA = 1;
@@ -121,17 +159,17 @@ public:
MatrixCoord problem_size,
int batch_count,
typename EpilogueOutputOp::Params output_op,
TensorRefA ref_A,
void const * ptr_B,
void const * ptr_C,
void * ptr_D,
int64_t inc_B,
int64_t inc_C,
int64_t inc_D,
int64_t batch_stride_A,
int64_t batch_stride_B,
int64_t batch_stride_C,
int64_t batch_stride_D
TensorRefA ref_A,
void const *ptr_B,
void const *ptr_C,
void *ptr_D,
int64_t inc_B,
int64_t inc_C,
int64_t inc_D,
int64_t batch_stride_A,
int64_t batch_stride_B,
int64_t batch_stride_C,
int64_t batch_stride_D
):
problem_size(problem_size),
batch_count(batch_count),
@@ -151,14 +189,44 @@ public:
Arguments(
MatrixCoord problem_size,
int batch_count,
typename EpilogueOutputOp::Params output_op,
TensorRefA ref_A,
void const * ptr_B,
void const * ptr_C,
void * ptr_D,
int64_t inc_B,
int64_t inc_C,
int64_t inc_D
TensorRefA ref_A,
void const *ptr_B,
void const *ptr_C,
void *ptr_D,
int64_t batch_stride_A,
int64_t batch_stride_B,
int64_t batch_stride_C,
int64_t batch_stride_D
):
Arguments(
problem_size,
batch_count,
output_op,
ref_A,
ptr_B,
ptr_C,
ptr_D,
1,
1,
1,
batch_stride_A,
batch_stride_B,
batch_stride_C,
batch_stride_D)
{ }
Arguments(
MatrixCoord problem_size,
typename EpilogueOutputOp::Params output_op,
TensorRefA ref_A,
void const *ptr_B,
void const *ptr_C,
void *ptr_D,
int64_t inc_B,
int64_t inc_C,
int64_t inc_D
):
Arguments(
problem_size,
@@ -206,7 +274,6 @@ public:
/// Determines whether kernel satisfies alignment
static Status can_implement(cutlass::MatrixCoord const & problem_size) {
return Status::kSuccess;
}
@@ -214,7 +281,7 @@ public:
return can_implement(args.problem_size);
}
/// Executes one GEMM
/// Executes one GEMV
CUTLASS_DEVICE
void operator()(Params const &params, SharedStorage &shared_storage) {
@@ -282,6 +349,288 @@ public:
/////////////////////////////////////////////////////////////////////////////////////////////////
// GEMV for row-major A matrix
template <
typename ElementA_,
typename ElementB_,
typename ElementC_,
typename ElementAccumulator_,
typename EpilogueOutputOp_,
int kElementsPerAccess_,
int kThreadCount_,
int kThreadsPerRow_
>
struct Gemv <
ElementA_,
layout::RowMajor,
ElementB_,
ElementC_,
ElementAccumulator_,
EpilogueOutputOp_,
kElementsPerAccess_,
kThreadCount_,
kThreadsPerRow_
>{
public:
using ElementA = ElementA_;
using LayoutA = layout::RowMajor;
using TensorRefA = TensorRef<ElementA, LayoutA>;
using ElementB = ElementB_;
using ElementC = ElementC_;
using ElementAccumulator = ElementAccumulator_;
using EpilogueOutputOp = EpilogueOutputOp_;
static ComplexTransform const kTransformA = ComplexTransform::kNone;
static ComplexTransform const kTransformB = ComplexTransform::kNone;
static FloatRoundStyle const Round = cutlass::FloatRoundStyle::round_to_nearest;
// number of return elements in a global access
static int const kElementsPerAccess = kElementsPerAccess_;
using FragmentA = Array<ElementA, kElementsPerAccess>;
using FragmentB = Array<ElementB, kElementsPerAccess>;
using FragmentCompute = Array<ElementAccumulator, kElementsPerAccess>;
// thread block shape (kThreadsPerRow, kThreadCount / kThreadsPerRow, 1)
static int const kThreadCount = (kThreadCount_ == 0) ? 128 : kThreadCount_;
static int const kThreadsPerRow = (kThreadsPerRow_ == 0) ?
std::min(static_cast<int>(kThreadCount / (kElementsPerAccess * sizeof(ElementA))), 16)
: kThreadsPerRow_;
//
// Structures
//
/// Argument structure
struct Arguments {
MatrixCoord problem_size;
int32_t batch_count;
typename EpilogueOutputOp::Params output_op;
TensorRefA ref_A;
ElementB const *ptr_B;
ElementC const *ptr_C;
ElementC *ptr_D;
int64_t batch_stride_A;
int64_t batch_stride_B;
int64_t batch_stride_C;
int64_t batch_stride_D;
//
// Methods
//
Arguments(): batch_count(0) { }
Arguments(
MatrixCoord problem_size,
int32_t batch_count,
typename EpilogueOutputOp::Params output_op,
TensorRefA ref_A,
void const *ptr_B,
void const *ptr_C,
void *ptr_D,
int64_t batch_stride_A,
int64_t batch_stride_B,
int64_t batch_stride_C,
int64_t batch_stride_D
):
problem_size(problem_size),
batch_count(batch_count),
output_op(output_op),
ref_A(ref_A),
ptr_B(static_cast<ElementB const *>(ptr_B)),
ptr_C(static_cast<ElementC const *>(ptr_C)),
ptr_D(static_cast<ElementC *>(ptr_D)),
batch_stride_A(batch_stride_A),
batch_stride_B(batch_stride_B),
batch_stride_C(batch_stride_C),
batch_stride_D(batch_stride_D)
{ }
Arguments(
MatrixCoord problem_size,
typename EpilogueOutputOp::Params output_op,
TensorRefA ref_A,
void const *ptr_B,
void const *ptr_C,
void *ptr_D
):
Arguments(
problem_size,
1,
output_op,
ref_A,
ptr_B,
ptr_C,
ptr_D,
1,
1,
1,
1)
{ }
Status update(Arguments const &args) {
problem_size = args.problem_size;
batch_count = args.batch_count;
output_op = args.output_op;
ref_A = ref_A;
ptr_B = args.ptr_B;
ptr_C = args.ptr_C;
ptr_D = args.ptr_D;
batch_stride_A = args.batch_stride_A;
batch_stride_B = args.batch_stride_B;
batch_stride_C = args.batch_stride_C;
batch_stride_D = args.batch_stride_D;
return Status::kSuccess;
}
};
using Params = Arguments;
/// Shared memory storage structure
union SharedStorage {
};
public:
//
// Methods
//
CUTLASS_DEVICE
Gemv() {}
/// Determines whether kernel satisfies alignment
static Status can_implement(cutlass::MatrixCoord const &problem_size) {
if (problem_size.column() % kElementsPerAccess != 0) {
return Status::kErrorMisalignedOperand;
}
return Status::kSuccess;
}
static Status can_implement(Arguments const &args) {
return can_implement(args.problem_size);
}
/// Executes one GEMV
CUTLASS_DEVICE
void operator()(Params const &params, SharedStorage &shared_storage) {
// Loop over batch indices
for (int batch_idx = blockIdx.z; batch_idx < params.batch_count; batch_idx += gridDim.z) {
int idx_col_k = threadIdx.x;
int idx_row_m = blockIdx.x * blockDim.y + threadIdx.y;
if (idx_row_m < params.problem_size.row()) {
// problem_size (row = m, column = k)
// matrix A (batch, m, k)
// vector B (batch, 1, k)
// vector C (batch, m, 1)
// vector D (batch, m, 1)
// move in the batch dimension
ElementA const *ptr_A = params.ref_A.data() + batch_idx * params.batch_stride_A;
ElementB const *ptr_B = params.ptr_B + batch_idx * params.batch_stride_B;
ElementC const *ptr_C = params.ptr_C + batch_idx * params.batch_stride_C;
ElementC *ptr_D = params.ptr_D + batch_idx * params.batch_stride_D;
// move in the k dimension
ptr_A += idx_col_k * kElementsPerAccess;
ptr_B += idx_col_k * kElementsPerAccess;
// move in the m dimension
ptr_A += idx_row_m * params.problem_size.column();
ptr_C += idx_row_m;
ptr_D += idx_row_m;
NumericArrayConverter<ElementAccumulator, ElementA, kElementsPerAccess, Round> srcA_converter;
NumericArrayConverter<ElementAccumulator, ElementB, kElementsPerAccess, Round> srcB_converter;
ElementAccumulator accum = 0.f;
FragmentB fragB;
FragmentA fragA;
int unroll_col_k = 0;
// rows of the rolling tile
int const tileA_k = kThreadsPerRow * kElementsPerAccess;
for (; unroll_col_k < params.problem_size.column() / tileA_k * tileA_k; unroll_col_k += tileA_k) {
// fetch from matrix A
arch::global_load<FragmentA,
sizeof(FragmentA),
arch::CacheOperation::LastUse>(fragA, (ptr_A + unroll_col_k), true);
// fetch from vector B
arch::global_load<FragmentB,
sizeof(FragmentB),
arch::CacheOperation::Always>(fragB, (ptr_B + unroll_col_k), true);
FragmentCompute fragB_Compute = srcB_converter(fragB);
FragmentCompute fragA_Compute = srcA_converter(fragA);
// Math
CUTLASS_PRAGMA_UNROLL
for (int e = 0; e < kElementsPerAccess; e++) {
accum += fragA_Compute.at(e) * fragB_Compute.at(e);
}
}
// calculate the rest of K elements
// each thread fetch 1 element each time
for (int k = unroll_col_k + idx_col_k; k < params.problem_size.column(); k += kThreadsPerRow) {
ElementB b = *(ptr_B - idx_col_k * kElementsPerAccess + k);
ElementA a = *(ptr_A - idx_col_k * kElementsPerAccess + k);
accum += ElementAccumulator(a) * ElementAccumulator(b);
}
EpilogueOutputOp output_op(params.output_op);
typename EpilogueOutputOp::FragmentOutput source_fragment;
// prefetch from source matrix C
if (output_op.is_source_needed()) {
source_fragment[0] = *(ptr_C);
}
typename EpilogueOutputOp::FragmentAccumulator accum_fragment;
typename EpilogueOutputOp::FragmentOutput output_fragment;
for (int mask = (kThreadsPerRow >> 1); mask > 0; mask >>= 1) {
accum += __shfl_xor_sync(0xFFFFFFFF, accum, mask, 32);
}
if (idx_col_k == 0) {
accum_fragment[0] = accum;
if (output_op.is_source_needed()) {
output_fragment = output_op(accum_fragment, source_fragment);
}
else {
output_fragment = output_op(accum_fragment);
}
*ptr_D = output_fragment[0];
}
}
}
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace kernel
} // namespace gemm
} // namespace cutlass
include/cutlass/gemm/kernel/gemv_strided_batched.h
-368
View File
@@ -1,368 +0,0 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/fast_math.h"
#include "cutlass/matrix_coord.h"
#include "cutlass/complex.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/arch/memory.h"
#include "cutlass/arch/cache_operation.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/numeric_conversion.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace kernel {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
typename ElementA_, /// matrix
typename LayoutA_,
typename ElementB_, /// vector
typename ElementC_,
typename ElementAccumulator_,
int kElementsPerAccess_,
typename EpilogueOutputOp_
>
struct GemvStridedBatched {
public:
using ElementA = ElementA_;
using LayoutA = layout::RowMajor;
using TensorRefA = TensorRef<ElementA, LayoutA>;
static_assert(std::is_same<LayoutA, LayoutA_>::value,
"Only supported for row-major A matrix");
using ElementB = ElementB_;
using ElementC = ElementC_;
using ElementAccumulator = ElementAccumulator_;
using EpilogueOutputOp = EpilogueOutputOp_;
static ComplexTransform const kTransformA = ComplexTransform::kNone;
static ComplexTransform const kTransformB = ComplexTransform::kNone;
static FloatRoundStyle const Round = cutlass::FloatRoundStyle::round_to_nearest;
// number of return elements in a global access
static int const kElementsPerAccess = kElementsPerAccess_;
using FragmentA = Array<ElementA, kElementsPerAccess>;
using FragmentB = Array<ElementB, kElementsPerAccess>;
using FragmentCompute = Array<ElementAccumulator, kElementsPerAccess>;
// thread block shape (kThreadCount, mThreadCount)
static int const kThreadCount = std::min(static_cast<int>(128 / (kElementsPerAccess * sizeof(ElementA))), 16);
static int const mThreadCount = 128 / kThreadCount;
// rolling tile shape
static int const kTileA = kThreadCount * kElementsPerAccess;
static int const mTileA = mThreadCount * 8;
//
// Structures
//
/// Argument structure
struct Arguments
{
MatrixCoord problem_size;
int32_t batch_count;
typename EpilogueOutputOp::Params output_op;
TensorRefA ref_A;
ElementB const *ptr_B;
ElementC const *ptr_C;
ElementC *ptr_D;
int64_t batch_stride_A;
int64_t batch_stride_B;
int64_t batch_stride_C;
int64_t batch_stride_D;
//
// Methods
//
Arguments() : batch_count(0) {}
Arguments(
MatrixCoord problem_size,
int32_t batch_count,
typename EpilogueOutputOp::Params output_op,
TensorRefA ref_A,
void const *ptr_B,
void const *ptr_C,
void *ptr_D,
int64_t batch_stride_A,
int64_t batch_stride_B,
int64_t batch_stride_C,
int64_t batch_stride_D) : problem_size(problem_size),
batch_count(batch_count),
output_op(output_op),
ref_A(ref_A),
ptr_B(static_cast<ElementB const *>(ptr_B)),
ptr_C(static_cast<ElementC const *>(ptr_C)),
ptr_D(static_cast<ElementC *>(ptr_D)),
batch_stride_A(batch_stride_A),
batch_stride_B(batch_stride_B),
batch_stride_C(batch_stride_C),
batch_stride_D(batch_stride_D)
{
}
Arguments(
MatrixCoord problem_size,
typename EpilogueOutputOp::Params output_op,
TensorRefA ref_A,
void const *ptr_B,
void const *ptr_C,
void *ptr_D) : Arguments(problem_size,
1,
1,
output_op,
ref_A,
ptr_B,
ptr_C,
ptr_D,
1,
1,
1,
1)
{
}
Status update(Arguments const &args)
{
problem_size = args.problem_size;
batch_count = args.batch_count;
output_op = args.output_op;
ref_A = ref_A;
ptr_B = args.ptr_B;
ptr_C = args.ptr_C;
ptr_D = args.ptr_D;
batch_stride_A = args.batch_stride_A;
batch_stride_B = args.batch_stride_B;
batch_stride_C = args.batch_stride_C;
batch_stride_D = args.batch_stride_D;
return Status::kSuccess;
}
};
using Params = Arguments;
/// Shared memory storage structure
union SharedStorage
{
};
public:
//
// Methods
//
CUTLASS_DEVICE
GemvStridedBatched() {}
/// Determines whether kernel satisfies alignment
static Status can_implement(cutlass::MatrixCoord const &problem_size)
{
if (problem_size.column() % kElementsPerAccess != 0)
return Status::kErrorMisalignedOperand;
return Status::kSuccess;
}
static Status can_implement(Arguments const &args)
{
return can_implement(args.problem_size);
}
/// Executes one GEMV
CUTLASS_DEVICE
void operator()(Params const &params, SharedStorage &shared_storage)
{
// Loop over batch indices
for (int batch_idx = blockIdx.z; batch_idx < params.batch_count; batch_idx += gridDim.z)
{
int k_col_id = threadIdx.x;
int m_row_id = threadIdx.y;
// problem_size (row = m, column = k)
// matrix A (batch, m, k)
// vector B (batch, 1, k)
// vector C (batch, m, 1)
// vector D (batch, m, 1)
// move in the batch dimension
ElementA const *ptr_A = params.ref_A.data() + batch_idx * params.batch_stride_A;
ElementB const *ptr_B = params.ptr_B + batch_idx * params.batch_stride_B;
ElementC const *ptr_C = params.ptr_C + batch_idx * params.batch_stride_C;
ElementC *ptr_D = params.ptr_D + batch_idx * params.batch_stride_D;
// move in the k dimension
ptr_A += k_col_id * kElementsPerAccess;
ptr_B += k_col_id * kElementsPerAccess;
// move in the m dimension
ptr_A += m_row_id * params.problem_size.column();
ptr_C += m_row_id;
ptr_D += m_row_id;
NumericArrayConverter<ElementAccumulator, ElementA, kElementsPerAccess, Round> srcA_converter;
NumericArrayConverter<ElementAccumulator, ElementB, kElementsPerAccess, Round> srcB_converter;
for (; m_row_id < params.problem_size.row(); m_row_id += mTileA)
{
ElementAccumulator accum[mTileA / mThreadCount] = {0.f};
FragmentB fragB;
FragmentA fragA[mTileA / mThreadCount];
int mElemCountPerTile = min(mTileA / mThreadCount, (params.problem_size.row() - m_row_id - 1) / mThreadCount + 1);
int kUnroll = 0;
for (; kUnroll < params.problem_size.column() / kTileA * kTileA; kUnroll += kTileA)
{
for (int m = 0; m < mElemCountPerTile; m++)
{
// fetch from matrix A
arch::global_load<FragmentA,
sizeof(FragmentA),
arch::CacheOperation::LastUse>(fragA[m], (ptr_A + kUnroll + m * mThreadCount * params.problem_size.column()), true);
}
// fetch from vector B
arch::global_load<FragmentB,
sizeof(FragmentB),
arch::CacheOperation::Always>(fragB, (ptr_B + kUnroll), true);
for (int m = 0; m < mElemCountPerTile; m++)
{
FragmentCompute fragB_Compute = srcB_converter(fragB);
FragmentCompute fragA_Compute = srcA_converter(fragA[m]);
// Math
CUTLASS_PRAGMA_UNROLL
for (int e = 0; e < kElementsPerAccess; e++)
{
accum[m] += fragA_Compute.at(e) * fragB_Compute.at(e);
}
}
}
// calculate the rest of K elements
// each thread fetch 1 element each time
for (int k = kUnroll + k_col_id; k < params.problem_size.column(); k += kThreadCount)
{
ElementB b = *(ptr_B - k_col_id * kElementsPerAccess + k);
for (int m = 0; m < mElemCountPerTile; m++)
{
ElementA a = *(ptr_A - k_col_id * kElementsPerAccess + k + m * mThreadCount * params.problem_size.column());
accum[m] += ElementAccumulator(a) * ElementAccumulator(b);
}
}
EpilogueOutputOp output_op(params.output_op);
typename EpilogueOutputOp::FragmentOutput source_fragment[mTileA / mThreadCount];
// prefetch from source matrix C
if (output_op.is_source_needed())
{
for (int m = 0; m < mElemCountPerTile; m++)
{
source_fragment[m][0] = *(ptr_C + m * mThreadCount);
}
}
typename EpilogueOutputOp::FragmentAccumulator accum_fragment;
typename EpilogueOutputOp::FragmentOutput output_fragment;
for (int m = 0; m < mElemCountPerTile; m++)
{
for (int mask = (kThreadCount >> 1); mask > 0; mask >>= 1)
{
accum[m] += __shfl_xor_sync(0xFFFFFFFF, accum[m], mask, 32);
}
if (k_col_id == 0)
{
accum_fragment[0] = accum[m];
if (output_op.is_source_needed())
{
output_fragment = output_op(accum_fragment, source_fragment[m]);
}
else
{
output_fragment = output_op(accum_fragment);
}
*(ptr_D + m * mThreadCount) = output_fragment[0];
}
}
ptr_A += mTileA * params.problem_size.column();
ptr_C += mTileA;
ptr_D += mTileA;
}
}
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace kernel
} // namespace gemm
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/kernel/sm70_gemm.hpp
+4 -8
View File
@@ -129,21 +129,18 @@ public:
};
}
static
bool
static bool
can_implement(Arguments const& args) {
return args.mode == GemmUniversalMode::kGemm or
(args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
}
static
int
static int
get_workspace_size(Arguments const& args) {
return 0;
}
static constexpr
dim3
static dim3
get_grid_shape(Params const& params) {
int batch_count = 1;
if constexpr (rank(ProblemShape{}) == 4) {
@@ -157,8 +154,7 @@ public:
);
}
static constexpr
dim3
static dim3
get_block_shape() {
return dim3(MaxThreadsPerBlock, 1, 1);
}
include/cutlass/gemm/kernel/sm90_gemm_tma.hpp
+8 -11
View File
@@ -172,20 +172,20 @@ public:
auto N = get<1>(args.problem_shape);
auto K = get<2>(args.problem_shape);
// Contiguous dimension for the TMA tensor should be 128b aligned
implementable = std::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> ?
implementable = std::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> ?
K % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0;
implementable = implementable && (std::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor> ?
implementable = implementable && (std::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor> ?
N % min_tma_aligned_elements == 0 : K % min_tma_aligned_elements == 0);
implementable = implementable && (!cutlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA<CollectiveEpilogue>::value ||
(cutlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA<CollectiveEpilogue>::value &&
std::is_same_v<gemm::detail::StrideToLayoutTagC_t<StrideC>, layout::RowMajor> ?
std::is_same_v<gemm::detail::StrideToLayoutTagC_t<StrideC>, layout::RowMajor> ?
N % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0));
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
return implementable;
}
constexpr bool is_beta_supported =
constexpr bool is_beta_supported =
CollectiveEpilogue::ThreadEpilogueOp::kScale == cutlass::epilogue::thread::ScaleType::Default;
implementable = is_beta_supported || (args.epilogue.thread.beta == 0 && args.epilogue.thread.beta_ptr == nullptr);
if (!implementable) {
@@ -196,15 +196,13 @@ public:
return implementable;
}
static
int
static int
get_workspace_size(Arguments const& args) {
return 0;
}
// Computes the kernel launch grid shape based on runtime parameters
static constexpr
dim3
static dim3
get_grid_shape(Params const& params) {
auto cluster_shape = ClusterShape{};
auto tile_shape = TileShape{};
@@ -213,8 +211,7 @@ public:
problem_shape_MNKL, tile_shape, cluster_shape);
}
static constexpr
dim3
static dim3
get_block_shape() {
return dim3(MaxThreadsPerBlock, 1, 1);
}
@@ -243,7 +240,7 @@ public:
int warp_idx = canonical_warp_idx();
int lane_predicate = cute::elect_one_sync();
// Issue Tma Descriptor Prefetch from a single thread
// Issue Tma Descriptor Prefetch from a single thread
if ((warp_idx == 0) && lane_predicate) {
CollectiveMainloop::prefetch_tma_descriptors(params.mainloop);
}
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp
+7 -9
View File
@@ -179,20 +179,20 @@ public:
auto N = get<1>(args.problem_shape);
auto K = get<2>(args.problem_shape);
// Contiguous dimension for the TMA tensor should be 128b aligned
implementable = std::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> ?
implementable = std::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> ?
K % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0;
implementable = implementable && (std::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor> ?
implementable = implementable && (std::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor> ?
N % min_tma_aligned_elements == 0 : K % min_tma_aligned_elements == 0);
implementable = implementable && (!cutlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA<CollectiveEpilogue>::value ||
(cutlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA<CollectiveEpilogue>::value &&
std::is_same_v<gemm::detail::StrideToLayoutTagC_t<StrideC>, layout::RowMajor> ?
std::is_same_v<gemm::detail::StrideToLayoutTagC_t<StrideC>, layout::RowMajor> ?
N % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0));
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
return implementable;
}
constexpr bool is_beta_supported =
constexpr bool is_beta_supported = not cute::is_void_v<ElementC> &&
CollectiveEpilogue::ThreadEpilogueOp::kScale == cutlass::epilogue::thread::ScaleType::Default;
implementable = is_beta_supported || (args.epilogue.thread.beta == 0 && args.epilogue.thread.beta_ptr == nullptr);
if (!implementable) {
@@ -210,8 +210,7 @@ public:
}
// Computes the kernel launch grid shape based on runtime parameters
static constexpr
dim3
static dim3
get_grid_shape(Params const& params) {
auto cluster_shape = ClusterShape{};
auto tile_shape = TileShape{};
@@ -220,8 +219,7 @@ public:
problem_shape_MNKL, tile_shape, cluster_shape);
}
static constexpr
dim3
static dim3
get_block_shape() {
return dim3(MaxThreadsPerBlock, 1, 1);
}
@@ -300,7 +298,7 @@ public:
typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state;
typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state;
// For the DMA Load (producer) we start with an opposite phase
// For the DMA Load (producer) we start with an opposite phase
// i.e., we skip all waits since we know that the buffer is indeed empty
PipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+7 -9
View File
@@ -202,20 +202,20 @@ public:
auto N = get<1>(args.problem_shape);
auto K = get<2>(args.problem_shape);
// Contiguous dimension for the TMA tensor should be 128b aligned
implementable = std::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> ?
implementable = std::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> ?
K % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0;
implementable = implementable && (std::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor> ?
implementable = implementable && (std::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor> ?
N % min_tma_aligned_elements == 0 : K % min_tma_aligned_elements == 0);
implementable = implementable && (!cutlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA<CollectiveEpilogue>::value ||
(cutlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA<CollectiveEpilogue>::value &&
std::is_same_v<gemm::detail::StrideToLayoutTagC_t<StrideC>, layout::RowMajor> ?
std::is_same_v<gemm::detail::StrideToLayoutTagC_t<StrideC>, layout::RowMajor> ?
N % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0));
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
return implementable;
}
constexpr bool is_beta_supported =
constexpr bool is_beta_supported =
CollectiveEpilogue::ThreadEpilogueOp::kScale == cutlass::epilogue::thread::ScaleType::Default;
implementable = is_beta_supported || (args.epilogue.thread.beta == 0 && args.epilogue.thread.beta_ptr == nullptr);
if (!implementable) {
@@ -233,15 +233,13 @@ public:
}
// Computes the kernel launch grid shape based on runtime parameters
static constexpr
dim3
static dim3
get_grid_shape(Params const& params) {
// Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently
return detail::PersistentTileSchedulerSm90::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info);
}
static constexpr
dim3
static dim3
get_block_shape() {
return dim3(MaxThreadsPerBlock, 1, 1);
}
@@ -333,7 +331,7 @@ public:
typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state;
typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state;
// For the DMA Load (producer) we start with an opposite phase
// For the DMA Load (producer) we start with an opposite phase
// i.e., we skip all waits since we know that the buffer is indeed empty
PipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
+12 -14
View File
@@ -110,7 +110,7 @@ public:
static constexpr uint32_t LoadRegisterRequirement = 40;
static constexpr uint32_t MmaRegisterRequirement = 232;
// Order Sequence barrier with two stages: one for Mainloop and one for Epilogue
// Order Sequence barrier with two stages: one for Mainloop and one for Epilogue
static constexpr uint32_t StagesPerMathWarpGroup = 2;
using MathWarpGroupOrderBarrier = cutlass::OrderedSequenceBarrier<
StagesPerMathWarpGroup, NumMmaWarpGroups>;
@@ -210,20 +210,20 @@ public:
auto N = get<1>(args.problem_shape);
auto K = get<2>(args.problem_shape);
// Contiguous dimension for the TMA tensor should be 128b aligned
implementable = std::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> ?
implementable = std::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> ?
K % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0;
implementable = implementable && (std::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor> ?
implementable = implementable && (std::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor> ?
N % min_tma_aligned_elements == 0 : K % min_tma_aligned_elements == 0);
implementable = implementable && (!cutlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA<CollectiveEpilogue>::value ||
(cutlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA<CollectiveEpilogue>::value &&
std::is_same_v<gemm::detail::StrideToLayoutTagC_t<StrideC>, layout::RowMajor> ?
std::is_same_v<gemm::detail::StrideToLayoutTagC_t<StrideC>, layout::RowMajor> ?
N % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0));
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
return implementable;
}
constexpr bool is_beta_supported =
constexpr bool is_beta_supported =
CollectiveEpilogue::ThreadEpilogueOp::kScale == cutlass::epilogue::thread::ScaleType::Default;
implementable = is_beta_supported || (args.epilogue.thread.beta == 0 && args.epilogue.thread.beta_ptr == nullptr);
if (!implementable) {
@@ -241,15 +241,13 @@ public:
}
// Computes the kernel launch grid shape based on runtime parameters
static constexpr
dim3
static dim3
get_grid_shape(Params const& params) {
// Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently
return detail::PersistentTileSchedulerSm90::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info);
}
static constexpr
dim3
static dim3
get_block_shape() {
return dim3(MaxThreadsPerBlock, 1, 1);
}
@@ -341,7 +339,7 @@ public:
typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state;
typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state;
// For the DMA Load (producer) we start with an opposite phase
// For the DMA Load (producer) we start with an opposite phase
// i.e., we skip all waits since we know that the buffer is indeed empty
PipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
@@ -389,9 +387,9 @@ public:
detail::PersistentTileSchedulerSm90 scheduler;
if (warp_group_role == WarpGroupRole::Consumer1) {
// Advance 2nd Math WG to the next work tile for the startup
// Advance 2nd Math WG to the next work tile for the startup
scheduler.advance_to_next_work();
// Advance 2nd Math WG pipeline states to the end of 1st Math WG
// Advance 2nd Math WG pipeline states to the end of 1st Math WG
mainloop_pipe_consumer_state.advance(k_tile_count);
epi_load_pipe_consumer_state.advance(c_tile_count);
epi_store_pipe_producer_state.advance(d_tile_count);
@@ -486,7 +484,7 @@ public:
params.mainloop
);
// Cue for next Math WG's MMA to start
// Cue for next Math WG's MMA to start
math_wg_order_barrier.arrive();
// Make sure the math instructions are done and free buffers before entering the epilogue
@@ -522,7 +520,7 @@ public:
// Wait for all TMA stores to complete
epi_store_pipeline.producer_tail(epi_store_pipe_producer_state);
// Cue for next Math WG's Epilogue to start
// Cue for next Math WG's Epilogue to start
math_wg_order_barrier.arrive();
// Get next work tile
include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp
+3 -3
View File
@@ -108,7 +108,7 @@ public:
return {work_idx_m, work_idx_n, static_cast<int32_t>(work_idx_l), current_work_linear_idx_ < scheduler_params.blocks_per_problem_};
}
CUTLASS_DEVICE
CUTLASS_DEVICE
void
advance_to_next_work(uint32_t advance_count = 1) {
current_work_linear_idx_ += grid_blocks_total_ * advance_count;
@@ -117,7 +117,7 @@ public:
// Given the inputs, computes the total number of output blocks this problem will compute over
// Note that this is only the logical size of our grid, not the physical grid we will actually launch.
template<class ProblemShapeMNKL, class BlockShape, class ClusterShape>
CUTLASS_HOST_DEVICE constexpr static
CUTLASS_HOST_DEVICE static
dim3
get_tiled_blk_shape_mnl(ProblemShapeMNKL problem_shape_mnkl, BlockShape blk_shape, ClusterShape cluster_shape) {
// Across M and N is our Cluster tile, so we must round up the blocks to the nearest whole number of Cluster tiles
@@ -135,7 +135,7 @@ public:
// Given the inputs, computes the physical grid we should launch.
template<class ProblemShapeMNKL, class BlockShape, class ClusterShape>
CUTLASS_HOST_DEVICE constexpr static
CUTLASS_HOST_DEVICE static
dim3
get_grid_shape(ProblemShapeMNKL problem_shape_mnk, BlockShape blk_shape, ClusterShape cluster_shape, KernelHardwareInfo hw_info) {
int const sm_count = hw_info.sm_count;
include/cutlass/gemm/threadblock/ell_mma_multistage.h
+6
View File
@@ -630,6 +630,12 @@ public:
accum = plus_accum(accum, tmp_accum);
}
// Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
};
include/cutlass/gemm/threadblock/mma_blas3_multistage.h
+5
View File
@@ -690,6 +690,11 @@ public:
__syncthreads();
}
// Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
};
include/cutlass/gemm/threadblock/mma_layernorm_mainloop_fusion_multistage.h
+4 -6
View File
@@ -846,12 +846,10 @@ public:
}
if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
// commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
// commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
};
include/cutlass/gemm/threadblock/mma_multistage.h
+4 -6
View File
@@ -660,12 +660,10 @@ public:
accum = plus_accum(accum, pipe_state.tmp_accum_);
}
// Optionally commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
// Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
include/cutlass/gemm/threadblock/mma_planar_complex_multistage.h
+6
View File
@@ -628,6 +628,12 @@ public:
}
// Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
};
include/cutlass/gemm/threadblock/mma_softmax_mainloop_fusion_multistage.h
+5
View File
@@ -739,6 +739,11 @@ public:
__syncthreads();
}
// Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
};
include/cutlass/gemm/threadblock/mma_sparse_multistage.h
+6
View File
@@ -650,6 +650,12 @@ public:
}
// Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
};
include/cutlass/gemm/threadblock/mma_with_reduction_multistage.h
+4 -6
View File
@@ -528,12 +528,10 @@ public:
}
if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
// commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
// commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
};