This is an automated email from the ASF dual-hosted git repository.
syfeng pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 38807e9322 [Unity][Transform] SimplifyNormInference (#14221)
38807e9322 is described below
commit 38807e93225c71919651c39adda57beaba51d48a
Author: Chaofan Lin <[email protected]>
AuthorDate: Wed Mar 8 14:20:22 2023 +0800
[Unity][Transform] SimplifyNormInference (#14221)
This PR adds a new pass SimplifyNormInference to unpack the norm operator
into a sequence of operators, which is same as the pass SimplifyInference in
Relay.
---------
Co-authored-by: Siyuan Feng <[email protected]>
---
include/tvm/relax/transform.h | 7 +
python/tvm/relax/transform/transform.py | 14 ++
src/relax/transform/simplify_norm_inference.cc | 131 ++++++++++++++++++
src/relax/transform/utils.h | 105 ++++++++++++++
.../test_transform_simpilify_norm_inference.py | 153 +++++++++++++++++++++
5 files changed, 410 insertions(+)
diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h
index 715c8e56ff..9838fe53b3 100644
--- a/include/tvm/relax/transform.h
+++ b/include/tvm/relax/transform.h
@@ -272,6 +272,13 @@ TVM_DLL Pass RemoveUnusedFunctions(Array<runtime::String>
entry_functions);
TVM_DLL Pass RunCodegen(Optional<Map<String, Map<String, ObjectRef>>>
target_options,
Array<runtime::String> entry_functions);
+/*!
+ * \brief Simplify normalization operators during inference. For example, the
result
+ * of a batch norm which is indexed at tuple index 0 will be unpacked into a
+ * number of simplified operators.
+ * \return The Pass.
+ */
+TVM_DLL Pass SimplifyNormInference();
} // namespace transform
} // namespace relax
} // namespace tvm
diff --git a/python/tvm/relax/transform/transform.py
b/python/tvm/relax/transform/transform.py
index a33ad63093..48560792e1 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -528,6 +528,20 @@ def MetaScheduleTuneIRMod(
return _ffi_api.MetaScheduleTuneIRMod(params, work_dir, max_trials_global)
# type: ignore
+def SimplifyNormInference() -> tvm.ir.transform.Pass:
+ """Simplify normalization operators during inference. For example, the
result
+ of a batch norm which is indexed at tuple index 0 will be unpacked into a
+ number of simplified operators.
+
+ Returns
+ -------
+ ret : tvm.transform.Pass
+ The registered pass
+ """
+
+ return _ffi_api.SimplifyNormInference() # type: ignore
+
+
def _wrap_class_function_pass(pass_cls, pass_info):
"""Wrap a python class as function pass."""
diff --git a/src/relax/transform/simplify_norm_inference.cc
b/src/relax/transform/simplify_norm_inference.cc
new file mode 100644
index 0000000000..545098db28
--- /dev/null
+++ b/src/relax/transform/simplify_norm_inference.cc
@@ -0,0 +1,131 @@
+/*
+ * 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/simplify_norm_inference.cc */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/nn.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);
+}
+
+Expr SimplifyBatchNorm(const CallNode* 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(static_cast<float>(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 out;
+}
+
+/*! \brief A mutator to simplify the normalization inference. */
+class NormInferenceSimplifier : public ExprMutator {
+ public:
+ static Expr Simplify(Expr expr) { return NormInferenceSimplifier()(expr); }
+
+ private:
+ using ExprMutator::VisitExpr_;
+ Expr VisitExpr_(const TupleGetItemNode* op) final {
+ Expr expr = ExprMutator::VisitExpr_(op);
+ op = expr.as<TupleGetItemNode>();
+ ICHECK_NOTNULL(op);
+
+ auto it = batch_norm_map_.find(op->tuple);
+ if (it != batch_norm_map_.end() && op->index == 0) {
+ return (*it).second;
+ } else {
+ return expr;
+ }
+ }
+
+ void VisitBinding_(const VarBindingNode* binding, const CallNode* val) final
{
+ ExprMutator::VisitBinding_(binding, val);
+ if (val->op == Op::Get("relax.nn.batch_norm")) {
+ // NOTE: we won't directly replace the batch_norm call since
+ // the following bindings may depend on the returned moving_mean and
moving_var.
+ // Instead, we will store the unpacked value in the batch_norm_map_, and
replace it
+ // at the TupleGetItemNode. And the original batch_norm call will be
removed in the
+ // follow-up pass `RemoveAllUnused`
+ batch_norm_map_.Set(binding->var, SimplifyBatchNorm(val));
+ }
+ }
+
+ private:
+ /*! \brief The mapping from binding var of batch_norm to the unpacked value.
*/
+ Map<Expr, Expr> batch_norm_map_;
+};
+
+namespace transform {
+Pass SimplifyNormInference() {
+ runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)>
pass_func =
+ [=](Function f, IRModule m, PassContext pc) {
+ f = Downcast<Function>(NormInferenceSimplifier::Simplify(f));
+ // Remove original batch_norm op if it's not used.
+ return RemoveAllUnused(f);
+ };
+ return CreateFunctionPass(/*pass_function=*/pass_func, //
+ /*opt_level=*/0, //
+ /*pass_name=*/"SimplifyNormInference", //
+ /*required=*/{});
+}
+
+TVM_REGISTER_GLOBAL("relax.transform.SimplifyNormInference").set_body_typed(SimplifyNormInference);
+
+} // namespace transform
+} // namespace relax
+} // namespace tvm
diff --git a/src/relax/transform/utils.h b/src/relax/transform/utils.h
index d94c1e3b3e..463e69d56c 100644
--- a/src/relax/transform/utils.h
+++ b/src/relax/transform/utils.h
@@ -24,14 +24,32 @@
#ifndef TVM_RELAX_TRANSFORM_UTILS_H_
#define TVM_RELAX_TRANSFORM_UTILS_H_
+#include <builtin_fp16.h>
#include <tvm/ir/module.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/expr_functor.h>
+#include <algorithm>
#include <string>
#include <unordered_map>
+#include <vector>
#include "../../relay/analysis/graph_partitioner.h"
+#include "../../support/array.h"
+#include "../op/nn/convolution.h"
+#include "../op/nn/nn.h"
+#include "../op/nn/pooling.h"
+#include "../op/tensor/binary.h"
+#include "../op/tensor/create.h"
+#include "../op/tensor/datatype.h"
+#include "../op/tensor/index.h"
+#include "../op/tensor/linear_algebra.h"
+#include "../op/tensor/manipulate.h"
+#include "../op/tensor/search.h"
+#include "../op/tensor/set.h"
+#include "../op/tensor/statistical.h"
+#include "../op/tensor/ternary.h"
+#include "../op/tensor/unary.h"
namespace tvm {
namespace relax {
@@ -116,6 +134,93 @@ IRModule MakeGroupedFunctions(
const std::unordered_map<const Object*, relay::GraphPartitioner::Group*>&
partition,
bool lift_constants = true);
+/*!
+ * \brief Check if the given StructInfo is a nested tensor StructInfo
satisfying the given
+ * condition f_condition.
+ * \param sinfo The StructInfo to be checked.
+ * \param f_condition The condition function for each leaf StructInfo with
signature
+ * `bool f_condition(TensorStructInfo)`.
+ * \tparam FType The condition function type.
+ * \return true if the given StructInfo is a nested tensor satisfying the
given f_condition.
+ */
+template <typename FType>
+bool IsNestedTensorConditioned(const StructInfo& sinfo, FType f_condition) {
+ if (const auto* tensor_sinfo = sinfo.as<TensorStructInfoNode>()) {
+ return f_condition(GetRef<TensorStructInfo>(tensor_sinfo));
+ } else if (const auto* tuple_sinfo = sinfo.as<TupleStructInfoNode>()) {
+ return !std::any_of(
+ tuple_sinfo->fields.begin(), tuple_sinfo->fields.end(),
+ [&](const StructInfo& field) { return
!IsNestedTensorConditioned(field, f_condition); });
+ }
+ return false;
+}
+
+/*!
+ * \brief Create a Constant with a scalar
+ *
+ * \param dtype The data type.
+ * \param value The value of the scalar.
+ * \return A Constant.
+ */
+template <typename T>
+inline Constant MakeConstantScalar(T value, DataType dtype) {
+ runtime::NDArray arr = runtime::NDArray::Empty({}, dtype, {kDLCPU, 0});
+ if (dtype == DataType::Float(32)) {
+ *static_cast<float*>(arr->data) = static_cast<float>(value);
+ } else if (dtype == DataType::Float(64)) {
+ *static_cast<double*>(arr->data) = static_cast<double>(value);
+ } else if (dtype == DataType::Int(32)) {
+ *static_cast<int32_t*>(arr->data) = static_cast<int32_t>(value);
+ } else if (dtype == DataType::Int(64)) {
+ *static_cast<int64_t*>(arr->data) = static_cast<int64_t>(value);
+ } else if (dtype == DataType::UInt(1)) {
+ *static_cast<bool*>(arr->data) = static_cast<bool>(value);
+ } else if (dtype == DataType::UInt(8)) {
+ *static_cast<uint8_t*>(arr->data) = static_cast<uint8_t>(value);
+ } else if (dtype == DataType::UInt(16)) {
+ *static_cast<uint16_t*>(arr->data) = static_cast<uint16_t>(value);
+ } else if (dtype == DataType::UInt(32)) {
+ *static_cast<uint32_t*>(arr->data) = static_cast<uint32_t>(value);
+ } else if (dtype == DataType::UInt(64)) {
+ *static_cast<uint64_t*>(arr->data) = static_cast<uint64_t>(value);
+ } else if (dtype == DataType::Int(8)) {
+ *static_cast<int8_t*>(arr->data) = static_cast<int8_t>(value);
+ } else if (dtype == DataType::Int(16)) {
+ *static_cast<int16_t*>(arr->data) = static_cast<int16_t>(value);
+ } else if (dtype == DataType::Int(32)) {
+ *static_cast<int32_t*>(arr->data) = static_cast<int32_t>(value);
+ } else if (dtype == DataType::Int(64)) {
+ *static_cast<int64_t*>(arr->data) = static_cast<int64_t>(value);
+ } else if (dtype == DataType::Float(16)) {
+ // convert to float16 storage is uint16_t
+ *static_cast<uint16_t*>(arr->data) =
+ __truncXfYf2__<float, uint32_t, 23, uint16_t, uint16_t,
10>(static_cast<float>(value));
+ } else if (dtype == DataType::BFloat(16)) {
+ // convert to bfloat16 storage is uint16_t
+ *static_cast<uint16_t*>(arr->data) =
+ __truncXfYf2__<float, uint32_t, 23, uint16_t, uint16_t,
7>(static_cast<float>(value));
+ } else {
+ LOG(FATAL) << "Unsupported dtype " << dtype;
+ }
+ return Constant(arr);
+}
+
+inline Array<Integer> GetOrderedPositiveAxes(const Array<Integer>& axes, int
ndim) {
+ std::vector<int64_t> ret;
+ ret.reserve(axes.size());
+ for (const auto& axis : axes) {
+ int64_t axis_val = axis->value;
+ if (axis_val < 0) {
+ axis_val += ndim;
+ }
+ ICHECK(axis_val >= 0 && axis_val < ndim) << "axis " << axis << " is out of
bounds for array of "
+ << "dimension " << ndim;
+ ret.push_back(axis_val);
+ }
+ std::sort(ret.begin(), ret.end());
+ return support::AsArray<int64_t, Integer>(ret);
+}
+
} // namespace relax
} // namespace tvm
diff --git a/tests/python/relax/test_transform_simpilify_norm_inference.py
b/tests/python/relax/test_transform_simpilify_norm_inference.py
new file mode 100644
index 0000000000..3c981ba035
--- /dev/null
+++ b/tests/python/relax/test_transform_simpilify_norm_inference.py
@@ -0,0 +1,153 @@
+# 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.
+
+from typing import Union
+
+import tvm
+import tvm.script
+import tvm.testing
+from tvm import IRModule, relax
+from tvm.relax import Function
+from tvm.script import relax as R
+
+
+def _check(before: Union[Function, IRModule], expected: Union[Function,
IRModule]):
+ if isinstance(before, Function):
+ before = IRModule({"main": before})
+ if isinstance(expected, Function):
+ expected = IRModule({"main": expected})
+ after = relax.transform.SimplifyNormInference()(before)
+ tvm.ir.assert_structural_equal(expected, after)
+
+
+def test_batch_norm_simple():
+ @R.function
+ def before(
+ x: R.Tensor((1, 64, 112, 112), "float32"),
+ gamma: R.Tensor((64,), "float32"),
+ beta: R.Tensor((64,), "float32"),
+ moving_mean: R.Tensor((64,), "float32"),
+ moving_var: R.Tensor((64,), "float32"),
+ ):
+ with R.dataflow():
+ bn = R.nn.batch_norm(
+ x,
+ gamma,
+ beta,
+ moving_mean,
+ moving_var,
+ axis=1,
+ epsilon=1e-5,
+ center=True,
+ scale=True,
+ )
+ gv = bn[0]
+ R.output(gv)
+ return gv
+
+ @R.function
+ def expected(
+ x: R.Tensor((1, 64, 112, 112), "float32"),
+ gamma: R.Tensor((64,), "float32"),
+ beta: R.Tensor((64,), "float32"),
+ moving_mean: R.Tensor((64,), "float32"),
+ moving_var: R.Tensor((64,), "float32"),
+ ):
+ with R.dataflow():
+ mean = R.expand_dims(moving_mean, axis=[0, 2, 3])
+ out = x - mean
+ var = R.expand_dims(moving_var, axis=[0, 2, 3])
+ var_eps = var + R.const(1e-05, "float32")
+ sqrt_var = R.sqrt(var_eps)
+ div = R.divide(out, sqrt_var)
+ new_gamma = R.expand_dims(gamma, axis=[0, 2, 3])
+ out = div * new_gamma
+ new_beta = R.expand_dims(beta, axis=[0, 2, 3])
+ out = out + new_beta
+ R.output(out)
+ return out
+
+ _check(before, expected)
+
+
+def test_batch_norm_complex():
+ @R.function
+ def before(
+ x: R.Tensor((1, 64, 112, 112), "float32"),
+ gamma: R.Tensor((64,), "float32"),
+ beta: R.Tensor((64,), "float32"),
+ moving_mean: R.Tensor((64,), "float32"),
+ moving_var: R.Tensor((64,), "float32"),
+ ):
+ with R.dataflow():
+ bn = R.nn.batch_norm(
+ x,
+ gamma,
+ beta,
+ moving_mean,
+ moving_var,
+ axis=1,
+ epsilon=1e-5,
+ center=True,
+ scale=True,
+ )
+ gv0 = bn[0]
+ gv1 = bn[1]
+ R.output(gv0, gv1)
+ return gv0, gv1
+
+ @R.function
+ def expected(
+ x: R.Tensor((1, 64, 112, 112), "float32"),
+ gamma: R.Tensor((64,), "float32"),
+ beta: R.Tensor((64,), "float32"),
+ moving_mean: R.Tensor((64,), "float32"),
+ moving_var: R.Tensor((64,), "float32"),
+ ):
+ with R.dataflow():
+ # bn[1] is used, so we need to keep the original batch_norm
+ # NOTE: It's a rare case, so that we don't optimize it for now
+ bn = R.nn.batch_norm(
+ x,
+ gamma,
+ beta,
+ moving_mean,
+ moving_var,
+ axis=1,
+ epsilon=1e-5,
+ center=True,
+ scale=True,
+ )
+ mean = R.expand_dims(moving_mean, axis=[0, 2, 3])
+ out = x - mean
+ var = R.expand_dims(moving_var, axis=[0, 2, 3])
+ var_eps = var + R.const(1e-05, "float32")
+ sqrt_var = R.sqrt(var_eps)
+ div = R.divide(out, sqrt_var)
+ new_gamma = R.expand_dims(gamma, axis=[0, 2, 3])
+ out = div * new_gamma
+ new_beta = R.expand_dims(beta, axis=[0, 2, 3])
+ out = out + new_beta
+ gv1 = bn[1]
+ R.output(out, gv1)
+ return out, gv1
+
+ _check(before, expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
