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

Updates for 3.1 (#932)

Browse Source
This commit is contained in:
ANIKET SHIVAM
2023-04-29 09:34:27 -04:00
committed by GitHub
parent 6f8596ce3f
commit 7c04f95415
51 changed files with 1796 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp
+3 -17
View File
@@ -190,7 +190,8 @@ struct CollectiveMma<
"SmemLayoutB K must be 128bytes to be transposed.");
static_assert(!transform::collective::detail::use_universal_transposition<InternalSmemLayoutAtomB, InternalElementB>(),
"Warp specialized ARF kernels have not supported universal B transposition yet.");
static_assert(!TransposeB || shape<0>(TileShape{}) == 64, "Optimized transpose RS kernel requires TileShape M = 64.");
static_assert(!TransposeB || !cute::is_same_v<KernelSchedule, KernelTmaWarpSpecializedCooperative>,
"Transpose RS kernel requires kernel schedule schmem is not KernelTmaWarpSpecializedCooperative.");
struct SharedStorage
{
@@ -294,7 +295,7 @@ struct CollectiveMma<
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
static constexpr int K_PIPE_MMAS = DispatchPolicy::PipelineAsyncMmaStages;
static_assert(K_PIPE_MMAS >= 1, "At least one MMA stage should be asynchronous for this mainloop.");
static_assert(K_PIPE_MMAS == 0, "no MMA stage should be asynchronous for this mainloop for now.");
static constexpr uint32_t TmaTransactionBytes =
(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof(InternalElementA)))+
(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof(InternalElementB)));
@@ -368,21 +369,6 @@ struct CollectiveMma<
}
}
// Issue the prologue loads
int k_tile_prologue = min(k_tile_count, K_PIPE_MAX);
CUTLASS_PRAGMA_UNROLL
for (int count = 0; count < k_tile_prologue; ++count) {
pipeline.producer_acquire(smem_pipe_write);
using BarrierType = typename MainloopPipeline::ProducerBarrierType;
BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write);
int write_stage = smem_pipe_write.index();
copy(tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
copy(tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
++k_tile_iter;
++smem_pipe_write;
}
k_tile_count -= k_tile_prologue;
// Mainloop
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp
-16
View File
@@ -303,22 +303,6 @@ struct CollectiveMma<
}
}
// Issue the prologue loads
int k_tile_prologue = min(k_tile_count, K_PIPE_MAX);
CUTLASS_PRAGMA_UNROLL
for (int count = 0; count < k_tile_prologue; ++count) {
pipeline.producer_acquire(smem_pipe_write);
using BarrierType = typename MainloopPipeline::ProducerBarrierType;
BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write);
int write_stage = smem_pipe_write.index();
copy(tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
copy(tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
++k_tile_iter;
++smem_pipe_write;
}
k_tile_count -= k_tile_prologue;
// Mainloop
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count)
include/cutlass/gemm/device/base_grouped.h
+8 -9
View File
@@ -301,18 +301,19 @@ public:
return 0;
}
result = cudaGetDeviceProperties(&properties, device_idx);
int multiprocessor_count;
result = cudaDeviceGetAttribute(&multiprocessor_count,
cudaDevAttrMultiProcessorCount, device_idx);
if (result != cudaSuccess) {
// Call cudaGetLastError() to clear the error bit
result = cudaGetLastError();
CUTLASS_TRACE_HOST(" cudaGetDeviceProperties() returned error "
<< cudaGetErrorString(result));
CUTLASS_TRACE_HOST(
" cudaDeviceGetAttribute() returned error "
<< cudaGetErrorString(result));
return 0;
}
bool override_sm_count = (available_sm_count < 0 || available_sm_count > properties.multiProcessorCount);
bool override_sm_count = (available_sm_count < 0 || available_sm_count > multiprocessor_count);
if (override_sm_count) {
available_sm_count = properties.multiProcessorCount;
available_sm_count = multiprocessor_count;
}
int max_active_blocks = maximum_active_blocks();
@@ -440,8 +441,6 @@ public:
cudaError_t result = cudaGetLastError();
if (result != cudaSuccess) {
// Call cudaGetLastError() to clear the error bit
result = cudaGetLastError();
CUTLASS_TRACE_HOST(" grid launch failed with error " << cudaGetErrorString(result));
return Status::kErrorInternal;
}
include/cutlass/gemm/device/default_gemm_configuration.h
+3 -3
View File
@@ -490,9 +490,9 @@ struct DefaultGemmConfiguration<arch::OpClassTensorOp, arch::Sm80, double,
static int const kAlignmentA = 1;
static int const kAlignmentB = 1;
using ThreadblockShape = GemmShape<128, 256, 64>;
using WarpShape = GemmShape<64, 64, 64>;
using InstructionShape = GemmShape<16, 8, 16>;
using ThreadblockShape = GemmShape<128, 128, 16>;
using WarpShape = GemmShape<32, 64, 16>;
using InstructionShape = GemmShape<8, 8, 4>;
static int const kStages = 3;
using EpilogueOutputOp = epilogue::thread::LinearCombination<
include/cutlass/gemm/device/gemm_universal_with_broadcast.h
+1 -1
View File
@@ -31,7 +31,7 @@
/*! \file
\brief Template for a GEMM kernel that can broadcast bias vector in the
epigloue.
epilogue.
*/
#pragma once
include/cutlass/gemm/device/gemv_strided_batched.h
+167
View File
@@ -0,0 +1,167 @@
/***************************************************************************************************
* 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/dispatch_policy.hpp
+1 -1
View File
@@ -150,7 +150,7 @@ template<
int Stages_,
class ClusterShape_ = Shape<_1,_1,_1>,
class KernelSchedule = KernelTmaWarpSpecialized,
int PipelineAsyncMmaStages_ = 1
int PipelineAsyncMmaStages_ = 0
>
struct MainloopSm90TmaGmmaRmemAWarpSpecialized {
constexpr static int Stages = Stages_;
include/cutlass/gemm/kernel/gemv_strided_batched.h
+368
View File
@@ -0,0 +1,368 @@
/***************************************************************************************************
* 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/sm90_gemm_tma.hpp
+12 -2
View File
@@ -166,8 +166,8 @@ public:
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Size don't meet the requirements.\n");
return implementable;
}
static constexpr int tma_alignment_bits = 128;
static constexpr int min_tma_aligned_elements = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
constexpr int tma_alignment_bits = 128;
constexpr int min_tma_aligned_elements = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
auto M = get<0>(args.problem_shape);
auto N = get<1>(args.problem_shape);
auto K = get<2>(args.problem_shape);
@@ -182,7 +182,17 @@ public:
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 =
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) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Scaling params don't meet ThreadEpilogueOp requirements.\n");
return implementable;
}
return implementable;
}
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp
+12 -2
View File
@@ -173,8 +173,8 @@ public:
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Size don't meet the requirements.\n");
return implementable;
}
static constexpr int tma_alignment_bits = 128;
static constexpr int min_tma_aligned_elements = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
constexpr int tma_alignment_bits = 128;
constexpr int min_tma_aligned_elements = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
auto M = get<0>(args.problem_shape);
auto N = get<1>(args.problem_shape);
auto K = get<2>(args.problem_shape);
@@ -189,7 +189,17 @@ public:
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 =
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) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Scaling params don't meet ThreadEpilogueOp requirements.\n");
return implementable;
}
return implementable;
}
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+12 -2
View File
@@ -196,8 +196,8 @@ public:
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Size don't meet the requirements.\n");
return implementable;
}
static constexpr int tma_alignment_bits = 128;
static constexpr int min_tma_aligned_elements = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
constexpr int tma_alignment_bits = 128;
constexpr int min_tma_aligned_elements = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
auto M = get<0>(args.problem_shape);
auto N = get<1>(args.problem_shape);
auto K = get<2>(args.problem_shape);
@@ -212,7 +212,17 @@ public:
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 =
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) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Scaling params don't meet ThreadEpilogueOp requirements.\n");
return implementable;
}
return implementable;
}
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
+12 -2
View File
@@ -204,8 +204,8 @@ public:
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Size don't meet the requirements.\n");
return implementable;
}
static constexpr int tma_alignment_bits = 128;
static constexpr int min_tma_aligned_elements = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
constexpr int tma_alignment_bits = 128;
constexpr int min_tma_aligned_elements = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
auto M = get<0>(args.problem_shape);
auto N = get<1>(args.problem_shape);
auto K = get<2>(args.problem_shape);
@@ -220,7 +220,17 @@ public:
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 =
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) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Scaling params don't meet ThreadEpilogueOp requirements.\n");
return implementable;
}
return implementable;
}
include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp
+6
View File
@@ -163,6 +163,12 @@ public:
int const min_num_gpc = sm_count < max_sm_per_gpc ? 1 : sm_count / max_sm_per_gpc;
int const max_blk_occupancy_per_gpc = max_sm_per_gpc - (max_sm_per_gpc % size(cluster_shape));
int blk_per_device = min_num_gpc * max_blk_occupancy_per_gpc;
// The calculation below allows for larger grid size launch for different GPUs.
int const num_gpc_residual = sm_count < max_sm_per_gpc ? 0 : sm_count % max_sm_per_gpc;
int const max_blk_occupancy_per_residual_gpc = num_gpc_residual - (num_gpc_residual % size(cluster_shape));
blk_per_device += max_blk_occupancy_per_residual_gpc;
blk_per_device = sm_count < blk_per_device ? sm_count : blk_per_device;
launch_grid.x = std::min(
include/cutlass/gemm/threadblock/threadblock_swizzle_streamk.h
+2 -7
View File
@@ -630,9 +630,6 @@ struct ThreadblockSwizzleStreamK {
}
// Guards needed for PyCUTLASS library generation
#if !defined(CUTLASS_PYTHON_HOST_CC)
//
// Device-side interface
//
@@ -692,7 +689,7 @@ struct ThreadblockSwizzleStreamK {
return GemmCoord(m, n, get_batch_idx());
}
/// Obtains the calling threadblock's tiled coordinates for the given tile index (row-major rastorization)
/// Obtains the calling threadblock's tiled coordinates for the given tile index (row-major rasterization)
CUTLASS_DEVICE
GemmCoord get_tile_offset_row_major(int tile_idx) const
{
@@ -740,7 +737,7 @@ struct ThreadblockSwizzleStreamK {
div_mod_sk_iters_per_region(region_idx, iter_in_region, iter);
int big_block_iters = (sk_big_blocks_per_region * sk_iters_per_normal_block()) + sk_big_blocks_per_region; // number of iterations in the region's big blocks
int normal_block_iters = iter_in_region - big_block_iters; // number of iterations in the region's normal bocks
int normal_block_iters = iter_in_region - big_block_iters; // number of iterations in the region's normal blocks
int big_block_idx_in_region = div_mod_sk_iters_per_big_block.div(iter_in_region);
int normal_block_idx_in_region = sk_big_blocks_per_region + div_mod_sk_iters_per_normal_block.div(normal_block_iters);
@@ -794,8 +791,6 @@ struct ThreadblockSwizzleStreamK {
return get_sk_block_idx(iter);
}
#endif // !defined(CUTLASS_PYTHON_HOST_CC)
};
/////////////////////////////////////////////////////////////////////////////////////////////////