packed_func.h 37.1 KB
Newer Older
1
/*!
2
 *  Copyright (c) 2017 by Contributors
tqchen committed
3
 * \file tvm/runtime/packed_func.h
4
 * \brief Type-erased function used across TVM API.
5
 */
6 7
#ifndef TVM_RUNTIME_PACKED_FUNC_H_
#define TVM_RUNTIME_PACKED_FUNC_H_
8

9
#include <dmlc/logging.h>
10 11
#include <functional>
#include <tuple>
12 13
#include <vector>
#include <string>
14 15 16
#include <limits>
#include <memory>
#include <type_traits>
17 18 19
#include "c_runtime_api.h"
#include "module.h"
#include "ndarray.h"
20
#include "node_base.h"
21

22
namespace HalideIR {
23 24 25 26 27 28
// Forward declare type for extensions
// The header works fine without depending on this.
struct Type;
struct Expr;
}

29 30 31 32 33
// Whether use TVM runtime in header only mode.
#ifndef TVM_RUNTIME_HEADER_ONLY
#define TVM_RUNTIME_HEADER_ONLY 0
#endif

34
namespace tvm {
35 36 37
// forward declarations
class Integer;

38
namespace runtime {
39 40 41 42 43
// forward declarations
class TVMArgs;
class TVMArgValue;
class TVMRetValue;
class TVMArgsSetter;
44 45

/*!
46 47
 * \brief Packed function is a type-erased function.
 *  The arguments are passed by packed format.
48
 *
49 50 51
 *  This is an useful unified interface to call generated functions,
 *  It is the unified function function type of TVM.
 *  It corresponds to TVMFunctionHandle in C runtime API.
52 53 54
 */
class PackedFunc {
 public:
55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73
  /*!
   * \brief The internal std::function
   * \param args The arguments to the function.
   * \param rv The return value.
   *
   * \code
   *   // Example code on how to implemented FType
   *   void MyPackedFunc(TVMArgs args, TVMRetValue* rv) {
   *     // automatically convert arguments to desired type.
   *     int a0 = args[0];
   *     float a1 = args[1];
   *     ...
   *     // automatically assign values to rv
   *     std::string my_return_value = "x";
   *     *rv = my_return_value;
   *   }
   * \endcode
   */
  using FType = std::function<void (TVMArgs args, TVMRetValue* rv)>;
74
  /*! \brief default constructor */
75
  PackedFunc() {}
76 77
  /*! \brief constructor from null */
  PackedFunc(std::nullptr_t null) {}  // NOLINT(*)
78 79 80 81
  /*!
   * \brief constructing a packed function from a std::function.
   * \param body the internal container of packed function.
   */
82 83
  explicit PackedFunc(FType body) : body_(body) {}
  /*!
84
   * \brief Call packed function by directly passing in unpacked format.
85 86
   * \param args Arguments to be passed.
   * \tparam Args arguments to be passed.
87 88 89 90 91 92 93 94 95
   *
   * \code
   *   // Example code on how to call packed function
   *   void CallPacked(PackedFunc f) {
   *     // call like normal functions by pass in arguments
   *     // return value is automatically converted back
   *     int rvalue = f(1, 2.0);
   *   }
   * \endcode
96 97
   */
  template<typename... Args>
98
  inline TVMRetValue operator()(Args&& ...args) const;
99 100 101
  /*!
   * \brief Call the function in packed format.
   * \param args The arguments
102
   * \param rv The return value.
103
   */
104
  inline void CallPacked(TVMArgs args, TVMRetValue* rv) const;
105
  /*! \return the internal body function */
106
  inline FType body() const;
107 108 109 110 111 112 113 114
  /*! \return Whether the packed function is nullptr */
  bool operator==(std::nullptr_t null) const {
    return body_ == nullptr;
  }
  /*! \return Whether the packed function is not nullptr */
  bool operator!=(std::nullptr_t null) const {
    return body_ != nullptr;
  }
115 116 117 118 119 120

