onnxscript 0.1.0__py3-none-any.whl

onnxscript/__init__.py

@@ -0,0 +1,131 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+__all__ = [
+    "script",
+    "graph",
+    "ir",
+    "optimizer",
+    "rewriter",
+    "export_onnx_lib",
+    "OnnxFunction",
+    "TracedOnnxFunction",
+    "proto2python",
+    "external_tensor",
+    "BFLOAT16",
+    "FLOAT16",
+    "FLOAT8E4M3FN",
+    "FLOAT8E4M3FNUZ",
+    "FLOAT8E5M2",
+    "FLOAT8E5M2FNUZ",
+    "FLOAT",
+    "DOUBLE",
+    "INT8",
+    "INT16",
+    "INT32",
+    "INT64",
+    "UINT8",
+    "UINT16",
+    "UINT32",
+    "UINT64",
+    "BOOL",
+    "STRING",
+    "COMPLEX64",
+    "COMPLEX128",
+    "opset1",
+    "opset2",
+    "opset3",
+    "opset4",
+    "opset5",
+    "opset6",
+    "opset7",
+    "opset8",
+    "opset9",
+    "opset10",
+    "opset11",
+    "opset12",
+    "opset13",
+    "opset14",
+    "opset15",
+    "opset16",
+    "opset17",
+    "opset18",
+    "opset19",
+    "opset20",
+    "opset_ai_onnx_ml1",
+    "opset_ai_onnx_ml2",
+    "opset_ai_onnx_ml3",
+    "opset_ai_onnx_ml4",
+    "DEBUG",
+]
+
+import importlib.metadata
+
+from .backend.onnx_export import export2python as proto2python
+from .main import export_onnx_lib, graph, script
+
+# isort: off
+from .onnx_opset import (
+    opset1,
+    opset2,
+    opset3,
+    opset4,
+    opset5,
+    opset6,
+    opset7,
+    opset8,
+    opset9,
+    opset10,
+    opset11,
+    opset12,
+    opset13,
+    opset14,
+    opset15,
+    opset16,
+    opset17,
+    opset18,
+    opset19,
+    opset20,
+    opset_ai_onnx_ml1,
+    opset_ai_onnx_ml2,
+    opset_ai_onnx_ml3,
+    opset_ai_onnx_ml4,
+)
+
+from .onnx_types import (
+    BFLOAT16,
+    FLOAT16,
+    FLOAT8E4M3FN,
+    FLOAT8E4M3FNUZ,
+    FLOAT8E5M2,
+    FLOAT8E5M2FNUZ,
+    FLOAT,
+    DOUBLE,
+    INT8,
+    INT16,
+    INT32,
+    INT64,
+    UINT8,
+    UINT16,
+    UINT32,
+    UINT64,
+    BOOL,
+    STRING,
+    COMPLEX64,
+    COMPLEX128,
+)
+
+# isort: on
+
+from . import ir, optimizer, rewriter
+from ._internal.utils import external_tensor
+from .values import OnnxFunction, TracedOnnxFunction
+
+# Set DEBUG to True to enable additional debug checks
+DEBUG: bool = False
+
+try:  # noqa: SIM105
+    __version__ = importlib.metadata.version("onnxscript")
+except importlib.metadata.PackageNotFoundError:
+    # package is not installed
+    pass

