You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/09/21 15:08:01 UTC
(tvm) branch tvmscript-generic-parser-builder updated: Promote canonical TVMScript parser and builder interfaces
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tvmscript-generic-parser-builder
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/tvmscript-generic-parser-builder by this push:
new 37b1cc6ab4 Promote canonical TVMScript parser and builder interfaces
37b1cc6ab4 is described below
commit 37b1cc6ab4c14cf9cfb1db3a04e5c3ffcee0d099
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Mon Sep 21 15:03:16 2026 +0000
Promote canonical TVMScript parser and builder interfaces
---
docs/reference/api/python/script/parser.rst | 24 +-
docs/tirx/api/script.rst | 10 +-
python/tvm/relax/script/__init__.py | 49 +-
python/tvm/relax/script/builder/__init__.py | 452 +++++++-
python/tvm/relax/script/builder_v2/__init__.py | 465 --------
python/tvm/relax/script/parser/__init__.py | 52 -
python/tvm/relax/script/parser/dist.py | 106 --
python/tvm/relax/script/parser/entry.py | 521 ---------
python/tvm/relax/script/parser/parser.py | 506 --------
python/tvm/relax/script/v2.py | 58 -
python/tvm/script/__init__.py | 62 +-
python/tvm/script/parser/__init__.py | 39 +-
python/tvm/script/parser/_core.py | 24 -
.../script/{parser_v2 => parser}/annotations.py | 0
python/tvm/script/parser/core/__init__.py | 20 -
python/tvm/script/parser/core/diagnostics.py | 324 ------
python/tvm/script/parser/core/dispatch.py | 157 ---
python/tvm/script/parser/core/doc.py | 320 -----
python/tvm/script/parser/core/doc_core.py | 1123 ------------------
python/tvm/script/parser/core/entry.py | 223 ----
python/tvm/script/parser/core/error.py | 27 -
python/tvm/script/parser/core/evaluator.py | 626 ----------
python/tvm/script/parser/core/parser.py | 1071 -----------------
python/tvm/script/parser/core/utils.py | 212 ----
.../script/{parser_v2 => parser}/diagnostics.py | 0
.../tvm/script/{parser_v2 => parser}/frontend.py | 0
.../tvm/script/{parser_v2 => parser}/functions.py | 0
python/tvm/script/{parser_v2 => parser}/ir.py | 0
python/tvm/script/parser/ir/__init__.py | 39 -
python/tvm/script/parser/ir/entry.py | 163 ---
python/tvm/script/parser/ir/parser.py | 212 ----
.../tvm/script/{parser_v2 => parser}/transform.py | 0
python/tvm/script/parser_v2/__init__.py | 34 -
python/tvm/tirx/script/__init__.py | 50 +-
python/tvm/tirx/script/builder/__init__.py | 571 ++++++++-
python/tvm/tirx/script/builder/ir.py | 3 -
python/tvm/tirx/script/builder_v2/__init__.py | 579 ----------
python/tvm/tirx/script/parser/__init__.py | 50 -
python/tvm/tirx/script/parser/entry.py | 581 ----------
python/tvm/tirx/script/parser/operation.py | 167 ---
python/tvm/tirx/script/parser/parser.py | 1217 --------------------
python/tvm/tirx/script/v2.py | 47 -
tests/python/relax/test_analysis_type_analysis.py | 2 -
tests/python/tirx/test_op_namespace_cleanup.py | 5 +-
.../{test_parser_v2.py => test_parser.py} | 8 +-
.../tvmscript/test_tvmscript_error_report.py | 47 -
.../tvmscript/test_tvmscript_parser_evaluator.py | 66 --
.../tvmscript/test_tvmscript_parser_source.py | 116 +-
48 files changed, 1202 insertions(+), 9226 deletions(-)
diff --git a/docs/reference/api/python/script/parser.rst b/docs/reference/api/python/script/parser.rst
index 583b27b276..e9bea20ae6 100644
--- a/docs/reference/api/python/script/parser.rst
+++ b/docs/reference/api/python/script/parser.rst
@@ -20,6 +20,11 @@ tvm.script.parser
tvm.script.parser
*****************
+The shared parser translates Python source into calls on registered construction
+namespaces. Dialect script packages provide their decorators and register the
+corresponding builder operations. Construction returns concrete IR values, with
+support for eager and postponed function annotations.
+
.. automodule:: tvm.script.parser
:members:
:imported-members:
@@ -28,10 +33,17 @@ tvm.script.parser.ir
********************
.. automodule:: tvm.script.parser.ir
-tvm.script.parser.relax
-***********************
-.. automodule:: tvm.script.parser.relax
+tvm.relax.script
+****************
+.. automodule:: tvm.relax.script
+ :members: function, macro
+ :imported-members:
+
+tvm.tirx.script
+***************
+.. automodule:: tvm.tirx.script
+ :members: prim_func, inline, macro
+ :imported-members:
-tvm.script.parser.tirx
-**********************
-.. automodule:: tvm.script.parser.tirx
+The public aliases ``tvm.script.parser.relax`` and ``tvm.script.parser.tirx``
+continue to resolve to these dialect script packages.
diff --git a/docs/tirx/api/script.rst b/docs/tirx/api/script.rst
index f30d508095..9ce118f7bf 100644
--- a/docs/tirx/api/script.rst
+++ b/docs/tirx/api/script.rst
@@ -30,10 +30,14 @@ authoring layers and TIRx IR, see :ref:`tirx-programming-model`.
Parser entry points
-------------------
-.. automodule:: tvm.tirx.script.parser.entry
- :members:
+.. automodule:: tvm.tirx.script
+ :members: prim_func, inline, macro
+ :imported-members:
:no-index:
- :exclude-members: TIRInline
+
+The entry points use the shared ``tvm.script.parser`` and the concrete
+construction operations in ``tvm.tirx.script.builder``. Function signatures
+support both eager and postponed Python annotations.
Core IR builder
---------------
diff --git a/python/tvm/relax/script/__init__.py b/python/tvm/relax/script/__init__.py
index b4603ea678..97cbd7277e 100644
--- a/python/tvm/relax/script/__init__.py
+++ b/python/tvm/relax/script/__init__.py
@@ -14,14 +14,45 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Relax-layer TVMScript pieces (parser, builder).
+"""TVMScript entry point using concrete Relax construction operations."""
-After the per-dialect TVMScript restructure, the Relax layer owns its own
-``script/{parser,builder}`` subpackages. ``tvm.script.relax`` resolves to
-this module via the dialect registry, so the public parser surface
-(``function``, ``Tensor``, ``match_cast``, etc.) is re-exported here.
-"""
+# pylint: disable=wildcard-import,unused-wildcard-import,redefined-builtin
+import sys as _sys
-# pylint: disable=redefined-builtin,wildcard-import,unused-wildcard-import
-from .v2 import *
-from .v2 import dist
+from tvm import relax as _relax
+from tvm.relax.base_py_module import PyModuleFactory as _PyModuleFactory
+from tvm.script.parser.frontend import make_decorator as _make_decorator
+from tvm.script.parser.frontend import make_helper as _make_helper
+from tvm.script.parser.frontend import register_namespace as _register_namespace
+from tvm.script.parser.functions import register_opaque_factory as _register_opaque_factory
+
+from . import builder as _builder
+from .builder import *
+
+function = _make_decorator(_builder, option_map={"pure": "is_pure", "private": "is_private"})
+macro = _make_helper(_builder, preserve_return=True)
+
+_register_namespace("R", _sys.modules[__name__])
+_register_namespace("relax", _sys.modules[__name__])
+
+
+def _opaque_function(name, function, source, span):
+ return _relax.ExternFunc(name, span=span).with_attrs(
+ {
+ "is_pyfunc": True,
+ "function_type": "python",
+ "python_function_name": name,
+ "python_source": source,
+ "python_packed_func": function,
+ }
+ )
+
+
+def _python_module(module, original, bases):
+ candidates = (original, *bases)
+ if any(isinstance(base, type) and issubclass(base, _relax.BasePyModule) for base in candidates):
+ return _PyModuleFactory(module, original_class=original)
+ return module
+
+
+_register_opaque_factory(_opaque_function, module_adapter=_python_module)
diff --git a/python/tvm/relax/script/builder/__init__.py b/python/tvm/relax/script/builder/__init__.py
index 021bcf9b1c..234f279b5e 100644
--- a/python/tvm/relax/script/builder/__init__.py
+++ b/python/tvm/relax/script/builder/__init__.py
@@ -14,12 +14,452 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Package tvm.relax.script.builder.
+"""Concrete Relax construction operations over the shared native builder stack."""
-Holds the per-dialect ir_builder API for Relax. The legacy path
-``tvm.script.ir_builder.relax`` resolves here via the dialect registry.
-"""
+# pylint: disable=wildcard-import,redefined-builtin,invalid-name
+import builtins as _python
+import numbers as _numbers
+import sys as _sys
+from functools import partial as _partial
-# pylint: disable=wildcard-import,redefined-builtin
-from . import distributed, frame, ir
+import tvm_ffi as _ffi
+
+from tvm import ir as _ir
+from tvm import relax as _relax
+from tvm import tirx as _tir
+from tvm.relax.distributed import DeviceMesh as _DeviceMesh
+from tvm.relax.distributed import DTensorType as _DTensorType
+from tvm.relax.distributed import Placement as _Placement
+from tvm.relax.distributed import device_mesh as device_mesh
+from tvm.script.ir_builder import IRBuilder as _IRBuilder
+from tvm.script.ir_builder import ir as _I
+from tvm.script.ir_builder import protocol as _protocol
+
+from . import _ffi_api, distributed, frame, ir
+from . import distributed as dist
+from . import frame as _frame
+from . import ir as _native
+from .distributed.ir import _lookup_device_mesh
from .ir import *
+
+
+@_protocol.expression_args("shape", introduce=True, dtype="int64", scalar_strings=False)
+def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
+ """Construct a concrete tensor type from already resolved dimensions."""
+ if isinstance(shape, _python.str) and dtype is None:
+ dtype, shape = shape, None
+ if isinstance(vdevice, _python.str):
+ target, *qualifiers = vdevice.split(":", 2)
+ vdevice = _I.lookup_vdevice(target, int(qualifiers[0]) if qualifiers else 0)
+ return _relax.TensorType(shape, dtype, vdevice, ndim, span)
+
+
+@_protocol.expression_args("shape", introduce=True, dtype="int64", scalar_strings=False)
+def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *, ndim=-1, span=None):
+ """Construct a concrete distributed tensor type with resolved shape symbols."""
+ if device_mesh is None:
+ device_mesh = _DeviceMesh([], _ir.Range(0, 1))
+ elif isinstance(device_mesh, _python.str):
+ device_mesh = _lookup_device_mesh(device_mesh)
+ if isinstance(placement, _python.str):
+ placement = _Placement.from_text(placement)
+ return _DTensorType(Tensor(shape, dtype, ndim=ndim), device_mesh, placement, span)
+
+
+Range = _ir.Range
+
+
+@_protocol.expression_args("values", introduce=True, dtype="int64")
+def Shape(values=None, ndim=-1, *, span=None):
+ """Construct a concrete shape type."""
+ return _relax.ShapeType(values, ndim, span)
+
+
+def _type(value):
+ if value is None:
+ return _ir.TupleType([])
+ if callable(value):
+ value = value()
+ if _ir.is_prim_expr(value):
+ value = value.ty
+ if not isinstance(value, _ir.Type):
+ raise TypeError(f"Expected a concrete type, got {type(value).__name__}")
+ return value
+
+
+def Callable(params=None, ret=None, purity=None, derive_func=None, *, span=None):
+ """Construct a concrete function type."""
+ if purity is None:
+ purity = params is not None
+ if params is None:
+ return _relax.FuncType.opaque_func(
+ ret=None if ret is None else _type(ret),
+ derive_func=derive_func,
+ purity=purity,
+ span=span,
+ )
+ if derive_func is not None:
+ raise ValueError("A derivation function requires an opaque callable")
+ if not isinstance(params, list | _python.tuple):
+ params = [params]
+ return _relax.FuncType([_type(param) for param in params], _type(ret), purity, span)
+
+
+def Tuple(*fields, span=None):
+ """Construct a concrete tuple type."""
+ if len(fields) == 1 and isinstance(fields[0], list | _python.tuple):
+ fields = fields[0]
+ return _ir.TupleType([_type(field) for field in fields], span)
+
+
+def Prim(dtype, *, span=None):
+ """Construct a primitive type."""
+ return _ir.PrimType(dtype)
+
+
+Prim.__tvm_parameter_dtype__ = "dtype"
+
+
+def Object(*, span=None):
+ """Construct the unconstrained Relax value type."""
+ return _relax.AnyType(span)
+
+
+Any = Object
+
+
+is_type_var = _ir.is_prim_var
+
+
+def type_var(name, *, dtype=None, span=None):
+ """Construct a signature symbol under Relax's default shape dtype policy."""
+ return _ir.Var(name, "int64" if dtype is None else dtype, span)
+
+
+class _Frame:
+ """Retain source metadata and exports around an existing native frame."""
+
+ def __init__(self, native, span=None):
+ self.native = native
+ self.span = span
+ self.result = {}
+
+ def __getattr__(self, name):
+ return getattr(self.native, name)
+
+ @property
+ def reference(self):
+ """Return the stable module or local function reference after declaration."""
+ if isinstance(self.native, _frame.FunctionFrame):
+ local_var = self.native.local_var
+ return local_var if local_var is not None else self.native.global_var
+ raise AttributeError("This frame does not declare a function")
+
+ def __enter__(self):
+ with _protocol.span_context(self.span):
+ self.native.__enter__()
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ with _protocol.span_context(self.span):
+ self.native.__exit__(exc_type, exc_value, traceback)
+ if exc_type is None:
+ if isinstance(self.native, _frame.BindingBlockFrame):
+ self.result = {var.name: var for var in self.native.output_vars}
+ elif (
+ isinstance(self.native, _frame.FunctionFrame) and self.native.local_var is not None
+ ):
+ self.result = {self.native.name: self.native.local_var}
+ elif isinstance(self.native, _frame.IfFrame):
+ self.result = {self.native.var_name: self.native.var}
+ return False
+
+
+def function(is_pure=True, is_private=False, *, local=False, reference=None, span=None):
+ """Enter a definition using the native Relax function frame."""
+ if local:
+ if reference is None:
+ raise ValueError("A local function requires its declared reference")
+ return _Frame(_ffi_api.LocalFunction(is_pure, reference), span)
+ return _Frame(_native.function(is_pure, is_private), span)
+
+
+def decl_function(is_pure=True, is_private=False, *, local=False, span=None):
+ """Declare a bodyless function with the same signature operations as a definition."""
+ return _Frame(_ffi_api.DeclFunction(is_pure, is_private, local), span)
+
+
+def arg(name, ty, *, span=None):
+ """Add a parameter, retaining a cached parameter's identity when supplied."""
+ with _protocol.span_context(span):
+ if isinstance(ty, _ir.Var):
+ return _ffi_api.ArgVar(name, ty)
+ return _protocol.at(span, _native.arg(name, _type(ty)))
+
+
+def func_ret_type(ret_ty):
+ """Set the concrete return type of the active declaration or definition."""
+ return _native.func_ret_type(_type(ret_ty))
+
+
+func_ret_ty = func_ret_type
+
+
+def dataflow(*, span=None):
+ """Create a dataflow region whose result maps exported names to finalized vars."""
+ return _Frame(_native.dataflow(), span)
+
+
+def If(condition, *, span=None):
+ """Create a conditional region with finalized named exports."""
+ return _Frame(_native.If(condition), span)
+
+
+def Then(*, span=None):
+ """Create the true branch of a conditional."""
+ return _Frame(_native.Then(), span)
+
+
+def Else(*, span=None):
+ """Create the false branch of a conditional."""
+ return _Frame(_native.Else(), span)
+
+
+def _check_unterminated():
+ for active_frame in reversed(_IRBuilder.current().frames):
+ if isinstance(active_frame, _frame.FunctionFrame):
+ if active_frame.output is not None:
+ raise ValueError("A Relax operation cannot follow an unconditional return")
+ break
+
+
+def _value(value, ty=None):
+ if isinstance(value, _python.tuple):
+ return _relax.utils.convert_to_expr(value)
+ if isinstance(value, _numbers.Number):
+ if isinstance(ty, _ir.PrimType):
+ return _relax.prim_value(value, dtype=ty.dtype)
+ return _relax.const(value)
+ return value
+
+
+def bind_(
+ value=_protocol.MISSING,
+ *,
+ ty=None,
+ name=None,
+ span=None,
+ name_span=None,
+ previous=_protocol.MISSING,
+ declaration=False,
+ frame_value=False,
+):
+ """Emit a binding, or name and preserve an existing frame-owned value."""
+ _check_unterminated()
+ if frame_value:
+ if isinstance(value, _python.list | _python.tuple | _ir.Array):
+ for index, item in enumerate(value):
+ bind_(
+ item,
+ name=None if name is None else f"{name}_{index}",
+ span=span,
+ name_span=name_span,
+ frame_value=True,
+ )
+ elif isinstance(value, _ir.Var):
+ if name is not None:
+ _IRBuilder.name(name, value)
+ _protocol.at(name_span if name_span is not None else span, value)
+ return value
+ if declaration:
+ if not _ir.is_prim_var(value):
+ raise TypeError("A symbol declaration requires a concrete primitive variable")
+ if ty is not None and not _ffi.structural_equal(_type(ty), value.ty):
+ raise TypeError("The symbol declaration has an incompatible type")
+ if previous is not _protocol.MISSING:
+ if not _ir.is_prim_var(previous) or not _ffi.structural_equal(previous.ty, value.ty):
+ raise TypeError("The symbol declaration has an incompatible signature dtype")
+ return previous
+ if name is not None:
+ _IRBuilder.name(name, value)
+ return _protocol.at(name_span if name_span is not None else span, value)
+ if value is _protocol.MISSING:
+ raise ValueError("Relax bindings require an initializer")
+ if isinstance(value, _I.meta_var):
+ return value.value
+ ty = None if ty is None else _type(ty)
+ value = _value(value, ty)
+ with _protocol.span_context(span):
+ if isinstance(value, _relax.MatchCast):
+ if ty is not None and not _ffi.structural_equal(ty, value.ty):
+ raise TypeError("The binding annotation differs from the match-cast type")
+ result = _ffi_api.EmitMatchCastV2(value.value, value.ty, name_span)
+ elif isinstance(value, _relax.Expr):
+ result = _ffi_api.EmitV2(value, ty, name_span)
+ else:
+ return value
+ if name is not None:
+ _IRBuilder.name(name, result)
+ return _protocol.at(name_span if name_span is not None else span, result)
+
+
+def emit_(value, *, span=None):
+ """Emit a void expression statement, retaining effect-only operations returning None."""
+ if value is None:
+ return
+ if not isinstance(value, _relax.Expr):
+ raise TypeError(f"Unsupported expression statement value: {type(value).__name__}")
+ result = bind_(value, name="_", span=span)
+ if not isinstance(result.ty, _ir.TupleType) or len(result.ty.fields) != 0:
+ raise ValueError(
+ "Non-void expressions must be bound to a variable; "
+ f"expression of type {result.ty} was used as a statement"
+ )
+
+
+def return_(value=None, *, span=None):
+ """Record the function result without exiting Python construction."""
+ _check_unterminated()
+ with _protocol.span_context(span):
+ if value is None:
+ value = _relax.Tuple([])
+ _native.func_ret_value(_value(value))
+
+
+def match_cast(value, ty, *, span=None):
+ """Construct a concrete match-cast binding for bind_ to consume."""
+ if value is None:
+ raise ValueError("The match-cast value cannot be None")
+ ty = _type(ty)
+ return _relax.MatchCast(_ir.Var("", ty), _value(value), ty, span)
+
+
+def unpack(value):
+ """Project an IR tuple with known arity; leave Python iteration unchanged."""
+ if isinstance(value, _relax.Tuple):
+ return _python.tuple(value.fields)
+ if isinstance(value, _relax.Expr) and isinstance(value.ty, _ir.TupleType):
+ return _python.tuple(_relax.TupleGetItem(value, i) for i in range(len(value.ty.fields)))
+ return value
+
+
+def assert_(condition, message="", *, span=None):
+ """Construct a runtime assertion with construction-time diagnostic text."""
+ if not isinstance(message, _python.str):
+ raise TypeError("An assertion message must be construction-time text")
+ with _protocol.span_context(span):
+ emit_(_protocol.at(span, _native.assert_op(condition, format=message)), span=span)
+
+
+def For(*args, span=None, **kwargs):
+ """Reject imperative loops in the Relax expression dialect."""
+ raise TypeError("Relax does not support imperative for loops")
+
+
+def break_(*, span=None):
+ """Reject imperative loop control in Relax."""
+ raise TypeError("Relax does not support break")
+
+
+def continue_(*, span=None):
+ """Reject imperative loop control in Relax."""
+ raise TypeError("Relax does not support continue")
+
+
+def setitem(target, index, value, *, span=None):
+ """Reject mutable stores, which are not a Relax binding operation."""
+ raise TypeError("Relax does not support indexed assignment")
+
+
+__all__ = [
+ *_native.__all__,
+ "Any",
+ "Callable",
+ "DTensor",
+ "For",
+ "Object",
+ "Prim",
+ "Range",
+ "Shape",
+ "Tensor",
+ "Tuple",
+ "assert_",
+ "break_",
+ "continue_",
+ "bind_",
+ "decl_function",
+ "device_mesh",
+ "dist",
+ "emit_",
+ "is_type_var",
+ "match_cast",
+ "return_",
+ "setitem",
+ "type_var",
+ "unpack",
+]
+
+
+def _logical_pair(lhs, rhs, operation, primitive, python_operation):
+ if not isinstance(lhs, _ir.Expr) and not isinstance(rhs, _ir.Expr):
+ return python_operation(lhs, rhs)
+ if _ir.is_prim_expr(lhs) or _ir.is_prim_expr(rhs):
+ return primitive(lhs, rhs)
+ return operation(_value(lhs), _value(rhs))
+
+
+def logical_and(*values):
+ """Construct conjunction of concrete tensor or primitive expressions."""
+ if not values:
+ raise TypeError("logical_and requires at least one operand")
+ result = values[0]
+ for value in values[1:]:
+ result = _logical_pair(result, value, _relax.op.logical_and, _tir.And, lambda a, b: a and b)
+ return result
+
+
+def logical_or(*values):
+ """Construct disjunction of concrete tensor or primitive expressions."""
+ if not values:
+ raise TypeError("logical_or requires at least one operand")
+ result = values[0]
+ for value in values[1:]:
+ result = _logical_pair(result, value, _relax.op.logical_or, _tir.Or, lambda a, b: a or b)
+ return result
+
+
+def logical_not(value):
+ """Construct negation without coercing an IR expression to Python bool."""
+ if _ir.is_prim_expr(value):
+ return _tir.Not(value)
+ if isinstance(value, _ir.Expr):
+ return _relax.op.logical_not(value)
+ return not value
+
+
+def select(condition, true_value, false_value):
+ """Construct an elementwise conditional or select ordinary Python values."""
+ if _ir.is_prim_expr(condition):
+ return _tir.Select(condition, true_value, false_value)
+ if isinstance(condition, _ir.Expr):
+ return _relax.op.where(condition, _value(true_value), _value(false_value))
+ return true_value if condition else false_value
+
+
+__all__ += ["logical_and", "logical_not", "logical_or", "select"]
+
+
+def _call_global(function, *args):
+ return _relax.Call(function, [_relax.utils.convert_to_expr(value) for value in args])
+
+
+def _global_callee(function):
+ return _partial(_call_global, function)
+
+
+_protocol.register_call_kind(_sys.modules[__name__], _ir.GlobalVar, _global_callee)
+
+
+from tvm.script.parser.annotations import enable_eager_constructors as _enable_eager_constructors
+
+_enable_eager_constructors(_sys.modules[__name__])
diff --git a/python/tvm/relax/script/builder_v2/__init__.py b/python/tvm/relax/script/builder_v2/__init__.py
deleted file mode 100644
index 52745a4642..0000000000
--- a/python/tvm/relax/script/builder_v2/__init__.py
+++ /dev/null
@@ -1,465 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Concrete Relax construction operations over the shared native builder stack."""
-
-# pylint: disable=wildcard-import,redefined-builtin,invalid-name
-import builtins as _python
-import numbers as _numbers
-import sys as _sys
-from functools import partial as _partial
-
-import tvm_ffi as _ffi
-
-from tvm import ir as _ir
-from tvm import relax as _relax
-from tvm import tirx as _tir
-from tvm.relax.distributed import DeviceMesh as _DeviceMesh
-from tvm.relax.distributed import DTensorType as _DTensorType
-from tvm.relax.distributed import Placement as _Placement
-from tvm.relax.distributed import device_mesh as device_mesh
-from tvm.script.ir_builder import IRBuilder as _IRBuilder
-from tvm.script.ir_builder import ir as _I
-from tvm.script.ir_builder import protocol as _protocol
-
-from .. import builder as _legacy
-from ..builder import *
-from ..builder import _ffi_api
-from ..builder import distributed as dist
-from ..builder import frame as _frame
-from ..builder.distributed.ir import _lookup_device_mesh
-
-
-@_protocol.expression_args("shape", introduce=True, dtype="int64", scalar_strings=False)
-def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
- """Construct a concrete tensor type from already resolved dimensions."""
- if isinstance(shape, _python.str) and dtype is None:
- dtype, shape = shape, None
- if isinstance(vdevice, _python.str):
- target, *qualifiers = vdevice.split(":", 2)
- vdevice = _I.lookup_vdevice(target, int(qualifiers[0]) if qualifiers else 0)
- return _relax.TensorType(shape, dtype, vdevice, ndim, span)
-
-
-@_protocol.expression_args("shape", introduce=True, dtype="int64", scalar_strings=False)
-def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *, ndim=-1, span=None):
- """Construct a concrete distributed tensor type with resolved shape symbols."""
- if device_mesh is None:
- device_mesh = _DeviceMesh([], _ir.Range(0, 1))
- elif isinstance(device_mesh, _python.str):
- device_mesh = _lookup_device_mesh(device_mesh)
- if isinstance(placement, _python.str):
- placement = _Placement.from_text(placement)
- return _DTensorType(Tensor(shape, dtype, ndim=ndim), device_mesh, placement, span)
-
-
-Range = _ir.Range
-
-
-@_protocol.expression_args("values", introduce=True, dtype="int64")
-def Shape(values=None, ndim=-1, *, span=None):
- """Construct a concrete shape type."""
- return _relax.ShapeType(values, ndim, span)
-
-
-def _type(value):
- if value is None:
- return _ir.TupleType([])
- if callable(value):
- value = value()
- if _ir.is_prim_expr(value):
- value = value.ty
- if not isinstance(value, _ir.Type):
- raise TypeError(f"Expected a concrete type, got {type(value).__name__}")
- return value
-
-
-def Callable(params=None, ret=None, purity=None, derive_func=None, *, span=None):
- """Construct a concrete function type."""
- if purity is None:
- purity = params is not None
- if params is None:
- return _relax.FuncType.opaque_func(
- ret=None if ret is None else _type(ret),
- derive_func=derive_func,
- purity=purity,
- span=span,
- )
- if derive_func is not None:
- raise ValueError("A derivation function requires an opaque callable")
- if not isinstance(params, list | _python.tuple):
- params = [params]
- return _relax.FuncType([_type(param) for param in params], _type(ret), purity, span)
-
-
-def Tuple(*fields, span=None):
- """Construct a concrete tuple type."""
- if len(fields) == 1 and isinstance(fields[0], list | _python.tuple):
- fields = fields[0]
- return _ir.TupleType([_type(field) for field in fields], span)
-
-
-def Prim(dtype, *, span=None):
- """Construct a primitive type."""
- return _ir.PrimType(dtype)
-
-
-Prim.__tvm_parameter_dtype__ = "dtype"
-
-
-def Object(*, span=None):
- """Construct the unconstrained Relax value type."""
- return _relax.AnyType(span)
-
-
-Any = Object
-
-
-is_type_var = _ir.is_prim_var
-
-
-def type_var(name, *, dtype=None, span=None):
- """Construct a signature symbol under Relax's default shape dtype policy."""
- return _ir.Var(name, "int64" if dtype is None else dtype, span)
-
-
-class _Frame:
- """Retain source metadata and exports around an existing native frame."""
-
- def __init__(self, native, span=None):
- self.native = native
- self.span = span
- self.result = {}
-
- def __getattr__(self, name):
- return getattr(self.native, name)
-
- @property
- def reference(self):
- """Return the stable module or local function reference after declaration."""
- if isinstance(self.native, _frame.FunctionFrame):
- local_var = self.native.local_var
- return local_var if local_var is not None else self.native.global_var
- raise AttributeError("This frame does not declare a function")
-
- def __enter__(self):
- with _protocol.span_context(self.span):
- self.native.__enter__()
- return self
-
- def __exit__(self, exc_type, exc_value, traceback):
- with _protocol.span_context(self.span):
- self.native.__exit__(exc_type, exc_value, traceback)
- if exc_type is None:
- if isinstance(self.native, _frame.BindingBlockFrame):
- self.result = {var.name: var for var in self.native.output_vars}
- elif (
- isinstance(self.native, _frame.FunctionFrame) and self.native.local_var is not None
- ):
- self.result = {self.native.name: self.native.local_var}
- elif isinstance(self.native, _frame.IfFrame):
- self.result = {self.native.var_name: self.native.var}
- return False
-
-
-def function(is_pure=True, is_private=False, *, local=False, reference=None, span=None):
- """Enter a definition using the native Relax function frame."""
- if local:
- if reference is None:
- raise ValueError("A local function requires its declared reference")
- return _Frame(_ffi_api.LocalFunction(is_pure, reference), span)
- return _Frame(_legacy.function(is_pure, is_private), span)
-
-
-def decl_function(is_pure=True, is_private=False, *, local=False, span=None):
- """Declare a bodyless function with the same signature operations as a definition."""
- return _Frame(_ffi_api.DeclFunction(is_pure, is_private, local), span)
-
-
-def arg(name, ty, *, span=None):
- """Add a parameter, retaining a cached parameter's identity when supplied."""
- with _protocol.span_context(span):
- if isinstance(ty, _ir.Var):
- return _ffi_api.ArgVar(name, ty)
- return _protocol.at(span, _legacy.arg(name, _type(ty)))
-
-
-def func_ret_type(ret_ty):
- """Set the concrete return type of the active declaration or definition."""
- return _legacy.func_ret_type(_type(ret_ty))
-
-
-func_ret_ty = func_ret_type
-
-
-def dataflow(*, span=None):
- """Create a dataflow region whose result maps exported names to finalized vars."""
- return _Frame(_legacy.dataflow(), span)
-
-
-def If(condition, *, span=None):
- """Create a conditional region with finalized named exports."""
- return _Frame(_legacy.If(condition), span)
-
-
-def Then(*, span=None):
- """Create the true branch of a conditional."""
- return _Frame(_legacy.Then(), span)
-
-
-def Else(*, span=None):
- """Create the false branch of a conditional."""
- return _Frame(_legacy.Else(), span)
-
-
-def _check_unterminated():
- for frame in reversed(_IRBuilder.current().frames):
- if isinstance(frame, _frame.FunctionFrame):
- if frame.output is not None:
- raise ValueError("A Relax operation cannot follow an unconditional return")
- break
-
-
-def _value(value, ty=None):
- if isinstance(value, _python.tuple):
- return _relax.utils.convert_to_expr(value)
- if isinstance(value, _numbers.Number):
- if isinstance(ty, _ir.PrimType):
- return _relax.prim_value(value, dtype=ty.dtype)
- return _relax.const(value)
- return value
-
-
-def bind_(
- value=_protocol.MISSING,
- *,
- ty=None,
- name=None,
- span=None,
- name_span=None,
- previous=_protocol.MISSING,
- declaration=False,
- frame_value=False,
-):
- """Emit a binding, or name and preserve an existing frame-owned value."""
- _check_unterminated()
- if frame_value:
- if isinstance(value, _python.list | _python.tuple | _ir.Array):
- for index, item in enumerate(value):
- bind_(
- item,
- name=None if name is None else f"{name}_{index}",
- span=span,
- name_span=name_span,
- frame_value=True,
- )
- elif isinstance(value, _ir.Var):
- if name is not None:
- _IRBuilder.name(name, value)
- _protocol.at(name_span if name_span is not None else span, value)
- return value
- if declaration:
- if not _ir.is_prim_var(value):
- raise TypeError("A symbol declaration requires a concrete primitive variable")
- if ty is not None and not _ffi.structural_equal(_type(ty), value.ty):
- raise TypeError("The symbol declaration has an incompatible type")
- if previous is not _protocol.MISSING:
- if not _ir.is_prim_var(previous) or not _ffi.structural_equal(previous.ty, value.ty):
- raise TypeError("The symbol declaration has an incompatible signature dtype")
- return previous
- if name is not None:
- _IRBuilder.name(name, value)
- return _protocol.at(name_span if name_span is not None else span, value)
- if value is _protocol.MISSING:
- raise ValueError("Relax bindings require an initializer")
- if isinstance(value, _I.meta_var):
- return value.value
- ty = None if ty is None else _type(ty)
- value = _value(value, ty)
- with _protocol.span_context(span):
- if isinstance(value, _relax.MatchCast):
- if ty is not None and not _ffi.structural_equal(ty, value.ty):
- raise TypeError("The binding annotation differs from the match-cast type")
- result = _ffi_api.EmitMatchCastV2(value.value, value.ty, name_span)
- elif isinstance(value, _relax.Expr):
- result = _ffi_api.EmitV2(value, ty, name_span)
- else:
- return value
- if name is not None:
- _IRBuilder.name(name, result)
- return _protocol.at(name_span if name_span is not None else span, result)
-
-
-def emit_(value, *, span=None):
- """Emit a void expression statement, retaining effect-only operations returning None."""
- if value is None:
- return
- if not isinstance(value, _relax.Expr):
- raise TypeError(f"Unsupported expression statement value: {type(value).__name__}")
- result = bind_(value, name="_", span=span)
- if not isinstance(result.ty, _ir.TupleType) or len(result.ty.fields) != 0:
- raise ValueError(
- "Non-void expressions must be bound to a variable; "
- f"expression of type {result.ty} was used as a statement"
- )
-
-
-def return_(value=None, *, span=None):
- """Record the function result without exiting Python construction."""
- _check_unterminated()
- with _protocol.span_context(span):
- if value is None:
- value = _relax.Tuple([])
- _legacy.func_ret_value(_value(value))
-
-
-def match_cast(value, ty, *, span=None):
- """Construct a concrete match-cast binding for bind_ to consume."""
- if value is None:
- raise ValueError("The match-cast value cannot be None")
- ty = _type(ty)
- return _relax.MatchCast(_ir.Var("", ty), _value(value), ty, span)
-
-
-def unpack(value):
- """Project an IR tuple with known arity; leave Python iteration unchanged."""
- if isinstance(value, _relax.Tuple):
- return _python.tuple(value.fields)
- if isinstance(value, _relax.Expr) and isinstance(value.ty, _ir.TupleType):
- return _python.tuple(_relax.TupleGetItem(value, i) for i in range(len(value.ty.fields)))
- return value
-
-
-def assert_(condition, message="", *, span=None):
- """Construct a runtime assertion with construction-time diagnostic text."""
- if not isinstance(message, _python.str):
- raise TypeError("An assertion message must be construction-time text")
- with _protocol.span_context(span):
- emit_(_protocol.at(span, _legacy.assert_op(condition, format=message)), span=span)
-
-
-def For(*args, span=None, **kwargs):
- """Reject imperative loops in the Relax expression dialect."""
- raise TypeError("Relax does not support imperative for loops")
-
-
-def break_(*, span=None):
- """Reject imperative loop control in Relax."""
- raise TypeError("Relax does not support break")
-
-
-def continue_(*, span=None):
- """Reject imperative loop control in Relax."""
- raise TypeError("Relax does not support continue")
-
-
-def setitem(target, index, value, *, span=None):
- """Reject mutable stores, which are not a Relax binding operation."""
- raise TypeError("Relax does not support indexed assignment")
-
-
-__all__ = [
- *_legacy.ir.__all__,
- "Any",
- "Callable",
- "DTensor",
- "For",
- "Object",
- "Prim",
- "Range",
- "Shape",
- "Tensor",
- "Tuple",
- "assert_",
- "break_",
- "continue_",
- "bind_",
- "decl_function",
- "device_mesh",
- "dist",
- "emit_",
- "is_type_var",
- "match_cast",
- "return_",
- "setitem",
- "type_var",
- "unpack",
-]
-
-
-def _logical_pair(lhs, rhs, operation, primitive, python_operation):
- if not isinstance(lhs, _ir.Expr) and not isinstance(rhs, _ir.Expr):
- return python_operation(lhs, rhs)
- if _ir.is_prim_expr(lhs) or _ir.is_prim_expr(rhs):
- return primitive(lhs, rhs)
- return operation(_value(lhs), _value(rhs))
-
-
-def logical_and(*values):
- """Construct conjunction of concrete tensor or primitive expressions."""
- if not values:
- raise TypeError("logical_and requires at least one operand")
- result = values[0]
- for value in values[1:]:
- result = _logical_pair(result, value, _relax.op.logical_and, _tir.And, lambda a, b: a and b)
- return result
-
-
-def logical_or(*values):
- """Construct disjunction of concrete tensor or primitive expressions."""
- if not values:
- raise TypeError("logical_or requires at least one operand")
- result = values[0]
- for value in values[1:]:
- result = _logical_pair(result, value, _relax.op.logical_or, _tir.Or, lambda a, b: a or b)
- return result
-
-
-def logical_not(value):
- """Construct negation without coercing an IR expression to Python bool."""
- if _ir.is_prim_expr(value):
- return _tir.Not(value)
- if isinstance(value, _ir.Expr):
- return _relax.op.logical_not(value)
- return not value
-
-
-def select(condition, true_value, false_value):
- """Construct an elementwise conditional or select ordinary Python values."""
- if _ir.is_prim_expr(condition):
- return _tir.Select(condition, true_value, false_value)
- if isinstance(condition, _ir.Expr):
- return _relax.op.where(condition, _value(true_value), _value(false_value))
- return true_value if condition else false_value
-
-
-__all__ += ["logical_and", "logical_not", "logical_or", "select"]
-
-
-def _call_global(function, *args):
- return _relax.Call(function, [_relax.utils.convert_to_expr(value) for value in args])
-
-
-def _global_callee(function):
- return _partial(_call_global, function)
-
-
-_protocol.register_call_kind(_sys.modules[__name__], _ir.GlobalVar, _global_callee)
-
-
-from tvm.script.parser_v2.annotations import enable_eager_constructors as _enable_eager_constructors
-
-_enable_eager_constructors(_sys.modules[__name__])
diff --git a/python/tvm/relax/script/parser/__init__.py b/python/tvm/relax/script/parser/__init__.py
deleted file mode 100644
index 9795764cb4..0000000000
--- a/python/tvm/relax/script/parser/__init__.py
+++ /dev/null
@@ -1,52 +0,0 @@
-# isort: skip_file
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# ruff: noqa: RUF005
-"""Initial impl of relax parser for sugars"""
-
-from typing import TYPE_CHECKING
-
-from tvm.relax.script.builder import * # pylint: disable=redefined-builtin
-from tvm.relax.script.builder import ir as _relax
-
-from . import parser as _parser
-from .entry import Any, Callable, Object, Prim, Shape, Tensor, Tuple, match_cast
-
-from . import dist
-from .dist import * # pylint: disable=wildcard-import,redefined-builtin
-
-if TYPE_CHECKING:
- # pylint: disable=invalid-name
- # Define prim_func and make it type check as static method
- # so most tvmscript won't trigger pylint error here.
- function = staticmethod
-else:
- from .entry import function, macro
-
-__all__ = _relax.__all__ + [
- "dist",
- "Any",
- "Callable",
- "Object",
- "Prim",
- "Shape",
- "Tensor",
- "Tuple",
- "function",
- "macro",
- "match_cast",
-]
diff --git a/python/tvm/relax/script/parser/dist.py b/python/tvm/relax/script/parser/dist.py
deleted file mode 100644
index 71030be905..0000000000
--- a/python/tvm/relax/script/parser/dist.py
+++ /dev/null
@@ -1,106 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=redefined-builtin,missing-docstring, invalid-name, unused-import, redefined-outer-name
-# ruff: noqa: F401
-
-from typing import Any, Optional, Union
-
-from tvm.ir import Range
-from tvm.relax import TensorType
-from tvm.relax.distributed import DeviceMesh, DTensorType, Placement, device_mesh
-from tvm.relax.script.builder.distributed import (
- annotate_sharding,
- call_tir,
- call_tir_local_view,
- const,
- redistribute,
- redistribute_replica_to_shard,
-)
-from tvm.script.ir_builder import IRBuilder
-from tvm.script.ir_builder.ir import IRModuleFrame
-from tvm.tirx import Expr
-
-from .entry import TensorProxy, TypeProxy
-
-############################### R.DTensor ###############################
-
-
-class DTensorProxy(TypeProxy):
- tensor_ty_proxy: TensorProxy
- device_mesh: DeviceMesh
- placement: Placement
-
- def __init__(
- self,
- tensor_ty_proxy: TensorProxy,
- device_mesh: DeviceMesh,
- placement: Placement,
- ) -> None:
- self.device_mesh = device_mesh
- self.placement = placement
- self.tensor_ty_proxy = tensor_ty_proxy
- super().__init__()
-
- def get_symbolic_vars(self) -> set[str]:
- return self.tensor_ty_proxy.get_symbolic_vars()
-
- def as_ty(self, dict_globals: dict[str, Any] | None = None) -> DTensorType:
- return DTensorType(
- self.tensor_ty_proxy.as_ty(dict_globals),
- self.device_mesh,
- self.placement,
- )
-
-
-def DTensor(
- shape: list[Expr | str] | None = None,
- dtype: str | None = None,
- device_mesh: DeviceMesh | str = DeviceMesh([], Range(0, 1)),
- placement: Placement | str = "",
- *,
- ndim: int = -1,
-) -> DTensorProxy:
- # scalar tensor case
- if shape is not None and len(shape) == 0:
- shape = []
- if isinstance(shape, str) and dtype is None:
- dtype = shape
- shape = None
-
- if shape is not None and not isinstance(shape, tuple | list):
- raise ValueError(f"shape must be a list or tuple, but got: {shape}")
- if isinstance(device_mesh, str):
- if not IRBuilder.is_in_scope():
- return (
- DTensorProxy(
- TensorProxy(shape, dtype, None, ndim), DeviceMesh([], Range(0, 1)), ""
- ),
- )
- name, index = device_mesh.split("[")
- index = int(index[:-1])
- frames = IRBuilder.current().frames
- for f in frames:
- if isinstance(f, IRModuleFrame):
- device_mesh = f.global_infos[name][index]
- break
- assert isinstance(device_mesh, DeviceMesh)
- if isinstance(placement, str):
- placement = Placement.from_text(placement)
- return DTensorProxy(TensorProxy(shape, dtype, None, ndim), device_mesh, placement)
-
-
-__all__ = ["DTensor", "device_mesh"]
diff --git a/python/tvm/relax/script/parser/entry.py b/python/tvm/relax/script/parser/entry.py
deleted file mode 100644
index 20605840f6..0000000000
--- a/python/tvm/relax/script/parser/entry.py
+++ /dev/null
@@ -1,521 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=missing-docstring, invalid-name
-import inspect
-from collections.abc import Callable as _Callable
-from typing import Any, TypeVar
-
-import tvm
-from tvm.ir import PrimType
-from tvm.relax import (
- AnyType,
- Expr,
- Function,
- FuncType,
- SeqExpr,
- ShapeExpr,
- ShapeType,
- TensorType,
- TupleType,
- Type,
-)
-from tvm.relax.expr import Var
-from tvm.relax.script import builder as R
-from tvm.runtime import ObjectConvertible
-from tvm.script.ir_builder.ir import lookup_vdevice
-from tvm.script.parser._core import doc, parse, utils
-from tvm.script.parser.core.entry import scan_macro
-from tvm.script.parser.core.parser import Parser, ScriptMacro
-
-FType = TypeVar("FType", bound=_Callable)
-
-############################## R.function ##############################
-
-
-# this formulation allows us to support having @R.function
-# appear as a decorator by itself or to have optional arguments
-# like @R.function(pure=False)
-def function(
- f: FType | None = None, pure: bool = True, private: bool = False, check_well_formed=True
-) -> Function | FType:
- # pylint: disable=unused-argument
- # (pure and private aren't used here, but are used later in parsing)
-
- # need to inspect the stack first because is_defined_in_class expects the outer class
- # to be in a particular position in the stack
- orig_stack = inspect.stack()
-
- def decorator_wrapper(f):
- if not inspect.isfunction(f):
- raise TypeError(f"Expect a function, but got: {f}")
- if utils.is_defined_in_class(orig_stack, f):
- return f
- return parse(f, utils.inspect_function_capture(f), check_well_formed=check_well_formed)
-
- if f is not None:
- # if there are no optional args given, this will directly invoke the wrapper
- return decorator_wrapper(f)
- else:
- # if there is a optional arg given, it returns the wrapper function
- # as a new decorator and applies it
- setattr(decorator_wrapper, "dispatch_token", "relax")
- return decorator_wrapper
-
-
-setattr(function, "dispatch_token", "relax")
-
-
-############################## R.macro ##############################
-
-
-class RelaxMacro(ScriptMacro):
- """Specialization of the ScriptMacro class for Relax."""
-
- def parse_macro(self, parser: Parser) -> Expr:
- macro_def = self.get_macro_def()
- ret_value = None
-
- with R.SeqExpr() as seq:
- for idx, stmt in enumerate(macro_def.body):
- # Normally, a "return" statement is only allowed in a R.function. We don't
- # want to parse the macro's body as if it was a body of a function, because
- # the latter imposes some constraints that we want to avoid.
- # At the same time, we want to use "return" to indicate the value of the
- # macro (since in Relax everything is an expression), so add special handling
- # of "return".
- if isinstance(stmt, doc.Return):
- ret_value = parser.eval_expr(stmt.value)
- if idx + 1 != len(macro_def.body):
- parser.report_error(macro_def, "'return' should be the last statement")
- break
- parser.visit(stmt)
-
- if ret_value is None:
- parser.report_error(macro_def, "Macros must end with a return statement")
-
- return SeqExpr(seq.binding_blocks, ret_value)
-
-
-def macro(*args, hygienic: bool = True) -> _Callable:
- """Decorator for macro definitions.
-
- Parameters
- ----------
- hygienic: bool
- Specifies whether the macro is hygienic or not.
- A macro is hygienic if all symbols used in the macro's body are resolved
- to values from the location of the macro definition. A non-hygienic macro
- will have its symbols resolved to values at the time of the macro's use.
- """
-
- def _decorator(func: _Callable) -> ScriptMacro:
- source, closure_vars = scan_macro(func, utils.inspect_function_capture(func))
- obj = RelaxMacro(source, closure_vars, func, hygienic)
-
- def wrapper(*args, **kwargs):
- return obj(*args, **kwargs)
-
- return wrapper
-
- if len(args) == 0:
- return _decorator
- if len(args) == 1 and inspect.isfunction(args[0]):
- return _decorator(args[0])
-
- raise ValueError(
- "Invalid use of R.macro. Usage: @R.macro, @R.macro(), @R.macro(hygienic=[True|False])"
- )
-
-
-############################# Type ##############################
-
-
-class TypeProxy(ObjectConvertible):
- def as_ty(self, dict_globals: dict[str, Any] | None = None) -> Type:
- raise NotImplementedError()
-
- def get_symbolic_vars(self) -> set[str]:
- return {}
-
- def asobject(self):
- return self.as_ty(None)
-
-
-############################### R.Any ################################
-
-
-class AnyProxy(TypeProxy):
- """The proxy for AnyType.
-
- Parameters
- ----------
- values : Optional[List[Expr]]
- The symbolic shape values if known.
-
- ndim : Optional[int]
- The size of the shape.
- """
-
- def __init__(self) -> None:
- pass
-
- def get_symbolic_vars(self) -> set[str]:
- return set()
-
- def as_ty(self, dict_globals: dict[str, Any] | None = None) -> AnyType:
- return AnyType()
-
-
-def Any() -> AnyProxy:
- return AnyProxy()
-
-
-ObjectProxy = AnyProxy
-
-
-def Object() -> AnyProxy:
- return AnyProxy()
-
-
-############################### R.Tensor ###############################
-
-
-def _eval_shape(expr: str | Expr, dict_globals: dict[str, Any] | None) -> Expr:
- if isinstance(expr, str):
- code = compile(expr, "<string>", "eval")
- return eval(code, dict_globals or {}) # pylint: disable=eval-used
- else:
- return expr
-
-
-class TensorProxy(TypeProxy):
- shape: list[str | Expr] | None
- dtype: str
- vdevice: str | None
- ndim: int
-
- def __init__(
- self,
- shape: list[Expr | str] | Expr | None = None,
- dtype: str | None = None,
- vdevice: str | None = None,
- ndim: int = -1,
- ) -> None:
- if isinstance(shape, Expr):
- if not isinstance(shape, ShapeExpr | Var):
- raise ValueError(
- "When the shape is an Expr, it must be a ShapeExpr or a Var with ShapeExpr "
- f"value. But got: {shape} with type: {type(shape)}"
- )
- if isinstance(shape, Var) and not isinstance(shape.ty, ShapeType):
- raise ValueError(
- "When the shape is a Var, it must have shape ty. But got "
- f"{shape} with ty: {shape.ty}"
- )
- self.shape = shape
- self.dtype = dtype
- self.vdevice = vdevice
- self.ndim = ndim
-
- def get_symbolic_vars(self) -> set[str]:
- if self.shape is None or isinstance(self.shape, Expr):
- return {}
- else:
- return {s for s in self.shape if isinstance(s, str) and s.isidentifier()}
-
- def as_ty(self, dict_globals: dict[str, Any] | None = None) -> TensorType:
- vdev = self.vdevice
- if isinstance(self.vdevice, str):
- if ":" in self.vdevice:
- split_vdev = self.vdevice.split(":")
- vdev = lookup_vdevice(split_vdev[0], int(split_vdev[1]))
- else:
- vdev = lookup_vdevice(self.vdevice, 0)
-
- if self.shape is None:
- return TensorType(None, self.dtype, vdev, self.ndim)
- elif isinstance(self.shape, ShapeExpr | Var):
- return TensorType(self.shape, self.dtype, vdev, self.ndim)
- else:
- if dict_globals is None and any([isinstance(s, str) for s in self.shape]):
- raise ValueError(
- "String-defined shape expr is only allowed when parsing function parameters "
- "and return annotations for TVMScript."
- )
- shape = [_eval_shape(s, dict_globals) for s in self.shape]
- return TensorType(shape, self.dtype, vdev, self.ndim)
-
-
-def Tensor(
- shape: list[Expr | str] | Expr | None = None,
- dtype: str | None = None,
- vdevice: str | None = None,
- ndim: int = -1,
-) -> TensorProxy:
- # scalar tensor case
- if shape is not None and not isinstance(shape, Var) and len(shape) == 0:
- shape = []
- if isinstance(shape, str) and dtype is None:
- dtype = shape
- shape = None
-
- if shape is not None and not isinstance(shape, tuple | list) and not isinstance(shape, Expr):
- raise ValueError(f"shape must be a list/tuple or an Expr, but got: {shape}")
- return TensorProxy(shape, dtype, vdevice, ndim)
-
-
-############################## R.Callable ##############################
-
-
-class CallableProxy(TypeProxy):
- params: list[TypeProxy]
- ret: TypeProxy
- purity: bool
- derive_func: str | tvm.ir.EnvFunc | None
-
- """Function type.
-
- A function type consists of a list of type parameters to enable
- the definition of generic functions,
- a set of type constraints which we omit for the time being,
- a sequence of argument types, the purity of the function, and a return type.
-
- Parameters
- ----------
- params : List[TypeProxy]
- The argument TypeProxy
-
- ret : TypeProxy
- The return TypeProxy.
-
- purity : bool
- Whether the callable is pure.
-
- derive_func: Optional[Union[str, tvm.ir.EnvFunc]]
- The derivation function to determine the output Type,
- based on the arguments provided to the function. The
- specified function should be accessible using
- `tvm.get_global_func`, and should have a signature
- `Callable[[relax.Call, relax.BlockBuilder], relax.Type]`.
-
- """
-
- def __init__(
- self,
- params: TypeProxy | list[TypeProxy] | None = None,
- ret: TypeProxy | None = None,
- purity: bool | None = None,
- derive_func: str | tvm.ir.EnvFunc | None = None,
- ) -> None:
- if params is None:
- self.params = params
- else:
- if not isinstance(params, list | tuple):
- params = [params]
- # convert `R.Callable` to `R.Callable()`
- self.params = [param() if callable(param) else param for param in params]
-
- # Mimic the C++ defaults, where an opaque function is assumed
- # to be impure, and a non-opaque function is assumed to be
- # pure.
- if purity is None:
- purity = params is not None
-
- self.ret = ret() if callable(ret) else ret
- self.purity = purity
- self.derive_func = derive_func
-
- def get_symbolic_vars(self) -> set[str]:
- if self.params is None:
- return set()
- else:
- return set().union(*[p.get_symbolic_vars() for p in self.params])
-
- def as_ty(self, dict_globals: dict[str, Any] | None = None) -> FuncType:
- if self.ret is None:
- ret = None
- else:
- ret = self.ret.as_ty(dict_globals)
-
- if self.params is None:
- params = None
- else:
- params = [param.as_ty(dict_globals) for param in self.params]
-
- if params is None:
- return FuncType.opaque_func(ret=ret, derive_func=self.derive_func, purity=self.purity)
- else:
- return FuncType(params, ret, purity=self.purity)
-
-
-def Callable(
- params: TypeProxy | list[TypeProxy] | None = None,
- ret: TypeProxy | None = None,
- purity: bool | None = None,
- derive_func: str | tvm.ir.EnvFunc | None = None,
-) -> CallableProxy:
- return CallableProxy(params, ret, purity=purity, derive_func=derive_func)
-
-
-############################### R.Tuple ################################
-
-
-class TupleProxy(TypeProxy):
- fields: list[TypeProxy]
- """The type of tuple values.
-
- Parameters
- ----------
- fields : List[TypeProxy]
- The fields in the tuple
- """
-
- def __init__(
- self,
- *fields: list[TypeProxy],
- ) -> None:
- if len(fields) == 1 and isinstance(fields[0], tuple | list):
- fields = fields[0]
- # convert `R.Tensor` to `R.Tensor()`
- self.fields = [field() if callable(field) else field for field in fields]
-
- def get_symbolic_vars(self) -> set[str]:
- return set().union(*[f.get_symbolic_vars() for f in self.fields])
-
- def as_ty(self, dict_globals: dict[str, Any] | None = None) -> TupleType:
- fields = [field.as_ty(dict_globals) for field in self.fields]
- return TupleType(fields)
-
-
-def Tuple(*fields: list[TypeProxy]) -> TupleProxy:
- return TupleProxy(*fields)
-
-
-############################### R.Shape ################################
-
-
-class ShapeProxy(TypeProxy):
- values: list[Expr] | None
- ndim: int
- """The type of shape values.
-
- Parameters
- ----------
- values : Optional[List[Expr]]
- The symbolic shape values if known.
-
- ndim : Optional[int]
- The size of the shape.
- """
-
- def __init__(
- self,
- values: list[Expr] | None = None,
- ndim: int = -1,
- ) -> None:
- self.values = values
- self.ndim = ndim
-
- def get_symbolic_vars(self) -> set[str]:
- if self.values is None:
- return set()
- else:
- return {v for v in self.values if isinstance(v, str) and v.isidentifier()}
-
- def as_ty(self, dict_globals: dict[str, Any] | None = None) -> ShapeType:
- values = [_eval_shape(v, dict_globals) for v in self.values] if self.values else None
- return ShapeType(values, self.ndim)
-
-
-def Shape(values: list[Expr] | None = None, ndim: int = -1) -> ShapeProxy:
- return ShapeProxy(values, ndim)
-
-
-################################ R.Prim ################################
-
-
-class PrimProxy(TypeProxy):
- dtype: str
-
- """The type of TIR-representable values.
-
- Parameters
- ----------
- dtype : str
- The data type.
-
- """
-
- def __init__(
- self,
- dtype: str,
- ) -> None:
- self.dtype = dtype
-
- def get_symbolic_vars(self) -> set[str]:
- return set()
-
- def as_ty(self, dict_globals: dict[str, Any] | None = None) -> PrimType:
- return PrimType(self.dtype)
-
-
-def Prim(
- dtype: str,
-) -> PrimProxy:
- return PrimProxy(dtype)
-
-
-############################ R.match_cast #############################
-class MatchCastPair:
- value: Expr
- ty: Type
-
- def __init__(self, value: Expr, ty: Type) -> None:
- self.value = value
- self.ty = ty
-
-
-def match_cast(value: Expr, ty: Type):
- ty = _normalize_ty(ty)
-
- if value is None:
- raise ValueError("value of match_cast cannot be None")
- if ty is None:
- raise ValueError("ty of match_cast cannot be None")
- return MatchCastPair(value, ty)
-
-
-def _normalize_ty_proxy(annotation) -> TypeProxy:
- if annotation is None:
- return TupleProxy([])
- elif callable(annotation):
- annotation = annotation()
- if tvm.ir.is_prim_expr(annotation):
- return PrimProxy(annotation.ty.dtype)
- return annotation
- elif isinstance(annotation, TypeProxy):
- return annotation
- else:
- raise TypeError(f"Expected TypeProxy but got {type(annotation)}.")
-
-
-def _normalize_ty(ty, dict_globals: dict[str, Any] | None = None) -> Type:
- if isinstance(ty, Type):
- return ty
- else:
- proxy = _normalize_ty_proxy(ty)
- return proxy.as_ty(dict_globals)
diff --git a/python/tvm/relax/script/parser/parser.py b/python/tvm/relax/script/parser/parser.py
deleted file mode 100644
index 083aa3ea93..0000000000
--- a/python/tvm/relax/script/parser/parser.py
+++ /dev/null
@@ -1,506 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=missing-docstring, unused-argument
-
-import functools
-import numbers
-from typing import Any
-
-import numpy as np
-import tvm_ffi
-
-import tvm
-from tvm import relax
-from tvm.ir import GlobalVar
-from tvm.relax import Expr, Type
-from tvm.relax.script import builder as R
-from tvm.relax.script.builder.frame import BindingBlockFrame
-from tvm.relax.utils import convert_to_expr
-from tvm.script.ir_builder import ir as I
-from tvm.script.ir_builder.base import IRBuilder
-from tvm.script.parser._core import Parser, collect_signature_type_vars, dispatch, doc
-from tvm.tirx.script import builder as T
-
-from .entry import (
- MatchCastPair,
- PrimProxy,
- TypeProxy,
- _normalize_ty,
- _normalize_ty_proxy,
-)
-
-relax.Expr._dispatch_type = relax.Expr # pylint: disable=protected-access
-dispatch.register_op(relax.Expr, doc.GtE, 0)(lambda lhs, rhs: lhs >= rhs)
-dispatch.register_op(relax.Expr, doc.Gt, 0)(lambda lhs, rhs: lhs > rhs)
-dispatch.register_op(relax.Expr, doc.LtE, 0)(lambda lhs, rhs: lhs <= rhs)
-dispatch.register_op(relax.Expr, doc.Lt, 0)(lambda lhs, rhs: lhs < rhs)
-
-
-def bind_assign_value(
- self: Parser,
- node: doc.expr,
- var_name: str,
- value: Any,
- anno_ty: Type | None = None,
- symbolic_declarations: set[str] | None = None,
-) -> Any:
- var_table = self.var_table.get()
-
- # Primitive assignments emit Relax bindings by default, reversing the old default.
- # Explicit I.meta_var values and literal T.dtype() declarations are the two
- # parser-time exceptions.
- if isinstance(value, I.meta_var):
- return value.value
-
- if tvm.ir.is_prim_var(value) and var_name in (symbolic_declarations or set()):
- if anno_ty is not None and not tvm_ffi.structural_equal(anno_ty, value.ty):
- self.report_error(
- node,
- f"Expected the declared symbolic variable to have type {anno_ty}, "
- f"but got {value.ty}",
- )
- if value.name and var_name != value.name:
- self.report_error(
- node,
- "Cannot define PrimType values with different names. The LHS of binding "
- "must have the same name provided in the RHS.",
- )
- if var_name in var_table:
- prev_value = var_table[var_name]
- if not tvm.ir.is_prim_var(prev_value):
- self.report_error(
- node,
- "Cannot redefine a non-PrimType object to a PrimType value. Please "
- "define the PrimType value with another name.",
- )
- if prev_value.ty != value.ty:
- self.report_error(
- node,
- f"Expected the same dtype for PrimType values but got "
- f"{value.ty} vs {prev_value.ty}",
- )
- value = prev_value
- IRBuilder.name(var_name, value)
- return value
-
- if isinstance(value, tuple):
- value = convert_to_expr(value)
- if isinstance(value, numbers.Number | np.bool_):
- if isinstance(anno_ty, tvm.ir.PrimType):
- value = relax.prim_value(value, dtype=anno_ty.dtype)
- else:
- value = R.const(value)
-
- if isinstance(value, relax.Expr):
- var = R.emit(value, anno_ty)
- elif isinstance(value, MatchCastPair):
- if anno_ty is not None and not tvm_ffi.structural_equal(anno_ty, value.ty):
- self.report_error(
- node, "Cannot specify inconsistent annotation for a match cast pair. "
- )
- var = R.emit_match_cast(value.value, value.ty)
- else:
- return value
-
- IRBuilder.name(var_name, var)
- return var
-
-
-def is_symbolic_var_declaration(node: doc.expr) -> bool:
- """Return whether an expression is literal ``T.dtype()`` declaration syntax.
-
- Detection is syntactic so declaration semantics do not depend on evaluating its RHS.
- """
- if not (
- isinstance(node, doc.Call)
- and not node.args
- and not node.keywords
- and isinstance(node.func, doc.Attribute)
- and isinstance(node.func.value, doc.Name)
- and node.func.value.id == "T"
- ):
- return False
- constructor = getattr(T, node.func.attr, None)
- return isinstance(constructor, T.DtypeConstructor) or constructor is T.bool
-
-
-def collect_symbolic_var_declaration_nodes(
- target: doc.expr, value: doc.expr
-) -> dict[str, doc.expr]:
- """Pair targets with literal ``T.dtype()`` declarations.
-
- Declaration-vs-binding is per target, allowing tuple unpacking to mix both.
- """
- if isinstance(target, doc.Name):
- return {target.id: value} if is_symbolic_var_declaration(value) else {}
- if isinstance(target, doc.Tuple | doc.List) and isinstance(value, doc.Tuple | doc.List):
- if len(target.elts) != len(value.elts):
- return {}
- declarations = {}
- for lhs, rhs in zip(target.elts, value.elts):
- declarations.update(collect_symbolic_var_declaration_nodes(lhs, rhs))
- return declarations
- return {}
-
-
-def collect_symbolic_var_declarations(target: doc.expr, value: doc.expr) -> set[str]:
- """Collect assignment targets whose matching RHS is ``T.dtype()``."""
- return set(collect_symbolic_var_declaration_nodes(target, value))
-
-
-def eval_ty_proxy(self: Parser, node: doc.expr) -> TypeProxy:
- try:
- annotation = self.eval_expr(node)
- return _normalize_ty_proxy(annotation)
- except Exception as err: # pylint: disable=broad-except
- self.report_error(node, err)
- raise
-
-
-def eval_ty(self: Parser, node: doc.expr, eval_str: bool = False) -> Type:
- var_table = self.var_table.get() if eval_str else None
- try:
- ty = self.eval_expr(node)
- return _normalize_ty(ty, var_table)
- except Exception as err: # pylint: disable=broad-except
- self.report_error(node, err)
- raise
-
-
-def is_called(node: Any, func_name: str) -> bool:
- # Check if it calls into a func
- if isinstance(node, doc.Call):
- # Recursive call was found
- if isinstance(node.func, doc.Name) and node.func.id == func_name:
- return True
- elif isinstance(node, list | tuple):
- for stmt in node:
- if is_called(stmt, func_name):
- return True
- elif isinstance(node, doc.AnnAssign | doc.Assign | doc.Return | doc.Expr):
- return is_called(node.value, func_name)
- elif isinstance(node, doc.With):
- return is_called(node.body, func_name)
- elif isinstance(node, doc.If):
- smts = []
- if node.body is not None:
- smts = smts + list(node.body)
- if node.orelse is not None:
- smts = smts + list(node.orelse)
- return is_called(smts, func_name)
- return False
-
-
-def is_recursive(node: doc.FunctionDef) -> bool:
- # Check if it is a recursive function
- for stmt in node.body:
- if is_called(stmt, node.name):
- return True
- return False
-
-
-def collect_symbolic_var_from_prelude(
- self: Parser, node: doc.FunctionDef, symbolic_vars: dict[str, tvm.ir.Var]
-) -> dict[str, tvm.ir.Var]:
- prelude_vars = {}
- for stmt in node.body:
- if isinstance(stmt, doc.Assign) and len(stmt.targets) == 1:
- declarations = collect_symbolic_var_declaration_nodes(stmt.targets[0], stmt.value)
- for name, value_node in declarations.items():
- if name not in symbolic_vars:
- continue
- declared_var = self.eval_expr(value_node)
- if tvm.ir.is_prim_var(declared_var) and declared_var.ty == symbolic_vars[name].ty:
- prelude_vars[name] = declared_var
-
- return {**symbolic_vars, **prelude_vars}
-
-
-def collect_symbolic_var_from_params(self: Parser, node: doc.FunctionDef) -> None:
- symbolic_vars = collect_signature_type_vars(self, node)
- prim_params = set()
- with self.var_table.with_frame():
- for var_name, var in symbolic_vars.items():
- self.var_table.add(var_name, var)
- for arg in node.args.args:
- if arg.annotation is None:
- self.report_error(arg, "Type annotation is required for function parameters.")
- param_ty_proxy = eval_ty_proxy(self, arg.annotation)
-
- for var_name in param_ty_proxy.get_symbolic_vars():
- if var_name not in prim_params and var_name not in symbolic_vars:
- symbolic_vars[var_name] = tvm.ir.Var(var_name, "int64")
-
- if isinstance(param_ty_proxy, PrimProxy):
- temp_param = tvm.ir.Var(arg.arg, param_ty_proxy.as_ty())
- self.var_table.add(arg.arg, temp_param)
- prim_params.add(arg.arg)
-
- # Prelude declarations select each scope symbol's canonical Var object.
- # Per-assignment declaration classification separately controls emission.
- symbolic_vars = collect_symbolic_var_from_prelude(self, node, symbolic_vars)
-
- # Define symbolic vars to the current var_table frame
- for var_name, var in symbolic_vars.items():
- self.var_table.add(var_name, var, allow_shadowing=False)
-
-
-def parse_function_params(self: Parser, node: doc.FunctionDef, param_factory) -> list[tvm.ir.Var]:
- collect_symbolic_var_from_params(self, node)
- params = []
- for arg in node.args.args:
- if arg.annotation is None:
- self.report_error(arg, "Type annotation is required for function parameters.")
- param_ty = eval_ty(self, arg.annotation, eval_str=True)
- param = param_factory(arg.arg, param_ty)
- self.var_table.add(arg.arg, param)
- params.append(param)
-
- return params
-
-
-@dispatch.register(token="relax", type_name="FunctionDef")
-def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
- is_inner_function = self.inside_function
- self.inside_function = True
-
- # reserve a var for local function
- func_val = self.var_table.get().get(node.name)
- if func_val is None and is_recursive(node):
- with self.var_table.with_frame():
- provisional_params = parse_function_params(self, node, tvm.ir.Var)
- if node.returns is None:
- ret_ty = relax.TupleType([])
- else:
- ret_ty = eval_ty(self, node.returns, eval_str=True)
- params_ty = [param.ty for param in provisional_params]
- # created a var for the local function, the same var could be used for recursive call
- local_func_var = tvm.ir.Var(node.name, relax.FuncType(params_ty, ret_ty))
- self.var_table.add(node.name, local_func_var)
-
- purity = find_decorator_annotation(node, "pure")
- # treat the function as private if we are inside another function
- # or if it has a privacy annotation
- privacy = is_inner_function or find_decorator_annotation(node, "private", default=False)
-
- with self.var_table.with_frame():
- with self.with_dispatch_token("relax"):
- with R.function(is_pure=purity, is_private=privacy):
- R.func_name(node.name)
- parse_function_params(self, node, R.arg)
-
- if node.returns is not None:
- ann_ty = eval_ty(self, node.returns, eval_str=True)
- R.func_ret_ty(ann_ty)
-
- for stmt in node.body:
- if isinstance(stmt, doc.FunctionDef):
- if not stmt.decorator_list:
- self.report_error(stmt, "Function must be decorated")
- dec = self.eval_expr(stmt.decorator_list[-1])
- # inline prim_func was found
- if dec.dispatch_token == "tirx":
- self.report_error(stmt, "inline prim_func is disallowed in Relax IR")
-
- self.visit_body(node.body)
- self.inside_function = is_inner_function
-
-
-def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool:
- """
- Check the value of given annotation (argument name) in the function decorator.
- Returns the value of the annotation if present, otherwise giving the default value.
- """
- # look for the named argument in the function decorator
- for dec in node.decorator_list:
- if not isinstance(dec, doc.Call) or dec.func.attr != "function":
- continue
- for keyword in dec.keywords:
- if keyword.arg == annotation:
- return keyword.value.value
- return default
-
-
-@dispatch.register(token="relax", type_name="tvm_declare_function")
-def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
- with self.var_table.with_frame():
- params = parse_function_params(self, node, tvm.ir.Var)
-
- if node.returns is None:
- # Use AnyType as unknown return type
- # NOTE: Cannot use VoidType here because the return type can be refined later.
- ret_ty = relax.AnyType()
- else:
- ret_ty = eval_ty(self, node.returns, eval_str=True)
- is_pure = find_decorator_annotation(node, "pure")
-
- func_signature = relax.Function.create_empty(params, ret_ty, is_pure=is_pure)
- return I.decl_function(node.name, func_signature)
-
-
-@dispatch.register(token="relax", type_name="pre_visit_local_function")
-def pre_visit_local_function(self: Parser, node: doc.Expr) -> None:
- ir_builder = IRBuilder()
- ir_builder.__enter__()
-
-
-@dispatch.register(token="relax", type_name="post_visit_local_function")
-def post_visit_local_function(self: Parser, node: doc.Expr) -> None:
- ir_builder = IRBuilder.current()
- result = ir_builder.get()
- ir_builder.__exit__(None, None, None)
- # reuse var if it is reserved
- reserved_var = self.var_table.get().get(node.name)
- if reserved_var is not None:
- var = R.emit_var_binding(relax.VarBinding(reserved_var, result))
- else:
- var = R.emit(result)
- IRBuilder.name(node.name, var)
- self.var_table.add(node.name, var, allow_shadowing=False)
-
-
-@dispatch.register(token="relax", type_name="Expr")
-def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
- value = self.eval_expr(node.value)
- if isinstance(value, relax.Expr):
- var = R.emit(value)
- IRBuilder.name("_", var)
- is_void_value = isinstance(var.ty, relax.TupleType) and len(var.ty.fields) == 0
-
- if not is_void_value:
- self.report_error(
- node,
- f"Non-void relax expressions must be bound to a variable, "
- f"but expression of type {var.ty} was used as a statement.",
- )
-
- elif value is not None:
- self.report_error(node, f"Unsupported Expr stmt type {value}.")
-
-
-@dispatch.register(token="relax", type_name="arguments")
-def visit_arguments(self: Parser, node: doc.arguments) -> None:
- arg: doc.arg
- for arg in node.args:
- if arg.annotation is None:
- self.report_error(arg, "Type annotation is required for function parameters.")
- param_ty = eval_ty(self, arg.annotation, eval_str=True)
- param = R.arg(arg.arg, param_ty)
-
- self.var_table.add(arg.arg, param)
-
-
-@dispatch.register(token="relax", type_name="tvm_annotation")
-def visit_tvm_annotation(self: Parser, node: doc.expr) -> Type:
- return eval_ty(self, node, eval_str=False)
-
-
-@dispatch.register(token="relax", type_name="With")
-def visit_with(self: Parser, node: doc.With) -> None:
- # Currently only `with R.dataflow()` is supported
- if len(node.items) != 1:
- self.report_error(node, "Only one item is allowed.")
- item = node.items[0]
- if item.optional_vars is not None:
- self.report_error(
- item.context_expr,
- "Relax syntax doesn't allow binding expressions in `with` to variables",
- )
- frame = self.eval_expr(item.context_expr)
- with self.var_table.with_frame():
- with frame:
- self.visit(node.body)
- if isinstance(frame, BindingBlockFrame) and frame.is_dataflow:
- output_vars = frame.output_vars
- for var in output_vars:
- self.var_table.add(var.name, var, allow_shadowing=True)
-
-
-@dispatch.register(token="relax", type_name="Assign")
-def visit_assign(self: Parser, node: doc.Assign) -> None:
- if len(node.targets) != 1:
- self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
- lhs = node.targets[0]
- rhs = self.eval_expr(node.value)
- self.eval_assign(
- target=lhs,
- source=rhs,
- bind_value=functools.partial(
- bind_assign_value,
- symbolic_declarations=collect_symbolic_var_declarations(lhs, node.value),
- ),
- allow_shadowing=True,
- )
-
-
-@dispatch.register(token="relax", type_name="AnnAssign")
-def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
- lhs = node.target
- rhs = self.eval_expr(node.value)
- anno_ty = self.visit_tvm_annotation(node.annotation)
- self.eval_assign(
- target=lhs,
- source=rhs,
- bind_value=functools.partial(
- bind_assign_value,
- anno_ty=anno_ty,
- symbolic_declarations=collect_symbolic_var_declarations(lhs, node.value),
- ),
- allow_shadowing=True,
- )
-
-
-@dispatch.register(token="relax", type_name="Return")
-def visit_return(self: Parser, node: doc.Assign) -> None:
- value = self.eval_expr(node.value)
- value = convert_to_expr(value)
- R.func_ret_value(value)
-
-
-@dispatch.register(token="relax", type_name="If")
-def visit_if(self: Parser, node: doc.If) -> None:
- if node.orelse is None:
- raise ValueError("Else statements are required for relax dialect.")
- with R.If(self.eval_expr(node.test)) as if_frame:
- with self.var_table.with_frame():
- with R.Then():
- self.visit_body(node.body)
- with self.var_table.with_frame():
- with R.Else():
- self.visit_body(node.orelse)
- self.var_table.add(if_frame.var_name, if_frame.var, allow_shadowing=True)
-
-
-@dispatch.register(token="relax", type_name="enter_token")
-def enter_token(self: Parser) -> dict[str, Any]:
- def relax_call(self, *args) -> Expr:
- args = [convert_to_expr(arg) if isinstance(arg, tuple) else arg for arg in args]
-
- if all(isinstance(x, Expr) for x in args):
- return relax.Call(self, args)
- arg_types = [type(x) for x in args]
- raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for types {arg_types}")
-
- context = {"GlobalVar.__call__": GlobalVar.__call__}
- GlobalVar.__call__ = relax_call
- return context
-
-
-@dispatch.register(token="relax", type_name="exit_token")
-def exit_token(self: Parser, context: dict[str, Any]) -> None:
- assert "GlobalVar.__call__" in context
- GlobalVar.__call__ = context.get("GlobalVar.__call__")
diff --git a/python/tvm/relax/script/v2.py b/python/tvm/relax/script/v2.py
deleted file mode 100644
index 1b627f04b5..0000000000
--- a/python/tvm/relax/script/v2.py
+++ /dev/null
@@ -1,58 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Opt-in TVMScript entry point using concrete Relax construction operations."""
-
-# pylint: disable=wildcard-import,unused-wildcard-import,redefined-builtin
-import sys as _sys
-
-from tvm import relax as _relax
-from tvm.relax.base_py_module import PyModuleFactory as _PyModuleFactory
-from tvm.script.parser_v2.frontend import make_decorator as _make_decorator
-from tvm.script.parser_v2.frontend import make_helper as _make_helper
-from tvm.script.parser_v2.frontend import register_namespace as _register_namespace
-from tvm.script.parser_v2.functions import register_opaque_factory as _register_opaque_factory
-
-from . import builder_v2 as _builder
-from .builder_v2 import * # noqa: F403
-
-function = _make_decorator(_builder, option_map={"pure": "is_pure", "private": "is_private"})
-macro = _make_helper(_builder, preserve_return=True)
-
-_register_namespace("R", _sys.modules[__name__])
-_register_namespace("relax", _sys.modules[__name__])
-
-
-def _opaque_function(name, function, source, span):
- return _relax.ExternFunc(name, span=span).with_attrs(
- {
- "is_pyfunc": True,
- "function_type": "python",
- "python_function_name": name,
- "python_source": source,
- "python_packed_func": function,
- }
- )
-
-
-def _python_module(module, original, bases):
- candidates = (original, *bases)
- if any(isinstance(base, type) and issubclass(base, _relax.BasePyModule) for base in candidates):
- return _PyModuleFactory(module, original_class=original)
- return module
-
-
-_register_opaque_factory(_opaque_function, module_adapter=_python_module)
diff --git a/python/tvm/script/__init__.py b/python/tvm/script/__init__.py
index bdf9b5378b..b198a0eb10 100644
--- a/python/tvm/script/__init__.py
+++ b/python/tvm/script/__init__.py
@@ -17,9 +17,9 @@
"""
TVMScript public namespace.
-Construction APIs return concrete IR values. Use
-``from __future__ import annotations`` for Python-defined script functions so
-signature symbols are resolved by the parser before constructors run.
+Construction APIs return concrete IR values. Python-defined script functions
+support both eager and postponed annotations. The shared parser resolves
+signature symbols through the registered construction protocol.
Dialect resolution mechanism
----------------------------
@@ -38,23 +38,24 @@ Dialect resolution mechanism
and cached as a normal module attribute. Subsequent accesses
skip ``__getattr__`` (cached in ``globals()``).
-- Subpackages ``tvm.script.parser``, ``tvm.script.ir_builder``, etc.
- each define their own ``__getattr__`` that consults the SAME
- ``_DIALECT_REGISTRY`` and appends their suffix. So
- ``tvm.script.parser.tirx`` resolves to ``tvm.tirx.script.parser`` via
- the dialect registry + ``.parser`` suffix.
+- Each dialect's script package exports its decorators and construction
+ operations, and registers source aliases with ``tvm.script.parser``.
+ The parser consumes those registrations without importing dialect packages.
+ Existing imports such as ``from tvm.script.parser import tirx`` resolve to
+ the registered script namespace. ``tvm.script.ir_builder`` separately uses
+ ``_DIALECT_REGISTRY`` to resolve a dialect's ``builder`` package.
-- For deep statement-form imports like
- ``from tvm.script.parser.tirx.entry import ObjectProxy``, PEP 562's
+- For statement-form imports like
+ ``import tvm.script.ir_builder.relax.ir``, PEP 562's
``__getattr__`` is not enough — it only handles one-level
``from X import Y``. A ``sys.meta_path`` finder (see
``_DialectRedirectFinder``) intercepts the import machinery to
register the real module under the legacy name in ``sys.modules``,
so subsequent attribute walks resolve correctly.
-Each dialect's ``tvm.<dialect>.script`` package MUST expose ``parser``,
-``ir_builder``, and (where applicable) ``printer`` as submodules. This
-convention is what makes the suffix-append redirect work uniformly.
+Each dialect's ``tvm.<dialect>.script`` package owns its script entry points
+and a ``builder`` subpackage. ``tvm.script.parser.<dialect>`` aliases the
+script package itself; ``tvm.script.ir_builder.<dialect>`` aliases its builder.
IR is foundational (script depends on ir) and is NOT a dialect; its
script handlers live in the shared core, not via this registry.
@@ -77,10 +78,10 @@ from typing import Any
_DIALECT_REGISTRY: dict[str, str] = {}
# Subpackages of `tvm.script` whose per-dialect children are redirected to a
-# matching subpackage under `tvm.<dialect>.script`. The values are the
-# subpackage name on the dialect side (e.g. `tvm.<dialect>.script.parser`).
+# matching package under `tvm.<dialect>.script`. An empty suffix selects the
+# script entry package itself; other values select its named subpackage.
_REDIRECTED_SUBPACKAGES = {
- "tvm.script.parser": "parser",
+ "tvm.script.parser": "",
"tvm.script.ir_builder": "builder",
}
@@ -90,11 +91,11 @@ def register_dialect(name: str, module_path: str) -> None:
Writes ``name -> module_path`` into ``_DIALECT_REGISTRY``. After
registration, ``tvm.script.<name>`` resolves to ``module_path`` via
- ``__getattr__``, and ``tvm.script.parser.<name>`` / ``tvm.script.ir_builder.<name>``
- resolve to ``module_path + ".parser"`` / ``module_path + ".builder"`` etc.
- via each subpackage's own ``__getattr__``. Deep statement-form imports
- (e.g., ``from tvm.script.parser.<name>.entry import X``) are handled
- by ``_DialectRedirectFinder`` on ``sys.meta_path``.
+ ``__getattr__``. Statement-form imports of ``tvm.script.parser.<name>``
+ and ``tvm.script.ir_builder.<name>`` resolve to ``module_path`` and
+ ``module_path + ".builder"`` through ``_DialectRedirectFinder``.
+ The script entry package separately registers its source aliases and
+ construction policies with the shared parser.
This function is idempotent — re-registering the same name with the same
path is harmless.
@@ -113,9 +114,8 @@ def register_dialect(name: str, module_path: str) -> None:
The short name exposed under ``tvm.script.<name>`` (e.g. ``"tirx"``).
module_path : str
The full dotted module path of the dialect's script package, e.g.
- ``"tvm.tirx.script"``. That package must expose ``parser`` and
- ``ir_builder`` as submodules (and ``printer`` where applicable) so
- that the suffix-append redirect works uniformly.
+ ``"tvm.tirx.script"``. That package exports the script entry points
+ and exposes its construction operations through a ``builder`` subpackage.
"""
_DIALECT_REGISTRY[name] = module_path
@@ -139,7 +139,9 @@ def _redirect_target(fullname: str) -> str | None:
rest = fullname[len(prefix) + 1 :]
head, _, tail = rest.partition(".")
if head in _DIALECT_REGISTRY:
- target = f"{_DIALECT_REGISTRY[head]}.{sub}"
+ target = _DIALECT_REGISTRY[head]
+ if sub:
+ target = f"{target}.{sub}"
return f"{target}.{tail}" if tail else target
return None
@@ -151,14 +153,14 @@ class _DialectRedirectFinder:
(``from tvm.script import tirx``). It cannot intercept deep
statement-form imports such as::
- from tvm.script.parser.tirx.entry import ObjectProxy
+ import tvm.script.parser.tirx
import tvm.script.ir_builder.relax.ir
This finder is installed on ``sys.meta_path`` to cover those cases.
When the import machinery asks for a module whose full name starts with
``tvm.script.<dialect>`` (or ``tvm.script.parser.<dialect>``, etc.) and
that dialect is in ``_DIALECT_REGISTRY``, :meth:`find_spec` imports the
- real target module (e.g. ``tvm.tirx.script.parser.entry``) and returns an
+ real target module (e.g. ``tvm.tirx.script``) and returns an
alias spec whose loader hands back that module, so the import machinery
registers it in ``sys.modules`` under the legacy name and all subsequent
imports and attribute walks resolve without going through the redirect
@@ -220,18 +222,18 @@ def __getattr__(name: str) -> Any:
return module
if name == "ir":
# IR is foundational and shares the generic parser entry point.
- ir_parser = importlib.import_module("tvm.script.parser_v2.ir")
+ ir_parser = importlib.import_module("tvm.script.parser.ir")
globals()["ir"] = ir_parser
return ir_parser
if name in ("from_source", "parse"):
- from .parser_v2 import parse # pylint: disable=import-outside-toplevel
+ from .parser import parse # pylint: disable=import-outside-toplevel
globals()["from_source"] = parse
globals()["parse"] = parse
return parse
if name == "ir_module":
# The foundational module entry does not use the dialect registry.
- ir_parser = importlib.import_module("tvm.script.parser_v2.ir")
+ ir_parser = importlib.import_module("tvm.script.parser.ir")
ir_module_value = ir_parser.ir_module
globals()["ir_module"] = ir_module_value
return ir_module_value
diff --git a/python/tvm/script/parser/__init__.py b/python/tvm/script/parser/__init__.py
index f4fa263211..787d6d82a2 100644
--- a/python/tvm/script/parser/__init__.py
+++ b/python/tvm/script/parser/__init__.py
@@ -14,24 +14,31 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Default TVMScript parser entries.
+"""Translate original Python ASTs into calls on a registered construction namespace.
-The generic implementation is retained under ``parser_v2`` until the migration
-cleanup. Dialect namespaces resolve through the shared script registry.
+Shared I infrastructure owns builder lifetime, spans and module identities; each
+context selects X from reverse-registered function metadata. Calls return concrete
+values, constructor metadata describes syntax, and compiled AST locations remain
+those of the original source. Entry modules register policies; this parser never
+imports their namespaces.
"""
-import importlib
-from typing import Any
+from . import ir
+from .frontend import (
+ _NAMESPACES,
+ from_source,
+ ir_module,
+ make_decorator,
+ make_helper,
+ parse,
+ pyfunc,
+ register_namespace,
+)
-from ..parser_v2 import ir, ir_module, parse
-
-def __getattr__(name: str) -> Any:
- # Lazy import to avoid loading tvm.script during dialect bootstrap.
- from tvm.script import _DIALECT_REGISTRY # pylint: disable=import-outside-toplevel
-
- if name in _DIALECT_REGISTRY:
- module = importlib.import_module(_DIALECT_REGISTRY[name])
- globals()[name] = module
- return module
- raise AttributeError(f"module 'tvm.script.parser' has no attribute {name!r}")
+def __getattr__(name):
+ # Entry modules register source aliases; parsing never imports their owners.
+ try:
+ return _NAMESPACES[name]
+ except KeyError:
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
diff --git a/python/tvm/script/parser/_core.py b/python/tvm/script/parser/_core.py
deleted file mode 100644
index d87214f22f..0000000000
--- a/python/tvm/script/parser/_core.py
+++ /dev/null
@@ -1,24 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the Licens.
-# ruff: noqa: F401
-"""The core parser infra"""
-
-# pylint: disable=unused-import
-from .core import dispatch, doc, utils
-from .core.dispatch import OpMethod, register_op
-from .core.entry import parse, scan_macro
-from .core.parser import Parser, collect_signature_type_vars
diff --git a/python/tvm/script/parser_v2/annotations.py b/python/tvm/script/parser/annotations.py
similarity index 100%
rename from python/tvm/script/parser_v2/annotations.py
rename to python/tvm/script/parser/annotations.py
diff --git a/python/tvm/script/parser/core/__init__.py b/python/tvm/script/parser/core/__init__.py
deleted file mode 100644
index 01bcd86ae4..0000000000
--- a/python/tvm/script/parser/core/__init__.py
+++ /dev/null
@@ -1,20 +0,0 @@
-# isort: skip_file
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""The core parser infra"""
-
-from . import diagnostics, dispatch, doc, doc_core, entry, evaluator, parser, utils
diff --git a/python/tvm/script/parser/core/diagnostics.py b/python/tvm/script/parser/core/diagnostics.py
deleted file mode 100644
index 974c5c593f..0000000000
--- a/python/tvm/script/parser/core/diagnostics.py
+++ /dev/null
@@ -1,324 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# ruff: noqa: E741
-"""TVM Script Parser Source and diagnostics"""
-
-import inspect
-import sys
-
-from tvm.error import DiagnosticError
-from tvm.ir import SourceName, Span
-
-from . import doc
-
-
-class Source:
- """Source code class for TVMScript.
-
- It is constructed by source code str or doc AST tree.
-
- Parameters
- ----------
- source_name : str
- The filename of the file where the source code locates.
-
- start_line : int
- The first line number of the source code.
-
- start_column : int
- The first column number of the first line of the source code.
-
- source : str
- The source code str of source code.
-
- full_source : str
- The complete source code of the file where the source code locates.
- """
-
- source_name: str
- start_line: int
- start_column: int
- source: str
- full_source: str
-
- def __init__(self, program: str | doc.AST):
- if isinstance(program, str):
- self.source_name = "<str>"
- self.start_line = 1
- self.start_column = 0
- self.source = program
- self.full_source = program
- return
-
- self.source_name = inspect.getsourcefile(program) # type: ignore
- lines, self.start_line = getsourcelines(program) # type: ignore
- if lines:
- self.start_column = len(lines[0]) - len(lines[0].lstrip())
- else:
- self.start_column = 0
- if self.start_column and lines:
- self.source = "\n".join([l[self.start_column :].rstrip() for l in lines])
- else:
- self.source = "".join(lines)
- try:
- # It will cause a problem when running in Jupyter Notebook.
- # `mod` will be <module '__main__'>, which is a built-in module
- # and `getsource` will throw a TypeError
- mod = inspect.getmodule(program)
- if mod:
- self.full_source = inspect.getsource(mod)
- else:
- self.full_source = self.source
- except TypeError:
- # It's a work around for Jupyter problem.
- # Since `findsource` is an internal API of inspect, we just use it
- # as a fallback method.
- src, _ = inspect.findsource(program) # type: ignore
- self.full_source = "".join(src)
-
- def as_ast(self) -> doc.AST:
- """Parse the source code into AST.
-
- Returns
- -------
- res : doc.AST
- The AST of source code.
- """
- return doc.parse(self.source)
-
- def location(self, node: doc.AST) -> tuple[int, int, int, int]:
- """Return the absolute 1-based source range of an AST node."""
- lineno = getattr(node, "lineno", 1) or 1
- col_offset = getattr(node, "col_offset", self.start_column)
- col_offset = self.start_column if col_offset is None else col_offset
- end_lineno = getattr(node, "end_lineno", lineno) or lineno
- end_col_offset = getattr(node, "end_col_offset", col_offset)
- end_col_offset = col_offset if end_col_offset is None else end_col_offset
- lineno += self.start_line - 1
- end_lineno += self.start_line - 1
- col_offset += self.start_column + 1
- end_col_offset += self.start_column + 1
- return lineno, col_offset, end_lineno, end_col_offset
-
- def to_span(self, node: doc.AST) -> Span:
- """Convert an AST node to the canonical IR source span."""
- lineno, col_offset, end_lineno, end_col_offset = self.location(node)
- return Span(
- SourceName(self.source_name or "<unknown>"),
- lineno,
- end_lineno,
- col_offset,
- end_col_offset,
- )
-
-
-_getfile = inspect.getfile # pylint: disable=invalid-name
-_findsource = inspect.findsource # pylint: disable=invalid-name
-
-
-def _patched_inspect_getfile(obj):
- """Work out which source or compiled file an object was defined in."""
- if not inspect.isclass(obj):
- return _getfile(obj)
- mod = getattr(obj, "__module__", None)
- if mod is not None:
- file = getattr(sys.modules[mod], "__file__", None)
- if file is not None:
- return file
- for _, member in inspect.getmembers(obj):
- if inspect.isfunction(member):
- if obj.__qualname__ + "." + member.__name__ == member.__qualname__:
- return inspect.getfile(member)
- raise TypeError(f"Source for {obj!r} not found")
-
-
-def findsource(obj):
- """Return the entire source file and starting line number for an object."""
- import linecache # pylint: disable=import-outside-toplevel
-
- if not inspect.isclass(obj):
- return _findsource(obj)
-
- file = inspect.getsourcefile(obj)
- if file:
- linecache.checkcache(file)
- else:
- file = inspect.getfile(obj)
- if not (file.startswith("<") and file.endswith(">")):
- raise OSError("source code not available")
-
- module = inspect.getmodule(obj, file)
- if module:
- lines = linecache.getlines(file, module.__dict__)
- else:
- lines = linecache.getlines(file)
- if not lines:
- raise OSError("could not get source code")
- qual_names = obj.__qualname__.replace(".<locals>", "<locals>").split(".")
- in_comment = 0
- scope_stack = []
- indent_info = {}
- for i, line in enumerate(lines):
- n_comment = line.count('"""')
- if n_comment:
- # update multi-line comments status
- in_comment = in_comment ^ (n_comment & 1)
- continue
- if in_comment:
- # skip lines within multi-line comments
- continue
- indent = len(line) - len(line.lstrip())
- tokens = line.split()
- if len(tokens) > 1:
- name = None
- if tokens[0] == "def":
- name = tokens[1].split(":")[0].split("(")[0].split("[")[0] + "<locals>"
- elif tokens[0] == "class":
- name = tokens[1].split(":")[0].split("(")[0]
- # pop scope if we are less indented
- while scope_stack and indent_info[scope_stack[-1]] >= indent:
- scope_stack.pop()
- if name:
- scope_stack.append(name)
- indent_info[name] = indent
- if scope_stack == qual_names:
- return lines, i
-
- raise OSError("could not find class definition")
-
-
-def getsourcelines(obj):
- """Extract the block of code at the top of the given list of lines."""
- obj = inspect.unwrap(obj)
- lines, l_num = findsource(obj)
- return inspect.getblock(lines[l_num:]), l_num + 1
-
-
-inspect.getfile = _patched_inspect_getfile
-
-
-def _format_source_snippet(
- source_lines: list,
- lineno: int,
- col_offset: int,
- end_lineno: int,
- end_col_offset: int,
-) -> str:
- """Format a source code snippet with a column/span marker.
-
- Renders every source line spanned by the diagnostic (``lineno`` through
- ``end_lineno``, inclusive) with a per-line gutter, followed by a caret
- underline covering the offending span. For a single-line span
- (``end_lineno == lineno``) only the columns ``col_offset`` (inclusive)
- through ``end_col_offset`` (exclusive) are underlined; for a multi-line
- span the underline covers the start column to end-of-line on the first
- line, the full text of interior lines, and the start of the final line up
- to ``end_col_offset``.
-
- Parameters
- ----------
- source_lines : list of str
- Lines of the source code.
-
- lineno : int
- 1-based starting line number in the source.
-
- col_offset : int
- 1-based starting column (inclusive) on ``lineno``.
-
- end_lineno : int
- 1-based ending line number in the source (>= ``lineno``).
-
- end_col_offset : int
- 1-based ending column (exclusive) on ``end_lineno``.
-
- Returns
- -------
- snippet : str
- Formatted source snippet with caret-marker line(s).
- """
- if end_lineno < lineno:
- end_lineno = lineno
-
- # Determine the gutter width so that all line numbers line up.
- header_width = len(f" {end_lineno} ")
- no_line_header = " " * header_width
-
- parts = [f"{no_line_header}| "]
- for cur_lineno in range(lineno, end_lineno + 1):
- idx = cur_lineno - 1
- if not 0 <= idx < len(source_lines):
- continue
- line_text = source_lines[idx].rstrip("\n")
- line_header = f" {cur_lineno} ".rjust(header_width)
-
- # Compute the underline span [start_col, stop_col) for this line.
- start_col = col_offset if cur_lineno == lineno else 1
- stop_col = end_col_offset if cur_lineno == end_lineno else len(line_text) + 1
-
- marker = ""
- for i in range(1, len(line_text) + 1):
- if start_col <= i < stop_col:
- marker += "^"
- else:
- marker += " "
- parts.append(f"{line_header}| {line_text}")
- parts.append(f"{no_line_header}| {marker}")
-
- if len(parts) == 1:
- return ""
- return "\n".join(parts)
-
-
-class Diagnostics:
- """Diagnostics class for error reporting in parser.
-
- Formats parse errors with source location context and raises directly,
- without going through DiagnosticContext.
-
- Parameters
- ----------
- source : Source
- The source code.
- """
-
- source: Source
-
- def __init__(self, source: Source):
- self.source = source
-
- def error(self, node: doc.AST, message: str) -> None:
- """Emit a diagnostic error by raising with source location context.
-
- Parameters
- ----------
- node : doc.AST
- The node with diagnostic error.
-
- message : str
- The diagnostic message.
- """
- lineno, col_offset, end_lineno, end_col_offset = self.source.location(node)
-
- source_lines = self.source.full_source.splitlines(keepends=True)
- snippet = _format_source_snippet(
- source_lines, lineno, col_offset, end_lineno, end_col_offset
- )
-
- location = f"{self.source.source_name}:{lineno}:{col_offset}"
- formatted = f"error: {message}\n --> {location}\n{snippet}"
- raise DiagnosticError(formatted)
diff --git a/python/tvm/script/parser/core/dispatch.py b/python/tvm/script/parser/core/dispatch.py
deleted file mode 100644
index 5f880be833..0000000000
--- a/python/tvm/script/parser/core/dispatch.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Parser dispatching infrastructure"""
-
-from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
-
-from .doc import AST
-
-if TYPE_CHECKING:
- from .parser import Parser
-
-
-ParseMethod = Callable[["Parser", AST], None]
-ParseVTable: dict[tuple[str, str], ParseMethod] = {}
-
-OpMethod = Callable[..., Any]
-OpVTable: dict[tuple[type, AST, int], OpMethod] = {}
-
-
-def register(token: str, type_name: str):
- """Register a method for a dispatch token and type name.
-
- Parameters
- ----------
- token : str
- The token for IR, e.g., T for TIR and R for Relax.
-
- type_name : str
- The type name of AST node, e.g., FunctionDef, With, For.
-
- Returns
- -------
- func : callable
- The function to register dispatched method of parsing
- corresponding token and AST node type.
- """
-
- def func(method: ParseMethod):
- """Register a method in parser virtual table.
-
- Parameters
- ----------
- method : ParseMethod
- The dispatched method to be registered in parser virtual table.
- """
- ParseVTable[(token, type_name)] = method
-
- return func
-
-
-def get(
- token: str,
- type_name: str,
- default: ParseMethod | None = None,
-) -> ParseMethod | None:
- """Get a registered method for a dispatch token and type name,
- or return a default method if no registered methods with this dispatch token and type name.
-
- Parameters
- ----------
- token : str
- The token for IR, e.g., T for TIR and R for Relax.
-
- type_name : str
- The type name of AST node, e.g., FunctionDef, With, For.
-
- default : Optional[ParseMethod]
- The default method when no registered methods with this dispatch token and type name.
-
- Returns
- -------
- func : Optional[ParseMethod]
- The dispatched method of parsing corresponding token and AST node type.
- """
- return ParseVTable.get((token, type_name), default)
-
-
-def register_op(operand_type: type, op_node_type: AST, operand_index: int):
- """Register a method for a operand type, AST operator node and operand index.
-
- Parameters
- ----------
- operand_type : Type
- The type of operands, e.g., tirx.Expr, tirx.IterVar.
-
- op_node_type : AST
- The doc AST operator node type, e.g., doc.Add, doc.Eq.
-
- operand_index : int
- The operand index, i.e., 0 for left operand and 1 for right operand.
-
- Returns
- -------
- func : callable
- The function to register dispatched method of parsing
- corresponding a operand type, AST operator node and operand index.
- """
-
- def func(method: OpMethod):
- """Register a method in parser operator virtual table.
-
- Parameters
- ----------
- method : ParseMethod
- The dispatched method to be registered in parser operator virtual table.
- """
- OpVTable[(operand_type, op_node_type, operand_index)] = method
-
- return func
-
-
-def get_op(
- operand_type: type,
- op_node_type: type,
- operand_index: int,
- default: OpMethod | None = None,
-) -> OpMethod | None:
- """Register a method for a operand type, AST operator node and operand index.
-
- Parameters
- ----------
- operand_type : Type
- The type of operands, e.g., tirx.Expr, tirx.IterVar.
-
- op_node_type : AST
- The doc AST operator node type, e.g., doc.Add, doc.Eq.
-
- operand_index : int
- The operand index, i.e., 0 for left operand and 1 for right operand.
-
-
- default : Optional[OpMethod]
- The default method when no registered methods with this operand type,
- AST operator node and operand index.
-
- Returns
- -------
- func : Optional[OpMethod]
- The function to register dispatched method of parsing
- corresponding a operand type, AST operator node and operand index.
- """
- return OpVTable.get((operand_type, op_node_type, operand_index), default)
diff --git a/python/tvm/script/parser/core/doc.py b/python/tvm/script/parser/core/doc.py
deleted file mode 100644
index 7d44dd50d6..0000000000
--- a/python/tvm/script/parser/core/doc.py
+++ /dev/null
@@ -1,320 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# ruff: noqa: E722, F403
-"""TVM Script Parser doc AST"""
-
-import ast
-import inspect
-import typing
-from collections import defaultdict
-
-from . import doc_core as doc
-from .doc_core import * # pylint: disable=unused-import,wildcard-import,redefined-builtin,W0614
-
-FnToDoc = typing.Callable[[ast.AST], doc.AST]
-FnFromDoc = typing.Callable[[doc.AST], ast.AST]
-
-
-class Entry:
- """Mapping entry between python AST node type str and doc AST.
-
- Parameters
- ----------
- to_doc : typing.Optional[FnToDoc]
- The callable methods for converting python AST node to doc AST.
-
- from_doc : typing.Optional[FnFromDoc]
- The callable methods for converting doc AST to python AST node.
- """
-
- to_doc: FnToDoc | None
- from_doc: FnFromDoc | None
-
- def __init__(self):
- self.to_doc = None
- self.from_doc = None
-
-
-class Registry:
- """Registration map from python AST node type str to methods of conversion
- between python AST node and doc AST node.
-
- Parameters
- ----------
- _inst : typing.Optional["Registry"]
- The instance of Registry.
-
- table : typing.Dict[str, Entry]
- The registration map from python AST node type str to methods of conversion
- between python AST node and doc AST node.
- """
-
- _inst: typing.Optional["Registry"] = None
- table: dict[str, Entry]
-
- def __init__(self):
- self.table = defaultdict(Entry)
-
-
-def register_to_doc(name: str):
- """Register the to_doc method for python AST node type.
-
- Parameters
- ----------
- name : str
- The type of python AST node.
-
- Returns
- -------
- f : Callable[[FnToDoc], None]
- The function of registering the to_doc method for python AST node type.
- """
-
- def f(to_doc: FnToDoc): # pylint: disable=redefined-outer-name
- reg = Registry._inst # pylint: disable=protected-access
- reg.table[name].to_doc = to_doc
-
- return f
-
-
-def register_from_doc(name: str):
- """Register the from_doc method for python AST node type.
-
- Parameters
- ----------
- name : str
- The type of python AST node.
-
- Returns
- -------
- f : Callable[[FnFromDoc], None]
- The function of registering the from_doc method for python AST node type.
- """
-
- def f(to_doc: FnFromDoc): # pylint: disable=redefined-outer-name
- reg = Registry._inst # pylint: disable=protected-access
- reg.table[name].from_doc = to_doc
-
- return f
-
-
-def _is_atomic_type(node):
- return (
- node is None
- or node in [..., True, False]
- or isinstance(
- node,
- int | float | str | bool | bytes | complex,
- )
- )
-
-
-def _get_registry_entry(cls_name, attr):
- cls_name = cls_name.split(".")[-1]
- reg = Registry._inst # pylint: disable=protected-access
- if cls_name in reg.table:
- entry = reg.table[cls_name]
- return getattr(entry, attr, None)
- return None
-
-
-def from_doc(node):
- """Get original python AST node from doc AST node.
-
- Parameters
- ----------
- node : doc.AST
- The doc AST node.
-
- Returns
- -------
- res : ast.AST
- The corresponding AST node.
- """
- if _is_atomic_type(node):
- return node
- if isinstance(node, tuple):
- return tuple(from_doc(n) for n in node)
- if isinstance(node, list):
- return [from_doc(n) for n in node]
- func = _get_registry_entry(node.__class__.__name__, "from_doc")
- if not func:
- raise NotImplementedError(f"from_doc is not implemented for: {node.__class__.__name__}")
- return func(node)
-
-
-def to_doc(node):
- """Get doc AST node from python AST node.
-
- Parameters
- ----------
- node : ast.AST
- The AST node.
-
- Returns
- -------
- res : doc.AST
- The corresponding doc AST node.
- """
- if _is_atomic_type(node):
- return node
- if isinstance(node, tuple):
- return tuple(to_doc(n) for n in node)
- if isinstance(node, list):
- return [to_doc(n) for n in node]
- func = _get_registry_entry(node.__class__.__name__, "to_doc")
- if not func:
- raise NotImplementedError(f"to_doc is not implemented for: {node.__class__.__name__}")
- return func(node)
-
-
-def parse(
- source: str,
- filename: str = "<unknown>",
- mode: str = "exec",
-) -> doc.AST:
- """Parse TVMScript source code str to doc AST.
-
- Its interface is consistent with python built-in ast.parse.
- And it will parse by python 3.8 first if possible,
- or it will parse with python version in current environment.
-
- Parameters
- ----------
- source : str
- The TVMScript source code.
-
- filename : str
- The optional filename of the file where source code locates.
-
- mode : str
- The parsing mode for ast.parse.
-
- Returns
- -------
- res : doc.AST
- The parsed doc AST.
- """
- try:
- program = ast.parse( # pylint: disable=unexpected-keyword-arg
- source=source,
- filename=filename,
- mode=mode,
- feature_version=(3, 8),
- )
- except: # pylint: disable=bare-except
- program = ast.parse(
- source=source,
- filename=filename,
- mode=mode,
- )
- return to_doc(program)
-
-
-class NodeVisitor:
- """Node visitor for doc AST"""
-
- def visit(self, node: doc.AST) -> None:
- if isinstance(node, list | tuple):
- for item in node:
- self.visit(item)
- return
- if not isinstance(node, doc.AST):
- return
- getattr(
- self,
- "visit_" + node.__class__.__name__.split(".")[-1],
- self.generic_visit,
- )(node)
-
- def generic_visit(self, node: doc.AST) -> None:
- for field in node.__class__._FIELDS: # pylint: disable=protected-access
- value = getattr(node, field, None)
- if value is None:
- pass
- elif isinstance(value, doc.AST | list | tuple):
- self.visit(value)
-
-
-class NodeTransformer:
- """Node transformer for doc AST"""
-
- def visit(self, node: doc.AST) -> doc.AST:
- if isinstance(node, list):
- return [self.visit(item) for item in node]
- if isinstance(node, tuple):
- return tuple(self.visit(item) for item in node)
- if not isinstance(node, doc.AST):
- return node
- return getattr(
- self,
- "visit_" + node.__class__.__name__.split(".")[-1],
- self.generic_visit,
- )(node)
-
- def generic_visit(self, node: doc.AST) -> doc.AST:
- kv: dict[str, typing.Any] = {}
- for field in node.__class__._FIELDS: # pylint: disable=protected-access
- value = getattr(node, field, None)
- if value is None:
- pass
- elif isinstance(value, doc.AST | list | tuple):
- value = self.visit(value)
- kv[field] = value
- return node.__class__(**kv)
-
-
-def _register_default():
- class DefaultTranslator:
- def __init__(self, doc_cls, func, fields):
- self.doc_cls = doc_cls # getattr(doc, name)
- self.func = func
- self.fields = fields
-
- def __call__(self, node):
- kv = {attr: self.func(getattr(node, attr, None)) for attr in self.fields}
- return self.doc_cls(**kv)
-
- Registry._inst = Registry() # pylint: disable=protected-access
- for cls_name in dir(doc):
- doc_cls = getattr(doc, cls_name)
- if not hasattr(ast, cls_name):
- continue
- if inspect.isclass(doc_cls) and issubclass(doc_cls, doc.AST):
- assert "." not in cls_name
- register_to_doc(cls_name)(
- DefaultTranslator(
- getattr(doc, cls_name),
- to_doc,
- doc_cls._FIELDS, # pylint: disable=protected-access
- )
- )
- register_from_doc(cls_name)(
- DefaultTranslator(
- getattr(ast, cls_name),
- from_doc,
- [
- field
- for field in doc_cls._FIELDS # pylint: disable=protected-access
- if field in getattr(getattr(ast, cls_name), "_fields", ())
- or field in {"lineno", "col_offset", "end_lineno", "end_col_offset"}
- ],
- )
- )
-
-
-_register_default()
diff --git a/python/tvm/script/parser/core/doc_core.py b/python/tvm/script/parser/core/doc_core.py
deleted file mode 100644
index 74b7e57e3d..0000000000
--- a/python/tvm/script/parser/core/doc_core.py
+++ /dev/null
@@ -1,1123 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=redefined-outer-name,missing-docstring,invalid-name
-# pylint: disable=useless-super-delegation,redefined-builtin
-# pylint: disable=too-few-public-methods,too-many-arguments
-# ruff: noqa: RUF012
-
-
-# minimal python version is 3.9
-class AST:
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class mod(AST):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Module(mod):
- _FIELDS = ["body"]
-
- def __init__(self, body):
- super().__init__()
- self.body = body
-
-
-class Interactive(mod):
- _FIELDS = ["body"]
-
- def __init__(self, body):
- super().__init__()
- self.body = body
-
-
-class Expression(mod):
- _FIELDS = ["body"]
-
- def __init__(self, body):
- super().__init__()
- self.body = body
-
-
-class stmt(AST):
- _FIELDS = ["lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__()
- self.lineno = lineno
- self.col_offset = col_offset
- self.end_lineno = end_lineno
- self.end_col_offset = end_col_offset
-
-
-class FunctionDef(stmt):
- _FIELDS = [
- "name",
- "args",
- "body",
- "decorator_list",
- "returns",
- "type_params",
- "lineno",
- "col_offset",
- "end_lineno",
- "end_col_offset",
- ]
-
- def __init__(
- self,
- name,
- args,
- body,
- decorator_list,
- returns,
- type_params,
- lineno,
- col_offset,
- end_lineno,
- end_col_offset,
- ):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.name = name
- self.args = args
- self.body = body
- self.decorator_list = decorator_list
- self.returns = returns
- self.type_params = type_params
-
-
-class ClassDef(stmt):
- _FIELDS = [
- "name",
- "bases",
- "keywords",
- "body",
- "decorator_list",
- "lineno",
- "col_offset",
- "end_lineno",
- "end_col_offset",
- ]
-
- def __init__(
- self,
- name,
- bases,
- keywords,
- body,
- decorator_list,
- lineno,
- col_offset,
- end_lineno,
- end_col_offset,
- ):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.name = name
- self.bases = bases
- self.keywords = keywords
- self.body = body
- self.decorator_list = decorator_list
-
-
-class Return(stmt):
- _FIELDS = ["value", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, value, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.value = value
-
-
-class Delete(stmt):
- _FIELDS = ["targets", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, targets, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.targets = targets
-
-
-class Assign(stmt):
- _FIELDS = ["targets", "value", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, targets, value, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.targets = targets
- self.value = value
-
-
-class AugAssign(stmt):
- _FIELDS = ["target", "op", "value", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, target, op, value, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.target = target
- self.op = op
- self.value = value
-
-
-class AnnAssign(stmt):
- _FIELDS = [
- "target",
- "annotation",
- "value",
- "simple",
- "lineno",
- "col_offset",
- "end_lineno",
- "end_col_offset",
- ]
-
- def __init__(
- self, target, annotation, value, simple, lineno, col_offset, end_lineno, end_col_offset
- ):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.target = target
- self.annotation = annotation
- self.value = value
- self.simple = simple
-
-
-class For(stmt):
- _FIELDS = [
- "target",
- "iter",
- "body",
- "orelse",
- "lineno",
- "col_offset",
- "end_lineno",
- "end_col_offset",
- ]
-
- def __init__(self, target, iter, body, orelse, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.target = target
- self.iter = iter
- self.body = body
- self.orelse = orelse
-
-
-class While(stmt):
- _FIELDS = ["test", "body", "orelse", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, test, body, orelse, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.test = test
- self.body = body
- self.orelse = orelse
-
-
-class If(stmt):
- _FIELDS = ["test", "body", "orelse", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, test, body, orelse, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.test = test
- self.body = body
- self.orelse = orelse
-
-
-class With(stmt):
- _FIELDS = ["items", "body", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, items, body, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.items = items
- self.body = body
-
-
-class Raise(stmt):
- _FIELDS = ["exc", "cause", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, exc, cause, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.exc = exc
- self.cause = cause
-
-
-class Try(stmt):
- _FIELDS = [
- "body",
- "handlers",
- "orelse",
- "finalbody",
- "lineno",
- "col_offset",
- "end_lineno",
- "end_col_offset",
- ]
-
- def __init__(
- self, body, handlers, orelse, finalbody, lineno, col_offset, end_lineno, end_col_offset
- ):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.body = body
- self.handlers = handlers
- self.orelse = orelse
- self.finalbody = finalbody
-
-
-class Assert(stmt):
- _FIELDS = ["test", "msg", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, test, msg, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.test = test
- self.msg = msg
-
-
-class Import(stmt):
- _FIELDS = ["names", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, names, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.names = names
-
-
-class ImportFrom(stmt):
- _FIELDS = ["module", "names", "level", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, module, names, level, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.module = module
- self.names = names
- self.level = level
-
-
-class Global(stmt):
- _FIELDS = ["names", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, names, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.names = names
-
-
-class Nonlocal(stmt):
- _FIELDS = ["names", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, names, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.names = names
-
-
-class Expr(stmt):
- _FIELDS = ["value", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, value, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.value = value
-
-
-class Pass(stmt):
- _FIELDS = ["lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
-
-
-class Break(stmt):
- _FIELDS = ["lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
-
-
-class Continue(stmt):
- _FIELDS = ["lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
-
-
-class expr(AST):
- _FIELDS = ["lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__()
- self.lineno = lineno
- self.col_offset = col_offset
- self.end_lineno = end_lineno
- self.end_col_offset = end_col_offset
-
-
-class type_param(AST):
- _FIELDS = ["lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__()
- self.lineno = lineno
- self.col_offset = col_offset
- self.end_lineno = end_lineno
- self.end_col_offset = end_col_offset
-
-
-class TypeVar(type_param):
- _FIELDS = [
- "name",
- "bound",
- "default_value",
- "lineno",
- "col_offset",
- "end_lineno",
- "end_col_offset",
- ]
-
- def __init__(self, name, bound, default_value, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.name = name
- self.bound = bound
- self.default_value = default_value
-
-
-class ParamSpec(type_param):
- _FIELDS = [
- "name",
- "default_value",
- "lineno",
- "col_offset",
- "end_lineno",
- "end_col_offset",
- ]
-
- def __init__(self, name, default_value, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.name = name
- self.default_value = default_value
-
-
-class TypeVarTuple(type_param):
- _FIELDS = [
- "name",
- "default_value",
- "lineno",
- "col_offset",
- "end_lineno",
- "end_col_offset",
- ]
-
- def __init__(self, name, default_value, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.name = name
- self.default_value = default_value
-
-
-class BoolOp(expr):
- _FIELDS = ["op", "values", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, op, values, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.op = op
- self.values = values
-
-
-class BinOp(expr):
- _FIELDS = ["left", "op", "right", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, left, op, right, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.left = left
- self.op = op
- self.right = right
-
-
-class UnaryOp(expr):
- _FIELDS = ["op", "operand", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, op, operand, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.op = op
- self.operand = operand
-
-
-class Lambda(expr):
- _FIELDS = ["args", "body", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, args, body, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.args = args
- self.body = body
-
-
-class IfExp(expr):
- _FIELDS = ["test", "body", "orelse", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, test, body, orelse, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.test = test
- self.body = body
- self.orelse = orelse
-
-
-class Dict(expr):
- _FIELDS = ["keys", "values", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, keys, values, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.keys = keys
- self.values = values
-
-
-class Set(expr):
- _FIELDS = ["elts", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, elts, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.elts = elts
-
-
-class ListComp(expr):
- _FIELDS = ["elt", "generators", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, elt, generators, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.elt = elt
- self.generators = generators
-
-
-class SetComp(expr):
- _FIELDS = ["elt", "generators", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, elt, generators, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.elt = elt
- self.generators = generators
-
-
-class DictComp(expr):
- _FIELDS = ["key", "value", "generators", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, key, value, generators, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.key = key
- self.value = value
- self.generators = generators
-
-
-class GeneratorExp(expr):
- _FIELDS = ["elt", "generators", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, elt, generators, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.elt = elt
- self.generators = generators
-
-
-class Yield(expr):
- _FIELDS = ["value", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, value, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.value = value
-
-
-class YieldFrom(expr):
- _FIELDS = ["value", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, value, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.value = value
-
-
-class Compare(expr):
- _FIELDS = ["left", "ops", "comparators", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, left, ops, comparators, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.left = left
- self.ops = ops
- self.comparators = comparators
-
-
-class Call(expr):
- _FIELDS = ["func", "args", "keywords", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, func, args, keywords, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.func = func
- self.args = args
- self.keywords = keywords
-
-
-class FormattedValue(expr):
- _FIELDS = [
- "value",
- "conversion",
- "format_spec",
- "lineno",
- "col_offset",
- "end_lineno",
- "end_col_offset",
- ]
-
- def __init__(
- self, value, conversion, format_spec, lineno, col_offset, end_lineno, end_col_offset
- ):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.value = value
- self.conversion = conversion
- self.format_spec = format_spec
-
-
-class JoinedStr(expr):
- _FIELDS = ["values", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, values, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.values = values
-
-
-class Constant(expr):
- _FIELDS = ["value", "kind", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, value, kind, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.value = value
- self.kind = kind
-
-
-class NamedExpr(expr):
- _FIELDS = ["target", "value", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, target, value, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.target = target
- self.value = value
-
-
-class Attribute(expr):
- _FIELDS = ["value", "attr", "ctx", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, value, attr, ctx, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.value = value
- self.attr = attr
- self.ctx = ctx
-
-
-class Subscript(expr):
- _FIELDS = ["value", "slice", "ctx", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, value, slice, ctx, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.value = value
- self.slice = slice
- self.ctx = ctx
-
-
-class Starred(expr):
- _FIELDS = ["value", "ctx", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, value, ctx, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.value = value
- self.ctx = ctx
-
-
-class Name(expr):
- _FIELDS = ["id", "ctx", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, id, ctx, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.id = id
- self.ctx = ctx
-
-
-class List(expr):
- _FIELDS = ["elts", "ctx", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, elts, ctx, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.elts = elts
- self.ctx = ctx
-
-
-class Tuple(expr):
- _FIELDS = ["elts", "ctx", "lineno", "col_offset", "end_lineno", "end_col_offset"]
-
- def __init__(self, elts, ctx, lineno, col_offset, end_lineno, end_col_offset):
- super().__init__(lineno, col_offset, end_lineno, end_col_offset)
- self.elts = elts
- self.ctx = ctx
-
-
-class Slice(AST):
- _FIELDS = ["lower", "upper", "step"]
-
- def __init__(self, lower, upper, step):
- super().__init__()
- self.lower = lower
- self.upper = upper
- self.step = step
-
-
-class expr_context(AST):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Del(expr_context):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Load(expr_context):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Store(expr_context):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class boolop(AST):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class And(boolop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Or(boolop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class operator(AST):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Add(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class BitAnd(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class BitOr(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class BitXor(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Div(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class FloorDiv(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class LShift(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Mod(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Mult(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class MatMult(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Pow(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class RShift(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Sub(operator):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class unaryop(AST):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Invert(unaryop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Not(unaryop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class UAdd(unaryop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class USub(unaryop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class cmpop(AST):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Eq(cmpop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Gt(cmpop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class GtE(cmpop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class In(cmpop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Is(cmpop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class IsNot(cmpop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class Lt(cmpop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class LtE(cmpop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class NotEq(cmpop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class NotIn(cmpop):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class comprehension(AST):
- _FIELDS = ["target", "iter", "ifs", "is_async"]
-
- def __init__(self, target, iter, ifs, is_async):
- super().__init__()
- self.target = target
- self.iter = iter
- self.ifs = ifs
- self.is_async = is_async
-
-
-class excepthandler(AST):
- _FIELDS = []
-
- def __init__(self):
- super().__init__()
-
-
-class ExceptHandler(excepthandler):
- _FIELDS = ["type", "name", "body"]
-
- def __init__(self, type, name, body):
- super().__init__()
- self.type = type
- self.name = name
- self.body = body
-
-
-class arguments(AST):
- _FIELDS = ["args", "vararg", "kwonlyargs", "kw_defaults", "kwarg", "defaults", "posonlyargs"]
-
- def __init__(
- self,
- args,
- vararg,
- kwonlyargs,
- kw_defaults,
- kwarg,
- defaults,
- posonlyargs,
- ):
- super().__init__()
- self.args = args
- self.vararg = vararg
- self.kwonlyargs = kwonlyargs
- self.kw_defaults = kw_defaults
- self.kwarg = kwarg
- self.defaults = defaults
- self.posonlyargs = posonlyargs
-
-
-class arg(AST):
- _FIELDS = ["arg", "annotation"]
-
- def __init__(self, arg, annotation):
- super().__init__()
- self.arg = arg
- self.annotation = annotation
-
-
-class keyword(AST):
- _FIELDS = ["arg", "value"]
-
- def __init__(self, arg, value):
- super().__init__()
- self.arg = arg
- self.value = value
-
-
-class alias(AST):
- _FIELDS = ["name", "asname"]
-
- def __init__(self, name, asname):
- super().__init__()
- self.name = name
- self.asname = asname
-
-
-class withitem(AST):
- _FIELDS = ["context_expr", "optional_vars"]
-
- def __init__(self, context_expr, optional_vars):
- super().__init__()
- self.context_expr = context_expr
- self.optional_vars = optional_vars
-
-
-__all__ = [
- "AST",
- "Add",
- "And",
- "AnnAssign",
- "Assert",
- "Assign",
- "Attribute",
- "AugAssign",
- "BinOp",
- "BitAnd",
- "BitOr",
- "BitXor",
- "BoolOp",
- "Break",
- "Call",
- "ClassDef",
- "Compare",
- "Constant",
- "Continue",
- "Del",
- "Delete",
- "Dict",
- "DictComp",
- "Div",
- "Eq",
- "ExceptHandler",
- "Expr",
- "Expression",
- "FloorDiv",
- "For",
- "FormattedValue",
- "FunctionDef",
- "GeneratorExp",
- "Global",
- "Gt",
- "GtE",
- "If",
- "IfExp",
- "Import",
- "ImportFrom",
- "In",
- "Interactive",
- "Invert",
- "Is",
- "IsNot",
- "JoinedStr",
- "LShift",
- "Lambda",
- "List",
- "ListComp",
- "Load",
- "Lt",
- "LtE",
- "MatMult",
- "Mod",
- "Module",
- "Mult",
- "Name",
- "NamedExpr",
- "Nonlocal",
- "Not",
- "NotEq",
- "NotIn",
- "Or",
- "ParamSpec",
- "Pass",
- "Pow",
- "RShift",
- "Raise",
- "Return",
- "Set",
- "SetComp",
- "Slice",
- "Starred",
- "Store",
- "Sub",
- "Subscript",
- "Try",
- "Tuple",
- "TypeVar",
- "TypeVarTuple",
- "UAdd",
- "USub",
- "UnaryOp",
- "While",
- "With",
- "Yield",
- "YieldFrom",
- "alias",
- "arg",
- "arguments",
- "boolop",
- "cmpop",
- "comprehension",
- "excepthandler",
- "expr",
- "expr_context",
- "keyword",
- "mod",
- "operator",
- "stmt",
- "type_param",
- "unaryop",
- "withitem",
-]
diff --git a/python/tvm/script/parser/core/entry.py b/python/tvm/script/parser/core/entry.py
deleted file mode 100644
index afd155f3be..0000000000
--- a/python/tvm/script/parser/core/entry.py
+++ /dev/null
@@ -1,223 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""The entry point of TVM parser."""
-
-import inspect
-from typing import Any, TypeVar
-
-import tvm
-
-from ....ir.module import IRModule
-from ...ir_builder import IRBuilder
-from . import doc
-from .diagnostics import Source
-from .error import ParserError
-from .parser import Parser
-
-WELL_FORMED_ERROR_MESSAGE = (
- "Program is not well-formed. If this is deliberate, consider "
- "setting check_well_formed in the top-level decorator to False "
- "(e.g., @I.ir_module(check_well_formed=False) or "
- "@R.function(check_well_formed=False))."
-)
-
-
-def _default_globals() -> dict[str, Any]:
- # lazy import here to avoid circular deps
- from tvm.script import tirx as _tirx_dsl # pylint: disable=import-outside-toplevel
- from tvm.script.parser import (
- ir, # pylint: disable=import-outside-toplevel
- relax, # pylint: disable=import-outside-toplevel
- )
- from tvm.script.parser import tirx as _tirx_parser # pylint: disable=import-outside-toplevel
- from tvm.script.tirx import tile as _tirx_tile # pylint: disable=import-outside-toplevel
- from tvm.tirx import layout as _tirx_layout # pylint: disable=import-outside-toplevel
-
- # Expose the layout `Axis` class so printed layout sugar like
- # `4 @ Axis.laneid` round-trips without per-script imports. Injecting just
- # `Axis` (one short symbol) avoids name collisions with common user shape
- # vars like `m`, `P`, `F` that registered axes happen to share names with.
- return {
- "tvm": tvm,
- "I": ir,
- "ir": ir,
- "T": _tirx_parser,
- "tir": _tirx_parser,
- "R": relax,
- "relax": relax,
- "Tx": _tirx_tile,
- "tirx": _tirx_dsl,
- "Axis": _tirx_layout.Axis,
- "TypeVar": TypeVar,
- }
-
-
-def scan_macro(program: Any | str, extra_vars: dict[str, Any] | None = None) -> Any:
- """Generate the AST, and the source code for __repr__."""
- # The AST will be converted into TIR at the time of expansion.
- source = Source(program)
- closure_vars = extra_vars or _default_globals()
- return source, closure_vars
-
-
-def parse(
- program: doc.AST | Any | str,
- extra_vars: dict[str, Any] | None = None,
- check_well_formed: bool = True,
- s_tir: bool = False,
- absent_params: dict[str, None] | None = None,
-) -> Any:
- """Register a method for a operand type, AST operator node and operand index.
-
- Parameters
- ----------
- program : Union[doc.AST, Any, str]
- The TVMScript code to parse.
-
- extra_vars : Dict[str, Any]
- The extra variable table for parsing.
-
- check_well_formed : bool
- Whether to check well-formedness after parsing.
-
- s_tir : bool
- Compatibility argument accepted by parse/from_source. It no longer
- selects verification: each PrimFunc's s_tir attribute determines its
- dialect checks, and common well-formedness checks cover the full module.
-
- absent_params : Optional[Dict[str, None]]
- Function parameters removed by a compile-time specialization. The
- dialect-specific function parser decides how to bind these names.
-
- Returns
- -------
- func : Any
- The parsed TVMScript program.
- """
- if extra_vars is None:
- extra_vars = _default_globals()
-
- ann = {}
- all_pyfuncs = {}
- if inspect.isfunction(program):
- ann = {program.__name__: program.__annotations__}
- elif inspect.isclass(program):
- for name, func in program.__dict__.items():
- if inspect.isfunction(func):
- ann[name] = func.__annotations__
- all_pyfuncs[name] = func
-
- source = Source(program)
- parser = Parser(source, ann, absent_params=absent_params)
- with IRBuilder() as builder:
- try:
- parser.parse(extra_vars=extra_vars)
- except ParserError as err:
- parser.report_error(err.node, err.args[0])
- ret = builder.get()
- # Attach pyfuncs to the IRModule
- if inspect.isclass(program) and isinstance(ret, IRModule):
- _attach_pyfuncs_to_irmodule(ret, all_pyfuncs)
-
- # check well-formedness in both Relax and TIR
- if check_well_formed:
- check_ret = ret
- if not isinstance(check_ret, IRModule):
- check_ret = IRModule.from_expr(ret)
-
- source_ast = source.as_ast()
-
- if isinstance(
- ret, IRModule | tvm.relax.Function
- ) and not tvm.relax.analysis.check_well_formed(ret):
- parser.report_error(source_ast, err=WELL_FORMED_ERROR_MESSAGE)
-
- try:
- # Keep shared-definition checks across both dialects, then apply
- # ordinary TIRX execution restrictions only to its own functions.
- tvm.s_tir.analysis.verify_well_formed(check_ret)
- for func in check_ret.functions.values():
- if isinstance(func, tvm.tirx.PrimFunc) and not func.attrs.get("s_tir", False):
- tvm.tirx.analysis.verify_tirx_well_formed(func)
- except Exception as err: # pylint: disable=broad-exception-caught
- parser.report_error(
- source_ast,
- err=f"{WELL_FORMED_ERROR_MESSAGE}\n\nTraceback: {err!s}",
- )
- return ret
-
-
-def _create_python_packed_func(pyfunc):
- """Create a PackedFunc wrapper for a Python function.
-
- This function creates a PackedFunc that can be called from TVM runtime
- and will execute the original Python function.
-
- Parameters
- ----------
- pyfunc : Callable
- The Python function to wrap.
-
- Returns
- -------
- PackedFunc
- A PackedFunc that wraps the Python function.
- """
-
- def packed_func_wrapper(*args, **kwargs):
- """Wrapper function that calls the original Python function."""
- try:
- result = pyfunc(*args, **kwargs)
- return result
- except Exception as error:
- print(f"Error calling Python function {pyfunc.__name__}: {error}")
- raise
-
- return packed_func_wrapper
-
-
-def _attach_pyfuncs_to_irmodule(irmodule, all_pyfuncs):
- """Attach Python functions to IRModule with reduced nesting."""
- if not all_pyfuncs:
- return
-
- if not hasattr(irmodule, "pyfuncs"):
- irmodule.pyfuncs = {}
-
- for global_var, func in irmodule.functions_items():
- if not isinstance(func, tvm.relax.ExternFunc):
- continue
- if not func.attrs.get("is_pyfunc", False):
- continue
-
- pyfunc_name = global_var.name_hint
- if pyfunc_name not in all_pyfuncs:
- continue
-
- pyfunc = all_pyfuncs[pyfunc_name]
- irmodule.pyfuncs[pyfunc_name] = pyfunc
-
- try:
- source_code = inspect.getsource(pyfunc)
- func = func.with_attr("python_source", source_code)
- except (OSError, TypeError):
- func = func.with_attr("python_source", f"# Source unavailable for {pyfunc_name}")
-
- packed_func = _create_python_packed_func(pyfunc)
- func = func.with_attr("python_packed_func", packed_func)
-
- irmodule[global_var] = func
diff --git a/python/tvm/script/parser/core/error.py b/python/tvm/script/parser/core/error.py
deleted file mode 100644
index a28a8daf74..0000000000
--- a/python/tvm/script/parser/core/error.py
+++ /dev/null
@@ -1,27 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Error classes for diagnostics."""
-
-from . import doc
-
-
-class ParserError(Exception):
- """Error class for diagnostics."""
-
- def __init__(self, node: doc.AST, msg: str):
- super().__init__(msg)
- self.node = node
diff --git a/python/tvm/script/parser/core/evaluator.py b/python/tvm/script/parser/core/evaluator.py
deleted file mode 100644
index 71a1b3bc8f..0000000000
--- a/python/tvm/script/parser/core/evaluator.py
+++ /dev/null
@@ -1,626 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""AST Evaluation"""
-
-import ast
-from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
-
-import tvm
-
-from . import dispatch, doc
-from .error import ParserError
-
-if TYPE_CHECKING:
- from .parser import Parser
-
-DEFAULT_OP: dict[type, Callable[..., Any]] = {
- doc.Add: lambda a, b: a + b,
- doc.Sub: lambda a, b: a - b,
- doc.Mult: lambda a, b: a * b,
- doc.Div: lambda a, b: a / b,
- doc.FloorDiv: lambda a, b: a // b,
- doc.Mod: lambda a, b: a % b,
- doc.LShift: lambda a, b: a << b,
- doc.RShift: lambda a, b: a >> b,
- doc.BitOr: lambda a, b: a | b,
- doc.BitXor: lambda a, b: a ^ b,
- doc.BitAnd: lambda a, b: a & b,
- doc.MatMult: lambda a, b: a @ b,
- doc.Pow: lambda a, b: a**b,
- doc.Eq: lambda a, b: a == b,
- doc.NotEq: lambda a, b: a != b,
- doc.Lt: lambda a, b: a < b,
- doc.LtE: lambda a, b: a <= b,
- doc.Gt: lambda a, b: a > b,
- doc.GtE: lambda a, b: a >= b,
- doc.Is: lambda a, b: a is b,
- doc.IsNot: lambda a, b: a is not b,
- doc.In: lambda a, b: a in b,
- doc.NotIn: lambda a, b: a not in b,
- doc.And: lambda a, b: a and b,
- doc.Or: lambda a, b: a or b,
- doc.Invert: lambda a: ~a,
- doc.Not: lambda a: not a,
- doc.UAdd: lambda a: +a,
- doc.USub: lambda a: -a,
-}
-
-
-def _get_builtin_or_none(name: str):
- builtins = globals().get("__builtins__")
- if not builtins:
- return None
- if not isinstance(builtins, dict) and hasattr(builtins, "__dict__"):
- builtins = builtins.__dict__
- if isinstance(builtins, dict):
- return builtins.get(name)
- return None
-
-
-class ExprEvaluator:
- """Expression evaluator for TVMScript parser.
-
- Parameters
- ----------
- parser : Parser
- The parser bound with the evaluator.
-
- value_table : Dict[str, Any]
- The value table for expression evaluation.
-
- new_value_count : int
- The count for intermediate result added during evaluation.
- """
-
- parser: "Parser"
- value_table: dict[str, Any]
- new_value_count: int
-
- def __init__(self, parser: "Parser", value_table: dict[str, Any]) -> None:
- super().__init__()
- self.parser = parser
- self.value_table = value_table
- self.new_value_count = 0
-
- @staticmethod
- def eval(parser: "Parser", value_table: dict[str, Any], node: doc.AST) -> Any:
- """Expression evaluation for TVMScript parser.
-
- Parameters
- ----------
- parser : Parser
- The parser bound with the evaluator.
-
- value_table : Dict[str, Any]
- The value table for expression evaluation.
-
- node : doc.AST
- The root node of AST tree node of expression to evaluate.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- self = ExprEvaluator(parser, value_table)
- result = self._visit(node) # pylint: disable=protected-access
- if isinstance(result, doc.Name):
- if result.id in self.value_table:
- return self.value_table[result.id]
- else:
- builtin = _get_builtin_or_none(result.id)
- if builtin:
- return builtin
- raise ParserError(result, f"Undefined variable: {result.id}")
- if isinstance(result, doc.Constant):
- return result.value
- raise TypeError(f"Unexpected result type: {type(result)}")
-
- def _add_intermediate_result(self, value: Any) -> doc.Name:
- """Add intermediate result during evaluation into value table.
-
- Parameters
- ----------
- value : Any
- The intermediate result.
-
- Returns
- -------
- name : doc.Name
- The doc AST name node with intermediate name for intermediate result.
- """
- if self.parser is not None:
- value = self.parser.annotate_current_source_span(value)
- name = f"__tvm_tmp_value_{self.new_value_count}"
- self.new_value_count += 1
- self.value_table[name] = value
- lineno = 0
- col_offset = 0
- return doc.Name(
- id=name,
- ctx=doc.Load(),
- lineno=lineno,
- col_offset=col_offset,
- end_lineno=None,
- end_col_offset=None,
- )
-
- def _visit(self, node: doc.AST) -> Any:
- """General doc AST node visiting method for expression evaluation.
-
- Parameters
- ----------
- node : doc.AST
- The root node of AST tree node of expression to evaluate.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- if isinstance(node, list):
- return [self._visit(n) for n in node]
- if isinstance(node, tuple):
- return tuple(self._visit(n) for n in node)
- assert isinstance(node, doc.AST)
- if self.parser is None:
- return self._visit_node(node)
- with self.parser.with_source_span(node):
- return self._visit_node(node)
-
- def _visit_node(self, node: doc.AST) -> Any:
- """Evaluate one AST node while its source span is active."""
- # Boolean operators and conditional expressions require AST-level
- # control over operand evaluation. Walking all fields first would
- # eagerly evaluate dead compile-time branches.
- if isinstance(node, doc.BoolOp):
- try:
- value = self._eval_bool_op(node)
- except Exception as err: # pylint: disable=broad-except
- self.parser.report_error(node, err)
- return self._add_intermediate_result(value)
- if isinstance(node, doc.IfExp):
- try:
- value = self._eval_if_exp(node)
- except Exception as err: # pylint: disable=broad-except
- self.parser.report_error(node, err)
- return self._add_intermediate_result(value)
-
- args = []
- if (
- isinstance(node, doc.Call)
- and hasattr(node.func, "attr")
- and node.func.attr not in ["reads", "writes", "match_buffer"]
- ) or isinstance(node, doc.BinOp | doc.UnaryOp | doc.Compare):
- if isinstance(node, doc.BinOp):
- args = [node.left, node.right]
- elif isinstance(node, doc.UnaryOp):
- args = [node.operand]
- elif isinstance(node, doc.Compare):
- args = [node.left, *node.comparators]
- elif isinstance(node, doc.Call):
- args = node.args
- for arg in args:
- if isinstance(arg, doc.Subscript) and isinstance(arg.slice, doc.Slice | doc.Tuple):
- if isinstance(arg.slice, doc.Slice):
- check_slices = [arg.slice]
- else:
- check_slices = []
- for p in arg.slice.elts:
- if isinstance(p, doc.Slice):
- check_slices.append(p)
- for s in check_slices:
- if not s.step and s.upper and s.lower:
- s.step = doc.Constant(
- 1,
- None,
- s.upper.lineno,
- s.upper.end_col_offset + 1,
- s.upper.lineno,
- s.upper.end_col_offset + 2,
- )
- if isinstance(node, doc.Name):
- if node.id not in self.value_table and not _get_builtin_or_none(node.id):
- raise ParserError(node, f"Undefined variable: {node.id}")
- return node
- if isinstance(
- node,
- doc.Constant | doc.expr_context | doc.operator | doc.boolop | doc.unaryop | doc.cmpop,
- ):
- return node
- if isinstance(node, doc.keyword):
- return doc.keyword(arg=node.arg, value=self._visit(node.value))
- if not isinstance(node, doc.expr | doc.Slice):
- return node
- if isinstance(node, doc.Lambda):
- return self._eval_lambda(node)
- if isinstance(node, doc.Starred):
- value = self._visit(node.value)
- return doc.Starred(
- value=value,
- ctx=node.ctx,
- lineno=node.lineno,
- col_offset=node.col_offset,
- end_lineno=node.end_lineno,
- end_col_offset=node.end_col_offset,
- )
-
- if isinstance(node, doc.ListComp | doc.SetComp | doc.DictComp):
- value = self._eval_expr(node)
- return self._add_intermediate_result(value)
-
- fields = {}
- for field in node.__class__._FIELDS: # pylint: disable=protected-access
- attr = getattr(node, field)
- if isinstance(attr, doc.AST | tuple | list):
- fields[field] = self._visit(attr)
- else:
- fields[field] = attr
- try:
- if isinstance(node, doc.Compare):
- value = self._eval_compare(fields)
- elif isinstance(node, doc.UnaryOp):
- value = self._eval_unary_op(fields)
- elif isinstance(node, doc.BinOp):
- value = self._eval_bin_op(fields)
- elif isinstance(node, doc.Slice):
- value = self._eval_slice(fields)
- else:
- value = self._eval_expr(node.__class__(**fields))
- except Exception as err: # pylint: disable=broad-except
- self.parser.report_error(node, err)
- return self._add_intermediate_result(value)
-
- def _eval_lambda(self, node: doc.Lambda) -> Any:
- """The doc AST lambda node evaluating method.
-
- Parameters
- ----------
- node : doc.Lambda
- The root node of AST tree node of expression to evaluate.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- try:
- value = self._eval_expr(node)
- except Exception as err: # pylint: disable=broad-except
- self.parser.report_error(node, err)
- return self._add_intermediate_result(value)
-
- def _eval_bool_op(self, node: doc.BoolOp) -> Any:
- """The doc AST boolean operator node evaluating method.
-
- Parameters
- ----------
- node : doc.BoolOp
- The boolean operation whose operands are evaluated in order.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- op = node.op
- if not isinstance(op, doc.And | doc.Or):
- raise TypeError(f"Unexpected operator: {op}")
- value = self._eval_expr(self._visit(node.values[0]))
- for rhs_node in node.values[1:]:
- if isinstance(value, bool):
- if isinstance(op, doc.And) and not value:
- return value
- if isinstance(op, doc.Or) and value:
- return value
- rhs = self._eval_expr(self._visit(rhs_node))
- value = _eval_op(op, values=[value, rhs])
- return value
-
- def _eval_compare(self, fields: dict[str, Any]) -> Any:
- """The doc AST comparison operation node evaluating method.
-
- Parameters
- ----------
- fields : Dict[str, Any]
- The dictionary of comparison operation information,
- e.g., operator types, operand values.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- values = [self._eval_expr(fields["left"])]
- values.extend([self._eval_expr(rhs) for rhs in fields["comparators"]])
- result = None
- assert len(fields["ops"]) == len(values) - 1
-
- for index, op in enumerate(fields["ops"]):
- sub_result = _eval_op(op, values=[values[index], values[index + 1]])
- if result is None:
- result = sub_result
- else:
- result = _eval_op(doc.And(), values=[result, sub_result])
- return result
-
- def _eval_unary_op(self, fields: dict[str, Any]) -> Any:
- """The doc AST unary operation node evaluating method.
-
- Parameters
- ----------
- fields : Dict[str, Any]
- The dictionary of unary operation information,
- e.g., operator types, operand values.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- value = self._eval_expr(fields["operand"])
- value = _eval_op(fields["op"], values=[value])
- return value
-
- def _eval_bin_op(self, fields: dict[str, Any]) -> Any:
- """The doc AST binary operation node evaluating method.
-
- Parameters
- ----------
- fields : Dict[str, Any]
- The dictionary of binary operation information,
- e.g., operator types, operand values.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- return _eval_op(
- fields["op"],
- values=[
- self._eval_expr(fields["left"]),
- self._eval_expr(fields["right"]),
- ],
- )
-
- def _eval_if_exp(self, node: doc.IfExp) -> Any:
- """The doc AST if-else expression node evaluating method.
-
- Parameters
- ----------
- node : doc.IfExp
- The conditional expression to evaluate.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- test = self._eval_expr(self._visit(node.test))
- if isinstance(test, bool):
- selected = node.body if test else node.orelse
- return self._eval_expr(self._visit(selected))
- elif tvm.ir.is_prim_expr(test) and test.ty.matches_code(tvm.DataTypeCode.BOOL):
- body = self._eval_expr(self._visit(node.body))
- orelse = self._eval_expr(self._visit(node.orelse))
- return tvm.tirx.op.if_then_else(test, body, orelse)
- else:
- raise TypeError(f"Expected Python bool or TIR bool, but got {type(test)}")
-
- def _eval_slice(self, fields: dict[str, Any]) -> slice:
- """The doc AST slice node evaluating method.
-
- Parameters
- ----------
- fields : Dict[str, Any]
- The dictionary of slice information,
- e.g., lower bound, upper bound, step.
-
- Returns
- -------
- res : slice
- The evaluation result.
- """
- lower, upper, step = fields["lower"], fields["upper"], fields["step"]
-
- lower = self._eval_expr(lower) if lower is not None else None
- upper = self._eval_expr(upper) if upper is not None else None
- step = self._eval_expr(step) if step is not None else None
-
- return slice(lower, upper, step)
-
- def _eval_expr(self, v: Any) -> Any:
- """The doc AST expression node evaluating method.
-
- Parameters
- ----------
- v : Any
- The root node of AST tree node of expression to evaluate.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- return _eval_expr(v, self.value_table)
-
-
-def eval_expr(
- parser: "Parser",
- node: doc.expr | doc.Expression,
- dict_globals: dict[str, Any] | None,
-) -> Any:
- """Expression evaluation for TVMScript parser.
-
- Parameters
- ----------
- parser : Parser
- The parser bound with the evaluator.
-
- node : Union[doc.expr, doc.Expression]
- The root node of AST tree node of expression to evaluate.
-
- dict_globals : Optional[Dict[str, Any]]
- The optional global value table for expression evaluation.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- value_table = {}
- if dict_globals is not None:
- value_table.update(dict_globals)
- return ExprEvaluator.eval(parser, value_table, node)
-
-
-def eval_assign(
- parser: "Parser",
- target: doc.expr,
- source: Any,
-) -> dict[str, Any]:
- """Expression assignment evaluation for TVMScript parser.
-
- Parameters
- ----------
- parser : Parser
- The parser bound with the evaluator.
-
- target : doc.expr
- The root node of AST tree node of assigned expression to evaluate.
-
- source : Any
- The source to be assigned with evaluated expression.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- try:
- return _eval_assign(target, source)
- except Exception as err: # pylint: disable=broad-except
- parser.report_error(target, err)
- raise
-
-
-def _eval_expr(
- node: doc.expr | doc.Expression,
- dict_globals: dict[str, Any] | None,
-) -> Any:
- """Expression evaluation implementation for TVMScript parser.
-
- Parameters
- ----------
- node : Union[doc.expr, doc.Expression]
- The root node of AST tree node of expression to evaluate.
-
- dict_globals : Optional[Dict[str, Any]]
- The optional global value table for expression evaluation.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- node = doc.from_doc(node)
- if isinstance(node, ast.expr):
- node = ast.Expression(body=node)
- assert isinstance(node, ast.Expression), "Expects an ast.Expression, but gets: " + str(node)
- if dict_globals is None:
- dict_globals = {}
- node = ast.fix_missing_locations(node)
- exe = compile(node, filename="<ast>", mode="eval")
- return eval(exe, dict_globals) # pylint: disable=eval-used
-
-
-def _eval_op(
- op: doc.AST,
- values: list[Any],
-):
- """Operation expression evaluation implementation for TVMScript parser.
-
- Parameters
- ----------
- op : doc.AST
- The root node of AST tree node of operation expression to evaluate.
-
- values : List[Any]
- The list of values of operands.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- op_type = type(op) # pylint: disable=protected-access
- for i, v in enumerate(values):
- v_type = getattr(type(v), "_dispatch_type", None)
- if v_type is None:
- continue
- f = dispatch.get_op(
- operand_type=v_type, op_node_type=op_type, operand_index=i, default=None
- )
- if f is not None:
- return f(*values)
- return DEFAULT_OP[op_type](*values)
-
-
-def _eval_assign(
- target: doc.expr,
- source: Any,
-) -> dict[str, Any]:
- """Expression assignment evaluation implementation for TVMScript parser.
-
- Parameters
- ----------
- target : doc.expr
- The root node of AST tree node of assigned expression to evaluate.
-
- source : Any
- The source to be assigned with evaluated expression.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- target = doc.from_doc(target)
- assert isinstance(target, ast.expr)
- RHS_VAR_NAME = "__tvm_rhs_var__" # pylint: disable=invalid-name
- rhs_var_name = RHS_VAR_NAME
- dict_locals = {rhs_var_name: source}
- mod = ast.fix_missing_locations(
- ast.Module(
- body=[
- ast.Assign(
- targets=[target],
- value=ast.Name(
- id=rhs_var_name,
- ctx=ast.Load(),
- ),
- )
- ],
- type_ignores=[],
- )
- )
- exe = compile(mod, filename="<ast>", mode="exec")
- exec(exe, {}, dict_locals) # pylint: disable=exec-used
- del dict_locals[rhs_var_name]
- return dict_locals
diff --git a/python/tvm/script/parser/core/parser.py b/python/tvm/script/parser/core/parser.py
deleted file mode 100644
index 6319eacb34..0000000000
--- a/python/tvm/script/parser/core/parser.py
+++ /dev/null
@@ -1,1071 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""The core parser"""
-
-import abc
-import ast
-import inspect
-from collections import defaultdict
-from collections.abc import Callable
-from contextlib import contextmanager
-from typing import Any, TypeVar
-
-import numpy as np
-
-from tvm.error import DiagnosticError
-from tvm.ir import GlobalVar, Var
-from tvm.runtime import Object
-
-from ...ir_builder import IRBuilder
-from . import dispatch, doc
-from .diagnostics import Diagnostics, Source
-from .evaluator import eval_assign, eval_expr
-
-DEFAULT_VISIT = {
- "Interactive",
- "Module",
- "Expression",
- "Pass",
-}
-
-
-def collect_signature_type_vars(parser: "Parser", node: doc.FunctionDef) -> dict[str, Var]:
- """Collect symbolic variables declared by function type-parameter syntax."""
-
- annotation_names = set()
- parameter_names = {arg.arg for arg in node.args.args}
-
- class AnnotationNameCollector(doc.NodeVisitor):
- def visit_Name(self, name): # pylint: disable=invalid-name
- annotation_names.add(name.id)
-
- def visit_Constant(self, constant): # pylint: disable=invalid-name
- if not isinstance(constant.value, str):
- return
- try:
- expression = ast.parse(constant.value, mode="eval")
- except SyntaxError:
- return
- annotation_names.update(
- child.id for child in ast.walk(expression) if isinstance(child, ast.Name)
- )
-
- name_collector = AnnotationNameCollector()
- for arg in node.args.args:
- name_collector.visit(arg.annotation)
- name_collector.visit(node.returns)
-
- symbolic_vars = {}
- for binding_name, value in parser.var_table.get().items():
- if (
- binding_name not in annotation_names
- or binding_name in parameter_names
- or not isinstance(value, TypeVar)
- ):
- continue
- if value.__name__ != binding_name:
- parser.report_error(
- node,
- f"TypeVar binding {binding_name!r} must match its declared name {value.__name__!r}",
- )
- if value.__constraints__ or value.__bound__ is not None:
- parser.report_error(
- node,
- f"Symbolic TypeVar {binding_name!r} must not have constraints or a bound",
- )
- symbolic_vars[binding_name] = Var(binding_name, "int64")
-
- for type_param in node.type_params or []:
- if not isinstance(type_param, doc.TypeVar):
- parser.report_error(type_param, "Only PEP 695 TypeVar parameters are supported")
- # Both a bare parameter and the optional ``int`` bound declare the
- # int64 PrimVar represented by this type parameter. The printer uses
- # the bare spelling, while accepting the bound as explicit input.
- if type_param.bound is not None and not (
- isinstance(type_param.bound, doc.Name) and type_param.bound.id == "int"
- ):
- parser.report_error(
- type_param,
- f"Symbolic TypeVar {type_param.name!r} must be unannotated or annotated as int",
- )
- if type_param.default_value is not None:
- parser.report_error(
- type_param,
- f"Symbolic TypeVar {type_param.name!r} must not have a default",
- )
- symbolic_vars[type_param.name] = Var(type_param.name, "int64")
- return symbolic_vars
-
-
-def _deferred(exit_f: Callable[[], None]):
- """Created context with certain exit function.
-
- Parameters
- ----------
- exit_f : Callable[[], None]
- The function to call when exiting the context.
-
- Returns
- -------
- res : Any
- The created context.
- """
-
- @contextmanager
- def context():
- try:
- yield
- finally:
- exit_f()
-
- return context()
-
-
-def _do_nothing(*args, **kwargs): # pylint: disable=unused-argument
- pass
-
-
-class ScriptMacro(abc.ABC):
- """Representation of a script macro.
-
- This is a callable object, intended to be called from the expression evaluator.
- The evaluator is expected to insert the current parser into the environment
- undef the name given by "parser_object_name".
-
- Once called, the ScriptMacro object will locate the current parser, and use it
- to parse the macro's body and produce the result.
-
- There were two major considerations for this design:
- 1. Implementing hygienic and non-hygienic macros.
- 2. Implementing macros that return values.
-
- Macro uses in TIR are only allowed at a statement-level, and they don't produce
- any values. Parsing of such macros could easily be done by intercepting doc.Call
- nodes in the TIR parser. If a macro is a value-producing expression, then there
- may not be a direct way to intercept calls to it if it's embedded in a complex
- expression. Because macros use function-call syntax, the evaluator will try to
- call the macro object, which this design relies on to parse and evaluate the macro.
- """
-
- parser_object_name = "__current_script_parser__"
-
- def __init__(
- self,
- source: Source,
- closure_vars: dict[str, Any],
- func: Callable,
- hygienic: bool,
- ) -> None:
- self.source = source
- self.closure_vars = closure_vars
- self.func = func
- self.hygienic = hygienic
-
- def __repr__(self):
- return self.source.source
-
- @abc.abstractmethod
- def parse_macro(self, parser: "Parser") -> Any:
- """The main macro parsing function. Different scripts may have different
- ways to parse a macro, and to return a value to the evaluator.
-
- Parameters
- ----------
- parser : Parser
- The parser with the appropriate frame already created and populated depending
- macro's hygiene settings,
-
- Returns
- -------
- The return value depends on the specifics of the particular script. It can be
- "None" or any other value or any type.
- """
-
- def _find_parser_def(self):
- outer_frame_infos = inspect.getouterframes(inspect.currentframe())
- for finfo in outer_frame_infos:
- parser = finfo.frame.f_globals.get(ScriptMacro.parser_object_name)
- if parser is not None:
- return parser
- raise RuntimeError(f"{ScriptMacro.parser_object_name} not available")
-
- def get_macro_def(self):
- ast_module = self.source.as_ast()
- for decl in ast_module.body:
- if isinstance(decl, doc.FunctionDef) and decl.name == self.func.__name__:
- return decl
- raise RuntimeError(f"cannot find macro definition for {self.func.__name__}")
-
- def __call__(self, *args, **kwargs):
- param_binding = inspect.signature(self.func).bind(*args, **kwargs)
- param_binding.apply_defaults()
- local_vars = param_binding.arguments
- parser = self._find_parser_def()
-
- with parser.with_diag_source(self.source):
- if self.hygienic:
- saved_var_table = parser.var_table
- parser.var_table = VarTable()
-
- with parser.var_table.with_frame():
- for k, v in self.closure_vars.items():
- parser.var_table.add(k, v)
- for k, v in local_vars.items():
- parser.var_table.add(k, v)
-
- parse_result = self.parse_macro(parser)
-
- parser.var_table = saved_var_table
-
- else:
- with parser.var_table.with_frame():
- for k, v in local_vars.items():
- parser.var_table.add(k, v)
-
- parse_result = self.parse_macro(parser)
-
- return parse_result
-
-
-class VarTableFrame:
- """The variable table frame.
- A frame of variable table stores the variables created in one block or scope.
-
- Parameters
- ----------
- vars : Set[str]
- The set of variable names in the variable table frame.
- """
-
- vars: set[str]
-
- def __init__(self):
- self.vars = set()
-
- def add(self, var: str):
- """Add a new variable into variable table frame.
-
- Parameters
- ----------
- var : str
- The name of new variable.
- """
- if var in self.vars:
- raise ValueError(f"Variable {var} already defined in current scope")
- self.vars.add(var)
-
- def pop_all(self, fn_pop: Callable[[str], None]):
- """Pop out all variable in variable table frame.
-
- Parameters
- ----------
- fn_pop : Callable[[str], None]
- The methods to call when popping each variable.
- """
- for var in self.vars:
- fn_pop(var)
- self.vars.clear()
-
-
-class VarTable:
- """The variable table.
- A variable table stores the all variables when parsing TVMScript.
-
- Parameters
- ----------
- frames : List[VarTableFrame]
- The list or stack of variable table frame.
-
- name2value : Dict[str, List[Any]]
- The dictionary for variable table name-based query.
- """
-
- frames: list[VarTableFrame]
- name2value: dict[str, list[Any]]
-
- def __init__(self):
- self.frames = []
- self.name2value = defaultdict(list)
-
- def with_frame(self):
- """Create a new variable table frame as with statement.
-
- Returns
- -------
- res : Any
- The context with new variable table frame.
- """
-
- def pop_frame():
- frame = self.frames.pop()
- frame.pop_all(lambda name: self.name2value[name].pop())
-
- self.frames.append(VarTableFrame())
- return _deferred(pop_frame)
-
- def add(self, var: str, value: Any, allow_shadowing: bool = False):
- """Add a new variable to variable table.
-
- Parameters
- ----------
- var : str
- The name of variable.
-
- value : Any
- The value of variable.
-
- allow_shadowing : bool
- The options of whether variable shadowing allowed for this variable.
- """
- # Skip if the key and value are equal to those in the var_table
- if self.name2value[var] and isinstance(self.name2value[var][-1], type(value)):
- old_value = self.name2value[var][-1]
- if isinstance(value, np.ndarray) and (old_value == value).all():
- return
- if isinstance(old_value, Object):
- if old_value.same_as(value):
- return
- elif old_value == value:
- return
- if allow_shadowing and var in self.frames[-1].vars:
- # Shadowing
- self.name2value[var][-1] = value
- else:
- self.frames[-1].add(var)
- self.name2value[var].append(value)
-
- def get(self) -> dict[str, Any]:
- """Get a variable dictionary of latest variables.
-
- Returns
- -------
- res : Any
- The variable dictionary copy of latest variables.
- """
- return {key: values[-1] for key, values in self.name2value.items() if values}
-
- def contains_in_current_frame(self, name: str) -> bool:
- """Check whether a variable name exists in the current frame."""
- return bool(self.frames) and name in self.frames[-1].vars
-
- def get_at_depth(self, depth: int) -> dict[str, Any]:
- """Get variables visible at the given frame depth, using current values.
-
- For each variable name that appears in frames 0..depth-1, count how many
- times it was pushed (to handle shadowing), then index into name2value at
- count-1 to retrieve the latest value visible at that depth.
-
- Parameters
- ----------
- depth : int
- The frame depth (number of frames visible).
-
- Returns
- -------
- res : dict[str, Any]
- Variable dictionary of values visible at the given depth.
- """
- result: dict[str, Any] = {}
- name_count: dict[str, int] = defaultdict(int)
- for frame_idx in range(min(depth, len(self.frames))):
- for name in self.frames[frame_idx].vars:
- name_count[name] += 1
- for name, count in name_count.items():
- if self.name2value[name]:
- result[name] = self.name2value[name][count - 1]
- return result
-
- def exist(self, value: Any) -> bool:
- """Check if any value exists in variable table.
-
- Parameters
- ----------
- value : Any
- The value of variable.
-
- Returns
- -------
- res : bool
- The existence of the value.
- """
- return any(
- value.same_as(known_value)
- for known_value_stack in self.name2value.values()
- for known_value in known_value_stack
- )
-
-
-def _dispatch_wrapper(func: dispatch.ParseMethod) -> dispatch.ParseMethod:
- def _wrapper(self: "Parser", node: doc.AST) -> None:
- try:
- return func(self, node)
- except Exception as err: # pylint: disable=broad-except
- self.report_error(node, err)
- raise
-
- return _wrapper
-
-
-def _dispatch(self: "Parser", type_name: str) -> dispatch.ParseMethod:
- for token in [self.dispatch_tokens[-1], "default"]:
- func = dispatch.get(token=token, type_name=type_name, default=None)
- if func is not None:
- return _dispatch_wrapper(func)
- return _dispatch_wrapper(lambda self, node: self.generic_visit(node))
-
-
-class Parser(doc.NodeVisitor):
- """The TVMScript parser
-
- Parameters
- ----------
- diag : Diagnostics
- The diagnostics for error reporting.
-
- dispatch_tokens : List[str]
- The list of dispatching tokens to dispatching parsing method
- of different IRs and different doc AST structure.
-
- var_table : VarTable
- The variable table for parsing.
-
- absent_params : Dict[str, None]
- Function parameters removed by a compile-time specialization.
- """
-
- diag: Diagnostics
- dispatch_tokens: list[str]
- function_annotations: dict[str, dict[str, Any]] | None
- absent_params: dict[str, None]
- var_table: VarTable
- inside_function: bool # whether we are within a function
- current_class: str | None = None # current class being parsed
- base_py_module_context: bool = False # whether current class inherits from BasePyModule
-
- def __init__(
- self,
- source: Source,
- function_annotations: dict[str, dict[str, Any]],
- absent_params: dict[str, None] | None = None,
- ) -> None:
- self.diag = Diagnostics(source)
- self.dispatch_tokens = ["default"]
- self.function_annotations = function_annotations
- self.absent_params = absent_params or {}
- self.var_table = VarTable()
- self.inside_function = False
-
- def parse(self, extra_vars: dict[str, Any] | None = None) -> Any:
- """The main parse method for parser.
-
- Parameters
- ----------
- extra_vars : Optional[Dict[str, Any]]
- The optional global value table for parsing.
-
- Returns
- -------
- res : Any
- The doc AST node visiting result.
- """
- if extra_vars is None:
- extra_vars = {}
- with self.var_table.with_frame():
- for k, v in extra_vars.items():
- self.var_table.add(k, v)
- node = self.diag.source.as_ast()
- self.visit(node)
-
- def get_dispatch_token(self, node: doc.FunctionDef) -> str:
- if not isinstance(node, doc.FunctionDef):
- self.report_error(node, "Only can get dispatch token for function.")
- if not node.decorator_list:
- self.report_error(node, "Function must be decorated")
- # TODO: only the last decorator is parsed
- decorator = self.eval_expr(node.decorator_list[-1])
- if not hasattr(decorator, "dispatch_token"):
- self.report_error(node, "The parser does not understand the decorator")
- return decorator.dispatch_token
-
- def with_dispatch_token(self, token: str):
- """Add a new dispatching token as with statement.
-
- Parameters
- ----------
- token : str
- The dispatching token.
-
- Returns
- -------
- res : Any
- The context with new dispatching token.
- """
-
- self.dispatch_tokens.append(token)
- enter_func = dispatch.get(token=token, type_name="enter_token", default=lambda *args: None)
- context = enter_func(self)
-
- def pop_token():
- exit_func = dispatch.get(
- token=token, type_name="exit_token", default=lambda *args: None
- )
- exit_func(self, context)
- self.dispatch_tokens.pop()
-
- return _deferred(pop_token)
-
- def set_class_context(self, class_name: str, is_base_py_module: bool = False):
- """Set the current class context for parsing.
-
- Parameters
- ----------
- class_name : str
- The name of the current class being parsed.
- is_base_py_module : bool
- Whether the current class inherits from BasePyModule.
- """
- self.current_class = class_name
- self.base_py_module_context = is_base_py_module
-
- def _get_current_class_context(self) -> str | None:
- """Get the current class context.
-
- Returns
- -------
- Optional[str]
- The name of the current class, or None if not in a class context.
- """
- return self.current_class
-
- def _is_base_py_module_context(self) -> bool:
- """Check if the current class context allows Python functions.
-
- Returns
- -------
- bool
- True if Python functions are allowed in the current context.
- """
- return self.base_py_module_context
-
- def with_diag_source(self, source: Source):
- """Add a new source as with statement.
-
- Parameters
- ----------
- source : Source
- The source for diagnostics.
-
- Returns
- -------
- res : Any
- The context with new source.
- """
-
- last_diag = self.diag
- self.diag = Diagnostics(source)
-
- def pop_source():
- self.diag = last_diag
-
- return _deferred(pop_source)
-
- @contextmanager
- def with_source_span(self, node: doc.AST):
- """Make ``node``'s source range active while constructing its IR."""
- if (
- not IRBuilder.is_in_scope()
- or getattr(node, "lineno", None) is None
- or getattr(node, "col_offset", None) is None
- ):
- yield
- return
- with IRBuilder.current().with_source_span(self.diag.source.to_span(node)):
- yield
-
- def annotate_current_source_span(self, value: Any) -> Any:
- """Attach the active parser span to an expression result, when applicable."""
- if isinstance(value, Object) and IRBuilder.is_in_scope():
- return IRBuilder.current()._set_current_source_span(value) # pylint: disable=protected-access
- return value
-
- def eval_expr(
- self,
- node: doc.Expression | doc.expr,
- extra_vars: dict[str, Any] | None = None,
- ) -> Any:
- """Expression evaluation when parsing.
-
- Parameters
- ----------
- node : Union[doc.expr, doc.Expression]
- The root node of AST tree node of expression to evaluate.
-
- extra_vars : Optional[Dict[str, Any]]
- The optional global value table for expression evaluation.
-
- Returns
- -------
- res : Any
- The evaluation result.
- """
- var_values = self.var_table.get()
- if extra_vars is not None:
- for k, v in extra_vars.items():
- var_values[k] = v
- var_values[ScriptMacro.parser_object_name] = self
- value = eval_expr(self, node, var_values)
-
- return self.annotate_current_source_span(value)
-
- def _duplicate_lhs_check(self, target: doc.expr) -> bool | set[str]:
- """Check whether duplicate lhs exists in assignment.
-
- Parameters
- ----------
- target : doc.expr
- The doc AST expr node for lhs.
-
- Returns
- -------
- res : Union[bool, Set[str]]
- The result of true if duplicate lhs exists,
- or the set of lhs names if no duplicate lhs exists.
- """
- if isinstance(target, doc.Tuple | doc.List):
- vars: set[str] = set() # pylint: disable=redefined-builtin
- for i in target.elts:
- res = self._duplicate_lhs_check(i)
- if isinstance(res, bool) and res:
- return True
- assert isinstance(res, set)
- if vars & res:
- return True
- vars = vars.union(res)
- return vars
- elif isinstance(target, doc.Name):
- return {target.id}
- elif isinstance(target, doc.Starred):
- return self._duplicate_lhs_check(target.value)
- elif isinstance(target, doc.Attribute):
- # Attribute assignment like packedB.data = ..., treated as rebinding.
- return {target.attr}
- else:
- self.report_error(target, "Invalid type in assign statement")
- raise NotImplementedError
-
- def eval_assign(
- self,
- target: doc.expr,
- source: Any,
- bind_value: Callable[["Parser", doc.expr, str, Any], Any],
- allow_shadowing: bool = False,
- ) -> dict[str, Any]:
- """Expression assignment evaluation when parsing.
-
- Parameters
- ----------
- target : doc.expr
- The root node of AST tree node of assigned expression to evaluate.
-
- source : Any
- The source to be assigned with evaluated expression.
-
- bind_value : Callable[["Parser", doc.expr, str, Any], Any]
- The value binding method when assigning the values to variables.
-
- allow_shadowing : bool
- The options of whether variable shadowing allowed for assignment.
-
- Returns
- -------
- res : Dict[str, Any]
- The dictionary of assignment result.
- """
- if self._duplicate_lhs_check(target) is True:
- self.report_error(target, "Duplicate vars assigned.")
- var_values = eval_assign(self, target, source)
- for k, v in var_values.items():
- var = bind_value(self, target, k, v)
- self.var_table.add(k, var, allow_shadowing)
- return var_values
-
- def report_error(self, node: doc.AST, err: Exception | str) -> None: # pylint: disable=no-self-use
- """The error reporting when parsing.
-
- Parameters
- ----------
- node : doc.AST
- The doc AST node with errors.
-
- err: Union[Exception, str]
- The error to report.
- """
-
- # If the error is already being raised as a DiagnosticError,
- # re-raise it without wrapping it in a DiagnosticContext.
- if isinstance(err, DiagnosticError):
- raise err
-
- # Only take the last line of the error message
- if isinstance(err, RuntimeError):
- lines = list(filter(None, str(err).split("\n")))
- msg = lines[-1] if lines else (str(err) or type(err).__name__)
- elif isinstance(err, KeyError):
- msg = "KeyError: " + str(err)
- else:
- msg = str(err)
-
- try:
- self.diag.error(node, msg)
- except Exception as diag_err:
- # Calling self.diag.error is guaranteed to throw an
- # exception. When shown to a user, this error should
- # reference the point of error within the provided
- # TVMScript. However, when caught in pdb, the full
- # traceback should be available for debugging.
- if isinstance(err, Exception):
- diag_err = diag_err.with_traceback(err.__traceback__)
- raise diag_err
-
- def visit(self, node: doc.AST) -> None:
- """The general visiting method.
-
- Parameters
- ----------
- node : doc.AST
- The doc AST node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- if isinstance(node, list | tuple):
- for item in node:
- self.visit(item)
- return
- if not isinstance(node, doc.AST):
- return
- name = node.__class__.__name__.split(".")[-1]
- if name in DEFAULT_VISIT:
- func = self.generic_visit
- else:
- func = getattr(self, "visit_" + name, None)
- if func is None:
- raise NotImplementedError(f"Visitor of AST node is not implemented: {name}")
- try:
- with self.with_source_span(node):
- func(node)
- except Exception as err: # pylint: disable=broad-except
- self.report_error(node, err)
-
- def visit_body(self, node: list[doc.stmt]) -> Any:
- """The general body visiting method.
-
- Parameters
- ----------
- node : List[doc.stmt]
- The list of statements in body.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- for stmt in node:
- self.visit(stmt)
-
- def visit_tvm_annotation(self, node: doc.expr) -> Any:
- """The general TVM annotation visiting method.
-
- Parameters
- ----------
- node : doc.expr
- The doc AST expr node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "tvm_annotation")(self, node)
-
- def visit_FunctionDef(self, node: doc.FunctionDef) -> None: # pylint: disable=invalid-name
- """The general function definition visit method.
-
- Parameters
- ----------
- node : doc.FunctionDef
- The doc FunctionDef node.
- """
- token = self.get_dispatch_token(node)
- func = dispatch.get(token=token, type_name="FunctionDef", default=None)
- if func is None:
- self.report_error(
- node,
- """The parser does not understand the decorator,
- or visit_FunctionDef is not implemented for the decorator with token: """
- + token,
- )
- _dispatch(self, "pre_visit_local_function")(self, node)
- _dispatch_wrapper(func)(self, node)
- _dispatch(self, "post_visit_local_function")(self, node)
-
- def visit_tvm_declare_function(self, node: doc.FunctionDef) -> GlobalVar:
- token = self.get_dispatch_token(node)
- with self.with_dispatch_token(token):
- return _dispatch(self, "tvm_declare_function")(self, node)
-
- def visit_ClassDef(self, node: doc.ClassDef) -> Any: # pylint: disable=invalid-name
- """The general class definition visiting method.
-
- Parameters
- ----------
- node : doc.ClassDef
- The doc AST class definition node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- func = dispatch.get(token="ir", type_name="ClassDef", default=None)
- if func is None:
- self.report_error(node, "The parser does not understand the decorator")
- _dispatch_wrapper(func)(self, node)
-
- def visit_ImportFrom(self, node: doc.ImportFrom) -> None: # pylint: disable=invalid-name
- """Accept postponed annotations emitted by the PEP 695 printer."""
- is_future_annotations = (
- node.module == "__future__"
- and node.level == 0
- and len(node.names) == 1
- and node.names[0].name == "annotations"
- and node.names[0].asname is None
- )
- if not is_future_annotations:
- self.report_error(node, "Only 'from __future__ import annotations' is supported")
-
- def visit_arguments(self, node: doc.arguments) -> Any:
- """The general arguments visiting method.
-
- Parameters
- ----------
- node : doc.arguments
- The doc AST arguments node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "arguments")(self, node)
-
- def visit_For(self, node: doc.For) -> Any: # pylint: disable=invalid-name
- """The general for visiting method.
-
- Parameters
- ----------
- node : doc.For
- The doc AST for node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "For")(self, node)
-
- def visit_While(self, node: doc.While) -> Any: # pylint: disable=invalid-name
- """The general while visiting method.
-
- Parameters
- ----------
- node : doc.While
- The doc AST while node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "While")(self, node)
-
- def visit_With(self, node: doc.With) -> Any: # pylint: disable=invalid-name
- """The general with visiting method.
-
- Parameters
- ----------
- node : doc.With
- The doc AST with node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "With")(self, node)
-
- def visit_Assign(self, node: doc.Assign) -> Any: # pylint: disable=invalid-name
- """The general assign visiting method.
-
- Parameters
- ----------
- node : doc.Assign
- The doc AST assign node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "Assign")(self, node)
-
- def visit_AnnAssign(self, node: doc.AnnAssign) -> Any: # pylint: disable=invalid-name
- """The general annotated assign visiting method.
-
- Parameters
- ----------
- node : doc.Assign
- The doc AST annotated assign node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "AnnAssign")(self, node)
-
- def visit_Expr(self, node: doc.Expr) -> Any: # pylint: disable=invalid-name
- """The general expression visiting method.
-
- Parameters
- ----------
- node : doc.Expr
- The doc AST expression node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "Expr")(self, node)
-
- def visit_If(self, node: doc.If) -> Any: # pylint: disable=invalid-name
- """The general if visiting method.
-
- Parameters
- ----------
- node : doc.If
- The doc AST if node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "If")(self, node)
-
- def visit_AugAssign(self, node: doc.AugAssign) -> Any: # pylint: disable=invalid-name
- """The general augmented assignment visiting method.
-
- Parameters
- ----------
- node : doc.AugAssign
- The doc AST augmented assignment node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "AugAssign")(self, node)
-
- def visit_Assert(self, node: doc.Assert) -> Any: # pylint: disable=invalid-name
- """The general assert visiting method.
-
- Parameters
- ----------
- node : doc.Assert
- The doc AST assert node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "Assert")(self, node)
-
- def visit_Return(self, node: doc.Return) -> Any: # pylint: disable=invalid-name
- """The general return visiting method.
-
- Parameters
- ----------
- node : doc.Return
- The doc AST return node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "Return")(self, node)
-
- def visit_Continue(self, node: doc.Continue) -> Any: # pylint: disable=invalid-name
- """The general continue visiting method.
-
- Parameters
- ----------
- node : doc.Continue
- The doc AST continue node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "Continue")(self, node)
-
- def visit_Break(self, node: doc.Break) -> Any: # pylint: disable=invalid-name
- """The general break visiting method.
-
- Parameters
- ----------
- node : doc.Break
- The doc AST break node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "Break")(self, node)
-
- def visit_Nonlocal(self, node: doc.Nonlocal) -> Any: # pylint: disable=invalid-name
- """The general nonlocal visiting method.
-
- Parameters
- ----------
- node : doc.Nonlocal
- The doc AST nonlocal node.
-
- Returns
- -------
- res : Any
- The visiting result.
- """
- return _dispatch(self, "Nonlocal")(self, node)
diff --git a/python/tvm/script/parser/core/utils.py b/python/tvm/script/parser/core/utils.py
deleted file mode 100644
index fc8a928e05..0000000000
--- a/python/tvm/script/parser/core/utils.py
+++ /dev/null
@@ -1,212 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""TVM Script Parser utils"""
-
-import inspect
-from collections.abc import Callable
-from types import FrameType
-from typing import Any
-
-from .diagnostics import findsource
-
-
-def get_func_nonlocals(func):
- """A modified version of `inspect.getclosurevars`"""
-
- if inspect.ismethod(func):
- func = func.__func__
-
- if not inspect.isfunction(func):
- raise TypeError(f"{func!r} is not a Python function")
-
- code = func.__code__
- # Nonlocal references are named in co_freevars and resolved
- # by looking them up in __closure__ by positional index
- nonlocal_vars = {}
- if func.__closure__ is not None:
- for var, cell in zip(code.co_freevars, func.__closure__):
- try:
- nonlocal_vars[var] = cell.cell_contents
- except ValueError as err:
- # cell_contents may raise ValueError if the cell is empty.
- if "empty" not in str(err):
- raise
- return nonlocal_vars
-
-
-def inspect_function_capture(func: Callable) -> dict[str, Any]:
- """Capture function non-locals and global variables.
-
- Parameters
- ----------
- func : Callable
- The function to inspect.
-
- Returns
- -------
- res : Dict[str, Any]
- The function variables map with non-local or global variables.
- """
- captured = {
- **func.__globals__, # type: ignore
- **get_func_nonlocals(func),
- }
- return captured
-
-
-def inspect_class_capture(cls: type) -> dict[str, Any]:
- """Capture class non-locals and global variables.
-
- Parameters
- ----------
- cls : type
- The class to inspect.
-
- Returns
- -------
- res : Dict[str, Any]
- The class variables map with non-local or global variables.
- """
- result: dict[str, Any] = {}
- for _, v in cls.__dict__.items():
- if inspect.isfunction(v):
- func_vars = inspect_function_capture(v)
- result.update(**func_vars)
- return result
-
-
-def _collect_annotation_names(source_obj: type | Callable) -> set[str]:
- """Parse source AST to find names used in function annotations.
-
- Returns the set of ``ast.Name`` identifiers found inside argument
- annotations and return annotations of any function definitions in
- *source_obj*.
- """
- import ast
- import textwrap
-
- try:
- source = textwrap.dedent(inspect.getsource(source_obj))
- tree = ast.parse(source)
- except (OSError, TypeError):
- return set()
-
- names: set[str] = set()
- for node in ast.walk(tree):
- if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
- for arg in node.args.args + node.args.posonlyargs + node.args.kwonlyargs:
- if arg.annotation:
- for n in ast.walk(arg.annotation):
- if isinstance(n, ast.Name):
- names.add(n.id)
- if node.returns:
- for n in ast.walk(node.returns):
- if isinstance(n, ast.Name):
- names.add(n.id)
- return names
-
-
-def _has_string_annotations(source_obj: type | Callable) -> bool:
- """Check if *source_obj* has stringified annotations (PEP 563)."""
- if inspect.isclass(source_obj):
- return any(
- isinstance(a, str)
- for v in source_obj.__dict__.values()
- if inspect.isfunction(v)
- for a in v.__annotations__.values()
- )
- return any(isinstance(a, str) for a in getattr(source_obj, "__annotations__", {}).values())
-
-
-def _get_enclosing_scope_names(qualname: str) -> set[str]:
- """Extract lexically enclosing scope names from ``__qualname__``.
-
- For ``outer.<locals>.inner.<locals>.func`` this returns ``{"outer", "inner"}``.
- """
- parts = qualname.split(".")
- return {p for p in parts[:-1] if p != "<locals>"}
-
-
-def resolve_closure_vars(
- source_obj: type | Callable, extra_vars: dict[str, Any], outer_stack: list
-) -> None:
- """Resolve closure variables hidden by PEP 563.
-
- With ``from __future__ import annotations``, variables used only in
- annotations are not captured in ``__closure__``. This function parses
- the source AST to find names used in function annotations, then looks
- them up in lexically enclosing scope frames identified via
- ``__qualname__``.
-
- Only triggered when annotations are actually strings (PEP 563 active).
- Only annotation-referenced names are added, and only from enclosing
- scopes — not from arbitrary caller frames.
-
- Works for both classes (``@I.ir_module``) and functions (``@T.prim_func``).
- """
- if not _has_string_annotations(source_obj):
- return
- ann_names = _collect_annotation_names(source_obj)
- enclosing = _get_enclosing_scope_names(source_obj.__qualname__)
- for name in ann_names:
- if name not in extra_vars:
- for frame_info in outer_stack[1:]:
- if frame_info.frame.f_code.co_name in enclosing:
- if name in frame_info.frame.f_locals:
- extra_vars[name] = frame_info.frame.f_locals[name]
- break
-
-
-def is_defined_in_class(frames: list[FrameType], obj: Any) -> bool:
- """Check whether a object is defined in a class scope.
-
- Parameters
- ----------
- frames : List[FrameType]
- The frame stack of the object, obtained by `inspect.stack()`.
-
- Returns
- -------
- res : bool
- The result if the object is defined in a class scope.
- """
-
- def _is_tvmscript_class_annotator(line: str) -> bool:
- """Checks if the line contains a TVMScript annotator for a class
-
- These match either `@I.ir_module` or `@R.rewriter`, or their
- imported names `@ir_module` or `@rewriter`.
- """
-
- return line.startswith("@") and ("ir_module" in line or "rewriter" in line)
-
- if len(frames) > 2:
- frame_info = frames[2]
- code_context = frame_info.code_context
- if code_context is None:
- return False
- line = code_context[0].strip()
- if _is_tvmscript_class_annotator(line):
- return True
- if line.startswith("class"):
- lineno = frame_info.lineno
- if lineno >= 2:
- source, _ = findsource(obj)
- line = source[lineno - 2].strip()
- if _is_tvmscript_class_annotator(line):
- return True
- return False
diff --git a/python/tvm/script/parser_v2/diagnostics.py b/python/tvm/script/parser/diagnostics.py
similarity index 100%
rename from python/tvm/script/parser_v2/diagnostics.py
rename to python/tvm/script/parser/diagnostics.py
diff --git a/python/tvm/script/parser_v2/frontend.py b/python/tvm/script/parser/frontend.py
similarity index 100%
rename from python/tvm/script/parser_v2/frontend.py
rename to python/tvm/script/parser/frontend.py
diff --git a/python/tvm/script/parser_v2/functions.py b/python/tvm/script/parser/functions.py
similarity index 100%
rename from python/tvm/script/parser_v2/functions.py
rename to python/tvm/script/parser/functions.py
diff --git a/python/tvm/script/parser_v2/ir.py b/python/tvm/script/parser/ir.py
similarity index 100%
rename from python/tvm/script/parser_v2/ir.py
rename to python/tvm/script/parser/ir.py
diff --git a/python/tvm/script/parser/ir/__init__.py b/python/tvm/script/parser/ir/__init__.py
deleted file mode 100644
index 063712f33d..0000000000
--- a/python/tvm/script/parser/ir/__init__.py
+++ /dev/null
@@ -1,39 +0,0 @@
-# isort: skip_file
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""The ir module parser"""
-
-from tvm.ir import Range, StringType, StringImm, GenericConst
-from ...ir_builder.ir import * # pylint: disable=redefined-builtin
-from . import parser as _parser
-from .entry import ir_module, pyfunc
-
-
-__all__ = [
- "GenericConst",
- "Range",
- "StringImm",
- "StringType",
- "dummy_global_info",
- "ir_module",
- "lookup_vdevice",
- "meta_var",
- "module_attrs",
- "module_global_infos",
- "pyfunc",
- "vdevice",
-]
diff --git a/python/tvm/script/parser/ir/entry.py b/python/tvm/script/parser/ir/entry.py
deleted file mode 100644
index 5191ce0f87..0000000000
--- a/python/tvm/script/parser/ir/entry.py
+++ /dev/null
@@ -1,163 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=import-outside-toplevel
-"""The entry point of TVM parser for ir module."""
-
-import inspect
-from collections.abc import Callable
-
-from tvm import cpu, ir
-from tvm.ir import GlobalVar, IRModule
-
-from .._core import parse, utils
-
-
-# this formulation allows us to support having @I.ir_module
-# appear as a decorator by itself or to have optional arguments
-# like @I.ir_module(check_well_formed=False)
-def ir_module(
- mod: type | None = None, check_well_formed: bool = True, s_tir: bool = False
-) -> IRModule:
- """The parsing method for ir module, by using `@ir_module` as decorator.
-
- Parameters
- ----------
- mod : Type
- The class to be parsed as ir module.
-
- check_well_formed : bool
- Whether to check well-formedness during parsing.
-
- Returns
- -------
- ir_module : IRModule
- The parsed ir module.
- """
-
- # Capture stack outside wrapper (wrapper adds to the stack)
- outer_stack = inspect.stack()
-
- def decorator_wrapper(mod):
- if not inspect.isclass(mod):
- raise TypeError(f"Expect a class, but got: {mod}")
-
- # Check BasePyModule inheritance
- base_py_module_inherited = any(base.__name__ == "BasePyModule" for base in mod.__bases__)
-
- extra_vars = utils.inspect_class_capture(mod)
- # Resolve closure variables hidden by PEP 563 (annotation-only names)
- utils.resolve_closure_vars(mod, extra_vars, outer_stack)
- m = parse(mod, extra_vars, check_well_formed=check_well_formed)
-
- if base_py_module_inherited:
- # Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
- # because tvm.script is loaded before tvm.relax during tvm initialization.
- from tvm.relax.base_py_module import BasePyModule
- from tvm.relax.expr import ExternFunc # pylint: disable=import-outside-toplevel
-
- # Collect pyfunc methods
- pyfunc_methods = [
- name
- for name, attr in mod.__dict__.items()
- if hasattr(attr, "dispatch_token") and attr.dispatch_token == "pyfunc"
- ]
-
- mod._pyfunc_methods = pyfunc_methods
-
- # Create ExternFunc nodes
-
- for method_name in pyfunc_methods:
- try:
- existing_gvars = [
- global_var
- for global_var in m.get_global_vars()
- if global_var.name_hint == method_name
- ]
-
- extern_func = ExternFunc(method_name)
- extern_func = extern_func.with_attr("is_pyfunc", True)
- extern_func = extern_func.with_attr("function_type", "python")
- extern_func = extern_func.with_attr("python_function_name", method_name)
- extern_func = extern_func.with_attr(
- "python_source", f"# Source for {method_name}"
- )
- extern_func = extern_func.with_attr("python_packed_func", None)
-
- if existing_gvars:
- m[existing_gvars[0]] = extern_func
- else:
- m[GlobalVar(method_name)] = extern_func
-
- except Exception: # pylint: disable=broad-exception-caught
- continue
-
- class ModuleFactory:
- """Factory class for creating BasePyModule instances with Python functions."""
-
- def __init__(self, module, pyfunc_methods, original_class):
- self.ir_module = module
- self.pyfunc_methods = pyfunc_methods
- self.original_class = original_class
-
- def __call__(self, device=None, target=None):
- if device is None:
- device = cpu(0)
-
- instance_ir_mod = ir.IRModule()
- for global_var, func in self.ir_module.functions_items():
- instance_ir_mod[global_var] = func
-
- instance = BasePyModule(instance_ir_mod, device, target)
-
- for method_name in self.pyfunc_methods:
- if hasattr(self.original_class, method_name):
- method = getattr(self.original_class, method_name)
- instance.add_python_function(method_name, method)
-
- return instance
-
- def __getattr__(self, name):
- if hasattr(self.ir_module, name):
- return getattr(self.ir_module, name)
- raise AttributeError(
- f"'{self.__class__.__name__}' object has no attribute '{name}'"
- )
-
- factory = ModuleFactory(m, pyfunc_methods, mod)
- setattr(factory, "__name__", mod.__name__)
- return factory
-
- setattr(m, "__name__", mod.__name__)
- return m
-
- if mod is not None:
- # if there are no optional args given, this will directly invoke the wrapper
- return decorator_wrapper(mod)
- else:
- # if there is a optional arg given, it returns the wrapper function
- # as a new decorator and applies it
- setattr(decorator_wrapper, "dispatch_token", "ir")
- return decorator_wrapper
-
-
-def pyfunc(func: Callable):
- # Set the dispatch_token on the decorated function
- setattr(func, "dispatch_token", "pyfunc")
- return func
-
-
-setattr(pyfunc, "dispatch_token", "pyfunc")
diff --git a/python/tvm/script/parser/ir/parser.py b/python/tvm/script/parser/ir/parser.py
deleted file mode 100644
index 848472404d..0000000000
--- a/python/tvm/script/parser/ir/parser.py
+++ /dev/null
@@ -1,212 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=unused-argument
-"""The base parser for ir module"""
-
-from tvm.ir import GlobalVar
-
-from ...ir_builder import ir as I
-from .._core import Parser, dispatch, doc
-
-
-class ModuleWithGlobalVars:
- """A Module that can add global vars during parsing, to support `Module.function` syntax."""
-
- def __getattr__(self, attr):
- # Customize the error message.
- # NOTE: `__getattr__` is only called when the attribute access fails with an AttributeError
- raise AttributeError(f"Cannot find the function `{attr}` in the current IRModule")
-
-
-@dispatch.register(token="ir", type_name="ClassDef")
-def _visit_class_def(self: Parser, node: doc.ClassDef) -> None:
- """The class definition visiting method for ir module.
-
- Parameters
- ----------
- self : Parser
- The visiting parser.
-
- node : doc.ClassDef
- The doc AST class definition node.
- """
-
- with self.var_table.with_frame():
- with I.ir_module():
- # Step 0. Add the class name to the var table
- fake_module = ModuleWithGlobalVars()
- self.var_table.add(node.name, fake_module)
-
- # Step 1: Check if this class inherits from BasePyModule
- is_base_py_module = _check_base_py_module_inheritance(node)
- if is_base_py_module:
- # Store this information in the IRModule for later use
- I.module_attrs({"base_py_module": True})
- # Set the parser context to allow Python functions
- self.set_class_context(node.name, True)
- else:
- # Set the parser context to disallow Python functions
- self.set_class_context(node.name, False)
-
- # Step 2. Visit non-function stmts, including but not limited to
- # 1. `I.module_attrs`
- # 2. `I.module_global_infos`
- with self.with_dispatch_token("ir"):
- for stmt in node.body:
- if not isinstance(stmt, doc.FunctionDef):
- self.visit(stmt)
-
- # Step 3. Visit function stmts to declare the global vars
- for stmt in node.body:
- if isinstance(stmt, doc.FunctionDef):
- global_var = self.visit_tvm_declare_function(stmt)
- fake_module.__setattr__(stmt.name, global_var)
-
- # Step 4. Visit and parse the functions
- with self.with_dispatch_token("ir"):
- for stmt in node.body:
- if isinstance(stmt, doc.FunctionDef):
- self.visit(stmt)
-
-
-@dispatch.register(token="ir", type_name="Assign")
-def _visit_assign(self: Parser, node: doc.Assign) -> None:
- """The assign visiting method for ir module.
-
- Parameters
- ----------
- self : Parser
- The visiting parser.
-
- node : doc.ClassDef
- The doc AST assign node.
- """
- if len(node.targets) != 1:
- self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
- lhs = node.targets[0].id
- rhs = self.eval_expr(node.value)
-
- I.decl_function(lhs, rhs)
- I.def_function(lhs, rhs)
-
-
-@dispatch.register(token="ir", type_name="Expr")
-def _visit_expr(self: Parser, node: doc.Expr) -> None:
- """The expression visiting method for ir module.
-
- Parameters
- ----------
- self : Parser
- The visiting parser.
-
- node : doc.ClassDef
- The doc AST expression node.
- """
- self.eval_expr(node.value)
-
-
-@dispatch.register(token="default", type_name="Assign")
-def visit_assign(self: Parser, node: doc.Assign) -> None:
- if len(node.targets) != 1:
- self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
- lhs = node.targets[0]
- rhs = self.eval_expr(node.value)
- self.eval_assign(
- target=lhs, source=rhs, bind_value=lambda _a, _b, _c, value: value, allow_shadowing=True
- )
-
-
-@dispatch.register(token="default", type_name="pre_visit_local_function")
-def pre_visit_local_function(self: Parser, node: doc.Expr) -> None:
- pass
-
-
-@dispatch.register(token="default", type_name="post_visit_local_function")
-def post_visit_local_function(self: Parser, node: doc.Expr) -> None:
- pass
-
-
-@dispatch.register(token="pyfunc", type_name="tvm_declare_function")
-def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
- """Declare a Python function as an ExternFunc in the IRModule."""
- # Check if Python functions are allowed in this context
- # We need to check if we're in a class that inherits from BasePyModule
- current_class = self._get_current_class_context()
- if current_class and not self._is_base_py_module_context():
- self.report_error(
- node,
- "@I.pyfunc are only allowed in classes that inherit from BasePyModule. "
- f"Class '{current_class}' does not inherit from BasePyModule.",
- )
-
- # Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
- # because tvm.script is loaded before tvm.relax during tvm initialization.
- from tvm.relax import ExternFunc # pylint: disable=import-outside-toplevel
-
- # Create ExternFunc with proper attributes for Python functions
- func = ExternFunc(node.name)
- func = func.with_attr("is_pyfunc", True)
- func = func.with_attr("function_type", "python")
- func = func.with_attr("python_function_name", node.name)
-
- # Add placeholder attributes that will be filled in later
- func = func.with_attr("python_source", f"# Source will be filled for {node.name}")
- func = func.with_attr("python_packed_func", None) # Will be filled in entry.py
-
- # Store the function name for later retrieval
- return I.decl_function(node.name, func)
-
-
-@dispatch.register(token="pyfunc", type_name="FunctionDef")
-def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
- """Visit Python function definition - no need to parse the body."""
- # Python function body is not parsed in TVMScript
-
-
-def _check_base_py_module_inheritance(node: doc.ClassDef) -> bool:
- """Check if a class inherits from BasePyModule.
-
- Parameters
- ----------
- node : doc.ClassDef
- The class definition node to check.
-
- Returns
- -------
- bool
- True if the class inherits from BasePyModule, False otherwise.
- """
- if not node.bases:
- return False
-
- # Check each base class
- for base in node.bases:
- if hasattr(base, "id"):
- if base.id == "BasePyModule":
- return True
- elif hasattr(base, "attr"):
- if base.attr == "BasePyModule":
- return True
- elif hasattr(base, "value") and hasattr(base.value, "id"):
- if (
- base.value.id in ["BasePyModule", "tvm", "relax"]
- and hasattr(base, "attr")
- and base.attr == "BasePyModule"
- ):
- return True
-
- return False
diff --git a/python/tvm/script/parser_v2/transform.py b/python/tvm/script/parser/transform.py
similarity index 100%
rename from python/tvm/script/parser_v2/transform.py
rename to python/tvm/script/parser/transform.py
diff --git a/python/tvm/script/parser_v2/__init__.py b/python/tvm/script/parser_v2/__init__.py
deleted file mode 100644
index 3128291335..0000000000
--- a/python/tvm/script/parser_v2/__init__.py
+++ /dev/null
@@ -1,34 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Translate original Python ASTs into calls on a registered construction namespace.
-
-Shared I infrastructure owns builder lifetime, spans and module identities; each
-context selects X from reverse-registered function metadata. Calls return concrete
-values, constructor metadata describes syntax, and compiled AST locations remain
-those of the original source. Entry modules register policies; this parser never
-imports their namespaces.
-"""
-
-from .frontend import (
- from_source,
- ir_module,
- make_decorator,
- make_helper,
- parse,
- pyfunc,
- register_namespace,
-)
diff --git a/python/tvm/tirx/script/__init__.py b/python/tvm/tirx/script/__init__.py
index f1566e0473..9127ad4160 100644
--- a/python/tvm/tirx/script/__init__.py
+++ b/python/tvm/tirx/script/__init__.py
@@ -14,24 +14,38 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""TIRX-layer TVMScript pieces (parser, builder).
-
-After the per-dialect TVMScript restructure, the TIRX layer owns its own
-``script/{parser,builder}`` subpackages. ``tvm.script.tirx`` resolves to
-this module via the dialect registry, so the public parser surface
-(``prim_func``, ``Buffer``, ``Ptr``, etc.) is re-exported here.
-"""
-
-# pylint: disable=redefined-builtin,wildcard-import,unused-wildcard-import
-from .v2 import *
-from .v2 import Buffer, Ptr, prim_func
-
-try:
- from .v2 import macro
-except ImportError:
- macro = None
+"""TVMScript entry point using concrete TIRx construction operations."""
+
+# pylint: disable=wildcard-import,unused-wildcard-import,redefined-builtin
+import sys as _sys
+
+from tvm.script.parser.frontend import make_decorator as _make_decorator
+from tvm.script.parser.frontend import make_helper as _make_helper
+from tvm.script.parser.frontend import register_namespace as _register_namespace
+from tvm.tirx.layout import Axis as _Axis
+
+from . import builder as _builder
+from . import tile as _tile
+from .builder import *
+from .tile import cluster as cluster
+from .tile import cta as cta
+from .tile import thread as thread
+from .tile import warp as warp
+from .tile import warpgroup as warpgroup
+from .tile import wg as wg
+
+tile = _tile
+prim_func = _make_decorator(
+ _builder, option_map={"private": "private", "s_tir": "s_tir", "persistent": "persistent"}
+)
+inline = _make_helper(_builder, preserve_return=True, late_binding=True)
+macro = _make_helper(_builder, preserve_return=False)
+
+_register_namespace("T", _sys.modules[__name__])
+_register_namespace("tirx", _sys.modules[__name__])
+
+_register_namespace("Axis", _Axis)
+
from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool
-from . import tile
from .builder.ir import TensorMap, meta_class
-from .tile import cluster, cta, thread, warp, warpgroup, wg
diff --git a/python/tvm/tirx/script/builder/__init__.py b/python/tvm/tirx/script/builder/__init__.py
index 769d13c317..4069fed4fa 100644
--- a/python/tvm/tirx/script/builder/__init__.py
+++ b/python/tvm/tirx/script/builder/__init__.py
@@ -1,4 +1,3 @@
-# isort: skip_file
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
@@ -15,12 +14,572 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Package tvm.script.ir_builder.tirx"""
+"""Concrete TIRx construction operations over the shared native IRBuilder stack."""
-from .ir import * # pylint: disable=wildcard-import,redefined-builtin
-from .ir import boolean as bool # pylint: disable=redefined-builtin
-from .ir import buffer as Buffer
-from .utils import buffer_proxy, frame_scope, seq_scope
+import builtins as _python
+import sys as _sys
+from functools import partial as _partial
+from functools import wraps as _wraps
+
+import tvm_ffi as _ffi
+
+from tvm import ir as _ir
+from tvm import tirx as _tir
+from tvm.script.ir_builder import IRBuilder as _IRBuilder
+from tvm.script.ir_builder import ir as _I
+from tvm.script.ir_builder.base import IRBuilderFrame as _NativeFrame
+from tvm.script.ir_builder.protocol import MISSING as _MISSING
+from tvm.script.ir_builder.protocol import at as _at
+from tvm.script.ir_builder.protocol import expression_args as _expression_args
+from tvm.script.ir_builder.protocol import register_call_kind as _register_call_kind
+from tvm.script.ir_builder.protocol import register_declaration as _register_declaration
+from tvm.script.ir_builder.protocol import span_context as _span_context
from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool
+
+from . import _ffi_api
+from . import frame as _frame
+from . import ir as _native
from . import tirx as tile
+from .ir import * # pylint: disable=wildcard-import,redefined-builtin
+from .ir import boolean as bool # pylint: disable=redefined-builtin
from .tirx import cluster, cta, thread, warp, warpgroup, wg
+from .utils import buffer_proxy, frame_scope, seq_scope
+
+is_type_var = _ir.is_prim_var
+
+
+def type_var(name, *, dtype=None, span=None):
+ """Construct a signature symbol; shape symbols default to int64."""
+ return _ir.Var(name, "int64" if dtype is None else dtype, span)
+
+
+@_expression_args(
+ "shape",
+ "strides",
+ "elem_offset",
+ "byte_offset",
+ introduce=True,
+ implicit_dtype="int32",
+ compound_declarations=True,
+)
+def Buffer(
+ shape,
+ dtype="float32",
+ data=None,
+ strides=None,
+ elem_offset=None,
+ byte_offset=None,
+ scope="global",
+ align=0,
+ offset_factor=0,
+ layout="default",
+ allocated_addr=None,
+ buffer_name="",
+ *,
+ span=None,
+):
+ """Construct a concrete buffer from resolved shape expressions."""
+ with _span_context(span):
+ return _at(
+ span,
+ _native.buffer(
+ shape,
+ dtype,
+ data,
+ strides,
+ elem_offset,
+ byte_offset,
+ scope,
+ align,
+ offset_factor,
+ layout,
+ allocated_addr,
+ buffer_name,
+ ),
+ )
+
+
+buffer = Buffer
+
+
+def Ptr(dtype, storage_scope="global", *, span=None):
+ """Construct a concrete pointer variable usable as a function annotation."""
+ if callable(dtype) and not isinstance(dtype, _ir.Expr):
+ dtype = dtype()
+ if isinstance(dtype, _ir.Expr):
+ dtype = dtype.ty
+ if isinstance(dtype, _ir.PrimType):
+ dtype = dtype.dtype
+ with _span_context(span):
+ return _at(span, _native.ptr(dtype, storage_scope))
+
+
+class _Frame:
+ """Preserve a frame's source range through native finalization."""
+
+ def __init__(self, native, span=None):
+ self.native = native
+ self.span = span
+ self.result = {}
+
+ def __enter__(self):
+ with _span_context(self.span):
+ value = self.native.__enter__()
+ return self if value is self.native else value
+
+ def __exit__(self, *exc):
+ with _span_context(self.span):
+ return self.native.__exit__(*exc)
+
+ @property
+ def reference(self):
+ """Return the stable module reference after signature finalization."""
+ return self.native.global_var
+
+ def __getattr__(self, name):
+ return getattr(self.native, name)
+
+
+def function(*, private=False, s_tir=False, persistent=False, span=None):
+ """Enter a native primitive-function definition frame."""
+ with _span_context(span):
+ return _Frame(_native.prim_func(private=private, s_tir=s_tir, persistent=persistent), span)
+
+
+def decl_function(*, private=False, s_tir=False, persistent=False, span=None):
+ """Declare a bodyless signature using the native function frame."""
+ with _span_context(span):
+ return _Frame(_ffi_api.DeclFunction(private, s_tir, persistent), span)
+
+
+def arg(name, annotation, *, span=None):
+ """Use the same concrete parameter object in declaration and definition."""
+ if callable(annotation) and not isinstance(annotation, _ir.Expr):
+ annotation = annotation()
+ if isinstance(annotation, _ir.Type):
+ annotation = _ir.Var(name, annotation)
+ with _span_context(span):
+ if _tir.is_buffer_var(annotation) and annotation.ty.layout is not None:
+ frames = _IRBuilder.current().frames
+ if _python.any(
+ isinstance(frame, _frame.PrimFuncFrame) and frame.s_tir for frame in frames
+ ):
+ ty = annotation.ty
+ annotation = _native.buffer(
+ ty.shape,
+ ty.dtype,
+ strides=ty.strides,
+ elem_offset=ty.elem_offset,
+ scope=ty.storage_scope,
+ align=ty.data_alignment,
+ offset_factor=ty.offset_factor,
+ layout=None,
+ allocated_addr=list(ty.allocated_addr),
+ buffer_name=name,
+ )
+ return _native.arg(name, _at(span, annotation))
+
+
+def func_ret_type(annotation, *, span=None):
+ """Set the signature's concrete return type."""
+ if callable(annotation) and not isinstance(annotation, _ir.Expr):
+ annotation = annotation()
+ if isinstance(annotation, _ir.Expr):
+ annotation = annotation.ty
+ with _span_context(span):
+ return _native.func_ret(annotation)
+
+
+def _name(value, name, span):
+ if name is not None:
+ _IRBuilder.name(name, value)
+ return _at(span, value)
+
+
+def _enter_concise(frame):
+ native = frame.native if isinstance(frame, _Frame) else frame
+ native.add_callback(_partial(frame.__exit__, None, None, None))
+ return frame.__enter__()
+
+
+def _as_expr(value):
+ if isinstance(value, _ffi.ObjectConvertible):
+ value = value.asobject()
+ if isinstance(value, _ir.Expr):
+ return value
+ if isinstance(value, str):
+ return _ir.StringImm(value)
+ if isinstance(value, list | tuple):
+ return _ir.Tuple([_as_expr(item) for item in value])
+ return _tir.const(value)
+
+
+def bind_(
+ value=_MISSING,
+ *,
+ ty=None,
+ name=None,
+ span=None,
+ name_span=None,
+ previous=_MISSING,
+ declaration=False,
+ frame_value=False,
+):
+ """Bind values, or name a frame-owned value without introducing new storage."""
+ name_span = span if name_span is None else name_span
+ with _span_context(span):
+ if frame_value:
+ if isinstance(value, _frame.SBlockFrame):
+ raise TypeError("A block does not introduce an as-target value")
+ if isinstance(value, _python.list | _python.tuple | _ir.Array):
+ for index, item in enumerate(value):
+ bind_(
+ item,
+ name=None if name is None else f"{name}_{index}",
+ span=span,
+ name_span=name_span,
+ frame_value=True,
+ )
+ elif isinstance(value, _ir.Var | _tir.IterVar | _tir.Layout):
+ _name(value, name, name_span)
+ elif isinstance(value, _ir.TensorLoad) and _tir.is_buffer_var(value.source):
+ _name(value.source, name, name_span)
+ return value
+ if previous is not _MISSING and _tir.is_buffer_var(previous):
+ shape = previous.ty.shape
+ if len(shape) == 1 and bool(shape[0] == 1):
+ if value is _MISSING:
+ raise ValueError("A reassignment requires an initializer")
+ buffer_store(previous, value, [0])
+ return previous
+ if previous is not _MISSING and isinstance(getattr(previous, "ty", None), _ir.PointerType):
+ raise ValueError(f"Pointer variable {name!r} cannot be reassigned")
+ if previous is not _MISSING and (
+ _tir.is_buffer_var(previous)
+ or isinstance(previous, _tir.IterVar)
+ or _python.any(
+ isinstance(frame, _frame.SBlockFrame)
+ and _python.any(axis.var.same_as(previous) for axis in frame.iter_vars)
+ for frame in _IRBuilder.current().frames
+ )
+ ):
+ raise ValueError(f"Cannot rebind buffer or block axis {name!r}")
+ if declaration:
+ if not _ir.is_prim_var(value):
+ raise TypeError("A symbol declaration requires a concrete primitive variable")
+ if ty is not None:
+ annotation = ty() if callable(ty) else ty
+ annotation = annotation.ty if isinstance(annotation, _ir.Expr) else annotation
+ if not _ffi.structural_equal(annotation, value.ty):
+ raise TypeError("The symbol declaration has an incompatible type")
+ if previous is not _MISSING:
+ if not _ir.is_prim_var(previous) or not _ffi.structural_equal(
+ previous.ty, value.ty
+ ):
+ raise TypeError("The symbol declaration has an incompatible signature dtype")
+ return previous
+ return _name(value, name, name_span)
+ if previous is not _MISSING and isinstance(previous, _ir.TensorLoad):
+ if value is _MISSING:
+ raise ValueError("A reassignment requires an initializer")
+ buffer_store(previous.source, value, list(previous.indices))
+ return previous
+ if isinstance(value, _I.meta_var):
+ return value.value
+ if isinstance(ty, _native.LocalVectorAnnotation):
+ if value is not _MISSING:
+ raise ValueError("Vector annotation does not support an initializer")
+ return _name(_native.alloc_local(ty.shape, ty.dtype), name, name_span)
+ if isinstance(ty, _native.LetAnnotation):
+ if value is _MISSING:
+ raise ValueError("An immutable binding requires an initializer")
+ value = _as_expr(value)
+ variable = _name(ty.as_var(rhs_dtype=value.ty), name, name_span)
+ _native.Bind(value, var=variable)
+ return variable
+ if ty is not None:
+ annotation = ty() if callable(ty) and not isinstance(ty, _ir.Expr) else ty
+ annotation = annotation.ty if isinstance(annotation, _ir.Expr) else annotation
+ if not isinstance(annotation, _ir.PrimType) or str(annotation) == "handle":
+ raise TypeError("Mutable scalar annotations require a primitive scalar type")
+ result = _native.local_scalar(str(annotation)).scalar
+ _name(result.source, name, name_span)
+ if value is not _MISSING:
+ buffer_store(result.source, value, [0])
+ return result
+ if value is _MISSING:
+ raise ValueError("An uninitialized binding requires a scalar type annotation")
+ if (
+ isinstance(value, _ir.TensorLoad)
+ and _tir.is_buffer_var(value.source)
+ and not value.source.name
+ and len(value.source.ty.shape) == 1
+ and isinstance(value.source.ty.shape[0], _tir.IntImm)
+ and value.source.ty.shape[0].value == 1
+ ):
+ _name(value.source, name, name_span)
+ return value
+ if isinstance(value, _native.scalar_wrapper):
+ _name(value.scalar.source, name, name_span)
+ return value.scalar
+ if isinstance(value, _NativeFrame | _Frame):
+ return _name(_enter_concise(value), name, name_span)
+ if isinstance(value, list | tuple):
+ for index, item in enumerate(value):
+ bind_(item, name=None if name is None else f"{name}_{index}", span=span)
+ return value
+ if getattr(type(value), "_is_meta_class", False):
+ if name is not None:
+ _native.name_meta_class_value(name, value)
+ return value
+ if _tir.is_buffer_var(value) or isinstance(value, _tir.IterVar | _tir.Layout):
+ return _name(value, name, name_span)
+ if isinstance(value, _ir.Var) and not value.name:
+ return _name(value, name, name_span)
+ if isinstance(value, _ir.TensorRegion):
+ return value
+ if not isinstance(value, _ir.Expr | _python.int | _python.float | _python.bool | str):
+ return value
+ value = _as_expr(value)
+ if _ir.is_prim_expr(value):
+ result = _native.local_scalar(str(value.ty.dtype)).scalar
+ _name(result.source, name, name_span)
+ buffer_store(result.source, value, [0])
+ return result
+ return _name(_native.Bind(value), name, name_span)
+
+
+def emit_(value, *, span=None):
+ """Consume one expression statement, including effect-only calls."""
+ if value is None or isinstance(value, str | _ir.Var):
+ return
+ with _span_context(span):
+ if isinstance(value, _NativeFrame | _Frame):
+ _enter_concise(value)
+ elif hasattr(value, "frames"):
+ for frame in value.frames:
+ _enter_concise(frame)
+ elif isinstance(value, _tir.BufferStore):
+ buffer_store(value.buffer, value.value, value.indices)
+ else:
+ _native.evaluate(value)
+
+
+def setitem(target, key, value, *, span=None):
+ """Construct an indexed store after the caller has evaluated its operands."""
+ with _span_context(span):
+ buffer_store(target, value, key)
+
+
+def setattr(target, name, value, *, span=None):
+ """Store through scalar attributes, or update ordinary Python metadata."""
+ if isinstance(value, _I.meta_var):
+ _python.setattr(target, name, value.value)
+ return
+ previous = getattr(target, name, _MISSING)
+ if isinstance(previous, _native.scalar_wrapper):
+ previous = previous.scalar
+ buffer = previous.source if isinstance(previous, _ir.TensorLoad) else previous
+ if _tir.is_buffer_var(buffer):
+ shape = buffer.ty.shape
+ if len(shape) == 1 and bool(shape[0] == 1):
+ bind_(value, previous=previous, span=span)
+ return
+ _python.setattr(target, name, value)
+
+
+def return_(value, *, span=None):
+ """Emit a native return; subsequent unreachable statements remain in the IR."""
+ if value is None:
+ raise TypeError("A primitive function return requires an expression")
+ with _span_context(span):
+ _native.Return(_as_expr(value))
+
+
+def _require_loop():
+ for frame in reversed(_IRBuilder.current().frames):
+ if isinstance(frame, _frame.ForFrame | _frame.WhileFrame):
+ return
+ if isinstance(frame, _frame.PrimFuncFrame):
+ break
+ raise ValueError("Loop control requires an enclosing primitive loop")
+
+
+def break_(*, span=None):
+ """Construct a break targeting the nearest primitive loop."""
+ _require_loop()
+ with _span_context(span):
+ _native.evaluate(_native.break_loop())
+
+
+def continue_(*, span=None):
+ """Construct a continue targeting the nearest primitive loop."""
+ _require_loop()
+ with _span_context(span):
+ _native.evaluate(_native.continue_loop())
+
+
+def assert_(condition, message="", *, span=None):
+ """Emit the native flat assertion with its own source range."""
+ kind = "RuntimeError"
+ if isinstance(message, tuple):
+ if len(message) != 2 or not isinstance(message[0], str):
+ raise TypeError("Assertion metadata must be (error_kind, message_parts)")
+ kind, message = message
+ if isinstance(message, list | tuple):
+ message = [str(part) for part in message]
+ if not isinstance(message, list | tuple):
+ message = [message]
+ with _span_context(span):
+ with _native.Assert(condition, message, error_kind=kind):
+ pass
+
+
+def If(condition, *, span=None):
+ with _span_context(span):
+ return _Frame(_native.If(condition), span)
+
+
+def Then(*, span=None):
+ with _span_context(span):
+ return _Frame(_native.Then(), span)
+
+
+def Else(*, span=None):
+ with _span_context(span):
+ return _Frame(_native.Else(), span)
+
+
+def For(iterable, *, span=None):
+ """Adapt a native loop frame, or a Python range, to construction scope."""
+ if isinstance(iterable, _python.range):
+ iterable = _native.serial(iterable.start, iterable.stop, step=iterable.step)
+ if not isinstance(iterable, _frame.ForFrame):
+ raise TypeError("A primitive for loop requires a native loop frame or range")
+ return _Frame(iterable, span)
+
+
+def While(condition, *, span=None):
+ with _span_context(span):
+ return _Frame(_native.While(condition), span)
+
+
+def unpack(value):
+ """Project a concrete IR tuple of known arity, preserving Python iteration."""
+ if isinstance(value, _ir.Tuple):
+ return _python.tuple(value.fields)
+ if isinstance(value, _ir.Expr) and isinstance(value.ty, _ir.TupleType):
+ return _python.tuple(_ir.TupleGetItem(value, i) for i in range(len(value.ty.fields)))
+ return value
+
+
+def alloc_scalar(dtype="float32", scope="global"):
+ """Allocate scalar storage and return its concrete load expression."""
+ value = _native.alloc_scalar(dtype, scope)
+ return value.scalar if isinstance(value, _native.scalar_wrapper) else value
+
+
+def local_scalar(dtype="float32"):
+ return alloc_scalar(dtype, "local")
+
+
+def shared_scalar(dtype="float32"):
+ return alloc_scalar(dtype, "shared")
+
+
+@_expression_args(
+ "shape",
+ "strides",
+ "elem_offset",
+ introduce=True,
+ implicit_dtype="int32",
+ compound_declarations=True,
+)
+@_wraps(_native.match_buffer)
+def match_buffer(*args, **kwargs):
+ """Construct a native buffer match with resolved symbolic shape fields."""
+ return _native.match_buffer(*args, **kwargs)
+
+
+# Constructor identities carry syntax policy; aliases share it without wrappers.
+for _constructor in vars(_native).values():
+ if isinstance(_constructor, _native.DtypeConstructor):
+ _register_declaration(_constructor, dtype=_constructor._dtype_str)
+del _constructor
+
+
+def range_(*args):
+ """Construct a serial loop from the source builtin range arguments."""
+ if len(args) in (1, 2):
+ return _native.serial(*args)
+ if len(args) != 3:
+ raise TypeError("range expects one to three arguments")
+ start, stop, step = args
+ if isinstance(step, _python.int) and step == 0:
+ raise ValueError("range step cannot be zero")
+ return _native.serial(start, stop, step=step)
+
+
+__tvm_call_overrides__ = {_python.range: range_}
+
+
+def logical_and(*values):
+ """Construct scalar/vector conjunction, preserving ordinary Python values."""
+ if not values:
+ raise TypeError("logical_and requires at least one operand")
+ values = [
+ value.asobject() if isinstance(value, _ffi.ObjectConvertible) else value for value in values
+ ]
+ result = values[0]
+ for value in values[1:]:
+ if not isinstance(result, _ir.Expr) and not isinstance(value, _ir.Expr):
+ result = result and value
+ else:
+ lhs, rhs = _as_expr(result), _as_expr(value)
+ result = _tir.And(lhs, rhs) if lhs.ty.is_scalar() and rhs.ty.is_scalar() else lhs & rhs
+ return result
+
+
+def logical_or(*values):
+ """Construct scalar/vector disjunction, preserving ordinary Python values."""
+ if not values:
+ raise TypeError("logical_or requires at least one operand")
+ values = [
+ value.asobject() if isinstance(value, _ffi.ObjectConvertible) else value for value in values
+ ]
+ result = values[0]
+ for value in values[1:]:
+ if not isinstance(result, _ir.Expr) and not isinstance(value, _ir.Expr):
+ result = result or value
+ else:
+ lhs, rhs = _as_expr(result), _as_expr(value)
+ result = _tir.Or(lhs, rhs) if lhs.ty.is_scalar() and rhs.ty.is_scalar() else lhs | rhs
+ return result
+
+
+def logical_not(value):
+ """Construct IR negation without coercing an IR expression to Python bool."""
+ if isinstance(value, _ffi.ObjectConvertible):
+ value = value.asobject()
+ return _tir.Not(value) if isinstance(value, _ir.Expr) else not value
+
+
+def select(condition, true_value, false_value):
+ """Construct a conditional expression or select an ordinary Python value."""
+ if isinstance(condition, _ffi.ObjectConvertible):
+ condition = condition.asobject()
+ if not isinstance(condition, _ir.Expr):
+ return true_value if condition else false_value
+ return _tir.if_then_else(condition, true_value, false_value)
+
+
+def _global_callee(function):
+ return _partial(_native._call_global, function)
+
+
+_register_call_kind(_sys.modules[__name__], _ir.GlobalVar, _global_callee)
+
+
+from tvm.script.parser.annotations import enable_eager_constructors as _enable_eager_constructors
+
+_enable_eager_constructors(_sys.modules[__name__], classes=("Buffer",))
diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py
index 4cb399c04f..3f4e2b600f 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -3098,9 +3098,6 @@ def register_script_namespace(name: str, namespace: object) -> object:
for module_name in [
"tvm.tirx.script.builder",
- "tvm.tirx.script.builder_v2",
- "tvm.tirx.script.v2",
- "tvm.tirx.script.parser",
"tvm.tirx.script",
"tvm.script.tirx",
]:
diff --git a/python/tvm/tirx/script/builder_v2/__init__.py b/python/tvm/tirx/script/builder_v2/__init__.py
deleted file mode 100644
index 21ac4bb50f..0000000000
--- a/python/tvm/tirx/script/builder_v2/__init__.py
+++ /dev/null
@@ -1,579 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Concrete TIRx construction operations over the shared native IRBuilder stack."""
-
-import builtins as _python
-import sys as _sys
-from functools import partial as _partial
-from functools import wraps as _wraps
-
-import tvm_ffi as _ffi
-
-from tvm import ir as _ir
-from tvm import tirx as _tir
-from tvm.script.ir_builder import IRBuilder as _IRBuilder
-from tvm.script.ir_builder import ir as _I
-from tvm.script.ir_builder.base import IRBuilderFrame as _NativeFrame
-from tvm.script.ir_builder.protocol import MISSING as _MISSING
-from tvm.script.ir_builder.protocol import at as _at
-from tvm.script.ir_builder.protocol import expression_args as _expression_args
-from tvm.script.ir_builder.protocol import register_call_kind as _register_call_kind
-from tvm.script.ir_builder.protocol import register_declaration as _register_declaration
-from tvm.script.ir_builder.protocol import span_context as _span_context
-from tvm.tirx.script import builder as _T
-from tvm.tirx.script.builder import * # pylint: disable=wildcard-import,unused-wildcard-import
-from tvm.tirx.script.builder import _ffi_api
-from tvm.tirx.script.builder import frame as _frame
-
-is_type_var = _ir.is_prim_var
-
-
-def type_var(name, *, dtype=None, span=None):
- """Construct a signature symbol; shape symbols default to int64."""
- return _ir.Var(name, "int64" if dtype is None else dtype, span)
-
-
-@_expression_args(
- "shape",
- "strides",
- "elem_offset",
- "byte_offset",
- introduce=True,
- implicit_dtype="int32",
- compound_declarations=True,
-)
-def Buffer(
- shape,
- dtype="float32",
- data=None,
- strides=None,
- elem_offset=None,
- byte_offset=None,
- scope="global",
- align=0,
- offset_factor=0,
- layout="default",
- allocated_addr=None,
- buffer_name="",
- *,
- span=None,
-):
- """Construct a concrete buffer from resolved shape expressions."""
- with _span_context(span):
- return _at(
- span,
- _T.Buffer(
- shape,
- dtype,
- data,
- strides,
- elem_offset,
- byte_offset,
- scope,
- align,
- offset_factor,
- layout,
- allocated_addr,
- buffer_name,
- ),
- )
-
-
-buffer = Buffer
-
-
-def Ptr(dtype, storage_scope="global", *, span=None):
- """Construct a concrete pointer variable usable as a function annotation."""
- if callable(dtype) and not isinstance(dtype, _ir.Expr):
- dtype = dtype()
- if isinstance(dtype, _ir.Expr):
- dtype = dtype.ty
- if isinstance(dtype, _ir.PrimType):
- dtype = dtype.dtype
- with _span_context(span):
- return _at(span, _T.ptr(dtype, storage_scope))
-
-
-class _Frame:
- """Preserve a frame's source range through native finalization."""
-
- def __init__(self, native, span=None):
- self.native = native
- self.span = span
- self.result = {}
-
- def __enter__(self):
- with _span_context(self.span):
- value = self.native.__enter__()
- return self if value is self.native else value
-
- def __exit__(self, *exc):
- with _span_context(self.span):
- return self.native.__exit__(*exc)
-
- @property
- def reference(self):
- """Return the stable module reference after signature finalization."""
- return self.native.global_var
-
- def __getattr__(self, name):
- return getattr(self.native, name)
-
-
-def function(*, private=False, s_tir=False, persistent=False, span=None):
- """Enter a native primitive-function definition frame."""
- with _span_context(span):
- return _Frame(_T.prim_func(private=private, s_tir=s_tir, persistent=persistent), span)
-
-
-def decl_function(*, private=False, s_tir=False, persistent=False, span=None):
- """Declare a bodyless signature using the native function frame."""
- with _span_context(span):
- return _Frame(_ffi_api.DeclFunction(private, s_tir, persistent), span)
-
-
-def arg(name, annotation, *, span=None):
- """Use the same concrete parameter object in declaration and definition."""
- if callable(annotation) and not isinstance(annotation, _ir.Expr):
- annotation = annotation()
- if isinstance(annotation, _ir.Type):
- annotation = _ir.Var(name, annotation)
- with _span_context(span):
- if _tir.is_buffer_var(annotation) and annotation.ty.layout is not None:
- frames = _IRBuilder.current().frames
- if _python.any(
- isinstance(frame, _frame.PrimFuncFrame) and frame.s_tir for frame in frames
- ):
- ty = annotation.ty
- annotation = _T.Buffer(
- ty.shape,
- ty.dtype,
- strides=ty.strides,
- elem_offset=ty.elem_offset,
- scope=ty.storage_scope,
- align=ty.data_alignment,
- offset_factor=ty.offset_factor,
- layout=None,
- allocated_addr=list(ty.allocated_addr),
- buffer_name=name,
- )
- return _T.arg(name, _at(span, annotation))
-
-
-def func_ret_type(annotation, *, span=None):
- """Set the signature's concrete return type."""
- if callable(annotation) and not isinstance(annotation, _ir.Expr):
- annotation = annotation()
- if isinstance(annotation, _ir.Expr):
- annotation = annotation.ty
- with _span_context(span):
- return _T.func_ret(annotation)
-
-
-def _name(value, name, span):
- if name is not None:
- _IRBuilder.name(name, value)
- return _at(span, value)
-
-
-def _enter_concise(frame):
- native = frame.native if isinstance(frame, _Frame) else frame
- native.add_callback(_partial(frame.__exit__, None, None, None))
- return frame.__enter__()
-
-
-def _as_expr(value):
- if isinstance(value, _ffi.ObjectConvertible):
- value = value.asobject()
- if isinstance(value, _ir.Expr):
- return value
- if isinstance(value, str):
- return _ir.StringImm(value)
- if isinstance(value, list | tuple):
- return _ir.Tuple([_as_expr(item) for item in value])
- return _tir.const(value)
-
-
-def bind_(
- value=_MISSING,
- *,
- ty=None,
- name=None,
- span=None,
- name_span=None,
- previous=_MISSING,
- declaration=False,
- frame_value=False,
-):
- """Bind values, or name a frame-owned value without introducing new storage."""
- name_span = span if name_span is None else name_span
- with _span_context(span):
- if frame_value:
- if isinstance(value, _frame.SBlockFrame):
- raise TypeError("A block does not introduce an as-target value")
- if isinstance(value, _python.list | _python.tuple | _ir.Array):
- for index, item in enumerate(value):
- bind_(
- item,
- name=None if name is None else f"{name}_{index}",
- span=span,
- name_span=name_span,
- frame_value=True,
- )
- elif isinstance(value, _ir.Var | _tir.IterVar | _tir.Layout):
- _name(value, name, name_span)
- elif isinstance(value, _ir.TensorLoad) and _tir.is_buffer_var(value.source):
- _name(value.source, name, name_span)
- return value
- if previous is not _MISSING and _tir.is_buffer_var(previous):
- shape = previous.ty.shape
- if len(shape) == 1 and bool(shape[0] == 1):
- if value is _MISSING:
- raise ValueError("A reassignment requires an initializer")
- _T.buffer_store(previous, value, [0])
- return previous
- if previous is not _MISSING and isinstance(getattr(previous, "ty", None), _ir.PointerType):
- raise ValueError(f"Pointer variable {name!r} cannot be reassigned")
- if previous is not _MISSING and (
- _tir.is_buffer_var(previous)
- or isinstance(previous, _tir.IterVar)
- or _python.any(
- isinstance(frame, _frame.SBlockFrame)
- and _python.any(axis.var.same_as(previous) for axis in frame.iter_vars)
- for frame in _IRBuilder.current().frames
- )
- ):
- raise ValueError(f"Cannot rebind buffer or block axis {name!r}")
- if declaration:
- if not _ir.is_prim_var(value):
- raise TypeError("A symbol declaration requires a concrete primitive variable")
- if ty is not None:
- annotation = ty() if callable(ty) else ty
- annotation = annotation.ty if isinstance(annotation, _ir.Expr) else annotation
- if not _ffi.structural_equal(annotation, value.ty):
- raise TypeError("The symbol declaration has an incompatible type")
- if previous is not _MISSING:
- if not _ir.is_prim_var(previous) or not _ffi.structural_equal(
- previous.ty, value.ty
- ):
- raise TypeError("The symbol declaration has an incompatible signature dtype")
- return previous
- return _name(value, name, name_span)
- if previous is not _MISSING and isinstance(previous, _ir.TensorLoad):
- if value is _MISSING:
- raise ValueError("A reassignment requires an initializer")
- _T.buffer_store(previous.source, value, list(previous.indices))
- return previous
- if isinstance(value, _I.meta_var):
- return value.value
- if isinstance(ty, _T.LocalVectorAnnotation):
- if value is not _MISSING:
- raise ValueError("Vector annotation does not support an initializer")
- return _name(_T.alloc_local(ty.shape, ty.dtype), name, name_span)
- if isinstance(ty, _T.LetAnnotation):
- if value is _MISSING:
- raise ValueError("An immutable binding requires an initializer")
- value = _as_expr(value)
- variable = _name(ty.as_var(rhs_dtype=value.ty), name, name_span)
- _T.Bind(value, var=variable)
- return variable
- if ty is not None:
- annotation = ty() if callable(ty) and not isinstance(ty, _ir.Expr) else ty
- annotation = annotation.ty if isinstance(annotation, _ir.Expr) else annotation
- if not isinstance(annotation, _ir.PrimType) or str(annotation) == "handle":
- raise TypeError("Mutable scalar annotations require a primitive scalar type")
- result = _T.local_scalar(str(annotation)).scalar
- _name(result.source, name, name_span)
- if value is not _MISSING:
- _T.buffer_store(result.source, value, [0])
- return result
- if value is _MISSING:
- raise ValueError("An uninitialized binding requires a scalar type annotation")
- if (
- isinstance(value, _ir.TensorLoad)
- and _tir.is_buffer_var(value.source)
- and not value.source.name
- and len(value.source.ty.shape) == 1
- and isinstance(value.source.ty.shape[0], _tir.IntImm)
- and value.source.ty.shape[0].value == 1
- ):
- _name(value.source, name, name_span)
- return value
- if isinstance(value, _T.scalar_wrapper):
- _name(value.scalar.source, name, name_span)
- return value.scalar
- if isinstance(value, _NativeFrame | _Frame):
- return _name(_enter_concise(value), name, name_span)
- if isinstance(value, list | tuple):
- for index, item in enumerate(value):
- bind_(item, name=None if name is None else f"{name}_{index}", span=span)
- return value
- if getattr(type(value), "_is_meta_class", False):
- if name is not None:
- _T.ir.name_meta_class_value(name, value)
- return value
- if _tir.is_buffer_var(value) or isinstance(value, _tir.IterVar | _tir.Layout):
- return _name(value, name, name_span)
- if isinstance(value, _ir.Var) and not value.name:
- return _name(value, name, name_span)
- if isinstance(value, _ir.TensorRegion):
- return value
- if not isinstance(value, _ir.Expr | _python.int | _python.float | _python.bool | str):
- return value
- value = _as_expr(value)
- if _ir.is_prim_expr(value):
- result = _T.local_scalar(str(value.ty.dtype)).scalar
- _name(result.source, name, name_span)
- _T.buffer_store(result.source, value, [0])
- return result
- return _name(_T.Bind(value), name, name_span)
-
-
-def emit_(value, *, span=None):
- """Consume one expression statement, including effect-only calls."""
- if value is None or isinstance(value, str | _ir.Var):
- return
- with _span_context(span):
- if isinstance(value, _NativeFrame | _Frame):
- _enter_concise(value)
- elif hasattr(value, "frames"):
- for frame in value.frames:
- _enter_concise(frame)
- elif isinstance(value, _tir.BufferStore):
- _T.buffer_store(value.buffer, value.value, value.indices)
- else:
- _T.evaluate(value)
-
-
-def setitem(target, key, value, *, span=None):
- """Construct an indexed store after the caller has evaluated its operands."""
- with _span_context(span):
- _T.buffer_store(target, value, key)
-
-
-def setattr(target, name, value, *, span=None):
- """Store through scalar attributes, or update ordinary Python metadata."""
- if isinstance(value, _I.meta_var):
- _python.setattr(target, name, value.value)
- return
- previous = getattr(target, name, _MISSING)
- if isinstance(previous, _T.scalar_wrapper):
- previous = previous.scalar
- buffer = previous.source if isinstance(previous, _ir.TensorLoad) else previous
- if _tir.is_buffer_var(buffer):
- shape = buffer.ty.shape
- if len(shape) == 1 and bool(shape[0] == 1):
- bind_(value, previous=previous, span=span)
- return
- _python.setattr(target, name, value)
-
-
-def return_(value, *, span=None):
- """Emit a native return; subsequent unreachable statements remain in the IR."""
- if value is None:
- raise TypeError("A primitive function return requires an expression")
- with _span_context(span):
- _T.Return(_as_expr(value))
-
-
-def _require_loop():
- for frame in reversed(_IRBuilder.current().frames):
- if isinstance(frame, _frame.ForFrame | _frame.WhileFrame):
- return
- if isinstance(frame, _frame.PrimFuncFrame):
- break
- raise ValueError("Loop control requires an enclosing primitive loop")
-
-
-def break_(*, span=None):
- """Construct a break targeting the nearest primitive loop."""
- _require_loop()
- with _span_context(span):
- _T.evaluate(_T.break_loop())
-
-
-def continue_(*, span=None):
- """Construct a continue targeting the nearest primitive loop."""
- _require_loop()
- with _span_context(span):
- _T.evaluate(_T.continue_loop())
-
-
-def assert_(condition, message="", *, span=None):
- """Emit the native flat assertion with its own source range."""
- kind = "RuntimeError"
- if isinstance(message, tuple):
- if len(message) != 2 or not isinstance(message[0], str):
- raise TypeError("Assertion metadata must be (error_kind, message_parts)")
- kind, message = message
- if isinstance(message, list | tuple):
- message = [str(part) for part in message]
- if not isinstance(message, list | tuple):
- message = [message]
- with _span_context(span):
- with _T.Assert(condition, message, error_kind=kind):
- pass
-
-
-def If(condition, *, span=None):
- with _span_context(span):
- return _Frame(_T.If(condition), span)
-
-
-def Then(*, span=None):
- with _span_context(span):
- return _Frame(_T.Then(), span)
-
-
-def Else(*, span=None):
- with _span_context(span):
- return _Frame(_T.Else(), span)
-
-
-def For(iterable, *, span=None):
- """Adapt a native loop frame, or a Python range, to construction scope."""
- if isinstance(iterable, _python.range):
- iterable = _T.serial(iterable.start, iterable.stop, step=iterable.step)
- if not isinstance(iterable, _frame.ForFrame):
- raise TypeError("A primitive for loop requires a native loop frame or range")
- return _Frame(iterable, span)
-
-
-def While(condition, *, span=None):
- with _span_context(span):
- return _Frame(_T.While(condition), span)
-
-
-def unpack(value):
- """Project a concrete IR tuple of known arity, preserving Python iteration."""
- if isinstance(value, _ir.Tuple):
- return _python.tuple(value.fields)
- if isinstance(value, _ir.Expr) and isinstance(value.ty, _ir.TupleType):
- return _python.tuple(_ir.TupleGetItem(value, i) for i in range(len(value.ty.fields)))
- return value
-
-
-def alloc_scalar(dtype="float32", scope="global"):
- """Allocate scalar storage and return its concrete load expression."""
- value = _T.alloc_scalar(dtype, scope)
- return value.scalar if isinstance(value, _T.scalar_wrapper) else value
-
-
-def local_scalar(dtype="float32"):
- return alloc_scalar(dtype, "local")
-
-
-def shared_scalar(dtype="float32"):
- return alloc_scalar(dtype, "shared")
-
-
-@_expression_args(
- "shape",
- "strides",
- "elem_offset",
- introduce=True,
- implicit_dtype="int32",
- compound_declarations=True,
-)
-@_wraps(_T.match_buffer)
-def match_buffer(*args, **kwargs):
- """Construct a native buffer match with resolved symbolic shape fields."""
- return _T.match_buffer(*args, **kwargs)
-
-
-# Constructor identities carry syntax policy; aliases share it without wrappers.
-for _constructor in vars(_T).values():
- if isinstance(_constructor, _T.DtypeConstructor):
- _register_declaration(_constructor, dtype=_constructor._dtype_str)
-del _constructor
-
-
-def range_(*args):
- """Construct a serial loop from the source builtin range arguments."""
- if len(args) in (1, 2):
- return _T.serial(*args)
- if len(args) != 3:
- raise TypeError("range expects one to three arguments")
- start, stop, step = args
- if isinstance(step, _python.int) and step == 0:
- raise ValueError("range step cannot be zero")
- return _T.serial(start, stop, step=step)
-
-
-__tvm_call_overrides__ = {_python.range: range_}
-
-
-def logical_and(*values):
- """Construct scalar/vector conjunction, preserving ordinary Python values."""
- if not values:
- raise TypeError("logical_and requires at least one operand")
- values = [
- value.asobject() if isinstance(value, _ffi.ObjectConvertible) else value for value in values
- ]
- result = values[0]
- for value in values[1:]:
- if not isinstance(result, _ir.Expr) and not isinstance(value, _ir.Expr):
- result = result and value
- else:
- lhs, rhs = _as_expr(result), _as_expr(value)
- result = _tir.And(lhs, rhs) if lhs.ty.is_scalar() and rhs.ty.is_scalar() else lhs & rhs
- return result
-
-
-def logical_or(*values):
- """Construct scalar/vector disjunction, preserving ordinary Python values."""
- if not values:
- raise TypeError("logical_or requires at least one operand")
- values = [
- value.asobject() if isinstance(value, _ffi.ObjectConvertible) else value for value in values
- ]
- result = values[0]
- for value in values[1:]:
- if not isinstance(result, _ir.Expr) and not isinstance(value, _ir.Expr):
- result = result or value
- else:
- lhs, rhs = _as_expr(result), _as_expr(value)
- result = _tir.Or(lhs, rhs) if lhs.ty.is_scalar() and rhs.ty.is_scalar() else lhs | rhs
- return result
-
-
-def logical_not(value):
- """Construct IR negation without coercing an IR expression to Python bool."""
- if isinstance(value, _ffi.ObjectConvertible):
- value = value.asobject()
- return _tir.Not(value) if isinstance(value, _ir.Expr) else not value
-
-
-def select(condition, true_value, false_value):
- """Construct a conditional expression or select an ordinary Python value."""
- if isinstance(condition, _ffi.ObjectConvertible):
- condition = condition.asobject()
- if not isinstance(condition, _ir.Expr):
- return true_value if condition else false_value
- return _tir.if_then_else(condition, true_value, false_value)
-
-
-def _global_callee(function):
- return _partial(_T.ir._call_global, function)
-
-
-_register_call_kind(_sys.modules[__name__], _ir.GlobalVar, _global_callee)
-
-
-from tvm.script.parser_v2.annotations import enable_eager_constructors as _enable_eager_constructors
-
-_enable_eager_constructors(_sys.modules[__name__], classes=("Buffer",))
diff --git a/python/tvm/tirx/script/parser/__init__.py b/python/tvm/tirx/script/parser/__init__.py
deleted file mode 100644
index 2b4301b0be..0000000000
--- a/python/tvm/tirx/script/parser/__init__.py
+++ /dev/null
@@ -1,50 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# ruff: noqa: RUF005
-"""The tirx parser"""
-
-from typing import TYPE_CHECKING
-
-from tvm.tirx.script.builder import * # pylint: disable=redefined-builtin
-from tvm.tirx.script.builder import ir as _tir
-
-from . import operation as _operation
-from . import parser as _parser
-from .entry import Buffer, Optional, Ptr, constexpr
-
-if TYPE_CHECKING:
- # pylint: disable=invalid-name
- # Define prim_func and make it type check as static method
- # so most tvmscript won't trigger pylint error here.
- prim_func = staticmethod
- jit = staticmethod
-else:
- from .entry import inline, jit, macro, prim_func
-
-__all__ = _tir.__all__ + [
- "Buffer",
- "Optional",
- "Ptr",
- "SMEMPool",
- "TMEMPool",
- "bool",
- "constexpr",
- "inline",
- "jit",
- "macro",
- "prim_func",
-]
diff --git a/python/tvm/tirx/script/parser/entry.py b/python/tvm/tirx/script/parser/entry.py
deleted file mode 100644
index 03bcc9d643..0000000000
--- a/python/tvm/tirx/script/parser/entry.py
+++ /dev/null
@@ -1,581 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""The entry point of TVM parser for tirx."""
-
-import ast
-import inspect
-from collections.abc import Callable
-from typing import Any
-
-from tvm.ir.base import deprecated
-from tvm.script.parser._core import parse, scan_macro, utils
-from tvm.script.parser.core.parser import Parser, ScriptMacro, VarTable
-from tvm.tirx import Buffer, PrimFunc
-from tvm.tirx.script.builder import block_name_suffix_context, buffer, ptr
-
-
-def prim_func(
- func: Callable | None = None,
- private: bool = False,
- check_well_formed=True,
- s_tir: bool = False,
- persistent: bool = False,
-) -> PrimFunc | Callable:
- """The parsing method for tirx prim func, by using `@prim_func` as decorator.
-
- Parameters
- ----------
- func : Callable
- The function to be parsed as prim func.
- (Listed as optional to allow the decorator to be used
- without arguments, like `@prim_func`,
- or with an argument, `@prim_func(private=True)`)
-
- private : bool, optional
- Whether the function should be treated as private.
- A private function has no global symbol attribute;
- if the function is not private, it will have a global symbol
- matching the function name.
-
- Returns
- -------
- res : Union[PrimFunc, Callable]
- The parsed tirx prim func.
- """
- # pylint: disable=unused-argument
- # (private will be used in the parser, but not immediately)
-
- # need to capture this var outside the wrapper because the wrapper
- # adds to the stack
- outer_stack = inspect.stack()
-
- def decorator_wrapper(func):
- if not inspect.isfunction(func):
- raise TypeError(f"Expect a function, but got: {func}")
- if utils.is_defined_in_class(outer_stack, func):
- return func
- extra_vars = utils.inspect_function_capture(func)
- utils.resolve_closure_vars(func, extra_vars, outer_stack)
- f = parse(func, extra_vars, check_well_formed=check_well_formed)
- setattr(f, "__name__", func.__name__)
- return f
-
- if func is not None:
- # no optional args given => use wrapper directly
- return decorator_wrapper(func)
- else:
- # if there is an optional arg given, return a new decorator
- # that will then be invoked
- setattr(decorator_wrapper, "dispatch_token", "tirx")
- return decorator_wrapper
-
-
-setattr(prim_func, "dispatch_token", "tirx")
-
-
-class TIRInline(ScriptMacro):
- """Specialization of ScriptMacro for TIR with Python LEGB scoping.
-
- Two definition paths:
- 1. Outside @T.prim_func (standalone @T.inline): definition_depth is None,
- closure_vars captured at definition time are used (module globals are
- effectively late-bound since they don't change during parsing).
- 2. Inside @T.prim_func (inline def in parsed body): definition_depth is set
- to the VarTable frame depth at definition time, and defining_var_table
- stores a reference to the VarTable that was active. At call time,
- defining_var_table.get_at_depth(definition_depth) reads current values
- from the lexically enclosing frames.
-
- Attributes
- ----------
- definition_depth : Optional[int]
- VarTable frame depth at definition time, or None for outside-prim_func.
- defining_var_table : Optional[VarTable]
- Reference to the VarTable that was active at definition time.
- call_count : int
- Counter for unique block name suffixes.
- """
-
- def __init__(
- self,
- source,
- closure_vars: dict[str, Any],
- func: Callable,
- definition_depth: int | None = None,
- defining_var_table: VarTable | None = None,
- ) -> None:
- # hygienic=True for the base class (field kept for compat but not used in dispatch)
- super().__init__(source, closure_vars, func, hygienic=True)
- self.definition_depth = definition_depth
- self.defining_var_table = defining_var_table
- self.call_count = 0
-
- def parse_macro(self, parser: Parser) -> None:
- macro_def = self.get_macro_def()
- suffix = f"_{self.call_count}" if self.call_count > 0 else ""
- self.call_count += 1
- with block_name_suffix_context(suffix):
- parser.visit_body(macro_def.body)
-
- def __call__(self, *args, **kwargs):
- param_binding = inspect.signature(self.func).bind(*args, **kwargs)
- param_binding.apply_defaults()
- local_vars = param_binding.arguments
- parser = self._find_parser_def()
-
- with parser.with_diag_source(self.source):
- if self.defining_var_table is not None:
- # Inside-prim_func path: LEGB late binding from the defining scope
- enclosing_vars = self.defining_var_table.get_at_depth(self.definition_depth)
- else:
- # Outside-prim_func path: use captured closure vars
- enclosing_vars = self.closure_vars
-
- saved_var_table = parser.var_table
- parser.var_table = VarTable()
-
- with parser.var_table.with_frame():
- for k, v in enclosing_vars.items():
- parser.var_table.add(k, v)
- with parser.var_table.with_frame():
- for k, v in local_vars.items():
- parser.var_table.add(k, v)
-
- parse_result = self.parse_macro(parser)
-
- parser.var_table = saved_var_table
-
- return parse_result
-
-
-def inline(*args, definition_depth: int | None = None, defining_var_table=None) -> Callable:
- """Decorator for inline function definitions with Python LEGB scoping.
-
- @T.inline follows Python's lexical scoping with late binding:
- - At definition time, record which scopes are visible.
- - At call time, read current values from those scopes.
-
- Example::
-
- import tvm
- from tvm.script import tirx as T
-
- x_value = 128
-
- @T.inline
- def capture(A, B):
- B[()] = A[x_value] # x_value resolved from enclosing scope
-
- @T.prim_func(s_tir=True)
- def use(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
- capture(A, B) # Produces B[()] = A[128]
- """
-
- def _decorator(func: Callable) -> Callable:
- source, closure_vars = scan_macro(func, utils.inspect_function_capture(func))
- obj = TIRInline(
- source,
- closure_vars,
- func,
- definition_depth=definition_depth,
- defining_var_table=defining_var_table,
- )
-
- def wrapper(*args, **kwargs):
- return obj(*args, **kwargs)
-
- return wrapper
-
- if len(args) == 0:
- setattr(_decorator, "dispatch_token", "tir.inline")
- return _decorator
- if len(args) == 1 and inspect.isfunction(args[0]):
- return _decorator(args[0])
-
- raise ValueError("Invalid use of T.inline. Usage: @T.inline or @T.inline()")
-
-
-setattr(inline, "dispatch_token", "tir.inline")
-
-
-class TIRJit:
- """Top-level kernel decorator with compile-time ``.specialize()`` params.
-
- Parses the function body lazily: parsing is deferred until ``.specialize()``
- supplies concrete values for the params annotated as ``T.constexpr``. The
- return type of ``.specialize()`` is a ``tvm.tirx.PrimFunc``, identical in
- type to what ``@T.prim_func`` produces today.
-
- Constexpr params are removed from the resulting PrimFunc's parameter list;
- their values are baked into the IR (e.g. into ``T.Buffer((M, K), ...)``
- shape annotations and into the body).
- """
-
- def __init__(
- self,
- func: Callable,
- check_well_formed: bool = True,
- is_stir: bool = False,
- persistent: bool = False,
- private: bool = False,
- ) -> None:
- self.func = func
- self.check_well_formed = check_well_formed
- self.is_stir = is_stir
- self.persistent = persistent # pylint: disable=unused-private-member
- self.private = private # pylint: disable=unused-private-member
- # Resolved closure vars (computed once; the function itself is the
- # capture point, so this never changes between specializations).
- self._closure_vars: dict[str, Any] = utils.inspect_function_capture(func)
- # Detect which params are marked T.constexpr or T.Optional. With PEP 563
- # (``from __future__ import annotations``), each annotation is a
- # string; we eval them one-by-one so a constexpr probe is not
- # blocked by sibling annotations that reference yet-undefined names
- # (e.g. ``A: T.Buffer((N,), ...)`` referencing constexpr ``N``).
- raw_anns = getattr(func, "__annotations__", {}) or {}
- eval_globals = {**func.__globals__, **self._closure_vars}
- sig = inspect.signature(func)
- constexpr_names: set[str] = set()
- constexpr_defaults: dict[str, Any] = {}
- optional_names: set[str] = set()
- for name, param in sig.parameters.items():
- raw_ann = raw_anns.get(name)
- ann = raw_ann
- if isinstance(ann, str):
- try:
- ann = eval(ann, eval_globals) # pylint: disable=eval-used
- except Exception: # pylint: disable=broad-except
- ann = None
- if ann is constexpr:
- constexpr_names.add(name)
- if param.default is not inspect.Parameter.empty:
- constexpr_defaults[name] = param.default
- if isinstance(ann, _OptionalAnnotation) or _is_explicit_optional_annotation(raw_ann):
- optional_names.add(name)
- self.constexpr_names: frozenset[str] = frozenset(constexpr_names)
- self.constexpr_defaults: dict[str, Any] = constexpr_defaults
- self.optional_names: frozenset[str] = frozenset(optional_names)
- self._cache: dict[tuple, PrimFunc] = {}
-
- def specialize(self, **specialization_kwargs) -> PrimFunc:
- """Build a PrimFunc by binding constexprs and absent optional params.
-
- Parameters
- ----------
- **specialization_kwargs
- One value per ``T.constexpr``-annotated parameter. A
- ``T.Optional`` parameter may additionally be supplied as ``None``
- to remove it from the resulting PrimFunc ABI. Omitting an
- optional parameter keeps it as a normal runtime parameter.
-
- Returns
- -------
- PrimFunc
- A concrete TIRx PrimFunc, identical in type to the output of
- ``@T.prim_func``.
- """
- specializable_names = self.constexpr_names | self.optional_names
- extra = specialization_kwargs.keys() - specializable_names
- if extra:
- raise TypeError(
- f"{self.func.__name__}.specialize() got unexpected arg(s): "
- f"{sorted(extra)} (specializable params are: {sorted(specializable_names)})"
- )
- invalid_optional = {
- name: specialization_kwargs[name]
- for name in self.optional_names & specialization_kwargs.keys()
- if specialization_kwargs[name] is not None
- }
- if invalid_optional:
- raise TypeError(
- f"{self.func.__name__}.specialize(): T.Optional parameters only accept None; "
- f"pass actual tensors when calling the compiled kernel (got: {invalid_optional!r})"
- )
-
- supplied_constexprs = {
- name: specialization_kwargs[name]
- for name in self.constexpr_names & specialization_kwargs.keys()
- }
- effective = {**self.constexpr_defaults, **supplied_constexprs}
- missing = self.constexpr_names - effective.keys()
- if missing:
- raise TypeError(
- f"{self.func.__name__}.specialize() missing constexpr arg(s) "
- f"(no default provided): {sorted(missing)}"
- )
-
- absent_params = {name: None for name in self.optional_names & specialization_kwargs.keys()}
- try:
- cache_key = (tuple(sorted(effective.items())), tuple(sorted(absent_params)))
- cached = self._cache.get(cache_key)
- except TypeError as err:
- raise TypeError(
- f"{self.func.__name__}.specialize(): all constexpr values must "
- f"be hashable (got: {effective!r})"
- ) from err
- if cached is not None:
- return cached
-
- extra_vars = {**self._closure_vars, **effective}
- prim_func = parse(
- self.func,
- extra_vars,
- check_well_formed=self.check_well_formed,
- absent_params=absent_params,
- )
- setattr(prim_func, "__name__", self.func.__name__)
- self._cache[cache_key] = prim_func
- return prim_func
-
-
-def jit(
- func: Callable | None = None,
- private: bool = False,
- check_well_formed: bool = True,
- is_stir: bool = False,
- persistent: bool = False,
-) -> "TIRJit | Callable":
- """Decorator: capture the kernel and defer parsing until ``.specialize()``.
-
- Use ``@T.jit`` (instead of ``@T.prim_func``) when the kernel takes
- compile-time parameters annotated with ``T.constexpr`` or runtime
- parameters that may be removed with ``T.Optional``. The resulting object
- exposes ``.specialize(**specialization_kwargs)``, which returns a
- ``tvm.tirx.PrimFunc``.
-
- Example::
-
- from tvm.script import tirx as T
-
- @T.jit
- def add(
- A: T.Buffer((N,), "float32"),
- B: T.Buffer((N,), "float32"),
- *,
- N: T.constexpr,
- ):
- ...
-
- kernel = add.specialize(N=1024) # returns a PrimFunc
-
- @T.jit
- def guarded(optional: T.Optional(T.handle), out: T.handle):
- if optional is not None:
- ...
-
- present = guarded.specialize()
- absent = guarded.specialize(optional=None)
- """
-
- def decorator_wrapper(func: Callable) -> TIRJit:
- if not inspect.isfunction(func):
- raise TypeError(f"Expect a function, but got: {func}")
- return TIRJit(
- func,
- check_well_formed=check_well_formed,
- is_stir=is_stir,
- persistent=persistent,
- private=private,
- )
-
- if func is not None:
- return decorator_wrapper(func)
- setattr(decorator_wrapper, "dispatch_token", "tirx")
- return decorator_wrapper
-
-
-setattr(jit, "dispatch_token", "tirx")
-
-
-class TIRMacro(ScriptMacro):
- """Specialization of the ScriptMacro class for TIR.
-
- Apache-compatible hygienic macro. Distinct from ``TIRInline`` (which
- uses Python LEGB late binding) so upstream code that relies on
- capture-at-definition-time semantics keeps working.
-
- Attributes
- ----------
- call_count : int
- Counter for the number of times this macro has been invoked.
- Used to generate unique block name suffixes.
- """
-
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.call_count = 0
-
- def parse_macro(self, parser: Parser) -> None:
- macro_def = self.get_macro_def()
- suffix = f"_{self.call_count}" if self.call_count > 0 else ""
- self.call_count += 1
- with block_name_suffix_context(suffix):
- parser.visit_body(macro_def.body)
-
-
-def macro(*args, hygienic: bool = True) -> Callable:
- """Decorator for macro definitions with hygienic capture.
-
- Parameters
- ----------
- hygienic: bool
- Specifies whether the macro is hygienic or not. A hygienic macro
- resolves symbols at definition time; a non-hygienic macro at use
- time. Defaults to ``True``.
- """
-
- def _decorator(func: Callable) -> TIRMacro:
- source, closure_vars = scan_macro(func, utils.inspect_function_capture(func))
- obj = TIRMacro(source, closure_vars, func, hygienic)
-
- def wrapper(*args, **kwargs):
- return obj(*args, **kwargs)
-
- return wrapper
-
- if len(args) == 0:
- return _decorator
- if len(args) == 1 and inspect.isfunction(args[0]):
- return _decorator(args[0])
-
- raise ValueError("Invalid use of T.macro. Usage: @T.macro or @T.macro()")
-
-
-setattr(macro, "dispatch_token", "tir.macro")
-
-
-class BufferProxy:
- """Buffer proxy class for constructing tirx buffer."""
-
- def __or__(self, other):
- """Support ``T.Buffer | None`` union syntax in annotations."""
- return self
-
- def __ror__(self, other):
- """Support ``None | T.Buffer`` union syntax in annotations."""
- return self
-
- def __call__(
- self,
- shape,
- dtype="float32",
- data=None,
- strides=None,
- elem_offset=None,
- byte_offset=None,
- scope="global",
- align=0,
- offset_factor=0,
- layout="default",
- ) -> Buffer:
- return buffer(
- shape,
- dtype=dtype,
- data=data,
- strides=strides,
- elem_offset=elem_offset,
- byte_offset=byte_offset,
- scope=scope,
- align=align,
- offset_factor=offset_factor,
- layout=layout,
- )
-
- @deprecated("T.Buffer[...]", "T.Buffer(...)")
- def __getitem__(self, keys) -> Buffer:
- if not isinstance(keys, tuple):
- return self(keys)
- if len(keys) >= 2 and not isinstance(keys[1], str):
- return self(keys)
- return self(*keys) # type: ignore[attr-defined] # pylint: disable=no-member
-
-
-class PtrProxy:
- """Ptr proxy class for constructing tirx pointer."""
-
- def __or__(self, other):
- """Support union syntax in annotations."""
- return self
-
- def __ror__(self, other):
- """Support union syntax in annotations."""
- return self
-
- @deprecated("T.Ptr(...)", "T.handle(...)")
- def __call__(self, dtype, storage_scope="global"):
- if callable(dtype):
- dtype = dtype().ty.dtype
- return ptr(dtype, storage_scope) # type: ignore[attr-defined] # pylint: disable=no-member
-
- @deprecated("T.Ptr[...]", "T.handle(...)")
- def __getitem__(self, keys):
- if not isinstance(keys, tuple):
- return self(keys)
- return self(*keys)
-
-
-class _ConstexprProxy:
- """Sentinel marker for compile-time (specialization-time) parameters.
-
- Used as a parameter annotation in ``@T.jit`` decorated functions to mark
- a parameter as constexpr — its value is supplied to ``.specialize(**kwargs)``
- rather than at call time, and it is removed from the generated PrimFunc's
- runtime parameter list.
- """
-
- def __or__(self, other):
- return self
-
- def __ror__(self, other):
- return self
-
-
-class _OptionalAnnotation:
- """Wrapper around the runtime annotation of an optional JIT parameter."""
-
- def __init__(self, annotation: Any) -> None:
- self.annotation = annotation
-
-
-class _OptionalProxy:
- """Construct an explicit ``T.Optional(...)`` JIT parameter annotation."""
-
- def __call__(self, annotation: Any) -> _OptionalAnnotation:
- return _OptionalAnnotation(annotation)
-
-
-def _is_explicit_optional_annotation(annotation: Any) -> bool:
- """Recognize deferred ``T.Optional(...)`` without accepting typing unions."""
-
- if not isinstance(annotation, str):
- return False
- try:
- node = ast.parse(annotation, mode="eval").body
- except SyntaxError:
- return False
- if not isinstance(node, ast.Call):
- return False
- func = node.func
- return (isinstance(func, ast.Attribute) and func.attr == "Optional") or (
- isinstance(func, ast.Name) and func.id == "Optional"
- )
-
-
-Buffer = BufferProxy() # pylint: disable=invalid-name
-Ptr = PtrProxy() # pylint: disable=invalid-name
-constexpr = _ConstexprProxy() # pylint: disable=invalid-name
-Optional = _OptionalProxy() # pylint: disable=invalid-name
diff --git a/python/tvm/tirx/script/parser/operation.py b/python/tvm/tirx/script/parser/operation.py
deleted file mode 100644
index fd67d6f125..0000000000
--- a/python/tvm/tirx/script/parser/operation.py
+++ /dev/null
@@ -1,167 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""The tirx expression operation registration"""
-
-import tvm
-from tvm import tirx
-from tvm.ir import PrimType
-from tvm.runtime import DataTypeCode
-from tvm.script.parser._core import OpMethod, doc, register_op
-from tvm.tirx import IntImm
-from tvm.tirx.expr import FloatImm
-
-
-def _register_expr_op(ty: type): # pylint: disable=invalid-name
- ty._dispatch_type = ty # pylint: disable=protected-access
-
- def _expr_ty(expr):
- ty = expr.ty if tvm.ir.is_prim_expr(expr) else None
- if not isinstance(ty, PrimType):
- ty = expr.expr_ty()
- if not isinstance(ty, PrimType):
- raise TypeError(f"Expected a PrimType expression, but got {ty}")
- return ty
-
- def _and(a, b):
- if isinstance(a, bool):
- a = IntImm("bool", a)
- if isinstance(b, bool):
- b = IntImm("bool", b)
- if not _expr_ty(a).is_scalar() or not _expr_ty(b).is_scalar():
- return a & b
- else:
- return tirx.And(a, b)
-
- def _or(a, b):
- if isinstance(a, bool):
- a = IntImm("bool", a)
- if isinstance(b, bool):
- b = IntImm("bool", b)
- if not _expr_ty(a).is_scalar() or not _expr_ty(b).is_scalar():
- return a | b
- else:
- return tirx.Or(a, b)
-
- def _get_type_str(ty: PrimType):
- dtype_str = str(ty.dtype)
- if ty.is_scalar():
- return dtype_str
- index = dtype_str.find("x")
- return dtype_str[0:index]
-
- def _auto_broadcast(a, b, op):
- if isinstance(a, int):
- if tvm.ir.is_prim_expr(b) or hasattr(b, "expr_ty"):
- b_ty = _expr_ty(b)
- if b_ty.matches_code(DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL):
- a = IntImm(_get_type_str(b_ty), a)
- elif b_ty.matches_code(DataTypeCode.FLOAT):
- a = FloatImm(_get_type_str(b_ty), a)
- elif isinstance(b, float):
- a = FloatImm("float32", a)
- else:
- a = IntImm("int32", a)
- elif isinstance(a, float):
- b_ty = _expr_ty(b)
- if b_ty.matches_code(DataTypeCode.FLOAT):
- a = FloatImm(_get_type_str(b_ty), a)
- else:
- a = FloatImm("float32", a)
-
- assert tvm.ir.is_prim_expr(a), "Operand should be a Expr."
- if isinstance(b, int):
- a_ty = _expr_ty(a)
- if a_ty.matches_code(DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL):
- b = IntImm(_get_type_str(a_ty), b)
- elif a_ty.matches_code(DataTypeCode.FLOAT):
- b = FloatImm(_get_type_str(a_ty), b)
- elif isinstance(b, float):
- b = FloatImm(_get_type_str(_expr_ty(a)), b)
-
- a_ty = _expr_ty(a)
- b_ty = _expr_ty(b)
- if a_ty.dtype.lanes == b_ty.dtype.lanes:
- return op(a, b)
- elif a_ty.is_scalar() and a_ty.dtype.lanes != b_ty.dtype.lanes:
- broadcast_a = tirx.Broadcast(a, b_ty.dtype.lanes)
- return op(broadcast_a, b)
- elif b_ty.is_scalar() and a_ty.dtype.lanes != b_ty.dtype.lanes:
- broadcast_b = tirx.Broadcast(b, a_ty.dtype.lanes)
- return op(a, broadcast_b)
- else:
- raise TypeError("do not know how to deal with it.")
-
- def _eq(a, b):
- return _auto_broadcast(a, b, tirx.EQ)
-
- def _ne(a, b):
- return _auto_broadcast(a, b, tirx.NE)
-
- def _lt(a, b):
- return _auto_broadcast(a, b, tirx.LT)
-
- def _le(a, b):
- return _auto_broadcast(a, b, tirx.LE)
-
- def _gt(a, b):
- return _auto_broadcast(a, b, tirx.GT)
-
- def _ge(a, b):
- return _auto_broadcast(a, b, tirx.GE)
-
- def r(op: type, i: int, m: OpMethod): # pylint: disable=invalid-name
- register_op(ty, op, i)(m)
-
- for i in [0, 1]:
- # Case 1. binop
- # doc.Add <-- is overloaded
- # doc.Sub <-- is overloaded
- # doc.Mult <-- is overloaded
- # doc.Div <-- is overloaded
- # doc.FloorDiv <-- is overloaded
- # doc.Mod <-- is overloaded
- # doc.LShift <-- is overloaded
- # doc.RShift <-- is overloaded
- # doc.BitOr <-- is overloaded
- # doc.BitXor <-- is overloaded
- # doc.BitAnd <-- is overloaded
- # doc.MatMult <-- not implemented
- # doc.Pow <-- not implemented
- # Case 2. cmpop
- r(doc.Eq, i, _eq)
- r(doc.NotEq, i, _ne)
- r(doc.Lt, i, _lt)
- r(doc.LtE, i, _le)
- r(doc.Gt, i, _gt)
- r(doc.GtE, i, _ge)
- # doc.Is <-- not implemented
- # doc.IsNot <-- not implemented
- # doc.In <-- not implemented
- # doc.NotIn <-- not implemented
- # Case 3. boolop
- r(doc.And, i, _and)
- r(doc.Or, i, _or)
- for i in [0]:
- # Case 4. unaryop
- # doc.Invert <-- is overloaded
- r(doc.Not, i, tirx.Not)
- # doc.UAdd <-- is overloaded
- # doc.USub <-- is overloaded
-
-
-_register_expr_op(tirx.Expr)
-_register_expr_op(tirx.IterVar)
diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py
deleted file mode 100644
index 0c13295533..0000000000
--- a/python/tvm/tirx/script/parser/parser.py
+++ /dev/null
@@ -1,1217 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""The base parser for tirx"""
-
-import ast
-import contextlib
-from copy import deepcopy
-from functools import partial
-from typing import Any, TypeVar
-
-import tvm
-from tvm.ir import Expr, GlobalVar, PointerType, PrimType, TensorLoad, TensorRegion
-from tvm.script.ir_builder import ir as I
-from tvm.script.ir_builder.base import IRBuilder
-from tvm.script.ir_builder.base import IRBuilderFrame as Frame
-from tvm.script.parser._core import Parser, collect_signature_type_vars, dispatch, doc
-from tvm.script.parser.core.doc import from_doc
-from tvm.tirx import Buffer, IterVar, Layout, buffer_data, is_buffer_var
-from tvm.tirx.script import builder as T
-from tvm.tirx.script.builder.ir import _call_global, name_meta_class_value
-
-from .entry import _OptionalAnnotation, inline
-from .entry import constexpr as _constexpr_sentinel
-
-
-@dispatch.register(token="tirx", type_name="enter_token")
-def enter_token(self: Parser) -> dict[str, Any]:
- context = {"GlobalVar.__call__": GlobalVar.__call__}
- GlobalVar.__call__ = _call_global
- return context
-
-
-@dispatch.register(token="tirx", type_name="exit_token")
-def exit_token(self: Parser, context: dict[str, Any]) -> None:
- GlobalVar.__call__ = context["GlobalVar.__call__"]
-
-
-def slice_buffer_from_region(br: TensorRegion) -> Buffer:
- """Create a matched DeclBuffer from a TensorRegion.
-
- Slices the layout (if present) or computes elem_offset for the sub-region,
- producing a DeclBuffer that views the same underlying data.
- """
- import functools # pylint: disable=import-outside-toplevel
-
- if not is_buffer_var(br.source):
- raise TypeError("A matched buffer region requires a BufferVar source")
- buf = br.source
- region = br.region
- new_shape = [r.extent for r in region]
- sliced_layout = None
- if buf.ty.layout is not None:
- range_pairs = [(r.min, r.min + r.extent) for r in region]
- sliced_layout = buf.ty.layout.slice(list(buf.ty.shape), range_pairs)
- if sliced_layout is not None:
- return T.decl_buffer(
- new_shape,
- buf.ty.dtype,
- buffer_data(buf),
- buf.ty.strides,
- buf.ty.elem_offset,
- None,
- buf.ty.storage_scope,
- buf.ty.data_alignment,
- buf.ty.offset_factor,
- sliced_layout,
- )
- # Fallback: compute elem_offset for default/no layout
- strides = []
- for i in range(len(buf.ty.shape)):
- stride = functools.reduce(
- lambda x, y: x * y, buf.ty.shape[i + 1 :], tvm.tirx.const(1, "int32")
- )
- strides.append(stride)
- offset = tvm.tirx.const(0, "int32")
- for i, r in enumerate(region):
- offset = offset + r.min * strides[i]
- new_elem_offset = buf.ty.elem_offset + offset
- return T.decl_buffer(
- new_shape,
- buf.ty.dtype,
- buffer_data(buf),
- buf.ty.strides,
- new_elem_offset,
- None,
- buf.ty.storage_scope,
- buf.ty.data_alignment,
- buf.ty.offset_factor,
- buf.ty.layout,
- )
-
-
-def bind_with_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
- """Value binding methods when parsing with statement.
- e.g. binding i, j, k with T.grid(128, 128, 128), when parsing
- with T.grid(128, 128, 18) as i, j, k.
-
- Parameters
- ----------
- self : Parser
- The current parser.
-
- node : doc.expr
- The doc AST expression node for error reporting.
-
- var_name : str
- The variable name.
-
- value : Any
- The value to be bound with.
-
- Returns
- -------
- res : Any
- The bound value.
- """
- if isinstance(value, list | tuple):
- for i, v in enumerate(value):
- bind_with_value(self, node, f"{var_name}_{i}", v)
- return value
- elif isinstance(value, tvm.ir.Var):
- IRBuilder.name(var_name, value)
- return value
- else:
- self.report_error(node, f"Do not know how to bind type: {type(value)} in with statement")
- raise NotImplementedError
-
-
-def bind_for_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
- """Value binding methods when parsing for statement.
- e.g. binding i, j, k with T.grid(128, 128, 128), when parsing
- for i, j, k in T.grid(128, 128, 128).
-
- Parameters
- ----------
- self : Parser
- The current parser.
-
- node : doc.expr
- The doc AST expression node for error reporting.
-
- var_name : str
- The variable name.
-
- value : Any
- The value to be bound with.
-
- Returns
- -------
- res : Any
- The bound value.
- """
- if isinstance(value, list | tuple | tvm.ir.Array):
- for i, v in enumerate(value):
- bind_for_value(self, node, f"{var_name}_{i}", v)
- return value
- elif isinstance(value, tvm.ir.Var):
- IRBuilder.name(var_name, value)
- return value
- else:
- self.report_error(node, f"Do not know how to bind type: {type(value)} in for statement")
- raise NotImplementedError
-
-
-def _convert_tuple_literal(self: Parser, node: doc.expr, value: Any) -> Any:
- """Convert an explicit list/tuple T.let value to a core IR Tuple.
-
- The generic evaluator deliberately keeps Python containers intact because
- many TIRx builder APIs use them structurally. Conversion is restricted to
- the internal immutable-binding boundary that explicitly opts in below.
- """
- if not isinstance(node, doc.List | doc.Tuple):
- return value
-
- def convert_field(field: Any) -> Expr:
- if isinstance(field, list | tuple):
- return tvm.ir.Tuple([convert_field(item) for item in field])
- if isinstance(field, Expr):
- return field
- if isinstance(field, str):
- return tvm.ir.StringImm(field)
- if isinstance(field, bool | int | float):
- return tvm.tirx.const(field)
- self.report_error(
- node,
- f"Tuple fields must be expressions or scalar literals, got {type(field).__name__}",
- )
- raise NotImplementedError
-
- return convert_field(value)
-
-
-def bind_assign_value(
- self: Parser,
- node: doc.expr,
- var_name: str,
- value: Any,
- prim_var_declarations: set[str] | None = None,
-) -> Any:
- """Value binding methods when parsing assign statement.
- e.g. binding vi, vj, vk with T.axis.remap("SSR", [i, j, k]), when parsing
- vi, vj, vk = T.axis.remap("SSR", [i, j, k]).
-
- Parameters
- ----------
- self : Parser
- The current parser.
-
- node : doc.expr
- The doc AST expression node for error reporting.
-
- var_name : str
- The variable name.
-
- value : Any
- The value to be bound with.
-
- Returns
- -------
- res : Any
- The bound value.
- """
- if var_name in (prim_var_declarations or set()):
- # A quoted Buffer shape may have already created this PrimVar. In that
- # case ``n = T.int32()`` is match-like syntax: bind the Python name to
- # the signature Var instead of emitting a second definition.
- if _get_signature_match_var(self, var_name) is not None:
- return _reuse_signature_match_var(self, node, var_name, value)
- if isinstance(value, T.scalar_wrapper): # pylint: disable=protected-access
- # special case for scalar, name the buffer, but the var is used as TensorLoad
- assert isinstance(value.scalar, TensorLoad)
- IRBuilder.name(var_name, value.scalar.source)
- return value.scalar
- if isinstance(value, I.meta_var):
- return value.value
- elif getattr(type(value), "_is_meta_class", False):
- name_meta_class_value(var_name, value)
- return value
- elif isinstance(value, list | tuple):
- # Tuple-unpacking with a starred target (e.g. ``vi, *vs = T.axis.remap(...)``)
- # collects multiple elements into a single list bound here. Recurse so each
- # element gets a per-index name; this matches apache's behavior.
- for i, v in enumerate(value):
- bind_assign_value(self, node, f"{var_name}_{i}", v)
- return value
- elif isinstance(value, TensorRegion):
- return value
- elif isinstance(value, Frame):
- value.add_callback(partial(value.__exit__, None, None, None))
- res = value.__enter__()
- IRBuilder.name(var_name, res)
- return res
- elif (
- is_buffer_var(value)
- or isinstance(value, IterVar | Layout)
- or (isinstance(value, tvm.ir.Var) and not self.var_table.exist(value))
- ):
- IRBuilder.name(var_name, value)
- return value
- else:
- is_pointer_expr = isinstance(value, tvm.ir.Expr) and isinstance(
- getattr(value, "ty", None), PointerType
- )
- if is_pointer_expr:
- if self.var_table.contains_in_current_frame(var_name):
- self.report_error(node, f"Pointer variable '{var_name}' cannot be reassigned")
- ann_var = T.Bind(value)
- IRBuilder.name(var_name, ann_var)
- return ann_var
- if not tvm.ir.is_prim_expr(value) and not isinstance(value, Expr):
- # Python scalar (int/float/bool) -> const prim expr
- value = tvm.tirx.const(value)
- if isinstance(value, tvm.ir.StringImm) or not tvm.ir.is_prim_expr(value):
- # StringImm or non-prim-expr (e.g. pointer Call): immutable Bind var
- ann_var = tvm.tirx.Var(var_name, value.ty)
- IRBuilder.name(var_name, ann_var)
- T.Bind(value, var=ann_var)
- return ann_var
- else:
- # x = expr -> scalar (auto-typed from value)
- scalar = T.local_scalar(dtype=str(value.ty.dtype))
- IRBuilder.name(var_name, scalar.scalar.source)
- T.buffer_store(scalar.scalar.source, value, [0])
- return scalar.scalar
-
-
-def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool:
- """
- Check the value of given annotation (argument name) in the prim_func decorator.
- Returns the value of the annotation if present, otherwise giving the default value.
- """
- # look for the named argument in the prim_func / jit decorator
- for dec in node.decorator_list:
- if not isinstance(dec, doc.Call) or dec.func.attr not in ("prim_func", "jit"):
- continue
- for keyword in dec.keywords:
- if keyword.arg == annotation:
- return keyword.value.value
- return default
-
-
-def _is_jit_function(node: doc.FunctionDef) -> bool:
- """Return whether the parsed source function is decorated with ``T.jit``."""
-
- for decorator in node.decorator_list:
- target = decorator.func if isinstance(decorator, doc.Call) else decorator
- if isinstance(target, doc.Attribute) and target.attr == "jit":
- return True
- if isinstance(target, doc.Name) and target.id == "jit":
- return True
- return False
-
-
-def _is_prim_var_declaration(node: doc.expr) -> bool:
- """Return whether an expression is literal ``T.dtype()`` declaration syntax."""
- if not (
- isinstance(node, doc.Call)
- and not node.args
- and not node.keywords
- and isinstance(node.func, doc.Attribute)
- and isinstance(node.func.value, doc.Name)
- and node.func.value.id == "T"
- ):
- return False
- constructor = getattr(T, node.func.attr, None)
- return isinstance(constructor, T.DtypeConstructor) or constructor is T.bool
-
-
-def _prim_var_annotation_dtype(node: doc.expr) -> str | None:
- """Return the dtype of a scalar ``T.dtype`` parameter annotation."""
- if not (
- isinstance(node, doc.Attribute)
- and isinstance(node.value, doc.Name)
- and node.value.id == "T"
- ):
- return None
- constructor = getattr(T, node.attr, None)
- if not isinstance(constructor, T.DtypeConstructor) and constructor is not T.bool:
- return None
- return str(constructor().ty.dtype)
-
-
-def _prim_var_declaration_names(target: doc.expr, value: doc.expr) -> set[str]:
- """Return targets paired with literal ``T.dtype()`` declarations.
-
- This classifies only the assignment currently being visited. In
- particular, it is not a predeclaration pass over the function body.
- """
- if isinstance(target, doc.Name):
- return {target.id} if _is_prim_var_declaration(value) else set()
- if isinstance(target, doc.Tuple | doc.List) and isinstance(value, doc.Tuple | doc.List):
- if len(target.elts) != len(value.elts):
- return set()
- declarations = set()
- for lhs, rhs in zip(target.elts, value.elts):
- declarations.update(_prim_var_declaration_names(lhs, rhs))
- return declarations
- return set()
-
-
-def _signature_prim_var_dtypes(node: doc.FunctionDef) -> dict[str, str]:
- """Collect scalar parameter dtypes without inspecting the function body."""
- return {
- arg.arg: dtype
- for arg in node.args.args
- if arg.annotation is not None
- if (dtype := _prim_var_annotation_dtype(arg.annotation)) is not None
- }
-
-
-@contextlib.contextmanager
-def _signature_match_var_scope(self: Parser):
- """Track PrimVars introduced by Buffer-shape match expressions."""
- previous = getattr(self, "_signature_match_vars", None)
- self._signature_match_vars = {}
- try:
- yield
- finally:
- if previous is None:
- del self._signature_match_vars
- else:
- self._signature_match_vars = previous
-
-
-def _get_signature_match_var(self: Parser, var_name: str) -> tvm.ir.Var | None:
- """Return the active PrimVar created by a signature match expression."""
- current = self.var_table.get().get(var_name)
- match_var = getattr(self, "_signature_match_vars", {}).get(var_name)
- if (
- isinstance(current, tvm.ir.Var)
- and isinstance(match_var, tvm.ir.Var)
- and current.same_as(match_var)
- ):
- return current
- return None
-
-
-def _reuse_signature_match_var(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any:
- """Match a later scalar declaration to a signature-defined PrimVar."""
- previous = _get_signature_match_var(self, var_name)
- if previous is None:
- return value
- if not (
- isinstance(previous.ty, PrimType)
- and isinstance(value, tvm.ir.Var)
- and isinstance(value.ty, PrimType)
- and not is_buffer_var(previous)
- and not is_buffer_var(value)
- ):
- self.report_error(node, f"{var_name} is already bound to a non-PrimVar value")
- if previous.ty != value.ty:
- self.report_error(
- node,
- f"Expected the same dtype for PrimVars but got {value.ty} vs {previous.ty}",
- )
- return previous
-
-
-def _eval_signature_annotation(
- self: Parser,
- node: doc.expr,
- signature_dtypes: dict[str, str],
- *,
- define_missing: bool = True,
-) -> Any:
- """Evaluate a function annotation using Buffer-shape match semantics.
-
- Function parameters are parsed from left to right. An ordinary name in an
- annotation is therefore a lookup and must already be bound. A quoted
- expression in ``T.Buffer``'s shape is the one exception: it is a match
- scope, so its first occurrence creates every missing PrimVar before the
- expression is evaluated. Later dimensions and parameters reuse those
- exact Var objects.
-
- ``signature_dtypes`` comes only from scalar parameter annotations, never
- from the function body. It lets a match expression that precedes a scalar
- parameter create the Var with that parameter's declared dtype. Otherwise
- TIR match variables default to int32. Return annotations pass
- ``define_missing=False`` and consequently cannot introduce variables.
- """
-
- class ShapeStringRewriter(ast.NodeTransformer):
- def visit_Name(self, name): # pylint: disable=invalid-name
- # Direct AST names follow normal definition-before-use ordering.
- if isinstance(name.ctx, ast.Load) and name.id not in self_parser.var_table.get():
- raise NameError(f"name '{name.id}' is not defined")
- return name
-
- def visit_Constant(self, constant): # pylint: disable=invalid-name
- if not isinstance(constant.value, str):
- return constant
- # Replace the quoted expression with its AST. Register its names
- # first so evaluating the rewritten annotation resolves all uses to
- # the canonical match-scope Vars.
- expression = ast.parse(constant.value, mode="eval").body
- for child in ast.walk(expression):
- if not isinstance(child, ast.Name) or not isinstance(child.ctx, ast.Load):
- continue
- current_value = self_parser.var_table.get().get(child.id)
- is_shadowed_type_var = child.id in signature_dtypes and isinstance(
- current_value, TypeVar
- )
- if define_missing and (current_value is None or is_shadowed_type_var):
- # TIR match-scope indices default to int32. A later scalar
- # parameter keeps its explicitly declared dtype. That
- # runtime parameter also shadows a module TypeVar with the
- # same name, such as one emitted for another function.
- var = tvm.tirx.Var(child.id, signature_dtypes.get(child.id, "int32"))
- self_parser.var_table.add(child.id, var, allow_shadowing=False)
- self_parser._signature_match_vars[child.id] = var
- for child in ast.walk(expression):
- ast.copy_location(child, constant)
- return expression
-
- class BufferShapeRewriter(ast.NodeTransformer):
- def visit_Call(self, call): # pylint: disable=invalid-name
- is_buffer = (
- isinstance(call.func, ast.Attribute)
- and call.func.attr == "Buffer"
- and isinstance(call.func.value, ast.Name)
- and call.func.value.id == "T"
- )
- if not is_buffer:
- return self.generic_visit(call)
- # Only Buffer.shape has match-scope semantics. Dtype and other
- # Buffer options remain ordinary expressions.
- call.func = self.visit(call.func)
- if call.args:
- call.args[0] = ShapeStringRewriter().visit(call.args[0])
- call.args[1:] = [self.visit(arg) for arg in call.args[1:]]
- for keyword in call.keywords:
- if keyword.arg == "shape":
- keyword.value = ShapeStringRewriter().visit(keyword.value)
- else:
- keyword.value = self.visit(keyword.value)
- return call
-
- self_parser = self
- python_node = from_doc(deepcopy(node))
- python_node = BufferShapeRewriter().visit(python_node)
- ast.fix_missing_locations(python_node)
- return self.eval_expr(doc.to_doc(python_node))
-
-
-@dispatch.register(token="tirx", type_name="For")
-def visit_for(self: Parser, node: doc.For) -> None:
- """The for visiting method for tirx.
-
- Parameters
- ----------
- self : Parser
- The visiting parser.
-
- node : doc.For
- The doc AST for node.
- """
- # Intercept range() at AST level so it works with both Python ints and Exprs.
- # In other contexts (e.g. list comprehensions), range remains Python's builtin.
- if (
- isinstance(node.iter, doc.Call)
- and isinstance(node.iter.func, doc.Name)
- and node.iter.func.id == "range"
- ):
- args = [self.eval_expr(a) for a in node.iter.args]
- kwargs = {kw.arg: self.eval_expr(kw.value) for kw in node.iter.keywords}
- if len(args) == 1:
- for_frame = T.serial(0, args[0], **kwargs)
- elif len(args) == 2:
- for_frame = T.serial(args[0], args[1], **kwargs)
- elif len(args) == 3:
- for_frame = T.serial(args[0], args[1], step=args[2], **kwargs)
- else:
- self.report_error(node.iter, "range() takes 1 to 3 arguments")
- else:
- for_frame = self.eval_expr(node.iter)
- if not isinstance(for_frame, T.frame.ForFrame):
- self.report_error(
- node.iter,
- "Expect the for loop to be one of the following: "
- "range, T.serial, T.grid, T.parallel, T.vectorized, T.unroll, T.thread_binding",
- )
- with self.var_table.with_frame():
- with for_frame as iters:
- self.eval_assign(target=node.target, source=iters, bind_value=bind_for_value)
- self.visit_body(node.body)
-
-
-@dispatch.register(token="tirx", type_name="While")
-def visit_while(self: Parser, node: doc.While) -> None:
- """The while visiting method for tirx.
-
- Parameters
- ----------
- self : Parser
- The visiting parser.
-
- node : doc.While
- The doc AST while node.
- """
- with self.var_table.with_frame():
- cond = self.eval_expr(node.test)
- with T.While(cond):
- self.visit_body(node.body)
-
-
-@dispatch.register(token="tirx", type_name="Break")
-def visit_break(self: Parser, node: doc.Break) -> None:
- """The break visiting method for tir.
-
- Parameters
- ----------
- self : Parser
- The visiting parser.
-
- node : doc.Break
- The doc AST break node.
- """
- T.evaluate(T.break_loop())
-
-
-@dispatch.register(token="tirx", type_name="Continue")
-def visit_continue(self: Parser, node: doc.Continue) -> None:
- """The continue visiting method for tir.
-
- Parameters
- ----------
- self : Parser
- The visiting parser.
-
- node : doc.Continue
- The doc AST continue node.
- """
- T.evaluate(T.continue_loop())
-
-
-@dispatch.register(token="tirx", type_name="Assign")
-def visit_assign(self: Parser, node: doc.Assign) -> None:
- """The assign visiting method for tirx.
-
- Parameters
- ----------
- self : Parser
- The visiting parser.
-
- node : doc.Assign
- The doc AST assign node.
- """
- if len(node.targets) != 1:
- self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
- lhs = node.targets[0]
-
- if isinstance(node.value, doc.Subscript):
- check_slices = []
- if isinstance(node.value.slice, doc.Slice):
- check_slices = [node.value.slice]
- elif isinstance(node.value.slice, doc.Tuple):
- for p in node.value.slice.elts:
- if isinstance(p, doc.Slice):
- check_slices.append(p)
- for s in check_slices:
- if not s.step and s.upper and s.lower:
- s.step = doc.Constant(
- 1,
- None,
- s.upper.lineno,
- s.upper.end_col_offset + 1,
- s.upper.lineno,
- s.upper.end_col_offset + 2,
- )
-
- rhs = self.eval_expr(node.value)
- if isinstance(lhs, doc.Subscript):
- if isinstance(lhs.slice, doc.Tuple):
- indices = []
- for index in lhs.slice.elts:
- if isinstance(index, doc.Starred):
- # x[*y]
- indices.extend(self.eval_expr(index.value))
- else:
- indices.append(self.eval_expr(index))
- else:
- indices = self.eval_expr(lhs.slice)
- T.buffer_store(self.eval_expr(lhs.value), rhs, indices)
... 1025 lines suppressed ...