python/docs/cutlass.html
Registry of elementwise epilogues
Elementwise epilogues can be added to many CUTLASS kernels in the CUTLAS Python interface via code like the following for GEMM:
plan=cutlass.op.Gemm(element=cutlass.DataType.f32,layout=cutlass.LayoutType.RowMajor)plan.activation=cutlass.epilogue.relu
cutlass.epilogue.get_activation_epilogue(activation, element_output, elements_per_access, element_accumulator, element_compute)[source]#
Return an epilogue corresponding to the activation function, data types, and alignment used in the kernel
Parameters:
activation – elementwise activation function to use
element_output – data type of the output
elements_per_access (int) – alignment of operand C of the kernel
element_accumulator – data type of the accumulated output C
element_compute – data type in which compute operations should be performed
Returns:
epilogue functor
cutlass.epilogue.get_activations()[source]#
Returns a list of available activation functions
Returns:
list of available activation functions
Return type:
list
Classes containing valid operations for a given compute capability and data types.
class cutlass.library_defaults.ArchOptions(target_cc, kernel_cc, operation_kind, gemm_kinds, allowed_math_operations=[<MathOperation.multiply_add: 1>, <MathOperation.multiply_add_saturate: 2>])[source]#
Bases: object
Structure for keeping track of kernels available on a given compute capability
Parameters:
target_cc (int) – compute capability of the device on which kernels will be run
kernel_cc (int) – compute capability of the kernels to generate
operation_kind (cutlass.OperationKind) – type of operation to register
gemm_kinds (list) – types of GEMM operations that can be included
allowed_math_operations (list) – types of primitive math operations allowed
opclass_supports_combination(op_class, datatype_comb, layout_comb)[source]#
Returns whether the provided operation class supports the provided data type and layout combination
Parameters:
op_class (cutlass.OpcodeClass) – operation class to consider
datatype_comb (tuple_[cutlass.DataType]_) – tuple of data types for (element_A, element_B, element_accumulator)
layout_comb (tuple_[cutlass.LayoutType]_) – tuple of data types for (layout_A, layout_B)
Returns:
set of operation classes that support the provided data type and layout combination
Return type:
set
operations(op_class, element_a, element_b, element_accumulator, layout_a, layout_b)[source]#
Returns whether the provided operation class supports the provided data type combination
Parameters:
op_class (cutlass.OpcodeClass) – operation class to consider
element_a (cutlass.DataType) – data type of operand A
element_b (cutlass.DataType) – data type of operand B
element_accumulator (cutlass.DataType) – data type of accumulator
layout_a (cutlass.LayoutType) – layout of operand A
layout_b (cutlass.LayoutType) – layout of operand B
Returns:
container of kernels by alignment supported by the provided combination of parameters
Return type:
supporting_opclasses(element_a, element_b, element_accumulator, layout_a, layout_b)[source]#
Returns a set of operation classes that support the provided data type combination
Parameters:
element_a (cutlass.DataType) – data type of operand A
element_b (cutlass.DataType) – data type of operand B
element_accumulator (cutlass.DataType) – data type of accumulator
layout_a (cutlass.LayoutType) – layout of operand A
layout_b (cutlass.LayoutType) – layout of operand B
Returns:
set of operation classes that support the provided data type combination
Return type:
set
class cutlass.library_defaults.KernelsForDataType(datatype_comb, layout_comb)[source]#
Bases: object
Container class for keeping track of kernels that correspond to a particular combination of data types for operands A, B, and accumulator
Parameters:
datatype_comb (tuple) –
layout_comb (tuple) –
Add an operation to the list of supported kernels
property alignments#
Returns an unsorted list of alignments supported by this data type combination
Returns:
unsorted list of alignments supported by this data type combination
Return type:
list
property all_operations#
Returns a list of all operations supported by this data type combination
Returns:
list of all operations supported by this data type combination
Return type:
list
find_alignment(shape, layout)[source]#
Returns the most preferable alignment for a given shape and layout
Parameters:
shape (tuple) – extent of each dimension of the tensor
layout (cutlass.LayoutType) – layout of the tensor
Returns:
maximum alignment supported by the data type combination and tensor size
Return type:
int
operations(alignment)[source]#
Returns operations satisfying the alignment constraint indicated by alignment
Parameters:
alignment (int) – alignment constraint of operations to return
Returns:
list of operations
Return type:
list
Sorts each list of kernels in kernels_by_alignment in descending order of threadblock shape
class cutlass.library_defaults.OptionRegistry(target_cc)[source]#
Bases: object
Container of all architecture-specific options
Parameters:
target_cc (int) – compute capability of the device on which operations will be run
options_for_cc(cc)[source]#Parameters:
cc (int) –
Return type:
Registry of swizzling functions