packed_func.h 36.6 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 35
namespace tvm {
namespace runtime {
36 37 38 39 40
// forward declarations
class TVMArgs;
class TVMArgValue;
class TVMRetValue;
class TVMArgsSetter;
41 42

/*!
43 44
 * \brief Packed function is a type-erased function.
 *  The arguments are passed by packed format.
45
 *
46 47 48
 *  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.
49 50 51
 */
class PackedFunc {
 public:
52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70
  /*!
   * \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)>;
71
  /*! \brief default constructor */
72
  PackedFunc() {}
73 74 75 76
  /*!
   * \brief constructing a packed function from a std::function.
   * \param body the internal container of packed function.
   */
77 78
  explicit PackedFunc(FType body) : body_(body) {}
  /*!
79
   * \brief Call packed function by directly passing in unpacked format.
80 81
   * \param args Arguments to be passed.
   * \tparam Args arguments to be passed.
82 83 84 85 86 87 88 89 90
   *
   * \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
91 92
   */
  template<typename... Args>
93
  inline TVMRetValue operator()(Args&& ...args) const;
94 95 96
  /*!
   * \brief Call the function in packed format.
   * \param args The arguments
97
   * \param rv The return value.
98
   */
99
  inline void CallPacked(TVMArgs args, TVMRetValue* rv) const;
100
  /*! \return the internal body function */
101
  inline FType body() const;
102 103 104 105 106 107 108 109
  /*! \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;
  }
110 111 112 113 114 115

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

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
/*!
 * \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() {}
  /*!
   * \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
   */
178 179 180 181 182 183 184 185 186 187 188
  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(*)
189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208
  /*!
   * \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>
209
  TypedPackedFunc(const FLambda& typed_lambda) {  // NOLINT(*)
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 264
    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_;
  }
265 266 267 268 269 270 271 272
  /*! \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;
  }
273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288

 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);
};

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
/*! \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);

331 332 333 334 335 336 337
/*!
 * \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);

338 339 340 341 342 343
// macro to check type code.
#define TVM_CHECK_TYPE_CODE(CODE, T)                           \
  CHECK_EQ(CODE, T) << " expected "                            \
  << TypeCode2Str(T) << " but get " << TypeCode2Str(CODE)      \

/*!
344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359
 * \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;
};

/*!
360
 * \brief Runtime function table about extension type.
361
 */
362 363
class ExtTypeVTable {
 public:
364 365 366 367
  /*! \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);
368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384
  /*!
   * \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);
385 386 387
};

/*!
388 389 390 391 392 393
 * \brief Internal base class to
 *  handle conversion to POD values.
 */
class TVMPODValue_ {
 public:
  operator double() const {
394 395 396 397 398 399
    // 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);
    }
400
    TVM_CHECK_TYPE_CODE(type_code_, kDLFloat);
401 402 403
    return value_.v_float64;
  }
  operator int64_t() const {
404
    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
405 406 407
    return value_.v_int64;
  }
  operator uint64_t() const {
408
    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
409 410 411
    return value_.v_int64;
  }
  operator int() const {
412
    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
413 414 415 416 417
    CHECK_LE(value_.v_int64,
             std::numeric_limits<int>::max());
    return static_cast<int>(value_.v_int64);
  }
  operator bool() const {
418
    TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
419 420 421 422 423 424 425 426
    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;
  }
427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442
  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));
443
  }
444 445 446 447
  operator TVMContext() const {
    TVM_CHECK_TYPE_CODE(type_code_, kTVMContext);
    return value_.v_ctx;
  }
448 449
  template<typename TExtension>
  const TExtension& AsExtension() const {
450 451
    CHECK_LT(type_code_, kExtEnd);
    return static_cast<TExtension*>(value_.v_handle)[0];
452
  }
453 454 455 456 457 458 459 460 461 462 463 464
  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);
  }
465 466 467 468 469 470 471 472

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

473 474 475 476 477 478 479 480 481 482 483 484 485 486
  /*! \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:
487 488
  /*! \brief default constructor */
  TVMArgValue() {}
