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 * \brief Helper utilities to implement hybrid_op.
22 * \file hybrid_op.h
23 */
24#ifndef TVM_TE_OPERATION_HYBRID_OP_H_
25#define TVM_TE_OPERATION_HYBRID_OP_H_
26
27#include <tvm/te/schedule.h>
28#include <tvm/tir/expr.h>
29
30#include <unordered_map>
31#include <unordered_set>
32#include <vector>
33
34#include "../../tir/transforms/arg_binder.h"
35#include "../../tir/transforms/ir_utils.h"
36#include "../schedule/message_passing.h"
37
38namespace tvm {
39namespace te {
40
41/*!
42 * \brief Find all the iteration variables in the given statement body.
43 * \param stmt The body to be inspected.
44 */
45std::vector<IterVar> GatherLoopVars(Stmt stmt);
46
47/*!
48 * \brief Replace the tensor reference (especially in Provide's) in stmt by the replace map.
49 * \param stmt The statement to be processed.
50 * \param replace The replacement rule.
51 */
52Stmt ReplaceProvideTensor(Stmt stmt, const std::unordered_map<Tensor, Tensor>& replace);
53
54/*!
55 * \brief Apply the schedule manipulation on the function body.
56 * \param stmt The statement to be processed.
57 * \param dom_map The extents of the iterative variables may be used.
58 * \param stage The schedule information to be applied.
59 */
60Stmt ApplySchedule(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
61 Stmt stmt);
62
63/*!
64 * \brief Apply loop splits and fuses in the schedule on the function body.
65 * \param stage The schedule information to be applied.
66 * \param dom_map The extents of the iterative variables may be used.
67 * \param stmt The statement to be processed.
68 */
69Stmt ApplyLoopShapes(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
70 Stmt stmt);
71
72/*!
73 * \brief Apply loop annotation in the schedule on the function body.
74 * \param stage The schedule information to be applied.
75 * \param rebased The map specifies the rebase, a.k.a rename, relationship of these variables.
76 * \param stmt The statement to be processed.
77 */
78Stmt ApplyLoopAnnotations(const Stage& stage, const std::unordered_map<IterVar, IterVar>& rebased,
79 Stmt stmt);
80
81/*!
82 * \brief Apply loop order in the schedule on the function body.
83 * \param stage The schedule information to be applied.
84 * \param dom_map The extents of the iterative variables may be used.
85 * \param rebased The map specifies the rebase, a.k.a rename, relationship of these variables.
86 * \param stmt The statement to be processed.
87 */
88Stmt ApplyLoopOrder(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
89 const std::unordered_map<IterVar, IterVar>& rebased, Stmt stmt);
90
91} // namespace te
92} // namespace tvm
93
94#endif // TVM_TE_OPERATION_HYBRID_OP_H_
95