CUTLASS 3.3.0 (#1167)
* Release 3.3.0 Adds support for mixed precision GEMMs On Hopper and Ampere Adds support for < 16B aligned GEMMs on Hopper Enhancements to EVT Enhancements to Python interface Enhancements to Sub-byte type handling in CuTe Several other bug-fixes and performance improvements. * minor doc update
This commit is contained in:
Pradeep Ramani
@@ -34,14 +34,20 @@
|
||||
Utilities for emitting GEMM kernels
|
||||
"""
|
||||
|
||||
import collections
|
||||
import enum
|
||||
import os.path
|
||||
import shutil
|
||||
import functools
|
||||
import operator
|
||||
import collections
|
||||
import os.path
|
||||
import shutil
|
||||
|
||||
from cutlass_library.library import *
|
||||
try:
|
||||
import builtins
|
||||
if hasattr(builtins, "CUTLASS_IGNORE_PACKAGE") and CUTLASS_IGNORE_PACKAGE == True:
|
||||
raise ImportError("Disabling attempt to import cutlass_library")
|
||||
from cutlass_library.library import *
|
||||
except ImportError:
|
||||
from library import *
|
||||
|
||||
###################################################################################################
|
||||
#
|
||||
@@ -55,9 +61,14 @@ class GemmOperation:
|
||||
def __init__(self, gemm_kind, arch, tile_description, A, B, C, element_epilogue, \
|
||||
epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity8, D = None,
|
||||
kernel_schedule = KernelScheduleType.ScheduleAuto, epilogue_schedule = EpilogueScheduleType.ScheduleAuto,
|
||||
tile_scheduler = TileSchedulerType.Default):
|
||||
tile_scheduler = TileSchedulerType.Default, extra_args = None):
|
||||
|
||||
self.prefix = "3x" if gemm_kind == GemmKind.Universal3x else ""
|
||||
kinds_3x = {
|
||||
GemmKind.Universal3x,
|
||||
GemmKind.SparseUniversal3x,
|
||||
}
|
||||
self.is_3x = gemm_kind in kinds_3x
|
||||
self.prefix = "3x" if self.is_3x else ""
|
||||
self.operation_kind = OperationKind.Gemm
|
||||
self.arch = arch
|
||||
self.tile_description = tile_description
|
||||
@@ -66,10 +77,11 @@ class GemmOperation:
|
||||
self.B = B
|
||||
self.C = C
|
||||
self.D = D
|
||||
|
||||
if self.D == None:
|
||||
self.D = self.C
|
||||
|
||||
if gemm_kind != GemmKind.Universal3x:
|
||||
if not self.is_3x:
|
||||
assert(kernel_schedule == KernelScheduleType.ScheduleAuto)
|
||||
assert(epilogue_schedule == EpilogueScheduleType.ScheduleAuto)
|
||||
self.kernel_schedule = kernel_schedule
|
||||
@@ -91,7 +103,7 @@ class GemmOperation:
|
||||
#
|
||||
def is_mixed_input(self):
|
||||
return self.A.element != self.B.element
|
||||
|
||||
|
||||
#
|
||||
def is_planar_complex(self):
|
||||
return self.gemm_kind in (GemmKind.PlanarComplex, GemmKind.PlanarComplexArray)
|
||||
@@ -125,13 +137,20 @@ class GemmOperation:
|
||||
MathOperation.and_popc: 'and'
|
||||
}
|
||||
|
||||
if self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp or \
|
||||
self.tile_description.math_instruction.opcode_class == OpcodeClass.WmmaTensorOp:
|
||||
tensor_ops = [
|
||||
OpcodeClass.TensorOp,
|
||||
OpcodeClass.WmmaTensorOp,
|
||||
OpcodeClass.SparseTensorOp,
|
||||
]
|
||||
|
||||
is_tensor_op = self.tile_description.math_instruction.opcode_class in tensor_ops
|
||||
|
||||
if is_tensor_op:
|
||||
|
||||
math_op = self.tile_description.math_instruction.math_operation
|
||||
math_op_string = math_operations_map[math_op] if math_op in math_operations_map.keys() else ''
|
||||
|
||||
if self.gemm_kind == GemmKind.Universal3x:
|
||||
if self.is_3x:
|
||||
inst_shape = "{0}x{1}x{2}".format(*tuple(self.tile_description.math_instruction.instruction_shape))
|
||||
else:
|
||||
inst_shape = "{0}{1}{2}".format(*tuple(self.tile_description.math_instruction.instruction_shape))
|
||||
@@ -183,6 +202,16 @@ class GemmOperation:
|
||||
core_name = self.core_name())
|
||||
return extended_name
|
||||
|
||||
def datatype_name_3x(self):
|
||||
'''Generates a string representing the MMA atom. Assumes accumulator type is C type.'''
|
||||
datatype_name = "{element_a}_{element_b}_{element_acc}_{element_c}_{element_d}".format(
|
||||
element_a = DataTypeNames[self.A.element],
|
||||
element_b = DataTypeNames[self.B.element],
|
||||
element_acc = DataTypeNames[self.tile_description.math_instruction.element_accumulator],
|
||||
element_c = DataTypeNames[self.C.element],
|
||||
element_d = DataTypeNames[self.D.element])
|
||||
return datatype_name
|
||||
|
||||
# Generates a short string representing the AB layout tags (e.g. nt or tn)
|
||||
def layout_name(self):
|
||||
if self.is_complex() or self.is_planar_complex():
|
||||
@@ -213,6 +242,10 @@ class GemmOperation:
|
||||
def epilogue_schedule_name_3x(self):
|
||||
return EpilogueScheduleSuffixes[self.epilogue_schedule]
|
||||
|
||||
# Generate a short string representing the operation class
|
||||
def opcode_class_name(self):
|
||||
return OpcodeClassNames[self.tile_description.math_instruction.opcode_class]
|
||||
|
||||
# Generates the full kernel function name
|
||||
def procedural_name(self):
|
||||
''' The full procedural name indicates architecture, extended name, tile size, and layout. '''
|
||||
@@ -661,7 +694,6 @@ ${compile_guard_end}
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class EmitGemmUniversal3xInstance:
|
||||
''' Responsible for emitting a CUTLASS 3.x template definition'''
|
||||
|
||||
@@ -687,10 +719,10 @@ class EmitGemmUniversal3xInstance:
|
||||
|
||||
using ${operation_name}_epilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
${arch}, ${opcode_class},
|
||||
${arch}, ${opcode_class_epi},
|
||||
cute::Shape<cute::_${tile_shape_m}, cute::_${tile_shape_n}, cute::_${tile_shape_k}>,
|
||||
cute::Shape<cute::_${cluster_m},cute::_${cluster_n},cute::_${cluster_k}>,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
${epi_tile_mn},
|
||||
${element_accumulator}, ${element_epilogue},
|
||||
${element_c}, ${layout_c}, ${align_c},
|
||||
${element_d}, ${layout_d}, ${align_d},
|
||||
@@ -699,7 +731,7 @@ using ${operation_name}_epilogue =
|
||||
|
||||
using ${operation_name}_mainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
${arch}, ${opcode_class},
|
||||
${arch}, ${opcode_class_main},
|
||||
${element_a}, ${layout_a}, ${align_a},
|
||||
${element_b}, ${layout_b}, ${align_b},
|
||||
${element_accumulator},
|
||||
@@ -743,6 +775,10 @@ ${compile_guard_end}
|
||||
stage_count_string = f"cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename {str(operation.procedural_name())}_epilogue::SharedStorage)>"
|
||||
warp_shape = [tile_shape[idx] // warp_count[idx] for idx in range(3)]
|
||||
|
||||
epi_tile_mn = "cutlass::epilogue::collective::EpilogueTileAuto"
|
||||
opcode_class_main = operation.tile_description.math_instruction.opcode_class
|
||||
opcode_class_epi = opcode_class_main
|
||||
|
||||
instance_layout_A, instance_layout_B, instance_layout_C , instance_layout_D = \
|
||||
(operation.A.layout, operation.B.layout, operation.C.layout, operation.D.layout)
|
||||
|
||||
@@ -760,20 +796,23 @@ ${compile_guard_end}
|
||||
else:
|
||||
epilogue_functor = self.epilogue_functor.emit_declaration()
|
||||
#
|
||||
|
||||
element_a = DataTypeTag[operation.A.element]
|
||||
element_b = DataTypeTag[operation.B.element]
|
||||
epilogue_schedule_type = EpilogueScheduleTag[operation.epilogue_schedule]
|
||||
values = {
|
||||
'operation_name': operation.procedural_name(),
|
||||
'operation_suffix': self.operation_suffix,
|
||||
'element_a': DataTypeTag[operation.A.element],
|
||||
'element_a': element_a,
|
||||
'layout_a': LayoutTag[instance_layout_A],
|
||||
'element_b': DataTypeTag[operation.B.element],
|
||||
'element_b': element_b,
|
||||
'layout_b': LayoutTag[instance_layout_B],
|
||||
'element_c': DataTypeTag[operation.C.element],
|
||||
'layout_c': LayoutTag[instance_layout_C],
|
||||
'element_d': DataTypeTag[operation.D.element],
|
||||
'layout_d': LayoutTag[instance_layout_D],
|
||||
'element_accumulator': DataTypeTag[operation.accumulator_type()],
|
||||
'opcode_class': OpcodeClassTag[operation.tile_description.math_instruction.opcode_class],
|
||||
'opcode_class_main': OpcodeClassTag[opcode_class_main],
|
||||
'opcode_class_epi': OpcodeClassTag[opcode_class_epi],
|
||||
'arch': "cutlass::arch::Sm%d" % operation.arch,
|
||||
'tile_shape_m': str(operation.tile_description.tile_shape[0]),
|
||||
'tile_shape_n': str(operation.tile_description.tile_shape[1]),
|
||||
@@ -788,7 +827,8 @@ ${compile_guard_end}
|
||||
'instruction_shape_n': str(operation.tile_description.math_instruction.instruction_shape[1]),
|
||||
'instruction_shape_k': str(operation.tile_description.math_instruction.instruction_shape[2]),
|
||||
'kernel_schedule' : str(KernelScheduleTag[operation.kernel_schedule]),
|
||||
'epilogue_schedule' : str(EpilogueScheduleTag[operation.epilogue_schedule]),
|
||||
'epilogue_schedule' : str(epilogue_schedule_type),
|
||||
'epi_tile_mn' : epi_tile_mn,
|
||||
'epilogue_functor': epilogue_functor,
|
||||
'stages': stage_count_string,
|
||||
'align_a': str(operation.A.alignment),
|
||||
@@ -800,7 +840,7 @@ ${compile_guard_end}
|
||||
'math_operation': MathOperationTag[operation.tile_description.math_instruction.math_operation],
|
||||
'epilogue_vector_length': str(epilogue_vector_length),
|
||||
'element_epilogue': str(DataTypeTag[operation.element_epilogue]),
|
||||
'tile_scheduler': str(TileSchedulerTag[operation.tile_scheduler])
|
||||
'tile_scheduler': str(TileSchedulerTag[operation.tile_scheduler]),
|
||||
}
|
||||
|
||||
return SubstituteTemplate(self.gemm_template, values)
|
||||
|
||||
Reference in New Issue
Block a user