This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 37f7b580c8 [REFACTOR][TVMScript] Resolve global-info arguments in
builders (#20439)
37f7b580c8 is described below
commit 37f7b580c8d6a58998b01d5f91079fa463cd271e
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 24 20:06:57 2026 -0400
[REFACTOR][TVMScript] Resolve global-info arguments in builders (#20439)
Resolve named global-info arguments through the shared
`resolve_global_info_args` builder decorator. Positional, keyword,
aliased and unpacked calls use ordinary Python argument evaluation,
including declared string defaults, while dialect resolvers own selector
syntax and module lookup. Remove the parser's global-info argument
policy and rewriting.
Require an active module for string selectors and discard completed or
failed module scopes. Preserve executable printed code with postponed
function annotations and existing virtual-device lookups for eager typed
external-function declarations.
---
docs/arch/tvmscript.rst | 19 +++-
python/tvm/relax/script/ir_builder/__init__.py | 34 +++----
.../tvm/relax/script/ir_builder/parser_protocol.py | 6 +-
python/tvm/script/ir_builder/__init__.py | 2 +
python/tvm/script/ir_builder/base.py | 67 ++++++++++++
python/tvm/script/ir_builder/frame.py | 8 ++
python/tvm/script/ir_builder/parser_protocol.py | 18 ++--
python/tvm/script/parser/protocol_registry.py | 6 +-
python/tvm/script/parser/transpile.py | 14 +--
src/relax/script/printer/dependent_type.cc | 36 ++++++-
src/relax/script/printer/distributed.cc | 9 ++
src/script/ir_builder/ir/frame.cc | 19 +++-
.../test_transform_fold_vdevice_scope_change.py | 8 +-
..._distributed_transform_legalize_redistribute.py | 12 +--
.../test_distributed_transform_lower_distir.py | 16 +--
...ributed_transform_lower_global_to_local_view.py | 64 ++++++------
...est_distributed_transform_propagate_sharding.py | 112 ++++++++++-----------
.../test_distributed_tvmscript_parser.py | 12 +--
.../test_distributed_tvmscript_printer.py | 6 +-
tests/python/relax/test_analysis_well_formed.py | 2 +-
.../relax/test_backend_dispatch_sort_scan.py | 20 ++--
tests/python/relax/test_transform_legalize_ops.py | 6 +-
.../python/relax/test_transform_realize_vdevice.py | 54 +++++-----
...nsform_specialize_primfunc_based_on_callsite.py | 12 +--
.../python/relax/test_transform_update_vdevice.py | 20 ++--
tests/python/relax/test_tvmscript_parser.py | 8 +-
tests/python/relax/test_vm_build.py | 2 +-
tests/python/script/minilang.py | 7 +-
.../python/script/test_special_parser_protocol.py | 30 +++---
.../test_tvmscript_ir_builder_irmodule.py | 1 +
30 files changed, 378 insertions(+), 252 deletions(-)
diff --git a/docs/arch/tvmscript.rst b/docs/arch/tvmscript.rst
index 2b1e1f9d6d..cead1d0582 100644
--- a/docs/arch/tvmscript.rst
+++ b/docs/arch/tvmscript.rst
@@ -173,9 +173,11 @@ TIR control-flow frames in a primitive function and Relax
control-flow frames in
function. The transpiler emits these operations without implementing dialect
IR semantics.
``parser/protocol_registry.py`` owns callable syntax metadata and
registration. Argument policies such
-as ``expr_str`` translate symbolic strings written directly in source
expressions, while ``global_info``
-preserves a module reference for builder-side lookup. Dtype and placement
strings remain
-literal. Captured or computed symbolic shapes must already contain explicit IR
variables;
+as ``expr_str`` translate symbolic strings written directly in source
expressions.
+Global-info selectors use the builder decorator
+:func:`tvm.script.ir_builder.resolve_global_info_args`, which resolves named
string arguments
+through an explicit dialect callback after ordinary Python argument
evaluation. Captured
+aliases and unpacked calls use the same decorator. Dtype and placement strings
remain literal. Captured or computed symbolic shapes must already contain
explicit IR variables;
the parser does not interpret expression strings found inside captured values.
Registration keeps syntax facts in simple dictionaries keyed by canonical
namespace paths.
Known namespace aliases normalize to the registered path; ordinary captured
callables,
@@ -263,8 +265,15 @@ DSL expression strings such as ``R.Tensor(("n", 4),
"float32")`` use
``type_var_map``. Repeated names refer to the same symbol across parameters,
return
annotations and the body. Quoted lookup does not introduce a Python name; an
explicit
symbol declaration does. A Python name in ``R.Tensor((n, 4), ...)`` follows
its source
-lexical scope. Module global-info references use ``I.resolve_global_info_``
against the
-active native module frame's existing map.
+lexical scope. Relax global-info selectors use ``R.resolve_global_info_``
against the
+nearest active native module frame's existing map. String selectors outside an
active module
+raise, including when an unrelated builder is active. Concrete metadata
objects retain identity.
+When defining Python functions before their module builder opens, quote the
whole annotation
+or use ``from __future__ import annotations``; the parser reconstructs these
annotations inside
+the module. Printed module function signatures containing selectors request
postponed annotations.
+Eager printed class assignments, such as typed ``R.ExternFunc`` declarations,
retrieve concrete
+virtual devices with ``R.lookup_vdevice`` because postponed annotations do not
defer assignment
+values. This explicit lookup preserves the module metadata object before
tensor construction.
During rewriting, each needed source range becomes a fixed native span entry
in an
injected table. ``_S[i](value)`` attaches that location while preserving
identity;
diff --git a/python/tvm/relax/script/ir_builder/__init__.py
b/python/tvm/relax/script/ir_builder/__init__.py
index f381cb31c5..29987e713f 100644
--- a/python/tvm/relax/script/ir_builder/__init__.py
+++ b/python/tvm/relax/script/ir_builder/__init__.py
@@ -27,7 +27,7 @@ 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 resolve_global_info_args as
_resolve_global_info_args
from tvm.script.ir_builder.base import at as _at
from tvm.script.ir_builder.base import source_span as _source_span
from tvm.script.parser.protocol_registry import ARGS_POLICIES as _ARGS_POLICIES
@@ -83,7 +83,8 @@ For = for_
supports_mutable_declarations = False
-@_args_policy("R.Tensor", {"shape": "expr_str", "vdevice": "global_info"},
scalar_strings=False)
+@_resolve_global_info_args("vdevice", resolver=resolve_global_info_)
+@_args_policy("R.Tensor", {"shape": "expr_str"}, scalar_strings=False)
def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
"""Construct a Relax tensor type.
@@ -97,7 +98,9 @@ def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *,
span=None):
Element type; None leaves the element type unknown.
vdevice : VDevice or str, optional
Concrete virtual device or a module metadata selector, such as
"cuda:0".
- None leaves the virtual device unspecified.
+ None leaves the virtual device unspecified. Strings require an active
module
+ builder; use a quoted whole annotation or postponed annotations when
defining
+ a Python function before its module builder opens.
ndim : int, optional
Rank when shape is unknown; -1 means unknown rank. Do not supply
an explicit rank together with a known shape.
@@ -107,22 +110,16 @@ def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1,
*, span=None):
Returns
-------
result : TensorType or Type
- The tensor type. A string device selector evaluated before the module
- builder opens returns a missing type for the generated declaration to
resolve.
+ The tensor type, or a missing type for an unresolved eager shape
annotation.
+ String selectors outside an active module always raise ValueError.
"""
if isinstance(shape, _python.str) and dtype is None:
dtype, shape = shape, None
- if isinstance(vdevice, _python.str) and not _IRBuilder.is_in_scope():
- # Python evaluates annotations before the module decorator opens its
- # frame. The generated declaration resolves this module-owned
reference.
- return _ir.Type.missing()
- vdevice = resolve_global_info_(vdevice)
return _relax.TensorType(shape, dtype, vdevice, ndim, _source_span(span))
-@_args_policy(
- "R.DTensor", {"shape": "expr_str", "device_mesh": "global_info"},
scalar_strings=False
-)
+@_resolve_global_info_args("device_mesh", resolver=resolve_global_info_)
+@_args_policy("R.DTensor", {"shape": "expr_str"}, scalar_strings=False)
def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *,
ndim=-1, span=None):
"""Construct a Relax distributed tensor type.
@@ -135,7 +132,8 @@ def DTensor(shape=None, dtype=None, device_mesh=None,
placement="", *, ndim=-1,
Element type; None leaves the element type unknown.
device_mesh : DeviceMesh or str, optional
Concrete mesh or module metadata selector. None creates an empty mesh
- placeholder. A selector is resolved in the enclosing module builder.
+ placeholder. A string selector requires an active module builder; use
a quoted
+ whole annotation or postponed annotations before the builder opens.
placement : Placement or str, optional
Distribution placement. Text, including the default empty string, is
parsed with Placement.from_text.
@@ -147,15 +145,11 @@ def DTensor(shape=None, dtype=None, device_mesh=None,
placement="", *, ndim=-1,
Returns
-------
result : DTensorType or Type
- The distributed type. A string mesh selector evaluated before the
- module builder opens returns a missing type for later resolution.
+ The distributed type, or a missing type for an unresolved eager shape
annotation.
+ String selectors outside an active module always raise ValueError.
"""
- if isinstance(device_mesh, _python.str) and not _IRBuilder.is_in_scope():
- return _ir.Type.missing()
if device_mesh is None:
device_mesh = _DeviceMesh([], _ir.Range(0, 1))
- else:
- device_mesh = resolve_global_info_(device_mesh)
if isinstance(placement, _python.str):
placement = _Placement.from_text(placement)
return _DTensorType(Tensor(shape, dtype, ndim=ndim), device_mesh,
placement, _source_span(span))
diff --git a/python/tvm/relax/script/ir_builder/parser_protocol.py
b/python/tvm/relax/script/ir_builder/parser_protocol.py
index 539d9506de..f7ba61e7c7 100644
--- a/python/tvm/relax/script/ir_builder/parser_protocol.py
+++ b/python/tvm/relax/script/ir_builder/parser_protocol.py
@@ -253,10 +253,10 @@ def resolve_global_info_(content: Any) -> Any:
.. code:: python
- # Source
+ # The constructor decorator calls this resolver for string selectors.
R.Tensor((n,), "float32", vdevice="cuda:0")
- # Generated builder
- R.Tensor((n,), "float32", vdevice=R.resolve_global_info_("cuda:0"))
+ # Direct resolution requires the same active module frame.
+ device = R.resolve_global_info_("cuda:0")
"""
if not isinstance(content, str):
return content
diff --git a/python/tvm/script/ir_builder/__init__.py
b/python/tvm/script/ir_builder/__init__.py
index 3906665490..93bcec4b8d 100644
--- a/python/tvm/script/ir_builder/__init__.py
+++ b/python/tvm/script/ir_builder/__init__.py
@@ -29,6 +29,7 @@ from .base import (
annotation_value_,
at_,
require_defined,
+ resolve_global_info_args,
with_at_group_,
)
from .frame import IRModuleFrame
@@ -70,6 +71,7 @@ __all__ = [
"module_member_",
"module_set_attr",
"require_defined",
+ "resolve_global_info_args",
"with_at_group_",
]
diff --git a/python/tvm/script/ir_builder/base.py
b/python/tvm/script/ir_builder/base.py
index 6411c821ad..14a960e9ca 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -19,6 +19,7 @@
from collections.abc import Callable
from contextlib import contextmanager, nullcontext
from functools import wraps
+from inspect import signature
from typing import Any, Generic, TypeVar
from tvm_ffi import register_object as _register_object
@@ -30,6 +31,72 @@ from tvm.runtime import Object as _Object
from . import _ffi_api
+def resolve_global_info_args(
+ *fields: str, resolver: Callable[[str], Any]
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+ """Resolve selected string arguments before calling a builder operation.
+
+ Parameters
+ ----------
+ *fields : str
+ Names of positional-only, positional-or-keyword, or keyword-only
parameters.
+ Repeated names are resolved once. Omitted arguments use their declared
defaults.
+ resolver : Callable[[str], Any]
+ Explicit callback receiving each selected string and returning its
replacement.
+ The callback owns selector syntax, lookup scope, and errors.
Non-string values
+ pass through with their identity preserved; containers are not decoded
recursively.
+
+ Returns
+ -------
+ Callable
+ Decorator preserving the callable's signature, name, and
documentation. The
+ signature is inspected once when decorating, then reused for argument
binding.
+
+ Raises
+ ------
+ ValueError
+ If a selected name is absent or names a variadic parameter.
+
+ Notes
+ -----
+ All argument expressions are evaluated once in ordinary Python order
before binding,
+ resolution, and the callable body. Positional, keyword, unpacked, and
aliased calls
+ share this behavior. Selected string defaults are resolved on every call.
Exceptions
+ from binding, the resolver, and the callable propagate unchanged.
+
+ .. code:: python
+
+ @resolve_global_info_args("device", resolver=lookup_device)
+ def tensor(shape, device="default"):
+ return make_tensor(shape, device)
+ """
+ fields = tuple(dict.fromkeys(fields))
+
+ def decorate(function: Callable[..., Any]) -> Callable[..., Any]:
+ call_signature = signature(function)
+ for field in fields:
+ parameter = call_signature.parameters.get(field)
+ if parameter is None or parameter.kind in (
+ parameter.VAR_POSITIONAL,
+ parameter.VAR_KEYWORD,
+ ):
+ raise ValueError(f"Unknown or variadic global-info argument:
{field!r}")
+
+ @wraps(function)
+ def invoke(*args, **kwargs):
+ bound = call_signature.bind(*args, **kwargs)
+ bound.apply_defaults()
+ for field in fields:
+ value = bound.arguments[field]
+ if isinstance(value, str):
+ bound.arguments[field] = resolver(value)
+ return function(*bound.args, **bound.kwargs)
+
+ return invoke
+
+ return decorate
+
+
@_register_object("script.ir_builder.IRBuilderFrame")
class IRBuilderFrame(_Object):
"""A stack frame of the IRBuilder used to keep track of the current scope.
diff --git a/python/tvm/script/ir_builder/frame.py
b/python/tvm/script/ir_builder/frame.py
index 778fe2fed3..367949c2f6 100644
--- a/python/tvm/script/ir_builder/frame.py
+++ b/python/tvm/script/ir_builder/frame.py
@@ -18,11 +18,19 @@
from tvm_ffi import register_object as _register_object
+from . import _ffi_api
from .base import IRBuilderFrame
@_register_object("script.ir_builder.IRModuleFrame")
class IRModuleFrame(IRBuilderFrame):
+ def __exit__(self, exc_type, exc_value, trace):
+ try:
+ return super().__exit__(exc_type, exc_value, trace)
+ finally:
+ # Body and exit-callback failures must not leave an active module
scope.
+ _ffi_api.IRModuleFrameAbort(self)
+
def __getattr__(self, name):
"""Expose a declared function through the native module's reference
map.
diff --git a/python/tvm/script/ir_builder/parser_protocol.py
b/python/tvm/script/ir_builder/parser_protocol.py
index 06000c78f0..120e97c18e 100644
--- a/python/tvm/script/ir_builder/parser_protocol.py
+++ b/python/tvm/script/ir_builder/parser_protocol.py
@@ -728,15 +728,19 @@ def resolve_global_info_(content: Any) -> Any:
Notes
-----
- A string argument marked ``global_info`` is evaluated once and passed to
- this hook without attaching a new span to the existing metadata object.
+ A builder may supply this hook to
+ :func:`tvm.script.ir_builder.resolve_global_info_args`. The decorator
resolves
+ selected string arguments after ordinary Python argument evaluation,
without
+ attaching a new span to the existing metadata object. The parser does not
rewrite
+ selector arguments.
.. code:: python
- # Source: an argument policy marks metadata as global_info.
- X.CustomType(metadata="mesh[0]")
- # Generated builder
- X.CustomType(metadata=X.resolve_global_info_("mesh[0]"))
+ @resolve_global_info_args("metadata", resolver=resolve_global_info_)
+ def CustomType(metadata):
+ return make_type(metadata)
+
+ CustomType(metadata="mesh[0]")
"""
raise NotImplementedError
@@ -868,7 +872,7 @@ def module_member_(name: str, value: Any) -> Any:
#
# Syntax markers live in tvm.script.parser.protocol_registry.
# ``constexpr(value)`` selects host evaluation in marked control flow.
-# ``args_policy(path, fields)`` marks expression-string or global-info
arguments.
+# ``args_policy(path, fields)`` marks expression-string arguments.
# ``register_type_var_decl(path, constructor, dtype=...)`` marks symbolic
declarations.
# ``mutable_cell_decl(path)`` marks mutable storage declarations.
# ``result_span(path)`` permits attaching a call's result span without a call
context.
diff --git a/python/tvm/script/parser/protocol_registry.py
b/python/tvm/script/parser/protocol_registry.py
index c336a17781..c9c0595102 100644
--- a/python/tvm/script/parser/protocol_registry.py
+++ b/python/tvm/script/parser/protocol_registry.py
@@ -59,7 +59,7 @@ class ExprStrPolicy(NamedTuple):
class ArgsPolicy(NamedTuple):
"""Argument policies and positional names computed once at registration.
- ``fields`` maps parameter names to ``expr_str`` or ``global_info``.
+ ``fields`` maps parameter names to ``expr_str``.
``expression`` describes the expression-string subset.
``positional_parameters``
lists positional-only and positional-or-keyword names in signature order;
keyword-only parameters remain available through ``fields``.
@@ -126,7 +126,7 @@ def args_policy(
namespace_path : str
Registered namespace alias followed by the exported member path.
fields : Mapping[str, str]
- Parameter names mapped to ``expr_str`` or ``global_info``.
+ Parameter names mapped to ``expr_str``.
scalar_strings : bool, optional
Whether bare strings in expression fields denote expressions. Defaults
to True.
dtype : object, optional
@@ -159,7 +159,7 @@ def args_policy(
# Builder: M.Tensor((M.resolve_type_var_("n"),))
"""
fields = dict(fields)
- unsupported = set(fields.values()).difference(("expr_str", "global_info"))
+ unsupported = set(fields.values()).difference(("expr_str",))
if unsupported:
raise ValueError(f"Unknown argument policies: {sorted(unsupported)}")
diff --git a/python/tvm/script/parser/transpile.py
b/python/tvm/script/parser/transpile.py
index 1519c9a4dc..a09b21ca19 100644
--- a/python/tvm/script/parser/transpile.py
+++ b/python/tvm/script/parser/transpile.py
@@ -610,18 +610,6 @@ class IRBuilderTranspiler(ast.NodeTransformer):
if kind == "expr_str" and (nested or policy.scalar_strings):
with self._use_string_policy(policy):
return self.visit(parse_expression_string(node,
self.module.filename))
- if kind == "global_info":
- # -------------------- Pattern --------------------
- # Python source:
- # X.Tensor(vdevice="cuda:1")
- #
- # Builder:
- # X.Tensor(vdevice=X.resolve_global_info_("cuda:1"))
- # -------------------------------------------------
- value = self.visit(node)
- return self._call(
- self.function.dialect_prefix, "resolve_global_info_",
[value], node
- )
return self.visit(node)
def _visit_direct_operand(self, node: ast.expr) -> ast.expr:
@@ -649,7 +637,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# X.Tensor(("n",), vdevice="cuda:0")
#
# Builder:
- # X.Tensor((X.resolve_type_var_("n"),),
vdevice=X.resolve_global_info_("cuda:0"))
+ # X.Tensor((X.resolve_type_var_("n"),), vdevice="cuda:0")
# -------------------------------------------------
binding_value = node is self.binding_expression
marker = self._read_constexpr_operand(node)
diff --git a/src/relax/script/printer/dependent_type.cc
b/src/relax/script/printer/dependent_type.cc
index 41d4beecab..cef8bee1f7 100644
--- a/src/relax/script/printer/dependent_type.cc
+++ b/src/relax/script/printer/dependent_type.cc
@@ -20,6 +20,7 @@
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/tirx/stmt_functor.h>
+#include "../../../script/printer/ir/utils.h"
#include "./utils.h"
namespace tvm {
@@ -117,12 +118,37 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kwargs_values.push_back(LiteralDoc::Int(n->ndim, n_p->Attr("ndim")));
}
if (n->vdevice.has_value() && n->vdevice.value()->target.defined()) {
+ // Function annotations defer module-owned selectors until
declaration.
+ bool has_relax_frame = false;
+ bool has_module_frame = false;
+ for (const Frame& frame : d->frames) {
+ if (const auto* relax_frame = frame.as<RelaxFrameNode>()) {
+ has_relax_frame = true;
+ if (relax_frame->func_vars != nullptr) {
+ d->ir_usage.insert("future_annotations");
+ }
+ } else if (const auto* ir_frame = frame.as<IRFrameNode>()) {
+ has_module_frame = ir_frame->global_infos != nullptr;
+ }
+ }
kwargs_keys.push_back("vdevice");
- std::string dev_kind = n->vdevice.value()->target->kind->name;
- int dev_index = FindVDeviceIndexByTargetKind(n->vdevice.value(), d);
- kwargs_values.push_back(LiteralDoc::Str(
- dev_kind + ":" + std::to_string(dev_index) + ":" +
n->vdevice.value()->memory_scope,
- n_p->Attr("vdevice")));
+ if (has_relax_frame) {
+ std::string dev_kind = n->vdevice.value()->target->kind->name;
+ int dev_index = FindVDeviceIndexByTargetKind(n->vdevice.value(),
d);
+ kwargs_values.push_back(LiteralDoc::Str(
+ dev_kind + ":" + std::to_string(dev_index) + ":" +
n->vdevice.value()->memory_scope,
+ n_p->Attr("vdevice")));
+ } else if (has_module_frame) {
+ // Class assignments execute eagerly; retrieve their existing
concrete metadata.
+ kwargs_values.push_back(
+ Relax(d, "lookup_vdevice")
+
->Call({LiteralDoc::Str(n->vdevice.value()->target->kind->name,
+ n_p->Attr("vdevice")),
+
LiteralDoc::Int(FindVDeviceIndexByTargetKind(n->vdevice.value(), d),
+ n_p->Attr("vdevice"))}));
+ } else {
+ kwargs_values.push_back(d->AsDoc<ExprDoc>(n->vdevice.value(),
n_p->Attr("vdevice")));
+ }
}
if (args.empty() && kwargs_keys.empty()) {
return Relax(d, "Tensor");
diff --git a/src/relax/script/printer/distributed.cc
b/src/relax/script/printer/distributed.cc
index 0e8810ea25..b7327bedfe 100644
--- a/src/relax/script/printer/distributed.cc
+++ b/src/relax/script/printer/distributed.cc
@@ -116,6 +116,15 @@ TVM_FFI_STATIC_INIT_BLOCK() {
for (const auto& kv : *f->global_infos) {
for (int i = 0; i < static_cast<int>(kv.second.size()); i++) {
if (kv.second[i].same_as(n)) {
+ // Module-owned selectors must be evaluated inside the
declaration frame.
+ for (const Frame& frame : d->frames) {
+ if (const auto* relax_frame = frame.as<RelaxFrameNode>()) {
+ if (relax_frame->func_vars != nullptr) {
+ d->ir_usage.insert("future_annotations");
+ break;
+ }
+ }
+ }
std::stringstream ss;
ss << kv.first << "[" << i << "]";
return d->AsDoc<Doc>(ffi::String(ss.str()), n_p);
diff --git a/src/script/ir_builder/ir/frame.cc
b/src/script/ir_builder/ir/frame.cc
index b216bd8334..284cd7e8d9 100644
--- a/src/script/ir_builder/ir/frame.cc
+++ b/src/script/ir_builder/ir/frame.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/module.h>
#include <tvm/script/ir_builder/ir/frame.h>
@@ -25,9 +26,25 @@ namespace script {
namespace ir_builder {
namespace ir {
-TVM_FFI_STATIC_INIT_BLOCK() { IRModuleFrameNode::RegisterReflection(); }
+TVM_FFI_STATIC_INIT_BLOCK() {
+ IRModuleFrameNode::RegisterReflection();
+ ffi::reflection::GlobalDef().def("script.ir_builder.ir.IRModuleFrameAbort",
+ [](const IRModuleFrame& frame) {
+ if (!IRBuilder::IsInScope()) return;
+ auto& frames =
IRBuilder::Current()->frames;
+ for (size_t i = frames.size(); i > 0;
--i) {
+ if (frames[i - 1].same_as(frame)) {
+ // Drop this failed scope without
callbacks or
+ // finalization.
+ while (frames.size() >= i)
frames.pop_back();
+ return;
+ }
+ }
+ });
+}
void IRModuleFrameNode::ExitWithScope() {
+ IRBuilderFrameNode::ExitWithScope();
ffi::Map<GlobalVar, BaseFunc> func_map;
TVM_FFI_ICHECK_EQ(functions.size(), global_var_map.size())
<< "All functions must be defined in the IRModule. Got " <<
global_var_map.size()
diff --git
a/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
b/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
index a416f38b36..e097727402 100644
---
a/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
+++
b/tests/python/relax/backend/adreno/test_transform_fold_vdevice_scope_change.py
@@ -127,8 +127,8 @@ def test_maxpool2d_scope_folding():
@R.function
def main(
- x: R.Tensor((2, 4, 26, 26), dtype="float32",
vdevice="opencl:1:global"),
- ) -> R.Tensor((2, 4, 13, 13), dtype="float32",
vdevice="opencl:1:global"):
+ x: 'R.Tensor((2, 4, 26, 26), dtype="float32",
vdevice="opencl:1:global")',
+ ) -> 'R.Tensor((2, 4, 13, 13), dtype="float32",
vdevice="opencl:1:global")':
cls = Input
with R.dataflow():
lv = R.call_tir(
@@ -252,8 +252,8 @@ def test_maxpool2d_scope_folding():
@R.function
def main(
- x: R.Tensor((2, 4, 26, 26), dtype="float32",
vdevice="opencl:1:global"),
- ) -> R.Tensor((2, 4, 13, 13), dtype="float32",
vdevice="opencl:1:global"):
+ x: 'R.Tensor((2, 4, 26, 26), dtype="float32",
vdevice="opencl:1:global")',
+ ) -> 'R.Tensor((2, 4, 13, 13), dtype="float32",
vdevice="opencl:1:global")':
cls = Expected
with R.dataflow():
lv = R.call_tir(
diff --git
a/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py
b/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py
index 337a19aa6f..a4ca009774 100644
---
a/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py
+++
b/tests/python/relax/distributed/test_distributed_transform_legalize_redistribute.py
@@ -31,8 +31,8 @@ def test_simple():
@R.function
def foo(
- x1: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
- x2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
+ x1: 'R.DTensor((128, 128), "float32", "mesh[0]", "R")',
+ x2: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[0]")',
):
R.func_attr({"num_input": 1})
# scatter
@@ -48,12 +48,12 @@ def test_simple():
@R.function
def foo(
- x1: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
- x2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
- ) -> R.Tuple(
+ x1: 'R.DTensor((128, 128), "float32", "mesh[0]", "R")',
+ x2: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[0]")',
+ ) -> """R.Tuple(
R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
- ):
+ )""":
R.func_attr({"num_input": 1})
lv0: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") = (
R.dist.redistribute_replica_to_shard(x1, num_workers=2, axis=1)
diff --git
a/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
b/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
index b7a89a121b..4279089793 100644
--- a/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
+++ b/tests/python/relax/distributed/test_distributed_transform_lower_distir.py
@@ -112,10 +112,10 @@ def test_mlp():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
- weight1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
- weight2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
- ) -> R.DTensor((128, 128), "float32", "mesh[0]", "R"):
+ x: 'R.DTensor((128, 128), "float32", "mesh[0]", "R")',
+ weight1: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[1]")',
+ weight2: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[0]")',
+ ) -> 'R.DTensor((128, 128), "float32", "mesh[0]", "R")':
R.func_attr({"num_input": 1})
cls = MLP
lv0: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.dist.call_tir_local_view(
@@ -292,12 +292,12 @@ def test_mlp_with_tuple():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
- weight_packed: R.Tuple(
+ x: 'R.DTensor((128, 128), "float32", "mesh[0]", "R")',
+ weight_packed: """R.Tuple(
R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
- ),
- ) -> R.DTensor((64, 128), "float32", "mesh[0]", "R"):
+ )""",
+ ) -> 'R.DTensor((64, 128), "float32", "mesh[0]", "R")':
cls = MLPWithTuple
weight1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
weight_packed[0]
lv0: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.dist.call_tir_local_view(
diff --git
a/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
b/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
index 1cf1f0bb92..2a825b3b54 100644
---
a/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
+++
b/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
@@ -95,10 +95,10 @@ def test_mlp():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
- weight1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
- weight2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
- ) -> R.DTensor((128, 128), "float32", "mesh[0]", "R"):
+ x: 'R.DTensor((128, 128), "float32", "mesh[0]", "R")',
+ weight1: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[1]")',
+ weight2: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[0]")',
+ ) -> 'R.DTensor((128, 128), "float32", "mesh[0]", "R")':
cls = MLP
lv0 = R.dist.call_tir(
cls.matmul,
@@ -201,10 +201,10 @@ def test_mlp():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
- weight1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
- weight2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
- ) -> R.DTensor((128, 128), "float32", "mesh[0]", "R"):
+ x: 'R.DTensor((128, 128), "float32", "mesh[0]", "R")',
+ weight1: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[1]")',
+ weight2: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[0]")',
+ ) -> 'R.DTensor((128, 128), "float32", "mesh[0]", "R")':
cls = Expected
lv0: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.dist.call_tir_local_view(
cls.matmul1,
@@ -647,19 +647,19 @@ def test_llama_attention():
@R.function(pure=False)
def foo(
- input_tokens: R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R"),
- mask: R.DTensor((1, 1, 256, 256), "float16", "mesh[0]", "R"),
- div_const: R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]"),
- maximum_const: R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]"),
+ input_tokens: 'R.DTensor((1, 256, 4096), "float16", "mesh[0]",
"R")',
+ mask: 'R.DTensor((1, 1, 256, 256), "float16", "mesh[0]", "R")',
+ div_const: 'R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]")',
+ maximum_const: 'R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]")',
kv_cache: R.Tuple(R.Any, R.Any),
- linear_weight: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight1: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight2: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight3: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[1]"),
- rms_norm_weight: R.DTensor((4096,), "float16", "mesh[0]", "R"),
- cos_cached: R.DTensor((2048, 128), "float16", "mesh[0]", "R"),
- sin_cached: R.DTensor((2048, 128), "float16", "mesh[0]", "R"),
- ) -> R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R"):
+ linear_weight: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight1: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight2: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight3: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[1]")',
+ rms_norm_weight: 'R.DTensor((4096,), "float16", "mesh[0]", "R")',
+ cos_cached: 'R.DTensor((2048, 128), "float16", "mesh[0]", "R")',
+ sin_cached: 'R.DTensor((2048, 128), "float16", "mesh[0]", "R")',
+ ) -> 'R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R")':
cls = LlamaAttentionLayer
lv6 = R.dist.call_tir(
cls.rms_norm,
@@ -1319,19 +1319,19 @@ def test_llama_attention():
@R.function(pure=False)
def foo(
- input_tokens: R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R"),
- mask: R.DTensor((1, 1, 256, 256), "float16", "mesh[0]", "R"),
- div_const: R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]"),
- maximum_const: R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]"),
+ input_tokens: 'R.DTensor((1, 256, 4096), "float16", "mesh[0]",
"R")',
+ mask: 'R.DTensor((1, 1, 256, 256), "float16", "mesh[0]", "R")',
+ div_const: 'R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]")',
+ maximum_const: 'R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]")',
kv_cache: R.Tuple(R.Any, R.Any),
- linear_weight: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight1: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight2: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight3: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[1]"),
- rms_norm_weight: R.DTensor((4096,), "float16", "mesh[0]", "R"),
- cos_cached: R.DTensor((2048, 128), "float16", "mesh[0]", "R"),
- sin_cached: R.DTensor((2048, 128), "float16", "mesh[0]", "R"),
- ) -> R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R"):
+ linear_weight: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight1: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight2: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight3: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[1]")',
+ rms_norm_weight: 'R.DTensor((4096,), "float16", "mesh[0]", "R")',
+ cos_cached: 'R.DTensor((2048, 128), "float16", "mesh[0]", "R")',
+ sin_cached: 'R.DTensor((2048, 128), "float16", "mesh[0]", "R")',
+ ) -> 'R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R")':
cls = Expected
lv6: R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R") =
R.dist.call_tir_local_view(
cls.rms_norm,
diff --git
a/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
b/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
index bcfe0d34c7..eb4baa09b3 100644
---
a/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
+++
b/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
@@ -61,10 +61,10 @@ def test_mlp():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
- weight1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
- weight2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
- ) -> R.DTensor((128, 128), "float32", "mesh[0]", "R"):
+ x: 'R.DTensor((128, 128), "float32", "mesh[0]", "R")',
+ weight1: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[1]")',
+ weight2: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[0]")',
+ ) -> 'R.DTensor((128, 128), "float32", "mesh[0]", "R")':
lv0: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.matmul(x, weight1)
lv1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.nn.gelu(lv0)
lv3: R.DTensor((128, 128), "float32", "mesh[0]", "R") =
R.matmul(lv1, weight2)
@@ -155,12 +155,12 @@ def test_mlp_with_tuple():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
- weight_packed: R.Tuple(
+ x: 'R.DTensor((128, 128), "float32", "mesh[0]", "R")',
+ weight_packed: """R.Tuple(
R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
- ),
- ) -> R.DTensor((64, 128), "float32", "mesh[0]", "R"):
+ )""",
+ ) -> 'R.DTensor((64, 128), "float32", "mesh[0]", "R")':
cls = ShardedMLPWithTuple
weight1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
weight_packed[0]
lv0: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.matmul(x, weight1)
@@ -217,10 +217,10 @@ def test_mlp_const():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
- weight1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
- weight2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
- ) -> R.DTensor((128, 128), "float32", "mesh[0]", "R"):
+ x: 'R.DTensor((128, 128), "float32", "mesh[0]", "R")',
+ weight1: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[1]")',
+ weight2: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[0]")',
+ ) -> 'R.DTensor((128, 128), "float32", "mesh[0]", "R")':
lv0: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.matmul(x, weight1)
lv1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.nn.gelu(lv0)
lv2: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") = R.add(
@@ -267,10 +267,10 @@ def test_mlp_dynamic_shape():
@R.function
def foo(
- x: R.DTensor(("m", "k0"), "float32", "mesh[0]", "R"),
- weight1: R.DTensor(("k0", "k1"), "float32", "mesh[0]", "S[1]"),
- weight2: R.DTensor(("k1", "n"), "float32", "mesh[0]", "S[0]"),
- ) -> R.DTensor(("m", "n"), "float32", "mesh[0]", "R"):
+ x: 'R.DTensor(("m", "k0"), "float32", "mesh[0]", "R")',
+ weight1: 'R.DTensor(("k0", "k1"), "float32", "mesh[0]", "S[1]")',
+ weight2: 'R.DTensor(("k1", "n"), "float32", "mesh[0]", "S[0]")',
+ ) -> 'R.DTensor(("m", "n"), "float32", "mesh[0]", "R")':
m = T.int64()
n = T.int64()
k0 = T.int64()
@@ -328,12 +328,12 @@ def test_mlp_pipeline_parallelism():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", "mesh[0]", "R"),
- weight1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]"),
- weight2: R.DTensor((128, 128), "float32", "mesh[0]", "S[0]"),
- weight3: R.DTensor((128, 128), "float32", "mesh[1]", "S[1]"),
- weight4: R.DTensor((128, 128), "float32", "mesh[1]", "S[0]"),
- ) -> R.DTensor((128, 128), "float32", "mesh[1]", "R"):
+ x: 'R.DTensor((128, 128), "float32", "mesh[0]", "R")',
+ weight1: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[1]")',
+ weight2: 'R.DTensor((128, 128), "float32", "mesh[0]", "S[0]")',
+ weight3: 'R.DTensor((128, 128), "float32", "mesh[1]", "S[1]")',
+ weight4: 'R.DTensor((128, 128), "float32", "mesh[1]", "S[0]")',
+ ) -> 'R.DTensor((128, 128), "float32", "mesh[1]", "R")':
lv0: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.matmul(x, weight1)
lv1: R.DTensor((128, 128), "float32", "mesh[0]", "S[1]") =
R.nn.gelu(lv0)
lv3: R.DTensor((128, 128), "float32", "mesh[0]", "R") =
R.matmul(lv1, weight2)
@@ -627,19 +627,19 @@ def test_decoder_layer():
@R.function(pure=False)
def foo(
- input_tokens: R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R"),
- mask: R.DTensor((1, 1, 256, 256), "float16", "mesh[0]", "R"),
- div_const: R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]"),
- maximum_const: R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]"),
+ input_tokens: 'R.DTensor((1, 256, 4096), "float16", "mesh[0]",
"R")',
+ mask: 'R.DTensor((1, 1, 256, 256), "float16", "mesh[0]", "R")',
+ div_const: 'R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]")',
+ maximum_const: 'R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]")',
kv_cache: R.Tuple(R.Any, R.Any),
- linear_weight: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight1: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight2: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight3: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[1]"),
- rms_norm_weight: R.DTensor((4096,), "float16", "mesh[0]", "R"),
- cos_cached: R.DTensor((2048, 128), "float16", "mesh[0]", "R"),
- sin_cached: R.DTensor((2048, 128), "float16", "mesh[0]", "R"),
- ) -> R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R"):
+ linear_weight: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight1: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight2: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight3: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[1]")',
+ rms_norm_weight: 'R.DTensor((4096,), "float16", "mesh[0]", "R")',
+ cos_cached: 'R.DTensor((2048, 128), "float16", "mesh[0]", "R")',
+ sin_cached: 'R.DTensor((2048, 128), "float16", "mesh[0]", "R")',
+ ) -> 'R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R")':
cls = ShardedLlamaAttentionLayer
lv6 = R.dist.call_tir(
cls.rms_norm,
@@ -1334,19 +1334,19 @@ def test_decoder_layer_tir():
@R.function(pure=False)
def foo(
- input_tokens: R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R"),
- mask: R.DTensor((1, 1, 256, 256), "float16", "mesh[0]", "R"),
- div_const: R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]"),
- maximum_const: R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]"),
+ input_tokens: 'R.DTensor((1, 256, 4096), "float16", "mesh[0]",
"R")',
+ mask: 'R.DTensor((1, 1, 256, 256), "float16", "mesh[0]", "R")',
+ div_const: 'R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]")',
+ maximum_const: 'R.DTensor((1, 32, 256, 256), "float16", "mesh[0]",
"S[1]")',
kv_cache: R.Tuple(R.Any, R.Any),
- linear_weight: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight1: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight2: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight3: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[1]"),
- rms_norm_weight: R.DTensor((4096,), "float16", "mesh[0]", "R"),
- cos_cached: R.DTensor((2048, 128), "float16", "mesh[0]", "R"),
- sin_cached: R.DTensor((2048, 128), "float16", "mesh[0]", "R"),
- ) -> R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R"):
+ linear_weight: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight1: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight2: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight3: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[1]")',
+ rms_norm_weight: 'R.DTensor((4096,), "float16", "mesh[0]", "R")',
+ cos_cached: 'R.DTensor((2048, 128), "float16", "mesh[0]", "R")',
+ sin_cached: 'R.DTensor((2048, 128), "float16", "mesh[0]", "R")',
+ ) -> 'R.DTensor((1, 256, 4096), "float16", "mesh[0]", "R")':
cls = ShardedLlamaAttentionLayerTIR
lv6 = R.dist.call_tir(
LlamaAttentionLayerTIR.get_global_var("rms_norm"),
@@ -1824,17 +1824,17 @@ def test_decoder_layer_dynamic_shape():
@R.function(pure=False)
def foo(
- input_tokens: R.DTensor((1, "n", 4096), "float16", "mesh[0]", "R"),
- mask: R.DTensor((1, 1, "n", "m"), "float16", "mesh[0]", "R"),
+ input_tokens: 'R.DTensor((1, "n", 4096), "float16", "mesh[0]",
"R")',
+ mask: 'R.DTensor((1, 1, "n", "m"), "float16", "mesh[0]", "R")',
kv_cache: R.Tuple(R.Any, R.Any),
- linear_weight: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight1: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight2: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]"),
- linear_weight3: R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[1]"),
- rms_norm_weight: R.DTensor((4096,), "float16", "mesh[0]", "R"),
- cos_cached: R.DTensor((2048, 128), "float16", "mesh[0]", "R"),
- sin_cached: R.DTensor((2048, 128), "float16", "mesh[0]", "R"),
- ) -> R.DTensor((1, "n", 4096), "float16", "mesh[0]", "R"):
+ linear_weight: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight1: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight2: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[0]")',
+ linear_weight3: 'R.DTensor((4096, 4096), "float16", "mesh[0]",
"S[1]")',
+ rms_norm_weight: 'R.DTensor((4096,), "float16", "mesh[0]", "R")',
+ cos_cached: 'R.DTensor((2048, 128), "float16", "mesh[0]", "R")',
+ sin_cached: 'R.DTensor((2048, 128), "float16", "mesh[0]", "R")',
+ ) -> 'R.DTensor((1, "n", 4096), "float16", "mesh[0]", "R")':
n = T.int64()
m = T.int64()
cls = ShardedLlamaAttentionLayerDynamicShape
diff --git
a/tests/python/relax/distributed/test_distributed_tvmscript_parser.py
b/tests/python/relax/distributed/test_distributed_tvmscript_parser.py
index 4a9a4e70d7..d6ffc96e84 100644
--- a/tests/python/relax/distributed/test_distributed_tvmscript_parser.py
+++ b/tests/python/relax/distributed/test_distributed_tvmscript_parser.py
@@ -69,8 +69,8 @@ def test_call_tir_dtensor():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R"),
- ) -> R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R"):
+ x: 'R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R")',
+ ) -> 'R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R")':
gv0 = R.dist.call_tir(
TestModule.tir_func,
x,
@@ -132,8 +132,8 @@ def test_explicit_device_id():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R"),
- ) -> R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R"):
+ x: 'R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R")',
+ ) -> 'R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R")':
gv0 = R.dist.call_tir(
TestModule.tir_func,
x,
@@ -172,8 +172,8 @@ def test_constant():
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R"),
- ) -> R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R"):
+ x: 'R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R")',
+ ) -> 'R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R")':
gv0 = R.dist.call_tir(
TestModule.tir_func,
x,
diff --git
a/tests/python/relax/distributed/test_distributed_tvmscript_printer.py
b/tests/python/relax/distributed/test_distributed_tvmscript_printer.py
index 050d0a9d1a..d6af70eb4d 100644
--- a/tests/python/relax/distributed/test_distributed_tvmscript_printer.py
+++ b/tests/python/relax/distributed/test_distributed_tvmscript_printer.py
@@ -90,8 +90,8 @@ class TestModule:
@R.function
def foo(
- x: R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R"),
- ) -> R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R"):
+ x: 'R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R")',
+ ) -> 'R.DTensor((128, 128), "float32", device_mesh="mesh[0]",
placement="S[0], R")':
gv0 = R.dist.call_tir(
TestModule.tir_func,
x,
@@ -120,6 +120,8 @@ def test_module():
_assert_print(
TestModule,
"""
+from __future__ import annotations
+
# from tvm.script import ir as I
# from tvm.script import tirx as T
# from tvm.tirx.layout import Axis
diff --git a/tests/python/relax/test_analysis_well_formed.py
b/tests/python/relax/test_analysis_well_formed.py
index d3593fd3da..5ebbb87b4c 100644
--- a/tests/python/relax/test_analysis_well_formed.py
+++ b/tests/python/relax/test_analysis_well_formed.py
@@ -1196,7 +1196,7 @@ def test_call_tir_with_dtensor_arguments():
I.module_global_infos({"mesh": [R.dist.device_mesh([4], I.Range(0,
4))]})
@R.function
- def main(A: R.dist.DTensor([8, 4], "float16", "mesh[0]", "S[0]")):
+ def main(A: 'R.dist.DTensor([8, 4], "float16", "mesh[0]", "S[0]")'):
B = R.dist.call_tir(
Module.flatten, A, out_ty=R.dist.DTensor([64], "float16",
"mesh[0]", "S[0]")
)
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
index 578d3c56b0..0884ca42c7 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -39,7 +39,7 @@ def test_dispatch_scanop():
I.module_global_infos({"vdevice": [R.vdevice("llvm", 0)]})
@R.function
- def foo(x: R.Tensor((2, 3), "float32", "llvm")):
+ def foo(x: 'R.Tensor((2, 3), "float32", "llvm")'):
with R.dataflow():
lv0 = R.cumsum(x, axis=1, dtype="float64", exclusive=False)
lv1 = R.cumprod(lv0, axis=1, dtype="float64", exclusive=False)
@@ -78,7 +78,7 @@ def test_dispatch_scanop_cuda():
I.module_global_infos({"vdevice": [R.vdevice("cuda", 0)]})
@R.function
- def main(x: R.Tensor(("m", 3), "float32", "cuda")):
+ def main(x: 'R.Tensor(("m", 3), "float32", "cuda")'):
with R.dataflow():
lv0 = R.cumsum(x, axis=1, exclusive=True)
lv1 = R.cumprod(lv0, axis=1)
@@ -124,7 +124,7 @@ def test_dispatch_sort():
I.module_global_infos({"vdevice": [R.vdevice("llvm", 0)]})
@R.function
- def foo(x: R.Tensor(("m", 3), "float32", "llvm")):
+ def foo(x: 'R.Tensor(("m", 3), "float32", "llvm")'):
m = T.int64()
with R.dataflow():
lv = R.sort(x, axis=1, descending=False)
@@ -156,7 +156,7 @@ def test_dispatch_sort_cuda():
I.module_global_infos({"vdevice": [R.vdevice("cuda")]})
@R.function
- def foo(x: R.Tensor((2, 3), "float32", "cuda")):
+ def foo(x: 'R.Tensor((2, 3), "float32", "cuda")'):
with R.dataflow():
lv = R.sort(x, axis=1, descending=False)
gv = lv
@@ -221,7 +221,7 @@ def test_dispatch_argsort():
I.module_global_infos({"vdevice": [R.vdevice("llvm", 0)]})
@R.function
- def foo(x: R.Tensor(("m", 3), "float32", "llvm")):
+ def foo(x: 'R.Tensor(("m", 3), "float32", "llvm")'):
m = T.int64()
with R.dataflow():
lv = R.argsort(x, axis=1, descending=False, dtype="int32")
@@ -252,7 +252,7 @@ def test_dispatch_argsort_cuda():
I.module_global_infos({"vdevice": [R.vdevice("cuda")]})
@R.function
- def foo(x: R.Tensor((2, 3), "float32", "cuda")):
+ def foo(x: 'R.Tensor((2, 3), "float32", "cuda")'):
with R.dataflow():
lv = R.argsort(x, axis=1, descending=False)
gv = lv
@@ -314,7 +314,7 @@ def test_dispatch_topk():
I.module_global_infos({"vdevice": [R.vdevice("llvm", 0)]})
@R.function
- def foo(x: R.Tensor(("m", 3), "float32", "llvm")):
+ def foo(x: 'R.Tensor(("m", 3), "float32", "llvm")'):
m = T.int64()
with R.dataflow():
lv = R.topk(x, k=2, axis=1, largest=True)
@@ -345,7 +345,7 @@ def test_dispatch_topk_cuda():
I.module_global_infos({"vdevice": [R.vdevice("cuda")]})
@R.function
- def foo(x: R.Tensor((2, 3), "float32", "cuda")):
+ def foo(x: 'R.Tensor((2, 3), "float32", "cuda")'):
with R.dataflow():
lv = R.topk(x, k=2, axis=1, largest=True)
gv = lv
@@ -379,7 +379,7 @@ def test_dispatch_topk_gpu():
I.module_global_infos({"vdevice": [R.vdevice("vulkan")]})
@R.function
- def foo(x: R.Tensor((2, 3), "float32", "vulkan")):
+ def foo(x: 'R.Tensor((2, 3), "float32", "vulkan")'):
with R.dataflow():
# Two same calls should have only one PrimFunc
lv0 = R.topk(x, k=2, axis=1, largest=True)
@@ -696,7 +696,7 @@ def
test_dispatch_cumsum_webgpu_symbolic_non_contiguous_axis():
I.module_global_infos({"vdevice": [R.vdevice("webgpu", 0)]})
@R.function
- def main(x: R.Tensor((1, "n", 9), "float32", "webgpu")):
+ def main(x: 'R.Tensor((1, "n", 9), "float32", "webgpu")'):
return R.cumsum(x, axis=1)
target = tvm.target.Target("webgpu", host="llvm")
diff --git a/tests/python/relax/test_transform_legalize_ops.py
b/tests/python/relax/test_transform_legalize_ops.py
index 6e76b5887d..78254923c4 100644
--- a/tests/python/relax/test_transform_legalize_ops.py
+++ b/tests/python/relax/test_transform_legalize_ops.py
@@ -375,7 +375,7 @@ def test_legalize_with_vdevice():
@R.function
def func_llvm(
- A: R.Tensor([32, 32], "float32", "llvm"), B: R.Tensor([32, 32],
"float32", "llvm")
+ A: 'R.Tensor([32, 32], "float32", "llvm")', B: 'R.Tensor([32, 32],
"float32", "llvm")'
):
C = R.add(A, B)
return C
@@ -407,8 +407,8 @@ def test_legalize_with_vdevice():
@R.function
def func_llvm(
- A: R.Tensor((32, 32), dtype="float32", vdevice="llvm"),
- B: R.Tensor((32, 32), dtype="float32", vdevice="llvm"),
+ A: 'R.Tensor((32, 32), dtype="float32", vdevice="llvm")',
+ B: 'R.Tensor((32, 32), dtype="float32", vdevice="llvm")',
):
cls = Expected
C = R.call_tir(
diff --git a/tests/python/relax/test_transform_realize_vdevice.py
b/tests/python/relax/test_transform_realize_vdevice.py
index d363d54226..f3a315f8c3 100644
--- a/tests/python/relax/test_transform_realize_vdevice.py
+++ b/tests/python/relax/test_transform_realize_vdevice.py
@@ -86,10 +86,10 @@ def test_dataflow_binding():
@R.function
def foo(
- x: R.Tensor((2, 3), "float32", "llvm"),
- y: R.Tensor((2, 3), "float32", "llvm"),
- z: R.Tensor((2, 3), "float32", "llvm"),
- ) -> R.Tensor((2, 3), "float32", "llvm"):
+ x: 'R.Tensor((2, 3), "float32", "llvm")',
+ y: 'R.Tensor((2, 3), "float32", "llvm")',
+ z: 'R.Tensor((2, 3), "float32", "llvm")',
+ ) -> 'R.Tensor((2, 3), "float32", "llvm")':
with R.dataflow():
x1: R.Tensor((2, 3), "float32", "llvm") = x
y1: R.Tensor((2, 3), "float32", "llvm") = y
@@ -144,10 +144,10 @@ def test_binding():
@R.function
def foo(
- x: R.Tensor((2, 3), "float32", "llvm"),
- y: R.Tensor((2, 3), "float32", "llvm"),
- z: R.Tensor((2, 3), "float32", "llvm"),
- ) -> R.Tensor((2, 3), "float32", "llvm"):
+ x: 'R.Tensor((2, 3), "float32", "llvm")',
+ y: 'R.Tensor((2, 3), "float32", "llvm")',
+ z: 'R.Tensor((2, 3), "float32", "llvm")',
+ ) -> 'R.Tensor((2, 3), "float32", "llvm")':
x1: R.Tensor((2, 3), "float32", "llvm") = x
y1: R.Tensor((2, 3), "float32", "llvm") = y
x2: R.Tensor((2, 3), "float32", "llvm") = x1
@@ -177,7 +177,7 @@ def test_func_ret():
x: R.Tensor((2, 3), "float32"),
y: R.Tensor((2, 3), "float32"),
z: R.Tensor((2, 3), "float32"),
- ) -> R.Tensor((2, 3), "float32", "cuda"):
+ ) -> 'R.Tensor((2, 3), "float32", "cuda")':
with R.dataflow():
lv0 = R.add(x, y)
gv = R.multiply(lv0, z)
@@ -197,10 +197,10 @@ def test_func_ret():
@R.function
def foo(
- x: R.Tensor((2, 3), "float32", "cuda"),
- y: R.Tensor((2, 3), "float32", "cuda"),
- z: R.Tensor((2, 3), "float32", "cuda"),
- ) -> R.Tensor((2, 3), "float32", "cuda"):
+ x: 'R.Tensor((2, 3), "float32", "cuda")',
+ y: 'R.Tensor((2, 3), "float32", "cuda")',
+ z: 'R.Tensor((2, 3), "float32", "cuda")',
+ ) -> 'R.Tensor((2, 3), "float32", "cuda")':
with R.dataflow():
lv0: R.Tensor((2, 3), "float32", "cuda") = R.add(x, y)
gv: R.Tensor((2, 3), "float32", "cuda") = R.multiply(lv0, z)
@@ -227,7 +227,7 @@ def test_tuple_func_ret():
x: R.Tensor((2, 3), "float32"),
y: R.Tensor((2, 3), "float32"),
z: R.Tensor((2, 3), "float32"),
- ) -> R.Tuple([R.Tensor((2, 3), "float32", "cuda"), R.Tensor((2, 3),
"float32", "cuda")]):
+ ) -> 'R.Tuple([R.Tensor((2, 3), "float32", "cuda"), R.Tensor((2, 3),
"float32", "cuda")])':
with R.dataflow():
lv0 = R.add(x, y)
gv = R.multiply(lv0, z)
@@ -247,10 +247,10 @@ def test_tuple_func_ret():
@R.function
def foo(
- x: R.Tensor((2, 3), "float32", "cuda"),
- y: R.Tensor((2, 3), "float32", "cuda"),
- z: R.Tensor((2, 3), "float32", "cuda"),
- ) -> R.Tuple([R.Tensor((2, 3), "float32", "cuda"), R.Tensor((2, 3),
"float32", "cuda")]):
+ x: 'R.Tensor((2, 3), "float32", "cuda")',
+ y: 'R.Tensor((2, 3), "float32", "cuda")',
+ z: 'R.Tensor((2, 3), "float32", "cuda")',
+ ) -> 'R.Tuple([R.Tensor((2, 3), "float32", "cuda"), R.Tensor((2, 3),
"float32", "cuda")])':
with R.dataflow():
lv0: R.Tensor((2, 3), "float32", "cuda") = R.add(x, y)
gv: R.Tensor((2, 3), "float32", "cuda") = R.multiply(lv0, z)
@@ -280,7 +280,7 @@ def test_multi_device():
x: R.Tensor((2, 3), "float32"),
y: R.Tensor((2, 3), "float32"),
z: R.Tensor((2, 3), "float32"),
- ) -> R.Tensor((2, 3), "float32", "cuda"):
+ ) -> 'R.Tensor((2, 3), "float32", "cuda")':
with R.dataflow():
lv0 = R.add(x, y)
lv0 = R.hint_on_device(lv0, tvm.cpu())
@@ -306,10 +306,10 @@ def test_multi_device():
@R.function
def foo(
- x: R.Tensor((2, 3), "float32", "llvm"),
- y: R.Tensor((2, 3), "float32", "llvm"),
- z: R.Tensor((2, 3), "float32", "cuda"),
- ) -> R.Tensor((2, 3), "float32", "cuda"):
+ x: 'R.Tensor((2, 3), "float32", "llvm")',
+ y: 'R.Tensor((2, 3), "float32", "llvm")',
+ z: 'R.Tensor((2, 3), "float32", "cuda")',
+ ) -> 'R.Tensor((2, 3), "float32", "cuda")':
with R.dataflow():
lv0: R.Tensor((2, 3), "float32", "llvm") = R.add(x, y)
lv0: R.Tensor((2, 3), "float32", "llvm") = lv0
@@ -369,10 +369,10 @@ def test_insert_to_vdevice():
@R.function
def foo(
- x: R.Tensor((2, 3), "float32", "llvm"),
- y: R.Tensor((2, 3), "float32", "llvm"),
- z: R.Tensor((2, 3), "float32", "cuda"),
- ) -> R.Tensor((2, 3), "float32", "cuda"):
+ x: 'R.Tensor((2, 3), "float32", "llvm")',
+ y: 'R.Tensor((2, 3), "float32", "llvm")',
+ z: 'R.Tensor((2, 3), "float32", "cuda")',
+ ) -> 'R.Tensor((2, 3), "float32", "cuda")':
with R.dataflow():
lv0: R.Tensor((2, 3), "float32", "llvm") = y
lv1: R.Tensor((2, 3), "float32", "llvm") = R.add(x, lv0)
diff --git
a/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py
b/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py
index 3b441191f2..5ebbb113c5 100644
--- a/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py
+++ b/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py
@@ -179,8 +179,8 @@ def test_single_arg_return():
@R.function
def main(
- x: R.Tensor((2, 4, 26, 26), dtype="float32",
vdevice="opencl:1:global"),
- ) -> R.Tensor((2, 4, 13, 13), dtype="float32",
vdevice="opencl:1:global"):
+ x: 'R.Tensor((2, 4, 26, 26), dtype="float32",
vdevice="opencl:1:global")',
+ ) -> 'R.Tensor((2, 4, 13, 13), dtype="float32",
vdevice="opencl:1:global")':
cls = Input
with R.dataflow():
lv = R.call_tir(
@@ -277,12 +277,12 @@ def test_multi_arg_return():
@R.function
def main(
- x: R.Tensor((2, 16, 28, 28), dtype="float32",
vdevice="opencl:1:global"),
- w: R.Tensor((4, 16, 3, 3), dtype="float32",
vdevice="opencl:1:global"),
- ) -> R.Tuple(
+ x: 'R.Tensor((2, 16, 28, 28), dtype="float32",
vdevice="opencl:1:global")',
+ w: 'R.Tensor((4, 16, 3, 3), dtype="float32",
vdevice="opencl:1:global")',
+ ) -> """R.Tuple(
R.Tensor((2, 4, 26, 26), dtype="float32",
vdevice="opencl:1:global"),
R.Tensor((2, 4, 26, 26), dtype="float32",
vdevice="opencl:1:global"),
- ):
+ )""":
cls = Input
with R.dataflow():
lv = R.call_tir(
diff --git a/tests/python/relax/test_transform_update_vdevice.py
b/tests/python/relax/test_transform_update_vdevice.py
index fc62430f1a..be7d01f80a 100644
--- a/tests/python/relax/test_transform_update_vdevice.py
+++ b/tests/python/relax/test_transform_update_vdevice.py
@@ -55,8 +55,8 @@ def test_update():
@R.function
def main(
- a: R.Tensor((128, 128), "float32", "cuda:1"),
- c: R.Tensor((128, 128), "float32", "vdevice:3"),
+ a: 'R.Tensor((128, 128), "float32", "cuda:1")',
+ c: 'R.Tensor((128, 128), "float32", "vdevice:3")',
) -> R.Tensor((128, 128), "float32"):
s = R.add(a, c)
return s
@@ -77,9 +77,9 @@ def test_update():
@R.function
def main(
- a: R.Tensor((128, 128), dtype="float32", vdevice="metal:1"),
- c: R.Tensor((128, 128), dtype="float32", vdevice="metal:1"),
- ) -> R.Tensor((128, 128), dtype="float32", vdevice="metal:1"):
+ a: 'R.Tensor((128, 128), dtype="float32", vdevice="metal:1")',
+ c: 'R.Tensor((128, 128), dtype="float32", vdevice="metal:1")',
+ ) -> 'R.Tensor((128, 128), dtype="float32", vdevice="metal:1")':
s: R.Tensor((128, 128), dtype="float32", vdevice="metal:1") =
R.add(a, c)
return s
@@ -97,8 +97,8 @@ def test_update():
@R.function
def main(
- a: R.Tensor((128, 128), "float32", "cuda:0"),
- c: R.Tensor((128, 128), "float32", "cuda:0"),
+ a: 'R.Tensor((128, 128), "float32", "cuda:0")',
+ c: 'R.Tensor((128, 128), "float32", "cuda:0")',
) -> R.Tensor((128, 128), "float32"):
s = R.add(a, c)
return s
@@ -117,9 +117,9 @@ def test_update():
@R.function
def main(
- a: R.Tensor((128, 128), "float32", "llvm:1"),
- c: R.Tensor((128, 128), "float32", "llvm:1"),
- ) -> R.Tensor((128, 128), "float32", "llvm:1"):
+ a: 'R.Tensor((128, 128), "float32", "llvm:1")',
+ c: 'R.Tensor((128, 128), "float32", "llvm:1")',
+ ) -> 'R.Tensor((128, 128), "float32", "llvm:1")':
s: R.Tensor((128, 128), "float32", "llvm:1") = R.add(a, c)
return s
diff --git a/tests/python/relax/test_tvmscript_parser.py
b/tests/python/relax/test_tvmscript_parser.py
index 210f79acaf..7b49fa40aa 100644
--- a/tests/python/relax/test_tvmscript_parser.py
+++ b/tests/python/relax/test_tvmscript_parser.py
@@ -816,10 +816,10 @@ def test_tensor_with_vdevice():
@R.function
def foo(
- a: R.Tensor((128, 128), "float32", "cuda:1"),
- b: R.Tensor((128, 128), "float32", "llvm"),
- c: R.Tensor((128, 128), "float32", "vdevice:3"),
- ) -> R.Tensor((128, 128), "float32", "cuda:1"):
+ a: 'R.Tensor((128, 128), "float32", "cuda:1")',
+ b: 'R.Tensor((128, 128), "float32", "llvm")',
+ c: 'R.Tensor((128, 128), "float32", "vdevice:3")',
+ ) -> 'R.Tensor((128, 128), "float32", "cuda:1")':
s = R.add(a, c)
return s
diff --git a/tests/python/relax/test_vm_build.py
b/tests/python/relax/test_vm_build.py
index 67a22136c7..7c532de00a 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -1221,7 +1221,7 @@ def test_relax_module_with_multiple_targets(exec_mode):
@R.function
def func_llvm(
- A: R.Tensor([32, 32], "float32", "llvm"), B: R.Tensor([32, 32],
"float32", "llvm")
+ A: 'R.Tensor([32, 32], "float32", "llvm")', B: 'R.Tensor([32, 32],
"float32", "llvm")'
):
C = R.add(A, B)
return C
diff --git a/tests/python/script/minilang.py b/tests/python/script/minilang.py
index e7c40271a1..de9c99b30e 100644
--- a/tests/python/script/minilang.py
+++ b/tests/python/script/minilang.py
@@ -28,7 +28,7 @@ from contextlib import contextmanager
from dataclasses import dataclass, field
from types import SimpleNamespace
-from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import IRBuilder, resolve_global_info_args
from tvm.script.ir_builder.base import MISSING, AlreadyEmitted
from tvm.script.parser import entry
from tvm.script.parser import protocol_registry as registry
@@ -215,9 +215,8 @@ class Language:
entry.make_decorator(self.M)
)
- @registry.args_policy(
- "M.Tensor", {"shape": "expr_str", "device": "global_info"},
scalar_strings=False
- )
+ @resolve_global_info_args("device", resolver=self.resolve_global_info)
+ @registry.args_policy("M.Tensor", {"shape": "expr_str"},
scalar_strings=False)
def Tensor(shape=None, dtype="float32", device=None, placement="S[0]"):
return Value("tensor", (shape, dtype, device, placement))
diff --git a/tests/python/script/test_special_parser_protocol.py
b/tests/python/script/test_special_parser_protocol.py
index ed14fe7eb1..1c1da9d8ef 100644
--- a/tests/python/script/test_special_parser_protocol.py
+++ b/tests/python/script/test_special_parser_protocol.py
@@ -24,28 +24,28 @@ import traceback
import pytest
from tvm.script import tirx as T
+from tvm.script.ir_builder import resolve_global_info_args
from tvm.script.parser import protocol_registry as registry
def test_constructor_policy_survives_a_failed_definition(language):
- # A registered argument uses the dialect lookup, even after a constructor
failure.
+ # Aliased calls use the same builder lookup, even after a constructor
failure.
M = language.M
calls = []
dialect_infos = {}
first, invalid, second = object(), object(), object()
failure = RuntimeError("constructor failed")
- @registry.args_policy("M.constructor", {"device": "global_info"})
+ def resolve(name):
+ return dialect_infos[name]
+
+ @resolve_global_info_args("device", resolver=resolve)
def constructor(device):
calls.append(device)
if device is invalid:
raise failure
return device
- def resolve(name):
- return dialect_infos[name]
-
- M.resolve_global_info_ = resolve
M.constructor = constructor
def build():
@@ -57,14 +57,14 @@ def
test_constructor_policy_survives_a_failed_definition(language):
return main
dialect_infos["mesh"] = first
- assert build().body == [("emit", first), ("emit", "mesh")]
+ assert build().body == [("emit", first), ("emit", first)]
dialect_infos["mesh"] = invalid
with pytest.raises(RuntimeError, match="constructor failed") as caught:
build()
assert caught.value is failure
dialect_infos["mesh"] = second
- assert build().body == [("emit", second), ("emit", "mesh")]
- assert calls == [first, "mesh", invalid, second, "mesh"]
+ assert build().body == [("emit", second), ("emit", second)]
+ assert calls == [first, first, invalid, second, second]
def test_argument_policy_preserves_expression_dtype(language):
@@ -88,9 +88,9 @@ def test_argument_policy_preserves_expression_dtype(language):
assert increment.args[1] == 1
-def test_tirx_rejects_global_info_at_the_source_argument(monkeypatch):
- # A custom global-info argument must report TIRx's unsupported policy at
its source.
- @registry.args_policy("T.global_annotation", {"device": "global_info"})
+def test_tirx_rejects_global_info_at_the_call_site(monkeypatch):
+ # A custom builder reports its resolver failure at the ordinary call site.
+ @resolve_global_info_args("device", resolver=T.resolve_global_info_)
def global_annotation(device):
return T.int32
@@ -103,7 +103,7 @@ def
test_tirx_rejects_global_info_at_the_source_argument(monkeypatch):
def main(value: T.global_annotation(device="cuda:0")):
T.evaluate(value)
- lines, first =
inspect.getsourcelines(test_tirx_rejects_global_info_at_the_source_argument)
+ lines, first =
inspect.getsourcelines(test_tirx_rejects_global_info_at_the_call_site)
index, line = next((i, line) for i, line in enumerate(lines) if "def
main(value:" in line)
location = first + index
frames = traceback.extract_tb(caught.value.__traceback__)
@@ -112,7 +112,7 @@ def
test_tirx_rejects_global_info_at_the_source_argument(monkeypatch):
]
assert source_frames
if getattr(source_frames[-1], "colno", None) is not None:
- column = line.index('"cuda:0"')
+ column = line.index("T.global_annotation(")
assert (
source_frames[-1].colno,
source_frames[-1].end_lineno,
@@ -120,5 +120,5 @@ def
test_tirx_rejects_global_info_at_the_source_argument(monkeypatch):
) == (
column,
location,
- column + len('"cuda:0"'),
+ column + len('T.global_annotation(device="cuda:0")'),
)
diff --git a/tests/python/tvmscript/test_tvmscript_ir_builder_irmodule.py
b/tests/python/tvmscript/test_tvmscript_ir_builder_irmodule.py
index 85c9359795..760d7e7b8e 100644
--- a/tests/python/tvmscript/test_tvmscript_ir_builder_irmodule.py
+++ b/tests/python/tvmscript/test_tvmscript_ir_builder_irmodule.py
@@ -30,6 +30,7 @@ def test_ir_builder_irmodule():
with IRBuilder() as ib: # pylint: disable=invalid-name
with I.ir_module():
pass
+ assert not ib.frames
# the ir_module generated by IRBuilder
ir_module_actual = ib.get()
