universal_onnxruntime / onnxruntime /tools /pytorch_export_contrib_ops.py

Upload 379 files

610d9c2 verified about 3 hours ago

5.21 kB

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

"""
Support for registering ONNX Runtime's built-in contrib ops with
PyTorch-ONNX exporter (torch.onnx.export).
"""

import contextlib
import typing

try:
    # TODO(justinchuby): Create a function to alert users when torch is not installed
    import torch
except ModuleNotFoundError:
    raise ModuleNotFoundError(  # noqa: B904
        "This module is only useful in combination with PyTorch. To install PyTorch see https://pytorch.org/."
    )

from torch.onnx import symbolic_helper

_OPSET_VERSION = 1
_registered_ops: typing.AbstractSet[str] = set()


def _reg(symbolic_fn: typing.Callable, namespace: str = "aten"):
    name = f"{namespace}::{symbolic_fn.__name__}"
    torch.onnx.register_custom_op_symbolic(name, symbolic_fn, _OPSET_VERSION)
    _registered_ops.add(name)


def register():
    """Register ONNX Runtime's built-in contrib ops.
    Should be run before torch.onnx.export().
    """

    def grid_sampler(g, input, grid, mode, padding_mode, align_corners):
        # mode
        #   'bilinear'      : onnx::Constant[value={0}]
        #   'nearest'       : onnx::Constant[value={1}]
        #   'bicubic'       : onnx::Constant[value={2}]
        # padding_mode
        #   'zeros'         : onnx::Constant[value={0}]
        #   'border'        : onnx::Constant[value={1}]
        #   'reflection'    : onnx::Constant[value={2}]
        mode = symbolic_helper._maybe_get_const(mode, "i")
        padding_mode = symbolic_helper._maybe_get_const(padding_mode, "i")
        mode_str = ["bilinear", "nearest", "bicubic"][mode]
        padding_mode_str = ["zeros", "border", "reflection"][padding_mode]
        align_corners = int(symbolic_helper._maybe_get_const(align_corners, "b"))

        return g.op(
            "com.microsoft::GridSample",
            input,
            grid,
            mode_s=mode_str,
            padding_mode_s=padding_mode_str,
            align_corners_i=align_corners,
        )

    _reg(grid_sampler)

    def inverse(g, self):
        return g.op("com.microsoft::Inverse", self).setType(self.type())

    _reg(inverse)
    torch.onnx.register_custom_op_symbolic("aten::linalg_inv", inverse, _OPSET_VERSION)
    _registered_ops.add("aten::linalg_inv")

    def gelu(g, self: torch._C.Value, approximate="none"):
        # PyTorch can emit aten::gelu with or without the optional approximate arg.
        if not isinstance(approximate, str):
            approximate = symbolic_helper._maybe_get_const(approximate, "s")

        # Use microsoft::Gelu for performance if possible. It only supports approximate == "none".
        if approximate == "none":
            return g.op("com.microsoft::Gelu", self).setType(self.type())
        return torch.onnx.symbolic_opset9.gelu(g, self, approximate)

    _reg(gelu)
    # Some PyTorch versions dispatch GELU symbolic lookup by exporter opset.
    # Registering across stable opsets keeps ORT Gelu fusion consistently enabled.
    for opset in range(9, 21):
        torch.onnx.register_custom_op_symbolic("aten::gelu", gelu, opset)

    def triu(g, self, diagonal):
        return g.op("com.microsoft::Trilu", self, diagonal, upper_i=1).setType(self.type())

    _reg(triu)

    def tril(g, self, diagonal):
        return g.op("com.microsoft::Trilu", self, diagonal, upper_i=0).setType(self.type())

    _reg(tril)

    @torch.onnx.symbolic_helper.parse_args("v")
    def DynamicTimeWarping(g, self):  # noqa: N802
        return g.op("com.microsoft::DynamicTimeWarping", self)

    _reg(DynamicTimeWarping, namespace="onnxruntime")

    def UnfoldTensor(g, self, dim, size, step):  # noqa: N802
        dim = int(symbolic_helper._maybe_get_const(dim, "i"))
        size = int(symbolic_helper._maybe_get_const(size, "i"))
        step = int(symbolic_helper._maybe_get_const(step, "i"))
        return g.op(
            "com.microsoft::UnfoldTensor",
            self,
            dim_i=dim,
            size_i=size,
            step_i=step,
        ).setType(self.type().with_sizes([None, None, None, None, size]))

    _reg(UnfoldTensor, namespace="onnxruntime")