489 490 491 492 493 494 495 496 497 498 499 500 501 502 503
  /*!
   * \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*;
504 505
  using TVMPODValue_::operator DLTensor*;
  using TVMPODValue_::operator NDArray;
506
  using TVMPODValue_::operator TVMContext;
507

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

/*!
 * \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_) {
584 585
    other.value_.v_handle = nullptr;
    other.type_code_ = kNull;
586 587 588 589 590 591 592 593 594 595 596 597
  }
  /*! \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*;
598
  using TVMPODValue_::operator DLTensor*;
599
  using TVMPODValue_::operator TVMContext;
600
  using TVMPODValue_::operator NDArray;
601 602 603 604 605 606
  // Disable copy and assign from another value, but allow move.
  TVMRetValue(const TVMRetValue& other) {
    this->Assign(other);
  }
  // conversion operators
  operator std::string() const {
607 608
    if (type_code_ == kTVMType) {
      return TVMType2String(operator TVMType());
609 610
    } else if (type_code_ == kBytes) {
      return *ptr<std::string>();
611
    }
612 613 614 615 616 617 618 619 620 621 622
    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 {
623
    if (type_code_ == kNull) return PackedFunc();
624 625 626
    TVM_CHECK_TYPE_CODE(type_code_, kFuncHandle);
    return *ptr<PackedFunc>();
  }
627 628 629 630
  template<typename FType>
  operator TypedPackedFunc<FType>() const {
    return TypedPackedFunc<FType>(operator PackedFunc());
  }
631 632 633 634
  operator Module() const {
    TVM_CHECK_TYPE_CODE(type_code_, kModuleHandle);
    return *ptr<Module>();
  }
635 636 637 638 639 640 641 642 643
  // 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) {
644
    this->SwitchToPOD(kDLFloat);
645 646 647 648 649 650 651 652 653 654 655 656 657 658
    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) {
659
    this->SwitchToPOD(kDLInt);
660 661 662 663
    value_.v_int64 = value;
    return *this;
  }
  TVMRetValue& operator=(int value) {
664
    this->SwitchToPOD(kDLInt);
665 666 667
    value_.v_int64 = value;
    return *this;
  }
668 669 670 671 672
  TVMRetValue& operator=(TVMContext value) {
    this->SwitchToPOD(kTVMContext);
    value_.v_ctx = value;
    return *this;
  }
673 674 675 676 677 678
  TVMRetValue& operator=(TVMType t) {
    this->SwitchToPOD(kTVMType);
    value_.v_type = t;
    return *this;
  }
  TVMRetValue& operator=(bool value) {
679
    this->SwitchToPOD(kDLInt);
680 681 682 683 684 685 686
    value_.v_int64 = value;
    return *this;
  }
  TVMRetValue& operator=(std::string value) {
    this->SwitchToClass(kStr, value);
    return *this;
  }
687 688 689 690
  TVMRetValue& operator=(TVMByteArray value) {
    this->SwitchToClass(kBytes, std::string(value.data, value.size));
    return *this;
  }
691 692 693 694 695 696 697
  TVMRetValue& operator=(NDArray other) {
    this->Clear();
    type_code_ = kNDArrayContainer;
    value_.v_handle = other.data_;
    other.data_ = nullptr;
    return *this;
  }
698 699 700 701
  TVMRetValue& operator=(PackedFunc f) {
    this->SwitchToClass(kFuncHandle, f);
    return *this;
  }
702 703 704 705
  template<typename FType>
  TVMRetValue& operator=(const TypedPackedFunc<FType>& f) {
    return operator=(f.packed());
  }
706 707 708 709
  TVMRetValue& operator=(Module m) {
    this->SwitchToClass(kModuleHandle, m);
    return *this;
  }
710 711 712 713
  TVMRetValue& operator=(const TVMRetValue& other) {  // NOLINT(*0
    this->Assign(other);
    return *this;
  }
714
  TVMRetValue& operator=(const TVMArgValue& other) {
715 716 717
    this->Assign(other);
    return *this;
  }
718 719 720 721 722 723 724 725
  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;
  }
726 727 728 729 730 731 732 733 734 735 736 737
  /*!
   * \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.
738
    CHECK(type_code_ != kStr && type_code_ != kBytes);
739 740 741 742
    *ret_value = value_;
    *ret_type_code = type_code_;
    type_code_ = kNull;
  }
743 744 745 746
  /*! \return The value field, if the data is POD */
  const TVMValue& value() const {
    CHECK(type_code_ != kNodeHandle &&
          type_code_ != kFuncHandle &&
747
          type_code_ != kModuleHandle &&
748 749 750
          type_code_ != kStr) << "TVMRetValue.value can only be used for POD data";
    return value_;
  }