 private:
  /*! \brief internal container of packed function */
  FType body_;
};

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
/*!
 * \brief Please refer to \ref TypedPackedFuncAnchor "TypedPackedFunc<R(Args..)>"
 */
template<typename FType>
class TypedPackedFunc;

/*!
 * \anchor TypedPackedFuncAnchor
 * \brief A PackedFunc wrapper to provide typed function signature.
 * It is backed by a PackedFunc internally.
 *
 * TypedPackedFunc enables compile time type checking.
 * TypedPackedFunc works with the runtime system:
 * - It can be passed as an argument of PackedFunc.
 * - It can be assigned to TVMRetValue.
 * - It can be directly converted to a type-erased PackedFunc.
 *
 * Developers should prefer TypedPackedFunc over PackedFunc in C++ code
 * as it enables compile time checking.
 * We can construct a TypedPackedFunc from a lambda function
 * with the same signature.
 *
 * \code
 *  // user defined lambda function.
 *  auto addone = [](int x)->int {
 *    return x + 1;
 *  };
 *  // We can directly convert
 *  // lambda function to TypedPackedFunc
 *  TypedPackedFunc<int(int)> ftyped(addone);
 *  // invoke the function.
 *  int y = ftyped(1);
 *  // Can be directly converted to PackedFunc
 *  PackedFunc packed = ftype;
 * \endcode
 * \tparam R The return value of the function.
 * \tparam Args The argument signature of the function.
 */
template<typename R, typename ...Args>
class TypedPackedFunc<R(Args...)> {
 public:
  /*! \brief short hand for this function type */
  using TSelf = TypedPackedFunc<R(Args...)>;
  /*! \brief default constructor */
  TypedPackedFunc() {}
166 167
  /*! \brief constructor from null */
  TypedPackedFunc(std::nullptr_t null) {}  // NOLINT(*)
168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184
  /*!
   * \brief construct by wrap a PackedFunc
   *
   * Example usage:
   * \code
   * PackedFunc packed([](TVMArgs args, TVMRetValue *rv) {
   *   int x = args[0];
   *   *rv = x + 1;
   *  });
   * // construct from packed function
   * TypedPackedFunc<int(int)> ftyped(packed);
   * // call the typed version.
   * CHECK_EQ(ftyped(1), 2);
   * \endcode
   *
   * \param packed The packed function
   */
185 186 187 188 189 190 191 192 193 194 195
  inline TypedPackedFunc(PackedFunc packed);  // NOLINT(*)
  /*!
   * \brief constructor from TVMRetValue
   * \param value The TVMRetValue
   */
  inline TypedPackedFunc(const TVMRetValue& value);  // NOLINT(*)
  /*!
   * \brief constructor from TVMArgValue
   * \param value The TVMArgValue
   */
  inline TypedPackedFunc(const TVMArgValue& value);  // NOLINT(*)
196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215
  /*!
   * \brief construct from a lambda function with the same signature.
   *
   * Example usage:
   * \code
   * auto typed_lambda = [](int x)->int { return x + 1; }
   * // construct from packed function
   * TypedPackedFunc<int(int)> ftyped(typed_lambda);
   * // call the typed version.
   * CHECK_EQ(ftyped(1), 2);
   * \endcode
   *
   * \param typed_lambda typed lambda function.
   * \tparam FLambda the type of the lambda function.
   */
  template<typename FLambda,
           typename = typename std::enable_if<
             std::is_convertible<FLambda,
                                 std::function<R(Args...)>
                                 >::value>::type>
216
  TypedPackedFunc(const FLambda& typed_lambda) {  // NOLINT(*)
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 264 265 266 267 268 269 270 271
    this->AssignTypedLambda(typed_lambda);
  }
  /*!
   * \brief copy assignment operator from typed lambda
   *
   * Example usage:
   * \code
   * // construct from packed function
   * TypedPackedFunc<int(int)> ftyped;
   * ftyped = [](int x) { return x + 1; }
   * // call the typed version.
   * CHECK_EQ(ftyped(1), 2);
   * \endcode
   *
   * \param typed_lambda typed lambda function.
   * \tparam FLambda the type of the lambda function.
   * \returns reference to self.
   */
  template<typename FLambda,
           typename = typename std::enable_if<
             std::is_convertible<FLambda,
                                 std::function<R(Args...)>
                                 >::value>::type>
  TSelf& operator=(FLambda typed_lambda) {  // NOLINT(*)
    this->AssignTypedLambda(typed_lambda);
    return *this;
  }
  /*!
   * \brief copy assignment operator from PackedFunc.
   * \param packed The packed function.
   * \returns reference to self.
   */
  TSelf& operator=(PackedFunc packed) {
    packed_ = packed;
    return *this;
  }
  /*!
   * \brief Invoke the operator.
   * \param args The arguments
   * \returns The return value.
   */
  inline R operator()(Args ...args) const;
  /*!
   * \brief convert to PackedFunc
   * \return the internal PackedFunc
   */
  operator PackedFunc() const {
    return packed();
  }
  /*!
   * \return reference the internal PackedFunc
   */
  const PackedFunc& packed() const {
    return packed_;
  }
272 273 274 275 276 277 278 279
  /*! \return Whether the packed function is nullptr */
  bool operator==(std::nullptr_t null) const {
    return packed_ == nullptr;
  }
  /*! \return Whether the packed function is not nullptr */
  bool operator!=(std::nullptr_t null) const {
    return packed_ != nullptr;
  }
280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295

 private:
  friend class TVMRetValue;
  /*! \brief The internal packed function */
  PackedFunc packed_;
  /*!
   * \brief Assign the packed field using a typed lambda function.
   *
   * \param flambda The lambda function.
   * \tparam FLambda The lambda function type.
   * \note We capture the lambda when possible for maximum efficiency.
   */
  template<typename FLambda>
  inline void AssignTypedLambda(FLambda flambda);
};

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
/*! \brief Arguments into TVM functions. */
class TVMArgs {
 public:
  const TVMValue* values;
  const int* type_codes;
  int num_args;
  /*!
   * \brief constructor
   * \param values The argument values
   * \param type_codes The argument type codes
   * \param num_args number of arguments.
   */
  TVMArgs(const TVMValue* values,
          const int* type_codes,
          int num_args)
      : values(values),
        type_codes(type_codes),
        num_args(num_args) { }
  /*! \return size of the arguments */
  inline int size() const;
  /*!
   * \brief Get i-th argument
   * \param i the index.
   * \return the ith argument.
   */
  inline TVMArgValue operator[](int i) const;
};

/*!
 * \brief Convert type code to its name
 * \param type_code The type code .
 * \return The name of type code.
 */
inline const char* TypeCode2Str(int type_code);

/*!
 * \brief convert a string to TVM type.
 * \param s The string to be converted.
 * \return The corresponding tvm type.
 */
inline TVMType String2TVMType(std::string s);

338 339 340 341 342 343 344
/*!
 * \brief convert a TVM type to string.
 * \param t The type to be converted.
 * \return The corresponding tvm type in string.
 */
inline std::string TVMType2String(TVMType t);

345 346 347 348 349 350
// macro to check type code.
#define TVM_CHECK_TYPE_CODE(CODE, T)                           \
  CHECK_EQ(CODE, T) << " expected "                            \
  << TypeCode2Str(T) << " but get " << TypeCode2Str(CODE)      \

/*!
351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366
 * \brief Type traits to mark if a class is tvm extension type.
 *
 * To enable extension type in C++ must be register () ed via marco.
 * TVM_REGISTER_EXT_TYPE(TypeName) after defining this with this traits.
 *
 * Extension class can be passed and returned via PackedFunc in all tvm runtime.
 * Internally extension class is stored as T*.
 *
 * \tparam T the typename
 */
template<typename T>
struct extension_class_info {
  static const int code = 0;
};

/*!
367
 * \brief Runtime function table about extension type.
368
 */
369 370
class ExtTypeVTable {
 public:
371 372 373 374
  /*! \brief function to be called to delete a handle */
  void (*destroy)(void* handle);
  /*! \brief function to be called when clone a handle */
  void* (*clone)(void* handle);
375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391
  /*!
   * \brief Register type
   * \tparam T The type to be register.
   * \return The registered vtable.
   */
  template <typename T>
  static inline ExtTypeVTable* Register_();
  /*!
   * \brief Get a vtable based on type code.
   * \param type_code The type code
   * \return The registered vtable.
   */
  TVM_DLL static ExtTypeVTable* Get(int type_code);

