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/10/19 22:11:18 UTC
(tvm) 02/15: [REFACTOR][IR] Unify raising well-formedness checks (#20793)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tirx-standalone-root
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 7844dd5ea84d02090abb8cecd81400b252a90744
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Mon Oct 19 11:42:27 2026 -0400
[REFACTOR][IR] Unify raising well-formedness checks (#20793)
Standardize well-formedness APIs on `check_well_formed` /
`CheckWellFormed`. S-TIR, Relax, and layout checks raise directly; TIRx
retains `assert_mode` and its boolean result contract. Callers propagate
the original diagnostics through analysis and passes.
Remove the script module-validator registry, dialect registration
callbacks, and generated whole-module dispatch. Module validation is an
explicit analysis operation, while standalone function decorators retain
validation and the `check_well_formed=False` option. Update affected
callers and concise standalone API documentation.
---
tirx/include/tvm/tirx/analysis.h | 35 ++++----------
tirx/include/tvm/tirx/layout.h | 6 +--
tirx/python/tvm_tirx/__init__.py | 2 +-
tirx/python/tvm_tirx/backend/_autoload_backends.py | 2 +-
tirx/python/tvm_tirx/script/__init__.py | 46 ++----------------
tirx/python/tvm_tirx/script/ir_builder/__init__.py | 2 -
.../tvm_tirx/script/ir_builder/parser_protocol.py | 43 -----------------
tirx/python/tvm_tirx/script/parser/entry.py | 3 +-
tirx/python/tvm_tirx/script/parser/transpile.py | 9 +---
tirx/python/tvm_tirx/tirx/__init__.py | 10 ----
tirx/python/tvm_tirx/tirx/analysis/analysis.py | 17 ++++---
tirx/python/tvm_tirx/tirx/layout.py | 24 +++++-----
.../tirx/script/ir_builder/parser_protocol.py | 9 +---
.../backend/rocm/codegen/llvm/codegen_amdgpu.cc | 2 +-
tirx/src/tirx/analysis/check_well_formed.cc | 18 ++++----
tirx/src/tirx/ir/layout/compose_layout.cc | 12 ++---
tirx/src/tirx/ir/layout/layout.cc | 2 +-
tirx/src/tirx/ir/layout/tile_core.cc | 9 ++--
tirx/tests/python/script/test_basic_usage.py | 15 +++---
.../tirx/analysis/test_check_well_defined.py | 54 +++++++++++-----------
tirx/tests/python/tirx/script/test_basic_usage.py | 29 +++++++++---
tirx/tests/python/tirx/test_layout.py | 36 +++++++--------
tirx/tests/python/tirx/test_verifier.py | 2 +-
.../tirx/transform/test_common_subexpr_elim.py | 2 +-
.../python/tirx/transform/test_convert_ssa.py | 2 +-
.../tirx/transform/test_lower_tensor_layouts.py | 2 +-
.../transform/test_pointer_value_type_rewrite.py | 2 +-
tirx/tests/python/tirx/transform/test_simplify.py | 2 +-
.../tirx/transform/test_split_host_device.py | 2 +-
.../python/tirx/transform/test_storage_rewrite.py | 4 +-
30 files changed, 151 insertions(+), 252 deletions(-)
diff --git a/tirx/include/tvm/tirx/analysis.h b/tirx/include/tvm/tirx/analysis.h
index d574282a5d..fa19ccdc67 100644
--- a/tirx/include/tvm/tirx/analysis.h
+++ b/tirx/include/tvm/tirx/analysis.h
@@ -57,39 +57,20 @@ TVM_TIRX_DLL bool VerifyMemory(const Function& func);
TVM_TIRX_DLL size_t CalculateExprComplexity(const PrimExpr& expr);
/*!
- * \brief Verify if the given TIR is well-formed. The verification includes:
- *
- * - All variables are defined prior to their point of use.
- *
- * - No variables are used outside of the scope of their definition.
- *
- * - Each variable has a single point of definition.
- *
- * - Tensor loads have consistent source, indices, and result types.
- *
- * - Loop control, tile-operation categories, and scope IDs are valid.
- *
- * Dialect statements require their dialect-specific verifier.
- *
+ * \brief Check whether a TIRx function is well formed.
* \param func The Function to be verified.
- * \param assert_mode The indicator if it raises an error when the function is not well-formed.
- * \return Whether it is a well-formed TIR function.
+ * \param assert_mode Whether to raise on invalid IR (default: true).
+ * \return True for valid IR; false for invalid IR when assert_mode is false.
*/
-TVM_TIRX_DLL bool VerifyWellFormed(const Function& func, bool assert_mode = true);
+TVM_TIRX_DLL bool CheckWellFormed(const Function& func, bool assert_mode = true);
/*!
- * \brief Verify if the TIR in the given IRMOdule is well-formed.
- *
- * In addition to the checks performed for each Function (see above),
- * the following checks are performed:
- *
- * - The same TIR variable may not be defined in more than one function
- *
+ * \brief Check whether a TIRx module is well formed.
* \param mod The IRModule to be verified.
- * \param assert_mode The indicator if it raises an error when the function is not well-formed.
- * \return Whether it is a well-formed TIR module.
+ * \param assert_mode Whether to raise on invalid IR (default: true).
+ * \return True for valid IR; false for invalid IR when assert_mode is false.
*/
-TVM_TIRX_DLL bool VerifyWellFormed(const IRModule& mod, bool assert_mode = true);
+TVM_TIRX_DLL bool CheckWellFormed(const IRModule& mod, bool assert_mode = true);
/*!
* \brief Find the entry function of the given IRModule, i.e, functions marked by
diff --git a/tirx/include/tvm/tirx/layout.h b/tirx/include/tvm/tirx/layout.h
index 18297a1a68..7c8baab9df 100644
--- a/tirx/include/tvm/tirx/layout.h
+++ b/tirx/include/tvm/tirx/layout.h
@@ -48,7 +48,7 @@ using ffi::Tuple;
class LayoutNode : public ffi::Object {
public:
/*! \brief Verify if the layout is well-formed */
- virtual bool VerifyWellFormed() const = 0;
+ virtual void CheckWellFormed() const = 0;
/*! \brief Get the size of the layout (of some axis) */
virtual PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const = 0;
@@ -164,7 +164,7 @@ class TileLayoutNode : public LayoutNode {
}
/*! \brief Verify if the layout is well-formed */
- bool VerifyWellFormed() const final;
+ void CheckWellFormed() const final;
/*! \brief Get the size of the layout (of some axis) */
PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
@@ -254,7 +254,7 @@ class ComposeLayoutNode : public LayoutNode {
}
/*! \brief Verify if the layout is well-formed */
- bool VerifyWellFormed() const final;
+ void CheckWellFormed() const final;
/*! \brief Get the size (of some axis) of the layout */
PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
diff --git a/tirx/python/tvm_tirx/__init__.py b/tirx/python/tvm_tirx/__init__.py
index f42166e8a4..9fcdf53125 100644
--- a/tirx/python/tvm_tirx/__init__.py
+++ b/tirx/python/tvm_tirx/__init__.py
@@ -96,7 +96,7 @@ if TYPE_CHECKING or not _is_config_mode():
sys.excepthook = tvm_wrap_excepthook(sys.excepthook)
# Autoload loads built-in and out-of-tree backends. Out-of-tree extensions opt
- # into being loaded automatically at ``import tvm`` time by declaring an entry
+ # into being loaded automatically at ``import tvm_tirx`` time by declaring an entry
# point in the ``tvm.backends`` group:
# [project.entry-points."tvm.backends"] tvm_foo = "tvm_foo:_autoload".
# Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``.
diff --git a/tirx/python/tvm_tirx/backend/_autoload_backends.py b/tirx/python/tvm_tirx/backend/_autoload_backends.py
index 600075640b..4e225f09a0 100644
--- a/tirx/python/tvm_tirx/backend/_autoload_backends.py
+++ b/tirx/python/tvm_tirx/backend/_autoload_backends.py
@@ -62,7 +62,7 @@ def _autoload_backends() -> None:
if not _RUNTIME_ONLY:
_load_builtin_backends()
- # Out-of-tree extensions opt into being loaded automatically at ``import tvm`` time
+ # Out-of-tree extensions opt into being loaded automatically at ``import tvm_tirx`` time
# by declaring an entry point in the ``tvm_tirx.backends`` group:
# [project.entry-points."tvm_tirx.backends"] tvm_foo = "tvm_foo:_autoload".
# Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``.
diff --git a/tirx/python/tvm_tirx/script/__init__.py b/tirx/python/tvm_tirx/script/__init__.py
index 138b7aba40..bcd71c957b 100644
--- a/tirx/python/tvm_tirx/script/__init__.py
+++ b/tirx/python/tvm_tirx/script/__init__.py
@@ -25,14 +25,10 @@ import importlib
import importlib.machinery
import importlib.util
import sys
-from collections.abc import Callable
from typing import Any
_DIALECT_REGISTRY: dict[str, str] = {}
_DIALECT_BUILDER_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]] = []
# Public parser and builder compatibility paths consume dialect-owned registrations.
_REDIRECTED_SUBPACKAGES = {
@@ -41,38 +37,6 @@ _REDIRECTED_SUBPACKAGES = {
}
-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, *, builder_path: str | None = None) -> None:
"""Register a dialect's script namespace and canonical builder package.
@@ -119,7 +83,7 @@ def _redirect_target(fullname: str) -> str | None:
Returns ``None`` if ``fullname`` is not a redirected name.
"""
if fullname.startswith("tvm_tirx.script."):
- # tvm.script.<language variant>[.subpath]
+ # tvm_tirx.script.<language variant>[.subpath]
rest = fullname[len("tvm_tirx.script.") :]
head, _, tail = rest.partition(".")
if head in _DIALECT_REGISTRY and "." not in head:
@@ -208,7 +172,7 @@ class _AliasLoader:
module.__loader__ = self._loader
-# Install the redirect finder once. Re-importing tvm.script (e.g. during a
+# Install the redirect finder once. Re-importing tvm_tirx.script (e.g. during a
# pytest reload) must not stack duplicates.
if not any(isinstance(f, _DialectRedirectFinder) for f in sys.meta_path):
# Handle deep aliases before PathFinder can execute a canonical source file
@@ -223,8 +187,8 @@ def __getattr__(name: str) -> Any:
return module
if name == "ir":
# IR is foundational — its parser is a real submodule under
- # tvm.script.parser.ir, exposed here as `tvm.script.ir` for the
- # legacy `from tvm.script import ir as I` pattern.
+ # tvm_tirx.script.parser.ir, exposed here as `tvm_tirx.script.ir` for the
+ # legacy `from tvm_tirx.script import ir as I` pattern.
ir_parser = importlib.import_module("tvm_tirx.script.parser.ir")
globals()["ir"] = ir_parser
return ir_parser
@@ -235,7 +199,7 @@ def __getattr__(name: str) -> Any:
globals()["parse"] = parse
return parse
if name == "ir_module":
- # ir_module lives in the IR parser at tvm.script.parser.ir; the IR
+ # ir_module lives in the IR parser at tvm_tirx.script.parser.ir; the IR
# layer is foundational, so we resolve it directly rather than via
# the language variant registry.
ir_parser = importlib.import_module("tvm_tirx.script.parser.ir")
diff --git a/tirx/python/tvm_tirx/script/ir_builder/__init__.py b/tirx/python/tvm_tirx/script/ir_builder/__init__.py
index 29a5dd6e35..bf2b887ab5 100644
--- a/tirx/python/tvm_tirx/script/ir_builder/__init__.py
+++ b/tirx/python/tvm_tirx/script/ir_builder/__init__.py
@@ -58,7 +58,6 @@ from .ir import (
)
from .ir import _get_dialect_builder as __getattr__
from .parser_protocol import (
- check_well_formed_,
decl_function,
ir_module,
module_member_,
@@ -88,7 +87,6 @@ __all__ = [
"Type",
"Var",
"at_",
- "check_well_formed_",
"constexpr",
"decl_function",
"dtype",
diff --git a/tirx/python/tvm_tirx/script/ir_builder/parser_protocol.py b/tirx/python/tvm_tirx/script/ir_builder/parser_protocol.py
index 704a0ab9f8..902803ab85 100644
--- a/tirx/python/tvm_tirx/script/ir_builder/parser_protocol.py
+++ b/tirx/python/tvm_tirx/script/ir_builder/parser_protocol.py
@@ -244,49 +244,6 @@ def func_ret_type_(annotation: Any, *, loc: _Loc = None) -> None:
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_tirx.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_tirx.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)
-
-
# --------------------------------------
# Bindings
# --------------------------------------
diff --git a/tirx/python/tvm_tirx/script/parser/entry.py b/tirx/python/tvm_tirx/script/parser/entry.py
index 00d529b920..e7999a824c 100644
--- a/tirx/python/tvm_tirx/script/parser/entry.py
+++ b/tirx/python/tvm_tirx/script/parser/entry.py
@@ -328,7 +328,8 @@ def make_decorator(builder: object, *, namespace_path: str) -> Callable[..., Any
global symbol matching its name. Defaults to False.
check_well_formed : bool, optional
Whether to check that the constructed function is well formed.
- Defaults to True.
+ Defaults to True. Module construction does not validate the completed
+ module; call the language analysis API explicitly.
persistent : bool, optional
For ``T.function``, mark the resulting function as a persistent
kernel. Defaults to False. See
diff --git a/tirx/python/tvm_tirx/script/parser/transpile.py b/tirx/python/tvm_tirx/script/parser/transpile.py
index 1980f67db0..1735b89374 100644
--- a/tirx/python/tvm_tirx/script/parser/transpile.py
+++ b/tirx/python/tvm_tirx/script/parser/transpile.py
@@ -2278,7 +2278,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
root,
)
)
- if check_well_formed:
+ if check_well_formed and not is_module:
# -------------------- Pattern --------------------
# Python source:
# @X.function
@@ -2289,12 +2289,7 @@ class IRBuilderTranspiler(ast.NodeTransformer):
# result = builder.get()
# X.check_well_formed_(result)
# -------------------------------------------------
- # Module syntax selects I.check_well_formed_ instead, after all bodies complete.
- namespace = (
- self.module.ir_prefix
- if is_module
- else self._function_namespace(root, self.read_function_metadata(root)[0])
- )
+ namespace = self._function_namespace(root, self.read_function_metadata(root)[0])
statements.append(
ast.copy_location(
ast.Expr(
diff --git a/tirx/python/tvm_tirx/tirx/__init__.py b/tirx/python/tvm_tirx/tirx/__init__.py
index 6ef8dfcc8f..79c23e4e88 100644
--- a/tirx/python/tvm_tirx/tirx/__init__.py
+++ b/tirx/python/tvm_tirx/tirx/__init__.py
@@ -99,16 +99,6 @@ if not _RUNTIME_ONLY_TIRX:
)
-def _check_script_module(module: "tvm_tirx.ir.IRModule") -> None:
- # Delay builder imports until validation, keeping dialect/runtime bootstrap safe.
- from tvm_tirx.tirx.script.ir_builder.parser_protocol import _check_module_well_formed
-
- _check_module_well_formed(module)
-
-
-tvm_tirx.script.register_module_validator(_check_script_module)
-
-
def _initialize_script_namespace() -> None:
from . import script
diff --git a/tirx/python/tvm_tirx/tirx/analysis/analysis.py b/tirx/python/tvm_tirx/tirx/analysis/analysis.py
index 4a0879724c..d0c6d9904f 100644
--- a/tirx/python/tvm_tirx/tirx/analysis/analysis.py
+++ b/tirx/python/tvm_tirx/tirx/analysis/analysis.py
@@ -61,10 +61,8 @@ def undefined_vars(node: Stmt | Expr, defs: list[Var] | None = None) -> list[Var
return _ffi_api.UndefinedVars(node, defs) # type: ignore # pylint: disable=no-member
-def verify_well_formed(obj: Function | IRModule, assert_mode: bool = True) -> bool:
- """Verify definitions, tensor-load types, loop control, tile ops, and scope IDs.
-
- Use ``tvm.s_tir.analysis.verify_well_formed`` for schedulable blocks.
+def check_well_formed(obj: Function | IRModule, assert_mode: bool = True) -> bool:
+ """Check whether a TIRx function is well formed.
Parameters
----------
@@ -72,11 +70,16 @@ def verify_well_formed(obj: Function | IRModule, assert_mode: bool = True) -> bo
The function or module to be verified.
assert_mode: bool
- The indicator if it raises an error when the function is not well-formed.
+ Whether to raise on invalid IR. Defaults to True.
Returns
-------
result: bool
- Whether it is a well-formed TIR function.
+ True for valid IR; False for invalid IR when assert_mode is False.
+
+ Raises
+ ------
+ RuntimeError
+ If the IR is invalid and assert_mode is True.
"""
- return _ffi_api.VerifyWellFormed(obj, assert_mode) # type: ignore # pylint: disable=no-member
+ return _ffi_api.CheckWellFormed(obj, assert_mode) # type: ignore # pylint: disable=no-member
diff --git a/tirx/python/tvm_tirx/tirx/layout.py b/tirx/python/tvm_tirx/tirx/layout.py
index 15f0e59424..219e15e400 100644
--- a/tirx/python/tvm_tirx/tirx/layout.py
+++ b/tirx/python/tvm_tirx/tirx/layout.py
@@ -32,15 +32,15 @@ from .thread_scope import ThreadScope
@tvm_ffi.register_object("tirx.Layout")
class Layout(tvm_ffi.Object):
- def verify_well_formed(self) -> bool:
- """Verify if the layout is well-formed.
+ def check_well_formed(self) -> None:
+ """Check whether the layout is well formed.
- Returns
- -------
- bool
- True if the layout is well-formed, False otherwise
+ Raises
+ ------
+ ValueError, RuntimeError
+ If the layout is not well formed.
"""
- return _ffi_api.LayoutVerifyWellFormed(self) # pylint: disable=no-member
+ _ffi_api.LayoutCheckWellFormed(self) # pylint: disable=no-member
def size(self, axis_name: str | None = None):
"""Get the size of the layout.
@@ -263,10 +263,10 @@ class Layout(tvm_ffi.Object):
# Set of axis names registered on the C++ side. Used for lazy resolution of
-# both module-level (`from tvm.tirx.layout import m`) and class-attribute
+# both module-level (`from tvm_tirx.tirx.layout import m`) and class-attribute
# (`Axis.m`) accesses. The actual FFI call to look up each axis is
-# deferred until first access — keeps `import tvm.tirx.layout` runtime-safe
-# (compiler-side FFI need not be present, matching apache's discipline).
+# deferred until first access — keeps `import tvm_tirx.tirx.layout` runtime-safe
+# (compiler-side FFI need not be present).
_AXIS_NAMES = ("m",)
@@ -326,9 +326,9 @@ class Axis(Enum, metaclass=_AxisMeta, type_key="tirx.Axis", init=False):
# ------------------------------------------------------------------
# 2) Lazy module-level axis lookup
# ------------------------------------------------------------------
-# PEP 562 module-level __getattr__ for `from tvm.tirx.layout import m`.
+# PEP 562 module-level __getattr__ for `from tvm_tirx.tirx.layout import m`.
# The FFI call to look up each axis is deferred until first access; bare
-# `import tvm.tirx.layout` performs zero compiler-side FFI calls.
+# `import tvm_tirx.tirx.layout` performs zero compiler-side FFI calls.
def __getattr__(name):
if name in _AXIS_NAMES:
return Axis.get(name)
diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py
index 7f297df198..42387cc3a2 100644
--- a/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py
+++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py
@@ -194,7 +194,7 @@ def func_attr(attrs: dict[str, Any]) -> None:
def check_well_formed_(function: _tir.Function) -> None:
"""Validate a completed TIRx function."""
try:
- _tir.analysis.verify_well_formed(function)
+ _tir.analysis.check_well_formed(function)
except Exception as error:
raise ValueError(
"Program is not well-formed. If this is deliberate, set "
@@ -202,13 +202,6 @@ def check_well_formed_(function: _tir.Function) -> None:
) from error
-def _check_module_well_formed(module: _ir.IRModule) -> None:
- """Validate completed functions belonging to the TIRx dialect."""
- for function in module.functions.values():
- if isinstance(function, _tir.Function) and not function.attrs.get("s_tir", False):
- check_well_formed_(function)
-
-
# --------------------------------------
# Bindings
# --------------------------------------
diff --git a/tirx/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc b/tirx/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
index 62141ad59d..9f6c191140 100644
--- a/tirx/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
+++ b/tirx/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
@@ -326,7 +326,7 @@ ffi::Module BuildAMDGPU(IRModule mod, Target target) {
auto flink = tvm::ffi::Function::GetGlobal("tvm_callback_rocm_link");
TVM_FFI_ICHECK(flink.has_value())
- << "Require tvm_callback_rocm_link to exist, do import tvm.support.rocm";
+ << "Require tvm_callback_rocm_link to exist, do import tvm_tirx.support.rocm";
TVMFFIByteArray arr;
arr.data = &obj[0];
diff --git a/tirx/src/tirx/analysis/check_well_formed.cc b/tirx/src/tirx/analysis/check_well_formed.cc
index 8b9261c368..7be7d44887 100644
--- a/tirx/src/tirx/analysis/check_well_formed.cc
+++ b/tirx/src/tirx/analysis/check_well_formed.cc
@@ -564,7 +564,7 @@ auto WellFormednessVerifier::Visit_(const TensorLoadNode* op) -> Result {
return this->Visit(op->indices);
}
-bool VerifyWellFormedCommon(const ffi::ObjectRef& node, bool assert_mode) {
+bool CheckWellFormedCommon(const ffi::ObjectRef& node, bool assert_mode) {
using Check = WellFormednessVerifier::Check;
for (Check check :
{Check::kVariables, Check::kTensors, Check::kTensorLoadType, Check::kLoopControl}) {
@@ -575,25 +575,25 @@ bool VerifyWellFormedCommon(const ffi::ObjectRef& node, bool assert_mode) {
} // namespace
-bool VerifyWellFormed(const Function& func, bool assert_mode) {
- return VerifyWellFormedCommon(func, assert_mode);
+bool CheckWellFormed(const Function& func, bool assert_mode) {
+ return CheckWellFormedCommon(func, assert_mode);
}
-bool VerifyWellFormed(const IRModule& mod, bool assert_mode) {
- return VerifyWellFormedCommon(mod, assert_mode);
+bool CheckWellFormed(const IRModule& mod, bool assert_mode) {
+ return CheckWellFormedCommon(mod, assert_mode);
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
- "tirx.analysis.VerifyWellFormed", [](const ffi::ObjectRef& obj, bool assert_mode) {
+ "tirx.analysis.CheckWellFormed", [](const ffi::ObjectRef& obj, bool assert_mode) {
if (auto opt = obj.as<Function>()) {
- return VerifyWellFormed(opt.value(), assert_mode);
+ return CheckWellFormed(opt.value(), assert_mode);
} else if (auto opt = obj.as<IRModule>()) {
- return VerifyWellFormed(opt.value(), assert_mode);
+ return CheckWellFormed(opt.value(), assert_mode);
} else {
TVM_FFI_THROW(InternalError)
- << "Expected VerifyWellFormed argument to be a Function or IRModule, but found "
+ << "Expected CheckWellFormed argument to be a Function or IRModule, but found "
<< obj->GetTypeKey();
TVM_FFI_UNREACHABLE();
}
diff --git a/tirx/src/tirx/ir/layout/compose_layout.cc b/tirx/src/tirx/ir/layout/compose_layout.cc
index 5f38dcbdf8..b09095aa8f 100644
--- a/tirx/src/tirx/ir/layout/compose_layout.cc
+++ b/tirx/src/tirx/ir/layout/compose_layout.cc
@@ -308,7 +308,7 @@ ComposeLayout::ComposeLayout(int per_element, int swizzle_len, int atom_len, Til
n->swizzle_len = swizzle_len;
n->atom_len = atom_len;
n->swizzle_inner = swizzle_inner;
- TVM_FFI_ICHECK(n->VerifyWellFormed()) << "ValueError: The compose layout is not well-formed";
+ n->CheckWellFormed();
int swizzle_mask = (1 << swizzle_len) - 1;
n->inner_mask = swizzle_mask;
n->outer_mask = swizzle_mask << atom_len;
@@ -323,11 +323,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-bool ComposeLayoutNode::VerifyWellFormed() const {
- if (!(per_element >= 0 && swizzle_len >= 0 && atom_len >= swizzle_len)) {
- return false;
- }
- return tile_layout->VerifyWellFormed();
+void ComposeLayoutNode::CheckWellFormed() const {
+ TVM_FFI_CHECK(per_element >= 0 && swizzle_len >= 0 && atom_len >= swizzle_len, ValueError)
+ << "A compose layout requires nonnegative per_element and swizzle_len, and "
+ << "atom_len >= swizzle_len";
+ tile_layout->CheckWellFormed();
}
PrimExpr ComposeLayoutNode::GetSize(ffi::Optional<ffi::String> axis_name) const {
diff --git a/tirx/src/tirx/ir/layout/layout.cc b/tirx/src/tirx/ir/layout/layout.cc
index e7c1e733f9..e1135fb4c5 100644
--- a/tirx/src/tirx/ir/layout/layout.cc
+++ b/tirx/src/tirx/ir/layout/layout.cc
@@ -33,7 +33,7 @@ ffi::Map<ffi::String, PrimExpr> LayoutNode::Apply(const ffi::Array<PrimExpr>& co
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
auto def = refl::GlobalDef();
- def.def("tirx.LayoutVerifyWellFormed", [](Layout layout) { return layout->VerifyWellFormed(); });
+ def.def("tirx.LayoutCheckWellFormed", [](Layout layout) { return layout->CheckWellFormed(); });
def.def("tirx.LayoutGetSize", [](Layout layout, ffi::Optional<ffi::String> axis_name) {
return layout->GetSize(axis_name);
});
diff --git a/tirx/src/tirx/ir/layout/tile_core.cc b/tirx/src/tirx/ir/layout/tile_core.cc
index e9ee7cd3b9..12c224a718 100644
--- a/tirx/src/tirx/ir/layout/tile_core.cc
+++ b/tirx/src/tirx/ir/layout/tile_core.cc
@@ -196,12 +196,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-bool TileLayoutNode::VerifyWellFormed() const {
- // 1. Check if the scope is connected
- if (!GetScope().has_value() && HasThreadAxis()) {
- return false;
- }
- return true;
+void TileLayoutNode::CheckWellFormed() const {
+ TVM_FFI_CHECK(GetScope().has_value() || !HasThreadAxis(), ValueError)
+ << "A tile layout with thread axes must have a connected scope";
}
PrimExpr TileLayoutNode::GetSize(ffi::Optional<ffi::String> axis_name) const {
diff --git a/tirx/tests/python/script/test_basic_usage.py b/tirx/tests/python/script/test_basic_usage.py
index 54183f0d33..515e589bec 100644
--- a/tirx/tests/python/script/test_basic_usage.py
+++ b/tirx/tests/python/script/test_basic_usage.py
@@ -291,9 +291,9 @@ def test_loop_control_validation_preserves_valid_and_unchecked_ir():
invalid = tirx.Function(params=[], body=tvm_tirx.ir.Break())
ir.assert_structural_equal(invalid.body, tvm_tirx.ir.SeqStmt([tvm_tirx.ir.Break()]))
- assert not tirx.analysis.verify_well_formed(invalid, assert_mode=False)
+ assert not tirx.analysis.check_well_formed(invalid, assert_mode=False)
with pytest.raises(RuntimeError, match="requires an enclosing loop"):
- tirx.analysis.verify_well_formed(invalid)
+ tirx.analysis.check_well_formed(invalid)
@T.function
def valid():
@@ -308,11 +308,14 @@ def test_loop_control_validation_preserves_valid_and_unchecked_ir():
bad = invalid
assert Unchecked["bad"].same_as(invalid)
- with pytest.raises(ValueError, match="requires an enclosing loop"):
- @I.ir_module(extra_vars={"invalid": invalid})
- class Rejected:
- bad = invalid
+ @I.ir_module(extra_vars={"invalid": invalid})
+ class Explicit:
+ bad = invalid
+
+ assert Explicit["bad"].same_as(invalid)
+ with pytest.raises(RuntimeError, match="requires an enclosing loop"):
+ tirx.analysis.check_well_formed(Explicit)
def test_written_comparison_order(primitive_language):
diff --git a/tirx/tests/python/tirx/analysis/test_check_well_defined.py b/tirx/tests/python/tirx/analysis/test_check_well_defined.py
index ad39f4cee5..bcf8e58f0f 100644
--- a/tirx/tests/python/tirx/analysis/test_check_well_defined.py
+++ b/tirx/tests/python/tirx/analysis/test_check_well_defined.py
@@ -40,8 +40,8 @@ def test_pass_simple():
for i, j in T.grid(128, 128):
C[i, j] = B[i, j] * 2.0
- assert tvm_tirx.tirx.analysis.verify_well_formed(element_wise)
- assert tvm_tirx.tirx.analysis.verify_well_formed(tvm_tirx.IRModule.from_expr(element_wise))
+ assert tvm_tirx.tirx.analysis.check_well_formed(element_wise)
+ assert tvm_tirx.tirx.analysis.check_well_formed(tvm_tirx.IRModule.from_expr(element_wise))
def test_error_for_out_of_scope_usage():
@@ -69,7 +69,7 @@ def test_error_for_out_of_scope_usage():
(ValueError, RuntimeError),
match="Invalid use of undefined variable i at .* no longer in-scope.",
):
- tvm_tirx.tirx.analysis.verify_well_formed(func)
+ tvm_tirx.tirx.analysis.check_well_formed(func)
def test_error_for_nested_rebind_usage():
@@ -87,7 +87,7 @@ def test_error_for_nested_rebind_usage():
(ValueError, RuntimeError),
match="ill-formed, due to multiple nested definitions of variable i",
):
- tvm_tirx.tirx.analysis.verify_well_formed(func)
+ tvm_tirx.tirx.analysis.check_well_formed(func)
def test_error_for_repeated_binding():
@@ -111,7 +111,7 @@ def test_error_for_repeated_binding():
(ValueError, RuntimeError),
match="multiple nested definitions of variable i",
):
- tvm_tirx.tirx.analysis.verify_well_formed(func)
+ tvm_tirx.tirx.analysis.check_well_formed(func)
def test_error_for_cross_function_reuse():
@@ -132,7 +132,7 @@ def test_error_for_cross_function_reuse():
T.evaluate(i)
with pytest.raises((ValueError, RuntimeError), match="multiple definitions of variable i"):
- tvm_tirx.tirx.analysis.verify_well_formed(mod)
+ tvm_tirx.tirx.analysis.check_well_formed(mod)
def test_sibling_launch_bindings_are_well_formed():
@@ -146,7 +146,7 @@ def test_sibling_launch_bindings_are_well_formed():
with T.launch_thread("threadIdx.x", 256) as threadIdx_x:
A[threadIdx_x] = A[threadIdx_x] + 2.0
- tvm_tirx.tirx.analysis.verify_well_formed(func)
+ tvm_tirx.tirx.analysis.check_well_formed(func)
def test_reuse_of_region_parameter_across_functions_is_ill_formed():
@@ -178,7 +178,7 @@ def test_reuse_of_region_parameter_across_functions_is_ill_formed():
(ValueError, RuntimeError),
match="multiple definitions of variable threadIdx_x",
):
- tvm_tirx.tirx.analysis.verify_well_formed(mod)
+ tvm_tirx.tirx.analysis.check_well_formed(mod)
def test_multiple_tensor_arguments_may_share_allocation():
@@ -195,7 +195,7 @@ def test_multiple_tensor_arguments_may_share_allocation():
def func(A: T.Tensor([256], "float32"), B: T.Tensor([256], "float32", data=A.data)): # noqa: F821
pass
- tvm_tirx.tirx.analysis.verify_well_formed(mod)
+ tvm_tirx.tirx.analysis.check_well_formed(mod)
def test_error_message_without_previous_definition_location():
@@ -221,7 +221,7 @@ def test_error_message_without_previous_definition_location():
T.evaluate(x)
with pytest.raises((ValueError, RuntimeError)) as exc_info:
- tvm_tirx.tirx.analysis.verify_well_formed(func, assert_mode=True)
+ tvm_tirx.tirx.analysis.check_well_formed(func, assert_mode=True)
error_msg = str(exc_info.value)
@@ -246,7 +246,7 @@ def test_error_message_with_previous_definition_location():
T.evaluate(x)
with pytest.raises((ValueError, RuntimeError)) as exc_info:
- tvm_tirx.tirx.analysis.verify_well_formed(func, assert_mode=True)
+ tvm_tirx.tirx.analysis.check_well_formed(func, assert_mode=True)
error_msg = str(exc_info.value)
@@ -277,7 +277,7 @@ def test_sequential_redefinition_with_location():
T.evaluate(x)
with pytest.raises((ValueError, RuntimeError)) as exc_info:
- tvm_tirx.tirx.analysis.verify_well_formed(func, assert_mode=True)
+ tvm_tirx.tirx.analysis.check_well_formed(func, assert_mode=True)
error_msg = str(exc_info.value)
@@ -295,7 +295,7 @@ def test_tensor_param_is_well_formed():
for i in T.grid(128):
B[i] = A[i] * 2.0
- tvm_tirx.tirx.analysis.verify_well_formed(func)
+ tvm_tirx.tirx.analysis.check_well_formed(func)
def test_decl_tensor_is_well_formed():
@@ -307,7 +307,7 @@ def test_decl_tensor_is_well_formed():
for i in T.grid(128):
B[i] = A[i] * 2.0
- tvm_tirx.tirx.analysis.verify_well_formed(func)
+ tvm_tirx.tirx.analysis.check_well_formed(func)
def test_alloc_tensor_is_well_formed():
@@ -321,7 +321,7 @@ def test_alloc_tensor_is_well_formed():
for i in T.grid(128):
B[i] = A[i] * 2.0
- tvm_tirx.tirx.analysis.verify_well_formed(mod)
+ tvm_tirx.tirx.analysis.check_well_formed(mod)
def test_tensor_load_asserted_type_matches_source_and_indices():
@@ -332,9 +332,9 @@ def test_tensor_load_asserted_type_matches_source_and_indices():
serialized = tvm_tirx.ir.save_json(func)
round_tripped = tvm_tirx.ir.load_json(serialized)
- assert tvm_tirx.tirx.analysis.verify_well_formed(round_tripped, assert_mode=False)
+ assert tvm_tirx.tirx.analysis.check_well_formed(round_tripped, assert_mode=False)
with ThreadPoolExecutor(max_workers=1) as executor:
- assert executor.submit(tvm_tirx.tirx.analysis.verify_well_formed, func, False).result()
+ assert executor.submit(tvm_tirx.tirx.analysis.check_well_formed, func, False).result()
graph = json.loads(serialized)
load = next(node for node in graph["nodes"] if node["type"] == "ir.TensorLoad")
@@ -346,9 +346,9 @@ def test_tensor_load_asserted_type_matches_source_and_indices():
load["data"]["ty"] = int32_type_index
malformed = tvm_tirx.ir.load_json(json.dumps(graph))
- assert not tvm_tirx.tirx.analysis.verify_well_formed(malformed, assert_mode=False)
+ assert not tvm_tirx.tirx.analysis.check_well_formed(malformed, assert_mode=False)
with pytest.raises(RuntimeError, match="asserts result type"):
- tvm_tirx.tirx.analysis.verify_well_formed(malformed)
+ tvm_tirx.tirx.analysis.check_well_formed(malformed)
def test_tensor_load_malformed_indices_return_false_without_asserting():
@@ -365,24 +365,24 @@ def test_tensor_load_malformed_indices_return_false_without_asserting():
rank_mismatch_indices = rank_mismatch_graph["nodes"][load_node["data"]["indices"]]["data"]
rank_mismatch_indices.pop()
rank_mismatch = tvm_tirx.ir.load_json(json.dumps(rank_mismatch_graph))
- assert not tvm_tirx.tirx.analysis.verify_well_formed(rank_mismatch, assert_mode=False)
+ assert not tvm_tirx.tirx.analysis.check_well_formed(rank_mismatch, assert_mode=False)
with pytest.raises(RuntimeError, match="indexes 2-dimensional tensor"):
- tvm_tirx.tirx.analysis.verify_well_formed(rank_mismatch)
+ tvm_tirx.tirx.analysis.check_well_formed(rank_mismatch)
indices[0], indices[1] = indices[1], indices[0]
non_final_vector = tvm_tirx.ir.load_json(json.dumps(graph))
- assert not tvm_tirx.tirx.analysis.verify_well_formed(non_final_vector, assert_mode=False)
+ assert not tvm_tirx.tirx.analysis.check_well_formed(non_final_vector, assert_mode=False)
with pytest.raises(RuntimeError, match="only the final index"):
- tvm_tirx.tirx.analysis.verify_well_formed(non_final_vector)
+ tvm_tirx.tirx.analysis.check_well_formed(non_final_vector)
def test_verify_unique_definitions():
x = tvm_tirx.tirx.Var("x", "int32")
y = tvm_tirx.tirx.Var("tindex", "int32")
z = tvm_tirx.ir.Evaluate(x + y)
- assert tvm_tirx.tirx.analysis.verify_well_formed(tvm_tirx.tirx.Function([x, y], z))
+ assert tvm_tirx.tirx.analysis.check_well_formed(tvm_tirx.tirx.Function([x, y], z))
- assert not tvm_tirx.tirx.analysis.verify_well_formed(
+ assert not tvm_tirx.tirx.analysis.check_well_formed(
tvm_tirx.tirx.Function([x, y], tvm_tirx.ir.SeqStmt([tvm_tirx.ir.Bind(x, 1), z])),
assert_mode=False,
)
@@ -394,10 +394,10 @@ def test_verify_repeated_let_definitions():
z2 = tvm_tirx.tirx.Let(x, 2, x + 2)
# Each occurrence must define a distinct variable, even for a shared Let.
- assert not tvm_tirx.tirx.analysis.verify_well_formed(
+ assert not tvm_tirx.tirx.analysis.check_well_formed(
tvm_tirx.tirx.Function([], tvm_tirx.ir.Evaluate(z1 + z1)), assert_mode=False
)
- assert not tvm_tirx.tirx.analysis.verify_well_formed(
+ assert not tvm_tirx.tirx.analysis.check_well_formed(
tvm_tirx.tirx.Function([], tvm_tirx.ir.Evaluate(z1 * z2)), assert_mode=False
)
diff --git a/tirx/tests/python/tirx/script/test_basic_usage.py b/tirx/tests/python/tirx/script/test_basic_usage.py
index fdde71cba9..58b590c7c3 100644
--- a/tirx/tests/python/tirx/script/test_basic_usage.py
+++ b/tirx/tests/python/tirx/script/test_basic_usage.py
@@ -579,9 +579,9 @@ def test_loop_control_validation_preserves_valid_and_unchecked_ir():
# Direct construction retains the native statement for an explicit verifier pass.
ir.assert_structural_equal(invalid.body, tvm_tirx.ir.SeqStmt([tvm_tirx.ir.Break()]))
- assert not tirx.analysis.verify_well_formed(invalid, assert_mode=False)
+ assert not tirx.analysis.check_well_formed(invalid, assert_mode=False)
with pytest.raises(RuntimeError, match="requires an enclosing loop"):
- tirx.analysis.verify_well_formed(invalid)
+ tirx.analysis.check_well_formed(invalid)
@T.function
def valid():
@@ -596,11 +596,28 @@ def test_loop_control_validation_preserves_valid_and_unchecked_ir():
bad = invalid
assert Unchecked["bad"].same_as(invalid)
- with pytest.raises(ValueError, match="requires an enclosing loop"):
- @I.ir_module(extra_vars={"invalid": invalid})
- class Rejected:
- bad = invalid
+ @I.ir_module(extra_vars={"invalid": invalid})
+ class Explicit:
+ bad = invalid
+
+ assert Explicit["bad"].same_as(invalid)
+ with pytest.raises(RuntimeError, match="requires an enclosing loop"):
+ tirx.analysis.check_well_formed(Explicit)
+
+
+def test_function_validation_opt_out():
+ source = """
+@T.function(check_well_formed=False)
+def unchecked():
+ break
+"""
+ invalid = tvm_tirx.script.from_source(source, extra_vars={"T": T})
+ assert not tvm_tirx.tirx.analysis.check_well_formed(invalid, assert_mode=False)
+ with pytest.raises(ValueError, match="requires an enclosing loop"):
+ tvm_tirx.script.from_source(
+ source.replace("check_well_formed=False", "check_well_formed=True"), extra_vars={"T": T}
+ )
def test_roundtrip_break_for():
diff --git a/tirx/tests/python/tirx/test_layout.py b/tirx/tests/python/tirx/test_layout.py
index 9f3db370ed..5dba9394e4 100644
--- a/tirx/tests/python/tirx/test_layout.py
+++ b/tirx/tests/python/tirx/test_layout.py
@@ -196,28 +196,28 @@ def test_spec_builder():
assert str(TileLayout(S[1:1] + R[4:"laneid"])) == str(TileLayout(S[1:1] + R[4 : 1 @ laneid]))
-def test_verify_well_formed():
+def test_check_well_formed():
def test_scope_connected():
layout = TileLayout(S[(8, 4, 2) : (4 @ laneid, 1 @ laneid, 1)])
res = layout.get_scope()
assert res is not None
assert res[0].name == "cuda.thread"
assert res[1].name == "cuda.warp"
- assert layout.verify_well_formed()
+ assert layout.check_well_formed() is None
layout = TileLayout(S[8 : 4 @ laneid] + R[4 : 1 @ laneid])
res = layout.get_scope()
assert res is not None
assert res[0].name == "cuda.thread"
assert res[1].name == "cuda.warp"
- assert layout.verify_well_formed()
+ layout.check_well_formed()
layout = TileLayout(S[(8, 4, 2) : (4 @ laneid, 1 @ laneid, 1)])
res = layout.get_scope()
assert res is not None
assert res[0].name == "cuda.thread"
assert res[1].name == "cuda.warp"
- assert layout.verify_well_formed()
+ layout.check_well_formed()
layout = TileLayout(
S[(2, 8, 2, 4, 2) : (2 @ warpid, 4 @ laneid, 1 @ warpid, 1 @ laneid, 1)]
@@ -226,7 +226,7 @@ def test_verify_well_formed():
assert res is not None
assert res[0].name == "cuda.thread"
assert res[1].name == "cuda.cta"
- assert layout.verify_well_formed()
+ layout.check_well_formed()
layout = TileLayout(
S[(2, 8, 2, 4, 2) : (2 @ wid_in_wg, 4 @ laneid, 1 @ wid_in_wg, 1 @ laneid, 1)]
@@ -235,11 +235,11 @@ def test_verify_well_formed():
assert res is not None
assert res[0].name == "cuda.thread"
assert res[1].name == "cuda.warpgroup"
- assert layout.verify_well_formed()
+ layout.check_well_formed()
layout = TileLayout(S[(2, 8, 2, 4, 2) : (2 @ wgid, 4 @ laneid, 1 @ wgid, 1 @ laneid, 1)])
- with pytest.raises(Exception):
- layout.verify_well_formed()
+ with pytest.raises(RuntimeError, match="disconnected scope chain"):
+ layout.check_well_formed()
test_scope_connected()
@@ -435,14 +435,14 @@ def test_normalize_tile_layout():
with tvm_tirx.target.Target("cuda"):
layout = TileLayout(S[(2, 8, 2, 4) : (2 @ warpid, 4 @ laneid, 1 @ warpid, 1 @ laneid)])
layout_expected = TileLayout(S[(2, 8, 2, 4) : (64 @ tx, 4 @ tx, 32 @ tx, 1 @ tx)])
- assert layout.verify_well_formed()
- assert layout_expected.verify_well_formed()
+ layout.check_well_formed()
+ layout_expected.check_well_formed()
assert_structural_equal(layout_expected, layout.canonicalize())
layout = TileLayout(S[(2, 2, 8, 4) : (2 @ warpid, 1 @ warpid, 4 @ laneid, 1 @ laneid)])
layout_expected = TileLayout(S[128 : 1 @ tx])
- assert layout.verify_well_formed()
- assert layout_expected.verify_well_formed()
+ layout.check_well_formed()
+ layout_expected.check_well_formed()
assert_structural_equal(layout_expected, layout.canonicalize())
layout = TileLayout(
@@ -460,8 +460,8 @@ def test_normalize_tile_layout():
layout_expected = TileLayout(
S[(2, 2, 8, 2, 2, 4) : (256 @ tx, 64 @ tx, 4 @ tx, 128 @ tx, 32 @ tx, 1 @ tx)]
)
- assert layout.verify_well_formed()
- assert layout_expected.verify_well_formed()
+ layout.check_well_formed()
+ layout_expected.check_well_formed()
assert_structural_equal(layout_expected, layout.canonicalize())
layout = TileLayout(
@@ -470,16 +470,16 @@ def test_normalize_tile_layout():
layout_expected = TileLayout(
S[(2, 8, 2, 4) : (64 @ tid_in_wg, 4 @ tid_in_wg, 32 @ tid_in_wg, 1 @ tid_in_wg)]
)
- assert layout.verify_well_formed()
- assert layout_expected.verify_well_formed()
+ layout.check_well_formed()
+ layout_expected.check_well_formed()
assert_structural_equal(layout_expected, layout.canonicalize())
layout = TileLayout(
S[(2, 2, 4, 32) : (2 @ wgid, 1 @ wgid, 32 @ tid_in_wg, 1 @ tid_in_wg)]
)
layout_expected = TileLayout(S[512 : 1 @ tx])
- assert layout.verify_well_formed()
- assert layout_expected.verify_well_formed()
+ layout.check_well_formed()
+ layout_expected.check_well_formed()
assert_structural_equal(layout_expected, layout.canonicalize())
case_fuse_axis()
diff --git a/tirx/tests/python/tirx/test_verifier.py b/tirx/tests/python/tirx/test_verifier.py
index c92aed876c..91fb02e9cc 100644
--- a/tirx/tests/python/tirx/test_verifier.py
+++ b/tirx/tests/python/tirx/test_verifier.py
@@ -17,7 +17,7 @@
import pytest
from tvm_tirx.script import tirx as T
-from tvm_tirx.tirx.analysis import verify_well_formed as verify
+from tvm_tirx.tirx.analysis import check_well_formed as verify
def test_root_scope():
diff --git a/tirx/tests/python/tirx/transform/test_common_subexpr_elim.py b/tirx/tests/python/tirx/transform/test_common_subexpr_elim.py
index 02010ce4e5..ecf25aef86 100644
--- a/tirx/tests/python/tirx/transform/test_common_subexpr_elim.py
+++ b/tirx/tests/python/tirx/transform/test_common_subexpr_elim.py
@@ -615,7 +615,7 @@ def test_shared_subtree_stays_ssa():
assert body[1].same_as(body[3])
After = tvm_tirx.tirx.transform.CommonSubexprElim()(unrolled)
tvm_tirx.ir.assert_structural_equal(After, Expected)
- assert tvm_tirx.tirx.analysis.verify_well_formed(After["main"])
+ assert tvm_tirx.tirx.analysis.check_well_formed(After["main"])
if __name__ == "__main__":
diff --git a/tirx/tests/python/tirx/transform/test_convert_ssa.py b/tirx/tests/python/tirx/transform/test_convert_ssa.py
index e550a6bb56..632adbb8ad 100644
--- a/tirx/tests/python/tirx/transform/test_convert_ssa.py
+++ b/tirx/tests/python/tirx/transform/test_convert_ssa.py
@@ -494,7 +494,7 @@ def test_reused_loop_var_in_decl_tensor_elem_offset():
assert Before["main"].body[1].same_as(Before["main"].body[2])
After = tvm_tirx.tirx.transform.ConvertSSA()(Before)
tvm_tirx.ir.assert_structural_equal(After, Expected)
- tvm_tirx.tirx.analysis.verify_well_formed(After["main"], assert_mode=True)
+ tvm_tirx.tirx.analysis.check_well_formed(After["main"], assert_mode=True)
if __name__ == "__main__":
diff --git a/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py b/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
index adc2526ea9..edae515ca5 100644
--- a/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
+++ b/tirx/tests/python/tirx/transform/test_lower_tensor_layouts.py
@@ -1200,7 +1200,7 @@ def test_lower_remaps_tensor_memory_address_metadata(offset):
lowered = LowerTensorLayouts()(Dispatched)["main"]
tvm_tirx.ir.assert_structural_equal(lowered, Expected["main"])
assert not tvm_tirx.tirx.analysis.undefined_vars(lowered.body, lowered.params)
- tvm_tirx.tirx.analysis.verify_well_formed(lowered)
+ tvm_tirx.tirx.analysis.check_well_formed(lowered)
bindings = []
tvm_ffi.structural_walk(lowered.body, (tvm_tirx.ir.Bind, bindings.append))
address = next(
diff --git a/tirx/tests/python/tirx/transform/test_pointer_value_type_rewrite.py b/tirx/tests/python/tirx/transform/test_pointer_value_type_rewrite.py
index 392ad98984..ff29b9778b 100644
--- a/tirx/tests/python/tirx/transform/test_pointer_value_type_rewrite.py
+++ b/tirx/tests/python/tirx/transform/test_pointer_value_type_rewrite.py
@@ -172,7 +172,7 @@ def test_decl_tensor_alias_chain_uses_flat_root_map():
After = transform(Before)
tvm_tirx.ir.assert_structural_equal(After, Expected)
- assert tvm_tirx.tirx.analysis.verify_well_formed(After)
+ assert tvm_tirx.tirx.analysis.check_well_formed(After)
if __name__ == "__main__":
diff --git a/tirx/tests/python/tirx/transform/test_simplify.py b/tirx/tests/python/tirx/transform/test_simplify.py
index 747293cfa1..4fdecdace8 100644
--- a/tirx/tests/python/tirx/transform/test_simplify.py
+++ b/tirx/tests/python/tirx/transform/test_simplify.py
@@ -838,7 +838,7 @@ def test_simplify_tensor_identity_well_formed():
after = _apply_simplify(before)
tvm_tirx.ir.assert_structural_equal(after, expected)
- tvm_tirx.tirx.analysis.verify_well_formed(after)
+ tvm_tirx.tirx.analysis.check_well_formed(after)
def test_tensor_shape_constraint():
diff --git a/tirx/tests/python/tirx/transform/test_split_host_device.py b/tirx/tests/python/tirx/transform/test_split_host_device.py
index 8477f9e2f5..7623b50ebd 100644
--- a/tirx/tests/python/tirx/transform/test_split_host_device.py
+++ b/tirx/tests/python/tirx/transform/test_split_host_device.py
@@ -361,7 +361,7 @@ def test_dynamic_launch_thread():
after = tvm_tirx.tirx.transform.SplitHostDevice()(before)
- tvm_tirx.tirx.analysis.verify_well_formed(after)
+ tvm_tirx.tirx.analysis.check_well_formed(after)
tvm_tirx.ir.assert_structural_equal(expected, after)
diff --git a/tirx/tests/python/tirx/transform/test_storage_rewrite.py b/tirx/tests/python/tirx/transform/test_storage_rewrite.py
index 3f85d07473..363c1f2038 100644
--- a/tirx/tests/python/tirx/transform/test_storage_rewrite.py
+++ b/tirx/tests/python/tirx/transform/test_storage_rewrite.py
@@ -849,7 +849,7 @@ def test_rewrite_decl_tensor():
D[i] = C[i]
After = tvm_tirx.tirx.transform.StorageRewrite()(Before)
- assert tvm_tirx.tirx.analysis.verify_well_formed(After)
+ assert tvm_tirx.tirx.analysis.check_well_formed(After)
tvm_tirx.ir.assert_structural_equal(After, Expected)
@@ -877,7 +877,7 @@ def test_decl_tensor_alias_chain_uses_flat_root():
D[0] = C[0]
After = tvm_tirx.tirx.transform.StorageRewrite()(Before)
- assert tvm_tirx.tirx.analysis.verify_well_formed(After)
+ assert tvm_tirx.tirx.analysis.check_well_formed(After)
tvm_tirx.ir.assert_structural_equal(After, Expected)