751
  // NodeRef related extenstions: in tvm/packed_func_ext.h
752 753 754 755 756 757
  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;
758
  inline TVMRetValue& operator=(const NodeRef& other);
759
  inline TVMRetValue& operator=(const NodePtr<Node>& other);
760
  // type related
761 762
  inline operator HalideIR::Type() const;
  inline TVMRetValue& operator=(const HalideIR::Type& other);
763 764 765 766 767

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

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

nhynes committed
874
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
875
inline std::ostream& operator<<(std::ostream& os, TVMType t) {  // NOLINT(*)
876 877 878
  os << TypeCode2Str(t.code);
  if (t.code == kHandle) return os;
  os << static_cast<int>(t.bits);
879 880 881 882 883
  if (t.lanes != 1) {
    os << 'x' << static_cast<int>(t.lanes);
  }
  return os;
}
nhynes committed
884
#endif
885 886

inline std::string TVMType2String(TVMType t) {
887
  if (t.bits == 0) return "";
nhynes committed
888
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
889 890 891
  std::ostringstream os;
  os << t;
  return os.str();
nhynes committed
892 893 894 895 896 897 898 899 900 901
#else
  std::string repr = "";
  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
902 903
}

904 905
inline TVMType String2TVMType(std::string s) {
  TVMType t;
906 907 908 909 910
  // handle None type
  if (s.length() == 0) {
    t.bits = 0; t.lanes = 0; t.code = kHandle;
    return t;
  }
911 912 913
  t.bits = 32; t.lanes = 1;
  const char* scan;
  if (s.substr(0, 3) == "int") {
914
    t.code = kDLInt;  scan = s.c_str() + 3;
915
  } else if (s.substr(0, 4) == "uint") {
916
    t.code = kDLUInt; scan = s.c_str() + 4;
917
  } else if (s.substr(0, 5) == "float") {
918
    t.code = kDLFloat; scan = s.c_str() + 5;
919
  } else if (s.substr(0, 6) == "handle") {
920 921 922 923 924 925 926
    t.code = kHandle;
    t.bits = 64;  // handle uses 64 bit by default.
    scan = s.c_str() + 6;
  } else {
    scan = s.c_str();
    LOG(FATAL) << "unknown type " << s;
  }
nhynes committed
927
  char* xdelim;  // emulate sscanf("%ux%u", bits, lanes)
928 929
  uint8_t bits = static_cast<uint8_t>(strtoul(scan, &xdelim, 10));
  if (bits != 0) t.bits = bits;
nhynes committed
930
  if (*xdelim == 'x') {
931
    t.lanes = static_cast<uint16_t>(strtoul(xdelim + 1, nullptr, 10));
nhynes committed
932
  }
933 934 935 936 937 938 939
  return t;
}

inline TVMArgValue TVMArgs::operator[](int i) const {
  CHECK_LT(i, num_args)
      << "not enough argument passed, "
      << num_args << " passed"
940
      << " but request arg[" << i << "].";
941 942 943 944 945 946 947 948 949
  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);
950 951
}

952 953 954 955
inline PackedFunc::FType PackedFunc::body() const {
  return body_;
}

956 957


958 959
// internal namespace
namespace detail {
960 961

template<bool stop, std::size_t I, typename F>
962
struct for_each_dispatcher {
963 964 965 966 967
  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)...);
968 969 970
  }
};

971 972 973
template<std::size_t I, typename F>
struct for_each_dispatcher<true, I, F>  {
  static void run(const F& f) {}  // NOLINT(*)
974 975 976
};

template<typename F, typename ...Args>
977 978 979
inline void for_each(const F& f, Args&&... args) {  // NOLINT(*)
  for_each_dispatcher<sizeof...(Args) == 0, 0, F>
      ::run(f, std::forward<Args>(args)...);
980
}
981
}  // namespace detail
982

983 984 985
/* \brief argument settter to PackedFunc */
class TVMArgsSetter {
 public:
986 987
  TVMArgsSetter(TVMValue* values, int* type_codes)
      : values_(values), type_codes_(type_codes) {}
988 989
  // setters for POD types
  template<typename T,
990 991
           typename = typename std::enable_if<
             std::is_integral<T>::value>::type>
992 993
  void operator()(size_t i, T value) const {
    values_[i].v_int64 = static_cast<int64_t>(value);
994
    type_codes_[i] = kDLInt;
995 996 997 998 999
  }
  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()));
