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 | |
38 | namespace tvm { |
39 | namespace te { |
40 | |
41 | /*! |
42 | * \brief Find all the iteration variables in the given statement body. |
43 | * \param stmt The body to be inspected. |
44 | */ |
45 | std::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 | */ |
52 | Stmt 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 | */ |
60 | Stmt 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 | */ |
69 | Stmt 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 | */ |
78 | Stmt 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 | */ |
88 | Stmt 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 | |