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
34namespace tvm {
35namespace te {
36
37/*!
38 * \brief To automatically inline the element-wise operations.
39 *
40 * \param sch The schedule to be inlined.
41 */
42void AutoInlineElemWise(Schedule sch);
43
44/*!
45 * \brief To automatically inline the broadcast operations.
46 *
47 * \param sch The schedule to be inlined.
48 */
49void 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 */
59TVM_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 */
67Map<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 */
76bool 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 */
89Stmt 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 */
107PrimFunc 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