leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 25 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
tools/library/include/cutlass/library/arch_mappings.h
+6
View File
@@ -102,6 +102,12 @@ template <typename OperatorClass> struct ArchMap<arch::Sm90, OperatorClass> {
static int const kMax = 1024;
};
// Arch conditional WGMMA
template <> struct ArchMap<arch::Sm90, arch::OpClassTensorOp> {
static int const kMin = 90;
static int const kMax = 90;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace library
tools/library/include/cutlass/library/handle.h
+11 -11
View File
@@ -178,7 +178,7 @@ public:
int K, /// GEMM K dimension
NumericTypeID element_compute, /// Data type of internal accumulation
NumericTypeID element_scalar, /// Data type of alpha/beta scalars
void const *alpha, /// Pointer to alpha scalar
@@ -186,29 +186,29 @@ public:
NumericTypeID element_A, /// Data type of A matrix elements
LayoutTypeID layout_A, /// Layout of A matrix
ComplexTransform transform_A, /// Complex transformation applied to A matrix - ignored for real-valued matrices
void const * ptr_A, /// Pointer to A matrix in Global Memory
int64_t lda, /// Leading dimension of A matrix
int64_t lda, /// Leading dimension of A matrix
NumericTypeID element_B, /// Data type of B matrix elements
LayoutTypeID layout_B, /// Layout of B matrix
ComplexTransform transform_B, /// Complex transformation applied to B matrix - ignored for real-valued matrices
void const * ptr_B, /// Pointer to B matrix in Global Memory
int64_t ldb, /// Leading dimension of B matrix
int64_t ldb, /// Leading dimension of B matrix
void const * beta, /// Pointer to beta scalar
NumericTypeID element_C, /// Data type of C and D matrices
NumericTypeID element_C, /// Data type of C matrix
LayoutTypeID layout_C, /// Layout of D matrix
void const * ptr_C, /// Pointer to C matrix
int64_t ldc, /// Leading dimension of C matrix
int64_t ldc, /// Leading dimension of C matrix
NumericTypeID element_D, /// Data type of D matrix
LayoutTypeID layout_D, /// Layout of D matrix
void * ptr_D, /// Pointer to D matrix
int64_t ldd, /// Leading dimension of D matrix
int64_t ldd, /// Leading dimension of D matrix
int batch_count = 1, /// Batch count or number of split-K slices
int64_t batch_stride_A = 0, /// Batch stride of A operand
int64_t batch_stride_B = 0, /// Batch stride of B operand
int64_t batch_stride_C = 0, /// Batch stride of C operand
tools/library/include/cutlass/library/library.h
+13 -2
View File
@@ -114,6 +114,8 @@ enum class NumericTypeID {
kS16,
kS32,
kS64,
kFE4M3,
kFE5M2,
kF16,
kBF16,
kTF32,
@@ -474,9 +476,12 @@ struct GemmDescription : public OperationDescription {
/// Describes the B operand
TensorDescription B;
/// Describes the source and destination matrices
/// Describes the source matrix
TensorDescription C;
/// Describes the destination matrix
TensorDescription D;
/// Describes the sparse meta matrices
TensorDescription E;
@@ -501,6 +506,7 @@ struct GemmDescription : public OperationDescription {
TensorDescription const &A = TensorDescription(),
TensorDescription const &B = TensorDescription(),
TensorDescription const &C = TensorDescription(),
TensorDescription const &D = TensorDescription(),
NumericTypeID element_epilogue = NumericTypeID::kInvalid,
SplitKMode split_k_mode = SplitKMode::kNone,
ComplexTransform transform_A = ComplexTransform::kNone,
@@ -510,6 +516,7 @@ struct GemmDescription : public OperationDescription {
A(A),
B(B),
C(C),
D(D),
element_epilogue(element_epilogue),
split_k_mode(split_k_mode),
transform_A(transform_A),
@@ -527,13 +534,14 @@ struct SparseGemmDescription : public GemmDescription {
TensorDescription const &A = TensorDescription(),
TensorDescription const &B = TensorDescription(),
TensorDescription const &C = TensorDescription(),
TensorDescription const &D = TensorDescription(),
TensorDescription const &E = TensorDescription(),
NumericTypeID element_epilogue = NumericTypeID::kInvalid,
SplitKMode split_k_mode = SplitKMode::kNone,
ComplexTransform transform_A = ComplexTransform::kNone,
ComplexTransform transform_B = ComplexTransform::kNone
):
GemmDescription(gemm_kind, A, B, C, element_epilogue, split_k_mode, transform_A, transform_B)
GemmDescription(gemm_kind, A, B, C, D, element_epilogue, split_k_mode, transform_A, transform_B)
{this->E = E;}
};
@@ -1019,6 +1027,9 @@ struct GemmUniversalArguments {
int64_t batch_stride_B;
int64_t batch_stride_C;
int64_t batch_stride_D;
// Needed for some 3.x kernels
int sm_count;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/include/cutlass/library/operation_table.h
+32 -14
View File
@@ -66,6 +66,9 @@ struct GemmFunctionalKey {
LayoutTypeID layout_B;
ComplexTransform transform_B;
NumericTypeID element_C;
LayoutTypeID layout_C;
NumericTypeID element_D;
LayoutTypeID layout_D;
//
// Methods
@@ -83,7 +86,10 @@ struct GemmFunctionalKey {
NumericTypeID element_B = NumericTypeID::kF16,
LayoutTypeID layout_B = LayoutTypeID::kColumnMajor,
ComplexTransform transform_B = ComplexTransform::kNone,
NumericTypeID element_C = NumericTypeID::kF16
NumericTypeID element_C = NumericTypeID::kF16,
LayoutTypeID layout_C = LayoutTypeID::kColumnMajor,
NumericTypeID element_D = NumericTypeID::kF16,
LayoutTypeID layout_D = LayoutTypeID::kColumnMajor
):
provider(provider),
gemm_kind(gemm_kind),
@@ -95,7 +101,10 @@ struct GemmFunctionalKey {
element_B(element_B),
layout_B(layout_B),
transform_B(transform_B),
element_C(element_C)
element_C(element_C),
layout_C(layout_C),
element_D(element_D),
layout_D(layout_D)
{ }
inline
@@ -111,7 +120,10 @@ struct GemmFunctionalKey {
(element_B == rhs.element_B) &&
(layout_B == rhs.layout_B) &&
(transform_B == rhs.transform_B) &&
(element_C == rhs.element_C);
(element_C == rhs.element_C) &&
(layout_C == rhs.layout_C) &&
(element_D == rhs.element_D) &&
(layout_D == rhs.layout_D);
}
inline
@@ -137,6 +149,9 @@ std::ostream & operator<<(std::ostream &out, cutlass::library::GemmFunctionalKey
<< " layout_B: " << to_string(k.layout_B) << "\n"
<< " transform_B: " << to_string(k.transform_B) << "\n"
<< " element_C: " << to_string(k.element_C) << "\n"
<< " layout_C: " << to_string(k.layout_C) << "\n"
<< " element_D: " << to_string(k.element_D) << "\n"
<< " layout_D: " << to_string(k.layout_D) << "\n"
<< "}";
return out;
@@ -157,18 +172,21 @@ struct GemmFunctionalKeyHasher {
size_t operator()(GemmFunctionalKey const &key) const {
IntHash hash;
return
rotl(hash(int(key.provider)), 1) ^
rotl(hash(int(key.gemm_kind)), 2) ^
return
rotl(hash(int(key.provider)), 1) ^
rotl(hash(int(key.gemm_kind)), 2) ^
rotl(hash(int(key.element_compute)), 3) ^
rotl(hash(int(key.element_scalar)), 4) ^
rotl(hash(int(key.element_A)), 5) ^
rotl(hash(int(key.layout_A)), 6) ^
rotl(hash(int(key.transform_A)), 7) ^
rotl(hash(int(key.element_B)), 8) ^
rotl(hash(int(key.layout_B)), 9) ^
rotl(hash(int(key.transform_B)), 10) ^
rotl(hash(int(key.element_C)), 11);
rotl(hash(int(key.element_scalar)), 4) ^
rotl(hash(int(key.element_A)), 5) ^
rotl(hash(int(key.layout_A)), 6) ^
rotl(hash(int(key.transform_A)), 7) ^
rotl(hash(int(key.element_B)), 8) ^
rotl(hash(int(key.layout_B)), 9) ^
rotl(hash(int(key.transform_B)), 10) ^
rotl(hash(int(key.element_C)), 11) ^
rotl(hash(int(key.layout_C)), 12) ^
rotl(hash(int(key.element_D)), 13) ^
rotl(hash(int(key.layout_D)), 14);
}
};