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 cf3c5aff50 [REFACTOR][IR] Move dialect and codegen registrations to
their owners (#20368)
cf3c5aff50 is described below
commit cf3c5aff50190ded12f68bb0e7658b50f1f66687
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 17 07:10:17 2026 -0400
[REFACTOR][IR] Move dialect and codegen registrations to their owners
(#20368)
Move TIRX-specific intrinsic-lowering registration and TensorMapType out
of core IR, and place runtime-module JSON callbacks beside their codegen
serializers. Remove unused higher-layer includes from IR
implementations.
TensorMapType now belongs to `tvm::tirx` in `tvm/tirx/type.h`, with
`tirx.TensorMapType` registration and `tvm.tirx.TensorMapType` Python
binding. Intrinsic-lowering registration uses
`tirx.RegisterOpLowerIntrinsic` and `tvm.tirx.register_intrin_lowering`.
GlobalVar calls use shared `ir.Call` construction, which preserves known
declared return types while leaving argument-dependent results for
normalization. Explicit constructor return-type overrides remain
supported. Host/device lowering preserves declared results and converts
void at external and packed-call ABI boundaries. Remove
`tvm.tirx.call_tir` and use ordinary function-call syntax in TVMScript.
Relax's distinct `call_tir` operator is unchanged.
---
include/tvm/ir/type.h | 24 --------
include/tvm/tirx/script/builder/ir.h | 3 +-
include/tvm/tirx/type.h | 56 +++++++++++++++++++
python/tvm/ir/__init__.py | 2 +-
python/tvm/ir/expr.py | 42 +++++++-------
python/tvm/ir/op.py | 37 -------------
python/tvm/ir/type.py | 17 ------
python/tvm/target/intrin.py | 3 +-
python/tvm/tirx/__init__.py | 3 +-
python/tvm/tirx/op.py | 59 +++++++++++++-------
python/tvm/tirx/script/parser/parser.py | 10 +---
python/tvm/tirx/type.py | 40 ++++++++++++++
src/ir/expr.cc | 2 -
src/ir/module.cc | 15 -----
src/ir/op.cc | 7 ---
src/ir/transform.cc | 1 -
src/ir/type.cc | 30 ----------
src/target/build_common.h | 3 +-
src/target/codegen.cc | 17 ++++++
src/target/llvm/codegen_llvm.cc | 5 +-
src/target/source/codegen_c.cc | 7 ++-
src/tirx/ir/type.cc | 64 ++++++++++++++++++++++
src/tirx/op/op.cc | 12 +++-
src/tirx/script/printer/expr.cc | 15 ++++-
src/tirx/script/printer/ir.cc | 3 +-
src/tirx/transform/split_host_device.cc | 8 +--
tests/python/tirx-base/test_tir_constructor.py | 1 -
tests/python/tvmscript/test_tvmscript_roundtrip.py | 5 +-
28 files changed, 283 insertions(+), 208 deletions(-)
diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h
index cc4fb6a711..efdd1c384b 100644
--- a/include/tvm/ir/type.h
+++ b/include/tvm/ir/type.h
@@ -179,29 +179,5 @@ class FuncType : public Type {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncType, Type, FuncTypeNode);
};
-/*!
- * \brief The type of tensor map.
- * \sa TensorMapType
- */
-class TensorMapTypeNode : public TypeNode {
- public:
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TensorMapTypeNode>();
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.TensorMapType", TensorMapTypeNode,
TypeNode);
-};
-
-/*!
- * \brief Managed reference to TensorMapTypeNode.
- * \sa TensorMapTypeNode
- */
-class TensorMapType : public Type {
- public:
- TVM_DLL TensorMapType(Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorMapType, Type,
TensorMapTypeNode);
-};
-
} // namespace tvm
#endif // TVM_IR_TYPE_H_
diff --git a/include/tvm/tirx/script/builder/ir.h
b/include/tvm/tirx/script/builder/ir.h
index 57a8d1098f..9e3e78a88c 100644
--- a/include/tvm/tirx/script/builder/ir.h
+++ b/include/tvm/tirx/script/builder/ir.h
@@ -28,6 +28,7 @@
#include <tvm/tirx/op.h>
#include <tvm/tirx/script/builder/frame.h>
#include <tvm/tirx/tile_primitive.h>
+#include <tvm/tirx/type.h>
namespace tvm {
namespace script {
@@ -524,7 +525,7 @@ inline Var Handle(ffi::Optional<PrimType> dtype =
std::nullopt,
return tvm::tirx::Var("", type_annotation);
}
-inline Var TensorMap() { return tvm::tirx::Var("",
PointerType(TensorMapType())); }
+inline Var TensorMap() { return tvm::tirx::Var("",
PointerType(tvm::tirx::TensorMapType())); }
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType)
\
inline PrimExpr FuncName(ffi::Optional<PrimExpr> expr = std::nullopt) {
\
diff --git a/include/tvm/tirx/type.h b/include/tvm/tirx/type.h
new file mode 100644
index 0000000000..20906eeabc
--- /dev/null
+++ b/include/tvm/tirx/type.h
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/tirx/type.h
+ * \brief Types specific to TIRX.
+ */
+#ifndef TVM_TIRX_TYPE_H_
+#define TVM_TIRX_TYPE_H_
+
+#include <tvm/ir/type.h>
+
+namespace tvm::tirx {
+
+/*!
+ * \brief The type of tensor map.
+ * \sa TensorMapType
+ */
+class TensorMapTypeNode : public TypeNode {
+ public:
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<TensorMapTypeNode>();
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TensorMapType", TensorMapTypeNode,
TypeNode);
+};
+
+/*!
+ * \brief Managed reference to TensorMapTypeNode.
+ * \sa TensorMapTypeNode
+ */
+class TensorMapType : public Type {
+ public:
+ TVM_DLL TensorMapType(Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorMapType, Type,
TensorMapTypeNode);
+};
+
+} // namespace tvm::tirx
+#endif // TVM_TIRX_TYPE_H_
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index a02e77cdd6..af039e0b9c 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -53,6 +53,6 @@ from . import prim
from .function import BaseFunc, CallingConv
from .global_info import GlobalInfo
from .module import IRModule
-from .op import Op, register_intrin_lowering, register_op_attr
+from .op import Op, register_op_attr
from tvm_ffi import Array, Map
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 2ae414f758..993d13bfc5 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -16,8 +16,6 @@
# under the License.
"""Common expressions data structures in the IR."""
-from numbers import Number
-
import tvm_ffi
import tvm
@@ -107,23 +105,7 @@ class GlobalVar(Expr):
call: Expr
A call taking the variable as a function.
"""
- from .type import PointerType
-
- def is_tir_arg(x):
- return (
- isinstance(x, Number)
- or is_prim_expr(x)
- or (isinstance(x, Expr) and isinstance(x.ty, PointerType))
- )
-
- if args and all(is_tir_arg(x) for x in args):
- return tvm.tirx.call_tir(self, *args)
-
- if all(isinstance(x, Expr) for x in args):
- return Call(self, args)
-
- arg_types = [type(x) for x in args]
- raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for
types {arg_types}")
+ return Call(self, args)
class ExprOperand:
@@ -467,7 +449,12 @@ class TensorLoad(_CallableExprWithOp):
@tvm_ffi.register_object("ir.Call")
class Call(_CallableExprWithOp):
- """Core function call node."""
+ """Core function call node.
+
+ When ``ret_ty`` is omitted, use the callee signature's declared return type
+ if available, or a missing type otherwise. Argument-dependent signatures
+ retain a missing type for subsequent normalization.
+ """
op: Expr
args: list[Expr]
@@ -487,14 +474,25 @@ class Call(_CallableExprWithOp):
# pylint: disable=import-outside-toplevel
from .attrs import DictAttrs
from .op import Op
- from .type import PointerType, PrimType, Type
+ from .type import PointerType, PrimType, TupleType, Type
if isinstance(op, str):
op = Op.get(op)
if attrs is not None and isinstance(attrs, dict):
attrs = DictAttrs(attrs)
if ret_ty is None:
- ret_ty = Type.missing()
+ # Reuse a declared signature without invoking dialect-specific
inference.
+ signature = getattr(op, "ty", None)
+ ret_ty = getattr(signature, "ret_type", None)
+ if not isinstance(ret_ty, Type):
+ ret_ty = getattr(signature, "ret", None)
+ # Rich signatures may specialize their result using arguments.
+ # Reuse only fixed shared scalar, pointer, or void results
here.
+ is_fixed_result = isinstance(ret_ty, PrimType | PointerType)
or (
+ isinstance(ret_ty, TupleType) and not ret_ty.fields
+ )
+ if not is_fixed_result or getattr(signature, "derive_func",
None) is not None:
+ ret_ty = Type.missing()
if isinstance(ret_ty, str) and ret_ty == "handle":
ret_ty = PointerType(PrimType("void"))
elif ret_ty is not None and not isinstance(ret_ty, Type):
diff --git a/python/tvm/ir/op.py b/python/tvm/ir/op.py
index 02034fdf99..e1de79d9e1 100644
--- a/python/tvm/ir/op.py
+++ b/python/tvm/ir/op.py
@@ -187,40 +187,3 @@ def register_op_attr(op_name, attr_key, value=None,
level=10):
return v
return _register(value) if value is not None else _register
-
-
-def register_intrin_lowering(
- op_name,
- target,
- *,
- f=None,
- level=10,
-):
- """Register Op lowering function
-
- Parameters
- ----------
- op_name : str
- The op name
-
- target : str
- The target string for given intrinsic lowering function
-
- f : function, optional
- The function to be registered.
-
- level : int
- The priority level
-
- Returns
- -------
- fregister : function
- Register op lowering function if f is not specified.
- """
-
- def _register(f):
- """internal register function"""
- _ffi_api.RegisterOpLowerIntrinsic(op_name, f, target, level)
- return f
-
- return _register(f) if f is not None else _register
diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py
index 015232963e..d9f6dbc707 100644
--- a/python/tvm/ir/type.py
+++ b/python/tvm/ir/type.py
@@ -160,20 +160,3 @@ class FuncType(Type):
arg_types,
ret_type,
)
-
-
-@tvm_ffi.register_object("ir.TensorMapType")
-class TensorMapType(Type):
- """TensorMapType used in the low-level TIR.
-
- Parameters
- ----------
- span : tvm.ir.Span
- The span information.
- """
-
- def __init__(self, span=None):
- self.__init_handle_by_constructor__(
- _ffi_api.TensorMapType,
- span, # pylint: disable=no-member
- )
diff --git a/python/tvm/target/intrin.py b/python/tvm/target/intrin.py
index 70c9a5fc2f..11ec8c33cc 100644
--- a/python/tvm/target/intrin.py
+++ b/python/tvm/target/intrin.py
@@ -16,8 +16,7 @@
# under the License.
"""Target dependent intrinsic registration."""
-from tvm.ir import register_intrin_lowering
-from tvm.tirx import call_pure_extern
+from tvm.tirx import call_pure_extern, register_intrin_lowering
def _rule_float_suffix(op):
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index e524a749c9..af4b7161b0 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -35,6 +35,7 @@ from .buffer import (
decl_buffer,
is_buffer_var,
)
+from .type import TensorMapType
from .expr import convert
from .expr import Var, Reduce, FloatImm, IntImm, StringImm, Cast
from .expr import Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod
@@ -57,7 +58,7 @@ from .tile_primitive import DispatchContext, LambdaExpr,
TilePrimitiveCall
from .function import PrimFunc, TensorIntrin, IndexMap
-from .op import call_packed_lowered, call_cpacked_lowered, call_tir
+from .op import call_packed_lowered, call_cpacked_lowered,
register_intrin_lowering
from .op import call_packed, call_cpacked, call_intrin, call_pure_extern,
call_extern
from .op import call_llvm_intrin, call_llvm_pure_intrin, all, any, min_value,
max_value, trace
from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index cfb875befc..1b321c9c28 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -29,12 +29,12 @@ from tvm.ir import Call, Expr, ExprWithOp, Op, PointerType,
PrimType, TensorLoad
from tvm.ir.base import Span
from tvm.ir.prim import clz as clz
from tvm.ir.prim import max_value, min_value
-from tvm.ir.type import TensorMapType
from tvm.runtime import const
from . import _ffi_api
from .buffer import Buffer, buffer_data, is_buffer_var
from .expr import BufferLoad, CommReducer, ExprOp, IntImm, Var
+from .type import TensorMapType
tir = tirx # alias for backward compat with upstream tir.convert() calls
@@ -50,6 +50,43 @@ _DEVICE_INTRIN_PREFIX_TO_NAMESPACE = {
}
+def register_intrin_lowering(
+ op_name,
+ target,
+ *,
+ f=None,
+ level=10,
+):
+ """Register Op lowering function
+
+ Parameters
+ ----------
+ op_name : str
+ The op name
+
+ target : str
+ The target string for given intrinsic lowering function
+
+ f : function, optional
+ The function to be registered.
+
+ level : int
+ The priority level
+
+ Returns
+ -------
+ fregister : function
+ Register op lowering function if f is not specified.
+ """
+
+ def _register(f):
+ """internal register function"""
+ _ffi_api.RegisterOpLowerIntrinsic(op_name, f, target, level)
+ return f
+
+ return _register(f) if f is not None else _register
+
+
def _canonical_device_intrin_name(func_name: str) -> str:
"""Return the canonical registry name for statically registered device
intrinsics."""
@@ -547,26 +584,6 @@ def undef():
return call_intrin("int32", "tirx.undef")
-def call_tir(global_var: tvm.ir.GlobalVar, *args):
- """Performs a call into another PrimFunc in the same IRModule
-
- Returns
- -------
- call : Expr
- The call expression.
- """
- assert isinstance(global_var, tvm.ir.GlobalVar)
- args = tuple(_reject_buffer_region(arg, "call_tir") for arg in args)
-
- dtype = "void"
- if global_var.ty is not None:
- ret_ty = global_var.ty.ret
- if isinstance(ret_ty, tvm.ir.PrimType):
- dtype = ret_ty
-
- return Call(op=global_var, args=args, ret_ty=dtype)
-
-
def start_profile_intrinsic(id):
"""Start profile intrinsic.
Parameters
diff --git a/python/tvm/tirx/script/parser/parser.py
b/python/tvm/tirx/script/parser/parser.py
index 607988a4d1..e9c9aafc43 100644
--- a/python/tvm/tirx/script/parser/parser.py
+++ b/python/tvm/tirx/script/parser/parser.py
@@ -1032,14 +1032,8 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) ->
None:
elif isinstance(res, int | bool):
T.evaluate(tvm.tirx.const(res))
elif isinstance(res, tvm.ir.Call) and not tvm.ir.is_prim_expr(res):
- if isinstance(res.op, tvm.ir.GlobalVar) and res.ty.is_missing():
- # GlobalVar calls with a missing return type are ambiguous, as
each IR has a
- # different function Call representation. Convert to the TIR
representation.
- T.evaluate(tvm.tirx.call_tir(res.op, *res.args))
- else:
- # Pointer-valued TIR calls are general Expr rather than Expr,
- # but are still valid standalone Evaluate statements.
- T.evaluate(res)
+ # Non-primitive calls are still valid standalone Evaluate statements.
+ T.evaluate(res)
elif isinstance(res, str):
# Ignore docstrings
pass
diff --git a/python/tvm/tirx/type.py b/python/tvm/tirx/type.py
new file mode 100644
index 0000000000..7b03990e89
--- /dev/null
+++ b/python/tvm/tirx/type.py
@@ -0,0 +1,40 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Types specific to TIRX."""
+
+import tvm_ffi
+
+from tvm.ir import Type
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("tirx.TensorMapType")
+class TensorMapType(Type):
+ """TensorMapType used in the low-level TIR.
+
+ Parameters
+ ----------
+ span : tvm.ir.Span
+ The span information.
+ """
+
+ def __init__(self, span=None):
+ self.__init_handle_by_constructor__(
+ _ffi_api.TensorMapType,
+ span, # pylint: disable=no-member
+ )
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 0f4502271c..cf2ffbeb26 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -21,7 +21,6 @@
* \file src/ir/expr.cc
* \brief The expression AST nodes for the common IR infra.
*/
-#include <tvm/arith/analyzer.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
@@ -32,7 +31,6 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
#include <tvm/ir/type.h>
-#include <tvm/te/tensor.h>
#include <cmath>
#include <utility>
diff --git a/src/ir/module.cc b/src/ir/module.cc
index 49f607f6c5..f2b63bf937 100644
--- a/src/ir/module.cc
+++ b/src/ir/module.cc
@@ -22,15 +22,12 @@
*/
#include <tvm/ffi/cast.h>
#include <tvm/ffi/container/variant.h>
-#include <tvm/ffi/extra/base64.h>
-#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/rvalue_ref.h>
#include <tvm/ir/module.h>
#include <tvm/ir/unique_name_supply.h>
-#include <tvm/target/codegen.h>
#include <algorithm>
#include <fstream>
@@ -258,18 +255,6 @@ IRModule IRModule::FromExpr(const Expr& expr,
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::TypeAttrDef<ffi::ModuleObj>()
- .def("__data_to_json__",
- [](const ffi::ModuleObj* node) {
- std::string bytes =
codegen::SerializeModuleToBytes(ffi::GetRef<ffi::Module>(node),
-
/*export_dso*/ false);
- return ffi::Base64Encode(ffi::Bytes(bytes));
- })
- .def("__data_from_json__", [](const ffi::String& base64_bytes) {
- ffi::Bytes bytes = ffi::Base64Decode(base64_bytes);
- ffi::Module rtmod = codegen::DeserializeModuleFromBytes(bytes.operator
std::string());
- return rtmod;
- });
refl::GlobalDef()
.def("ir.Module_Clone",
[](IRModule mod) -> IRModule {
diff --git a/src/ir/op.cc b/src/ir/op.cc
index 4b50f06ab1..8f7897e502 100644
--- a/src/ir/op.cc
+++ b/src/ir/op.cc
@@ -27,7 +27,6 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
#include <tvm/ir/type.h>
-#include <tvm/tirx/op_attr_types.h>
#include <memory>
@@ -57,7 +56,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ArgumentInfoNode::RegisterReflection(); }
using ffi::Any;
using ffi::Function;
using ffi::PackedArgs;
-using tirx::FLowerIntrinsic;
using OpRegistry = AttrRegistry<OpRegEntry, Op>;
@@ -182,11 +180,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
} else {
reg.set_attr(attr_key, value, plevel);
}
- })
- .def("ir.RegisterOpLowerIntrinsic",
- [](ffi::String name, ffi::Function f, ffi::String target, int
plevel) {
- tvm::OpRegEntry::RegisterOrGet(name).set_attr<FLowerIntrinsic>(
- target + ".FLowerIntrinsic", f, plevel);
});
// override OpNode to use name as the repr
}
diff --git a/src/ir/transform.cc b/src/ir/transform.cc
index 42a5741b0c..d33895c997 100644
--- a/src/ir/transform.cc
+++ b/src/ir/transform.cc
@@ -28,7 +28,6 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/rvalue_ref.h>
#include <tvm/ir/transform.h>
-#include <tvm/relax/expr.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/logging.h>
diff --git a/src/ir/type.cc b/src/ir/type.cc
index dd28b98e55..6e0ce2ab8d 100644
--- a/src/ir/type.cc
+++ b/src/ir/type.cc
@@ -229,18 +229,6 @@ TVMFFIAny
TupleTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
return ffi::Unchanged().CopyToTVMFFIAny();
}
-TVMFFIAny TensorMapTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView)
noexcept {
- return ffi::AnyView(nullptr).CopyToTVMFFIAny();
-}
-
-TVMFFIAny TensorMapTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView)
noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
-}
-
-TVMFFIAny TensorMapTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*,
ffi::AnyView) noexcept {
- return ffi::Unchanged().CopyToTVMFFIAny();
-}
-
} // namespace
Type Type::Missing() {
@@ -428,22 +416,4 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TupleType TupleType::Empty() { return TupleType(ffi::Array<Type>()); }
-TensorMapType::TensorMapType(Span span) : Type(ffi::UnsafeInit{}) {
- ffi::ObjectPtr<TensorMapTypeNode> n = ffi::make_object<TensorMapTypeNode>();
- n->span = std::move(span);
- data_ = std::move(n);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- TensorMapTypeNode::RegisterReflection();
- refl::TypeAttrDef<TensorMapTypeNode>()
- .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&TensorMapTypeVisit))
- .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&TensorMapTypeMutate))
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- reinterpret_cast<void*>(&TensorMapTypeMaybeInplaceMutate));
-
- refl::GlobalDef().def("ir.TensorMapType", [](Span span) { return
TensorMapType(span); });
-}
-
} // namespace tvm
diff --git a/src/target/build_common.h b/src/target/build_common.h
index aba8f0796d..a12cca310a 100644
--- a/src/target/build_common.h
+++ b/src/target/build_common.h
@@ -32,6 +32,7 @@
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt.h>
+#include <tvm/tirx/type.h>
#include <string>
@@ -64,7 +65,7 @@ inline ffi::Map<ffi::String, runtime::FunctionInfo>
ExtractFuncInfo(const IRModu
if (type == nullptr) {
return false;
}
- return type->element_type.as<TensorMapTypeNode>() != nullptr;
+ return type->element_type.as<tirx::TensorMapTypeNode>() != nullptr;
};
arg_extra_tags.push_back(is_tensormap(f->params[i]) ?
runtime::ArgExtraTags::kTensorMap
:
runtime::ArgExtraTags::kNone);
diff --git a/src/target/codegen.cc b/src/target/codegen.cc
index 4437667c70..de16eb1e02 100644
--- a/src/target/codegen.cc
+++ b/src/target/codegen.cc
@@ -21,6 +21,7 @@
* \file codegen.cc
* \brief Common utilities to generated C style code.
*/
+#include <tvm/ffi/extra/base64.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
@@ -347,6 +348,22 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("target.Build", Build);
}
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::TypeAttrDef<ffi::ModuleObj>()
+ .def("__data_to_json__",
+ [](const ffi::ModuleObj* node) {
+ std::string bytes =
codegen::SerializeModuleToBytes(ffi::GetRef<ffi::Module>(node),
+
/*export_dso*/ false);
+ return ffi::Base64Encode(ffi::Bytes(bytes));
+ })
+ .def("__data_from_json__", [](const ffi::String& base64_bytes) {
+ ffi::Bytes bytes = ffi::Base64Decode(base64_bytes);
+ ffi::Module rtmod = codegen::DeserializeModuleFromBytes(bytes.operator
std::string());
+ return rtmod;
+ });
+}
+
// Export a few auxiliary function to the runtime namespace.
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index f70a4f9ed9..92892ec7bf 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -82,6 +82,7 @@
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/op.h>
+#include <tvm/tirx/type.h>
#include <algorithm>
#include <functional>
@@ -643,14 +644,14 @@ llvm::Type* CodeGenLLVM::GetLLVMType(const Type& type)
const {
if (PrimType(primtype->dtype).IsVoid()) {
return t_void_p_;
}
- } else if (ptr->element_type->IsInstance<TensorMapTypeNode>()) {
+ } else if (ptr->element_type->IsInstance<tirx::TensorMapTypeNode>()) {
return llvmGetPointerTo(t_tvm_tensormap_, 0);
}
// TODO(tvm-team) consider put storage scope into the pointer type.
return llvmGetPointerTo(GetLLVMType(ptr->element_type),
GetGlobalAddressSpace());
} else if (IsVoidType(type)) {
return t_void_;
- } else if (type->IsInstance<TensorMapTypeNode>()) {
+ } else if (type->IsInstance<tirx::TensorMapTypeNode>()) {
return t_tvm_tensormap_;
} else {
TVM_FFI_THROW(InternalError) << "Type " << type << " does not have a
corresponding LLVM Type";
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 5086dbe0aa..4cca7b6255 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -25,6 +25,7 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
#include <tvm/ir/unique_name_supply.h>
+#include <tvm/tirx/type.h>
#include <cctype>
#include <iomanip>
@@ -99,7 +100,7 @@ void CodeGenC::PrintFunctionSignature(const ffi::String&
function_name, const Pr
auto is_tensormap_ptr = [&]() -> bool {
if (auto* ptr = v->ty.as<PointerTypeNode>()) {
- return ptr->element_type.as<TensorMapTypeNode>();
+ return ptr->element_type.as<tirx::TensorMapTypeNode>();
}
return false;
};
@@ -112,7 +113,7 @@ void CodeGenC::PrintFunctionSignature(const ffi::String&
function_name, const Pr
bool no_alias = func->HasNonzeroAttr(tirx::attr::kNoAlias);
bool is_handle = v->ty.as<PointerTypeNode>();
auto* ptr = v->ty.as<PointerTypeNode>();
- if (ptr && ptr->element_type.as<TensorMapTypeNode>()) {
+ if (ptr && ptr->element_type.as<tirx::TensorMapTypeNode>()) {
is_handle = false;
}
if (no_alias && is_handle) {
@@ -799,7 +800,7 @@ void CodeGenC::Dispatch_(const CallNode* op, std::ostream&
os) { // NOLINT(*)
<< "Builtin address_of() expects the argument to be a TensorLoad
or Var, but "
<< "received argument " << op->args[0];
if (auto* ptr = var->ty.as<PointerTypeNode>()) {
- if (ptr->element_type.as<TensorMapTypeNode>()) {
+ if (ptr->element_type.as<tirx::TensorMapTypeNode>()) {
os << "((unsigned long long)(&(";
this->PrintExpr(op->args[0], os);
os << ")))";
diff --git a/src/tirx/ir/type.cc b/src/tirx/ir/type.cc
new file mode 100644
index 0000000000..704c60a40f
--- /dev/null
+++ b/src/tirx/ir/type.cc
@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tirx/ir/type.cc
+ * \brief Types specific to TIRX.
+ */
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/tirx/type.h>
+
+namespace tvm::tirx {
+namespace {
+
+TVMFFIAny TensorMapTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView)
noexcept {
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny TensorMapTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView)
noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny TensorMapTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*,
ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+} // namespace
+
+TensorMapType::TensorMapType(Span span) : Type(ffi::UnsafeInit{}) {
+ ffi::ObjectPtr<TensorMapTypeNode> n = ffi::make_object<TensorMapTypeNode>();
+ n->span = std::move(span);
+ data_ = std::move(n);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ TensorMapTypeNode::RegisterReflection();
+ refl::TypeAttrDef<TensorMapTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&TensorMapTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&TensorMapTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&TensorMapTypeMaybeInplaceMutate));
+
+ refl::GlobalDef().def("tirx.TensorMapType", [](Span span) { return
TensorMapType(span); });
+}
+
+} // namespace tvm::tirx
diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc
index 828568d16f..8c521b88c3 100644
--- a/src/tirx/op/op.cc
+++ b/src/tirx/op/op.cc
@@ -34,6 +34,7 @@
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/op_attr_types.h>
+#include <tvm/tirx/type.h>
#include <tvm/tirx/var.h>
#include <cmath>
@@ -104,7 +105,7 @@ Type GetType(const PrimExpr& expr) {
if (auto var = address_of->args[0].as<Var>()) {
if (auto* ptr = var.value()->ty.as<PointerTypeNode>()) {
- if (ptr->element_type.as<TensorMapTypeNode>()) {
+ if (ptr->element_type.as<tirx::TensorMapTypeNode>()) {
return PrimType::UInt(64);
}
}
@@ -149,6 +150,15 @@ PrimExpr break_loop(Span span) {
.as_or_throw<PrimExpr>();
}
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::GlobalDef().def("tirx.RegisterOpLowerIntrinsic",
+ [](ffi::String name, ffi::Function f, ffi::String
target, int plevel) {
+
OpRegEntry::RegisterOrGet(name).set_attr<tirx::FLowerIntrinsic>(
+ target + ".FLowerIntrinsic", f, plevel);
+ });
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index c9174c2040..ab29eb1623 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -19,6 +19,7 @@
#include <tvm/ir/prim/builtin.h>
#include <tvm/te/operation.h>
#include <tvm/tirx/builtin.h>
+#include <tvm/tirx/type.h>
#include "./utils.h"
@@ -72,7 +73,7 @@ ExprDoc PrintVarCreation(const tirx::Var& var, const
AccessPath& var_p, const IR
kwargs_keys, kwargs_values);
}
}
- } else if (ptr_type->element_type->IsInstance<TensorMapTypeNode>()) {
+ } else if (ptr_type->element_type->IsInstance<tirx::TensorMapTypeNode>()) {
rhs = TIR(d, "TensorMap")->Call({}, {}, {});
}
} else {
@@ -348,6 +349,16 @@ Doc PrintTIRCall(Call call, AccessPath call_p, IRDocsifier
d) {
"types, but got "
<< call->ty;
};
+ auto get_call_return_type_doc = [&]() -> ExprDoc {
+ if (call->ty.IsMissing()) {
+ return IdDoc("tvm")->Attr("ir")->Attr("Type")->Attr("missing")->Call({});
+ }
+ if (call_prim_type || call->ty.as<PointerTypeNode>()) {
+ return get_call_type_doc(call_p->Attr("ty"));
+ }
+ // Annotation spellings such as None for an empty tuple are not type
values.
+ return d->AddMetadata(call->ty);
+ };
if (call->attrs.defined()) {
ffi::Array<ExprDoc> call_args;
int n_args = call->args.size();
@@ -358,7 +369,7 @@ Doc PrintTIRCall(Call call, AccessPath call_p, IRDocsifier
d) {
ExprDoc op_doc = call->op.as<Op>()
? LiteralDoc::Str(call->op.as<Op>().value()->name,
call_p->Attr("op"))
: d->AsDoc<ExprDoc>(call->op, call_p->Attr("op"));
- ExprDoc ret_ty_doc = get_call_type_doc(call_p->Attr("ty"));
+ ExprDoc ret_ty_doc = get_call_return_type_doc();
return TIR(d, "Call")->Call(
{op_doc, ListDoc(call_args)}, {"attrs", "ret_ty"},
{d->AsDoc<ExprDoc>(call->attrs, call_p->Attr("attrs")), ret_ty_doc});
diff --git a/src/tirx/script/printer/ir.cc b/src/tirx/script/printer/ir.cc
index 462382864a..e0aa4f8202 100644
--- a/src/tirx/script/printer/ir.cc
+++ b/src/tirx/script/printer/ir.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/target/target.h>
+#include <tvm/tirx/type.h>
#include "./utils.h"
@@ -89,7 +90,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
element_type = LiteralDoc::DataType(prim_type->dtype, //
ty_p->Attr("element_type")->Attr("dtype"));
- } else if (ty->element_type.as<TensorMapTypeNode>()) {
+ } else if (ty->element_type.as<tirx::TensorMapTypeNode>()) {
return TIR(d, "TensorMap")->Call({});
} else {
element_type = d->AsDoc<ExprDoc>(ty->element_type,
ty_p->Attr("element_type"));
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index 9804f2d610..51acd55019 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -337,8 +337,7 @@ class HostDeviceSplitter : public StmtExprMutator {
assert_success});
} else {
- return Evaluate(
- Call(PrimType::Void(), kernel_symbol_global,
call_args).as_or_throw<PrimExpr>());
+ return Evaluate(Call(kernel_ret_type, kernel_symbol_global, call_args));
}
}
@@ -811,7 +810,8 @@ class DeviceKernelMutator : public StmtExprMutator {
for (const Expr& arg : node->args) {
args.push_back(arg);
}
- return Call(node->ty, builtin::call_extern(), args);
+ Type ret_ty = IsVoidType(node->ty) ? PrimType::Void() : node->ty;
+ return Call(ret_ty, builtin::call_extern(), args);
}
}
@@ -856,7 +856,7 @@ class DeviceKernelMutator : public StmtExprMutator {
.as_or_throw<PrimExpr>());
}
- PrimType node_ty = node->ty.as_or_throw<PrimType>();
+ PrimType node_ty = IsVoidType(node->ty) ? PrimType::Void() :
node->ty.as_or_throw<PrimType>();
PrimType ret_ty = node_ty.IsVoid() ? PrimType::Int(32) : node_ty;
return Call(ret_ty, builtin::tvm_call_packed(),
call_args).as_or_throw<PrimExpr>();
diff --git a/tests/python/tirx-base/test_tir_constructor.py
b/tests/python/tirx-base/test_tir_constructor.py
index 05e91a2591..88d0f36964 100644
--- a/tests/python/tirx-base/test_tir_constructor.py
+++ b/tests/python/tirx-base/test_tir_constructor.py
@@ -217,7 +217,6 @@ def test_buffer_region_call_wrappers_reject():
lambda: tvm.tirx.call_cpacked("consume", region, 0),
lambda: tvm.tirx.call_packed_lowered("consume", region),
lambda: tvm.tirx.call_cpacked_lowered("consume", region, 0),
- lambda: tvm.tirx.call_tir(tvm.ir.GlobalVar("callee"), region),
lambda: tvm.tirx.trace([region]),
lambda: T.evaluate(region),
]
diff --git a/tests/python/tvmscript/test_tvmscript_roundtrip.py
b/tests/python/tvmscript/test_tvmscript_roundtrip.py
index b55d5279f9..e41755a9a2 100644
--- a/tests/python/tvmscript/test_tvmscript_roundtrip.py
+++ b/tests/python/tvmscript/test_tvmscript_roundtrip.py
@@ -3003,10 +3003,7 @@ def subroutine_call_without_arguments():
class mod:
@T.prim_func(s_tir=True)
def main():
- # Should be equivalent to the bare "mod.subroutine()", but
- # that relies on `GlobalVar.__call__` returning the
- # correct IR type.
- tirx.call_tir(mod.subroutine)
+ mod.subroutine()
@T.prim_func(s_tir=True)
def subroutine():
