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 e689b275156fcbe7812c720d391d92f61039be33 Author: Tianqi Chen <[email protected]> AuthorDate: Mon Oct 19 11:41:30 2026 -0400 Support standard optional JIT annotations and NoneType parameters (#20794) JIT optional annotations accept `Tensor(...) | None` and imported or qualified `typing.Optional[...]`, alongside `T.Optional(...)`. Selecting `None` preserves the argument's position with shared `ir.NoneType`; the packed API requires actual `None` in that slot, including when the body does not use it. Deferred annotations, compile-time guards, and specialization caching retain their behavior. Rename the deferred kernel class to `JITFunction` while preserving `@T.jit` and `.specialize()`. --- tirx/include/tvm/ir/type.h | 14 +++++ tirx/python/tvm_tirx/ir/__init__.py | 1 + tirx/python/tvm_tirx/ir/type.py | 11 ++++ tirx/python/tvm_tirx/script/parser/jit_support.py | 36 ++++++++++++- tirx/python/tvm_tirx/script/parser/transpile.py | 43 +++++++++++++-- tirx/python/tvm_tirx/tirx/script/jit.py | 46 ++++++++++------ tirx/src/ir/type.cc | 16 ++++++ tirx/src/script/printer/ir/prim_type.cc | 10 ++++ tirx/src/script/printer/script_printer.cc | 1 + tirx/src/tirx/transform/tvm_ffi_binder.cc | 7 +++ tirx/tests/python/ir/test_type.py | 5 ++ .../tests/python/script/test_jit_specialization.py | 20 ++++--- .../python/tirx/script/test_jit_specialization.py | 3 +- tirx/tests/python/tirx/test_jit.py | 62 +++++++++++++++++----- 14 files changed, 232 insertions(+), 43 deletions(-) diff --git a/tirx/include/tvm/ir/type.h b/tirx/include/tvm/ir/type.h index 14b96a443b..b961b9757c 100644 --- a/tirx/include/tvm/ir/type.h +++ b/tirx/include/tvm/ir/type.h @@ -35,6 +35,20 @@ namespace tvm { +/*! \brief The type whose only value is None. */ +class NoneTypeNode : public TypeNode { + public: + static void RegisterReflection() { ffi::reflection::ObjectDef<NoneTypeNode>(); } + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.NoneType", NoneTypeNode, TypeNode); +}; + +/*! \brief Managed reference to NoneTypeNode. */ +class NoneType : public Type { + public: + TVM_TIRX_DLL NoneType(); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(NoneType, Type, NoneTypeNode); +}; + /*! * \brief The top type, which admits any value. * diff --git a/tirx/python/tvm_tirx/ir/__init__.py b/tirx/python/tvm_tirx/ir/__init__.py index f51dae0f41..d10809eab8 100644 --- a/tirx/python/tvm_tirx/ir/__init__.py +++ b/tirx/python/tvm_tirx/ir/__init__.py @@ -29,6 +29,7 @@ from .type import ( AnyType, FuncType, MissingType, + NoneType, OpaqueType, PtrType, PrimType, diff --git a/tirx/python/tvm_tirx/ir/type.py b/tirx/python/tvm_tirx/ir/type.py index 2626e92f1e..ffba9b5b94 100644 --- a/tirx/python/tvm_tirx/ir/type.py +++ b/tirx/python/tvm_tirx/ir/type.py @@ -16,6 +16,7 @@ # under the License. """Unified type system in the project.""" +import keyword import sys from collections.abc import Callable from types import ModuleType @@ -75,6 +76,14 @@ class MissingType(Type): self.__init_handle_by_constructor__(_ffi_api.MissingType) +@tvm_ffi.register_object("ir.NoneType") +class NoneType(Type): + """The concrete type whose only value is None.""" + + def __init__(self): + self.__init_handle_by_constructor__(_ffi_api.NoneType) + + @tvm_ffi.register_object("ir.AnyType") class AnyType(Type): """The top type, which admits any value.""" @@ -281,6 +290,8 @@ def _init_type_api(namespace: str, target_module: str | ModuleType) -> None: if root_index not in info.type_ancestors: continue alias = name[:-4] if name.endswith("Type") else name + if keyword.iskeyword(alias): + alias = name if prefix == namespace: # Reserve unbound dialect names so a later repeat can fill them. dialect_aliases.add(alias) diff --git a/tirx/python/tvm_tirx/script/parser/jit_support.py b/tirx/python/tvm_tirx/script/parser/jit_support.py index 143a5e85a5..5bcc03c7ec 100644 --- a/tirx/python/tvm_tirx/script/parser/jit_support.py +++ b/tirx/python/tvm_tirx/script/parser/jit_support.py @@ -24,11 +24,45 @@ state, and ordinary parsing does not import the TIRx JIT entry point. from __future__ import annotations +import ast +import typing from collections.abc import Iterator, Mapping from contextlib import contextmanager from contextvars import ContextVar from typing import Any + +def optional_annotation_body(node: ast.expr, environment: Mapping[str, Any]) -> ast.expr | None: + """Recognize standard optional syntax without evaluating its contained annotation. + + Only the actual typing.Optional object or typing module is accepted. Attribute + lookup on arbitrary receivers stays opaque, just as for language namespaces. + """ + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr): + if isinstance(node.right, ast.Constant) and node.right.value is None: + inner = node.left + elif isinstance(node.left, ast.Constant) and node.left.value is None: + inner = node.right + else: + return None + if not (isinstance(inner, ast.BinOp) and isinstance(inner.op, ast.BitOr)): + return inner + if isinstance(node, ast.Subscript): + marker = node.value + if isinstance(marker, ast.Name): + optional = environment.get(marker.id) is typing.Optional + else: + optional = ( + isinstance(marker, ast.Attribute) + and isinstance(marker.value, ast.Name) + and environment.get(marker.value.id) is typing + and marker.attr == "Optional" + ) + if optional: + return node.slice + return None + + # Per-execution inputs, never a second owner of native function or parameter state. # Host/JIT values and callable annotation adapters have open-ended Python types. _SPECIALIZATION: ContextVar[tuple[str | None, dict[str, Any]] | None] = ContextVar( @@ -47,7 +81,7 @@ def use_specialization(name: str | None, const_args: Mapping[str, Any] | None) - exact name. None is used when the parsed root is not a function. const_args : Mapping[str, Any] or None Parameter names mapped to fixed values, including explicit None values for - omitted optional parameters. None selects ordinary parsing; an empty + None-typed optional parameters. None selects ordinary parsing; an empty mapping still selects specialization with no compile-time values. The mapping is copied on entry, without copying its values. diff --git a/tirx/python/tvm_tirx/script/parser/transpile.py b/tirx/python/tvm_tirx/script/parser/transpile.py index a6e3a2558f..1980f67db0 100644 --- a/tirx/python/tvm_tirx/script/parser/transpile.py +++ b/tirx/python/tvm_tirx/script/parser/transpile.py @@ -45,6 +45,7 @@ from typing import Any, NamedTuple, NoReturn, TypeVar from . import protocol_registry from .inspect_source import _AnnotationScope +from .jit_support import optional_annotation_body from .prescan import ( BindingKind, PrescanContext, @@ -113,7 +114,7 @@ def _unwrap_optional_annotation( Notes ----- Generated specialization calls this after checking selected values and - absences, so omitted parameters never evaluate their annotation. Ordinary + absences, so None-specialized parameters never evaluate their annotation. Ordinary argument construction delegates annotation validation to the language variant. Adapter lookup and execution exceptions propagate unchanged. """ @@ -1847,6 +1848,10 @@ class IRBuilderTranspiler(ast.NodeTransformer): [], ) else: + if const_args is not None: + inner = optional_annotation_body(annotation, self.module.environment) + if inner is not None: + annotation = inner if name in forward: translated = ast.Name(forward[name], ast.Load()) elif ( @@ -1875,8 +1880,8 @@ class IRBuilderTranspiler(ast.NodeTransformer): ) # x = X.arg_("x", annotation, loc=_L[i]) fallback = self._call_dialect("arg_", [ast.Constant(name), translated], parameter) - # Specialized parameters have no runtime ABI slot. The annotation - # thunk is absent from the selected generated Python branch. + # Constexprs have no ABI slot. Optional None parameters retain a slot + # while binding Python None for later annotations and body guards. value = ( fallback if const_args is None @@ -1884,7 +1889,26 @@ class IRBuilderTranspiler(ast.NodeTransformer): name, fallback, parameter, const_args=const_args ) ) - declaration.append(self._assign(alias, value, parameter)) + if const_args is not None and not is_constexpr: + from tvm_tirx.ir import NoneType + + none_arg = self._call_dialect( + "arg_", + [ast.Constant(name), ast.Call(self._inject(NoneType), [], [])], + parameter, + ) + declaration.append( + ast.copy_location( + ast.If( + value.test, + [ast.Expr(none_arg), self._assign(alias, value.body, parameter)], + [self._assign(alias, value.orelse, parameter)], + ), + parameter, + ) + ) + else: + declaration.append(self._assign(alias, value, parameter)) unbound.discard(name) if alias != name: declaration.append(self._assign(name, ast.Name(alias, ast.Load()), parameter)) @@ -2002,12 +2026,21 @@ class IRBuilderTranspiler(ast.NodeTransformer): ) for parameter in parameters: name = parameter.arg + if name not in constexpr_aliases: + # Every runtime parameter, including NoneType, owns a slot. + body.append( + self._assign( + name, + ast.Call(self._inject(next), [ast.Name(iterator, ast.Load())], []), + parameter, + ) + ) value = ( ast.Name(constexpr_aliases[name], ast.Load()) if name in constexpr_aliases else self._select_specialized_value( name, - ast.Call(self._inject(next), [ast.Name(iterator, ast.Load())], []), + ast.Name(name, ast.Load()), parameter, const_args=const_args, ) diff --git a/tirx/python/tvm_tirx/tirx/script/jit.py b/tirx/python/tvm_tirx/tirx/script/jit.py index 4d9fcdf4dc..0ca3a8bfc8 100644 --- a/tirx/python/tvm_tirx/tirx/script/jit.py +++ b/tirx/python/tvm_tirx/tirx/script/jit.py @@ -20,17 +20,19 @@ from __future__ import annotations import ast import inspect +import typing from collections.abc import Callable, Mapping -from types import FunctionType +from types import FunctionType, UnionType from typing import Any from tvm_tirx.script.ir_builder.ir import constexpr from tvm_tirx.script.parser import protocol_registry +from tvm_tirx.script.parser.jit_support import optional_annotation_body from tvm_tirx.tirx import Function class OptionalAnnotation: - """Retain the runtime annotation of a removable JIT parameter.""" + """Retain the runtime annotation of an optional JIT parameter.""" def __init__(self, annotation: Any) -> None: self.annotation = annotation @@ -66,12 +68,12 @@ def make_jit(builder: object, *, namespace_path: str) -> Callable[..., Any]: private: bool = False, check_well_formed: bool = True, persistent: bool = False, - ) -> TIRJit | Callable[[FunctionType], TIRJit]: + ) -> JITFunction | Callable[[FunctionType], JITFunction]: """Decorator: capture the kernel and defer parsing until ``.specialize()``. Use ``@T.jit`` (instead of ``@T.function``) when the kernel takes - compile-time parameters annotated with ``T.constexpr`` or runtime - parameters that may be removed with ``T.Optional``. The resulting object + compile-time parameters annotated with ``T.constexpr`` or optional runtime + parameters that may specialize to ``ir.NoneType``. The resulting object exposes ``.specialize(**const_args)``, which returns a ``tvm_tirx.tirx.Function``. @@ -90,7 +92,7 @@ def make_jit(builder: object, *, namespace_path: str) -> Callable[..., Any]: Returns ------- - TIRJit or Callable[[types.FunctionType], TIRJit] + JITFunction or Callable[[types.FunctionType], JITFunction] Deferred kernel, or its definition-site decorator when func is None. Raises @@ -125,7 +127,7 @@ def make_jit(builder: object, *, namespace_path: str) -> Callable[..., Any]: @T.jit def guarded( - optional: T.Optional(T.Tensor((1,), "int32")), + optional: T.Tensor((1,), "int32") | None, output: T.Tensor((1,), "int32"), ): if T.constexpr(optional is not None): @@ -137,7 +139,7 @@ def make_jit(builder: object, *, namespace_path: str) -> Callable[..., Any]: absent = guarded.specialize(optional=None) """ - def apply(function: FunctionType) -> TIRJit: + def apply(function: FunctionType) -> JITFunction: from tvm_tirx.script.parser.inspect_source import ( capture_definition_scope, require_definition_site, @@ -153,7 +155,7 @@ def make_jit(builder: object, *, namespace_path: str) -> Callable[..., Any]: definition_scope = capture_definition_scope(frame) finally: del frame - return TIRJit( + return JITFunction( function, check_well_formed, persistent, @@ -167,7 +169,7 @@ def make_jit(builder: object, *, namespace_path: str) -> Callable[..., Any]: return jit -class TIRJit: +class JITFunction: """Top-level kernel decorator with compile-time ``.specialize()`` params. Parses the function body lazily: parsing is deferred until ``.specialize()`` @@ -178,6 +180,11 @@ class TIRJit: Constexpr params are removed from the resulting Function's parameter list; their values are baked into the IR (e.g. into ``T.Tensor((M, K), ...)`` shape annotations and into the body). + + Optional parameters accept ``annotation | None``, ``typing.Optional[annotation]`` + and ``T.Optional(annotation)``. Selecting None preserves their parameter slots + with ``ir.NoneType``; the body sees Python None for compile-time guards. + Use ``from __future__ import annotations`` for deferred tensor annotations. """ def __init__( @@ -256,6 +263,9 @@ class TIRJit: node = ast.parse(ann, mode="eval").body if isinstance(node, ast.Constant) and isinstance(node.value, str): node = ast.parse(node.value, mode="eval").body + if optional_annotation_body(node, annotation_scope) is not None: + optional_names.add(name) + continue key = ( resolve_namespace_key(node, annotation_scope) if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) @@ -279,6 +289,10 @@ class TIRJit: constexpr_defaults[name] = param.default if ann is OptionalAnnotation or isinstance(ann, OptionalAnnotation): optional_names.add(name) + elif typing.get_origin(ann) in (typing.Union, UnionType): + args = typing.get_args(ann) + if len(args) == 2 and type(None) in args: + optional_names.add(name) self.constexpr_names: frozenset[str] = frozenset(constexpr_names) self.constexpr_defaults: dict[str, Any] = constexpr_defaults self.optional_names: frozenset[str] = frozenset(optional_names) @@ -290,9 +304,9 @@ class TIRJit: Parameters ---------- **const_args : Any - One hashable value per ``T.constexpr``-annotated parameter. A - ``T.Optional`` parameter may additionally be supplied as ``None`` - to remove it from the resulting Function ABI. Omitting an + One hashable value per ``T.constexpr``-annotated parameter. An + optional parameter may additionally be supplied as ``None`` + to retain its ABI slot with ``ir.NoneType``. Omitting an optional parameter keeps it as a normal runtime parameter. Returns @@ -311,8 +325,8 @@ class TIRJit: Notes ----- No supplied kwargs still selects JIT construction with an empty map when - no constexpr defaults exist. Selected None optional values omit parameters - and skip their annotations. Source and builder exceptions propagate unchanged. + no constexpr defaults exist. Selected None optional values retain parameters + and skip their original annotations. Source and builder exceptions propagate unchanged. """ from tvm_tirx.script.parser.entry import parse @@ -330,7 +344,7 @@ class TIRJit: } if invalid_optional: raise TypeError( - f"{self.func.__name__}.specialize(): T.Optional parameters only accept None; " + f"{self.func.__name__}.specialize(): optional parameters only accept None; " f"pass actual tensors when calling the compiled kernel (got: {invalid_optional!r})" ) diff --git a/tirx/src/ir/type.cc b/tirx/src/ir/type.cc index d10c33a3fa..1bc6bda096 100644 --- a/tirx/src/ir/type.cc +++ b/tirx/src/ir/type.cc @@ -275,6 +275,22 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("ir.MissingType", []() { return MissingType(); }); } +NoneType::NoneType() : Type(ffi::UnsafeInit{}) { + static const auto singleton = ffi::make_object<NoneTypeNode>(); + data_ = singleton; +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + NoneTypeNode::RegisterReflection(); + refl::TypeAttrDef<NoneTypeNode>() + .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&TypeVisit>()) + .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&TypeMutate>()) + .attr(ffi::type_attr::kStructuralMaybeInplaceMutate, + ffi::FStructuralMutate::FromNative<&TypeMaybeInplaceMutate>()); + refl::GlobalDef().def("ir.NoneType", []() { return NoneType(); }); +} + AnyType::AnyType(Location loc) : Type(ffi::UnsafeInit{}) { ffi::ObjectPtr<AnyTypeNode> n = ffi::make_object<AnyTypeNode>(); n->loc = loc; diff --git a/tirx/src/script/printer/ir/prim_type.cc b/tirx/src/script/printer/ir/prim_type.cc index d1f8f06bcf..610c0ae494 100644 --- a/tirx/src/script/printer/ir/prim_type.cc +++ b/tirx/src/script/printer/ir/prim_type.cc @@ -48,6 +48,16 @@ TVM_FFI_STATIC_INIT_BLOCK() { FDocTranslate::FromNative<&MemoryScopeDocTranslate>()); } +ffi::Optional<ExprDoc> NoneTypeDocTranslate(DocTranslatorObj*, ffi::AnyView, const ffi::Object*) { + return NamespaceDoc("ir")->Attr("NoneType")->Call({}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + ffi::reflection::TypeAttrDef<NoneTypeNode>().attr( + tvm::script::printer::type_attr::kDocTranslate.name, + FDocTranslate::FromNative<&NoneTypeDocTranslate>()); +} + bool HasPrimTypeShorthand(const PrimType& type) { bool supported_element = (type.MatchesCode(kDLInt, kDLUInt) && diff --git a/tirx/src/script/printer/script_printer.cc b/tirx/src/script/printer/script_printer.cc index df4b2d5515..7cd4aa82c2 100644 --- a/tirx/src/script/printer/script_printer.cc +++ b/tirx/src/script/printer/script_printer.cc @@ -44,6 +44,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { RegisterScriptRepr<DataTypeImmNode>(); RegisterScriptRepr<GenericConstNode>(); RegisterScriptRepr<AnyTypeNode>(); + RegisterScriptRepr<NoneTypeNode>(); RegisterScriptRepr<CallNode>(); RegisterScriptRepr<DictAttrsNode>(); RegisterScriptRepr<FloatImmNode>(); diff --git a/tirx/src/tirx/transform/tvm_ffi_binder.cc b/tirx/src/tirx/transform/tvm_ffi_binder.cc index b86ffa5e46..cf64f59c4d 100644 --- a/tirx/src/tirx/transform/tvm_ffi_binder.cc +++ b/tirx/src/tirx/transform/tvm_ffi_binder.cc @@ -538,6 +538,13 @@ void TVMFFIABIBuilder::DecodeParam(int param_index) { ffi::reflection::AccessPath param_path = ffi::reflection::AccessPath::Root()->Extend(AccessStep::ArrayItem(param_index)); + if (param->ty.as<NoneTypeNode>()) { + // A None parameter still occupies an ABI slot, even when the body ignores it. + EmitTypeIndexCheck(param_index, + type_index.as_or_throw<PrimExpr>() == ffi::TypeIndex::kTVMFFINone, "None"); + return; + } + if (param->ty.as<TensorTypeNode>()) { Var handle(param->name + ".handle", PtrType::VoidPtrType()); Expr handle_value = DecodeParamOpaqueHandle(param_index, type_index.as_or_throw<PrimExpr>()); diff --git a/tirx/tests/python/ir/test_type.py b/tirx/tests/python/ir/test_type.py index 1e5b87497e..d0119083fa 100644 --- a/tirx/tests/python/ir/test_type.py +++ b/tirx/tests/python/ir/test_type.py @@ -36,6 +36,11 @@ def test_missing_type(): assert isinstance(missing, tvm_tirx.ir.Type) assert isinstance(missing, tvm_tirx.ir.MissingType) + none = tvm_tirx.ir.NoneType() + assert isinstance(none, tvm_tirx.ir.Type) + assert not tvm_ffi.structural_equal(none, missing) + assert not tvm_ffi.structural_equal(none, tvm_tirx.ir.TupleType([])) + check_json_roundtrip(none) def test_prim_type(): diff --git a/tirx/tests/python/script/test_jit_specialization.py b/tirx/tests/python/script/test_jit_specialization.py index dbc4c5a2de..c9805d5aca 100644 --- a/tirx/tests/python/script/test_jit_specialization.py +++ b/tirx/tests/python/script/test_jit_specialization.py @@ -25,10 +25,11 @@ from __future__ import annotations import gc import weakref from types import SimpleNamespace +from typing import Optional import pytest -from tvm_tirx.ir import assert_structural_equal +from tvm_tirx.ir import NoneType, assert_structural_equal from tvm_tirx.script import ir as I from tvm_tirx.script import tirx as T from tvm_tirx.script.ir_builder import IRBuilder @@ -64,7 +65,6 @@ def test_annotation_constructor_executes_in_required_builder_context(jit_languag def test_optional_annotation_is_evaluated_only_for_present_parameters(jit_language): # A selected absence skips its annotation entirely; a present optional unwraps once. M = jit_language.M - M.Optional = T.Optional calls = [] def annotation(): @@ -72,11 +72,12 @@ def test_optional_annotation_is_evaluated_only_for_present_parameters(jit_langua return M.Tensor((5,)) @M.jit - def function(value: M.Optional(annotation())): + def function(value: Optional[annotation()]): # noqa: UP045 M.record(value) absent = function.specialize(value=None) - assert absent.params == [] and absent.body == [("emit", None)] + assert len(absent.params) == 1 and absent.body == [("emit", None)] + assert isinstance(absent.params[0].args[0], NoneType) assert calls == [] present = function.specialize() assert calls == ["annotation"] @@ -213,14 +214,14 @@ def test_tirx_jit_specializes_captured_shape_and_value(): def test_optional_missing_annotation_stays_lazy(jit_language): M = jit_language.M - M.Optional = T.Optional @M.jit - def function(value: M.Optional(missing_annotation)): # noqa: F821 + def function(value: missing_annotation | None): # noqa: F821 M.record(value) absent = function.specialize(value=None) - assert absent.params == [] and absent.body == [("emit", None)] + assert len(absent.params) == 1 and absent.body == [("emit", None)] + assert isinstance(absent.params[0].args[0], NoneType) with pytest.raises(NameError, match="missing_annotation"): function.specialize() @@ -276,7 +277,8 @@ def kernel(output: Script.Tensor((1,), "int32"), *, value: Script.constexpr, assert kernel.constexpr_names == {"value"} assert kernel.optional_names == {"optional"} result = kernel.specialize(value=3, optional=None) - assert len(result.params) == 1 + assert len(result.params) == 2 + assert isinstance(result.params[1].ty, NoneType) assert int(result.body[0].value) == 3 @@ -301,6 +303,8 @@ def test_postponed_jit_uses_registered_constexpr_syntax(jit_language, monkeypatc def ordinary( value: Ordinary.constexpr, optional: Ordinary.Optional(unavailable()), + standard: Ordinary.Optional[unavailable()], + union: M.int32 | M.float32, alias: M.marker_alias, ): M.record(value) diff --git a/tirx/tests/python/tirx/script/test_jit_specialization.py b/tirx/tests/python/tirx/script/test_jit_specialization.py index 26ba0f0b82..b35b2855b3 100644 --- a/tirx/tests/python/tirx/script/test_jit_specialization.py +++ b/tirx/tests/python/tirx/script/test_jit_specialization.py @@ -19,6 +19,7 @@ from __future__ import annotations from tvm_tirx.ir import assert_structural_equal +from tvm_tirx.script import ir as I from tvm_tirx.script import tirx as T @@ -43,7 +44,7 @@ def test_jit_optional_tensor(): T.evaluate(0) @T.function(private=True) - def absent(): + def absent(value: I.NoneType()): T.evaluate(0) @T.function(private=True) diff --git a/tirx/tests/python/tirx/test_jit.py b/tirx/tests/python/tirx/test_jit.py index 8614952bcd..10275a54ef 100644 --- a/tirx/tests/python/tirx/test_jit.py +++ b/tirx/tests/python/tirx/test_jit.py @@ -19,6 +19,7 @@ from __future__ import annotations +import ctypes import traceback import typing @@ -26,7 +27,9 @@ import pytest import tvm_tirx from tvm_tirx.ir import assert_structural_equal +from tvm_tirx.script import ir as I from tvm_tirx.script import tirx as T +from tvm_tirx.tirx.script.jit import JITFunction def test_int_constexpr_specializes_loop_bound(): @@ -259,17 +262,22 @@ def test_optional_param_present_and_absent_ir(): out[0] = A[0] @T.function(private=True) - def expected_absent(out: T.Tensor((1,), "int32")): + def expected_absent(a: I.NoneType(), out: T.Tensor((1,), "int32")): out[0] = -1 present = kernel.specialize() absent = kernel.specialize(a=None) + assert isinstance(kernel, JITFunction) assert_structural_equal(present, expected_present, map_free_vars=True) assert_structural_equal(absent, expected_absent, map_free_vars=True) assert [param.name for param in present.params] == ["a", "out"] - assert [param.name for param in absent.params] == ["out"] + assert [param.name for param in absent.params] == ["a", "out"] assert all(tvm_tirx.tirx.is_tensor_var(param) for param in present.params) - assert all(tvm_tirx.tirx.is_tensor_var(param) for param in absent.params) + assert isinstance(absent.params[0].ty, tvm_tirx.ir.NoneType) + assert tvm_tirx.tirx.is_tensor_var(absent.params[1]) + assert_structural_equal( + tvm_tirx.script.from_source(absent.script(), extra_vars={"I": I, "T": T}), absent + ) def test_optional_specialization_cache_includes_presence(): @@ -311,13 +319,16 @@ def test_multiple_optional_params_preserve_runtime_order(): ] assert [param.name for param in kernel.specialize(a=None).params] == [ "first", + "a", "scale", "b", "out", ] assert [param.name for param in kernel.specialize(a=None, b=None).params] == [ "first", + "a", "scale", + "b", "out", ] @@ -333,19 +344,28 @@ def test_optional_only_accepts_none_at_specialization_time(): kernel.specialize(a=object()) -def test_only_explicit_t_optional_is_specializable(): +def test_standard_optional_annotations_are_specializable(): + from typing import Optional + @T.jit(private=True) - def typing_optional(a: typing.Optional[T.handle], out_h: T.handle): # noqa: UP045 + def typing_optional( + a: typing.Optional[T.Tensor((1,), "int32")], # noqa: UP045 + out: Optional[T.Tensor((1,), "int32")], # noqa: UP045 + ): T.evaluate(0) @T.jit(private=True) - def union_optional(a: T.handle | None, out_h: T.handle): + def union_optional(a: T.Tensor((1,), "int32") | None, out: T.Tensor((1,), "int32")): T.evaluate(0) - with pytest.raises(TypeError, match="unexpected"): - typing_optional.specialize(a=None) - with pytest.raises(TypeError, match="unexpected"): - union_optional.specialize(a=None) + for kernel in (typing_optional, union_optional): + present = kernel.specialize() + absent = kernel.specialize(a=None) + assert all(tvm_tirx.tirx.is_tensor_var(param) for param in present.params) + assert [param.name for param in absent.params] == ["a", "out"] + assert isinstance(absent.params[0].ty, tvm_tirx.ir.NoneType) + assert tvm_tirx.tirx.is_tensor_var(absent.params[1]) + assert isinstance(typing_optional.specialize(out=None).params[1].ty, tvm_tirx.ir.NoneType) def test_t_optional_is_restricted_to_jit(): @@ -368,7 +388,7 @@ def test_compile_time_if_binding_uses_python_scope(): present = kernel.specialize() absent = kernel.specialize(a=None) assert len(present.params) == 2 - assert len(absent.params) == 1 + assert len(absent.params) == 2 assert sum(tvm_tirx.tirx.is_tensor_var(param) for param in present.params) == 2 assert sum(tvm_tirx.tirx.is_tensor_var(param) for param in absent.params) == 1 @@ -386,7 +406,7 @@ def test_compile_time_bool_ops_and_if_expression_short_circuit(): out[0] = 3 if T.constexpr(a is None) else fail_if_evaluated() absent = kernel.specialize(a=None) - assert [param.name for param in absent.params] == ["out"] + assert [param.name for param in absent.params] == ["a", "out"] def test_runtime_tir_if_cannot_guard_absent_optional_param(): @@ -431,5 +451,23 @@ def test_present_optional_param_still_rejects_ffi_none(): executable(None) +def test_none_specialization_checks_unused_packed_argument(): + @T.jit + def kernel(before: T.int32, value: T.Tensor((1,), "int32") | None, after: T.int32): + T.evaluate(before + after) + + executable = tvm_tirx.tirx.compile( + kernel.specialize(value=None), target="llvm", pipeline="tirx" + ) + executable(1, None, 2) + tensor = tvm_tirx.runtime.empty((1,), "int32") + for value in (0, False, 1.0, "not None", tensor, ctypes.c_void_p()): + with pytest.raises( + TypeError, + match=r"Mismatched type on argument #1[\s\S]*kernel[\s\S]*value:[\s\S]*expected None", + ): + executable(1, value, 2) + + if __name__ == "__main__": pytest.main([__file__, "-v"])
