leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 27 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/src/gemm_operation.h
+18 -3
View File
@@ -64,6 +64,8 @@ public:
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementD = ElementC;
using LayoutD = LayoutC;
// assuming all tensors use same type for StrideIndex
using StrideIndex = typename Operator::LayoutA::Index;
using ElementAccumulator = typename Operator::ElementAccumulator;
@@ -121,6 +123,7 @@ public:
description_.A = make_TensorDescription<ElementA, LayoutA>(Operator::kAlignmentA);
description_.B = make_TensorDescription<ElementB, LayoutB>(Operator::kAlignmentB);
description_.C = make_TensorDescription<ElementC, LayoutC>(Operator::kAlignmentC);
description_.D = make_TensorDescription<ElementD, LayoutD>(Operator::kAlignmentC);
description_.element_epilogue = NumericTypeMap<ElementCompute>::kId;
description_.split_k_mode = SplitKMode::kNone;
@@ -147,6 +150,8 @@ public:
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementD = ElementC;
using LayoutD = LayoutC;
using ElementAccumulator = typename Operator::ElementAccumulator;
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
using OperatorArguments = typename Operator::Arguments;
@@ -204,7 +209,7 @@ protected:
operator_args.ref_A.reset(static_cast<ElementA const *>(arguments->A));
operator_args.ref_B.reset(static_cast<ElementB const *>(arguments->B));
operator_args.ref_C.reset(static_cast<ElementC const *>(arguments->C));
operator_args.ref_D.reset(static_cast<ElementC *>(arguments->D));
operator_args.ref_D.reset(static_cast<ElementD *>(arguments->D));
return Status::kSuccess;
}
@@ -345,6 +350,8 @@ public:
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementD = ElementC;
using LayoutD = LayoutC;
using ElementE = typename Operator::ElementE;
using LayoutE = typename Operator::LayoutE;
using ElementAccumulator = typename Operator::ElementAccumulator;
@@ -405,7 +412,7 @@ protected:
operator_args.ref_A.reset(static_cast<ElementA const *>(arguments->A));
operator_args.ref_B.reset(static_cast<ElementB const *>(arguments->B));
operator_args.ref_C.reset(static_cast<ElementC const *>(arguments->C));
operator_args.ref_D.reset(static_cast<ElementC *>(arguments->D));
operator_args.ref_D.reset(static_cast<ElementD *>(arguments->D));
operator_args.ref_E.reset(static_cast<ElementE const *>(arguments->E));
return Status::kSuccess;
@@ -547,6 +554,8 @@ public:
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementD = ElementC;
using LayoutD = LayoutC;
using ElementAccumulator = typename Operator::ElementAccumulator;
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
@@ -751,6 +760,8 @@ public:
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementD = ElementC;
using LayoutD = LayoutC;
using ElementAccumulator = typename Operator::ElementAccumulator;
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
@@ -958,6 +969,8 @@ public:
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementD = ElementC;
using LayoutD = LayoutC;
using ElementAccumulator = typename Operator::ElementAccumulator;
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
@@ -1159,6 +1172,8 @@ public:
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementD = ElementC;
using LayoutD = LayoutC;
using ElementAccumulator = typename Operator::ElementAccumulator;
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
@@ -1218,7 +1233,7 @@ protected:
op_args.ptr_A = static_cast<ElementA **>(arguments->ptr_A);
op_args.ptr_B = static_cast<ElementB **>(arguments->ptr_B);
op_args.ptr_C = static_cast<ElementC **>(arguments->ptr_C);
op_args.ptr_D = static_cast<ElementC **>(arguments->ptr_D);
op_args.ptr_D = static_cast<ElementD **>(arguments->ptr_D);
op_args.lda = arguments->lda;
op_args.ldb = arguments->ldb;
tools/library/src/gemm_operation_3x.hpp
+27 -14
View File
@@ -34,7 +34,6 @@
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/kernel_hardware_info.hpp"
#include "cutlass/library/library.h"
#include "library_internal.h"
@@ -56,6 +55,8 @@ public:
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementD = typename Operator::ElementD;
using LayoutD = typename Operator::LayoutD;
// assuming all tensors use same type for StrideIndex
using StrideIndex = typename Operator::LayoutA::Index;
using ElementAccumulator = typename Operator::ElementAccumulator;
@@ -117,6 +118,7 @@ public:
description_.A = make_TensorDescription<ElementA, LayoutA>(Operator::kAlignmentA);
description_.B = make_TensorDescription<ElementB, LayoutB>(Operator::kAlignmentB);
description_.C = make_TensorDescription<ElementC, LayoutC>(Operator::kAlignmentC);
description_.D = make_TensorDescription<ElementD, LayoutD>(Operator::kAlignmentD);
description_.element_epilogue = NumericTypeMap<ElementCompute>::kId;
description_.split_k_mode = SplitKMode::kNone;
@@ -144,6 +146,8 @@ public:
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementD = typename Operator::ElementD;
using LayoutD = typename Operator::LayoutD;
using ElementAccumulator = typename Operator::ElementAccumulator;
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
@@ -167,9 +171,6 @@ protected:
// Do nothing here and construct kernel arguments in update_arguments_ instead
// We also cannot construct TMA descriptors without all the arguments available
if (operator_args.hw_info.sm_count <= 0) {
operator_args.hw_info.sm_count = KernelHardwareInfo::query_device_multiprocessor_count();
}
operator_args.mode = configuration->mode;
return Status::kSuccess;
}
@@ -181,13 +182,13 @@ protected:
typename ThreadEpilogueOp::Params params(
*static_cast<ElementCompute const *>(arguments->alpha),
*static_cast<ElementCompute const *>(arguments->beta));
operator_args.epilogue_params.thread_params = params;
operator_args.epilogue.thread = params;
}
else if (arguments->pointer_mode == ScalarPointerMode::kDevice) {
typename ThreadEpilogueOp::Params params(
static_cast<ElementCompute const *>(arguments->alpha),
static_cast<ElementCompute const *>(arguments->beta));
operator_args.epilogue_params.thread_params = params;
operator_args.epilogue.thread = params;
}
else {
return Status::kErrorInvalidProblem;
@@ -201,18 +202,21 @@ protected:
arguments->batch_count);
// update arguments
operator_args.ptr_A = static_cast<ElementA const *>(arguments->A);
operator_args.ptr_B = static_cast<ElementB const *>(arguments->B);
operator_args.epilogue_params.ptr_C = static_cast<ElementC const *>(arguments->C);
operator_args.epilogue_params.ptr_D = static_cast<ElementC *>(arguments->D);
operator_args.mainloop.ptr_A = static_cast<ElementA const *>(arguments->A);
operator_args.mainloop.ptr_B = static_cast<ElementB const *>(arguments->B);
operator_args.epilogue.ptr_C = static_cast<ElementC const *>(arguments->C);
operator_args.epilogue.ptr_D = static_cast<ElementD *>(arguments->D);
operator_args.dA = cute::make_int_tuple_from<typename Operator::GemmKernel::StrideA>(
operator_args.mainloop.dA = cute::make_int_tuple_from<typename Operator::GemmKernel::StrideA>(
arguments->lda, arguments->batch_stride_A);
operator_args.dB = cute::make_int_tuple_from<typename Operator::GemmKernel::StrideB>(
operator_args.mainloop.dB = cute::make_int_tuple_from<typename Operator::GemmKernel::StrideB>(
arguments->ldb, arguments->batch_stride_B);
operator_args.epilogue_params.dC = cute::make_int_tuple_from<typename Operator::GemmKernel::StrideC>(
operator_args.epilogue.dC = cute::make_int_tuple_from<typename Operator::GemmKernel::StrideC>(
arguments->ldc, arguments->batch_stride_C);
operator_args.epilogue_params.dD = operator_args.epilogue_params.dC;
operator_args.epilogue.dD = operator_args.epilogue.dC;
/* Query device SM count to pass onto the kernel as an argument, where needed */
operator_args.hw_info.sm_count = arguments->sm_count;
return Status::kSuccess;
}
@@ -223,6 +227,8 @@ public:
Status can_implement(
void const *configuration_ptr, void const *arguments_ptr) const override {
GemmUniversalConfiguration const *configuration =
static_cast<GemmUniversalConfiguration const *>(configuration_ptr);
GemmUniversalArguments const *arguments =
static_cast<GemmUniversalArguments const *>(arguments_ptr);
@@ -232,6 +238,13 @@ public:
return status;
}
// can_implement rules may need access to problem shape
args.problem_shape = cute::make_shape(
configuration->problem_size.m(),
configuration->problem_size.n(),
configuration->problem_size.k(),
configuration->batch_count);
return Operator::can_implement(args);
}
tools/library/src/handle.cu
+29 -14
View File
@@ -379,7 +379,10 @@ Status Handle::gemm(
element_B,
layout_B,
transform_B,
element_C
element_C, // C/D are same type and col major default
LayoutTypeID::kColumnMajor,
element_C,
LayoutTypeID::kColumnMajor
);
auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
@@ -498,26 +501,26 @@ Status Handle::gemm_universal(
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, /// Batch count or number of split-K slices
@@ -542,7 +545,10 @@ Status Handle::gemm_universal(
element_B,
layout_B,
transform_B,
element_C
element_C,
layout_C,
element_D,
layout_D
);
auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
@@ -741,7 +747,10 @@ Status Handle::gemm_planar_complex(
element_B,
layout_B,
transform_B,
element_C
element_C, // C/D are same type
LayoutTypeID::kColumnMajor,
element_C,
LayoutTypeID::kColumnMajor
);
auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
@@ -935,7 +944,10 @@ Status Handle::gemm_planar_complex_array(
element_B,
layout_B,
transform_B,
element_C
element_C, // C/D are same type
LayoutTypeID::kColumnMajor,
element_C,
LayoutTypeID::kColumnMajor
);
auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
@@ -1121,7 +1133,7 @@ Operation const* find_gemm_operation_for_parallel_reduction(Operation const *ope
static_cast<GemmDescription const &>(operation->description());
// if the curren gemm operation accumulator and output data type match return operation
if(gemm_desc.tile_description.math_instruction.element_accumulator == gemm_desc.C.element) {
if(gemm_desc.tile_description.math_instruction.element_accumulator == gemm_desc.D.element) {
return operation;
}
@@ -1137,7 +1149,10 @@ Operation const* find_gemm_operation_for_parallel_reduction(Operation const *ope
gemm_desc.B.element,
gemm_desc.B.layout,
gemm_desc.transform_B,
gemm_desc.tile_description.math_instruction.element_accumulator);
gemm_desc.tile_description.math_instruction.element_accumulator, // C/D are same type
LayoutTypeID::kColumnMajor,
gemm_desc.tile_description.math_instruction.element_accumulator,
LayoutTypeID::kColumnMajor);
// gemm operation table
auto gemm_operations = Singleton::get().operation_table.gemm_operations;
tools/library/src/library_internal.h
+8
View File
@@ -96,6 +96,14 @@ template <> struct NumericTypeMap<uint8_t> {
static NumericTypeID const kId = NumericTypeID::kU8;
};
template <> struct NumericTypeMap<cutlass::float_e4m3_t> {
static NumericTypeID const kId = NumericTypeID::kFE4M3;
};
template <> struct NumericTypeMap<cutlass::float_e5m2_t> {
static NumericTypeID const kId = NumericTypeID::kFE5M2;
};
template <> struct NumericTypeMap<uint16_t> {
static NumericTypeID const kId = NumericTypeID::kU16;
};
tools/library/src/manifest.cpp
+1 -1
View File
@@ -57,7 +57,7 @@ Status Manifest::initialize() {
// initialize procedurally generated cutlass op in manifest object
initialize_all(*this);
// initialize manually instanced conv3d reference op in manifest object
// initialize manually instanced reference op in manifest object
initialize_reference_operations(*this);
// initialize manually instanced reduction reference op in manifest object
tools/library/src/operation_table.cu
+4 -1
View File
@@ -66,7 +66,10 @@ void OperationTable::append(Manifest const &manifest) {
gemm_desc.B.element,
gemm_desc.B.layout,
gemm_desc.transform_B,
gemm_desc.C.element
gemm_desc.C.element,
gemm_desc.C.layout,
gemm_desc.D.element,
gemm_desc.D.layout
);
Operation const *op = operation.get();
tools/library/src/reduction/init_reduction_operations.cu
-1
View File
@@ -58,7 +58,6 @@ void initialize_all_reduction_op(Manifest &manifest) {
initialize_reduce_add_linear_combination_f32_f32_f32(manifest);
initialize_reduce_add_linear_combination_f64_f64_f64(manifest);
initialize_reduce_add_linear_combination_cf32_cf32_cf32(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reduction/reduction_device.cu
+11 -10
View File
@@ -43,9 +43,10 @@ namespace library {
// naming convention initialize_reduce_[ReductionOp]_[EpilogueOp]_[ElementWorkspace]_[ElementAccumulator]_[ElementOutput]
void initialize_reduce_add_linear_combination_f16_f16_f16(Manifest &manifest) {
using ElementWorkspace = cutlass::half_t;
using ElementWorkspace = cutlass::half_t;
using ElementAccumulator = cutlass::half_t;
using ElementOutput = cutlass::half_t;
using ElementCompute = cutlass::half_t;
@@ -58,7 +59,7 @@ void initialize_reduce_add_linear_combination_f16_f16_f16(Manifest &manifest) {
>;
using ReductionOp = cutlass::reduction::thread::ReduceAdd<
ElementAccumulator,
ElementAccumulator,
typename EpilogueOutputOp::ElementAccumulator,
EpilogueOutputOp::kCount
>;
@@ -79,7 +80,7 @@ void initialize_reduce_add_linear_combination_f16_f16_f16(Manifest &manifest) {
void initialize_reduce_add_linear_combination_f32_f32_f16(Manifest &manifest) {
using ElementWorkspace = float;
using ElementWorkspace = float;
using ElementAccumulator = float;
using ElementOutput = cutlass::half_t;
using ElementCompute = float;
@@ -92,7 +93,7 @@ void initialize_reduce_add_linear_combination_f32_f32_f16(Manifest &manifest) {
>;
using ReductionOp = cutlass::reduction::thread::ReduceAdd<
ElementAccumulator,
ElementAccumulator,
typename EpilogueOutputOp::ElementAccumulator,
EpilogueOutputOp::kCount
>;
@@ -114,7 +115,7 @@ void initialize_reduce_add_linear_combination_f32_f32_f16(Manifest &manifest) {
void initialize_reduce_add_linear_combination_f32_f32_f32(Manifest &manifest) {
using ElementWorkspace = float;
using ElementWorkspace = float;
using ElementAccumulator = float;
using ElementOutput = float;
using ElementCompute = float;
@@ -127,7 +128,7 @@ void initialize_reduce_add_linear_combination_f32_f32_f32(Manifest &manifest) {
>;
using ReductionOp = cutlass::reduction::thread::ReduceAdd<
ElementAccumulator,
ElementAccumulator,
typename EpilogueOutputOp::ElementAccumulator,
EpilogueOutputOp::kCount
>;
@@ -148,7 +149,7 @@ void initialize_reduce_add_linear_combination_f32_f32_f32(Manifest &manifest) {
void initialize_reduce_add_linear_combination_f64_f64_f64(Manifest &manifest) {
using ElementWorkspace = double;
using ElementWorkspace = double;
using ElementAccumulator = double;
using ElementOutput = double;
using ElementCompute = double;
@@ -161,7 +162,7 @@ void initialize_reduce_add_linear_combination_f64_f64_f64(Manifest &manifest) {
>;
using ReductionOp = cutlass::reduction::thread::ReduceAdd<
ElementAccumulator,
ElementAccumulator,
typename EpilogueOutputOp::ElementAccumulator,
EpilogueOutputOp::kCount
>;
@@ -182,7 +183,7 @@ void initialize_reduce_add_linear_combination_f64_f64_f64(Manifest &manifest) {
void initialize_reduce_add_linear_combination_cf32_cf32_cf32(Manifest &manifest) {
using ElementWorkspace = cutlass::complex<float>;
using ElementWorkspace = cutlass::complex<float>;
using ElementAccumulator = cutlass::complex<float>;
using ElementOutput = cutlass::complex<float>;
using ElementCompute = cutlass::complex<float>;
@@ -195,7 +196,7 @@ void initialize_reduce_add_linear_combination_cf32_cf32_cf32(Manifest &manifest)
>;
using ReductionOp = cutlass::reduction::thread::ReduceAdd<
ElementAccumulator,
ElementAccumulator,
typename EpilogueOutputOp::ElementAccumulator,
EpilogueOutputOp::kCount
>;
tools/library/src/reference/conv_reference_operation.h
+1
View File
@@ -146,6 +146,7 @@ struct ConvReferenceDispatcher<
LayoutC,
ElementCompute,
ElementAccumulator,
ElementC,
ConvertOp,
InnerProductOp
>(
tools/library/src/reference/gemm.cu
+367
View File
@@ -137,6 +137,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int8_t,
float,
int32_t,
int8_t,
NumericConverterClamp<int8_t, float>
>(manifest);
@@ -146,6 +147,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
@@ -163,6 +165,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int8_t,
float,
int32_t,
int8_t,
NumericConverterClamp<int8_t, float>
>(manifest);
@@ -172,6 +175,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
@@ -191,6 +195,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
@@ -201,6 +206,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int8_t,
float,
int32_t,
int8_t,
NumericConverterClamp<int8_t, float>
>(manifest);
@@ -220,6 +226,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
@@ -230,6 +237,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
uint8_t,
float,
int32_t,
uint8_t,
NumericConverterClamp<uint8_t, float>
>(manifest);
@@ -240,6 +248,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int8_t,
float,
int32_t,
int8_t,
NumericConverterClamp<int8_t, float>
>(manifest);
@@ -259,6 +268,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
@@ -269,6 +279,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int4b_t,
float,
int32_t,
int4b_t,
NumericConverterClamp<int4b_t, float>
>(manifest);
@@ -288,6 +299,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
@@ -298,6 +310,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
uint4b_t,
float,
int32_t,
uint4b_t,
NumericConverterClamp<uint4b_t, float>
>(manifest);
@@ -308,6 +321,7 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
int4b_t,
float,
int32_t,
int4b_t,
NumericConverterClamp<int4b_t, float>
>(manifest);
@@ -330,6 +344,359 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
complex<double>,
complex<double>
>(manifest);
//
// FP8 GEMMs
//
//////////////////////////////////
/// ElementC: half_t
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float , // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float , // ElementAccumulator
half_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
half_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
/// ElementC: bfloat16_t
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
bfloat16_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
bfloat16_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
bfloat16_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
bfloat16_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
bfloat16_t, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
/// ElementC: float
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
//////////////////////////////////
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
float, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_reference_operation.h
+77 -8
View File
@@ -67,7 +67,8 @@ template <
typename LayoutC_,
typename ElementCompute_,
typename ElementAccumulator_ = ElementCompute_,
typename ConvertOp_ = NumericConverter<ElementC_, ElementCompute_>,
typename ElementD_ = ElementC_,
typename ConvertOp_ = NumericConverter<ElementD_, ElementCompute_>,
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
>
class GemmReferenceOperation : public Operation {
@@ -84,7 +85,9 @@ public:
static cutlass::ComplexTransform const kTransformB = TransformB;
using ElementC = ElementC_;
using LayoutC = LayoutC_;
using ElementD = ElementD_;
using TensorRefC = TensorRef<ElementC, LayoutC>;
using TensorRefD = TensorRef<ElementD, LayoutC>;
using ElementCompute = ElementCompute_;
using ElementAccumulator = ElementAccumulator_;
using ConvertOp = ConvertOp_;
@@ -114,6 +117,7 @@ public:
description_.B = make_TensorDescription<ElementB, LayoutB>();
description_.transform_B = ComplexTransformMap<kTransformB>::kId;
description_.C = make_TensorDescription<ElementC, LayoutC>();
description_.D = make_TensorDescription<ElementD, LayoutC>();
// Epilogue compute and accumulator type description
description_.element_epilogue = NumericTypeMap<ElementCompute>::kId;
@@ -196,7 +200,7 @@ public:
TensorRefA ref_A{static_cast<ElementA *>(const_cast<void *>(args.A)), LayoutA(int(config.lda))};
TensorRefB ref_B{static_cast<ElementB *>(const_cast<void *>(args.B)), LayoutB(int(config.ldb))};
TensorRefC ref_C{static_cast<ElementC *>(const_cast<void *>(args.C)), LayoutC(int(config.ldc))};
TensorRefC ref_D{static_cast<ElementC *>(args.D), LayoutC(int(config.ldd))};
TensorRefD ref_D{static_cast<ElementD *>(args.D), LayoutC(int(config.ldd))};
if (kProvider == Provider::kReferenceHost) {
@@ -209,6 +213,7 @@ public:
LayoutC,
ElementCompute,
ElementAccumulator,
ElementD,
ConvertOp,
InnerProductOp
>(
@@ -242,6 +247,7 @@ public:
LayoutC,
ElementCompute,
ElementAccumulator,
ElementD,
ConvertOp,
InnerProductOp
>(
@@ -282,7 +288,8 @@ template <
typename LayoutC_,
typename ElementCompute_,
typename ElementAccumulator_ = ElementCompute_,
typename ConvertOp_ = NumericConverter<ElementC_, ElementCompute_>,
typename ElementD_ = ElementC_,
typename ConvertOp_ = NumericConverter<ElementD_, ElementCompute_>,
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
>
void make_gemm(Manifest &manifest) {
@@ -294,6 +301,7 @@ void make_gemm(Manifest &manifest) {
ElementC_, LayoutC_,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>);
@@ -305,6 +313,7 @@ void make_gemm(Manifest &manifest) {
ElementC_, LayoutC_,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>);
@@ -317,37 +326,42 @@ template <
typename ElementC_,
typename ElementCompute_,
typename ElementAccumulator_ = ElementCompute_,
typename ConvertOp_ = NumericConverter<ElementC_, ElementCompute_>,
typename ElementD_ = ElementC_,
typename ConvertOp_ = NumericConverter<ElementD_, ElementCompute_>,
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
>
void make_gemm_canonical_layouts(Manifest &manifest) {
// M Major outputs
make_gemm<
ElementA_, cutlass::layout::ColumnMajor, TransformA,
ElementB_, cutlass::layout::ColumnMajor, TransformB,
ElementC_, cutlass::layout::ColumnMajor,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
make_gemm<
ElementA_, cutlass::layout::ColumnMajor, TransformA,
ElementB_, cutlass::layout::RowMajor, TransformB,
ElementC_, cutlass::layout::ColumnMajor,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
make_gemm<
ElementA_, cutlass::layout::RowMajor, TransformA,
ElementB_, cutlass::layout::ColumnMajor, TransformB,
ElementC_, cutlass::layout::ColumnMajor,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
@@ -358,6 +372,52 @@ void make_gemm_canonical_layouts(Manifest &manifest) {
ElementC_, cutlass::layout::ColumnMajor,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
// N Major outputs
make_gemm<
ElementA_, cutlass::layout::ColumnMajor, TransformA,
ElementB_, cutlass::layout::ColumnMajor, TransformB,
ElementC_, cutlass::layout::RowMajor,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
make_gemm<
ElementA_, cutlass::layout::ColumnMajor, TransformA,
ElementB_, cutlass::layout::RowMajor, TransformB,
ElementC_, cutlass::layout::RowMajor,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
make_gemm<
ElementA_, cutlass::layout::RowMajor, TransformA,
ElementB_, cutlass::layout::ColumnMajor, TransformB,
ElementC_, cutlass::layout::RowMajor,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
make_gemm<
ElementA_, cutlass::layout::RowMajor, TransformA,
ElementB_, cutlass::layout::RowMajor, TransformB,
ElementC_, cutlass::layout::RowMajor,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
@@ -372,6 +432,7 @@ template <
typename ElementC_,
typename ElementCompute_,
typename ElementAccumulator_ = ElementCompute_,
typename ElementD_ = ElementC_,
typename ConvertOp_ = NumericConverter<ElementC_, ElementCompute_>,
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
>
@@ -383,6 +444,7 @@ void make_gemm_interleaved_layouts(Manifest &manifest) {
ElementC_, cutlass::layout::ColumnMajor,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
@@ -396,7 +458,8 @@ template <
typename ElementC_,
typename ElementCompute_,
typename ElementAccumulator_ = ElementCompute_,
typename ConvertOp_ = NumericConverter<ElementC_, ElementCompute_>,
typename ElementD_ = ElementC_,
typename ConvertOp_ = NumericConverter<ElementD_, ElementCompute_>,
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
>
void make_gemm_real_canonical_layouts(Manifest &manifest) {
@@ -406,6 +469,7 @@ void make_gemm_real_canonical_layouts(Manifest &manifest) {
ElementC_,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
@@ -418,7 +482,8 @@ template <
typename ElementC_,
typename ElementCompute_,
typename ElementAccumulator_ = ElementCompute_,
typename ConvertOp_ = NumericConverter<ElementC_, ElementCompute_>,
typename ElementD_ = ElementC_,
typename ConvertOp_ = NumericConverter<ElementD_, ElementCompute_>,
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
>
void make_gemm_complex_canonical_layouts(Manifest &manifest) {
@@ -429,6 +494,7 @@ void make_gemm_complex_canonical_layouts(Manifest &manifest) {
ElementC_,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
@@ -439,6 +505,7 @@ void make_gemm_complex_canonical_layouts(Manifest &manifest) {
ElementC_,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
@@ -449,6 +516,7 @@ void make_gemm_complex_canonical_layouts(Manifest &manifest) {
ElementC_,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
@@ -459,6 +527,7 @@ void make_gemm_complex_canonical_layouts(Manifest &manifest) {
ElementC_,
ElementCompute_,
ElementAccumulator_,
ElementD_,
ConvertOp_,
InnerProductOp_
>(manifest);
tools/library/src/util.cu
+64
View File
@@ -443,6 +443,8 @@ NumericTypeID_enumerants[] = {
{"s16", "S16", NumericTypeID::kS16},
{"s32", "S32", NumericTypeID::kS32},
{"s64", "S64", NumericTypeID::kS64},
{"fe4m3", "FE4M3", NumericTypeID::kFE4M3},
{"fe5m2", "FE5M2", NumericTypeID::kFE5M2},
{"f16", "F16", NumericTypeID::kF16},
{"bf16", "BF16", NumericTypeID::kBF16},
{"f32", "F32", NumericTypeID::kF32},
@@ -504,6 +506,8 @@ NumericTypeID from_string<NumericTypeID>(std::string const &str) {
/// Returns the size of a data type in bits
int sizeof_bits(NumericTypeID type) {
switch (type) {
case NumericTypeID::kFE4M3: return 8;
case NumericTypeID::kFE5M2: return 8;
case NumericTypeID::kF16: return 16;
case NumericTypeID::kBF16: return 16;
case NumericTypeID::kTF32: return 32;
@@ -581,6 +585,8 @@ bool is_integer_type(NumericTypeID type) {
/// Returns true if numeric type is signed
bool is_signed_type(NumericTypeID type) {
switch (type) {
case NumericTypeID::kFE4M3: return true;
case NumericTypeID::kFE5M2: return true;
case NumericTypeID::kF16: return true;
case NumericTypeID::kBF16: return true;
case NumericTypeID::kTF32: return true;
@@ -610,6 +616,8 @@ bool is_unsigned_integer(NumericTypeID type) {
/// Returns true if numeric type is floating-point type
bool is_float_type(NumericTypeID type) {
switch (type) {
case NumericTypeID::kFE4M3: return true;
case NumericTypeID::kFE5M2: return true;
case NumericTypeID::kF16: return true;
case NumericTypeID::kBF16: return true;
case NumericTypeID::kTF32: return true;
@@ -1050,6 +1058,20 @@ bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string c
ss >> *reinterpret_cast<int64_t *>(bytes.data());
}
break;
case NumericTypeID::kFE4M3:
{
float tmp;
ss >> tmp;
*reinterpret_cast<float_e4m3_t *>(bytes.data()) = static_cast<float_e4m3_t>(tmp);
}
break;
case NumericTypeID::kFE5M2:
{
float tmp;
ss >> tmp;
*reinterpret_cast<float_e5m2_t *>(bytes.data()) = static_cast<float_e5m2_t>(tmp);
}
break;
case NumericTypeID::kF16:
{
float tmp;
@@ -1187,6 +1209,18 @@ std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type) {
ss << *reinterpret_cast<int64_t *>(bytes.data());
}
break;
case NumericTypeID::kFE4M3:
{
float tmp = *reinterpret_cast<float_e4m3_t *>(bytes.data());
ss << tmp;
}
break;
case NumericTypeID::kFE5M2:
{
float tmp = *reinterpret_cast<float_e5m2_t *>(bytes.data());
ss << tmp;
}
break;
case NumericTypeID::kF16:
{
float tmp = *reinterpret_cast<half_t *>(bytes.data());
@@ -1329,6 +1363,16 @@ bool cast_from_int64(std::vector<uint8_t> &bytes, NumericTypeID type, int64_t sr
*reinterpret_cast<int64_t *>(bytes.data()) = static_cast<int64_t>(src);
}
break;
case NumericTypeID::kFE4M3:
{
*reinterpret_cast<float_e4m3_t *>(bytes.data()) = static_cast<float_e4m3_t>(float(src));
}
break;
case NumericTypeID::kFE5M2:
{
*reinterpret_cast<float_e5m2_t *>(bytes.data()) = static_cast<float_e5m2_t>(float(src));
}
break;
case NumericTypeID::kF16:
{
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
@@ -1429,6 +1473,16 @@ bool cast_from_uint64(std::vector<uint8_t> &bytes, NumericTypeID type, uint64_t
*reinterpret_cast<int64_t *>(bytes.data()) = static_cast<int64_t>(src);
}
break;
case NumericTypeID::kFE4M3:
{
*reinterpret_cast<float_e4m3_t *>(bytes.data()) = static_cast<float_e4m3_t>(float(src));
}
break;
case NumericTypeID::kFE5M2:
{
*reinterpret_cast<float_e5m2_t *>(bytes.data()) = static_cast<float_e5m2_t>(float(src));
}
break;
case NumericTypeID::kF16:
{
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
@@ -1530,6 +1584,16 @@ bool cast_from_double(std::vector<uint8_t> &bytes, NumericTypeID type, double sr
*reinterpret_cast<int64_t *>(bytes.data()) = static_cast<int64_t>(src);
}
break;
case NumericTypeID::kFE4M3:
{
*reinterpret_cast<float_e4m3_t *>(bytes.data()) = static_cast<float_e4m3_t>(float(src));
}
break;
case NumericTypeID::kFE5M2:
{
*reinterpret_cast<float_e5m2_t *>(bytes.data()) = static_cast<float_e5m2_t>(float(src));
}
break;
case NumericTypeID::kF16:
{
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));