device_annotation.cc 19.5 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 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37
/*!
 * Copyright (c) 2018 by Contributors
 *
 * \file deivce_annotation.cc
 * \brief Passes to rewrite annotated program and retrieve the device allocation
 * of expression.
 *
 * The following passes are performed:
 *  1. Validate the unnecessary and redundant annotation.
 *  2. Rewrite the annotated program and insert data copy operators.
 *  3. Collect the device allocation of each expression.
 */

#include <tvm/relay/attrs/device_copy.h>
#include <tvm/relay/attrs/annotation.h>
#include <tvm/relay/expr.h>
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/pass.h>
38
#include <tvm/relay/transform.h>
39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70

#include <memory>
#include <unordered_map>
#include <unordered_set>

namespace tvm {
namespace relay {

namespace {

bool IsOnDeviceNode(const ExprNode* node) {
  const auto* call_node = dynamic_cast<const CallNode*>(node);
  return call_node != nullptr && call_node->attrs.as<OnDeviceAttrs>();
}

bool IsDeviceCopyNode(const ExprNode* node) {
  const auto* call_node = dynamic_cast<const CallNode*>(node);
  return call_node != nullptr && call_node->attrs.as<DeviceCopyAttrs>();
}

}  // namespace

class ValidateAnnotation : private ExprVisitor {
 public:
  static std::unordered_map<const ExprNode*, int> Validate(const Expr& expr) {
    ValidateAnnotation valid;
    valid(expr);
    return valid.annotation_map_;
  }

 private:
  void VisitExpr_(const CallNode* call_node) final {
71
    ExprVisitor::VisitExpr_(call_node);
72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89
    if (IsOnDeviceNode(call_node)) {
      int device_type = GetDeviceId(call_node);
      if (annotation_map_.count(call_node)) {
        CHECK_EQ(annotation_map_.at(call_node), device_type)
            << "An expression node can only be annotated to one device.";
      } else {
        annotation_map_.insert({call_node, GetDeviceId(call_node)});
      }

      CHECK_EQ(call_node->args.size(), 1U);
      const auto* node = call_node->args[0].operator->();
      if (annotation_map_.count(node)) {
        CHECK_EQ(annotation_map_.at(node), device_type)
            << "An expression node can only be annotated to one device.";
      } else {
        annotation_map_.insert({node, GetDeviceId(call_node)});
      }
    }
90 91 92 93 94 95 96 97
  }

  void VisitExpr_(const TupleGetItemNode* get_elem) final {
    ExprVisitor::VisitExpr_(get_elem);
    const auto* tn = get_elem->tuple.operator->();
    if (annotation_map_.count(tn)) {
      annotation_map_.insert({get_elem, annotation_map_.at(tn)});
    }
98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187
  }

  /*
   * \brief Get the device type of the annotation node.
   * \param call_node The on_device annotation call node.
   * \return The device type.
   */
  int GetDeviceId(const CallNode* call_node) {
    CHECK(IsOnDeviceNode(call_node))
        << "The input call node must be on_device node.";
    const OnDeviceAttrs* on_device_attr = call_node->attrs.as<OnDeviceAttrs>();
    return on_device_attr->device_type;
  }

  std::unordered_map<const ExprNode*, int> annotation_map_;
};

// Replace the use of an expression with the output of a `copy_device` operator
// if the `on_device` operator takes the annotated expr as an input.
//
// This actually replaces annotation ops with device copy ops and connects any
// two dependent expressions with a `device_copy` op when needed. Note that the
// device type of a `device_copy` op is identical to that of the destination op
// since it is where the data should be copied to.
class RewriteAnnotation : public ExprMutator {
 public:
  Expr Rewrite(const Expr& expr, int fallback_device) {
    fallback_device_ = fallback_device;
    annotation_map_ = ValidateAnnotation::Validate(expr);
    return this->VisitExpr(expr);
  }

  Expr VisitExpr_(const LetNode* op) final {
    Expr value = GetDeviceCopyExpr(op->value, op);
    Expr body = GetDeviceCopyExpr(op->body, op);

    if (value.same_as(op->value) && body.same_as(op->body)) {
      return ExprMutator::VisitExpr_(op);
    } else {
      Expr new_let = LetNode::make(op->var, value, body);
      UpdateAnnotationMap(op, new_let.operator->());
      return this->VisitExpr(new_let);
    }
  }

