You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/04/08 18:28:41 UTC
[GitHub] [tvm] zxybazh opened a new pull request #7809: [WIP][Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
zxybazh opened a new pull request #7809:
URL: https://github.com/apache/tvm/pull/7809
This PR updated the intrinsic lowering pass to support the new op registry and avoid overloading the global tvm registry. Meanwhile, it kept the fallback mechanism to find the most suitable lower intrinsic function, e.g., `llvm.FLowerIntrinsic` vs. `default.FLowerIntrinsic`. All previous op registration are ported to new functions, and some missing ops would be added.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r619199095
##########
File path: src/tir/transforms/lower_intrin.cc
##########
@@ -42,28 +42,38 @@ class IntrinInjecter : public tvm::arith::IRMutatorWithAnalyzer {
IntrinInjecter(arith::Analyzer* analyzer, std::string target, std::string mtriple = "")
: IRMutatorWithAnalyzer(analyzer) {
- patterns_.push_back("tvm.intrin.rule." + target + ".");
+ patterns_.push_back(target + ".FLowerIntrinsic");
bool is_llvm_aarch64 = (mtriple.find("aarch64") != std::string::npos);
if (is_llvm_aarch64) {
- patterns_.push_back("tvm.intrin.rule." + target + "." + "aarch64.");
+ patterns_.push_back(target + ".aarch64.FLowerIntrinsic");
}
- patterns_.push_back("tvm.intrin.rule.default.");
- fma_ = runtime::Registry::Get(patterns_[0] + "fma");
+ patterns_.push_back("default.FLowerIntrinsic");
+ fma_ = runtime::Registry::Get("tvm.intrin.rule." + target + ".fma");
if (target == "stackvm") {
support_bitwise_op_ = false;
}
}
PrimExpr VisitExpr_(const CallNode* op) final {
if (auto* ptr_op = op->op.as<OpNode>()) {
- // Still use legacy string based rewriting
- // TODO(tvm-team): migrate the pattern application from global function look up
- // to an OpAttrMap<PackedFunc>
- std::string name = ptr_op->name;
- PrimExpr r = ApplyPattern(name, GetRef<PrimExpr>(op));
- if (r.defined()) return r;
+ for (const std::string& pattern : patterns_)
+ if (Op::HasAttrMap(pattern)) {
+ auto f_lower_intrin_map = Op::GetAttrMap<FLowerIntrinsic>(pattern);
Review comment:
Let us restructure this. Because the list of patterns are known in constructor time, it is better to cache the attr_map early. So we do not involve string based lookup in this internal CallNode(can be quite frequent).
We can cache a `vector<OpAttrMap<FLowerIntrinsic>>` or related kind (perhaps unique_ptr needed depends on the ability of copy construct etc) and directly iterate over that during op lookup.
The string based lookup is done once during constructor time, then the rest of lookups are all index based. cc @junrushao1994
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r618662368
##########
File path: include/tvm/ir/op.h
##########
@@ -270,14 +270,17 @@ class OpRegEntry {
* an higher priority level attribute
* will replace lower priority level attribute.
* Must be bigger than 0.
+ * \param can_override Whether to explicitly allow
+ * overriding the attribute, any non-zero value
+ * implies allowance and 0 means disallowance.
*
* Cannot set with same plevel twice in the code.
*
* \tparam ValueType The type of the value to be set.
*/
template <typename ValueType>
inline OpRegEntry& set_attr(const std::string& attr_name, // NOLINT(*)
- const ValueType& value, int plevel = 10);
+ const ValueType& value, int plevel = 10, int can_override = 0);
Review comment:
In terms of API interface additional override flag can cause confusion(because the user might ask how does it interact with the plevel). So I still think in this case we should not add this argument(for API conciseness's pov). It is not as frequent when a developer want to override an intrinsic, and when they do, assuming the knowledge of plevel may not be a bad thing(or they can just put a very big number)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r618661232
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>));
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
Review comment:
The original API was built when we did not have `PrimExpr(PrimExpr)`. Given that I don;t think the FLowerIntrinsic signature will change in the future, it would be great to switch to the typed version
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r617093473
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,43 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ target,
+ f=None,
+ level=10,
+ override=False,
+):
Review comment:
Thanks for pointing that out. I thought about this issue again and I think here `f` is used overwhelmingly more than `level` and `override`. In most cases, it's just `op_name`, `target` and `f`. Therefore, I insist that we keep the current order for most intuitive function usage.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r618664584
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>));
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
Review comment:
I see. In that case typed version would be better. Thanks for pointing that out.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on pull request #7809: [WIP][Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#issuecomment-819034462
Thanks for the contribution! A few high-level comments:
1. Do not add any C API - we can use packed func to expose the interface. Op is a compile-time concept so please move the registration logic out of the runtime :-)
2. Please change `register_op` to some name like `register_intrin_lowering`, and consider accept a target kind like `default`/`llvm`/`cuda`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r617075762
##########
File path: python/tvm/ir/op.py
##########
@@ -17,9 +17,21 @@
# pylint: disable=invalid-name
"""Primitive operators in the TVM IR."""
import tvm._ffi
+from tvm._ffi.base import _FFI_MODE
from .expr import RelayExpr
from . import _ffi_api
+try:
+ # pylint: disable=wrong-import-position,unused-import
+ if _FFI_MODE == "ctypes":
+ raise ImportError()
+ from tvm._ffi._cy3.core import convert_to_tvm_func, _get_global_func, PackedFuncBase
+except (RuntimeError, ImportError) as error:
+ # pylint: disable=wrong-import-position,unused-import
+ if _FFI_MODE == "cython":
+ raise error
+ from tvm._ffi._ctypes.packed_func import convert_to_tvm_func, _get_global_func, PackedFuncBase
Review comment:
Removed, thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r614585193
##########
File path: src/target/intrin_rule.cc
##########
@@ -24,179 +24,234 @@
#include "intrin_rule.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace intrin {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.exp").set_body(DispatchPureExtern<FloatSuffix>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>("default.FLowerIntrinsic",
Review comment:
shall we remove this file to, like, `lower_intrinsic_default.cc`? CC: @tqchen
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on pull request #7809: [WIP][Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#issuecomment-818262055
Python side op lowering function has been changed from `register_func` to `register_op`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r616164032
##########
File path: include/tvm/tir/op_attr_types.h
##########
@@ -28,11 +28,13 @@
#ifndef TVM_TIR_OP_ATTR_TYPES_H_
#define TVM_TIR_OP_ATTR_TYPES_H_
+#include <tvm/ir/expr.h>
Review comment:
This is kept for usage of `tvm::Integer` class in the file. Not sure why it was not required in the orginal file though.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 merged pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 merged pull request #7809:
URL: https://github.com/apache/tvm/pull/7809
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r618637227
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>));
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
Review comment:
It seems that in order to use typed version of `PackedFunc` the signature of the later function must be exactly the same, i.e., changing signatures of all functions to `PrimExpr(PrimExpr)`. I think the current functions are pretty much in good shape and if we want to change the function type of `FLowerIntrinsic` in the future we don't have to modify the code. What do you think?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r618633054
##########
File path: src/target/llvm/intrin_rule_hexagon.cc
##########
@@ -19,44 +19,54 @@
#ifdef TVM_LLVM_VERSION
+#include <tvm/tir/op_attr_types.h>
+
#include "intrin_rule_llvm.h"
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.hexagon.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "hexagon.FLowerIntrinsic", PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>));
Review comment:
Thanks! Fixed in the new commit.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r614576312
##########
File path: include/tvm/tir/op_attr_types.h
##########
@@ -28,11 +28,13 @@
#ifndef TVM_TIR_OP_ATTR_TYPES_H_
#define TVM_TIR_OP_ATTR_TYPES_H_
+#include <tvm/ir/expr.h>
Review comment:
No need for this
##########
File path: include/tvm/tir/op_attr_types.h
##########
@@ -28,11 +28,13 @@
#ifndef TVM_TIR_OP_ATTR_TYPES_H_
#define TVM_TIR_OP_ATTR_TYPES_H_
+#include <tvm/ir/expr.h>
#include <tvm/runtime/container.h>
+#include <tvm/runtime/packed_func.h>
namespace tvm {
namespace tir {
-
+using namespace runtime;
Review comment:
No need for this
##########
File path: include/tvm/tir/op_attr_types.h
##########
@@ -43,6 +45,11 @@ using TGlobalSymbol = String;
*/
using TVectorizable = bool;
+/*!
+ * \brief The intrinsic lowering function for given OP.
+ */
+using FLowerIntrinsic = PackedFunc;
Review comment:
Using `runtime::TypedPackedFunc` when possible, which is a simple wrapper of `PackedFunc`, but with the clear signature like `std::function`, so it could help readability :-)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r614581355
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,50 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ f=None,
+ target="default",
+ level=10,
+ override=False,
+):
+ """Register Op lowering function
+
+ Parameters
+ ----------
+ op_name : str or function
+ The op name
+
+ f : function, optional
+ The function to be registered.
+
+ target : str
+ The target string for given intrinsic lowering function
+
+ level : int
+ The priority level
+
+ override: boolean optional
+ Whether override existing entry.
+
+ Returns
+ -------
+ fregister : function
+ Register op lowering function if f is not specified.
+ """
+ if not isinstance(op_name, str):
+ raise ValueError("expect string op name")
+
+ def _register(myf):
+ """internal intrinsic lowering registration function"""
+ assert isinstance(target, str)
+ if not isinstance(myf, PackedFuncBase):
+ myf = convert_to_tvm_func(myf)
+ _ffi_api.RegisterOpLowerIntrinsic(op_name, myf.handle, target, level, override)
+ return myf
Review comment:
I assume the logic could be a bit simpler? Perhaps you might want to refer the some previous implementations, like "def register_compute" :-)
##########
File path: src/ir/op.cc
##########
@@ -122,6 +124,19 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
}
});
+TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
Review comment:
I assume this should be put on `src/tir/op.cc`?
##########
File path: src/ir/op.cc
##########
@@ -122,6 +124,19 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
}
});
+TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
+ .set_body_typed([](String name, TVMFunctionHandle f, String target = "default", int plevel = 10,
+ int can_override = 0) {
Review comment:
Forget about TVMFunctionHandle, which is only used in the C API but nowhere else :-)
Let's embrace the PackedFunc magic! It is a nicer wrapper of a function from either C++ or python (or potentially other languages).
BTW, don't need to set default arguments of a lambda because they are not used any way.
```suggestion
.set_body_typed([](String name, PackedFunc f, String target, int plevel, bool override_) {
```
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,50 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ f=None,
+ target="default",
+ level=10,
+ override=False,
+):
Review comment:
Let's avoid providing too many defaults :-)
```suggestion
def register_op_intrin_lowering(
op_name,
f,
target,
level=10,
override=False,
):
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#issuecomment-824532034
@tqchen would you like to take a second look? If there is other comments, I will get it merged the other day
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r616225480
##########
File path: src/ir/op.cc
##########
@@ -122,6 +124,19 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
}
});
+TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
Review comment:
Yeah I have the same concern but I checked both files and found most similar functions like `RegisterOpAttr` function in the same file. Therefore, I think it better to keep the function here or we would need to migrate all related functions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r617207042
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,43 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ target,
+ f=None,
+ level=10,
+ override=False,
+):
Review comment:
This is an interesting topic to discuss :-)
In our particular case, we could still provide default values for some arguments, like:
```python
>>> def f(a, *, b, c=1):
... print(f"a = {a}, b = {b}, c = {c}")
...
>>> f(1, 2)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: f() takes 1 positional argument but 2 were given
>>> f(1, b=2)
a = 1, b = 2, c = 1
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#issuecomment-824270527
CC: @comaniac @yzhliu @icemelon9
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r619494697
##########
File path: src/tir/transforms/lower_intrin.cc
##########
@@ -42,28 +42,38 @@ class IntrinInjecter : public tvm::arith::IRMutatorWithAnalyzer {
IntrinInjecter(arith::Analyzer* analyzer, std::string target, std::string mtriple = "")
: IRMutatorWithAnalyzer(analyzer) {
- patterns_.push_back("tvm.intrin.rule." + target + ".");
+ patterns_.push_back(target + ".FLowerIntrinsic");
bool is_llvm_aarch64 = (mtriple.find("aarch64") != std::string::npos);
if (is_llvm_aarch64) {
- patterns_.push_back("tvm.intrin.rule." + target + "." + "aarch64.");
+ patterns_.push_back(target + ".aarch64.FLowerIntrinsic");
}
- patterns_.push_back("tvm.intrin.rule.default.");
- fma_ = runtime::Registry::Get(patterns_[0] + "fma");
+ patterns_.push_back("default.FLowerIntrinsic");
+ fma_ = runtime::Registry::Get("tvm.intrin.rule." + target + ".fma");
if (target == "stackvm") {
support_bitwise_op_ = false;
}
}
PrimExpr VisitExpr_(const CallNode* op) final {
if (auto* ptr_op = op->op.as<OpNode>()) {
- // Still use legacy string based rewriting
- // TODO(tvm-team): migrate the pattern application from global function look up
- // to an OpAttrMap<PackedFunc>
- std::string name = ptr_op->name;
- PrimExpr r = ApplyPattern(name, GetRef<PrimExpr>(op));
- if (r.defined()) return r;
+ for (const std::string& pattern : patterns_)
+ if (Op::HasAttrMap(pattern)) {
+ auto f_lower_intrin_map = Op::GetAttrMap<FLowerIntrinsic>(pattern);
Review comment:
Good point, fixed in the new commit.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r618592045
##########
File path: include/tvm/ir/op.h
##########
@@ -270,14 +270,17 @@ class OpRegEntry {
* an higher priority level attribute
* will replace lower priority level attribute.
* Must be bigger than 0.
+ * \param can_override Whether to explicitly allow
+ * overriding the attribute, any non-zero value
+ * implies allowance and 0 means disallowance.
*
* Cannot set with same plevel twice in the code.
*
* \tparam ValueType The type of the value to be set.
*/
template <typename ValueType>
inline OpRegEntry& set_attr(const std::string& attr_name, // NOLINT(*)
- const ValueType& value, int plevel = 10);
+ const ValueType& value, int plevel = 10, int can_override = 0);
Review comment:
Thanks for pointing it out. In this case, new attributes would be accepted with a different `plevel` but in the cases of overriding, it becomes a bit troublesome without `can_override` - we would need to either query the current `plevel` or try to use a different number somehow. Therefore, I think it would be more concise to simply add this flag for convenience but definitely there could be better solutions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 edited a comment on pull request #7809: [WIP][Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 edited a comment on pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#issuecomment-819034462
Thanks for the contribution! A few high-level comments:
1. Do not add any C API - we can use packed func to expose the interface. Op is a compile-time concept so please move the registration logic out of the runtime :-)
2. Please change `register_op` to some name like `register_intrin_lowering`, and consider accepting a target kind like `default`/`llvm`/`cuda`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r616227265
##########
File path: src/target/intrin_rule.cc
##########
@@ -24,179 +24,234 @@
#include "intrin_rule.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace intrin {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.exp").set_body(DispatchPureExtern<FloatSuffix>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>("default.FLowerIntrinsic",
Review comment:
I think the naming is just fine because the `lower_intrinsic.cc` file kind of implies it's default rules.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r618664064
##########
File path: include/tvm/ir/op.h
##########
@@ -270,14 +270,17 @@ class OpRegEntry {
* an higher priority level attribute
* will replace lower priority level attribute.
* Must be bigger than 0.
+ * \param can_override Whether to explicitly allow
+ * overriding the attribute, any non-zero value
+ * implies allowance and 0 means disallowance.
*
* Cannot set with same plevel twice in the code.
*
* \tparam ValueType The type of the value to be set.
*/
template <typename ValueType>
inline OpRegEntry& set_attr(const std::string& attr_name, // NOLINT(*)
- const ValueType& value, int plevel = 10);
+ const ValueType& value, int plevel = 10, int can_override = 0);
Review comment:
Yeah that make sense. Thanks for the explanation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r616950245
##########
File path: python/tvm/ir/op.py
##########
@@ -17,9 +17,21 @@
# pylint: disable=invalid-name
"""Primitive operators in the TVM IR."""
import tvm._ffi
+from tvm._ffi.base import _FFI_MODE
from .expr import RelayExpr
from . import _ffi_api
+try:
+ # pylint: disable=wrong-import-position,unused-import
+ if _FFI_MODE == "ctypes":
+ raise ImportError()
+ from tvm._ffi._cy3.core import convert_to_tvm_func, _get_global_func, PackedFuncBase
+except (RuntimeError, ImportError) as error:
+ # pylint: disable=wrong-import-position,unused-import
+ if _FFI_MODE == "cython":
+ raise error
+ from tvm._ffi._ctypes.packed_func import convert_to_tvm_func, _get_global_func, PackedFuncBase
Review comment:
๐ i suppose these lines become useless given you have addressed the comments i left before?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r616979231
##########
File path: src/tir/transforms/lower_intrin.cc
##########
@@ -42,28 +42,36 @@ class IntrinInjecter : public tvm::arith::IRMutatorWithAnalyzer {
IntrinInjecter(arith::Analyzer* analyzer, std::string target, std::string mtriple = "")
: IRMutatorWithAnalyzer(analyzer) {
- patterns_.push_back("tvm.intrin.rule." + target + ".");
+ patterns_.push_back(target + ".FLowerIntrinsic");
bool is_llvm_aarch64 = (mtriple.find("aarch64") != std::string::npos);
if (is_llvm_aarch64) {
- patterns_.push_back("tvm.intrin.rule." + target + "." + "aarch64.");
+ patterns_.push_back(target + ".aarch64.FLowerIntrinsic");
}
- patterns_.push_back("tvm.intrin.rule.default.");
- fma_ = runtime::Registry::Get(patterns_[0] + "fma");
+ patterns_.push_back("default.FLowerIntrinsic");
+ fma_ = runtime::Registry::Get("tvm.intrin.rule." + target + ".fma");
if (target == "stackvm") {
support_bitwise_op_ = false;
}
}
PrimExpr VisitExpr_(const CallNode* op) final {
if (auto* ptr_op = op->op.as<OpNode>()) {
- // Still use legacy string based rewriting
- // TODO(tvm-team): migrate the pattern application from global function look up
- // to an OpAttrMap<PackedFunc>
- std::string name = ptr_op->name;
- PrimExpr r = ApplyPattern(name, GetRef<PrimExpr>(op));
- if (r.defined()) return r;
+ for (size_t i = 0; i < patterns_.size(); ++i)
Review comment:
```suggestion
for (const std::string& pattern : patterns_)
```
##########
File path: src/tir/transforms/lower_intrin.cc
##########
@@ -42,28 +42,36 @@ class IntrinInjecter : public tvm::arith::IRMutatorWithAnalyzer {
IntrinInjecter(arith::Analyzer* analyzer, std::string target, std::string mtriple = "")
: IRMutatorWithAnalyzer(analyzer) {
- patterns_.push_back("tvm.intrin.rule." + target + ".");
+ patterns_.push_back(target + ".FLowerIntrinsic");
bool is_llvm_aarch64 = (mtriple.find("aarch64") != std::string::npos);
if (is_llvm_aarch64) {
- patterns_.push_back("tvm.intrin.rule." + target + "." + "aarch64.");
+ patterns_.push_back(target + ".aarch64.FLowerIntrinsic");
}
- patterns_.push_back("tvm.intrin.rule.default.");
- fma_ = runtime::Registry::Get(patterns_[0] + "fma");
+ patterns_.push_back("default.FLowerIntrinsic");
+ fma_ = runtime::Registry::Get("tvm.intrin.rule." + target + ".fma");
if (target == "stackvm") {
support_bitwise_op_ = false;
}
}
PrimExpr VisitExpr_(const CallNode* op) final {
if (auto* ptr_op = op->op.as<OpNode>()) {
- // Still use legacy string based rewriting
- // TODO(tvm-team): migrate the pattern application from global function look up
- // to an OpAttrMap<PackedFunc>
- std::string name = ptr_op->name;
- PrimExpr r = ApplyPattern(name, GetRef<PrimExpr>(op));
- if (r.defined()) return r;
+ for (size_t i = 0; i < patterns_.size(); ++i)
+ if (Op::HasAttrMap(patterns_[i])) {
+ auto default_intrin = Op::GetAttrMap<FLowerIntrinsic>(patterns_[i]);
Review comment:
it is not necessarily default intrin, so let's find a more accurate name here? like, `f_lower_intrin_map`?
##########
File path: src/tir/transforms/lower_intrin.cc
##########
@@ -42,28 +42,36 @@ class IntrinInjecter : public tvm::arith::IRMutatorWithAnalyzer {
IntrinInjecter(arith::Analyzer* analyzer, std::string target, std::string mtriple = "")
: IRMutatorWithAnalyzer(analyzer) {
- patterns_.push_back("tvm.intrin.rule." + target + ".");
+ patterns_.push_back(target + ".FLowerIntrinsic");
bool is_llvm_aarch64 = (mtriple.find("aarch64") != std::string::npos);
if (is_llvm_aarch64) {
- patterns_.push_back("tvm.intrin.rule." + target + "." + "aarch64.");
+ patterns_.push_back(target + ".aarch64.FLowerIntrinsic");
}
- patterns_.push_back("tvm.intrin.rule.default.");
- fma_ = runtime::Registry::Get(patterns_[0] + "fma");
+ patterns_.push_back("default.FLowerIntrinsic");
+ fma_ = runtime::Registry::Get("tvm.intrin.rule." + target + ".fma");
if (target == "stackvm") {
support_bitwise_op_ = false;
}
}
PrimExpr VisitExpr_(const CallNode* op) final {
if (auto* ptr_op = op->op.as<OpNode>()) {
- // Still use legacy string based rewriting
- // TODO(tvm-team): migrate the pattern application from global function look up
- // to an OpAttrMap<PackedFunc>
- std::string name = ptr_op->name;
- PrimExpr r = ApplyPattern(name, GetRef<PrimExpr>(op));
- if (r.defined()) return r;
+ for (size_t i = 0; i < patterns_.size(); ++i)
+ if (Op::HasAttrMap(patterns_[i])) {
+ auto default_intrin = Op::GetAttrMap<FLowerIntrinsic>(patterns_[i]);
+ FLowerIntrinsic f = default_intrin.get(GetRef<Op>(ptr_op), nullptr);
+ PrimExpr e = GetRef<PrimExpr>(op);
+ if (f != nullptr) {
+ PrimExpr r = f(e);
Review comment:
```suggestion
if (f != nullptr) {
PrimExpr e = GetRef<PrimExpr>(op);
PrimExpr r = f(e);
```
##########
File path: src/tir/transforms/lower_intrin.cc
##########
@@ -42,28 +42,36 @@ class IntrinInjecter : public tvm::arith::IRMutatorWithAnalyzer {
IntrinInjecter(arith::Analyzer* analyzer, std::string target, std::string mtriple = "")
: IRMutatorWithAnalyzer(analyzer) {
- patterns_.push_back("tvm.intrin.rule." + target + ".");
+ patterns_.push_back(target + ".FLowerIntrinsic");
bool is_llvm_aarch64 = (mtriple.find("aarch64") != std::string::npos);
if (is_llvm_aarch64) {
- patterns_.push_back("tvm.intrin.rule." + target + "." + "aarch64.");
+ patterns_.push_back(target + ".aarch64.FLowerIntrinsic");
}
- patterns_.push_back("tvm.intrin.rule.default.");
- fma_ = runtime::Registry::Get(patterns_[0] + "fma");
+ patterns_.push_back("default.FLowerIntrinsic");
+ fma_ = runtime::Registry::Get("tvm.intrin.rule." + target + ".fma");
if (target == "stackvm") {
support_bitwise_op_ = false;
}
}
PrimExpr VisitExpr_(const CallNode* op) final {
if (auto* ptr_op = op->op.as<OpNode>()) {
- // Still use legacy string based rewriting
- // TODO(tvm-team): migrate the pattern application from global function look up
- // to an OpAttrMap<PackedFunc>
- std::string name = ptr_op->name;
- PrimExpr r = ApplyPattern(name, GetRef<PrimExpr>(op));
- if (r.defined()) return r;
+ for (size_t i = 0; i < patterns_.size(); ++i)
+ if (Op::HasAttrMap(patterns_[i])) {
+ auto default_intrin = Op::GetAttrMap<FLowerIntrinsic>(patterns_[i]);
+ FLowerIntrinsic f = default_intrin.get(GetRef<Op>(ptr_op), nullptr);
+ PrimExpr e = GetRef<PrimExpr>(op);
+ if (f != nullptr) {
+ PrimExpr r = f(e);
+ ICHECK(r.defined()) << "intrinsic rule must always return valid Expr";
+ if (!r.same_as(e)) {
+ r = this->VisitExpr(r);
+ if (r.defined()) return r;
Review comment:
```suggestion
if (r.defined()) {
return r;
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r616227265
##########
File path: src/target/intrin_rule.cc
##########
@@ -24,179 +24,234 @@
#include "intrin_rule.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace intrin {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.default.exp").set_body(DispatchPureExtern<FloatSuffix>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>("default.FLowerIntrinsic",
Review comment:
I think the naming is just fine because the `lower_intrinsic.cc`file kind of implies it's default rules.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r617085138
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>));
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
Review comment:
oh i see...let's keep it then :-)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r616224603
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,50 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ f=None,
+ target="default",
+ level=10,
+ override=False,
+):
+ """Register Op lowering function
+
+ Parameters
+ ----------
+ op_name : str or function
+ The op name
+
+ f : function, optional
+ The function to be registered.
+
+ target : str
+ The target string for given intrinsic lowering function
+
+ level : int
+ The priority level
+
+ override: boolean optional
+ Whether override existing entry.
+
+ Returns
+ -------
+ fregister : function
+ Register op lowering function if f is not specified.
+ """
+ if not isinstance(op_name, str):
+ raise ValueError("expect string op name")
+
+ def _register(myf):
+ """internal intrinsic lowering registration function"""
+ assert isinstance(target, str)
+ if not isinstance(myf, PackedFuncBase):
+ myf = convert_to_tvm_func(myf)
+ _ffi_api.RegisterOpLowerIntrinsic(op_name, myf.handle, target, level, override)
+ return myf
Review comment:
Sure it could be simplified but I'll still keep the `_register` function for consistensy in function implementation compared to `register_op_attr`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r617083997
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>));
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
Review comment:
Unfortunately if I remove the type conversion I get the following error.
`tvm/src/target/llvm/intrin_rule_llvm.cc:196:7: note: cannot convert โ<lambda closure object>tvm::codegen::llvm::<lambda(const tvm::runtime::TVMArgs&, tvm::runtime::TVMRetValue*)>{}โ (type โtvm::codegen::llvm::<lambda(const tvm::runtime::TVMArgs&, tvm::runtime::TVMRetValue*)>โ) to type โconst tvm::runtime::TypedPackedFunc<tvm::PrimExpr(tvm::PrimExpr)>&โ
}));`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r619197339
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(), 2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs, 1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
Review comment:
[](PrimExpr expr) -> PrimExpr
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
Review comment:
Consider convert this to TypedPackedFunc as well, by using a `[](PrimExpr expr) -> PrimExpr` signature
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(), 2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs, 1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr one = make_const(x.dtype(), 1);
+ PrimExpr two = make_const(x.dtype(), 2);
+ PrimExpr neg_two = make_const(x.dtype(), -2);
+
+ PrimExpr exp_neg2x = exp(neg_two * x);
+ PrimExpr exp_pos2x = exp(two * x);
+
+ PrimExpr tanh_pos = (one - exp_neg2x) / (one + exp_neg2x);
+ PrimExpr tanh_neg = (exp_pos2x - one) / (exp_pos2x + one);
+ *rv = tir::Select(x >= make_zero(x.dtype()), tanh_pos, tanh_neg);
+ }));
+
+TVM_REGISTER_OP("tir.pow").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
+
+TVM_REGISTER_OP("tir.popcount")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
+
+TVM_REGISTER_OP("tir.tan").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
Review comment:
NOTE: we should change these to FLegalize later once we start to introduce legalize registries logics. And perhaps move to generic legalize
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(), 2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs, 1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
Review comment:
NOTE: we should change these to FLegalize later once we start to introduce legalize registries logics. And perhaps move to generic legalize
##########
File path: src/tir/transforms/lower_intrin.cc
##########
@@ -42,28 +42,38 @@ class IntrinInjecter : public tvm::arith::IRMutatorWithAnalyzer {
IntrinInjecter(arith::Analyzer* analyzer, std::string target, std::string mtriple = "")
: IRMutatorWithAnalyzer(analyzer) {
- patterns_.push_back("tvm.intrin.rule." + target + ".");
+ patterns_.push_back(target + ".FLowerIntrinsic");
bool is_llvm_aarch64 = (mtriple.find("aarch64") != std::string::npos);
if (is_llvm_aarch64) {
- patterns_.push_back("tvm.intrin.rule." + target + "." + "aarch64.");
+ patterns_.push_back(target + ".aarch64.FLowerIntrinsic");
}
- patterns_.push_back("tvm.intrin.rule.default.");
- fma_ = runtime::Registry::Get(patterns_[0] + "fma");
+ patterns_.push_back("default.FLowerIntrinsic");
+ fma_ = runtime::Registry::Get("tvm.intrin.rule." + target + ".fma");
if (target == "stackvm") {
support_bitwise_op_ = false;
}
}
PrimExpr VisitExpr_(const CallNode* op) final {
if (auto* ptr_op = op->op.as<OpNode>()) {
- // Still use legacy string based rewriting
- // TODO(tvm-team): migrate the pattern application from global function look up
- // to an OpAttrMap<PackedFunc>
- std::string name = ptr_op->name;
- PrimExpr r = ApplyPattern(name, GetRef<PrimExpr>(op));
- if (r.defined()) return r;
+ for (const std::string& pattern : patterns_)
+ if (Op::HasAttrMap(pattern)) {
+ auto f_lower_intrin_map = Op::GetAttrMap<FLowerIntrinsic>(pattern);
Review comment:
Let us restructure this. Because the list of patterns are known in constructor time, it is better to cache the attr_map early. So we do not involve string based lookup in this internal CallNode.
The string based lookup is done once during constructor time, then the rest of lookups are all index based. cc @junrushao1994
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(), 2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs, 1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr one = make_const(x.dtype(), 1);
+ PrimExpr two = make_const(x.dtype(), 2);
+ PrimExpr neg_two = make_const(x.dtype(), -2);
+
+ PrimExpr exp_neg2x = exp(neg_two * x);
+ PrimExpr exp_pos2x = exp(two * x);
+
+ PrimExpr tanh_pos = (one - exp_neg2x) / (one + exp_neg2x);
+ PrimExpr tanh_neg = (exp_pos2x - one) / (exp_pos2x + one);
+ *rv = tir::Select(x >= make_zero(x.dtype()), tanh_pos, tanh_neg);
+ }));
+
+TVM_REGISTER_OP("tir.pow").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
+
+TVM_REGISTER_OP("tir.popcount")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
+
+TVM_REGISTER_OP("tir.tan").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
Review comment:
[](PrimExpr expr) -> PrimExpr
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(), 2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs, 1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr one = make_const(x.dtype(), 1);
+ PrimExpr two = make_const(x.dtype(), 2);
+ PrimExpr neg_two = make_const(x.dtype(), -2);
+
+ PrimExpr exp_neg2x = exp(neg_two * x);
+ PrimExpr exp_pos2x = exp(two * x);
+
+ PrimExpr tanh_pos = (one - exp_neg2x) / (one + exp_neg2x);
+ PrimExpr tanh_neg = (exp_pos2x - one) / (exp_pos2x + one);
+ *rv = tir::Select(x >= make_zero(x.dtype()), tanh_pos, tanh_neg);
+ }));
+
+TVM_REGISTER_OP("tir.pow").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
+
+TVM_REGISTER_OP("tir.popcount")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
+
+TVM_REGISTER_OP("tir.tan").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
PrimExpr e = targs[0];
const tir::CallNode* call = e.as<tir::CallNode>();
ICHECK(call != nullptr);
const PrimExpr& x = call->args[0];
- PrimExpr ln10 = make_const(x.dtype(), 2.302585093);
- PrimExpr ret = exp(x * ln10);
- *rv = ret;
- });
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.fma")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 3>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log10")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.sqrt")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.floor")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.ceil")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.trunc")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.fabs")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.round")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.nearbyint")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.tanh")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
- PrimExpr e = targs[0];
- const tir::CallNode* call = e.as<tir::CallNode>();
- ICHECK(call != nullptr);
- const PrimExpr& x = call->args[0];
- PrimExpr one = make_const(x.dtype(), 1);
- PrimExpr two = make_const(x.dtype(), 2);
- PrimExpr neg_two = make_const(x.dtype(), -2);
-
- PrimExpr exp_neg2x = exp(neg_two * x);
- PrimExpr exp_pos2x = exp(two * x);
-
- PrimExpr tanh_pos = (one - exp_neg2x) / (one + exp_neg2x);
- PrimExpr tanh_neg = (exp_pos2x - one) / (exp_pos2x + one);
- *rv = tir::Select(x >= make_zero(x.dtype()), tanh_pos, tanh_neg);
- });
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.pow")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.popcount")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.tan").set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- PrimExpr e = targs[0];
- const tir::CallNode* call = e.as<tir::CallNode>();
- ICHECK(call != nullptr);
- const PrimExpr& x = call->args[0];
- PrimExpr tan_x = sin(x) / cos(x);
- *rv = tan_x;
-});
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.cos")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::cos, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.cosh")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+ PrimExpr tan_x = sin(x) / cos(x);
+ *rv = tan_x;
+ }));
+
+TVM_REGISTER_OP("tir.cos").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::cos, 1>);
+
+TVM_REGISTER_OP("tir.cosh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr two = make_const(x.dtype(), 2);
+ PrimExpr neg_one = make_const(x.dtype(), -1);
+ PrimExpr exp_negx = exp(neg_one * x);
+ PrimExpr exp_posx = exp(x);
+ PrimExpr ret = (exp_posx + exp_negx) / two;
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.sin").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::sin, 1>);
+
+TVM_REGISTER_OP("tir.sinh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
Review comment:
[](PrimExpr expr) -> PrimExpr
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(), 2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs, 1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call = e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr one = make_const(x.dtype(), 1);
+ PrimExpr two = make_const(x.dtype(), 2);
+ PrimExpr neg_two = make_const(x.dtype(), -2);
+
+ PrimExpr exp_neg2x = exp(neg_two * x);
+ PrimExpr exp_pos2x = exp(two * x);
+
+ PrimExpr tanh_pos = (one - exp_neg2x) / (one + exp_neg2x);
+ PrimExpr tanh_neg = (exp_pos2x - one) / (exp_pos2x + one);
+ *rv = tir::Select(x >= make_zero(x.dtype()), tanh_pos, tanh_neg);
+ }));
+
+TVM_REGISTER_OP("tir.pow").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
+
+TVM_REGISTER_OP("tir.popcount")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
+
+TVM_REGISTER_OP("tir.tan").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
PrimExpr e = targs[0];
const tir::CallNode* call = e.as<tir::CallNode>();
ICHECK(call != nullptr);
const PrimExpr& x = call->args[0];
- PrimExpr ln10 = make_const(x.dtype(), 2.302585093);
- PrimExpr ret = exp(x * ln10);
- *rv = ret;
- });
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.fma")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 3>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log10")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.sqrt")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.floor")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.ceil")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.trunc")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.fabs")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.round")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.nearbyint")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.tanh")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
- PrimExpr e = targs[0];
- const tir::CallNode* call = e.as<tir::CallNode>();
- ICHECK(call != nullptr);
- const PrimExpr& x = call->args[0];
- PrimExpr one = make_const(x.dtype(), 1);
- PrimExpr two = make_const(x.dtype(), 2);
- PrimExpr neg_two = make_const(x.dtype(), -2);
-
- PrimExpr exp_neg2x = exp(neg_two * x);
- PrimExpr exp_pos2x = exp(two * x);
-
- PrimExpr tanh_pos = (one - exp_neg2x) / (one + exp_neg2x);
- PrimExpr tanh_neg = (exp_pos2x - one) / (exp_pos2x + one);
- *rv = tir::Select(x >= make_zero(x.dtype()), tanh_pos, tanh_neg);
- });
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.pow")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.popcount")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.tan").set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- PrimExpr e = targs[0];
- const tir::CallNode* call = e.as<tir::CallNode>();
- ICHECK(call != nullptr);
- const PrimExpr& x = call->args[0];
- PrimExpr tan_x = sin(x) / cos(x);
- *rv = tan_x;
-});
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.cos")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::cos, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.cosh")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+ PrimExpr tan_x = sin(x) / cos(x);
+ *rv = tan_x;
+ }));
+
+TVM_REGISTER_OP("tir.cos").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::cos, 1>);
+
+TVM_REGISTER_OP("tir.cosh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
Review comment:
[](PrimExpr expr) -> PrimExpr
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
tqchen commented on pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#issuecomment-822878863
note, please do not checkin vta-hw submodule change(do a git submodule update)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#issuecomment-823503782
> note, please do not checkin vta-hw submodule change(do a git submodule update)
Noted with thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r617170972
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,43 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ target,
+ f=None,
+ level=10,
+ override=False,
+):
Review comment:
In that case, it would be awkward to use the function as decorators as we would need to mark all the argument. What do you think?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r616970743
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,43 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ target,
+ f=None,
+ level=10,
+ override=False,
+):
Review comment:
oh BTW, the order of the arguments matter, and we need to double check if the following syntax works
```
@tvm.ir.op.register_op_intrin_lowering("tir.exp", target="cuda", override=True)
def my_exp(expr) -> expr:
pass
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r618561537
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>));
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
Review comment:
Would be great use the typed version of PackedFunc
##########
File path: include/tvm/ir/op.h
##########
@@ -270,14 +270,17 @@ class OpRegEntry {
* an higher priority level attribute
* will replace lower priority level attribute.
* Must be bigger than 0.
+ * \param can_override Whether to explicitly allow
+ * overriding the attribute, any non-zero value
+ * implies allowance and 0 means disallowance.
*
* Cannot set with same plevel twice in the code.
*
* \tparam ValueType The type of the value to be set.
*/
template <typename ValueType>
inline OpRegEntry& set_attr(const std::string& attr_name, // NOLINT(*)
- const ValueType& value, int plevel = 10);
+ const ValueType& value, int plevel = 10, int can_override = 0);
Review comment:
Wju do we need to introduce can_override when we already have plevel to specify overriding rule.
Previously we check the plevel, if the plevel have higher priority level then we can override the attribute
##########
File path: src/target/llvm/intrin_rule_hexagon.cc
##########
@@ -19,44 +19,54 @@
#ifdef TVM_LLVM_VERSION
+#include <tvm/tir/op_attr_types.h>
+
#include "intrin_rule_llvm.h"
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.hexagon.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "hexagon.FLowerIntrinsic", PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>));
Review comment:
would be nice if we can modify the signature of DispatchLLVMPureIntrin to be PrimExpr->PrimExpr, so we do not have to do the convertion to PackedFunc
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#issuecomment-820744866
@junrushao1994 Please review, thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on pull request #7809: [WIP][Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#issuecomment-819124827
Thanks for the advice! I have changed the name of function in python-side to `register_op_intrin_lowering` and removed change to C API. Instead, I utilized the global function registration and packedfunc to register op intrinsic lowering function.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r617150823
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,43 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ target,
+ f=None,
+ level=10,
+ override=False,
+):
Review comment:
An alternative approach is to force kwargs, as demonstrated below, so that users are required to explicitly mark which argument is override and which is level
```python
def (a, b, *, c, d)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r616961994
##########
File path: src/ir/op.cc
##########
@@ -122,6 +124,18 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
}
});
+TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
+ .set_body_typed([](String name, PackedFunc f, String target, int plevel, int can_override) {
+ if (Op::HasAttrMap(target + ".FLowerIntrinsic") &&
+ OpRegistry::Global()->Get(name) != nullptr &&
+ Op::GetAttrMap<FLowerIntrinsic>(target + ".FLowerIntrinsic").count(Op::Get(name))) {
+ ICHECK(can_override) << "Op " << name << "'s intrinsic lowering function " << target
Review comment:
Use `CHECK` here because it is a user-facing error message.
```suggestion
CHECK(can_override) << "Op " << name << "'s intrinsic lowering function " << target
```
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>));
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>));
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs, TVMRetValue* rv) {
Review comment:
IIRC without this `PackedFunc` conversion the code still compiles (because we do implicit conversion), so in this particular case (registering with a lambda), we can remove this explicit conversion
##########
File path: src/ir/op.cc
##########
@@ -122,6 +124,18 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
}
});
+TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
+ .set_body_typed([](String name, PackedFunc f, String target, int plevel, int can_override) {
+ if (Op::HasAttrMap(target + ".FLowerIntrinsic") &&
+ OpRegistry::Global()->Get(name) != nullptr &&
+ Op::GetAttrMap<FLowerIntrinsic>(target + ".FLowerIntrinsic").count(Op::Get(name))) {
+ ICHECK(can_override) << "Op " << name << "'s intrinsic lowering function " << target
+ << ".FlowerIntrinsic is already registered";
+ }
+ tvm::OpRegEntry::RegisterOrGet(name).set_name().set_attr<FLowerIntrinsic>(
Review comment:
No need to set_name in this case
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,43 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ target,
+ f=None,
+ level=10,
+ override=False,
+):
+ """Register Op lowering function
+
+ Parameters
+ ----------
+ op_name : str or function
+ The op name
Review comment:
Let's only allow the string case
##########
File path: python/tvm/target/intrin.py
##########
@@ -112,14 +68,18 @@ def _rule_float_direct(op):
See Also
--------
- register_intrin_rule : The registration function for intrin rule.
+ tvm.ir.op.register_op_intrin_lowering : The registration function for intrinsic lowering rule.
"""
if str(op.dtype).startswith("float"):
return call_pure_extern(op.dtype, op.op.name[4:], *op.args)
return None
# opencl pattern for exp
-register_intrin_rule("opencl", "exp", _rule_float_direct, override=True)
+tvm.ir.op.register_op_intrin_lowering(
Review comment:
we should do like, `from tvm.ir.op import register_op_intrin_lowering`, in the beginning of this file?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r617093473
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,43 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ target,
+ f=None,
+ level=10,
+ override=False,
+):
Review comment:
Thanks for point that out. I thought about this issue again and I think here `f` is used overwhelmingly more than `level` and `override`. In most cases, it's just `op_name`, `target` and `f`. Therefore, I insist that we keep the current order for most intuitive function usage.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#issuecomment-826041090
Thanks @zxybazh @tqchen!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r618558723
##########
File path: include/tvm/ir/op.h
##########
@@ -270,14 +270,17 @@ class OpRegEntry {
* an higher priority level attribute
* will replace lower priority level attribute.
* Must be bigger than 0.
+ * \param can_override Whether to explicitly allow
+ * overriding the attribute, any non-zero value
+ * implies allowance and 0 means disallowance.
*
* Cannot set with same plevel twice in the code.
*
* \tparam ValueType The type of the value to be set.
*/
template <typename ValueType>
inline OpRegEntry& set_attr(const std::string& attr_name, // NOLINT(*)
- const ValueType& value, int plevel = 10);
+ const ValueType& value, int plevel = 10, int can_override = 0);
Review comment:
Not sure if can_override flag is needed here.
Previously we check the plevel, if the plevel have higher priority level then we can override the attribute
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a change in pull request #7809: [Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass
Posted by GitBox <gi...@apache.org>.
zxybazh commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r617210922
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,43 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ target,
+ f=None,
+ level=10,
+ override=False,
+):
Review comment:
Ah I see. With default values things could be a lot easier. I would prefer to use `def register_intrin_lowering(op_name, target, *, f=None, level=10, override=False,)`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org