Function parse_paddle_kernels

tools/parse_kernel_info.py:91–135

(lib="phi", kernel_type="function", print_detail=False)

def parse_paddle_kernels(lib="phi", kernel_type="function", print_detail=False):
    assert lib in ["fluid", "phi"]

    if lib == "phi":
        assert kernel_type in ["function", "structure", "all"]
        # phi kernel type can be: function, structure, all
        kernel_infos = paddle.base.core._get_registered_phi_kernels(kernel_type)
    else:
        # fluid, phi, all
        assert kernel_type in ["fluid", "phi", "all"]
        kernel_infos = paddle.base.core._get_all_register_op_kernels(
            kernel_type
        )

    max_op_type_lengths = 0
    stats = KernelRegistryStatistics()

    kernel_info_dict = {}
    for key, value in kernel_infos.items():
        info = KernelInfo(key)
        if lib == "phi":
            info.parse_phi_dtypes(value, device="GPU")
        else:
            info.parse_fluid_dtypes(value, device="gpu")
        kernel_info_dict[key] = info
        if len(info.op_type) > max_op_type_lengths:
            max_op_type_lengths = len(info.op_type)
        stats.update(info.supported_dtypes)

    if print_detail:
        print(
            f"==================== lib={lib}, kernel_type={kernel_type} ===================="
        )
        print(
            "{} : {}".format(
                "op_type".ljust(max_op_type_lengths + 4),
                "supported_dtypes for GPU",
            )
        )
        for key, value in sorted(kernel_info_dict.items()):
            print(
                f"{value.op_type.ljust(max_op_type_lengths + 4)} : {value.supported_dtypes}"
            )
        print()
    return stats


def main(lib):