More updates for 3.1 (#958)
* Updates for 3.1 * Minor change * doc link fix * Minor updates
This commit is contained in:
@@ -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):
|
||||
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user