This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 116bb48ff5 Support standard optional JIT annotations and NoneType
parameters (#20794)
116bb48ff5 is described below
commit 116bb48ff5a601a113ec0d3072df6603d61b3621
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()`.
---
docs/tirx/native_basics/cuda/functions.rst | 13 ++++-
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 +++++++++++++++++-----
15 files changed, 244 insertions(+), 44 deletions(-)
diff --git a/docs/tirx/native_basics/cuda/functions.rst
b/docs/tirx/native_basics/cuda/functions.rst
index d0c9507706..4b0941b125 100644
--- a/docs/tirx/native_basics/cuda/functions.rst
+++ b/docs/tirx/native_basics/cuda/functions.rst
@@ -184,7 +184,7 @@ device checks (e.g. asserting ``B.shape[0] == n``)::
- ``@Tx.function`` parses the function immediately into a ``Function``. Sizes
are
whatever you wrote — concrete ints, or runtime-symbolic vars (above).
-- ``@Tx.jit`` **defers** parsing until you call ``.specialize(**constexpr)``:
+- ``@Tx.jit`` returns a ``JITFunction`` and **defers** parsing until you call
``.specialize(**constexpr)``:
parameters annotated ``Tx.constexpr`` are baked in as compile-time constants
and
the result is an ordinary ``Function``. Use it when you want sizes/flags
fixed at
compile time (so the compiler can unroll, statically size shared memory,
etc.).
@@ -216,6 +216,17 @@ So: a **symbolic shape** is one kernel whose size is
resolved at run time; a
**constexpr + jit** produces a specialized kernel per value, resolved at
compile
time.
+Optional annotations may use ``Tx.Tensor((N,), "float32") | None`` or
+``typing.Optional[Tx.Tensor((N,), "float32")]`` (including an imported
+``Optional``). The existing ``Tx.Optional(...)`` spelling remains supported.
+Use ``from __future__ import annotations`` for these deferred tensor
annotations;
+they are interpreted as script syntax rather than evaluated as Python unions.
+Calling ``kernel.specialize(optional=None)`` skips that parameter's original
+annotation and binds Python ``None`` in compile-time guards. The resulting
+``Function`` retains the parameter in its original position with
``ir.NoneType``;
+the compiled function requires ``None`` in that argument slot. Omitting the
+specialization keyword keeps the parameter's ordinary runtime type.
+
Launch parameters
-----------------
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"])
