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/24 18:25:25 UTC
(tvm) branch main updated: [REFACTOR][TVMSCRIPT] Introduce a new transpiler based parser (#20423)
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 0db97a5a2c [REFACTOR][TVMSCRIPT] Introduce a new transpiler based parser (#20423)
0db97a5a2c is described below
commit 0db97a5a2cf041f7d83db34325c4901db8ee0380
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Sep 24 14:25:19 2026 -0400
[REFACTOR][TVMSCRIPT] Introduce a new transpiler based parser (#20423)
This PR introduces a modular, customizable parser infrastructure for
TVMScript through transpilation. The key idea is to deterministically
transpile language syntax into a corresponding sequence of builder calls
while preserving source locations and providing span context.
Key design elements include:
- Acquire source and its definition scope, collect syntax facts once,
rewrite the owned AST, and execute the resulting builder program with
the original lexical bindings.
- Each language variant customizes behavior through a
`parser_protocol.py` file, which defines special handling for syntax
such as `for` and `if`.
- Create reusable span entries during rewriting. Attach locations to
meaningful expressions and pass spans directly to builder operations
that own them.
- Reuse existing builder and IR APIs, so new types and functions can be
added without changing the parser, provided they have well-defined
builder behavior.
- Generate specialization handling only when specialization is
requested.
For example, this kernel:
```python
@T.prim_func
def add_one(A: T.Buffer((128,), "float32"), B: T.Buffer((16, 8), "float32")):
for i in range(128):
B[i // 8, i % 8] = A[i] + 1
```
is lowered to the following builder calls. This presentation wraps the
emitted program in `build_add_one` and gives its local scaffolding
readable aliases. `I` supplies shared builders, `T` is the language
variant namespace, and `_S` holds source-span entries.
```python
def build_add_one():
with I.IRBuilder() as builder:
with T.function_(span=_S[6]) as fn:
T.func_name('add_one')
A = T.arg('A', _S[1](T.Buffer(_S[0]((128,)), 'float32')), span=_S[2])
B = T.arg('B', _S[4](T.Buffer(_S[3]((16, 8)), 'float32')), span=_S[5])
def build_body():
A, B = fn.params
with T.for_(T.range_(128), span=_S[7], names='i') as i:
T.setitem_(value=_S[9](_S[8](A[i]) + 1), target=B, key=(_S[10](i // 8), _S[11](i % 8)), span=_S[12])
build_body()
result = fn.function
result.__name__ = 'add_one'
T.check_well_formed_(result)
return fn.function
```
---
.github/workflows/main.yml | 4 +-
docs/arch/tvmscript.rst | 413 ++--
.../tutorials/mix_python_and_tvm_with_pymodule.py | 16 +-
docs/reference/api/python/relax/relax.rst | 1 +
docs/reference/api/python/script/ir_builder.rst | 6 +
docs/reference/api/python/script/parser.rst | 3 +
docs/reference/api/python/script/script.rst | 15 +
docs/tirx/api/script.rst | 6 +-
docs/tirx/tile_primitives/copy_async/dsmem.rst | 4 +-
include/tvm/ir/global_info.h | 1 +
include/tvm/relax/global_info.h | 1 -
include/tvm/relax/script/builder/frame.h | 17 +-
include/tvm/relax/script/builder/ir.h | 23 +
include/tvm/script/ir_builder/base.h | 74 +-
include/tvm/tirx/script/builder/frame.h | 49 +-
include/tvm/tirx/script/builder/ir.h | 32 +-
python/tvm/backend/cuda/lang/clc.py | 2 +-
python/tvm/backend/cuda/lang/warp_role.py | 2 +-
.../cuda/tile_primitive/copy_async/dsmem.py | 8 +-
python/tvm/backend/trn/script.py | 10 +-
.../backend/trn/tile_primitive/binary/default.py | 2 +-
.../trn/tile_primitive/compose_op/binary_reduce.py | 4 +-
.../trn/tile_primitive/compose_op/unary_reduce.py | 4 +-
.../tvm/backend/trn/tile_primitive/copy/default.py | 12 +-
.../tvm/backend/trn/tile_primitive/gemm/default.py | 8 +-
.../backend/trn/tile_primitive/select/default.py | 3 +-
.../tvm/backend/trn/tile_primitive/unary/utils.py | 6 +-
python/tvm/ir/module.py | 2 +-
python/tvm/relax/__init__.py | 21 +
python/tvm/relax/backend/gpu_generic/sampling.py | 2 +-
python/tvm/relax/base_py_module.py | 19 +-
python/tvm/relax/frontend/nn/llm/_page_kernels.py | 12 +-
.../relax/frontend/nn/llm/position_embedding.py | 4 +-
python/tvm/relax/frontend/nn/llm/tree_attn.py | 8 +-
.../tvm/relax/frontend/tflite/tflite_frontend.py | 4 +-
python/tvm/relax/global_info.py | 10 +-
python/tvm/relax/relax_to_pyfunc_converter.py | 6 +-
python/tvm/relax/script/__init__.py | 209 +-
python/tvm/relax/script/builder/__init__.py | 306 ++-
python/tvm/relax/script/builder/distributed/ir.py | 32 +-
python/tvm/relax/script/builder/frame.py | 10 +-
python/tvm/relax/script/builder/ir.py | 124 +-
python/tvm/relax/script/builder/parser_protocol.py | 546 +++++
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/transform/transform.py | 4 +-
python/tvm/s_tir/tensor_intrin/cuda.py | 54 +-
.../tvm/s_tir/tensor_intrin/dot_product_common.py | 2 +-
python/tvm/s_tir/tensor_intrin/metal.py | 4 +-
python/tvm/s_tir/tensor_intrin/rocm.py | 52 +-
python/tvm/s_tir/tensor_intrin/x86.py | 10 +-
python/tvm/script/__init__.py | 146 +-
python/tvm/script/ir_builder/__init__.py | 21 +-
python/tvm/script/ir_builder/base.py | 273 ++-
python/tvm/script/ir_builder/ir/__init__.py | 8 +-
python/tvm/script/ir_builder/ir/frame.py | 12 +-
python/tvm/script/ir_builder/ir/ir.py | 149 +-
python/tvm/script/ir_builder/ir/parser_protocol.py | 1490 ++++++++++++
python/tvm/script/parser/__init__.py | 88 +-
python/tvm/script/parser/_core.py | 24 -
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/evaluator.py | 626 -----
python/tvm/script/parser/core/parser.py | 1071 ---------
python/tvm/script/parser/core/utils.py | 212 --
python/tvm/script/parser/entry.py | 821 +++++++
python/tvm/script/parser/expr_str_handling.py | 136 ++
python/tvm/script/parser/inspect_source.py | 269 +++
python/tvm/script/parser/{core/error.py => ir.py} | 15 +-
python/tvm/script/parser/ir/__init__.py | 39 -
python/tvm/script/parser/ir/entry.py | 163 --
python/tvm/script/parser/ir/parser.py | 212 --
python/tvm/script/parser/jit_support.py | 83 +
python/tvm/script/parser/prescan.py | 728 ++++++
python/tvm/script/parser/protocol_registry.py | 314 +++
python/tvm/script/parser/transpile.py | 2520 ++++++++++++++++++++
python/tvm/script/printer/doc.py | 2 +-
python/tvm/tirx/__init__.py | 23 +
python/tvm/tirx/buffer.py | 4 +
python/tvm/tirx/layout.py | 9 +
python/tvm/tirx/op.py | 16 +
python/tvm/tirx/script/__init__.py | 87 +-
python/tvm/tirx/script/builder/__init__.py | 418 +++-
python/tvm/tirx/script/builder/frame.py | 52 +-
python/tvm/tirx/script/builder/ir.py | 396 ++-
python/tvm/tirx/script/builder/parser_protocol.py | 601 +++++
python/tvm/tirx/script/builder/tirx.py | 45 +-
python/tvm/tirx/script/builder/triton.py | 8 +-
python/tvm/tirx/script/jit.py | 278 +++
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 ----------
src/{relax => }/ir/global_info.cc | 30 +-
src/relax/ir/global_info.cc | 11 +-
src/relax/script/builder/frame.cc | 120 +-
src/relax/script/builder/ir.cc | 98 +-
src/relax/script/builder/utils.h | 13 +-
src/relax/{ir => script/printer}/global_info.cc | 48 +-
.../schedule/primitive/loop_transformation.cc | 40 +-
src/script/ir_builder/base.cc | 119 +-
src/script/ir_builder/ir/ir.cc | 44 +-
.../printer/doc_printer/python_doc_printer.cc | 1 -
src/script/printer/ir/ir.cc | 20 -
src/script/printer/utils.h | 26 +-
src/sym/modular_set.cc | 4 +-
src/tirx/analysis/verify_well_formed.h | 60 +-
src/tirx/script/builder/frame.cc | 139 +-
src/tirx/script/builder/ir.cc | 108 +-
src/tirx/script/builder/utils.h | 17 +-
src/tirx/script/printer/buffer.cc | 80 +-
src/tirx/script/printer/function.cc | 45 +-
src/tirx/script/printer/utils.h | 4 +-
src/tirx/transform/tile_primitive_dispatch.cc | 113 +-
tests/python/codegen/test_target_codegen_vulkan.py | 20 +-
.../python/contrib/test_tir_triton_integration.py | 14 +-
tests/python/disco/test_callback.py | 4 +-
.../test_transform_fold_vdevice_scope_change.py | 8 +-
.../test_distributed_tvmscript_printer.py | 8 +-
tests/python/relax/test_analysis_type_analysis.py | 8 +-
tests/python/relax/test_analysis_well_formed.py | 6 +-
.../relax/test_backend_dispatch_sort_scan.py | 38 +-
tests/python/relax/test_base_py_module_printer.py | 82 +-
.../relax/test_base_py_module_symbolic_shape.py | 5 +-
tests/python/relax/test_bind_symbolic_vars.py | 3 +-
tests/python/relax/test_blockbuilder_core.py | 10 +-
tests/python/relax/test_blockbuilder_emit_te.py | 2 +-
tests/python/relax/test_codegen_cutlass.py | 6 +-
tests/python/relax/test_contrib_vllm.py | 4 +-
tests/python/relax/test_dataflow_pattern.py | 1 -
tests/python/relax/test_dataflow_rewriter.py | 4 +-
tests/python/relax/test_frontend_nn_op.py | 6 +-
tests/python/relax/test_op_binary.py | 12 +-
tests/python/relax/test_op_index.py | 4 +-
tests/python/relax/test_op_inspect.py | 7 +-
tests/python/relax/test_op_view.py | 25 +-
tests/python/relax/test_pytorch_integration.py | 2 +-
tests/python/relax/test_relax_operators.py | 6 +-
.../python/relax/test_relax_to_pyfunc_converter.py | 120 +-
.../python/relax/test_runtime_builtin_rnn_state.py | 4 +-
tests/python/relax/test_tir_call_source_kernel.py | 3 +-
tests/python/relax/test_transform.py | 14 +-
.../relax/test_transform_bundle_model_params.py | 4 +-
.../relax/test_transform_canonicalize_bindings.py | 4 +-
.../relax/test_transform_compute_prim_value.py | 4 +-
.../python/relax/test_transform_convert_layout.py | 66 +-
tests/python/relax/test_transform_cse.py | 12 +-
tests/python/relax/test_transform_fuse_ops.py | 4 +-
.../relax/test_transform_fuse_ops_by_pattern.py | 14 +-
tests/python/relax/test_transform_fuse_tir.py | 14 +-
.../relax/test_transform_lazy_transform_params.py | 16 +-
tests/python/relax/test_transform_legalize_ops.py | 8 +-
.../relax/test_transform_legalize_ops_binary.py | 60 +-
..._transform_legalize_ops_index_linear_algebra.py | 4 +-
.../python/relax/test_transform_legalize_ops_nn.py | 13 +-
.../relax/test_transform_lift_transform_params.py | 2 +-
.../relax/test_transform_meta_schedule_tuning.py | 1 -
.../python/relax/test_transform_realize_vdevice.py | 62 +-
.../test_transform_remove_unused_parameters.py | 4 +-
.../test_transform_rewrite_dataflow_reshape.py | 4 +-
...nsform_specialize_primfunc_based_on_callsite.py | 8 +-
.../python/relax/test_transform_update_vdevice.py | 24 +-
tests/python/relax/test_tvmscript_parser.py | 196 +-
.../relax/test_tvmscript_parser_op_manipulate.py | 4 +-
tests/python/relax/test_tvmscript_printer_relax.py | 6 +-
tests/python/relax/test_tvmscript_pyfunc.py | 14 +-
tests/python/relax/test_tvmscript_type_vars.py | 4 +-
tests/python/relax/test_utils.py | 13 +-
tests/python/relax/test_vm_build.py | 2 +-
tests/python/relax/test_vm_multi_device.py | 14 +-
.../s_tir/dlight/test_gpu_matmul_tensorize.py | 180 +-
...est_meta_schedule_postproc_rewrite_tensorize.py | 10 +-
.../test_meta_schedule_trace_apply.py | 36 +-
.../schedule/test_tir_schedule_cache_read_write.py | 8 +-
.../s_tir/schedule/test_tir_schedule_pad_einsum.py | 8 +-
.../s_tir/schedule/test_tir_schedule_reindex.py | 3 +-
.../s_tir/schedule/test_tir_schedule_set_scope.py | 7 +-
.../schedule/test_tir_schedule_transform_layout.py | 31 +-
tests/python/s_tir/test_arith_domain_touched.py | 7 +-
.../test_s_tir_transform_inject_ptx_async_copy.py | 28 +-
.../test_s_tir_transform_lower_opaque_block.py | 5 +-
...st_s_tir_transform_memhammer_lower_auto_copy.py | 3 +-
tests/python/script/conftest.py | 66 +
tests/python/script/minilang.py | 348 +++
tests/python/script/test_basic_usage.py | 1046 ++++++++
tests/python/script/test_cross_function_call.py | 83 +
tests/python/script/test_error_handling.py | 223 ++
tests/python/script/test_jit_specialization.py | 210 ++
tests/python/script/test_meta_programming.py | 500 ++++
tests/python/script/test_parser_entry.py | 87 +
tests/python/script/test_source_locations.py | 486 ++++
.../python/script/test_special_parser_protocol.py | 124 +
tests/python/script/test_symbolic_shape.py | 191 ++
tests/python/te/test_te_create_primfunc.py | 6 +-
.../test_tir_analysis_verify_well_formed.py | 6 +-
.../test_tir_transform_bf16_legalize.py | 1 +
...test_tir_transform_force_narrow_index_to_i32.py | 8 +-
.../tirx-transform/test_tir_transform_vectorize.py | 65 +-
.../python/tirx/codegen/test_codegen_blackwell.py | 10 +-
tests/python/tirx/codegen/test_codegen_cuda.py | 4 +-
tests/python/tirx/codegen/test_ptx_addr.py | 24 +-
tests/python/tirx/codegen/test_ptx_dialect.py | 132 +-
tests/python/tirx/iket/test_iket_profiler.py | 10 +-
.../operator/tile_primitive/trn/test_binary_trn.py | 38 +-
.../operator/tile_primitive/trn/test_unary_trn.py | 12 +-
tests/python/tirx/test_inline.py | 2 +-
tests/python/tirx/test_jit.py | 43 +-
tests/python/tirx/test_op_namespace_cleanup.py | 7 +-
tests/python/tirx/test_parser_printer.py | 303 +--
tests/python/tirx/test_tvmscript_type_vars.py | 4 +-
.../tirx/transform/test_transform_lower_tirx.py | 48 +-
.../tvmscript/test_tvmscript_error_report.py | 274 +--
.../tvmscript/test_tvmscript_ir_builder_tir.py | 2 +-
tests/python/tvmscript/test_tvmscript_ops.py | 2 +-
.../tvmscript/test_tvmscript_parser_evaluator.py | 13 +-
.../tvmscript/test_tvmscript_parser_source.py | 6 +-
.../python/tvmscript/test_tvmscript_parser_tir.py | 31 +-
.../tvmscript/test_tvmscript_printer_annotation.py | 9 +-
.../python/tvmscript/test_tvmscript_printer_tir.py | 19 +-
tests/python/tvmscript/test_tvmscript_roundtrip.py | 17 +-
.../tvmscript/test_tvmscript_syntax_sugar.py | 51 +-
227 files changed, 15600 insertions(+), 10277 deletions(-)
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 6ff75d1aea..2827ed917a 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -76,7 +76,7 @@ jobs:
- name: Test
shell: bash -l {0}
run: >-
- python -m pytest -v tests/python/all-platform-minimal-test
+ python -m pytest -v tests/python/all-platform-minimal-test tests/python/script
- name: Minimal Metal Compile-Only
shell: bash -l {0}
run: |
@@ -116,4 +116,4 @@ jobs:
- name: Test
shell: cmd /C call {0}
run: >-
- python -m pytest -v tests/python/all-platform-minimal-test
+ python -m pytest -v tests/python/all-platform-minimal-test tests/python/script
diff --git a/docs/arch/tvmscript.rst b/docs/arch/tvmscript.rst
index 0dc8fc63cb..ce3f3f86f4 100644
--- a/docs/arch/tvmscript.rst
+++ b/docs/arch/tvmscript.rst
@@ -22,22 +22,22 @@ TVMScript
TVMScript is a Python-based domain-specific language (DSL) for writing TVM IR. It lets users
define ``IRModule``\ s — containing both Relax functions and TIR ``PrimFunc``\ s — using
-familiar Python syntax. Although TVMScript *looks* like Python, it is **not executed by the
-Python interpreter**. Instead, Python decorators extract the AST from the source code and
-transform it into TVM IR through a dedicated parser and IR builder pipeline.
+familiar Python syntax. Decorators acquire the source and its definition context. A syntax
+transpiler translates the Python AST into a Python builder program; executing that program
+constructs TVM IR. The scripted function body is therefore translated before execution.
TVMScript serves two roles in the TVM stack:
- **Authoring**: users write TIR kernels and Relax programs directly in TVMScript.
-- **Roundtrip**: every ``IRModule`` can be printed back to TVMScript via ``mod.script()`` and
- re-parsed to produce an equivalent module. This makes TVMScript the primary tool for
- inspecting, debugging, and serializing IR.
+- **Roundtrip**: ``mod.script()`` prints IR as TVMScript, which can be re-parsed with any
+ required external symbols to reconstruct equivalent IR. This supports inspection,
+ debugging, and serialization of printable IR.
Overview
--------
-The TVMScript system has three components:
+TVMScript separates source acquisition, syntax translation, IR construction, and printing:
.. code-block:: text
@@ -45,13 +45,17 @@ The TVMScript system has three components:
Python source (TVMScript)
│
- ▼ ast.parse + convert
+ ▼ Frontend: source + definition context
│
- Doc AST (mirror of Python AST)
+ Python AST
│
- ▼ Parser (dispatch by token: ir / tirx / relax)
+ ▼ Syntax transpiler + construction protocol
│
- ▼ IR Builder (frame stack)
+ Generated Python builder program
+ │
+ ▼ Frontend: recompose and execute callable
+ │
+ ▼ IR builders (frame stack + dialect policy)
│
TVM IR (IRModule, PrimFunc, relax.Function)
@@ -68,14 +72,17 @@ The TVMScript system has three components:
│
TVMScript text
-- **Parser** (Python): reads Python source, converts it to a ``Doc AST`` (a mirror of
- Python's ``ast`` module), then walks the tree using dialect-specific handlers that call
- into the IR builder.
-- **IR Builder** (Python + C++): provides a frame-stack API where each ``with`` block or
- decorator pushes a frame. When the frame exits, the constructed IR is finalized. The builder
- is shared across dialects — TIR and Relax each register their own frame types.
-- **Printer** (C++): converts TVM IR objects to a ``Doc`` tree (an intermediate representation
- of Python syntax), then formats the tree into valid TVMScript text.
+- **Frontend** (Python): captures definition context, acquires source, and composes the
+ generated callable with its lexical environment. It owns decorators, helper entry points,
+ and execution of the builder program.
+- **Syntax transpiler** (Python): rewrites Python AST nodes into calls on the selected
+ construction namespace. It preserves source scopes, evaluation order, and locations.
+- **IR builders** (Python + C++): own typed values, symbol identity, dialect policy, and
+ construction frames. Exiting a frame finalizes its IR and attaches it to its parent.
+- **Printer** (C++): converts IR to a ``Doc`` tree and formats it as Python syntax. The
+ printer's Doc tree is separate from the Python AST used for parsing.
+
+All source entry points use the canonical ``tvm.script.parser`` implementation.
Decorators
@@ -88,6 +95,12 @@ TVMScript uses three import aliases by convention:
from tvm.script import ir as I # module-level constructs
from tvm.script import tirx as T # TIR constructs
from tvm.script import relax as R # Relax constructs
+ from tvm.script.parser.entry import make_macro_decorator
+ from tvm.target import Target
+
+These are public authoring APIs. ``Target`` configures compilation targets; it remains in
+``tvm.target``. Public dialect namespaces expose decorators and source constructors, so
+scripted programs do not need direct imports from builder implementation packages.
The primary decorators are:
@@ -102,7 +115,7 @@ These can be composed:
@I.ir_module
class MyModule:
- @T.prim_func
+ @T.prim_func(s_tir=True)
def add_kernel(A: T.Buffer((128,), "float32"),
B: T.Buffer((128,), "float32"),
C: T.Buffer((128,), "float32")):
@@ -115,82 +128,228 @@ These can be composed:
def main(x: R.Tensor((128,), "float32"),
y: R.Tensor((128,), "float32")) -> R.Tensor((128,), "float32"):
with R.dataflow():
- out = R.call_tir(cls.add_kernel, (x, y),
+ out = R.call_tir(MyModule.add_kernel, (x, y),
out_ty=R.Tensor((128,), "float32"))
R.output(out)
return out
-When Python encounters ``@I.ir_module``, the decorator does **not** execute the class body.
-Instead, it calls ``tvm.script.parse()`` which extracts the source code of the class,
-builds a Doc AST, and hands it to the parser.
-
-
-Parser Architecture
--------------------
+Python creates the class before invoking ``@I.ir_module``. Function decorators inside an
+IR-module class retain method definitions for module construction. The module frontend then
+acquires the source and constructs declarations and bodies in the module's builder context.
+Standalone ``@T.prim_func`` and ``@R.function`` definitions construct their IR immediately.
-The parser lives in ``python/tvm/script/parser/``.
+A helper can construct expressions or statements in the caller's active builder frames:
-Dispatch mechanism
-~~~~~~~~~~~~~~~~~~
-
-Different IR dialects (TIR, Relax) need different handling for the same Python syntax. For
-example, ``if ... else`` inside ``@T.prim_func`` creates a TIR ``If`` branch, while the same
-syntax inside ``@R.function`` creates a Relax ``If`` node with different semantics.
-
-The parser maintains a **dispatch token** stack (``["default"]`` initially). When it encounters
-a decorated function, it inspects the decorator to determine the token — ``"tirx"`` for
-``@T.prim_func``, ``"relax"`` for ``@R.function`` — and pushes it onto the stack.
+.. code-block:: python
-Each AST node type is dispatched via a virtual table:
+ @make_macro_decorator(T)
+ def add_one(value):
+ return value + 1
-.. code-block:: text
+Unlike a function decorator, ``make_macro_decorator`` does not create an IR function for
+each helper. Its default return behavior is ordinary Python return from the generated
+helper program.
+``T.inline`` and the dialect macro decorators are supplied through the same frontend helper
+mechanism.
- ParseVTable[(token, node_type)] → handler function
- Lookup order:
- 1. (current_token, node_type) e.g. ("tirx", "For")
- 2. ("default", node_type) e.g. ("default", "For")
- 3. generic_visit fallback
+Parser Architecture
+-------------------
-Dialect-specific parsers (``parser/tirx/parser.py``, ``parser/relax/parser.py``) register
-handlers using ``@dispatch.register(token, type_name)`` decorators.
+The parser lives in ``python/tvm/script/parser/``. ``inspect_source.py`` owns source
+acquisition, the ``Source`` representation and its coordinate helpers, and definition-context
+capture. Documented functions in ``entry.py`` manage the parse sequence, callable composition
+and execution.
+The transpiler handles syntax; the builders interpret concrete values. TIRx JIT validation,
+defaults and specialization caching live in ``tvm.tirx.script.jit``; shared parser utilities
+remain independent of that policy.
+
+Syntax and construction protocol
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+A registered function decorator selects a construction namespace. The same Python syntax
+then lowers to operations supplied by that namespace: for example, unmarked ``if`` creates
+TIR control-flow frames in a primitive function and Relax control-flow frames in a Relax
+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;
+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,
+local aliases, methods and closures execute as Python without marker inference. The
+dictionaries retain no source functions or parse captures. Dialects own their namespace
+exports, entry points and aliases and register them with the generic parser.
+The registry selects a call's argument policy before traversal and reuses positional
+parameter names computed once at registration;
+``parser/expr_str_handling.py`` decodes annotation/expression strings and maps their
+escaped, multiline and UTF-8 source ranges. The main visitor rewrites the decoded syntax.
+
+Assignments become binding operations, standalone expressions become emission operations,
+and loops and scopes become builder contexts. An ordinary assignment passes its evaluated
+RHS and a separate ``value_span`` to ``bind_``, without first stamping the returned value.
+The binding target has its own span. TIRx preserves returned Vars, including BufferVars,
+and ordinary metadata objects; other expressions follow the dialect's binding rules.
+Construction still runs in the source context, and nested expressions retain their locations.
+Concrete binding, type checking, comparison construction, and frame finalization belong to
+the builders. Ordinary host calls and
+operator overloads execute as part of the generated Python program.
+
+Loop frames create their named variables during construction. Entering a frame returns
+the variable directly for one dimension and the variable sequence for multiple dimensions,
+so ordinary Python uses ``with T.serial(0, n) as i``. Generated scalar loops use the same
+entry form. Generated tuple, list and starred targets unpack the entered frame's stable
+``vars`` sequence, preserving explicit one-dimensional unpacking and native variable identity.
+
+Parser syntax restrictions raise ``SyntaxError`` with the original filename and source
+range. Python helper and builder exceptions propagate unchanged, preserving their identity,
+type and source-mapped traceback. Temporary parser contexts are released on both successful
+and exceptional exits.
Parse flow
~~~~~~~~~~
-The entry point is ``parse(program, extra_vars)``:
-
-1. **Source extraction**: the program's source code is extracted (from a class, function, or
- string) and converted to a Doc AST via Python's ``ast`` module.
+The public entry points include ``tvm.script.parse`` and ``tvm.script.from_source``:
+
+1. **Acquire source and context**: the frontend accepts a function, class, or source string,
+ obtains its Python AST, and retains source locations. Decorators capture the definition
+ context temporarily through ``definition_scope``; source-string callers can supply
+ external bindings through ``extra_vars``. Source inspection belongs to
+ ``parser.inspect_source``.
+2. **Prescan and translate syntax**: acquisition supplies a fresh owned AST. One
+ ``PrescanCollector`` produces read-only ``PrescanContext`` facts for reserved names,
+ scoped declarations and conditional outputs. ``IRBuilderTranspiler`` rewrites that owned
+ tree with the standard Python AST visitor, using fixed namespace metadata and registered
+ argument policies without copying or reparsing the tree. One shared name allocator
+ avoids collisions with user identifiers through rewriting and recomposition.
+3. **Compose the callable**: the private frontend helper ``_recompose_builder`` compiles
+ the generated program with the source's globals, closure bindings, and necessary
+ annotation context. Explicit temporary phase inputs connect generated helpers to
+ their source functions without rediscovering them through a generated-AST walk.
+4. **Execute construction**: the callable enters native builder frames and reconstructs
+ declarations and annotations in their required context. When forward declarations are
+ needed, it declares all signatures first. Inside each retained native function frame,
+ it defines and calls a lexical body helper without an explicit frame argument. Re-entry
+ uses the existing parameters instead of adding them again. Native frames own module
+ references, results and function-local symbols.
+5. **Validate and return IR**: generated code calls the dialect's
+ ``X.check_well_formed_`` for a completed function or the shared
+ ``I.check_well_formed_`` for a completed module. Public parsing validates by default.
+ Parsing owns execution and releases its temporary builder and captures afterward.
+
+Definition and symbol scopes
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Python lexical bindings and DSL symbolic declarations have different owners. The frontend
+captures the exact definition scope during decoration, including enclosing names used only
+in annotations that Python may omit from the function's closure. Ordinary body references
+retain their source globals and closure bindings. Unrelated callers' local variables do not
+become part of this environment.
+
+Annotations are reconstructed as builder operations within the module/function context.
+The generated program preserves their definition scope and original names, using a
+distinct compiler binding only when a source-local binding would shadow the annotation.
+One prescan-seeded name allocator serves rewriting and callable recomposition. Body
+globals and closures retain their original lookup, including a method closure whose name
+also appears as a class member. Callable recomposition and
+lexical environment setup belong to the frontend. The transpiler preserves source scopes
+through Python syntax; it does not classify each name by membership in a runtime environment
+or search caller stacks for values.
+
+DSL expression strings such as ``R.Tensor(("n", 4), "float32")`` use
+``X.resolve_type_var_("n")``, which delegates to the nearest native function frame's
+``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.
+
+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;
+``_S[i].ctx(lambda: helper(value))`` evaluates the callee and arguments once within
+the source context and restores it on success or failure. Entries live in
+``ir_builder.base`` and retain no AST or active caller context. Dynamic caller/definition
+provenance is composed when an entry is used. Ordinary names, Python literals, fixed
+namespace/callee lookup and structural indexing tuples need no entries. Arithmetic and
+buffer loads retain their own expression ranges.
+
+Generated hooks receive entries explicitly: ``X.for_(X.range_(3), names="i", span=_S[i])``
+and ``X.emit_(value, span=_S[j])`` own their result locations. A callable registered with
+``result_span`` promises that its complete IR effect is represented by its returned node
+or emission receipt, allowing result attachment without a construction context. Opaque
+helpers still require ``.ctx`` to locate internal effects. At an ordinary assignment boundary,
+``.ctx(..., attach_result=False)`` preserves construction context while leaving the returned
+value's attribution to ``bind_`` through its separate ``value_span``.
+The parser-owned ``parser.protocol_registry`` module defines registration APIs, metadata and
+persistent registration state. Dialects apply these policies beside concrete definitions,
+constructor creation or necessary exposure sites, without namespace scans or bulk
+registration inventories. Dialect ``parser_protocol`` modules implement construction hooks. The shared
+``ir_builder.ir.parser_protocol`` documents the detailed builder-hook contract and links
+to the registry for special syntax policies. Dialect hooks retain typed signatures and
+document only their specific behavior. For example, dialects import ``args_policy``
+from ``tvm.script.parser.protocol_registry`` to register literal argument policies.
+
+Ordinary assignments pass returned values to ``bind_``. TIRx Vars keep their exact identity,
+producer names and spans, including scope variables, block axes and explicit ``T.bind``
+results. This behavior belongs to the builder's value handling and requires no parser
+producer classification. Self-emitting builders return ``AlreadyEmitted[T]`` so
+source-location handling can annotate the emitted object without emitting it twice.
+Layout, meta_class and view results likewise follow dialect value handling for assignment
+and standalone statements. Unsupported host values still fail under the dialect's emission
+rules. Expression-facing thread APIs return native Vars; IterVar remains native metadata.
+
+Simple comparison chains such as ``0 < i < 10`` lower to
+``X.and_(X.lt_(0, i), X.lt_(i, 10))``. Only names and numeric literals, including signed
+literals, are supported in a chain. Complex chain operands fail with a source-located
+syntax error before they evaluate. Single comparisons retain general operands, and the
+explicit constexpr boundary retains Python comparison semantics.
+
+A bare loop target receives the native entry result: ``for iters in T.grid(m, n)`` binds
+the original variable sequence, while ``for i in T.serial(n)`` binds one Var. Explicit
+tuple, list and starred targets unpack the native frame's stable ``vars`` field.
+
+Function decorators pass their construction namespace and options explicitly to parsing.
+Generated ``X.function_(**options)`` creates the frame; the dialect hook owns option defaults.
+Declaration creates stable references before function bodies; generated code reads the
+native ``global_var`` or ``local_var`` field directly, without a separate reservation pass.
+There is no function-registration record or copied source-function metadata. The root
+retains ``pyfunc`` behavior and original Python callables for module assembly.
+
+Completed modules use generic root coordination over dialect-owned validation hooks.
+Each dialect checks its members and cross-function requirements against the whole module,
+including captured members. Shared GlobalInfo construction and lookup belong to ``tvm.ir``;
+existing Relax class exports and serialized type keys remain compatible.
+
+Each parse rewrites the fresh AST returned by source acquisition directly, without
+copying or reparsing it. Ordinary functions construct arguments directly and bind their
+known body parameters from ``frame.params``. Specialization setup and conditional
+parameter selection are generated only for an explicit specialization request; an empty
+mapping still requests that path, while ``None`` selects ordinary parsing.
+
+Explicit host control flow
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Mark compile-time Python selection with ``I.constexpr`` or the identical dialect alias:
-2. **AST walking**: the ``Parser`` (a subclass of ``doc.NodeVisitor``) walks the Doc AST.
- For each node, it looks up the handler in the dispatch table.
-
-3. **Expression evaluation**: expressions like ``T.grid(128, 128)`` are evaluated by the
- ``ExprEvaluator``, which resolves names against the variable table and the ``T.``/``R.``
- module namespaces.
-
-4. **Value binding**: assignment statements (``A = T.match_buffer(...)`` in TIR,
- ``lv = R.add(x, y)`` in Relax) go through dialect-specific ``bind_*_value()`` functions
- that register the resulting TVM objects in the parser's ``VarTable``.
-
-5. **Scoping**: the ``VarTable`` maintains a stack of frames. Entering a ``with`` block,
- ``for`` loop, or function body pushes a new frame; exiting pops it. This ensures variables
- are scoped correctly.
-
-Variable table
-~~~~~~~~~~~~~~
+.. code-block:: python
-The ``VarTable`` is the parser's symbol table:
+ if I.constexpr(enabled):
+ T.evaluate(1)
+ else:
+ T.evaluate(0)
-.. code-block:: text
+ value = T.ramp(0, 1, lanes) if T.constexpr(lanes > 1) else 0
- VarTable
- ├── frames: [VarTableFrame, ...] ← stack of scopes
- └── name2value: {str: [Any, ...]} ← name → value stack (for shadowing)
+The marker applies to the controlling value. The generated program evaluates it once and
+executes only the selected Python branch in the existing builder scope. Branch bodies and
+result expressions still translate normally. Marked ``and`` and ``or`` decisions retain
+Python short-circuiting; unmarked conditionals and logical expressions call dialect IR
+constructors. A Python boolean alone does not request compile-time branch selection.
-When a name is looked up, the most recent binding wins. When a frame is popped, all bindings
-introduced in that frame are removed.
+JIT supplies validated specialization bindings and optional-argument absence to builder
+execution. Syntax translation does not inspect those values to select a branch.
IR Builder Architecture
@@ -208,7 +367,7 @@ represents and attaches it to the parent frame.
.. code-block:: text
- IRBuilder (thread-local singleton)
+ IRBuilder (current construction context)
└── frame stack:
├── IRModuleFrame ← @I.ir_module
│ ├── PrimFuncFrame ← @T.prim_func
@@ -219,19 +378,19 @@ represents and attaches it to the parent frame.
│ └── BindingBlockFrame ← R.dataflow()
└── ...
-This design means the parser never needs to build a complete IR tree in memory — it
-constructs IR top-down by entering and exiting frames, and each frame handles its own
-finalization.
+The generated program constructs IR by entering and exiting frames. Each frame owns its
+construction policy and finalization; the syntax transpiler does not maintain an IR tree.
TIR builder
~~~~~~~~~~~
-The TIR builder (``ir_builder/tirx/ir.py``) provides functions that map directly to TVMScript
-syntax. Key categories:
+The TIR builder (``python/tvm/tirx/script/builder/``) implements the construction operations
+behind public ``T`` syntax. Builder modules also support direct programmatic construction;
+public decorators and internal frame-opening operations have distinct roles. Key categories:
**Function and block**:
-- ``T.prim_func()`` → ``PrimFuncFrame``
+- ``@T.prim_func`` selects TIR construction; its builder ``function()`` opens ``PrimFuncFrame``
- ``T.sblock(name)`` → ``SBlockFrame`` (spatial block)
- ``T.init()`` → ``BlockInitFrame`` (reduction initialization)
- ``T.reads(...)``, ``T.writes(...)`` → declare buffer access regions
@@ -257,11 +416,11 @@ syntax. Key categories:
Relax builder
~~~~~~~~~~~~~
-The Relax builder (``ir_builder/relax/ir.py``) provides:
+The Relax builder (``python/tvm/relax/script/builder/``) implements:
**Function and dataflow**:
-- ``R.function()`` → ``FunctionFrame``
+- ``@R.function`` selects Relax construction; its builder ``function()`` opens ``FunctionFrame``
- ``R.dataflow()`` → ``BindingBlockFrame``
- ``R.output(*vars)`` → expose variables from a dataflow block
@@ -270,6 +429,9 @@ The Relax builder (``ir_builder/relax/ir.py``) provides:
- ``R.emit(value)`` → emit a binding, returns a ``Var``
- ``R.emit_match_cast(value, ty)`` → emit with type assertion
+Generated standalone-statement emission uses the construction protocol's separate ``emit_``
+operation. Public ``R.emit`` retains its binding-returning imperative API.
+
**Type annotations**:
- ``R.Tensor(shape, dtype)`` — tensor type
@@ -325,17 +487,17 @@ It maintains:
Each IR dialect registers its own converters:
-- ``src/script/printer/tirx/`` — converts PrimFunc, Buffer, SBlock, loops, expressions.
-- ``src/script/printer/relax/`` — converts relax.Function, bindings, types, operators.
+- ``src/tirx/script/printer/`` — converts PrimFunc, Buffer, SBlock, loops, expressions.
+- ``src/relax/script/printer/`` — converts relax.Function, bindings, types, operators.
- ``src/script/printer/ir/`` — converts IRModule, shared types.
The final step calls ``DocToPythonScript()`` (``src/script/printer/doc_printer/python_doc_printer.cc``)
to format the Doc tree into properly indented Python text.
-Roundtrip guarantee
-~~~~~~~~~~~~~~~~~~~
+Roundtrip
+~~~~~~~~~
-For any ``IRModule`` constructed through the compiler:
+For printable IR with the required source context available:
.. code-block:: python
@@ -344,9 +506,9 @@ For any ``IRModule`` constructed through the compiler:
tvm.ir.assert_structural_equal(mod, reparsed)
This roundtrip property is relied upon by testing infrastructure and serialization workflows.
-Note that the printed text may differ from hand-written TVMScript — the printer uses canonical
-forms (e.g., explicit ``R.emit`` calls, fully qualified buffer annotations) that are not required
-in hand-written code.
+Printed text uses canonical forms and may differ from hand-written TVMScript. External
+objects or Python attachments require their corresponding context; text alone does not
+serialize arbitrary Python state.
Supported Python Syntax
@@ -364,7 +526,7 @@ and how each construct is interpreted:
- Relax
* - ``for i in range(n)``
- Serial loop nest
- - Not supported (no Relax-level ``for`` handler)
+ - No Relax IR loop construction
* - ``with T.sblock(...)``
- Spatial block scope
- N/A
@@ -372,8 +534,8 @@ and how each construct is interpreted:
- N/A
- Dataflow block
* - ``if ... else``
- - TIR ``If`` branch (PrimExpr condition) or static eval (Python bool)
- - Relax ``If`` node (plain Python ``if cond:`` syntax)
+ - TIR ``IfThenElse``; explicit ``I.constexpr`` selects a Python branch
+ - Relax ``If``; explicit ``I.constexpr`` selects a Python branch
* - ``while``
- ``T.While`` loop
- Not supported
@@ -385,9 +547,9 @@ and how each construct is interpreted:
- N/A
* - ``x: R.Tensor(...)``
- N/A
- - Struct info annotation
+ - Type annotation
* - ``return``
- - Not used
+ - Primitive return expression
- Function return value
* - ``A[i, j]``
- Buffer load
@@ -402,9 +564,9 @@ and how each construct is interpreted:
- ``T.*`` intrinsics
- ``R.*`` operators or ``call_tir`` / ``call_packed``
-**Not supported**: ``class`` definitions (except for ``@I.ir_module``), ``try/except``,
-``yield``, ``async/await``, list comprehensions, ``lambda``, ``import``, and ``global``
-statements.
+Inside scripted function bodies, unsupported statement forms include ``class``,
+``try/except``, ``import``, and ``global``. ``yield`` and ``async/await`` are also unsupported.
+Place imports and ordinary Python class definitions in the surrounding Python program.
TIR Syntax Reference
@@ -417,6 +579,8 @@ Function definition
@T.prim_func
def func_name(a: T.handle, b: T.handle):
+ m = T.int32()
+ n = T.int32()
A = T.match_buffer(a, (m, n), "float32")
B = T.match_buffer(b, (m,), "float32")
# function body
@@ -428,6 +592,8 @@ Function definition
Block and axes
~~~~~~~~~~~~~~
+Use ``@T.prim_func(s_tir=True)`` for functions with scheduled-TIR blocks and block axes.
+
.. code-block:: python
for i, j in T.grid(128, 128):
@@ -472,7 +638,7 @@ Common intrinsics
T.cast(x, "float16") # type cast
T.if_then_else(cond, true_val, false_val) # conditional expression
T.min(a, b), T.max(a, b) # min/max
- T.call_extern("func_name", *args) # external function call
+ T.call_extern("int32", "func_name", *args) # external function call
T.call_packed("func_name", *args) # packed function call
T.tvm_storage_sync("shared") # GPU memory fence
@@ -514,9 +680,9 @@ Calling TIR functions
.. code-block:: python
- out = R.call_tir(cls.my_kernel, (x, y), out_ty=R.Tensor((128,), "float32"))
+ out = R.call_tir(MyModule.my_kernel, (x, y), out_ty=R.Tensor((128,), "float32"))
-- ``cls.my_kernel`` — references a TIR ``PrimFunc`` in the same module.
+- ``MyModule.my_kernel`` — references a TIR ``PrimFunc`` in the same module.
- ``out_ty`` — the type (shape and dtype) of the output tensor.
Control flow
@@ -534,42 +700,3 @@ boolean type. Both branches are required.
else:
result = R.multiply(x, x)
return result
-
-
-Source Code Map
----------------
-
-.. list-table::
- :header-rows: 1
- :widths: 50 50
-
- * - Path
- - Contents
- * - ``python/tvm/script/parser/core/``
- - Core parser: dispatch, expression evaluator, variable table, Doc AST
- * - ``python/tvm/script/parser/tirx/``
- - TIR-specific parser handlers and value binding
- * - ``python/tvm/script/parser/relax/``
- - Relax-specific parser handlers and value binding
- * - ``python/tvm/script/parser/ir/``
- - ``@I.ir_module`` entry point and module-level parsing
- * - ``python/tvm/script/ir_builder/base.py``
- - IRBuilder base class and frame stack mechanism
- * - ``python/tvm/script/ir_builder/tirx/``
- - TIR frame types and builder functions (``T.*``)
- * - ``python/tvm/script/ir_builder/relax/``
- - Relax frame types and builder functions (``R.*``)
- * - ``python/tvm/script/ir_builder/ir/``
- - IRModule builder (``I.*``)
- * - ``src/script/printer/``
- - C++ printer: Doc tree, IRDocsifier, Python code generation
- * - ``src/script/printer/tirx/``
- - TIR-specific IR-to-Doc converters
- * - ``src/script/printer/relax/``
- - Relax-specific IR-to-Doc converters
- * - ``src/script/ir_builder/``
- - C++ backend for frame stack and IR construction
- * - ``include/tvm/script/printer/``
- - C++ headers: Doc classes, IRDocsifier, dispatch functor
- * - ``include/tvm/script/ir_builder/``
- - C++ headers: builder base, dialect-specific frame types
diff --git a/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py b/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py
index 9515f6874d..bf42154bfb 100644
--- a/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py
+++ b/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py
@@ -71,7 +71,7 @@ RUN_EXAMPLE = HAS_TORCH and not IS_IN_CI
######################################################################
# Step 1: Your First Hybrid Module
# ----------------------------------
-# The core idea: decorate a class with ``@I.ir_module``, inherit from ``BasePyModule``, and use
+# The core idea: decorate a class with ``@R.py_module``, inherit from ``BasePyModule``, and use
# three decorators for three kinds of functions:
#
# - ``@T.prim_func`` — low-level TIR kernel (JIT-compiled on instantiation)
@@ -83,7 +83,7 @@ RUN_EXAMPLE = HAS_TORCH and not IS_IN_CI
if RUN_EXAMPLE:
- @I.ir_module
+ @R.py_module
class MyFirstModule(BasePyModule):
@T.prim_func(s_tir=True)
def add_tir(
@@ -129,7 +129,7 @@ if RUN_EXAMPLE:
if RUN_EXAMPLE:
- @I.ir_module
+ @R.py_module
class DebugModule(BasePyModule):
@T.prim_func(s_tir=True)
def matmul_tir(var_A: T.handle, var_B: T.handle, var_C: T.handle):
@@ -207,7 +207,7 @@ if RUN_EXAMPLE:
out_np = x_np + b_np
out[:] = out_np
- @I.ir_module
+ @R.py_module
class PipelineModule(BasePyModule):
@T.prim_func(s_tir=True)
def matmul_tir(var_A: T.handle, var_B: T.handle, var_C: T.handle):
@@ -304,7 +304,7 @@ if RUN_EXAMPLE:
w = torch.randn(4, 4)
b = torch.randn(4)
- py_result_early = converted_early.pyfuncs["main"](x, w, b)
+ py_result_early = converted_early.__pyfuncs__["main"](x, w, b)
expected = F.relu(x @ w + b)
print("Before optimization:")
@@ -320,7 +320,7 @@ if RUN_EXAMPLE:
converter_late = RelaxToPyFuncConverter(optimized_mod)
converted_late = converter_late.convert(["main"])
- py_result_late = converted_late.pyfuncs["main"](x, w, b)
+ py_result_late = converted_late.__pyfuncs__["main"](x, w, b)
print("\nAfter CanonicalizeBindings pass:")
print(" Converted result:", py_result_late)
@@ -345,7 +345,7 @@ if RUN_EXAMPLE:
if RUN_EXAMPLE:
- @I.ir_module
+ @R.py_module
class HybridVMModule(BasePyModule):
@I.pyfunc
def silu(self, x):
@@ -399,7 +399,7 @@ if RUN_EXAMPLE:
if RUN_EXAMPLE:
- @I.ir_module
+ @R.py_module
class DynamicModule(BasePyModule):
@T.prim_func(s_tir=True)
def scale_tir(var_x: T.handle, var_out: T.handle):
diff --git a/docs/reference/api/python/relax/relax.rst b/docs/reference/api/python/relax/relax.rst
index 28cdc5e32b..7ae3359987 100644
--- a/docs/reference/api/python/relax/relax.rst
+++ b/docs/reference/api/python/relax/relax.rst
@@ -17,6 +17,7 @@
tvm.relax
---------
+
.. automodule:: tvm.relax
:members:
:imported-members:
diff --git a/docs/reference/api/python/script/ir_builder.rst b/docs/reference/api/python/script/ir_builder.rst
index 201090a0e4..150d4a43de 100644
--- a/docs/reference/api/python/script/ir_builder.rst
+++ b/docs/reference/api/python/script/ir_builder.rst
@@ -23,17 +23,23 @@ tvm.script.ir_builder
.. automodule:: tvm.script.ir_builder
:members:
:imported-members:
+ :exclude-members: Any
tvm.script.ir_builder.ir
************************
+The :class:`~tvm.script.ir_builder.IRBuilder` class is also available as
+``tvm.script.ir_builder.ir.IRBuilder``.
+
.. automodule:: tvm.script.ir_builder.ir
:members:
:imported-members:
+ :exclude-members: GenericConst, IRBuilder, Range, StringImm, StringType
tvm.script.ir_builder.relax
***************************
.. automodule:: tvm.script.ir_builder.relax
:members:
+ :exclude-members: ExternFunc, ShapeExpr, TupleGetItem, Range
tvm.script.ir_builder.relax.distributed
***************************************
diff --git a/docs/reference/api/python/script/parser.rst b/docs/reference/api/python/script/parser.rst
index 583b27b276..716738c747 100644
--- a/docs/reference/api/python/script/parser.rst
+++ b/docs/reference/api/python/script/parser.rst
@@ -24,6 +24,9 @@ tvm.script.parser
:members:
:imported-members:
+The language variant aliases below share the public construction namespaces documented
+in :doc:`script`. Parser entry points above use the canonical frontend.
+
tvm.script.parser.ir
********************
.. automodule:: tvm.script.parser.ir
diff --git a/docs/reference/api/python/script/script.rst b/docs/reference/api/python/script/script.rst
index 5e819bd1bb..096b4413e8 100644
--- a/docs/reference/api/python/script/script.rst
+++ b/docs/reference/api/python/script/script.rst
@@ -23,13 +23,28 @@ tvm.script
.. automodule:: tvm.script
:members:
:imported-members:
+ :exclude-members: Any
tvm.script.relax
****************
.. automodule:: tvm.script.relax
:members:
+ :exclude-members: ExternFunc, ShapeExpr, TupleGetItem, Range, function, macro
+
+.. autofunction:: tvm.script.relax.function
+
+.. autofunction:: tvm.script.relax.macro
tvm.script.tirx
***************
.. automodule:: tvm.script.tirx
:members:
+ :exclude-members: Range, meta_var, Var, Call, CommReducer, Reduce, SMEMPool, TMEMPool, FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod, LShift, RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min, Max, EQ, NE, LT, LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle, CallEffectKind, IterVar, prim_func, jit, inline, macro
+
+.. autofunction:: tvm.script.tirx.prim_func
+
+.. autofunction:: tvm.script.tirx.jit
+
+.. autofunction:: tvm.script.tirx.inline
+
+.. autofunction:: tvm.script.tirx.macro
diff --git a/docs/tirx/api/script.rst b/docs/tirx/api/script.rst
index f30d508095..8d905e11f5 100644
--- a/docs/tirx/api/script.rst
+++ b/docs/tirx/api/script.rst
@@ -30,10 +30,10 @@ authoring layers and TIRx IR, see :ref:`tirx-programming-model`.
Parser entry points
-------------------
-.. automodule:: tvm.tirx.script.parser.entry
- :members:
+.. automodule:: tvm.script.tirx
+ :members: prim_func, jit, inline, macro
+ :undoc-members:
:no-index:
- :exclude-members: TIRInline
Core IR builder
---------------
diff --git a/docs/tirx/tile_primitives/copy_async/dsmem.rst b/docs/tirx/tile_primitives/copy_async/dsmem.rst
index 0a59fcffe9..6681573e51 100644
--- a/docs/tirx/tile_primitives/copy_async/dsmem.rst
+++ b/docs/tirx/tile_primitives/copy_async/dsmem.rst
@@ -141,12 +141,12 @@ chunk's offsets each step:
.. code-block:: python
- if not outer_extents: # one contiguous chunk
+ if Tx.constexpr(not outer_extents): # one contiguous chunk
Tx.ptx["cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"](
Tx.cast(mapped[1], "uint32"), src_buf.ptr_to(src_st),
Tx.cast(chunk_bytes, "uint32"), Tx.cast(mapped[0], "uint32"))
else:
- for loop_vars in Tx.grid(*outer_extents): # one chunk per outer coord
+ for loop_vars in T.grid(*outer_extents): # one chunk per outer coord
... # re-decl src/dst views at the per-chunk offset
Tx.ptx["cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"](
Tx.cast(mapped[1], "uint32"), src_ptr,
diff --git a/include/tvm/ir/global_info.h b/include/tvm/ir/global_info.h
index 46d7cce46c..a4ee836ff4 100644
--- a/include/tvm/ir/global_info.h
+++ b/include/tvm/ir/global_info.h
@@ -25,6 +25,7 @@
#define TVM_IR_GLOBAL_INFO_H_
#include <tvm/ffi/object.h>
+#include <tvm/ffi/reflection/registry.h>
namespace tvm {
diff --git a/include/tvm/relax/global_info.h b/include/tvm/relax/global_info.h
index 4cb49064b8..e08ffe4f45 100644
--- a/include/tvm/relax/global_info.h
+++ b/include/tvm/relax/global_info.h
@@ -24,7 +24,6 @@
#ifndef TVM_RELAX_GLOBAL_INFO_H_
#define TVM_RELAX_GLOBAL_INFO_H_
-#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/global_info.h>
#include <tvm/target/target.h>
diff --git a/include/tvm/relax/script/builder/frame.h b/include/tvm/relax/script/builder/frame.h
index a2f663fb50..2d3cd6cf8f 100644
--- a/include/tvm/relax/script/builder/frame.h
+++ b/include/tvm/relax/script/builder/frame.h
@@ -97,6 +97,8 @@ class FunctionFrameNode : public SeqExprFrameNode {
* However, we must specify the name by `R.func_name` before exit this frame.
*/
ffi::Optional<ffi::String> name;
+ /*! \brief Function-local symbols retained across declaration and body entry. */
+ ffi::Map<ffi::String, tvm::Var> type_var_map;
/*! \brief The function params. */
ffi::Array<tvm::Var> params;
/*!
@@ -117,15 +119,26 @@ class FunctionFrameNode : public SeqExprFrameNode {
ffi::Map<ffi::String, Any> attrs;
/*! \brief The block builder to create Relax function. */
tvm::relax::BlockBuilder block_builder;
+ /*! \brief Whether this frame constructs only a function signature. */
+ bool declaration = false;
+ bool local = false;
+ ffi::Optional<tvm::Var> local_var;
+ /*! \brief Finalized function and its stable module reference. */
+ ffi::Optional<tvm::relax::Function> function;
+ ffi::Optional<tvm::GlobalVar> global_var;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FunctionFrameNode>()
+ .def_ro("type_var_map", &FunctionFrameNode::type_var_map)
.def_ro("name", &FunctionFrameNode::name)
.def_ro("params", &FunctionFrameNode::params)
.def_ro("ret_ty", &FunctionFrameNode::ret_ty)
.def_ro("is_pure", &FunctionFrameNode::is_pure)
- .def_ro("attrs", &FunctionFrameNode::attrs);
+ .def_ro("attrs", &FunctionFrameNode::attrs)
+ .def_ro("function", &FunctionFrameNode::function)
+ .def_ro("global_var", &FunctionFrameNode::global_var)
+ .def_ro("local_var", &FunctionFrameNode::local_var);
// `binding_blocks` and `output` are inherited from SeqExprFrameNode.
// `block_builder` is not registered as it's not visited.
}
@@ -163,6 +176,8 @@ class BindingBlockFrameNode : public RelaxFrameNode {
* \note Only used for a dataflow block.
*/
ffi::Array<tvm::Var> output_vars;
+ /*! \brief Statement ranges for explicitly emitted bindings. */
+ ffi::Map<tvm::Var, Span> binding_spans;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
diff --git a/include/tvm/relax/script/builder/ir.h b/include/tvm/relax/script/builder/ir.h
index 2516dc0134..be008474c2 100644
--- a/include/tvm/relax/script/builder/ir.h
+++ b/include/tvm/relax/script/builder/ir.h
@@ -29,6 +29,9 @@ namespace script {
namespace ir_builder {
namespace relax {
+/*! \brief Find a module virtual device by target kind and index. */
+TVM_DLL tvm::relax::VDevice LookupVDevice(ffi::String target_kind, int device_index);
+
/////////////////////////////// Function ////////////////////////////////
/*!
@@ -39,6 +42,15 @@ namespace relax {
*/
TVM_DLL FunctionFrame Function(bool is_pure, bool is_private);
+/*! \brief Start a bodyless declaration using the normal signature operations. */
+TVM_DLL FunctionFrame DeclFunction(bool is_pure, bool is_private, bool local);
+
+/*! \brief Define a local function under its cached declaration identity. */
+TVM_DLL FunctionFrame LocalFunction(bool is_pure, const tvm::Var& reference);
+
+/*! \brief Add a cached parameter without changing its identity. */
+TVM_DLL tvm::Var ArgVar(const ffi::String& name, const tvm::Var& var);
+
/*!
* \brief Add a parameter to the last function frame.
* \param name The name of the parameter.
@@ -117,6 +129,17 @@ TVM_DLL tvm::Var EmitMatchCast(const tvm::relax::Expr& value, const tvm::Type& t
*/
TVM_DLL tvm::Var EmitVarBinding(const tvm::relax::VarBinding& binding);
+/*! \brief Emit a binding with separate statement and variable-name ranges. */
+TVM_DLL tvm::Var EmitWithSpan(const tvm::relax::Expr& value,
+ const ffi::Optional<tvm::Type>& annotate_ty,
+ const ffi::Optional<Span>& name_span,
+ const ffi::Optional<Span>& span = std::nullopt);
+
+/*! \brief Emit a match cast with separate statement and variable-name ranges. */
+TVM_DLL tvm::Var EmitMatchCastWithSpan(const tvm::relax::Expr& value, const tvm::Type& ty,
+ const ffi::Optional<Span>& name_span,
+ const ffi::Optional<Span>& span = std::nullopt);
+
///////////////////////////// If Then Else /////////////////////////////
/*!
diff --git a/include/tvm/script/ir_builder/base.h b/include/tvm/script/ir_builder/base.h
index 417e9b0c38..ff9c798618 100644
--- a/include/tvm/script/ir_builder/base.h
+++ b/include/tvm/script/ir_builder/base.h
@@ -37,39 +37,21 @@ namespace ir_builder {
* Furthermore, the information stored in each stack frame can be useful for context-dependent
* IR construction.
*
- * \example
- *
- * The `T::MatchBuffer` below annotates a PrimFunc parameter with BufferType:
- *
- * \code {.cpp}
- *
- * using T = tvm::script::ir_builder::tirx;
- * With <PrimFuncFrame> _(...);
- * Buffer buffer = T::MatchBuffer(...);
- *
- * \endcode
- *
- * The `T::MatchBuffer` below instead generates `s_tir::MatchBufferRegion` in a TIR block:
- *
- * \code {.cpp}
- *
- * using T = tvm::script::ir_builder::tirx;
- * With <PrimFuncFrame> _(...);
- * {
- * With<SBlockFrame> _2(...);
- * Buffer buffer = T::MatchBuffer(...);
- * }
- *
- * \endcode
+ * Frames capture construction context and nested scopes. Language variant builders define
+ * concrete frame subclasses and decide how their contents contribute to the result.
*/
class IRBuilderFrameNode : public ffi::Object {
public:
+ /*! \brief Capture the source-call context when this frame is constructed. */
+ IRBuilderFrameNode();
/*! \brief A list of callbacks used when exiting the frame. */
std::vector<ffi::TypedFunction<void()>> callbacks;
+ /*! \brief Source context retained until this frame constructs its result. */
+ mutable Span source_span;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<IRBuilderFrameNode>();
+ refl::ObjectDef<IRBuilderFrameNode>().def_ro("source_span", &IRBuilderFrameNode::source_span);
// `callbacks` is not registered as it's not visited.
}
@@ -129,28 +111,29 @@ class IRBuilderFrame : public ffi::ObjectRef {
}
};
+/*! \brief Resolve one symbol in a native function frame's retained map. */
+TVM_DLL tvm::Var ResolveTypeVar(ffi::Map<ffi::String, tvm::Var>* symbols, const ffi::String& name,
+ ffi::Optional<PrimType> dtype, ffi::Optional<tvm::Var> value,
+ Span span);
+
////////////////////////////// IRBuilder //////////////////////////////
/*!
- * \brief A dialect-agnostic IRBuilder that constructs any IR of TVM.
+ * \brief A language variant agnostic IRBuilder that constructs any IR of TVM.
* An idiomatic use of this class is to put this inside the RAII with-scope,
- * call dialect-specific methods accordingly. Upon exiting the scope.
+ * call language variant specific methods accordingly. Upon exiting the scope.
*
* \code
*
- * PrimFunc ConstructPrimFunc() {
- * using tvm::script::ir_builder::IRBuilder;
- * using T = tvm::script::ir_builder::tirx;
+ * IRModule ConstructModule() {
+ * using namespace tvm::script::ir_builder;
* IRBuilder builder;
- * // Step 1. Place IRBuilder inside the with-scope.
* {
- * With<IRBuilder> _(builder);
- * // Step 2. Call dialect-specific methods.
- * With<T::PrimFuncFrame> _2(...);
- * T::MatchBuffer(...);
+ * With<IRBuilder> scope(builder);
+ * With<ir::IRModuleFrame> module(ir::IRModule());
+ * // Add function declarations and definitions with the chosen language variant builder.
* }
- * // Step 3. Return the constructed PrimFunc.
- * return builder->Get<PrimFunc>();
+ * return builder->Get<IRModule>();
* }
*
* \endcode
@@ -183,7 +166,7 @@ class IRBuilderNode : public ffi::Object {
template <typename TFrame>
inline ffi::Optional<TFrame> FindFrame() const;
/*!
- * \brief Get the frame on top of the stack `this->frames` if its type is `TFrame`.
+ * \brief Get the top frame if its type is `TFrame`.
* \tparam TFrame The assumed type of the last frame on stack.
* \return The frame if the stack is non-empty and the top of the stack is of type `TFrame`.
* Otherwise std::nullopt.
@@ -202,9 +185,11 @@ class IRBuilderNode : public ffi::Object {
/*! \brief Pop the innermost frontend source span. */
void PopSourceSpan();
/*! \brief Return the normalized active source span, including expansion history. */
- Span GetCurrentSourceSpan() const;
- /*! \brief Attach the active source span to an expression that has no span yet. */
+ Span GetCurrentSourceSpan(Span location = Span()) const;
+ /*! \brief Compose active source context onto a supported node or construction frame. */
ffi::ObjectRef SetCurrentSourceSpan(ffi::ObjectRef obj) const;
+ /*! \brief Attach an explicit location composed with the active source-call context. */
+ ffi::ObjectRef SetSourceSpan(ffi::ObjectRef obj, Span span) const;
};
/*!
@@ -267,6 +252,13 @@ class IRBuilder : public ffi::ObjectRef {
namespace details {
+/*! \brief Language variant extensible access to an object's mutable source span. */
+class SourceSpanAccessor {
+ public:
+ using FType = ObjectFunctor<Span*(const ffi::ObjectRef&)>;
+ static FType& vtable();
+};
+
class Namer {
public:
using FType = ObjectFunctor<void(const ffi::ObjectRef&, ffi::String)>;
@@ -297,7 +289,7 @@ template <typename TFrame>
inline ffi::Optional<TFrame> IRBuilderNode::GetLastFrame() const {
using TFrameNode = typename TFrame::ContainerType;
if (!frames.empty() && frames.back()->IsInstance<TFrameNode>()) {
- return frames.back().as_or_throw<TFrame>();
+ return frames.back().template as_or_throw<TFrame>();
}
return std::nullopt;
}
diff --git a/include/tvm/tirx/script/builder/frame.h b/include/tvm/tirx/script/builder/frame.h
index 0427b9def3..c28613e97a 100644
--- a/include/tvm/tirx/script/builder/frame.h
+++ b/include/tvm/tirx/script/builder/frame.h
@@ -70,6 +70,8 @@ class TIRFrame : public IRBuilderFrame {
*/
class PrimFuncFrameNode : public TIRFrameNode {
public:
+ /*! \brief Function-local symbols retained across declaration and body entry. */
+ ffi::Map<ffi::String, tvm::Var> type_var_map;
/*! \brief The name of the block. */
ffi::Optional<ffi::String> name;
/*! \brief Function parameters. */
@@ -93,10 +95,16 @@ class PrimFuncFrameNode : public TIRFrameNode {
bool s_tir;
/*! \brief Whether it is a persistent kernel. */
bool persistent;
+ /*! \brief Whether this frame declares a bodyless signature. */
+ bool is_declaration{false};
+ /*! \brief Finalized function and its module identity. */
+ ffi::Optional<tvm::tirx::PrimFunc> function;
+ ffi::Optional<GlobalVar> global_var;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PrimFuncFrameNode>()
+ .def_ro("type_var_map", &PrimFuncFrameNode::type_var_map)
.def_ro("name", &PrimFuncFrameNode::name)
.def_ro("args", &PrimFuncFrameNode::args)
.def_ro("is_private", &PrimFuncFrameNode::is_private)
@@ -106,7 +114,10 @@ class PrimFuncFrameNode : public TIRFrameNode {
.def_ro("env_threads", &PrimFuncFrameNode::env_threads)
.def_ro("root_alloc_buffers", &PrimFuncFrameNode::root_alloc_buffers)
.def_ro("s_tir", &PrimFuncFrameNode::s_tir)
- .def_ro("persistent", &PrimFuncFrameNode::persistent);
+ .def_ro("persistent", &PrimFuncFrameNode::persistent)
+ .def_ro("is_declaration", &PrimFuncFrameNode::is_declaration)
+ .def_ro("function", &PrimFuncFrameNode::function)
+ .def_ro("global_var", &PrimFuncFrameNode::global_var);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.PrimFuncFrame", PrimFuncFrameNode,
TIRFrameNode);
@@ -264,7 +275,7 @@ class ForFrameNode : public TIRFrameNode {
*/
using FMakeForLoop = ffi::TypedFunction<tvm::tirx::Stmt(
ffi::Array<tvm::tirx::Var> loop_vars, ffi::Array<Range> loop_extents,
- ffi::Array<ffi::Optional<PrimExpr>> loop_steps, tvm::tirx::Stmt loop_body)>;
+ ffi::Array<ffi::Optional<PrimExpr>> loop_steps, tvm::tirx::Stmt loop_body, Span span)>;
/*! \brief The loop variable. */
ffi::Array<tvm::tirx::Var> vars;
/*! \brief The domains of iteration. */
@@ -284,10 +295,9 @@ class ForFrameNode : public TIRFrameNode {
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ForFrame", ForFrameNode, TIRFrameNode);
public:
- /*!
- * \brief The method called when exiting RAII scope.
- * \sa tvm::support::With
- */
+ /*! \brief Apply source target names before entry, preserving variable identity. */
+ void SetNames(ffi::Optional<ffi::Variant<ffi::String, ffi::Array<ffi::String>>> names);
+ /*! \brief Construct the loop nest with this frame's stored source location. */
void ExitWithScope() final;
};
@@ -639,33 +649,6 @@ class DeclBufferFrame : public TIRFrame {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DeclBufferFrame, TIRFrame, DeclBufferFrameNode);
};
-class AllocBufferFrameNode : public TIRFrameNode {
- public:
- /*! \brief The allocated buffer. */
- tvm::tirx::BufferVar buffer;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AllocBufferFrameNode>().def_ro("buffer", &AllocBufferFrameNode::buffer);
- }
-
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.AllocBufferFrame", AllocBufferFrameNode,
- TIRFrameNode);
-
- public:
- void ExitWithScope() final;
-};
-
-class AllocBufferFrame : public TIRFrame {
- public:
- explicit AllocBufferFrame(ffi::ObjectPtr<AllocBufferFrameNode> data)
- : TIRFrame(ffi::UnsafeInit{}) {
- TVM_FFI_ICHECK(data != nullptr);
- data_ = std::move(data);
- }
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AllocBufferFrame, TIRFrame, AllocBufferFrameNode);
-};
-
/*!
* \brief A frame that represents a hint directive for the sketch language.
*
diff --git a/include/tvm/tirx/script/builder/ir.h b/include/tvm/tirx/script/builder/ir.h
index 537c23fb2b..7e3e7283a9 100644
--- a/include/tvm/tirx/script/builder/ir.h
+++ b/include/tvm/tirx/script/builder/ir.h
@@ -68,6 +68,9 @@ BufferVar BufferDecl(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String buf
*/
PrimFuncFrame PrimFunc(bool is_private, bool s_tir = false, bool persistent = false);
+/*! \brief Construct a bodyless function signature using the ordinary signature operations. */
+PrimFuncFrame DeclFunction(bool is_private = false, bool s_tir = false, bool persistent = false);
+
/*!
* \brief The PrimFunc variable arguments adding function.
* \param name The name of the variable.
@@ -114,13 +117,16 @@ Type FuncRet(Type ret_type);
* \param storage_scope The optional storage scope of buffer data pointer.
* \param align The alignment requirement of data pointer in bytes.
* \param offset_factor The factor of elem_offset field.
+ * \param layout The buffer layout.
+ * \param allocated_addr Addresses assigned to the buffer allocation.
* \return The matched buffer.
*/
BufferVar MatchBuffer(ffi::ObjectRef param, ffi::Array<PrimExpr> shape,
PrimType dtype = PrimType::Float(32), ffi::Optional<Expr> data = std::nullopt,
ffi::Array<PrimExpr> strides = {}, PrimExpr elem_offset = PrimExpr(),
ffi::String storage_scope = "global", int align = -1, int offset_factor = 0,
- ffi::Optional<Layout> layout = std::nullopt);
+ ffi::Optional<Layout> layout = std::nullopt,
+ ffi::Array<PrimExpr> allocated_addr = {});
/*!
* \brief The block declaration statement.
@@ -206,15 +212,14 @@ void BlockAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
* \param offset_factor The factor of elem_offset field.
* \param layout The layout of the buffer.
* \param allocated_addr The allocated address of the buffer. Might be multi-dimensional.
- * \return The allocated buffer or the AllocBufferFrame if the function is called under
- * T.prim_func(tirx=True).
+ * \return The buffer attached to its enclosing block or function allocation list.
*/
-ffi::Variant<BufferVar, AllocBufferFrame> SBlockAllocBuffer(
- ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
- ffi::Optional<Expr> data = std::nullopt, ffi::Array<PrimExpr> strides = {},
- PrimExpr elem_offset = PrimExpr(), ffi::String storage_scope = "", int align = -1,
- int offset_factor = 0, ffi::Optional<Layout> layout = std::nullopt,
- ffi::Array<PrimExpr> allocated_addr = {});
+BufferVar SBlockAllocBuffer(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
+ ffi::Optional<Expr> data = std::nullopt,
+ ffi::Array<PrimExpr> strides = {}, PrimExpr elem_offset = PrimExpr(),
+ ffi::String storage_scope = "", int align = -1, int offset_factor = 0,
+ ffi::Optional<Layout> layout = std::nullopt,
+ ffi::Array<PrimExpr> allocated_addr = {});
namespace axis {
@@ -394,8 +399,9 @@ WhileFrame While(PrimExpr condition);
/*!
* \brief Create a return statement.
* \param value The value to return.
+ * \return The same statement that was added to the parent frame.
*/
-void Return(Expr value);
+tvm::tirx::Stmt Return(Expr value);
/*!
* \brief Create a break statement.
@@ -488,14 +494,16 @@ Var EnvThread(ffi::String thread_tag, PrimType dtype = PrimType::Int(32));
* \param buffer The buffer.
* \param value The value to be stored.
* \param indices The indices location to be stored.
+ * \return The same statement that was added to the parent frame.
*/
-void BufferStore(BufferVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices);
+tvm::tirx::Stmt BufferStore(BufferVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices);
/*!
* \brief Evaluate the input expression.
* \param value The input expression to evaluate.
+ * \return The same statement that was added to the parent frame.
*/
-void Evaluate(Expr value);
+tvm::tirx::Stmt Evaluate(Expr value);
/*!
* \brief Create a TIR var that represents a pointer
diff --git a/python/tvm/backend/cuda/lang/clc.py b/python/tvm/backend/cuda/lang/clc.py
index bc79b85ec1..65577963a7 100644
--- a/python/tvm/backend/cuda/lang/clc.py
+++ b/python/tvm/backend/cuda/lang/clc.py
@@ -41,7 +41,7 @@ def query_cancel_first_ctaid_x(first_ctaid_x, handle, *, use_ld_acquire=True):
T.ptx[f"ld{'.acquire.cta' if use_ld_acquire else ''}.shared.b128"](response, handle)
T.ptx.clusterlaunchcontrol.query_cancel.is_canceled.pred.b128(canceled, response)
- first_ctaid_x = T.uint32(0xFFFFFFFF)
+ T.buffer_store(first_ctaid_x.source, T.uint32(0xFFFFFFFF), first_ctaid_x.indices)
T.ptx.clusterlaunchcontrol.query_cancel.get_first_ctaid__x.b32.b128(
first_ctaid_x, response, pred=canceled
)
diff --git a/python/tvm/backend/cuda/lang/warp_role.py b/python/tvm/backend/cuda/lang/warp_role.py
index b61cd8822d..51974e2ee9 100644
--- a/python/tvm/backend/cuda/lang/warp_role.py
+++ b/python/tvm/backend/cuda/lang/warp_role.py
@@ -132,7 +132,7 @@ class WarpgroupRole:
def __enter__(self):
if isinstance(self.wg_id_val, tuple):
start, stop = self.wg_id_val
- self._if_frame = T.If(start <= self.wg_id_var and self.wg_id_var < stop)
+ self._if_frame = T.If(T.And(T.LE(start, self.wg_id_var), self.wg_id_var < stop))
else:
self._if_frame = T.If(self.wg_id_var == self.wg_id_val)
self._if_frame.__enter__()
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
index 35a0d2570d..468e64734b 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
@@ -134,13 +134,9 @@ def copy_dsmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFu
# Helper to compute element offsets from loop variables (called via T.meta_var)
def compute_offsets(loop_vars):
- if len(outer_extents) == 1:
- lvs = [loop_vars]
- else:
- lvs = list(loop_vars)
src_off = 0
dst_off = 0
- for j, v in enumerate(lvs):
+ for j, v in enumerate(loop_vars):
src_off = src_off + v * outer_src_strides[j]
dst_off = dst_off + v * outer_dst_strides[j]
return src_off, dst_off
@@ -158,7 +154,7 @@ def copy_dsmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFu
T.ptx.mapa.u64(mapped[0], mbar, T.uint32(remote_cta_id))
remote_mbar = mapped[0]
- if not outer_extents:
+ if T.constexpr(not outer_extents):
# Single contiguous chunk — no iteration needed
src_ptr = src_buf.ptr_to(src_st)
T.ptx.mapa.u64(mapped[1], dst_buf.ptr_to(dst_st), T.uint32(remote_cta_id))
diff --git a/python/tvm/backend/trn/script.py b/python/tvm/backend/trn/script.py
index 5d507f462b..4f9d9e436d 100644
--- a/python/tvm/backend/trn/script.py
+++ b/python/tvm/backend/trn/script.py
@@ -21,22 +21,26 @@ from __future__ import annotations
from collections.abc import Callable
from typing import Any
+from tvm.script.parser.protocol_registry import result_span as _result_span
+
from . import op as _trn_op
OpWrapper = Callable[[Callable[..., Any]], Callable[..., Any]]
-def _default_op_wrapper() -> OpWrapper:
+def _default_op_wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
+ """Create the default facade for a returned-node NKI intrinsic."""
from tvm.tirx.script.builder.ir import _op_wrapper # pylint: disable=import-outside-toplevel
- return _op_wrapper
+ # NKI producers return a single call_intrin node and never emit separately.
+ return _result_span(f"T.nki.{func.__name__.removeprefix('nki_')}")(_op_wrapper(func))
class NKINamespace:
"""The NKI instructions submodule."""
def __init__(self, op_wrapper: OpWrapper | None = None):
- wrap = op_wrapper or _default_op_wrapper()
+ wrap = _default_op_wrapper if op_wrapper is None else op_wrapper
self.load = wrap(_trn_op.nki_load)
self.store = wrap(_trn_op.nki_store)
self.tensor_copy = wrap(_trn_op.nki_tensor_copy)
diff --git a/python/tvm/backend/trn/tile_primitive/binary/default.py b/python/tvm/backend/trn/tile_primitive/binary/default.py
index 85f19d11c5..9cf669d892 100644
--- a/python/tvm/backend/trn/tile_primitive/binary/default.py
+++ b/python/tvm/backend/trn/tile_primitive/binary/default.py
@@ -84,7 +84,7 @@ def binary_trn(
if inst_gen.make_guard(_dst):
dst_indices = T.meta_var(inst_gen.generate_indices(_dst))
src1_indices = T.meta_var(inst_gen.generate_indices(_src1))
- if CONST is None:
+ if T.constexpr(CONST is None):
src2_indices = T.meta_var(inst_gen.generate_indices(_src2))
T.evaluate(
func(
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
index a3a83bb3a4..fd00bfc0a1 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
@@ -113,7 +113,7 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc
vec_dst_idx = T.meta_var(inst_gen.generate_indices(binary_output))
reduce_dst_idx = T.meta_var(inst_gen.generate_indices(reduce_output))
if inst_gen.make_guard(binary_output):
- if CONST is None:
+ if T.constexpr(CONST is None):
src_2_indices = T.meta_var(inst_gen.generate_indices(binary_input2)) # noqa: E501
T.nki.tensorscalar_reduce(dst2[tuple(reduce_dst_idx)], dst1[tuple(vec_dst_idx)], src1[tuple(src_1_indices)], src2[tuple(src_2_indices)], binary_opcode, reduce_opcode, reverse[0]) # noqa: E501
else:
@@ -134,7 +134,7 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc
if inst_gen.make_guard(binary_output):
src_1_indices = T.meta_var(inst_gen.generate_indices(binary_input1)) # noqa: E501
vec_dst_idx = T.meta_var(inst_gen.generate_indices(binary_output)) # noqa: E501
- if CONST is None:
+ if T.constexpr(CONST is None):
src_2_indices = T.meta_var(inst_gen.generate_indices(binary_input2)) # noqa: E501
T.nki.tensorscalar_reduce(intermediate_buffer[p_loop, reduction_b_loop], dst1[tuple(vec_dst_idx)], src1[tuple(src_1_indices)], src2[tuple(src_2_indices)], binary_opcode, reduce_opcode, reverse[0]) # noqa: E501
else:
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
index 2cc80e57c2..8b56df56b1 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
@@ -112,7 +112,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
dst_1_indices = T.meta_var(inst_gen.generate_indices(unary_output))
dst_2_indices = T.meta_var(inst_gen.generate_indices(reduce_output))
if inst_gen.make_guard(unary_output):
- if isinstance(bias, TensorRegion):
+ if T.constexpr(isinstance(bias, TensorRegion)):
src_bias_indices = T.meta_var(inst_gen.generate_indices(bias))
T.evaluate(T.nki.activation_reduce(dst2[tuple(dst_2_indices)], dst1[tuple(dst_1_indices)], src[tuple(src_1_indices)], unary_opcode, reduce_opcode, bias_buffer[tuple(src_bias_indices)], scale)) # noqa: E501
else:
@@ -138,7 +138,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
src_1_indices = T.meta_var(inst_gen.generate_indices(unary_input))
dst_1_indices = T.meta_var(inst_gen.generate_indices(unary_output))
if inst_gen.make_guard(unary_output):
- if isinstance(bias, TensorRegion):
+ if T.constexpr(isinstance(bias, TensorRegion)):
src_bias_indices = T.meta_var(inst_gen.generate_indices(bias)) # noqa: E501
T.evaluate(T.nki.activation_reduce(intermediate_buffer[p_loop, reduction_b_loop], dst1[tuple(dst_1_indices)], src[tuple(src_1_indices)], unary_opcode, reduce_opcode, bias_buffer[tuple(src_bias_indices)], scale)) # noqa: E501
else:
diff --git a/python/tvm/backend/trn/tile_primitive/copy/default.py b/python/tvm/backend/trn/tile_primitive/copy/default.py
index 561fca20c9..7b1fae5c50 100644
--- a/python/tvm/backend/trn/tile_primitive/copy/default.py
+++ b/python/tvm/backend/trn/tile_primitive/copy/default.py
@@ -133,9 +133,9 @@ def transpose_schedule(
)
src_indices = T.meta_var(inst_gen.generate_indices(src_region))
dst_indices = T.meta_var(inst_gen.generate_indices(dst_region))
- src_guard = T.meta_var(inst_gen.make_guard(src_region))
- dst_guard = T.meta_var(inst_gen.make_guard(dst_region))
- if src_guard and dst_guard:
+ if inst_gen.make_guard(src_region) and inst_gen.make_guard(
+ dst_region
+ ):
T.evaluate(
T.nki.matmul(
dst_buffer[tuple(dst_indices)],
@@ -176,16 +176,14 @@ def transpose_schedule(
for rhs_f_loop in T.serial(0, rhs_f_size, annotations={nki_dim: "rhs_F"}): # noqa: E501
inst_gen.set_bind_map(src_region, {b_var: b_loop, lhs_f: lhs_f_loop, lhs_p: p_loop, extend_b: extend_b_loop}) # noqa: E501
src_indices = T.meta_var(inst_gen.generate_indices(src_region))
- src_guard = T.meta_var(inst_gen.make_guard(src_region))
- if src_guard:
+ if inst_gen.make_guard(src_region):
T.evaluate(T.nki.matmul(acc_psum[b_loop % max_psum_slots, lhs_f_loop,extend_b_loop * rhs_f_size + rhs_f_loop], src_buffer[tuple(src_indices)], identity_tensor[p_loop, rhs_f_loop])) # noqa: E501
with T.attr(0, "tensorized_nki_instruction", 1):
for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
for f_loop in T.serial(0, rhs_f_size * extend_len, annotations={nki_dim: "F"}):
inst_gen.set_bind_map(dst_region, {b_var: b_loop, lhs_f: p_loop, dst_f: f_loop % rhs_f_size, extend_b: f_loop // rhs_f_size}) # noqa: E501
- dst_guard = T.meta_var(inst_gen.make_guard(dst_region))
dst_indices = T.meta_var(inst_gen.generate_indices(dst_region))
- if dst_guard:
+ if inst_gen.make_guard(dst_region):
T.evaluate(T.nki.tensor_copy(dst_buffer[tuple(dst_indices)], acc_psum[b_loop % max_psum_slots, p_loop, f_loop])) # noqa: E501
# fmt: on
return transpose_sbuf_output
diff --git a/python/tvm/backend/trn/tile_primitive/gemm/default.py b/python/tvm/backend/trn/tile_primitive/gemm/default.py
index 9935a38970..8cb4af4e07 100644
--- a/python/tvm/backend/trn/tile_primitive/gemm/default.py
+++ b/python/tvm/backend/trn/tile_primitive/gemm/default.py
@@ -227,7 +227,6 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
for p_loop in T.serial(0, p_size, annotations={"nki_dim": "P"}):
for lhs_f_loop in T.serial(0, lhs_f_size, annotations={"nki_dim": "lhs_F"}):
for rhs_f_loop in T.serial(0, inst_repr.size, annotations={"nki_dim": "rhs_F"}):
- b_idx = T.meta_var(lhs_b_loop * rhs_b_extent + rhs_b_loop)
inst_gen.set_bind_map(A_buffer_region, {lhs_b: lhs_b_loop, lhs_f: lhs_f_loop, p: p_loop, reduction_b: reduction_b_loop}) # noqa: E501
inst_gen.set_bind_map(B_buffer_region, {rhs_b: rhs_b_loop, rhs_f: rhs_f_loop, p: p_loop, reduction_b: reduction_b_loop}) # noqa: E501
inst_gen.set_bind_map(C_buffer_region, {lhs_f: lhs_f_loop, rhs_f: rhs_f_loop, lhs_b: lhs_b_loop, rhs_b: rhs_b_loop}) # noqa: E501
@@ -235,10 +234,10 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
rhs_indices = T.meta_var(inst_gen.generate_indices(B_buffer_region))
C_indices = T.meta_var(inst_gen.generate_indices(C_buffer_region))
if inst_gen.make_guard(A_buffer_region) and inst_gen.make_guard(B_buffer_region): # noqa: E501
- if C_as_output:
+ if T.constexpr(C_as_output):
T.evaluate(T.nki.matmul(acc[C_indices], A[lhs_indices], B[rhs_indices])) # noqa: E501
else:
- T.evaluate(T.nki.matmul(acc[b_idx % max_psum_slots, lhs_f_loop, rhs_f_loop], A[lhs_indices], B[rhs_indices])) # noqa: E501
+ T.evaluate(T.nki.matmul(acc[(lhs_b_loop * rhs_b_extent + rhs_b_loop) % max_psum_slots, lhs_f_loop, rhs_f_loop], A[lhs_indices], B[rhs_indices])) # noqa: E501
if C.scope() == "trn.psum":
# This fragment captures buffers and indices from its insertion scope.
@@ -280,11 +279,10 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
with T.attr(0, "tensorized_nki_instruction", 1):
for lhs_f_loop in T.serial(0, lhs_f_size, annotations={"nki_dim": "P"}):
for rhs_f_loop in T.serial(0, inst_repr.size, annotations={"nki_dim": "F"}):
- b_idx = T.meta_var(lhs_b_loop * rhs_b_extent + rhs_b_loop)
inst_gen.set_bind_map(C_buffer_region, {lhs_f: lhs_f_loop, rhs_f: rhs_f_loop, lhs_b: lhs_b_loop, rhs_b: rhs_b_loop}) # noqa: E501
if inst_gen.make_guard(C_buffer_region):
acc_indices = T.meta_var(inst_gen.generate_indices(C_buffer_region))
- T.evaluate(T.nki.tensor_copy(C[acc_indices], acc_psum[b_idx % max_psum_slots, lhs_f_loop, rhs_f_loop])) # noqa: E501
+ T.evaluate(T.nki.tensor_copy(C[acc_indices], acc_psum[(lhs_b_loop * rhs_b_extent + rhs_b_loop) % max_psum_slots, lhs_f_loop, rhs_f_loop])) # noqa: E501
# fmt: on
return impl_C_sbuf
diff --git a/python/tvm/backend/trn/tile_primitive/select/default.py b/python/tvm/backend/trn/tile_primitive/select/default.py
index 8e281c8e9f..e8f40f4507 100644
--- a/python/tvm/backend/trn/tile_primitive/select/default.py
+++ b/python/tvm/backend/trn/tile_primitive/select/default.py
@@ -120,8 +120,7 @@ def select_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
if inst_gen.make_guard(dst):
dst_indices = T.meta_var(inst_gen.generate_indices(dst))
true_value_indices = T.meta_var(inst_gen.generate_indices(true_value))
- pred = T.meta_var(analyzer.simplify(op.predicate.apply(inst_gen.generate_axes(dst)))) # noqa: E501
- T.evaluate(T.nki.affine_select(dst_buffer[tuple(dst_indices)], pred, true_value_buffer[tuple(true_value_indices)], false_value)) # noqa: E501
+ T.evaluate(T.nki.affine_select(dst_buffer[tuple(dst_indices)], analyzer.simplify(op.predicate.apply(inst_gen.generate_axes(dst))), true_value_buffer[tuple(true_value_indices)], false_value)) # noqa: E501
# fmt: on
return impl
diff --git a/python/tvm/backend/trn/tile_primitive/unary/utils.py b/python/tvm/backend/trn/tile_primitive/unary/utils.py
index 106648b17d..ff635350a0 100644
--- a/python/tvm/backend/trn/tile_primitive/unary/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/unary/utils.py
@@ -177,13 +177,13 @@ def generate_unary_func(
inst_gen.set_bind_map_all({p_var: p_loop, f_var: f_loop, b_var: b_loop})
dst_indices = T.meta_var(inst_gen.generate_indices(dst_buffer_region))
if inst_gen.make_guard(dst_buffer_region):
- if unary_op == MapOpType.FILL:
+ if T.constexpr(unary_op == MapOpType.FILL):
T.evaluate(T.nki.memset(dst[tuple(dst_indices)], _src))
else:
src_indices = T.meta_var(inst_gen.generate_indices(_src))
- if unary_op == MapOpType.RECIPROCAL:
+ if T.constexpr(unary_op == MapOpType.RECIPROCAL):
T.evaluate(T.nki.reciprocal(dst[tuple(dst_indices)], src[tuple(src_indices)])) # noqa: E501
- elif isinstance(bias, TensorRegion):
+ elif T.constexpr(isinstance(bias, TensorRegion)):
bias_indices = T.meta_var(inst_gen.generate_indices(bias))
T.evaluate(T.nki.activation(dst[tuple(dst_indices)], src[tuple(src_indices)], opcode, scale=scale, bias=bias_buffer[tuple(bias_indices)])) # noqa: E501
else:
diff --git a/python/tvm/ir/module.py b/python/tvm/ir/module.py
index 1ce93e3700..fd5ebc59aa 100644
--- a/python/tvm/ir/module.py
+++ b/python/tvm/ir/module.py
@@ -64,7 +64,7 @@ class IRModule(Node, Scriptable):
attrs,
global_infos,
)
- self.pyfuncs = {}
+ self.__pyfuncs__ = {}
def clone(self) -> "IRModule":
return _ffi_api.Module_Clone(self)
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index a9c933b7d1..955c988906 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -113,3 +113,24 @@ from .binding_rewrite import DataflowBlockRewrite
import tvm.script
tvm.script.register_dialect("relax", "tvm.relax.script")
+
+
+def _check_script_module(module: "tvm.ir.IRModule") -> None:
+ # Delay builder imports until validation, keeping dialect/runtime bootstrap safe.
+ from .script.builder.parser_protocol import _check_module_well_formed
+
+ _check_module_well_formed(module)
+
+
+tvm.script.register_module_validator(_check_script_module, prepend=True)
+
+
+def _initialize_script_namespace() -> None:
+ from . import script
+
+ script._initialize()
+
+
+from tvm.script.parser import register_namespace_initializer as _register_namespace_initializer
+
+_register_namespace_initializer(_initialize_script_namespace, aliases=("R", "relax"))
diff --git a/python/tvm/relax/backend/gpu_generic/sampling.py b/python/tvm/relax/backend/gpu_generic/sampling.py
index 487027ce7b..302b8b7e8c 100644
--- a/python/tvm/relax/backend/gpu_generic/sampling.py
+++ b/python/tvm/relax/backend/gpu_generic/sampling.py
@@ -174,7 +174,7 @@ def gpu_multinomial_from_uniform(
local_sum[()] = T.Cast(dtype, init_value)
for i in T.unroll(thread_elem):
- if mask_local is not None:
+ if T.constexpr(mask_local is not None):
if mask_local[i]:
local_sum[()] = reduce_op(local_sum[()], data_local[i])
else:
diff --git a/python/tvm/relax/base_py_module.py b/python/tvm/relax/base_py_module.py
index 825b532364..cba22b5dba 100644
--- a/python/tvm/relax/base_py_module.py
+++ b/python/tvm/relax/base_py_module.py
@@ -52,7 +52,8 @@ class BasePyModule:
3. Wrapping Relax functions for easy Python calling.
4. Cross-function calls between Python, TIR, and Relax functions.
- Only IRModules that inherit from this class are allowed to contain Python functions.
+ Shared IRModules collect Python functions in ``__pyfuncs__``. Decorating a
+ subclass with ``R.py_module`` adds this executable runtime interface.
"""
def __del__(self):
@@ -168,8 +169,8 @@ class BasePyModule:
def _create_relax_wrapper(name):
def wrapper(*args, **kwargs):
"""Wrapper for Relax function with automatic tensor conversion."""
- if hasattr(self.ir_mod, "pyfuncs") and name in self.ir_mod.pyfuncs:
- return self.ir_mod.pyfuncs[name](*args, **kwargs)
+ if hasattr(self.ir_mod, "__pyfuncs__") and name in self.ir_mod.__pyfuncs__:
+ return self.ir_mod.__pyfuncs__[name](*args, **kwargs)
if self.relax_vm is not None:
converted_args = self._convert_pytorch_to_tvm(list(args))
@@ -191,7 +192,7 @@ class BasePyModule:
def _register_python_functions(self):
"""Register Python functions with the VM runtime for call_py_func support."""
- if not hasattr(self.ir_mod, "pyfuncs") or not self.ir_mod.pyfuncs:
+ if not hasattr(self.ir_mod, "__pyfuncs__") or not self.ir_mod.__pyfuncs__:
return
try:
@@ -199,7 +200,7 @@ class BasePyModule:
except ValueError:
return
- for func_name, py_func in self.ir_mod.pyfuncs.items():
+ for func_name, py_func in self.ir_mod.__pyfuncs__.items():
def create_py_func_wrapper(name, original_func):
def wrapper(*args, **kwargs):
@@ -514,7 +515,7 @@ class BasePyModule:
"""Print TVM IR into TVMScript text format with Python function support.
This method extends the standard IRModule script() method to handle
- Python functions stored in the IRModule's pyfuncs attribute.
+ Python functions stored in the IRModule's ``__pyfuncs__`` attribute.
"""
# First get the standard IRModule script
base_script = self.ir_mod.script(
@@ -535,7 +536,7 @@ class BasePyModule:
)
# If there are no Python functions, return the base script
- if not hasattr(self.ir_mod, "pyfuncs") or not self.ir_mod.pyfuncs:
+ if not hasattr(self.ir_mod, "__pyfuncs__") or not self.ir_mod.__pyfuncs__:
return base_script
# Insert Python functions into the script
@@ -559,8 +560,8 @@ class BasePyModule:
class_indent = len(line) - len(line.lstrip())
# Insert Python functions after the class definition
- if hasattr(self.ir_mod, "pyfuncs") and self.ir_mod.pyfuncs:
- for func_name, func in self.ir_mod.pyfuncs.items():
+ if hasattr(self.ir_mod, "__pyfuncs__") and self.ir_mod.__pyfuncs__:
+ for func_name, func in self.ir_mod.__pyfuncs__.items():
# Get the function source code
func_source = self._get_function_source(func)
if func_source:
diff --git a/python/tvm/relax/frontend/nn/llm/_page_kernels.py b/python/tvm/relax/frontend/nn/llm/_page_kernels.py
index 6682e5f18c..882de88e45 100644
--- a/python/tvm/relax/frontend/nn/llm/_page_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_page_kernels.py
@@ -48,7 +48,7 @@ def _kv_cache_transpose_append(num_key_value_heads, head_dim, dtype, page_size:
var_position_map: T.handle,
):
T.func_attr({"tirx.noalias": True})
- ntoken = T.Var("num_tokens_excluding_cache", "int64")
+ ntoken = T.int64()
num_pages = T.int64()
pages_elem_offset = T.int64()
position_map_elem_offset = T.int32()
@@ -84,7 +84,7 @@ def _kv_cache_transpose_append_mla(d_qk: int, dtype, page_size: int = 16):
var_position_map: T.handle,
):
T.func_attr({"tirx.noalias": True})
- ntoken = T.Var("num_tokens_excluding_cache", "int64")
+ ntoken = T.int64()
num_pages = T.int64()
pages_elem_offset = T.int64()
position_map_elem_offset = T.int32()
@@ -115,8 +115,8 @@ def _kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, head_dim, dty
layer_id: T.int64,
):
T.func_attr({"tirx.noalias": True})
- seqlen = T.Var("num_tokens_including_cache", "int64")
- page_size = T.Var("page_size", "int64")
+ seqlen = T.int64()
+ page_size = T.int64()
num_pages = T.int64()
pages_elem_offset = T.int64()
position_map_elem_offset = T.int64()
@@ -147,8 +147,8 @@ def _kv_cache_debug_get_kv_mla(num_hidden_layers, d_qk, dtype):
layer_id: T.int64,
):
T.func_attr({"tirx.noalias": True})
- seqlen = T.Var("num_tokens_including_cache", "int64")
- page_size = T.Var("page_size", "int64")
+ seqlen = T.int64()
+ page_size = T.int64()
num_pages = T.int64()
pages_elem_offset = T.int64()
position_map_elem_offset = T.int64()
diff --git a/python/tvm/relax/frontend/nn/llm/position_embedding.py b/python/tvm/relax/frontend/nn/llm/position_embedding.py
index 9af71b1730..1ff7e9c66d 100644
--- a/python/tvm/relax/frontend/nn/llm/position_embedding.py
+++ b/python/tvm/relax/frontend/nn/llm/position_embedding.py
@@ -499,7 +499,7 @@ def llama_rope_with_position_map( # pylint: disable=too-many-arguments
ext_factors: T.Buffer | None = None,
):
kwargs = {}
- if ext_factors:
+ if ext_factors is not None:
kwargs["ext_factors"] = ext_factors
cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
pos * scale, d, rotary_dim, theta, "float32", **kwargs
@@ -725,7 +725,7 @@ def llama4_rope_with_position_map( # pylint: disable=too-many-arguments
ext_factors: T.Buffer | None = None,
):
kwargs = {}
- if ext_factors:
+ if ext_factors is not None:
kwargs["ext_factors"] = ext_factors
cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
pos * scale, d, rotary_dim, theta, "float32", **kwargs
diff --git a/python/tvm/relax/frontend/nn/llm/tree_attn.py b/python/tvm/relax/frontend/nn/llm/tree_attn.py
index 9cf0ccb678..237cebc917 100644
--- a/python/tvm/relax/frontend/nn/llm/tree_attn.py
+++ b/python/tvm/relax/frontend/nn/llm/tree_attn.py
@@ -416,10 +416,10 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
with T.sblock():
for li, lj, lk in T.grid(tile_x, tile_z, tile_y):
with T.sblock("S_gemm"):
- i, j, k = T.axis.remap("SSR", [li, lj, lk])
+ i, j, k_axis = T.axis.remap("SSR", [li, lj, lk])
with T.init():
S_local[i, j] = 0.0
- S_local[i, j] += T.cast(Q_smem[i, k], "float32") * T.cast(K_smem[j, k], "float32") * sm_scale * math.log2(math.exp(1))
+ S_local[i, j] += T.cast(Q_smem[i, k_axis], "float32") * T.cast(K_smem[j, k_axis], "float32") * sm_scale * math.log2(math.exp(1))
T.tvm_storage_sync("shared")
for li, lj in T.grid(tile_x, tile_z):
with T.sblock("S_store"):
@@ -482,10 +482,10 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
with T.sblock():
for li, lj, lk in T.grid(tile_x, tile_y, tile_z):
with T.sblock("O_gemm"):
- i, j, k = T.axis.remap("SSR", [li, lj, lk])
+ i, j, k_axis = T.axis.remap("SSR", [li, lj, lk])
with T.init():
O_local[i, j] *= T.exp2(m_prev_smem[i] - m_smem[i])
- O_local[i, j] += S_smem[i, k] * T.cast(V_smem[k, j], "float32")
+ O_local[i, j] += S_smem[i, k_axis] * T.cast(V_smem[k_axis, j], "float32")
# Store O from smem to gmem
for li, lj in T.grid(tile_x, tile_y):
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index b6704660de..7129198aa4 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -8343,8 +8343,8 @@ def _build_tflite_rfft2d_primfunc(input_shape, output_pair_shape):
for b_idx, out_y, out_x in T.grid(batch, height, out_width):
with T.sblock("rfft2d"):
v_b, v_oy, v_ox = T.axis.remap("SSS", [b_idx, out_y, out_x])
- real_sum = T.float32(0)
- imag_sum = T.float32(0)
+ real_sum: T.float32 = T.float32(0)
+ imag_sum: T.float32 = T.float32(0)
input_base = v_b * height * width
for in_y, in_x in T.grid(height, width):
phase_y = T.Cast("float32", v_oy) * T.Cast("float32", in_y) / T.float32(height)
diff --git a/python/tvm/relax/global_info.py b/python/tvm/relax/global_info.py
index 2eacc8b072..ecbf99c9dd 100644
--- a/python/tvm/relax/global_info.py
+++ b/python/tvm/relax/global_info.py
@@ -14,13 +14,13 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Global information used by Relax."""
+"""Concrete global information used by Relax."""
import tvm_ffi
-import tvm
from tvm.ir import GlobalInfo
-from tvm.runtime import Device
+from tvm.runtime import Device, convert
+from tvm.target import Target
from . import _ffi_api
@@ -46,7 +46,7 @@ class VDevice(GlobalInfo):
memory_scope: str = "global",
) -> None:
if isinstance(target, dict | str):
- target = tvm.target.Target(tvm.runtime.convert(target))
+ target = Target(convert(target))
if isinstance(target, Device):
- target = tvm.target.Target.from_device(target)
+ target = Target.from_device(target)
self.__init_handle_by_constructor__(_ffi_api.VDevice, target, vdevice_id, memory_scope)
diff --git a/python/tvm/relax/relax_to_pyfunc_converter.py b/python/tvm/relax/relax_to_pyfunc_converter.py
index 39667dda92..4efaa74181 100644
--- a/python/tvm/relax/relax_to_pyfunc_converter.py
+++ b/python/tvm/relax/relax_to_pyfunc_converter.py
@@ -88,8 +88,8 @@ class RelaxToPyFuncConverter:
new_ir_mod = self.ir_module.clone()
# Initialize pyfuncs if not exists
- if not hasattr(new_ir_mod, "pyfuncs"):
- new_ir_mod.pyfuncs = {}
+ if not hasattr(new_ir_mod, "__pyfuncs__"):
+ new_ir_mod.__pyfuncs__ = {}
# Get Relax function names from IRModule
relax_func_names = []
@@ -116,7 +116,7 @@ class RelaxToPyFuncConverter:
py_func = self._convert_relax_func_to_python(relax_func, func_name)
# Store in pyfuncs
- new_ir_mod.pyfuncs[func_name] = py_func
+ new_ir_mod.__pyfuncs__[func_name] = py_func
return new_ir_mod
diff --git a/python/tvm/relax/script/__init__.py b/python/tvm/relax/script/__init__.py
index acbd9d2435..172af4eece 100644
--- a/python/tvm/relax/script/__init__.py
+++ b/python/tvm/relax/script/__init__.py
@@ -14,14 +14,205 @@
# 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).
+"""Public canonical TVMScript dialect namespace."""
-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.
-"""
+from __future__ import annotations
-# pylint: disable=redefined-builtin,wildcard-import,unused-wildcard-import
-from .parser import *
-from .parser import dist
+import ast as _ast
+import importlib as _importlib
+import inspect as _inspect
+import sys as _sys
+from collections.abc import Callable as _Callable
+from typing import TYPE_CHECKING
+from typing import Any as _Any
+
+from tvm.script.parser.protocol_registry import module_decorator as _module_decorator
+
+if TYPE_CHECKING:
+ from tvm.ir import IRModule
+ from tvm.relax.base_py_module import BasePyModule
+ from tvm.runtime import Device
+ from tvm.target import Target
+
+
+_initialized = False
+_initializing = False
+
+
+def _initialize() -> None:
+ global _initialized, _initializing
+ if _initialized or _initializing:
+ return
+ _initializing = True
+ try:
+ from tvm.script.parser import entry, register_namespace
+ from tvm.script.parser.protocol_registry import declaration_kind
+
+ from . import builder
+
+ globals().update(
+ (name, value) for name, value in vars(builder).items() if not name.startswith("_")
+ )
+ globals()["__tvm_value_if__"] = builder.__tvm_value_if__
+ globals()["function"] = declaration_kind("R.function", "function")(
+ entry.make_decorator(builder)
+ )
+ globals()["macro"] = declaration_kind("R.macro", "helper")(
+ entry.make_macro_decorator(builder, preserve_return=True)
+ )
+ namespace = _sys.modules[__name__]
+ register_namespace("R", namespace)
+ register_namespace("relax", namespace)
+ globals()["__all__"] = sorted(name for name in globals() if not name.startswith("_"))
+ _initialized = True
+ finally:
+ _initializing = False
+
+
+def __getattr__(name: str) -> _Any:
+ if name in ("builder", "tile"):
+ return _importlib.import_module(__name__ + "." + name)
+ if name.startswith("_") and name != "__all__":
+ raise AttributeError(name)
+ _initialize()
+ if name in globals():
+ return globals()[name]
+ return getattr(_importlib.import_module(__name__ + ".builder"), name)
+
+
+class _PyModuleFactory:
+ """Instantiate an executable module without sharing its runtime registry."""
+
+ def __init__(self, module: IRModule, original_class: type) -> None:
+ self.ir_module = module
+ self.original_class = original_class
+ self.__name__ = original_class.__name__
+
+ def __call__(self, device: Device | None = None, target: Target | None = None) -> BasePyModule:
+ from tvm import cpu, ir
+ from tvm.relax.base_py_module import BasePyModule
+
+ source = self.ir_module
+ instance_module = ir.IRModule(
+ source.functions, attrs=source.attrs, global_infos=source.global_infos
+ )
+ instance = BasePyModule(instance_module, device or cpu(0), target)
+ for name, function in source.__pyfuncs__.items():
+ instance.add_python_function(name, function)
+ return instance
+
+ def __getattr__(self, name: str) -> _Any:
+ return getattr(self.ir_module, name)
+
+
+def py_module(
+ module: type | None = None, **options: _Any
+) -> IRModule | _PyModuleFactory | _Callable[[type], IRModule | _PyModuleFactory]:
+ """Parse a class and attach executable Relax Python-function metadata.
+
+ Parameters
+ ----------
+ module : type, optional
+ Class containing registered script functions and ``I.pyfunc`` members.
+ None, the default, returns a decorator. Python bodies are retained as
+ their original callables and are not executed during construction.
+ **options : _Any
+ Options forwarded once to shared ``parse``, including
+ ``check_well_formed`` (True by default) and ``track_span`` (True).
+
+ Returns
+ -------
+ IRModule or callable
+ An IRModule with original Python callables in ``__pyfuncs__`` and
+ corresponding Relax ExternFunc metadata. For a BasePyModule subclass,
+ return a factory accepting optional device and target, creating a fresh
+ runtime instance on each call. With no class, return the decorator.
+
+ Raises
+ ------
+ TypeError
+ If the decorated value is not a class, or an unsupported parser option
+ is supplied.
+ SyntaxError
+ If the source violates a parser syntax restriction. Construction and
+ enabled validation propagate their original exceptions unchanged.
+ OSError
+ If source inspection cannot recover a Python function's definition.
+
+ Notes
+ -----
+ The decorator captures the original class-definition scope and passes it
+ explicitly to shared parsing. It retains no Python frame or scope snapshot.
+ ExternFunc metadata uses each original function's source coordinates when
+ span tracking is enabled. Runtime compilation/registration is deferred to
+ BasePyModule construction; plain IRModule construction registers no runtime
+ functions. Ordinary shared ``I.ir_module`` never performs these Relax steps.
+
+ .. code:: python
+
+ # Source
+ @R.py_module
+ class Module:
+ @I.pyfunc
+ def twice(value):
+ return value * 2
+
+ # Relax adaptation after one shared parse
+ assert Module.__pyfuncs__["twice"](3) == 6
+ assert isinstance(Module["twice"], relax.ExternFunc)
+ """
+
+ def apply(source_class: type) -> IRModule | _PyModuleFactory:
+ from tvm import ir, relax
+ from tvm.relax.base_py_module import BasePyModule
+ from tvm.script.parser.entry import parse
+ from tvm.script.parser.inspect_source import acquire_source, capture_definition_scope
+
+ if not _inspect.isclass(source_class):
+ raise TypeError(f"Expect a class, but got: {source_class}")
+ frame = _inspect.currentframe().f_back
+ try:
+ if frame.f_code is py_module.__code__:
+ frame = frame.f_back
+ definition_scope = capture_definition_scope(frame)
+ definition_source = (frame.f_code.co_filename, frame.f_lineno)
+ finally:
+ del frame
+ result = parse(
+ source_class,
+ definition_scope=definition_scope,
+ _definition_source=definition_source,
+ **options,
+ )
+ result.__pyfuncs__ = getattr(result, "__pyfuncs__", {})
+ for name, function in result.__pyfuncs__.items():
+ tree, filename, _ = acquire_source(function)
+ node = tree.body[-1]
+ span = None
+ if options.get("track_span", True):
+ span = ir.Span(
+ ir.SourceName(filename),
+ node.lineno,
+ node.end_lineno,
+ node.col_offset + 1,
+ node.end_col_offset + 1,
+ )
+ result[name] = relax.ExternFunc(name, span=span).with_attrs(
+ {
+ "is_pyfunc": True,
+ "function_type": "python",
+ "python_function_name": name,
+ "python_source": _ast.unparse(node),
+ "python_packed_func": function,
+ }
+ )
+ result.__name__ = source_class.__name__
+ if issubclass(source_class, BasePyModule):
+ return _PyModuleFactory(result, source_class)
+ return result
+
+ return apply(module) if module is not None else apply
+
+
+# Member decorators may defer to this shared-parser entry without importing Relax.
+_module_decorator("R.py_module")(py_module)
diff --git a/python/tvm/relax/script/builder/__init__.py b/python/tvm/relax/script/builder/__init__.py
index 021bcf9b1c..b94c54cce9 100644
--- a/python/tvm/relax/script/builder/__init__.py
+++ b/python/tvm/relax/script/builder/__init__.py
@@ -14,12 +14,306 @@
# 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
-# pylint: disable=wildcard-import,redefined-builtin
-from . import distributed, frame, ir
+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.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
+from tvm.script.parser.protocol_registry import args_policy as _args_policy
+from tvm.script.parser.protocol_registry import constexpr as constexpr
+
+from . import distributed as dist
+from . import ir as _native
+from .distributed.ir import _lookup_device_mesh
from .ir import *
+from .parser_protocol import (
+ and_,
+ arg,
+ assert_,
+ bind_,
+ break_,
+ call_global_var_,
+ check_well_formed_,
+ continue_,
+ decl_mutable_cell_,
+ else_,
+ emit_,
+ eq_,
+ for_,
+ func_name,
+ func_ret_type,
+ function_,
+ ge_,
+ gt_,
+ if_,
+ if_then_else_,
+ le_,
+ lt_,
+ ne_,
+ not_,
+ or_,
+ range_,
+ resolve_global_info_,
+ resolve_type_var_,
+ return_,
+ set_mutable_cell_,
+ setattr_,
+ setitem_,
+ then_,
+ unpack,
+ while_,
+)
+
+If = if_
+For = for_
+
+# Syntax capability: mutable declaration policies apply only in this dialect.
+supports_mutable_declarations = False
+
+
+@_args_policy("R.Tensor", {"shape": "expr_str", "vdevice": "global_info"}, scalar_strings=False)
+def Tensor(shape=None, dtype=None, vdevice=None, ndim=-1, *, span=None):
+ """Construct a tensor type from concrete shape dimensions."""
+ 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
+)
+def DTensor(shape=None, dtype=None, device_mesh=None, placement="", *, ndim=-1, span=None):
+ """Construct a distributed tensor type from concrete dimensions."""
+ 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))
+
+
+# The distributed source spelling shares concrete constructors and argument policy.
+dist.DTensor = DTensor
+_ARGS_POLICIES["R.dist.DTensor"] = _ARGS_POLICIES["R.DTensor"]
+dist.device_mesh = device_mesh
+
+Range = _ir.Range
+
+# Syntax metadata: Relax statements return one same-named branch output.
+__tvm_value_if__ = True
+
+
+@_args_policy("R.Shape", {"values": "expr_str"}, dtype="int64")
+def Shape(values=None, ndim=-1, *, span=None):
+ """Construct a shape type from concrete dimensions."""
+ return _relax.ShapeType(values, ndim, _source_span(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 or opaque Relax 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=_source_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, _source_span(span)
+ )
+
+
+def Tuple(*fields, span=None):
+ """Construct a tuple type while preserving missing component types."""
+ if len(fields) == 1 and isinstance(fields[0], list | _python.tuple):
+ fields = fields[0]
+ return _ir.TupleType([_type(field) for field in fields], _source_span(span))
+
+
+def Object(*, span=None):
+ """Construct the unconstrained Relax value type."""
+ return _relax.AnyType(_source_span(span))
+
+
+Any = Object
+
+
+is_type_var = _ir.is_prim_var
+
+
+def type_var(name, *, dtype=None, span=None):
+ """Construct a fresh standalone primitive symbol."""
+ return _ir.Var(name, "int64" if dtype is None else dtype, _source_span(span))
+
+
+func_ret_ty = func_ret_type
+
+
+def dataflow(*, span=None):
+ """Create a dataflow context with explicit finalized exports.
+
+ Parameters
+ ----------
+ span : Span or source location, optional
+ Source location attached to the constructed IR.
+
+ Returns
+ -------
+ res : frame.BindingBlockFrame
+ The constructed frame, retaining source metadata.
+ """
+ return _at(span, _native.dataflow())
+
+
+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 match_cast(value, ty, *, span=None):
+ """Construct a match-cast descriptor 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, _source_span(span))
+
+
+__all__ = [
+ "then_",
+ "else_",
+ "while_",
+ "setattr_",
+ *_native.__all__,
+ "Any",
+ "Callable",
+ "DTensor",
+ "For",
+ "for_",
+ "function_",
+ "if_",
+ "check_well_formed_",
+ "resolve_type_var_",
+ "resolve_global_info_",
+ "call_global_var_",
+ "decl_mutable_cell_",
+ "set_mutable_cell_",
+ "Object",
+ "Range",
+ "Shape",
+ "Tensor",
+ "Tuple",
+ "assert_",
+ "break_",
+ "continue_",
+ "bind_",
+ "device_mesh",
+ "dist",
+ "emit_",
+ "eq_",
+ "ge_",
+ "gt_",
+ "le_",
+ "lt_",
+ "ne_",
+ "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 host, primitive, or tensor values."""
+ 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 host, primitive, or tensor values."""
+ 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):
+ """Negate a host or IR boolean without coercing IR 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 elementwise selection or select already-built host 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"]
+
+
+__all__ += ["and_", "if_then_else_", "not_", "or_"]
diff --git a/python/tvm/relax/script/builder/distributed/ir.py b/python/tvm/relax/script/builder/distributed/ir.py
index 6192862473..827dea67fd 100644
--- a/python/tvm/relax/script/builder/distributed/ir.py
+++ b/python/tvm/relax/script/builder/distributed/ir.py
@@ -45,6 +45,7 @@ from tvm.relax.utils import convert_to_expr
from tvm.runtime import _tensor
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder.ir import IRModuleFrame
+from tvm.script.ir_builder.ir.ir import lookup_global_info
from . import _ffi_api
@@ -93,24 +94,26 @@ def const(
value: bool | int | float | _np.ndarray | tvm.runtime.Tensor,
ty: DTensorType,
) -> GenericConst:
- """Create a constant value.
+ """Create a distributed constant value with the specified tensor type.
Parameters
----------
- value: Union[bool, int, float, numpy.ndarray, tvm.runtime.Tensor]
+ value : bool, int, float, numpy.ndarray or tvm.runtime.Tensor
The constant value.
- dtype: Optional[str]
- The data type of the resulting constant.
+ ty : DTensorType
+ The distributed tensor type, including its tensor dtype, device mesh
+ and placement.
- Note
- ----
- When dtype is None, we use the following rule:
+ Returns
+ -------
+ res : GenericConst
+ The constant carrying the supplied distributed tensor type.
- - int maps to "int32"
- - float maps to "float32"
- - bool maps to "bool"
- - other using the same default rule as numpy.
+ Notes
+ -----
+ Python scalars and NumPy values are converted to the dtype specified by
+ ``ty.tensor_ty``. The dtype is not inferred from the Python value.
"""
ty = tvm.runtime.convert(ty)
if not isinstance(ty, DTensorType):
@@ -131,10 +134,13 @@ def const(
def _lookup_device_mesh(device_mesh_str: py_str) -> DeviceMesh:
- if not IRBuilder.is_in_scope():
- raise ValueError("device_mesh cannot be found in global info")
name, index_str = device_mesh_str.split("[")
index = int(index_str[:-1])
+ if not IRBuilder.is_in_scope():
+ device_mesh = lookup_global_info(name, index)
+ if not isinstance(device_mesh, DeviceMesh):
+ raise TypeError("The device_mesh global info must be a DeviceMesh.")
+ return device_mesh
frames = IRBuilder.current().frames
for f in frames:
if isinstance(f, IRModuleFrame):
diff --git a/python/tvm/relax/script/builder/frame.py b/python/tvm/relax/script/builder/frame.py
index fec8a88af9..e81d11808b 100644
--- a/python/tvm/relax/script/builder/frame.py
+++ b/python/tvm/relax/script/builder/frame.py
@@ -18,7 +18,9 @@
from tvm_ffi import register_object as _register_object
-from tvm.script.ir_builder.base import IRBuilderFrame
+from tvm.script.ir_builder.base import IRBuilderFrame, _resolve_type_var
+
+from . import _ffi_api
@_register_object("script.ir_builder.relax.RelaxFrame")
@@ -32,7 +34,11 @@ class SeqExprFrame(RelaxFrame): ...
@_register_object("script.ir_builder.relax.FunctionFrame")
class FunctionFrame(SeqExprFrame):
- """The ir_builder frame for the relax function."""
+ """Native function frame retaining signature, symbols and finalized results."""
+
+ def resolve_type_var(self, name, dtype=None, *, value=None, span=None):
+ """Resolve a primitive symbol in this function's native map."""
+ return _resolve_type_var(self, _ffi_api.ResolveTypeVar, name, dtype, value=value, span=span)
@_register_object("script.ir_builder.relax.BindingBlockFrame")
diff --git a/python/tvm/relax/script/builder/ir.py b/python/tvm/relax/script/builder/ir.py
index a3d7d3c038..a406c402de 100644
--- a/python/tvm/relax/script/builder/ir.py
+++ b/python/tvm/relax/script/builder/ir.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=redefined-builtin, wrong-import-order, no-member, invalid-name
-"""IRBuilder for Relax dialect"""
+"""IRBuilder for Relax language variant"""
import builtins
import functools
@@ -28,7 +28,7 @@ from tvm import DataType, relax
from tvm.ir import IRModule, StringImm
from tvm.relax import Call, Expr, ExternFunc, ShapeExpr, TupleGetItem, Var, VarBinding, const
from tvm.relax.dpl import PatternMatchingRewriter
-from tvm.relax.global_info import VDevice
+from tvm.relax.global_info import DummyGlobalInfo, VDevice
############################### Operators ###############################
from tvm.relax.op import (
@@ -207,7 +207,10 @@ from tvm.runtime._tensor import (
vulkan,
webgpu,
)
-from tvm.script.ir_builder.ir import decl_function, lookup_vdevice
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder.ir import decl_function
+from tvm.script.ir_builder.ir.ir import _class_global_infos
+from tvm.script.parser.protocol_registry import module_decorator as _module_decorator
from . import _ffi_api, frame
@@ -218,7 +221,72 @@ py_tuple = tuple # pylint: disable=used-before-assignment
py_str = str # pylint: disable=used-before-assignment
-################################ Device ################################
+############################### GlobalInfo ###############################
+
+
+def dummy_global_info() -> "DummyGlobalInfo":
+ """Create a dummy global info expression.
+
+ Returns
+ -------
+ res : DummyGlobalInfo
+ The result dummy global info.
+ """
+ return DummyGlobalInfo() # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def vdevice(target=None, vdevice_id: int = 0, memory_scope: py_str = "global") -> "VDevice":
+ """Create a virtual device global info.
+ Parameters
+ ----------
+ target
+ The target.
+ vdevice_id: int
+ The virtual device index.
+ memory_scope: py_str
+ The memory scope, default is "global"
+
+ Returns
+ -------
+ res : VDevice
+ The result virtual device.
+ """
+ return VDevice(target, vdevice_id, memory_scope) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def lookup_vdevice(target_kind: py_str | None = None, device_index: int = -1) -> "VDevice":
+ """Retrieve a virtual device from the globalinfo vdevice list.
+
+ Parameters
+ ----------
+ target_kind: str
+ The target device kind, for example 'llvm' or 'cuda'. Use 'vdevice'
+ to index the complete virtual-device list.
+ device_index: int
+ The zero-based index among devices of the selected target kind, or
+ among all devices when target_kind is 'vdevice'.
+
+ Returns
+ -------
+ res : VDevice
+ The result virtual device.
+ """
+ if IRBuilder.is_in_scope():
+ return _ffi_api.LookupVDevice(target_kind, device_index)
+ infos = _class_global_infos()
+ if not infos:
+ raise ValueError("The GlobalInfos in the IRModule is not defined.")
+ vdevices = infos["vdevice"]
+ if device_index < 0 or device_index >= len(vdevices):
+ raise ValueError("The target VDevice in the GlobalInfos was not found.")
+ if not all(isinstance(value, VDevice) for value in vdevices):
+ raise TypeError("The vdevice global infos must contain VDevice values.")
+ if target_kind == "vdevice":
+ return vdevices[device_index]
+ matches = [value for value in vdevices if value.target.kind.name == target_kind]
+ if device_index >= len(matches):
+ raise ValueError("The target VDevice in the GlobalInfos was not found.")
+ return matches[device_index]
def to_vdevice(data: Expr, dst_vdevice: py_str | VDevice) -> Expr:
@@ -229,7 +297,7 @@ def to_vdevice(data: Expr, dst_vdevice: py_str | VDevice) -> Expr:
data : Expr
The tensor to be copied.
- dst_device : Union[py_str, VDevice]
+ dst_vdevice : Union[py_str, VDevice]
The destination device where the data is copied to.
Returns
@@ -252,6 +320,7 @@ def to_vdevice(data: Expr, dst_vdevice: py_str | VDevice) -> Expr:
def function(is_pure: bool = True, is_private: bool = False) -> frame.FunctionFrame:
"""Start a function frame.
+
Parameters
----------
is_pure: bool
@@ -272,6 +341,7 @@ def function(is_pure: bool = True, is_private: bool = False) -> frame.FunctionFr
def arg(name: py_str, ty: Type) -> Var:
"""Add a parameter to the last function frame.
+
Parameters
----------
name: str
@@ -290,6 +360,7 @@ def arg(name: py_str, ty: Type) -> Var:
def func_name(name: py_str) -> None:
"""Specify the name of the last function frame.
+
Parameters
----------
name: str
@@ -300,6 +371,7 @@ def func_name(name: py_str) -> None:
def func_attr(attrs: dict[py_str, tvm_Object]) -> None:
"""Specify the attrs of the last function frame.
+
Parameters
----------
attrs: Dict[str, Object]
@@ -310,6 +382,7 @@ def func_attr(attrs: dict[py_str, tvm_Object]) -> None:
def func_ret_type(ret_ty: Type) -> None:
"""Specify the return type of the last function frame.
+
Parameters
----------
ret_ty: Type
@@ -325,6 +398,7 @@ def func_ret_ty(ret_ty: Type) -> None:
def func_ret_value(value: Expr) -> None:
"""Specify the return value of the last function frame.
+
Parameters
----------
value: Expr
@@ -368,18 +442,49 @@ def rewriter(rewriter_mod: IRModule | type) -> PatternMatchingRewriter:
A rewriter object, which can be applied either to a Relax
function or to an entire IRModule.
+ Notes
+ -----
+ Class members are parsed together after the class body completes. Their
+ annotations use the decorator's original definition scope, which is released
+ after parsing. An existing IRModule is used directly.
+
"""
if not isinstance(rewriter_mod, IRModule):
- rewriter_mod = tvm.script.ir_module(rewriter_mod)
+ from tvm.script.parser.entry import parse
+ from tvm.script.parser.inspect_source import capture_definition_scope
+
+ if not inspect.isclass(rewriter_mod):
+ raise TypeError(f"Expect a class, but got: {rewriter_mod}")
+ frame = inspect.currentframe().f_back
+ try:
+ definition_scope = capture_definition_scope(frame)
+ definition_source = (frame.f_code.co_filename, frame.f_lineno)
+ finally:
+ del frame
+ try:
+ module = parse(
+ rewriter_mod,
+ definition_scope=definition_scope,
+ _definition_source=definition_source,
+ )
+ finally:
+ del definition_scope
+ module.__name__ = rewriter_mod.__name__
+ rewriter_mod = module
return PatternMatchingRewriter.from_module(rewriter_mod)
+# Member decorators defer to the same shared module construction boundary.
+_module_decorator("R.rewriter")(rewriter)
+
+
############################# BindingBlock ##############################
def dataflow() -> frame.BindingBlockFrame:
"""Start a dataflow binding block frame.
+
Returns
-------
frame: frame.BindingBlockFrame
@@ -390,6 +495,7 @@ def dataflow() -> frame.BindingBlockFrame:
def output(*vars: tuple[Var]) -> None:
"""Expose the dataflow block output variables as global ones.
+
Parameters
----------
vars: Tuple[Var]
@@ -649,6 +755,7 @@ def emit_with_ty(
def SeqExpr() -> frame.SeqExprFrame: # pylint: disable=invalid-name
"""Create a SeqExpr frame.
+
Returns
-------
res : frame.SeqExprFrame
@@ -684,6 +791,7 @@ def If(condition: Expr) -> frame.IfFrame: # pylint: disable=invalid-name
def Then() -> frame.ThenFrame: # pylint: disable=invalid-name
"""Create a then frame.
+
Returns
-------
res : frame.ThenFrame
@@ -694,6 +802,7 @@ def Then() -> frame.ThenFrame: # pylint: disable=invalid-name
def Else() -> frame.ElseFrame: # pylint: disable=invalid-name
"""Create an else frame.
+
Returns
-------
res : frame.ElseFrame
@@ -846,6 +955,7 @@ __all__ = [
"device",
"divide",
"dtype",
+ "dummy_global_info",
"dynamic_strided_slice",
"einsum",
"emit",
@@ -902,6 +1012,7 @@ __all__ = [
"logical_not",
"logical_or",
"logical_xor",
+ "lookup_vdevice",
"make_closure",
"matmul",
"max",
@@ -978,6 +1089,7 @@ __all__ = [
"tuple",
"unique",
"variance",
+ "vdevice",
"vision",
"vm",
"vpi",
diff --git a/python/tvm/relax/script/builder/parser_protocol.py b/python/tvm/relax/script/builder/parser_protocol.py
new file mode 100644
index 0000000000..e759767195
--- /dev/null
+++ b/python/tvm/relax/script/builder/parser_protocol.py
@@ -0,0 +1,546 @@
+# 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.
+"""Relax implementation of the shared source-to-builder protocol.
+
+Hooks delegate construction to this language variant's native builder and IR APIs.
+For example, generated ``X.if_(condition)`` creates the native conditional frame;
+``X.then_()`` and ``X.else_()`` enter its branches. See the corresponding shared
+``tvm.script.ir_builder.ir.parser_protocol`` hooks for operand and span contracts.
+"""
+
+from __future__ import annotations
+
+import builtins as _python
+import numbers as _numbers
+import re as _re
+from collections.abc import Sequence
+from typing import Any, NoReturn
+
+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.ir.prim import _ffi_api as _prim_ffi
+from tvm.script.ir_builder import IRBuilder as _IRBuilder
+from tvm.script.ir_builder import base as _base
+from tvm.script.ir_builder.ir.frame import IRModuleFrame as _IRModuleFrame
+
+from .. import builder as _builder
+from . import _ffi_api
+from . import frame as _frame
+from . import ir as _native
+
+_Span = _base.SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int] | None
+
+
+# ------------
+# Section: control flow
+# ------------
+#
+# ``with X.if_(cond):`` opens a conditional frame.
+# ``with X.then_():`` enters its true branch.
+# ``with X.else_():`` enters its false branch.
+# ``with X.for_(loop) as i:`` binds a loop iterator.
+# ``with X.while_(cond):`` builds a while loop.
+# ``X.range_(start, stop)`` constructs a loop descriptor.
+# ``X.break_()`` emits a loop exit.
+# ``X.continue_()`` emits a loop continuation.
+
+
+def if_(condition: Any, *, span: _Span = None) -> _frame.IfFrame:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.if_`."""
+ return _base.at_(span, _native.If(condition))
+
+
+def then_(*, span: _Span = None) -> _frame.ThenFrame:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.then_`."""
+ return _base.at_(span, _native.Then())
+
+
+def else_(*, span: _Span = None) -> _frame.ElseFrame:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.else_`."""
+ return _base.at_(span, _native.Else())
+
+
+def for_(
+ iterable: Any, *, names: str | Sequence[str] | None = None, span: _Span = None
+) -> NoReturn:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.for_`.
+
+ Relax rejects source loops; use supported functional control flow.
+ """
+ raise TypeError("Relax does not support imperative for loops")
+
+
+def while_(condition: Any, *, span: _Span = None) -> NoReturn:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.while_`."""
+ raise TypeError("Relax does not support imperative while loops")
+
+
+def range_(*args: Any, annotations: dict[str, Any] | None = None) -> NoReturn:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.range_`."""
+ raise TypeError("Relax does not support imperative for loops")
+
+
+def break_(*, span: _Span = None) -> NoReturn:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.break_`.
+
+ Relax rejects loop-control statements.
+ """
+ raise TypeError("Relax does not support break")
+
+
+def continue_(*, span: _Span = None) -> NoReturn:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.continue_`.
+
+ Relax rejects loop-control statements.
+ """
+ raise TypeError("Relax does not support continue")
+
+
+# ------------
+# Section: operator overloading
+# ------------
+#
+# ``X.if_then_else_(c, a, b)`` selects an expression.
+# ``X.and_(a, b)`` constructs conjunction.
+# ``X.or_(a, b)`` constructs disjunction.
+# ``X.not_(a)`` negates a condition.
+# ``X.lt_(a, b)`` lowers ``a < b``.
+# ``X.le_(a, b)`` lowers ``a <= b``.
+# ``X.gt_(a, b)`` lowers ``a > b``.
+# ``X.ge_(a, b)`` lowers ``a >= b``.
+# ``X.eq_(a, b)`` lowers ``a == b``.
+# ``X.ne_(a, b)`` lowers ``a != b``.
+
+
+def if_then_else_(condition: Any, true_value: Any, false_value: Any) -> Any:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.if_then_else_`."""
+ if isinstance(condition, _ffi.ObjectConvertible):
+ condition = condition.asobject()
+ if not isinstance(condition, _ir.Expr):
+ return true_value if condition else false_value
+ true_value = (
+ true_value.asobject() if isinstance(true_value, _ffi.ObjectConvertible) else true_value
+ )
+ false_value = (
+ false_value.asobject() if isinstance(false_value, _ffi.ObjectConvertible) else false_value
+ )
+ if _ir.is_prim_expr(condition) and all(
+ _ir.is_prim_expr(value)
+ if isinstance(value, _ir.Expr)
+ else isinstance(value, _numbers.Number)
+ for value in (true_value, false_value)
+ ):
+ return _tir.if_then_else(condition, true_value, false_value)
+ return _relax.If(condition, _builder._value(true_value), _builder._value(false_value))
+
+
+def and_(*values: Any) -> Any:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.and_`."""
+ return _builder.logical_and(*values)
+
+
+def or_(*values: Any) -> Any:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.or_`."""
+ return _builder.logical_or(*values)
+
+
+def not_(value: Any) -> Any:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.not_`."""
+ return _builder.logical_not(value)
+
+
+def lt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.lt_`."""
+ if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
+ lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
+ rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
+ return _base.at_(span, _relax.op.less(lhs, rhs))
+ return _prim_ffi._OpLT(lhs, rhs, _base.source_span(span))
+
+
+def le_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.le_`."""
+ if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
+ lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
+ rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
+ return _base.at_(span, _relax.op.less_equal(lhs, rhs))
+ return _prim_ffi._OpLE(lhs, rhs, _base.source_span(span))
+
+
+def gt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.gt_`."""
+ if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
+ lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
+ rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
+ return _base.at_(span, _relax.op.greater(lhs, rhs))
+ return _prim_ffi._OpGT(lhs, rhs, _base.source_span(span))
+
+
+def ge_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.ge_`."""
+ if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
+ lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
+ rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
+ return _base.at_(span, _relax.op.greater_equal(lhs, rhs))
+ return _prim_ffi._OpGE(lhs, rhs, _base.source_span(span))
+
+
+def eq_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.eq_`."""
+ if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
+ lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
+ rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
+ return _base.at_(span, _relax.op.equal(lhs, rhs))
+ return _prim_ffi._OpEQ(lhs, rhs, _base.source_span(span))
+
+
+def ne_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.ne_`."""
+ if any(isinstance(value, _ir.Expr) and not _ir.is_prim_expr(value) for value in (lhs, rhs)):
+ lhs = _relax.const(lhs) if isinstance(lhs, _numbers.Number) else lhs
+ rhs = _relax.const(rhs) if isinstance(rhs, _numbers.Number) else rhs
+ return _base.at_(span, _relax.op.not_equal(lhs, rhs))
+ return _prim_ffi._OpNE(lhs, rhs, _base.source_span(span))
+
+
+# ------------
+# Section: context lookup and resolution
+# ------------
+#
+# ``X.resolve_global_info_(key)`` resolves module metadata.
+# ``X.resolve_type_var_("n")`` resolves a symbolic dimension.
+# ``X.call_global_var_(f, args)`` calls a module function.
+
+
+def resolve_global_info_(content: Any) -> Any:
+ """Resolve a module-owned selector or retain a concrete object.
+
+ Parameters
+ ----------
+ content : str or Any
+ Original global-info selector or concrete value. "mesh[0]" indexes a named list;
+ "cuda:1" selects the second CUDA vdevice; "vdevice:0" selects by absolute index.
+ A trailing memory-scope suffix is accepted without changing device selection.
+
+ Returns
+ -------
+ Any
+ The exact registered global-info object, or the unchanged non-string input.
+
+ Notes
+ -----
+ String lookup requires the nearest active native module frame and creates no metadata.
+ Missing context, malformed selectors or unmatched devices raise ValueError; missing map
+ entries or out-of-range indices propagate KeyError/IndexError. Non-string values require
+ no frame. No source span is attached to an existing metadata object.
+
+ .. code:: python
+
+ # Source
+ R.Tensor((n,), "float32", vdevice="cuda:0")
+ # Generated builder
+ R.Tensor((n,), "float32", vdevice=R.resolve_global_info_("cuda:0"))
+ """
+ if not isinstance(content, str):
+ return content
+ if not _IRBuilder.is_in_scope():
+ raise ValueError("Global-info lookup requires an enclosing module frame")
+ for frame in reversed(_IRBuilder.current().frames):
+ if isinstance(frame, _IRModuleFrame):
+ break
+ else:
+ raise ValueError("Global-info lookup requires an enclosing module frame")
+ match = _re.fullmatch(r"([^\[\]]+)\[(\d+)\]", content)
+ if match:
+ name, index = match.groups()
+ return frame.global_infos[name][int(index)]
+ selector = _re.fullmatch(r"([^:\[\]]+)(?::(\d+)(?::([^:]+))?)?", content)
+ if selector is None:
+ raise ValueError(f"Invalid global-info reference: {content!r}")
+ target, index, _scope = selector.groups()
+ ordinal = int(index) if index is not None else 0
+ devices = frame.global_infos.get("vdevice", ())
+ if target == "vdevice":
+ return devices[ordinal]
+ for device in devices:
+ if device.target.kind.name == target:
+ if ordinal == 0:
+ return device
+ ordinal -= 1
+ raise ValueError(f"Global-info device reference was not found: {content!r}")
+
+
+def resolve_type_var_(
+ name: str,
+ dtype: str | _ir.Type | _ir.Var | None = None,
+ *,
+ value: _ir.Var | None = None,
+ span: _Span = None,
+) -> _ir.Var:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.resolve_type_var_`."""
+ return _base._current_function_frame().resolve_type_var(name, dtype, value=value, span=span)
+
+
+def call_global_var_(function: _ir.GlobalVar, args: Sequence[Any]) -> _ir.Expr:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.call_global_var_`."""
+ return _relax.Call(function, [_relax.utils.convert_to_expr(value) for value in args])
+
+
+# ------------
+# Section: special protocol
+# ------------
+#
+# Syntax markers live in tvm.script.parser.protocol_registry.
+# ``with X.function_(...) as fn:`` builds a function frame.
+# ``a = X.arg("a", ty)`` declares a parameter.
+# ``X.func_name("main")`` sets the function name.
+# ``X.func_ret_type(ty)`` sets the result annotation.
+# ``X.check_well_formed_(result)`` validates completed IR.
+
+
+def function_(
+ pure: bool = True,
+ private: bool = False,
+ *,
+ decl: bool = False,
+ local: bool = False,
+ reference: _ir.Var | None = None,
+ span: _Span = None,
+) -> _frame.FunctionFrame:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.function_`.
+
+ Public pure/private options pass to native purity and visibility controls.
+ Local function bodies require their previously declared reference.
+ """
+ if decl:
+ return _base.at_(span, _ffi_api.DeclFunction(pure, private, local))
+ if local:
+ if reference is None:
+ raise ValueError("A local function requires its declared reference")
+ return _base.at_(span, _ffi_api.LocalFunction(pure, reference))
+ return _base.at_(span, _native.function(pure, private))
+
+
+def arg(name: str, ty: Any, *, span: _Span = None) -> _ir.Var:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.arg`."""
+ if not isinstance(ty, _ir.Var):
+ ty = _builder._type(ty)
+ if isinstance(ty, _ir.PrimType) or _ir.is_prim_var(ty):
+ ty = resolve_type_var_(name, ty, span=span)
+ if isinstance(ty, _ir.Var):
+ return _ffi_api.ArgVar(name, ty)
+ return _base.at_(span, _native.arg(name, _builder._type(ty)))
+
+
+def func_name(name: str) -> None:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.func_name`."""
+ return _native.func_name(name)
+
+
+def func_ret_type(annotation: Any, *, span: _Span = None) -> None:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.func_ret_type`."""
+ return _native.func_ret_type(_builder._type(_base._return_annotation(annotation)))
+
+
+def check_well_formed_(function: _relax.Function) -> None:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.check_well_formed_`.
+
+ Validate the completed Relax function. Whole-module cross-function
+ validation is registered separately with the root coordinator.
+ """
+ from tvm import s_tir
+
+ message = (
+ "Program is not well-formed. If this is deliberate, set "
+ "check_well_formed=False in the top-level decorator."
+ )
+ if not _relax.analysis.check_well_formed(function):
+ raise ValueError(message)
+ try:
+ s_tir.analysis.verify_well_formed(_ir.IRModule.from_expr(function))
+ except Exception as error:
+ raise ValueError(f"{message}\n{error}") from error
+
+
+def _check_module_well_formed(module: _ir.IRModule) -> None:
+ """Own Relax whole-module validation, including cross-function references."""
+ if not _relax.analysis.check_well_formed(module):
+ raise ValueError(
+ "Program is not well-formed. If this is deliberate, set "
+ "check_well_formed=False in the top-level decorator."
+ )
+
+
+# ------------
+# Section: binding
+# ------------
+#
+# ``a = X.bind_(value, name="a")`` binds a source name.
+# ``X.decl_mutable_cell_(value, ty=ty)`` declares mutable storage.
+# ``X.set_mutable_cell_(cell, value)`` updates mutable storage.
+# ``a, b = X.unpack(value)`` destructures a binding.
+
+
+def bind_(
+ value: Any = _base.MISSING,
+ *,
+ ty: Any = None,
+ name: str | None = None,
+ span: _Span = None,
+ value_span: _Span = None,
+ name_span: _Span = None,
+ frame_value: bool = False,
+) -> Any:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.bind_`.
+
+ Relax values emit a binding or MatchCast; value_span belongs to the actual
+ RHS and span to the binding target. Metadata passes through unchanged.
+ """
+ name_span = _base.source_span(span if name_span is None else name_span)
+ 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=_base.source_span(span),
+ name_span=name_span,
+ frame_value=True,
+ )
+ elif isinstance(value, _ir.Var):
+ if name is not None:
+ _IRBuilder.name(name, value)
+ _base.at_(name_span if name_span is not None else span, value)
+ return value
+ if value is _base.MISSING:
+ raise ValueError("Relax bindings require an initializer")
+ ty = None if ty is None else _builder._type(ty)
+ if isinstance(value, _base.AlreadyEmitted):
+ return _base.at_(value_span, value)
+ value = _builder._value(value, ty)
+ if isinstance(value, _relax.MatchCast):
+ _base.at_(value_span, value.value)
+ 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.EmitMatchCastWithSpan(
+ value.value, value.ty, name_span, _base.source_span(span)
+ )
+ elif isinstance(value, _relax.Expr):
+ _base.at_(value_span, value)
+ result = _ffi_api.EmitWithSpan(value, ty, name_span, _base.source_span(span))
+ else:
+ return value
+ if name is not None:
+ _IRBuilder.name(name, result)
+ return _base.at_(name_span if name_span is not None else span, result)
+
+
+def decl_mutable_cell_(
+ value: Any = _base.MISSING,
+ *,
+ ty: Any = None,
+ name: str | None = None,
+ span: _Span = None,
+ name_span: _Span = None,
+) -> NoReturn:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.decl_mutable_cell_`.
+
+ Relax has immutable bindings and rejects mutable storage declarations.
+ """
+ raise TypeError("Relax does not support mutable local storage")
+
+
+def set_mutable_cell_(target: Any, value: Any, *, span: _Span = None) -> NoReturn:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.set_mutable_cell_`.
+
+ Relax has immutable bindings and rejects mutable storage updates.
+ """
+ raise TypeError("Relax does not support mutable local storage")
+
+
+def unpack(value: Any) -> Any:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.unpack`."""
+ 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
+
+
+# ------------
+# Section: statement
+# ------------
+#
+# ``X.emit_(value)`` emits an expression statement.
+# ``X.return_(value)`` emits a function return.
+# ``X.setitem_(a, i, value)`` emits an indexed store.
+# ``X.setattr_(a, "field", value)`` emits an attribute store.
+# ``X.assert_(condition, message)`` emits an assertion.
+
+
+def emit_(value: Any, *, span: _Span = None) -> None:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.emit_`.
+
+ Only void expressions may be discarded. Receipts and None emit nothing;
+ non-void expressions raise ValueError and unsupported host values raise TypeError.
+ """
+ if isinstance(value, _base.AlreadyEmitted):
+ _base.at_(span, value)
+ return None
+ if value is None:
+ return
+ if not isinstance(value, _relax.Expr):
+ raise TypeError(f"Unsupported expression statement value: {type(value).__name__}")
+ result = bind_(_base.at_(span, 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: Any = None, *, span: _Span = None) -> None:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.return_`."""
+ if value is None:
+ value = _relax.Tuple([])
+ # Normalization may emit bindings, but an existing result keeps its own span.
+ _base.with_at_group_(span, lambda: _native.func_ret_value(_builder._value(value)))
+
+
+def setitem_(target: Any, key: Any, value: Any, *, span: _Span = None) -> NoReturn:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.setitem_`."""
+ raise TypeError("Relax does not support indexed assignment")
+
+
+def setattr_(target: Any, name: str, value: Any, *, span: _Span = None) -> NoReturn:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.setattr_`."""
+ raise TypeError("Relax does not support attribute assignment")
+
+
+def assert_(
+ condition: Any,
+ message: str | tuple[str, Sequence[Any]] | Sequence[Any] = "",
+ *,
+ span: _Span = None,
+) -> None:
+ """Implements :func:`tvm.script.ir_builder.ir.parser_protocol.assert_`."""
+ if not isinstance(message, _python.str):
+ raise TypeError("An assertion message must be construction-time text")
+ emit_(_base.at_(span, _native.assert_op(condition, format=message)), span=span)
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/transform/transform.py b/python/tvm/relax/transform/transform.py
index e7056d393c..c954fa09b4 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -335,7 +335,7 @@ def LazyGetInput() -> tvm.ir.transform.Pass:
...
@R.function
- def after(fget_param: R.Callable([R.Prim('int64'), R.Any], R.Any)):
+ def after(fget_param: R.Callable([T.int64, R.Any], R.Any)):
A_untyped = fget_param(0, R.str('A'))
A = R.match_cast(A_untyped, R.Tensor([16,32], "float32")
...
@@ -373,7 +373,7 @@ def LazySetOutput() -> tvm.ir.transform.Pass:
return (A, B)
@R.function
- def after(args, fset_param: R.Callable([R.Prim('int64'), R.Any])):
+ def after(args, fset_param: R.Callable([T.int64, R.Any])):
...
fset_param(0, A)
...
diff --git a/python/tvm/s_tir/tensor_intrin/cuda.py b/python/tvm/s_tir/tensor_intrin/cuda.py
index fb28a0c359..a755c4c9da 100644
--- a/python/tvm/s_tir/tensor_intrin/cuda.py
+++ b/python/tvm/s_tir/tensor_intrin/cuda.py
@@ -177,9 +177,9 @@ def get_ldmatrix_intrin(
v0, v1 = T.axis.remap("SS", [ax0, ax1])
T.reads(shared[v0, v1])
- thread_id, local_id = T.meta_var(index_map(v0, v1))
- T.writes(warp[thread_id, local_id])
- warp[thread_id, local_id] = shared[v0, v1]
+ warp_indices = index_map(v0, v1)
+ T.writes(warp[warp_indices[0], warp_indices[1]])
+ warp[warp_indices[0], warp_indices[1]] = shared[v0, v1]
@T.prim_func(s_tir=True)
def ldmatrix_impl(warp_handle: T.handle, shared_handle: T.handle) -> None:
@@ -376,23 +376,23 @@ def get_mma_intrin(
for i, j, k in T.grid(M_DIM, N_DIM, k_dim):
with T.sblock("C"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
- a_row_ind, a_col_ind = T.meta_var(swap_if_flag(vi, vk, a_transposed))
- b_row_ind, b_col_ind = T.meta_var(swap_if_flag(vk, vj, b_transposed))
+ a_indices = swap_if_flag(vi, vk, a_transposed)
+ b_indices = swap_if_flag(vk, vj, b_transposed)
- thread_id_C, local_id_C = T.meta_var(index_map_C(vi, vj))
- thread_id_A, local_id_A = T.meta_var(index_map_A(a_row_ind, a_col_ind))
- thread_id_B, local_id_B = T.meta_var(index_map_B(b_row_ind, b_col_ind))
+ c_warp_indices = index_map_C(vi, vj)
+ a_warp_indices = index_map_A(a_indices[0], a_indices[1])
+ b_warp_indices = index_map_B(b_indices[0], b_indices[1])
T.reads(
- C[thread_id_C, local_id_C],
- A[thread_id_A, local_id_A],
- B[thread_id_B, local_id_B],
+ C[c_warp_indices[0], c_warp_indices[1]],
+ A[a_warp_indices[0], a_warp_indices[1]],
+ B[b_warp_indices[0], b_warp_indices[1]],
)
- T.writes(C[thread_id_C, local_id_C])
+ T.writes(C[c_warp_indices[0], c_warp_indices[1]])
- C[thread_id_C, local_id_C] += cast_to_out_dtype(
- A[thread_id_A, local_id_A]
- ) * cast_to_out_dtype(B[thread_id_B, local_id_B])
+ C[c_warp_indices[0], c_warp_indices[1]] += cast_to_out_dtype(
+ A[a_warp_indices[0], a_warp_indices[1]]
+ ) * cast_to_out_dtype(B[b_warp_indices[0], b_warp_indices[1]])
@T.prim_func(s_tir=True)
def mma_sync_impl(a: T.handle, b: T.handle, c: T.handle) -> None:
@@ -564,10 +564,10 @@ def get_mma_fill_intrin(dtype, local_size):
for i0, i1 in T.grid(M_DIM, N_DIM):
with T.sblock("C_warp"):
i, j = T.axis.remap("SS", [i0, i1])
- thread_id, local_id = T.meta_var(index_map(i, j))
+ warp_indices = index_map(i, j)
T.reads()
- T.writes(C_warp[thread_id, local_id])
- C_warp[thread_id, local_id] = zero
+ T.writes(C_warp[warp_indices[0], warp_indices[1]])
+ C_warp[warp_indices[0], warp_indices[1]] = zero
@T.prim_func(s_tir=True)
def mma_fill_impl(a: T.handle) -> None:
@@ -613,10 +613,10 @@ def get_mma_store_intrin(dtype, local_size, scope="global", use_mma_store_intrin
for i0, i1 in T.grid(M_DIM, N_DIM):
with T.sblock("C_warp"):
v0, v1 = T.axis.remap("SS", [i0, i1])
- thread_id, local_id = T.meta_var(index_map(v0, v1))
- T.reads(C_warp[thread_id, local_id])
+ warp_indices = index_map(v0, v1)
+ T.reads(C_warp[warp_indices[0], warp_indices[1]])
T.writes(C[v0, v1])
- C[v0, v1] = C_warp[thread_id, local_id]
+ C[v0, v1] = C_warp[warp_indices[0], warp_indices[1]]
if use_mma_store_intrinic:
@@ -669,8 +669,8 @@ def get_mma_store_intrin(dtype, local_size, scope="global", use_mma_store_intrin
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
for local_id in T.serial(local_size):
- row, col = T.meta_var(index_map_rev(tx, local_id))
- C[row, col] = C_warp[tx, local_id]
+ matrix_indices = index_map_rev(tx, local_id)
+ C[matrix_indices[0], matrix_indices[1]] = C_warp[tx, local_id]
return mma_store_desc, mma_store_impl
@@ -1081,9 +1081,9 @@ def get_wmma_sync_intrin(
for i, j, k in T.grid(m_dim, n_dim, k_dim):
with T.sblock(""):
vii, vjj, vkk = T.axis.remap("SSR", [i, j, k])
- B_index_0, B_index_1 = T.meta_var(maybe_swap(vkk, vjj))
+ b_indices = maybe_swap(vkk, vjj)
C[vii, vjj] = C[vii, vjj] + maybe_cast(A[vii, vkk]) * maybe_cast(
- B[B_index_0, B_index_1]
+ B[b_indices[0], b_indices[1]]
)
@T.prim_func(s_tir=True)
@@ -1633,9 +1633,9 @@ def get_mma_sync_intrin(
for i, j, k in T.grid(m_dim, n_dim, k_dim):
with T.sblock("m16n8k8_sync"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
- B_index_0, B_index_1 = T.meta_var(maybe_swap(vk, vj))
+ b_indices = maybe_swap(vk, vj)
C[vi, vj] = C[vi, vj] + maybe_cast(A[vi, vk]) * maybe_cast(
- B[B_index_0, B_index_1]
+ B[b_indices[0], b_indices[1]]
)
@T.prim_func(s_tir=True)
diff --git a/python/tvm/s_tir/tensor_intrin/dot_product_common.py b/python/tvm/s_tir/tensor_intrin/dot_product_common.py
index 7272477406..74b1acebf0 100644
--- a/python/tvm/s_tir/tensor_intrin/dot_product_common.py
+++ b/python/tvm/s_tir/tensor_intrin/dot_product_common.py
@@ -56,7 +56,7 @@ def get_dp4a_intrin(dtype_a, dtype_b, dtype_c):
"__dp4a",
A.vload([0], vec_type_a),
B.vload([0], vec_type_b),
- T.uint32(0) if dtype_c == "uint32" else T.int32(0),
+ T.uint32(0) if T.constexpr(dtype_c == "uint32") else T.int32(0),
dtype=dtype_c,
)
diff --git a/python/tvm/s_tir/tensor_intrin/metal.py b/python/tvm/s_tir/tensor_intrin/metal.py
index 1750338044..8c63e9b46f 100644
--- a/python/tvm/s_tir/tensor_intrin/metal.py
+++ b/python/tvm/s_tir/tensor_intrin/metal.py
@@ -93,7 +93,7 @@ def get_simdgroup_load_intrin(
for i, j in T.grid(col, row):
with T.sblock("load"):
vii, vjj = T.axis.remap("SS", [i, j])
- if transpose_matrix:
+ if T.constexpr(transpose_matrix):
# C[vii, vjj] = A[vjj, vii]
C[vjj, vii] = A[vii, vjj]
else:
@@ -157,7 +157,7 @@ def get_simdgroup_store_intrin(
for i, j in T.grid(col, row):
with T.sblock("store"):
vii, vjj = T.axis.remap("SS", [i, j])
- if transpose_matrix:
+ if T.constexpr(transpose_matrix):
C[vjj, vii] = A[vii, vjj]
else:
C[vii, vjj] = A[vii, vjj]
diff --git a/python/tvm/s_tir/tensor_intrin/rocm.py b/python/tvm/s_tir/tensor_intrin/rocm.py
index 29749dd443..2646c0f497 100644
--- a/python/tvm/s_tir/tensor_intrin/rocm.py
+++ b/python/tvm/s_tir/tensor_intrin/rocm.py
@@ -131,10 +131,10 @@ def get_mma_fill_intrin(dtype, local_size):
for i0, i1 in T.grid(M_DIM, N_DIM):
with T.sblock("C_warp"):
i, j = T.axis.remap("SS", [i0, i1])
- thread_id, local_id = T.meta_var(index_map(i, j))
+ warp_indices = index_map(i, j)
T.reads()
- T.writes(C_warp[thread_id, local_id])
- C_warp[thread_id, local_id] = zero
+ T.writes(C_warp[warp_indices[0], warp_indices[1]])
+ C_warp[warp_indices[0], warp_indices[1]] = zero
@T.prim_func(s_tir=True)
def mma_fill_impl(a: T.handle) -> None:
@@ -221,9 +221,9 @@ def get_mfma_load_intrin(
v0, v1 = T.axis.remap("SS", [ax0, ax1])
T.reads(memory[v0, v1])
- thread_id, local_id = T.meta_var(index_map(v0, v1))
- T.writes(reg[thread_id, local_id])
- reg[thread_id, local_id] = memory[v0, v1]
+ warp_indices = index_map(v0, v1)
+ T.writes(reg[warp_indices[0], warp_indices[1]])
+ reg[warp_indices[0], warp_indices[1]] = memory[v0, v1]
@T.prim_func(s_tir=True)
def mfma_load_impl(reg_handle: T.handle, memory_handle: T.handle) -> None:
@@ -248,9 +248,9 @@ def get_mfma_load_intrin(
T.writes(reg[0:WARP_SIZE, 0:local_size])
tx = T.env_thread("threadIdx.x")
for local_id in T.serial(0, local_size):
- row, col = T.meta_var(reverse_index_map(tx, local_id))
+ matrix_indices = reverse_index_map(tx, local_id)
T.launch_thread(tx, WARP_SIZE)
- reg[tx, local_id] = memory[row, col]
+ reg[tx, local_id] = memory[matrix_indices[0], matrix_indices[1]]
return mfma_load_desc, mfma_load_impl
@@ -302,22 +302,22 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
for i, j, k in T.grid(M_DIM, N_DIM, k_dim):
with T.sblock("C"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
- b_row_ind, b_col_ind = T.meta_var(maybe_swap(vk, vj))
+ b_indices = maybe_swap(vk, vj)
- thread_id_C, local_id_C = T.meta_var(index_map_C(vi, vj))
- thread_id_A, local_id_A = T.meta_var(index_map_A(vi, vk))
- thread_id_B, local_id_B = T.meta_var(index_map_B(b_row_ind, b_col_ind))
+ c_warp_indices = index_map_C(vi, vj)
+ a_warp_indices = index_map_A(vi, vk)
+ b_warp_indices = index_map_B(b_indices[0], b_indices[1])
T.reads(
- C[thread_id_C, local_id_C],
- A[thread_id_A, local_id_A],
- B[thread_id_B, local_id_B],
+ C[c_warp_indices[0], c_warp_indices[1]],
+ A[a_warp_indices[0], a_warp_indices[1]],
+ B[b_warp_indices[0], b_warp_indices[1]],
)
- T.writes(C[thread_id_C, local_id_C])
+ T.writes(C[c_warp_indices[0], c_warp_indices[1]])
- C[thread_id_C, local_id_C] += maybe_cast(
- A[thread_id_A, local_id_A]
- ) * maybe_cast(B[thread_id_B, local_id_B])
+ C[c_warp_indices[0], c_warp_indices[1]] += maybe_cast(
+ A[a_warp_indices[0], a_warp_indices[1]]
+ ) * maybe_cast(B[b_warp_indices[0], b_warp_indices[1]])
@T.prim_func(s_tir=True)
def mfma_sync_impl_float(a: T.handle, b: T.handle, c: T.handle) -> None:
@@ -336,8 +336,8 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
T.launch_thread(tx, WARP_SIZE)
C[tx, T.ramp(0, 1, local_size_out)] = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id(mfma_intrin),
- A[tx, T.ramp(0, 1, local_size) if local_size > 1 else 0],
- B[tx, T.ramp(0, 1, local_size) if local_size > 1 else 0],
+ A[tx, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
+ B[tx, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
C[tx, T.ramp(0, 1, local_size_out)],
T.int32(0),
T.int32(0),
@@ -366,12 +366,12 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
T.call_intrin(
"int32",
"tirx.reinterpret",
- A[tx, T.ramp(0, 1, local_size) if local_size > 1 else 0],
+ A[tx, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
),
T.call_intrin(
"int32",
"tirx.reinterpret",
- B[tx, T.ramp(0, 1, local_size) if local_size > 1 else 0],
+ B[tx, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
),
C[tx, T.ramp(0, 1, local_size_out)],
T.int32(0),
@@ -401,10 +401,10 @@ def get_mfma_store_intrin(local_size=4, dtype="float32", scope="global"):
for i0, i1 in T.grid(M_DIM, N_DIM):
with T.sblock("C_warp"):
v0, v1 = T.axis.remap("SS", [i0, i1])
- thread_id, local_id = T.meta_var(index_map(v0, v1))
- T.reads(C_warp[thread_id, local_id])
+ warp_indices = index_map(v0, v1)
+ T.reads(C_warp[warp_indices[0], warp_indices[1]])
T.writes(C[v0, v1])
- C[v0, v1] = C_warp[thread_id, local_id]
+ C[v0, v1] = C_warp[warp_indices[0], warp_indices[1]]
@T.prim_func(s_tir=True)
def mfma_store_impl(a: T.handle, c: T.handle) -> None:
diff --git a/python/tvm/s_tir/tensor_intrin/x86.py b/python/tvm/s_tir/tensor_intrin/x86.py
index 2fad505104..1db0cc51e3 100644
--- a/python/tvm/s_tir/tensor_intrin/x86.py
+++ b/python/tvm/s_tir/tensor_intrin/x86.py
@@ -51,12 +51,12 @@ def dot_product_16x4_u8i8i32_vnni(
T.reads(C[0:16], A[0:4], B[0:16, 0:4])
T.writes(C[0:16])
- A_u8x4 = A.vload([0], "uint8x4")
- A_i32 = T.reinterpret(A_u8x4, dtype="int32")
+ A_u8x4: T.uint8x4 = A.vload([0], "uint8x4")
+ A_i32: T.int32 = T.reinterpret(A_u8x4, dtype="int32")
- B_i8x64 = B.vload([0, 0], dtype="int8x64")
- B_i32x16 = T.reinterpret(B_i8x64, dtype="int32x16")
- C_i32x16 = C.vload([0], dtype="int32x16")
+ B_i8x64: T.int8x64 = B.vload([0, 0], dtype="int8x64")
+ B_i32x16: T.int32x16 = T.reinterpret(B_i8x64, dtype="int32x16")
+ C_i32x16: T.int32x16 = C.vload([0], dtype="int32x16")
C[T.ramp(T.int32(0), 1, 16)] = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id("llvm.x86.avx512.vpdpbusd.512"),
diff --git a/python/tvm/script/__init__.py b/python/tvm/script/__init__.py
index 940739df07..a71ce59180 100644
--- a/python/tvm/script/__init__.py
+++ b/python/tvm/script/__init__.py
@@ -14,93 +14,84 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""
-TVMScript public namespace.
-
-Dialect resolution mechanism
-----------------------------
-
-``tvm.script`` is a virtual namespace: dialect names like ``tirx`` and
-``relax`` are not bound as static attributes here. Instead:
-
-- ``register_dialect(name, module_path)`` writes an entry to
- ``_DIALECT_REGISTRY: dict[str, str]``. Each in-tree dialect's
- ``__init__.py`` calls this on import (e.g., ``tvm.tirx.__init__.py``
- calls ``tvm.script.register_dialect("tirx", "tvm.tirx.script")``).
- Out-of-tree dialects can register themselves the same way.
-
-- ``__getattr__(name)`` (PEP 562) fires on missing attribute access.
- If ``name`` is in ``_DIALECT_REGISTRY``, the listed module is imported
- 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.
-
-- For deep statement-form imports like
- ``from tvm.script.parser.tirx.entry import ObjectProxy``, 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.
-IR is foundational (script depends on ir) and is NOT a dialect; its
-script handlers live in the shared core, not via this registry.
-
-Bootstrap order
----------------
-
-``python/tvm/__init__.py`` imports ``tvm.script`` BEFORE importing any
-dialect package (``tvm.tirx``, ``tvm.relax``, …). This guarantees that
-``tvm.script.register_dialect`` is reachable the moment a dialect's own
-``__init__.py`` runs and calls it. The ``tvm.script`` module itself
-stays dialect-agnostic at load time (no dialect submodules are eagerly
-imported here), so there is no circular dependency.
+"""TVMScript public namespace and registered language variant exports.
+
+All source parsing uses :mod:`tvm.script.parser`; language variant packages expose source
+constructors and concrete builders from the same canonical implementation.
+The registry permits out-of-tree language variants to expose the same public spellings.
"""
import importlib
import importlib.util
import sys
+from collections.abc import Callable
from typing import Any
_DIALECT_REGISTRY: dict[str, str] = {}
+# Language variants register whole-module checks at import, independently of source decorators.
+# Callbacks own concrete eligibility; this list retains no source or construction state.
+_MODULE_VALIDATORS: list[Callable[[Any], None]] = []
-# 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`).
+# Subpackages of `tvm.script` whose per-language variant children are redirected to a
+# canonical package under `tvm.<language variant>.script`. An empty suffix names
+# the public language variant namespace itself.
_REDIRECTED_SUBPACKAGES = {
- "tvm.script.parser": "parser",
+ "tvm.script.parser": "",
"tvm.script.ir_builder": "builder",
}
+def register_module_validator(validator: Callable[[Any], None], *, prepend: bool = False) -> None:
+ """Register an opaque language variant check over a completed module.
+
+ Parameters
+ ----------
+ validator : Callable[[Any], None]
+ Whole-module validator. It owns concrete IR eligibility and raises on
+ failure; shared coordination forwards its exception unchanged.
+ prepend : bool, optional
+ Run this check before existing callbacks. False appends it instead.
+ Re-registering the same callback is a no-op and preserves its position.
+
+ Returns
+ -------
+ None
+ Registration changes no IR and returns no source-visible result.
+
+ Notes
+ -----
+ Register beside language variant initialization so captured and preexisting functions
+ are checked even when no corresponding source decorator appears. A callback
+ may lazily import its language variant implementation to preserve import order. The
+ registry owns static hooks only, never source scopes, frames or IR results.
+ """
+ if any(existing is validator for existing in _MODULE_VALIDATORS):
+ return
+ if prepend:
+ _MODULE_VALIDATORS.insert(0, validator)
+ else:
+ _MODULE_VALIDATORS.append(validator)
+
+
def register_dialect(name: str, module_path: str) -> None:
- """Register a dialect's script package path.
+ """Register a language variant's script package path.
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``.
+ resolve to the same script namespace and ``module_path + ".builder"``
+ respectively. Deep builder imports are handled by
+ ``_DialectRedirectFinder`` on ``sys.meta_path``.
This function is idempotent — re-registering the same name with the same
path is harmless.
- Each in-tree dialect calls this from its own ``__init__.py``::
+ Each in-tree language variant calls this from its own ``__init__.py``::
import tvm.script
tvm.script.register_dialect("tirx", "tvm.tirx.script")
- Out-of-tree dialects do the same in their own package init without
+ Out-of-tree language variants do the same in their own package init without
editing any in-tree file.
Parameters
@@ -108,10 +99,9 @@ def register_dialect(name: str, module_path: str) -> None:
name : str
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.
+ The full dotted module path of the language variant's script package, e.g.
+ ``"tvm.tirx.script"``. That package exposes its public construction
+ entry points directly and its imperative APIs in a ``builder`` submodule.
"""
_DIALECT_REGISTRY[name] = module_path
@@ -122,41 +112,41 @@ def _redirect_target(fullname: str) -> str | None:
Returns ``None`` if ``fullname`` is not a redirected name.
"""
if fullname.startswith("tvm.script."):
- # tvm.script.<dialect>[.subpath]
+ # tvm.script.<language variant>[.subpath]
rest = fullname[len("tvm.script.") :]
head, _, tail = rest.partition(".")
if head in _DIALECT_REGISTRY and "." not in head:
target = _DIALECT_REGISTRY[head]
return f"{target}.{tail}" if tail else target
- # tvm.script.parser.<dialect>[.subpath] / tvm.script.ir_builder.<dialect>[.subpath]
+ # Registered language variants may redirect parser and builder subpackages.
for prefix, sub in _REDIRECTED_SUBPACKAGES.items():
if fullname == prefix or not fullname.startswith(prefix + "."):
continue
rest = fullname[len(prefix) + 1 :]
head, _, tail = rest.partition(".")
if head in _DIALECT_REGISTRY:
- target = f"{_DIALECT_REGISTRY[head]}.{sub}"
+ target = _DIALECT_REGISTRY[head] + (f".{sub}" if sub else "")
return f"{target}.{tail}" if tail else target
return None
class _DialectRedirectFinder:
- """``sys.meta_path`` finder that redirects ``tvm.script.<dialect>`` import paths.
+ """``sys.meta_path`` finder that redirects ``tvm.script.<language variant>`` import paths.
PEP 562 ``__getattr__`` only handles one-level attribute lookups
(``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.tirx.tile
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
+ ``tvm.script.<language variant>`` (or ``tvm.script.parser.<language variant>``, etc.) and
+ that language variant is in ``_DIALECT_REGISTRY``, :meth:`find_spec` imports the
+ real target module (e.g. ``tvm.tirx.script.tile``) 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
+ registers it in ``sys.modules`` under the public alias and all subsequent
imports and attribute walks resolve without going through the redirect
again.
"""
@@ -222,7 +212,7 @@ def __getattr__(name: str) -> Any:
globals()["ir"] = ir_parser
return ir_parser
if name in ("from_source", "parse"):
- from .parser._core import parse # pylint: disable=import-outside-toplevel
+ from .parser import parse # pylint: disable=import-outside-toplevel
globals()["from_source"] = parse
globals()["parse"] = parse
@@ -230,9 +220,15 @@ def __getattr__(name: str) -> Any:
if name == "ir_module":
# ir_module lives in the IR parser at tvm.script.parser.ir; the IR
# layer is foundational, so we resolve it directly rather than via
- # the dialect registry.
+ # the language variant registry.
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
raise AttributeError(f"module 'tvm.script' has no attribute {name!r}")
+
+
+# Shared source entry points participate in the same opaque namespace boundary.
+from .parser import register_namespace as _register_namespace
+
+_register_namespace("script", sys.modules[__name__])
diff --git a/python/tvm/script/ir_builder/__init__.py b/python/tvm/script/ir_builder/__init__.py
index ee114fd1d8..2f1769459c 100644
--- a/python/tvm/script/ir_builder/__init__.py
+++ b/python/tvm/script/ir_builder/__init__.py
@@ -14,31 +14,16 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The ir_builder subpackage of TVMScript.
-
-Per-dialect builder submodules (``tvm.script.ir_builder.tirx``, etc.) are
-resolved lazily via :data:`tvm.script._DIALECT_REGISTRY`. When a dialect
-is accessed (e.g. ``tvm.script.ir_builder.tirx``), this subpackage's
-``__getattr__`` looks up the dialect in ``_DIALECT_REGISTRY`` and imports
-``<dialect_module_path>.builder`` (e.g. ``tvm.tirx.script.builder``),
-caching the result so subsequent accesses skip ``__getattr__``.
-
-The IR layer is foundational and is NOT registered as a dialect — its
-builder lives as a real submodule ``tvm.script.ir_builder.ir``.
-
-See :mod:`tvm.script` for a full description of the dialect resolution
-mechanism, including the ``_DialectRedirectFinder`` that handles
-deep statement-form imports.
-"""
+"""The ir_builder subpackage of TVMScript."""
import importlib
from typing import Any
-from .base import IRBuilder
+from .base import AlreadyEmitted, IRBuilder
def __getattr__(name: str) -> Any:
- # Lazy import to avoid loading tvm.script during dialect bootstrap.
+ # Lazy import to avoid loading tvm.script during language variant bootstrap.
from tvm.script import _DIALECT_REGISTRY # pylint: disable=import-outside-toplevel
if name in _DIALECT_REGISTRY:
diff --git a/python/tvm/script/ir_builder/base.py b/python/tvm/script/ir_builder/base.py
index 56e78b9072..815bee7be2 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -17,11 +17,14 @@
"""A generic IRBuilder across the TVM stack"""
from collections.abc import Callable
-from contextlib import contextmanager
-from typing import Any
+from contextlib import contextmanager, nullcontext
+from functools import wraps
+from typing import Any, Generic, TypeVar
from tvm_ffi import register_object as _register_object
+from tvm_ffi.dataclasses import MISSING
+from tvm import ir
from tvm.runtime import Object as _Object
from . import _ffi_api
@@ -31,39 +34,11 @@ from . import _ffi_api
class IRBuilderFrame(_Object):
"""A stack frame of the IRBuilder used to keep track of the current scope.
- Furthermore, the information stored in each stack frame can be useful for context-dependent
- IR construction.
-
- Examples
- --------
-
- The `T.match_buffer` below instead an element in the buffer map of `PrimFuncFrame`:
-
- .. code-block:: python
-
- from tvm.script.ir_builder import tirx as T
- from tvm.script.ir_builder import IRBuilder
-
- with IRBuilder() as builder:
- with T.prim_func(...): # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
- # to `builder`'s stack of frames
- buffer = T.match_buffer(...)
-
-
- The `T.match_buffer` below instead generates `MatchBufferRegion` in a TIR block:
-
- .. code-block:: python
-
- from tvm.script.ir_builder import tirx as T
- from tvm.script.ir_builder import IRBuilder
-
- with IRBuilder() as builder:
- with T.prim_func(...): # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
- # to `builder`'s stack of frames
- with T.sblock(...): # pushes a BlockFrame (subclass of IRBuilderFrame)
- # to `builder`'s stack of frames
- buffer = T.match_buffer(...)
-
+ A language variant supplies frame subclasses that retain the information
+ needed for context-dependent construction. Entering a frame pushes it onto
+ the active builder's stack; nested operations can inspect that stack to
+ find their enclosing scope. Normal exit finalizes the frame and runs its
+ registered callbacks.
"""
def __enter__(self) -> "IRBuilderFrame":
@@ -90,24 +65,22 @@ class IRBuilderFrame(_Object):
@_register_object("script.ir_builder.IRBuilder")
class IRBuilder(_Object):
- """A dialect-agnostic IRBuilder that constructs any IR of TVM.
+ """A shared construction context for any language variant's IR.
Examples
--------
- An idiomatic use of this class is to put this inside the with-scope,
- call dialect-specific methods accordingly. Upon exiting the scope.
+ Enter the builder before using a language variant's construction frames.
+ Each completed frame contributes to the result returned by ``get``.
.. code-block:: python
- from tvm.script.ir_builder import tirx as T
from tvm.script.ir_builder import IRBuilder
- with IRBuilder() as builder:
- with T.prim_func(...): # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
- # to `builder`'s stack of frames
- buffer = T.match_buffer(...)
-
- return builder.get() # returns the constructed IR, i.e. tirx.PrimFunc
+ def build(frame, emit_body):
+ with IRBuilder() as builder:
+ with frame:
+ emit_body()
+ return builder.get()
"""
def __init__(self) -> None:
@@ -185,7 +158,7 @@ class IRBuilder(_Object):
)
def _set_current_source_span(self, value):
- """Attach the active source span to an expression without one."""
+ """Compose the active source span onto the same supported IR node or frame."""
return _ffi_api.IRBuilderSetCurrentSourceSpan( # type: ignore[attr-defined] # pylint: disable=no-member
self, value
)
@@ -229,3 +202,211 @@ class IRBuilder(_Object):
"""
assert len(s) == len(vs)
return [IRBuilder.name(i, v) for i, v in zip(s, vs)]
+
+
+_T = TypeVar("_T")
+
+
+class SpanEntry:
+ """A materialized source range shared by generated builder operations.
+
+ Entries retain only fixed source metadata. Calling an entry attaches its
+ span to the same result; ``ctx(thunk)`` additionally supplies call provenance
+ during evaluation. ``ctx(thunk, attach_result=False)`` supplies only the
+ evaluation context, leaving result attachment to the binding operation.
+ Both compose the active caller context at invocation.
+ Builders accepting an explicit span normalize an entry with ``source_span``.
+ """
+
+ __slots__ = ("span",)
+
+ def __init__(self, span: ir.Span) -> None:
+ self.span = span
+
+ def __call__(self, value: _T) -> _T:
+ """Attach this range to the same value, receipt, or native frame."""
+ return at(self.span, value)
+
+ def ctx(self, thunk: Callable[[], _T], *, attach_result: bool = True) -> _T:
+ """Evaluate once under this range, optionally attaching it to the result.
+
+ Context is restored even on failure. Frames constructed during the call
+ retain their native construction span regardless of result attachment.
+ """
+ return with_at_group_(self.span, thunk, attach_result=attach_result)
+
+
+def source_span(
+ location: SpanEntry | ir.Span | tuple[str | ir.SourceName, int, int, int, int] | None,
+) -> ir.Span | None:
+ """Normalize an entry or materialize a range without retaining source-unit state."""
+ if isinstance(location, SpanEntry):
+ return location.span
+ if location is None or isinstance(location, ir.Span | ir.SequentialSpan):
+ return location
+ source_name, line, end_line, column, end_column = location
+ if isinstance(source_name, str):
+ source_name = ir.SourceName(source_name)
+ return ir.Span(source_name, line, end_line, column, end_column)
+
+
+class AlreadyEmitted(Generic[_T]):
+ """Hold the exact emitted value; location handling preserves this receipt."""
+
+ __slots__ = ("value",)
+
+ def __init__(self, value: _T) -> None:
+ self.value = value
+
+
+def at(
+ span: SpanEntry | ir.Span | tuple[str | ir.SourceName, int, int, int, int] | None, value: _T
+) -> _T:
+ """Attach source context to the same IR node, emission receipt, or frame.
+
+ Native mutation annotates the statement held by the builder itself. Keep
+ the original Python facade as well, including callable objects and frames.
+ Unsupported objects and ordinary Python values pass through unchanged.
+ """
+ if span is None or not IRBuilder.is_in_scope():
+ return value
+ target = value.value if isinstance(value, AlreadyEmitted) else value
+ targets = target if isinstance(target, list | tuple) else (target,)
+ for item in targets:
+ if isinstance(item, _Object):
+ _ffi_api.IRBuilderSetSourceSpan(IRBuilder.current(), item, source_span(span))
+ return value
+
+
+def require_defined(value, name):
+ """Report a source name whose designated region output was not produced."""
+ if value is MISSING:
+ raise NameError(f"name {name!r} is not defined")
+ return value
+
+
+def with_at_group_(
+ location: SpanEntry | ir.Span | tuple[str | ir.SourceName, int, int, int, int] | None,
+ thunk: Callable[[], _T],
+ *,
+ attach_result: bool = True,
+) -> _T:
+ """Evaluate once under a location, optionally attaching it to the same result."""
+ span = source_span(location)
+ context = (
+ IRBuilder.current().with_source_span(span)
+ if span is not None and IRBuilder.is_in_scope()
+ else nullcontext()
+ )
+ with context:
+ value = thunk()
+ return at(location, value) if attach_result else value
+
+
+at_ = at
+
+
+def _resolve_type_var(frame, ffi_resolver, name, dtype=None, *, value=None, span=None):
+ """Validate a native function resolver's inputs, leaving its map owned by C++."""
+ if not isinstance(name, str) or not name:
+ raise ValueError("A symbolic variable requires a nonempty string name")
+ if isinstance(dtype, ir.Var):
+ value, dtype = dtype, dtype.ty
+ if isinstance(dtype, str):
+ dtype = ir.PrimType(dtype)
+ if dtype is not None and not isinstance(dtype, ir.PrimType):
+ raise TypeError("A symbolic variable requires a primitive type")
+ if value is not None and not ir.is_prim_var(value):
+ raise TypeError("A symbolic binding requires a primitive Var")
+ return ffi_resolver(frame, name, dtype, value, source_span(span))
+
+
+def _current_function_frame():
+ """Find the nearest function for eager shared annotation constructors."""
+ if IRBuilder.is_in_scope():
+ for frame in reversed(IRBuilder.current().frames):
+ if callable(getattr(frame, "resolve_type_var", None)):
+ return frame
+ raise ValueError("Symbol resolution requires an active function frame")
+
+
+def wrap_expression_constructor(constructor, call_signature, policy, *, as_type=False):
+ """Adapt an eager constructor using parser-owned expression-string metadata."""
+ fields = policy.fields
+
+ def unresolved(value, nested=False):
+ if isinstance(value, str):
+ return nested or policy.scalar_strings
+ if isinstance(value, TypeVar):
+ return True
+ if isinstance(value, tuple | list):
+ return any(unresolved(item, True) for item in value)
+ return False
+
+ @wraps(constructor)
+ def invoke(*args, **kwargs):
+ bound = call_signature.bind(*args, **kwargs)
+ if IRBuilder.is_in_scope():
+ # typing.TypeVar is ordinary eager Python metadata. Resolve it
+ # here, never in the syntax-only transpiler.
+ def resolve(value):
+ if isinstance(value, TypeVar):
+ if value.__bound__ is not None or value.__constraints__:
+ raise TypeError("A symbolic TypeVar cannot have constraints or a bound")
+ return _current_function_frame().resolve_type_var(value.__name__)
+ if isinstance(value, tuple):
+ return tuple(resolve(item) for item in value)
+ if isinstance(value, list):
+ return [resolve(item) for item in value]
+ return value
+
+ for field in fields:
+ if field in bound.arguments:
+ bound.arguments[field] = resolve(bound.arguments[field])
+ if any(unresolved(bound.arguments[field]) for field in fields if field in bound.arguments):
+ if IRBuilder.is_in_scope():
+ raise TypeError(
+ "Builder expression arguments require concrete symbols, not strings"
+ )
+ return ir.Type.missing()
+ return constructor(*bound.args, **bound.kwargs)
+
+ result = invoke
+ if as_type:
+ # The class is an annotation surface, not an IR or proxy type.
+ # __new__ returns the concrete construction result (or MissingType).
+ result = type(
+ constructor.__name__,
+ (),
+ {
+ "__new__": lambda cls, *args, **kwargs: invoke(*args, **kwargs),
+ "__signature__": call_signature,
+ "__doc__": constructor.__doc__,
+ "__module__": constructor.__module__,
+ },
+ )
+ return result
+
+
+def _return_annotation(annotation):
+ """Evaluate a return annotation without introducing return-only symbols."""
+ if not callable(annotation) or isinstance(annotation, ir.Expr | ir.Type):
+ return annotation
+ frame = _current_function_frame()
+ declared = set(frame.type_var_map)
+ annotation = annotation()
+ introduced = set(frame.type_var_map) - declared
+ if introduced:
+ raise ValueError(f"Return annotation introduces unbound symbol {sorted(introduced)[0]!r}")
+ return annotation
+
+
+def annotation_value_(name, value):
+ """Adapt a real definition-context symbol using the native function map."""
+ if isinstance(value, TypeVar):
+ if value.__bound__ is not None or value.__constraints__:
+ raise TypeError("A symbolic TypeVar cannot have constraints or a bound")
+ return _current_function_frame().resolve_type_var(name)
+ if ir.is_prim_var(value):
+ return _current_function_frame().resolve_type_var(name, value=value)
+ return value
diff --git a/python/tvm/script/ir_builder/ir/__init__.py b/python/tvm/script/ir_builder/ir/__init__.py
index 86c6f1bf27..59ccf8e395 100644
--- a/python/tvm/script/ir_builder/ir/__init__.py
+++ b/python/tvm/script/ir_builder/ir/__init__.py
@@ -17,6 +17,11 @@
# under the License.
"""Package tvm.script.ir_builder.ir"""
+from tvm.ir import GenericConst, Range, StringImm, StringType
+from tvm.script.parser.protocol_registry import constexpr
+
+from ..base import IRBuilder, MISSING, annotation_value_, at_, with_at_group_, require_defined
+from .parser_protocol import check_well_formed_, module_member_
from .frame import IRModuleFrame
from .ir import (
decl_function,
@@ -26,9 +31,6 @@ from .ir import (
module_get_attr,
module_set_attr,
module_global_infos,
- lookup_vdevice,
lookup_name,
meta_var,
- vdevice,
- dummy_global_info,
)
diff --git a/python/tvm/script/ir_builder/ir/frame.py b/python/tvm/script/ir_builder/ir/frame.py
index b96c8c9366..346ec00c62 100644
--- a/python/tvm/script/ir_builder/ir/frame.py
+++ b/python/tvm/script/ir_builder/ir/frame.py
@@ -22,4 +22,14 @@ from ..base import IRBuilderFrame
@_register_object("script.ir_builder.IRModuleFrame")
-class IRModuleFrame(IRBuilderFrame): ...
+class IRModuleFrame(IRBuilderFrame):
+ def __getattr__(self, name):
+ """Expose a declared function through the native module's reference map.
+
+ Reflected native fields and Python methods resolve before this fallback.
+ Source aliases therefore retain this frame and its plain GlobalVars.
+ """
+ try:
+ return self.global_vars[name]
+ except KeyError:
+ raise AttributeError(f"IRModuleFrame has no attribute {name!r}") from None
diff --git a/python/tvm/script/ir_builder/ir/ir.py b/python/tvm/script/ir_builder/ir/ir.py
index e2810ce698..efe38fc5a7 100644
--- a/python/tvm/script/ir_builder/ir/ir.py
+++ b/python/tvm/script/ir_builder/ir/ir.py
@@ -16,48 +16,45 @@
# under the License.
"""Package tvm.script.ir_builder.ir.ir"""
-from typing import TYPE_CHECKING, Any, TypeVar
+import inspect
+from typing import TypeVar
from tvm.ir import BaseFunc, GlobalInfo, GlobalVar
from tvm.runtime import Object as tvm_Object
+from ..base import IRBuilder
from . import _ffi_api
from .frame import IRModuleFrame
-if TYPE_CHECKING:
- from tvm.relax import DummyGlobalInfo, VDevice
+T = TypeVar("T")
- T = TypeVar("T")
- def meta_var(value: T) -> T:
- """Mark a value for parser-time metaprogramming."""
- return value
+def meta_var(value: T) -> T:
+ """Return a Python metadata value without binding, naming or relocating it.
-else:
-
- class meta_var: # pylint: disable=invalid-name
- """A value used only for TVMScript parser-time metaprogramming.
-
- Assignments unwrap this object without emitting an IR binding. The
- shared wrapper is exposed as ``I.meta_var``; dialect namespaces may
- provide compatibility aliases to the same implementation.
- For Relax, this is the explicit opt-out from default primitive binding emission.
+ Parameters
+ ----------
+ value : T
+ Any host or IR object, including an unpackable sequence.
- Parameters
- ----------
- value : Any
- The parser-time value.
- """
+ Returns
+ -------
+ T
+ The exact input object. No frame is required, no IR is emitted, and
+ existing names and source locations are retained.
- def __init__(self, value: Any) -> None:
- self.value = value
+ .. code:: python
- def __iter__(self):
- return (meta_var(item) for item in self.value)
+ # Source and generated Python (the value retains its identity)
+ value = I.meta_var(existing_value)
+ a, b = I.meta_var((left, right))
+ """
+ return value
def ir_module() -> IRModuleFrame:
"""Start a ir_module frame.
+
Returns
-------
frame: IRModuleFrame
@@ -153,65 +150,63 @@ def module_set_attr(
def module_global_infos(global_infos: dict[str, list[GlobalInfo]]) -> None:
"""Specify the global infos of the ir_module frame.
+
Parameters
----------
global_infos: Dict[str, List[GlobalInfo]]
The module global infos.
"""
- return _ffi_api.ModuleGlobalInfos(global_infos) # type: ignore[attr-defined] # pylint: disable=no-member
-
-
-############################### GlobalInfo ###############################
-
-
-def dummy_global_info() -> "DummyGlobalInfo":
- """Create a dummy global info expression.
- Returns
- -------
- res : DummyGlobalInfo
- The result dummy global info.
- """
- from tvm.relax import DummyGlobalInfo # pylint: disable=import-outside-toplevel
-
- return DummyGlobalInfo() # type: ignore[attr-defined] # pylint: disable=no-member
-
-
-def vdevice(target=None, vdevice_id: int = 0, memory_scope: str = "global") -> "VDevice":
- """Create a virtual device global info.
- Parameters
- ----------
- target
- The target.
- vdevice_id: int
- The virtual device index.
- memory_scope: str
- The memory scope, default is "global"
-
- Returns
- -------
- res : VDevice
- The result virtual device.
- """
- from tvm.relax import VDevice # pylint: disable=import-outside-toplevel
-
- return VDevice(target, vdevice_id, memory_scope) # type: ignore[attr-defined] # pylint: disable=no-member
-
+ if IRBuilder.is_in_scope():
+ return _ffi_api.ModuleGlobalInfos(global_infos)
+ # Keep native argument validation even before Python has applied the module decorator.
+ _ffi_api.ModuleGlobalInfos(global_infos)
+ frame = inspect.currentframe().f_back
+ try:
+ if _is_class_frame(frame):
+ previous = frame.f_locals.get("__tvm_script_global_infos__")
+ if previous:
+ raise ValueError(f"Duplicate module global_infos, previous one is:\n{previous}")
+ frame.f_locals["__tvm_script_global_infos__"] = {
+ name: tuple(values) for name, values in global_infos.items()
+ }
+ finally:
+ del frame
+
+
+def _is_class_frame(frame):
+ return (
+ not frame.f_code.co_flags & inspect.CO_NEWLOCALS
+ and frame.f_locals is not frame.f_globals
+ and "__module__" in frame.f_locals
+ and frame.f_locals.get("__qualname__", "").split(".")[-1] == frame.f_code.co_name
+ )
-def lookup_vdevice(target_kind: str | None = None, device_index: int = -1) -> "VDevice":
- """Retrieve a virtual device from the globalinfo vdevice list.
- Parameters
- ----------
- target_kind: str
- The target device kind, for example 'llvm' or 'cuda'.
- device_index: int
- The virtual device index.
- Returns
- -------
- res : VDevice
- The result virtual device.
- """
- return _ffi_api.LookupVDevice(target_kind, device_index) # type: ignore[attr-defined] # pylint: disable=no-member
+def _class_global_infos():
+ # Only a still-executing class owns eager signature context. Never retain its frame,
+ # and do not fall through an inner class to an unrelated outer module declaration.
+ frame = inspect.currentframe().f_back
+ try:
+ while frame is not None:
+ if _is_class_frame(frame):
+ return frame.f_locals.get("__tvm_script_global_infos__", {})
+ frame = frame.f_back
+ finally:
+ del frame
+ return {}
+
+
+def lookup_global_info(name: str, index: int) -> GlobalInfo:
+ """Resolve a concrete global info in a builder or an executing module class."""
+ if IRBuilder.is_in_scope():
+ for frame in reversed(IRBuilder.current().frames):
+ if isinstance(frame, IRModuleFrame):
+ return frame.global_infos[name][index]
+ raise ValueError("The GlobalInfos in the IRModule is not defined.")
+ infos = _class_global_infos()
+ if not infos:
+ raise ValueError("The GlobalInfos in the IRModule is not defined.")
+ return infos[name][index]
def lookup_name(name: str) -> bool:
diff --git a/python/tvm/script/ir_builder/ir/parser_protocol.py b/python/tvm/script/ir_builder/ir/parser_protocol.py
new file mode 100644
index 0000000000..48c2f1594d
--- /dev/null
+++ b/python/tvm/script/ir_builder/ir/parser_protocol.py
@@ -0,0 +1,1490 @@
+# 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.
+"""Shared structural contract for source-to-builder rewriting.
+
+The source namespace selects language variant ``X``; shared module support uses ``I``.
+For example, ``target()[index()] = value()`` lowers to
+``X.setitem_(value=value(), target=target(), key=index())``. Python evaluates the
+keywords in that order, preserving the assignment's RHS-first semantics.
+
+Native frames own construction state. Language variant implementations live in
+``builder.parser_protocol``; shared contract stubs here raise NotImplementedError.
+Special syntax markers are documented in the special protocol section below.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+from tvm import ir as _ir
+
+from ..base import MISSING, AlreadyEmitted, IRBuilderFrame, SpanEntry
+
+_Span = SpanEntry | _ir.Span | tuple[_ir.SourceName, int, int, int, int] | None
+
+
+# ------------
+# Section: control flow
+# ------------
+#
+# ``with X.if_(cond):`` opens a conditional frame.
+# ``with X.then_():`` enters its true branch.
+# ``with X.else_():`` enters its false branch.
+# ``with X.for_(loop) as i:`` binds a loop iterator.
+# ``with X.while_(cond):`` builds a while loop.
+# ``X.range_(start, stop)`` constructs a loop descriptor.
+# ``X.break_()`` emits a loop exit.
+# ``X.continue_()`` emits a loop continuation.
+
+
+def if_(condition: Any, *, span: _Span = None) -> IRBuilderFrame:
+ """Create the native conditional frame for a source if statement.
+
+ Parameters
+ ----------
+ condition : Expr or bool
+ Already-evaluated predicate; both branch bodies construct IR without testing it
+ in Python.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ IRBuilderFrame
+ A native conditional context manager. Relax exposes its same-named merged output
+ as frame.var after exit.
+
+ Notes
+ -----
+ Requires an active function region. Enter then_/else_ frames inside it; each branch's
+ lexical helper is defined and called inside that branch frame. TIRx permits a missing
+ else; value-producing Relax branches must agree on their final binding name/type.
+ Invalid context/condition/outputs produce builder diagnostics. Explicit constexpr
+ conditions remain ordinary Python if statements.
+
+ .. code:: python
+
+ # Source
+ if condition:
+ T.evaluate(1)
+ # Generated builder
+ with X.if_(condition):
+ with X.then_():
+ def branch():
+ X.emit_(X.evaluate(1))
+ branch()
+ """
+ raise NotImplementedError
+
+
+def then_(*, span: _Span = None) -> IRBuilderFrame:
+ """Create the true branch of the active conditional.
+
+ Parameters
+ ----------
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ IRBuilderFrame
+ The native branch context manager; its exit finalizes its region.
+
+ Notes
+ -----
+ Requires an active ``if_`` frame and records the true branch. Invalid ordering or repeated
+ branches raise native builder errors. The frame stores its source location; statements
+ keep their individual locations. Lexical helpers take zero explicit arguments and run
+ inside the entered frame.
+
+ .. code:: python
+
+ # Source
+ if condition:
+ pass
+ else:
+ pass
+ # Generated builder, inside X.if_(condition)
+ with X.then_():
+ def branch():
+ pass
+ branch()
+ """
+ raise NotImplementedError
+
+
+def else_(*, span: _Span = None) -> IRBuilderFrame:
+ """Create the false branch of the active conditional.
+
+ Parameters
+ ----------
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ IRBuilderFrame
+ The native branch context manager; its exit finalizes its region.
+
+ Notes
+ -----
+ Requires an active ``if_`` frame and records the false branch. Invalid ordering or repeated
+ branches raise native builder errors. The frame stores its source location; statements
+ keep their individual locations. Lexical helpers take zero explicit arguments and run
+ inside the entered frame.
+
+ .. code:: python
+
+ # Source
+ if condition:
+ pass
+ else:
+ pass
+ # Generated builder, inside X.if_(condition)
+ with X.else_():
+ def branch():
+ pass
+ branch()
+ """
+ raise NotImplementedError
+
+
+def for_(
+ iterable: Any, *, names: str | Sequence[str] | None = None, span: _Span = None
+) -> IRBuilderFrame:
+ """Configure and return the native iteration frame for a source for loop.
+
+ Parameters
+ ----------
+ iterable : ForFrame or range
+ Already-evaluated iteration specification. TIRx converts a Python range to its
+ native serial frame.
+ names : str or sequence of str, optional
+ Source target names, including at most one ``*starred`` group. None (the default)
+ retains constructor defaults. Native configuration expands/validates them before
+ entry; names never control the entry return shape.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ IRBuilderFrame
+ The same configured ForFrame. Entry returns the native variable for one
+ dimension, or the original variable sequence otherwise.
+
+ Notes
+ -----
+ TIRx requires an active primitive function before entry. Variables already have final
+ names at entry. Simple scalar targets use the entry result; generated tuple,
+ list or starred targets use the stable frame.vars sequence for unpacking.
+ Invalid iterable/names raise TypeError, ValueError or native errors. Relax rejects imperative
+ loops. A frame stores its location before deferred body finalization.
+
+ .. code:: python
+
+ # Source
+ for i in range(n):
+ T.evaluate(i)
+ # Generated builder
+ with X.for_(X.range_(n), names=("i",)) as i:
+ X.emit_(X.evaluate(i))
+ """
+ raise NotImplementedError
+
+
+def while_(condition: Any, *, span: _Span = None) -> IRBuilderFrame:
+ """Create a native while-loop frame.
+
+ Parameters
+ ----------
+ condition : Expr or bool
+ Loop predicate expression, constructed once and evaluated by the IR at runtime.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ IRBuilderFrame
+ A context manager whose exit completes the loop body.
+
+ Notes
+ -----
+ TIRx requires an active function region. Invalid predicate/context produces native
+ errors; Relax raises TypeError. This does not repeatedly execute the Python body. The
+ frame owns its stored source span.
+
+ .. code:: python
+
+ # Source
+ while condition:
+ T.evaluate(1)
+ # Generated builder
+ with X.while_(condition):
+ X.emit_(X.evaluate(1))
+ """
+ raise NotImplementedError
+
+
+def range_(*args: Any, annotations: dict[str, Any] | None = None) -> IRBuilderFrame:
+ """Normalize builtin range syntax to a native serial loop frame.
+
+ Parameters
+ ----------
+ args : Expr or int
+ One to three already-evaluated bounds: stop; start, stop; or start, stop, step.
+ Omitted start is zero and omitted step uses the native default.
+ annotations : dict, optional
+ Native loop annotations. None (the default) adds none.
+
+ Returns
+ -------
+ IRBuilderFrame
+ An unentered serial loop frame; ``for_`` configures names and source span.
+
+ Notes
+ -----
+ TIRx validates arity and rejects a literal zero step (TypeError/ValueError). Native
+ bound/type diagnostics propagate. Relax raises TypeError. Only the resolved builtin
+ range is normalized; unrelated host callables named range retain ordinary call behavior.
+
+ .. code:: python
+
+ # Source
+ for i in range(2, n, 2):
+ pass
+ # Generated builder iteration specification
+ X.range_(2, n, 2)
+ """
+ raise NotImplementedError
+
+
+def break_(*, span: _Span = None) -> AlreadyEmitted[Any]:
+ """Emit break for the enclosing language variant loop.
+
+ Parameters
+ ----------
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ AlreadyEmitted[Any]
+ Receipt retaining the exact emitted statement; consuming it emits nothing again.
+
+ Notes
+ -----
+ TIRx requires an enclosing native ForFrame or WhileFrame inside the current primitive
+ function, emits break for that loop and raises ValueError otherwise. Relax raises
+ TypeError. Explicit constexpr loops retain ordinary Python control flow.
+
+ .. code:: python
+
+ # Source
+ break
+ # Generated builder
+ X.break_()
+ """
+ raise NotImplementedError
+
+
+def continue_(*, span: _Span = None) -> AlreadyEmitted[Any]:
+ """Emit continue for the enclosing language variant loop.
+
+ Parameters
+ ----------
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ AlreadyEmitted[Any]
+ Receipt retaining the exact emitted statement; consuming it emits nothing again.
+
+ Notes
+ -----
+ TIRx requires an enclosing native ForFrame or WhileFrame inside the current primitive
+ function, emits continue for that loop and raises ValueError otherwise. Relax raises
+ TypeError. Explicit constexpr loops retain ordinary Python control flow.
+
+ .. code:: python
+
+ # Source
+ continue
+ # Generated builder
+ X.continue_()
+ """
+ raise NotImplementedError
+
+
+# ------------
+# Section: operator overloading
+# ------------
+#
+# ``X.if_then_else_(c, a, b)`` selects an expression.
+# ``X.and_(a, b)`` constructs conjunction.
+# ``X.or_(a, b)`` constructs disjunction.
+# ``X.not_(a)`` negates a condition.
+# ``X.lt_(a, b)`` lowers ``a < b``.
+# ``X.le_(a, b)`` lowers ``a <= b``.
+# ``X.gt_(a, b)`` lowers ``a > b``.
+# ``X.ge_(a, b)`` lowers ``a >= b``.
+# ``X.eq_(a, b)`` lowers ``a == b``.
+# ``X.ne_(a, b)`` lowers ``a != b``.
+
+
+def if_then_else_(condition: Any, true_value: Any, false_value: Any) -> Any:
+ """Construct conditional evaluation from eagerly constructed operands.
+
+ Parameters
+ ----------
+ condition : Any
+ Host boolean or scalar IR predicate.
+ true_value : Any
+ Already-constructed true arm.
+ false_value : Any
+ Already-constructed false arm.
+
+ Returns
+ -------
+ Any
+ The selected host value or native conditional expression.
+
+ Notes
+ -----
+ No statement emission or frame entry occurs. Host predicates select an existing value;
+ scalar IR uses conditional evaluation. Relax also accepts its expression arms. Operand
+ construction remains eager; runtime evaluation follows language variant conditional semantics.
+ Invalid combinations raise native type errors. Source-result handling supplies spans.
+
+ .. code:: python
+
+ # Source
+ yes if condition else no
+ # Generated builder
+ X.if_then_else_(condition, yes, no)
+ """
+ raise NotImplementedError
+
+
+def and_(*values: Any) -> Any:
+ """Construct a conjunction from ordinary boolean or comparison operands.
+
+ Parameters
+ ----------
+ *values : Any
+ Eagerly constructed host or IR operands in source order. TIRx constructs
+ primitive conjunctions; Relax selects primitive or tensor operations.
+
+ Returns
+ -------
+ Any
+ The host boolean or native conjunction expression.
+
+ Notes
+ -----
+ Empty operands raise TypeError. Simple comparison chains accept only names
+ and numeric literals, including signed literals. They lower directly to
+ adjacent comparisons with this conjunction; complex chain operands are a
+ source error before evaluation. There are no temporary operand bindings or
+ substitution. Explicit constexpr code retains Python short-circuit behavior.
+
+ .. code:: python
+
+ # Source
+ 0 < i < 10
+ # Generated builder
+ X.and_(X.lt_(0, i), X.lt_(i, 10))
+ """
+ raise NotImplementedError
+
+
+def or_(*values: Any) -> Any:
+ """Construct a disjunction from already-evaluated operands.
+
+ Parameters
+ ----------
+ values : Any
+ One or more host or language variant boolean operands in source order.
+
+ Returns
+ -------
+ Any
+ Host result or language variant disjunction expression.
+
+ Notes
+ -----
+ No statement is emitted or frame entered. TIRx supports scalar/vector predicates; Relax
+ also supports tensor predicates. Empty operands raise TypeError; native type checks
+ propagate. Explicit constexpr operands retain host short-circuit syntax in the generated
+ program. Result spans are attached by shared source handling.
+
+ .. code:: python
+
+ # Source
+ a or b
+ # Generated builder
+ X.or_(a, b)
+ """
+ raise NotImplementedError
+
+
+def not_(value: Any) -> Any:
+ """Negate a host or language variant boolean without coercing IR truth in Python.
+
+ Parameters
+ ----------
+ value : Any
+ Once-evaluated host boolean or primitive/tensor boolean expression.
+
+ Returns
+ -------
+ Any
+ Host boolean or native logical negation expression.
+
+ Notes
+ -----
+ No statement is emitted or frame entered. TIRx constructs primitive negation; Relax also
+ supports tensor negation. Invalid operand types propagate native errors. Shared source
+ handling attaches the result span.
+
+ .. code:: python
+
+ # Source
+ not value
+ # Generated builder
+ X.not_(value)
+ """
+ raise NotImplementedError
+
+
+def lt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Construct the < comparison in written operand order.
+
+ Parameters
+ ----------
+ lhs : Expr or scalar
+ Already-evaluated left operand.
+ rhs : Expr or scalar
+ Already-evaluated right operand.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ _ir.Expr
+ A native comparison expression.
+
+ Notes
+ -----
+ No statement is emitted or frame entered. TIRx compares primitive operands;
+ Relax selects tensor comparison if either operand is a nonprimitive IR
+ expression, otherwise primitive comparison. Native dtype/shape errors propagate. Do not
+ use Python object equality to compare these operands.
+
+ .. code:: python
+
+ # Source
+ lhs < rhs
+ # Generated builder
+ X.lt_(lhs, rhs)
+ """
+ raise NotImplementedError
+
+
+def le_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Construct the <= comparison in written operand order.
+
+ Parameters
+ ----------
+ lhs : Expr or scalar
+ Already-evaluated left operand.
+ rhs : Expr or scalar
+ Already-evaluated right operand.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ _ir.Expr
+ A native comparison expression.
+
+ Notes
+ -----
+ No statement is emitted or frame entered. TIRx compares primitive operands;
+ Relax selects tensor comparison if either operand is a nonprimitive IR
+ expression, otherwise primitive comparison. Native dtype/shape errors propagate. Do not
+ use Python object equality to compare these operands.
+
+ .. code:: python
+
+ # Source
+ lhs <= rhs
+ # Generated builder
+ X.le_(lhs, rhs)
+ """
+ raise NotImplementedError
+
+
+def gt_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Construct the > comparison in written operand order.
+
+ Parameters
+ ----------
+ lhs : Expr or scalar
+ Already-evaluated left operand.
+ rhs : Expr or scalar
+ Already-evaluated right operand.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ _ir.Expr
+ A native comparison expression.
+
+ Notes
+ -----
+ No statement is emitted or frame entered. TIRx compares primitive operands;
+ Relax selects tensor comparison if either operand is a nonprimitive IR
+ expression, otherwise primitive comparison. Native dtype/shape errors propagate. Do not
+ use Python object equality to compare these operands.
+
+ .. code:: python
+
+ # Source
+ lhs > rhs
+ # Generated builder
+ X.gt_(lhs, rhs)
+ """
+ raise NotImplementedError
+
+
+def ge_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Construct the >= comparison in written operand order.
+
+ Parameters
+ ----------
+ lhs : Expr or scalar
+ Already-evaluated left operand.
+ rhs : Expr or scalar
+ Already-evaluated right operand.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ _ir.Expr
+ A native comparison expression.
+
+ Notes
+ -----
+ No statement is emitted or frame entered. TIRx compares primitive operands;
+ Relax selects tensor comparison if either operand is a nonprimitive IR
+ expression, otherwise primitive comparison. Native dtype/shape errors propagate. Do not
+ use Python object equality to compare these operands.
+
+ .. code:: python
+
+ # Source
+ lhs >= rhs
+ # Generated builder
+ X.ge_(lhs, rhs)
+ """
+ raise NotImplementedError
+
+
+def eq_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Construct the == comparison in written operand order.
+
+ Parameters
+ ----------
+ lhs : Expr or scalar
+ Already-evaluated left operand.
+ rhs : Expr or scalar
+ Already-evaluated right operand.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ _ir.Expr
+ A native comparison expression.
+
+ Notes
+ -----
+ No statement is emitted or frame entered. TIRx compares primitive operands;
+ Relax selects tensor comparison if either operand is a nonprimitive IR
+ expression, otherwise primitive comparison. Native dtype/shape errors propagate. Do not
+ use Python object equality to compare these operands.
+
+ .. code:: python
+
+ # Source
+ lhs == rhs
+ # Generated builder
+ X.eq_(lhs, rhs)
+ """
+ raise NotImplementedError
+
+
+def ne_(lhs: Any, rhs: Any, *, span: _Span = None) -> _ir.Expr:
+ """Construct the != comparison in written operand order.
+
+ Parameters
+ ----------
+ lhs : Expr or scalar
+ Already-evaluated left operand.
+ rhs : Expr or scalar
+ Already-evaluated right operand.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ _ir.Expr
+ A native comparison expression.
+
+ Notes
+ -----
+ No statement is emitted or frame entered. TIRx compares primitive operands;
+ Relax selects tensor comparison if either operand is a nonprimitive IR
+ expression, otherwise primitive comparison. Native dtype/shape errors propagate. Do not
+ use Python object equality to compare these operands.
+
+ .. code:: python
+
+ # Source
+ lhs != rhs
+ # Generated builder
+ X.ne_(lhs, rhs)
+ """
+ raise NotImplementedError
+
+
+# ------------
+# Section: context lookup and resolution
+# ------------
+#
+# ``X.resolve_global_info_(key)`` resolves module metadata.
+# ``X.resolve_type_var_("n")`` resolves a symbolic dimension.
+# ``X.call_global_var_(f, args)`` calls a module function.
+# ``I.module_member_("f", f)`` installs a module member.
+
+
+def resolve_global_info_(content: Any) -> Any:
+ """Look up language-owned global metadata selected by a source argument.
+
+ Parameters
+ ----------
+ content : str or Any
+ A selector or concrete value interpreted by the active language.
+
+ Returns
+ -------
+ Any
+ The language's metadata value. Selector syntax, lookup context and
+ unsupported-value errors belong to the language variant implementation.
+
+ 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.
+
+ .. 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]"))
+ """
+ raise NotImplementedError
+
+
+def resolve_type_var_(
+ name: str,
+ dtype: str | _ir.Type | _ir.Var | None = None,
+ *,
+ value: _ir.Var | None = None,
+ span: _Span = None,
+) -> _ir.Var:
+ """Resolve or declare a symbolic variable in the nearest native function.
+
+ Parameters
+ ----------
+ name : str
+ Function-local lookup key; quoted symbols do not create a Python binding.
+ dtype : str, Type or Var, optional
+ Explicit primitive type or supplied variable. None defaults new symbols to
+ int64; existing symbols must agree with an explicit dtype.
+ value : Var, optional
+ Existing primitive variable to register without replacement. None creates a
+ variable only if the name is not already registered.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ _ir.Var
+ The exact existing or newly registered variable.
+
+ Notes
+ -----
+ Requires an active native function frame. The signature and resumed body use the same
+ symbol map; nested functions use separate maps. Conflicting dtype/variable declarations
+ or missing context raise a builder error. No statement is emitted.
+
+ .. code:: python
+
+ # Source
+ n = T.int64()
+ # Generated builder
+ n = X.resolve_type_var_("n", dtype="int64")
+ """
+ raise NotImplementedError
+
+
+def call_global_var_(function: _ir.GlobalVar, args: Sequence[Any]) -> _ir.Expr:
+ """Construct a call to a declared module function.
+
+ Parameters
+ ----------
+ function : GlobalVar
+ Native callee reference reserved by the module declaration phase.
+ args : sequence of Any
+ Positional operands evaluated once, in source order. Generated global calls do
+ not accept keyword arguments.
+
+ Returns
+ -------
+ _ir.Expr
+ The caller language variant's call expression; constructing it does not emit a statement.
+
+ Notes
+ -----
+ The callee declaration must be available in the active module context. TIRx preserves
+ its exact declared return type; Relax converts operands using its expression conversion.
+ Invalid types or an unavailable declaration produce native builder errors. Source-call
+ handling attaches locations to the returned expression.
+
+ .. code:: python
+
+ # Source
+ Module.callee(x)
+ # Generated builder
+ X.call_global_var_(callee_reference, [x])
+ """
+ raise NotImplementedError
+
+
+def module_member_(name: str, value: Any) -> Any:
+ """Register a concrete function member or retain ordinary class setup.
+
+ Parameters
+ ----------
+ name : str
+ Source class member identifier.
+ value : Any
+ Already-evaluated class member value; a BaseFunc is declared and defined as a
+ module function.
+
+ Returns
+ -------
+ Any
+ The reserved GlobalVar for a concrete BaseFunc, otherwise the exact original
+ value.
+
+ Notes
+ -----
+ Concrete function registration requires an active module frame and follows native
+ duplicate/type checks. Other values need no frame and cause no IR emission or renaming.
+ Shared source class setup runs before function signatures so module attributes/global
+ info are available.
+
+ .. code:: python
+
+ # Source
+ class Module:
+ helper = existing_function
+ # Generated builder, inside I.ir_module()
+ helper = I.module_member_("helper", existing_function)
+ """
+ from tvm.ir import BaseFunc
+
+ from .ir import decl_function, def_function
+
+ if isinstance(value, BaseFunc):
+ reference = decl_function(name, value)
+ def_function(name, value)
+ return reference
+ return value
+
+
+# ------------
+# Section: special protocol
+# ------------
+#
+# 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.
+# ``register_type_var_decl(path, 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.
+# ``module_decorator(path)`` marks module declaration decorators.
+# ``declaration_kind(path, kind)`` marks function and helper declarations.
+# ``with X.function_(...) as fn:`` builds a function frame.
+# ``a = X.arg("a", ty)`` declares a parameter.
+# ``X.func_name("main")`` sets the function name.
+# ``X.func_ret_type(ty)`` sets the result annotation.
+# ``X.check_well_formed_(result)`` validates completed IR.
+
+
+def function_(*, decl: bool = False, span: _Span = None, **options: Any) -> IRBuilderFrame:
+ """Create the native function frame used for signature and body construction.
+
+ Parameters
+ ----------
+ decl : bool, optional
+ False (default) constructs a complete function on one entry. True collects a
+ signature on the first entry and retains this frame for body re-entry.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+ options : Any
+ Language variant options: TIRx private, s_tir and persistent default to False. Relax
+ is_pure defaults to True, is_private/local to False; local=True requires a
+ declared reference when building its body.
+
+ Returns
+ -------
+ IRBuilderFrame
+ The native context manager. It owns params, result type, symbol map, reference
+ and completed result.
+
+ Notes
+ -----
+ Requires an active IRBuilder; module definitions also require its module frame. Declare
+ every sibling signature before any body for forward references. Re-enter the same frame,
+ define and invoke a zero-argument lexical helper inside it, then exit before validation.
+ Invalid options/context raise builder errors. The frame stores its source location
+ independently of exit-time ambient context.
+
+ .. code:: python
+
+ # Source
+ @T.prim_func
+ def f(a: T.int32):
+ T.evaluate(a)
+ # Generated builder
+ with X.function_() as fn:
+ X.func_name("f")
+ a = X.arg("a", X.int32)
+ X.emit_(X.evaluate(a))
+ """
+ raise NotImplementedError
+
+
+def arg(name: str, annotation: Any, *, span: _Span = None) -> _ir.Var:
+ """Add a parameter to the active native function signature.
+
+ Parameters
+ ----------
+ name : str
+ Source parameter name.
+ annotation : Type, Var, Buffer or callable
+ Concrete rewritten annotation or existing native parameter. A callable
+ annotation is evaluated; an existing variable retains identity.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ _ir.Var
+ The created or retained parameter variable (buffers are native variables).
+
+ Notes
+ -----
+ Requires an active function signature frame. Primitive symbols use that frame's
+ resolver, preserving annotation/body identity. TIRx handles buffer layout according to
+ its function policy; Relax converts its type annotation. Invalid annotations or context
+ raise TypeError/native builder errors.
+
+ .. code:: python
+
+ # Source
+ def f(a: T.int32):
+ pass
+ # Generated builder, inside the signature frame
+ a = X.arg("a", X.int32)
+ """
+ raise NotImplementedError
+
+
+def func_name(name: str) -> None:
+ """Set the active native function's source name.
+
+ Parameters
+ ----------
+ name : str
+ Function identifier; used for the module reference and public symbol according
+ to language variant privacy options.
+
+ Returns
+ -------
+ None
+ No value.
+
+ Notes
+ -----
+ Requires a function frame. Mutates its signature metadata and emits no statement;
+ invalid or repeated naming follows native diagnostics. No source span argument is needed
+ because the frame owns its location.
+
+ .. code:: python
+
+ # Source
+ def f():
+ pass
+ # Generated builder, inside the function frame
+ X.func_name("f")
+ """
+ raise NotImplementedError
+
+
+def func_ret_type(annotation: Any, *, span: _Span = None) -> None:
+ """Set the active native function's return annotation.
+
+ Parameters
+ ----------
+ annotation : Type, Expr or callable
+ Rewritten return annotation; expression annotations supply their type. None
+ denotes a void/empty tuple return as supported by the language variant.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ None
+ No value.
+
+ Notes
+ -----
+ Requires a function signature frame. Resolves callable annotations and records the type;
+ no body is emitted. Unsupported or conflicting types raise native builder errors. The
+ completed function retains its frame span.
+
+ .. code:: python
+
+ # Source
+ def f() -> T.int32:
+ return 1
+ # Generated builder, inside the signature frame
+ X.func_ret_type(X.int32)
+ """
+ raise NotImplementedError
+
+
+def check_well_formed_(module: _ir.IRModule) -> None:
+ """Validate a completed module after every body has been finalized.
+
+ Parameters
+ ----------
+ module : IRModule
+ Completed native module, including all resolved forward references and mixed-
+ language variant members.
+
+ Returns
+ -------
+ None
+ No value or mutation; success means the existing native validators accepted the
+ program.
+
+ Notes
+ -----
+ Requires completed IR, with no active construction frame. Root coordination
+ invokes opaque whole-module hooks supplied by language variant initialization through
+ tvm.script.register_module_validator. Each hook owns concrete types, eligibility
+ and cross-function validation, including captured/preexisting members. No source
+ decorator inventory selects the hooks. Validator exceptions propagate unchanged;
+ an empty hook list raises RuntimeError instead of establishing validity.
+ check_well_formed=False omits the generated call entirely.
+
+ .. code:: python
+
+ # Source
+ @I.ir_module
+ class Module:
+ pass
+ # Generated builder, after module frame exit
+ I.check_well_formed_(module)
+ """
+ from tvm.script import _MODULE_VALIDATORS
+
+ validators = tuple(_MODULE_VALIDATORS)
+ if not validators:
+ raise RuntimeError("No completed-module validators are registered")
+ for validator in validators:
+ validator(module)
+
+
+# ------------
+# Section: binding
+# ------------
+#
+# ``a = X.bind_(value, name="a")`` binds a source name.
+# ``X.decl_mutable_cell_(value, ty=ty)`` declares mutable storage.
+# ``X.set_mutable_cell_(cell, value)`` updates mutable storage.
+# ``a, b = X.unpack(value)`` destructures a binding.
+
+
+def bind_(
+ value: Any = MISSING,
+ *,
+ ty: Any = None,
+ name: str | None = None,
+ span: _Span = None,
+ value_span: _Span = None,
+ name_span: _Span = None,
+ frame_value: bool = False,
+) -> Any:
+ """Apply the language variant's ordinary assignment policy.
+
+ Parameters
+ ----------
+ value : Any, optional
+ Once-evaluated RHS. MISSING (the default) denotes an omitted initializer and is
+ rejected unless the language variant supports the annotation-only form.
+ ty : Type or annotation callable, optional
+ Already-rewritten source annotation. None (the default) lets the language variant infer
+ the binding type.
+ name : str, optional
+ Source target name. None (the default) requests no source-derived name.
+ span : SpanEntry, Span or source-location tuple, optional
+ Binding-target location for a newly constructed binding. None (the default)
+ leaves it unspecified; this is separate from the RHS location.
+ value_span : SpanEntry, Span or source-location tuple, optional
+ RHS source location, passed separately without first stamping the returned value.
+ None (the default) leaves explicit RHS attribution unspecified. The language variant
+ applies it only when binding/conversion requires value attribution; TIRx
+ variable and metadata passthrough retain producer names and spans.
+ name_span : SpanEntry, Span or source-location tuple, optional
+ Location of the target identifier. None (the default) uses span; it can differ
+ from the emitted statement location.
+ frame_value : bool, optional
+ False by default. True names an already-entered with-target without constructing
+ another binding or entering its frame.
+
+ Returns
+ -------
+ Any
+ The value bound to the Python target; usually a native variable, or an unchanged
+ host/frame-owned value.
+
+ Notes
+ -----
+ Requires the language variant construction context when producing IR. TIRx emits immutable Bind
+ statements; Relax emits normalized bindings and match-casts. Unsupported
+ values/annotations raise TypeError or ValueError. Ordinary TIRx Vars (including
+ BufferVars) pass through before general Expr binding, without naming, stamping
+ or another binding. Other Expr values retain ordinary language variant binding. Non-Expr
+ metadata such as Layout and ordinary meta_class instances passes through
+ unchanged without inspecting or naming its resources. Explicit typed declarations,
+ mutable storage and frame targets retain their separate contracts. AlreadyEmitted
+ receipts retain RHS attribution without another emission. A DSL may opt into
+ concise scope entry: register the returned child
+ frame's exit callback on the active parent before child entry, then return its entered
+ value. Later statements enter that child, and parent exit closes it. This is language variant
+ policy; it adds no parser-owned scope state. Direct-call results and source module
+ aliases bypass this operation.
+
+ .. code:: python
+
+ # Source
+ x = value
+ # Generated builder
+ x = X.bind_(value, name="x", span=_S[0], value_span=_S[1])
+
+ # TIRx concise scope entry
+ tid = T.launch_thread("threadIdx.x", 128)
+ # Generated builder
+ tid = X.bind_(X.launch_thread("threadIdx.x", 128), name="tid")
+ """
+ raise NotImplementedError
+
+
+def decl_mutable_cell_(
+ value: Any = MISSING,
+ *,
+ ty: Any = None,
+ name: str | None = None,
+ span: _Span = None,
+ name_span: _Span = None,
+) -> Any:
+ """Introduce explicitly declared mutable storage.
+
+ Parameters
+ ----------
+ value : Any, optional
+ Once-evaluated storage handle for a call declaration, or initializer for an
+ annotation declaration. MISSING (the default) means no initializer.
+ ty : Type or annotation callable, optional
+ Scalar or vector storage annotation. None (the default) identifies an already-
+ created storage handle.
+ name : str, optional
+ Source target name. None (the default) requests no source-derived name.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+ name_span : SpanEntry, Span or source-location tuple, optional
+ Location of the target identifier. None (the default) uses span; it can differ
+ from the emitted statement location.
+
+ Returns
+ -------
+ Any
+ The same declared storage handle, or the handle allocated for the annotation.
+
+ Notes
+ -----
+ TIRx requires an active primitive function; scalar primitive annotations allocate local
+ storage and may initialize it. Vector annotations allocate storage but reject an
+ initializer. Invalid handle/type combinations raise TypeError or ValueError. Relax
+ always rejects mutable storage with TypeError. Declaration syntax takes precedence over
+ any outer mutable target name.
+
+ .. code:: python
+
+ # Source
+ x: T.int32 = 1
+ # Generated builder
+ x = X.decl_mutable_cell_(1, ty=X.int32, name="x")
+ """
+ raise NotImplementedError
+
+
+def set_mutable_cell_(target: Any, value: Any, *, span: _Span = None) -> AlreadyEmitted[Any]:
+ """Emit an update through an existing mutable handle without rebinding it.
+
+ Parameters
+ ----------
+ target : TensorLoad, scalar wrapper or one-element buffer Var
+ Storage handle returned by an explicit mutable declaration. Its identity is
+ retained; this argument is not a source name or a new declaration.
+ value : Expr or scalar convertible to Expr
+ Once-evaluated value to store. Native store checking validates its type and
+ indices against the target.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location of the emitted store. None (the default) leaves explicit location
+ unspecified; existing source-call provenance is retained.
+
+ Returns
+ -------
+ AlreadyEmitted[Any]
+ Receipt retaining the exact emitted statement; consuming it emits nothing again.
+
+ Notes
+ -----
+ TIRx requires an active primitive function/statement region and appends a native store
+ to that region. Scalar wrappers are unwrapped, TensorLoad indices are preserved, and
+ one-element buffer targets use index zero. Unsupported targets raise TypeError; native
+ type/index checks propagate. Relax always raises TypeError because its bindings are
+ immutable. No new allocation or immutable binding is created.
+
+ .. code:: python
+
+ # Source
+ x: T.int32 = 0
+ x = value
+ # Generated builder
+ x = X.decl_mutable_cell_(0, ty=X.int32, name="x")
+ X.set_mutable_cell_(x, value)
+ """
+ raise NotImplementedError
+
+
+def unpack(value: Any) -> Any:
+ """Expose elements for ordinary Python target unpacking.
+
+ Parameters
+ ----------
+ value : Any
+ Concrete IR tuple, typed tuple expression, or ordinary host iterable.
+
+ Returns
+ -------
+ Any
+ A Python tuple of IR fields/projections for an IR tuple; otherwise the original
+ value.
+
+ Notes
+ -----
+ No statement is emitted and no builder frame is entered. Known IR tuple types supply
+ arity; ordinary Python performs target-count/starred-unpacking checks. Field identities
+ are retained for concrete tuples. Direct-call results keep ordinary Python unpacking
+ without this hook.
+
+ .. code:: python
+
+ # Source
+ a, b = value
+ # Generated builder
+ left, right = X.unpack(value)
+ a = X.bind_(left, name="a")
+ b = X.bind_(right, name="b")
+ """
+ raise NotImplementedError
+
+
+# ------------
+# Section: statement
+# ------------
+#
+# ``X.emit_(value)`` emits an expression statement.
+# ``X.return_(value)`` emits a function return.
+# ``X.setitem_(a, i, value)`` emits an indexed store.
+# ``X.setattr_(a, "field", value)`` emits an attribute store.
+# ``X.assert_(condition, message)`` emits an assertion.
+
+
+def emit_(value: Any, *, span: _Span = None) -> None:
+ """Consume a source expression statement.
+
+ Parameters
+ ----------
+ value : Any
+ Once-evaluated expression result. AlreadyEmitted receipts and None produce no
+ additional emission.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location of the emitted statement and expression, or the existing node in an
+ AlreadyEmitted receipt. None leaves explicit attribution unspecified. Active
+ caller provenance is composed without adding a construction context; native
+ frames retain the location for finalization.
+
+ Returns
+ -------
+ None
+ No source-visible value.
+
+ Notes
+ -----
+ Requires an active language variant function/region when emitting IR. TIRx adds statements,
+ evaluates expressions, and can enter concise frames. Variables, text, layouts,
+ and meta_class instances are inert; sequences are consumed
+ elementwise. Relax accepts only void expressions (or
+ None/AlreadyEmitted); unsupported values raise TypeError and non-void expressions raise
+ ValueError. An explicit span annotates the exact previously emitted node in a receipt
+ without emitting it again. Known builder results need no separate result wrapper;
+ opaque source calls keep their scoped provenance before this hook.
+
+ .. code:: python
+
+ # Source
+ T.evaluate(1)
+ # Generated builder
+ X.emit_(X.evaluate(1), span=_S[0])
+ """
+ raise NotImplementedError
+
+
+def return_(value: Any = None, *, span: _Span = None) -> AlreadyEmitted[Any] | None:
+ """Record a language variant function return while continuing Python construction.
+
+ Parameters
+ ----------
+ value : Any, optional
+ Return operand. None (the default) means an empty tuple in Relax; TIRx requires
+ an expression.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ AlreadyEmitted[Any] | None
+ Receipt retaining the exact emitted statement; consuming it emits nothing again.
+ Relax records a function result and returns None.
+
+ Notes
+ -----
+ Requires an active function/region. TIRx emits its native return operation; Relax
+ records the function result. Unsupported missing/type/context combinations raise
+ TypeError or native errors. A source macro retaining ordinary Python return does not
+ call this hook.
+
+ .. code:: python
+
+ # Source
+ return value
+ # Generated builder
+ X.return_(value)
+ """
+ raise NotImplementedError
+
+
+def setitem_(target: Any, key: Any, value: Any, *, span: _Span = None) -> AlreadyEmitted[Any]:
+ """Apply an indexed assignment using already-evaluated operands.
+
+ Parameters
+ ----------
+ target : buffer Var
+ Destination buffer.
+ key : Expr, int, slice or sequence
+ Indices in written order; native buffer-store rules validate supported forms.
+ value : Expr or scalar
+ Once-evaluated stored value.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ AlreadyEmitted[Any]
+ Receipt retaining the exact emitted statement; consuming it emits nothing again.
+
+ Notes
+ -----
+ TIRx requires an active statement region and preserves target/index evaluation order.
+ Native shape/type/index errors propagate. Relax raises TypeError because indexed
+ mutation is unsupported.
+
+ .. code:: python
+
+ # Source
+ A[i] = value
+ # Generated builder
+ X.setitem_(A, i, value)
+ """
+ raise NotImplementedError
+
+
+def setattr_(
+ target: Any, name: str, value: Any, *, span: _Span = None
+) -> AlreadyEmitted[Any] | None:
+ """Apply an attribute assignment using already-evaluated operands.
+
+ Parameters
+ ----------
+ target : Any
+ Object containing a scalar storage attribute or ordinary mutable Python
+ metadata.
+ name : str
+ Attribute identifier, evaluated by source syntax before this hook.
+ value : Any
+ Once-evaluated replacement or stored value.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ AlreadyEmitted[Any] | None
+ Receipt retaining the exact emitted statement; consuming it emits nothing again.
+ Ordinary host attribute updates return None because they emit no IR statement.
+
+ Notes
+ -----
+ TIRx stores through scalar storage attributes without replacing their handles; other
+ attributes use ordinary Python setattr. Native store checks and Python attribute errors
+ propagate. Relax raises TypeError. A native store needs an active function; ordinary
+ host metadata updates do not.
+
+ .. code:: python
+
+ # Source
+ state.count = value
+ # Generated builder
+ X.setattr_(state, "count", value)
+ """
+ raise NotImplementedError
+
+
+def assert_(
+ condition: Any,
+ message: str | tuple[str, Sequence[Any]] | Sequence[Any] = "",
+ *,
+ span: _Span = None,
+) -> None:
+ """Emit a runtime assertion.
+
+ Parameters
+ ----------
+ condition : Expr or bool
+ Already-evaluated predicate.
+ message : str or assertion metadata, optional
+ Empty text by default. Relax requires construction-time text. TIRx also accepts
+ message parts or an (error_kind, parts) pair.
+ span : SpanEntry, Span or source-location tuple, optional
+ Location for the constructed result. None (the default) leaves explicit location
+ unspecified; active source-call provenance is composed by the builder. Frames
+ retain their location until finalization.
+
+ Returns
+ -------
+ None
+ No value; emits a native assertion.
+
+ Notes
+ -----
+ Requires an active function/statement region. Malformed diagnostic metadata raises
+ TypeError; native predicate checking propagates. Construction does not test an IR
+ predicate as a host boolean.
+
+ .. code:: python
+
+ # Source
+ assert condition, "failed"
+ # Generated builder
+ X.assert_(condition, "failed")
+ """
+ raise NotImplementedError
diff --git a/python/tvm/script/parser/__init__.py b/python/tvm/script/parser/__init__.py
index d9911322a0..c67be91e9d 100644
--- a/python/tvm/script/parser/__init__.py
+++ b/python/tvm/script/parser/__init__.py
@@ -14,38 +14,72 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The parser subpackage of TVMScript.
+"""Canonical TVMScript parser with language variant namespace initialization."""
-Per-dialect parser submodules (``tvm.script.parser.tirx``, etc.) are
-resolved lazily via :data:`tvm.script._DIALECT_REGISTRY`. When a dialect
-is accessed (e.g. ``tvm.script.parser.tirx``), this subpackage's
-``__getattr__`` looks up the dialect in ``_DIALECT_REGISTRY`` and imports
-``<dialect_module_path>.parser`` (e.g. ``tvm.tirx.script.parser``),
-caching the result so subsequent accesses skip ``__getattr__``.
+from __future__ import annotations
-The IR layer is foundational and is NOT registered as a dialect — its
-parser lives as a real submodule ``tvm.script.parser.ir``, with
-``ir_module`` re-exported at this level for convenience.
+import importlib
+from collections.abc import Callable
+from typing import Any, TypeVar
-See :mod:`tvm.script` for a full description of the dialect resolution
-mechanism, including the ``_DialectRedirectFinder`` that handles
-deep statement-form imports.
-"""
+_NAMESPACES: dict[str, object] = {}
+_NAMESPACE_INITIALIZERS: list[Callable[[], None]] = []
+_NAMESPACE_ALIASES: set[str] = {"I", "ir"}
+_ENTRY_EXPORTS = (
+ "from_source",
+ "ir_module",
+ "make_decorator",
+ "make_macro_decorator",
+ "parse",
+ "pyfunc",
+)
+__all__ = [*_ENTRY_EXPORTS, "register_namespace", "register_namespace_initializer"]
+_initialized = False
+_initializing = False
+
+
+def register_namespace(alias: str, namespace: object) -> None:
+ """Register an opaque fixed namespace; its first alias is the syntax key root.
+
+ Each parse borrows these namespaces. Replacing an alias affects future parses
+ only; registration enters no builder frame and inspects no namespace members.
+ """
+ _NAMESPACES[alias] = namespace
+
+
+def register_namespace_initializer(
+ initializer: Callable[[], None], *, aliases: tuple[str, ...] = ()
+) -> None:
+ """Register a lazy bootstrap callback without importing the language variant."""
+ _NAMESPACE_ALIASES.update(aliases)
+ if any(existing is initializer for existing in _NAMESPACE_INITIALIZERS):
+ return
+ _NAMESPACE_INITIALIZERS.append(initializer)
+ if _initialized:
+ initializer()
-import importlib
-from typing import Any
-from . import _core, ir, tirx
-from ._core import parse
-from .ir import ir_module
+def _initialize() -> None:
+ global _initialized, _initializing
+ if _initialized or _initializing:
+ return
+ _initializing = True
+ try:
+ importlib.import_module(__name__ + ".ir")
+ register_namespace("TypeVar", TypeVar)
+ for initialize in _NAMESPACE_INITIALIZERS:
+ initialize()
+ _initialized = True
+ finally:
+ _initializing = False
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(f"{_DIALECT_REGISTRY[name]}.parser")
- globals()[name] = module
- return module
- raise AttributeError(f"module 'tvm.script.parser' has no attribute {name!r}")
+ if name not in _ENTRY_EXPORTS and name not in _NAMESPACES and name not in _NAMESPACE_ALIASES:
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+ _initialize()
+ if name in _ENTRY_EXPORTS:
+ return getattr(importlib.import_module(__name__ + ".entry"), name)
+ if name in _NAMESPACES:
+ return _NAMESPACES[name]
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
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/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/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,
... 24796 lines suppressed ...