def unregister():
    """Unregister ONNX Runtime's built-in contrib ops."""
    for name in _registered_ops:
        try:
            torch.onnx.unregister_custom_op_symbolic(name, _OPSET_VERSION)
        except AttributeError:
            # The symbolic_registry module was removed in PyTorch 1.13.
            # We are importing it here for backwards compatibility
            # because unregister_custom_op_symbolic is not available before PyTorch 1.12
            from torch.onnx import symbolic_registry  # noqa: PLC0415

            namespace, kind = name.split("::")
            for version in symbolic_helper._onnx_stable_opsets:
                if version >= _OPSET_VERSION and symbolic_registry.is_registered_op(kind, namespace, version):
                    del symbolic_registry._registry[(namespace, version)][kind]

    # Also clean up gelu's multi-opset registrations (see register()).
    for opset in range(9, 21):
        with contextlib.suppress(Exception):
            torch.onnx.unregister_custom_op_symbolic("aten::gelu", opset)

	# Copyright (c) Microsoft Corporation. All rights reserved.
	# Licensed under the MIT License.

	"""
	Support for registering ONNX Runtime's built-in contrib ops with
	PyTorch-ONNX exporter (torch.onnx.export).
	"""

	import contextlib
	import typing

	try:
	# TODO(justinchuby): Create a function to alert users when torch is not installed
	import torch
	except ModuleNotFoundError:
	raise ModuleNotFoundError( # noqa: B904
	"This module is only useful in combination with PyTorch. To install PyTorch see https://pytorch.org/."
	)

	from torch.onnx import symbolic_helper

	_OPSET_VERSION = 1
	_registered_ops: typing.AbstractSet[str] = set()


	def _reg(symbolic_fn: typing.Callable, namespace: str = "aten"):
	name = f"{namespace}::{symbolic_fn.__name__}"
	torch.onnx.register_custom_op_symbolic(name, symbolic_fn, _OPSET_VERSION)
	_registered_ops.add(name)


	def register():
	"""Register ONNX Runtime's built-in contrib ops.

	Should be run before torch.onnx.export().
	"""

	def grid_sampler(g, input, grid, mode, padding_mode, align_corners):
	# mode
	# 'bilinear' : onnx::Constant[value={0}]
	# 'nearest' : onnx::Constant[value={1}]
	# 'bicubic' : onnx::Constant[value={2}]
	# padding_mode
	# 'zeros' : onnx::Constant[value={0}]
	# 'border' : onnx::Constant[value={1}]
	# 'reflection' : onnx::Constant[value={2}]
	mode = symbolic_helper._maybe_get_const(mode, "i")
	padding_mode = symbolic_helper._maybe_get_const(padding_mode, "i")
	mode_str = ["bilinear", "nearest", "bicubic"][mode]
	padding_mode_str = ["zeros", "border", "reflection"][padding_mode]
	align_corners = int(symbolic_helper._maybe_get_const(align_corners, "b"))

	return g.op(
	"com.microsoft::GridSample",
	input,
	grid,
	mode_s=mode_str,
	padding_mode_s=padding_mode_str,
	align_corners_i=align_corners,
	)

	_reg(grid_sampler)

	def inverse(g, self):
	return g.op("com.microsoft::Inverse", self).setType(self.type())

	_reg(inverse)
	torch.onnx.register_custom_op_symbolic("aten::linalg_inv", inverse, _OPSET_VERSION)
	_registered_ops.add("aten::linalg_inv")

	def gelu(g, self: torch._C.Value, approximate="none"):
	# PyTorch can emit aten::gelu with or without the optional approximate arg.
	if not isinstance(approximate, str):
	approximate = symbolic_helper._maybe_get_const(approximate, "s")

	# Use microsoft::Gelu for performance if possible. It only supports approximate == "none".
	if approximate == "none":
	return g.op("com.microsoft::Gelu", self).setType(self.type())
	return torch.onnx.symbolic_opset9.gelu(g, self, approximate)

	_reg(gelu)
	# Some PyTorch versions dispatch GELU symbolic lookup by exporter opset.
	# Registering across stable opsets keeps ORT Gelu fusion consistently enabled.
	for opset in range(9, 21):
	torch.onnx.register_custom_op_symbolic("aten::gelu", gelu, opset)

	def triu(g, self, diagonal):
	return g.op("com.microsoft::Trilu", self, diagonal, upper_i=1).setType(self.type())

	_reg(triu)

	def tril(g, self, diagonal):
	return g.op("com.microsoft::Trilu", self, diagonal, upper_i=0).setType(self.type())

	_reg(tril)

	@torch.onnx.symbolic_helper.parse_args("v")
	def DynamicTimeWarping(g, self): # noqa: N802
	return g.op("com.microsoft::DynamicTimeWarping", self)

	_reg(DynamicTimeWarping, namespace="onnxruntime")

	def UnfoldTensor(g, self, dim, size, step): # noqa: N802
	dim = int(symbolic_helper._maybe_get_const(dim, "i"))
	size = int(symbolic_helper._maybe_get_const(size, "i"))
	step = int(symbolic_helper._maybe_get_const(step, "i"))
	return g.op(
	"com.microsoft::UnfoldTensor",
	self,
	dim_i=dim,
	size_i=size,
	step_i=step,
	).setType(self.type().with_sizes([None, None, None, None, size]))

	_reg(UnfoldTensor, namespace="onnxruntime")


	def unregister():
	"""Unregister ONNX Runtime's built-in contrib ops."""
	for name in _registered_ops:
	try:
	torch.onnx.unregister_custom_op_symbolic(name, _OPSET_VERSION)
	except AttributeError:
	# The symbolic_registry module was removed in PyTorch 1.13.
	# We are importing it here for backwards compatibility
	# because unregister_custom_op_symbolic is not available before PyTorch 1.12
	from torch.onnx import symbolic_registry # noqa: PLC0415

	namespace, kind = name.split("::")
	for version in symbolic_helper._onnx_stable_opsets:
	if version >= _OPSET_VERSION and symbolic_registry.is_registered_op(kind, namespace, version):
	del symbolic_registry._registry[(namespace, version)][kind]

	# Also clean up gelu's multi-opset registrations (see register()).
	for opset in range(9, 21):
	with contextlib.suppress(Exception):
	torch.onnx.unregister_custom_op_symbolic("aten::gelu", opset)