 private:
  // Internal registration function.
  TVM_DLL static ExtTypeVTable* RegisterInternal(int type_code, const ExtTypeVTable& vt);
392 393 394
};

/*!
395 396 397 398 399 400
 * \brief Internal base class to
 *  handle conversion to POD values.
 */
class TVMPODValue_ {
 public:
  operator double() const {
401 402 403 404 405 406
    // Allow automatic conversion from int to float
    // This avoids errors when user pass in int from
    // the frontend while the API expects a float.
    if (type_code_ == kDLInt) {
      return static_cast<double>(value_.v_int64);
    }
407
    TVM_CHECK_TYPE_CODE(type_code_, kDLFloat);
408 409 410
    return value_.v_float64;
  }
  operator int64_t() const {
411
    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
412 413 414
    return value_.v_int64;
  }
  operator uint64_t() const {
415
    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
416 417 418
    return value_.v_int64;
  }
  operator int() const {
419
    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
420 421 422 423 424
    CHECK_LE(value_.v_int64,
             std::numeric_limits<int>::max());
    return static_cast<int>(value_.v_int64);
  }
  operator bool() const {
425
    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
426 427 428 429 430 431 432 433
    return value_.v_int64 != 0;
  }
  operator void*() const {
    if (type_code_ == kNull) return nullptr;
    if (type_code_ == kArrayHandle) return value_.v_handle;
    TVM_CHECK_TYPE_CODE(type_code_, kHandle);
    return value_.v_handle;
  }
434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449
  operator DLTensor*() const {
    if (type_code_ == kArrayHandle ||
        type_code_ == kNDArrayContainer) {
      return static_cast<DLTensor*>(value_.v_handle);
    } else {
      if (type_code_ == kNull) return nullptr;
      LOG(FATAL) << "Expected "
                 << "DLTensor* or NDArray but get "
                 << TypeCode2Str(type_code_);
      return nullptr;
    }
  }
  operator NDArray() const {
    if (type_code_ == kNull) return NDArray();
    TVM_CHECK_TYPE_CODE(type_code_, kNDArrayContainer);
    return NDArray(static_cast<NDArray::Container*>(value_.v_handle));
450
  }
451 452 453 454
  operator TVMContext() const {
    TVM_CHECK_TYPE_CODE(type_code_, kTVMContext);
    return value_.v_ctx;
  }
455 456
  template<typename TExtension>
  const TExtension& AsExtension() const {
457 458
    CHECK_LT(type_code_, kExtEnd);
    return static_cast<TExtension*>(value_.v_handle)[0];
459
  }
460 461 462 463 464 465 466 467 468 469 470 471
  int type_code() const {
    return type_code_;
  }
  /*!
   * \brief return handle as specific pointer type.
   * \tparam T the data type.
   * \return The pointer type.
   */
  template<typename T>
  T* ptr() const {
    return static_cast<T*>(value_.v_handle);
  }
472 473 474 475 476 477 478 479

