Function CreateConv2dOperator

dnn/scripts/cutlass_generator/generator.py:107–181

(
    manifest,
    layout,
    tile_descriptions,
    data_type,
    alignment,
    conv_kinds=[ConvKind.Fprop, ConvKind.Dgrad, ConvKind.Wgrad],
    epilogue_functor=EpilogueFunctor.LinearCombination,
)

###########################################################################################################
# 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,
):

    element_a, element_b, element_c, element_epilogue = data_type

    # one exceptional case
    alignment_c = min(8, alignment)

    # iterator algorithm (analytic and optimized)
    iterator_algorithms = [IteratorAlgorithm.Analytic, IteratorAlgorithm.Optimized]

    # by default, only generate the largest tile size
    if manifest.args.kernels == "":
        tile_descriptions = [tile_descriptions[0]]

    operations = []

    for tile in tile_descriptions:
        for conv_kind in conv_kinds:
            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)

                # 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)

                # strided dgrad is not supported by Optimized Dgrad
                if (iterator_algorithm == IteratorAlgorithm.Optimized) and (
                    conv_kind == ConvKind.Dgrad
                ):
                    continue

                # strided support for Fprop (Analytic/Optimized), Dgrad (Analytic), and Wgrad (Analytic)