# # \file generator.py # # \brief Generates the CUTLASS Library's instances # import enum import os.path import shutil import functools import operator from library import * ################################################################################################### # # Data structure modeling a GEMM operation # ################################################################################################### # class GemmOperation: # def __init__(self, gemm_kind, arch, tile_description, A, B, C, element_epilogue, \ epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity8, \ required_cuda_ver_major = 9, required_cuda_ver_minor = 2): self.operation_kind = OperationKind.Gemm self.arch = arch self.tile_description = tile_description self.gemm_kind = gemm_kind self.A = A self.B = B self.C = C self.element_epilogue = element_epilogue self.epilogue_functor = epilogue_functor self.swizzling_functor = swizzling_functor self.required_cuda_ver_major = required_cuda_ver_major self.required_cuda_ver_minor = required_cuda_ver_minor # 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 # def is_split_k_parallel(self): return self.gemm_kind == GemmKind.SplitKParallel # def is_planar_complex(self): return self.gemm_kind in (GemmKind.PlanarComplex, GemmKind.PlanarComplexArray) # def accumulator_type(self): accum = self.tile_description.math_instruction.element_accumulator if self.is_complex(): return get_complex_from_real(accum) return accum # 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()] # def core_name(self): ''' The basic operation kind is prefixed with a letter indicating the accumulation type. ''' inst_shape = '' inst_operation = '' intermediate_type = '' math_operations_map = { MathOperation.xor_popc: 'xor', } if self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp or \ self.tile_description.math_instruction.opcode_class == OpcodeClass.WmmaTensorOp: 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 '' inst_shape = "%d%d%d" % tuple(self.tile_description.math_instruction.instruction_shape) inst_shape += math_op_string if self.tile_description.math_instruction.element_a != self.A.element and \ self.tile_description.math_instruction.element_a != self.tile_description.math_instruction.element_accumulator: intermediate_type = DataTypeNames[self.tile_description.math_instruction.element_a] return "%s%s%s%s" % (self.short_math_name(), inst_shape, intermediate_type, GemmKindNames[self.gemm_kind]) # def extended_name(self): ''' Append data types if they differ from compute type. ''' if self.is_complex(): extended_name = "${core_name}" else: if self.C.element != self.tile_description.math_instruction.element_accumulator and \ self.A.element != self.tile_description.math_instruction.element_accumulator: extended_name = "${element_c}_${core_name}_${element_a}" elif self.C.element == self.tile_description.math_instruction.element_accumulator and \ self.A.element != self.tile_description.math_instruction.element_accumulator: extended_name = "${core_name}_${element_a}" else: extended_name = "${core_name}" extended_name = SubstituteTemplate(extended_name, { 'element_a': DataTypeNames[self.A.element], 'element_c': DataTypeNames[self.C.element], 'core_name': self.core_name() }) return extended_name # def layout_name(self): if self.is_complex() or self.is_planar_complex(): return "%s%s" % ( ShortComplexLayoutNames[(self.A.layout, self.A.complex_transform)], ShortComplexLayoutNames[(self.B.layout, self.B.complex_transform)] ) return "%s%s" % (ShortLayoutTypeNames[self.A.layout], ShortLayoutTypeNames[self.B.layout]) # def procedural_name(self): ''' The full procedural name indicates architecture, extended name, tile size, and layout. ''' threadblock = self.tile_description.procedural_name() opcode_class_name = OpcodeClassNames[self.tile_description.math_instruction.opcode_class] alignment = max([self.A.alignment, self.B.alignment, self.C.alignment]) return SubstituteTemplate( "cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}_align${alignment}", { 'opcode_class': opcode_class_name, 'extended_name': self.extended_name(), 'threadblock': threadblock, 'layout': self.layout_name(), 'alignment': "%d" % self.A.alignment, } ) # def configuration_name(self): ''' The full procedural name indicates architecture, extended name, tile size, and layout. ''' return self.procedural_name() ################################################################################################### # # Data structure modeling a GEMV Batched Strided operation # ################################################################################################### # class GemvBatchedStridedOperation: # def __init__(self, gemm_kind, arch, math_inst, threadblock_shape, thread_shape, A, B, C, \ required_cuda_ver_major = 9, required_cuda_ver_minor = 2): self.operation_kind = OperationKind.Gemm self.arch = arch self.gemm_kind = gemm_kind self.math_instruction = math_inst self.threadblock_shape = threadblock_shape self.thread_shape = thread_shape self.A = A self.B = B self.C = C self.required_cuda_ver_major = required_cuda_ver_major self.required_cuda_ver_minor = required_cuda_ver_minor # def accumulator_type(self): accum = self.math_instruction.element_accumulator return accum # def short_math_name(self): return ShortDataTypeNames[self.accumulator_type()] # def core_name(self): ''' The basic operation kind is prefixed with a letter indicating the accumulation type. ''' return "%s%s" % (self.short_math_name(), \ GemmKindNames[self.gemm_kind]) # def extended_name(self): ''' Append data types if they differ from compute type. ''' if self.C.element != self.math_instruction.element_accumulator and \ self.A.element != self.math_instruction.element_accumulator: extended_name = "${element_c}_${core_name}_${element_a}" elif self.C.element == self.math_instruction.element_accumulator and \ self.A.element != self.math_instruction.element_accumulator: extended_name = "${core_name}_${element_a}" else: extended_name = "${core_name}" extended_name = SubstituteTemplate(extended_name, { 'element_a': DataTypeNames[self.A.element], 'element_c': DataTypeNames[self.C.element], 'core_name': self.core_name() }) return extended_name # def layout_name(self): return "%s%s" % (ShortLayoutTypeNames[self.A.layout], ShortLayoutTypeNames[self.B.layout]) # def procedural_name(self): ''' The full procedural name indicates architecture, extended name, tile size, and layout. ''' threadblock = "%dx%d_%d" % (self.threadblock_shape[0], self.threadblock_shape[1], self.threadblock_shape[2]) opcode_class_name = OpcodeClassNames[self.math_instruction.opcode_class] alignment_a = self.A.alignment alignment_b = self.B.alignment return SubstituteTemplate( "cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}_align${alignment_a}x${alignment_b}", { 'opcode_class': opcode_class_name, 'extended_name': self.extended_name(), 'threadblock': threadblock, 'layout': self.layout_name(), 'alignment_a': "%d" % alignment_a, 'alignment_b': "%d" % alignment_b, } ) # def configuration_name(self): ''' The full procedural name indicates architecture, extended name, tile size, and layout. ''' return self.procedural_name() # def GeneratesGemm(tile, data_type, layout_a, layout_b, layout_c, min_cc, align_a = 32, align_b = 32, align_c = 32, required_cuda_ver_major = 9, required_cuda_ver_minor = 2): operations = [] swizzling_functor = SwizzlingFunctor.Identity1 element_a, element_b, element_c, element_epilogue = data_type if tile.math_instruction.element_accumulator == DataType.s32: epilogues = [EpilogueFunctor.LinearCombinationClamp] else: assert tile.math_instruction.element_accumulator == DataType.f32 or \ tile.math_instruction.element_accumulator == DataType.f16 epilogues = [EpilogueFunctor.LinearCombination] for epilogue in epilogues: A = TensorDescription(element_a, layout_a, int(align_a//DataTypeSize[element_a])) B = TensorDescription(element_b, layout_b, int(align_b//DataTypeSize[element_b])) C = TensorDescription(element_c, layout_c, int(align_c//DataTypeSize[element_c])) operations.append(GemmOperation(GemmKind.Gemm, min_cc, tile, A, B, C, \ element_epilogue, epilogue, swizzling_functor, \ required_cuda_ver_major, required_cuda_ver_minor)) operations.append(GemmOperation(GemmKind.SplitKParallel, min_cc, tile, A, B, C, \ element_epilogue, epilogue, swizzling_functor, \ required_cuda_ver_major, required_cuda_ver_minor)) return operations def GeneratesGemv(math_inst, threadblock_shape, thread_shape, data_type, layout_a, layout_b, layout_c, min_cc, \ align_a = 32, align_b = 32, align_c = 32, \ required_cuda_ver_major = 9, required_cuda_ver_minor = 2): element_a, element_b, element_c, element_epilogue = data_type A = TensorDescription(element_a, layout_a, int(align_a//DataTypeSize[element_a])) B = TensorDescription(element_b, layout_b, int(align_b//DataTypeSize[element_b])) C = TensorDescription(element_c, layout_c, int(align_c//DataTypeSize[element_c])) return GemvBatchedStridedOperation(GemmKind.GemvBatchedStrided, min_cc, math_inst, threadblock_shape, thread_shape, \ A, B, C, required_cuda_ver_major, required_cuda_ver_minor) ################################################################################################### # # Emits single instances of a CUTLASS device-wide operator # ################################################################################################### # class EmitGemmInstance: ''' Responsible for emitting a CUTLASS template definition''' def __init__(self): self.gemm_template = """ // Gemm operator ${operation_name} using Operation_${operation_name} = cutlass::gemm::device::Gemm< ${element_a}, ${layout_a}, ${element_b}, ${layout_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}, ${align_a}, ${align_b}, false, ${math_operation} ${residual} >; """ self.gemm_complex_template = """ // Gemm operator ${operation_name} using Operation_${operation_name} = cutlass::gemm::device::GemmComplex< ${element_a}, ${layout_a}, ${element_b}, ${layout_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}, ${transform_a}, ${transform_b}, ${math_operation} ${residual} >; """ def emit(self, operation): warp_shape = [operation.tile_description.threadblock_shape[idx] // operation.tile_description.warp_count[idx] for idx in range(3)] epilogue_vector_length = int(min(operation.C.alignment * DataTypeSize[operation.C.element], 128) / DataTypeSize[operation.C.element]) residual = '' values = { 'operation_name': operation.procedural_name(), 'element_a': DataTypeTag[operation.A.element], 'layout_a': LayoutTag[operation.A.layout], 'element_b': DataTypeTag[operation.B.element], 'layout_b': LayoutTag[operation.B.layout], 'element_c': DataTypeTag[operation.C.element], 'layout_c': LayoutTag[operation.C.layout], '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], 'residual': residual } template = self.gemm_complex_template if operation.is_complex() else self.gemm_template return SubstituteTemplate(template, values) # class EmitGemvBatchedStridedInstance: ''' Responsible for emitting a CUTLASS template definition''' def __init__(self): self.template = """ // Gemm operator ${operation_name} using Operation_${operation_name} = cutlass::gemm::kernel::DefaultGemv< cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>, cutlass::gemm::GemmShape<${thread_shape_m}, ${thread_shape_n}, ${thread_shape_k}>, ${element_a}, ${layout_a}, ${element_b}, ${layout_b}, ${element_c}, ${layout_c} >; """ def emit(self, operation): values = { 'operation_name': operation.procedural_name(), 'element_a': DataTypeTag[operation.A.element], 'layout_a': LayoutTag[operation.A.layout], 'element_b': DataTypeTag[operation.B.element], 'layout_b': LayoutTag[operation.B.layout], 'element_c': DataTypeTag[operation.C.element], 'layout_c': LayoutTag[operation.C.layout], 'threadblock_shape_m': str(operation.threadblock_shape[0]), 'threadblock_shape_n': str(operation.threadblock_shape[1]), 'threadblock_shape_k': str(operation.threadblock_shape[2]), 'thread_shape_m': str(operation.thread_shape[0]), 'thread_shape_n': str(operation.thread_shape[1]), 'thread_shape_k': str(operation.thread_shape[2]), } return SubstituteTemplate(self.template, values) ################################################################################################### class EmitSparseGemmInstance: ''' Responsible for emitting a CUTLASS template definition''' def __init__(self): self.gemm_template = """ // Gemm operator ${operation_name} using Operation_${operation_name} = cutlass::gemm::device::SparseGemm< ${element_a}, ${layout_a}, ${element_b}, ${layout_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}, ${align_a}, ${align_b}, false, ${math_operation} ${residual} >; """ def emit(self, operation): warp_shape = [operation.tile_description.threadblock_shape[idx] // operation.tile_description.warp_count[idx] for idx in range(3)] epilogue_vector_length = int(min(operation.C.alignment * DataTypeSize[operation.C.element], 128) / DataTypeSize[operation.C.element]) residual = '' values = { 'operation_name': operation.procedural_name(), 'element_a': DataTypeTag[operation.A.element], 'layout_a': LayoutTag[operation.A.layout], 'element_b': DataTypeTag[operation.B.element], 'layout_b': LayoutTag[operation.B.layout], 'element_c': DataTypeTag[operation.C.element], 'layout_c': LayoutTag[operation.C.layout], '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], 'residual': residual } template = self.gemm_template return SubstituteTemplate(template, values) ################################################################################################### # 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)] 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''' def __init__(self): self.template = """ // Gemm operator ${operation_name} using Operation_${operation_name} = typename cutlass::gemm::kernel::DefaultGemmPlanarComplexUniversal< ${element_a}, ${layout_a}, ${transform_a}, ${alignment_a}, ${element_b}, ${layout_b}, ${transform_b}, ${alignment_b}, ${element_c}, cutlass::layout::RowMajor, ${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}>, cutlass::epilogue::thread::LinearCombinationPlanarComplex< ${element_c}, ${alignment_c}, ${element_accumulator}, ${element_epilogue} >, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, ${stages}, ${math_operator} >::GemmKernel; struct ${operation_name} : public Operation_${operation_name} { }; """ def emit(self, operation): warp_shape = [operation.tile_description.threadblock_shape[idx] // operation.tile_description.warp_count[idx] for idx in range(3)] # exchange and transpose A and B types, layouts, and complex transforms since the C layout is row-major transposed_layout_A = TransposedLayout[operation.A.layout] transposed_layout_B = TransposedLayout[operation.B.layout] values = { 'operation_name': operation.procedural_name(), 'element_a': DataTypeTag[operation.B.element], 'layout_a': LayoutTag[transposed_layout_B], 'transform_a': ComplexTransformTag[operation.B.complex_transform], 'alignment_a': str(operation.B.alignment), 'element_b': DataTypeTag[operation.A.element], 'layout_b': LayoutTag[transposed_layout_A], 'transform_b': ComplexTransformTag[operation.A.complex_transform], 'alignment_b': str(operation.A.alignment), 'element_c': DataTypeTag[operation.C.element], 'layout_c': LayoutTag[operation.C.layout], 'element_accumulator': DataTypeTag[operation.tile_description.math_instruction.element_accumulator], '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]), 'alignment_c': str(operation.C.alignment), 'element_epilogue': str(DataTypeTag[operation.element_epilogue]), 'stages': str(operation.tile_description.stages), 'math_operator': 'cutlass::arch::OpMultiplyAdd' } return SubstituteTemplate(self.template, values) ################################################################################################### # class EmitGemmPlanarComplexArrayInstance: ''' Responsible for emitting a CUTLASS template definition''' def __init__(self): self.template = """ // Gemm operator ${operation_name} using Operation_${operation_name} = typename cutlass::gemm::kernel::DefaultGemmPlanarComplexUniversal< ${element_a}, ${layout_a}, ${transform_a}, ${alignment_a}, ${element_b}, ${layout_b}, ${transform_b}, ${alignment_b}, ${element_c}, cutlass::layout::RowMajor, ${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}>, cutlass::epilogue::thread::LinearCombinationPlanarComplex< ${element_c}, ${alignment_c}, ${element_accumulator}, ${element_epilogue} >, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, ${stages}, ${math_operator} >::GemmArrayKernel; struct ${operation_name} : public Operation_${operation_name} { }; """ def emit(self, operation): warp_shape = [operation.tile_description.threadblock_shape[idx] // operation.tile_description.warp_count[idx] for idx in range(3)] # exchange and transpose A and B types, layouts, and complex transforms since the C layout is row-major transposed_layout_A = TransposedLayout[operation.A.layout] transposed_layout_B = TransposedLayout[operation.B.layout] values = { 'operation_name': operation.procedural_name(), 'element_a': DataTypeTag[operation.B.element], 'layout_a': LayoutTag[transposed_layout_B], 'transform_a': ComplexTransformTag[operation.B.complex_transform], 'alignment_a': str(operation.B.alignment), 'element_b': DataTypeTag[operation.A.element], 'layout_b': LayoutTag[transposed_layout_A], 'transform_b': ComplexTransformTag[operation.A.complex_transform], 'alignment_b': str(operation.A.alignment), 'element_c': DataTypeTag[operation.C.element], 'layout_c': LayoutTag[operation.C.layout], 'element_accumulator': DataTypeTag[operation.tile_description.math_instruction.element_accumulator], '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]), 'alignment_c': str(operation.C.alignment), 'element_epilogue': str(DataTypeTag[operation.element_epilogue]), 'stages': str(operation.tile_description.stages), 'math_operator': 'cutlass::arch::OpMultiplyAdd' } return SubstituteTemplate(self.template, values) # class EmitGemmSplitKParallelInstance: ''' Responsible for emitting a CUTLASS template definition''' def __init__(self): self.template = """ // Gemm operator ${operation_name} using Operation_${operation_name} = cutlass::gemm::device::GemmSplitKParallel< ${element_a}, ${layout_a}, ${element_b}, ${layout_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} >, cutlass::epilogue::thread::Convert< ${element_accumulator}, ${epilogue_vector_length}, ${element_accumulator} >, cutlass::reduction::thread::ReduceAdd< ${element_accumulator}, ${element_accumulator}, ${epilogue_vector_length} >, cutlass::gemm::threadblock::GemmSplitKHorizontalThreadblockSwizzle, ${stages}, ${align_a}, ${align_b}, ${math_operation} >; """ def emit(self, operation): warp_shape = [operation.tile_description.threadblock_shape[idx] // operation.tile_description.warp_count[idx] for idx in range(3)] epilogue_vector_length = int(min(operation.C.alignment * DataTypeSize[operation.C.element], 128) / DataTypeSize[operation.C.element]) values = { 'operation_name': operation.procedural_name(), 'element_a': DataTypeTag[operation.A.element], 'layout_a': LayoutTag[operation.A.layout], 'element_b': DataTypeTag[operation.B.element], 'layout_b': LayoutTag[operation.B.layout], 'element_c': DataTypeTag[operation.C.element], 'layout_c': LayoutTag[operation.C.layout], '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], 'stages': str(operation.tile_description.stages), 'math_operation': MathOperationTag[operation.tile_description.math_instruction.math_operation], 'align_a': str(operation.A.alignment), 'align_b': str(operation.B.alignment), } return SubstituteTemplate(self.template, values) ################################################################################################### ################################################################################################### # # Emitters functions for all targets # ################################################################################################### class EmitGemmConfigurationLibrary: def __init__(self, operation_path, configuration_name): self.configuration_name = configuration_name self.configuration_path = os.path.join(operation_path, "%s.cu" % configuration_name).replace('\\', '/') self.instance_emitter = { GemmKind.Gemm: EmitGemmInstance, GemmKind.Sparse: EmitSparseGemmInstance, GemmKind.Universal: EmitGemmUniversalInstance, GemmKind.PlanarComplex: EmitGemmPlanarComplexInstance, GemmKind.PlanarComplexArray: EmitGemmPlanarComplexArrayInstance } self.gemm_kind_wrappers = { GemmKind.Gemm: 'GemmOperation', GemmKind.Sparse: 'GemmSparseOperation', GemmKind.Universal: 'GemmUniversalOperation', GemmKind.PlanarComplex: 'GemmPlanarComplexOperation', GemmKind.PlanarComplexArray: 'GemmPlanarComplexArrayOperation' } self.wmma_guard_start = "#if defined(CUTLASS_ARCH_WMMA_SM${sm_number}_ENABLED)" self.instance_template = { GemmKind.Gemm: """ ${compile_guard_start} manifest.append(new ${gemm_kind}("${operation_name}")); ${compile_guard_end} """, GemmKind.Sparse: """ ${compile_guard_start} manifest.append(new ${gemm_kind}("${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} manifest.append(new ${gemm_kind}< cutlass::gemm::device::GemmUniversalAdapter<${operation_name}> >("${operation_name}")); ${compile_guard_end} """, GemmKind.PlanarComplexArray: """ ${compile_guard_start} manifest.append(new ${gemm_kind}< cutlass::gemm::device::GemmUniversalAdapter<${operation_name}> >("${operation_name}")); ${compile_guard_end} """ } self.header_template = """ /* Generated by gemm_operation.py - Do not edit. */ /////////////////////////////////////////////////////////////////////////////////////////////////// #include "cutlass/arch/wmma.h" #include "cutlass/cutlass.h" #include "cutlass/library/library.h" #include "cutlass/library/manifest.h" #include "library_internal.h" #include "gemm_operation.h" /////////////////////////////////////////////////////////////////////////////////////////////////// """ self.initialize_function_template = """ /////////////////////////////////////////////////////////////////////////////////////////////////// namespace cutlass { namespace library { /////////////////////////////////////////////////////////////////////////////////////////////////// void initialize_${configuration_name}(Manifest &manifest) { """ self.epilogue_template = """ } /////////////////////////////////////////////////////////////////////////////////////////////////// } // namespace library } // namespace cutlass /////////////////////////////////////////////////////////////////////////////////////////////////// """ def __enter__(self): self.configuration_file = open(self.configuration_path, "w") self.configuration_file.write(self.header_template) self.instance_definitions = [] self.instance_wrappers = [] self.operations = [] return self def emit(self, operation): emitter = self.instance_emitter[operation.gemm_kind]() self.operations.append(operation) self.instance_definitions.append(emitter.emit(operation)) self.instance_wrappers.append(SubstituteTemplate(self.instance_template[operation.gemm_kind], { 'configuration_name': self.configuration_name, 'operation_name': operation.procedural_name(), 'gemm_kind': self.gemm_kind_wrappers[operation.gemm_kind], 'compile_guard_start': SubstituteTemplate(self.wmma_guard_start, {'sm_number': str(operation.arch)}) \ if operation.tile_description.math_instruction.opcode_class == OpcodeClass.WmmaTensorOp else "", 'compile_guard_end': "#endif" \ if operation.tile_description.math_instruction.opcode_class == OpcodeClass.WmmaTensorOp else "" })) def __exit__(self, exception_type, exception_value, traceback): # Write instance definitions in top-level namespace for instance_definition in self.instance_definitions: self.configuration_file.write(instance_definition) # Add wrapper objects within initialize() function self.configuration_file.write(SubstituteTemplate(self.initialize_function_template, { 'configuration_name': self.configuration_name })) for instance_wrapper in self.instance_wrappers: self.configuration_file.write(instance_wrapper) self.configuration_file.write(self.epilogue_template) self.configuration_file.close() ################################################################################################### ################################################################################################### class EmitGemmSingleKernelWrapper: def __init__(self, kernel_path, gemm_operation, short_path=False): self.short_path = short_path self.kernel_path = kernel_path self.operation = gemm_operation instance_emitters = { GemmKind.Gemm: EmitGemmInstance(), GemmKind.SplitKParallel: EmitGemmSplitKParallelInstance(), } self.instance_emitter = instance_emitters[self.operation.gemm_kind] self.header_template = """ #if __CUDACC_VER_MAJOR__ > ${required_cuda_ver_major} || (__CUDACC_VER_MAJOR__ == ${required_cuda_ver_major} && __CUDACC_VER_MINOR__ >= ${required_cuda_ver_minor}) // ignore warning of cutlass #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wunused-parameter" #pragma GCC diagnostic ignored "-Wstrict-aliasing" #pragma GCC diagnostic ignored "-Wuninitialized" #pragma GCC diagnostic ignored "-Wmaybe-uninitialized" #include "cutlass/gemm/device/gemm.h" #include "cutlass/gemm/device/gemm_splitk_parallel.h" #include "src/cuda/cutlass/manifest.h" #include "src/cuda/cutlass/gemm_operation.h" """ self.instance_template = """ ${operation_instance} """ self.manifest_template = """ namespace cutlass { namespace library { void initialize_${operation_name}(Manifest &manifest) { manifest.append(new GemmOperation< Operation_${operation_name} >("${operation_name}")); } } // namespace library } // namespace cutlass """ self.epilogue_template = """ #pragma GCC diagnostic pop #endif """ # def __enter__(self): if self.short_path: self.kernel_path = os.path.join(self.kernel_path, "%s.cu" % GlobalCnt.cnt) GlobalCnt.cnt += 1 else: self.kernel_path = os.path.join(self.kernel_path, "%s.cu" % self.operation.procedural_name()) self.kernel_file = open(self.kernel_path, "w") self.kernel_file.write(SubstituteTemplate(self.header_template, { 'required_cuda_ver_major': str(self.operation.required_cuda_ver_major), 'required_cuda_ver_minor': str(self.operation.required_cuda_ver_minor), })) return self # def emit(self): self.kernel_file.write(SubstituteTemplate(self.instance_template, { 'operation_instance': self.instance_emitter.emit(self.operation), })) # emit manifest helper manifest = SubstituteTemplate(self.manifest_template, { 'operation_name': self.operation.procedural_name(), }) self.kernel_file.write(manifest) # def __exit__(self, exception_type, exception_value, traceback): self.kernel_file.write(self.epilogue_template) self.kernel_file.close() ################################################################################################### ################################################################################################### class EmitGemvSingleKernelWrapper: def __init__(self, kernel_path, gemm_operation, wrapper_path, short_path=False): self.kernel_path = kernel_path self.wrapper_path = wrapper_path self.operation = gemm_operation self.short_path = short_path self.wrapper_template = """ template void megdnn::cuda::cutlass_wrapper:: cutlass_vector_matrix_mul_batched_strided_wrapper( BatchedGemmCoord const& problem_size, const typename Operation_${operation_name}::ElementA* d_A, size_t lda, size_t batch_stride_a, const typename Operation_${operation_name}::ElementB* d_B, size_t ldb, size_t batch_stride_b, typename Operation_${operation_name}::ElementCD* d_C, size_t ldc, size_t batch_stride_c, cudaStream_t stream); """ self.instance_emitter = EmitGemvBatchedStridedInstance() self.header_template = """ #if __CUDACC_VER_MAJOR__ > ${required_cuda_ver_major} || (__CUDACC_VER_MAJOR__ == ${required_cuda_ver_major} && __CUDACC_VER_MINOR__ >= ${required_cuda_ver_minor}) // ignore warning of cutlass #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wunused-parameter" #pragma GCC diagnostic ignored "-Wstrict-aliasing" #pragma GCC diagnostic ignored "-Wuninitialized" #pragma GCC diagnostic ignored "-Wmaybe-uninitialized" #include "${wrapper_path}" """ self.instance_template = """ ${operation_instance} """ self.epilogue_template = """ #pragma GCC diagnostic pop #endif """ # def __enter__(self): if self.short_path: self.kernel_path = os.path.join(self.kernel_path, "%s.cu" % GlobalCnt.cnt) GlobalCnt.cnt += 1 else: self.kernel_path = os.path.join(self.kernel_path, "%s.cu" % self.operation.procedural_name()) self.kernel_file = open(self.kernel_path, "w") self.kernel_file.write(SubstituteTemplate(self.header_template, { 'wrapper_path': self.wrapper_path, 'required_cuda_ver_major': str(self.operation.required_cuda_ver_major), 'required_cuda_ver_minor': str(self.operation.required_cuda_ver_minor), })) return self # def emit(self): self.kernel_file.write(SubstituteTemplate(self.instance_template, { 'operation_instance': self.instance_emitter.emit(self.operation), })) # emit wrapper wrapper = SubstituteTemplate(self.wrapper_template, { 'operation_name': self.operation.procedural_name(), }) self.kernel_file.write(wrapper) # def __exit__(self, exception_type, exception_value, traceback): self.kernel_file.write(self.epilogue_template) self.kernel_file.close() ################################################################################################### ###################################################################################################