You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "SiriusNEO (via GitHub)" <gi...@apache.org> on 2023/04/05 17:44:34 UTC
[GitHub] [tvm] SiriusNEO opened a new pull request, #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
SiriusNEO opened a new pull request, #14512:
URL: https://github.com/apache/tvm/pull/14512
Prior to this PR, we first have a pass named `SimplifyNormInference` (https://github.com/apache/tvm/pull/14221), which is used to decompose `batch_norm` into simple operators for optimization during inference. And there are some follow-up changes:
- https://github.com/apache/tvm/pull/14282 This PR introduces a policy to decompose `tensor_to_shape` op.
- https://github.com/mlc-ai/relax/pull/170 In mlc, we make `SimplifyNormInference` support op simplification during training (Because `batch_norm` behaves differently in inference and training).
- And in the future maybe there will be more operators needed to be decomposed.
So this PR polishes the changes in https://github.com/apache/tvm/pull/14282 mainly in the following aspects:
- Use one mutator (instead of two) to visit and decompose ops.
- Separate inference and training by registering two passes (Adviced by TQ).
- Shorten the name (DecomposeCompositeOps -> DecomposeOps)
Now the code is clear and easy to read. Now we use `if-else if-else if ...` pattern to recognize the op we want to decompose. In the future maybe we can introduce a map or register the decomposition policy in the op attribute, just like what we do in the `LegalizeOps` pass.
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Hzfengsy commented on a diff in pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "Hzfengsy (via GitHub)" <gi...@apache.org>.
Hzfengsy commented on code in PR #14512:
URL: https://github.com/apache/tvm/pull/14512#discussion_r1159363592
##########
tests/python/relax/test_op_nn.py:
##########
@@ -272,7 +272,7 @@ def test_batch_norm_infer_struct_info():
_check_inference(
bb,
- relax.op.nn.batch_norm(x0, gamma0, beta0, moving_mean0, moving_var0, axis=1),
Review Comment:
Can we use the default value for momentum? Then we can keep using the current codebase.
##########
src/relax/transform/decompose_ops.cc:
##########
@@ -0,0 +1,267 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*! \file src/relax/transform/decompose_ops.cc */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/nn.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/relax/transform.h>
+
+#include "utils.h"
+
+namespace tvm {
+namespace relax {
+
+TensorStructInfo MatchTensorStructInfo(Expr data) {
+ auto _sinfo = MatchStructInfo<TensorStructInfo>(data);
+ ICHECK(_sinfo.defined()) << "Expect data to be a tensor, but get " << GetStructInfo(data);
+ return _sinfo.value();
+}
+
+Expr ExpandToMatchInput(Expr data, int ndim, Array<Integer> axes) {
+ axes = GetOrderedPositiveAxes(axes, ndim);
+ Array<Integer> expand_axes;
+ for (int i = 0, j = 0; i < ndim; ++i) {
+ if (j < static_cast<int>(axes.size()) && i == axes[j]->value) {
+ ++j;
+ } else {
+ expand_axes.push_back(i);
+ }
+ }
+ return expand_dims(data, expand_axes);
+}
+
+Tuple SimplifyBatchNormInference(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Expr moving_mean = ExpandToMatchInput(call->args[3], sinfo->ndim, {attrs->axis});
+ Expr moving_var = ExpandToMatchInput(call->args[4], sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(moving_var, epsilon));
+ Expr out = divide(subtract(data, moving_mean), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ return Tuple({out, call->args[3], call->args[4]});
+}
+
+Tuple SimplifyBatchNormTraining(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Array<Integer> reduce_axes;
+ for (int i = 0; i < sinfo->ndim; ++i) {
+ if (i != attrs->axis) {
+ reduce_axes.push_back(i);
+ }
+ }
+
+ Expr data_mean = mean(data, reduce_axes, false);
+ Expr data_mean_rs = ExpandToMatchInput(data_mean, sinfo->ndim, {attrs->axis});
+ Expr data_var = variance(data, reduce_axes, false);
+ Expr data_var_rs = ExpandToMatchInput(data_var, sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(data_var_rs, epsilon));
+ Expr out = divide(subtract(data, data_mean_rs), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ Expr moving_mean = call->args[3];
+ Expr moving_var = call->args[4];
+ Expr momentum = MakeConstantScalar(attrs->momentum, sinfo->dtype);
+ Expr one_minus_mom = MakeConstantScalar(1 - attrs->momentum, sinfo->dtype);
+
+ return Tuple({
+ out,
+ add(multiply(one_minus_mom, moving_mean), multiply(momentum, data_mean)),
+ add(multiply(one_minus_mom, moving_var), multiply(momentum, data_var)),
+ });
+}
+
+Expr SimplifyLayerNorm(const Call& call) {
Review Comment:
As decomposing `LayerNorm` is used for training and may influence the inference perf, let's only decompose it in training mode
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] SiriusNEO commented on a diff in pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "SiriusNEO (via GitHub)" <gi...@apache.org>.
SiriusNEO commented on code in PR #14512:
URL: https://github.com/apache/tvm/pull/14512#discussion_r1159298963
##########
src/relax/transform/decompose_ops.cc:
##########
@@ -0,0 +1,267 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*! \file src/relax/transform/decompose_ops.cc */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/nn.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/relax/transform.h>
+
+#include "utils.h"
+
+namespace tvm {
+namespace relax {
+
+TensorStructInfo MatchTensorStructInfo(Expr data) {
+ auto _sinfo = MatchStructInfo<TensorStructInfo>(data);
+ ICHECK(_sinfo.defined()) << "Expect data to be a tensor, but get " << GetStructInfo(data);
+ return _sinfo.value();
+}
+
+Expr ExpandToMatchInput(Expr data, int ndim, Array<Integer> axes) {
+ axes = GetOrderedPositiveAxes(axes, ndim);
+ Array<Integer> expand_axes;
+ for (int i = 0, j = 0; i < ndim; ++i) {
+ if (j < static_cast<int>(axes.size()) && i == axes[j]->value) {
+ ++j;
+ } else {
+ expand_axes.push_back(i);
+ }
+ }
+ return expand_dims(data, expand_axes);
+}
+
+Tuple SimplifyBatchNormInference(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Expr moving_mean = ExpandToMatchInput(call->args[3], sinfo->ndim, {attrs->axis});
+ Expr moving_var = ExpandToMatchInput(call->args[4], sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(moving_var, epsilon));
+ Expr out = divide(subtract(data, moving_mean), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ return Tuple({out, call->args[3], call->args[4]});
+}
+
+Tuple SimplifyBatchNormTraining(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Array<Integer> reduce_axes;
+ for (int i = 0; i < sinfo->ndim; ++i) {
+ if (i != attrs->axis) {
+ reduce_axes.push_back(i);
+ }
+ }
+
+ Expr data_mean = mean(data, reduce_axes, false);
+ Expr data_mean_rs = ExpandToMatchInput(data_mean, sinfo->ndim, {attrs->axis});
+ Expr data_var = variance(data, reduce_axes, false);
+ Expr data_var_rs = ExpandToMatchInput(data_var, sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(data_var_rs, epsilon));
+ Expr out = divide(subtract(data, data_mean_rs), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ Expr moving_mean = call->args[3];
+ Expr moving_var = call->args[4];
+ Expr momentum = MakeConstantScalar(attrs->momentum, sinfo->dtype);
+ Expr one_minus_mom = MakeConstantScalar(1 - attrs->momentum, sinfo->dtype);
+
+ return Tuple({
+ out,
+ add(multiply(one_minus_mom, moving_mean), multiply(momentum, data_mean)),
+ add(multiply(one_minus_mom, moving_var), multiply(momentum, data_var)),
+ });
+}
+
+Expr SimplifyLayerNorm(const Call& call) {
Review Comment:
Ah here we simplify `LayerNorm` for now just because it is easy to get its gradient. If it is more efficient as a single op, maybe in the future we can remove this from the pass after we finish the gradient function of the whole `LayerNorm`
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14512:
URL: https://github.com/apache/tvm/pull/14512#discussion_r1160029376
##########
src/relax/transform/decompose_ops.cc:
##########
@@ -0,0 +1,267 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*! \file src/relax/transform/decompose_ops.cc */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/nn.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/relax/transform.h>
+
+#include "utils.h"
+
+namespace tvm {
+namespace relax {
+
+TensorStructInfo MatchTensorStructInfo(Expr data) {
+ auto _sinfo = MatchStructInfo<TensorStructInfo>(data);
+ ICHECK(_sinfo.defined()) << "Expect data to be a tensor, but get " << GetStructInfo(data);
+ return _sinfo.value();
+}
+
+Expr ExpandToMatchInput(Expr data, int ndim, Array<Integer> axes) {
+ axes = GetOrderedPositiveAxes(axes, ndim);
+ Array<Integer> expand_axes;
+ for (int i = 0, j = 0; i < ndim; ++i) {
+ if (j < static_cast<int>(axes.size()) && i == axes[j]->value) {
+ ++j;
+ } else {
+ expand_axes.push_back(i);
+ }
+ }
+ return expand_dims(data, expand_axes);
+}
+
+Tuple SimplifyBatchNormInference(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Expr moving_mean = ExpandToMatchInput(call->args[3], sinfo->ndim, {attrs->axis});
+ Expr moving_var = ExpandToMatchInput(call->args[4], sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(moving_var, epsilon));
+ Expr out = divide(subtract(data, moving_mean), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ return Tuple({out, call->args[3], call->args[4]});
+}
+
+Tuple SimplifyBatchNormTraining(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Array<Integer> reduce_axes;
+ for (int i = 0; i < sinfo->ndim; ++i) {
+ if (i != attrs->axis) {
+ reduce_axes.push_back(i);
+ }
+ }
+
+ Expr data_mean = mean(data, reduce_axes, false);
+ Expr data_mean_rs = ExpandToMatchInput(data_mean, sinfo->ndim, {attrs->axis});
+ Expr data_var = variance(data, reduce_axes, false);
+ Expr data_var_rs = ExpandToMatchInput(data_var, sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(data_var_rs, epsilon));
+ Expr out = divide(subtract(data, data_mean_rs), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ Expr moving_mean = call->args[3];
+ Expr moving_var = call->args[4];
+ Expr momentum = MakeConstantScalar(attrs->momentum, sinfo->dtype);
+ Expr one_minus_mom = MakeConstantScalar(1 - attrs->momentum, sinfo->dtype);
+
+ return Tuple({
+ out,
+ add(multiply(one_minus_mom, moving_mean), multiply(momentum, data_mean)),
+ add(multiply(one_minus_mom, moving_var), multiply(momentum, data_var)),
+ });
+}
+
+Expr SimplifyLayerNorm(const Call& call) {
Review Comment:
> Is there any complication for fusion if we have a composite ops in general?
Yes, we need to be aware of the implementation when setting op type (`kOpaque` or `kInjective`). This will cause different fusion results.
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] SiriusNEO commented on a diff in pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "SiriusNEO (via GitHub)" <gi...@apache.org>.
SiriusNEO commented on code in PR #14512:
URL: https://github.com/apache/tvm/pull/14512#discussion_r1159299613
##########
src/relax/transform/decompose_ops.cc:
##########
@@ -0,0 +1,267 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*! \file src/relax/transform/decompose_ops.cc */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/nn.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/relax/transform.h>
+
+#include "utils.h"
+
+namespace tvm {
+namespace relax {
+
+TensorStructInfo MatchTensorStructInfo(Expr data) {
+ auto _sinfo = MatchStructInfo<TensorStructInfo>(data);
+ ICHECK(_sinfo.defined()) << "Expect data to be a tensor, but get " << GetStructInfo(data);
+ return _sinfo.value();
+}
+
+Expr ExpandToMatchInput(Expr data, int ndim, Array<Integer> axes) {
+ axes = GetOrderedPositiveAxes(axes, ndim);
+ Array<Integer> expand_axes;
+ for (int i = 0, j = 0; i < ndim; ++i) {
+ if (j < static_cast<int>(axes.size()) && i == axes[j]->value) {
+ ++j;
+ } else {
+ expand_axes.push_back(i);
+ }
+ }
+ return expand_dims(data, expand_axes);
+}
+
+Tuple SimplifyBatchNormInference(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Expr moving_mean = ExpandToMatchInput(call->args[3], sinfo->ndim, {attrs->axis});
+ Expr moving_var = ExpandToMatchInput(call->args[4], sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(moving_var, epsilon));
+ Expr out = divide(subtract(data, moving_mean), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ return Tuple({out, call->args[3], call->args[4]});
+}
+
+Tuple SimplifyBatchNormTraining(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Array<Integer> reduce_axes;
+ for (int i = 0; i < sinfo->ndim; ++i) {
+ if (i != attrs->axis) {
+ reduce_axes.push_back(i);
+ }
+ }
+
+ Expr data_mean = mean(data, reduce_axes, false);
+ Expr data_mean_rs = ExpandToMatchInput(data_mean, sinfo->ndim, {attrs->axis});
+ Expr data_var = variance(data, reduce_axes, false);
+ Expr data_var_rs = ExpandToMatchInput(data_var, sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(data_var_rs, epsilon));
+ Expr out = divide(subtract(data, data_mean_rs), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ Expr moving_mean = call->args[3];
+ Expr moving_var = call->args[4];
+ Expr momentum = MakeConstantScalar(attrs->momentum, sinfo->dtype);
+ Expr one_minus_mom = MakeConstantScalar(1 - attrs->momentum, sinfo->dtype);
+
+ return Tuple({
+ out,
+ add(multiply(one_minus_mom, moving_mean), multiply(momentum, data_mean)),
+ add(multiply(one_minus_mom, moving_var), multiply(momentum, data_var)),
+ });
+}
+
+Expr SimplifyLayerNorm(const Call& call) {
Review Comment:
Or maybe we can add some arguments for the pass to let the user specify what operators he want to simplify?
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on PR #14512:
URL: https://github.com/apache/tvm/pull/14512#issuecomment-1499307950
@sunggg please take another look
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] sunggg commented on a diff in pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "sunggg (via GitHub)" <gi...@apache.org>.
sunggg commented on code in PR #14512:
URL: https://github.com/apache/tvm/pull/14512#discussion_r1158952771
##########
src/relax/transform/decompose_ops.cc:
##########
@@ -0,0 +1,267 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*! \file src/relax/transform/decompose_ops.cc */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/nn.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/relax/transform.h>
+
+#include "utils.h"
+
+namespace tvm {
+namespace relax {
+
+TensorStructInfo MatchTensorStructInfo(Expr data) {
+ auto _sinfo = MatchStructInfo<TensorStructInfo>(data);
+ ICHECK(_sinfo.defined()) << "Expect data to be a tensor, but get " << GetStructInfo(data);
+ return _sinfo.value();
+}
+
+Expr ExpandToMatchInput(Expr data, int ndim, Array<Integer> axes) {
+ axes = GetOrderedPositiveAxes(axes, ndim);
+ Array<Integer> expand_axes;
+ for (int i = 0, j = 0; i < ndim; ++i) {
+ if (j < static_cast<int>(axes.size()) && i == axes[j]->value) {
+ ++j;
+ } else {
+ expand_axes.push_back(i);
+ }
+ }
+ return expand_dims(data, expand_axes);
+}
+
+Tuple SimplifyBatchNormInference(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Expr moving_mean = ExpandToMatchInput(call->args[3], sinfo->ndim, {attrs->axis});
+ Expr moving_var = ExpandToMatchInput(call->args[4], sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(moving_var, epsilon));
+ Expr out = divide(subtract(data, moving_mean), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ return Tuple({out, call->args[3], call->args[4]});
+}
+
+Tuple SimplifyBatchNormTraining(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Array<Integer> reduce_axes;
+ for (int i = 0; i < sinfo->ndim; ++i) {
+ if (i != attrs->axis) {
+ reduce_axes.push_back(i);
+ }
+ }
+
+ Expr data_mean = mean(data, reduce_axes, false);
+ Expr data_mean_rs = ExpandToMatchInput(data_mean, sinfo->ndim, {attrs->axis});
+ Expr data_var = variance(data, reduce_axes, false);
+ Expr data_var_rs = ExpandToMatchInput(data_var, sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(data_var_rs, epsilon));
+ Expr out = divide(subtract(data, data_mean_rs), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ Expr moving_mean = call->args[3];
+ Expr moving_var = call->args[4];
+ Expr momentum = MakeConstantScalar(attrs->momentum, sinfo->dtype);
+ Expr one_minus_mom = MakeConstantScalar(1 - attrs->momentum, sinfo->dtype);
+
+ return Tuple({
+ out,
+ add(multiply(one_minus_mom, moving_mean), multiply(momentum, data_mean)),
+ add(multiply(one_minus_mom, moving_var), multiply(momentum, data_var)),
+ });
+}
+
+Expr SimplifyLayerNorm(const Call& call) {
+ auto attrs = call->attrs.as<LayerNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Expr data_mean = mean(data, attrs->axes, true);
+ Expr data_var = variance(data, attrs->axes, true);
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(data_var, epsilon));
+ Expr out = divide(subtract(data, data_mean), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, gamma);
+ }
+ if (attrs->center) {
+ out = add(out, beta);
+ }
+
+ return out;
+}
+
+Expr TensorToShape(const Call& call_node, const BlockBuilder& builder) {
+ ICHECK(call_node->struct_info_.defined());
+ Expr expr = call_node->args[0];
+ const ShapeStructInfoNode* sinfo = GetStructInfoAs<ShapeStructInfoNode>(call_node);
+ ICHECK(sinfo);
+ // call builtin function that converts tensor to shape tuple
+ // TODO(@sunggg): Register operator for "vm.builtin.tensor_to_shape"
+ Var call = builder->Emit(
+ Call(ExternFunc("vm.builtin.tensor_to_shape"), {expr}, {}, {GetRef<ShapeStructInfo>(sinfo)}));
+
+ // Operators like reshape take the output of `TensorToShape` as their output shape.
+ // Because TOPI expects to have such output shape in symbolic shape at least (i.e.,
+ // Array<PrimExpr>), we define symbolic variables and returns them as a ShapeExpr.
+ Array<PrimExpr> shape_var;
+ for (int i = 0; i < sinfo->ndim; i++) {
+ shape_var.push_back(tir::Var("x", DataType::Int(64)));
+ }
+ // bind symbolic variables to the shape tuple
+ relax::Var var("y", ShapeStructInfo(shape_var));
+ builder->EmitNormalized(MatchCast(var, call, ShapeStructInfo(shape_var)));
+ return ShapeExpr(shape_var);
+}
+
+class OpDecomposer : public ExprMutator {
+ public:
+ constexpr static const char* kModeEval = "eval";
+ constexpr static const char* kModeTraining = "training";
+
+ explicit OpDecomposer(String mode) : ExprMutator(), mode_(mode) {
+ CHECK(mode == kModeEval || mode == kModeTraining)
+ << "The argument mode must be one of the following values: \"eval\", \"training\".";
+ }
+
+ private:
+ using ExprMutator::VisitExpr_;
+
+ Expr VisitExpr_(const CallNode* call_node) final {
+ Call call = Downcast<Call>(VisitExprPostOrder_(call_node));
+ if (call->op == batch_norm_op_) {
+ if (mode_ == kModeEval) {
+ return SimplifyBatchNormInference(call);
+ } else {
+ ICHECK_EQ(mode_, kModeTraining);
+ return SimplifyBatchNormTraining(call);
+ }
+ } else if (call->op == layer_norm_op_) {
+ return SimplifyLayerNorm(call);
+ } else if (call->op == tensor_to_shape_op_) {
+ return TensorToShape(call, builder_);
+ }
+ return call;
+ }
+
+ const String mode_;
+
+ /* composite opeartor list */
+ const Op& batch_norm_op_ = Op::Get("relax.nn.batch_norm");
+ const Op& layer_norm_op_ = Op::Get("relax.nn.layer_norm");
+ const Op& tensor_to_shape_op_ = Op::Get("relax.tensor_to_shape");
+};
+
+IRModule Decompose(IRModule mod, Optional<String> func_name, String mode) {
+ auto op_decomposer = OpDecomposer(mode);
+
+ IRModuleNode* new_module = mod.CopyOnWrite();
+
+ if (!func_name.defined()) { // simplify all functions
+ Map<GlobalVar, BaseFunc> functions = mod->functions;
+ for (const auto& func_pr : functions) {
+ if (const auto* relax_f = func_pr.second.as<FunctionNode>()) {
+ Function f = Downcast<Function>(op_decomposer(GetRef<Function>(relax_f)));
+ new_module->Update(func_pr.first, f);
+ }
+ }
+ } else { // simplify specified function
+ auto* func_ptr = mod->Lookup(func_name.value()).as<FunctionNode>();
+ CHECK(func_ptr) << func_name.value() << "is not a Relax Function";
+ auto gvar = mod->GetGlobalVar(func_name.value());
+ auto func = GetRef<Function>(func_ptr);
+ func = Downcast<Function>(op_decomposer(func));
+ new_module->Update(gvar, func);
+ }
+
+ return GetRef<IRModule>(new_module);
+}
+
+namespace transform {
+Pass DecomposeOpsForInference(Optional<String> func_name) {
+ runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func = [=](IRModule mod,
+ PassContext pc) {
+ return Decompose(mod, func_name, OpDecomposer::kModeEval);
Review Comment:
Nit: To be consistent with the pass name, can we rename `kModeEval` to `kModeInference`?
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] SiriusNEO commented on pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "SiriusNEO (via GitHub)" <gi...@apache.org>.
SiriusNEO commented on PR #14512:
URL: https://github.com/apache/tvm/pull/14512#issuecomment-1497882881
cc @tqchen @sunggg @Hzfengsy
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tvm-bot commented on pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "tvm-bot (via GitHub)" <gi...@apache.org>.
tvm-bot commented on PR #14512:
URL: https://github.com/apache/tvm/pull/14512#issuecomment-1497880888
<!---bot-comment-->
Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from [Reviewers](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @-ing them in a comment.
<!--bot-comment-ccs-start-->
* cc @quic-sanirudh <sub>See [#10317](https://github.com/apache/tvm/issues/10317) for details</sub><!--bot-comment-ccs-end-->
<sub>Generated by [tvm-bot](https://github.com/apache/tvm/blob/main/ci/README.md#github-actions)</sub>
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14512:
URL: https://github.com/apache/tvm/pull/14512#discussion_r1158899871
##########
src/relax/transform/decompose_ops.cc:
##########
@@ -0,0 +1,267 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*! \file src/relax/transform/decompose_ops.cc */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/nn.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/relax/transform.h>
+
+#include "utils.h"
+
+namespace tvm {
+namespace relax {
+
+TensorStructInfo MatchTensorStructInfo(Expr data) {
+ auto _sinfo = MatchStructInfo<TensorStructInfo>(data);
+ ICHECK(_sinfo.defined()) << "Expect data to be a tensor, but get " << GetStructInfo(data);
+ return _sinfo.value();
+}
+
+Expr ExpandToMatchInput(Expr data, int ndim, Array<Integer> axes) {
+ axes = GetOrderedPositiveAxes(axes, ndim);
+ Array<Integer> expand_axes;
+ for (int i = 0, j = 0; i < ndim; ++i) {
+ if (j < static_cast<int>(axes.size()) && i == axes[j]->value) {
+ ++j;
+ } else {
+ expand_axes.push_back(i);
+ }
+ }
+ return expand_dims(data, expand_axes);
+}
+
+Tuple SimplifyBatchNormInference(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Expr moving_mean = ExpandToMatchInput(call->args[3], sinfo->ndim, {attrs->axis});
+ Expr moving_var = ExpandToMatchInput(call->args[4], sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(moving_var, epsilon));
+ Expr out = divide(subtract(data, moving_mean), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ return Tuple({out, call->args[3], call->args[4]});
+}
+
+Tuple SimplifyBatchNormTraining(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Array<Integer> reduce_axes;
+ for (int i = 0; i < sinfo->ndim; ++i) {
+ if (i != attrs->axis) {
+ reduce_axes.push_back(i);
+ }
+ }
+
+ Expr data_mean = mean(data, reduce_axes, false);
+ Expr data_mean_rs = ExpandToMatchInput(data_mean, sinfo->ndim, {attrs->axis});
+ Expr data_var = variance(data, reduce_axes, false);
+ Expr data_var_rs = ExpandToMatchInput(data_var, sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(data_var_rs, epsilon));
+ Expr out = divide(subtract(data, data_mean_rs), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ Expr moving_mean = call->args[3];
+ Expr moving_var = call->args[4];
+ Expr momentum = MakeConstantScalar(attrs->momentum, sinfo->dtype);
+ Expr one_minus_mom = MakeConstantScalar(1 - attrs->momentum, sinfo->dtype);
+
+ return Tuple({
+ out,
+ add(multiply(one_minus_mom, moving_mean), multiply(momentum, data_mean)),
+ add(multiply(one_minus_mom, moving_var), multiply(momentum, data_var)),
+ });
+}
+
+Expr SimplifyLayerNorm(const Call& call) {
Review Comment:
We don't need to decompose `LayerNorm` since it can be implemented more efficiently as a single op
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14512:
URL: https://github.com/apache/tvm/pull/14512#discussion_r1160029376
##########
src/relax/transform/decompose_ops.cc:
##########
@@ -0,0 +1,267 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*! \file src/relax/transform/decompose_ops.cc */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/nn.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/relax/transform.h>
+
+#include "utils.h"
+
+namespace tvm {
+namespace relax {
+
+TensorStructInfo MatchTensorStructInfo(Expr data) {
+ auto _sinfo = MatchStructInfo<TensorStructInfo>(data);
+ ICHECK(_sinfo.defined()) << "Expect data to be a tensor, but get " << GetStructInfo(data);
+ return _sinfo.value();
+}
+
+Expr ExpandToMatchInput(Expr data, int ndim, Array<Integer> axes) {
+ axes = GetOrderedPositiveAxes(axes, ndim);
+ Array<Integer> expand_axes;
+ for (int i = 0, j = 0; i < ndim; ++i) {
+ if (j < static_cast<int>(axes.size()) && i == axes[j]->value) {
+ ++j;
+ } else {
+ expand_axes.push_back(i);
+ }
+ }
+ return expand_dims(data, expand_axes);
+}
+
+Tuple SimplifyBatchNormInference(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Expr moving_mean = ExpandToMatchInput(call->args[3], sinfo->ndim, {attrs->axis});
+ Expr moving_var = ExpandToMatchInput(call->args[4], sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(moving_var, epsilon));
+ Expr out = divide(subtract(data, moving_mean), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ return Tuple({out, call->args[3], call->args[4]});
+}
+
+Tuple SimplifyBatchNormTraining(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Array<Integer> reduce_axes;
+ for (int i = 0; i < sinfo->ndim; ++i) {
+ if (i != attrs->axis) {
+ reduce_axes.push_back(i);
+ }
+ }
+
+ Expr data_mean = mean(data, reduce_axes, false);
+ Expr data_mean_rs = ExpandToMatchInput(data_mean, sinfo->ndim, {attrs->axis});
+ Expr data_var = variance(data, reduce_axes, false);
+ Expr data_var_rs = ExpandToMatchInput(data_var, sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(data_var_rs, epsilon));
+ Expr out = divide(subtract(data, data_mean_rs), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ Expr moving_mean = call->args[3];
+ Expr moving_var = call->args[4];
+ Expr momentum = MakeConstantScalar(attrs->momentum, sinfo->dtype);
+ Expr one_minus_mom = MakeConstantScalar(1 - attrs->momentum, sinfo->dtype);
+
+ return Tuple({
+ out,
+ add(multiply(one_minus_mom, moving_mean), multiply(momentum, data_mean)),
+ add(multiply(one_minus_mom, moving_var), multiply(momentum, data_var)),
+ });
+}
+
+Expr SimplifyLayerNorm(const Call& call) {
Review Comment:
> Is there any complication for fusion if we have a composite ops in general?
Yes, we need to be aware of the implementation when setting op type (`kOpaque` or `kInjective`). This will cause different fusion results.
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] sunggg commented on a diff in pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "sunggg (via GitHub)" <gi...@apache.org>.
sunggg commented on code in PR #14512:
URL: https://github.com/apache/tvm/pull/14512#discussion_r1158951241
##########
src/relax/transform/decompose_ops.cc:
##########
@@ -0,0 +1,267 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*! \file src/relax/transform/decompose_ops.cc */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/nn.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/relax/transform.h>
+
+#include "utils.h"
+
+namespace tvm {
+namespace relax {
+
+TensorStructInfo MatchTensorStructInfo(Expr data) {
+ auto _sinfo = MatchStructInfo<TensorStructInfo>(data);
+ ICHECK(_sinfo.defined()) << "Expect data to be a tensor, but get " << GetStructInfo(data);
+ return _sinfo.value();
+}
+
+Expr ExpandToMatchInput(Expr data, int ndim, Array<Integer> axes) {
+ axes = GetOrderedPositiveAxes(axes, ndim);
+ Array<Integer> expand_axes;
+ for (int i = 0, j = 0; i < ndim; ++i) {
+ if (j < static_cast<int>(axes.size()) && i == axes[j]->value) {
+ ++j;
+ } else {
+ expand_axes.push_back(i);
+ }
+ }
+ return expand_dims(data, expand_axes);
+}
+
+Tuple SimplifyBatchNormInference(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Expr moving_mean = ExpandToMatchInput(call->args[3], sinfo->ndim, {attrs->axis});
+ Expr moving_var = ExpandToMatchInput(call->args[4], sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(moving_var, epsilon));
+ Expr out = divide(subtract(data, moving_mean), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ return Tuple({out, call->args[3], call->args[4]});
+}
+
+Tuple SimplifyBatchNormTraining(const Call& call) {
+ auto attrs = call->attrs.as<BatchNormAttrs>();
+ ICHECK_NOTNULL(attrs);
+
+ Expr data = call->args[0];
+ TensorStructInfo sinfo = MatchTensorStructInfo(data);
+ Expr gamma = call->args[1];
+ Expr beta = call->args[2];
+
+ Array<Integer> reduce_axes;
+ for (int i = 0; i < sinfo->ndim; ++i) {
+ if (i != attrs->axis) {
+ reduce_axes.push_back(i);
+ }
+ }
+
+ Expr data_mean = mean(data, reduce_axes, false);
+ Expr data_mean_rs = ExpandToMatchInput(data_mean, sinfo->ndim, {attrs->axis});
+ Expr data_var = variance(data, reduce_axes, false);
+ Expr data_var_rs = ExpandToMatchInput(data_var, sinfo->ndim, {attrs->axis});
+
+ // output = (x - mean) / sqrt(var + epsilon) * gamma + beta
+ Expr epsilon = MakeConstantScalar(attrs->epsilon, sinfo->dtype);
+ Expr sqrt_var = sqrt(add(data_var_rs, epsilon));
+ Expr out = divide(subtract(data, data_mean_rs), sqrt_var);
+
+ if (attrs->scale) {
+ out = multiply(out, ExpandToMatchInput(gamma, sinfo->ndim, {attrs->axis}));
+ }
+ if (attrs->center) {
+ out = add(out, ExpandToMatchInput(beta, sinfo->ndim, {attrs->axis}));
+ }
+
+ Expr moving_mean = call->args[3];
+ Expr moving_var = call->args[4];
+ Expr momentum = MakeConstantScalar(attrs->momentum, sinfo->dtype);
+ Expr one_minus_mom = MakeConstantScalar(1 - attrs->momentum, sinfo->dtype);
+
+ return Tuple({
+ out,
+ add(multiply(one_minus_mom, moving_mean), multiply(momentum, data_mean)),
+ add(multiply(one_minus_mom, moving_var), multiply(momentum, data_var)),
+ });
+}
+
+Expr SimplifyLayerNorm(const Call& call) {
Review Comment:
Is there any complication for fusion if we have a composite ops in general? (not specifically about `LayerNorm`). i.e., Is there a case where we cannot fuse the composite op but we can when we decompose?
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 merged pull request #14512: [Unity][Transform] Some Improvements on pass DecomposeOps
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 merged PR #14512:
URL: https://github.com/apache/tvm/pull/14512
--
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.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org