Skip to content

Instantly share code, notes, and snippets.

@ruslo ruslo/cpp_deploy.cc
Created Dec 20, 2018

Embed
What would you like to do?
/*!
* Copyright (c) 2017 by Contributors
* \brief Example code on load and run TVM module.s
* \file cpp_deploy_example.cc
*/
#include <cstdio>
#include <dlpack/dlpack.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/registry.h>
#include <tvm/runtime/packed_func.h>
void Verify(tvm::runtime::Module mod, std::string fname) {
// Get the function from the module.
tvm::runtime::PackedFunc f = mod.GetFunction(fname);
CHECK(f != nullptr);
// Allocate the DLPack data structures.
//
// Note that we use TVM runtime API to allocate the DLTensor in this example.
// TVM accept DLPack compatible DLTensors, so function can be invoked
// as long as we pass correct pointer to DLTensor array.
//
// For more information please refer to dlpack.
// One thing to notice is that DLPack contains alignment requirement for
// the data pointer and TVM takes advantage of that.
// If you plan to use your customized data container, please
// make sure the DLTensor you pass in meet the alignment requirement.
//
DLTensor* x;
DLTensor* y;
int ndim = 1;
int dtype_code = kDLFloat;
int dtype_bits = 32;
int dtype_lanes = 1;
int device_type = kDLOpenCL;
int device_id = 0;
int64_t shape[1] = {10};
TVMArrayAlloc(shape, ndim, dtype_code, dtype_bits, dtype_lanes,
device_type, device_id, &x);
TVMArrayAlloc(shape, ndim, dtype_code, dtype_bits, dtype_lanes,
device_type, device_id, &y);
for (int i = 0; i < shape[0]; ++i) {
static_cast<float*>(x->data)[i] = i;
}
// Invoke the function
// PackedFunc is a function that can be invoked via positional argument.
// The signature of the function is specified in tvm.build
f(x, y);
// Print out the output
for (int i = 0; i < shape[0]; ++i) {
CHECK_EQ(static_cast<float*>(y->data)[i], i + 1.0f);
}
LOG(INFO) << "Finish verification...";
}
int main(void) {
tvm::runtime::Module mod_dylib =
tvm::runtime::Module::LoadFromFile("/full/path/to/lib/test_addone.so");
LOG(INFO) << "Verify dynamic loading from test_addone.so";
Verify(mod_dylib, "addone");
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.