Skip to content

Commit 6c985db

Browse files
vinx13vinx13
committed
[Relay][Pass] CanonicalizeExpr
1 parent 7c1c97d commit 6c985db

File tree

5 files changed

+210
-0
lines changed

5 files changed

+210
-0
lines changed

‎include/tvm/relay/transform.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -531,6 +531,13 @@ TVM_DLL Pass CanonicalizeOps();
531531
*/
532532
TVM_DLL Pass AlterOpLayout();
533533

534+
/*!
535+
* \brief Canonicalize an expression to make operator fusion more efficient.
536+
*
537+
* \return The pass.
538+
*/
539+
TVM_DLL Pass CanonicalizeExpr();
540+
534541
} // namespace transform
535542
} // namespace relay
536543
} // namespace tvm

‎python/tvm/relay/transform.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -593,3 +593,14 @@ def PartialEvaluate():
593593
The registered pass that performs partial evaluation on an expression.
594594
"""
595595
return _transform.PartialEvaluate()
596+
597+
def CanonicalizeExpr():
598+
"""
599+
Canonicalize an expression to make operator fusion more efficient.
600+
601+
Returns
602+
-------
603+
ret : tvm.relay.Pass
604+
The registered pass that canonicalizes an expression.
605+
"""
606+
return _transform.CanonicalizeExpr()

‎src/relay/backend/build_module.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,7 @@ class RelayBuildModule : public runtime::ModuleNode {
299299
pass_seqs.push_back(transform::CombineParallelConv2D(3));
300300
pass_seqs.push_back(transform::FoldConstant());
301301
pass_seqs.push_back(transform::FoldScaleAxis());
302+
pass_seqs.push_back(transform::CanonicalizeExpr());
302303
pass_seqs.push_back(transform::CanonicalizeOps());
303304

304305
// Alter layout transformation is only applied to homogeneous execution yet.

‎src/relay/pass/canonicalize_expr.cc‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
/*!
21+
* Copyright (c) 2019 by Contributors
22+
* \file canonicalize_expr.cc
23+
* \brief Canonicalize an expression to make operator fusion more efficient.
24+
*/
25+
#include <tvm/relay/pass.h>
26+
#include <tvm/relay/expr_functor.h>
27+
#include <tvm/relay/attrs/nn.h>
28+
#include <tvm/relay/transform.h>
29+
#include "pattern_util.h"
30+
#include "pass_util.h"
31+
32+
namespace tvm {
33+
namespace relay {
34+
35+
// This pass finds upcast that is referred by multiple elemwise/broadcast operators, and creates a
36+
// copy of it in each branch such that after fusion the previous function have output with fewer
37+
// bits.
38+
class ExprCanonicalizer : public ExprMutator {
39+
public:
40+
Expr VisitExpr_(const CallNode* call) {
41+
static auto fpattern = Op::GetAttr<TOpPattern>("TOpPattern");
42+
43+
if (const OpNode* opnode = call->op.as<OpNode>()) {
44+
auto pattern = fpattern[GetRef<Op>(opnode)];
45+
if (pattern <= kBroadcast) {
46+
Array<Expr> call_args = call->args;
47+
bool unchanged = true;
48+
for (size_t i = 0; i < call_args.size(); ++i) {
49+
Expr arg = call_args[i];
50+
Expr new_arg = GetNewCallArg(arg);
51+
if (!arg.same_as(new_arg)) {
52+
call_args.Set(i, new_arg);
53+
unchanged = false;
54+
}
55+
}
56+
if (unchanged) {
57+
return GetRef<Expr>(call);
58+
}
59+
return CallNode::make(call->op, call_args, call->attrs, call->type_args);
60+
}
61+
}
62+
63+
Expr new_expr = ExprMutator::VisitExpr_(call);
64+
return new_expr;
65+
}
66+
67+
private:
68+
std::unordered_map<const Node*, size_t> ref_counter_;
69+
70+
Expr GetNewCallArg(const Expr& e) {
71+
// if e is a upcast and ref count > 1, create an copy; otherwise call the default visitor
72+
73+
static auto& cast = Op::Get("cast");
74+
Expr new_expr = this->VisitExpr(e);
75+
76+
if (const CallNode* call = e.as<CallNode>()) {
77+
if (call->op.same_as(cast)) {
78+
auto attrs = call->attrs.as<CastAttrs>();
79+
const auto* from_type = call->args[0]->type_as<TensorTypeNode>();
80+
CHECK(from_type);
81+
82+
if (from_type->dtype.bits() < attrs->dtype.bits()) {
83+
if (++ref_counter_[call] > 1) {
84+
const CallNode* new_call = new_expr.as<CallNode>();
85+
CHECK(new_call);
86+
CHECK(new_call->op.same_as(cast));
87+
return CallNode::make(new_call->op, new_call->args, new_call->attrs,
88+
new_call->type_args);
89+
}
90+
}
91+
}
92+
}
93+
return new_expr;
94+
}
95+
};
96+
97+
Expr CanonicalizeExpr(const Expr& e) {
98+
return ExprCanonicalizer().Mutate(e);
99+
}
100+
101+
TVM_REGISTER_API("relay._ir_pass.canonicalize_expr")
102+
.set_body_typed(CanonicalizeExpr);
103+
104+
namespace transform {
105+
106+
Pass CanonicalizeExpr() {
107+
runtime::TypedPackedFunc<Function(Function, Module, PassContext)> pass_func =
108+
[=](Function f, Module m, PassContext pc) {
109+
return Downcast<Function>(CanonicalizeExpr(f));
110+
};
111+
return CreateFunctionPass(pass_func, 3, "CanonicalizeExpr",
112+
{ir::StringImm::make("InferType")});
113+
}
114+
115+
TVM_REGISTER_API("relay._transform.CanonicalizeExpr")
116+
.set_body_typed(CanonicalizeExpr);
117+
118+
} // namespace transform
119+
120+
} // namespace relay
121+
} // namespace tvm

‎tests/python/relay/test_pass_canonicalize_expr.py‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
import tvm
19+
import tvm.relay as relay
20+
import tvm.relay.module as _module
21+
import tvm.relay.transform as _transform
22+
23+
24+
def test_canonicalize_cast():
25+
def before(data, conv_weight, bias1, bias2):
26+
x = relay.nn.conv2d(data, conv_weight,
27+
channels=16,
28+
kernel_size=(3, 3),
29+
padding=(1, 1),
30+
out_dtype="int8")
31+
x1 = relay.cast(x, dtype="int32")
32+
y1 = relay.add(x1, bias1)
33+
y2 = relay.add(x1, bias2)
34+
y = relay.add(y1, y2)
35+
return relay.Function([data, conv_weight, bias1, bias2], y)
36+
37+
def expected(data, conv_weight, bias1, bias2):
38+
x = relay.nn.conv2d(data, conv_weight,
39+
channels=16,
40+
kernel_size=(3, 3),
41+
padding=(1, 1),
42+
out_dtype="int8")
43+
x1 = relay.cast(x, dtype="int32")
44+
x2 = relay.cast(x, dtype="int32")
45+
y1 = relay.add(x1, bias1)
46+
y2 = relay.add(x2, bias2)
47+
y = relay.add(y1, y2)
48+
return relay.Function([data, conv_weight, bias1, bias2], y)
49+
50+
def check(shape):
51+
data = relay.var("data", shape=shape, dtype="int8")
52+
conv_weight = relay.var("weight")
53+
bias1 = relay.var("bias1", shape=(16, 1, 1), dtype="int32")
54+
bias2 = relay.var("bias2", shape=(16, 1, 1), dtype="int32")
55+
y = before(data, conv_weight, bias1, bias2)
56+
mod = _module.Module.from_expr(y)
57+
seq = _transform.Sequential([_transform.InferType(), _transform.CanonicalizeExpr(),
58+
_transform.InferType()])
59+
with _transform.PassContext(opt_level=3):
60+
mod = seq(mod)
61+
y = mod[mod.entry_func.name_hint]
62+
y_expected = expected(data, conv_weight, bias1, bias2)
63+
y_expected = relay.ir_pass.infer_type(y_expected)
64+
assert relay.ir_pass.alpha_equal(y, y_expected)
65+
66+
check((1, 16, 7, 7))
67+
68+
69+
if __name__ == '__main__':
70+
test_canonicalize_cast()

0 commit comments

Comments
 (0)