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 * \file x86/default.h
22 * \brief default x86 schedule
23 */
24#ifndef TVM_TOPI_X86_DEFAULT_H_
25#define TVM_TOPI_X86_DEFAULT_H_
26
27#include <tvm/target/generic_func.h>
28#include <tvm/te/operation.h>
29#include <tvm/te/schedule_pass.h>
30#include <tvm/topi/detail/fuse.h>
31#include <tvm/topi/tags.h>
32
33namespace tvm {
34namespace topi {
35
36using namespace tvm::te;
37
38namespace x86 {
39/*!
40 * \brief Helper to create a default x86 schedule for the given ops.
41 *
42 * \param target The target to generate a schedule for.
43 * \param outs The output tensors.
44 * \param auto_inline Whether to apply the auto inline step.
45 *
46 * \return A schedule for the given ops.
47 */
48inline Schedule MakeDefaultSchedule(const Target& target, const Array<Tensor>& outs,
49 bool auto_inline) {
50 Array<Operation> out_ops;
51 for (auto t : outs) {
52 out_ops.push_back(t->op);
53 }
54 auto s = create_schedule(out_ops);
55 auto x = outs[0];
56 auto axis = s[x]->op.as<ComputeOpNode>()->axis;
57
58 if (auto_inline) {
59 tvm::te::AutoInlineInjective(s);
60 if (axis.size() > 0) {
61 detail::Fuse(s[x], axis);
62 }
63 return s;
64 }
65
66 if (axis.size() == 4) {
67 auto n = axis[0];
68 auto c = axis[1];
69 auto fused = detail::Fuse(s[x], {n, c}); // for nhwc layout, fuse n and h
70 s[x].parallel(fused);
71 } else {
72 s[x].parallel(axis[0]);
73 }
74
75 return s;
76}
77
78/*!
79 * \brief Create a default x86 schedule for the given ops.
80 *
81 * \param target The target to generate a schedule for.
82 * \param outs The output tensors.
83 *
84 * \return A schedule for the given ops.
85 */
86inline Schedule default_schedule(const Target& target, const Array<Tensor>& outs) {
87 return MakeDefaultSchedule(target, outs, false);
88}
89
90/*!
91 * \brief Create a default x86 schedule for the given ops, with auto inline
92 *
93 * \param target The target to generate a schedule for.
94 * \param outs The output tensors.
95 *
96 * \return A schedule for the given ops.
97 */
98inline Schedule default_schedule_auto_inline(const Target& target, const Array<Tensor>& outs) {
99 return MakeDefaultSchedule(target, outs, true);
100}
101
102} // namespace x86
103} // namespace topi
104} // namespace tvm
105#endif // TVM_TOPI_X86_DEFAULT_H_
106