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 <[email protected]> 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 = """ [email protected](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)
