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:
[email protected]
