Function GenerateConv2d_Simt

dnn/scripts/cutlass_generator/generator.py:188–279

(args)

def GenerateConv2d_Simt(args):
    operations = []

    layouts = [(LayoutType.TensorNC4HW4, LayoutType.TensorC4RSK4)]

    math_instructions = [
        MathInstruction(
            [1, 1, 4],
            DataType.s8,
            DataType.s8,
            DataType.s32,
            OpcodeClass.Simt,
            MathOperation.multiply_add,
        )
    ]

    dst_layouts = [
        LayoutType.TensorNC4HW4,
        LayoutType.TensorNC32HW32,
        LayoutType.TensorNHWC,
        LayoutType.TensorNHWC,
        LayoutType.TensorNCHW,
    ]

    dst_types = [DataType.s8, DataType.s8, DataType.u4, DataType.s4, DataType.f32]

    max_cc = 1024

    for math_inst in math_instructions:
        for layout in layouts:
            for dst_type, dst_layout in zip(dst_types, dst_layouts):
                if dst_type == DataType.s4 or dst_type == DataType.u4:
                    min_cc = 75
                    use_special_optimization = SpecialOptimizeDesc.NoneSpecialOpt
                else:
                    min_cc = 61
                    use_special_optimization = SpecialOptimizeDesc.ConvFilterUnity
                tile_descriptions = [
                    TileDescription(
                        [128, 128, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc
                    ),
                    TileDescription(
                        [128, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc
                    ),
                    TileDescription(
                        [64, 128, 32], 2, [1, 4, 1], math_inst, min_cc, max_cc
                    ),
                    TileDescription(
                        [128, 32, 32], 2, [2, 1, 1], math_inst, min_cc, max_cc
                    ),
                    TileDescription(
                        [32, 128, 32], 2, [1, 2, 1], math_inst, min_cc, max_cc
                    ),
                    TileDescription(
                        [32, 64, 32], 2, [1, 1, 1], math_inst, min_cc, max_cc
                    ),
                    TileDescription(
                        [64, 32, 32], 2, [1, 1, 1], math_inst, min_cc, max_cc