Пример #1
0
def _load_ops_as_dict():
    # Returns a list of dicts, each with a name that is a tuple of the form:
    #   (kernel_signature, variant)
    # The kernel signature is a reified form of the argument type signature
    # used throughout PyTorch:
    #   namespace::kernel_name(type1,type2)
    def reify_signature(reg_op):
        kernel_name, unused_variant = reg_op["name"]
        arg_types = [arg["type"] for arg in reg_op["arguments"]]
        return f"{kernel_name}({','.join(arg_types)})"

    reg_ops_list = get_registered_ops()
    return {reify_signature(reg_op): reg_op for reg_op in reg_ops_list}
Пример #2
0
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See frontends/pytorch/LICENSE for license information.
# RUN: %PYTHON %s | FileCheck %s

import _torch_mlir

# This check is just for a built-in op that is unlikely to change (and is
# otherwise insignificant).
# CHECK: {'name': ('aten::mul', 'Tensor'), 'is_c10_op': True, 'is_vararg': False, 'is_varret': False, 'is_mutable': False, 'arguments': [{'name': 'self', 'type': 'Tensor', 'pytype': 'Tensor'}, {'name': 'other', 'type': 'Tensor', 'pytype': 'Tensor'}], 'returns': [{'name': '', 'type': 'Tensor', 'pytype': 'Tensor'}]}
print('\n\n'.join([repr(r) for r in _torch_mlir.get_registered_ops()]))
Пример #3
0
def load_registry() -> Registry:
    return Registry([JitOperator(op_info) for op_info in get_registered_ops()])