  Expr VisitExp_(const TupleNode* op) {
    Array<Expr> fields;
    bool annotated = false;
    for (const auto& field : fields) {
      annotated |= NeedDeviceCopy(field.operator->(), op);
      fields.push_back(GetDeviceCopyExpr(field, op));
    }

    if (annotated) {
      Expr new_tuple = TupleNode::make(fields);
      UpdateAnnotationMap(op, new_tuple.operator->());
      return this->VisitExpr(new_tuple);
    } else {
      return ExprMutator::VisitExpr_(op);
    }
  }

  Expr VisitExpr_(const TupleGetItemNode* op) final {
    Expr tuple = op->tuple;
    if (NeedDeviceCopy(tuple.operator->(), op)) {
      Expr new_expr =
          TupleGetItemNode::make(GetDeviceCopyExpr(tuple, op), op->index);
      UpdateAnnotationMap(op, new_expr.operator->());
      return this->VisitExpr(new_expr);
    } else {
      return ExprMutator::VisitExpr_(op);
    }
  }

  Expr VisitExpr_(const IfNode* if_node) final {
    Expr cond = GetDeviceCopyExpr(if_node->cond, if_node);
    Expr true_br = GetDeviceCopyExpr(if_node->true_branch, if_node);
    Expr false_br = GetDeviceCopyExpr(if_node->false_branch, if_node);

    if (if_node->cond.same_as(cond) && if_node->true_branch.same_as(true_br) &&
        if_node->false_branch.same_as(false_br)) {
      return ExprMutator::VisitExpr_(if_node);
    } else {
      Expr new_if = IfNode::make(cond, true_br, false_br);
      UpdateAnnotationMap(if_node, new_if.operator->());
      return this->VisitExpr(new_if);
    }
  }

  Expr VisitExpr_(const CallNode* call_node) final {
188 189 190 191 192
    if (IsOnDeviceNode(call_node)) {
      return this->VisitExpr(call_node->args[0]);
    }

    if (IsDeviceCopyNode(call_node)) {
193 194 195 196 197 198 199 200 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 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263
      return ExprMutator::VisitExpr_(call_node);
    }

    Array<Expr> new_args;
    bool annotated = false;
    for (const auto& arg : call_node->args) {
      annotated |= NeedDeviceCopy(arg.operator->(), call_node);
      new_args.push_back(GetDeviceCopyExpr(arg, call_node));
    }

    if (annotated) {
      Call new_call = CallNode::make(call_node->op, new_args, call_node->attrs,
                                     call_node->type_args);

      UpdateAnnotationMap(call_node, new_call.operator->());
      return this->VisitExpr(new_call);
    } else {
      return ExprMutator::VisitExpr_(call_node);
    }
  }

 private:
  void UpdateAnnotationMap(const ExprNode* old_node, const ExprNode* new_node) {
    const auto it = annotation_map_.find(old_node);
    if (it == annotation_map_.end()) {
      annotation_map_.insert({new_node, fallback_device_});
    } else {
      annotation_map_.insert({new_node, it->second});
    }
    this->memo_[GetRef<Expr>(old_node)] = GetRef<Expr>(new_node);
  }

  Expr GetDeviceCopyExpr(const Expr& src, const ExprNode* dst) {
    const auto* src_node = src.operator->();
    if (!NeedDeviceCopy(src_node, dst)) return src;

    const auto sit = annotation_map_.find(src_node);
    if (sit == annotation_map_.end()) {
      const auto dit = annotation_map_.find(dst);
      CHECK(dit != annotation_map_.end())
          << "Device copy op is not required when both src and dst ops are not "
             "annotated.";
      return CreateDeviceCopy(src, fallback_device_, dit->second);
    } else {
      const auto dit = annotation_map_.find(dst);
      int dst_dev_type =
          dit == annotation_map_.end() ? fallback_device_ : dit->second;
      return CreateDeviceCopy(src, sit->second, dst_dev_type);
    }
  }

