CUTLASS 2.2 (#96)

Adds support for NVIDIA Ampere Architecture features. CUDA 11 Toolkit recommended.
This commit is contained in:
Andrew Kerr
2020-06-08 16:17:35 -07:00
committed by GitHub
parent e33d90b361
commit 86931fef85
584 changed files with 51080 additions and 3373 deletions
+3 -3
View File
@@ -1,4 +1,4 @@
# Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
# Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification, are permitted
# provided that the following conditions are met:
@@ -22,7 +22,7 @@
include(GNUInstallDirs)
find_package(Python3 3.6 COMPONENTS Interpreter REQUIRED)
find_package(Python3 3.5 COMPONENTS Interpreter REQUIRED)
add_library(cutlass_library_includes INTERFACE)
add_library(nvidia::cutlass::library::includes ALIAS cutlass_library_includes)
@@ -59,7 +59,7 @@ cutlass_add_library(
src/operation_table.cu
src/singleton.cu
src/util.cu
)
file(GLOB_RECURSE GENERATOR_PYTHON_SOURCES CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/scripts/*.py)
+59 -1
View File
@@ -45,6 +45,9 @@ private:
/// Host workspace
static int const kHostWorkspaceSize = (4 << 10);
/// Provider of operations
Provider provider_;
/// CUDA device properties
cudaDeviceProp device_;
@@ -90,6 +93,12 @@ public:
/// Gets the current CUDA stream
cudaStream_t get_stream() const;
/// Gets the current provider
Provider get_provider() const;
/// Sets the provider of operations
void set_provider(Provider provider);
/// Gets the device workspace size
size_t get_workspace_size() const;
@@ -149,6 +158,56 @@ public:
void * ptr_D, /// Pointer to D matrix
int ldd /// Leading dimension of D matrix
);
/// Executes a GEMM computation: D <= alpha * A*B + beta * C.
//
// Supports batched-strided, batched array or split-K serial or split-K parallel.
//
Status gemm_universal(
GemmUniversalMode mode, /// indicates the mode in which the kUniversal GEMM is launched
int M, /// GEMM M dimension
int N, /// GEMM N dimension
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
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
int 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
int ldb, /// Leading dimension of B matrix
void const * beta, /// Pointer to beta scalar
NumericTypeID element_C, /// Data type of C and D matrices
void const * ptr_C, /// Pointer to C matrix
int ldc, /// Leading dimension of C matrix
void * ptr_D, /// Pointer to D matrix
int 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
int64_t batch_stride_D = 0 /// Batch stride of D operand
);
/// Planar complex GEMM
///
@@ -276,7 +335,6 @@ public:
using HandlePtr = std::unique_ptr<Handle>;
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace library
} // namespace cutlass
+105 -7
View File
@@ -44,6 +44,7 @@
#include <vector>
#include <string>
#include <cstdint>
#include <stdexcept>
#include <cuda_runtime.h>
#include "cutlass/cutlass.h"
@@ -93,10 +94,14 @@ enum class NumericTypeID {
kS32,
kS64,
kF16,
kBF16,
kTF32,
kF32,
kF64,
kCF16,
kCBF16,
kCF32,
kCTF32,
kCF64,
kCS4,
kCS8,
@@ -120,6 +125,7 @@ enum class ComplexTransform {
/// Providers
enum class Provider {
kNone,
kCUTLASS,
kReferenceHost,
kReferenceDevice,
@@ -132,6 +138,8 @@ enum class Provider {
/// Enumeration indicating the kind of operation
enum class OperationKind {
kGemm,
kEqGemm,
kReduction,
kInvalid
};
@@ -160,9 +168,11 @@ enum class OpcodeClassID {
};
enum class MathOperationID {
kAdd,
kMultiplyAdd,
kMultiplyAddSaturate,
kMultiplyAddComplex,
kMultiplyAddGaussianComplex,
kXorPopc,
kInvalid
};
@@ -180,12 +190,17 @@ enum class GemmKind {
kInvalid
};
/// Mode of GEMM
enum class GemmUniversalMode {
kGemm,
kGemmSplitKParallel,
kBatched,
kArray,
/// Mode of Universal GEMM
using GemmUniversalMode = cutlass::gemm::GemmUniversalMode;
enum class EpilogueKind {
kUnknown,
kConversion,
kLinearCombination,
kLinearCombinationClamp,
kLinearCombinationPlanarComplex,
kLinearCombinationRelu,
kLinearCombinationSigmoid,
kInvalid
};
@@ -220,6 +235,22 @@ struct MathInstructionDescription {
opcode_class(opcode_class),
math_operation(math_operation) {}
// Equality operator
inline
bool operator==(MathInstructionDescription const& rhs) const{
return (
(instruction_shape == rhs.instruction_shape) &&
(element_accumulator == rhs.element_accumulator) &&
(opcode_class == rhs.opcode_class) &&
(math_operation == rhs.math_operation));
}
// Inequality operator
inline
bool operator!=(MathInstructionDescription const& rhs) const {
return !(*this == rhs);
}
};
/// Structure describing the tiled structure of a GEMM-like computation
@@ -261,6 +292,24 @@ struct TileDescription {
math_instruction(math_instruction),
minimum_compute_capability(minimum_compute_capability),
maximum_compute_capability(maximum_compute_capability) { }
// Equality operator
inline
bool operator==(TileDescription const& rhs) const{
return (
(threadblock_shape == rhs.threadblock_shape) &&
(threadblock_stages == rhs.threadblock_stages) &&
(warp_count == rhs.warp_count) &&
(math_instruction == rhs.math_instruction) &&
(minimum_compute_capability == rhs.minimum_compute_capability) &&
(maximum_compute_capability == rhs.maximum_compute_capability));
}
// Inequality operator
inline
bool operator!=(TileDescription const& rhs) const {
return !(*this == rhs);
}
};
/// High-level description of an operation
@@ -379,6 +428,20 @@ struct GemmDescription : public OperationDescription {
transform_B(transform_B) {}
};
/// Description of all Reduction operations
struct ReductionDescription : public OperationDescription {
/// Describes the data type of workspace
NumericTypeID element_workspace;
/// Describes the data type of final output
NumericTypeID element_output;
/// Describes the data type of the scalars passed to the epilogue
NumericTypeID element_epilogue;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -549,6 +612,42 @@ struct GemmArrayArguments {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Universal GEMM supporting multiple split-K modes, multiple batched modes, real and complex
//
// OperationKind: Gemm
// GemmKind: Universal
struct GemmUniversalConfiguration {
GemmUniversalMode mode;
gemm::GemmCoord problem_size;
int batch_count;
int64_t lda;
int64_t ldb;
int64_t ldc;
int64_t ldd;
};
struct GemmUniversalArguments {
void const *A;
void const *B;
void const *C;
void *D;
void const *alpha;
void const *beta;
ScalarPointerMode pointer_mode;
int64_t batch_stride_A;
int64_t batch_stride_B;
int64_t batch_stride_C;
int64_t batch_stride_D;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Complex valued GEMM in which real and imaginary parts are separated by a stride
//
// OperationKind: Gemm
@@ -648,7 +747,6 @@ struct GemmPlanarComplexArrayArguments {
ScalarPointerMode pointer_mode;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace library
@@ -45,6 +45,13 @@ namespace cutlass {
namespace library {
///////////////////////////////////////////////////////////////////////////////////////////////////
// Forward declaration
class Manifest;
// init and insert all cutlass gemm and conv2d op in manifest object (procedurally generated using generator.py)
void initialize_all(Manifest &manifest);
/////////////////////////////////////////////////////////////////////////////////////////////////////////
/// List of operations
using OperationVector = std::vector<std::unique_ptr<Operation>>;
@@ -29,24 +29,28 @@
*/
#pragma once
#include <fstream>
#include <iosfwd>
#include <unordered_map>
#include <algorithm>
#include "cutlass/library/library.h"
#include "cutlass/library/manifest.h"
#include "cutlass/library/util.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace library {
/////////////////////////////////////////////////////////////////////////////////////////////////
// Data Structures for Gemm Functional Maps
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Tuple uniquely identifying functional behavior
/// Tuple uniquely identifying Gemm functional behavior
struct GemmFunctionalKey {
Provider provider;
GemmKind gemm_kind;
NumericTypeID element_compute;
NumericTypeID element_scalar;
NumericTypeID element_A;
@@ -63,6 +67,8 @@ struct GemmFunctionalKey {
inline
GemmFunctionalKey(
Provider provider,
GemmKind gemm_kind = GemmKind::kGemm,
NumericTypeID element_compute = NumericTypeID::kF32,
NumericTypeID element_scalar = NumericTypeID::kF32,
NumericTypeID element_A = NumericTypeID::kF16,
@@ -73,6 +79,8 @@ struct GemmFunctionalKey {
ComplexTransform transform_B = ComplexTransform::kNone,
NumericTypeID element_C = NumericTypeID::kF16
):
provider(provider),
gemm_kind(gemm_kind),
element_compute(element_compute),
element_scalar(element_scalar),
element_A(element_A),
@@ -87,6 +95,8 @@ struct GemmFunctionalKey {
inline
bool operator==(GemmFunctionalKey const &rhs) const {
return
(provider == rhs.provider) &&
(gemm_kind == rhs.gemm_kind) &&
(element_compute == rhs.element_compute) &&
(element_scalar == rhs.element_scalar) &&
(element_A == rhs.element_A) &&
@@ -104,6 +114,28 @@ struct GemmFunctionalKey {
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
inline
std::ostream & operator<<(std::ostream &out, cutlass::library::GemmFunctionalKey const &k) {
out << "{\n"
<< " provider: " << to_string(k.provider) << "\n"
<< " gemm_kind: " << to_string(k.gemm_kind) << "\n"
<< " element_compute: " << to_string(k.element_compute) << "\n"
<< " element_scalar: " << to_string(k.element_scalar) << "\n"
<< " element_A: " << to_string(k.element_A) << "\n"
<< " layout_A: " << to_string(k.layout_A) << "\n"
<< " transform_A: " << to_string(k.transform_A) << "\n"
<< " element_B: " << to_string(k.element_B) << "\n"
<< " layout_B: " << to_string(k.layout_B) << "\n"
<< " transform_B: " << to_string(k.transform_B) << "\n"
<< " element_C: " << to_string(k.element_C) << "\n"
<< "}";
return out;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Hash function for GemmFunctionalKey
@@ -120,15 +152,17 @@ struct GemmFunctionalKeyHasher {
IntHash hash;
return
rotl(hash(int(key.element_compute)), 2) ^
rotl(hash(int(key.element_scalar)), 3) ^
rotl(hash(int(key.element_A)), 4) ^
rotl(hash(int(key.layout_A)), 5) ^
rotl(hash(int(key.transform_A)), 6) ^
rotl(hash(int(key.element_B)), 7) ^
rotl(hash(int(key.layout_B)), 8) ^
rotl(hash(int(key.transform_B)), 9) ^
rotl(hash(int(key.element_C)), 10);
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);
}
};
@@ -172,6 +206,7 @@ using GemmOperationFunctionalMap = std::unordered_map<
GemmOperationVectorMap,
GemmFunctionalKeyHasher
>;
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -179,15 +214,10 @@ using GemmOperationFunctionalMap = std::unordered_map<
class OperationTable {
public:
/// Map of all operations of type kGemm and gemm_kind of type kGemm
/// Map of all operations of type kGemm
// provider (kCUTLASS)
GemmOperationFunctionalMap gemm_operations;
/// Map of all operations of type kGemm and gemm_kind of type kPlanarComplex
GemmOperationFunctionalMap gemm_planar_complex_operations;
/// Map of all operations of type kGemm and gemm_kind of type kPlanarComplexArray
GemmOperationFunctionalMap gemm_planar_complex_array_operations;
public:
void append(Manifest const &manifest);
@@ -202,4 +232,3 @@ public:
/////////////////////////////////////////////////////////////////////////////////////////////////
std::ostream & operator<<(std::ostream &out, cutlass::library::GemmFunctionalKey const &k);
@@ -49,6 +49,9 @@ char const *to_string(Provider provider, bool pretty = false);
/// Parses a Provider enumerant from a string
template <> Provider from_string<Provider>(std::string const &str);
/// Converts a GemmKind enumerant to a string
char const *to_string(GemmKind type, bool pretty = false);
/// Converts a NumericType enumerant to a string
char const *to_string(OperationKind type, bool pretty = false);
@@ -111,6 +114,14 @@ char const *to_string(ComplexTransform type, bool pretty = false);
template <>
ComplexTransform from_string<ComplexTransform>(std::string const &str);
/// Converts a SplitKMode enumerant to a string
char const *to_string(SplitKMode split_k_mode, bool pretty = false);
/// Converts a SplitKMode enumerant from a string
template <>
SplitKMode from_string<SplitKMode>(std::string const &str);
/// Lexical cast from int64_t to string
std::string lexical_cast(int64_t int_value);
+147 -4
View File
@@ -23,7 +23,7 @@ from library import *
class GemmOperation:
#
def __init__(self, gemm_kind, arch, tile_description, A, B, C, element_epilogue, \
epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Cohort):
epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity8):
self.operation_kind = OperationKind.Gemm
self.arch = arch
@@ -40,6 +40,7 @@ class GemmOperation:
def is_complex(self):
complex_operators = [
MathOperation.multiply_add_complex,
MathOperation.multiply_add_complex_gaussian
]
return self.tile_description.math_instruction.math_operation in complex_operators
@@ -58,6 +59,8 @@ class GemmOperation:
#
def short_math_name(self):
if self.tile_description.math_instruction.math_operation == MathOperation.multiply_add_complex_gaussian:
return "g%s" % ShortDataTypeNames[self.accumulator_type()]
return ShortDataTypeNames[self.accumulator_type()]
@@ -259,6 +262,135 @@ class EmitGemmInstance:
###################################################################################################
#
class EmitGemmUniversalInstance:
''' Responsible for emitting a CUTLASS template definition'''
def __init__(self):
self.gemm_template = """
// Gemm operator ${operation_name}
using ${operation_name}_base =
typename cutlass::gemm::kernel::DefaultGemmUniversal<
${element_b}, ${layout_b}, ${transform_b}, ${align_b}, // transposed B operand
${element_a}, ${layout_a}, ${transform_a}, ${align_a}, // transposed A operand
${element_c}, ${layout_c},
${element_accumulator},
${opcode_class},
${arch},
cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,
cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
${epilogue_functor}<
${element_c},
${epilogue_vector_length},
${element_accumulator},
${element_epilogue}
>,
${swizzling_functor},
${stages},
${math_operation}
>::GemmKernel;
// Define named type
struct ${operation_name} :
public ${operation_name}_base { };
"""
self.gemm_template_interleaved = """
// Gemm operator ${operation_name}
using ${operation_name}_base =
typename cutlass::gemm::kernel::DefaultGemmUniversal<
${element_a}, ${layout_a}, ${transform_a}, ${align_a},
${element_b}, ${layout_b}, ${transform_b}, ${align_b},
${element_c}, ${layout_c},
${element_accumulator},
${opcode_class},
${arch},
cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,
cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
${epilogue_functor}<
${element_c},
${epilogue_vector_length},
${element_accumulator},
${element_epilogue}
>,
${swizzling_functor},
${stages},
${math_operation}
>::GemmKernel;
// Define named type
struct ${operation_name} :
public ${operation_name}_base { };
"""
def emit(self, operation):
threadblock_shape = operation.tile_description.threadblock_shape
warp_count = operation.tile_description.warp_count
warp_shape = [threadblock_shape[idx] // warp_count[idx] for idx in range(3)]
warp_shape[2] = operation.tile_description.threadblock_shape[2]
epilogue_vector_length = int(min(operation.C.alignment * DataTypeSize[operation.C.element], 128) / DataTypeSize[operation.C.element])
transpose_layouts = {
LayoutType.ColumnMajor: LayoutType.RowMajor,
LayoutType.RowMajor: LayoutType.ColumnMajor
}
if operation.A.layout in transpose_layouts.keys() and \
operation.B.layout in transpose_layouts.keys() and \
operation.C.layout in transpose_layouts.keys():
instance_layout_A = transpose_layouts[operation.A.layout]
instance_layout_B = transpose_layouts[operation.B.layout]
instance_layout_C = transpose_layouts[operation.C.layout]
gemm_template = self.gemm_template
else:
instance_layout_A, instance_layout_B, instance_layout_C = \
(operation.A.layout, operation.B.layout, operation.C.layout)
gemm_template = self.gemm_template_interleaved
#
values = {
'operation_name': operation.procedural_name(),
'element_a': DataTypeTag[operation.A.element],
'layout_a': LayoutTag[instance_layout_A],
'element_b': DataTypeTag[operation.B.element],
'layout_b': LayoutTag[instance_layout_B],
'element_c': DataTypeTag[operation.C.element],
'layout_c': LayoutTag[instance_layout_C],
'element_accumulator': DataTypeTag[operation.accumulator_type()],
'opcode_class': OpcodeClassTag[operation.tile_description.math_instruction.opcode_class],
'arch': "cutlass::arch::Sm%d" % operation.arch,
'threadblock_shape_m': str(operation.tile_description.threadblock_shape[0]),
'threadblock_shape_n': str(operation.tile_description.threadblock_shape[1]),
'threadblock_shape_k': str(operation.tile_description.threadblock_shape[2]),
'warp_shape_m': str(warp_shape[0]),
'warp_shape_n': str(warp_shape[1]),
'warp_shape_k': str(warp_shape[2]),
'instruction_shape_m': str(operation.tile_description.math_instruction.instruction_shape[0]),
'instruction_shape_n': str(operation.tile_description.math_instruction.instruction_shape[1]),
'instruction_shape_k': str(operation.tile_description.math_instruction.instruction_shape[2]),
'epilogue_vector_length': str(epilogue_vector_length),
'element_epilogue': str(DataTypeTag[operation.element_epilogue]),
'epilogue_functor': EpilogueFunctorTag[operation.epilogue_functor],
'swizzling_functor': SwizzlingFunctorTag[operation.swizzling_functor],
'stages': str(operation.tile_description.stages),
'align_a': str(operation.A.alignment),
'align_b': str(operation.B.alignment),
'transform_a': ComplexTransformTag[operation.A.complex_transform],
'transform_b': ComplexTransformTag[operation.B.complex_transform],
'math_operation': MathOperationTag[operation.tile_description.math_instruction.math_operation]
}
return SubstituteTemplate(gemm_template, values)
###################################################################################################
#
class EmitGemmPlanarComplexInstance:
''' Responsible for emitting a CUTLASS template definition'''
@@ -282,12 +414,13 @@ class EmitGemmPlanarComplexInstance:
${element_accumulator},
${element_epilogue}
>,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
${stages},
${math_operator}
>::GemmKernel;
struct ${operation_name} : public Operation_${operation_name} { };
struct ${operation_name} :
public Operation_${operation_name} { };
"""
def emit(self, operation):
@@ -355,7 +488,7 @@ class EmitGemmPlanarComplexArrayInstance:
${element_accumulator},
${element_epilogue}
>,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
${stages},
${math_operator}
>::GemmArrayKernel;
@@ -419,12 +552,14 @@ class EmitGemmConfigurationLibrary:
self.instance_emitter = {
GemmKind.Gemm: EmitGemmInstance,
GemmKind.Universal: EmitGemmUniversalInstance,
GemmKind.PlanarComplex: EmitGemmPlanarComplexInstance,
GemmKind.PlanarComplexArray: EmitGemmPlanarComplexArrayInstance
}
self.gemm_kind_wrappers = {
GemmKind.Gemm: 'GemmOperation',
GemmKind.Universal: 'GemmUniversalOperation',
GemmKind.PlanarComplex: 'GemmPlanarComplexOperation',
GemmKind.PlanarComplexArray: 'GemmPlanarComplexArrayOperation'
}
@@ -436,6 +571,13 @@ class EmitGemmConfigurationLibrary:
${compile_guard_start}
manifest.append(new ${gemm_kind}<Operation_${operation_name}>("${operation_name}"));
${compile_guard_end}
""",
GemmKind.Universal: """
${compile_guard_start}
manifest.append(new ${gemm_kind}<
cutlass::gemm::device::GemmUniversalAdapter<${operation_name}>
>("${operation_name}"));
${compile_guard_end}
""",
GemmKind.PlanarComplex: """
${compile_guard_start}
@@ -542,3 +684,4 @@ void initialize_${configuration_name}(Manifest &manifest) {
###################################################################################################
###################################################################################################
+858 -10
View File
@@ -18,7 +18,7 @@ from gemm_operation import *
def CudaToolkitVersionSatisfies(semantic_ver_string, major, minor, patch = 0):
# by default, use the latest CUDA Toolkit version
cuda_version = [10, 2, 82]
cuda_version = [11, 0, 132]
# Update cuda_version based on parsed string
if semantic_ver_string != '':
@@ -36,7 +36,7 @@ def CudaToolkitVersionSatisfies(semantic_ver_string, major, minor, patch = 0):
#
def CreateGemmOperator(manifest, layouts, tile_descriptions, data_type, \
alignment_constraints, complex_transforms = None, epilogue_functor = EpilogueFunctor.LinearCombination, \
swizzling_functor = SwizzlingFunctor.Cohort):
swizzling_functor = SwizzlingFunctor.Identity8):
if complex_transforms is None:
complex_transforms = [(ComplexTransform.none, ComplexTransform.none),]
@@ -61,7 +61,7 @@ def CreateGemmOperator(manifest, layouts, tile_descriptions, data_type, \
B = TensorDescription(element_b, layout[1], alignment, complex_transform[1])
C = TensorDescription(element_c, layout[2], alignment_c)
new_operation = GemmOperation(GemmKind.Gemm, tile_description.minimum_compute_capability, \
new_operation = GemmOperation(GemmKind.Universal, tile_description.minimum_compute_capability, \
tile_description, A, B, C, element_epilogue, epilogue_functor, swizzling_functor)
manifest.append(new_operation)
@@ -466,6 +466,9 @@ def GenerateSM70_WmmaTensorOp_161616(manifest, args):
def GenerateSM70(manifest, args):
GenerateSM70_TensorOp_884(manifest, args)
GenerateSM70_PlanarComplexTensorOp_884(manifest, args)
# To limit build size, WMMA GEMMs are disabled for now.
#
#GenerateSM70_WmmaTensorOp_161616(manifest, args)
###################################################################################################
@@ -621,6 +624,11 @@ def GenerateSM75_TensorOp_8816_TN(manifest, args):
DataType.s8, DataType.s8, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
MathInstruction( \
[8, 8, 16], \
DataType.u8, DataType.u8, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
]
min_cc = 75
@@ -654,7 +662,7 @@ def GenerateSM75_TensorOp_8816_TN(manifest, args):
data_type_mixed = [
math_inst.element_a,
math_inst.element_b,
math_inst.element_a,
DataType.s8,
DataType.f32,
]
@@ -687,6 +695,11 @@ def GenerateSM75_TensorOp_8816_Interleaved(manifest, args):
DataType.s8, DataType.s8, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
MathInstruction( \
[8, 8, 16], \
DataType.u8, DataType.u8, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
]
min_cc = 75
@@ -712,8 +725,7 @@ def GenerateSM75_TensorOp_8816_Interleaved(manifest, args):
]
operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp, \
SwizzlingFunctor.Identity)
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
for op in operations:
op.C.alignment = 8
@@ -736,6 +748,11 @@ def GenerateSM75_TensorOp_8832_TN(manifest, args):
DataType.s4, DataType.s4, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
MathInstruction( \
[8, 8, 32], \
DataType.u4, DataType.u4, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
]
min_cc = 75
@@ -769,7 +786,7 @@ def GenerateSM75_TensorOp_8832_TN(manifest, args):
data_type_mixed = [
math_inst.element_a,
math_inst.element_b,
math_inst.element_a,
DataType.s4,
DataType.f32,
]
@@ -804,6 +821,11 @@ def GenerateSM75_TensorOp_8832_Interleaved(manifest, args):
DataType.s4, DataType.s4, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
MathInstruction( \
[8, 8, 32], \
DataType.u4, DataType.u4, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
]
min_cc = 75
@@ -832,8 +854,7 @@ def GenerateSM75_TensorOp_8832_Interleaved(manifest, args):
]
operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp, \
SwizzlingFunctor.Identity)
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
for op in operations:
op.C.alignment = 16
@@ -911,6 +932,831 @@ def GenerateSM75(manifest, args):
###################################################################################################
###################################################################################################
#
def GenerateSM80_TensorOp_16816(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
]
math_instructions = [
MathInstruction( \
[16, 8, 16], \
DataType.f16, DataType.f16, DataType.f32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add),
MathInstruction( \
[16, 8, 16], \
DataType.f16, DataType.f16, DataType.f16, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add),
MathInstruction( \
[16, 8, 16], \
DataType.bf16, DataType.bf16, DataType.f32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add),
]
min_cc = 80
max_cc = 1024
alignment_constraints = [8, 4, 2]
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc),
TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 64], 3, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([128, 64, 64], 3, [2, 1, 2], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 64], 4, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([128, 64, 64], 4, [2, 1, 2], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 32], 10, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 64], 4, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 64], 5, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 256, 64], 3, [1, 4, 1], math_inst, min_cc, max_cc),
]
data_type = [
math_inst.element_a,
math_inst.element_b,
math_inst.element_accumulator,
math_inst.element_accumulator,
]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints)
# Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation)
if math_inst.element_a != math_inst.element_accumulator:
data_type_mixed = [
math_inst.element_a,
math_inst.element_b,
math_inst.element_a,
math_inst.element_accumulator,
]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type_mixed, alignment_constraints)
#
#
def GenerateSM80_PlanarComplexTensorOp_16816(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
]
complex_transforms = [
(ComplexTransform.none, ComplexTransform.none),
(ComplexTransform.conj, ComplexTransform.none),
(ComplexTransform.none, ComplexTransform.conj),
(ComplexTransform.conj, ComplexTransform.conj)
]
math_instructions = [
MathInstruction( \
[16, 8, 16], \
DataType.f16, DataType.f16, DataType.f32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add),
MathInstruction( \
[16, 8, 16], \
DataType.bf16, DataType.bf16, DataType.f32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add),
MathInstruction( \
[16, 8, 16], \
DataType.f16, DataType.f16, DataType.f16, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add),
]
min_cc = 80
max_cc = 1024
alignment_constraints = [8, ]
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([ 64, 128, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc),
]
data_type = [
math_inst.element_a,
math_inst.element_b,
math_inst.element_accumulator,
math_inst.element_accumulator,
]
CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints, complex_transforms)
# Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation)
if math_inst.element_a != math_inst.element_accumulator:
data_type_mixed = [
math_inst.element_a,
math_inst.element_b,
math_inst.element_a,
math_inst.element_accumulator,
]
CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, \
data_type_mixed, alignment_constraints, complex_transforms)
#
def GenerateSM80_TensorOp_16832_TN(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
]
math_instructions = [
MathInstruction( \
[16, 8, 32], \
DataType.s8, DataType.s8, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
MathInstruction( \
[16, 8, 32], \
DataType.u8, DataType.u8, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
]
min_cc = 80
max_cc = 1024
alignment_constraints = [16,]
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc),
TileDescription([64, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc),
TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([256, 64, 128], 3, [4, 1, 1], math_inst, min_cc, max_cc),
TileDescription([64, 256, 128], 3, [1, 4, 1], math_inst, min_cc, max_cc),
]
data_type = [math_inst.element_a, math_inst.element_b, DataType.s32, DataType.s32]
data_type_mixed = [math_inst.element_a, math_inst.element_b, DataType.s8, DataType.f32]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
operations = []
operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
for op in operations:
if op.tile_description.threadblock_shape[1] >= 128:
op.C.alignment = 16
else:
op.C.alignment = 8
#
#
def GenerateSM80_TensorOp_16832_Interleaved(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.ColumnMajorInterleaved32, LayoutType.RowMajorInterleaved32, LayoutType.ColumnMajorInterleaved32),
]
math_instructions = [
MathInstruction( \
[16, 8, 32], \
DataType.s8, DataType.s8, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
MathInstruction( \
[16, 8, 32], \
DataType.u8, DataType.u8, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
]
min_cc = 80
max_cc = 1024
alignment_constraints = [16,]
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc),
]
data_type_mixed = [math_inst.element_a, math_inst.element_b, DataType.s8, DataType.f32]
operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
for op in operations:
op.C.alignment = 8
#
#
def GenerateSM80_TensorOp_16864_TN(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
]
math_instructions = [
MathInstruction( \
[16, 8, 64], \
DataType.s4, DataType.s4, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
MathInstruction( \
[16, 8, 64], \
DataType.u4, DataType.u4, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
]
min_cc = 80
max_cc = 1024
alignment_constraints = [32,]
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([256, 128, 256], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 256], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 256], 5, [2, 2, 1], math_inst, min_cc, max_cc),
]
data_type = [math_inst.element_a, math_inst.element_b, DataType.s32, DataType.s32]
data_type_mixed = [math_inst.element_a, math_inst.element_b, DataType.s4, DataType.f32]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
operations = []
operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
for op in operations:
if op.tile_description.threadblock_shape[1] >= 128:
op.C.alignment = 8
elif op.tile_description.threadblock_shape[1] == 64:
op.C.alignment = 8
else:
op.C.alignment = 4
#
#
def GenerateSM80_TensorOp_16864_Interleaved(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.ColumnMajorInterleaved64, LayoutType.RowMajorInterleaved64, LayoutType.ColumnMajorInterleaved64),
]
math_instructions = [
MathInstruction( \
[16, 8, 64], \
DataType.s4, DataType.s4, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
MathInstruction( \
[16, 8, 64], \
DataType.u4, DataType.u4, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_saturate),
]
min_cc = 80
max_cc = 1024
alignment_constraints = [32,]
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc),
]
data_type_mixed = [math_inst.element_a, math_inst.element_b, DataType.s4, DataType.f32]
operations = []
operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
for op in operations:
op.C.alignment = 16
#
#
def GenerateSM80_TensorOp_168256(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
]
math_instructions = [
MathInstruction( \
[16, 8, 256], \
DataType.b1, DataType.b1, DataType.s32, \
OpcodeClass.TensorOp, \
MathOperation.xor_popc),
]
min_cc = 80
max_cc = 1024
alignment_constraints = [128,]
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([256, 128, 512], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 512], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 512], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 512], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 512], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 512], 5, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([256, 128, 1024], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 1024], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 1024], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 1024], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 1024], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 1024], 5, [2, 2, 1], math_inst, min_cc, max_cc),
]
data_type = [DataType.b1, DataType.b1, DataType.s32, DataType.s32]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
#
#
def GenerateSM80_TensorOp_1688(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
]
math_instructions = [
MathInstruction( \
[16, 8, 8], \
DataType.tf32, DataType.tf32, DataType.f32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add)
]
min_cc = 80
max_cc = 1024
alignment_constraints = [4, 2, 1]
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc),
TileDescription([256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 32], 3, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([128, 64, 32], 3, [2, 1, 2], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 32], 4, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([128, 64, 32], 4, [2, 1, 2], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 32], 4, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 32], 5, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 256, 32], 3, [1, 4, 1], math_inst, min_cc, max_cc),
]
data_type = [
math_inst.element_a,
math_inst.element_b,
math_inst.element_accumulator,
math_inst.element_accumulator,
]
data_type_mixed = [
math_inst.element_a,
math_inst.element_b,
math_inst.element_a,
math_inst.element_accumulator,
]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints)
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type_mixed, alignment_constraints)
#
#
def GenerateSM80_TensorOp_1688_fast_math(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
]
math_instructions = [
MathInstruction( \
[16, 8, 8], \
DataType.tf32, DataType.tf32, DataType.f32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add),
MathInstruction( \
[16, 8, 8], \
DataType.f16, DataType.f16, DataType.f32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_fast_f16),
MathInstruction( \
[16, 8, 8], \
DataType.bf16, DataType.bf16, DataType.f32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_fast_bf16)
]
min_cc = 80
max_cc = 1024
alignment_constraints = [4, 2, 1]
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc),
TileDescription([256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 32], 3, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([128, 64, 32], 3, [2, 1, 2], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 32], 4, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([128, 64, 32], 4, [2, 1, 2], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 32], 4, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 32], 5, [1, 2, 2], math_inst, min_cc, max_cc),
TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 256, 32], 3, [1, 4, 1], math_inst, min_cc, max_cc),
]
data_type = [DataType.f32, DataType.f32, DataType.f32, DataType.f32]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints)
#
#
def GenerateSM80_TensorOp_1688_complex(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
]
math_inst = MathInstruction( \
[16, 8, 8], \
DataType.f32, DataType.f32, DataType.f32, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_complex)
min_cc = 80
max_cc = 1024
tile_descriptions = [
TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([64, 128, 16], 4, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([64, 32, 16], 4, [2, 1, 1], math_inst, min_cc, max_cc),
TileDescription([32, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
]
data_type = [
DataType.cf32, DataType.cf32, DataType.cf32, DataType.cf32
]
alignment_constraints = [1,]
complex_transforms = [
(ComplexTransform.none, ComplexTransform.none),
(ComplexTransform.conj, ComplexTransform.none),
(ComplexTransform.none, ComplexTransform.conj),
(ComplexTransform.conj, ComplexTransform.conj)
]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints, complex_transforms)
#
#
def GenerateSM80_TensorOp_884(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
]
math_inst = \
MathInstruction( \
[8, 8, 4], \
DataType.f64, DataType.f64, DataType.f64, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add)
min_cc = 80
max_cc = 1024
alignment_constraints = [1,]
tile_descriptions = [
TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([32, 32, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([16, 32, 16], 5, [1, 2, 1], math_inst, min_cc, max_cc),
TileDescription([32, 16, 16], 5, [2, 1, 1], math_inst, min_cc, max_cc),
]
data_type = [DataType.f64, DataType.f64, DataType.f64, DataType.f64]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints)
#
#
def GenerateSM80_TensorOp_884_complex(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
]
math_inst = \
MathInstruction( \
[8, 8, 4], \
DataType.f64, DataType.f64, DataType.f64, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_complex)
min_cc = 80
max_cc = 1024
alignment_constraints = [1,]
tile_descriptions = [
TileDescription([128, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([64, 128, 8], 3, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([64, 64, 8], 3, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc),
TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc),
]
data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64]
complex_transforms = [
(ComplexTransform.none, ComplexTransform.none),
(ComplexTransform.conj, ComplexTransform.none),
(ComplexTransform.none, ComplexTransform.conj),
(ComplexTransform.conj, ComplexTransform.conj)
]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints, complex_transforms)
#
def GenerateSM80_TensorOp_884_complex_gaussian(manifest, args):
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
return
layouts = [
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
]
math_inst = \
MathInstruction( \
[8, 8, 4], \
DataType.f64, DataType.f64, DataType.f64, \
OpcodeClass.TensorOp, \
MathOperation.multiply_add_complex_gaussian)
min_cc = 80
max_cc = 1024
alignment_constraints = [1,]
tile_descriptions = [
TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc),
TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc),
]
data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64]
complex_transforms = [
(ComplexTransform.none, ComplexTransform.none),
(ComplexTransform.conj, ComplexTransform.none),
(ComplexTransform.none, ComplexTransform.conj),
(ComplexTransform.conj, ComplexTransform.conj)
]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints, complex_transforms)
#
###################################################################################################
#
def GenerateSM80_Simt(manifest, args):
layouts = [
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
]
math_instructions = [
MathInstruction( \
[1, 1, 1], \
DataType.f32, DataType.f32, DataType.f32, \
OpcodeClass.Simt, \
MathOperation.multiply_add),
]
min_cc = 80
max_cc = 1024
alignment_constraints = [1,]
for math_inst in math_instructions:
tile_descriptions = [
TileDescription([256, 128, 8], 5, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 8], 5, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 8], 5, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([256, 128, 8], 4, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 256, 8], 4, [2, 4, 1], math_inst, min_cc, max_cc),
TileDescription([128, 128, 8], 4, [4, 2, 1], math_inst, min_cc, max_cc),
TileDescription([128, 64, 8], 5, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 128, 8], 5, [2, 2, 1], math_inst, min_cc, max_cc),
TileDescription([ 64, 64, 8], 5, [2, 1, 1], math_inst, min_cc, max_cc),
TileDescription([128, 32, 8], 5, [2, 1, 1], math_inst, min_cc, max_cc),
TileDescription([ 32, 128, 8], 5, [1, 2, 1], math_inst, min_cc, max_cc),
]
data_type = [
math_inst.element_a,
math_inst.element_b,
math_inst.element_accumulator,
math_inst.element_accumulator,
]
CreateGemmOperator(manifest, layouts, tile_descriptions, \
data_type, alignment_constraints)
#
###################################################################################################
#
def GenerateSM80(manifest, args):
GenerateSM80_TensorOp_16816(manifest, args)
GenerateSM80_PlanarComplexTensorOp_16816(manifest, args)
GenerateSM80_TensorOp_1688(manifest, args)
GenerateSM80_TensorOp_1688_fast_math(manifest, args)
GenerateSM80_TensorOp_1688_complex(manifest, args)
GenerateSM80_TensorOp_884(manifest, args)
GenerateSM80_TensorOp_884_complex(manifest, args)
GenerateSM80_TensorOp_884_complex_gaussian(manifest, args)
GenerateSM80_TensorOp_16832_TN(manifest, args)
GenerateSM80_TensorOp_16832_Interleaved(manifest, args)
GenerateSM80_TensorOp_16864_TN(manifest, args)
GenerateSM80_TensorOp_16864_Interleaved(manifest, args)
GenerateSM80_TensorOp_168256(manifest, args)
GenerateSM80_Simt(manifest, args)
#
###################################################################################################
if __name__ == "__main__":
@@ -920,7 +1766,7 @@ if __name__ == "__main__":
parser.add_argument("--build-dir", default=".", required=False, help="CUTLASS top-level build directory")
parser.add_argument("--curr-build-dir", default=".", help="CUTLASS current build directory. cmake files will be emitted in this directory")
parser.add_argument("--generator-target", default='library', help="Target of CUTLASS Library Generator.")
parser.add_argument("--architectures", default='50;60;61;75', help="Target compute architectures")
parser.add_argument("--architectures", default='53;60;61;70;75;80', help="Target compute architectures")
parser.add_argument("--kernels", default='', help='Comma delimited list to filter kernels by name.')
parser.add_argument("--cuda-version", default="11.0.0", help="Semantic version string of CUDA Toolkit")
@@ -933,6 +1779,8 @@ if __name__ == "__main__":
GenerateSM61(manifest, args)
GenerateSM70(manifest, args)
GenerateSM75(manifest, args)
GenerateSM80(manifest, args)
if 'library' in args.generator_target.split(','):
manifest.emit(GeneratorTarget.Library)
+112 -66
View File
@@ -4,14 +4,32 @@
# \brief Generates the CUTLASS Library's instances
#
import enum
import re
###################################################################################################
import enum
# The following block implements enum.auto() for Python 3.5 variants that don't include it such
# as the default 3.5.2 on Ubuntu 16.04.
#
# https://codereview.stackexchange.com/questions/177309/reimplementing-pythons-enum-auto-for-compatibility
try:
from enum import auto as enum_auto
except ImportError:
__cutlass_library_auto_enum = 0
def enum_auto() -> int:
global __cutlass_library_auto_enum
i = __cutlass_library_auto_enum
__cutlass_library_auto_enum += 1
return i
###################################################################################################
#
class GeneratorTarget(enum.Enum):
Library = enum.auto()
Library = enum_auto()
#
GeneratorTargetNames = {
GeneratorTarget.Library: 'library'
@@ -22,33 +40,37 @@ GeneratorTargetNames = {
#
class DataType(enum.Enum):
b1 = enum.auto()
u4 = enum.auto()
u8 = enum.auto()
u16 = enum.auto()
u32 = enum.auto()
u64 = enum.auto()
s4 = enum.auto()
s8 = enum.auto()
s16 = enum.auto()
s32 = enum.auto()
s64 = enum.auto()
f16 = enum.auto()
f32 = enum.auto()
f64 = enum.auto()
cf16 = enum.auto()
cf32 = enum.auto()
cf64 = enum.auto()
cs4 = enum.auto()
cs8 = enum.auto()
cs16 = enum.auto()
cs32 = enum.auto()
cs64 = enum.auto()
cu4 = enum.auto()
cu8 = enum.auto()
cu16 = enum.auto()
cu32 = enum.auto()
cu64 = enum.auto()
b1 = enum_auto()
u4 = enum_auto()
u8 = enum_auto()
u16 = enum_auto()
u32 = enum_auto()
u64 = enum_auto()
s4 = enum_auto()
s8 = enum_auto()
s16 = enum_auto()
s32 = enum_auto()
s64 = enum_auto()
f16 = enum_auto()
bf16 = enum_auto()
f32 = enum_auto()
tf32 = enum_auto()
f64 = enum_auto()
cf16 = enum_auto()
cbf16 = enum_auto()
cf32 = enum_auto()
ctf32 = enum_auto()
cf64 = enum_auto()
cs4 = enum_auto()
cs8 = enum_auto()
cs16 = enum_auto()
cs32 = enum_auto()
cs64 = enum_auto()
cu4 = enum_auto()
cu8 = enum_auto()
cu16 = enum_auto()
cu32 = enum_auto()
cu64 = enum_auto()
#
ShortDataTypeNames = {
@@ -74,10 +96,14 @@ DataTypeNames = {
DataType.s32: "s32",
DataType.s64: "s64",
DataType.f16: "f16",
DataType.bf16: "bf16",
DataType.f32: "f32",
DataType.tf32: "tf32",
DataType.f64: "f64",
DataType.cf16: "cf16",
DataType.cbf16: "cbf16",
DataType.cf32: "cf32",
DataType.ctf32: "ctf32",
DataType.cf64: "cf64",
DataType.cu4: "cu4",
DataType.cu8: "cu8",
@@ -104,10 +130,14 @@ DataTypeTag = {
DataType.s32: "int32_t",
DataType.s64: "int64_t",
DataType.f16: "cutlass::half_t",
DataType.bf16: "cutlass::bfloat16_t",
DataType.f32: "float",
DataType.tf32: "cutlass::tfloat32_t",
DataType.f64: "double",
DataType.cf16: "cutlass::complex<cutlass::half_t>",
DataType.cbf16: "cutlass::complex<cutlass::bfloat16_t>",
DataType.cf32: "cutlass::complex<float>",
DataType.ctf32: "cutlass::complex<cutlass::tfloat32_t>",
DataType.cf64: "cutlass::complex<double>",
DataType.cu4: "cutlass::complex<cutlass::uint4b_t>",
DataType.cu8: "cutlass::complex<cutlass::uint8_t>",
@@ -134,10 +164,14 @@ DataTypeSize = {
DataType.s32: 32,
DataType.s64: 64,
DataType.f16: 16,
DataType.bf16: 16,
DataType.f32: 32,
DataType.tf32: 32,
DataType.f64: 64,
DataType.cf16: 32,
DataType.cbf16: 32,
DataType.cf32: 64,
DataType.ctf32: 32,
DataType.cf64: 128,
DataType.cu4: 8,
DataType.cu8: 16,
@@ -155,8 +189,8 @@ DataTypeSize = {
#
class ComplexTransform(enum.Enum):
none = enum.auto()
conj = enum.auto()
none = enum_auto()
conj = enum_auto()
#
ComplexTransformTag = {
@@ -194,40 +228,47 @@ def get_real_from_complex(complex_type):
#
class ComplexMultiplyOp(enum.Enum):
multiply_add = enum.auto()
gaussian = enum.auto()
multiply_add = enum_auto()
gaussian = enum_auto()
###################################################################################################
#
class MathOperation(enum.Enum):
multiply_add = enum.auto()
multiply_add_saturate = enum.auto()
xor_popc = enum.auto()
multiply_add_complex = enum.auto()
multiply_add = enum_auto()
multiply_add_saturate = enum_auto()
xor_popc = enum_auto()
multiply_add_fast_bf16 = enum_auto()
multiply_add_fast_f16 = enum_auto()
multiply_add_complex = enum_auto()
multiply_add_complex_gaussian = enum_auto()
#
MathOperationTag = {
MathOperation.multiply_add: 'cutlass::arch::OpMultiplyAdd',
MathOperation.multiply_add_saturate: 'cutlass::arch::OpMultiplyAddSaturate',
MathOperation.xor_popc: 'cutlass::arch::OpXorPopc',
MathOperation.multiply_add_fast_bf16: 'cutlass::arch::OpMultiplyAddFastBF16',
MathOperation.multiply_add_fast_f16: 'cutlass::arch::OpMultiplyAddFastF16',
MathOperation.multiply_add_complex: 'cutlass::arch::OpMultiplyAddComplex',
MathOperation.multiply_add_complex_gaussian: 'cutlass::arch::OpMultiplyAddGaussianComplex',
}
###################################################################################################
#
class LayoutType(enum.Enum):
ColumnMajor = enum.auto()
RowMajor = enum.auto()
ColumnMajorInterleaved32 = enum.auto()
RowMajorInterleaved32 = enum.auto()
ColumnMajorInterleaved64 = enum.auto()
RowMajorInterleaved64 = enum.auto()
TensorNHWC = enum.auto()
TensorNCHW = enum.auto()
TensorNGHWC = enum.auto()
TensorNCxHW32 = enum.auto()
TensorNCxHW64 = enum.auto()
ColumnMajor = enum_auto()
RowMajor = enum_auto()
ColumnMajorInterleaved32 = enum_auto()
RowMajorInterleaved32 = enum_auto()
ColumnMajorInterleaved64 = enum_auto()
RowMajorInterleaved64 = enum_auto()
TensorNHWC = enum_auto()
TensorNCHW = enum_auto()
TensorNGHWC = enum_auto()
TensorNCxHW32 = enum_auto()
TensorNCxHW64 = enum_auto()
#
LayoutTag = {
@@ -282,9 +323,9 @@ ShortComplexLayoutNames = {
#
class OpcodeClass(enum.Enum):
Simt = enum.auto()
TensorOp = enum.auto()
WmmaTensorOp = enum.auto()
Simt = enum_auto()
TensorOp = enum_auto()
WmmaTensorOp = enum_auto()
OpcodeClassNames = {
OpcodeClass.Simt: 'simt',
@@ -302,7 +343,7 @@ OpcodeClassTag = {
#
class OperationKind(enum.Enum):
Gemm = enum.auto()
Gemm = enum_auto()
#
OperationKindNames = {
OperationKind.Gemm: 'gemm'
@@ -310,7 +351,7 @@ OperationKindNames = {
#
class Target(enum.Enum):
library = enum.auto()
library = enum_auto()
ArchitectureNames = {
50: 'maxwell',
@@ -318,6 +359,7 @@ ArchitectureNames = {
61: 'pascal',
70: 'volta',
75: 'turing',
80: 'ampere',
}
###################################################################################################
@@ -340,27 +382,27 @@ def SubstituteTemplate(template, values):
#
class GemmKind(enum.Enum):
Gemm = enum.auto()
Batched = enum.auto()
Array = enum.auto()
Universal = enum.auto()
PlanarComplex = enum.auto()
PlanarComplexArray = enum.auto()
Gemm = enum_auto()
Batched = enum_auto()
Array = enum_auto()
Universal = enum_auto()
PlanarComplex = enum_auto()
PlanarComplexArray = enum_auto()
#
GemmKindNames = {
GemmKind.Gemm: "gemm",
GemmKind.Batched: "gemm_batched",
GemmKind.Array: "gemm_array",
GemmKind.Universal: "gemm_universal",
GemmKind.Universal: "gemm",
GemmKind.PlanarComplex: "gemm_planar_complex",
GemmKind.PlanarComplexArray: "gemm_planar_complex_array",
}
#
class EpilogueFunctor(enum.Enum):
LinearCombination = enum.auto()
LinearCombinationClamp = enum.auto()
LinearCombination = enum_auto()
LinearCombinationClamp = enum_auto()
#
EpilogueFunctorTag = {
@@ -370,13 +412,17 @@ EpilogueFunctorTag = {
#
class SwizzlingFunctor(enum.Enum):
Cohort = enum.auto()
Identity = enum.auto()
Identity1 = enum_auto()
Identity2 = enum_auto()
Identity4 = enum_auto()
Identity8 = enum_auto()
#
SwizzlingFunctorTag = {
SwizzlingFunctor.Cohort: 'cutlass::gemm::threadblock::GemmCohortThreadblockSwizzle<${layout_a}, ${layout_b}>',
SwizzlingFunctor.Identity: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle',
SwizzlingFunctor.Identity1: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>',
SwizzlingFunctor.Identity2: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<2>',
SwizzlingFunctor.Identity4: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<4>',
SwizzlingFunctor.Identity8: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>',
}
###################################################################################################
+1 -1
View File
@@ -127,7 +127,7 @@ class Manifest:
if args.kernels == 'all':
self.kernel_names = []
else:
self.kernel_names = args.kernels.split(',')
self.kernel_names = [x for x in args.kernels.split(',') if x != '']
self.operation_count = 0
self.operations_by_name = {}
+200 -4
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
@@ -29,13 +29,14 @@
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/gemm/kernel/default_gemm_planar_complex_universal.h"
#include "cutlass/gemm/device/gemm.h"
#include "cutlass/gemm/device/gemm_complex.h"
#include "cutlass/gemm/device/gemm_batched.h"
#include "cutlass/gemm/device/gemm_array.h"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/default_gemm_universal.h"
#include "cutlass/gemm/kernel/default_gemm_planar_complex_universal.h"
#include "cutlass/library/library.h"
#include "library_internal.h"
@@ -104,10 +105,10 @@ public:
MathOperationMap<typename Operator::Operator>::kId;
description_.tile_description.minimum_compute_capability =
ArchMap<typename Operator::ArchTag>::kMin;
ArchMap<typename Operator::ArchTag, typename Operator::OperatorClass>::kMin;
description_.tile_description.maximum_compute_capability =
ArchMap<typename Operator::ArchTag>::kMax;
ArchMap<typename Operator::ArchTag, typename Operator::OperatorClass>::kMax;
description_.A = make_TensorDescription<ElementA, LayoutA>(Operator::kAlignmentA);
description_.B = make_TensorDescription<ElementB, LayoutB>(Operator::kAlignmentB);
@@ -698,6 +699,201 @@ public:
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename Operator_>
class GemmUniversalOperation : public GemmOperationBase<Operator_> {
public:
using Operator = Operator_;
using ElementA = typename Operator::ElementA;
using LayoutA = typename Operator::LayoutA;
using ElementB = typename Operator::ElementB;
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementAccumulator = typename Operator::ElementAccumulator;
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
using OperatorArguments = typename Operator::Arguments;
public:
/// Constructor
GemmUniversalOperation(char const *name = "unknown_gemm"):
GemmOperationBase<Operator_>(name) {
this->description_.gemm_kind = GemmKind::kUniversal;
}
protected:
/// Constructs the arguments structure given the configuration and arguments
static Status construct_arguments_(
OperatorArguments &operator_args,
GemmUniversalConfiguration const *configuration) {
operator_args.mode = configuration->mode;
operator_args.problem_size = configuration->problem_size;
operator_args.batch_count = configuration->batch_count;
operator_args.lda = int(configuration->lda);
operator_args.ldb = int(configuration->ldb);
operator_args.ldc = int(configuration->ldc);
operator_args.ldd = int(configuration->ldd);
return Status::kSuccess;
}
/// Constructs the arguments structure given the configuration and arguments
static Status update_arguments_(
OperatorArguments &operator_args,
GemmUniversalArguments const *arguments) {
if (arguments->pointer_mode == ScalarPointerMode::kHost) {
typename Operator::EpilogueOutputOp::Params params(
*static_cast<ElementCompute const *>(arguments->alpha),
*static_cast<ElementCompute const *>(arguments->beta)
);
operator_args.epilogue = params;
}
else if (arguments->pointer_mode == ScalarPointerMode::kDevice){
typename Operator::EpilogueOutputOp::Params params(
static_cast<ElementCompute const *>(arguments->alpha),
static_cast<ElementCompute const *>(arguments->beta)
);
operator_args.epilogue = params;
}
else {
return Status::kErrorInvalidProblem;
}
// update arguments
operator_args.ptr_A = arguments->A;
operator_args.ptr_B = arguments->B;
operator_args.ptr_C = arguments->C;
operator_args.ptr_D = arguments->D;
operator_args.batch_stride_A = arguments->batch_stride_A;
operator_args.batch_stride_B = arguments->batch_stride_B;
operator_args.batch_stride_C = arguments->batch_stride_C;
operator_args.batch_stride_D = arguments->batch_stride_D;
return Status::kSuccess;
}
public:
/// Returns success if the operation can proceed
virtual Status can_implement(
void const *configuration_ptr,
void const *arguments_ptr) const {
GemmUniversalConfiguration const *configuration =
static_cast<GemmUniversalConfiguration const *>(configuration_ptr);
GemmUniversalArguments const *arguments =
static_cast<GemmUniversalArguments const *>(arguments_ptr);
OperatorArguments args;
Status status = construct_arguments_(args, configuration);
if (status != Status::kSuccess) {
return status;
}
status = update_arguments_(args, arguments);
if (status != Status::kSuccess) {
return status;
}
return Operator::can_implement(args);
}
/// Gets the host-side workspace
virtual uint64_t get_host_workspace_size(
void const *configuration) const {
return sizeof(Operator);
}
/// Gets the device-side workspace
virtual uint64_t get_device_workspace_size(
void const *configuration_ptr) const {
OperatorArguments args;
Status status = construct_arguments_(
args,
static_cast<GemmUniversalConfiguration const *>(configuration_ptr));
if (status != Status::kSuccess) {
return 0;
}
uint64_t size = Operator::get_workspace_size(args);
return size;
}
/// Initializes the workspace
virtual Status initialize(
void const *configuration_ptr,
void *host_workspace,
void *device_workspace,
cudaStream_t stream = nullptr) const {
OperatorArguments args;
Status status = construct_arguments_(
args,
static_cast<GemmUniversalConfiguration const *>(configuration_ptr));
if (status != Status::kSuccess) {
return status;
}
Operator *op = new (host_workspace) Operator;
status = op->initialize(args, device_workspace, stream);
return status;
}
/// Runs the kernel
virtual Status run(
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
OperatorArguments args;
Status status = update_arguments_(
args,
static_cast<GemmUniversalArguments const *>(arguments_ptr));
if (status != Status::kSuccess) {
return status;
}
Operator *op = static_cast<Operator *>(host_workspace);
status = op->update(args, device_workspace);
if (status != Status::kSuccess) {
return status;
}
status = op->run(stream);
return status;
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
template <typename Operator_>
+206 -7
View File
@@ -26,7 +26,7 @@
/*! \file
\brief CUTLASS Library handle.
*/
#include <iostream>
#include <stdexcept>
#include <cstdint>
@@ -43,7 +43,8 @@ namespace library {
Handle::Handle(
cudaStream_t stream,
size_t workspace_size
):
):
provider_(Provider::kCUTLASS),
stream_(stream),
workspace_(nullptr),
workspace_size_(0),
@@ -95,6 +96,7 @@ Handle::Handle(Handle && handle) {
/// Move assignment operator
Handle & Handle::operator=(Handle && handle) {
provider_ = handle.provider_;
device_ = handle.device_;
workspace_size_ = handle.workspace_size_;
workspace_ = handle.workspace_;
@@ -121,6 +123,16 @@ cudaStream_t Handle::get_stream() const {
return stream_;
}
/// Gets the current provider
Provider Handle::get_provider() const {
return provider_;
}
/// Sets the provider of operations
void Handle::set_provider(Provider provider) {
provider_ = provider;
}
/// Gets the device workspace size
size_t Handle::get_workspace_size() const {
return workspace_size_;
@@ -351,6 +363,8 @@ Status Handle::gemm(
//
GemmFunctionalKey key(
provider_,
GemmKind::kGemm,
element_compute,
element_scalar,
element_A,
@@ -457,6 +471,188 @@ Status Handle::gemm(
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Executes a GEMM computation: D <= alpha * A*B + beta * C.
//
// Supports batched-strided, batched array or split-K serial or split-K parallel.
//
Status Handle::gemm_universal(
GemmUniversalMode mode, /// indicates the mode in which the kUniversal GEMM is launched
int M, /// GEMM M dimension
int N, /// GEMM N dimension
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
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
int 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
int ldb, /// Leading dimension of B matrix
void const * beta, /// Pointer to beta scalar
NumericTypeID element_C, /// Data type of C and D matrices
void const * ptr_C, /// Pointer to C matrix
int ldc, /// Leading dimension of C matrix
void * ptr_D, /// Pointer to D matrix
int ldd, /// Leading dimension of D matrix
int batch_count, /// Batch count or number of split-K slices
int64_t batch_stride_A, /// Batch stride of A operand
int64_t batch_stride_B, /// Batch stride of B operand
int64_t batch_stride_C, /// Batch stride of C operand
int64_t batch_stride_D /// Batch stride of D operand
) {
//
// Find the operation
//
GemmFunctionalKey key(
provider_,
GemmKind::kUniversal,
element_compute,
element_scalar,
element_A,
layout_A,
transform_A,
element_B,
layout_B,
transform_B,
element_C
);
auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
if (operators_it == Singleton::get().operation_table.gemm_operations.end()) {
return cutlass::Status::kErrorNotSupported;
}
if (operators_it->second.empty()) {
return cutlass::Status::kErrorNotSupported;
}
//
// Compute the largest alignment restriction the kernel can satisfy.
//
// Maximum alignment expectation among all kernels (in units of bytes)
int const kMaximumAlignmentSize = 16;
void const *ptr_A_check = ptr_A;
void const *ptr_B_check = ptr_B;
void const *ptr_C_check = ptr_C;
void * ptr_D_check = ptr_D;
// Ignore alignment of pointers to pointers. We can't check this from the host,
// as each batch index has its own pointer in device memory.
if (mode == GemmUniversalMode::kArray) {
ptr_A_check = nullptr;
ptr_B_check = nullptr;
ptr_C_check = nullptr;
ptr_D_check = nullptr;
}
int alignment = gemm_problem_alignment(
M, N, K,
element_A, ptr_A_check, lda, 0,
element_B, ptr_B_check, ldb, 0,
element_C, ptr_C_check, ldc, 0,
ptr_D_check, ldd, 0, kMaximumAlignmentSize
);
//
// Find the best kernel in descending order of preference.
//
GemmPreferenceKey preference_key(compute_capability(), alignment);
Operation const *operation = find_gemm_operation(operators_it, preference_key);
if (!operation) {
return cutlass::Status::kErrorNotSupported;
}
last_operation_ = operation;
//
// Configure operation
//
GemmUniversalConfiguration configuration{
mode,
{M, N, K},
batch_count,
lda,
ldb,
ldc,
ldd
};
// Query host work space size
uint64_t host_workspace_size_needed = operation->get_host_workspace_size(&configuration);
if (uint64_t(kHostWorkspaceSize) < host_workspace_size_needed) {
return cutlass::Status::kErrorNotSupported;
}
char host_workspace[kHostWorkspaceSize];
// Query device workspace size
uint64_t device_workspace_size_needed = operation->get_device_workspace_size(&configuration);
if (uint64_t(workspace_size_) < device_workspace_size_needed) {
return cutlass::Status::kErrorNotSupported;
}
// Initialize host and device workspaces
Status status = operation->initialize(
&configuration,
host_workspace,
workspace_,
stream_);
if (status != cutlass::Status::kSuccess) {
return status;
}
// Run the operator
GemmUniversalArguments arguments{
ptr_A,
ptr_B,
ptr_C,
ptr_D,
alpha,
beta,
scalar_pointer_mode_,
batch_stride_A,
batch_stride_B,
batch_stride_C,
batch_stride_D
};
return operation->run(&arguments, host_workspace, workspace_, stream_);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Planar complex GEMM
Status Handle::gemm_planar_complex(
@@ -522,6 +718,8 @@ Status Handle::gemm_planar_complex(
//
GemmFunctionalKey key(
provider_,
GemmKind::kPlanarComplex,
element_compute,
element_scalar,
element_A,
@@ -533,9 +731,9 @@ Status Handle::gemm_planar_complex(
element_C
);
auto operators_it = Singleton::get().operation_table.gemm_planar_complex_operations.find(key);
auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
if (operators_it == Singleton::get().operation_table.gemm_planar_complex_operations.end()) {
if (operators_it == Singleton::get().operation_table.gemm_operations.end()) {
return cutlass::Status::kErrorNotSupported;
}
@@ -714,6 +912,8 @@ Status Handle::gemm_planar_complex_array(
//
GemmFunctionalKey key(
provider_,
GemmKind::kPlanarComplexArray,
element_compute,
element_scalar,
element_A,
@@ -725,9 +925,9 @@ Status Handle::gemm_planar_complex_array(
element_C
);
auto operators_it = Singleton::get().operation_table.gemm_planar_complex_array_operations.find(key);
auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
if (operators_it == Singleton::get().operation_table.gemm_planar_complex_array_operations.end()) {
if (operators_it == Singleton::get().operation_table.gemm_operations.end()) {
return cutlass::Status::kErrorNotSupported;
}
@@ -837,7 +1037,6 @@ Status Handle::gemm_planar_complex_array(
}
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace library
} // namespace cutlass
+29 -7
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
@@ -125,6 +125,14 @@ template <> struct NumericTypeMap<cutlass::complex<double> > {
static NumericTypeID const kId = NumericTypeID::kCF64;
};
template <> struct NumericTypeMap<cutlass::bfloat16_t> {
static NumericTypeID const kId = NumericTypeID::kBF16;
};
template <> struct NumericTypeMap<cutlass::tfloat32_t> {
static NumericTypeID const kId = NumericTypeID::kTF32;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T> struct MathOperationMap {
@@ -143,6 +151,10 @@ template <> struct MathOperationMap<cutlass::arch::OpMultiplyAddComplex> {
static MathOperationID const kId = MathOperationID::kMultiplyAddComplex;
};
template <> struct MathOperationMap<cutlass::arch::OpMultiplyAddGaussianComplex> {
static MathOperationID const kId = MathOperationID::kMultiplyAddGaussianComplex;
};
template <> struct MathOperationMap<cutlass::arch::OpXorPopc> {
static MathOperationID const kId = MathOperationID::kXorPopc;
};
@@ -217,33 +229,43 @@ template <> struct ComplexTransformMap<cutlass::ComplexTransform::kConjugate> {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T> struct ArchMap;
template <typename ArchTag, typename OperatorClass> struct ArchMap;
template <> struct ArchMap<arch::Sm50> {
template <> struct ArchMap<arch::Sm50, arch::OpClassSimt> {
static int const kMin = 50;
static int const kMax = 1024;
};
template <> struct ArchMap<arch::Sm60> {
template <> struct ArchMap<arch::Sm60, arch::OpClassSimt> {
static int const kMin = 60;
static int const kMax = 1024;
};
template <> struct ArchMap<arch::Sm61> {
template <> struct ArchMap<arch::Sm61, arch::OpClassSimt> {
static int const kMin = 61;
static int const kMax = 1024;
};
template <> struct ArchMap<arch::Sm70> {
template <> struct ArchMap<arch::Sm70, arch::OpClassWmmaTensorOp> {
static int const kMin = 70;
static int const kMax = 1024;
};
template <> struct ArchMap<arch::Sm70, arch::OpClassTensorOp> {
static int const kMin = 70;
static int const kMax = 75;
};
template <> struct ArchMap<arch::Sm75> {
template <typename OperatorClass> struct ArchMap<arch::Sm75, OperatorClass> {
static int const kMin = 75;
static int const kMax = 1024;
};
template <typename OperatorClass> struct ArchMap<arch::Sm80, OperatorClass> {
static int const kMin = 80;
static int const kMax = 1024;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename Element, typename Layout>
+3 -13
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
@@ -37,11 +37,6 @@ namespace library {
//////////////////////////////////////////////////////////////////////////////////////////////////////////
// init and insert all cutlass op in manifest object (procedurally generated using generator.py)
void initialize_all(Manifest &manifest);
/////////////////////////////////////////////////////////////////////////////////////////////////////////
/// Top-level initialization
Status Manifest::initialize() {
@@ -49,13 +44,8 @@ Status Manifest::initialize() {
operations_.clear();
}
switch(provider_) {
case Provider::kCUTLASS:
initialize_all(*this); break;
default:
break;
}
// initialize procedurally generated cutlass op in manifest object
initialize_all(*this);
return Status::kSuccess;
}
+22 -92
View File
@@ -28,30 +28,7 @@
instances may be queried.
*/
#include <fstream>
#include "cutlass/library/library.h"
#include "cutlass/library/operation_table.h"
#include "cutlass/library/util.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
std::ostream & operator<<(std::ostream &out, cutlass::library::GemmFunctionalKey const &k) {
out << "{\n"
<< " element_compute: " << to_string(k.element_compute) << "\n"
<< " element_scalar: " << to_string(k.element_scalar) << "\n"
<< " element_A: " << to_string(k.element_A) << "\n"
<< " layout_A: " << to_string(k.layout_A) << "\n"
<< " transform_A: " << to_string(k.transform_A) << "\n"
<< " element_B: " << to_string(k.element_B) << "\n"
<< " layout_B: " << to_string(k.layout_B) << "\n"
<< " transform_B: " << to_string(k.transform_B) << "\n"
<< " element_C: " << to_string(k.element_C) << "\n"
<< "}";
return out;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -67,85 +44,38 @@ void OperationTable::append(Manifest const &manifest) {
OperationDescription const &desc = operation->description();
// insert all gemm operation into operation table
if (desc.kind == OperationKind::kGemm) {
GemmDescription const &gemm_desc = static_cast<GemmDescription const &>(desc);
if (gemm_desc.gemm_kind == GemmKind::kGemm) {
GemmFunctionalKey functional_key(
gemm_desc.tile_description.math_instruction.element_accumulator,
gemm_desc.element_epilogue,
gemm_desc.A.element,
gemm_desc.A.layout,
gemm_desc.transform_A,
gemm_desc.B.element,
gemm_desc.B.layout,
gemm_desc.transform_B,
gemm_desc.C.element
);
GemmFunctionalKey functional_key(
gemm_desc.provider,
gemm_desc.gemm_kind,
gemm_desc.tile_description.math_instruction.element_accumulator,
gemm_desc.element_epilogue,
gemm_desc.A.element,
gemm_desc.A.layout,
gemm_desc.transform_A,
gemm_desc.B.element,
gemm_desc.B.layout,
gemm_desc.transform_B,
gemm_desc.C.element
);
Operation const *op = operation.get();
Operation const *op = operation.get();
int cc = gemm_desc.tile_description.minimum_compute_capability;
int cc = gemm_desc.tile_description.minimum_compute_capability;
int alignment = std::max(std::max(
gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
int alignment = std::max(std::max(
gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
GemmPreferenceKey preference_key(cc, alignment);
GemmPreferenceKey preference_key(cc, alignment);
gemm_operations[functional_key][preference_key].push_back(op);
}
else if (gemm_desc.gemm_kind == GemmKind::kPlanarComplex) {
GemmFunctionalKey functional_key(
gemm_desc.tile_description.math_instruction.element_accumulator,
gemm_desc.element_epilogue,
gemm_desc.A.element,
gemm_desc.A.layout,
gemm_desc.transform_A,
gemm_desc.B.element,
gemm_desc.B.layout,
gemm_desc.transform_B,
gemm_desc.C.element
);
Operation const *op = operation.get();
int cc = gemm_desc.tile_description.minimum_compute_capability;
int alignment = std::max(std::max(
gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
GemmPreferenceKey preference_key(cc, alignment);
gemm_planar_complex_operations[functional_key][preference_key].push_back(op);
}
else if (gemm_desc.gemm_kind == GemmKind::kPlanarComplexArray) {
GemmFunctionalKey functional_key(
gemm_desc.tile_description.math_instruction.element_accumulator,
gemm_desc.element_epilogue,
gemm_desc.A.element,
gemm_desc.A.layout,
gemm_desc.transform_A,
gemm_desc.B.element,
gemm_desc.B.layout,
gemm_desc.transform_B,
gemm_desc.C.element
);
Operation const *op = operation.get();
int cc = gemm_desc.tile_description.minimum_compute_capability;
int alignment = std::max(std::max(
gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
GemmPreferenceKey preference_key(cc, alignment);
gemm_planar_complex_array_operations[functional_key][preference_key].push_back(op);
}
gemm_operations[functional_key][preference_key].push_back(op);
}
}
}
+203 -1
View File
@@ -45,6 +45,7 @@ static struct {
Provider enumerant;
}
Provider_enumerants[] = {
{"none", "None", Provider::kNone},
{"cutlass", "CUTLASS", Provider::kCUTLASS},
{"host", "reference_host", Provider::kReferenceHost},
{"device", "reference_device", Provider::kReferenceDevice},
@@ -83,6 +84,38 @@ Provider from_string<Provider>(std::string const &str) {
}
///////////////////////////////////////////////////////////////////////////////////////////////////
static struct {
char const *text;
char const *pretty;
GemmKind enumerant;
}
GemmKind_enumerants[] = {
{"gemm", "<Gemm>", GemmKind::kGemm},
{"batched", "<Batched>", GemmKind::kBatched},
{"array", "<Array>", GemmKind::kArray},
{"universal", "<Universal>", GemmKind::kUniversal},
{"planar_complex", "<PlanarComplex>", GemmKind::kPlanarComplex},
{"planar_complex_array", "<PlanarComplexArray>", GemmKind::kPlanarComplexArray},
};
/// Converts a ConvKind enumerant to a string
char const *to_string(GemmKind type, bool pretty) {
for (auto const & possible : GemmKind_enumerants) {
if (type == possible.enumerant) {
if (pretty) {
return possible.pretty;
}
else {
return possible.text;
}
}
}
return pretty ? "Invalid" : "invalid";
}
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -92,6 +125,7 @@ static struct {
OperationKind enumerant;
}
OperationKind_enumerants[] = {
{"eq_gemm", "EqGemm", OperationKind::kEqGemm},
{"gemm", "Gemm", OperationKind::kGemm},
};
@@ -194,10 +228,14 @@ NumericTypeID_enumerants[] = {
{"s32", "S32", NumericTypeID::kS32},
{"s64", "S64", NumericTypeID::kS64},
{"f16", "F16", NumericTypeID::kF16},
{"bf16", "BF16", NumericTypeID::kBF16},
{"f32", "F32", NumericTypeID::kF32},
{"tf32", "TF32", NumericTypeID::kTF32},
{"f64", "F64", NumericTypeID::kF64},
{"cf16", "CF16", NumericTypeID::kCF16},
{"cbf16", "CBF16", NumericTypeID::kCBF16},
{"cf32", "CF32", NumericTypeID::kCF32},
{"ctf32", "CTF32", NumericTypeID::kCTF32},
{"cf64", "CF64", NumericTypeID::kCF64},
{"cu4", "CU4", NumericTypeID::kCU4},
{"cu8", "CU8", NumericTypeID::kCU8},
@@ -249,10 +287,14 @@ NumericTypeID from_string<NumericTypeID>(std::string const &str) {
int sizeof_bits(NumericTypeID type) {
switch (type) {
case NumericTypeID::kF16: return 16;
case NumericTypeID::kBF16: return 16;
case NumericTypeID::kTF32: return 32;
case NumericTypeID::kF32: return 32;
case NumericTypeID::kF64: return 64;
case NumericTypeID::kCF16: return 32;
case NumericTypeID::kCBF16: return 32;
case NumericTypeID::kCF32: return 64;
case NumericTypeID::kCTF32: return 64;
case NumericTypeID::kCF64: return 128;
case NumericTypeID::kS4: return 4;
case NumericTypeID::kS8: return 8;
@@ -276,6 +318,8 @@ bool is_complex_type(NumericTypeID type) {
case NumericTypeID::kCF16: return true;
case NumericTypeID::kCF32: return true;
case NumericTypeID::kCF64: return true;
case NumericTypeID::kCBF16: return true;
case NumericTypeID::kCTF32: return true;
default: break;
}
return false;
@@ -287,6 +331,8 @@ NumericTypeID get_real_type(NumericTypeID type) {
case NumericTypeID::kCF16: return NumericTypeID::kF16;
case NumericTypeID::kCF32: return NumericTypeID::kF32;
case NumericTypeID::kCF64: return NumericTypeID::kF64;
case NumericTypeID::kCBF16: return NumericTypeID::kBF16;
case NumericTypeID::kCTF32: return NumericTypeID::kTF32;
default: break;
}
return type;
@@ -314,6 +360,8 @@ bool is_integer_type(NumericTypeID type) {
bool is_signed_type(NumericTypeID type) {
switch (type) {
case NumericTypeID::kF16: return true;
case NumericTypeID::kBF16: return true;
case NumericTypeID::kTF32: return true;
case NumericTypeID::kF32: return true;
case NumericTypeID::kF64: return true;
case NumericTypeID::kS4: return true;
@@ -340,9 +388,13 @@ bool is_unsigned_integer(NumericTypeID type) {
bool is_float_type(NumericTypeID type) {
switch (type) {
case NumericTypeID::kF16: return true;
case NumericTypeID::kBF16: return true;
case NumericTypeID::kTF32: return true;
case NumericTypeID::kF32: return true;
case NumericTypeID::kF64: return true;
case NumericTypeID::kCF16: return true;
case NumericTypeID::kCBF16: return true;
case NumericTypeID::kCTF32: return true;
case NumericTypeID::kCF32: return true;
case NumericTypeID::kCF64: return true;
default: break;
@@ -431,7 +483,7 @@ OpcodeClassID_enumerants[] = {
{"simt", "<simt>", OpcodeClassID::kSimt},
{"tensorop", "<tensorop>", OpcodeClassID::kTensorOp},
{"wmmatensorop", "<wmmatensorop>", OpcodeClassID::kWmmaTensorOp},
{"wmma", "<wmma>", OpcodeClassID::kWmmaTensorOp}
{"wmma", "<wmma>", OpcodeClassID::kWmmaTensorOp},
};
/// Converts a OpcodeClassID enumerant to a string
@@ -509,6 +561,47 @@ ComplexTransform from_string<ComplexTransform>(std::string const &str) {
}
static struct {
char const *text;
char const *pretty;
SplitKMode enumerant;
}
SplitKMode_enumerants[] = {
{"serial", "<serial>", SplitKMode::kSerial},
{"parallel", "<parallel>", SplitKMode::kParallel},
};
/// Converts a SplitKMode enumerant to a string
char const *to_string(SplitKMode type, bool pretty) {
for (auto const & possible : SplitKMode_enumerants) {
if (type == possible.enumerant) {
if (pretty) {
return possible.pretty;
}
else {
return possible.text;
}
}
}
return pretty ? "Invalid" : "invalid";
}
/// Converts a SplitKMode enumerant from a string
template <>
SplitKMode from_string<SplitKMode>(std::string const &str) {
for (auto const & possible : SplitKMode_enumerants) {
if ((str.compare(possible.text) == 0) ||
(str.compare(possible.pretty) == 0)) {
return possible.enumerant;
}
}
return SplitKMode::kInvalid;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Lexical cast a string to a byte array. Returns true if cast is successful or false if invalid.
bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string const &str) {
@@ -570,6 +663,20 @@ bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string c
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(tmp);
}
break;
case NumericTypeID::kBF16:
{
float tmp;
ss >> tmp;
*reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(tmp);
}
break;
case NumericTypeID::kTF32:
{
float tmp;
ss >> tmp;
*reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(tmp);
}
break;
case NumericTypeID::kF32:
{
ss >> *reinterpret_cast<float *>(bytes.data());
@@ -589,11 +696,29 @@ bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string c
x->imag() = static_cast<half_t>(std::imag(tmp));
}
break;
case NumericTypeID::kCBF16:
{
std::complex<float> tmp;
ss >> tmp;
cutlass::complex<cutlass::bfloat16_t> *x = reinterpret_cast<cutlass::complex<bfloat16_t> *>(bytes.data());
x->real() = static_cast<bfloat16_t>(std::real(tmp));
x->imag() = static_cast<bfloat16_t>(std::imag(tmp));
}
break;
case NumericTypeID::kCF32:
{
ss >> *reinterpret_cast<std::complex<float>*>(bytes.data());
}
break;
case NumericTypeID::kCTF32:
{
std::complex<float> tmp;
ss >> tmp;
cutlass::complex<cutlass::tfloat32_t> *x = reinterpret_cast<cutlass::complex<tfloat32_t> *>(bytes.data());
x->real() = static_cast<tfloat32_t>(std::real(tmp));
x->imag() = static_cast<tfloat32_t>(std::imag(tmp));
}
break;
case NumericTypeID::kCF64:
{
ss >> *reinterpret_cast<std::complex<double>*>(bytes.data());
@@ -674,6 +799,18 @@ std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type) {
ss << tmp;
}
break;
case NumericTypeID::kBF16:
{
float tmp = *reinterpret_cast<bfloat16_t *>(bytes.data());;
ss << tmp;
}
break;
case NumericTypeID::kTF32:
{
float tmp = *reinterpret_cast<tfloat32_t *>(bytes.data());;
ss << tmp;
}
break;
case NumericTypeID::kF32:
{
ss << *reinterpret_cast<float *>(bytes.data());
@@ -696,6 +833,18 @@ std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type) {
}
}
break;
case NumericTypeID::kCBF16:
{
cutlass::complex<bfloat16_t> const *x =
reinterpret_cast<cutlass::complex<bfloat16_t> const *>(bytes.data());
ss << float(x->real());
if (x->imag() != cutlass::bfloat16_t()) {
ss << "+i" << float(x->imag());
}
}
break;
case NumericTypeID::kCF32:
{
cutlass::complex<float> const * x = reinterpret_cast<cutlass::complex<float> const *>(bytes.data());
@@ -707,6 +856,17 @@ std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type) {
}
}
break;
case NumericTypeID::kCTF32:
{
cutlass::complex<tfloat32_t> const * x = reinterpret_cast<cutlass::complex<tfloat32_t> const *>(bytes.data());
ss << float(x->real());
if (x->imag() != tfloat32_t()) {
ss << "+i" << float(x->imag());
}
}
break;
case NumericTypeID::kCF64:
{
cutlass::complex<double> const * x = reinterpret_cast<cutlass::complex<double> const *>(bytes.data());
@@ -780,6 +940,16 @@ bool cast_from_int64(std::vector<uint8_t> &bytes, NumericTypeID type, int64_t sr
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
}
break;
case NumericTypeID::kBF16:
{
*reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(float(src));
}
break;
case NumericTypeID::kTF32:
{
*reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(float(src));
}
break;
case NumericTypeID::kF32:
{
*reinterpret_cast<float *>(bytes.data()) = static_cast<float>(src);
@@ -870,6 +1040,16 @@ bool cast_from_uint64(std::vector<uint8_t> &bytes, NumericTypeID type, uint64_t
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
}
break;
case NumericTypeID::kBF16:
{
*reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(float(src));
}
break;
case NumericTypeID::kTF32:
{
*reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(float(src));
}
break;
case NumericTypeID::kF32:
{
*reinterpret_cast<float *>(bytes.data()) = static_cast<float>(src);
@@ -961,6 +1141,16 @@ bool cast_from_double(std::vector<uint8_t> &bytes, NumericTypeID type, double sr
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
}
break;
case NumericTypeID::kBF16:
{
*reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(float(src));
}
break;
case NumericTypeID::kTF32:
{
*reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(float(src));
}
break;
case NumericTypeID::kF32:
{
*reinterpret_cast<float *>(bytes.data()) = static_cast<float>(src);
@@ -978,11 +1168,23 @@ bool cast_from_double(std::vector<uint8_t> &bytes, NumericTypeID type, double sr
x->imag() = static_cast<half_t>(float(0));
}
break;
case NumericTypeID::kCBF16:
{
cutlass::complex<cutlass::bfloat16_t> *x = reinterpret_cast<cutlass::complex<bfloat16_t> *>(bytes.data());
x->real() = static_cast<bfloat16_t>(bfloat16_t(src));
x->imag() = static_cast<bfloat16_t>(bfloat16_t(0));
}
break;
case NumericTypeID::kCF32:
{
*reinterpret_cast<std::complex<float>*>(bytes.data()) = std::complex<float>(float(src), float(0));
}
break;
case NumericTypeID::kCTF32:
{
*reinterpret_cast<std::complex<tfloat32_t>*>(bytes.data()) = std::complex<tfloat32_t>(tfloat32_t(src), tfloat32_t(0));
}
break;
case NumericTypeID::kCF64:
{
*reinterpret_cast<std::complex<double>*>(bytes.data()) = std::complex<double>(src, double(0));