 protected:
  friend class TVMArgsSetter;
  friend class TVMRetValue;
  TVMPODValue_() : type_code_(kNull) {}
  TVMPODValue_(TVMValue value, int type_code)
      : value_(value), type_code_(type_code) {}

480 481 482 483 484 485 486 487 488 489 490 491 492 493
  /*! \brief The value */
  TVMValue value_;
  /*! \brief the type code */
  int type_code_;
};

/*!
 * \brief A single argument value to PackedFunc.
 *  Containing both type_code and TVMValue
 *
 *  Provides utilities to do type cast into other types.
 */
class TVMArgValue : public TVMPODValue_ {
 public:
494 495
  /*! \brief default constructor */
  TVMArgValue() {}
496 497 498 499 500 501 502 503 504 505 506 507 508 509 510
  /*!
   * \brief constructor
   * \param value of the function
   * \param type_code The type code.
   */
  TVMArgValue(TVMValue value, int type_code)
      : TVMPODValue_(value, type_code) {
  }
  // reuse converter from parent
  using TVMPODValue_::operator double;
  using TVMPODValue_::operator int64_t;
  using TVMPODValue_::operator uint64_t;
  using TVMPODValue_::operator int;
  using TVMPODValue_::operator bool;
  using TVMPODValue_::operator void*;
511 512
  using TVMPODValue_::operator DLTensor*;
  using TVMPODValue_::operator NDArray;
513
  using TVMPODValue_::operator TVMContext;
514

515 516
  // conversion operator.
  operator std::string() const {
517 518
    if (type_code_ == kTVMType) {
      return TVMType2String(operator TVMType());
519 520 521 522 523 524
    } else if (type_code_ == kBytes) {
      TVMByteArray* arr = static_cast<TVMByteArray*>(value_.v_handle);
      return std::string(arr->data, arr->size);
    } else {
      TVM_CHECK_TYPE_CODE(type_code_, kStr);
      return std::string(value_.v_str);
525
    }
526 527 528 529 530
  }
  operator TVMType() const {
    if (type_code_ == kStr) {
      return String2TVMType(operator std::string());
    }
531 532 533 534 535 536
    // None type
    if (type_code_ == kNull) {
      TVMType t;
      t.code = kHandle; t.bits = 0; t.lanes = 0;
      return t;
    }
537 538 539 540
    TVM_CHECK_TYPE_CODE(type_code_, kTVMType);
    return value_.v_type;
  }
  operator PackedFunc() const {
541
    if (type_code_ == kNull) return PackedFunc();
542 543 544
    TVM_CHECK_TYPE_CODE(type_code_, kFuncHandle);
    return *ptr<PackedFunc>();
  }
545 546 547 548
  template<typename FType>
  operator TypedPackedFunc<FType>() const {
    return TypedPackedFunc<FType>(operator PackedFunc());
  }
549 550 551 552
  operator Module() const {
    TVM_CHECK_TYPE_CODE(type_code_, kModuleHandle);
    return *ptr<Module>();
  }
553 554 555
  const TVMValue& value() const {
    return value_;
  }
556 557 558 559
  // Deferred extension handler.
  template<typename TNodeRef>
  inline TNodeRef AsNodeRef() const;
  template<typename T,
560
           typename = typename std::enable_if<
561 562
             std::is_class<T>::value>::type>
  inline operator T() const;
563 564 565 566
  template<typename TNodeRef,
           typename = typename std::enable_if<
             std::is_class<TNodeRef>::value>::type>
  inline bool IsNodeType() const;
567 568
  inline operator HalideIR::Type() const;
  inline operator HalideIR::Expr() const;
569
  inline operator tvm::Integer() const;
570
  // get internal node ptr, if it is node
571
  inline NodePtr<Node>& node_sptr();
572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591
};

/*!
 * \brief Return Value container,
 *  Unlike TVMArgValue, which only holds reference and do not delete
 *  the underlying container during destruction.
 *
 *  TVMRetValue holds value and will manage the underlying containers
 *  when it stores a complicated data type.
 */
class TVMRetValue : public TVMPODValue_ {
 public:
  /*! \brief default constructor */
  TVMRetValue() {}
  /*!
   * \brief move constructor from anoter return value.
   * \param other The other return value.
   */
  TVMRetValue(TVMRetValue&& other)
      : TVMPODValue_(other.value_, other.type_code_) {
592 593
    other.value_.v_handle = nullptr;
    other.type_code_ = kNull;
594 595 596 597 598 599 600 601 602 603 604 605
  }
  /*! \brief destructor */
  ~TVMRetValue() {
    this->Clear();
  }
  // reuse converter from parent
  using TVMPODValue_::operator double;
  using TVMPODValue_::operator int64_t;
  using TVMPODValue_::operator uint64_t;
  using TVMPODValue_::operator int;
  using TVMPODValue_::operator bool;
  using TVMPODValue_::operator void*;
606
  using TVMPODValue_::operator DLTensor*;
607
  using TVMPODValue_::operator TVMContext;
608
  using TVMPODValue_::operator NDArray;
609
  TVMRetValue(const TVMRetValue& other) : TVMPODValue_() {
610 611 612 613
    this->Assign(other);
  }
  // conversion operators
  operator std::string() const {
614 615
    if (type_code_ == kTVMType) {
      return TVMType2String(operator TVMType());
616 617
    } else if (type_code_ == kBytes) {
      return *ptr<std::string>();
618
    }
619 620 621 622 623 624 625 626 627 628 629
    TVM_CHECK_TYPE_CODE(type_code_, kStr);
    return *ptr<std::string>();
  }
  operator TVMType() const {
    if (type_code_ == kStr) {
      return String2TVMType(operator std::string());
    }
    TVM_CHECK_TYPE_CODE(type_code_, kTVMType);
    return value_.v_type;
  }
  operator PackedFunc() const {
630
    if (type_code_ == kNull) return PackedFunc();
631 632 633
    TVM_CHECK_TYPE_CODE(type_code_, kFuncHandle);
    return *ptr<PackedFunc>();
  }
634 635 636 637
  template<typename FType>
  operator TypedPackedFunc<FType>() const {
    return TypedPackedFunc<FType>(operator PackedFunc());
  }
638 639 640 641
  operator Module() const {
    TVM_CHECK_TYPE_CODE(type_code_, kModuleHandle);
    return *ptr<Module>();
  }
642 643 644 645 646 647 648 649 650
  // Assign operators
  TVMRetValue& operator=(TVMRetValue&& other) {
    this->Clear();
    value_ = other.value_;
    type_code_ = other.type_code_;
    other.type_code_ = kNull;
    return *this;
  }
  TVMRetValue& operator=(double value) {
651
    this->SwitchToPOD(kDLFloat);
652 653 654 655 656 657 658 659 660 661 662 663 664 665
    value_.v_float64 = value;
    return *this;
  }
  TVMRetValue& operator=(std::nullptr_t value) {
    this->SwitchToPOD(kNull);
    value_.v_handle = value;
    return *this;
  }
  TVMRetValue& operator=(void* value) {
    this->SwitchToPOD(kHandle);
    value_.v_handle = value;
    return *this;
  }
  TVMRetValue& operator=(int64_t value) {
666
    this->SwitchToPOD(kDLInt);
667 668 669 670
    value_.v_int64 = value;
    return *this;
  }
  TVMRetValue& operator=(int value) {
671
    this->SwitchToPOD(kDLInt);
672 673 674
    value_.v_int64 = value;
    return *this;
  }
675 676 677 678 679
  TVMRetValue& operator=(TVMContext value) {
    this->SwitchToPOD(kTVMContext);
    value_.v_ctx = value;
    return *this;
  }
680 681 682 683 684 685
  TVMRetValue& operator=(TVMType t) {
    this->SwitchToPOD(kTVMType);
    value_.v_type = t;
    return *this;
  }
  TVMRetValue& operator=(bool value) {
686
    this->SwitchToPOD(kDLInt);
687 688 689 690 691 692 693
    value_.v_int64 = value;
    return *this;
  }
  TVMRetValue& operator=(std::string value) {
    this->SwitchToClass(kStr, value);
    return *this;
  }
694 695 696 697
  TVMRetValue& operator=(TVMByteArray value) {
    this->SwitchToClass(kBytes, std::string(value.data, value.size));
    return *this;
  }
698 699 700 701 702 703 704
  TVMRetValue& operator=(NDArray other) {
    this->Clear();
    type_code_ = kNDArrayContainer;
    value_.v_handle = other.data_;
    other.data_ = nullptr;
    return *this;
  }
705 706 707 708
  TVMRetValue& operator=(PackedFunc f) {
    this->SwitchToClass(kFuncHandle, f);
    return *this;
  }
709 710 711 712
  template<typename FType>
  TVMRetValue& operator=(const TypedPackedFunc<FType>& f) {
    return operator=(f.packed());
  }
713 714 715 716
  TVMRetValue& operator=(Module m) {
    this->SwitchToClass(kModuleHandle, m);
    return *this;
  }
717 718 719 720
  TVMRetValue& operator=(const TVMRetValue& other) {  // NOLINT(*0
    this->Assign(other);
    return *this;
  }
721
  TVMRetValue& operator=(const TVMArgValue& other) {
722 723 724
    this->Assign(other);
    return *this;
  }
725 726 727 728 729 730 731 732
  template<typename T,
           typename = typename std::enable_if<
             extension_class_info<T>::code != 0>::type>
  TVMRetValue& operator=(const T& other) {
    this->SwitchToClass<T>(
        extension_class_info<T>::code, other);
    return *this;
  }
733 734 735 736 737 738 739 740 741 742 743 744
  /*!
   * \brief Move the value back to front-end via C API.
   *  This marks the current container as null.
   *  The managed resources is moved to front-end and
   *  the front end should take charge in managing them.
   *
   * \param ret_value The return value.
   * \param ret_type_code The return type code.
   */
  void MoveToCHost(TVMValue* ret_value,
                   int* ret_type_code) {
    // cannot move str; need specially handle.
745
    CHECK(type_code_ != kStr && type_code_ != kBytes);
746 747 748 749
    *ret_value = value_;
    *ret_type_code = type_code_;
    type_code_ = kNull;
  }
750 751 752 753
  /*! \return The value field, if the data is POD */
  const TVMValue& value() const {
    CHECK(type_code_ != kNodeHandle &&
          type_code_ != kFuncHandle &&
754
          type_code_ != kModuleHandle &&
755 756 757
          type_code_ != kStr) << "TVMRetValue.value can only be used for POD data";
    return value_;
  }
758
  // NodeRef related extenstions: in tvm/packed_func_ext.h
759 760 761 762 763 764
  template<typename T,
           typename = typename std::enable_if<
             std::is_class<T>::value>::type>
  inline operator T() const;
  template<typename TNodeRef>
  inline TNodeRef AsNodeRef() const;
765
  inline TVMRetValue& operator=(const NodeRef& other);
766
  inline TVMRetValue& operator=(const NodePtr<Node>& other);
767
  // type related
768 769
  inline operator HalideIR::Type() const;
  inline TVMRetValue& operator=(const HalideIR::Type& other);
770 771 772 773 774