  // Check if a device copy op is need between two ops.
  bool NeedDeviceCopy(const ExprNode* src, const ExprNode* dst) {
    if (annotation_map_.count(src)) {
      int src_dev_type = annotation_map_.at(src);
      if (annotation_map_.count(dst)) {
        return src_dev_type != annotation_map_.at(dst);
      } else {
        return src_dev_type != fallback_device_;
      }
    } else {
      if (annotation_map_.count(dst)) {
        // Though data copy op could be inserted whenever the `src` and `dst`
        // ops are annotated to different devices, it leads to high overhead.
        //
        // Here we need across device data transferring only when `src` is a
        // CallNode or FunctionNode and the `dst` is annotated with any device
        // id other than fallback_device_.
        if (src->is_type<CallNode>() || src->is_type<FunctionNode>()) {
          return annotation_map_.at(dst) != fallback_device_;
        } else {
264 265 266
          // There shouldn't be any copy nodes between var/constant and another
          // expression.
          return !(src->is_type<VarNode>() || src->is_type<ConstantNode>());
267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351
        }
      } else {
        return false;
      }
    }
  }

  /*
   * \brief Create an operator to copy data from the source device to the
   * destination device.
   * \param src The source expression that produces data to be copied.
   * \param src_dev_type The device type where the data is copied from.
   * \param dst_dev_type The device type where the data is copied to.
   * \return The created call node.
   */
  Call CreateDeviceCopy(const Expr& src, int src_dev_type, int dst_dev_type) {
    auto attrs = make_node<DeviceCopyAttrs>();
    attrs->src_dev_type = src_dev_type;
    attrs->dst_dev_type = dst_dev_type;
    static const Op& op = Op::Get("device_copy");
    Call device_copy = CallNode::make(op, {src}, Attrs(attrs), {});
    annotation_map_.insert({device_copy.operator->(), dst_dev_type});
    return device_copy;
  }

  std::unordered_map<const ExprNode*, int> annotation_map_;
  int fallback_device_;
};

// Get all annotation expressions.
class AnnotatationVisitor : private ExprVisitor {
 public:
  static Map<Expr, Integer> GetAnnotations(const Expr& expr) {
    AnnotatationVisitor visitor;
    visitor(expr);
    return visitor.annotations_;
  }
 private:
  void VisitExpr_(const CallNode* call_node) {
    if (IsOnDeviceNode(call_node)) {
      const auto* attr = call_node->attrs.as<OnDeviceAttrs>();
      annotations_.Set(GetRef<Expr>(call_node), attr->device_type);
    }
    ExprVisitor::VisitExpr_(call_node);
  }
  Map<Expr, Integer> annotations_;
};

/*
 * \brief Return device allocation map based on the post order traversed graph.
 * For the following program:
 * .. code-block:: python
 *     x = relay.var("x")
 *     y = relay.var("y")
 *     add = relay.add(x, y)
 *     sqrt = relay.sqrt(add)
 *     log = relay.log(add)
 *     subtract = relay.subtract(sqrt, log)
 *     exp = relay.exp(subtract)
 *
 * Suppose we have annotated add, sqrt, and log with device 1, 2, and 3,
 * respectively. The fallback/default device is 4. After Rewriting the
 * program, we can have the following graph, where each copy op has both
 * source and destination device type denoting which device the data should be
 * copied from and to.
 *
 *         x     y
 *          \   /
 *          add/1
 *          /   \
 *       copy1  copy2
 *         |     |
 *      sqrt/2 log/3
 *         |     |
 *       copy3 copy4
 *          \   /
 *        subtract
 *            |
 *           exp
 *
 * To Get the device mapping of each expression, we need to propagate the
 * device information from the copy ops. This can be done in two passes.
 *  -Pass 1: Propagating the source device type to ops in a bottom-up way to the
 *           ancestors until encountering another copy op. For example, this way
 *           provides add, x, and y device types from the copy operator, `copy1`.
352 353 354
 *  -Pass 2: Propagating the destination device type of "the last" copy op to the
 *           remain nodes. For instance, this offers `subtract` and `exp` the 
 *           same device type as `copy3`.
355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373
 */

class DeviceInfo {
 public:
  static Map<Expr, Integer> GetDeviceMap(const Expr& expr) {
    DeviceInfo device_info;
    device_info.post_visitor_ = PostDfsOrderVisitor();
    device_info.post_visitor_.Visit(expr);
    if (device_info.post_visitor_.num_device_copy_ops_ > 0) {
      device_info.PropagateDeviceId();
      return device_info.device_map_;
    } else {
      return Map<Expr, Integer>();
    }
  }

