1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 *
22 * \file src/relay/op/nn/upsampling.h
23 * \brief implementation of the InferCorrectLayout pass for upsampling
24 */
25
26#ifndef TVM_RELAY_OP_NN_UPSAMPLING_H_
27#define TVM_RELAY_OP_NN_UPSAMPLING_H_
28
29#include <tvm/relay/attrs/nn.h>
30#include <tvm/tir/data_layout.h>
31
32#include "../op_common.h"
33
34namespace tvm {
35namespace relay {
36
37template <typename T>
38InferCorrectLayoutOutput UpsamplingInferCorrectLayout(const Attrs& attrs,
39 const Array<Layout>& new_in_layouts,
40 const Array<Layout>& old_in_layouts,
41 const Array<tvm::relay::Type>& old_in_types) {
42 const auto* attrs_ptr = attrs.as<T>();
43 ICHECK(attrs_ptr);
44 ObjectPtr<T> params = make_object<T>(*attrs_ptr);
45
46 if (new_in_layouts.defined()) {
47 ICHECK_EQ(new_in_layouts.size(), 1);
48
49 Layout raw_layout(params->layout);
50 Layout input = new_in_layouts[0];
51 if (input.IndexOf(LayoutAxis::Get('W')) == raw_layout.IndexOf(LayoutAxis::Get('W')) &&
52 input.IndexOf(LayoutAxis::Get('H')) == raw_layout.IndexOf(LayoutAxis::Get('H')) &&
53 !input.Contains(LayoutAxis::Get('w')) && !input.Contains(LayoutAxis::Get('h')) &&
54 (input.IndexOf(LayoutAxis::Get('D')) == -1 ||
55 (input.IndexOf(LayoutAxis::Get('D')) == raw_layout.IndexOf(LayoutAxis::Get('D')) &&
56 !input.Contains(LayoutAxis::Get('d'))))) {
57 params->layout = input.name(); // modify self to follow the input layout
58 }
59 }
60
61 return InferCorrectLayoutOutput({params->layout}, {params->layout}, Attrs(params));
62}
63
64} // namespace relay
65} // namespace tvm
66
67#endif // TVM_RELAY_OP_NN_UPSAMPLING_H_
68