其他 C++ 打包函数的例子

其他 C++ 打包函数的例子#

加法偏函数#

#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>

using namespace tvm::runtime;

namespace tvm_ext {
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]); });
});

} // namespace tvm_ext

编译：

!cd ../cpp/bind_add && make clean && make

rm -rf outputs/*
g++ -std=c++17 -O2 -fPIC -I/media/pc/data/lxw/ai/tvm/include -I/media/pc/data/lxw/ai/tvm/3rdparty/dmlc-core/include -I/media/pc/data/lxw/ai/tvm/3rdparty/dlpack/include -Iinclude -DDMLC_USE_LOGGING_LIBRARY=\<tvm/runtime/logging.h\> -shared -o outputs/libs/libtvm_ext.so src/tvm_ext.cc -ldl -pthread -L/media/pc/data/lxw/ai/tvm/build

from tvm_book.tvm_ext.libinfo import load_lib
import tvm
_LIB_EXT, _LIB_EXT_NAME = load_lib(name="libtvm_ext.so", search_path=["../cpp/bind_add/outputs/libs"])
tvm._ffi._init_api("tvm_ext", __name__)

bind_add

<tvm.runtime.packed_func.PackedFunc at 0x7f7e8240d8d0>

def add(a, b):
    return a + b

f = bind_add(add, 7)
assert f(2) == 9

C++ 外部设备的例子#

#include <tvm/runtime/device_api.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
// #include <tvm/tir/op.h>

using namespace tvm::runtime;

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

} // namespace tvm_ext

编译：

!cd ../cpp/device_api && make clean && make

rm -rf outputs/*
g++ -std=c++17 -O2 -fPIC -I/media/pc/data/lxw/ai/tvm/include -I/media/pc/data/lxw/ai/tvm/3rdparty/dmlc-core/include -I/media/pc/data/lxw/ai/tvm/3rdparty/dlpack/include -Iinclude -DDMLC_USE_LOGGING_LIBRARY=\<tvm/runtime/logging.h\> -shared -o outputs/libs/libtvm_ext.so src/tvm_ext.cc -ldl -pthread -L/media/pc/data/lxw/ai/tvm/build

import tvm
from tvm_book.tvm_ext.libinfo import load_lib

_LIB_EXT, _LIB_EXT_NAME = load_lib(name="libtvm_ext.so", search_path=["../cpp/device_api/outputs/libs"])
tvm._ffi._init_api("tvm_ext", __name__)

import numpy as np
from tvm import te
n = 10
A = te.placeholder((n,), name="A")
B = te.compute((n,), lambda *i: A(*i) + 1.0, name="B")
s = te.create_schedule(B.op)

def check_llvm():
    f = tvm.build(s, [A, B], tvm.target.Target("ext_dev", "llvm"))
    dev = tvm.ext_dev(0)
    # launch the kernel.
    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
    b = tvm.nd.array(np.zeros(n, dtype=B.dtype), dev)
    f(a, b)
    np.testing.assert_allclose(b.numpy(), a.numpy() + 1)

check_llvm()

回调 C++ 端外部函数#

#include <tvm/runtime/registry.h>

// 暴露给运行时的外部函数
extern "C" float TVMTestAddOne(float y) { return y + 1; }

编译：

!cd ../cpp/extern_func && make clean && make

rm -rf outputs/*
g++ -std=c++17 -O2 -fPIC -I/media/pc/data/lxw/ai/tvm/include -I/media/pc/data/lxw/ai/tvm/3rdparty/dmlc-core/include -I/media/pc/data/lxw/ai/tvm/3rdparty/dlpack/include -Iinclude -DDMLC_USE_LOGGING_LIBRARY=\<tvm/runtime/logging.h\> -shared -o outputs/libs/libtvm_ext.so src/tvm_ext.cc -ldl -pthread -L/media/pc/data/lxw/ai/tvm/build

import tvm
from tvm_book.tvm_ext.libinfo import load_lib

_LIB_EXT, _LIB_EXT_NAME = load_lib(name="libtvm_ext.so", search_path=["../cpp/extern_func/outputs/libs"])
tvm._ffi._init_api("tvm_ext", __name__)

import numpy as np
from tvm import te
n = 10
A = te.placeholder((n,), name="A")
B = te.compute(
    (n,), lambda *i: tvm.tir.call_extern("float32", "TVMTestAddOne", A(*i)), name="B"
)
s = te.create_schedule(B.op)

def check_llvm():
    f = tvm.build(s, [A, B], "llvm")
    dev = tvm.cpu(0)
    # launch the kernel.
    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
    b = tvm.nd.array(np.zeros(n, dtype=B.dtype), dev)
    f(a, b)
    np.testing.assert_allclose(b.numpy(), a.numpy() + 1)

check_llvm()

提取外部 C++ 函数#

#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/registry.h>

using namespace tvm::runtime;
// 这个回调方法允许扩展，使 TVM 能够提取。
// 当想要使用仅包含头文件的最小版本的 TVM 运行时，这种方法会很有帮助。
extern "C" int TVMExtDeclare(TVMFunctionHandle pregister) {
  const PackedFunc& fregister = GetRef<PackedFunc>(static_cast<PackedFuncObj*>(pregister));
  // 等价于 const PackedFunc& fregister = *static_cast<PackedFunc*>(pregister);
  auto mul = [](TVMArgs args, TVMRetValue* rv) {
    int x = args[0];
    int y = args[1];
    *rv = x * y;
  };
  fregister("mul", PackedFunc(mul));
  return 0;
}

编译：

!cd ../cpp/mini_runtime && make clean && make

rm -rf outputs/*
g++ -std=c++17 -O2 -fPIC -I/media/pc/data/lxw/ai/tvm/include -I/media/pc/data/lxw/ai/tvm/3rdparty/dmlc-core/include -I/media/pc/data/lxw/ai/tvm/3rdparty/dlpack/include -Iinclude -DDMLC_USE_LOGGING_LIBRARY=\<tvm/runtime/logging.h\> -shared -o outputs/libs/libtvm_ext.so src/tvm_ext.cc -ldl -pthread -L/media/pc/data/lxw/ai/tvm/build

import tvm
from tvm_book.tvm_ext.libinfo import load_lib

_LIB, _LIB_NAME = load_lib(name="libtvm_ext.so", search_path=["../cpp/mini_runtime/outputs/libs"])
tvm._ffi._init_api("tvm_ext", __name__)

fdict = tvm._ffi.registry.extract_ext_funcs(_LIB.TVMExtDeclare)
assert fdict["mul"](3, 4) == 12