op_util.h 3.51 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

20 21
/*!
 *  Copyright (c) 2017 by Contributors
22 23
 * \file op_util.h
 * \brief Common utility used in operator construction.
24
 */
25 26
#ifndef TVM_OP_OP_UTIL_H_
#define TVM_OP_OP_UTIL_H_
27 28 29 30 31 32 33

#include <tvm/expr.h>
#include <tvm/schedule.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "../pass/ir_util.h"
34
#include "../pass/arg_binder.h"
35
#include "../schedule/message_passing.h"
36 37 38 39 40 41 42 43 44 45 46 47 48 49 50

namespace tvm {
namespace op {

using ir::MergeNest;

/*!
 * \brief Build loop nest for stage.
 *
 * \param stage The stage to create a loop nest.
 * \param dom_map The range of each iter var.
 * \param begin_iter_pos The beginning position of leaf_iter_vars to generate loop.
 * \param new_loop_var Whether create new loop variable.
 * \param skip_iter Whether skip certain iteration.
 * \param p_value_map The result value of each IterVar.
51
 * \param debug_keep_trivial_loop Whether keep trivial loops with extent of 1
52 53 54 55 56 57 58
 */
std::vector<std::vector<Stmt> >
MakeLoopNest(const Stage& stage,
             const std::unordered_map<IterVar, Range>& dom_map,
             size_t begin_iter_pos,
             bool new_loop_var,
             const std::unordered_set<IterVar>& skip_iter,
59
             std::unordered_map<IterVar, Expr>* p_value_map,
60
             bool debug_keep_trivial_loop);
61

62 63 64 65 66 67 68
/*!
 * \brief Create a nest of if checking the predicates.
 *
 * \param predicates The predicates to be checked.
 * \return List of If nest that checks the predicates.
 */
std::vector<Stmt> MakeIfNest(const std::vector<Expr>& predicates);
69 70

/*!
71
 * \brief Replace the tensor reference (especially in Call's) in stmt by the replace map.
72 73 74 75 76 77
 * \param stmt The statement to be processed.
 * \param replace The replacement rule.
 */
Stmt ReplaceTensor(Stmt stmt,
                   const std::unordered_map<Tensor, Tensor>& replace);
/*!
78
 * \brief Replace the tensor reference (especially in Call's) in stmt by the replace map.
79 80 81 82 83 84
 * \param expr The expression to be processed.
 * \param replace The replacement rule.
 */
Expr ReplaceTensor(Expr expr,
                   const std::unordered_map<Tensor, Tensor>& replace);

85 86 87 88 89 90 91 92 93
/*!
 * \brief Substitute the variables of stmt by value map.
 * \param stmt the statment
 * \param value_map The value map.
 * \return Substituted result.
 */
Stmt Substitute(Stmt stmt,
                const std::unordered_map<IterVar, Expr>& value_map);

94 95 96 97 98 99 100 101 102 103 104 105
/*!
 * \brief Converts Halide ForType to its corresponding IterVarType
 * \param for_type The ForType to be converted
 */
IterVarType ForTypeToIterVarType(ir::ForType for_type);

/*!
 * \brief Converts IterVarType to its corresponding Halide ForType
 * \param iter_type The IterVarType to be converted
 */
ir::ForType IterVarTypeToForType(IterVarType iter_type);

106 107
}  // namespace op
}  // namespace tvm
108
#endif  // TVM_OP_OP_UTIL_H_