CUTLASS 2.0 (#62)
CUTLASS 2.0 Substantially refactored for - Better performance, particularly for native Turing Tensor Cores - Robust and durable templates spanning the design space - Encapsulated functionality embodying modern C++11 programming techniques - Optimized containers and data types for efficient, generic, portable device code Updates to: - Quick start guide - Documentation - Utilities - CUTLASS Profiler Native Turing Tensor Cores - Efficient GEMM kernels targeting Turing Tensor Cores - Mixed-precision floating point, 8-bit integer, 4-bit integer, and binarized operands Coverage of existing CUTLASS functionality: - GEMM kernels targeting CUDA and Tensor Cores in NVIDIA GPUs - Volta Tensor Cores through native mma.sync and through WMMA API - Optimizations such as parallel reductions, threadblock rasterization, and intra-threadblock reductions - Batched GEMM operations - Complex-valued GEMMs Note: this commit and all that follow require a host compiler supporting C++11 or greater.
This commit is contained in:
Andrew Kerr
@@ -0,0 +1,450 @@
|
||||
#
|
||||
# \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):
|
||||
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
|
||||
|
||||
#
|
||||
def core_name(self):
|
||||
''' The basic operation kind is prefixed with a letter indicating the accumulation type. '''
|
||||
if self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp or \
|
||||
self.tile_description.math_instruction.opcode_class == OpcodeClass.WmmaTensorOp:
|
||||
inst_shape = "%d%d%d" % tuple(self.tile_description.math_instruction.instruction_shape)
|
||||
else:
|
||||
inst_shape = ''
|
||||
|
||||
return "%s%s%s" % (ShortDataTypeNames[self.tile_description.math_instruction.element_accumulator], inst_shape, GemmKindNames[self.gemm_kind])
|
||||
|
||||
#
|
||||
def extended_name(self):
|
||||
''' Append data types if they differ from compute type. '''
|
||||
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 procedural_name(self):
|
||||
''' The full procedural name indicates architecture, extended name, tile size, and layout. '''
|
||||
if self.tile_description.stages > 2:
|
||||
threadblock = "%dx%d_%dx%d" % (
|
||||
self.tile_description.threadblock_shape[0],
|
||||
self.tile_description.threadblock_shape[1],
|
||||
self.tile_description.threadblock_shape[2],
|
||||
self.tile_description.stages
|
||||
)
|
||||
else:
|
||||
threadblock = "%dx%d" % (self.tile_description.threadblock_shape[0], self.tile_description.threadblock_shape[1])
|
||||
|
||||
opcode_class_name = OpcodeClassNames[self.tile_description.math_instruction.opcode_class]
|
||||
|
||||
return SubstituteTemplate(
|
||||
"cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}",
|
||||
{
|
||||
'opcode_class': opcode_class_name,
|
||||
'extended_name': self.extended_name(),
|
||||
'threadblock': threadblock,
|
||||
'layout': "%s%s" % (ShortLayoutTypeNames[self.A.layout], ShortLayoutTypeNames[self.B.layout]),
|
||||
}
|
||||
)
|
||||
|
||||
#
|
||||
def configuration_name(self):
|
||||
''' The full procedural name indicates architecture, extended name, tile size, and layout. '''
|
||||
return self.procedural_name()
|
||||
|
||||
###################################################################################################
|
||||
#
|
||||
# Emits single instances of a CUTLASS device-wide operator
|
||||
#
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class EmitGemmInstance:
|
||||
''' Responsible for emitting a CUTLASS template definition'''
|
||||
|
||||
def __init__(self):
|
||||
self.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}>,
|
||||
cutlass::epilogue::thread::LinearCombination<
|
||||
${element_c},
|
||||
${epilogue_vector_length},
|
||||
${element_accumulator},
|
||||
${element_epilogue}
|
||||
>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle,
|
||||
${stages}
|
||||
>;
|
||||
"""
|
||||
|
||||
def emit(self, operation):
|
||||
|
||||
warp_shape = [operation.tile_description.threadblock_shape[idx] // operation.tile_description.warp_count[idx] for idx in range(3)]
|
||||
#warp_shape[2] = operation.tile_description.math_instruction.instruction_shape[2]
|
||||
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])
|
||||
|
||||
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.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]),
|
||||
'epilogue_vector_length': str(epilogue_vector_length),
|
||||
'element_epilogue': str(DataTypeTag[operation.element_epilogue]),
|
||||
'stages': str(operation.tile_description.stages)
|
||||
}
|
||||
|
||||
return SubstituteTemplate(self.template, values)
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class EmitGemmBatchedInstance:
|
||||
''' Responsible for emitting a CUTLASS template definition'''
|
||||
|
||||
def __init__(self):
|
||||
self.template = """
|
||||
// Gemm operator ${operation_name}
|
||||
using Operation_${operation_name} = cutlass::gemm::device::GemmBatched<
|
||||
${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}>,
|
||||
cutlass::epilogue::thread::LinearCombination<
|
||||
${element_c},
|
||||
${epilogue_vector_length},
|
||||
${element_accumulator},
|
||||
${element_epilogue}
|
||||
>,
|
||||
cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
|
||||
${stages},
|
||||
${align_a},
|
||||
${align_b}
|
||||
>;
|
||||
"""
|
||||
|
||||
def emit(self, operation):
|
||||
|
||||
warp_shape = [operation.tile_description.threadblock_shape[idx] // operation.tile_description.warp_count[idx] for idx in range(3)]
|
||||
#warp_shape[2] = operation.tile_description.math_instruction.instruction_shape[2]
|
||||
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])
|
||||
|
||||
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.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]),
|
||||
'epilogue_vector_length': str(epilogue_vector_length),
|
||||
'element_epilogue': str(DataTypeTag[operation.element_epilogue]),
|
||||
'stages': str(operation.tile_description.stages),
|
||||
'align_a': str(operation.A.alignment),
|
||||
'align_b': str(operation.B.alignment),
|
||||
}
|
||||
|
||||
return SubstituteTemplate(self.template, values)
|
||||
|
||||
###################################################################################################
|
||||
#
|
||||
# Generator functions for all layouts
|
||||
#
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def GenerateGemmSimt(gemm_kind, manifest, tile_descriptions, min_cc):
|
||||
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),
|
||||
]
|
||||
|
||||
# for each tile configuration, emit a GEMM
|
||||
for tile in tile_descriptions:
|
||||
for layout in layouts:
|
||||
|
||||
A = TensorDescription(tile.math_instruction.element_a, layout[0], 1)
|
||||
B = TensorDescription(tile.math_instruction.element_b, layout[1], 1)
|
||||
C = TensorDescription(tile.math_instruction.element_accumulator, layout[2], 1)
|
||||
|
||||
manifest.append(GemmOperation(gemm_kind, 50, tile, A, B, C, tile.math_instruction.element_accumulator))
|
||||
|
||||
#
|
||||
def GenerateGemmTensorOp(gemm_kind, manifest, tile_descriptions, min_cc, minimum_alignment = [128,]):
|
||||
|
||||
# Canonical matrix layouts
|
||||
canonical_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),
|
||||
]
|
||||
|
||||
# Interleaved matrix layouts
|
||||
interleaved_layouts = {
|
||||
8: [
|
||||
#(LayoutType.ColumnMajorInterleaved32, LayoutType.RowMajorInterleaved32, LayoutType.ColumnMajorInterleaved32),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
],
|
||||
4: [
|
||||
#(LayoutType.ColumnMajorInterleaved64, LayoutType.RowMajorInterleaved64, LayoutType.ColumnMajorInterleaved64),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
}
|
||||
|
||||
# for each tile configuration, emit a GEMM
|
||||
for align in minimum_alignment:
|
||||
for tile in tile_descriptions:
|
||||
|
||||
min_input_size = min(DataTypeSize[tile.math_instruction.element_a], DataTypeSize[tile.math_instruction.element_a])
|
||||
|
||||
# If the data type is large enough, use canonical layouts.
|
||||
if min_input_size >= 16:
|
||||
layouts = canonical_layouts
|
||||
else:
|
||||
layouts = interleaved_layouts[min_input_size]
|
||||
|
||||
for layout in layouts:
|
||||
|
||||
#
|
||||
output_types = [tile.math_instruction.element_a, tile.math_instruction.element_accumulator] \
|
||||
if DataTypeSize[tile.math_instruction.element_accumulator] == 32 \
|
||||
else [tile.math_instruction.element_accumulator,]
|
||||
|
||||
align_a = align // DataTypeSize[tile.math_instruction.element_a]
|
||||
align_b = align // DataTypeSize[tile.math_instruction.element_b]
|
||||
|
||||
|
||||
for output_type in output_types:
|
||||
|
||||
rows_per_warp = 8 // tile.warp_count[1]
|
||||
align_c = min(int(align / DataTypeSize[output_type]), tile.threadblock_shape[1] * rows_per_warp // 32)
|
||||
|
||||
A = TensorDescription(tile.math_instruction.element_a, layout[0], align_a)
|
||||
B = TensorDescription(tile.math_instruction.element_b, layout[1], align_b)
|
||||
C = TensorDescription(output_type, layout[2], max(1, align_c))
|
||||
|
||||
element_epilogue = DataType.f32 if tile.math_instruction.element_accumulator == DataType.s32 \
|
||||
else tile.math_instruction.element_accumulator
|
||||
|
||||
manifest.append(GemmOperation(gemm_kind, min_cc, tile, A, B, C, element_epilogue))
|
||||
|
||||
|
||||
#
|
||||
def GenerateGemmWmmaTensorOp(gemm_kind, manifest, tile_descriptions, min_cc, minimum_alignment = [128,]):
|
||||
|
||||
# Wmma supported matrix layouts
|
||||
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),
|
||||
]
|
||||
|
||||
# for each tile configuration, emit a GEMM
|
||||
for align in minimum_alignment:
|
||||
for tile in tile_descriptions:
|
||||
for layout in layouts:
|
||||
|
||||
#
|
||||
output_types = [tile.math_instruction.element_a, tile.math_instruction.element_accumulator] \
|
||||
if DataTypeSize[tile.math_instruction.element_accumulator] == 32 \
|
||||
else [tile.math_instruction.element_accumulator,]
|
||||
|
||||
align_a = align // DataTypeSize[tile.math_instruction.element_a]
|
||||
align_b = align // DataTypeSize[tile.math_instruction.element_b]
|
||||
|
||||
|
||||
for output_type in output_types:
|
||||
|
||||
rows_per_warp = 8 // tile.warp_count[1]
|
||||
align_c = min(int(align / DataTypeSize[output_type]), tile.threadblock_shape[1] * rows_per_warp // 32)
|
||||
|
||||
A = TensorDescription(tile.math_instruction.element_a, layout[0], align_a)
|
||||
B = TensorDescription(tile.math_instruction.element_b, layout[1], align_b)
|
||||
C = TensorDescription(output_type, layout[2], max(1, align_c))
|
||||
|
||||
element_epilogue = DataType.f32 if tile.math_instruction.element_accumulator == DataType.s32 \
|
||||
else tile.math_instruction.element_accumulator
|
||||
|
||||
manifest.append(GemmOperation(gemm_kind, min_cc, tile, A, B, C, element_epilogue))
|
||||
|
||||
###################################################################################################
|
||||
#
|
||||
# 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.Batched: EmitGemmBatchedInstance
|
||||
}
|
||||
|
||||
self.gemm_kind_wrappers = {
|
||||
GemmKind.Gemm: 'GemmOperation',
|
||||
GemmKind.Batched: 'GemmBatchedOperation',
|
||||
}
|
||||
|
||||
self.wmma_guard_start = "#if defined(CUTLASS_ARCH_WMMA_SM${sm_number}_ENABLED)"
|
||||
|
||||
self.instance_template = """
|
||||
${compile_guard_start}
|
||||
manifest.append(new ${gemm_kind}<Operation_${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"
|
||||
|
||||
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(SubstituteTemplate(self.header_template, {
|
||||
'configuration_name': self.configuration_name
|
||||
}))
|
||||
self.operations = []
|
||||
return self
|
||||
|
||||
def emit(self, operation):
|
||||
emitter = self.instance_emitter[operation.gemm_kind]()
|
||||
|
||||
self.operations.append(operation)
|
||||
self.configuration_file.write(emitter.emit(operation))
|
||||
self.configuration_file.write(SubstituteTemplate(self.instance_template, {
|
||||
'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):
|
||||
self.configuration_file.write(self.epilogue_template)
|
||||
self.configuration_file.close()
|
||||
|
||||
###################################################################################################
|
||||
###################################################################################################
|
||||
@@ -0,0 +1,235 @@
|
||||
#
|
||||
# \file generator.py
|
||||
#
|
||||
# \brief Generates the CUTLASS Library's instances
|
||||
#
|
||||
|
||||
import enum
|
||||
import os.path
|
||||
import shutil
|
||||
import argparse
|
||||
|
||||
from library import *
|
||||
from manifest import *
|
||||
from gemm_operation import *
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def CudaToolkitVersionSatisfies(semantic_ver_string, major, minor, patch = 0):
|
||||
if semantic_ver_string == '':
|
||||
cuda_version = [10, 2, 0]
|
||||
else:
|
||||
cuda_version = [int(x) for x in semantic_ver_string.split('.')]
|
||||
|
||||
return cuda_version >= [major, minor, patch]
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def GenerateSM50(manifest, args):
|
||||
|
||||
min_cc = 50
|
||||
max_cc = 1024
|
||||
stages = 2
|
||||
|
||||
# single-precision
|
||||
inst = MathInstruction([1, 1, 1], DataType.f32, DataType.f32, DataType.f32, OpcodeClass.Simt)
|
||||
tile_descriptions = [
|
||||
TileDescription([128, 128, 8], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 8], stages, [2, 4, 1], inst, min_cc, max_cc),
|
||||
TileDescription([256, 128, 8], stages, [4, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, 8], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 8], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 32, 8], stages, [4, 1, 1], inst, min_cc, max_cc),
|
||||
TileDescription([32, 128, 8], stages, [1, 4, 1], inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
GenerateGemmSimt(GemmKind.Gemm, manifest, tile_descriptions, min_cc)
|
||||
GenerateGemmSimt(GemmKind.Batched, manifest, tile_descriptions, min_cc)
|
||||
|
||||
# double precision
|
||||
inst = MathInstruction([1, 1, 1], DataType.f64, DataType.f64, DataType.f64, OpcodeClass.Simt)
|
||||
tile_descriptions = [
|
||||
TileDescription([128, 128, 8], stages, [4, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, 8], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 8], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 32, 8], stages, [4, 1, 1], inst, min_cc, max_cc),
|
||||
TileDescription([32, 128, 8], stages, [1, 4, 1], inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
GenerateGemmSimt(GemmKind.Gemm, manifest, tile_descriptions, min_cc)
|
||||
GenerateGemmSimt(GemmKind.Batched, manifest, tile_descriptions, min_cc)
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def GenerateSM60(manifest, args):
|
||||
|
||||
min_cc = 60
|
||||
max_cc = 1024
|
||||
stages = 2
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction([1, 1, 1], DataType.f16, DataType.f16, DataType.f16, OpcodeClass.Simt),
|
||||
]
|
||||
|
||||
tile_descriptions = []
|
||||
|
||||
for inst in math_instructions:
|
||||
tile_descriptions += [
|
||||
TileDescription([256, 256, 8], stages, [4, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 8], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 8], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, 8], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([32, 128, 8], stages, [1, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 32, 8], stages, [2, 1, 1], inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
GenerateGemmSimt(GemmKind.Gemm, manifest, tile_descriptions, min_cc)
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def GenerateSM61(manifest, args):
|
||||
|
||||
min_cc = 61
|
||||
max_cc = 1024
|
||||
stages = 2
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction([1, 1, 4], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.Simt),
|
||||
]
|
||||
|
||||
tile_descriptions = []
|
||||
|
||||
for inst in math_instructions:
|
||||
tile_descriptions += [
|
||||
TileDescription([128, 256, 32], stages, [2, 4, 1], inst, min_cc, max_cc),
|
||||
TileDescription([256, 128, 32], stages, [4, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 32], stages, [2, 4, 1], inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, 32], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 32], stages, [4, 1, 1], inst, min_cc, max_cc),
|
||||
TileDescription([32, 128, 32], stages, [1, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 32, 32], stages, [2, 1, 1], inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
GenerateGemmSimt(GemmKind.Gemm, manifest, tile_descriptions, min_cc)
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def GenerateSM70(manifest, args):
|
||||
|
||||
min_cc = 70
|
||||
max_cc = 75
|
||||
stages = 2
|
||||
k_groups = 8
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction([8, 8, 4], DataType.f16, DataType.f16, DataType.f16, OpcodeClass.TensorOp),
|
||||
MathInstruction([8, 8, 4], DataType.f16, DataType.f16, DataType.f32, OpcodeClass.TensorOp),
|
||||
]
|
||||
|
||||
tile_descriptions = []
|
||||
|
||||
for inst in math_instructions:
|
||||
kblock = k_groups * inst.instruction_shape[2]
|
||||
tile_descriptions += [
|
||||
TileDescription([256, 128, kblock], stages, [4, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, kblock], stages, [2, 4, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([64, 64, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
if CudaToolkitVersionSatisfies(args.cuda_version, 10, 1):
|
||||
GenerateGemmTensorOp(GemmKind.Gemm, manifest, tile_descriptions, min_cc)
|
||||
GenerateGemmTensorOp(GemmKind.Batched, manifest, tile_descriptions, min_cc)
|
||||
|
||||
# wmma tensor op SM70 Gemm kernels
|
||||
stages = 2
|
||||
k_groups = 2
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction([16, 16, 16], DataType.f16, DataType.f16, DataType.f16, OpcodeClass.WmmaTensorOp),
|
||||
MathInstruction([16, 16, 16], DataType.f16, DataType.f16, DataType.f32, OpcodeClass.WmmaTensorOp),
|
||||
]
|
||||
|
||||
tile_descriptions = []
|
||||
|
||||
for inst in math_instructions:
|
||||
kblock = k_groups * inst.instruction_shape[2]
|
||||
tile_descriptions += [
|
||||
TileDescription([128, 128, kblock], stages, [2, 4, 1], inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([64, 64, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
GenerateGemmWmmaTensorOp(GemmKind.Gemm, manifest, tile_descriptions, min_cc)
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def GenerateSM75(manifest, args):
|
||||
|
||||
min_cc = 75
|
||||
max_cc = 1024
|
||||
stages = 2
|
||||
k_groups = 4
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction([16, 8, 8], DataType.f16, DataType.f16, DataType.f16, OpcodeClass.TensorOp),
|
||||
MathInstruction([16, 8, 8], DataType.f16, DataType.f16, DataType.f32, OpcodeClass.TensorOp),
|
||||
MathInstruction([8, 8, 16], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.TensorOp),
|
||||
MathInstruction([8, 8, 32], DataType.s4, DataType.s4, DataType.s32, OpcodeClass.TensorOp)
|
||||
]
|
||||
|
||||
tile_descriptions = []
|
||||
|
||||
for inst in math_instructions:
|
||||
kblock = k_groups * inst.instruction_shape[2]
|
||||
tile_descriptions += [
|
||||
TileDescription([256, 128, kblock], stages, [4, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, kblock], stages, [2, 4, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
TileDescription([64, 64, kblock], stages, [2, 2, 1], inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
if CudaToolkitVersionSatisfies(args.cuda_version, 10, 2):
|
||||
GenerateGemmTensorOp(GemmKind.Gemm, manifest, tile_descriptions, min_cc)
|
||||
GenerateGemmTensorOp(GemmKind.Batched, manifest, tile_descriptions, min_cc)
|
||||
|
||||
|
||||
###################################################################################################
|
||||
###################################################################################################
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
parser = argparse.ArgumentParser(description="Generates device kernel registration code for CUTLASS Kernels")
|
||||
parser.add_argument("--operations", default="gemm", help="Specifies the operation to generate (gemm, all)")
|
||||
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("--kernels", default='', help='Comma delimited list to filter kernels by name.')
|
||||
parser.add_argument("--cuda-version", default="10.2.0", help="Semantic version string of CUDA Toolkit")
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
manifest = Manifest(args)
|
||||
|
||||
GenerateSM50(manifest, args)
|
||||
GenerateSM60(manifest, args)
|
||||
GenerateSM61(manifest, args)
|
||||
GenerateSM70(manifest, args)
|
||||
GenerateSM75(manifest, args)
|
||||
if 'library' in args.generator_target.split(','):
|
||||
manifest.emit(GeneratorTarget.Library)
|
||||
|
||||
#
|
||||
###################################################################################################
|
||||
@@ -0,0 +1,307 @@
|
||||
#
|
||||
# \file generator.py
|
||||
#
|
||||
# \brief Generates the CUTLASS Library's instances
|
||||
#
|
||||
|
||||
import enum
|
||||
import re
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class GeneratorTarget(enum.Enum):
|
||||
Library = enum.auto()
|
||||
#
|
||||
GeneratorTargetNames = {
|
||||
GeneratorTarget.Library: 'library'
|
||||
}
|
||||
#
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
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()
|
||||
|
||||
#
|
||||
ShortDataTypeNames = {
|
||||
DataType.s32: 'i',
|
||||
DataType.f16: 'h',
|
||||
DataType.f32: 's',
|
||||
DataType.f64: 'd',
|
||||
DataType.cf32: 'c',
|
||||
DataType.cf64: 'z',
|
||||
}
|
||||
|
||||
#
|
||||
DataTypeNames = {
|
||||
DataType.b1: "b1",
|
||||
DataType.u4: "u4",
|
||||
DataType.u8: "u8",
|
||||
DataType.u16: "u16",
|
||||
DataType.u32: "u32",
|
||||
DataType.u64: "u64",
|
||||
DataType.s4: "s4",
|
||||
DataType.s8: "s8",
|
||||
DataType.s16: "s16",
|
||||
DataType.s32: "s32",
|
||||
DataType.s64: "s64",
|
||||
DataType.f16: "f16",
|
||||
DataType.f32: "f32",
|
||||
DataType.f64: "f64",
|
||||
DataType.cf16: "cf16",
|
||||
DataType.cf32: "cf32",
|
||||
DataType.cf64: "cf64",
|
||||
DataType.cu4: "cu4",
|
||||
DataType.cu8: "cu8",
|
||||
DataType.cu16: "cu16",
|
||||
DataType.cu32: "cu32",
|
||||
DataType.cu64: "cu64",
|
||||
DataType.cs4: "cs4",
|
||||
DataType.cs8: "cs8",
|
||||
DataType.cs16: "cs16",
|
||||
DataType.cs32: "cs32",
|
||||
DataType.cs64: "cs64",
|
||||
}
|
||||
|
||||
DataTypeTag = {
|
||||
DataType.b1: "cutlass::uint1b_t",
|
||||
DataType.u4: "cutlass::uint4b_t",
|
||||
DataType.u8: "uint8_t",
|
||||
DataType.u16: "uint16_t",
|
||||
DataType.u32: "uint32_t",
|
||||
DataType.u64: "uint64_t",
|
||||
DataType.s4: "cutlass::int4b_t",
|
||||
DataType.s8: "int8_t",
|
||||
DataType.s16: "int16_t",
|
||||
DataType.s32: "int32_t",
|
||||
DataType.s64: "int64_t",
|
||||
DataType.f16: "cutlass::half_t",
|
||||
DataType.f32: "float",
|
||||
DataType.f64: "double",
|
||||
DataType.cf16: "cutlass::complex<cutlass::half_t>",
|
||||
DataType.cf32: "cutlass::complex<float>",
|
||||
DataType.cf64: "cutlass::complex<double>",
|
||||
DataType.cu4: "cutlass::complex<cutlass::uint4b_t>",
|
||||
DataType.cu8: "cutlass::complex<cutlass::uint8_t>",
|
||||
DataType.cu16: "cutlass::complex<cutlass::uint16_t>",
|
||||
DataType.cu32: "cutlass::complex<cutlass::uint32_t>",
|
||||
DataType.cu64: "cutlass::complex<cutlass::uint64_t>",
|
||||
DataType.cs4: "cutlass::complex<cutlass::int4b_t>",
|
||||
DataType.cs8: "cutlass::complex<cutlass::int8_t>",
|
||||
DataType.cs16: "cutlass::complex<cutlass::int16_t>",
|
||||
DataType.cs32: "cutlass::complex<cutlass::int32_t>",
|
||||
DataType.cs64: "cutlass::complex<cutlass::int64_t>",
|
||||
}
|
||||
|
||||
DataTypeSize = {
|
||||
DataType.b1: 1,
|
||||
DataType.u4: 4,
|
||||
DataType.u8: 4,
|
||||
DataType.u16: 16,
|
||||
DataType.u32: 32,
|
||||
DataType.u64: 64,
|
||||
DataType.s4: 4,
|
||||
DataType.s8: 8,
|
||||
DataType.s16: 16,
|
||||
DataType.s32: 32,
|
||||
DataType.s64: 64,
|
||||
DataType.f16: 16,
|
||||
DataType.f32: 32,
|
||||
DataType.f64: 64,
|
||||
DataType.cf16: 32,
|
||||
DataType.cf32: 64,
|
||||
DataType.cf64: 128,
|
||||
DataType.cu4: 8,
|
||||
DataType.cu8: 16,
|
||||
DataType.cu16: 32,
|
||||
DataType.cu32: 64,
|
||||
DataType.cu64: 128,
|
||||
DataType.cs4: 8,
|
||||
DataType.cs8: 16,
|
||||
DataType.cs16: 32,
|
||||
DataType.cs32: 64,
|
||||
DataType.cs64: 128,
|
||||
}
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
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()
|
||||
|
||||
#
|
||||
LayoutTag = {
|
||||
LayoutType.ColumnMajor: 'cutlass::layout::ColumnMajor',
|
||||
LayoutType.RowMajor: 'cutlass::layout::RowMajor',
|
||||
LayoutType.ColumnMajorInterleaved32: 'cutlass::layout::ColumnMajorInterleaved<32>',
|
||||
LayoutType.RowMajorInterleaved32: 'cutlass::layout::RowMajorInterleaved<32>',
|
||||
LayoutType.ColumnMajorInterleaved64: 'cutlass::layout::ColumnMajorInterleaved<64>',
|
||||
LayoutType.RowMajorInterleaved64: 'cutlass::layout::RowMajorInterleaved<64>',
|
||||
LayoutType.TensorNHWC: 'cutlass::layout::TensorNHWC',
|
||||
LayoutType.TensorNCHW: 'cutlass::layout::TensorNCHW',
|
||||
LayoutType.TensorNGHWC: 'cutlass::layout::TensorNGHWC',
|
||||
LayoutType.TensorNCxHW32: 'cutlass::layout::TensorNCxHW32',
|
||||
LayoutType.TensorNCxHW64: 'cutlass::layout::TensorNCxHW64'
|
||||
}
|
||||
|
||||
#
|
||||
ShortLayoutTypeNames = {
|
||||
LayoutType.ColumnMajor: 'n',
|
||||
LayoutType.ColumnMajorInterleaved32: 'n32',
|
||||
LayoutType.ColumnMajorInterleaved64: 'n64',
|
||||
LayoutType.RowMajor: 't',
|
||||
LayoutType.RowMajorInterleaved32: 't32',
|
||||
LayoutType.RowMajorInterleaved64: 't64',
|
||||
LayoutType.TensorNHWC: 'nhwc',
|
||||
LayoutType.TensorNCHW: 'nchw',
|
||||
LayoutType.TensorNGHWC: 'nghwc',
|
||||
LayoutType.TensorNCxHW32: 'ncxhw32',
|
||||
LayoutType.TensorNCxHW64: 'ncxhw64'
|
||||
}
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class OpcodeClass(enum.Enum):
|
||||
Simt = enum.auto()
|
||||
TensorOp = enum.auto()
|
||||
WmmaTensorOp = enum.auto()
|
||||
|
||||
OpcodeClassNames = {
|
||||
OpcodeClass.Simt: 'simt',
|
||||
OpcodeClass.TensorOp: 'tensorop',
|
||||
OpcodeClass.WmmaTensorOp: 'wmma_tensorop',
|
||||
}
|
||||
|
||||
OpcodeClassTag = {
|
||||
OpcodeClass.Simt: 'cutlass::arch::OpClassSimt',
|
||||
OpcodeClass.TensorOp: 'cutlass::arch::OpClassTensorOp',
|
||||
OpcodeClass.WmmaTensorOp: 'cutlass::arch::OpClassWmmaTensorOp',
|
||||
}
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class OperationKind(enum.Enum):
|
||||
Gemm = enum.auto()
|
||||
#
|
||||
OperationKindNames = {
|
||||
OperationKind.Gemm: 'gemm'
|
||||
}
|
||||
|
||||
#
|
||||
class Target(enum.Enum):
|
||||
library = enum.auto()
|
||||
|
||||
ArchitectureNames = {
|
||||
50: 'maxwell',
|
||||
60: 'pascal',
|
||||
61: 'pascal',
|
||||
70: 'volta',
|
||||
75: 'turing',
|
||||
}
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def SubstituteTemplate(template, values):
|
||||
text = template
|
||||
for key, value in values.items():
|
||||
regex = "\\$\\{%s\\}" % key
|
||||
text = re.sub(regex, value, text)
|
||||
return text
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class GemmKind(enum.Enum):
|
||||
Gemm = enum.auto()
|
||||
Batched = enum.auto()
|
||||
Array = enum.auto()
|
||||
PlanarComplex = enum.auto()
|
||||
PlanarComplexBatched = enum.auto()
|
||||
|
||||
#
|
||||
GemmKindNames = {
|
||||
GemmKind.Gemm: "gemm",
|
||||
GemmKind.Batched: "gemm_batched",
|
||||
GemmKind.Array: "gemm_array",
|
||||
GemmKind.PlanarComplex: "gemm_planar_complex",
|
||||
GemmKind.PlanarComplexBatched: "gemm_planar_complex_batched",
|
||||
}
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class MathInstruction:
|
||||
def __init__(self, instruction_shape, element_a, element_b, element_accumulator, opcode_class):
|
||||
self.instruction_shape = instruction_shape
|
||||
self.element_a = element_a
|
||||
self.element_b = element_b
|
||||
self.element_accumulator = element_accumulator
|
||||
self.opcode_class = opcode_class
|
||||
|
||||
|
||||
#
|
||||
class TileDescription:
|
||||
|
||||
def __init__(self, threadblock_shape, stages, warp_count, math_instruction, min_compute, max_compute):
|
||||
self.threadblock_shape = threadblock_shape
|
||||
self.stages = stages
|
||||
self.warp_count = warp_count
|
||||
self.math_instruction = math_instruction
|
||||
self.minimum_compute_capability = min_compute
|
||||
self.maximum_compute_capability = max_compute
|
||||
|
||||
def procedural_name(self):
|
||||
if self.stages == 2:
|
||||
return "%dx%dx%d" % self.threadblock_shape
|
||||
elif self.stages > 2:
|
||||
return "%dx%d_%dx%d" % (self.threadblock_shape[0], self.threadblock_shape[1], self.threadblock_shape[2], self.stages)
|
||||
|
||||
#
|
||||
class TensorDescription:
|
||||
def __init__(self, element, layout, alignment = 1):
|
||||
self.element = element
|
||||
self.layout = layout
|
||||
self.alignment = alignment
|
||||
|
||||
###################################################################################################
|
||||
@@ -0,0 +1,272 @@
|
||||
#
|
||||
# \file generator.py
|
||||
#
|
||||
# \brief Generates the CUTLASS Library's instances
|
||||
#
|
||||
|
||||
import enum
|
||||
import os.path
|
||||
import shutil
|
||||
|
||||
from library import *
|
||||
from gemm_operation import *
|
||||
###################################################################################################
|
||||
|
||||
class EmitOperationKindLibrary:
|
||||
def __init__(self, generated_path, kind, args):
|
||||
self.generated_path = generated_path
|
||||
self.kind = kind
|
||||
self.args = args
|
||||
|
||||
self.emitters = {
|
||||
OperationKind.Gemm: EmitGemmConfigurationLibrary
|
||||
}
|
||||
|
||||
self.configurations = [];
|
||||
|
||||
self.header_template ="""
|
||||
/*
|
||||
Generated by manifest.py - Do not edit.
|
||||
*/
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/library/library.h"
|
||||
#include "cutlass/library/manifest.h"
|
||||
|
||||
namespace cutlass {
|
||||
namespace library {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
"""
|
||||
self.entry_template = """
|
||||
|
||||
//
|
||||
// Entry point to construct operations
|
||||
//
|
||||
void initialize_all_${operation_name}_operations(Manifest &manifest) {
|
||||
"""
|
||||
self.configuration_prototype_template = "void initialize_${configuration_name}(Manifest &manifest);\n"
|
||||
self.configuration_template =" initialize_${configuration_name}(manifest);\n"
|
||||
|
||||
self.epilogue_template ="""
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace library
|
||||
} // namespace cutlass
|
||||
|
||||
"""
|
||||
|
||||
#
|
||||
def __enter__(self):
|
||||
self.operation_path = os.path.join(self.generated_path, OperationKindNames[self.kind])
|
||||
os.mkdir(self.operation_path)
|
||||
|
||||
self.top_level_path = os.path.join(self.operation_path, "all_%s_operations.cu" % OperationKindNames[self.kind])
|
||||
|
||||
self.top_level_file = open(self.top_level_path, "w")
|
||||
self.top_level_file.write(self.header_template)
|
||||
|
||||
self.source_files = [self.top_level_path,]
|
||||
|
||||
return self
|
||||
|
||||
#
|
||||
def emit(self, configuration_name, operations):
|
||||
|
||||
with self.emitters[self.kind](self.operation_path, configuration_name) as configuration_emitter:
|
||||
for operation in operations:
|
||||
configuration_emitter.emit(operation)
|
||||
|
||||
self.source_files.append(configuration_emitter.configuration_path)
|
||||
|
||||
self.configurations.append(configuration_name)
|
||||
self.top_level_file.write(SubstituteTemplate(self.configuration_prototype_template, {'configuration_name': configuration_name} ))
|
||||
|
||||
#
|
||||
def __exit__(self, exception_type, exception_value, traceback):
|
||||
self.top_level_file.write(SubstituteTemplate(self.entry_template, {'operation_name': OperationKindNames[self.kind]}))
|
||||
|
||||
for configuration_name in self.configurations:
|
||||
self.top_level_file.write(SubstituteTemplate(self.configuration_template, {'configuration_name': configuration_name}))
|
||||
|
||||
self.top_level_file.write(self.epilogue_template)
|
||||
self.top_level_file.close()
|
||||
|
||||
###################################################################################################
|
||||
###################################################################################################
|
||||
|
||||
class Options:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class Manifest:
|
||||
|
||||
#
|
||||
def __init__(self, args):
|
||||
self.operations = {}
|
||||
self.args = args
|
||||
self.compute_capabilities = [int(x) for x in args.architectures.split(';')]
|
||||
|
||||
if args.kernels == 'all':
|
||||
self.kernel_names = []
|
||||
else:
|
||||
self.kernel_names = args.kernels.split(',')
|
||||
|
||||
self.operation_count = 0
|
||||
self.operations_by_name = {}
|
||||
self.top_level_prologue = '''
|
||||
|
||||
#include "cutlass/library/library.h"
|
||||
#include "cutlass/library/manifest.h"
|
||||
|
||||
namespace cutlass {
|
||||
namespace library {
|
||||
|
||||
${prototypes}
|
||||
|
||||
void initialize_all(Manifest &manifest) {
|
||||
|
||||
'''
|
||||
self.top_level_reserve = ' manifest.reserve(${operation_count});\n\n'
|
||||
self.top_level_epilogue = '''
|
||||
}
|
||||
|
||||
} // namespace library
|
||||
} // namespace cutlass
|
||||
|
||||
'''
|
||||
|
||||
#
|
||||
def filter(self, operation):
|
||||
''' Filtering operations based on various criteria'''
|
||||
|
||||
# filter based on compute capability
|
||||
enabled = False
|
||||
for cc in self.compute_capabilities:
|
||||
if cc >= operation.tile_description.minimum_compute_capability and \
|
||||
cc <= operation.tile_description.maximum_compute_capability:
|
||||
|
||||
enabled = True
|
||||
break
|
||||
|
||||
if not enabled:
|
||||
return False
|
||||
|
||||
# eliminate duplicates
|
||||
if operation.procedural_name() in self.operations_by_name.keys():
|
||||
return False
|
||||
|
||||
# Filter based on list of valid substrings
|
||||
if len(self.kernel_names):
|
||||
name = operation.procedural_name()
|
||||
enabled = False
|
||||
for name_substr in self.kernel_names:
|
||||
if name_substr in name:
|
||||
enabled = True
|
||||
break
|
||||
|
||||
# todo: filter based on operation kind
|
||||
# todo: filter based on compute data type
|
||||
return enabled
|
||||
#
|
||||
|
||||
#
|
||||
def append(self, operation):
|
||||
'''
|
||||
Inserts the operation.
|
||||
|
||||
operation_kind -> configuration_name -> []
|
||||
'''
|
||||
|
||||
if self.filter(operation):
|
||||
|
||||
self.operations_by_name[operation.procedural_name()] = operation
|
||||
|
||||
# add the configuration
|
||||
configuration_name = operation.configuration_name()
|
||||
|
||||
if operation.operation_kind not in self.operations.keys():
|
||||
self.operations[operation.operation_kind] = {}
|
||||
|
||||
if configuration_name not in self.operations[operation.operation_kind].keys():
|
||||
self.operations[operation.operation_kind][configuration_name] = []
|
||||
|
||||
self.operations[operation.operation_kind][configuration_name].append(operation)
|
||||
self.operation_count += 1
|
||||
#
|
||||
|
||||
#
|
||||
def emit(self, target = GeneratorTarget.Library):
|
||||
|
||||
operation_emitters = {
|
||||
GeneratorTarget.Library: EmitOperationKindLibrary
|
||||
}
|
||||
|
||||
generated_path = os.path.join(self.args.curr_build_dir, 'generated')
|
||||
|
||||
# create generated/
|
||||
if os.path.exists(generated_path):
|
||||
shutil.rmtree(generated_path)
|
||||
|
||||
os.mkdir(generated_path)
|
||||
|
||||
source_files = []
|
||||
|
||||
top_level_path = os.path.join(generated_path, 'initialize_all.cpp')
|
||||
with open(top_level_path, 'w') as top_level_file:
|
||||
|
||||
if target == GeneratorTarget.Library:
|
||||
source_files.append(top_level_path)
|
||||
|
||||
prototypes = []
|
||||
for operation_kind, configurations in self.operations.items():
|
||||
prototypes.append(SubstituteTemplate(
|
||||
"void initialize_all_${operation_kind}_operations(Manifest &manifest);",
|
||||
{'operation_kind': OperationKindNames[operation_kind]}))
|
||||
|
||||
top_level_file.write(SubstituteTemplate(self.top_level_prologue,
|
||||
{'prototypes': "\n".join(prototypes)}))
|
||||
|
||||
top_level_file.write(SubstituteTemplate(
|
||||
self.top_level_reserve, {'operation_count': str(self.operation_count)}))
|
||||
|
||||
# for each operation kind, emit initializer for all configurations
|
||||
for operation_kind, configurations in self.operations.items():
|
||||
with operation_emitters[target](generated_path, operation_kind, self.args) as operation_kind_emitter:
|
||||
for configuration_name, operations in configurations.items():
|
||||
operation_kind_emitter.emit(configuration_name, operations)
|
||||
|
||||
source_files += operation_kind_emitter.source_files
|
||||
|
||||
top_level_file.write(SubstituteTemplate(
|
||||
" initialize_all_${operation_kind}_operations(manifest);\n",
|
||||
{'operation_kind': OperationKindNames[operation_kind]}))
|
||||
|
||||
top_level_file.write(self.top_level_epilogue)
|
||||
|
||||
# write the manifest.cmake file containing paths from all targets
|
||||
manifest_path = os.path.join(generated_path, "manifest.cmake")
|
||||
with open(manifest_path, "w") as manifest_file:
|
||||
|
||||
target_name = 'cutlass_lib'
|
||||
|
||||
target_text = SubstituteTemplate("""cutlass_target_sources(
|
||||
${target_name}
|
||||
PRIVATE
|
||||
""", { 'target_name': target_name})
|
||||
|
||||
manifest_file.write(target_text)
|
||||
|
||||
for source_file in source_files:
|
||||
manifest_file.write(" %s\n" % str(source_file.replace('\\', '/')))
|
||||
manifest_file.write(")")
|
||||
#
|
||||
|
||||
###################################################################################################
|
||||
Reference in New Issue
Block a user