 private:
  class PostDfsOrderVisitor : private ExprVisitor {
   public:
374 375
    void Visit(const Expr& expr) {
      if (const auto* fn = expr.as<FunctionNode>()) {
376 377 378
        for (const auto& param : fn->params) {
          this->VisitExpr(param);
        }
379 380 381 382 383
        this->VisitExpr(fn->body);
      } else {
        this->VisitExpr(expr);
      }
    }
384 385 386 387 388 389 390 391

   private:
    // Post order traversal.
    void VisitExpr_(const FunctionNode* fn) final {
      // TODO(zhiics) Skip annotation of function node for now.
    }

    void VisitExpr_(const ConstantNode* cn) final {
392
      post_dfs_order_.push_back(std::make_pair(cn, has_copy_));
393 394 395 396 397
    }

    void VisitExpr_(const CallNode* call) final {
      // Skip annotation nodes.
      if (!IsOnDeviceNode(call)) {
398
        if (GetDeviceCopyNode(call)) {
399
          num_device_copy_ops_++;
400 401 402 403 404 405 406 407
          bool has_copy_prev = has_copy_;
          has_copy_ = true;
          ExprVisitor::VisitExpr_(call);
          post_dfs_order_.push_back(std::make_pair(call, has_copy_));
          has_copy_ = has_copy_prev;
        } else {
          ExprVisitor::VisitExpr_(call);
          post_dfs_order_.push_back(std::make_pair(call, has_copy_));
408 409 410 411 412 413 414 415 416 417 418
        }
      }
    }

    void VisitExpr_(const TupleNode* tn) final {
      ExprVisitor::VisitExpr_(tn);
      // TODO(zhiics) Skip annotation of tuple node for now.
    }

    void VisitExpr_(const TupleGetItemNode* op) final {
      ExprVisitor::VisitExpr_(op);
419
      std::make_pair(op, has_copy_);
420 421
    }

422
    void VisitExpr_(const VarNode* vn) final {
423
      post_dfs_order_.push_back(std::make_pair(vn, has_copy_));
424
    }
425 426 427

    void VisitExpr_(const LetNode* ln) final {
      ExprVisitor::VisitExpr_(ln);
428
      post_dfs_order_.push_back(std::make_pair(ln, has_copy_));
429 430 431 432
    }

