You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by sy...@apache.org on 2023/02/08 15:22:47 UTC
[tvm] branch unity-staging updated (c1ccd5bd07 -> 63aaefa671)
This is an automated email from the ASF dual-hosted git repository.
syfeng pushed a change to branch unity-staging
in repository https://gitbox.apache.org/repos/asf/tvm.git
discard c1ccd5bd07 [Unity] Relax TVMScript Parser. (#13932)
new 63aaefa671 [Unity] Relax TVMScript Parser. (#13932)
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (c1ccd5bd07)
\
N -- N -- N refs/heads/unity-staging (63aaefa671)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
tests/python/relax/test_tvmscript_parser.py | 19 +------------------
1 file changed, 1 insertion(+), 18 deletions(-)
[tvm] 01/01: [Unity] Relax TVMScript Parser. (#13932)
Posted by sy...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
syfeng pushed a commit to branch unity-staging
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 63aaefa671b21c66564cfa6a0992620cc965eeaa
Author: Siyuan Feng <Hz...@sjtu.edu.cn>
AuthorDate: Wed Feb 8 22:31:47 2023 +0800
[Unity] Relax TVMScript Parser. (#13932)
This PR adds the TVMScript parser/ir_builder support based on the blockbuilder.
Co-authored-by: Ruihang Lai <ru...@cs.cmu.edu>
Co-authored-by: Junru Shao <ju...@gmail.com>
Co-authored-by: Tianqi Chen <ti...@gmail.com>
Co-authored-by: Yuchen Jin <yu...@cs.washington.edu>
Co-authored-by: Steven S. Lyubomirsky <sl...@gmail.com>
Co-authored-by: Yong Wu <yo...@gmail.com>
---
include/tvm/ir/expr.h | 1 +
include/tvm/script/ir_builder/ir/frame.h | 11 +-
include/tvm/script/ir_builder/ir/ir.h | 17 +
include/tvm/script/ir_builder/relax/frame.h | 293 ++++++
include/tvm/script/ir_builder/relax/ir.h | 137 +++
python/tvm/ir/expr.py | 50 +-
python/tvm/script/ir_builder/base.py | 6 +-
python/tvm/script/ir_builder/ir/__init__.py | 2 +-
python/tvm/script/ir_builder/ir/ir.py | 45 +
.../script/ir_builder/{ir => relax}/__init__.py | 7 +-
.../__init__.py => ir_builder/relax/_ffi_api.py} | 7 +-
python/tvm/script/ir_builder/relax/frame.py | 55 +
python/tvm/script/ir_builder/relax/ir.py | 407 ++++++++
python/tvm/script/parser/__init__.py | 3 +-
python/tvm/script/parser/core/diagnostics.py | 2 +-
python/tvm/script/parser/core/entry.py | 5 +
python/tvm/script/parser/core/evaluator.py | 2 +-
python/tvm/script/parser/core/parser.py | 50 +-
python/tvm/script/parser/ir/parser.py | 15 +
python/tvm/script/parser/relax/__init__.py | 17 +-
python/tvm/script/parser/relax/entry.py | 327 ++++++
python/tvm/script/parser/relax/parser.py | 276 +++++
python/tvm/script/parser/tir/entry.py | 4 +-
python/tvm/script/parser/tir/operation.py | 12 +-
python/tvm/script/parser/tir/parser.py | 26 +
src/ir/module.cc | 16 +-
src/script/ir_builder/ir/frame.cc | 12 +-
src/script/ir_builder/ir/ir.cc | 41 +-
src/script/ir_builder/ir/{frame.cc => utils.h} | 30 +-
src/script/ir_builder/relax/frame.cc | 273 +++++
src/script/ir_builder/relax/ir.cc | 236 +++++
src/script/ir_builder/relax/utils.h | 119 +++
src/script/ir_builder/tir/frame.cc | 15 +-
src/script/ir_builder/tir/utils.h | 2 +-
tests/python/relax/test_tvmscript_ir_builder.py | 153 +++
tests/python/relax/test_tvmscript_parser.py | 1062 ++++++++++++++++++++
36 files changed, 3667 insertions(+), 69 deletions(-)
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index d4ba628d36..c662067a04 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -462,6 +462,7 @@ class GlobalVarNode : public RelayExprNode {
v->Visit("virtual_device_", &virtual_device_);
v->Visit("span", &span);
v->Visit("_checked_type_", &checked_type_);
+ v->Visit("struct_info_", &struct_info_);
}
bool SEqualReduce(const GlobalVarNode* other, SEqualReducer equal) const {
diff --git a/include/tvm/script/ir_builder/ir/frame.h b/include/tvm/script/ir_builder/ir/frame.h
index 887981ccff..dacfc361a6 100644
--- a/include/tvm/script/ir_builder/ir/frame.h
+++ b/include/tvm/script/ir_builder/ir/frame.h
@@ -38,12 +38,17 @@ namespace ir {
*/
class IRModuleFrameNode : public IRBuilderFrameNode {
public:
- Array<GlobalVar> global_vars;
- Array<BaseFunc> functions;
+ /*! \brief A map from string names to global variables that ensures global uniqueness. */
+ Map<String, GlobalVar> global_var_map;
+ /*!
+ * \brief A map from GlobalVar to all global functions.
+ * \note Only defined functions are in the map, while declared functions are not included.
+ */
+ Map<GlobalVar, BaseFunc> functions;
void VisitAttrs(tvm::AttrVisitor* v) {
IRBuilderFrameNode::VisitAttrs(v);
- v->Visit("global_vars", &global_vars);
+ v->Visit("global_vars", &global_var_map);
v->Visit("functions", &functions);
}
diff --git a/include/tvm/script/ir_builder/ir/ir.h b/include/tvm/script/ir_builder/ir/ir.h
index f0e7cc6f5c..49bdcf60e6 100644
--- a/include/tvm/script/ir_builder/ir/ir.h
+++ b/include/tvm/script/ir_builder/ir/ir.h
@@ -37,6 +37,23 @@ namespace ir {
*/
TVM_DLL IRModuleFrame IRModule();
+/*!
+ * \brief Declare a Function without given the specific function implementation.
+ * \note It is usually used in cross-function call. And we can specify the function by `DefFunction`
+ * \param func_name The function unique name.
+ * \param func_signature A Function w/o body, which used to specify the function signature
+ * (i.e. func params and func return type/shape).
+ * \return The corresponding GlobalVar.
+ */
+TVM_DLL GlobalVar DeclFunction(const String& func_name, const BaseFunc& func_signature);
+
+/*!
+ * \brief Define the function which is declared before.
+ * \param func_name The function unique name.
+ * \param func The given function implementation
+ */
+TVM_DLL void DefFunction(const String& func_name, const BaseFunc& func);
+
} // namespace ir
} // namespace ir_builder
} // namespace script
diff --git a/include/tvm/script/ir_builder/relax/frame.h b/include/tvm/script/ir_builder/relax/frame.h
new file mode 100644
index 0000000000..0f544d3abc
--- /dev/null
+++ b/include/tvm/script/ir_builder/relax/frame.h
@@ -0,0 +1,293 @@
+/*
+ * 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.
+ */
+#ifndef TVM_SCRIPT_IR_BUILDER_RELAX_FRAME_H_
+#define TVM_SCRIPT_IR_BUILDER_RELAX_FRAME_H_
+
+#include <tvm/relax/block_builder.h>
+#include <tvm/relax/expr.h>
+#include <tvm/script/ir_builder/base.h>
+#include <tvm/script/ir_builder/ir/frame.h>
+#include <tvm/script/ir_builder/ir/ir.h>
+
+namespace tvm {
+namespace script {
+namespace ir_builder {
+namespace relax {
+
+/*! \brief The base ir_builder frame for the relax dialect. */
+class RelaxFrameNode : public IRBuilderFrameNode {
+ public:
+ void VisitAttrs(tvm::AttrVisitor* v) { IRBuilderFrameNode::VisitAttrs(v); }
+
+ static constexpr const char* _type_key = "script.ir_builder.relax.RelaxFrame";
+ TVM_DECLARE_BASE_OBJECT_INFO(RelaxFrameNode, IRBuilderFrameNode);
+};
+
+class RelaxFrame : public IRBuilderFrame {
+ public:
+ TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(RelaxFrame, IRBuilderFrame, RelaxFrameNode);
+
+ protected:
+ RelaxFrame() = default;
+};
+
+/*! \brief The base ir_builder frame for frames with SeqExpr
+ i.e. Functions, If branches
+ */
+class SeqExprFrameNode : public RelaxFrameNode {
+ public:
+ /*! \brief The binding blocks inside the frame. */
+ Array<tvm::relax::BindingBlock> binding_blocks;
+ /*! \brief The frame output expr. `NullOpt` when undefined. */
+ Optional<tvm::relax::Expr> output;
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ RelaxFrameNode::VisitAttrs(v);
+ v->Visit("binding_blocks", &binding_blocks);
+ v->Visit("output", &output);
+ }
+
+ static constexpr const char* _type_key = "script.ir_builder.relax.SeqExprFrame";
+ TVM_DECLARE_BASE_OBJECT_INFO(SeqExprFrameNode, RelaxFrameNode);
+
+ public:
+ void EnterWithScope() override;
+ void ExitWithScope() override;
+};
+
+class SeqExprFrame : public RelaxFrame {
+ public:
+ TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(SeqExprFrame, RelaxFrame, SeqExprFrameNode);
+};
+
+/*! \brief The ir_builder frame for the relax function. */
+class FunctionFrameNode : public SeqExprFrameNode {
+ public:
+ /*!
+ * \brief The function name.
+ * \note The name will not be specified in constructor, so it is "Optional",
+ * However, we must specify the name by `R.func_name` before exit this frame.
+ */
+ Optional<String> name;
+ /*! \brief The function params. */
+ Array<tvm::relax::Var> params;
+ /*!
+ * \brief The function return struct info.
+ * \note Usually the function return type can be deduced by the function body.
+ * But we can use this field to specify a more "accurate" return type.
+ * i.e. If the `ret_struct_info` is None, try to use the deduced type from body
+ * If the `ret_struct_info` is not None, we can still take body.struct_info
+ * if we ret_struct_info is base of body.struct_info. If not, we will
+ * take the specified `ret_struct_info`.
+ */
+ Optional<tvm::relax::StructInfo> ret_struct_info;
+
+ /*! \brief The function attributes. */
+ Map<String, ObjectRef> attrs;
+ /*! \brief The block builder to create Relax function. */
+ tvm::relax::BlockBuilder block_builder;
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ SeqExprFrameNode::VisitAttrs(v);
+ v->Visit("name", &name);
+ v->Visit("params", ¶ms);
+ v->Visit("ret_struct_info", &ret_struct_info);
+ v->Visit("attrs", &attrs);
+ v->Visit("binding_blocks", &binding_blocks);
+ v->Visit("output", &output);
+ // `block_builder` is not visited.
+ }
+
+ static constexpr const char* _type_key = "script.ir_builder.relax.FunctionFrame";
+ TVM_DECLARE_FINAL_OBJECT_INFO(FunctionFrameNode, SeqExprFrameNode);
+
+ public:
+ void ExitWithScope() final;
+};
+
+class FunctionFrame : public SeqExprFrame {
+ public:
+ TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(FunctionFrame, SeqExprFrame, FunctionFrameNode);
+};
+
+/*! \brief The ir_builder frame for relax binding blocks. */
+class BlockFrameNode : public RelaxFrameNode {
+ public:
+ /*! \brief The flag that indicates whether the block is a dataflow block. */
+ bool is_dataflow;
+ /*! \brief The variables emitted in this block. */
+ Array<tvm::relax::Var> emitted_vars;
+ /*!
+ * \brief A boolean indicating if the dataflow block is ended of construction.
+ * If it is true, any new binding trying to be emitted into this block will cause an error.
+ * \note Only used for a dataflow block.
+ */
+ bool block_ended;
+ /*!
+ * \brief The output vars of the dataflow block.
+ * \note Only used for a dataflow block.
+ */
+ Array<tvm::relax::Var> output_vars;
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ RelaxFrameNode::VisitAttrs(v);
+ v->Visit("is_dataflow", &is_dataflow);
+ v->Visit("emitted_vars", &emitted_vars);
+ v->Visit("output_vars", &output_vars);
+ // `block_ended` is not visited.
+ }
+
+ static constexpr const char* _type_key = "script.ir_builder.relax.BlockFrame";
+ TVM_DECLARE_FINAL_OBJECT_INFO(BlockFrameNode, RelaxFrameNode);
+
+ public:
+ void EnterWithScope() final;
+ void ExitWithScope() final;
+};
+
+class BlockFrame : public RelaxFrame {
+ public:
+ TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(BlockFrame, RelaxFrame, BlockFrameNode);
+};
+
+/*!
+ * \brief A frame that represents if statement.
+ *
+ * \sa IfFrame
+ */
+class IfFrameNode : public RelaxFrameNode {
+ public:
+ /*! \brief The condition of the if statement. */
+ tvm::relax::Expr condition;
+ /*! \brief The Bindings in the true branch. */
+ Optional<tvm::relax::Expr> then_expr;
+ /*! \brief The Bindings in the false branch. */
+ Optional<tvm::relax::Expr> else_expr;
+ /*! \brief The Binding var. */
+ tvm::relax::Var var;
+ /*! \brief The binding var name. */
+ String var_name;
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ RelaxFrameNode::VisitAttrs(v);
+ v->Visit("condition", &condition);
+ v->Visit("then_expr", &then_expr);
+ v->Visit("else_expr", &else_expr);
+ v->Visit("var", &var);
+ v->Visit("var_name", &var_name);
+ }
+
+ static constexpr const char* _type_key = "script.ir_builder.relax.IfFrame";
+ TVM_DECLARE_FINAL_OBJECT_INFO(IfFrameNode, RelaxFrameNode);
+
+ public:
+ /*!
+ * \brief The method called when entering RAII scope.
+ * \sa tvm::support::With
+ */
+ void EnterWithScope() final;
+ /*!
+ * \brief The method called when exiting RAII scope.
+ * \sa tvm::support::With
+ */
+ void ExitWithScope() final;
+};
+
+/*!
+ * \brief Managed reference to IfFrameNode.
+ *
+ * \sa IfFrameNode
+ */
+class IfFrame : public RelaxFrame {
+ public:
+ TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IfFrame, RelaxFrame, IfFrameNode);
+};
+
+/*!
+ * \brief A frame that represents then.
+ *
+ * \sa ThenFrame
+ */
+class ThenFrameNode : public SeqExprFrameNode {
+ public:
+ static constexpr const char* _type_key = "script.ir_builder.relax.ThenFrame";
+ TVM_DECLARE_FINAL_OBJECT_INFO(ThenFrameNode, SeqExprFrameNode);
+
+ public:
+ /*!
+ * \brief The method called when entering RAII scope.
+ * \sa tvm::support::With
+ */
+ void EnterWithScope() final;
+ /*!
+ * \brief The method called when exiting RAII scope.
+ * \sa tvm::support::With
+ */
+ void ExitWithScope() final;
+};
+
+/*!
+ * \brief Managed reference to ThenFrameNode.
+ *
+ * \sa ThenFrameNode
+ */
+class ThenFrame : public SeqExprFrame {
+ public:
+ TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(ThenFrame, SeqExprFrame, ThenFrameNode);
+};
+
+/*!
+ * \brief A frame that represents else.
+ *
+ * \sa ElseFrame
+ */
+class ElseFrameNode : public SeqExprFrameNode {
+ public:
+ static constexpr const char* _type_key = "script.ir_builder.relax.ElseFrame";
+ TVM_DECLARE_FINAL_OBJECT_INFO(ElseFrameNode, SeqExprFrameNode);
+
+ public:
+ /*!
+ * \brief The method called when entering RAII scope.
+ * \sa tvm::support::With
+ */
+ void EnterWithScope() final;
+ /*!
+ * \brief The method called when exiting RAII scope.
+ * \sa tvm::support::With
+ */
+ void ExitWithScope() final;
+};
+
+/*!
+ * \brief Managed reference to ElseFrameNode.
+ *
+ * \sa ElseFrameNode
+ */
+class ElseFrame : public SeqExprFrame {
+ public:
+ TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(ElseFrame, SeqExprFrame, ElseFrameNode);
+};
+
+} // namespace relax
+} // namespace ir_builder
+} // namespace script
+} // namespace tvm
+
+#endif // TVM_SCRIPT_IR_BUILDER_RELAX_FRAME_H_
diff --git a/include/tvm/script/ir_builder/relax/ir.h b/include/tvm/script/ir_builder/relax/ir.h
new file mode 100644
index 0000000000..72aab6684e
--- /dev/null
+++ b/include/tvm/script/ir_builder/relax/ir.h
@@ -0,0 +1,137 @@
+/*
+ * 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.
+ */
+#ifndef TVM_SCRIPT_IR_BUILDER_RELAX_IR_H_
+#define TVM_SCRIPT_IR_BUILDER_RELAX_IR_H_
+
+#include <tvm/relax/expr.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/script/ir_builder/base.h>
+#include <tvm/script/ir_builder/relax/frame.h>
+
+namespace tvm {
+namespace script {
+namespace ir_builder {
+namespace relax {
+
+/////////////////////////////// Function ////////////////////////////////
+
+/*!
+ * \brief Start a function frame.
+ * \return The created ir_builder Function frame.
+ */
+TVM_DLL FunctionFrame Function();
+
+/*!
+ * \brief Add a parameter to the last function frame.
+ * \param name The name of the parameter.
+ * \param struct_info The struct_info of the parameter.
+ * \return The created function parameter var.
+ */
+TVM_DLL tvm::relax::Var Arg(const String& name, const tvm::relax::StructInfo& struct_info);
+
+/*!
+ * \brief Specify the name of the last function frame.
+ * \param name The function name.
+ */
+TVM_DLL void FuncName(const String& name);
+
+/*!
+ * \brief Specify the attrs of the last function frame.
+ * \param attrs The function attrs.
+ */
+TVM_DLL void FuncAttrs(Map<String, ObjectRef> attrs);
+
+/*!
+ * \brief Specify the return struct info of the last function frame.
+ * \param ret_sinfo The return struct info.
+ */
+TVM_DLL void FuncRetStructInfo(const tvm::relax::StructInfo& ret_sinfo);
+
+/*!
+ * \brief Specify the return value of the last function frame.
+ * \param value The return value.
+ */
+TVM_DLL void FuncRetValue(const tvm::relax::Expr& value);
+
+///////////////////////////// BindingBlock //////////////////////////////
+
+/*!
+ * \brief Start a binding block frame.
+ * \return The created ir_builder Block frame.
+ */
+TVM_DLL BlockFrame BindingBlock();
+
+/*!
+ * \brief Start a dataflow binding block frame.
+ * \return The created ir_builder Block frame.
+ */
+TVM_DLL BlockFrame Dataflow();
+
+/*!
+ * \brief Expose the dataflow block output variables as global ones
+ * \param vars The output variables of a dataflow block
+ */
+TVM_DLL void DataflowBlockOutput(const Array<tvm::relax::Var>& vars);
+
+////////////////////////////// Bindings ////////////////////////////////
+
+/*!
+ * \brief Emit a binding to the last binding block frame.
+ * \param value The right side value of the bindings to be emitted.
+ * \param annotate_struct_info The optional struct info annotation for the emitted value.
+ * \return The left side var of the emitted binding.
+ */
+TVM_DLL tvm::relax::Var Emit(
+ const tvm::relax::Expr& value,
+ const Optional<tvm::relax::StructInfo>& annotate_struct_info = NullOpt);
+
+/*!
+ * \brief Emit a match_cast binding to the last binding block frame.
+ * \param value The value of the MatchCast to be emitted.
+ * \param struct_info The struct info of the MatchCast to be emitted.
+ * \return The left side var of the emitted binding.
+ */
+TVM_DLL tvm::relax::Var EmitMatchCast(const tvm::relax::Expr& value,
+ const tvm::relax::StructInfo& struct_info);
+
+///////////////////////////// If Then Else /////////////////////////////
+
+/*!
+ * \brief Create an if statement.
+ * \param condition The condition of if statement.
+ * \return The result IfFrame.
+ */
+IfFrame If(tvm::relax::Expr condition);
+/*!
+ * \brief Create a then.
+ * \return The result ThenFrame.
+ */
+ThenFrame Then();
+/*!
+ * \brief Create an else.
+ * \return The result ElseFrame.
+ */
+ElseFrame Else();
+
+} // namespace relax
+} // namespace ir_builder
+} // namespace script
+} // namespace tvm
+
+#endif // TVM_SCRIPT_IR_BUILDER_RELAX_IR_H_
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index f90468de66..721e12e7f8 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -93,10 +93,17 @@ class GlobalVar(RelayExpr):
A call taking the variable as a function.
"""
# pylint: disable=import-outside-toplevel
+
+ # TODO(@relax-team): replace with Relax base class after it's introduced
if all(isinstance(x, RelayExpr) for x in args):
- from tvm import relay
+ if all(is_relax_expr(x) for x in args):
+ from tvm import relax
+
+ return relax.Call(self, args)
+ else:
+ from tvm import relay
- return relay.Call(self, args)
+ return relay.Call(self, args)
arg_types = [type(x) for x in args]
raise RuntimeError(
"Do not know how to handle GlobalVar.__call__ for types {}".format(arg_types)
@@ -185,3 +192,42 @@ class Range(Node, Scriptable):
The constructed range.
"""
return _ffi_api.Range_from_min_extent(min_value, extent, span)
+
+
+# TODO(@relax-team): remove when we have a RelaxExpr base class
+def is_relax_expr(expr: RelayExpr) -> bool:
+ """check if a RelayExpr is a Relax expresssion.
+
+ Parameters
+ ----------
+ expr : RelayExpr
+ The expression to check.
+
+ Returns
+ -------
+ res : bool
+ If the expression is Relax expression, return True; otherwise return False.
+ """
+ from tvm import relax # pylint: disable=import-outside-toplevel
+
+ if isinstance(
+ expr,
+ (
+ relax.Call,
+ relax.Constant,
+ relax.Tuple,
+ relax.TupleGetItem,
+ relax.If,
+ relax.Var,
+ relax.DataflowVar,
+ relax.ShapeExpr,
+ relax.SeqExpr,
+ relax.Function,
+ relax.ExternFunc,
+ relax.PrimValue,
+ relax.StringImm,
+ relax.DataTypeImm,
+ ),
+ ):
+ return True
+ return False
diff --git a/python/tvm/script/ir_builder/base.py b/python/tvm/script/ir_builder/base.py
index 7aa33ee49c..b35bbd0a7d 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -64,8 +64,10 @@ class IRBuilderFrame(_Object):
_ffi_api.IRBuilderFrameEnter(self) # type: ignore[attr-defined] # pylint: disable=no-member
return self
- def __exit__(self, ptype, value, trace) -> None: # pylint: disable=unused-argument
- _ffi_api.IRBuilderFrameExit(self) # type: ignore[attr-defined] # pylint: disable=no-member
+ def __exit__(self, exc_type, exc_value, trace) -> None: # pylint: disable=unused-argument
+ if exc_type is None and exc_value is None:
+ # Do not execute `FrameExit` if the with scope exits because of exceptions
+ _ffi_api.IRBuilderFrameExit(self) # type: ignore[attr-defined] # pylint: disable=no-member
def add_callback(self, callback: Callable[[], None]) -> None:
"""Add a callback method invoked when exiting the with-scope.
diff --git a/python/tvm/script/ir_builder/ir/__init__.py b/python/tvm/script/ir_builder/ir/__init__.py
index ebb9728737..946be263a7 100644
--- a/python/tvm/script/ir_builder/ir/__init__.py
+++ b/python/tvm/script/ir_builder/ir/__init__.py
@@ -16,4 +16,4 @@
# under the License.
"""Package tvm.script.ir_builder.ir"""
from .frame import IRModuleFrame
-from .ir import ir_module
+from .ir import decl_function, def_function, ir_module
diff --git a/python/tvm/script/ir_builder/ir/ir.py b/python/tvm/script/ir_builder/ir/ir.py
index 213180463c..796d6f3aad 100644
--- a/python/tvm/script/ir_builder/ir/ir.py
+++ b/python/tvm/script/ir_builder/ir/ir.py
@@ -16,9 +16,54 @@
# under the License.
"""Package tvm.script.ir_builder.ir.ir"""
+from tvm.ir import BaseFunc, GlobalVar
+
from . import _ffi_api
from .frame import IRModuleFrame
def ir_module() -> IRModuleFrame:
+ """Start a ir_module frame.
+ Returns
+ -------
+ frame: IRModuleFrame
+ The constructed frame.
+ """
return _ffi_api.IRModule() # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def decl_function(func_name: str, func_signature: BaseFunc) -> GlobalVar:
+ """Declare a Function without given the specific function implementation.
+ Parameters
+ ----------
+ func_name : str
+ The function unique name.
+
+ func_signature: Optional[BaseFunc]
+ A Function w/o body, which used to specify the function signature
+ (i.e. func params and func return type/shape).
+
+ Note
+ ----
+ It is usually used in cross-function call. And we can specify the function by `DefFunction`
+ Returns
+ -------
+ gv : GlobalVar
+ The corresponding GlobalVar.
+ """
+
+ return _ffi_api.DeclFunction( # type: ignore[attr-defined] # pylint: disable=no-member
+ func_name, func_signature
+ )
+
+
+def def_function(func_name: str, func: BaseFunc) -> None:
+ """Define the function which is declared before.
+ Parameters
+ ----------
+ func_name : str
+ The function unique name.
+ func: BaseFunc
+ The given function implementation
+ """
+ return _ffi_api.DefFunction(func_name, func) # type: ignore[attr-defined] # pylint: disable=no-member
diff --git a/python/tvm/script/ir_builder/ir/__init__.py b/python/tvm/script/ir_builder/relax/__init__.py
similarity index 82%
copy from python/tvm/script/ir_builder/ir/__init__.py
copy to python/tvm/script/ir_builder/relax/__init__.py
index ebb9728737..f0905acf34 100644
--- a/python/tvm/script/ir_builder/ir/__init__.py
+++ b/python/tvm/script/ir_builder/relax/__init__.py
@@ -14,6 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Package tvm.script.ir_builder.ir"""
-from .frame import IRModuleFrame
-from .ir import ir_module
+# pylint: disable=unused-import
+"""Package tvm.script.ir_builder.relax"""
+from . import frame
+from .ir import * # pylint: disable=wildcard-import,redefined-builtin
diff --git a/python/tvm/script/parser/relax/__init__.py b/python/tvm/script/ir_builder/relax/_ffi_api.py
similarity index 83%
copy from python/tvm/script/parser/relax/__init__.py
copy to python/tvm/script/ir_builder/relax/_ffi_api.py
index feb8e68340..6e2098cf88 100644
--- a/python/tvm/script/parser/relax/__init__.py
+++ b/python/tvm/script/ir_builder/relax/_ffi_api.py
@@ -14,8 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Initial impl of relax parser for sugars"""
-from tvm.relax import TensorStructInfo, ShapeStructInfo
+"""FFI APIs for tvm.script.ir_builder.relax"""
+import tvm._ffi
-Tensor = TensorStructInfo
-Shape = ShapeStructInfo
+tvm._ffi._init_api("script.ir_builder.relax", __name__) # pylint: disable=protected-access
diff --git a/python/tvm/script/ir_builder/relax/frame.py b/python/tvm/script/ir_builder/relax/frame.py
new file mode 100644
index 0000000000..97e181fbe4
--- /dev/null
+++ b/python/tvm/script/ir_builder/relax/frame.py
@@ -0,0 +1,55 @@
+# 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.
+"""IR Builder Frame for Relax dialect"""
+from tvm._ffi import register_object as _register_object
+
+from ..base import IRBuilderFrame
+
+
+@_register_object("script.ir_builder.relax.RelaxFrame")
+class RelaxFrame(IRBuilderFrame):
+ """The base ir_builder frame for the relax dialect."""
+
+
+@_register_object("script.ir_builder.relax.SeqExprFrame")
+class SeqExprFrame(RelaxFrame):
+ ...
+
+
+@_register_object("script.ir_builder.relax.FunctionFrame")
+class FunctionFrame(SeqExprFrame):
+ """The ir_builder frame for the relax function."""
+
+
+@_register_object("script.ir_builder.relax.BlockFrame")
+class BlockFrame(RelaxFrame):
+ """The ir_builder frame for relax binding blocks."""
+
+
+@_register_object("script.ir_builder.relax.IfFrame")
+class IfFrame(RelaxFrame):
+ ...
+
+
+@_register_object("script.ir_builder.relax.ThenFrame")
+class ThenFrame(SeqExprFrame):
+ ...
+
+
+@_register_object("script.ir_builder.relax.ElseFrame")
+class ElseFrame(SeqExprFrame):
+ ...
diff --git a/python/tvm/script/ir_builder/relax/ir.py b/python/tvm/script/ir_builder/relax/ir.py
new file mode 100644
index 0000000000..647ef8f25a
--- /dev/null
+++ b/python/tvm/script/ir_builder/relax/ir.py
@@ -0,0 +1,407 @@
+# 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.
+# pylint: disable=redefined-builtin, wrong-import-order, no-member, invalid-name
+"""IRBuilder for Relax dialect"""
+
+import builtins
+import functools
+import inspect
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import tvm
+from tvm import DataType, relax
+from tvm.ir import PrimExpr
+from tvm.relax import Call, Expr, ExternFunc, TupleGetItem, Var, const
+
+############################### Operators ###############################
+from tvm.relax.op import (
+ add,
+ assert_op,
+ call_builtin_with_ctx,
+ call_tir,
+ invoke_closure,
+ make_closure,
+ multiply,
+ null_value,
+ print,
+ shape_of,
+)
+from tvm.relax.struct_info import StructInfo
+from tvm.relax.utils import args_converter
+from tvm.runtime import Object as tvm_Object
+from tvm.runtime import ObjectGeneric
+
+from . import _ffi_api, frame
+
+##################### Python Native Function Alias ######################
+
+py_print = builtins.print
+py_tuple = tuple
+py_str = str
+
+
+############################### Function ################################
+
+
+def function() -> frame.FunctionFrame:
+ """Start a function frame.
+ Returns
+ -------
+ frame: FunctionFrame
+ The constructed function frame.
+ """
+ return _ffi_api.Function() # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def arg(name: py_str, struct_info: StructInfo) -> Var:
+ """Add a parameter to the last function frame.
+ Parameters
+ ----------
+ name: str
+ The name of the parameter.
+ struct_info: StructInfo
+ The Struct Info of the parameter
+
+ Returns
+ -------
+ var: Var
+ The created function parameter var.
+ """
+
+ return _ffi_api.Arg(name, struct_info) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def func_name(name: py_str) -> None:
+ """Specify the name of the last function frame.
+ Parameters
+ ----------
+ name: str
+ The function name.
+ """
+ return _ffi_api.FuncName(name) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def func_attr(attrs: Dict[py_str, tvm_Object]) -> None:
+ """Specify the attrs of the last function frame.
+ Parameters
+ ----------
+ attrs: Dict[str, Object]
+ The function attrs.
+ """
+ return _ffi_api.FuncAttrs(attrs) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def func_ret_struct_info(ret_sinfo: StructInfo) -> None:
+ """Specify the return struct info of the last function frame.
+ Parameters
+ ----------
+ ret_type: StructInfo
+ The function return struct info.
+ """
+ return _ffi_api.FuncRetStructInfo(ret_sinfo) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def func_ret_value(value: Expr) -> None:
+ """Specify the return value of the last function frame.
+ Parameters
+ ----------
+ value: Expr
+ The function return value.
+ """
+ return _ffi_api.FuncRetValue(value) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+############################# BindingBlock ##############################
+
+
+def dataflow() -> frame.BlockFrame:
+ """Start a dataflow binding block frame.
+ Returns
+ -------
+ frame: frame.BlockFrame
+ The created ir_builder Block frame.
+ """
+ return _ffi_api.Dataflow() # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def output(*vars: Tuple[Var]) -> None:
+ """Expose the dataflow block output variables as global ones.
+ Parameters
+ ----------
+ vars: Tuple[Var]
+ The output variables of a dataflow block.
+ """
+ return _ffi_api.DataflowBlockOutput(vars) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+################################## Ops #################################
+
+
+@args_converter.auto
+def call_packed(
+ func: py_str,
+ *args: Expr,
+ sinfo_args: Union[StructInfo, List[StructInfo]],
+ **kwargs: Any,
+) -> Call:
+ """Create a relax Call, which calls a packed function.
+ Parameters
+ ----------
+ func: str
+ The name of extern function.
+ *args : Expr
+ The arguments.
+ sinfo_args: Union[StructInfo, List[StructInfo]]
+ The list of structure info arguments.
+ kwargs: Expr
+ The keyword arguments.
+
+ Returns
+ -------
+ call: Call
+ The created Relax Call
+ """
+ op = ExternFunc(func)
+ if sinfo_args is None:
+ raise ValueError("R.call_packed is required to have type_args")
+ if isinstance(sinfo_args, py_tuple): # type: ignore
+ sinfo_args = list(sinfo_args)
+ elif not isinstance(sinfo_args, list):
+ sinfo_args = [sinfo_args]
+ for i, sinfo_arg in enumerate(sinfo_args):
+ if callable(sinfo_arg):
+ sinfo_arg = sinfo_arg()
+ # Convert possible StructInfoProxy to StructInfo
+ if isinstance(sinfo_arg, ObjectGeneric):
+ sinfo_arg = sinfo_arg.asobject()
+ sinfo_args[i] = sinfo_arg
+
+ is_default = False
+ if "attrs_type_key" in kwargs:
+ attrs_type_key = kwargs["attrs_type_key"]
+ kwargs.pop("attrs_type_key")
+ else:
+ attrs_type_key = "DictAttrs"
+ is_default = True
+ attrs = None
+ if kwargs or not is_default:
+ attrs = tvm.ir.attrs.make_node(attrs_type_key, **kwargs)
+
+ return Call(op, args, attrs=attrs, sinfo_args=sinfo_args)
+
+
+def _sinfo_arg_wrapper(func):
+ """A wrapper to convert StructInfoProxies to StructInfo for builtin operators with sinfo_args"""
+
+ def _convert_tensor_type(args):
+ if isinstance(args, (list, py_tuple)): # type: ignore
+ new_args = [_convert_tensor_type(x) for x in args]
+ return type(args)(new_args)
+ if isinstance(args, dict):
+ return {_convert_tensor_type(k): _convert_tensor_type(v) for k, v in args.items()}
+ if inspect.isfunction(args):
+ args = args()
+ if isinstance(args, ObjectGeneric):
+ args = args.asobject()
+ return args
+
+ @functools.wraps(func)
+ def wrapped(*args, **kwargs):
+ return func(*_convert_tensor_type(args), **_convert_tensor_type(kwargs))
+
+ return wrapped # type: ignore
+
+
+invoke_closure = _sinfo_arg_wrapper(invoke_closure) # pylint: disable=invalid-name
+
+call_builtin_with_ctx = _sinfo_arg_wrapper(call_builtin_with_ctx) # pylint: disable=invalid-name
+
+############################### Bindings ###############################
+
+
+def emit(value: Expr, annotate_struct_info: Optional[StructInfo] = None) -> Var:
+ """Emit a binding to the last binding block frame.
+ Parameters
+ ----------
+ value: Expr
+ The right side value of the bindings to be emitted.
+
+ annotate_struct_info: Optional[StructInfo]
+ The optional struct info annotation for the emitted value.
+
+ Returns
+ -------
+ var: Var
+ The left side var of the emitted binding.
+ """
+ return _ffi_api.Emit(value, annotate_struct_info) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def emit_match_cast(value: Expr, struct_info: StructInfo) -> Var:
+ """Emit a match_cast binding to the last binding block frame.
+ Parameters
+ ----------
+ value: Expr
+ The value of the MatchCast to be emitted.
+ struct_info: StructInfo
+ The struct_info of the MatchCast to be emitted.
+
+ Returns
+ -------
+ var: Var
+ The left side var of the emitted binding.
+ """
+ return _ffi_api.EmitMatchCast(value, struct_info) # type: ignore
+
+
+############################# If Then Else #############################
+
+
+def If(condition: Expr) -> frame.IfFrame: # pylint: disable=invalid-name
+ """Create an if frame.
+ Parameters
+ ----------
+ condition : Expr
+ The condition of if statement, executes the true branch if the condition is true,
+ otherwise jump into the false branch.
+ Returns
+ -------
+ res : frame.IfFrame
+ The result IfFrame.
+ """
+ return _ffi_api.If(condition) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def Then() -> frame.ThenFrame: # pylint: disable=invalid-name
+ """Create a then frame.
+ Returns
+ -------
+ res : frame.ThenFrame
+ The result ThenFrame.
+ """
+ return _ffi_api.Then() # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def Else() -> frame.ElseFrame: # pylint: disable=invalid-name
+ """Create an else frame.
+ Returns
+ -------
+ res : frame.ElseFrame
+ The result ElseFrame.
+ """
+ return _ffi_api.Else() # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+############################### R.tuple ################################
+
+
+def tuple(*fields: Expr) -> Expr:
+ """Create a tuple expression.
+ Parameters
+ ----------
+ *fields : Expr
+ The fields of the tuple.
+ Returns
+ -------
+ res : Expr
+ The result tuple.
+ """
+ if len(fields) == 0:
+ fields = py_tuple()
+
+ return relax.Tuple(fields) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+############################### PrimValue ##############################
+
+
+def prim_value(value: PrimExpr) -> Expr:
+ """Create a prim value expression.
+ Parameters
+ ----------
+ value : PrimExpr
+ The value of the prim value.
+ Returns
+ -------
+ res : Expr
+ The result prim value.
+ """
+ return relax.PrimValue(value) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def str(value: py_str) -> Expr:
+ """Create a string imm expression.
+ Parameters
+ ----------
+ value : str
+ The value of the str.
+ Returns
+ -------
+ res : Expr
+ The result str.
+ """
+ return relax.StringImm(value) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+def dtype(value: Union[py_str, DataType]) -> Expr:
+ """Create a dtype imm expression.
+ Parameters
+ ----------
+ value : dtype
+ The value of the dtype.
+ Returns
+ -------
+ res : Expr
+ The result dtype.
+ """
+ return relax.DataTypeImm(value) # type: ignore[attr-defined] # pylint: disable=no-member
+
+
+############################### Importer ###############################
+
+__all__ = [
+ "Else",
+ "If",
+ "Then",
+ "TupleGetItem",
+ "add",
+ "arg",
+ "assert_op",
+ "call_packed",
+ "call_tir",
+ "call_builtin_with_ctx",
+ "const",
+ "dataflow",
+ "dtype",
+ "emit",
+ "emit_match_cast",
+ "func_attr",
+ "func_name",
+ "func_ret_struct_info",
+ "func_ret_value",
+ "function",
+ "invoke_closure",
+ "make_closure",
+ "multiply",
+ "null_value",
+ "output",
+ "prim_value",
+ "print",
+ "shape_of",
+ "str",
+ "tuple",
+]
diff --git a/python/tvm/script/parser/__init__.py b/python/tvm/script/parser/__init__.py
index 5161a2601c..678297799e 100644
--- a/python/tvm/script/parser/__init__.py
+++ b/python/tvm/script/parser/__init__.py
@@ -15,7 +15,8 @@
# specific language governing permissions and limitations
# under the Licens.
"""The parser"""
-from . import _core, ir, tir
+from . import _core, ir, tir, relax
from ._core import parse
from .ir import ir_module
from .tir import prim_func
+from .relax import function
diff --git a/python/tvm/script/parser/core/diagnostics.py b/python/tvm/script/parser/core/diagnostics.py
index ad7ae50347..2767a97f60 100644
--- a/python/tvm/script/parser/core/diagnostics.py
+++ b/python/tvm/script/parser/core/diagnostics.py
@@ -220,7 +220,7 @@ class Diagnostics:
level : diagnostics.DiagnosticLevel
The diagnostic level.
"""
- lineno = node.lineno or self.source.start_line
+ lineno = node.lineno or 1
col_offset = node.col_offset or self.source.start_column
end_lineno = node.end_lineno or lineno
end_col_offset = node.end_col_offset or col_offset
diff --git a/python/tvm/script/parser/core/entry.py b/python/tvm/script/parser/core/entry.py
index bf6a118672..a98dba91a0 100644
--- a/python/tvm/script/parser/core/entry.py
+++ b/python/tvm/script/parser/core/entry.py
@@ -41,7 +41,9 @@ def parse(program: Union[doc.AST, Any, str], extra_vars: Dict[str, Any] = None)
The parsed TVMScript program.
"""
if extra_vars is None:
+ import tvm # pylint: disable=import-outside-toplevel
from tvm.script.parser import ir # pylint: disable=import-outside-toplevel
+ from tvm.script.parser import relax # pylint: disable=import-outside-toplevel
from tvm.script.parser import tir # pylint: disable=import-outside-toplevel
extra_vars = {
@@ -49,6 +51,9 @@ def parse(program: Union[doc.AST, Any, str], extra_vars: Dict[str, Any] = None)
"ir": ir,
"T": tir,
"tir": tir,
+ "relax": relax,
+ "R": relax,
+ "tvm": tvm,
}
source = Source(program)
diff --git a/python/tvm/script/parser/core/evaluator.py b/python/tvm/script/parser/core/evaluator.py
index 3a72a3c331..075aedd891 100644
--- a/python/tvm/script/parser/core/evaluator.py
+++ b/python/tvm/script/parser/core/evaluator.py
@@ -203,7 +203,7 @@ class ExprEvaluator:
else:
value = self._eval_expr(node.__class__(**fields))
except Exception as e: # pylint: disable=broad-except,invalid-name
- self.parser.report_error(node, str(e))
+ self.parser.report_error(node, e)
return self._add_intermediate_result(value)
def _eval_lambda(self, node: doc.Lambda) -> Any:
diff --git a/python/tvm/script/parser/core/parser.py b/python/tvm/script/parser/core/parser.py
index 7c699c42ae..105164ed5f 100644
--- a/python/tvm/script/parser/core/parser.py
+++ b/python/tvm/script/parser/core/parser.py
@@ -60,6 +60,10 @@ def _deferred(exit_f: Callable[[], None]):
return context()
+def _do_nothing(*args, **kwargs): # pylint: disable=unused-argument
+ pass
+
+
class VarTableFrame:
"""The variable table frame.
A frame of variable table stores the variables created in one block or scope.
@@ -259,6 +263,17 @@ class Parser(doc.NodeVisitor):
node = self.diag.source.as_ast()
self.visit(node)
+ def get_dispatch_token(self, node: doc.FunctionDef) -> str:
+ if not isinstance(node, doc.FunctionDef):
+ self.report_error(node, "Only can get dispatch token for function.")
+ if not node.decorator_list:
+ self.report_error(node, "Function must be decorated")
+ # TODO: only the last decorator is parsed
+ decorator = self.eval_expr(node.decorator_list[-1])
+ if not hasattr(decorator, "dispatch_token"):
+ self.report_error(node, "The parser does not understand the decorator")
+ return decorator.dispatch_token
+
def with_dispatch_token(self, token: str):
"""Add a new dispatching token as with statement.
@@ -388,6 +403,8 @@ class Parser(doc.NodeVisitor):
# Only take the last line of the error message
if isinstance(err, TVMError):
msg = list(filter(None, str(err).split("\n")))[-1]
+ elif isinstance(err, KeyError):
+ msg = "KeyError: " + str(err)
else:
msg = str(err)
self.diag.error(node, msg)
@@ -457,30 +474,33 @@ class Parser(doc.NodeVisitor):
"""
return _dispatch(self, "tvm_annotation")(self, node)
- def visit_FunctionDef(self, node: doc.FunctionDef) -> Any: # pylint: disable=invalid-name
- """The general function definition visiting method.
+ def visit_FunctionDef(self, node: doc.FunctionDef) -> None: # pylint: disable=invalid-name
+ """The general function definition visit method.
Parameters
----------
node : doc.FunctionDef
- The doc AST function definition node.
-
- Returns
- -------
- res : Any
- The visiting result.
+ The doc FunctionDef node.
"""
- if not node.decorator_list:
- self.report_error(node, "Function must be decorated")
- # TODO: only the last decorator is parsed
- decorator = self.eval_expr(node.decorator_list[-1])
- if not hasattr(decorator, "dispatch_token"):
- self.report_error(node, "The parser does not understand the decorator")
- token = decorator.dispatch_token
+ token = self.get_dispatch_token(node)
+ current_token = self.dispatch_tokens[-1]
func = dispatch.get(token=token, type_name="FunctionDef", default=None)
if func is None:
self.report_error(node, "The parser does not understand the decorator")
+ pre_func = dispatch.get(
+ token=current_token, type_name="pre_token_switch", default=_do_nothing
+ )
+ post_func = dispatch.get(
+ token=current_token, type_name="post_token_switch", default=_do_nothing
+ )
+ pre_func(self, node)
_dispatch_wrapper(func)(self, node)
+ post_func(self, node)
+
+ def visit_tvm_declare_function(self, node: doc.FunctionDef) -> None:
+ token = self.get_dispatch_token(node)
+ with self.with_dispatch_token(token):
+ _dispatch(self, "tvm_declare_function")(self, node)
def visit_ClassDef(self, node: doc.ClassDef) -> Any: # pylint: disable=invalid-name
"""The general class definition visiting method.
diff --git a/python/tvm/script/parser/ir/parser.py b/python/tvm/script/parser/ir/parser.py
index 9532e7e32c..13b3e29859 100644
--- a/python/tvm/script/parser/ir/parser.py
+++ b/python/tvm/script/parser/ir/parser.py
@@ -32,8 +32,12 @@ def _visit_class_def(self: Parser, node: doc.ClassDef) -> None:
node : doc.ClassDef
The doc AST class definition node.
"""
+
with self.var_table.with_frame():
with I.ir_module():
+ for stmt in node.body:
+ if isinstance(stmt, doc.FunctionDef):
+ self.visit_tvm_declare_function(stmt)
with self.with_dispatch_token("ir"):
self.visit_body(node.body)
@@ -64,3 +68,14 @@ def _visit_expr(_self: Parser, _node: doc.Expr) -> None:
node : doc.ClassDef
The doc AST expression node.
"""
+
+
+@dispatch.register(token="default", type_name="Assign")
+def visit_assign(self: Parser, node: doc.Assign) -> None:
+ if len(node.targets) != 1:
+ self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
+ lhs = node.targets[0]
+ rhs = self.eval_expr(node.value)
+ self.eval_assign(
+ target=lhs, source=rhs, bind_value=lambda _a, _b, _c, value: value, allow_shadowing=True
+ )
diff --git a/python/tvm/script/parser/relax/__init__.py b/python/tvm/script/parser/relax/__init__.py
index feb8e68340..04f3fea21c 100644
--- a/python/tvm/script/parser/relax/__init__.py
+++ b/python/tvm/script/parser/relax/__init__.py
@@ -15,7 +15,18 @@
# specific language governing permissions and limitations
# under the License.
"""Initial impl of relax parser for sugars"""
-from tvm.relax import TensorStructInfo, ShapeStructInfo
+from ...ir_builder.relax import * # pylint: disable=redefined-builtin
+from ...ir_builder.relax import ir as _relax
+from . import parser as _parser
+from .entry import Callable, Object, Prim, Shape, Tensor, Tuple, function, match_cast
-Tensor = TensorStructInfo
-Shape = ShapeStructInfo
+__all__ = _relax.__all__ + [
+ "Callable",
+ "Object",
+ "Prim",
+ "Shape",
+ "Tensor",
+ "Tuple",
+ "function",
+ "match_cast",
+]
diff --git a/python/tvm/script/parser/relax/entry.py b/python/tvm/script/parser/relax/entry.py
new file mode 100644
index 0000000000..d93f9a2826
--- /dev/null
+++ b/python/tvm/script/parser/relax/entry.py
@@ -0,0 +1,327 @@
+# 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.
+# pylint: disable=missing-docstring, invalid-name
+import inspect
+from typing import Any
+from typing import Callable as _Callable
+from typing import Dict, List, Optional, Set, TypeVar, Union
+
+from tvm.relax import (
+ Expr,
+ FuncStructInfo,
+ Function,
+ ObjectStructInfo,
+ PrimStructInfo,
+ ShapeStructInfo,
+ StructInfo,
+ TensorStructInfo,
+ TupleStructInfo,
+)
+from tvm.runtime import ObjectGeneric
+from tvm.tir import PrimExpr
+
+from .._core import parse, utils
+
+FType = TypeVar("FType", bound=_Callable)
+
+############################## R.function ##############################
+
+
+def function(f: FType) -> Union[Function, FType]:
+ if not inspect.isfunction(f):
+ raise TypeError(f"Expect a function, but got: {f}")
+ if utils.is_defined_in_class(inspect.stack(), f):
+ return f
+ return parse(f, utils.inspect_function_capture(f))
+
+
+setattr(function, "dispatch_token", "relax")
+
+
+############################# Struct Info ##############################
+
+
+class StructInfoProxy(ObjectGeneric):
+ def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> StructInfo:
+ raise NotImplementedError()
+
+ def get_symbolic_vars(self) -> Set[str]:
+ return {}
+
+ def asobject(self):
+ return self.as_struct_info(None)
+
+
+############################### R.Tensor ###############################
+
+
+def _eval_shape(expr: Union[str, PrimExpr], dict_globals: Optional[Dict[str, Any]]) -> PrimExpr:
+ if isinstance(expr, str):
+ code = compile(expr, "<string>", "eval")
+ return eval(code, dict_globals or {}) # pylint: disable=eval-used
+ else:
+ return expr
+
+
+class TensorProxy(StructInfoProxy):
+ shape: Optional[List[Union[str, PrimExpr]]]
+ dtype: str
+ ndim: int
+
+ def __init__(
+ self,
+ shape: Optional[List[Union[PrimExpr, str]]] = None,
+ dtype: Optional[str] = None,
+ ndim: int = -1,
+ ) -> None:
+ self.shape = shape
+ self.dtype = dtype
+ self.ndim = ndim
+ super().__init__()
+
+ def get_symbolic_vars(self) -> Set[str]:
+ if self.shape is None:
+ return {}
+ else:
+ return {s for s in self.shape if isinstance(s, str) and s.isidentifier()}
+
+ def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> TensorStructInfo:
+ if self.shape is None:
+ return TensorStructInfo(None, self.dtype, self.ndim)
+ else:
+ if dict_globals is None and any([isinstance(s, str) for s in self.shape]):
+ raise ValueError(
+ "String-defined shape expr is only allowed when parsing function parameters "
+ "and return annotations for TVMScript."
+ )
+ shape = [_eval_shape(s, dict_globals) for s in self.shape]
+ return TensorStructInfo(shape, self.dtype, self.ndim)
+
+
+def Tensor(
+ shape: Optional[List[Union[PrimExpr, str]]] = None,
+ dtype: Optional[str] = None,
+ ndim: int = -1,
+) -> TensorProxy:
+ # scalar tensor case
+ if shape is not None and len(shape) == 0:
+ shape = []
+ if isinstance(shape, str) and dtype is None:
+ dtype = shape
+ shape = None
+
+ if shape is not None and not isinstance(shape, (tuple, list)):
+ raise ValueError(f"shape must be a list or tuple, but got: {shape}")
+ return TensorProxy(shape, dtype, ndim)
+
+
+############################## R.Callable ##############################
+
+
+class CallableProxy(StructInfoProxy):
+ params: List[StructInfoProxy]
+ ret: StructInfoProxy
+ """Function type.
+
+ A function type consists of a list of type parameters to enable
+ the definition of generic functions,
+ a set of type constraints which we omit for the time being,
+ a sequence of argument types, and a return type.
+
+ Parameters
+ ----------
+ params : List[StructInfoProxy]
+ The argument StructInfoProxy
+
+ ret : StructInfoProxy
+ The return StructInfoProxy.
+
+ """
+
+ def __init__(
+ self,
+ params: Union[StructInfoProxy, List[StructInfoProxy]],
+ ret: StructInfoProxy,
+ ) -> None:
+ if not isinstance(params, (list, tuple)):
+ params = [params]
+ # convert `R.Tensor` to `R.Tensor()`
+ self.params = [param() if callable(param) else param for param in params]
+ self.ret = ret() if callable(ret) else ret
+
+ def get_symbolic_vars(self) -> Set[str]:
+ return set().union(*[p.get_symbolic_vars() for p in self.params])
+
+ def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> FuncStructInfo:
+ params = [param.as_struct_info(dict_globals) for param in self.params]
+ ret = self.ret.as_struct_info(dict_globals)
+ return FuncStructInfo(params, ret)
+
+
+def Callable(
+ params: Union[StructInfoProxy, List[StructInfoProxy]],
+ ret: StructInfoProxy,
+) -> CallableProxy:
+ return CallableProxy(params, ret)
+
+
+############################### R.Tuple ################################
+
+
+class TupleProxy(StructInfoProxy):
+ fields: List[StructInfoProxy]
+ """The type of tuple values.
+
+ Parameters
+ ----------
+ fields : List[StructInfoProxy]
+ The fields in the tuple
+ """
+
+ def __init__(
+ self,
+ *fields: List[StructInfoProxy],
+ ) -> None:
+ if len(fields) == 1 and isinstance(fields[0], (tuple, list)):
+ fields = fields[0]
+ # convert `R.Tensor` to `R.Tensor()`
+ self.fields = [field() if callable(field) else field for field in fields]
+
+ def get_symbolic_vars(self) -> Set[str]:
+ return set().union(*[f.get_symbolic_vars() for f in self.fields])
+
+ def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> TupleStructInfo:
+ fields = [field.as_struct_info(dict_globals) for field in self.fields]
+ return TupleStructInfo(fields)
+
+
+def Tuple(*fields: List[StructInfoProxy]) -> TupleProxy:
+ return TupleProxy(*fields)
+
+
+############################### R.Shape ################################
+
+
+class ShapeProxy(StructInfoProxy):
+ values: Optional[List[PrimExpr]]
+ ndim: int
+ """The type of shape values.
+
+ Parameters
+ ----------
+ values : Optional[List[PrimExpr]]
+ The symbolic shape values if known.
+
+ ndim : Optional[int]
+ The size of the shape.
+ """
+
+ def __init__(
+ self,
+ values: Optional[List[PrimExpr]] = None,
+ ndim: int = -1,
+ ) -> None:
+ self.values = values
+ self.ndim = ndim
+
+ def get_symbolic_vars(self) -> Set[str]:
+ if self.values is None:
+ return {}
+ else:
+ return {v for v in self.values if isinstance(v, str) and v.isidentifier()}
+
+ def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> ShapeStructInfo:
+ values = [_eval_shape(v, dict_globals) for v in self.values] if self.values else None
+ return ShapeStructInfo(values, self.ndim)
+
+
+def Shape(values: Optional[List[PrimExpr]] = None, ndim: int = -1) -> ShapeProxy:
+ return ShapeProxy(values, ndim)
+
+
+############################### R.Object ################################
+
+
+class ObjectProxy(StructInfoProxy):
+ """The proxy fo ObjectStructInfo.
+
+ Parameters
+ ----------
+ values : Optional[List[PrimExpr]]
+ The symbolic shape values if known.
+
+ ndim : Optional[int]
+ The size of the shape.
+ """
+
+ def __init__(self) -> None:
+ pass
+
+ def get_symbolic_vars(self) -> Set[str]:
+ return set()
+
+ def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> ShapeStructInfo:
+ return ObjectStructInfo()
+
+
+def Object() -> ObjectProxy:
+ return ObjectProxy()
+
+
+################################ R.Prim ################################
+
+
+class PrimProxy(StructInfoProxy):
+ dtype: str
+ """The type of shape values.
+
+ Parameters
+ ----------
+ dtype : str
+ The data type.
+ """
+
+ def __init__(self, dtype: str) -> None:
+ self.dtype = dtype
+
+ def get_symbolic_vars(self) -> Set[str]:
+ return set()
+
+ def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> ShapeStructInfo:
+ return PrimStructInfo(self.dtype)
+
+
+def Prim(dtype: str) -> PrimProxy:
+ return PrimProxy(dtype)
+
+
+############################ R.match_cast #############################
+class MatchCastPair:
+ value: Expr
+ struct_info: StructInfo
+
+ def __init__(self, value: Expr, struct_info: StructInfo) -> None:
+ self.value = value
+ self.struct_info = struct_info
+
+
+def match_cast(value: Expr, struct_info: StructInfo):
+ if value is None:
+ raise ValueError("value of match_cast cannot be None")
+ if struct_info is None:
+ raise ValueError("struct_info of match_cast cannot be None")
+ return MatchCastPair(value, struct_info)
diff --git a/python/tvm/script/parser/relax/parser.py b/python/tvm/script/parser/relax/parser.py
new file mode 100644
index 0000000000..ef26ddd6e9
--- /dev/null
+++ b/python/tvm/script/parser/relax/parser.py
@@ -0,0 +1,276 @@
+# 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.
+# pylint: disable=missing-docstring
+
+import functools
+import numbers
+from typing import Any, Optional
+
+from tvm import relax, tir
+from tvm.ir import structural_equal
+from tvm.relax import StructInfo
+from tvm.relax.utils import convert_to_expr
+from tvm.script.ir_builder.relax.frame import BlockFrame
+
+from ...ir_builder import ir as I
+from ...ir_builder import relax as R
+from ...ir_builder.base import IRBuilder
+from .._core import Parser, dispatch, doc
+from .entry import MatchCastPair, StructInfoProxy, TupleProxy
+
+
+def bind_assign_value(
+ self: Parser,
+ node: doc.expr,
+ var_name: str,
+ value: Any,
+ anno_sinfo: Optional[StructInfo] = None,
+) -> Any:
+ var_table = self.var_table.get()
+
+ if isinstance(value, tir.Var):
+ if value.name and var_name != value.name:
+ self.report_error(
+ node,
+ "Cannot define TIR variables with different names. The LHS of binding should "
+ "has the same name provided in RHS.",
+ )
+ if var_name in var_table:
+ prev_value = var_table[var_name]
+ if not isinstance(prev_value, tir.Var):
+ self.report_error(
+ node,
+ "Cannot redefine a non-TIR-variable object to a TIR variable. Please "
+ "define the TIR variable with another name.",
+ )
+ if prev_value.dtype != value.dtype:
+ self.report_error(
+ node,
+ "Expected the same dtype for TIR vars "
+ f"but got {value.dtype} vs {prev_value.dtype}",
+ )
+ return prev_value
+ IRBuilder.name(var_name, value)
+ return value
+
+ if isinstance(value, tuple):
+ value = convert_to_expr(value)
+ if isinstance(value, numbers.Number):
+ value = R.const(value)
+
+ if isinstance(value, relax.Expr):
+ var = R.emit(value, anno_sinfo)
+ elif isinstance(value, MatchCastPair):
+ if anno_sinfo is not None and not structural_equal(anno_sinfo, value.struct_info):
+ self.report_error(
+ node, "Cannot specify inconsistent annotation for a match cast pair. "
+ )
+ var = R.emit_match_cast(value.value, value.struct_info)
+ else:
+ raise TypeError(f"Unsupported type {type(value)} in assignment")
+
+ IRBuilder.name(var_name, var)
+ return var
+
+
+def eval_struct_info_proxy(self: Parser, node: doc.expr) -> StructInfoProxy:
+ try:
+ annotation = self.eval_expr(node)
+ if annotation is None:
+ return TupleProxy([])
+ if callable(annotation):
+ annotation = annotation()
+ if isinstance(annotation, StructInfoProxy):
+ return annotation
+ else:
+ raise TypeError(f"Expected StructInfoProxy but got {type(annotation)}.")
+ except Exception as err:
+ self.report_error(node, str(err))
+ raise err
+
+
+def eval_struct_info(self: Parser, node: doc.expr, eval_str: bool = False) -> StructInfo:
+ var_table = self.var_table.get() if eval_str else None
+ try:
+ return eval_struct_info_proxy(self, node).as_struct_info(var_table)
+ except Exception as err:
+ self.report_error(node, str(err))
+ raise err
+
+
+def collect_symbolic_var_from_params(self: Parser, node: doc.FunctionDef) -> None:
+ # Collect symbolic vars from parameters
+ symbolic_vars = set()
+ for arg in node.args.args:
+ if arg.annotation is None:
+ self.report_error(arg, "Type annotation is required for function parameters.")
+ param_sinfo_proxy = eval_struct_info_proxy(self, arg.annotation)
+ symbolic_vars.update(param_sinfo_proxy.get_symbolic_vars())
+
+ # Define symbolic vars to the current var_table frame
+ for var_name in symbolic_vars:
+ self.var_table.add(var_name, tir.Var(var_name, "int64"), allow_shadowing=False)
+
+
+@dispatch.register(token="relax", type_name="FunctionDef")
+def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
+ with self.var_table.with_frame():
+ with self.with_dispatch_token("relax"):
+ with R.function():
+ R.func_name(node.name)
+ collect_symbolic_var_from_params(self, node)
+
+ if node.returns is not None:
+ ann_sinfo = eval_struct_info(self, node.returns, eval_str=True)
+ R.func_ret_struct_info(ann_sinfo)
+
+ self.visit(node.args)
+ self.visit_body(node.body)
+
+
+@dispatch.register(token="relax", type_name="tvm_declare_function")
+def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> None:
+ with self.var_table.with_frame():
+ collect_symbolic_var_from_params(self, node)
+
+ if node.returns is None:
+ # Use ObjectStructInfo as unknown return type
+ # NOTE: Cannot use VoidStructInfo here because the return type can be refined later.
+ ret_sinfo = relax.ObjectStructInfo()
+ else:
+ ret_sinfo = eval_struct_info(self, node.returns, eval_str=True)
+ params = []
+ params_sinfo = []
+ for arg in node.args.args:
+ if arg.annotation is None:
+ self.report_error(arg, "Type annotation is required for function parameters.")
+ param_sinfo = eval_struct_info(self, arg.annotation, eval_str=True)
+ params_sinfo.append(param_sinfo)
+ params.append(relax.Var(arg.arg, param_sinfo))
+
+ func_signature = relax.Function.create_empty(params, ret_sinfo)
+ global_var = I.decl_function(node.name, func_signature)
+ self.var_table.add(node.name, global_var)
+
+
+@dispatch.register(token="relax", type_name="pre_token_switch")
+def pre_token_switch(self: Parser, node: doc.Expr) -> None: # pylint: disable=unused-argument
+ ir_builder = IRBuilder()
+ ir_builder.__enter__()
+
+
+@dispatch.register(token="relax", type_name="post_token_switch")
+def post_token_switch(self: Parser, node: doc.Expr) -> None:
+ ir_builder = IRBuilder.current()
+ result = ir_builder.get()
+ ir_builder.__exit__(None, None, None)
+ var = R.emit(result)
+ IRBuilder.name(node.name, var)
+ self.var_table.add(node.name, var, allow_shadowing=False)
+
+
+@dispatch.register(token="relax", type_name="Expr")
+def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
+ value = self.eval_expr(node.value)
+ if value is not None:
+ self.report_error(node, f"Unsupported Expr stmt type {value}.")
+
+
+@dispatch.register(token="relax", type_name="arguments")
+def visit_arguments(self: Parser, node: doc.arguments) -> None:
+ arg: doc.arg
+ for arg in node.args:
+ if arg.annotation is None:
+ self.report_error(arg, "Type annotation is required for function parameters.")
+ param_sinfo = eval_struct_info(self, arg.annotation, eval_str=True)
+ param = R.arg(arg.arg, param_sinfo)
+
+ self.var_table.add(arg.arg, param)
+
+
+@dispatch.register(token="relax", type_name="tvm_annotation")
+def visit_tvm_annotation(self: Parser, node: doc.expr) -> StructInfo:
+ return eval_struct_info(self, node, eval_str=False)
+
+
+@dispatch.register(token="relax", type_name="With")
+def visit_with(self: Parser, node: doc.With) -> None:
+ # Currently only `with R.dataflow()` is supported
+ if len(node.items) != 1:
+ self.report_error(node, "Only one item is allowed.")
+ item = node.items[0]
+ if item.optional_vars is not None:
+ self.report_error(
+ item.context_expr,
+ "Relax syntax doesn't allow binding expressions in `with` to variables",
+ )
+ frame = self.eval_expr(item.context_expr)
+ with self.var_table.with_frame():
+ with frame:
+ self.visit(node.body)
+ if isinstance(frame, BlockFrame) and frame.is_dataflow:
+ output_vars = frame.output_vars
+ for var in output_vars:
+ self.var_table.add(var.name_hint, var, allow_shadowing=True)
+
+
+@dispatch.register(token="relax", type_name="Assign")
+def visit_assign(self: Parser, node: doc.Assign) -> None:
+ if len(node.targets) != 1:
+ self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
+ lhs = node.targets[0]
+ rhs = self.eval_expr(node.value)
+ self.eval_assign(
+ target=lhs,
+ source=rhs,
+ bind_value=bind_assign_value,
+ allow_shadowing=True,
+ )
+
+
+@dispatch.register(token="relax", type_name="AnnAssign")
+def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
+ lhs = node.target
+ rhs = self.eval_expr(node.value)
+ anno_sinfo = self.visit_tvm_annotation(node.annotation)
+ self.eval_assign(
+ target=lhs,
+ source=rhs,
+ bind_value=functools.partial(bind_assign_value, anno_sinfo=anno_sinfo),
+ allow_shadowing=True,
+ )
+
+
+@dispatch.register(token="relax", type_name="Return")
+def visit_return(self: Parser, node: doc.Assign) -> None:
+ value = self.eval_expr(node.value)
+ value = convert_to_expr(value)
+ R.func_ret_value(value)
+
+
+@dispatch.register(token="relax", type_name="If")
+def visit_if(self: Parser, node: doc.If) -> None:
+ if node.orelse is None:
+ raise ValueError("Else statements are required for relax dialect.")
+ with R.If(self.eval_expr(node.test)) as if_frame:
+ with self.var_table.with_frame():
+ with R.Then():
+ self.visit_body(node.body)
+ with self.var_table.with_frame():
+ with R.Else():
+ self.visit_body(node.orelse)
+ self.var_table.add(if_frame.var_name, if_frame.var, allow_shadowing=True)
diff --git a/python/tvm/script/parser/tir/entry.py b/python/tvm/script/parser/tir/entry.py
index e7ec7cf886..b38cbb09cf 100644
--- a/python/tvm/script/parser/tir/entry.py
+++ b/python/tvm/script/parser/tir/entry.py
@@ -83,7 +83,7 @@ class BufferProxy:
return self(keys)
if len(keys) >= 2 and not isinstance(keys[1], str):
return self(keys)
- return self(*keys) # pylint: disable=no-member # type: ignore
+ return self(*keys) # type: ignore[attr-defined] # pylint: disable=no-member
class PtrProxy:
@@ -94,7 +94,7 @@ class PtrProxy:
def __call__(self, dtype, storage_scope="global"):
if callable(dtype):
dtype = dtype().dtype
- return ptr(dtype, storage_scope) # pylint: disable=no-member # type: ignore
+ return ptr(dtype, storage_scope) # type: ignore[attr-defined] # pylint: disable=no-member
def __getitem__(self, keys):
if not isinstance(keys, tuple):
diff --git a/python/tvm/script/parser/tir/operation.py b/python/tvm/script/parser/tir/operation.py
index f0c04f47cd..ed8f07a063 100644
--- a/python/tvm/script/parser/tir/operation.py
+++ b/python/tvm/script/parser/tir/operation.py
@@ -46,12 +46,12 @@ def _register_expr_op(ty: Type): # pylint: disable=invalid-name
for i in [0, 1]:
# Case 1. binop
- r(doc.Add, i, tir.Add)
- r(doc.Sub, i, tir.Sub)
- r(doc.Mult, i, tir.Mul)
- r(doc.Div, i, tir.Div)
- r(doc.FloorDiv, i, tir.FloorDiv)
- r(doc.Mod, i, tir.FloorMod)
+ r(doc.Add, i, lambda a, b: a + b)
+ r(doc.Sub, i, lambda a, b: a - b)
+ r(doc.Mult, i, lambda a, b: a * b)
+ r(doc.Div, i, lambda a, b: a / b)
+ r(doc.FloorDiv, i, lambda a, b: a // b)
+ r(doc.Mod, i, lambda a, b: a % b)
r(doc.LShift, i, lambda a, b: a << b)
r(doc.RShift, i, lambda a, b: a >> b)
r(doc.BitOr, i, lambda a, b: a | b)
diff --git a/python/tvm/script/parser/tir/parser.py b/python/tvm/script/parser/tir/parser.py
index 0e74114ba2..07bf052d41 100644
--- a/python/tvm/script/parser/tir/parser.py
+++ b/python/tvm/script/parser/tir/parser.py
@@ -24,6 +24,7 @@ import tvm
from tvm.ir import PrimType
from tvm.tir import Buffer, IterVar, PrimExpr, Var
+from ...ir_builder import ir as I
from ...ir_builder import tir as T
from ...ir_builder.base import IRBuilder
from ...ir_builder.base import IRBuilderFrame as Frame
@@ -471,3 +472,28 @@ def visit_return(self: Parser, node: doc.Return) -> None:
The doc AST return node.
"""
self.report_error(node, "Return is not allowed.")
+
+
+@dispatch.register(token="tir", type_name="tvm_declare_function")
+def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> None:
+ """The function declaration step for tir
+
+ Parameters
+ ----------
+ self : Parser
+ The visiting parser.
+
+ node : doc.Return
+ The doc AST return node.
+ """
+
+ ret_type = None
+ if node.returns is not None:
+ ret_type = self.eval_expr(node.returns)
+ if callable(ret_type):
+ ret_type = PrimType(ret_type().dtype)
+
+ # Only ret_type is needed for func_signature.
+ func_signature = tvm.tir.PrimFunc([], None, ret_type=ret_type)
+ global_var = I.decl_function(node.name, func_signature)
+ self.var_table.add(node.name, global_var)
diff --git a/src/ir/module.cc b/src/ir/module.cc
index 22c6faf3d6..4a09bdaaf7 100644
--- a/src/ir/module.cc
+++ b/src/ir/module.cc
@@ -63,15 +63,25 @@ IRModule::IRModule(tvm::Map<GlobalVar, BaseFunc> functions,
}
bool IRModuleNode::SEqualReduce(const IRModuleNode* other, SEqualReducer equal) const {
- if (functions.size() != other->functions.size()) return false;
if (!equal(this->attrs, other->attrs)) return false;
+
+ if (functions.size() != other->functions.size()) return false;
+ // Update GlobalVar remap
+ for (const auto& gv : this->GetGlobalVars()) {
+ if (!other->ContainGlobalVar(gv->name_hint)) return false;
+ if (!equal.DefEqual(gv, other->GetGlobalVar(gv->name_hint))) return false;
+ }
for (const auto& kv : this->functions) {
- if (!other->ContainGlobalVar(kv.first->name_hint)) return false;
if (!equal(kv.second, other->Lookup(kv.first->name_hint))) return false;
}
+
if (type_definitions.size() != other->type_definitions.size()) return false;
+ // Update GlobalTypeVar remap
+ for (const auto& gtv : this->GetGlobalTypeVars()) {
+ if (!other->ContainGlobalTypeVar(gtv->name_hint)) return false;
+ if (!equal.DefEqual(gtv, other->GetGlobalTypeVar(gtv->name_hint))) return false;
+ }
for (const auto& kv : this->type_definitions) {
- if (!other->ContainGlobalTypeVar(kv.first->name_hint)) return false;
if (!equal(kv.second, other->LookupTypeDef(kv.first->name_hint))) return false;
}
return true;
diff --git a/src/script/ir_builder/ir/frame.cc b/src/script/ir_builder/ir/frame.cc
index a81c56922d..addf129284 100644
--- a/src/script/ir_builder/ir/frame.cc
+++ b/src/script/ir_builder/ir/frame.cc
@@ -26,11 +26,15 @@ namespace ir_builder {
namespace ir {
void IRModuleFrameNode::ExitWithScope() {
- ICHECK_EQ(functions.size(), global_vars.size());
- int n = functions.size();
Map<GlobalVar, BaseFunc> func_map;
- for (int i = 0; i < n; ++i) {
- func_map.Set(global_vars[i], functions[i]);
+ CHECK_EQ(functions.size(), global_var_map.size())
+ << "All functions must be defined in the IRModule. Got " << global_var_map.size()
+ << "declared function(s), but only " << functions.size() << "defined function(s).";
+ for (const auto& kv : functions) {
+ const GlobalVar& gv = kv.first;
+ const BaseFunc& func = kv.second;
+ CHECK(func.defined()) << "ValueError: function " << gv->name_hint << " is not defined";
+ func_map.Set(gv, func);
}
IRBuilder builder = IRBuilder::Current();
ICHECK(!builder->result.defined()) << "ValueError: Builder.result has already been set";
diff --git a/src/script/ir_builder/ir/ir.cc b/src/script/ir_builder/ir/ir.cc
index a8cc452e4f..da2330b577 100644
--- a/src/script/ir_builder/ir/ir.cc
+++ b/src/script/ir_builder/ir/ir.cc
@@ -17,9 +17,12 @@
* under the License.
*/
#include <tvm/ir/module.h>
+#include <tvm/relax/analysis.h>
#include <tvm/runtime/registry.h>
#include <tvm/script/ir_builder/ir/ir.h>
+#include "./utils.h"
+
namespace tvm {
namespace script {
namespace ir_builder {
@@ -27,12 +30,48 @@ namespace ir {
IRModuleFrame IRModule() {
ObjectPtr<IRModuleFrameNode> n = make_object<IRModuleFrameNode>();
- n->global_vars.clear();
+ n->global_var_map.clear();
n->functions.clear();
return IRModuleFrame(n);
}
+GlobalVar DeclFunction(const String& func_name, const BaseFunc& func_signature) {
+ IRModuleFrame frame = FindModuleFrame("I.DeclFunction");
+ CHECK(!frame->global_var_map.count(func_name))
+ << "ValueError: function " << func_name << " already exists";
+ GlobalVar gv = GlobalVar(func_name);
+ if (func_signature->struct_info_.defined()) {
+ gv->struct_info_ = tvm::relax::GetStructInfo(func_signature);
+ } else if (const auto* prim_func = func_signature.as<tvm::tir::PrimFuncNode>()) {
+ gv->struct_info_ =
+ tvm::relax::FuncStructInfo::OpaqueFunc(tvm::relax::StructInfoFromType(prim_func->ret_type));
+ } else {
+ LOG(FATAL) << "Unsupported function type: " << func_signature->GetTypeKey();
+ }
+ CHECK(frame->functions.find(gv) == frame->functions.end())
+ << "ValueError: function " << func_name << " has already been defined.";
+ frame->global_var_map.Set(func_name, gv);
+ if (func_signature.defined()) {
+ frame->functions.Set(gv, func_signature);
+ }
+ return gv;
+}
+
+void DefFunction(const String& func_name, const BaseFunc& func) {
+ IRModuleFrame frame = FindModuleFrame("I.DefFunction");
+ auto it = frame->global_var_map.find(func_name);
+ CHECK(it != frame->global_var_map.end())
+ << "ValueError: function " << func_name << " does not exist, please declare it first.";
+ const GlobalVar& gv = (*it).second;
+ frame->functions.Set(gv, func);
+ if (func->checked_type_.defined()) {
+ gv->checked_type_ = func->checked_type_;
+ }
+}
+
TVM_REGISTER_GLOBAL("script.ir_builder.ir.IRModule").set_body_typed(IRModule);
+TVM_REGISTER_GLOBAL("script.ir_builder.ir.DeclFunction").set_body_typed(DeclFunction);
+TVM_REGISTER_GLOBAL("script.ir_builder.ir.DefFunction").set_body_typed(DefFunction);
} // namespace ir
} // namespace ir_builder
diff --git a/src/script/ir_builder/ir/frame.cc b/src/script/ir_builder/ir/utils.h
similarity index 59%
copy from src/script/ir_builder/ir/frame.cc
copy to src/script/ir_builder/ir/utils.h
index a81c56922d..58d5e53f70 100644
--- a/src/script/ir_builder/ir/frame.cc
+++ b/src/script/ir_builder/ir/utils.h
@@ -16,8 +16,9 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include <tvm/ir/module.h>
-#include <tvm/runtime/registry.h>
+#ifndef TVM_SCRIPT_IR_BUILDER_IR_UTILS_H_
+#define TVM_SCRIPT_IR_BUILDER_IR_UTILS_H_
+
#include <tvm/script/ir_builder/ir/frame.h>
namespace tvm {
@@ -25,21 +26,24 @@ namespace script {
namespace ir_builder {
namespace ir {
-void IRModuleFrameNode::ExitWithScope() {
- ICHECK_EQ(functions.size(), global_vars.size());
- int n = functions.size();
- Map<GlobalVar, BaseFunc> func_map;
- for (int i = 0; i < n; ++i) {
- func_map.Set(global_vars[i], functions[i]);
- }
+inline IRModuleFrame FindModuleFrame(const String& method) {
IRBuilder builder = IRBuilder::Current();
- ICHECK(!builder->result.defined()) << "ValueError: Builder.result has already been set";
- builder->result = tvm::IRModule(func_map);
+ if (Optional<IRModuleFrame> frame = builder->FindFrame<IRModuleFrame>()) {
+ const Optional<IRModuleFrame>& last_module_frame = builder->GetLastFrame<IRModuleFrame>();
+ if (last_module_frame.defined() && last_module_frame.value() == frame) {
+ return frame.value();
+ }
+ } else {
+ LOG(FATAL) << "ValueError: IRModule frame not find. Please ensure '" << method
+ << "' is called under I.ir_module()";
+ }
+ LOG(FATAL) << "ValueError: '" << method << "' must be called immediately under I.ir_module()";
+ throw;
}
-TVM_REGISTER_NODE_TYPE(IRModuleFrameNode);
-
} // namespace ir
} // namespace ir_builder
} // namespace script
} // namespace tvm
+
+#endif // TVM_SCRIPT_IR_BUILDER_IR_UTILS_H_
diff --git a/src/script/ir_builder/relax/frame.cc b/src/script/ir_builder/relax/frame.cc
new file mode 100644
index 0000000000..c78b9e73c5
--- /dev/null
+++ b/src/script/ir_builder/relax/frame.cc
@@ -0,0 +1,273 @@
+/*
+ * 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.
+ */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/script/ir_builder/relax/frame.h>
+#include <tvm/script/ir_builder/relax/ir.h>
+
+#include "./utils.h"
+
+namespace tvm {
+namespace script {
+namespace ir_builder {
+namespace relax {
+
+void SeqExprFrameNode::ExitWithScope() {
+ // At this moment, there should be at most one BlockFrame which hasn't ended. In this case, call
+ // its `ExitBlockFrame` and check if there is any more unended BlockFrame.
+ if (Optional<BlockFrame> block_frame = IRBuilder::Current()->GetLastFrame<BlockFrame>()) {
+ block_frame.value()->ExitWithScope();
+ ICHECK(!IRBuilder::Current()->GetLastFrame<BlockFrame>().defined())
+ << "ValueError: There is some remaining BlockFrame that is not properly popped out.";
+ }
+ RelaxFrameNode::ExitWithScope();
+}
+
+void SeqExprFrameNode::EnterWithScope() {
+ RelaxFrameNode::EnterWithScope();
+ BindingBlock()->EnterWithScope();
+}
+
+void FunctionFrameNode::ExitWithScope() {
+ using ir::IRModuleFrame;
+ using tvm::relax::Expr;
+ IRBuilder builder = IRBuilder::Current();
+ SeqExprFrameNode::ExitWithScope();
+ // Step 1: Create the function.
+ CHECK(output.defined()) << "ValueError: A Relax function must have a return value. Please use "
+ "`return` to return an Expr";
+ this->block_builder->BeginScope(params);
+ Expr body = this->block_builder->Normalize(tvm::relax::SeqExpr(binding_blocks, output.value()));
+ auto dict_attrs = attrs.empty() ? NullValue<DictAttrs>() : DictAttrs(attrs);
+ this->block_builder->EndScope();
+ tvm::relax::Function func(/*params=*/params,
+ /*body=*/body,
+ /*ret_struct_info=*/ret_struct_info,
+ /*attrs=*/dict_attrs);
+ // Step 2: Update IRModule.
+ if (builder->frames.empty()) {
+ // Case 0. No outer frame, return function directly
+ ICHECK(!builder->result.defined()) << "ValueError: Builder.result has already been set";
+ builder->result = func;
+ } else if (Optional<IRModuleFrame> opt_frame = builder->FindFrame<IRModuleFrame>()) {
+ // Case 1. A global function of an IRModule
+ CHECK(name.defined()) << "ValueError: The function name must be defined before exiting the "
+ "function scope, if it's defined in a Module";
+ const IRModuleFrame& frame = opt_frame.value();
+ const String& func_name = name.value_or("");
+ if (!frame->global_var_map.count(func_name)) {
+ // First time visiting the function.
+ ir::DeclFunction(func_name, func);
+ }
+ // Define the function.
+ // Note we do checks to disallow redefinition of functions inside the `DefFunction`.
+ ir::DefFunction(func_name, func);
+ } else {
+ LOG(FATAL) << "ValueError: Cannot find where to insert Relax.Function";
+ }
+}
+
+void BlockFrameNode::EnterWithScope() {
+ // Step 1. If the last frame is a block frame. The start of a new block frame marks the end of the
+ // last block frame.
+ Optional<BlockFrame> block_frame = IRBuilder::Current()->GetLastFrame<BlockFrame>();
+ if (block_frame.defined()) {
+ block_frame.value()->ExitWithScope();
+ // Block frames cannot appear consecutively.
+ ICHECK(!IRBuilder::Current()->GetLastFrame<BlockFrame>());
+ }
+ // Step 2. Deal with the new block frame.
+ RelaxFrameNode::EnterWithScope();
+ Optional<FunctionFrame> func_frame = IRBuilder::Current()->FindFrame<FunctionFrame>();
+ CHECK(func_frame.defined())
+ << "ValueError: Cannot find FunctionFrame when creating BindingBlocks, Please ensure "
+ "creating the block under Relax function scope.";
+ const tvm::relax::BlockBuilder& block_builder = func_frame.value()->block_builder;
+ if (is_dataflow) {
+ block_builder->BeginDataflowBlock();
+ } else {
+ block_builder->BeginBindingBlock();
+ }
+}
+
+class DataflowBlockRewriter : public tvm::relax::ExprMutator {
+ public:
+ static tvm::relax::DataflowBlock Rewrite(const tvm::relax::DataflowBlock& block,
+ const Array<tvm::relax::Var>& output_vars) {
+ DataflowBlockRewriter rewriter(output_vars);
+ return Downcast<tvm::relax::DataflowBlock>(rewriter.VisitBindingBlock(block));
+ }
+
+ private:
+ explicit DataflowBlockRewriter(const Array<tvm::relax::Var>& output_vars) {
+ for (const tvm::relax::Var& var : output_vars) {
+ output_var_set_.insert(var.get());
+ }
+ }
+
+ tvm::relax::Var VisitVarDef_(const tvm::relax::DataflowVarNode* op) final {
+ auto it = output_var_set_.find(op);
+ if (it != output_var_set_.end()) {
+ // Rewrite dataflow vars to global vars
+ auto n = make_object<tvm::relax::VarNode>(*op);
+ tvm::relax::Var new_var(n);
+ this->var_remap_[op->vid] = new_var;
+ return new_var;
+ } else {
+ return GetRef<tvm::relax::Var>(op);
+ }
+ }
+
+ private:
+ std::unordered_set<const tvm::relax::VarNode*> output_var_set_;
+};
+
+void BlockFrameNode::ExitWithScope() {
+ // Step 1. Pop the current frame out of the frame stack.
+ RelaxFrameNode::ExitWithScope();
+
+ // Step 2. Get the constructed binding block from the block builder. The block should have at
+ // lease one binding - otherwise, the block is not supposed to be created.
+ const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
+ tvm::relax::BindingBlock block = block_builder->EndBlock();
+ if (block->bindings.empty()) {
+ return;
+ }
+
+ // Step 3. Rewrite the dataflow block.
+ if (is_dataflow) {
+ // Step 3.1. Rewrite block binding
+ block = DataflowBlockRewriter::Rewrite(Downcast<tvm::relax::DataflowBlock>(block), output_vars);
+
+ // Step 3.2. Collect global vars' reference in bindings
+ Map<tvm::relax::Id, tvm::relax::Var> new_global_vars;
+ for (const tvm::relax::Binding& binding : block->bindings) {
+ if (!binding->var->IsInstance<tvm::relax::DataflowVarNode>()) {
+ new_global_vars.Set(binding->var->vid, binding->var);
+ }
+ }
+
+ // Step 3.3. Rewrite output vars
+ Array<tvm::relax::Var> new_output_vars;
+ for (const auto& var : output_vars) {
+ auto it = new_global_vars.find(var->vid);
+ ICHECK(it != new_global_vars.end());
+ new_output_vars.push_back((*it).second);
+ }
+ output_vars = std::move(new_output_vars);
+ }
+
+ // Step 3. Get the last frame from the IRBuilder frame stack.
+ Optional<RelaxFrame> opt_last_frame = IRBuilder::Current()->GetLastFrame<RelaxFrame>();
+ ICHECK(opt_last_frame.defined());
+ RelaxFrame last_frame = opt_last_frame.value();
+
+ // Step 4. Since we popped out any possible block frame when entering the "with" scope of the
+ // current frame, the last frame cannot be a block frame.
+ ICHECK(!last_frame->IsInstance<BlockFrameNode>());
+
+ // Step 5. Push the block frame into the corresponding field of the last frame.
+ if (const auto* seq_frame = last_frame.as<SeqExprFrameNode>()) {
+ ICHECK(!seq_frame->output.defined())
+ << "The function is not expected to have output values when emitting blocks.";
+ auto frame = GetRef<SeqExprFrame>(seq_frame);
+ frame->binding_blocks.push_back(block);
+ } else {
+ LOG(FATAL) << "ValueError: Currently the last frame is supposed to be either a function frame "
+ "or a block frame. However, the last frame is \""
+ << last_frame->GetTypeKey() << "\".";
+ }
+
+ // Step 6. Start another binding block when a dataflow block ended.
+ if (is_dataflow) {
+ BindingBlock()->EnterWithScope();
+ }
+}
+
+void IfFrameNode::EnterWithScope() {
+ const Array<IRBuilderFrame>& frames = IRBuilder::Current()->frames;
+ for (const IRBuilderFrame& frame : frames) {
+ const auto* block_frame = frame.as<BlockFrameNode>();
+ if (block_frame && block_frame->is_dataflow) {
+ LOG(FATAL) << "ValueError: Cannot create an IfFrame inside a dataflow block.";
+ }
+ }
+ RelaxFrameNode::EnterWithScope();
+}
+
+void IfFrameNode::ExitWithScope() {
+ RelaxFrameNode::ExitWithScope();
+ CHECK(then_expr.defined())
+ << "ValueError: The body of then part is expected to be defined before exiting.";
+ CHECK(then_expr.defined())
+ << "ValueError: The body of else part is expected to be defined before exiting.";
+ auto body = tvm::relax::If(condition, then_expr.value(), else_expr.value());
+ var = Emit(body);
+ IRBuilder::Name(var_name, var);
+}
+
+void ThenFrameNode::EnterWithScope() {
+ IfFrame frame = FindIfFrame("R.Then");
+ CHECK(!frame->then_expr.defined())
+ << "ValueError: Duplicate then branch declaration, previous one is "
+ << frame->then_expr.value();
+ SeqExprFrameNode::EnterWithScope();
+}
+
+void ThenFrameNode::ExitWithScope() {
+ SeqExprFrameNode::ExitWithScope();
+ String var_name;
+ output = GetSeqExprForBranch(GetRef<ThenFrame>(this), &var_name);
+ IfFrame frame = FindIfFrame("R.Then");
+ frame->then_expr = output;
+ frame->var_name = var_name;
+}
+
+void ElseFrameNode::EnterWithScope() {
+ IfFrame frame = FindIfFrame("R.Else");
+ CHECK(frame->then_expr.defined()) << "The else branch should follow then branch";
+ CHECK(!frame->else_expr.defined())
+ << "ValueError: Duplicate else branch declaration, previous one is "
+ << frame->else_expr.value();
+ SeqExprFrameNode::EnterWithScope();
+}
+
+void ElseFrameNode::ExitWithScope() {
+ SeqExprFrameNode::ExitWithScope();
+ String var_name;
+ output = GetSeqExprForBranch(GetRef<ElseFrame>(this), &var_name);
+ IfFrame frame = FindIfFrame("R.Else");
+ frame->else_expr = output;
+ CHECK(frame->var_name == var_name)
+ << "This last binding of both branches must have the same variable.";
+}
+
+TVM_REGISTER_NODE_TYPE(FunctionFrameNode);
+TVM_REGISTER_NODE_TYPE(SeqExprFrameNode);
+TVM_REGISTER_NODE_TYPE(BlockFrameNode);
+TVM_REGISTER_NODE_TYPE(IfFrameNode);
+TVM_REGISTER_NODE_TYPE(ThenFrameNode);
+TVM_REGISTER_NODE_TYPE(ElseFrameNode);
+
+} // namespace relax
+} // namespace ir_builder
+} // namespace script
+} // namespace tvm
diff --git a/src/script/ir_builder/relax/ir.cc b/src/script/ir_builder/relax/ir.cc
new file mode 100644
index 0000000000..ece645243c
--- /dev/null
+++ b/src/script/ir_builder/relax/ir.cc
@@ -0,0 +1,236 @@
+/*
+ * 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.
+ */
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/script/ir_builder/relax/ir.h>
+#include <tvm/tir/op.h>
+
+#include "./utils.h"
+
+namespace tvm {
+namespace script {
+namespace ir_builder {
+namespace relax {
+
+///////////////////////////////// Vars //////////////////////////////////
+
+using tvm::script::ir_builder::details::Namer;
+
+TVM_STATIC_IR_FUNCTOR(Namer, vtable)
+ .set_dispatch<tvm::relax::VarNode>([](const ObjectRef& node, String name) -> void {
+ using tvm::relax::VarNode;
+ using tvm::relax::IdNode;
+ const VarNode* var = node.as<VarNode>();
+ IdNode* vid = const_cast<IdNode*>(var->vid.get());
+ vid->name_hint = name;
+ });
+
+TVM_STATIC_IR_FUNCTOR(Namer, vtable)
+ .set_dispatch<tvm::relax::DataflowVarNode>([](const ObjectRef& node, String name) -> void {
+ using tvm::relax::DataflowVarNode;
+ using tvm::relax::IdNode;
+ const DataflowVarNode* var = node.as<DataflowVarNode>();
+ IdNode* vid = const_cast<IdNode*>(var->vid.get());
+ vid->name_hint = name;
+ });
+
+/////////////////////////////// Function ////////////////////////////////
+
+FunctionFrame Function() {
+ ObjectPtr<FunctionFrameNode> n = make_object<FunctionFrameNode>();
+ const IRBuilder& ir_builder = IRBuilder::Current();
+ Optional<tvm::IRModule> mod = NullOpt;
+ if (const Optional<ir::IRModuleFrame> mod_frame = ir_builder->GetLastFrame<ir::IRModuleFrame>()) {
+ mod = tvm::IRModule(mod_frame.value()->functions);
+ }
+ n->block_builder = tvm::relax::BlockBuilder::Create(/*mod=*/mod);
+ return FunctionFrame(n);
+}
+
+tvm::relax::Var Arg(const String& name, const tvm::relax::StructInfo& struct_info) {
+ FunctionFrame frame = FindFunctionFrame("R.Arg");
+ tvm::relax::Var var(name, struct_info);
+ frame->params.push_back(var);
+ return var;
+}
+
+void FuncName(const String& name) {
+ FunctionFrame frame = FindFunctionFrame("R.func_name");
+ if (frame->name.defined()) {
+ LOG(FATAL) << "ValueError: Duplicate function name, previous one is: \"" << frame->name.value()
+ << "\"";
+ }
+ frame->name = name;
+}
+
+void FuncAttrs(Map<String, ObjectRef> attrs) {
+ FunctionFrame frame = FindFunctionFrame("R.func_attr");
+ if (!frame->attrs.empty()) {
+ LOG(FATAL) << "ValueError: Duplicate function attrs, previous one is:\n" << frame->attrs;
+ }
+ frame->attrs = attrs;
+}
+
+void FuncRetStructInfo(const tvm::relax::StructInfo& ret_sinfo) {
+ FunctionFrame frame = FindFunctionFrame("R.func_ret_struct_info");
+ if (frame->ret_struct_info.defined()) {
+ LOG(FATAL) << "ValueError: Duplicate function return struct info, previous one is:\n "
+ << frame->ret_struct_info.value();
+ }
+ frame->ret_struct_info = ret_sinfo;
+}
+
+void FuncRetValue(const tvm::relax::Expr& value) {
+ // Step 0. Normalize the value.
+ const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
+ tvm::relax::Expr normalized_value = block_builder->Normalize(value);
+
+ // Step 1. The current Relax TVMScript syntax only allows function return appearing at the end of
+ // a function body. Therefore if there is any unended block frame when dealing with function
+ // return, we should end the block frame.
+ Optional<BlockFrame> block_frame = IRBuilder::Current()->GetLastFrame<BlockFrame>();
+ if (block_frame.defined()) {
+ block_frame.value()->ExitWithScope();
+ ICHECK(!IRBuilder::Current()->FindFrame<BlockFrame>())
+ << "All block frame are supposed to be popped out already";
+ }
+ // Step 2. Add the output value to the function frame.
+ FunctionFrame frame = FindFunctionFrame("return");
+ CHECK(!frame->output.defined())
+ << "ValueError: Relax functions don't support multiple return statement. Please make sure "
+ "the return statement appears at the end of function.";
+
+ frame->output = std::move(normalized_value);
+}
+
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.Function").set_body_typed(Function);
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.Arg").set_body_typed(Arg);
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.FuncName").set_body_typed(FuncName);
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.FuncAttrs").set_body_typed(FuncAttrs);
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.FuncRetStructInfo").set_body_typed(FuncRetStructInfo);
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.FuncRetValue").set_body_typed(FuncRetValue);
+
+///////////////////////////// BindingBlock //////////////////////////////
+
+BlockFrame Dataflow() {
+ ObjectPtr<BlockFrameNode> n = make_object<BlockFrameNode>();
+ n->is_dataflow = true;
+ n->block_ended = false;
+ return BlockFrame(n);
+}
+
+BlockFrame BindingBlock() {
+ ObjectPtr<BlockFrameNode> n = make_object<BlockFrameNode>();
+ n->is_dataflow = false;
+ n->block_ended = false;
+ return BlockFrame(n);
+}
+
+void DataflowBlockOutput(const Array<tvm::relax::Var>& vars) {
+ // Step 1. Check that we're in a Dataflow block that is not ended.
+ Optional<BlockFrame> block_frame = IRBuilder::Current()->GetLastFrame<BlockFrame>();
+ CHECK(block_frame.defined() && block_frame.value()->is_dataflow)
+ << "ValueError: `R.output` should appear inside a dataflow block. However, the current "
+ "innermost block is not a dataflow block.";
+ CHECK(!block_frame.value()->block_ended)
+ << "ValueError: It is not allowed for a dataflow block to have multiple output operation.";
+
+ // Step 2. Mark the block frame ended of construction, so that any followup binding after this
+ // mark in the dataflow block will lead to an error.
+ block_frame.value()->block_ended = true;
+
+ // Step 3. All the output variables must be global variables and must be emitted by this dataflow
+ // block.
+ const Array<tvm::relax::Var>& emitted_vars = block_frame.value()->emitted_vars;
+ for (const tvm::relax::Var& var : vars) {
+ CHECK(std::find(emitted_vars.begin(), emitted_vars.end(), var) != emitted_vars.end())
+ << "ValueError: An output variable is not emitted by this dataflow block. Please make sure "
+ "all dataflow block output variables are emitted exactly by this block.";
+ block_frame.value()->output_vars.push_back(var);
+ }
+}
+
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.Dataflow").set_body_typed(Dataflow);
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.BindingBlock").set_body_typed(BindingBlock);
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.DataflowBlockOutput")
+ .set_body_typed(DataflowBlockOutput);
+
+/////////////////////////////// Bindings ///////////////////////////////
+
+tvm::relax::Var Emit(const tvm::relax::Expr& expr,
+ const Optional<tvm::relax::StructInfo>& annotate_struct_info) {
+ using tvm::relax::GetStructInfo;
+ BlockFrame block_frame = CheckBlockFrameExistAndUnended();
+ const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
+ if (annotate_struct_info.defined()) {
+ const auto& sinfo = annotate_struct_info.value();
+ if (!expr->struct_info_.defined()) {
+ UpdateStructInfo(expr, sinfo);
+ } else {
+ CHECK(StructInfoBaseCheck(sinfo, GetStructInfo(expr)) != tvm::relax::BaseCheckResult::kFailL0)
+ << "Invalid annotation. Got rhs value struct info: " << GetStructInfo(expr)
+ << ", given struct info: " << sinfo;
+ }
+ }
+ tvm::relax::Var var = block_builder->Emit(expr);
+ block_frame->emitted_vars.push_back(var);
+ return var;
+}
+
+tvm::relax::Var EmitMatchCast(const tvm::relax::Expr& value,
+ const tvm::relax::StructInfo& struct_info) {
+ BlockFrame block_frame = CheckBlockFrameExistAndUnended();
+ const tvm::relax::BlockBuilder& block_builder = GetBlockBuilder();
+
+ tvm::relax::Var var = block_builder->EmitMatchCast(value, struct_info);
+ block_frame->emitted_vars.push_back(var);
+ return var;
+}
+
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.Emit").set_body_typed(Emit);
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.EmitMatchCast").set_body_typed(EmitMatchCast);
+
+///////////////////////////// If Then Else /////////////////////////////
+
+IfFrame If(tvm::relax::Expr condition) {
+ ObjectPtr<IfFrameNode> n = make_object<IfFrameNode>();
+ n->condition = condition;
+ n->then_expr = NullOpt;
+ n->else_expr = NullOpt;
+ return IfFrame(n);
+}
+
+ThenFrame Then() {
+ ObjectPtr<ThenFrameNode> n = make_object<ThenFrameNode>();
+ return ThenFrame(n);
+}
+
+ElseFrame Else() {
+ ObjectPtr<ElseFrameNode> n = make_object<ElseFrameNode>();
+ return ElseFrame(n);
+}
+
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.If").set_body_typed(If);
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.Then").set_body_typed(Then);
+TVM_REGISTER_GLOBAL("script.ir_builder.relax.Else").set_body_typed(Else);
+
+} // namespace relax
+} // namespace ir_builder
+} // namespace script
+} // namespace tvm
diff --git a/src/script/ir_builder/relax/utils.h b/src/script/ir_builder/relax/utils.h
new file mode 100644
index 0000000000..ae91d05769
--- /dev/null
+++ b/src/script/ir_builder/relax/utils.h
@@ -0,0 +1,119 @@
+/*
+ * 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.
+ */
+#ifndef TVM_SCRIPT_IR_BUILDER_RELAX_UTILS_H_
+#define TVM_SCRIPT_IR_BUILDER_RELAX_UTILS_H_
+
+#include <tvm/relax/struct_info_functor.h>
+#include <tvm/script/ir_builder/relax/frame.h>
+#include <tvm/script/ir_builder/relax/ir.h>
+
+#include <string>
+
+namespace tvm {
+namespace script {
+namespace ir_builder {
+namespace relax {
+
+inline FunctionFrame FindFunctionFrame(const String& method) {
+ if (Optional<FunctionFrame> frame = IRBuilder::Current()->FindFrame<FunctionFrame>()) {
+ return frame.value();
+ }
+ LOG(FATAL) << "ValueError: Function frame not find. Please ensure '" << method
+ << "' is called under R.function()";
+ throw;
+}
+
+inline IfFrame FindIfFrame(const String& method) {
+ if (Optional<IfFrame> frame = IRBuilder::Current()->GetLastFrame<IfFrame>()) {
+ return frame.value();
+ } else {
+ LOG(FATAL) << "ValueError: IfThenElse frame not find. Please ensure '" << method
+ << "' is called under R.if_()";
+ }
+ throw;
+}
+
+inline tvm::relax::BlockBuilder GetBlockBuilder() {
+ Optional<FunctionFrame> frame = IRBuilder::Current()->FindFrame<FunctionFrame>();
+ CHECK(frame.defined()) << "ValueError: Relax Function frame not find. Please ensure "
+ "assignment is called under R.function()";
+ return frame.value()->block_builder;
+}
+
+inline BlockFrame CheckBlockFrameExistAndUnended() {
+ // We check if the current block is "ended" - if a block is ended, it is not allowed to emit new
+ // bindings into this block, and we should throw exceptions.
+
+ Optional<BlockFrame> block_frame = IRBuilder::Current()->GetLastFrame<BlockFrame>();
+ CHECK(block_frame.defined()) << "ValueError: Block frame not find";
+ CHECK(!block_frame.value()->block_ended)
+ << "ValueError: New binding is not allowed after dataflow block output.";
+ return block_frame.value();
+}
+
+inline tvm::relax::SeqExpr GetSeqExprForBranch(const SeqExprFrame& frame, String* var_name) {
+ // Step 0. Check frame type
+ std::string method;
+ if (frame->IsInstance<ThenFrameNode>()) {
+ method = "R.Then";
+ } else if (frame->IsInstance<ElseFrameNode>()) {
+ method = "R.Else";
+ } else {
+ ICHECK(false) << "TypeError: Unsupported frame type: " << frame->GetTypeKey();
+ }
+
+ // Step 1. Check non-empty block and last binding is non-dataflow
+ CHECK(!frame->binding_blocks.empty())
+ << "Empty body is not allowed for '" << method << "' statements.";
+ const tvm::relax::BindingBlock& last_block = frame->binding_blocks.back();
+ CHECK(!last_block->bindings.empty()) << "Blocks are expected to be non-empty.";
+
+ // Step 2. Collect body from the last binding.
+ tvm::relax::Expr body;
+ const tvm::relax::Binding& last_binding = last_block->bindings.back();
+ if (const auto* var_binding = last_binding.as<tvm::relax::VarBindingNode>()) {
+ CHECK(!var_binding->var->IsInstance<tvm::relax::DataflowVarNode>())
+ << "A non-dataflow var is expected in the last binding of '" << method << "'.";
+ body = var_binding->value;
+ *var_name = var_binding->var->name_hint();
+ } else if (const auto* match_cast = last_binding.as<tvm::relax::MatchCastNode>()) {
+ CHECK(!match_cast->var->IsInstance<tvm::relax::DataflowVarNode>())
+ << "A non-dataflow var is expected in the last binding of '" << method << "'.";
+ body = var_binding->value;
+ *var_name = match_cast->var->name_hint();
+ } else {
+ ICHECK(false) << "TypeError: Unsupported binding type: " << last_binding->GetTypeKey();
+ }
+
+ // Step 3. Re-collect binding blocks to remove the last binding.
+ Array<tvm::relax::BindingBlock> new_blocks(frame->binding_blocks.begin(),
+ frame->binding_blocks.end() - 1);
+ Array<tvm::relax::Binding> last_block_bindings(last_block->bindings.begin(),
+ last_block->bindings.end() - 1);
+ new_blocks.push_back(tvm::relax::BindingBlock(last_block_bindings));
+
+ return tvm::relax::SeqExpr(new_blocks, body);
+}
+
+} // namespace relax
+} // namespace ir_builder
+} // namespace script
+} // namespace tvm
+
+#endif // TVM_SCRIPT_IR_BUILDER_RELAX_UTILS_H_
diff --git a/src/script/ir_builder/tir/frame.cc b/src/script/ir_builder/tir/frame.cc
index 1e63201a40..dd8d3c2ed3 100644
--- a/src/script/ir_builder/tir/frame.cc
+++ b/src/script/ir_builder/tir/frame.cc
@@ -16,6 +16,7 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/script/ir_builder/ir/ir.h>
#include <tvm/script/ir_builder/tir/frame.h>
#include <tvm/tir/function.h>
@@ -41,9 +42,17 @@ void PrimFuncFrameNode::ExitWithScope() {
ICHECK(!builder->result.defined()) << "ValueError: Builder.result has already been set";
builder->result = func;
} else if (Optional<ir::IRModuleFrame> opt_frame = builder->FindFrame<ir::IRModuleFrame>()) {
- ir::IRModuleFrame frame = opt_frame.value();
- frame->global_vars.push_back(GlobalVar(name.value_or("")));
- frame->functions.push_back(func);
+ CHECK(name.defined()) << "ValueError: The function name must be defined before exiting the "
+ "function scope, if it's defined in a Module";
+ const ir::IRModuleFrame& frame = opt_frame.value();
+ const String& func_name = name.value_or("");
+ if (!frame->global_var_map.count(func_name)) {
+ // Case. First time visiting the function.
+ ir::DeclFunction(func_name, func);
+ }
+ // Define the function.
+ // Note we do checks to disallow redefinition of functions inside the `DefFunction`.
+ ir::DefFunction(func_name, func);
} else {
LOG(FATAL) << "ValueError: Cannot find where to insert PrimFunc";
}
diff --git a/src/script/ir_builder/tir/utils.h b/src/script/ir_builder/tir/utils.h
index 733c975fad..80a92f923c 100644
--- a/src/script/ir_builder/tir/utils.h
+++ b/src/script/ir_builder/tir/utils.h
@@ -80,7 +80,7 @@ inline PrimFuncFrame FindPrimFuncFrame(const String& method) {
* \return The top frame of BlockFrame.
*/
inline BlockFrame FindBlockFrame(const String& method) {
- if (Optional<BlockFrame> frame = IRBuilder::Current()->GetLastFrame<BlockFrame>()) {
+ if (Optional<BlockFrame> frame = IRBuilder::Current()->FindFrame<BlockFrame>()) {
return frame.value();
}
LOG(FATAL) << "ValueError: Block frame not find. Please ensure '" << method
diff --git a/tests/python/relax/test_tvmscript_ir_builder.py b/tests/python/relax/test_tvmscript_ir_builder.py
new file mode 100644
index 0000000000..12d8b114b8
--- /dev/null
+++ b/tests/python/relax/test_tvmscript_ir_builder.py
@@ -0,0 +1,153 @@
+# 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.
+import tvm
+import tvm.testing
+from tvm import relax, tir
+from tvm.script.ir_builder import relax as R
+from tvm.script.ir_builder.base import IRBuilder
+
+
+def test_function_simple():
+ """
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
+ out = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
+ return out
+ """
+ # create with Script IRBuilder
+ with IRBuilder() as ir_builder:
+ with R.function():
+ R.func_name("foo")
+ R.func_attr({"Primitive": 1})
+ x = R.arg("x", relax.TensorStructInfo((128, 128), "float32"))
+ R.func_ret_struct_info(relax.TensorStructInfo(dtype="float32", ndim=2))
+ out = R.emit(
+ R.call_tir("extern_func", x, relax.TensorStructInfo((128, 128), dtype="float32"))
+ )
+ IRBuilder.name("out", out)
+ R.func_ret_value(out)
+ func = ir_builder.get()
+ # create with BlockBuilder
+ x = relax.Var("x", relax.TensorStructInfo((128, 128), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,), attrs={"Primitive": 1}):
+ out = bb.emit(
+ relax.call_tir("extern_func", x, relax.TensorStructInfo((128, 128), dtype="float32"))
+ )
+ bb.emit_func_output(out)
+ mod = bb.get()
+
+ tvm.ir.assert_structural_equal(func, mod["foo"])
+ # check names
+ assert func.params[0].name_hint == "x"
+ assert func.body.body.name_hint == "out"
+
+
+def test_match_cast():
+ """
+ @R.function
+ def foo(x: R.Tensor(None, "float32"), y: R.Tensor(None, "float32")):
+ m = T.var("int64")
+ n = T.var("int64")
+ _ = R.match_cast(x, R.Tensor((m,), "float32"))
+ y1 = R.match_cast(x, R.Tensor((n,), "float32"))
+ return (m, n * 2)
+ """
+ # create with Script IRBuilder
+ with IRBuilder() as ir_builder:
+ with R.function():
+ R.func_name("foo")
+ x = R.arg("x", relax.TensorStructInfo(ndim=-1, dtype="float32"))
+ y = R.arg("y", relax.TensorStructInfo(ndim=-1, dtype="float32"))
+ m = tir.Var("m", dtype="int64")
+ n = tir.Var("n", dtype="int64")
+ _ = R.emit_match_cast(x, relax.TensorStructInfo((m,), "float32"))
+ y1 = R.emit_match_cast(y, relax.TensorStructInfo((n,), "float32"))
+ IRBuilder.name("y1", y1)
+ R.func_ret_value(relax.ShapeExpr([m, n * 2]))
+ func = ir_builder.get()
+
+ # create with BlockBuilder
+ x = relax.Var("x", relax.TensorStructInfo(dtype="float32", ndim=-1))
+ y = relax.Var("y", relax.TensorStructInfo(dtype="float32", ndim=-1))
+ m = tir.Var("m", dtype="int64")
+ n = tir.Var("n", dtype="int64")
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x, y)):
+ _ = bb.match_cast(x, relax.TensorStructInfo((m,), "float32"))
+ y1 = bb.match_cast(y, relax.TensorStructInfo((n,), "float32"))
+ bb.emit_func_output(relax.ShapeExpr([m, n * 2]))
+ mod = bb.get()
+
+ tvm.ir.assert_structural_equal(func, mod["foo"])
+
+
+def test_dataflow_block():
+ """
+ @R.function
+ def foo(x: Tensor((128, 128), "float32")) -> Tensor(None, "float32", ndim = 2):
+ # block 0
+ with R.dataflow():
+ lv0 = R.call_tir("extern_func", (x,), R.Tensor((128, 128), dtype="float32"))
+ gv: Tensor((128, 128), "float32") = lv0
+ R.output(gv)
+ return gv
+ """
+ # create with Script IRBuilder
+ with IRBuilder() as ir_builder:
+ with R.function():
+ R.func_name("foo")
+ x = R.arg("x", relax.TensorStructInfo((128, 128), "float32"))
+ with R.dataflow() as df:
+ lv0 = R.emit(
+ R.call_tir(
+ "extern_func", x, relax.TensorStructInfo((128, 128), dtype="float32")
+ )
+ )
+ IRBuilder.name("lv0", lv0)
+ gv = R.emit(lv0)
+ IRBuilder.name("gv", gv)
+ R.output(gv)
+ (gv,) = df.output_vars
+ R.func_ret_value(gv)
+ func = ir_builder.get()
+
+ # create with BlockBuilder
+ x = relax.Var("x", relax.TensorStructInfo((128, 128), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ with bb.dataflow():
+ lv0 = bb.emit(
+ relax.call_tir(
+ "extern_func", x, relax.TensorStructInfo((128, 128), dtype="float32")
+ )
+ )
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+
+ tvm.ir.assert_structural_equal(func, bb.get()["foo"])
+
+
+def test_regression_py_print():
+ # Test that the py_print directs to python builtin print
+ from tvm.script.ir_builder.relax.ir import py_print # pylint: disable=import-outside-toplevel
+
+ assert py_print == print
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_tvmscript_parser.py b/tests/python/relax/test_tvmscript_parser.py
new file mode 100644
index 0000000000..34b02fdbb8
--- /dev/null
+++ b/tests/python/relax/test_tvmscript_parser.py
@@ -0,0 +1,1062 @@
+# 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.
+
+from typing import Optional, Union
+
+import pytest
+import tvm
+import tvm.script
+import tvm.testing
+from tvm import IRModule, relax, tir, topi
+from tvm.relax import DynTensorType
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tir as T
+
+
+def _check(
+ parsed: Union[relax.Function, IRModule],
+ expect: Optional[Union[relax.Function, IRModule]] = None,
+):
+ # TODO(relax-team): enable roundtrip testing when printer is ready
+ # test = parsed.script(show_meta=True)
+ # roundtrip_mod = tvm.script.parse(test)
+ # tvm.ir.assert_structural_equal(parsed, roundtrip_mod)
+ if expect:
+ tvm.ir.assert_structural_equal(parsed, expect)
+
+
+def test_simple_func():
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
+ R.func_attr({"Primitive": 1})
+ gv0 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
+ return gv0
+
+ x = relax.Var("x", R.Tensor((128, 128), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,), attrs={"Primitive": 1}):
+ out = bb.emit(relax.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32")))
+ bb.emit_func_output(out)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_error_report():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
+ # error: a = b = c is not allowed.
+ gv0 = gv1 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
+ return gv0
+
+
+def test_mismatch_cast_dims_and_ndim():
+ with pytest.raises(Exception):
+
+ @R.function
+ def f(
+ x: R.Tensor((2, 3), "float32", ndim=3)
+ ): # error: ndim and the shape dims are mismatch
+ return x
+
+
+def test_unexpected_num_kw_args():
+ with pytest.raises(Exception):
+
+ @R.function
+ def f(x: R.Tensor(dtype="float32", ndim=1, foo=2)): # error: unexpected kw args foo
+ return x
+
+
+def test_unexpected_ndim():
+ with pytest.raises(Exception):
+
+ @R.function
+ # error: dim is expected to be non-negative int or -1 for unknown
+ def f(x: R.Tensor(dtype="float32", ndim=-2)):
+ return x
+
+
+def test_unexpected_ndim_type():
+ with pytest.raises(Exception):
+
+ @R.function
+ def f(x: R.Tensor(dtype="float32", ndim="1")): # error: dim is expected to be int
+ return x
+
+
+def test_unexpected_tir_cast_args():
+
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def f(x: R.Tensor(("m",), "float32")):
+ m = T.var("int64")
+ # tir.cast expects 2 arguments, but got 3
+ return R.call_tir("foo", (x,), R.Tensor((T.cast("int32", m, 1),), dtype="float32"))
+
+
+def test_unexpected_tir_max_args():
+
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def f(x: R.Tensor(("m", "n"), "float32")):
+ m = T.var("int64")
+ # tir.max expects 2 arguments, but got 1
+ return relax.call_tir("foo", (x,), R.Tensor((T.max(m),), dtype="float32"))
+
+
+def test_func_type_annotation_fail():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def f(x, y): # error: the parameter type annotation is missing
+ z = R.add(x, y)
+ y = z
+ return y
+
+
+def test_if_mismatch_var_fail():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def f(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")):
+ if cond:
+ w = R.add(x, x)
+ y = R.multiply(w, w)
+ else:
+ w = R.multiply(x, x)
+ z = R.add(w, w) # error: The binding var is expected to `y`
+ return z
+
+
+def test_unassigned_call_fail():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def f(x: R.Tensor):
+ R.add(x, x)
+ return x
+
+
+def test_simple_module():
+ @I.ir_module
+ class TestModule:
+ @T.prim_func
+ def tir_func(
+ x: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+ y: T.Buffer((T.int64(128), T.int64(128)), "float32"),
+ ):
+ T.func_attr({"tir.noalias": True})
+ for i, j in T.grid(T.int64(128), T.int64(128)):
+ with T.block():
+ vi, vj = T.axis.remap("SS", [i, j])
+ y[vi, vj] = x[vi, vj] + 1.0
+
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
+ # TODO(Siyuan): Need to change to `TestModule.tir_func`
+ gv0 = R.call_tir(tir_func, x, R.Tensor((128, 128), dtype="float32"))
+ return gv0
+
+ x = relax.Var("x", R.Tensor((128, 128), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ out = bb.emit_te(lambda x: x + 1, x, primfunc_name_hint="tir_func")
+ bb.emit_func_output(out)
+
+ _check(TestModule, bb.get())
+
+
+def test_relax_tensor_op():
+ @R.function
+ def foo(x: R.Tensor((4, 4), "float32")) -> R.Tensor((4, 4), "float32"):
+ y = R.add(x, x)
+ z = R.multiply(x, y)
+ return z
+
+ x = relax.Var("x", R.Tensor((4, 4), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ y = bb.emit(relax.op.add(x, x))
+ z = bb.emit(relax.op.multiply(x, y))
+ bb.emit_func_output(z)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_symbolic_shape():
+ @R.function
+ def foo(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64", "m")
+ n = T.var("int64", "n")
+ gv0 = R.call_tir("extern_func", x, R.Tensor((m, n), dtype="float32"))
+ return gv0
+
+ @R.function
+ def bar(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv0 = R.call_tir("extern_func", x, R.Tensor((m, n), dtype="float32"))
+ return gv0
+
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def mismatch_dtype(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(None, "float32", ndim=2):
+ m = T.var("int64")
+ n = T.var("int32") # The shape dtype should be int64
+ gv0 = R.call_tir("extern_func", x, R.Tensor((m, n), dtype="float32"))
+ return gv0
+
+ def _expected(name: str):
+ n, m = tir.Var("n", "int64"), tir.Var("m", "int64")
+ x = relax.Var("x", R.Tensor([m, n], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function(name, (x,)):
+ out = bb.emit(relax.call_tir("extern_func", x, R.Tensor((m, n), dtype="float32")))
+ bb.emit_func_output(out)
+ return bb.get()[name]
+
+ _check(foo, _expected("foo"))
+ _check(bar, _expected("bar"))
+
+
+def test_shadowing():
+ @R.function
+ def foo(x: R.Tensor((4, 4), "float32")):
+ y = R.add(x, x)
+ z = R.multiply(x, y)
+ y = R.add(x, y)
+ y = z
+ y = R.multiply(y, x)
+ z = y
+ return z
+
+ x = relax.Var("x", R.Tensor((4, 4), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ y = bb.emit(relax.op.add(x, x))
+ z = bb.emit(relax.op.multiply(x, y))
+ y = bb.emit(relax.op.add(x, y))
+ y = bb.emit(z)
+ y = bb.emit(relax.op.multiply(y, x))
+ z = bb.emit(y)
+ bb.emit_func_output(z)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_match_cast():
+ @R.function
+ def foo(x: R.Tensor("float32"), y: R.Tensor("float32")):
+ m = T.var("int64")
+ n = T.var("int64")
+ x0 = R.match_cast(x, R.Tensor([m], "float32"))
+ with R.dataflow():
+ y0 = R.match_cast(y, R.Tensor([n], "float32"))
+ gv = y0
+ R.output(gv)
+ return (x0, (m, n * 2))
+
+ x = relax.Var("x", R.Tensor("float32"))
+ y = relax.Var("y", R.Tensor("float32"))
+ m = tir.Var("m", dtype="int64")
+ n = tir.Var("n", dtype="int64")
+ y2 = relax.Var("y", R.Tensor([n], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x, y)):
+ x0 = bb.match_cast(x, R.Tensor([m], "float32"))
+ with bb.dataflow():
+ y0 = bb.match_cast(y, R.Tensor([n], "float32"))
+ bb.emit_output(y0)
+ bb.emit_func_output(relax.Tuple([x0, relax.ShapeExpr([m, n * 2])]))
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_tuple_return():
+ @R.function
+ def foo(x: R.Tensor((4, 4), "float32")):
+ gv0 = R.call_tir("extern_func_0", x, R.Tensor((4, 4), dtype="float32"))
+ gv1 = R.call_tir("extern_func_1", x, R.Tensor((4, 4), dtype="float32"))
+ return (gv0, gv1)
+
+ x = relax.Var("x", R.Tensor((4, 4), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ gv0 = bb.emit(relax.call_tir("extern_func_0", x, R.Tensor((4, 4), dtype="float32")))
+ gv1 = bb.emit(relax.call_tir("extern_func_1", x, R.Tensor((4, 4), dtype="float32")))
+ bb.emit_func_output(relax.Tuple((gv0, gv1)))
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_tuple_return_2():
+ @R.function
+ def foo(x: R.Tensor("float32", ndim=2)):
+ n, m = T.var("int64"), T.var("int64")
+ x0 = R.match_cast(x, R.Tensor((n, m), "float32"))
+ return (x0, (n + 1, m, 1))
+
+ x = relax.Var("x", R.Tensor("float32", ndim=2))
+ n, m = tir.Var("n", "int64"), tir.Var("m", "int64")
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ x0 = bb.match_cast(x, R.Tensor((n, m), "float32"))
+ bb.emit_func_output(relax.Tuple([x0, relax.ShapeExpr([n + 1, m, 1])]))
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_tuple_binding():
+ @R.function
+ def foo(x: R.Tensor("float32", ndim=2)):
+ n, m = T.var("int64"), T.var("int64")
+ x0 = R.match_cast(x, R.Tensor((n, m), "float32"))
+ t0 = (x, x0)
+ t1 = (x, (n, m), t0)
+ return t1
+
+ x = relax.Var("x", R.Tensor("float32", ndim=2))
+ n, m = tir.Var("n", "int64"), tir.Var("m", "int64")
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ x0 = bb.match_cast(x, R.Tensor((n, m), "float32"))
+ t0 = bb.emit(relax.Tuple([x, x0]))
+ t1 = bb.emit(relax.Tuple([x, relax.ShapeExpr([n, m]), t0]))
+ bb.emit_func_output(t1)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_tuple_get_item():
+ @R.function
+ def foo(x: R.Tensor, y: R.Tensor):
+ t1 = R.tuple(x, y)
+ t2 = (x, y)
+ a = t1[0]
+ b = R.TupleGetItem(t2, 1)
+ c = R.add(a, b)
+ return c
+
+ x = relax.Var("x", R.Tensor())
+ y = relax.Var("y", R.Tensor())
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x, y)):
+ t1 = bb.emit(relax.Tuple([x, y]))
+ t2 = bb.emit(relax.Tuple([x, y]))
+ a = bb.emit(relax.TupleGetItem(t1, 0))
+ b = bb.emit(relax.TupleGetItem(t2, 1))
+ c = bb.emit(relax.op.add(a, b))
+ bb.emit_func_output(c)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_dataflow_block():
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
+ with R.dataflow():
+ lv0 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
+ lv1 = R.call_tir("extern_func", lv0, R.Tensor((128, 128), dtype="float32"))
+ gv = lv1
+ R.output(gv)
+ return gv
+
+ x = relax.Var("x", R.Tensor((128, 128), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32")))
+ lv1 = bb.emit(relax.call_tir("extern_func", lv0, R.Tensor((128, 128), dtype="float32")))
+ gv = bb.emit_output(lv1)
+ bb.emit_func_output(gv)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_dataflow_block_advanced():
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
+ gv0 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
+ gv1 = R.call_tir("extern_func", gv0, R.Tensor((128, 128), dtype="float32"))
+ with R.dataflow():
+ m = T.var("int64")
+ n = T.var("int64")
+ lv0 = R.call_tir("extern_func", gv1, R.Tensor((128, 128), dtype="float32"))
+ lv1 = R.match_cast(lv0, R.Tensor((m, n), "float32"))
+ gv2 = R.call_tir("extern_func", lv0, R.Tensor((128, 128), dtype="float32"))
+ gv2 = R.call_tir("extern_func", gv2, R.Tensor((128, 128), dtype="float32"))
+ gv3 = R.match_cast(gv2, R.Tensor((m, n), "float32"))
+ gv3 = R.match_cast(lv0, R.Tensor((m, n), "float32"))
+ gv4 = gv3
+ gv5 = gv2
+ R.output(gv5, gv4)
+ gv6 = R.call_tir("extern_func", gv5, R.Tensor((128, 128), dtype="float32"))
+ gv7 = R.call_tir("extern_func", gv6, R.Tensor((128, 128), dtype="float32"))
+ return gv7
+
+ x = relax.Var("x", R.Tensor((128, 128), "float32"))
+ bb = relax.BlockBuilder()
+ m = tir.Var("m", dtype="int64")
+ n = tir.Var("n", dtype="int64")
+ with bb.function("foo", (x,)):
+ gv0 = bb.emit(relax.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32")))
+ gv1 = bb.emit(relax.call_tir("extern_func", gv0, R.Tensor((128, 128), dtype="float32")))
+ with bb.dataflow():
+ lv0 = bb.emit(relax.call_tir("extern_func", gv1, R.Tensor((128, 128), dtype="float32")))
+ lv1 = bb.match_cast(lv0, R.Tensor((m, n), "float32"))
+ gv2 = bb.emit(relax.call_tir("extern_func", lv0, R.Tensor((128, 128), dtype="float32")))
+ gv21 = bb.emit(
+ relax.call_tir("extern_func", gv2, R.Tensor((128, 128), dtype="float32"))
+ )
+ gv3 = bb.match_cast(gv21, R.Tensor((m, n), "float32"))
+ gv31 = bb.match_cast(lv0, R.Tensor((m, n), "float32"))
+ gv32 = bb.emit_output(gv31)
+ gv22 = bb.emit_output(gv21)
+ gv4 = bb.emit(relax.call_tir("extern_func", gv22, R.Tensor((128, 128), dtype="float32")))
+ gv5 = bb.emit(relax.call_tir("extern_func", gv4, R.Tensor((128, 128), dtype="float32")))
+ bb.emit_func_output(gv5)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_dataflow_binding_after_output():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
+ with R.dataflow():
+ gv = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
+ R.output(gv)
+ lv = R.call_tir("extern_func", gv, R.Tensor((128, 128), dtype="float32"))
+ return gv
+
+
+def test_dataflow_output_global_var():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
+ gv0 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
+ with R.dataflow():
+ gv1 = R.call_tir("extern_func", gv0, R.Tensor((128, 128), dtype="float32"))
+ R.output(gv0, gv1)
+ return gv1
+
+
+def test_dataflow_multiple_output():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
+ with R.dataflow():
+ gv = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
+ R.output(gv)
+ R.output(gv)
+ return gv
+
+
+def test_dataflow_output_outside_dataflow_block():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
+ gv = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
+ R.output(gv)
+ return gv
+
+
+def test_dataflow_scope_fail():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def f(x: R.Tensor(ndim=2)):
+ with R.dataflow():
+ y = R.add(x, x)
+ z = R.multiply(y, x)
+ w = R.add(z, x)
+ R.output(y, w)
+ t = R.multiply(y, z) # z is not in the outer scope
+ return t
+
+
+def test_return_without_binding():
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")):
+ return x
+
+ x = relax.Var("x", R.Tensor((128, 128), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ bb.emit_func_output(x)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_multiple_return():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")):
+ return x
+ return x
+
+
+def test_function_without_return():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def foo(x: R.Tensor((128, 128), "float32")):
+ gv0 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
+
+
+def test_tensor_type_without_args():
+ @R.function
+ def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
+ v = R.call_tir("tir_relu", x, R.Tensor((32, 32), dtype="float32"))
+ return v
+
+ x = relax.Var("x", R.Tensor((32, 32), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x)):
+ v = bb.emit(relax.call_tir("tir_relu", x, R.Tensor((32, 32), dtype="float32")))
+ bb.emit_func_output(v)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_direct_return():
+ @R.function
+ def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor((32, 32), "float32"):
+ return x
+
+ x = relax.Var("x", R.Tensor((32, 32), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x)):
+ bb.emit_func_output(x)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_call_packed():
+ @R.function
+ def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
+ z = R.call_packed("vm.builtin.copy", x, sinfo_args=R.Tensor((32, 32), "float32"))
+ return z
+
+ x = relax.Var("x", R.Tensor((32, 32), "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x)):
+ z = bb.emit(
+ relax.Call(
+ relax.ExternFunc("vm.builtin.copy"),
+ (x,),
+ None,
+ sinfo_args=[R.Tensor((32, 32), "float32")],
+ )
+ )
+ bb.emit_func_output(z)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_annotation():
+ @R.function
+ def foo(
+ x: R.Tensor((32, "m"), "float32"),
+ y: R.Tensor(("m",), "float32"),
+ r: R.Tensor(dtype="int64"),
+ ) -> R.Object:
+ m = T.var("int64", "m")
+ z: R.Tensor((32, m), "float32") = R.multiply(x, y)
+ w: R.Tensor = R.multiply(z, z)
+ q: R.Tensor(ndim=2) = R.add(w, w)
+ t = R.add(w, z)
+ sh: R.Shape = R.call_packed("shape_of", x, sinfo_args=R.Shape)
+ o: R.Object = R.call_packed("contrib.tensor_array_stack", x, y, sinfo_args=R.Object)
+ return o
+
+ def _check_struct_info(binding, expected_sinfo):
+ tvm.ir.assert_structural_equal(binding.var.struct_info, expected_sinfo)
+ tvm.ir.assert_structural_equal(binding.value.struct_info, expected_sinfo)
+
+ # Cannot use block builder here because we need to check the annotated type,
+ # which may be inconsistent with deduced type.
+ assert isinstance(foo.ret_struct_info, relax.ObjectStructInfo)
+ m = relax.get_shape_of(foo.params[0])[1]
+ bindings = foo.body.blocks[0].bindings
+
+ _check_struct_info(bindings[0], relax.TensorStructInfo([32, m], "float32"))
+ _check_struct_info(bindings[1], relax.TensorStructInfo(dtype="", ndim=-1))
+ _check_struct_info(bindings[2], relax.TensorStructInfo(dtype="", ndim=2))
+ _check_struct_info(bindings[3], relax.TensorStructInfo(dtype="", ndim=-1))
+ _check_struct_info(bindings[4], relax.ShapeStructInfo(ndim=-1))
+ _check_struct_info(bindings[5], relax.ObjectStructInfo())
+
+
+def test_annotate_override():
+ @R.function
+ def foo(x: R.Tensor):
+ y = x
+ # z will be treated as object type even though it's a tensor
+ z: R.Object = R.add(x, y)
+ return z
+
+ assert isinstance(foo.ret_struct_info, relax.ObjectStructInfo)
+ y_bind, z_bind = foo.body.blocks[0].bindings
+ assert isinstance(y_bind.var.struct_info, relax.TensorStructInfo)
+ assert isinstance(z_bind.var.struct_info, relax.ObjectStructInfo)
+
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def test(x: R.Tensor):
+ # Error: x is of Tensor StructInfo, which can not annotate to R.Shape.
+ z: R.Shape = x
+ return z
+
+ @R.function
+ def bar(x: R.Tensor):
+ # x is of Tensor StructInfo, the annotation of `z` is ignored.
+ z: R.Object = x
+ return z
+
+ assert isinstance(bar.ret_struct_info, relax.TensorStructInfo)
+ (z_bind,) = bar.body.blocks[0].bindings
+ assert isinstance(z_bind.var.struct_info, relax.TensorStructInfo)
+
+
+def test_call_tir_empty_shape():
+ @R.function
+ def foo(x: R.Tensor((), "float32")):
+ z = R.call_tir("scalar_add", x, R.Tensor((), dtype="float32"))
+ return z
+
+ (z_bind,) = foo.body.blocks[0].bindings
+ shape_expr = z_bind.value.sinfo_args[0].shape
+
+ assert isinstance(shape_expr, relax.ShapeExpr)
+ assert len(shape_expr.values) == 0
+
+
+def test_call_tir_empty_tuple_arg():
+ bb = relax.BlockBuilder()
+ dummy_param = relax.Var("dummy_param", R.Tensor(()))
+ with bb.function("foo", [dummy_param]):
+ output = bb.emit_te(topi.full, shape=(16, 32), dtype="float32", fill_value=1.0)
+ bb.emit_func_output(output)
+
+ _check(bb.get())
+
+
+def test_call_tir_with_tir_var():
+ @I.ir_module
+ class Module:
+ @R.function
+ def main(
+ dumb_param: R.Tensor(("n",), "float32"), x: R.Tensor(("n * 2", "float32"))
+ ) -> R.Tensor(("n * 2",), "float32"):
+ n = T.var("int64")
+ y = R.call_tir(copy, (x,), R.Tensor(((n * 2,)), dtype="float32"), tir_vars=(n,))
+ return y
+
+ @T.prim_func
+ def copy(var_x: T.handle, var_y: T.handle, n: T.int64):
+ X = T.match_buffer(var_x, (n * 2,), dtype="float32")
+ Y = T.match_buffer(var_y, (n * 2,), dtype="float32")
+ for i in T.grid(n * 2):
+ with T.block("block"):
+ vi = T.axis.remap("S", [i])
+ Y[vi] = X[vi]
+
+ _check(Module)
+
+
+def test_local_function():
+ @R.function
+ def main(
+ x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")
+ ) -> R.Tensor((2, 3), "float32"):
+ @R.function
+ def outer_func(
+ c1: R.Tensor((2, 3), "float32")
+ ) -> R.Callable((R.Tensor(None, "float32", ndim=2),), R.Tensor(None, "float32", ndim=2)):
+ @R.function
+ def inner_func(x1: R.Tensor((2, 3), "float32")):
+ s: R.Tensor((2, 3), "float32") = R.add(x1, c1)
+ return s
+
+ return inner_func
+
+ in_call = outer_func(x)
+ res = in_call(y)
+ return res
+
+ main_bindings = main.body.blocks[0].bindings
+ assert len(main_bindings) == 3
+ outer_func = main_bindings[0].value
+ assert isinstance(outer_func, relax.Function)
+
+ outer_func_bindings = outer_func.body.blocks[0].bindings
+ assert len(outer_func_bindings) == 1
+ inner_func = outer_func_bindings[0].value
+ assert isinstance(inner_func, relax.Function)
+
+ @I.ir_module
+ class TestModule:
+ @R.function
+ def f(x: R.Tensor((128, 128), "float32"), y: R.Tensor((128, 128), "float32")):
+ @T.prim_func
+ def my_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
+ A = T.match_buffer(a, (128, 128))
+ B = T.match_buffer(b, (128, 128))
+ C = T.match_buffer(c, (128, 128))
+
+ for i, j, k in T.grid(128, 128, 128):
+ with T.block():
+ vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+ with T.init():
+ C[vi, vj] = 0.0
+ C[vi, vj] += A[vi, vk] * B[vj, vk]
+
+ z = relax.call_tir(my_matmul, (x, y), R.Tensor((128, 128), dtype="float32"))
+ return z
+
+ bindings = TestModule["f"].body.blocks[0].bindings
+ assert len(bindings) == 2
+ tir_func = bindings[0].value
+ assert isinstance(tir_func, tir.PrimFunc)
+
+
+def test_cross_function_call():
+ @I.ir_module
+ class Mod0:
+ @R.function
+ def foo(x: R.Tensor((10, 5), "float32")):
+ s = R.add(x, x)
+ return s
+
+ @R.function
+ def main(x: R.Tensor((10, 5), "float32")):
+ inner = foo
+ gv1 = inner(x)
+ gv2 = foo(x)
+ return (inner, gv1, gv2)
+
+ @I.ir_module
+ class Mod1:
+ @R.function
+ def main(x: R.Tensor((10, 5), "float32")):
+ inner = foo
+ gv1 = inner(x)
+ gv2 = foo(x)
+ return (inner, gv1, gv2)
+
+ @R.function
+ def foo(x: R.Tensor((10, 5), "float32")) -> R.Tensor((10, 5), "float32"):
+ s = R.add(x, x)
+ return s
+
+
+def test_if_branch():
+ @R.function
+ def foo(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")) -> R.Tensor((1,), "float32"):
+ if cond:
+ w = R.add(x, x)
+ y = R.multiply(w, w)
+ else:
+ w = R.multiply(x, x)
+ y = R.add(w, w)
+ return y
+
+ cond, x = foo.params
+ y_bind = foo.body.blocks[0].bindings[0]
+ y, ite = y_bind.var, y_bind.value
+
+ assert isinstance(y, relax.Var)
+ assert y.name_hint == "y"
+
+ assert isinstance(ite, relax.If)
+ assert isinstance(ite.true_branch, relax.SeqExpr)
+ assert isinstance(ite.false_branch, relax.SeqExpr)
+
+ def check_call(call, op, args):
+ assert isinstance(call, relax.Call)
+ if isinstance(op, str):
+ assert call.op.name == op
+ else:
+ assert call.op == op
+ tvm.ir.assert_structural_equal(call.args, args)
+
+ w_bind = ite.true_branch.blocks[0].bindings[0]
+ # the seq exprts in the branches are normalized to bind any call
+ # in the seq expr "body" to a var
+ y_bind = ite.true_branch.blocks[-1].bindings[-1]
+ assert w_bind.var.name_hint == "w"
+ check_call(w_bind.value, "relax.add", [x, x])
+ check_call(y_bind.value, "relax.multiply", [w_bind.var, w_bind.var])
+
+ w_bind = ite.false_branch.blocks[0].bindings[0]
+ y_bind = ite.false_branch.blocks[-1].bindings[-1]
+ assert w_bind.var.name_hint == "w"
+ check_call(w_bind.value, "relax.multiply", [x, x])
+ check_call(y_bind.value, "relax.add", [w_bind.var, w_bind.var])
+
+
+def test_if_inside_dataflow():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def foo(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")):
+ with R.dataflow():
+ if cond:
+ w = R.add(x, x)
+ y = R.multiply(w, w)
+ else:
+ w = R.multiply(x, x)
+ y = R.add(w, w)
+ R.output(y)
+ return y
+
+
+def test_var_if_scoping_fail():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def f(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")):
+ if cond:
+ w = R.add(x, x)
+ y = R.multiply(w, w)
+ else:
+ w = R.multiply(x, x)
+ y = R.add(w, w)
+ return w # error: The w is not defined in the outer scope
+
+
+def test_if_branch_var_scope():
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def foo(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")):
+ if cond:
+ w = R.add(x, x)
+ y = R.multiply(w, w)
+ else:
+ w = R.multiply(x, x)
+ y = R.add(w, w)
+ return w
+
+
+def test_erase_to_well_defined():
+ @R.function
+ def foo(x: R.Tensor):
+ q = x
+ m, n = T.var("int64"), T.var("int64")
+ z = R.match_cast(q, R.Tensor((m, n)))
+ w = z
+ return w
+
+ tvm.ir.assert_structural_equal(foo.ret_struct_info, R.Tensor(ndim=2))
+ _check(foo)
+
+
+def test_empty_tuple():
+ @R.function
+ def foo(x: R.Tuple()):
+ y: R.Tuple() = R.tuple()
+ return y
+
+ x = relax.Var("x", relax.TupleStructInfo([]))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x,)):
+ y = bb.emit(relax.Tuple([]))
+ bb.emit_func_output(y)
+
+ _check(foo, bb.get()["foo"])
+
+
+def test_symbolic_shape_computing():
+ # Tensor Case 1
+ @R.function
+ def foo(x: R.Tensor(("m + 1",), "float32"), y: R.Tensor(("m", 1), "float32")):
+ z = R.add(x, y)
+ return z
+
+ m = tir.Var("m", "int64")
+ x = relax.Var("x", relax.TensorStructInfo([m + 1], "float32"))
+ y = relax.Var("y", relax.TensorStructInfo([m, 1], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("foo", (x, y)):
+ z = bb.emit(relax.op.add(x, y))
+ bb.emit_func_output(z)
+
+ _check(foo, bb.get()["foo"])
+
+ # Tensor Case 2
+ @R.function
+ def bar(
+ x: R.Tensor(("m",), "float32"), y: R.Tensor(("T.max(m, 20)",), "float32")
+ ) -> R.Tensor(("T.max(m, 20) + 1",), "float32"):
+ m = T.var("int64")
+ z = R.call_tir("test_intrin", (x, y), R.Tensor((T.max(m, 20) + 1,), dtype="float32"))
+ return z
+
+ m = tir.Var("m", "int64")
+ x = relax.Var("x", relax.TensorStructInfo([m], "float32"))
+ y = relax.Var("y", relax.TensorStructInfo([tir.max(m, 20)], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("bar", (x, y)):
+ z = bb.emit(
+ relax.call_tir("test_intrin", (x, y), R.Tensor((tir.max(m, 20) + 1,), dtype="float32"))
+ )
+ bb.emit_func_output(z)
+
+ _check(bar, bb.get()["bar"])
+
+ # Shape Case
+ @R.function
+ def baz(x: R.Shape(("m",)), y: R.Tensor(("m * 2",), "float32")):
+ m = T.var("int64")
+ z = R.call_tir("test_intrin", y, R.Tensor((m * 2,), dtype="float32"))
+ return z
+
+ m = tir.Var("m", "int64")
+ x = relax.Var("x", relax.ShapeStructInfo([m]))
+ y = relax.Var("y", relax.TensorStructInfo([m * 2], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("baz", (x, y)):
+ z = bb.emit(relax.call_tir("test_intrin", (y), R.Tensor((m * 2,), dtype="float32")))
+ bb.emit_func_output(z)
+
+ _check(baz, bb.get()["baz"])
+
+ # Error Case
+ with pytest.raises(tvm.error.DiagnosticError):
+
+ @R.function
+ def foo(x: R.Tensor(("m + 1", "m * 2"), "float32")): # name 'm' is not defined
+ z = R.add(x, x)
+ return z
+
+
+# TODO(relax-team): enable this when vm ops are ready
+@pytest.mark.xfail
+def test_vm_ops():
+ @R.function
+ def foo(x: R.Tensor(("m", "n"), dtype="float32")):
+ m = T.var("int64")
+ n = T.var("int64")
+ storage = R.vm.alloc_storage((4 * m * n,), dtype="float32", runtime_device_index=0)
+ alloc = R.vm.alloc_tensor(storage, (m, n), offset=0, dtype="float32")
+ tensor = R.builtin.alloc_tensor((m, n), dtype="float32", runtime_device_index=0)
+ _ = R.vm.call_tir_dyn("te_func", (x, tensor, (m, n)))
+ gv = tensor
+ return alloc, gv
+
+
+def test_prim_value():
+ @R.function
+ def foo():
+ gv = R.call_packed("test", 1, sinfo_args=R.Tensor((32, 32), "float32"))
+ return gv
+
+ _check(foo)
+
+
+def test_string_imm():
+ @R.function
+ def foo():
+ gv = R.call_packed("test", "hello", sinfo_args=R.Tensor((32, 32), "float32"))
+ return gv
+
+ _check(foo)
+
+
+def test_datatype_imm():
+ @R.function
+ def foo():
+ gv = R.call_packed("test", R.dtype("float32"), sinfo_args=R.Tensor((32, 32), "float32"))
+ return gv
+
+ _check(foo)
+
+
+def test_function_void_return_type():
+ @tvm.script.ir_module
+ class Foo:
+ @R.function
+ def main(x: R.Tensor((3, 3), dtype="float32")):
+ res = mul(x)
+ return res
+
+ @R.function
+ def mul(x: R.Tensor((3, 3), dtype="float32")):
+ res = R.multiply(x, x)
+ return res
+
+ _check(Foo)
+ # Since the return type of function `mul` is not annotated,
+ # the function `main` regards it as a generic return type.
+ assert isinstance(Foo["main"].ret_struct_info, relax.ObjectStructInfo)
+ assert isinstance(Foo["mul"].ret_struct_info, relax.TensorStructInfo)
+
+ @tvm.script.ir_module
+ class Bar:
+ @R.function
+ def main(x1: R.Tensor((3, 3), dtype="float32")):
+ res1 = mul(x1)
+ return res1
+
+ @R.function
+ def mul(x: R.Tensor((3, 3), dtype="float32")) -> None:
+ res = R.multiply(x, x)
+ return res
+
+ # Since the return type of function `mul` is not annotated,
+ # the function `main` regards it as a generic return type.
+ _check(Bar)
+ tvm.ir.assert_structural_equal(Bar["main"].ret_struct_info, relax.TupleStructInfo([]))
+ tvm.ir.assert_structural_equal(Bar["mul"].ret_struct_info, relax.TupleStructInfo([]))
+
+
+def test_class_normalize():
+ @tvm.script.ir_module
+ class InputModule:
+ @R.function
+ def mul_add(x: R.Tensor) -> R.Tensor:
+ return R.multiply(R.add(x, x), R.add(x, x))
+
+ # The parser automatically normalizes the input AST to the following ANF form
+ @tvm.script.ir_module
+ class OutputModule:
+ @R.function
+ def mul_add(x: R.Tensor) -> R.Tensor:
+ gv = R.add(x, x)
+ gv1 = R.add(x, x)
+ return R.multiply(gv, gv1)
+
+ _check(InputModule, OutputModule)
+
+
+if __name__ == "__main__":
+ test_cross_function_call()
+ tvm.testing.main()