@@ -17,7 +17,7 @@ from library import *
|
||||
class Conv2dOperation:
|
||||
#
|
||||
def __init__(self, conv_kind, iterator_algorithm, arch, tile_description, A, B, C, element_epilogue, \
|
||||
stride_support, epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity4):
|
||||
stride_support, epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity1):
|
||||
|
||||
self.operation_kind = OperationKind.Conv2d
|
||||
self.arch = arch
|
||||
|
||||
@@ -141,18 +141,19 @@ def CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, data_t
|
||||
###########################################################################################################
|
||||
# ConvolutionOperator support variations
|
||||
# ____________________________________________________________________
|
||||
# ConvolutionalOperator | Analytic | Optimized
|
||||
# ConvolutionalOperator | Analytic | Optimized
|
||||
# ____________________________________________________________________
|
||||
# | Fprop | (strided) | (strided)
|
||||
# | Dgrad | (strided, unity*) | (unity)
|
||||
# | Wgrad | (strided) | (strided)
|
||||
# | Fprop | (strided) | (strided)
|
||||
# | Dgrad | (strided, unity*) | (strided, unity)
|
||||
# | Wgrad | (strided) | (strided)
|
||||
# ____________________________________________________________________
|
||||
#
|
||||
# Note : Operator marked (*) are supported but not generated to keep the instantiated kernel count low
|
||||
###########################################################################################################
|
||||
# Convolution for 2D operations
|
||||
def CreateConv2dOperator(manifest, layout, tile_descriptions, data_type, alignment, \
|
||||
conv_kinds = [ConvKind.Fprop, ConvKind.Dgrad, ConvKind.Wgrad], epilogue_functor = EpilogueFunctor.LinearCombination):
|
||||
conv_kinds = [ConvKind.Fprop, ConvKind.Dgrad, ConvKind.Wgrad], \
|
||||
epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity4):
|
||||
|
||||
element_a, element_b, element_c, element_epilogue = data_type
|
||||
|
||||
@@ -169,33 +170,66 @@ def CreateConv2dOperator(manifest, layout, tile_descriptions, data_type, alignme
|
||||
operations = []
|
||||
|
||||
for tile in tile_descriptions:
|
||||
for conv_kind in conv_kinds:
|
||||
A = TensorDescription(element_a, layout[0], alignment)
|
||||
B = TensorDescription(element_b, layout[1], alignment)
|
||||
C = TensorDescription(element_c, layout[2], alignment_c)
|
||||
|
||||
swizzling_functor_ = swizzling_functor
|
||||
|
||||
#
|
||||
# Conv2d Fprop
|
||||
#
|
||||
if ConvKind.Fprop in conv_kinds:
|
||||
|
||||
# Strided support for Analytic and Optimized Fprop
|
||||
for iterator_algorithm in iterator_algorithms:
|
||||
A = TensorDescription(element_a, layout[0], alignment)
|
||||
B = TensorDescription(element_b, layout[1], alignment)
|
||||
C = TensorDescription(element_c, layout[2], alignment_c)
|
||||
new_operation = Conv2dOperation(ConvKind.Fprop, iterator_algorithm, tile.minimum_compute_capability, tile,\
|
||||
A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor, swizzling_functor_)
|
||||
|
||||
# unity stride only for Optimized Dgrad
|
||||
if (iterator_algorithm == IteratorAlgorithm.Optimized) and (conv_kind == ConvKind.Dgrad):
|
||||
new_operation = Conv2dOperation(conv_kind, iterator_algorithm, tile.minimum_compute_capability, tile,\
|
||||
A, B, C, element_epilogue, StrideSupport.Unity, epilogue_functor)
|
||||
manifest.append(new_operation)
|
||||
operations.append(new_operation)
|
||||
|
||||
manifest.append(new_operation)
|
||||
operations.append(new_operation)
|
||||
#
|
||||
# Conv2d Dgrad
|
||||
#
|
||||
if ConvKind.Dgrad in conv_kinds:
|
||||
|
||||
# strided dgrad is not supported by Optimized Dgrad
|
||||
if (iterator_algorithm == IteratorAlgorithm.Optimized) and (conv_kind == ConvKind.Dgrad):
|
||||
continue
|
||||
# Unity stride for Analytic and Optimized Dgrad
|
||||
for iterator_algorithm in iterator_algorithms:
|
||||
new_operation = Conv2dOperation(ConvKind.Dgrad, iterator_algorithm, tile.minimum_compute_capability, tile,\
|
||||
A, B, C, element_epilogue, StrideSupport.Unity, epilogue_functor, swizzling_functor_)
|
||||
|
||||
# strided support for Fprop (Analytic/Optimized), Dgrad (Analytic), and Wgrad (Analytic)
|
||||
new_operation = Conv2dOperation(conv_kind, iterator_algorithm, tile.minimum_compute_capability, tile,\
|
||||
A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor)
|
||||
manifest.append(new_operation)
|
||||
operations.append(new_operation)
|
||||
|
||||
# Strided support for Analytic Dgrad
|
||||
# strided dgrad uses a special threadblock swizzle
|
||||
# note that SwizzlingFunctor.StridedDgradHorizontal might be
|
||||
# better for problem sizes with large activation channel count
|
||||
swizzling_functor_strided_dgrad_ = SwizzlingFunctor.StridedDgradIdentity1
|
||||
|
||||
new_operation = Conv2dOperation(ConvKind.Dgrad, IteratorAlgorithm.Analytic, tile.minimum_compute_capability, tile,\
|
||||
A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor, swizzling_functor_strided_dgrad_)
|
||||
|
||||
manifest.append(new_operation)
|
||||
operations.append(new_operation)
|
||||
|
||||
#
|
||||
# Conv2d Wgrad
|
||||
#
|
||||
if ConvKind.Wgrad in conv_kinds:
|
||||
|
||||
# Strided support for Analytic and Optimized Wgrad
|
||||
for iterator_algorithm in iterator_algorithms:
|
||||
new_operation = Conv2dOperation(ConvKind.Wgrad, iterator_algorithm, tile.minimum_compute_capability, tile,\
|
||||
A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor, swizzling_functor_)
|
||||
|
||||
manifest.append(new_operation)
|
||||
operations.append(new_operation)
|
||||
|
||||
return operations
|
||||
|
||||
|
||||
# Convolution for 3D operations
|
||||
def CreateConv3dOperator(manifest, layout, tile_descriptions, data_type, alignment, \
|
||||
conv_kinds = [ConvKind.Fprop, ConvKind.Dgrad, ConvKind.Wgrad], epilogue_functor = EpilogueFunctor.LinearCombination):
|
||||
@@ -315,6 +349,11 @@ def GenerateSM50_Simt_complex(manifest, args):
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([128, 64, 8], 2, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 8], 2, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 8], 2, [2, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 32, 8], 2, [2, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 32, 128, 8], 2, [1, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 8], 2, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
@@ -1272,6 +1311,7 @@ def GenerateSM80_TensorOp_16816(manifest, args):
|
||||
TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
@@ -1698,9 +1738,10 @@ def GenerateSM80_TensorOp_16864_TN(manifest, args):
|
||||
TileDescription([256, 64, 256], 4, [4, 1, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([ 64, 256, 256], 4, [1, 4, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([128, 128, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([128, 128, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 256], 5, [2, 2, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([ 64, 64, 256], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, DataType.s32]
|
||||
@@ -1713,14 +1754,14 @@ def GenerateSM80_TensorOp_16864_TN(manifest, args):
|
||||
|
||||
operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
|
||||
conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC)
|
||||
CreateConv2dOperator(manifest, conv_layout, tile_descriptions,
|
||||
data_type, 32, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
|
||||
operations += CreateConv2dOperator(manifest, conv_layout, tile_descriptions,
|
||||
data_type_mixed, 32, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
|
||||
for op in operations:
|
||||
if op.tile_description.threadblock_shape[1] >= 128:
|
||||
op.C.alignment = 8
|
||||
@@ -1934,6 +1975,7 @@ def GenerateSM80_TensorOp_1688(manifest, args):
|
||||
TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
@@ -1993,7 +2035,7 @@ def GenerateSM80_TensorOp_1688_fast_math(manifest, args):
|
||||
[16, 8, 8], \
|
||||
DataType.bf16, DataType.bf16, DataType.f32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_fast_bf16)
|
||||
MathOperation.multiply_add_fast_bf16),
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
@@ -2017,6 +2059,7 @@ def GenerateSM80_TensorOp_1688_fast_math(manifest, args):
|
||||
TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
@@ -2031,6 +2074,7 @@ def GenerateSM80_TensorOp_1688_fast_math(manifest, args):
|
||||
CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, 4)
|
||||
#
|
||||
|
||||
#
|
||||
#
|
||||
def GenerateSM80_SparseTensorOp_16816_fast_math(manifest, args):
|
||||
|
||||
@@ -2155,9 +2199,9 @@ def GenerateSM80_TensorOp_884(manifest, args):
|
||||
alignment_constraints = [1,]
|
||||
|
||||
tile_descriptions = [
|
||||
TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc_smem_limited),
|
||||
TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
@@ -2463,6 +2507,7 @@ if __name__ == "__main__":
|
||||
parser.add_argument('--kernel-filter-file', type=str, default=None, required=False, help='Full path of filter file')
|
||||
parser.add_argument('--selected-kernel-list', type=str, default=None, required=False,
|
||||
help='Specify the output log file containing all enabled kernels in this build')
|
||||
parser.add_argument("--interface-dir", default=None, required=False, help="Interface header to kernels")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -437,6 +437,10 @@ class SwizzlingFunctor(enum.Enum):
|
||||
Identity2 = enum_auto()
|
||||
Identity4 = enum_auto()
|
||||
Identity8 = enum_auto()
|
||||
Horizontal = enum_auto()
|
||||
StridedDgradIdentity1 = enum_auto()
|
||||
StridedDgradIdentity4 = enum_auto()
|
||||
StridedDgradHorizontal = enum_auto()
|
||||
|
||||
#
|
||||
SwizzlingFunctorTag = {
|
||||
@@ -444,6 +448,10 @@ SwizzlingFunctorTag = {
|
||||
SwizzlingFunctor.Identity2: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<2>',
|
||||
SwizzlingFunctor.Identity4: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<4>',
|
||||
SwizzlingFunctor.Identity8: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>',
|
||||
SwizzlingFunctor.Horizontal: 'cutlass::gemm::threadblock::GemmHorizontalThreadblockSwizzle',
|
||||
SwizzlingFunctor.StridedDgradIdentity1: 'cutlass::conv::threadblock::StridedDgradIdentityThreadblockSwizzle<1>',
|
||||
SwizzlingFunctor.StridedDgradIdentity4: 'cutlass::conv::threadblock::StridedDgradIdentityThreadblockSwizzle<4>',
|
||||
SwizzlingFunctor.StridedDgradHorizontal: 'cutlass::conv::threadblock::StridedDgradHorizontalThreadblockSwizzle',
|
||||
}
|
||||
|
||||
###################################################################################################
|
||||
|
||||
@@ -101,6 +101,70 @@ void initialize_all_${operation_name}_operations(Manifest &manifest) {
|
||||
self.top_level_file.write(self.epilogue_template)
|
||||
self.top_level_file.close()
|
||||
|
||||
class EmitInterfaceLibrary:
|
||||
def __init__(self, generated_path, operation_count, args):
|
||||
self.generated_path = generated_path
|
||||
self.args = args
|
||||
|
||||
|
||||
self.prototypes = []
|
||||
self.fn_calls = []
|
||||
self.operation_count = str(operation_count)
|
||||
|
||||
self.top_level_hdr_template = '''
|
||||
/*
|
||||
Generated by manifest.py - Do not edit.
|
||||
*/
|
||||
'''
|
||||
self.top_level_prologue = '''
|
||||
|
||||
#include "cutlass/library/library.h"
|
||||
#include "cutlass/library/manifest.h"
|
||||
|
||||
namespace cutlass {
|
||||
\tnamespace library {
|
||||
|
||||
${prototypes}
|
||||
|
||||
\t\tvoid initialize_all(Manifest &manifest) {
|
||||
\t\t\tmanifest.reserve(${operation_count});\n\n
|
||||
${fn_calls}
|
||||
\t\t\t}
|
||||
|
||||
\t} // namespace library
|
||||
} // namespace cutlass
|
||||
|
||||
'''
|
||||
|
||||
#
|
||||
def __enter__(self):
|
||||
self.top_level_path = os.path.join(self.generated_path, 'initialize_all.cpp')
|
||||
|
||||
self.top_level_file = open(self.top_level_path, "w")
|
||||
self.top_level_file.write(self.top_level_hdr_template)
|
||||
|
||||
self.source_files = [self.top_level_path,]
|
||||
|
||||
return self
|
||||
|
||||
#
|
||||
def emit(self, operation_name):
|
||||
self.prototypes.append(SubstituteTemplate(
|
||||
"\t\tvoid initialize_all_${operation_kind}_operations(Manifest &manifest);",
|
||||
{'operation_kind': operation_name}))
|
||||
self.fn_calls.append(SubstituteTemplate(
|
||||
"\t\t\tinitialize_all_${operation_kind}_operations(manifest);",
|
||||
{'operation_kind': operation_name}))
|
||||
|
||||
|
||||
|
||||
#
|
||||
def __exit__(self, exception_type, exception_value, traceback):
|
||||
self.top_level_file.write(SubstituteTemplate(self.top_level_prologue, {'prototypes':"\n".join(self.prototypes),
|
||||
'fn_calls':"\n".join(self.fn_calls),
|
||||
'operation_count': self.operation_count}))
|
||||
self.top_level_file.close()
|
||||
|
||||
###################################################################################################
|
||||
###################################################################################################
|
||||
|
||||
@@ -150,27 +214,6 @@ class Manifest:
|
||||
|
||||
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 get_kernel_filters (self, kernelListFile):
|
||||
@@ -288,6 +331,9 @@ void initialize_all(Manifest &manifest) {
|
||||
operation_emitters = {
|
||||
GeneratorTarget.Library: EmitOperationKindLibrary
|
||||
}
|
||||
interface_emitters = {
|
||||
GeneratorTarget.Library: EmitInterfaceLibrary
|
||||
}
|
||||
|
||||
generated_path = os.path.join(self.args.curr_build_dir, 'generated')
|
||||
|
||||
@@ -299,38 +345,20 @@ void initialize_all(Manifest &manifest) {
|
||||
|
||||
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 = []
|
||||
with interface_emitters[target](generated_path, self.operation_count, self.args) as iface_emitter:
|
||||
for operation_kind, configurations in self.operations.items():
|
||||
prototypes.append(SubstituteTemplate(
|
||||
"void initialize_all_${operation_kind}_operations(Manifest &manifest);",
|
||||
{'operation_kind': OperationKindNames[operation_kind]}))
|
||||
iface_emitter.emit(OperationKindNames[operation_kind])
|
||||
|
||||
top_level_file.write(SubstituteTemplate(self.top_level_prologue,
|
||||
{'prototypes': "\n".join(prototypes)}))
|
||||
source_files += iface_emitter.source_files
|
||||
|
||||
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)
|
||||
# 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)
|
||||
source_files += operation_kind_emitter.source_files
|
||||
|
||||
# write the manifest.cmake file containing paths from all targets
|
||||
manifest_path = os.path.join(generated_path, "manifest.cmake")
|
||||
|
||||
Reference in New Issue
Block a user