scan_op.cc 10.3 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
/*!
 * \brief Scan Operator.
 * \file scan_op.cc
 */
24
#include <tvm/runtime/registry.h>
25
#include <tvm/te/operation.h>
26 27
#include <tvm/tir/expr.h>
#include <tvm/tir/ir_pass.h>
28
#include "op_util.h"
29 30 31
#include "../schedule/graph.h"

namespace tvm {
32
namespace te {
33
using namespace tir;
34

35 36
TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
.set_dispatch<ScanOpNode>([](const ObjectRef& node, ReprPrinter* p) {
37
    auto* op = static_cast<const ScanOpNode*>(node.get());
38 39 40 41
    p->stream << "scan(" << op->name << ", " << op << ")";
});
TVM_REGISTER_NODE_TYPE(ScanOpNode);

42
inline bool prove_equal(PrimExpr lhs, PrimExpr rhs) {
43
  return is_zero(tir::Simplify(lhs - rhs));
44 45 46 47 48 49 50 51 52 53 54 55 56
}

int ScanOpNode::num_outputs() const {
  return static_cast<int>(update.size());
}
Array<IterVar> ScanOpNode::root_iter_vars() const {
  Array<IterVar> ret{scan_axis};
  for (IterVar iv : spatial_axis_) {
    ret.push_back(iv);
  }
  return ret;
}

57
DataType ScanOpNode::output_dtype(size_t i) const {
58 59 60
  return update[i]->dtype;
}

61
Array<PrimExpr> ScanOpNode::output_shape(size_t i) const {
62 63 64 65 66
  CHECK_LT(i, state_placeholder.size());
  return state_placeholder[i]->shape;
}

Operation ScanOpNode::make(std::string name,
67
                           std::string tag,
68
                           Map<std::string, ObjectRef> attrs,
69 70 71
                           IterVar axis,
                           Array<Tensor> init,
                           Array<Tensor> update,
72 73
                           Array<Tensor> state_placeholder,
                           Array<Tensor> inputs) {
74
  if (!attrs.defined()) {
75
    attrs = Map<std::string, ObjectRef>();
76
  }
77
  auto n = make_object<ScanOpNode>();
78 79 80 81 82 83
  CHECK_EQ(init.size(), update.size());
  CHECK_EQ(init.size(), state_placeholder.size());

  for (size_t i = 0; i < init.size(); ++i) {
    CHECK_EQ(init[i]->dtype, state_placeholder[i]->dtype);
    CHECK_EQ(init[i]->dtype, update[i]->dtype);
84
    CHECK(prove_equal(init[i]->shape[0], axis->dom->min))
85 86 87
        << "init.shape[0] need to match scan_axis.dom.min";
    CHECK(prove_equal(
        state_placeholder[i]->shape[0], axis->dom->min + axis->dom->extent))
88
        << "state_placeholder.shape[0] need to match"
89 90 91 92 93 94 95 96 97 98 99 100 101 102
        << " scan_axis.dom.min + scan_axis.dom.extent";
    CHECK_EQ(state_placeholder[i].ndim(), init[i].ndim())
        << "The dimension of init need to match state_placeholder";
    CHECK_EQ(update[i].ndim(), state_placeholder[i].ndim())
        << "The update.ndim need to be state_placeholder.ndim - 1";
    for (size_t k = 0;  k < update[i].ndim(); ++k) {
      CHECK(prove_equal(
          update[i]->shape[k], state_placeholder[i]->shape[k]));
      if (k != 0) {
        // setup spatial axis
        std::ostringstream spatial_name;
        spatial_name << name << ".out" << i << ".i" << k;
        n->spatial_axis_.push_back(
            IterVarNode::make(
103
                Range::make_by_min_extent(0, update[i]->shape[k]),
104 105 106 107 108 109 110 111 112
                Var(spatial_name.str()), kOpaque));
      }
    }

    for (size_t k = 1;  k < init[i].ndim(); ++k) {
      CHECK(prove_equal(
          init[i]->shape[k], state_placeholder[i]->shape[k]));
    }
  }
113 114 115 116 117 118 119 120
  n->name = std::move(name);
  n->tag = std::move(tag);
  n->attrs = std::move(attrs);
  n->scan_axis = std::move(axis);
  n->init = std::move(init);
  n->update = std::move(update);
  n->state_placeholder = std::move(state_placeholder);
  n->inputs = std::move(inputs);
121 122 123
  return Operation(n);
}

124 125 126 127
TVM_REGISTER_GLOBAL("te.ScanOp")
.set_body_typed(ScanOpNode::make);


128 129 130
Array<Tensor> scan(Array<Tensor> init,
                   Array<Tensor> update,
                   Array<Tensor> state_placeholder,
131
                   Array<Tensor> inputs,
132
                   std::string name,
133
                   std::string tag,
134
                   Map<std::string, ObjectRef> attrs) {
135 136
  IterVar scan_axis =
      IterVarNode::make(
137
          Range::make_by_min_extent(
138 139 140
              init[0]->shape[0], update[0]->shape[0] - init[0]->shape[0]),
          Var(name + ".idx"), kOrdered);
  Operation op = ScanOpNode::make(
141 142
      name, tag, attrs, scan_axis,
      init, update, state_placeholder, inputs);
143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164
  Array<Tensor> res;
  for (int i = 0; i < op->num_outputs(); ++i) {
    res.push_back(op.output(i));
  }
  return res;
}

Array<Tensor> ScanOpNode::InputTensors() const {
  Array<Tensor> ret;
  for (Tensor t : init) {
    ret.push_back(t);
  }
  for (Tensor t : update) {
    ret.push_back(t);
  }
  return ret;
}

Operation ScanOpNode::ReplaceInputs(
    const Operation& self,
    const std::unordered_map<Tensor, Tensor>& rmap) const {
  CHECK_EQ(self.operator->(), this);
165
  auto n = make_object<ScanOpNode>(*this);
166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183
  for (size_t i = 0; i < n->init.size(); ++i) {
    if (rmap.count(n->init[i])) {
      n->init.Set(i, rmap.at(n->init[i]));
    }
    if (rmap.count(n->update[i])) {
      n->update.Set(i, rmap.at(n->update[i]));
    }
  }
  if (!n->init.same_as(init) ||
      !n->update.same_as(update)) {
    return Operation(n);
  } else {
    return self;
  }
}

void ScanOpNode::PropBoundToInputs(
    const Operation& self,
184
    arith::Analyzer* analyzer,
185
    const std::unordered_map<const VarNode*, IntSet>& dom_map,
186 187 188 189 190 191 192 193 194 195 196 197 198 199
    std::unordered_map<Tensor, TensorDom>* out_dom_map) const {
  CHECK_EQ(self.operator->(), this);
  for (size_t i = 0, sp_idx = 0; i < this->init.size(); ++i) {
    TensorDom* init_dom = nullptr;
    TensorDom* update_dom = nullptr;
    if (out_dom_map->count(this->init[i])) {
      init_dom = &out_dom_map->at(this->init[i]);
    }
    if (out_dom_map->count(this->update[i])) {
      update_dom = &out_dom_map->at(this->update[i]);
    }
    // first dimension, always needed.
    if (init_dom) {
      init_dom->data[0].push_back(IntSet::range(
200
          Range::make_by_min_extent(0, this->init[i]->shape[0])));
201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236
    }
    if (update_dom) {
      update_dom->data[0].push_back(dom_map.at(this->scan_axis->var.get()));
    }
    // The update dimensions
    for (size_t k = 1; k < this->update[i]->shape.size(); ++k, ++sp_idx) {
      IterVar sp_ax = this->spatial_axis_[sp_idx];
      if (init_dom) {
        init_dom->data[k].push_back(dom_map.at(sp_ax->var.get()));
      }
      if (update_dom) {
        update_dom->data[k].push_back(dom_map.at(sp_ax->var.get()));
      }
    }
  }
}

void ScanOpNode::GatherBound(
    const Operation& self,
    const std::unordered_map<Tensor, TensorDom>& tensor_dom,
    std::unordered_map<IterVar, Range>* out_dom_map) const {
  CHECK_EQ(self.operator->(), this);
  CHECK(!out_dom_map->count(this->scan_axis));
  std::vector<Tensor> output(this->num_outputs());
  for (size_t i = 0; i < output.size(); ++i) {
    output[i] = self.output(i);
  }
  // Update for time axis.
  std::vector<IntSet> time_dom;
  for (size_t i = 0; i < output.size(); ++i) {
    const TensorDom& d = tensor_dom.at(output[i]);
    time_dom.insert(time_dom.end(), d.data[0].begin(), d.data[0].end());
  }
  CHECK(!out_dom_map->count(this->scan_axis));
  Range sdom = this->scan_axis->dom;
  Range r = arith::Union(time_dom).cover_range(sdom);
237
  (*out_dom_map)[this->scan_axis] = Range::make_by_min_extent(
238
      sdom->min, tir::Simplify(r->extent + r->min - sdom->min));
239
  Map<IterVar, PrimExpr> fix_pt = ScanFixPointAnalysis(self);
240 241 242 243 244 245 246 247
  // Update for spatial axis.
  size_t sp_idx = 0;
  for (size_t i = 0; i < output.size(); ++i) {
    const TensorDom& d = tensor_dom.at(output[i]);
    for (size_t k = 1; k < this->update[i]->shape.size(); ++k, ++sp_idx) {
      IterVar sp_ax = this->spatial_axis_[sp_idx];
      CHECK(!out_dom_map->count(sp_ax));
      CHECK(fix_pt.count(sp_ax));
248
      if (fix_pt[sp_ax].as<tir::IntImmNode>()->value) {
249
        // fix point, we can slice it.
250
        (*out_dom_map)[sp_ax] = arith::Union(d.data[k]).cover_range(sp_ax->dom);
251 252 253 254 255 256 257 258 259
      } else {
        // not a fix point, need to include everything.
        (*out_dom_map)[sp_ax] = sp_ax->dom;
      }
    }
  }
}

Stmt ScanOpNode::BuildRealize(
260
    const Stage& stage,
261 262
    const std::unordered_map<IterVar, Range>& dom_map,
    const Stmt& body) const {
263
  CHECK_EQ(stage->op.get(), this);
264
  Range sdom = dom_map.at(this->scan_axis);
265
  Range tdom = Range::make_by_min_extent(
266
      0, tir::Simplify(sdom->extent + sdom->min));
267 268 269
  Stmt ret = body;
  size_t sp_idx = 0;
  for (size_t i = 0; i < update.size(); ++i) {
270
    Tensor t = stage->op.output(i);
271
    CHECK_EQ(static_cast<size_t>(t->value_index), i);
272
    Region bounds;
273 274 275 276 277
    bounds.push_back(tdom);
    for (size_t k = 1; k < this->update[i]->shape.size(); ++k, ++sp_idx) {
      IterVar sp_ax = this->spatial_axis_[sp_idx];
      bounds.push_back(dom_map.at(sp_ax));
    }
278
    ret = tir::RealizeNode::make(t->op, t->value_index, t->dtype,
279 280 281 282 283 284 285
                            bounds, const_true(), ret);
  }
  return ret;
}

Stmt ScanOpNode::BuildProvide(
    const Stage& stage,
286
    const std::unordered_map<IterVar, Range>& dom_map,
287
    bool debug_keep_trivial_loop) const {
288
  CHECK_EQ(stage->op.operator->(), this);
289
  Stmt provide = AttrStmtNode::make(
290
      stage->op, tir::attr::scan_update_scope, this->scan_axis->var,
291 292
      EvaluateNode::make(0));
  Stmt init = AttrStmtNode::make(
293
      stage->op, tir::attr::scan_init_scope, 0,
294
      EvaluateNode::make(0));
295 296 297 298 299 300 301
  size_t begin_scan = 0;
  for (size_t  i = 0; i < stage->leaf_iter_vars.size(); ++i) {
    if (stage->leaf_iter_vars[i]->iter_type == kThreadIndex) {
      CHECK_EQ(begin_scan, i);
      begin_scan = i + 1;
    }
  }
302
  std::unordered_map<IterVar, PrimExpr> vmap;
303
  std::unordered_set<IterVar> empty;
304
  auto nest = MakeLoopNest(
305
      stage, dom_map, 0, false, empty, &vmap, debug_keep_trivial_loop);
306 307
  nest[begin_scan].push_back(init);
  nest.push_back(
308 309
      MakeIfNest(
          MakeBoundCheck(stage, dom_map, vmap, false, empty)));
310 311
  return MergeNest(nest, provide);
}
312
}  // namespace te
313
}  // namespace tvm