1000
    type_codes_[i] = kDLInt;
1001 1002 1003
  }
  void operator()(size_t i, double value) const {
    values_[i].v_float64 = value;
1004
    type_codes_[i] = kDLFloat;
1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017
  }
  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;
  }
1018
  void operator()(size_t i, DLTensor* value) const {
1019 1020 1021
    values_[i].v_handle = value;
    type_codes_[i] = kArrayHandle;
  }
1022 1023 1024 1025
  void operator()(size_t i, TVMContext value) const {
    values_[i].v_ctx = value;
    type_codes_[i] = kTVMContext;
  }
1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036
  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)
1037
  void operator()(size_t i, const std::string& value) const {  // NOLINT(*)
1038 1039 1040
    values_[i].v_str = value.c_str();
    type_codes_[i] = kStr;
  }
1041 1042
  void operator()(size_t i, const TVMByteArray& value) const {  // NOLINT(*)
    values_[i].v_handle = const_cast<TVMByteArray*>(&value);
1043 1044
    type_codes_[i] = kBytes;
  }
1045 1046
  void operator()(size_t i, const PackedFunc& value) const {  // NOLINT(*)
    values_[i].v_handle = const_cast<PackedFunc*>(&value);
1047 1048
    type_codes_[i] = kFuncHandle;
  }
1049 1050 1051 1052
  template<typename FType>
  void operator()(size_t i, const TypedPackedFunc<FType>& value) const {  // NOLINT(*)
    operator()(i, value.packed());
  }
1053 1054
  void operator()(size_t i, const Module& value) const {  // NOLINT(*)
    values_[i].v_handle = const_cast<Module*>(&value);
1055 1056
    type_codes_[i] = kModuleHandle;
  }
1057 1058 1059 1060
  void operator()(size_t i, const NDArray& value) const {  // NOLINT(*)
    values_[i].v_handle = value.data_;
    type_codes_[i] = kNDArrayContainer;
  }
1061
  void operator()(size_t i, const TVMRetValue& value) const {  // NOLINT(*)
1062 1063 1064 1065
    if (value.type_code() == kStr) {
      values_[i].v_str = value.ptr<std::string>()->c_str();
      type_codes_[i] = kStr;
    } else {
1066
      CHECK_NE(value.type_code(), kBytes) << "not handled.";
1067 1068 1069 1070
      values_[i] = value.value_;
      type_codes_[i] = value.type_code();
    }
  }
1071 1072 1073 1074 1075
  // 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;
1076
  // NodeRef related extenstions: in tvm/packed_func_ext.h
1077
  inline void operator()(size_t i, const NodeRef& other) const;  // NOLINT(*)
1078
  inline void operator()(size_t i, const HalideIR::Type& t) const;
1079 1080 1081 1082 1083 1084 1085 1086

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

1087
template<typename... Args>
1088
inline TVMRetValue PackedFunc::operator()(Args&& ...args) const {
1089
  const int kNumArgs = sizeof...(Args);
1090 1091 1092
  const int kArraySize = kNumArgs > 0 ? kNumArgs : 1;
  TVMValue values[kArraySize];
  int type_codes[kArraySize];
1093
  detail::for_each(TVMArgsSetter(values, type_codes),
1094
                   std::forward<Args>(args)...);
1095 1096 1097
  TVMRetValue rv;
  body_(TVMArgs(values, type_codes, kNumArgs), &rv);
  return rv;
1098
}
1099

1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 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
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>
1165 1166 1167 1168
TypedPackedFunc<R(Args...)>::TypedPackedFunc(PackedFunc packed)
  : packed_(packed) {}

template<typename R, typename ...Args>
1169 1170 1171 1172 1173 1174 1175 1176
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>
1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189
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)...);
}

1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220
// 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);
}

1221 1222 1223 1224 1225 1226 1227 1228
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);
}

1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239
// 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));
  }
};

1240 1241 1242 1243 1244 1245 1246 1247 1248
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);
1249
}
1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263

// 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;
}
1264 1265
}  // namespace runtime
}  // namespace tvm
1266
#endif  // TVM_RUNTIME_PACKED_FUNC_H_