#include <tvm/relay/transform.h>
#include <tvm/tir/data_layout.h>
+#include "../op/tensor/transform.h"
#include "pass_util.h"
#include "pattern_util.h"
*
* Use namespace to reduce potential naming conflict.
*/
+
namespace fold_scale_axis {
using runtime::TypedPackedFunc;
}
};
+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) {
+ if (axis.size() > 1) {
+ return ReshapeToMatchAxis(scale, shape, axis);
+ } else {
+ return ExpandBiasToMatchAxis(scale, shape.size(), axis);
+ }
+}
+
//----------------------------------------------
// Per operator defs for FScaleAxisForward
//----------------------------------------------
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();
+ }
Expr rhs = Divide(new_args[1], scale);
rnode->value = Call(ref_call->op, {slhs->value, rhs}, ref_call->attrs, ref_call->type_args);
rnode->scale = slhs->scale;
} 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();
+ }
Expr lhs = Divide(new_args[0], scale);
rnode->value = Call(ref_call->op, {lhs, srhs->value}, ref_call->attrs, ref_call->type_args);
rnode->scale = srhs->scale;
CHECK_GE(c_big_axis, 0);
Message none = NullValue<Message>();
- AxesSet data_axes = NullValue<AxesSet>();
// For now, we only support simple pattern (no folded weight/data)
// More general layout can be supported under the current framework.
// By using a unified layout transformation.
// only handle depthwise or full conv2d.
// TODO(tvm-team) handle grouped conv by reshape + bcast
bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, kernel_layout);
- if (kernel_layout.IndexOf(LayoutAxis::Get('i')) < 0 && c_small_axis < 0 &&
- (param->groups == 1 || is_depthwise_conv2d)) {
- data_axes = {c_big_axis};
- }
- if (data_axes.defined()) {
- return {Message(data_axes, false), none};
+ if (param->groups == 1 || is_depthwise_conv2d) {
+ auto ko_small_axis = kernel_layout.IndexOf(LayoutAxis::Get('o'));
+ auto ki_small_axis = kernel_layout.IndexOf(LayoutAxis::Get('i'));
+ if ((ko_small_axis < 0 && ki_small_axis < 0 && c_small_axis < 0) || // simple layout
+ (ko_small_axis >= 0 && ki_small_axis >= 0 && c_small_axis >= 0)) { // blocked layout
+ Array<Integer> arr{c_big_axis};
+ if (c_small_axis >= 0) {
+ arr.push_back(c_small_axis);
+ }
+ return {Message(arr, false), none};
+ }
}
return {none, none};
}
Layout kernel_layout(param->kernel_layout);
int c_big_axis = data_layout.IndexOf(LayoutAxis::Get('C'));
CHECK_GE(c_big_axis, 0);
- // For now, we only support simple pattern (no folded weight/data)
- // TODO(tvm-team) support general data layout
- CHECK_EQ(kernel_layout.IndexOf(LayoutAxis::Get('i')), -1);
- CHECK(sdata->axes.size() == 1 && c_big_axis == sdata->axes[0]->value);
- int big_oc_axis = kernel_layout.IndexOf(LayoutAxis::Get('O'));
- int big_ic_axis = kernel_layout.IndexOf(LayoutAxis::Get('I'));
+ int small_ko_axis = kernel_layout.IndexOf(LayoutAxis::Get('o'));
+ int small_ki_axis = kernel_layout.IndexOf(LayoutAxis::Get('i'));
+ int big_ki_axis = kernel_layout.IndexOf(LayoutAxis::Get('I'));
+ int big_ko_axis = kernel_layout.IndexOf(LayoutAxis::Get('O'));
+
+ bool is_simple = (small_ko_axis < 0 && small_ki_axis < 0 && big_ki_axis >= 0);
+ bool is_blocking = (small_ko_axis >= 0 && small_ki_axis >= 0 && big_ki_axis >= 0);
+ CHECK(is_simple || is_blocking);
// Check it must be depthwise or full conv2d.
bool is_depthwise_conv2d = IsDepthwiseConv2D(ref_call, param, kernel_layout);
// match the ic_axis
if (is_depthwise_conv2d) {
- Expr scale = ExpandBiasToMatchAxis(sdata->scale, kernel_layout.ndim(), {big_oc_axis});
- weight = Multiply(weight, scale);
+ if (is_simple) {
+ Expr scale = ExpandBiasToMatchAxis(sdata->scale, kernel_layout.ndim(), {big_ko_axis});
+ weight = Multiply(weight, scale);
+ } else {
+ weight = Multiply(weight,
+ ReshapeToMatchAxis(sdata->scale, weight->type_as<TensorTypeNode>()->shape,
+ {big_ko_axis, small_ko_axis}));
+ if (!weight.defined()) return Expr();
+ }
+
} else {
- Expr scale = ExpandBiasToMatchAxis(sdata->scale, kernel_layout.ndim(), {big_ic_axis});
- weight = Multiply(weight, scale);
+ if (is_simple) {
+ Expr scale = ExpandBiasToMatchAxis(sdata->scale, kernel_layout.ndim(), {big_ki_axis});
+ weight = Multiply(weight, scale);
+ } else {
+ weight = Multiply(weight,
+ ReshapeToMatchAxis(sdata->scale, weight->type_as<TensorTypeNode>()->shape,
+ {big_ki_axis, small_ki_axis}));
+ if (!weight.defined()) return Expr();
+ }
}
// return transformed conv2d
return Call(ref_call->op, {sdata->value, weight}, ref_call->attrs, ref_call->type_args);
CHECK(equal(message->axes, lhs_message->axes));
Expr lhs = transformer->Transform(call->args[0], message, scale);
Expr rhs = transformer->Transform(call->args[1], NullValue<Message>(), NullValue<Expr>());
- Expr rhs_scale = ExpandBiasToMatchAxis(scale, tlhs->shape.size(), message->axes);
+ Expr rhs_scale = ReshapeOrExpandToMatchAxis(scale, tlhs->shape, message->axes);
+ if (!rhs_scale.defined()) {
+ return transformer->NormalCallTransform(call.operator->());
+ }
rhs = Multiply(rhs, rhs_scale);
return Call(call->op, {lhs, rhs}, call->attrs, call->type_args);
} else if (rhs_message.defined()) {
CHECK(equal(message->axes, rhs_message->axes));
Expr lhs = transformer->Transform(call->args[0], NullValue<Message>(), NullValue<Expr>());
Expr rhs = transformer->Transform(call->args[1], message, scale);
- Expr lhs_scale = ExpandBiasToMatchAxis(scale, trhs->shape.size(), message->axes);
+ Expr lhs_scale = ReshapeOrExpandToMatchAxis(scale, trhs->shape, message->axes);
+ if (!lhs_scale.defined()) {
+ return transformer->NormalCallTransform(call.operator->());
+ }
lhs = Multiply(lhs, lhs_scale);
return Call(call->op, {lhs, rhs}, call->attrs, call->type_args);
} else {
// only handle depthwise or full conv2d.
// TODO(tvm-team) handle grouped conv by reshape + bcast
bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, kernel_layout);
- if (kernel_layout.IndexOf(LayoutAxis::Get('o')) < 0 &&
- kernel_layout.IndexOf(LayoutAxis::Get('i')) < 0 && c_small_axis < 0 &&
- (param->groups == 1 || is_depthwise_conv2d)) {
- return Message({c_big_axis}, false);
- } else {
- return NullValue<Message>();
+ if (param->groups == 1 || is_depthwise_conv2d) {
+ auto ko_small_axis = kernel_layout.IndexOf(LayoutAxis::Get('o'));
+ auto ki_small_axis = kernel_layout.IndexOf(LayoutAxis::Get('i'));
+ if ((ko_small_axis < 0 && ki_small_axis < 0 && c_small_axis < 0) || // simple layout
+ (ko_small_axis >= 0 && ki_small_axis >= 0 && c_small_axis >= 0)) { // blocked layout
+ Array<Integer> arr{c_big_axis};
+ if (c_small_axis >= 0) {
+ arr.push_back(c_small_axis);
+ }
+ return Message(arr, false);
+ }
}
+ return NullValue<Message>();
}
// Conv2D consumes the scale axis during transformation.
CHECK_GE(c_big_axis, 0);
// For now, we only support simple pattern (no folded weight/data)
// TODO(tvm-team) support general data layout
- CHECK_EQ(kernel_layout.IndexOf(LayoutAxis::Get('o')), -1);
- CHECK_EQ(kernel_layout.IndexOf(LayoutAxis::Get('i')), -1);
- CHECK(message->axes.size() == 1 && c_big_axis == message->axes[0]->value);
-
- int big_oc_axis = kernel_layout.IndexOf(LayoutAxis::Get('O'));
+ int small_ko_axis = kernel_layout.IndexOf(LayoutAxis::Get('o'));
+ int small_ki_axis = kernel_layout.IndexOf(LayoutAxis::Get('i'));
+ int big_ki_axis = kernel_layout.IndexOf(LayoutAxis::Get('I'));
+ int big_ko_axis = kernel_layout.IndexOf(LayoutAxis::Get('O'));
// Check it must be depthwise or full conv2d.
bool is_depthwise_conv2d = IsDepthwiseConv2D(call, param, kernel_layout);
CHECK(param->groups == 1 || is_depthwise_conv2d);
+ bool is_simple = (small_ko_axis < 0 && small_ki_axis < 0 && big_ki_axis >= 0);
+ bool is_blocking = (small_ko_axis >= 0 && small_ki_axis >= 0 && big_ki_axis >= 0);
+ CHECK(is_simple || is_blocking);
Expr data = transformer->Transform(call->args[0], NullValue<Message>(), NullValue<Expr>());
Expr weight = transformer->Transform(call->args[1], NullValue<Message>(), NullValue<Expr>());
// scale on input for deptwise.
- Expr wscale = ExpandBiasToMatchAxis(scale, kernel_layout.ndim(), {big_oc_axis});
+ Expr wscale;
+ if (is_simple) {
+ wscale = ExpandBiasToMatchAxis(scale, kernel_layout.ndim(), {big_ko_axis});
+ } else {
+ wscale = ReshapeToMatchAxis(scale, weight->type_as<TensorTypeNode>()->shape,
+ {big_ko_axis, small_ko_axis});
+ if (!wscale.defined()) return transformer->NormalCallTransform(call.operator->());
+ }
weight = Multiply(weight, wscale);
return Call(call->op, {data, weight}, call->attrs, call->type_args);
}
def test_fold_fwd_simple():
"""Simple testcase."""
- def before(x, conv_weight, in_bias, in_scale, channels):
+ def before(x, conv_weight, in_bias, in_scale, channels, blocking):
args = [x, conv_weight, in_bias]
- in_bias = relay.expand_dims(in_bias, axis=1, num_newaxis=2)
x = relay.multiply(x, in_scale)
x = relay.nn.relu(x)
x = relay.add(x, in_bias)
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW2i{}o".format(blocking[1]) if blocking else "OIHW")
return relay.Function(args, y)
- def expected(x, conv_weight, in_bias, in_scale, channels):
+ def expected(x, conv_weight, in_bias, in_scale, in_channels, channels, blocking):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight, in_bias]
- in_bias = relay.expand_dims(in_bias, axis=1, num_newaxis=2)
- squeezed_scale = relay.squeeze(in_scale, axis=[1,2])
- x = relay.nn.relu(x)
- in_bias = relay.divide(in_bias, relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
- x = relay.add(x, in_bias)
- conv_weight = relay.multiply(
- conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
+ if blocking:
+ squeezed_scale = relay.squeeze(in_scale, axis=[0,2,3])
+ x = relay.nn.relu(x)
+ in_bias = relay.divide(in_bias,
+ relay.reshape(squeezed_scale, (1, in_channels // blocking[0], 1, 1, blocking[0]))) #NCHWc
+ x = relay.add(x, in_bias)
+ conv_weight = relay.multiply(conv_weight,
+ relay.reshape(squeezed_scale, (1, in_channels//2, 1, 1, 2, 1))) #OIHWio
+ else:
+ squeezed_scale = relay.squeeze(in_scale, axis=[1,2])
+ x = relay.nn.relu(x)
+ in_bias = relay.divide(in_bias, relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
+ x = relay.add(x, in_bias)
+ conv_weight = relay.multiply(
+ conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
+
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW2i{}o".format(blocking[1]) if blocking else "OIHW")
return relay.Function(args, y)
- def check(shape, channels):
+ def check(shape, channels, blocking):
x = relay.var("x", shape=shape)
- in_channels = shape[1]
weight = relay.var("weight")
- in_bias = relay.var("in_bias", shape=(in_channels,))
- in_scale = relay.const(_get_positive_scale((in_channels, 1, 1)))
- y1 = before(x, weight, in_bias, in_scale, channels)
+ if blocking:
+ in_channels = shape[1] * shape[4]
+ in_bias = relay.var("in_bias", shape=(1, in_channels // blocking[0], 1, 1, blocking[0]))
+ in_scale = relay.const(_get_positive_scale((1, in_channels // blocking[0], 1, 1, blocking[0])))
+ else:
+ in_channels = shape[1]
+ in_bias = relay.var("in_bias", shape=(in_channels, 1, 1))
+ in_scale = relay.const(_get_positive_scale((in_channels, 1, 1)))
+ y1 = before(x, weight, in_bias, in_scale, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
type_dict = {x.name_hint:x.checked_type for x in y1.params}
weight = relay.var("weight", type_dict["weight"])
y1_folded = run_opt_pass(y1, transform.ForwardFoldScaleAxis())
- y1_expected = expected(x, weight, in_bias, in_scale, channels)
+ y1_expected = expected(x, weight, in_bias, in_scale, in_channels, channels, blocking)
y1_folded = run_opt_pass(y1_folded, transform.InferType())
y1_expected = run_opt_pass(y1_expected, transform.InferType())
assert tvm.ir.structural_equal(y1_folded, y1_expected)
- check((2, 4, 10, 10), 2)
-
+ check((2, 4, 10, 10), 2, None)
+ check((2, 2, 10, 10, 2), 8, (2, 4))
def test_fold_fwd_dual_path():
"""scale axis being consumed by two consumers"""
- def before(x, conv_weight, in_bias, in_scale, channels):
+ def before(x, conv_weight, in_bias, in_scale, channels, blocking):
args = [x, conv_weight, in_bias]
x = relay.multiply(in_scale, x)
x = relay.nn.relu(x)
y1 = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- data_layout="NHWC",
- kernel_layout="HWIO",
+ data_layout="NHWC{}c".format(blocking[0]) if blocking else "NHWC",
+ kernel_layout="HWIO1i{}o".format(blocking[1]) if blocking else "HWIO",
groups=channels,
padding=(1, 1))
y2 = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- data_layout="NHWC",
- kernel_layout="HWIO",
+ data_layout="NHWC{}c".format(blocking[0]) if blocking else "NHWC",
+ kernel_layout="HWIO1i{}o".format(blocking[1]) if blocking else "HWIO",
groups=channels,
padding=(1, 1))
z = relay.add(y1, y2)
return relay.Function(args, z)
- def expected(x, conv_weight, in_bias, in_scale, channels):
+ def expected(x, conv_weight, in_bias, in_scale, channels, blocking):
args = [x, conv_weight, in_bias]
x = relay.nn.relu(x)
- in_bias = relay.divide(in_bias, in_scale)
+ if blocking:
+ _in_scale = relay.reshape(in_scale, (1, 1, 1, channels//blocking[0], blocking[0])) #NHWCc
+ else:
+ _in_scale = in_scale
+ in_bias = relay.divide(in_bias, _in_scale)
x = relay.subtract(x, in_bias)
+ if blocking:
+ _in_scale = relay.reshape(in_scale, (1, 1, 1, channels//blocking[0], 1, blocking[0])) #HWIOio
y1 = relay.nn.conv2d(x,
- relay.multiply(conv_weight, in_scale),
+ relay.multiply(conv_weight, _in_scale),
channels=channels,
kernel_size=(3, 3),
- data_layout="NHWC",
- kernel_layout="HWIO",
+ data_layout="NHWC{}c".format(blocking[0]) if blocking else "NHWC",
+ kernel_layout="HWIO1i{}o".format(blocking[1]) if blocking else "HWIO",
groups=channels,
padding=(1, 1))
+ if blocking:
+ _in_scale = relay.reshape(in_scale, (1, 1, 1, channels//blocking[0], 1, blocking[0])) #HWIOio
y2 = relay.nn.conv2d(x,
- relay.multiply(conv_weight, in_scale),
+ relay.multiply(conv_weight, _in_scale),
channels=channels,
kernel_size=(3, 3),
- data_layout="NHWC",
- kernel_layout="HWIO",
+ data_layout="NHWC{}c".format(blocking[0]) if blocking else "NHWC",
+ kernel_layout="HWIO1i{}o".format(blocking[1]) if blocking else "HWIO",
groups=channels,
padding=(1, 1))
z = relay.add(y1, y2)
return relay.Function(args, z)
- def check(dshape, channels):
+ def check(dshape, channels, blocking):
x = relay.var("x", shape=dshape)
- in_channels = dshape[-1]
+ if blocking:
+ in_channels = dshape[3] * dshape[4]
+ wshape = (3, 3, 1, channels//blocking[1], 1, blocking[1]) # HWIOio
+ weight = relay.var("weight", shape=wshape)
+ in_bias = relay.var("in_bias", shape=(in_channels//blocking[0],blocking[0]))
+ in_scale = relay.const(_get_positive_scale((in_channels//blocking[0],blocking[0])))
+ else:
+ in_channels = dshape[-1]
+ wshape = (3, 3, 1, channels) # HWIO
+ weight = relay.var("weight", shape=wshape)
+ in_bias = relay.var("in_bias", shape=(in_channels,))
+ in_scale = relay.const(_get_positive_scale(in_channels,))
+
# test depthwise
assert in_channels == channels
- wshape = (3, 3, 1, channels) # HWIO
- weight = relay.var("weight", shape=wshape)
- in_bias = relay.var("in_bias", shape=(in_channels,))
- in_scale = relay.const(_get_positive_scale(in_channels,))
- y1 = before(x, weight, in_bias, in_scale, channels)
+
+ y1 = before(x, weight, in_bias, in_scale, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
y1_folded = run_opt_pass(y1, transform.ForwardFoldScaleAxis())
type_dict = {x.name_hint:x.checked_type for x in y1.params}
weight = relay.var("weight", type_dict["weight"])
- y1_expected = expected(x, weight, in_bias, in_scale, channels)
+ y1_expected = expected(x, weight, in_bias, in_scale, channels, blocking)
y1_expected = run_opt_pass(y1_expected, transform.InferType())
assert tvm.ir.structural_equal(y1_folded, y1_expected)
- check((2, 4, 10, 3), 3)
-
+ check((2, 4, 10, 3), 3, None)
+ check((2, 4, 10, 2, 2), 4, (2, 2))
def test_fold_fwd_fail():
"""testcase where we canont fold"""
- def before(x, conv_weight, in_bias, in_scale, channels):
+ def before(x, conv_weight, in_bias, in_scale, channels, blocking):
x = relay.multiply(x, in_scale)
xx = relay.nn.leaky_relu(x, alpha=0.1)
y1 = relay.nn.conv2d(xx, conv_weight,
channels=channels,
kernel_size=(3, 3),
- data_layout="NHWC",
+ data_layout="NHWC{}c".format(blocking[0]) if blocking else "NHWC",
+ kernel_layout="HWIO1i{}o".format(blocking[1]) if blocking else "HWIO",
padding=(1, 1))
z = relay.add(y1, x)
return relay.Function(relay.analysis.free_vars(z), z)
- def check(shape, channels):
+ def check(shape, channels, blocking):
x = relay.var("x", shape=shape)
- in_channels = shape[-1]
+ if blocking:
+ in_channels = shape[3] * shape[4]
+ in_bias = relay.var("in_bias", shape=(in_channels//blocking[0],blocking[0]))
+ in_scale = relay.const(_get_positive_scale((in_channels//blocking[0],blocking[0])))
+ else:
+ in_channels = shape[-1]
+ in_bias = relay.var("in_bias", shape=(in_channels,))
+ in_scale = relay.const(_get_positive_scale(size=(in_channels,)))
# test depthwise
assert in_channels == channels
weight = relay.var("weight")
- in_bias = relay.var("in_bias", shape=(in_channels,))
- in_scale = relay.const(_get_positive_scale(size=(in_channels,)))
- y1 = before(x, weight, in_bias, in_scale, channels)
+ y1 = before(x, weight, in_bias, in_scale, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
y1_folded = run_opt_pass(y1, transform.ForwardFoldScaleAxis())
assert tvm.ir.structural_equal(y1, y1_folded)
- check((2, 11, 10, 4), 4)
-
+ check((2, 11, 10, 4), 4, None)
+ check((2, 11, 10, 2, 2), 4, (2,2))
def test_fold_fwd_relu_fail():
"""testcase where we canont fold because scale can not pass relu"""
- def before(x, conv_weight, in_bias, in_scale, channels):
+ def before(x, conv_weight, in_bias, in_scale, channels, blocking):
x = relay.multiply(x, in_scale)
xx = relay.nn.relu(x)
y1 = relay.nn.conv2d(xx, conv_weight,
channels=channels,
kernel_size=(3, 3),
- data_layout="NHWC",
+ data_layout="NHWC{}c".format(blocking[0]) if blocking else "NHWC",
+ kernel_layout="HWIO1i{}o".format(blocking[1]) if blocking else "HWIO",
padding=(1, 1))
z = relay.add(y1, x)
return relay.Function(relay.analysis.free_vars(z), z)
- def check(shape, channels, in_scale):
+ def check(shape, channels, blocking, in_scale):
x = relay.var("x", shape=shape)
- in_channels = shape[-1]
- # test depthwise
- assert in_channels == channels
weight = relay.var("weight")
- in_bias = relay.var("in_bias", shape=(in_channels,))
- y1 = before(x, weight, in_bias, in_scale, channels)
+ if blocking:
+ in_channels = shape[3] * shape[4]
+ in_bias = relay.var("in_bias", shape=(1, in_channels // blocking[0], 1, 1, blocking[0]))
+ else:
+ in_channels = shape[-1]
+ in_bias = relay.var("in_bias", shape=(in_channels,))
+
+ assert in_channels == channels
+ y1 = before(x, weight, in_bias, in_scale, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
y1_folded = run_opt_pass(y1, transform.ForwardFoldScaleAxis())
assert tvm.ir.structural_equal(y1, y1_folded)
in_scale = relay.var("in_scale", shape=(4,))
- check((2, 11, 10, 4), 4, in_scale)
+ check((2, 11, 10, 4), 4, None, in_scale)
in_scale = relay.const(-_get_positive_scale((4,)))
- check((2, 11, 10, 4), 4, in_scale)
+ check((2, 11, 10, 4), 4, None, in_scale)
+
+ in_scale = relay.var("in_scale", shape=(1,1,1,2,2))
+ check((2, 11, 10, 2, 2), 4, (2, 2), in_scale)
+ in_scale = relay.const(-_get_positive_scale((1,1,1,2,2)))
+ check((2, 11, 10, 2, 2), 4, (2, 2), in_scale)
+
+
def test_fold_fwd_negative_scale():
"""Testcase of folding negative scale"""
- def before(x, conv_weight, in_scale, channels):
+ def before(x, conv_weight, in_scale, channels, blocking):
args = [x, conv_weight]
x = relay.multiply(x, in_scale)
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW4i{}o".format(blocking[1]) if blocking else "OIHW")
return relay.Function(args, y)
- def expected(x, conv_weight, in_scale, channels):
+ def expected(x, conv_weight, in_scale, in_channels, channels, blocking):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight]
- squeezed_scale = relay.squeeze(in_scale, axis=[1,2])
- conv_weight = relay.multiply(
- conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
+ if blocking:
+ squeezed_scale = relay.squeeze(in_scale, axis=[0,2,3])
+ conv_weight = relay.multiply(
+ conv_weight , relay.reshape(squeezed_scale, (1, in_channels//4, 1, 1, 4, 1)))
+ #blocking by "i" in OIHWio
+ else:
+ squeezed_scale = relay.squeeze(in_scale, axis=[1,2])
+ conv_weight = relay.multiply(
+ conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
y = relay.nn.conv2d(x,
conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW4i{}o".format(blocking[1]) if blocking else "OIHW")
return relay.Function(args, y)
- def check(shape, channels):
+ def check(shape, channels, blocking):
x = relay.var("x", shape=shape)
- in_channels = shape[1]
- in_scale = relay.const(-_get_positive_scale((in_channels, 1, 1)))
+ if blocking:
+ in_channels = shape[1] * shape[4]
+ in_scale = relay.const(-_get_positive_scale((1, shape[1], 1, 1, shape[4])))
+ else:
+ in_channels = shape[1]
+ in_scale = relay.const(-_get_positive_scale((in_channels, 1, 1)))
weight = relay.var("weight")
- y1 = before(x, weight, in_scale, channels)
+ y1 = before(x, weight, in_scale, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
type_dict = {x.name_hint:x.checked_type for x in y1.params}
weight = relay.var("weight", type_dict["weight"])
y1_folded = run_opt_pass(y1, transform.ForwardFoldScaleAxis())
- y1_expected = expected(x, weight, in_scale, channels)
+ y1_expected = expected(x, weight, in_scale, in_channels, channels, blocking)
y1_expected = run_opt_pass(y1_expected, transform.InferType())
assert tvm.ir.structural_equal(y1_folded, y1_expected)
- check((2, 4, 10, 10), 4)
-
+ check((2, 4, 10, 10), 4, None)
+ check((2, 2, 10, 10, 2), 8, (2, 2))
def test_fold_bwd_simple():
"""Simple testcase."""
- def before(x, conv_weight, out_bias, out_scale, channels):
+ def before(x, conv_weight, out_bias, out_scale, in_channels, channels, blocking):
args = [x, conv_weight, out_bias]
- out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
+ if blocking:
+ out_bias = relay.reshape(out_bias, (1, channels//blocking[1], 1, 1, blocking[1]))
+ else:
+ out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y = relay.add(y, out_bias)
y = relay.nn.relu(y)
+ if blocking:
+ out_scale = relay.reshape(out_scale, (1, channels//blocking[1], 1, 1, blocking[1]))
y = relay.multiply(y, out_scale)
return relay.Function(args, y)
- def expected(x, conv_weight, out_bias, out_scale, channels):
+ def expected(x, conv_weight, out_bias, out_scale, in_channels, channels, blocking):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight, out_bias]
- out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
- squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
- conv_weight = relay.multiply(
- conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
+ if blocking:
+ out_bias = relay.reshape(out_bias, (1, channels//blocking[1], 1, 1, blocking[1]))
+ out_scale = relay.reshape(out_scale, (1, channels//blocking[1], 1, 1, blocking[1]))
+ squeezed_scale = relay.squeeze(out_scale, axis=[0, 2, 3])
+ conv_weight = relay.multiply(
+ conv_weight , relay.reshape(squeezed_scale, (channels//blocking[1], 1, 1, 1, 1, blocking[1])))
+ else:
+ out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
+ squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
+ conv_weight = relay.multiply(
+ conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
- out_bias = relay.multiply(out_bias,
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
+ if blocking:
+ out_bias = relay.multiply(out_bias,
+ relay.reshape(squeezed_scale, (1, channels//blocking[1], 1, 1, blocking[1])))
+ else:
+ out_bias = relay.multiply(out_bias,
relay.expand_dims(squeezed_scale, axis=1, num_newaxis=2))
y = relay.add(y, out_bias)
y = relay.nn.relu(y)
return relay.Function(args, y)
- def check(shape, channels):
+ def check(shape, in_channels, channels, blocking):
x = relay.var("x", shape=shape)
- in_channels = shape[1]
weight = relay.var("weight")
out_bias = relay.var("out_bias", shape=(channels,))
- out_scale = relay.const(_get_positive_scale((channels, 1, 1)))
-
- y1 = before(x, weight, out_bias, out_scale, channels)
+ if blocking:
+ out_scale = relay.const(_get_positive_scale((channels,)))
+ else:
+ out_scale = relay.const(_get_positive_scale((channels,1, 1)))
+ y1 = before(x, weight, out_bias, out_scale, in_channels, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
type_dict = {x.name_hint:x.checked_type for x in y1.params}
weight = relay.var("weight", type_dict["weight"])
y1_folded = run_opt_pass(y1, transform.BackwardFoldScaleAxis())
- y1_expected = expected(x, weight, out_bias, out_scale, channels)
+ y1_expected = expected(x, weight, out_bias, out_scale, in_channels, channels, blocking)
y1_expected = run_opt_pass(y1_expected, transform.InferType())
assert tvm.ir.structural_equal(y1_folded, y1_expected)
- check((2, 4, 10, 10), 8)
+ check((2, 4, 10, 10), 4, 8, None)
+ check((2, 2, 10, 10, 16), 32, 64, (16, 16))
def test_fold_bwd_dual_path():
"""Dual path testcase."""
- def before(x, conv_weight, out_bias, out_scale, channels):
+ def before(x, conv_weight, out_bias, out_scale, in_channels, channels, blocking):
args = [x, conv_weight, out_bias]
y1 = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y1 = relay.nn.relu(y1)
y2 = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y2 = relay.nn.relu(y2)
y = relay.add(y1, y2)
y = relay.multiply(y, out_scale)
return relay.Function(args, y)
- def expected(x, conv_weight, out_bias, out_scale, channels):
+ def expected(x, conv_weight, out_bias, out_scale, in_channels, channels, blocking):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight, out_bias]
- out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
+ if not blocking:
+ out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
def fold_conv_weight():
- return relay.multiply(
- conv_weight ,
- relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
+ if blocking:
+ return relay.multiply(
+ conv_weight ,
+ relay.reshape(squeezed_scale, (channels//blocking[1], 1, 1, 1, 1, blocking[1])))
+ else:
+ return relay.multiply(
+ conv_weight ,
+ relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
y1 = relay.nn.conv2d(x, fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y1 = relay.nn.relu(y1)
y2 = relay.nn.conv2d(x, fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y2 = relay.nn.relu(y2)
y = relay.add(y1, y2)
return relay.Function(args, y)
- def check(shape, channels):
+ def check(shape, in_channels, channels, blocking):
x = relay.var("x", shape=shape)
- in_channels = shape[1]
weight = relay.var("weight")
- out_bias = relay.var("out_bias", shape=(channels,))
- out_scale = relay.const(_get_positive_scale((channels, 1, 1)))
-
- y1 = before(x, weight, out_bias, out_scale, channels)
+ if blocking:
+ out_bias = relay.var("out_bias", shape=(channels // blocking[1], 1, 1, blocking[1]))
+ out_scale = relay.const(_get_positive_scale((channels // blocking[1], 1, 1, blocking[1])))
+ else:
+ out_bias = relay.var("out_bias", shape=(channels,))
+ out_scale = relay.const(_get_positive_scale((channels, 1, 1)))
+
+ y1 = before(x, weight, out_bias, out_scale, in_channels, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
type_dict = {x.name_hint:x.checked_type for x in y1.params}
weight = relay.var("weight", type_dict["weight"])
y1_folded = run_opt_pass(y1, transform.BackwardFoldScaleAxis())
- y1_expected = expected(x, weight, out_bias, out_scale, channels)
+ y1_expected = expected(x, weight, out_bias, out_scale, in_channels, channels, blocking)
y1_expected = run_opt_pass(y1_expected, transform.InferType())
assert tvm.ir.structural_equal(y1_folded, y1_expected)
- check((2, 4, 10, 10), 8)
-
+ check((2, 4, 10, 10), 4, 8, None)
+ check((2, 2, 10, 10, 2), 4, 8, (2, 2))
def test_fold_bwd_dual_consumer():
- def before(x, conv_weight, out_bias, out_scale, channels):
+ def before(x, conv_weight, out_bias, out_scale, in_channels, channels, blocking):
args = [x, conv_weight, out_bias]
y0 = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y0 = relay.multiply(y0, out_scale)
y0 = relay.nn.relu(y0)
y1 = relay.nn.conv2d(y0, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y1 = relay.multiply(y1, out_scale)
y1 = relay.nn.relu(y1)
y2 = relay.nn.conv2d(y0, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y2 = relay.multiply(y2, out_scale)
y2 = relay.nn.relu(y2)
y = relay.add(y1, y2)
return relay.Function(args, y)
- def expected(x, conv_weight, out_bias, out_scale, channels):
+ def expected(x, conv_weight, out_bias, out_scale, in_channels, channels, blocking):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight, out_bias]
def fold_conv_weight():
squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
- return relay.multiply(
- conv_weight ,
- relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
+ if blocking:
+ return relay.multiply(
+ conv_weight ,
+ relay.reshape(squeezed_scale, (channels//blocking[1], 1, 1, 1, 1, blocking[1])))
+ else:
+ return relay.multiply(
+ conv_weight ,
+ relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
y0 = relay.nn.conv2d(x, fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y0 = relay.nn.relu(y0)
y1 = relay.nn.conv2d(y0, fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y1 = relay.nn.relu(y1)
y2 = relay.nn.conv2d(y0, fold_conv_weight(),
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y2 = relay.nn.relu(y2)
y = relay.add(y1, y2)
return relay.Function(args, y)
- def check(shape, channels):
+ def check(shape, in_channels, channels, blocking):
x = relay.var("x", shape=shape)
- in_channels = shape[1]
weight = relay.var("weight")
- out_bias = relay.var("out_bias", shape=(channels,))
- out_scale = relay.const(_get_positive_scale((channels,1, 1)))
-
- y1 = before(x, weight, out_bias, out_scale, channels)
+ if blocking:
+ out_bias = relay.var("out_bias", shape=(channels // blocking[1], 1, 1, blocking[1]))
+ out_scale = relay.const(_get_positive_scale((channels // blocking[1], 1, 1, blocking[1])))
+ else:
+ out_bias = relay.var("out_bias", shape=(channels,))
+ out_scale = relay.const(_get_positive_scale((channels, 1, 1)))
+
+ y1 = before(x, weight, out_bias, out_scale, in_channels, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
type_dict = {x.name_hint:x.checked_type for x in y1.params}
weight = relay.var("weight", type_dict["weight"])
y1_folded = run_opt_pass(y1, transform.BackwardFoldScaleAxis())
- y1_expected = expected(x, weight, out_bias, out_scale, channels)
+ y1_expected = expected(x, weight, out_bias, out_scale, in_channels, channels, blocking)
y1_expected = run_opt_pass(y1_expected, transform.InferType())
assert tvm.ir.structural_equal(y1_folded, y1_expected)
- check((2, 4, 10, 10), 4)
-
+ check((2, 4, 10, 10), 4, 4, None)
+ check((2, 2, 10, 10, 2), 4, 4, (2, 2))
def test_fold_bwd_fail():
"""Dual path testcase."""
- def fail1(x, conv_weight, out_bias, out_scale, channels):
+ def fail1(x, conv_weight, out_bias, out_scale, in_channels, channels, blocking):
args = [x, conv_weight, out_bias]
- out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
y1 = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y1 = relay.nn.relu(y1)
y2 = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
padding=(1, 1),
- out_layout="CNHW")
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW",
+ out_layout="CNHW{}c".format(blocking[1]) if blocking else "CNHW")
# fold will fail because the axis from two path
# differs from each other.
y2 = relay.nn.relu(y2)
y = relay.multiply(y, out_scale)
return relay.Function(args, y)
- def fail2(x, conv_weight, out_bias, out_scale, channels):
+ def fail2(x, conv_weight, out_bias, out_scale, in_channels, channels, blocking):
args = [x, conv_weight, out_bias]
- out_bias = relay.expand_dims(out_bias, axis=1, num_newaxis=2)
y1 = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y2 = relay.nn.relu(y1)
# fold will fail because y1 is referred also by y2
y1 = relay.multiply(y1, out_scale)
y = relay.add(y1, y2)
return relay.Function(args, y)
- def check(shape, channels, fbefore):
+ def check(shape, in_channels, channels, blocking, fbefore):
x = relay.var("x", shape=shape)
- in_channels = shape[1]
weight = relay.var("weight")
- out_bias = relay.var("out_bias", shape=(channels,))
- out_scale = relay.const(_get_positive_scale((channels, 1, 1)))
- y1 = fbefore(x, weight, out_bias, out_scale, channels)
+ if blocking:
+ out_bias = relay.var("out_bias", shape=(channels // blocking[1], 1, 1, blocking[1]))
+ out_scale = relay.const(_get_positive_scale((channels // blocking[1], 1, 1, blocking[1])))
+ else:
+ out_bias = relay.var("out_bias", shape=(channels, 1, 1))
+ out_scale = relay.const(_get_positive_scale((channels, 1, 1)))
+ y1 = fbefore(x, weight, out_bias, out_scale, in_channels, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
y1_folded = run_opt_pass(y1, transform.BackwardFoldScaleAxis())
assert tvm.ir.structural_equal(y1_folded, y1)
- check((4, 4, 10, 10), 4, fail1)
- check((4, 4, 10, 10), 4, fail2)
+ check((4, 4, 10, 10), 4, 4, None, fail1)
+ check((2, 2, 10, 10, 2), 4, 4, (2, 2), fail1)
+ check((4, 4, 10, 10), 4, 4, None, fail2)
+ check((4, 2, 10, 10, 2), 4, 4, (2, 2), fail2)
def test_fold_bwd_relu_fail():
"""testcase where we canont fold because scale can not pass relu"""
- def before(x, conv_weight, out_scale, channels):
+ def before(x, conv_weight, out_scale, channels, blocking):
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- data_layout="NCHW",
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y = relay.nn.relu(y)
y = relay.multiply(x, out_scale)
return relay.Function(relay.analysis.free_vars(y), y)
- def check(shape, channels, out_scale):
+ def check(shape, channels, blocking, out_scale):
x = relay.var("x", shape=shape)
in_channels = shape[1]
weight = relay.var("weight")
- y1 = before(x, weight, out_scale, channels)
+ y1 = before(x, weight, out_scale, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
y1_folded = run_opt_pass(y1, transform.BackwardFoldScaleAxis())
assert tvm.ir.structural_equal(y1, y1_folded)
out_scale = relay.var("in_scale", shape=(4, 1, 1))
- check((4, 4, 10, 10), 4, out_scale)
+ check((4, 4, 10, 10), 4, None, out_scale)
out_scale = relay.const(np.random.uniform(size=(4, 1, 1), low=-1.0, high=0.0)).astype("float32")
- check((4, 4, 10, 10), 4, out_scale)
+ check((4, 4, 10, 10), 4, None, out_scale)
+
+ out_scale = relay.var("in_scale", shape=(1, 2, 1, 1, 2))
+ check((4, 2, 10, 10, 2), 4, (2, 2), out_scale)
+ out_scale = relay.const(np.random.uniform(size=(1, 2, 1, 1, 2), low=-1.0, high=0.0)).astype("float32")
+ check((4, 2, 10, 10, 2), 4, (2, 2), out_scale)
def test_fold_bwd_negative_scale():
"""Testcase of folding negative scale"""
- def before(x, conv_weight, out_scale, channels):
+ def before(x, conv_weight, out_scale, channels, blocking):
args = [x, conv_weight]
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
y = relay.multiply(y, out_scale)
return relay.Function(args, y)
- def expected(x, conv_weight, out_scale, channels):
+ def expected(x, conv_weight, out_scale, channels, blocking):
# use a fixed order of args so alpha equal check can pass
args = [x, conv_weight]
- squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
- conv_weight = relay.multiply(
- conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
+ if blocking:
+ squeezed_scale = relay.squeeze(out_scale, axis=[0,2,3])
+ conv_weight = relay.multiply(
+ conv_weight , relay.reshape(squeezed_scale, (channels//blocking[1], 1, 1, 1, 1, blocking[1])))
+ else:
+ squeezed_scale = relay.squeeze(out_scale, axis=[1,2])
+ conv_weight = relay.multiply(
+ conv_weight , relay.expand_dims(squeezed_scale, axis=1, num_newaxis=3))
y = relay.nn.conv2d(x, conv_weight,
channels=channels,
kernel_size=(3, 3),
- padding=(1, 1))
+ padding=(1, 1),
+ data_layout="NCHW{}c".format(blocking[0]) if blocking else "NCHW",
+ kernel_layout="OIHW1i{}o".format(blocking[1]) if blocking else "OIHW")
return relay.Function(args, y)
- def check(shape, channels):
+ def check(shape, channels, blocking):
x = relay.var("x", shape=shape)
weight = relay.var("weight")
- out_scale = relay.const(-_get_positive_scale((channels, 1, 1)))
- y1 = before(x, weight, out_scale, channels)
+ if blocking:
+ out_scale = relay.const(-_get_positive_scale((1,channels//blocking[1], 1, 1, blocking[1])))
+ else:
+ out_scale = relay.const(-_get_positive_scale((channels, 1, 1)))
+ y1 = before(x, weight, out_scale, channels, blocking)
y1 = run_opt_pass(y1, transform.InferType())
type_dict = {x.name_hint:x.checked_type for x in y1.params}
weight = relay.var("weight", type_dict["weight"])
y1_folded = run_opt_pass(y1, transform.BackwardFoldScaleAxis())
- y1_expected = expected(x, weight, out_scale, channels)
+ y1_expected = expected(x, weight, out_scale, channels, blocking)
y1_expected = run_opt_pass(y1_expected, transform.InferType())
assert tvm.ir.structural_equal(y1_folded, y1_expected)
- check((2, 4, 10, 10), 8)
-
+ check((2, 4, 10, 10), 8, None)
+ check((2, 2, 10, 10, 2), 8, (2, 2))
if __name__ == "__main__":
test_fold_fwd_simple()