diff --git a/src/relay/transforms/fold_scale_axis.cc b/src/relay/transforms/fold_scale_axis.cc index ca9a286ff6a7..4b94159fe3d1 100644 --- a/src/relay/transforms/fold_scale_axis.cc +++ b/src/relay/transforms/fold_scale_axis.cc @@ -589,6 +589,28 @@ RELAY_REGISTER_OP("nn.conv2d").set_attr("FScaleAxisForwardPrep", C RELAY_REGISTER_OP("nn.conv2d") .set_attr("FScaleAxisForwardRewrite", Conv2DForwardRewrite); +// Dense send out requirement of axis folding. +Array DenseForwardPrep(const Call& call, const Message& out_message) { + return {Message({1}, false), NullValue()}; +} + +// Dense consumes the scale axis during transformation. +Expr DenseForwardRewrite(const Call& ref_call, const Array& new_args, + const Message& message) { + const auto* sdata = new_args[0].as(); + const auto* sweight = new_args[1].as(); + if (sdata == nullptr) return Expr(); + if (sweight != nullptr) return Expr(); + + Expr weight = Multiply(new_args[1], sdata->scale); + return Call(ref_call->op, {sdata->value, weight}, ref_call->attrs, ref_call->type_args); +} + +RELAY_REGISTER_OP("nn.dense").set_attr("FScaleAxisForwardPrep", DenseForwardPrep); + +RELAY_REGISTER_OP("nn.dense") + .set_attr("FScaleAxisForwardRewrite", DenseForwardRewrite); + Expr ForwardFoldScaleAxis(const Expr& data) { auto message = ForwardPrep().Prepare(data); for (const auto& m : message) { @@ -996,6 +1018,67 @@ RELAY_REGISTER_OP("nn.conv2d") RELAY_REGISTER_OP("nn.conv2d") .set_attr("FScaleAxisBackwardTransform", Conv2DBackwardTransform); +Message BiasAddBackwardPrep(const Call& call, const Array& in_messages) { + const BiasAddAttrs* attrs = call->attrs.as(); + ICHECK(attrs); + if (in_messages[0].defined() && in_messages[0]->axes.size() == 1 && + attrs->axis == static_cast(in_messages[0]->axes[0]->value)) { + return in_messages[0]; + } else { + return NullValue(); + } +} + +Expr BiasAddBackwardTransform(const Call& call, const Message& message, const Expr& scale, + const BackwardTransformer& transformer) { + if (!message.defined()) { + return transformer->NormalCallTransform(call.operator->()); + } + Message lhs_message = transformer->GetMessage(call->args[0]); + Message rhs_message = transformer->GetMessage(call->args[1]); + StructuralEqual equal; + + if (lhs_message.defined()) { + ICHECK(equal(message->axes, lhs_message->axes)); + Expr lhs = transformer->Transform(call->args[0], message, scale); + Expr rhs = transformer->Transform(call->args[1], NullValue(), NullValue()); + rhs = Multiply(rhs, scale); + return Call(call->op, {lhs, rhs}, call->attrs, call->type_args); + } else { + LOG(FATAL) << "outstanding scale"; + return Expr(); + } +} + +RELAY_REGISTER_OP("nn.bias_add") + .set_attr("FScaleAxisBackwardPrep", BiasAddBackwardPrep); + +RELAY_REGISTER_OP("nn.bias_add") + .set_attr("FScaleAxisBackwardTransform", BiasAddBackwardTransform); + +// Dense send out requirement of axis folding. +Message DenseBackwardPrep(const Call& call, const Array& in_messages) { + return Message({1}, false); +} + +// Dense consumes the sacle axis during trasformation. +Expr DenseBackwardTransform(const Call& call, const Message& message, const Expr& scale, + const BackwardTransformer& transformer) { + if (!message.defined()) { + return transformer->NormalCallTransform(call.operator->()); + } + Expr data = transformer->Transform(call->args[0], NullValue(), NullValue()); + Expr weight = transformer->Transform(call->args[1], NullValue(), NullValue()); + Expr wscale = ExpandBiasToMatchAxis(scale, 2, {0}); + weight = Multiply(weight, wscale); + return Call(call->op, {data, weight}, call->attrs, call->type_args); +} + +RELAY_REGISTER_OP("nn.dense").set_attr("FScaleAxisBackwardPrep", DenseBackwardPrep); + +RELAY_REGISTER_OP("nn.dense") + .set_attr("FScaleAxisBackwardTransform", DenseBackwardTransform); + Expr BackwardFoldScaleAxis(const Expr& data) { return make_object()->Fold(data); } diff --git a/tests/python/relay/test_pass_fold_scale_axis.py b/tests/python/relay/test_pass_fold_scale_axis.py index 3c2dc82cb07b..b5e7b1c81645 100644 --- a/tests/python/relay/test_pass_fold_scale_axis.py +++ b/tests/python/relay/test_pass_fold_scale_axis.py @@ -413,6 +413,46 @@ def check(shape, channels, blocking): check((2, 2, 10, 10, 2), 8, (2, 2)) +def test_fold_fwd_dense(): + """dense testcase.""" + + def before(x, weight, in_bias, in_scale): + args = [x, weight, in_bias] + x = relay.multiply(x, in_scale) + x = relay.nn.relu(x) + x = relay.add(x, in_bias) + y = relay.nn.dense(x, weight) + return relay.Function(args, y) + + def expected(x, weight, in_bias, in_scale): + # use a fixed order of args so alpha equal check can pass + args = [x, weight, in_bias] + x = relay.nn.relu(x) + in_bias = relay.divide(in_bias, in_scale) + x = relay.add(x, in_bias) + weight = relay.multiply(weight, in_scale) + y = relay.nn.dense(x, weight) + return relay.Function(args, y) + + def check(data_shape, weight_shape): + x = relay.var("x", shape=data_shape) + weight = relay.var("weight", shape=weight_shape) + in_channels = data_shape[1] + 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) + y1 = run_opt_pass(y1, transform.InferType()) + y1_folded = run_opt_pass(y1, transform.ForwardFoldScaleAxis()) + y1_expected = expected(x, weight, in_bias, in_scale) + + 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), (3, 4)) + check((3, 5), (4, 5)) + + def test_fold_bwd_simple(): """Simple testcase.""" @@ -888,15 +928,118 @@ def check(shape, channels, blocking): check((2, 2, 10, 10, 2), 8, (2, 2)) +def test_fold_bwd_dense(): + """dense testcase.""" + + def before(x, weight, in_bias, in_scale): + args = [x, weight, in_bias] + x = relay.nn.dense(x, weight) + x = relay.add(x, in_bias) + x = relay.nn.relu(x) + y = relay.multiply(x, in_scale) + return relay.Function(args, y) + + def expected(x, weight, in_bias, in_scale): + # use a fixed order of args so alpha equal check can pass + args = [x, weight, in_bias] + scale = relay.expand_dims(in_scale, axis=1) + weight = relay.multiply(weight, scale) + x = relay.nn.dense(x, weight) + bias = relay.multiply(in_bias, in_scale) + x = relay.add(x, bias) + y = relay.nn.relu(x) + return relay.Function(args, y) + + def check(data_shape, weight_shape): + x = relay.var("x", shape=data_shape) + weight = relay.var("weight", shape=weight_shape) + out_channels = weight_shape[0] + in_bias = relay.var("in_bias", shape=(out_channels,)) + in_scale = relay.const(_get_positive_scale((out_channels,))) + y1 = before(x, weight, in_bias, in_scale) + y1 = run_opt_pass(y1, transform.InferType()) + y1_folded = run_opt_pass(y1, transform.BackwardFoldScaleAxis()) + y1_expected = expected(x, weight, in_bias, in_scale) + + 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), (3, 4)) + check((3, 5), (4, 5)) + + +def test_fold_bwd_bias_add(): + """bias add testcase.""" + + def before(x, conv_weight, out_bias, out_scale, channels): + args = [x, conv_weight, out_bias] + y = relay.nn.conv2d( + x, + conv_weight, + channels=channels, + kernel_size=(3, 3), + padding=(1, 1), + data_layout="NCHW", + kernel_layout="OIHW", + ) + y = relay.nn.bias_add(y, out_bias) + y = relay.nn.relu(y) + y = relay.multiply(y, out_scale) + return relay.Function(args, y) + + def expected(x, conv_weight, out_bias, out_scale, channels): + # use a fixed order of args so alpha equal check can pass + args = [x, conv_weight, out_bias] + 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), + data_layout="NCHW", + kernel_layout="OIHW", + ) + + out_bias = relay.multiply(out_bias, squeezed_scale) + y = relay.nn.bias_add(y, out_bias) + y = relay.nn.relu(y) + return relay.Function(args, y) + + def check(shape, channels): + x = relay.var("x", shape=shape) + 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) + 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 = run_opt_pass(y1_expected, transform.InferType()) + assert tvm.ir.structural_equal(y1_folded, y1_expected) + + check((2, 4, 10, 10), 4) + + if __name__ == "__main__": test_fold_fwd_simple() test_fold_fwd_dual_path() test_fold_fwd_fail() test_fold_fwd_relu_fail() test_fold_fwd_negative_scale() + test_fold_fwd_dense() test_fold_bwd_simple() test_fold_bwd_dual_path() test_fold_bwd_dual_consumer() test_fold_bwd_fail() test_fold_bwd_relu_fail() test_fold_bwd_negative_scale() + test_fold_bwd_dense() + test_fold_bwd_bias_add()