domain_touched.cc 3.73 KB
Newer Older
1 2 3 4 5 6 7 8
/*
 * 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
9
 *
10
 *   http://www.apache.org/licenses/LICENSE-2.0
11
 *
12 13 14 15 16 17 18 19
 * 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 22 23
/*!
 * \file bound_deducer.cc
 * \brief Utility to deduce bound of expression
 */
24 25 26
#include <tvm/tir/expr.h>
#include <tvm/tir/ir_pass.h>
#include <tvm/tir/stmt_functor.h>
27
#include <tvm/te/tensor.h>
28
#include <tvm/runtime/registry.h>
29 30 31 32 33 34 35

#include <unordered_set>
#include <unordered_map>

namespace tvm {
namespace arith {

36
using namespace tir;
37 38

// Find Read region of the tensor in the stmt.
39
class FuncTouchedDomain final : public StmtExprVisitor {
40
 public:
41
  FuncTouchedDomain(const te::Tensor &tensor, bool consider_calls, bool consider_provides)
42 43 44
    : tensor_(tensor), consider_calls_(consider_calls), consider_provides_(consider_provides)  {}

  Domain Find(const Stmt& stmt) {
45
    operator()(stmt);
46 47 48 49 50 51 52 53
    Domain ret;
    Range none;
    for (size_t i = 0; i < bounds_.size(); ++i) {
      ret.push_back(arith::Union(bounds_[i]).cover_range(none));
    }
    return ret;
  }

54 55
  void VisitStmt_(const ForNode *op) final {
    const VarNode* var = op->loop_var.get();
56 57
    dom_map_[var] = IntSet::range(
        Range::make_by_min_extent(op->min, op->extent));
58
    StmtExprVisitor::VisitStmt_(op);
59 60 61
    dom_map_.erase(var);
  }

62
  void VisitStmt_(const LetStmtNode* op) final {
63 64
    dom_map_[op->var.get()] =
        arith::EvalSet(op->value, dom_map_);
65
    StmtExprVisitor::VisitStmt_(op);
66 67 68 69
    dom_map_.erase(op->var.get());
  }

  /* TODO: Thread extent unitest not generated.*/
70
  void VisitStmt_(const AttrStmtNode* op) final {
71 72 73
    if (op->attr_key == attr::thread_extent) {
      const IterVarNode* thread_axis = op->node.as<IterVarNode>();
      CHECK(thread_axis);
74
      const VarNode* var = thread_axis->var.get();
75
      dom_map_[var] = IntSet::range(Range(make_zero(op->value.dtype()), op->value));
76
      StmtExprVisitor::VisitStmt_(op);
77 78
      dom_map_.erase(var);
    } else {
79
      StmtExprVisitor::VisitStmt_(op);
80 81 82
    }
  }

83
  void VisitExpr_(const CallNode* op) final {
84 85 86 87
    if (consider_calls_ && tensor_->op.same_as(op->func)
        && tensor_->value_index == op->value_index) {
      Touch(op->args);
    }
88
    StmtExprVisitor::VisitExpr_(op);
89 90
  }

91
  void VisitStmt_(const ProvideNode* op) final {
92 93 94 95
    if (consider_provides_ && tensor_->op.same_as(op->func)
        && tensor_->value_index == op->value_index) {
      Touch(op->args);
    }
96
    StmtExprVisitor::VisitStmt_(op);
97 98 99
  }

 private:
100
  void Touch(const Array<PrimExpr>& args) {
101 102 103 104 105 106 107 108
    if (args.size() > bounds_.size()) {
      bounds_.resize(args.size());
    }
    for (size_t i = 0; i < args.size(); ++i) {
      bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
    }
  }

109
  const te::Tensor &tensor_;
110 111
  bool consider_calls_, consider_provides_;
  std::vector<std::vector<IntSet> > bounds_;
112
  std::unordered_map<const VarNode*, IntSet> dom_map_;
113 114
};

115
Domain DomainTouched(Stmt stmt,
116
                     const te::Tensor &tensor,
117 118
                     bool consider_calls,
                     bool consider_provides) {
119 120 121 122 123
  return FuncTouchedDomain(tensor, consider_calls, consider_provides).Find(stmt);
}

}  // namespace arith
}  // namespace tvm