onnxscript/_framework_apis/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Semi-private stable APIs for framework-specific usage only."""

onnxscript/_framework_apis/torch_2_5.py

@@ -0,0 +1,117 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Stable APIs for PyTorch 2.5."""
+
+from __future__ import annotations
+
+__all__ = [
+    "check_model",
+    "convert_version",
+    "get_torchlib_ops",
+    "optimize",
+    "save_model_with_external_data",
+]
+
+import dataclasses
+import os
+import pathlib
+from typing import Callable
+
+from onnxscript import ir, optimizer, version_converter
+from onnxscript.function_libs.torch_lib import registration
+
+
+@dataclasses.dataclass(frozen=True)
+class _OnnxFunctionMeta:
+    """A wrapper of onnx-script function with additional metadata.
+
+    qualified_name: The qualified name of the aten operator.
+    function: The onnx-script function.
+    domain: The domain of the function.
+    name: The name of the function.
+    is_complex: Whether the function is a complex function.
+    """
+
+    qualified_name: str
+    function: Callable
+    domain: str
+    name: str
+    is_complex: bool = False
+
+
+def optimize(model: ir.Model) -> ir.Model:
+    """Optimize the model."""
+    # Internal flag. Will go away.
+    enabled = os.getenv("TORCH_ONNX_ENABLE_OPTIMIZATION") == "1"
+    if enabled:
+        optimizer.optimize_ir(model)
+    return model
+
+
+def convert_version(model: ir.Model, target_version: int) -> ir.Model:
+    """Convert the model to the specified ONNX opset version."""
+    # Internal flag. Will go away.
+    enabled = os.getenv("TORCH_ONNX_ENABLE_VERSION_CONVERSION") == "1"
+    if enabled:
+        version_converter.convert_version(model, target_version)
+    return model
+
+
+def check_model(model: ir.Model) -> None:
+    """Check the model."""
+
+    del model  # Unused yet
+
+
+def save_model_with_external_data(model: ir.Model, model_path: str | os.PathLike) -> None:
+    """Save the model with external data. The model is unchanged after saving."""
+
+    # TODO(#1835): Decide if we want to externalize large attributes as well
+    for value in model.graph.initializers.values():
+        if value.const_value is None:
+            raise ValueError(
+                "The model contains uninitialized initializer values. "
+                "Please make sure all initializer values are initialized."
+            )
+    destination_path = pathlib.Path(model_path)
+    data_path = f"{destination_path.name}.data"
+
+    ir.save(model, model_path, external_data=data_path)
+
+
+def get_torchlib_ops() -> list[_OnnxFunctionMeta]:
+    # Trigger op registration
+    from onnxscript.function_libs.torch_lib import (  # pylint: disable=import-outside-toplevel
+        ops,
+    )
+
+    del ops  # Unused
+
+    torchlib_registry = registration.default_registry
+    function_metas = []
+
+    for qualified_name, aten_overloads_func in torchlib_registry.items():
+        if qualified_name.startswith("internal::"):
+            # Skip the custom defined internal functions
+            continue
+
+        for overload_func in aten_overloads_func.overloads:
+            function_meta = _OnnxFunctionMeta(
+                qualified_name=qualified_name,
+                function=overload_func,
+                domain=overload_func.function_ir.domain,
+                name=overload_func.name,
+                is_complex=False,
+            )
+            function_metas.append(function_meta)
+        for complex_func in aten_overloads_func.complex:
+            function_meta = _OnnxFunctionMeta(
+                qualified_name=qualified_name,
+                function=complex_func,
+                domain=complex_func.function_ir.domain,
+                name=complex_func.name,
+                is_complex=True,
+            )
+            function_metas.append(function_meta)
+
+    return function_metas

