icemelon9 commented on a change in pull request #5357:
URL: https://github.com/apache/incubator-tvm/pull/5357#discussion_r418430603
##########
File path: python/tvm/relay/op/strategy/x86.py
##########
@@ -18,13 +18,17 @@
# pylint:
disable=invalid-name,unused-argument,wildcard-import,unused-wildcard-import
import logging
+import re
import topi
from tvm.te import SpecializedCondition
from .generic import *
from .. import op as _op
logger = logging.getLogger('strategy')
+_NCHWc_matcher = re.compile("^NCHW[-+]?[0-9]+c$")
Review comment:
Why do we need `[-+]?` here?
##########
File path: python/tvm/relay/op/strategy/x86.py
##########
@@ -84,8 +88,13 @@ def conv2d_strategy_cpu(attrs, inputs, out_type, target):
raise ValueError("dilation should be positive value")
if groups == 1:
- if layout == "NCHW":
- assert kernel_layout == "OIHW"
+ if layout.startswith("NCHW"):
Review comment:
Could you separate "NCHW" and "NCHWc" into two different if branches?
##########
File path: src/relay/op/tensor/transform.h
##########
@@ -36,6 +36,9 @@
namespace tvm {
namespace relay {
+extern Expr MakeReshape(Expr data,
+ Array<Integer> newshape);
Review comment:
fix indent
##########
File path: src/relay/transforms/fold_scale_axis.cc
##########
@@ -390,8 +430,10 @@ Expr AddSubForwardRewrite(const Call& ref_call,
} else {
CHECK(srhs != nullptr);
CHECK(MatchBroadcastToLeftAxes(trhs, tlhs, srhs->axes));
- Expr scale = ExpandBiasToMatchAxis(
- srhs->scale, trhs->shape.size(), srhs->axes);
+ Expr scale = ReshapeOrExpandToMatchAxis(
+ srhs->scale, trhs->shape, srhs->axes);
+ if (!scale.defined())
+ return Expr();
Review comment:
same here
##########
File path: src/relay/transforms/fold_scale_axis.cc
##########
@@ -380,8 +418,10 @@ Expr AddSubForwardRewrite(const Call& ref_call,
if (slhs != nullptr) {
CHECK(srhs == nullptr);
CHECK(MatchBroadcastToLeftAxes(tlhs, trhs, slhs->axes));
- Expr scale = ExpandBiasToMatchAxis(
- slhs->scale, tlhs->shape.size(), slhs->axes);
+ Expr scale = ReshapeOrExpandToMatchAxis(
+ slhs->scale, tlhs->shape, slhs->axes);
+ if (!scale.defined())
+ return Expr();
Review comment:
still use { } to scope the if condition
##########
File path: src/relay/transforms/fold_scale_axis.cc
##########
@@ -314,6 +316,42 @@ class ForwardPrep : private ExprVisitor {
}
};
+static bool IsIntInArray(const Array<Integer>& axis, int v) {
+ for (size_t i = 0; i < axis.size(); i++) {
+ if (axis[i] == v)
+ return true;
+ }
+ return false;
+}
+
+static Expr ReshapeToMatchAxis(Expr scale, const Array<PrimExpr>& shape,
+ const Array<Integer>& axis) {
+ Array<Integer> arr;
+ for (size_t i = 0; i < shape.size(); i++) {
+ if (IsIntInArray(axis, i)) {
+ auto node = shape[i].as<IntImmNode>();
+ if (!node) {
+ // if the shape is not a constant, use normal transform
+ return Expr();
+ }
+ arr.push_back(node->value);
+ } else {
+ arr.push_back(1);
+ }
+ }
+ return MakeReshape(scale, std::move(arr));
+}
+
+// if only one axis, use expand dim. Else, use reshape
+static Expr ReshapeOrExpandToMatchAxis(Expr scale, const Array<PrimExpr>&
shape,
+ const Array<Integer>& axis) {
Review comment:
fix indent
----------------------------------------------------------------
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]
