def extract_config(
    func_attrs,
    dtype="float16",
    skip_simt_kernels=False,
    f_apply_special_config=None,
    op_kind=None,
    op_layout=None,
):
    """Extracts cutlass config for conv kernels."""
    import copy

    import cutlass_lib

    spec = CUDASpec()
    lib_dtype = spec.dtype_to_lib_type(dtype)

    if lib_dtype == "float":
        data_type = cutlass_lib.library.DataType.f32
        acc_type = cutlass_lib.library.DataType.f32
    elif "half" in lib_dtype:
        data_type = cutlass_lib.library.DataType.f16
        acc_type = cutlass_lib.library.DataType.f32
        # check target use fp16 acc
        if "use_fp16_acc" in Target.current()._kwargs:
            if Target.current()._kwargs["use_fp16_acc"]:
                acc_type = cutlass_lib.library.DataType.f16
    elif "bfloat16" in lib_dtype:
        data_type = cutlass_lib.library.DataType.bf16
        acc_type = cutlass_lib.library.DataType.f32
    else:
        raise RuntimeError(f"Unsupported dtype {lib_dtype}")

    def f_proc_op(op):
        ret = []
        if (
            skip_simt_kernels
            and op.tile_description.math_instruction.opcode_class
            == cutlass_lib.library.OpcodeClass.Simt
        ):
            return ret

        if (
            op.A.element == data_type
            and op.B.element == data_type
            and op.C.element == data_type
            and op.iterator_algorithm == cutlass_lib.library.IteratorAlgorithm.Optimized
            and op.tile_description.math_instruction.element_accumulator == acc_type
        ):
            op = copy.deepcopy(op)

            # set epilogue
            epilogue_name = func_attrs["epilogue"]
            op.epilogue_functor = cutlass_lib.library.EpilogueFunctorName[epilogue_name]
            op.element_epilogue = acc_type

            # apply special config if required
            if f_apply_special_config is not None:
                op = f_apply_special_config(func_attrs, op)

            # set C alignment depending on the dtype
            for i in alignment.get_alignments(dtype):
                op = copy.deepcopy(op)
                op.C.alignment = i
                ret.append(op)

        return ret

    if op_kind is None:
        op_kind = cutlass_lib.library.OperationKind.Conv2d
    extract_ops = list(Target.current()._operators[op_kind].items())
    conv_kind = cutlass_lib.library.ConvKind.Fprop

    conv_ops = OrderedDict()
    for _, value in extract_ops:
        op = value[0]
        if op.conv_kind == conv_kind:
            ret = f_proc_op(op)
            if len(ret) > 0:
                for op_inst in ret:
                    key = kernel_name(op_inst, layout=op_layout)
                    conv_ops[key] = op_inst
    return conv_ops
Extracts cutlass config for conv kernels.
python/aitemplate/backend/cuda/conv2d/common.py