## Motivation and context
A TIRx pointer carries two pieces of information that later lowering
needs: the pointee element type and the storage scope. Both must survive
when a pointer-producing expression is assigned to a Python name and
then used as the backing storage of a buffer.
A concrete example is accessing an mbarrier in another CTA through
distributed shared memory:
```python
ptr_ty = PointerType(PrimType("uint64"), "shared")
remote_ptr = T.reinterpret(
ptr_ty,
T.ptx.map_shared_rank(mbar.ptr_to([0]), T.int32(0)),
)
remote_mbar = T.decl_buffer(
[1], "uint64", data=remote_ptr, scope="shared"
)
```
`map_shared_rank` returns the raw `uint64` address produced by PTX
`mapa`, and `reinterpret` gives that address the intended
`PointerType(uint64, shared)`. Because `decl_buffer(data=...)` requires
a pointer `Var`, assigning the expression to `remote_ptr` should create
an immutable typed pointer binding.
Before this PR, an unannotated assignment such as `remote_ptr = <pointer
expression>` followed the same parser path as a numeric assignment. That
path allocates a mutable local scalar and therefore cannot represent a
`PointerType`. The pointer expression could not be carried as a
correctly typed `Var` into `decl_buffer` and CUDA lowering.
This PR makes an unannotated pointer-valued assignment emit a TIRx
`Bind`. The bound `Var` has exactly the type of the right-hand side,
including its element type and storage scope. Pointer bindings are
immutable, so reassignment in the same scope is diagnosed; shadowing a
name supplied through `extra_vars` remains valid. Numeric assignments
keep their existing mutable-local behavior.
## Type propagation fixes
The parser fix exposed several other boundaries where pointer type
information must remain consistent:
| Boundary | Previous behavior | Behavior after this PR |
| --- | --- | --- |
| Unannotated pointer assignment | Tried to materialize the value as a
local scalar | Emits an immutable `Bind` with the RHS `PointerType` |
| `address_of(buffer)` / `buffer.ptr_to(...)` | Reused the raw backing
pointer type | Returns a pointer to `buffer.dtype` while preserving the
backing pointer storage scope |
| `tvm_access_ptr` / `ptr_byte_offset` | Accepted strings or annotation
expressions, but not a `PrimType` object directly | Accepts `PrimType`
and produces the corresponding typed pointer |
| Printed `T.ptx.mapa` call | The printer emits all intrinsic attributes
positionally, but the Python helper required keyword-only arguments |
Accepts the canonical printed form so pointer code round-trips through
TVMScript |
The `address_of` distinction matters for typed views over byte-addressed
storage. For example, if a `float32` buffer is backed by a `uint8*`
allocation in `shared.dyn`, the address of a buffer element must be
`PointerType(float32, shared.dyn)`, not `PointerType(uint8,
shared.dyn)`.
With these changes, the DSMEM example above round-trips through
TVMScript and CUDA codegen declares the remote buffer pointer as
`uint64_t*`.
## TMA dtype normalization
This PR also contains a small, separate type-representation fix in TMA
descriptor construction. `TmaPlan.elem_dtype` is a string consumed by
the host-side `runtime.cuTensorMapEncodeTiled` packed call, but
`_assemble_plan` stored `g_buf.dtype`, which is a `PrimType`. Converting
it with `str(g_buf.dtype)` ensures that the generated packed-call
argument is `StringImm("float16")` rather than an IR type object. This
does not change the pointer-binding semantics described above.
## Testing
- Verify that an unannotated pointer expression creates a `Bind` whose
`Var` type matches the RHS type.
- Verify that pointer reassignment is rejected while shadowing an
`extra_vars` name is allowed.
- Verify parser/printer structural round-tripping for the pointer
binding and canonical `T.ptx.mapa` call.
- Verify that `address_of` uses the logical buffer element type and
preserves the storage scope for byte-backed buffer views.
- Verify that `tvm_access_ptr` and `ptr_byte_offset` accept `PrimType`
inputs.
- Compile the DSMEM `map_shared_rank` example through the CUDA TIRx
pipeline and check for a typed `uint64_t*` remote buffer pointer.
- Verify that the TMA host initialization passes the descriptor dtype as
a `StringImm`.
## Motivation and context
The TVMScript parser already tracks Python AST locations for
diagnostics, but TIRx statements and expression results emitted through
`IRBuilder` did not retain those locations. After parsing, a direct
intrinsic call, an inlined helper body, or a `TilePrimitiveCall`
therefore could not be traced back to the source range that produced it.
Inline expansion also needs more than a single flat location. The
generated IR should retain both the caller location and the
helper-definition location, while ordinary nested AST evaluation within
one source should not accumulate redundant enclosing spans.
## Changes
- Add an active source-span stack to `IRBuilder`, with scoped push/pop
support.
- Make the parser activate the current AST source range while visiting
statements and evaluating expressions.
- Attach the active span to emitted TIRx statements and to expression
results that do not already carry an explicit span.
- Normalize nested spans from the same source to the innermost relevant
range.
- Preserve cross-source inline expansion history as a `SequentialSpan`,
ordered from the call site to the expanded definition.
- Reuse the same source-coordinate calculation for diagnostics and IR
spans so their line and column conventions remain consistent.
Source spans remain diagnostic metadata: functions parsed from different
source locations keep the same structural hash and remain structurally
equal.
## Testing
- Verify exact parser source coordinates against diagnostic coordinates.
- Verify spans on direct intrinsic calls and `TilePrimitiveCall` nodes.
- Verify that inline expansion produces a `SequentialSpan` containing
caller and callee ranges.
- Verify direct `IRBuilder.with_source_span` behavior.
- Verify that source spans do not affect structural identity.
- Run changed-files pre-commit checks, including clang-format.
Focused result: 9 tests passed.
## Summary
- Make `PrimExpr` a typed C++ view over `Expr` values whose
`ExprNode::ty` is `PrimType`, instead of using a separate runtime node
class as the proof of primitive-ness.
- Use the shared `ir::Call` node for Relax, TIRX, and primitive-valued
calls, while keeping primitive-only APIs explicit at their semantic
boundaries.
- Keep Python on the general `Expr` surface for primitive-typed values
so `isinstance` behavior does not imply a nominal primitive-expression
subclass.
## Design Rationale
The main advantage of this change is that common expression nodes such
as `Call` can be unified without specializing each one to `PrimType`. A
single `ir::Call` can represent a Relax tensor call, a Relax scalar
call, or a primitive-valued intrinsic call; the result type stored in
`ExprNode::ty` determines whether that particular value can be viewed as
`PrimExpr`.
This keeps the IR node hierarchy focused on expression structure rather
than result-type categories. Nodes that are intrinsically primitive,
such as integer and floating-point literals or TIRX primitive operators,
still have strongly typed C++ APIs and data structures. General nodes
whose result type may vary, such as `Call`, remain general `Expr` nodes
and are narrowed to `PrimExpr` only where primitive-only semantics are
required.
The PR also keeps the compatibility surface practical: C++
primitive-only APIs continue to accept `PrimExpr`, Python exposes a
compatibility predicate for checking the primitive typed category, and
visitors/printers use one natural `Call` path rather than duplicating
Relax and primitive call handling. Missing expression types are
represented explicitly with `Type::Missing()` so constructors can leave
type inference to later analysis without relying on nullable `Type`
values.
In the past we have been using `DataType` in PrimExpr.dtype field to
check type information for PrimExpr while still having BaseExpr.ty for
richer type information. DataType is also used both in runtime and
compiler. This PR streamlines the boundary:
- PrimExpr.ty now carries PrimType that replaces original use of
`DataType`
- Runtime use will now favor DLPack DLDataType, removing one layer of
indirection.
- Constants attributes where values are usually runtime values, will use
`DLDataType`
- DataType will be phased out after this PR
We also brings up helper functions in PrimType, but also limits them to
a more concise set so the functions do not grow with the data type codes
in DLPack.
This is a major refactor that changes the IR primitive. It helps to
bring possible future benefits:
- Unified type mechanism through Expr.ty
- Possibility of carry future Type nodes
Migration Guide:
- Use `PrimType` when code reasons about compiler expression types,
tensor element compiler types, or constructs a `PrimExpr`/compiler type.
- Use existing source types such as `expr.ty()`, `ExprOp.expr_ty()`, or
TE tensor element `dtype` where possible instead of rebuilding a type
from dtype text.
- Use raw `DLDataType` for runtime constants, ABI paths, dtype-valued
attrs, and storage/runtime helper logic.
- Prefer direct `PrimType` equality, `MatchesCode(...)`,
`MatchesElementType(...)`, and `WithCode(...)` over local wrappers or
string dtype checks.
Performance:
Using Object type instead of DLDataType would indeed bring some
performance impact to the IR. We have done the following performance
optimizations:
- Make sure most of the outputs reuse one of the PrimType from inputs
- Cache a thread local PrimType based on input so we don't repeatly
realloc
We did benchmarks show that rewrite simplify operation stays within
+-10% overhead of original one. Which merits the refactor given the
benefit the unfication brings
This PR adds an autoload mechanism for out-of-tree backends, simplifies
TVM's Python library loading, and removes `TVMError` in favor of native
Python errors.
## Autoload out-of-tree backends
Out-of-tree packages can register an autoload callable under the
`tvm.backends` entry-point group (mirroring torch's device-backend
autoload). At `import tvm` startup each entry point is discovered and
its callable invoked once, after the core runtime and the `tvm`
namespace are fully initialized, so an extension can register
ops/targets/funcs or load extra libraries.
```toml
[project.entry-points."tvm.backends"]
tvm_foo = "tvm_foo:_autoload"
```
A failing extension is caught and surfaced via `warnings.warn` so it
cannot break `import tvm`. Autoload can be disabled with
`TVM_DEVICE_BACKEND_AUTOLOAD=0`.
## Simplify library loading
The library-loading path in `base.py` is consolidated around a single
`_LOADED_LIBS` dict (basename to ctypes handle) so downstream and
autoloaded extensions can skip already-loaded libraries; the per-backend
runtime DSO list is folded into `load_backend_libs`. Accumulated cruft
is removed: the Python-3.9 check, the readline shim, the `_FFI_MODE`
ctypes check, the `base.__version__` re-export, and `py_str` (call sites
inline `.decode("utf-8")`).
## Remove TVMError in favor of native Python errors
`TVMError` added a layer atop `RuntimeError` that downstream code had to
import and learn. It is removed; the registered FFI error kinds
(`InternalError`, `RPCError`, `OpError`, `DiagnosticError`,
`ScheduleError`) now subclass `RuntimeError` directly while staying
registered, so the FFI keeps throwing the right kinds. All `TVMError`
imports, `except`/`raise`/`isinstance` uses, and
`pytest.raises(tvm.TVMError)` sites move to the `RuntimeError` builtin.
Replace TVM's `Diagnostic` / `DiagnosticContext` machinery with the
tvm-ffi
`visit_error_context` mechanism. Validators throw an `ffi::Error` seeded
with the
offending node; leaf pass executors (`ModulePass` / relax `Function` /
`DataflowBlock`) catch and rethrow `EnrichPassErrorWithContext`, which
appends the
failing pass name and a TVMScript-rendered, underlined source location.
`relax.analysis.well_formed` now throws on the first violation; a new
`check_well_formed` returns a bool, and all C++/Python/test callers are
routed
accordingly. `include/tvm/ir/diagnostic.h` and `src/ir/diagnostic.cc`
are deleted.
The enrichment renders with `num_context_lines=10` so a small function
shows in
full with no skipped-lines marker, while a large module stays bounded.
The TVMScript parser diagnostics
(`python/tvm/script/parser/core/diagnostics.py`)
stay self-contained pure-Python with no `DiagnosticContext` dependency,
and
restore multi-line source rendering: a diagnostic whose offending AST
node spans
multiple source lines now renders every spanned line with its gutter
line number
and an underline covering the span. `tvm.error.DiagnosticError` (used by
the
TVMScript parser) is retained.
A rendered end-to-end enriched-error example is posted as a comment
below.
## Summary
This PR adds the initial TIRx support needed for low-level programming
of Blackwell-class GPU architectures. As part of the ongoing TIRx
refactor, it introduces TVMScript support for directly scripting
advanced hardware features without relying on scheduling as the primary
programming interface.
The change keeps existing `s_tir` script support intact while making
direct scripting a first-class path for TIRx programs.
## Main Changes
- Add TIRx operator dispatch and layout infrastructure.
- Add TVMScript support for new low-level TIRx operations.
- Add analysis, transform, and lowering support for TIRx IR nodes.
- Add CUDA/Blackwell-oriented codegen and intrinsic coverage.
- Add Python and C++ integration points for TIRx scripting and runtime
support.
## Validation
- `pre-commit run --all-files`
- `ninja -C build -j32`
- `CUDA_VISIBLE_DEVICES=2 pytest tests/python/tirx/ -n 16`
- `1723 passed, 47 skipped, 32 warnings`
- `CUDA_VISIBLE_DEVICES=2 python -m pytest -v
tests/python/all-platform-minimal-test`
- `37 passed, 105 skipped`
- `TVM_TEST_TARGETS=llvm python -m pytest -v tests/python/tirx-analysis
tests/python/tirx-base tests/python/tirx-transform -n 16`
- `664 passed, 25 skipped, 9 xfailed, 1 xpassed`
## Local CI Notes
Some full CI-equivalent jobs were not locally reproducible because this
machine is missing parts of the Apache TVM CI environment, including
`llvm-config-15/17`, Vulkan, ROCm, Maven, Sphinx, Doxygen, Emscripten,
and ARM/QEMU cross-toolchain components. Metal-specific tests were
skipped locally because no Metal runtime is available.
## Problem
On Python 3.14, any use of TVMScript raises a `TypeError` before the
module body is even parsed:
```
TypeError: unsupported format string passed to type.__format__
```
The traceback points to
`python/tvm/script/parser/core/diagnostics.py:120`:
```python
raise TypeError(f"Source for {obj:!r} not found")
```
## Root Cause
`{obj:!r}` is an invalid f-string expression. The `:` introduces a
`format_spec`, so `!r` is passed to `type.__format__` as a format string
— which it does not support.
The intended syntax for a `repr()` conversion is `{obj!r}` (no colon).
Python 3.14 re-implemented f-string parsing under [PEP
701](https://peps.python.org/pep-0701/) and now strictly validates
format specs, surfacing this latent bug. Python 3.10–3.13 silently
passed the invalid spec to `__format__` and happened not to raise in
most code paths, so the bug went unnoticed.
## Fix
```diff
- raise TypeError(f"Source for {obj:!r} not found")
+ raise TypeError(f"Source for {obj!r} not found")
```
One character change. Valid across all Python versions >= 3.6.
## Testing
Verified on Python 3.14.2 (darwin/arm64):
- TVMScript `ir_module` + `prim_func` parses and compiles correctly
after the fix
- Full TVMScript test suite: **628 passed, 1 xfailed** (the 1 failure in
`test_tvmscript_roundtrip.py::test_roundtrip[relax_symbolic_size_var]`
is pre-existing and unrelated to this change)
Add handling for `doc.keyword` nodes in `ExprEvaluator._visit` to ensure
expressions (e.g. `BoolOp`) in keyword arguments are processed with
correct evaluation methods.
Fix#18972 . For more details, please refer to this issue.
This PR brings up the tirx namespace. We have been spliting out the
original tir namespace to include high-level component s_tir and this PR
updates the remaining low-level part as tirx namespace
## Summary
Rename `LetStmtNode`/`LetStmt` to `BindNode`/`Bind` and remove the
`body` field.
The variable defined by `Bind(var, value)` is now visible in all
subsequent
statements within the same enclosing scope, rather than being scoped to
a nested body.
This flattens deeply nested let-chains into sequential
`SeqStmt([Bind(...), Bind(...), ...])`,
making the IR easier to read, transform, and analyze.
## Key Changes
- **New `BindNode`**: `{var, value}` — no body field. Variable scope is
the enclosing
statement's body (For, IfThenElse, AllocBuffer, etc.)
- **ScopeStack pattern**: Passes that need scope-aware cleanup
(ConvertSSA, CSE,
tir_visitor_with_path) use `ScopeStack` instead of manual save/restore
or RAII wrappers
- **All passes migrated**: 89 files updated across codegen backends, TIR
transforms,
S-TIR transforms, analyses, TVMScript printer/parser/ir_builder
With `from __future__ import annotations`, Python stores annotations as
strings
and does not capture annotation-only variables in `__closure__`. This
broke
TVMScript when buffer shapes/dtypes referenced closure variables.
Fix: wrap `extra_vars` in a `collections.ChainMap` with snapshots of all
live
caller-frame locals (from `inspect.stack()`) as fallback layers in both
`tir/entry.py` (`prim_func`) and `ir/entry.py` (`ir_module`). The
`ir_module`
function now also captures `outer_stack = inspect.stack()` at its entry
point,
mirroring the existing pattern in `prim_func`. Lookup falls back to
frame locals
only on cache miss, preserving existing behavior for non-PEP-563 code.
Add `tests/python/tvmscript/test_tvmscript_pep563_closure.py` (requires
`from __future__ import annotations` at the top) covering closure
variables in
buffer shapes, dtypes, nested scopes, ir_module, and mixed
annotation+body use.
This PR enables ruff pyupgrade (UP) rules with py310 target, auto-fixing
~5600 annotation modernizations (PEP 585 generics, PEP 604 unions,
deprecated typing imports).
Also removes from __future__ import annotations from ir/module.py and
rmsnorm.py, bumps requires-python to >=3.10, and removes absolute_import
aliases from topi/contrib files.
Implement proper parsing and evaluation of chained comparison operators
(e.g., `0 < i < 128`) in TVMScript. The sequence comparisons are now
correctly expanded to their logical equivalents (e.g., `(0 < i and i < 128)`).
Changes:
- Updated expression evaluator to handle sequence comparisons correctly
- Added test case to verify sequence comparison functionality
Add support for conditional expressions in TVMScript
This PR adds support for conditional expressions in TVMScript parser,
which allows developers to use Python-style conditional expressions
```python
@T.prim_func
def func(A: T.buffer((128, 128), "float32")):
for i, j in T.grid(128, 128):
A[i, j] = i if i < j else j
@T.prim_func
def expected(A: T.buffer((128, 128), "float32")):
for i, j in T.grid(128, 128):
A[i, j] = T.if_then_else(i < j, i, j)
```
* support continue and break in tvmscript
* fix black format
* fix pylint issue
* Update tests/python/tvmscript/test_tvmscript_syntax_sugar.py
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
* add printer/parser test, fix lint
* Fit to latest ffi update
* Skip i386 numpy-related test
* Introduce AnnotateIrregularLoop before any lowering loop expansions.
---------
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
### **Overview**
This PR implements native Python function support in TVM Relax through
the `@I.pyfunc` decorator and `BasePyModule`, which enable seamless
integration between TVM's compilation pipeline and Python/PyTorch runtime
environments. This enhancement allows users to write Python functions
directly in TVMScript that can interoperate with Relax and TIR functions
that provides enhanced debugging capabilities and leveraging existing
PyTorch operator libraries.
### **Key Features**
**TVMScript Parser Enhancement**
- `@I.pyfunc` decorator: Marks Python functions for integration into IRModules
- Dual storage format: Stores both raw string representation (for TVMScript
printing) and captured PackedFunc (for runtime execution)
- ExternFunc representation: Each Python function is represented as an
ExternFunc node with attributes storing source code and runtime wrapper
**Complete BasePyModule Implementation**
- DLPack-based tensor conversion: Seamless conversion between PyTorch
tensors and TVM NDArrays
- Cross-function interoperability: Python functions can call Relax/TIR
functions and vice versa
- JIT compilation: Delays compilation until module instantiation for flexible
late-stage modifications
- Dynamic function registration: Supports runtime addition of Python functions
### Future Work
- TVMScript printer for IRModules with Python functions: Print IRModules
in proper format with high-level operator mapping from Relax ops to PyTorch
ops, handling symbolic shapes
- R.call_py_func primitive: Introduce Relax primitive to invoke corresponding
PackedFunc of specified Python functions at runtime
- Set Python minimal version to 3.9 as Python 3.8 has reached end of life
- Update AST handling to match Python 3.9+ API changes
- Fix deprecation warning: "Support for arbitrary keyword arguments is deprecated and will be removed in Python 3.15"
- Update TVMScript parser components to use modern AST interfaces
- Adjust linting and test infrastructure for Python 3.9+ compatibility
[ARITH] Remove deprecated attributes from Constant AST node
Remove the deprecated `s` and `n` attributes from the `Constant` AST node
class in the script parser. These attributes were previously used for
string and numeric constants but are no longer needed in the current
AST implementation.
Updated all code that creates `Constant` objects to remove the
corresponding parameters:
- Removed `s` and `n` parameters from `Constant.__init__`
- Updated `_FIELDS` list to exclude deprecated attributes
- Fixed calls in `evaluator.py` and `parser.py` to remove extra arguments
This change simplifies the AST structure and removes unused legacy code.
This PR phases out tvm._ffi redirections in favor of new FFI
new functions are now called via tvm.ffi.
We also enabled limited API support for python 3.12+
so the compiled binary can be forward compatible to future
python versions.
* [TVMScript] Enable T.macro decorateing class method
This PR refactors the implementation of `T.macro`, so that the `self` argument can be passed through the TVMScript parser. Then we can decroate the class methods with `T.macro`.
* update test
* [TVMScript] Avoid segfault from invalid TVMScript
Prior to this commit, after the `DiagnosticContext` prints its error,
it overwrites the `DiagnosticRenderer` with a NULL renderer. If a
second call to `DiagnosticContext::Render` occurs, it will segfault.
This appears to be intended to prevent double-printing of error
messages, but double-printing error messages is much worse than a
segfault.
In addition, `DiagnosticContext::Render` should only be called once.
There's a common pattern in the parser where it will wrap exceptions
in `DiagnosticError`, but re-raise exceptions that are already a
`DiagnosticError`. This requires every such location to include
`except DiagnosticError: raise`, and can easily be missed.
This PR makes two changes: First, the `DiagnosticRenderer` is updated
to have a no-op callback rather than a NULL callback. Second, the
re-raising of `DiagnosticError` is moved to `Parser.report_error`, so
that it does not need to be handled separately at several independent
locations in the TVMScript parser.
* [TVMScript][Bugfix] Normalize relax::If with function's TIR var
Prior to this commit, the branches of `relax::If` were normalized
using `EraseToWellDefinedInScope`, using a fresh variable scope.
While this had the intended behavior of preventing variables defined
in a single branch from being usable outside of the conditional, it
also caused the conditional's branches to treat function-scope
symbolic variables as if they were undefined.
This commit updates the `tvm::relax::Normalizer` so that `relax::If`
is normalized within an inherited scope. This preserves the previous
behavior for symbolic variables defined within a branch, but allows
shapes within a branch to use symbolic variables defined outside of
the branch.
* [Relax] Canonicalize known symbolic shapes in Relax expressions
Prior to this commit, known constants in Relax functions would be
inlined by the `CanonicalizeBindings` pass, but only if they appeared as Relax
expressions (e.g. `R.const` or `R.prim_value`). Known constants that
appeared as TIR variables (e.g. symbolic shapes) would be kept as
dynamic parameters, even if they were known at compile time.
This commit updates the `CanonicalizeBindings` pass to identify known
values of symbolic shapes, and to use these known values in shape
expressions.
* [Relax][Refactor] Reorganize pattern-matching
A follow-up to https://github.com/apache/tvm/pull/16730. Now that the
implementations for `rewrite_call` and `rewrite_bindings` are in
separate classes, they can be further split out into separate files.
* [Relax][Refactor] Implement Rewriter class for pattern-rewrite
Prior to this commit, the pattern to be matched and the rewrite to be
performed were provided as separate arguments. This commit introduces
a new class `ExprRewriter`, which contains both parts.
This abstraction will make it easier to combine multiple different
rewrite rules, applying them in a single pass.
* lint fixes
* Remove unnecessary change which broke a unit test
* lint fix for import order
* Add docstrings
* lint fix
* Lint fix
* lint fixes
* lint fix
* Update based on review comments
* Add test case for matching against arbitrary dtype
* Fix breakage in unit tests
One unit test that had been relying on invalid shape propagation.
Another unit test that required constructed an ill-formed output to
test against.
* Updated base class name from ExprRewriter to PatternMatchingRewriter
* lint fix
* [Analysis] Allow calls to GlobalVar in @R.function
Prior to this commit, the post-parsing well-formed check performed by
TVMScript allowed a call to `GlobalVar` in a `@R.function`, but only
if it occurred within the context of a `@I.ir_module`. If
`@R.function` appeared on its own, calls to a `GlobalVar` would be
treated as calls to an undefined function.
* Use approrpirate well-formed checks TIR/Relax functions
* Lint fix
* Import order fix
* Check well-formedness in the parser
* Correct packed funcs in NN frontend
* Support the check_well_formed optional argument to I.ir_module
* Also check well-formedness in TIR
* Enable normalization for individual Relax functions and PrimFuncs
* Use the error raised by the TIR well-formed checker for the message
* Fix tvmscript test failures
* Whitespace
* Fix errors in verify_well_formed test
* Include a more helpful error message
* Fix TIR test failures
* Address well-formed failures in test_tir_specialize
* Correct well-formedness error in test_tir_analysis_oob
* Correct further well-formedness failures
* Remove __tvm_meta__ from test case to avoid parsing error
* Avoid circular import in entryy.py
* Formatting fixes
* lint fix
* Add pylint exceptions
* Fix whitespace
* Fix more failed test cases
* Catch inappropriate use of decl_function instead of segfaulting
* Fix test_lower.py
* Mark purity in test_relax_2d_buffer_allocation.py
* Mark purity in test_dma_builtin.py
* Remove __tvm_meta___ from test_tir_usmp_analysis_extract_bufferinfo.py
* Suppress well-formed check in test_tir_transform_convert_blocks_to_opaque.py
* Remove __tvm_meta__ in test_tir_usmp_algo.py
* Remove __tvm_meta__ from more USMP tests
* Fix incorrect var in test_tir_transform_storage_flatten.py
* Remove all remaining instances of __tvm_meta__
* Fix purity error in test_dataflow_pattern.py
* Fix purity error in test_ast_printer
* Fix test_arith_domain_touched example
* Okay to set check_well_formed to True in test_tir_analysis_identify_mcmcpy
* Define variable in test_tir_analysis_oob
* Typo fix
* Add explanatory comment to test case
* Define the undefined vars in test_tir_transform_common_subexpr_elim
* Exception no longer necessary in test_tir_transform_inject_rolling_buffer
* Remove unnecessary check exemption in test_tir_transform_convert_ssa
* Avoid checking exemption in test_inject_ptx_ldg32
* Note special case in test_distributed_transform_propagate_sharding
* Exempt well-formed error in dlight/test_benchmark
* Exempt well-formedness errors in test_ethosu/, mostly uninitialized vars
* Whitespace
* Include non-CUDA GPUs in IsScheduledOnGPU
* Fix thread binding bug by changing thread binding var dtype
* Include overrides in test_runtime_builtin_paged_attention_kv_cache.py
* add exemptions in test_ethosu/test_replace_conv2d
* Add more ethosu exemptions
* More exemptions for ethosu tests
* Remove unused reference
* Indicate purity in test_transform_rewrite_cuda_graph
* Indicate purity in test_transform_normalize
* Reorder MergeSharedMemoryAllocations in GPU codegen
* Add target parameter for FP8StorageLegalize and FP8ComputeLegalize
* Don't re-import Target in tvm/tir/transform/transform.py
Prior to this commit, exceptions raised during the parsing of
TVMScript would be caught and replaced with a new exception. While
this does allow the TVMScript location of the error to be included in
the exception, it also removes the stack trace of the original error.
This commit updates the `Parser.report_error` function to provide the
original stack trace alongside the updated exception object.
The builtins are already supported by `eval` (they are automatically
injected in the global scope), but they are not recognized by the
evaluator's checks. When the evaluator sees doc.Name, it looks it up
in the current `var_table`, and flags an error if it's not there.
Make the evaluator also consult the current builtins before erroring
out.
* [Script] Be more careful when generating ast.ExtSlice for Subscript
The ast.ExtSlice expects a non-empty list, otherwise evaluation
fails with "error: empty dims on ExtSlice". Also, each element
in "dims" list of ExtSlice must be either Slice or Index.
In python3.8 an expression A[()] is parsed (by ast) as Subscript
with slice being Index(value=Tuple(elts=[])). When we translate a
subscript from doc.AST to ast, we unconditionally convert every
tuple to ast.ExtSlice, which in this case is incorrect.
The fix is to map empty tuple back to the Index(Tuple[])) instead
of ExtSlice. In other cases, ensure that members of ExtSlice are
of correct types.
* Fix lint #1
* [TVMScript] Support starred indices in for-loop
An extension of https://github.com/apache/tvm/pull/15404, which
allowed starred expressions in the rhs of
assignments (e.g. `T.decl_buffer(shape=[*dim, 128])`), this PR also
enables starred expressions in the lhs of
assignments (e.g. `for *spatial,reduction in T.grid(*A.shape)`).
* Fix single-argument indices
* Updated test case for T.grid()
As a background info---the script parser works by visiting a "statement"
(or top-level expression) at a time. The expression parts of the state-
ment are evaluated, and then the IR corresponding to the statement is
constructed if necessary.
In TIR, macro calls can only occur at the statement level, and they don't
produce any values. This means that the statement visitor (visit_expr_stmt)
can see these calls directly in its node parameter. At this point it could
simply visit the body of the macro instead, which is the basis of the
existing implementation.
In other dialects there may be a need for macros to produce values. This
means that macro calls can occur in the middle of complex expressions.
As a result, these calls will not be present at the statement level, and
the TIR approach by intercepting them in visit_expr_stmt will no longer
work. Instead, these macros delay the visiting of the macro body to the
evaluation time. A macro is represented by an ScriptMacro (TIRMacro in
the current implementation) object (created via macro decorator). When the
evaluator evaluates an expression with a macro call, it will call the
macro object (since macro calls use function call syntax). It is in the
macro object's __call__ function where the macro parsing picks up. The
remaining issue was to pass the Parser object to the __call__ function.
This is done by injecting it into the global dictionary under a reserved
name.
It turns out that the same approach also works for TIR, and the macro
processing can be generalized, leaving only language-specific details to
the language-specific language macro objects.
* [TIR] Implement TIR macros
This patch introduces two new symbols: `T.macro` and `T.insert`.
`T.macro` is a decorator that, when applied to a function, turns the
body of that function into a piece of TIR that can be inserted via
`T.insert` into a PrimFunc.
For example:
```python
@T.macro
def copy_backwards(dst, src, size):
with T.block("backwards"):
for i in T.serial(size):
ai = T.axis.remap("S", [i])
T.reads(src[0:size])
T.writes(dst[0:size])
dst[ai] = src[size - ai - 1]
@T.prim_func
def foo_int32(A: T.Buffer((128,), "int32"), B: T.Buffer((128,), "int32")):
T.insert(copy_backwards, A, B, 128)
@T.prim_func
def foo_int8(A: T.Buffer((128,), "int8"), B: T.Buffer((128,), "int8")):
T.insert(copy_backwards, A, B, 128)
```
The above will generate two PrimFuncs that do the same backwards copy,
but applied to buffers with different data types.
Semantics:
- Function that is decorated with @T.macro can have any parameters that
follow Python syntax, i.e. positional, keyword, etc. Type annotations
are not required, but are allowed.
- The arguments to `T.insert` are macro name followed by the argument
list.
For `T.insert(arg1, arg2, arg3, ...)`, the values are substituted into
the body of the macro as in the call `arg1(arg2, arg3, ...)`.
The body with the substituted values is then inserted at the point
where the `T.insert` is located.
* Fix linter
* Fix linter again
One linter suggested something that the other didn't like...
* Get rid of T.insert, apply macro via function-call syntax
* Store closure vars in TIRMacro
* ast.parse always returns ast.Module, hence doc is doc.Module
* Simplify `expand_macro`, capture environment variables
* Implement macro hygiene
* Fix linter
* Make T.macro work same as T.macro()
The previous commit inadvertently made T.macro (without parentheses)
illegal, only abbreviated form allowed was T.macro(). Restore T.macro
as a valid decorator use.
* Edit comment: insertion -> expansion
* Add import pytest
* One more typo...
* Remove stale testcase
This PR implements the privacy annotation proposal. Namely, the @R.function decorator now has an optional private attribute. If a function is marked as private, then it will not have a global symbol attached to it and thus will not be externally accessible. By default, functions are not private, so the parser does insert a global symbol for them.
Hongyi Jin