    void VisitExpr_(const IfNode* in) final {
      ExprVisitor::VisitExpr_(in);
433
      post_dfs_order_.push_back(std::make_pair(in, has_copy_));
434 435
    }

436

437
    int num_device_copy_ops_{0};
438 439
    bool has_copy_ = false;
    std::vector<std::pair<const ExprNode*, bool>> post_dfs_order_;
440 441 442
    friend DeviceInfo;
  };

443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462
  /*
   * \brief Returns a device copy node based on the current expr node. It
   * returns a device copy node either the current expr node is a device copy
   * node or the current expr node is a function node whose body is a device
   * copy node (i.e. the fused function of a device copy call node).
   */
  static const ExprNode* GetDeviceCopyNode(const ExprNode* node) {
    if (IsDeviceCopyNode(node)) {
      return node;
    } else if (const auto* call_node = dynamic_cast<const CallNode*>(node)) {
      if (const auto* fn = call_node->op.as<FunctionNode>()) {
        const ExprNode* body = fn->body.operator->();
        if (IsDeviceCopyNode(body)) {
          return body;
        }
      }
    }
    return nullptr;
  }

463 464
  void PropagateDeviceId() {
    // Bottom-up propagation.
465 466 467
    int out_dev_type = BottomUpPropagation();
    // propagation for remained nodes.
    FillPropagation(out_dev_type);
468 469
  }

470
  int BottomUpPropagation() {
471 472
    const CallNode* last_copy_node = nullptr;
    int cur_dev_type = -1;
473
    int out_dev_type = -1;
474 475
    for (auto it = post_visitor_.post_dfs_order_.crbegin();
         it != post_visitor_.post_dfs_order_.crend(); ++it) {
476
      if (const auto* node = GetDeviceCopyNode(it->first)) {
477
        last_copy_node = dynamic_cast<const CallNode*>(node);
478 479
        const auto* attrs = last_copy_node->attrs.as<DeviceCopyAttrs>();
        cur_dev_type = attrs->src_dev_type;
480 481 482
        if (out_dev_type == -1) out_dev_type = attrs->dst_dev_type;
        if (it->second) device_map_.Set(GetRef<Expr>(it->first),
                                        attrs->dst_dev_type);
483
      } else if (last_copy_node) {
484
        Expr expr = GetRef<Expr>(it->first);
485
        CHECK_EQ(device_map_.count(expr), 0U);
486
        if (it->second) device_map_.Set(expr, cur_dev_type);
487 488
      }
    }
489
      return out_dev_type;
490 491
  }

492
  void FillPropagation(int out_dev_type) {
493
    for (const auto& it : post_visitor_.post_dfs_order_) {
494 495
        Expr expr = GetRef<Expr>(it.first);
        if (!it.second) device_map_.Set(expr, out_dev_type);
496 497 498
    }
  }

499

500 501 502 503 504 505
  PostDfsOrderVisitor post_visitor_;
  Map<Expr, Integer> device_map_;
};

Expr RewriteAnnotatedOps(const Expr& expr, int fallback_device) {
  RewriteAnnotation rewrote = RewriteAnnotation();
506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551
  Expr new_expr = rewrote.Rewrite(expr, fallback_device);

  // Remove OnDevice operators. Note that these operators are only present at the
  // leaves after annotation. Therefore, we can simply reconstruct the
  // Function/Expr by removing them directly.
  if (const FunctionNode* fn = new_expr.as<FunctionNode>()) {
    auto params = fn->params;
    auto body = fn->body;
    std::vector<Expr> new_body;
    if (const TupleNode* tuple = body.as<TupleNode>()) {
      for (const auto& field : tuple->fields) {
        if (!IsOnDeviceNode(field.operator->())) {
          new_body.push_back(field);
        }
      }
      CHECK_GT(new_body.size(), 0U);
      if (new_body.size() == 1) {
        return FunctionNode::make(params, new_body[0], Type(nullptr),
                                  fn->type_params, fn->attrs);
      } else if (tuple->fields.size() == new_body.size()) {
          return new_expr;
      } else {
        Tuple tuple_body = TupleNode::make(new_body);
        return FunctionNode::make(params, tuple_body, Type(nullptr),
                                  fn->type_params, fn->attrs);
      }
    } else {
      return new_expr;
    }
  } else if (const TupleNode* tuple = new_expr.as<TupleNode>()) {
    std::vector<Expr> new_fields;
    for (const auto& field : tuple->fields) {
      if (!IsOnDeviceNode(field.operator->())) {
        new_fields.push_back(field);
      }
    }
    CHECK_GT(new_fields.size(), 0U);
    if (tuple->fields.size() == new_fields.size()) {
      return new_fields.size() == 1 ? new_fields[0] : new_expr;
    } else {
      return new_fields.size() == 1 ? new_fields[0]
                                    : TupleNode::make(new_fields);
    }
  } else {
    return new_expr;
  }
552 553 554 555 556 557 558 559 560 561 562
}

Map<Expr, Integer> CollectDeviceInfo(const Expr& expr) {
  return DeviceInfo::GetDeviceMap(expr);
}

Map<Expr, Integer> CollectDeviceAnnotationOps(const Expr& expr) {
  return AnnotatationVisitor::GetAnnotations(expr);
}

TVM_REGISTER_API("relay._ir_pass.CollectDeviceInfo")
563
.set_body_typed(CollectDeviceInfo);
564 565

TVM_REGISTER_API("relay._ir_pass.RewriteDeviceAnnotation")
566
.set_body_typed(RewriteAnnotatedOps);
567 568

TVM_REGISTER_API("relay._ir_pass.CollectDeviceAnnotationOps")
569
.set_body_typed(CollectDeviceAnnotationOps);
570

571 572 573 574 575 576 577
namespace transform {

Pass RewriteAnnotatedOps(int fallback_device) {
  runtime::TypedPackedFunc<Function(Function, Module, PassContext)> pass_func =
    [=](Function f, Module m, PassContext pc) {
    return Downcast<Function>(RewriteAnnotatedOps(f, fallback_device));
  };
578 579
  return CreateFunctionPass(pass_func, 1, "RewriteAnnotatedOps",
                            {ir::StringImm::make("InferType")});
580 581
}

582 583 584
TVM_REGISTER_API("relay._transform.RewriteDeviceAnnotation")
.set_body_typed(RewriteAnnotatedOps);

585 586
}  // namespace transform

587 588
}  // namespace relay
}  // namespace tvm