cross_thread_reduction.cc 3.95 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 24
/*!
 * \brief Logics related to cross thread reduction, used by ComputeOpNode.
 * \file cross_thread_reduction.cc
 */
#include <tvm/ir_pass.h>
25 26
#include "compute_op.h"
#include "op_util.h"
27 28 29 30 31 32 33

namespace tvm {
using namespace ir;

Stmt MakeCrossThreadReduction(
    const ComputeOpNode* self,
    const Stage& stage,
34
    const std::unordered_map<IterVar, Range>& dom_map,
35
    bool debug_keep_trivial_loop) {
36 37 38 39 40 41
  Array<Expr>  args;
  for (IterVar iv : self->axis) {
    args.push_back(iv->var);
  }
  std::unordered_map<IterVar, Expr> value_map;
  auto nest = op::MakeLoopNest(
42
      stage, dom_map, 0, false, std::unordered_set<IterVar>(), &value_map, debug_keep_trivial_loop);
43 44 45
  auto conds = schedule::MakeBoundCheck(
      stage, dom_map, value_map, false,
      std::unordered_set<IterVar>());
46 47 48 49 50 51 52 53 54 55 56 57 58 59

  size_t size = self->body.size();
  CHECK_GT(size, 0);
  std::vector<const Reduce*> reduces(size);
  for (size_t i = 0; i < size; ++i) {
    const Reduce* reduce = self->body[i].as<Reduce>();
    CHECK(reduce);
    reduces[i] = reduce;
  }
  Expr cond = reduces[0]->condition;
  for (Expr v : conds) {
    cond = cond && v;
  }
  Array<Expr> freduce_args;
60
  freduce_args.push_back(make_const(DataType::UInt(32), static_cast<uint32_t>(size)));
61 62 63 64 65 66
  for (size_t i = 0; i < size; ++i) {
    freduce_args.push_back(reduces[0]->source[i]);
  }
  freduce_args.push_back(cond);
  std::vector<Var> res_handles(size);
  for (size_t idx = 0; idx < size; ++idx) {
67
    res_handles[idx] = Var("reduce_temp" + std::to_string(idx), DataType::Handle());
68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87
    freduce_args.push_back(res_handles[idx]);
  }

  for (IterVar iv : stage->leaf_iter_vars) {
    if (iv->iter_type == kCommReduce) {
      auto it = stage->iter_var_attrs.find(iv);
      if (it != stage->iter_var_attrs.end() &&
          (*it).second->bind_thread.defined()) {
        IterVar tv = (*it).second->bind_thread;
        freduce_args.push_back(tv->var);
      }
    }
  }
  // Checks for the thread.
  std::vector<Expr> thread_head_check;
  if (stage->store_predicate.defined()) {
    thread_head_check.emplace_back(stage->store_predicate);
  }

  Stmt reduce_body = Evaluate::make(Call::make(
88
      DataType::Handle(),
89 90 91 92 93
      ir::intrinsic::tvm_thread_allreduce,
      freduce_args, Call::Intrinsic));
  reduce_body = AttrStmt::make(
      reduces[0]->combiner,
      attr::reduce_scope,
94
      make_zero(DataType::Handle()),
95 96 97
      reduce_body);
  std::vector<Stmt> assigns(size);
  for (size_t idx = 0; idx < size; ++idx) {
98
    DataType t = reduces[idx]->dtype;
99 100 101 102
    assigns[idx] = Provide::make(
      stage->op, idx,
      Load::make(t, res_handles[idx], 0, const_true(t.lanes())), args);
  }
103
  Stmt assign_body = SeqStmt::Flatten(assigns);
104 105
  assign_body = MergeNest(op::MakeIfNest(thread_head_check), assign_body);
  assign_body = MergeNest(op::MakeIfNest(conds), assign_body);
106
  Stmt body = SeqStmt::Flatten(reduce_body, assign_body);
107 108
  for (size_t idx = size; idx != 0; --idx) {
    body = Allocate::make(
109
      res_handles[idx - 1], reduces[idx - 1]->dtype, {1}, const_true(), body);
110 111 112 113 114 115 116
    body = AttrStmt::make(
      res_handles[idx - 1], attr::storage_scope, StringImm::make("local"), body);
  }
  body = op::Substitute(body, value_map);
  return MergeNest(nest, body);
}
}  // namespace tvm