CUTLASS#
Subpackages#
Epilogue#
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 
 
Library Defaults#
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) – 
 
 - 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 
 
 
Swizzle#
Registry of swizzling functions