Review of TVM's codebase

Table of Contents

TE

TIR/Transform

Runtime

PackedFunc

PackedFunc 是 TVM 中的一种数据结构,为 C++和其它语言交互及部署提供了便利。

如下可以在 C++中定义一个 ~PackedFunc~,并在 Python 中执行:

#include <tvm/runtime/packed_func.h>

void MyAdd(TVMArgs args, TVMRetValue* rv) {
  // automatically convert arguments to desired type.
  int a = args[0];
  int b = args[1];
  // automatically assign value return to rv
  *rv = a + b;
}

void CallPacked() {
  PackedFunc myadd = PackedFunc(MyAdd);
  // get back 3
  int c = myadd(1, 2);
}

// register a global packed function in c++
TVM_REGISTER_GLOBAL("myadd")
.set_body(MyAdd);
import tvm

myadd = tvm.get_global_func("myadd")
# prints 3
print(myadd(1, 2))

同理,也可以在 Python 中定义函数,并在 C++中使用 PackedFunc 调用:

import tvm

def callback(msg):
  print(msg)

# convert to PackedFunc
f = tvm.convert(callback)
callhello = tvm.get_global_func("callhello")
# prints hello world
callhello(f)
TVM_REGISTER_GLOBAL("callhello")
.set_body([](TVMArgs args, TVMRetValue* rv) {
  PackedFunc f = args[0];
  f("hello world");
 });

TODO 代码结构

~TVMValue~是几种常见的数据类型的 union:

typedef union {
  int64_t v_int64;
  double v_float64;
  void* v_handle;
  const char* v_str;
  DLDataType v_type;
  TVMContext v_ctx;
} TVMValue;

TVMArgs1 包装了可能的参数,其中每一个值都是 TVMValue 类型,同时储存了 type_codes 保存了每个参数的类型。

TVMPODValue_2 包装了 TVMValue 到各种数据类型的显示转换。

TVMArgValue3

PackedFunc4包装一个 std::function 类。

TypedPackedFunc5

Registry6用于注册 global function,

Module

Footnotes:

Author: Zihao Ye

Email: expye@outlook.com

Date:

Last modified: 2021-02-19 Fri 21:14

Licensed under CC BY-NC 4.0