 private:
  template<typename T>
  void Assign(const T& other) {
    switch (other.type_code()) {
775
      case kStr: {
776 777 778
        SwitchToClass<std::string>(kStr, other);
        break;
      }
779 780 781 782
      case kBytes: {
        SwitchToClass<std::string>(kBytes, other);
        break;
      }
783 784 785 786
      case kFuncHandle: {
        SwitchToClass<PackedFunc>(kFuncHandle, other);
        break;
      }
787
      case kModuleHandle: {
788
        SwitchToClass<Module>(kModuleHandle, other);
789 790
        break;
      }
791 792 793 794
      case kNDArrayContainer: {
        *this = other.operator NDArray();
        break;
      }
795
      case kNodeHandle: {
796 797
        SwitchToClass<NodePtr<Node> >(
            kNodeHandle, *other.template ptr<NodePtr<Node> >());
798 799 800
        break;
      }
      default: {
801 802 803 804
        if (other.type_code() < kExtBegin) {
          SwitchToPOD(other.type_code());
          value_ = other.value_;
        } else {
805 806 807
#if TVM_RUNTIME_HEADER_ONLY
          LOG(FATAL) << "Header only mode do not support ext type";
#else
808 809 810 811 812
          this->Clear();
          type_code_ = other.type_code();
          value_.v_handle =
              (*(ExtTypeVTable::Get(other.type_code())->clone))(
                  other.value().v_handle);
813
#endif
814
        }
815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840
        break;
      }
    }
  }
  // get the internal container.
  void SwitchToPOD(int type_code) {
    if (type_code_ != type_code) {
      this->Clear();
      type_code_ = type_code;
    }
  }
  template<typename T>
  void SwitchToClass(int type_code, T v) {
    if (type_code_ != type_code) {
      this->Clear();
      type_code_ = type_code;
      value_.v_handle = new T(v);
    } else {
      *static_cast<T*>(value_.v_handle) = v;
    }
  }
  void Clear() {
    if (type_code_ == kNull) return;
    switch (type_code_) {
      case kStr: delete ptr<std::string>(); break;
      case kFuncHandle: delete ptr<PackedFunc>(); break;
841
      case kModuleHandle: delete ptr<Module>(); break;
842
      case kNodeHandle: delete ptr<NodePtr<Node> >(); break;
843 844 845 846
      case kNDArrayContainer: {
        static_cast<NDArray::Container*>(value_.v_handle)->DecRef();
        break;
      }
847
    }
848
    if (type_code_ > kExtBegin) {
849 850 851
#if TVM_RUNTIME_HEADER_ONLY
          LOG(FATAL) << "Header only mode do not support ext type";
#else
852
      (*(ExtTypeVTable::Get(type_code_)->destroy))(value_.v_handle);
853
#endif
854
    }
855 856 857 858 859 860 861
    type_code_ = kNull;
  }
};

// implementation details
inline const char* TypeCode2Str(int type_code) {
  switch (type_code) {
862 863 864
    case kDLInt: return "int";
    case kDLUInt: return "uint";
    case kDLFloat: return "float";
865
    case kStr: return "str";
866
    case kBytes: return "bytes";
867
    case kHandle: return "handle";
868 869 870 871
    case kNull: return "NULL";
    case kNodeHandle: return "NodeHandle";
    case kArrayHandle: return "ArrayHandle";
    case kTVMType: return "TVMType";
872
    case kTVMContext: return "TVMContext";
873
    case kFuncHandle: return "FunctionHandle";
874
    case kModuleHandle: return "ModuleHandle";
875
    case kNDArrayContainer: return "NDArrayContainer";
876 877 878 879 880
    default: LOG(FATAL) << "unknown type_code="
                        << static_cast<int>(type_code); return "";
  }
}

nhynes committed
881
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
882
inline std::ostream& operator<<(std::ostream& os, TVMType t) {  // NOLINT(*)
883 884 885
  if (t.bits == 1 && t.lanes == 1 && t.code == kDLUInt) {
    os << "bool"; return os;
  }
886 887 888
  os << TypeCode2Str(t.code);
  if (t.code == kHandle) return os;
  os << static_cast<int>(t.bits);
889 890 891 892 893
  if (t.lanes != 1) {
    os << 'x' << static_cast<int>(t.lanes);
  }
  return os;
}
894

nhynes committed
895
#endif
896 897

inline std::string TVMType2String(TVMType t) {
898
  if (t.bits == 0) return "";
nhynes committed
899
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
900 901 902
  std::ostringstream os;
  os << t;
  return os.str();
nhynes committed
903
#else
904 905 906
  if (t.bits == 1 && t.lanes == 1 && t.code == kDLUInt) {
    return "bool";
  }
nhynes committed
907 908 909 910 911 912 913 914
  repr += TypeCode2Str(t.code);
  if (t.code == kHandle) return repr;
  repr += std::to_string(static_cast<int>(t.bits));
  if (t.lanes != 1) {
    repr += "x" + std::to_string(static_cast<int>(t.lanes));
  }
  return repr;
#endif
915 916
}

917 918
inline TVMType String2TVMType(std::string s) {
  TVMType t;
919 920 921 922 923
  // handle None type
  if (s.length() == 0) {
    t.bits = 0; t.lanes = 0; t.code = kHandle;
    return t;
  }
924 925 926
  t.bits = 32; t.lanes = 1;
  const char* scan;
  if (s.substr(0, 3) == "int") {
927
    t.code = kDLInt;  scan = s.c_str() + 3;
928
  } else if (s.substr(0, 4) == "uint") {
929
    t.code = kDLUInt; scan = s.c_str() + 4;
930
  } else if (s.substr(0, 5) == "float") {
931
    t.code = kDLFloat; scan = s.c_str() + 5;
932
  } else if (s.substr(0, 6) == "handle") {
933 934 935
    t.code = kHandle;
    t.bits = 64;  // handle uses 64 bit by default.
    scan = s.c_str() + 6;
936 937 938 939 940
  } else if (s == "bool") {
    t.code = kDLUInt;
    t.bits = 1;
    t.lanes = 1;
    return t;
941 942 943 944
  } else {
    scan = s.c_str();
    LOG(FATAL) << "unknown type " << s;
  }
nhynes committed
945
  char* xdelim;  // emulate sscanf("%ux%u", bits, lanes)
946 947
  uint8_t bits = static_cast<uint8_t>(strtoul(scan, &xdelim, 10));
  if (bits != 0) t.bits = bits;
948
  char* endpt = xdelim;
nhynes committed
949
  if (*xdelim == 'x') {
950
    t.lanes = static_cast<uint16_t>(strtoul(xdelim + 1, &endpt, 10));
nhynes committed
951
  }
952
  CHECK(endpt == s.c_str() + s.length()) << "unknown type " << s;
953 954 955 956 957 958 959
  return t;
}

inline TVMArgValue TVMArgs::operator[](int i) const {
  CHECK_LT(i, num_args)
      << "not enough argument passed, "
      << num_args << " passed"
960
      << " but request arg[" << i << "].";
961 962 963 964 965 966 967 968 969
  return TVMArgValue(values[i], type_codes[i]);
}

inline int TVMArgs::size() const {
  return num_args;
}

inline void PackedFunc::CallPacked(TVMArgs args, TVMRetValue* rv) const {
  body_(args, rv);
970 971
}

972 973 974 975
inline PackedFunc::FType PackedFunc::body() const {
  return body_;
}

976 977


978 979
// internal namespace
namespace detail {
980 981

template<bool stop, std::size_t I, typename F>
982
struct for_each_dispatcher {
983 984 985 986 987
  template<typename T, typename ...Args>
  static void run(const F& f, T&& value, Args&&... args) {  // NOLINT(*)
    f(I, std::forward<T>(value));
    for_each_dispatcher<sizeof...(Args) == 0, (I+1), F>
        ::run(f, std::forward<Args>(args)...);
988 989 990
  }
};

991 992 993
template<std::size_t I, typename F>
struct for_each_dispatcher<true, I, F>  {
  static void run(const F& f) {}  // NOLINT(*)
994 995 996
};

template<typename F, typename ...Args>
997 998 999
inline void for_each(const F& f, Args&&... args) {  // NOLINT(*)
  for_each_dispatcher<sizeof...(Args) == 0, 0, F>
      ::run(f, std::forward<Args>(args)...);
1000
}
1001
}  // namespace detail
1002

1003 1004 1005
/* \brief argument settter to PackedFunc */
class TVMArgsSetter {
 public:
1006 1007
  TVMArgsSetter(TVMValue* values, int* type_codes)
      : values_(values), type_codes_(type_codes) {}
1008 1009
  // setters for POD types
  template<typename T,
1010 1011
           typename = typename std::enable_if<
             std::is_integral<T>::value>::type>
1012 1013
  void operator()(size_t i, T value) const {
    values_[i].v_int64 = static_cast<int64_t>(value);
1014
    type_codes_[i] = kDLInt;
1015 1016 1017 1018 1019
  }
  void operator()(size_t i, uint64_t value) const {
    values_[i].v_int64 = static_cast<int64_t>(value);
    CHECK_LE(value,
             static_cast<uint64_t>(std::numeric_limits<int64_t>::max()));
1020
    type_codes_[i] = kDLInt;
1021 1022 1023
  }
  void operator()(size_t i, double value) const {
    values_[i].v_float64 = value;
1024
    type_codes_[i] = kDLFloat;
1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037
  }
  void operator()(size_t i, std::nullptr_t value) const {
    values_[i].v_handle = value;
    type_codes_[i] = kNull;
  }
  void operator()(size_t i, const TVMArgValue& value) const {
    values_[i] = value.value_;
    type_codes_[i] = value.type_code_;
  }
  void operator()(size_t i, void* value) const {
    values_[i].v_handle = value;
    type_codes_[i] = kHandle;
  }
1038
  void operator()(size_t i, DLTensor* value) const {
1039 1040 1041
    values_[i].v_handle = value;
    type_codes_[i] = kArrayHandle;
  }
1042 1043 1044 1045
  void operator()(size_t i, TVMContext value) const {
    values_[i].v_ctx = value;
    type_codes_[i] = kTVMContext;
  }
1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056
  void operator()(size_t i, TVMType value) const {
    values_[i].v_type = value;
    type_codes_[i] = kTVMType;
  }
  void operator()(size_t i, const char* value) const {
    values_[i].v_str = value;
    type_codes_[i] = kStr;
  }
  // setters for container type
  // They must be reference(instead of const ref)
  // to make sure they are alive in the tuple(instead of getting converted)
1057
  void operator()(size_t i, const std::string& value) const {  // NOLINT(*)
1058 1059 1060
    values_[i].v_str = value.c_str();
    type_codes_[i] = kStr;
  }
1061 1062
  void operator()(size_t i, const TVMByteArray& value) const {  // NOLINT(*)
    values_[i].v_handle = const_cast<TVMByteArray*>(&value);
1063 1064
    type_codes_[i] = kBytes;
  }
1065 1066
  void operator()(size_t i, const PackedFunc& value) const {  // NOLINT(*)
    values_[i].v_handle = const_cast<PackedFunc*>(&value);
1067 1068
    type_codes_[i] = kFuncHandle;
  }
1069 1070 1071 1072
  template<typename FType>
  void operator()(size_t i, const TypedPackedFunc<FType>& value) const {  // NOLINT(*)
    operator()(i, value.packed());
  }
1073 1074
  void operator()(size_t i, const Module& value) const {  // NOLINT(*)
    values_[i].v_handle = const_cast<Module*>(&value);
1075 1076
    type_codes_[i] = kModuleHandle;
  }
1077 1078 1079 1080
  void operator()(size_t i, const NDArray& value) const {  // NOLINT(*)
    values_[i].v_handle = value.data_;
    type_codes_[i] = kNDArrayContainer;
  }
1081
  void operator()(size_t i, const TVMRetValue& value) const {  // NOLINT(*)
1082 1083 1084 1085
    if (value.type_code() == kStr) {
      values_[i].v_str = value.ptr<std::string>()->c_str();
      type_codes_[i] = kStr;
    } else {
1086
      CHECK_NE(value.type_code(), kBytes) << "not handled.";
1087 1088 1089 1090
      values_[i] = value.value_;
      type_codes_[i] = value.type_code();
    }
  }
1091 1092 1093 1094 1095
  // extension
  template<typename T,
           typename = typename std::enable_if<
             extension_class_info<T>::code != 0>::type>
  inline void operator()(size_t i, const T& value) const;
1096
  // NodeRef related extenstions: in tvm/packed_func_ext.h
1097
  inline void operator()(size_t i, const NodeRef& other) const;  // NOLINT(*)
1098
  inline void operator()(size_t i, const HalideIR::Type& t) const;
1099 1100 1101 1102 1103 1104 1105 1106

