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

CUTLASS 3.1 (#915)

Browse Source 
Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2023-04-14 23:19:34 -04:00
committed by GitHub
co-authored by Aniket Shivam
parent 9b8166e3f0
commit d572cc1aab
482 changed files with 37175 additions and 16410 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/gemm/device/base_grouped.h
+1 -1
View File
@@ -157,7 +157,7 @@ private:
static void reorder_array(T* data, const std::vector<size_t>& indices) {
// For now, simply create a copy of the data and then copy over to the original.
std::vector<T> copy(indices.size());
for (unsigned i = 0; i < indices.size(); ++i) {
for (size_t i = 0; i < indices.size(); ++i) {
copy.at(i) = data[indices[i]];
}
include/cutlass/gemm/device/default_gemm_configuration.h
-3
View File
@@ -763,9 +763,6 @@ struct DefaultGemmConfiguration<
};
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
template <typename ElementC,
typename ElementAccumulator>
struct DefaultGemmConfiguration<arch::OpClassTensorOp, arch::Sm90, double,
include/cutlass/gemm/device/gemm_universal.h
+27 -6
View File
@@ -132,7 +132,11 @@ template <
/// Scatter result D by using an index array
bool ScatterD = false,
/// Permute result D
typename PermuteDLayout = layout::NoPermute
typename PermuteDLayout_ = layout::NoPermute,
/// Permute operand A
typename PermuteALayout_ = layout::NoPermute,
/// Permute operand B
typename PermuteBLayout_ = layout::NoPermute
>
class GemmUniversal :
public GemmUniversalBase<
@@ -161,7 +165,9 @@ class GemmUniversal :
GatherA,
GatherB,
ScatterD,
PermuteDLayout
PermuteDLayout_,
PermuteALayout_,
PermuteBLayout_
>::GemmKernel
> {
@@ -176,6 +182,9 @@ class GemmUniversal :
using EpilogueOutputOp = EpilogueOutputOp_;
using ThreadblockSwizzle = ThreadblockSwizzle_;
using Operator = Operator_;
using PermuteDLayout = PermuteDLayout_;
using PermuteALayout = PermuteALayout_;
using PermuteBLayout = PermuteBLayout_;
static int const kStages = Stages;
static int const kAlignmentA = AlignmentA;
static int const kAlignmentB = AlignmentB;
@@ -209,7 +218,9 @@ class GemmUniversal :
GatherA,
GatherB,
ScatterD,
PermuteDLayout
PermuteDLayout_,
PermuteALayout_,
PermuteBLayout_
>::GemmKernel
>;
@@ -268,14 +279,19 @@ template <
/// Scatter result D by using an index array
bool ScatterD,
/// Permute result D
typename PermuteDLayout
typename PermuteDLayout_,
/// Permute operand A
typename PermuteALayout_,
/// Permute operand B
typename PermuteBLayout_
>
class GemmUniversal<ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_,
layout::ColumnMajor, // partially specialized on LayoutC
ElementAccumulator_, OperatorClass_, ArchTag_, ThreadblockShape_,
WarpShape_, InstructionShape_, EpilogueOutputOp_,
ThreadblockSwizzle_, Stages, AlignmentA, AlignmentB,
Operator_, TransformA, TransformB, GatherA, GatherB, ScatterD, PermuteDLayout> {
Operator_, TransformA, TransformB, GatherA, GatherB, ScatterD,
PermuteDLayout_, PermuteALayout_, PermuteBLayout_> {
public:
using ElementA = ElementA_;
@@ -297,6 +313,9 @@ class GemmUniversal<ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_,
using EpilogueOutputOp = EpilogueOutputOp_;
using ThreadblockSwizzle = ThreadblockSwizzle_;
using Operator = Operator_;
using PermuteDLayout = PermuteDLayout_;
using PermuteALayout = PermuteALayout_;
using PermuteBLayout = PermuteBLayout_;
static int const kStages = Stages;
static int const kAlignmentA = AlignmentA;
static int const kAlignmentB = AlignmentB;
@@ -327,7 +346,9 @@ class GemmUniversal<ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_,
GatherB,
GatherA,
ScatterD,
PermuteDLayout
PermuteDLayout,
PermuteBLayout,
PermuteALayout
>::Base;
using GemmKernel = typename UnderlyingOperator::GemmKernel;
include/cutlass/gemm/device/gemm_universal_adapter.h
+15 -10
View File
@@ -80,7 +80,7 @@ class GemmUniversalAdapter;
template <class GemmKernel_>
class GemmUniversalAdapter<
GemmKernel_,
std::enable_if_t<gemm::detail::IsCutlass3GemmKernel<GemmKernel_>::value>>
cute::enable_if_t<gemm::detail::IsCutlass3GemmKernel<GemmKernel_>::value>>
{
public:
using GemmKernel = GemmKernel_;
@@ -88,6 +88,7 @@ public:
using ElementA = typename GemmKernel::ElementA;
using ElementB = typename GemmKernel::ElementB;
using ElementC = typename GemmKernel::ElementC;
using ElementD = typename GemmKernel::ElementD;
using ElementAccumulator = typename GemmKernel::TiledMma::ValTypeC;
using DispatchPolicy = typename GemmKernel::DispatchPolicy;
using CollectiveMainloop = typename GemmKernel::CollectiveMainloop;
@@ -107,14 +108,14 @@ public:
using MathOperator = cutlass::arch::OpMultiplyAdd;
// If our TiledMMA's instruction thread layout size is larger than 1, we know its a tensorop!
using OperatorClass = std::conditional_t<
using OperatorClass = cute::conditional_t<
(cute::size(typename GemmKernel::TiledMma::AtomThrID{}) > 1),
cutlass::arch::OpClassTensorOp, cutlass::arch::OpClassSimt>;
using ArchTag = typename GemmKernel::ArchTag;
// NOTE: Assume identity swizzle for now
static_assert(std::is_void_v<typename GemmKernel::GridSwizzle>,
static_assert(cute::is_void_v<typename GemmKernel::GridSwizzle>,
"CUTLASS 3.x kernel types do not support grid swizzle functors yet.");
using ThreadblockSwizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
@@ -142,7 +143,7 @@ public:
// But we can best approximate it by inspecting the TiledMma::TiledShape_MNK
// For this, we make the assumption that we always have 4 warps along M, and rest along N, none along K
// We also always round up the warp count to 4 if the tiled mma is smaller than 128 threads
static constexpr int WarpsInMma = std::max(4, cute::size(typename GemmKernel::TiledMma{}) / 32);
static constexpr int WarpsInMma = cute::max(4, cute::size(typename GemmKernel::TiledMma{}) / 32);
static constexpr int WarpsInMmaM = 4;
static constexpr int WarpsInMmaN = cute::ceil_div(WarpsInMma, WarpsInMmaM);
using WarpCount = cutlass::gemm::GemmShape<WarpsInMmaM, WarpsInMmaN, 1>;
@@ -166,7 +167,7 @@ public:
using EpilogueOutputOp = typename CollectiveEpilogue::ThreadEpilogueOp;
// Split-K preserves splits that are 128b aligned
static int constexpr kSplitKAlignment = std::max(
static int constexpr kSplitKAlignment = cute::max(
128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value);
/// Argument structure: User API
@@ -208,8 +209,8 @@ public:
/// Computes the grid shape
static dim3
get_grid_shape(Arguments const& args) {
auto tmp_params = GemmKernel::to_underlying_arguments(args);
get_grid_shape(Arguments const& args, void* workspace = nullptr) {
auto tmp_params = GemmKernel::to_underlying_arguments(args, workspace);
return GemmKernel::get_grid_shape(tmp_params);
}
@@ -397,14 +398,14 @@ public:
template <typename GemmKernel_>
class GemmUniversalAdapter<
GemmKernel_,
std::enable_if_t<not gemm::detail::IsCutlass3GemmKernel<GemmKernel_>::value>>
cute::enable_if_t<not gemm::detail::IsCutlass3GemmKernel<GemmKernel_>::value>>
{
public:
using GemmKernel = GemmKernel_;
static bool const kInternalTranspose =
platform::is_same<typename GemmKernel::LayoutC, cutlass::layout::RowMajor>::value;
cute::is_same<typename GemmKernel::LayoutC, cutlass::layout::RowMajor>::value;
using ThreadblockShape = typename GemmKernel::Mma::Shape;
using WarpShape = typename GemmKernel::WarpShape;
@@ -447,11 +448,15 @@ public:
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>;
using TensorRefD = TensorRef<ElementC, LayoutC>;
using TensorRefD = TensorRef<ElementD, LayoutD>;
static int const kStages = GemmKernel::Mma::kStages;
include/cutlass/gemm/device/gemm_universal_base.h
+4
View File
@@ -36,7 +36,11 @@
#pragma once
#if defined(__CUDACC_RTC__)
#include <cuda/std/limits>
#else
#include <limits>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"