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 tvm/te/schedule_pass.h |
22 | * \brief Collection of Schedule pass functions. |
23 | * |
24 | * These passes works on the schedule hyper-graph |
25 | * and infers information such as bounds, check conditions |
26 | * read/write dependencies between the IterVar |
27 | */ |
28 | #ifndef TVM_TE_SCHEDULE_PASS_H_ |
29 | #define TVM_TE_SCHEDULE_PASS_H_ |
30 | |
31 | #include <tvm/te/schedule.h> |
32 | #include <tvm/tir/function.h> |
33 | |
34 | namespace tvm { |
35 | namespace te { |
36 | |
37 | /*! |
38 | * \brief To automatically inline the element-wise operations. |
39 | * |
40 | * \param sch The schedule to be inlined. |
41 | */ |
42 | void AutoInlineElemWise(Schedule sch); |
43 | |
44 | /*! |
45 | * \brief To automatically inline the broadcast operations. |
46 | * |
47 | * \param sch The schedule to be inlined. |
48 | */ |
49 | void AutoInlineBroarcast(Schedule sch); |
50 | |
51 | /*! |
52 | * \brief To automatically inline operations with injective writes |
53 | * (i.e. writes without reduction or sequential loops). Note |
54 | * that in this case, guarantees about contiguity, transpose, stride, |
55 | * alignemnt and memory footprint in general do not hold. |
56 | * |
57 | * \param sch The schedule to be inlined. |
58 | */ |
59 | TVM_DLL void AutoInlineInjective(Schedule sch); |
60 | |
61 | /*! |
62 | * \brief Infer the bound of all iteration variables relates to the schedule. |
63 | * |
64 | * \param sch The root schedule to infer all the bounds. |
65 | * \return the result bound of the iteration Variable |
66 | */ |
67 | Map<IterVar, Range> InferBound(const Schedule& sch); |
68 | |
69 | /*! |
70 | * \brief Verify if there is any argument bound to compact buffer. |
71 | * |
72 | * \param stmt The stmt to be verified. |
73 | * \return true if there is any buffer_bind_scope attribute found, |
74 | * otherwise, false. |
75 | */ |
76 | bool VerifyCompactBuffer(const Stmt& stmt); |
77 | |
78 | /*! |
79 | * \brief Schedule s' dependent operations. |
80 | * |
81 | * \param s The schedule to be realized |
82 | * \param dom_map The domain of each iter vars. |
83 | * \param debug_keep_trivial_loop Whether keep trivial loops with extent of 1 during lowering. |
84 | * This is a debug feature for dataflow/axis analysis. |
85 | * Note: If this is true, The lowered IR may be incorrect, |
86 | * because we will also delete the init part of reduction |
87 | * \return the result Stmt |
88 | */ |
89 | Stmt ScheduleOps(Schedule s, Map<IterVar, Range> dom_map, bool debug_keep_trivial_loop); |
90 | |
91 | /*! |
92 | * \brief Postprocessing the Stmt generated by ScheduleOps to create |
93 | * a PrimFunc that can then be used for further TIR optimizations. |
94 | * |
95 | * Perform this translation before running any TIR optimizations. |
96 | * |
97 | * List of actions taken by the function: |
98 | * - Remove occurrences of te::Tensor, te::Operation in the IR |
99 | * and replace them by corresponding IR nodes via tir::Buffer. |
100 | * - Add annotation of extern buffers using the buffer_map field |
101 | * in the PrimFunc type. |
102 | * |
103 | * \param arg_list Array of Tensor/Var/Buffer arguments to the function. |
104 | * \param body The body of the function. |
105 | * \param bindings potential Tensor to Buffer bindings for the Tensors in the body. |
106 | */ |
107 | PrimFunc SchedulePostProcToPrimFunc(Array<ObjectRef> arg_list, Stmt body, |
108 | Optional<Map<Tensor, Buffer>> bindings); |
109 | |
110 | } // namespace te |
111 | } // namespace tvm |
112 | #endif // TVM_TE_SCHEDULE_PASS_H_ |
113 | |