 private:
  /*! \brief The values fields */
  TVMValue* values_;
  /*! \brief The type code fields */
  int* type_codes_;
};

1107
template<typename... Args>
1108
inline TVMRetValue PackedFunc::operator()(Args&& ...args) const {
1109
  const int kNumArgs = sizeof...(Args);
1110 1111 1112
  const int kArraySize = kNumArgs > 0 ? kNumArgs : 1;
  TVMValue values[kArraySize];
  int type_codes[kArraySize];
1113
  detail::for_each(TVMArgsSetter(values, type_codes),
1114
                   std::forward<Args>(args)...);
1115 1116 1117
  TVMRetValue rv;
  body_(TVMArgs(values, type_codes, kNumArgs), &rv);
  return rv;
1118
}
1119

1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184
namespace detail {
template<typename R, int nleft, int index, typename F>
struct unpack_call_dispatcher {
  template<typename ...Args>
  static void run(const F& f,
                  const TVMArgs& args_pack,
                  TVMRetValue* rv,
                  Args&&... unpacked_args) {
    unpack_call_dispatcher<R, nleft - 1, index + 1, F>
        ::run(f, args_pack, rv,
              std::forward<Args>(unpacked_args)...,
              args_pack[index]);
  }
};

template<typename R, int index, typename F>
struct unpack_call_dispatcher<R, 0, index, F> {
  template<typename ...Args>
  static void run(const F& f,
                  const TVMArgs& args_pack,
                  TVMRetValue* rv,
                  Args&&... unpacked_args) {
    *rv = R(f(std::forward<Args>(unpacked_args)...));
  }
};

template<int index, typename F>
struct unpack_call_dispatcher<void, 0, index, F> {
  template<typename ...Args>
  static void run(const F& f,
                  const TVMArgs& args_pack,
                  TVMRetValue* rv,
                  Args&&... unpacked_args) {
    f(std::forward<Args>(unpacked_args)...);
  }
};

template<typename R, int nargs, typename F>
inline void unpack_call(const F& f, const TVMArgs& args, TVMRetValue* rv) {
  unpack_call_dispatcher<R, nargs, 0, F>::run(f, args, rv);
}

template<typename R, typename ...Args>
inline R call_packed(const PackedFunc& pf, Args&& ...args) {
  return R(pf(std::forward<Args>(args)...));
}

template<typename R>
struct typed_packed_call_dispatcher {
  template<typename ...Args>
  static inline R run(const PackedFunc& pf, Args&& ...args) {
    return pf(std::forward<Args>(args)...);
  }
};

template<>
struct typed_packed_call_dispatcher<void> {
  template<typename ...Args>
  static inline void run(const PackedFunc& pf, Args&& ...args) {
    pf(std::forward<Args>(args)...);
  }
};
}  // namespace detail

template<typename R, typename ...Args>
1185 1186 1187 1188
TypedPackedFunc<R(Args...)>::TypedPackedFunc(PackedFunc packed)
  : packed_(packed) {}

template<typename R, typename ...Args>
1189 1190 1191 1192 1193 1194 1195 1196
TypedPackedFunc<R(Args...)>::TypedPackedFunc(const TVMRetValue& value)
    : packed_(value.operator PackedFunc()) {}

template<typename R, typename ...Args>
TypedPackedFunc<R(Args...)>::TypedPackedFunc(const TVMArgValue& value)
    : packed_(value.operator PackedFunc()) {}

template<typename R, typename ...Args>
1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209
template<typename FType>
inline void TypedPackedFunc<R(Args...)>::AssignTypedLambda(FType flambda) {
  packed_ = PackedFunc([flambda](const TVMArgs& args, TVMRetValue* rv) {
      detail::unpack_call<R, sizeof...(Args)>(flambda, args, rv);
    });
}

template<typename R, typename ...Args>
inline R TypedPackedFunc<R(Args...)>::operator()(Args... args) const {
  return detail::typed_packed_call_dispatcher<R>
      ::run(packed_, std::forward<Args>(args)...);
}

1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240
// extension and node type handling
namespace detail {
template<typename T, typename TSrc, bool is_ext>
struct TVMValueCast {
  static T Apply(const TSrc* self) {
    return self->template AsNodeRef<T>();
  }
};

template<typename T, typename TSrc>
struct TVMValueCast<T, TSrc, true> {
  static T Apply(const TSrc* self) {
    return self->template AsExtension<T>();
  }
};
}  // namespace detail

template<typename T, typename>
inline TVMArgValue::operator T() const {
  return detail::
      TVMValueCast<T, TVMArgValue, extension_class_info<T>::code != 0>
      ::Apply(this);
}

template<typename T, typename>
inline TVMRetValue::operator T() const {
  return detail::
      TVMValueCast<T, TVMRetValue, extension_class_info<T>::code != 0>
      ::Apply(this);
}

1241 1242 1243 1244 1245 1246 1247 1248
template<typename T, typename>
inline void TVMArgsSetter::operator()(size_t i, const T& value) const {
  static_assert(extension_class_info<T>::code != 0,
                "Need to have extesion code");
  type_codes_[i] = extension_class_info<T>::code;
  values_[i].v_handle = const_cast<T*>(&value);
}

1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259
// extension type handling
template<typename T>
struct ExtTypeInfo {
  static void destroy(void* handle) {
    delete static_cast<T*>(handle);
  }
  static void* clone(void* handle) {
    return new T(*static_cast<T*>(handle));
  }
};

1260 1261 1262 1263 1264 1265 1266 1267 1268
template<typename T>
inline ExtTypeVTable* ExtTypeVTable::Register_() {
  const int code = extension_class_info<T>::code;
  static_assert(code != 0,
                "require extension_class_info traits to be declared with non-zero code");
  ExtTypeVTable vt;
  vt.clone = ExtTypeInfo<T>::clone;
  vt.destroy = ExtTypeInfo<T>::destroy;
  return ExtTypeVTable::RegisterInternal(code, vt);
1269
}
1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280 1281 1282 1283

// Implement Module::GetFunction
// Put implementation in this file so we have seen the PackedFunc
inline PackedFunc Module::GetFunction(const std::string& name, bool query_imports) {
  PackedFunc pf = node_->GetFunction(name, node_);
  if (pf != nullptr) return pf;
  if (query_imports) {
    for (const Module& m : node_->imports_) {
      pf = m.node_->GetFunction(name, m.node_);
      if (pf != nullptr) return pf;
    }
  }
  return pf;
}
1284 1285
}  // namespace runtime
}  // namespace tvm
1286
#endif  // TVM_RUNTIME_PACKED_FUNC_H_