More updates for 3.1 (#958)

* 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
+57 -25
View File
@@ -4079,32 +4079,36 @@ def GenerateSM90_TensorOp_16b_WGMMA_gemm(manifest, cuda_version):
max_cc = 90
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([math_inst.instruction_shape[0]*4, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [2,1,1]),
tile_descriptions_small = [
# Not compatible with TmaWarpSpecializedCooperative
TileDescription([math_inst.instruction_shape[0], math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [2,1,1]),
TileDescription([math_inst.instruction_shape[0], math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,2,1]),
TileDescription([math_inst.instruction_shape[0], math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,1,1]),
]
tile_descriptions_medium = [
TileDescription([math_inst.instruction_shape[0]*2, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [2,1,1]),
TileDescription([math_inst.instruction_shape[0]*2, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,2,1]),
TileDescription([math_inst.instruction_shape[0]*2, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,1,1]),
]
tile_descriptions_large = [
TileDescription([math_inst.instruction_shape[0]*4, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [2,1,1]),
TileDescription([math_inst.instruction_shape[0]*4, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,2,1]),
TileDescription([math_inst.instruction_shape[0]*4, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,1,1]),
TileDescription([math_inst.instruction_shape[0]*2, math_inst.instruction_shape[1]*2, math_inst.instruction_shape[2]*4],
0, [4, 2, 1], math_inst, min_cc, max_cc, [2,1,1]),
#TileDescription([math_inst.instruction_shape[0], math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
# 0, [4, 1, 1], math_inst, min_cc, max_cc, [2,1,1]), - Not compatible with TmaWarpSpecializedCooperative
TileDescription([math_inst.instruction_shape[0]*4, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,2,1]),
TileDescription([math_inst.instruction_shape[0]*4, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,2,1]),
TileDescription([math_inst.instruction_shape[0]*2, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,2,1]),
TileDescription([math_inst.instruction_shape[0]*2, math_inst.instruction_shape[1]*2, math_inst.instruction_shape[2]*4],
0, [4, 2, 1], math_inst, min_cc, max_cc, [1,2,1]),
#TileDescription([math_inst.instruction_shape[0], math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
# 0, [4, 1, 1], math_inst, min_cc, max_cc, [1,2,1]),- Not compatible with TmaWarpSpecializedCooperative
TileDescription([math_inst.instruction_shape[0]*4, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,1,1]),
TileDescription([math_inst.instruction_shape[0]*2, math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
0, [4, 1, 1], math_inst, min_cc, max_cc, [1,1,1]),
#TileDescription([math_inst.instruction_shape[0], math_inst.instruction_shape[1], math_inst.instruction_shape[2]*4],
# 0, [4, 1, 1], math_inst, min_cc, max_cc, [1,1,1]), - Not compatible with TmaWarpSpecializedCooperative
]
tile_descriptions = tile_descriptions_medium + tile_descriptions_large
data_type = {
"a_type" : math_inst.element_a,
@@ -4139,11 +4143,17 @@ def GenerateSM90_TensorOp_16b_WGMMA_gemm(manifest, cuda_version):
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, schedules)
# persistent kernels with TMA epilogues
if data_type["c_type"] in [DataType.f16, DataType.bf16] and CudaToolkitVersionSatisfies(cuda_version, 12, 1):
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[KernelScheduleType.TmaWarpSpecializedPingpong, EpilogueScheduleType.TmaWarpSpecialized],
[KernelScheduleType.TmaWarpSpecializedCooperative, EpilogueScheduleType.TmaWarpSpecializedCooperative]])
# Emit instance without C allocation+load
if CudaToolkitVersionSatisfies(cuda_version, 12, 1):
# not enough smem for 256x128 f32 out with C allocation
if data_type["d_type"] == DataType.f32:
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions_medium, data_type,
[[KernelScheduleType.TmaWarpSpecializedPingpong, EpilogueScheduleType.TmaWarpSpecialized],
[KernelScheduleType.TmaWarpSpecializedCooperative, EpilogueScheduleType.TmaWarpSpecializedCooperative]])
else:
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[KernelScheduleType.TmaWarpSpecializedPingpong, EpilogueScheduleType.TmaWarpSpecialized],
[KernelScheduleType.TmaWarpSpecializedCooperative, EpilogueScheduleType.TmaWarpSpecializedCooperative]])
# Emit instance without C allocation + load
data_type["c_type"] = DataType.void
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[KernelScheduleType.TmaWarpSpecializedPingpong, EpilogueScheduleType.TmaWarpSpecialized],
@@ -4170,7 +4180,7 @@ def GenerateSM90_TensorOp_16b_WGMMA_gemm(manifest, cuda_version):
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type_mixed, schedules)
# persistent kernels with TMA epilogues
if data_type_mixed["c_type"] in [DataType.f16, DataType.bf16] and CudaToolkitVersionSatisfies(cuda_version, 12, 1):
if CudaToolkitVersionSatisfies(cuda_version, 12, 1):
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type_mixed,
[[KernelScheduleType.TmaWarpSpecializedPingpong, EpilogueScheduleType.TmaWarpSpecialized],
[KernelScheduleType.TmaWarpSpecializedCooperative, EpilogueScheduleType.TmaWarpSpecializedCooperative]])
@@ -4331,6 +4341,28 @@ def GenerateSM90_TensorOp_int8_WGMMA_gemm(manifest, cuda_version):
for data_type in data_types:
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type)
# persistent kernels with TMA epilogues
if CudaToolkitVersionSatisfies(cuda_version, 12, 1):
# Emit instance without C allocation+load
data_types += [
{
"a_type" : math_inst.element_a,
"b_type" : math_inst.element_b,
"c_type" : DataType.void,
"d_type" : math_inst.element_accumulator,
"acc_type" : math_inst.element_accumulator,
"epi_type" : math_inst.element_accumulator
}
]
for data_type in data_types:
# Set alignment d based on Destination format.
for layout in layouts:
layout[2][1] = 128 // DataTypeSize[data_type["d_type"]]
CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type,
[[KernelScheduleType.TmaWarpSpecializedPingpong, EpilogueScheduleType.TmaWarpSpecialized],
[KernelScheduleType.TmaWarpSpecializedCooperative, EpilogueScheduleType.TmaWarpSpecializedCooperative]])
#
def GenerateSM90_TensorOp_1684(manifest, cuda_version):
-352
View File
@@ -345,358 +345,6 @@ void initialize_gemm_reference_operations(Manifest &manifest) {
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);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
+15 -2
View File
@@ -103,14 +103,20 @@ bool get_cublas_transpose_operation(
/// Maps a CUTLASS numeric type to a cuBLAS data type enumeration
bool get_cublas_datatype(cublasDataType_t &data_type, library::NumericTypeID element_type) {
switch (element_type) {
switch (element_type) {
case library::NumericTypeID::kFE4M3:
#if (__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8))
data_type = CUDA_R_8F_E4M3;
return true;
#endif
break;
case library::NumericTypeID::kFE5M2:
#if (__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8))
data_type = CUDA_R_8F_E5M2;
return true;
#endif
break;
case library::NumericTypeID::kF16:
data_type = CUDA_R_16F;
@@ -139,7 +145,7 @@ bool get_cublas_datatype(cublasDataType_t &data_type, library::NumericTypeID ele
return true;
case library::NumericTypeID::kS16:
break;
break;
case library::NumericTypeID::kS32:
data_type = CUDA_R_32I;
@@ -260,6 +266,13 @@ Status cublas_satisfies(library::GemmDescription const &desc) {
return Status::kErrorNotSupported;
}
// input type BF16 and TF32 not supported in cuBLAS
if (desc.A.element == library::NumericTypeID::kBF16 ||
desc.A.element == library::NumericTypeID::kTF32) {
return Status::kErrorNotSupported;
}
return Status::kSuccess;
}