From 118d87a12916ec99924c636e277dc511f25269de Mon Sep 17 00:00:00 2001 From: Jinkun Lin Date: Thu, 3 Feb 2022 00:04:45 +0000 Subject: [PATCH 1/9] Move function body to .cc file. --- src/relay/transforms/infer_layout_utils.cc | 134 +++++++++++++++++++++ src/relay/transforms/infer_layout_utils.h | 115 +----------------- 2 files changed, 136 insertions(+), 113 deletions(-) create mode 100644 src/relay/transforms/infer_layout_utils.cc diff --git a/src/relay/transforms/infer_layout_utils.cc b/src/relay/transforms/infer_layout_utils.cc new file mode 100644 index 000000000000..47366b6e9455 --- /dev/null +++ b/src/relay/transforms/infer_layout_utils.cc @@ -0,0 +1,134 @@ +#include "infer_layout_utils.h" + +#include +#include +#include + +#include +#include +#include + +#include "pattern_utils.h" +#include "tvm/runtime/logging.h" + +namespace tvm { +namespace relay { + +inline Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& old_layout, + const Array& old_shape) { + // For each subordinate axis + // 1) Find the corresponding dual axis. + // 2) Find the Index of this dual axis in old_layout. + // 3) Find the shape of the that axis in old_shape. + // 4) a) Adjust factor to 1, if that shape is 1. b) Else retain the factor. + std::string new_layout; + for (auto axis : src_layout->axes) { + if (!LayoutAxis::Get(axis).IsPrimal()) { + // 1) Find the corresponding dual axis + const auto& dual_axis = LayoutAxis::Get(axis).ToPrimal(); + + // 2) Find the index of this dual axis in old_layout + int old_axis = old_layout.IndexOf(dual_axis); + + // 3) Find the shape of this index in old_shape + auto shape_val = old_shape[old_axis]; + + // 4) a) Check if this shape element is 1. + bool is_shape_one = false; + if (auto* shape_int = shape_val.as()) { + if (shape_int->value == 1) { + new_layout += "1"; + is_shape_one = true; + } + } + + // 4) b) If shape is not 1, retain the factor. + if (!is_shape_one) { + auto new_shape_val = src_layout.FactorOf(dual_axis); + new_layout += std::to_string(new_shape_val); + } + } + new_layout += LayoutAxis::Get(axis).name(); + } + return Layout(new_layout); +} +inline std::pair, Array> BinaryBroadcastLayoutHelper( + const Attrs& attrs, const Array& new_in_layouts, const Array& old_in_layouts, + const Array& old_in_types) { + Array layouts; + Array> old_in_shapes; + for (auto old_in_t : old_in_types) { + ICHECK(old_in_t.as()); + old_in_shapes.push_back(old_in_t.as()->shape); + } + + if (new_in_layouts.defined()) { + layouts.Assign(new_in_layouts.begin(), new_in_layouts.end()); + } else { + layouts.Assign(old_in_layouts.begin(), old_in_layouts.end()); + } + + std::pair, Array> out_default{{Layout::Undef(), Layout::Undef()}, + {Layout::Undef()}}; + + if (!layouts[0].defined() && !layouts[1].defined()) { + // both undefined, infer fails + return out_default; + } else if (!layouts[0].defined() || !layouts[1].defined()) { + // only one is defined, use shape information to help infer + int defined_idx = layouts[0].defined() ? 0 : 1; + int undef_idx = 1 - defined_idx; + + if (old_in_shapes[defined_idx].size() >= old_in_shapes[undef_idx].size()) { + layouts.Set(undef_idx, layouts[defined_idx].SubLayout(old_in_shapes[defined_idx].size() - + old_in_shapes[undef_idx].size(), + old_in_shapes[undef_idx].size())); + return {layouts, {layouts[defined_idx]}}; + } else { + // only know the tensor with smaller dimensions, + // so we cannot infer the final broadcasted output. + // fails in this case. + return out_default; + } + } else if (layouts[0].defined() && layouts[1].defined() && + (layouts[0].ndim() == 0 || layouts[1].ndim() == 0)) { + int scalar = layouts[0].ndim() == 0 ? 0 : 1; + return {layouts, {layouts[1 - scalar]}}; + } else { + // Set the layout of the larger dimension. If one dimension size is lower, we call expand dims + // while transforming layout. + int large_idx = layouts[0].ndim_primal() >= layouts[1].ndim_primal() ? 0 : 1; + int small_idx = 1 - large_idx; + Layout ret = layouts[large_idx]; + + if (old_in_layouts[0].Equals(old_in_layouts[1])) { + // Support scenarios where original operands were of type [N, H, W, C] and [N, H, W, 1] + // In this case, we might have NCHW16c coming for 1 operand. However, the other operand does + // not have enough C dimension. To reuse broadcasting, we would want to use NCHW1c for the + // second operand. The following section of code walks through the layouts and shapes to + // perform that operation. + // a in NCHWC16c + // b in NHW1 + // b = layout_transform(b) from NHW1 -> NCHW1c + // add(a, b) + auto old_small_shape = old_in_shapes[small_idx]; + auto old_small_layout = old_in_layouts[small_idx]; + auto new_small_layout = + AdjustSubordinateFactors(layouts[large_idx], old_small_layout, old_small_shape); + layouts.Set(small_idx, new_small_layout); + } else { + // Support scenarios where original operands were of type [N, H, W, C] and [C]. In this case, + // while transforming the layout, we expand dims to make C go to NHWC, and then use the + // modified layout of the first operator to call the layout transform. E.g. + // a in NCHWC16c + // b in C + // b = expand_dims(b) from C -> NHWC + // b = layout_transform(b) from NHWC -> NCHW16c + // add(a, b) + layouts.Set(small_idx, ret); + } + return {layouts, {ret}}; + } +} +} // namespace relay +} // namespace tvm \ No newline at end of file diff --git a/src/relay/transforms/infer_layout_utils.h b/src/relay/transforms/infer_layout_utils.h index 76d6aa646f4c..ba0c5783105e 100644 --- a/src/relay/transforms/infer_layout_utils.h +++ b/src/relay/transforms/infer_layout_utils.h @@ -48,43 +48,7 @@ namespace relay { * \return The adjusted Layout. */ inline Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& old_layout, - const Array& old_shape) { - // For each subordinate axis - // 1) Find the corresponding dual axis. - // 2) Find the Index of this dual axis in old_layout. - // 3) Find the shape of the that axis in old_shape. - // 4) a) Adjust factor to 1, if that shape is 1. b) Else retain the factor. - std::string new_layout; - for (auto axis : src_layout->axes) { - if (!LayoutAxis::Get(axis).IsPrimal()) { - // 1) Find the corresponding dual axis - const auto& dual_axis = LayoutAxis::Get(axis).ToPrimal(); - - // 2) Find the index of this dual axis in old_layout - int old_axis = old_layout.IndexOf(dual_axis); - - // 3) Find the shape of this index in old_shape - auto shape_val = old_shape[old_axis]; - - // 4) a) Check if this shape element is 1. - bool is_shape_one = false; - if (auto* shape_int = shape_val.as()) { - if (shape_int->value == 1) { - new_layout += "1"; - is_shape_one = true; - } - } - - // 4) b) If shape is not 1, retain the factor. - if (!is_shape_one) { - auto new_shape_val = src_layout.FactorOf(dual_axis); - new_layout += std::to_string(new_shape_val); - } - } - new_layout += LayoutAxis::Get(axis).name(); - } - return Layout(new_layout); -} + const Array& old_shape); /* * \brief An output structure to hold results from FInferCorrectLayout calls. @@ -159,82 +123,7 @@ inline InferCorrectLayoutOutput ElemwiseArbitraryLayout( inline std::pair, Array> BinaryBroadcastLayoutHelper( const Attrs& attrs, const Array& new_in_layouts, const Array& old_in_layouts, - const Array& old_in_types) { - Array layouts; - Array> old_in_shapes; - for (auto old_in_t : old_in_types) { - ICHECK(old_in_t.as()); - old_in_shapes.push_back(old_in_t.as()->shape); - } - - if (new_in_layouts.defined()) { - layouts.Assign(new_in_layouts.begin(), new_in_layouts.end()); - } else { - layouts.Assign(old_in_layouts.begin(), old_in_layouts.end()); - } - - std::pair, Array> out_default{{Layout::Undef(), Layout::Undef()}, - {Layout::Undef()}}; - - if (!layouts[0].defined() && !layouts[1].defined()) { - // both undefined, infer fails - return out_default; - } else if (!layouts[0].defined() || !layouts[1].defined()) { - // only one is defined, use shape information to help infer - int defined_idx = layouts[0].defined() ? 0 : 1; - int undef_idx = 1 - defined_idx; - - if (old_in_shapes[defined_idx].size() >= old_in_shapes[undef_idx].size()) { - layouts.Set(undef_idx, layouts[defined_idx].SubLayout(old_in_shapes[defined_idx].size() - - old_in_shapes[undef_idx].size(), - old_in_shapes[undef_idx].size())); - return {layouts, {layouts[defined_idx]}}; - } else { - // only know the tensor with smaller dimensions, - // so we cannot infer the final broadcasted output. - // fails in this case. - return out_default; - } - } else if (layouts[0].defined() && layouts[1].defined() && - (layouts[0].ndim() == 0 || layouts[1].ndim() == 0)) { - int scalar = layouts[0].ndim() == 0 ? 0 : 1; - return {layouts, {layouts[1 - scalar]}}; - } else { - // Set the layout of the larger dimension. If one dimension size is lower, we call expand dims - // while transforming layout. - int large_idx = layouts[0].ndim_primal() >= layouts[1].ndim_primal() ? 0 : 1; - int small_idx = 1 - large_idx; - Layout ret = layouts[large_idx]; - - if (old_in_layouts[0].Equals(old_in_layouts[1])) { - // Support scenarios where original operands were of type [N, H, W, C] and [N, H, W, 1] - // In this case, we might have NCHW16c coming for 1 operand. However, the other operand does - // not have enough C dimension. To reuse broadcasting, we would want to use NCHW1c for the - // second operand. The following section of code walks through the layouts and shapes to - // perform that operation. - // a in NCHWC16c - // b in NHW1 - // b = layout_transform(b) from NHW1 -> NCHW1c - // add(a, b) - auto old_small_shape = old_in_shapes[small_idx]; - auto old_small_layout = old_in_layouts[small_idx]; - auto new_small_layout = - AdjustSubordinateFactors(layouts[large_idx], old_small_layout, old_small_shape); - layouts.Set(small_idx, new_small_layout); - } else { - // Support scenarios where original operands were of type [N, H, W, C] and [C]. In this case, - // while transforming the layout, we expand dims to make C go to NHWC, and then use the - // modified layout of the first operator to call the layout transform. E.g. - // a in NCHWC16c - // b in C - // b = expand_dims(b) from C -> NHWC - // b = layout_transform(b) from NHWC -> NCHW16c - // add(a, b) - layouts.Set(small_idx, ret); - } - return {layouts, {ret}}; - } -} + const Array& old_in_types); /*! \brief Infer layout for binary broadcast operators */ inline InferCorrectLayoutOutput BinaryBroadcastLayout(const Attrs& attrs, From 53bb9242a6be404457e3e0ba5b2509b1efc33c0a Mon Sep 17 00:00:00 2001 From: Jinkun Lin Date: Thu, 3 Feb 2022 04:35:40 +0000 Subject: [PATCH 2/9] fix broadcast infer layout --- src/relay/transforms/infer_layout_utils.cc | 242 ++++++++++++++++----- src/relay/transforms/infer_layout_utils.h | 22 +- 2 files changed, 202 insertions(+), 62 deletions(-) diff --git a/src/relay/transforms/infer_layout_utils.cc b/src/relay/transforms/infer_layout_utils.cc index 47366b6e9455..21eae6034281 100644 --- a/src/relay/transforms/infer_layout_utils.cc +++ b/src/relay/transforms/infer_layout_utils.cc @@ -1,9 +1,29 @@ +/* + * 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. + */ + #include "infer_layout_utils.h" #include #include #include +#include #include #include #include @@ -14,31 +34,39 @@ namespace tvm { namespace relay { -inline Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& old_layout, - const Array& old_shape) { +Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& old_layout, + const Array& old_shape) { // For each subordinate axis // 1) Find the corresponding dual axis. // 2) Find the Index of this dual axis in old_layout. // 3) Find the shape of the that axis in old_shape. // 4) a) Adjust factor to 1, if that shape is 1. b) Else retain the factor. + DLOG(INFO) << "AdjustSubordinateFactors" + << "src_layout: " << src_layout << " old_layout: " << old_layout + << " old_shape: " << old_shape << std::endl; std::string new_layout; for (auto axis : src_layout->axes) { if (!LayoutAxis::Get(axis).IsPrimal()) { + bool is_shape_one = false; // 1) Find the corresponding dual axis const auto& dual_axis = LayoutAxis::Get(axis).ToPrimal(); // 2) Find the index of this dual axis in old_layout int old_axis = old_layout.IndexOf(dual_axis); - // 3) Find the shape of this index in old_shape - auto shape_val = old_shape[old_axis]; + if (old_axis == -1) { + new_layout += "1"; + is_shape_one = true; + } else { + // 3) Find the shape of this index in old_shape + auto shape_val = old_shape[old_axis]; - // 4) a) Check if this shape element is 1. - bool is_shape_one = false; - if (auto* shape_int = shape_val.as()) { - if (shape_int->value == 1) { - new_layout += "1"; - is_shape_one = true; + // 4) a) Check if this shape element is 1. + if (auto* shape_int = shape_val.as()) { + if (shape_int->value == 1) { + new_layout += "1"; + is_shape_one = true; + } } } @@ -52,83 +80,179 @@ inline Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& o } return Layout(new_layout); } -inline std::pair, Array> BinaryBroadcastLayoutHelper( +bool Isomorphic(const Layout& lhs, const Layout& rhs) { + DLOG(INFO) << "Isomorphic: " + << "lhs: " << lhs << " rhs: " << rhs << std::endl; + ICHECK(lhs.defined()); + ICHECK(rhs.defined()); + if (lhs->axes.size() != rhs->axes.size()) return false; + std::map map_to, map_back; + for (size_t i = 0; i < lhs->axes.size(); ++i) { + auto& lhs_axis = LayoutAxis::Get(lhs->axes[i]); + auto& rhs_axis = LayoutAxis::Get(rhs->axes[i]); + std::string name_lhs = lhs_axis.name(); + std::string name_rhs = rhs_axis.name(); + if (lhs_axis.IsPrimal() != rhs_axis.IsPrimal()) return false; + + auto it = map_to.find(name_lhs); + if (it == map_to.end()) + map_to[name_lhs] = name_rhs; + else if (it->second != name_rhs) + return false; + + it = map_back.find(name_rhs); + if (it == map_back.end()) + map_back[name_rhs] = name_lhs; + else if (it->second != name_lhs) + return false; + if (!lhs_axis.IsPrimal() && lhs.FactorOf(lhs_axis) != rhs.FactorOf(rhs_axis)) return false; + } + return true; +} +Layout TryTransformLike(const Layout& old, const Layout& ref_old, const Layout& ref_new) { + DLOG(INFO) << "transform_layout: old = " << old << ", ref_new = " << ref_new + << ", ref_old = " << ref_old << std::endl; + ICHECK(ref_old.defined()); + ICHECK(ref_new.defined()); + ICHECK(old.defined()); + + { // check if old and ref_old are similar enough such that it's + // compatible for the transform ref_old -> ref_new + const Layout& large = ref_old.ndim() > old.ndim() ? ref_old : old; + const Layout& small = large == ref_old ? old : ref_old; + Layout large_sublayout = large.SubLayout(large.ndim() - small.ndim(), small.ndim()), + rest_sublayout = large.SubLayout(0, large.ndim() - small.ndim()); + bool orthorgonal = true; + for (auto i : rest_sublayout->axes) + if (large_sublayout.IndexOf(LayoutAxis::Get(i).ToPrimal()) != -1 || + large_sublayout.IndexOf(LayoutAxis::Get(i).ToSubordinate()) != -1) { + orthorgonal = false; + break; + } + if (!orthorgonal || !Isomorphic(large_sublayout, small)) + return Layout::Undef(); // For now this case is not supported. + } + + // `old` is compatible. Now learn the axis name mapping between `old` and `ref_old` + int mapping[26]; + bool used[26]; + memset(mapping, -1, sizeof mapping); + memset(used, 0, sizeof used); + auto find_unused = [&](char preference) -> char { + if (!used[preference - 'A']) return preference; // preference unused + for (int i = 0; i < 26; ++i) + if (!used[i]) return 'A' + i; + LOG(FATAL) << "All letters are used"; + return 0; + }; + for (int j = old->axes.size() - 1, i = ref_old->axes.size() - 1; j >= 0; --i, --j) { + char name_ref = LayoutAxis::Get(ref_old->axes[i]).ToPrimal().name()[0]; + char name = LayoutAxis::Get(old->axes[j]).ToPrimal().name()[0]; + mapping[name_ref - 'A'] = name - 'A'; + used[name - 'A'] = true; + } + for (int i = ref_old->axes.size() - 1; i >= 0; --i) { + char name_ref = LayoutAxis::Get(ref_old->axes[i]).ToPrimal().name()[0]; + int name = mapping[name_ref - 'A']; + if (name == -1) { + mapping[name_ref - 'A'] = find_unused(name_ref) - 'A'; + used[mapping[name_ref - 'A']] = true; + } + } + + // apply the mapping to rename `ref_new` + std::string new_layout; + for (auto c : std::string(ref_new->name)) { + if (c >= 'A' && c <= 'Z') { + ICHECK(mapping[c - 'A'] != -1); + new_layout += mapping[c - 'A'] + 'A'; + } else if (c >= 'a' && c <= 'z') { + ICHECK(mapping[c - 'a'] != -1); + new_layout += mapping[c - 'a'] + 'a'; + } else { + new_layout += c; + } + } + DLOG(INFO) << "new_layout = " << new_layout << std::endl; + return Layout(new_layout); +} + +std::pair, Array> BinaryBroadcastLayoutHelper( const Attrs& attrs, const Array& new_in_layouts, const Array& old_in_layouts, const Array& old_in_types) { + // Two steps. Step (2) only executes if the function is called after rewrite. + // (1) infer input layouts before rewrite + // (2) if some input layouts are changed by its producer after rewrite, rewrite the other + // layout to make sure it's changed in the same way, so that they are still broadcastable. Array layouts; Array> old_in_shapes; for (auto old_in_t : old_in_types) { ICHECK(old_in_t.as()); old_in_shapes.push_back(old_in_t.as()->shape); } + int old_large_idx = old_in_shapes[0].size() >= old_in_shapes[1].size() ? 0 : 1; - if (new_in_layouts.defined()) { - layouts.Assign(new_in_layouts.begin(), new_in_layouts.end()); - } else { - layouts.Assign(old_in_layouts.begin(), old_in_layouts.end()); - } + layouts.Assign(old_in_layouts.begin(), old_in_layouts.end()); + // always operate on the original layouts first for consistency - std::pair, Array> out_default{{Layout::Undef(), Layout::Undef()}, - {Layout::Undef()}}; + std::pair, Array> out, + out_default{{Layout::Undef(), Layout::Undef()}, {Layout::Undef()}}; if (!layouts[0].defined() && !layouts[1].defined()) { // both undefined, infer fails - return out_default; + out = out_default; } else if (!layouts[0].defined() || !layouts[1].defined()) { // only one is defined, use shape information to help infer int defined_idx = layouts[0].defined() ? 0 : 1; int undef_idx = 1 - defined_idx; if (old_in_shapes[defined_idx].size() >= old_in_shapes[undef_idx].size()) { + // TODO(lazycal): handle the case when the sublayout contains subcoordinate of factor one but + // the other tensor has the corresponding dimension size other than one. + // E.g. defined's shape = [x, x, x, x, 1] in NCHW1c and undefined's shape = [3] layouts.Set(undef_idx, layouts[defined_idx].SubLayout(old_in_shapes[defined_idx].size() - old_in_shapes[undef_idx].size(), old_in_shapes[undef_idx].size())); - return {layouts, {layouts[defined_idx]}}; + out = {layouts, {layouts[defined_idx]}}; } else { // only know the tensor with smaller dimensions, // so we cannot infer the final broadcasted output. // fails in this case. - return out_default; + out = out_default; } - } else if (layouts[0].defined() && layouts[1].defined() && - (layouts[0].ndim() == 0 || layouts[1].ndim() == 0)) { - int scalar = layouts[0].ndim() == 0 ? 0 : 1; - return {layouts, {layouts[1 - scalar]}}; } else { - // Set the layout of the larger dimension. If one dimension size is lower, we call expand dims - // while transforming layout. - int large_idx = layouts[0].ndim_primal() >= layouts[1].ndim_primal() ? 0 : 1; - int small_idx = 1 - large_idx; - Layout ret = layouts[large_idx]; - - if (old_in_layouts[0].Equals(old_in_layouts[1])) { - // Support scenarios where original operands were of type [N, H, W, C] and [N, H, W, 1] - // In this case, we might have NCHW16c coming for 1 operand. However, the other operand does - // not have enough C dimension. To reuse broadcasting, we would want to use NCHW1c for the - // second operand. The following section of code walks through the layouts and shapes to - // perform that operation. - // a in NCHWC16c - // b in NHW1 - // b = layout_transform(b) from NHW1 -> NCHW1c - // add(a, b) - auto old_small_shape = old_in_shapes[small_idx]; - auto old_small_layout = old_in_layouts[small_idx]; - auto new_small_layout = - AdjustSubordinateFactors(layouts[large_idx], old_small_layout, old_small_shape); - layouts.Set(small_idx, new_small_layout); - } else { - // Support scenarios where original operands were of type [N, H, W, C] and [C]. In this case, - // while transforming the layout, we expand dims to make C go to NHWC, and then use the - // modified layout of the first operator to call the layout transform. E.g. - // a in NCHWC16c - // b in C - // b = expand_dims(b) from C -> NHWC - // b = layout_transform(b) from NHWC -> NCHW16c - // add(a, b) - layouts.Set(small_idx, ret); - } - return {layouts, {ret}}; + // when both are defined, return the larger one + out = {layouts, {layouts[old_large_idx]}}; } + if (!new_in_layouts.defined()) return out; + // Step (2) rewrite the layouts to make them broadcastable again. + Layout ret = new_in_layouts[old_large_idx]; + int large_idx = new_in_layouts[0].ndim_primal() >= new_in_layouts[1].ndim_primal() ? 0 : 1; + int small_idx = 1 - large_idx; + // start adjusting + + // Apply a greedy strategy that always transform the small layout in the same way as the + // large layout is transformed, if possible. + Layout tgt_layout = + TryTransformLike(layouts[small_idx], layouts[large_idx], new_in_layouts[large_idx]); + if (!tgt_layout.defined()) return out_default; // fallback + + // Support scenarios where original operands were of type [N, H, W, C] and [N, H, W, 1] + // In this case, we might have NCHW16c coming for 1 operand. However, the other operand does + // not have enough C dimension. To reuse broadcasting, we would want to use NCHW1c for the + // second operand. The following section of code walks through the layouts and shapes to + // perform that operation. + // a in NCHWC16c + // b in NHW1 + // b = layout_transform(b) from NHW1 -> NCHW1c + // add(a, b) + auto old_small_shape = old_in_shapes[small_idx]; + auto old_small_layout = layouts[small_idx]; + auto new_small_layout = AdjustSubordinateFactors(tgt_layout, old_small_layout, old_small_shape); + layouts.Set(large_idx, new_in_layouts[large_idx]); + layouts.Set(small_idx, new_small_layout); + return {layouts, {ret}}; } + } // namespace relay -} // namespace tvm \ No newline at end of file +} // namespace tvm diff --git a/src/relay/transforms/infer_layout_utils.h b/src/relay/transforms/infer_layout_utils.h index ba0c5783105e..68c407d70a98 100644 --- a/src/relay/transforms/infer_layout_utils.h +++ b/src/relay/transforms/infer_layout_utils.h @@ -47,9 +47,25 @@ namespace relay { * \param old_shape The shape of the original tensor. * \return The adjusted Layout. */ -inline Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& old_layout, - const Array& old_shape); +Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& old_layout, + const Array& old_shape); +bool Isomorphic(const Layout& lhs, const Layout& rhs); +/*! + * \brief Try transforming `old` in as the smae way as how`ref_old` is transformed to `ref_new`. + * `old` and `ref_old` are expected to describe two broadcastable tensors. Layout with fewer rank + * will be expanded. For example, + * if old = 'NW', ref_old = 'NC', ref_new = 'NC1c', then the result is 'NW1w'; + * if old = 'W', ref_old = 'NC', ref_new = 'NC1c', then the result is 'NW1w'. + * When `old` and `ref_old` are isomorphic (same structure, only differ in naming), the transform + * is guaranteed to succeed, in which case the function is simply renaming the axes of `ref_new` + * to conform to `old`'s naming. + * \param old The layout to be transformed. + * \param ref_old The reference layout before transform. + * \param ref_new The reference layout after transform. + * \return The transformed layout. + */ +Layout TryTransformLike(const Layout& old, const Layout& ref_old, const Layout& ref_new); /* * \brief An output structure to hold results from FInferCorrectLayout calls. * \tparam input_layouts Inferred input layouts. @@ -121,7 +137,7 @@ inline InferCorrectLayoutOutput ElemwiseArbitraryLayout( return InferCorrectLayoutOutput(Array(old_in_layouts.size(), ret), {ret}, attrs); } -inline std::pair, Array> BinaryBroadcastLayoutHelper( +std::pair, Array> BinaryBroadcastLayoutHelper( const Attrs& attrs, const Array& new_in_layouts, const Array& old_in_layouts, const Array& old_in_types); From b2ce09cf8692c8c98ee7861ac08903c3ed210016 Mon Sep 17 00:00:00 2001 From: Jinkun Lin Date: Thu, 3 Feb 2022 04:35:49 +0000 Subject: [PATCH 3/9] add unittest --- .../python/relay/test_pass_alter_op_layout.py | 164 ++++++++++++++++++ 1 file changed, 164 insertions(+) diff --git a/tests/python/relay/test_pass_alter_op_layout.py b/tests/python/relay/test_pass_alter_op_layout.py index ea7fe0bd7871..c877c0a4cea0 100644 --- a/tests/python/relay/test_pass_alter_op_layout.py +++ b/tests/python/relay/test_pass_alter_op_layout.py @@ -1471,5 +1471,169 @@ def test_conv2d_reduce_channels(): relay.build(mod, params=params, target="llvm") +def test_alter_layout_nonscalar_broadcast(): + """Test boradcast operators""" + + def before(): + x = relay.var("x", shape=(1, 16, 3, 3)) + weight = relay.var("weight", shape=(16, 16, 1, 1)) + y = relay.nn.conv2d(x, weight, channels=16, + kernel_size=(1, 1), padding=(0, 0), data_layout="NCHW") + z = relay.var("z", shape=(1, 3, 3)) + y = y + z + y = relay.Function(analysis.free_vars(y), y) + return y + + def expected(): + x = relay.var("x", shape=(1, 16, 3, 3)) + weight = relay.var("weight", shape=(16, 16, 1, 1)) + x = relay.layout_transform(x, src_layout="NCHW", dst_layout="NCHW4c") + weight = relay.layout_transform(weight, src_layout="OIHW", dst_layout="OIHW4i4o") + y = relay.nn.conv2d( + x, + weight, + channels=16, + kernel_size=(1, 1), + padding=(0, 0), + data_layout="NCHW4c", + kernel_layout="OIHW4i4o", + ) + z = relay.var("z", shape=(1, 3, 3)) + z = relay.expand_dims(z, 0) + z = relay.layout_transform(z, src_layout="NCHW", dst_layout="NCHW1c") + y = y + z + y = relay.layout_transform(y, src_layout="NCHW4c", dst_layout="NCHW") + y = relay.Function(analysis.free_vars(y), y) + return y + + def alter_conv2d(attrs, inputs, tinfos, out_type): + data, weight = inputs + new_attrs = dict(attrs) + new_attrs["data_layout"] = "NCHW4c" + new_attrs["kernel_layout"] = "OIHW4i4o" + return relay.nn.conv2d(data, weight, **new_attrs) + + with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): + a = run_opt_pass(before(), transform.AlterOpLayout()) + b = run_opt_pass(expected(), transform.InferType()) + assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) +'\nExpected = \n' + str(b) + + inp = np.random.uniform(size=(1, 16, 3, 3)).astype(np.float32) + weight = np.random.uniform(size=(16, 16, 1, 1)).astype(np.float32) + z = np.random.uniform(size=(1, 3, 3)).astype(np.float32) + mod = tvm.IRModule.from_expr(before()) + with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): + with tvm.transform.PassContext(opt_level=4): + res = relay.build_module.create_executor( + 'graph', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight, z) + with tvm.transform.PassContext(opt_level=0): + res1 = relay.build_module.create_executor( + 'debug', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight, z) + np.testing.assert_allclose(res.numpy(), res1.numpy()) + + +def test_broadcast_non_adaptable(): + """NCHW4c + [x, x, 4] and NCHW4c is being altered to NCHW""" + + def before(): + x = relay.var("x", shape=(1, 4, 3, 3, 4)) + weight = relay.var("weight", shape=(4, 4, 1, 1, 4, 4)) + y = relay.nn.conv2d( + x, + weight, + channels=16, + kernel_size=(1, 1), + padding=(0, 0), + data_layout="NCHW4c", + kernel_layout="OIHW4i4o", + ) + z = relay.var("z", shape=(3, 3, 4)) + y = y + z + y = relay.Function(analysis.free_vars(y), y) + return y + + def expected(): + x = relay.var("x", shape=(1, 4, 3, 3, 4)) + weight = relay.var("weight", shape=(4, 4, 1, 1, 4, 4)) + x = relay.layout_transform(x, src_layout="NCHW4c", dst_layout="NCHW") + weight = relay.layout_transform(weight, src_layout="OIHW4i4o", dst_layout="OIHW") + y = relay.nn.conv2d( + x, + weight, + channels=16, + kernel_size=(1, 1), + padding=(0, 0), + data_layout="NCHW", + kernel_layout="OIHW", + ) + z = relay.var("z", shape=(3, 3, 4)) + y = relay.layout_transform(y, src_layout="NCHW", dst_layout="NCHW4c") + y = y + z + y = relay.Function(analysis.free_vars(y), y) + return y + + def alter_conv2d(attrs, inputs, tinfos, out_type): + data, weight = inputs + new_attrs = dict(attrs) + new_attrs["data_layout"] = "NCHW" + new_attrs["kernel_layout"] = "OIHW" + return relay.nn.conv2d(data, weight, **new_attrs) + + with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): + a = run_opt_pass(before(), transform.AlterOpLayout()) + b = run_opt_pass(expected(), transform.InferType()) + assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) +'\nExpected = \n' + str(b) + + inp = np.random.uniform(size=(1, 4, 3, 3, 4)).astype(np.float32) + weight = np.random.uniform(size=(4, 4, 1, 1, 4, 4)).astype(np.float32) + z = np.random.uniform(size=(3, 3, 4)).astype(np.float32) + mod = tvm.IRModule.from_expr(before()) + with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): + with tvm.transform.PassContext(opt_level=4): + res = relay.build_module.create_executor( + 'graph', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight, z) + with tvm.transform.PassContext(opt_level=0): + res1 = relay.build_module.create_executor( + 'debug', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight, z) + np.testing.assert_allclose(res.numpy(), res1.numpy()) + + +def test_broadcast_respect_input_layouts(): + def before(): + x = relay.var("x", shape=(1, 16, 1, 1)) + w = relay.var("w", shape=(16, 16, 1, 1)) + x = relay.nn.conv2d( + x, + w, + kernel_size=(1, 1), + padding=(0, 0), + channels=16, + ) + y1 = relay.min(x, axis=[2]) + y2 = relay.min(x, axis=[3]) + z = y1 + y2 + z = relay.Function(analysis.free_vars(z), z) + return z + + def alter_conv2d(attrs, inputs, tinfos, out_type): + data, weight = inputs + new_attrs = dict(attrs) + new_attrs["data_layout"] = "NCHW4c" + new_attrs["kernel_layout"] = "OIHW4i4o" + return relay.nn.conv2d(data, weight, **new_attrs) + + inp = np.random.uniform(size=(1, 16, 1, 1)).astype(np.float32) + weight = np.random.uniform(size=(16, 16, 1, 1)).astype(np.float32) + mod = tvm.IRModule.from_expr(before()) + with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): + with tvm.transform.PassContext(opt_level=4): + res = relay.build_module.create_executor( + 'graph', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight) + with tvm.transform.PassContext(opt_level=0): + res1 = relay.build_module.create_executor( + 'debug', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight) + np.testing.assert_allclose(res.asnumpy(), res1.asnumpy()) + + if __name__ == "__main__": pytest.main([__file__]) From cc13a51cb3580ae2f60ac43c0d6d23beb625b758 Mon Sep 17 00:00:00 2001 From: Jinkun Lin Date: Thu, 3 Feb 2022 05:32:00 +0000 Subject: [PATCH 4/9] backward-compat: optimize for scalar layout --- src/relay/transforms/infer_layout_utils.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/relay/transforms/infer_layout_utils.cc b/src/relay/transforms/infer_layout_utils.cc index 21eae6034281..aad8d61c2bd6 100644 --- a/src/relay/transforms/infer_layout_utils.cc +++ b/src/relay/transforms/infer_layout_utils.cc @@ -78,7 +78,7 @@ Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& old_layo } new_layout += LayoutAxis::Get(axis).name(); } - return Layout(new_layout); + return new_layout != "" ? Layout(new_layout) : Layout("H").SubLayout(0, 0); // hack to create a scalar layout } bool Isomorphic(const Layout& lhs, const Layout& rhs) { DLOG(INFO) << "Isomorphic: " @@ -134,6 +134,7 @@ Layout TryTransformLike(const Layout& old, const Layout& ref_old, const Layout& } // `old` is compatible. Now learn the axis name mapping between `old` and `ref_old` + if (old.ndim() == 0) return old; // an optmization for scalar: no-op int mapping[26]; bool used[26]; memset(mapping, -1, sizeof mapping); From 0cbf0dec0242af51becc7814996adc0099b417e0 Mon Sep 17 00:00:00 2001 From: Jinkun Lin Date: Thu, 3 Feb 2022 05:34:12 +0000 Subject: [PATCH 5/9] fix lint --- src/relay/transforms/infer_layout_utils.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/relay/transforms/infer_layout_utils.cc b/src/relay/transforms/infer_layout_utils.cc index aad8d61c2bd6..86ad89716ab1 100644 --- a/src/relay/transforms/infer_layout_utils.cc +++ b/src/relay/transforms/infer_layout_utils.cc @@ -78,7 +78,8 @@ Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& old_layo } new_layout += LayoutAxis::Get(axis).name(); } - return new_layout != "" ? Layout(new_layout) : Layout("H").SubLayout(0, 0); // hack to create a scalar layout + return new_layout != "" ? Layout(new_layout) + : Layout("H").SubLayout(0, 0); // hack to create a scalar layout } bool Isomorphic(const Layout& lhs, const Layout& rhs) { DLOG(INFO) << "Isomorphic: " From c6bac19ea324c57ae5a2fa38009fefeff4a65848 Mon Sep 17 00:00:00 2001 From: Jinkun Lin Date: Thu, 3 Feb 2022 06:10:51 +0000 Subject: [PATCH 6/9] fix lint and warning --- .../python/relay/test_pass_alter_op_layout.py | 29 ++++++++++++------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/tests/python/relay/test_pass_alter_op_layout.py b/tests/python/relay/test_pass_alter_op_layout.py index e13208f63f19..4ce1286eb508 100644 --- a/tests/python/relay/test_pass_alter_op_layout.py +++ b/tests/python/relay/test_pass_alter_op_layout.py @@ -1477,8 +1477,9 @@ def test_alter_layout_nonscalar_broadcast(): def before(): x = relay.var("x", shape=(1, 16, 3, 3)) weight = relay.var("weight", shape=(16, 16, 1, 1)) - y = relay.nn.conv2d(x, weight, channels=16, - kernel_size=(1, 1), padding=(0, 0), data_layout="NCHW") + y = relay.nn.conv2d( + x, weight, channels=16, kernel_size=(1, 1), padding=(0, 0), data_layout="NCHW" + ) z = relay.var("z", shape=(1, 3, 3)) y = y + z y = relay.Function(analysis.free_vars(y), y) @@ -1516,7 +1517,7 @@ def alter_conv2d(attrs, inputs, tinfos, out_type): with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = run_opt_pass(before(), transform.AlterOpLayout()) b = run_opt_pass(expected(), transform.InferType()) - assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) +'\nExpected = \n' + str(b) + assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) + "\nExpected = \n" + str(b) inp = np.random.uniform(size=(1, 16, 3, 3)).astype(np.float32) weight = np.random.uniform(size=(16, 16, 1, 1)).astype(np.float32) @@ -1525,10 +1526,12 @@ def alter_conv2d(attrs, inputs, tinfos, out_type): with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): with tvm.transform.PassContext(opt_level=4): res = relay.build_module.create_executor( - 'graph', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight, z) + "graph", mod, target="llvm", device=tvm.cpu() + ).evaluate()(inp, weight, z) with tvm.transform.PassContext(opt_level=0): res1 = relay.build_module.create_executor( - 'debug', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight, z) + "debug", mod, target="llvm", device=tvm.cpu() + ).evaluate()(inp, weight, z) np.testing.assert_allclose(res.numpy(), res1.numpy()) @@ -1582,7 +1585,7 @@ def alter_conv2d(attrs, inputs, tinfos, out_type): with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): a = run_opt_pass(before(), transform.AlterOpLayout()) b = run_opt_pass(expected(), transform.InferType()) - assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) +'\nExpected = \n' + str(b) + assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) + "\nExpected = \n" + str(b) inp = np.random.uniform(size=(1, 4, 3, 3, 4)).astype(np.float32) weight = np.random.uniform(size=(4, 4, 1, 1, 4, 4)).astype(np.float32) @@ -1591,10 +1594,12 @@ def alter_conv2d(attrs, inputs, tinfos, out_type): with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): with tvm.transform.PassContext(opt_level=4): res = relay.build_module.create_executor( - 'graph', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight, z) + "graph", mod, target="llvm", device=tvm.cpu() + ).evaluate()(inp, weight, z) with tvm.transform.PassContext(opt_level=0): res1 = relay.build_module.create_executor( - 'debug', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight, z) + "debug", mod, target="llvm", device=tvm.cpu() + ).evaluate()(inp, weight, z) np.testing.assert_allclose(res.numpy(), res1.numpy()) @@ -1628,11 +1633,13 @@ def alter_conv2d(attrs, inputs, tinfos, out_type): with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d): with tvm.transform.PassContext(opt_level=4): res = relay.build_module.create_executor( - 'graph', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight) + "graph", mod, target="llvm", device=tvm.cpu() + ).evaluate()(inp, weight) with tvm.transform.PassContext(opt_level=0): res1 = relay.build_module.create_executor( - 'debug', mod, target='llvm', device=tvm.cpu()).evaluate()(inp, weight) - np.testing.assert_allclose(res.asnumpy(), res1.asnumpy()) + "debug", mod, target="llvm", device=tvm.cpu() + ).evaluate()(inp, weight) + np.testing.assert_allclose(res.numpy(), res1.numpy()) def test_axis_semantic_change(): From a5c816189f06c293fa55e8f8c5bcab7e726ccceb Mon Sep 17 00:00:00 2001 From: Jinkun Lin Date: Fri, 4 Feb 2022 20:23:12 +0000 Subject: [PATCH 7/9] Add newlines; Use std::vector --- src/relay/transforms/infer_layout_utils.cc | 12 ++++++++---- src/relay/transforms/infer_layout_utils.h | 2 ++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/src/relay/transforms/infer_layout_utils.cc b/src/relay/transforms/infer_layout_utils.cc index 86ad89716ab1..a2317c7eb418 100644 --- a/src/relay/transforms/infer_layout_utils.cc +++ b/src/relay/transforms/infer_layout_utils.cc @@ -81,6 +81,7 @@ Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& old_layo return new_layout != "" ? Layout(new_layout) : Layout("H").SubLayout(0, 0); // hack to create a scalar layout } + bool Isomorphic(const Layout& lhs, const Layout& rhs) { DLOG(INFO) << "Isomorphic: " << "lhs: " << lhs << " rhs: " << rhs << std::endl; @@ -110,6 +111,7 @@ bool Isomorphic(const Layout& lhs, const Layout& rhs) { } return true; } + Layout TryTransformLike(const Layout& old, const Layout& ref_old, const Layout& ref_new) { DLOG(INFO) << "transform_layout: old = " << old << ", ref_new = " << ref_new << ", ref_old = " << ref_old << std::endl; @@ -136,10 +138,9 @@ Layout TryTransformLike(const Layout& old, const Layout& ref_old, const Layout& // `old` is compatible. Now learn the axis name mapping between `old` and `ref_old` if (old.ndim() == 0) return old; // an optmization for scalar: no-op - int mapping[26]; - bool used[26]; - memset(mapping, -1, sizeof mapping); - memset(used, 0, sizeof used); + std::vector mapping(26, -1); + std::vector used(26, false); + auto find_unused = [&](char preference) -> char { if (!used[preference - 'A']) return preference; // preference unused for (int i = 0; i < 26; ++i) @@ -147,12 +148,14 @@ Layout TryTransformLike(const Layout& old, const Layout& ref_old, const Layout& LOG(FATAL) << "All letters are used"; return 0; }; + for (int j = old->axes.size() - 1, i = ref_old->axes.size() - 1; j >= 0; --i, --j) { char name_ref = LayoutAxis::Get(ref_old->axes[i]).ToPrimal().name()[0]; char name = LayoutAxis::Get(old->axes[j]).ToPrimal().name()[0]; mapping[name_ref - 'A'] = name - 'A'; used[name - 'A'] = true; } + for (int i = ref_old->axes.size() - 1; i >= 0; --i) { char name_ref = LayoutAxis::Get(ref_old->axes[i]).ToPrimal().name()[0]; int name = mapping[name_ref - 'A']; @@ -175,6 +178,7 @@ Layout TryTransformLike(const Layout& old, const Layout& ref_old, const Layout& new_layout += c; } } + DLOG(INFO) << "new_layout = " << new_layout << std::endl; return Layout(new_layout); } diff --git a/src/relay/transforms/infer_layout_utils.h b/src/relay/transforms/infer_layout_utils.h index 68c407d70a98..3b1cb29951e4 100644 --- a/src/relay/transforms/infer_layout_utils.h +++ b/src/relay/transforms/infer_layout_utils.h @@ -51,6 +51,7 @@ Layout AdjustSubordinateFactors(const Layout& src_layout, const Layout& old_layo const Array& old_shape); bool Isomorphic(const Layout& lhs, const Layout& rhs); + /*! * \brief Try transforming `old` in as the smae way as how`ref_old` is transformed to `ref_new`. * `old` and `ref_old` are expected to describe two broadcastable tensors. Layout with fewer rank @@ -66,6 +67,7 @@ bool Isomorphic(const Layout& lhs, const Layout& rhs); * \return The transformed layout. */ Layout TryTransformLike(const Layout& old, const Layout& ref_old, const Layout& ref_new); + /* * \brief An output structure to hold results from FInferCorrectLayout calls. * \tparam input_layouts Inferred input layouts. From d45c47679729fa0ddcf192b7ce97bfb74a1c8d0d Mon Sep 17 00:00:00 2001 From: Jinkun Lin Date: Fri, 4 Feb 2022 20:28:27 +0000 Subject: [PATCH 8/9] fix lint --- src/relay/transforms/infer_layout_utils.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/src/relay/transforms/infer_layout_utils.cc b/src/relay/transforms/infer_layout_utils.cc index a2317c7eb418..32838e09a441 100644 --- a/src/relay/transforms/infer_layout_utils.cc +++ b/src/relay/transforms/infer_layout_utils.cc @@ -27,6 +27,7 @@ #include #include #include +#include #include "pattern_utils.h" #include "tvm/runtime/logging.h" From f7cefe1a2b79c1049a475ea2ac8ddffa9292d373 Mon Sep 17 00:00:00 2001 From: Jinkun Lin Date: Sat, 5 Feb 2022 02:35:38 +0000 Subject: [PATCH 9/9] jostle ci