tvm_ext.cc 1.54 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12

/*!
 *  Copyright (c) 2017 by Contributors
 * \brief Example package that uses TVM.
 * \file tvm_ext.cc
 */
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/registry.h>
#include <tvm/packed_func_ext.h>

namespace tvm_ext {
13 14 15 16 17 18 19 20 21 22 23 24 25 26 27
using IntVector = std::vector<int>;
}  // namespace tvm_ext

namespace tvm {
namespace runtime {
template<>
struct extension_class_info<tvm_ext::IntVector> {
  static const int code = 17;
};
}  // namespace tvm
}  // namespace runtime


namespace tvm_ext {

28 29 30
using namespace tvm;
using namespace tvm::runtime;

31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47
TVM_REGISTER_EXT_TYPE(IntVector);

TVM_REGISTER_GLOBAL("tvm_ext.ivec_create")
.set_body([](TVMArgs args, TVMRetValue *rv) {
    IntVector vec;
    for (int i = 0; i < args.size(); ++i) {
      vec.push_back(args[i].operator int());
    }
    *rv = vec;
  });

TVM_REGISTER_GLOBAL("tvm_ext.ivec_get")
.set_body([](TVMArgs args, TVMRetValue *rv) {
    *rv = args[0].AsExtension<IntVector>()[args[1].operator int()];
  });


48 49 50 51 52 53 54 55 56 57 58 59 60 61 62
TVM_REGISTER_GLOBAL("tvm_ext.bind_add")
.set_body([](TVMArgs args_, TVMRetValue *rv_) {
    PackedFunc pf = args_[0];
    int b = args_[1];
    *rv_ = PackedFunc([pf, b](TVMArgs args, TVMRetValue *rv) {
        *rv = pf(b, args[0]);
      });
  });

TVM_REGISTER_GLOBAL("tvm_ext.sym_add")
.set_body([](TVMArgs args, TVMRetValue *rv) {
    Var a = args[0];
    Var b = args[1];
    *rv = a + b;
  });
63 64 65 66 67

TVM_REGISTER_GLOBAL("device_api.ext_dev")
.set_body([](TVMArgs args, TVMRetValue *rv) {
    *rv = (*tvm::runtime::Registry::Get("device_api.cpu"))();
  });
68
}  // namespace tvm_ext