onnxscript/_framework_apis/torch_2_6.py

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Stable APIs for PyTorch 2.6."""
+
+from __future__ import annotations
+
+__all__ = [
+    "check_model",
+    "convert_version",
+    "get_torchlib_ops",
+    "optimize",
+    "save_model_with_external_data",
+    "torchlib_opset",
+]
+from typing import TYPE_CHECKING
+
+from onnxscript import ir, optimizer
+from onnxscript._framework_apis.torch_2_5 import (
+    check_model,
+    convert_version,
+    get_torchlib_ops,
+    save_model_with_external_data,
+)
+
+if TYPE_CHECKING:
+    from onnxscript.onnx_opset._impl.opset18 import Opset18
+
+
+def optimize(model: ir.Model) -> ir.Model:
+    """Optimize the model."""
+    optimizer.optimize_ir(model)
+    return model
+
+
+def torchlib_opset() -> Opset18:
+    """Return the default opset for torchlib."""
+    import onnxscript  # pylint: disable=import-outside-toplevel
+
+    return onnxscript.opset18  # type: ignore
+
+
+def torchlib_opset_version() -> int:
+    """Return the default opset version for torchlib."""
+
+    return torchlib_opset().version

onnxscript/_internal/analysis.py

@@ -0,0 +1,229 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import ast
+from typing import Any, Optional, Sequence, Set
+
+from onnxscript import sourceinfo
+from onnxscript._internal import ast_utils
+
+
+def _get_loop_var(for_stmt: ast.For, formatter: sourceinfo.Formatter) -> str:
+    if not isinstance(for_stmt.target, ast.Name):
+        raise TypeError(formatter(for_stmt, "For loop target must be a single variable."))
+    return for_stmt.target.id
+
+
+def _used_vars(expr: Optional[ast.expr]) -> Set[str]:
+    """Return set of all variables used, including function names, in an expression."""
+    if expr is None:
+        return set()
+    if isinstance(expr, ast.Name):
+        return {expr.id}
+    result = set()
+    if isinstance(expr, ast.Call):
+        # The callee-expression is not visited
+        children = expr.args
+        for keyword in expr.keywords:
+            if isinstance(keyword.value, ast.Name):
+                result.add(keyword.value.id)
+    else:
+        children = ast.iter_child_nodes(expr)  # type: ignore[assignment]
+    for c in children:
+        result = result | _used_vars(c)
+    return result
+
+
+def _lhs_vars(lhs: ast.expr) -> Set[str]:
+    """Return set of assigned variables in the lhs of an assignment statement."""
+
+    def get_id(e):
+        assert isinstance(e, ast.Name), "Only simple assignments supported."
+        return e.id
+
+    if isinstance(lhs, ast.Tuple):
+        return {get_id(x) for x in lhs.elts}
+    return {get_id(lhs)}
+
+
+def assigned_vars(
+    stmt: ast.stmt | list[ast.stmt], formatter: sourceinfo.Formatter
+) -> Set[str]:
+    """Return the set of all variables that may be assigned to in an execution of input stmt
+    or sequence of statements.
+    """
+
+    def assigned_in_block(block: Sequence[ast.stmt]) -> Set[str]:
+        result: set[Any] = set()
+        for s in block:
+            result = result | assigned_vars(s, formatter)
+        return result
+
+    if isinstance(stmt, ast.Assign):
+        return _lhs_vars(stmt.targets[0])
+    if isinstance(stmt, ast.AnnAssign):
+        return _lhs_vars(stmt.target)
+    if isinstance(stmt, ast.Return):
+        return set()
+    if isinstance(stmt, ast.If):
+        return assigned_in_block(stmt.body) | assigned_in_block(stmt.orelse)
+    if isinstance(stmt, ast.For):
+        return assigned_in_block(stmt.body) | {_get_loop_var(stmt, formatter)}
+    if isinstance(stmt, ast.While):
+        return assigned_in_block(stmt.body)
+    if isinstance(stmt, list):
+        return assigned_in_block(stmt)
+    if isinstance(stmt, ast.Break):
+        return set()
+    if ast_utils.is_print_call(stmt):
+        return set()
+    if ast_utils.is_doc_string(stmt):
+        return set()
+    error_message = formatter(stmt, f"Unsupported statement type {type(stmt)!r}.")
+    raise ValueError(error_message)
+
+
+def do_liveness_analysis(fun: ast.FunctionDef, formatter: sourceinfo.Formatter):
+    """Perform liveness analysis of the given function-ast. The results of the
+    analysis are stored directly with each statement-ast `s` as attributes `s.live_in`
+    and `s.live_out`.
+    """
+
+    def visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
+        stmt.live_out = live_out  # type: ignore[attr-defined]
+        live = do_visit(stmt, live_out)
+        stmt.live_in = live  # type: ignore[attr-defined]
+        return live
+
+    def do_visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
+        def visitBlock(block: Sequence[ast.stmt], live_out: Set[str]) -> Set[str]:
+            for s in reversed(block):
+                live_out = visit(s, live_out)
+            return live_out
+
+        if isinstance(stmt, ast.Assign):
+            return live_out.difference(_lhs_vars(stmt.targets[0])) | _used_vars(stmt.value)
+        if isinstance(stmt, ast.AnnAssign):
+            return live_out.difference(_lhs_vars(stmt.target)) | _used_vars(stmt.value)
+        if isinstance(stmt, ast.Return):
+            return _used_vars(stmt.value)
+        if isinstance(stmt, ast.If):
+            live1 = visitBlock(stmt.body, live_out)
+            live2 = visitBlock(stmt.orelse, live_out)
+            return live1 | live2 | _used_vars(stmt.test)
+        if isinstance(stmt, ast.For):
+            p_loop_var = _get_loop_var(stmt, formatter)
+            prev = None
+            curr = live_out
+            while curr != prev:
+                prev = curr
+                curr = visitBlock(stmt.body, prev).difference({p_loop_var})
+            return curr
+        if isinstance(stmt, ast.While):
+            cond_vars = _used_vars(stmt.test)
+            prev = None
+            curr = live_out | cond_vars
+            while curr != prev:
+                prev = curr
+                curr = visitBlock(stmt.body, prev) | cond_vars
+            return curr
+        if isinstance(stmt, ast.Break):
+            # The following is sufficient for the current restricted usage, where
+            # a (conditional) break is allowed only as the last statement of a loop.
+            # Break statements in the middle of the loop, however, will require
+            # a generalization.
+            return live_out
+        if ast_utils.is_doc_string(stmt):
+            return live_out
+        if isinstance(stmt, ast.FunctionDef):
+            return live_out
+        if ast_utils.is_print_call(stmt):
+            return live_out
+        raise ValueError(formatter(stmt, f"Unsupported statement type {type(stmt)!r}."))
+
+    assert isinstance(fun, ast.FunctionDef)
+    live: set[Any] = set()
+    for s in reversed(fun.body):
+        live = visit(s, live)
+
+
+def exposed_uses(stmts: Sequence[ast.stmt], formatter: sourceinfo.Formatter):
+    """Return the set of variables that are used before being defined by given block.
+    In essence, this identifies the "inputs" to a given code-block.
+    For example, consider the following code-block:
+    ::
+
+       x = x + 10
+       y = 20
+       z = x + y
+       x = 30
+
+    The exposed_uses of this code-block is { x }. The value of z is not used within
+    the block. Even though the value of y is used within the block, it is assigned
+    a value before it is used. However, in contrast, the incoming value of x is used
+    (in the first statement). Hence x is included in the exposed_uses.
+    """
+
+    def visitBlock(block: Sequence[ast.stmt], live_out: Set[str]) -> Set[str]:
+        for stmt in reversed(block):
+            live_out = visit(stmt, live_out)
+        return live_out
+
+    def visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
+        if isinstance(stmt, ast.Assign):
+            return live_out.difference(_lhs_vars(stmt.targets[0])) | _used_vars(stmt.value)
+        if isinstance(stmt, ast.AnnAssign):
+            return live_out.difference(_lhs_vars(stmt.target)) | _used_vars(stmt.value)
+        if isinstance(stmt, ast.Return):
+            return _used_vars(stmt.value)
+        if isinstance(stmt, ast.If):
+            live1 = visitBlock(stmt.body, live_out)
+            live2 = visitBlock(stmt.orelse, live_out)
+            return (live1 | live2) | _used_vars(stmt.test)
+        if ast_utils.is_print_call(stmt):
+            return live_out
+        if ast_utils.is_doc_string(stmt):
+            return live_out
+        if isinstance(stmt, ast.For):
+            # Analysis assumes loop may execute zero times. Results can be improved
+            # for loops that execute at least once.
+            loop_var_set = {_get_loop_var(stmt, formatter)}
+            used_after_loop = live_out.difference(loop_var_set)
+            used_inside_loop = visitBlock(stmt.body, set()).difference(loop_var_set)
+            used_in_loop_header = _used_vars(stmt.iter)
+            return used_inside_loop | used_in_loop_header | used_after_loop
+        if isinstance(stmt, ast.While):
+            # Analysis assumes loop may execute zero times. Results can be improved
+            # for loops that execute at least once.
+            used_inside_loop = visitBlock(stmt.body, set())
+            used_in_loop_header = _used_vars(stmt.test)
+            return used_inside_loop | used_in_loop_header | live_out
+        if isinstance(stmt, ast.Break):
+            # Currently, we assume that break statements are only allowed as the last
+            # statement in a loop, as "if cond: break".
+            return live_out
+        if isinstance(stmt, ast.FunctionDef):
+            if stmt.name in live_out:
+                live_out.remove(stmt.name)
+                live_out = live_out | outer_scope_variables(stmt, formatter)
+            return live_out
+        raise ValueError(formatter(stmt, f"Unsupported statement type {type(stmt)!r}."))
+
+    return visitBlock(stmts, set())
+
+
+def outer_scope_variables(fun: ast.FunctionDef, formatter: sourceinfo.Formatter):
+    """Return the set of outer-scope variables used in a nested function.
+
+    Args:
+        fun: The function-ast to analyze.
+        formatter: The formatter object.
+
+    Returns:
+        A set of variable names (strings).
+    """
+    assert isinstance(fun, ast.FunctionDef)
+    used_vars_ = exposed_uses(fun.body, formatter)
+    inputs = [x.arg for x in fun.args.args]
+    return used_vars_.difference(inputs)

onnxscript/_internal/ast_utils.py

@@ -0,0 +1,64 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Utilities for working with Python ASTs."""
+
+from __future__ import annotations
+
+import ast
+import inspect
+import sys
+import textwrap
+from typing import Callable
+
+PY_VERSION_GE_39 = sys.version_info >= (3, 9)
+
+
+def get_src_and_ast(func: Callable, /) -> tuple[str, ast.FunctionDef]:
+    try:
+        src = inspect.getsource(func)
+    except OSError as e:
+        raise RuntimeError(
+            f"Decorator script does not work on dynamically compiled function {func.__name__}."
+        ) from e
+    src = textwrap.dedent(src)
+    top_level_ast = ast.parse(src)
+    assert isinstance(top_level_ast, ast.Module)
+    assert len(top_level_ast.body) == 1
+    f_ast = top_level_ast.body[0]
+    assert isinstance(f_ast, ast.FunctionDef)
+    return src, f_ast
+
+
+def normalize_subscript_expr(expr: ast.Subscript):
+    # Normalizes the representation of a subscripted expression, handling python version
+    # differences as well as variations between A[x] (single-index) and A[x, y] (multiple indices)
+    # Returns a list of expressions, denoting the indices, after stripping the extraneous "Index"
+    # wrapper present in python versions before 3.9
+    index_expr = expr.slice
+    if PY_VERSION_GE_39:
+        if isinstance(index_expr, ast.Tuple):
+            return index_expr.elts  # multiple indices
+        else:
+            return [index_expr]  # single index
+    else:
+        if isinstance(index_expr, ast.ExtSlice):
+            indices = index_expr.dims  # type: ignore[attr-defined]
+        else:
+            indices = [index_expr]  # single slice-index
+        return [x.value if isinstance(x, ast.Index) else x for x in indices]  # type: ignore[attr-defined]
+
+
+def is_print_call(stmt: ast.stmt) -> bool:
+    """Return True if the statement is a call to the print function."""
+    if isinstance(stmt, ast.Expr):
+        if isinstance(stmt.value, ast.Call):
+            if isinstance(stmt.value.func, ast.Name):
+                return stmt.value.func.id == "print"
+    return False
+
+
+def is_doc_string(stmt: ast.stmt) -> bool:
+    """Return True if the statement is a docstring."""
+    if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant):
+        return isinstance(stmt.value.value, str)
+    return False