This is an automated email from the ASF dual-hosted git repository.
masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 62a5e7acf5 [Relay] Improve the "clip" op optimization in simplify expr
pass (#15068)
62a5e7acf5 is described below
commit 62a5e7acf537554e078a6d5924616be3d273a1dd
Author: 电线杆 <[email protected]>
AuthorDate: Thu Jun 15 12:02:50 2023 +0800
[Relay] Improve the "clip" op optimization in simplify expr pass (#15068)
* improve the simplification of clip op in simplify_expr
* 1. the arg of clip is forced to not be qnn.requantize and ann.add. This
is to avoid destroying the structure required by LegalizeQnnOpForDnnl; 2. code
reformatted.
* Fix the treatment of bfloat16 in function CheckDataTypeMaxMinValue. Add a
test case of bfloat16 in test_pass_simplify_expr.
* code reformatted
* fix a typo
---
src/relay/transforms/simplify_expr.cc | 54 ++++++++++--------
tests/python/relay/test_pass_simplify_expr.py | 82 +++++++++++++++++++++------
2 files changed, 96 insertions(+), 40 deletions(-)
diff --git a/src/relay/transforms/simplify_expr.cc
b/src/relay/transforms/simplify_expr.cc
index a557f2496b..fa3348b95a 100644
--- a/src/relay/transforms/simplify_expr.cc
+++ b/src/relay/transforms/simplify_expr.cc
@@ -161,17 +161,15 @@ class SimplifyConsecutiveCast : public DFPatternRewrite {
};
bool CheckDataTypeMaxMinValue(DataType dtype, double min_value, double
max_value) {
+ double lbound{}, ubound{};
if (dtype.is_int() || dtype.is_uint()) {
- double ubound =
static_cast<double>(Downcast<IntImm>(tvm::max_value(dtype))->value);
- double lbound =
static_cast<double>(Downcast<IntImm>(tvm::min_value(dtype))->value);
- return ubound == max_value && lbound == min_value;
- } else if (dtype.is_float()) {
- double ubound = Downcast<FloatImm>(tvm::max_value(dtype))->value;
- double lbound = Downcast<FloatImm>(tvm::min_value(dtype))->value;
- return ubound == max_value && lbound == min_value;
+ ubound =
static_cast<double>(Downcast<IntImm>(tvm::max_value(dtype))->value);
+ lbound =
static_cast<double>(Downcast<IntImm>(tvm::min_value(dtype))->value);
+ } else if (dtype.is_float() || dtype.is_bfloat16()) {
+ ubound = Downcast<FloatImm>(tvm::max_value(dtype))->value;
+ lbound = Downcast<FloatImm>(tvm::min_value(dtype))->value;
}
-
- return false;
+ return max_value >= ubound && min_value <= lbound;
}
/*!
@@ -224,8 +222,8 @@ class SimplifyClipAndConsecutiveCast : public
DFPatternRewrite {
};
/*!
- * \brief SimplifyCastClip matches the pattern cast->clip and remove redundant
Cast based on Clip
- * min/max values and min/max values of Cast target data type.
+ * \brief SimplifyClip removes redundant Clip based on its a_min/a_max values
and the min/max values
+ * of the data type.
*
* Example:
* %1 = cast(%0, dtype="uint8") [type=uint8]
@@ -234,30 +232,38 @@ class SimplifyClipAndConsecutiveCast : public
DFPatternRewrite {
* Optimized to (remove Clip):
* %1 = cast(%0, dtype="uint8") [type=uint8]
*/
-class SimplifyCastClip : public DFPatternRewrite {
+class SimplifyClip : public DFPatternRewrite {
public:
- SimplifyCastClip() {
- cast_ = IsOp("cast")({IsWildcard()});
- pattern_ = IsOp("clip")({cast_});
+ SimplifyClip() {
+ x_ = IsWildcard();
+ pattern_ = IsOp("clip")({x_});
}
Expr Callback(const Expr& pre, const Expr& post,
const Map<DFPattern, Array<Expr>>& node_map) const override {
- auto cast = Downcast<Call>(node_map[cast_][0]);
- DataType cast_dtype = Downcast<TensorType>(cast->checked_type())->dtype;
+ DataType cast_dtype = Downcast<TensorType>(pre->checked_type())->dtype;
- auto clip = Downcast<Call>(post);
- const CallNode* clip_node = clip.as<CallNode>();
+ const CallNode* clip_node = post.as<CallNode>();
const ClipAttrs* clip_attrs = clip_node->attrs.as<ClipAttrs>();
+ // TODO(kfeng123): For now, the arg of "clip" is forced to not be
"qnn.requantize" and
+ // "qnn.add". This is to avoid destroying the structure required by
LegalizeQnnOpForDnnl
+ auto child{post.as<CallNode>()->args[0].as<CallNode>()};
+ if (child && child->op.as<OpNode>()) {
+ String op_name{child->op.as<OpNode>()->name};
+ if (op_name == "qnn.requantize" || op_name == "qnn.add") {
+ return post;
+ }
+ }
+
if (CheckDataTypeMaxMinValue(cast_dtype, clip_attrs->a_min,
clip_attrs->a_max)) {
- return node_map[cast_][0];
+ return node_map[x_][0];
}
return post;
}
protected:
- DFPattern clip_, cast_;
+ DFPattern x_;
};
/*!
@@ -992,7 +998,7 @@ class SimplifyBinomial : public DFPatternRewrite {
DFPattern y_;
};
-/*! \brief Simplifying x/sqrt to x*sqrt */
+/*! \brief Simplifying x/sqrt to x*rsqrt */
class SimplifyRSqrt : public DFPatternRewrite {
public:
SimplifyRSqrt() {
@@ -1085,7 +1091,7 @@ Expr SimplifyExpr(const Expr& expr, const IRModule& mod) {
composer.AddRewrite<SimplifyDQArgMin>();
composer.AddRewrite<SimplifyDQArgSort>();
composer.AddRewrite<SimplifyClipAndConsecutiveCast>();
- composer.AddRewrite<SimplifyCastClip>();
+ composer.AddRewrite<SimplifyClip>();
composer.AddRewrite<SimplifyBinomial>();
return RewritePatterns(composer.MakeCallbacks(), expr, mod);
}
@@ -1099,7 +1105,7 @@ Expr SimplifyExprPostAlterOp(const Expr& expr, const
IRModule& mod) {
composer.AddRewrite<SimplifySameCast>();
composer.AddRewrite<SimplifyConsecutiveCast>();
composer.AddRewrite<SimplifyClipAndConsecutiveCast>();
- composer.AddRewrite<SimplifyCastClip>();
+ composer.AddRewrite<SimplifyClip>();
return RewritePatterns(composer.MakeCallbacks(), expr, mod);
}
diff --git a/tests/python/relay/test_pass_simplify_expr.py
b/tests/python/relay/test_pass_simplify_expr.py
index 4edb85f2d7..b117c91d1c 100644
--- a/tests/python/relay/test_pass_simplify_expr.py
+++ b/tests/python/relay/test_pass_simplify_expr.py
@@ -722,34 +722,84 @@ def test_simplify_dq_argsort():
def test_simplify_clip_cast():
- x = relay.var("x", shape=(4, 8), dtype="int32")
+ def before1():
+ x = relay.var("x", shape=(4, 8), dtype="int32")
+ clip = relay.clip(x, a_min=0.0, a_max=255.0)
+ cast = relay.cast(clip, "uint8")
+ cast = relay.cast(cast, "int32")
+ return relay.Function([x], cast)
- def before():
+ def expected1():
+ x = relay.var("x", shape=(4, 8), dtype="int32")
+ clip = relay.clip(x, a_min=0.0, a_max=255.0)
+ return relay.Function([x], clip)
+
+ def before2():
+ x = relay.var("x", shape=(4, 8), dtype="int32")
clip = relay.clip(x, a_min=0.0, a_max=255.0)
cast = relay.cast(clip, "uint8")
- return relay.cast(cast, "int32")
+ cast = relay.cast(cast, "int32")
+ return relay.Function([x], cast)
- def expected():
- return relay.clip(x, a_min=0.0, a_max=255.0)
+ def expected2():
+ x = relay.var("x", shape=(4, 8), dtype="int32")
+ clip = relay.clip(x, a_min=0.0, a_max=255.0)
+ return relay.Function([x], clip)
- opt = run_opt_pass(before(), transform.SimplifyExpr())
- ref = run_infer_type(expected())
- assert tvm.ir.structural_equal(opt, ref)
+ for before, expected in [
+ [before1(), expected1()],
+ [before2(), expected2()],
+ ]:
+ after = run_opt_pass(before, transform.SimplifyExpr())
+ expected = run_opt_pass(expected, transform.InferType())
+ assert tvm.ir.structural_equal(after, expected), "\nafter: {}
\nexpected: {}".format(
+ after, expected
+ )
def test_simplify_cast_clip():
- x = relay.var("x", shape=(4, 8), dtype="int32")
+ def before1():
+ x = relay.var("x", shape=(4, 8), dtype="int32")
+ cast = relay.cast(x, "uint8")
+ clip = relay.clip(cast, a_min=0.0, a_max=255.0)
+ return relay.Function([x], clip)
- def before():
+ def expected1():
+ x = relay.var("x", shape=(4, 8), dtype="int32")
cast = relay.cast(x, "uint8")
- return relay.clip(cast, a_min=0.0, a_max=255.0)
+ return relay.Function([x], cast)
- def expected():
- return relay.cast(x, "uint8")
+ def before2():
+ x = relay.var("x", shape=(4, 8), dtype="uint8")
+ clip = relay.clip(x, a_min=0.0, a_max=255.0)
+ return relay.Function([x], clip)
- opt = run_opt_pass(before(), transform.SimplifyExpr())
- ref = run_infer_type(expected())
- assert tvm.ir.structural_equal(opt, ref)
+ def expected2():
+ x = relay.var("x", shape=(4, 8), dtype="uint8")
+ return relay.Function([x], x)
+
+ def before3():
+ x = relay.var("x", shape=(4, 8), dtype="float32")
+ cast = relay.cast(x, "bfloat16")
+ clip = relay.clip(cast, a_min=-0.2, a_max=0.4)
+ return relay.Function([x], clip)
+
+ def expected3():
+ x = relay.var("x", shape=(4, 8), dtype="float32")
+ cast = relay.cast(x, "bfloat16")
+ clip = relay.clip(cast, a_min=-0.2, a_max=0.4)
+ return relay.Function([x], clip)
+
+ for before, expected in [
+ [before1(), expected1()],
+ [before2(), expected2()],
+ [before3(), expected3()],
+ ]:
+ after = run_opt_pass(before, transform.SimplifyExpr())
+ expected = run_opt_pass(expected, transform.InferType())
+ assert tvm.ir.structural_equal(after, expected), "\nafter: {}
\nexpected: {}".format(
+ after, expected
+ )
def test_simplify_add():
