Skip to content

T
tic
Commit 7f7dc073 by Zhi Committed by Yizhi Liu
Options

use packed func macro for external codegen (#4710)

parent ce807fe8
Showing with 81 additions and 104 deletions
python/tvm/_ffi/libinfo.py
......@@ -179,13 +179,20 @@ def find_include_path(name=None, search_path=None, optional=False):
else:
tvm_include_path = [os.path.join(p, name) for p in header_path]
dlpack_include_path = []
dmlc_include_path = []
else:
tvm_include_path = [os.path.join(p, 'include') for p in header_path]
dlpack_include_path = [os.path.join(p, 'dlpack/include') for p in header_path]
dlpack_include_path = [os.path.join(p, 'dlpack/include') for p in
header_path]
dmlc_include_path = [os.path.join(p, 'dmlc-core/include') for p in
header_path]
# try to find include path
include_found = [p for p in tvm_include_path if os.path.exists(p) and os.path.isdir(p)]
include_found += [p for p in dlpack_include_path if os.path.exists(p) and os.path.isdir(p)]
include_found += [p for p in dlpack_include_path if os.path.exists(p)
and os.path.isdir(p)]
include_found += [p for p in dmlc_include_path if os.path.exists(p)
and os.path.isdir(p)]
if not include_found:
message = ('Cannot find the files.\n' +
......
src/relay/backend/contrib/codegen_c/codegen.cc
......@@ -154,12 +154,9 @@ class CSourceCodegen : public CSourceModuleCodegenBase {
runtime::Module CreateCSourceModule(const ObjectRef& ref) override {
// Create headers
code_stream_ << "#include <cstdint>\n";
code_stream_ << "#include <iostream>\n";
code_stream_ << "#include <cstdlib>\n";
code_stream_ << "#include <stdio.h>\n";
code_stream_ << "#include <cstring>\n";
code_stream_ << "#include <tvm/runtime/c_runtime_api.h>\n";
code_stream_ << "#include <tvm/runtime/packed_func.h>\n";
code_stream_ << "#include <dlpack/dlpack.h>\n";
// Append some common macro for operator definition.
......
src/relay/backend/contrib/codegen_c/codegen_c.h
......@@ -99,63 +99,52 @@ class CodegenCBase {
* \code
*
* // An example code for the generated C function.
* extern "C" void foo(TVMValue* value, int* type_code, int nargs) {
* if (nargs != 3) {
* printf("foo expects 3 args, but received %d\n", nargs);
* return 1;
* }
*
* DLTensor* arg0 = static_cast<DLTensor*>(value[0].v_handle);
* DLTensor* arg1 = static_cast<DLTensor*>(value[1].v_handle);
* DLTensor* out = static_cast<DLTensor*>(value[2].v_handle);
*
* extern "C" void foo_wrapper_(DLTensor* arg0,
* DLTensor* arg1,
* DLTensor* out) {
* foo_(static_cast<float*>(arg0->data),
* static_cast<float*>(arg1->data),
* static_cast<float*>(out->data));
* return 0;
* }
*
* TVM_DLL_EXPORT_TYPED_FUNC(foo, foo_wrapper_);
*
* \endcode
*/
void GenerateBackendCFunc(const std::string& func_name, int arg_cnt) {
// Print signature
code_stream_ << "\n";
code_stream_ << "extern \"C\" int " << func_name;
code_stream_ << "(TVMValue* value, int* type_code, int nargs) {\n";
EnterScope();
// Print guard
PrintIndents();
code_stream_ << "if (nargs != " << arg_cnt << "){\n";
code_stream_ << "extern \"C\" int " << func_name << "_wrapper_(";
for (int i = 0; i < arg_cnt - 1; i++) {
code_stream_ << "DLTensor* arg" << i << ",\n";
code_stream_ << "\t";
}
if (arg_cnt > 0) {
code_stream_ << "DLTensor* arg" << arg_cnt - 1 << ") {\n";
}
EnterScope();
PrintIndents();
code_stream_ << "printf(\"" << func_name << " expects " << arg_cnt
<< " arguments, but received %d\\n\", nargs);\n";
PrintIndents();
code_stream_ << "return 1;\n";
ExitScope();
PrintIndents();
code_stream_ << "}\n";
// According to TVM's calling convention, the last one is output.
for (int i = 0; i < arg_cnt; i++) {
PrintIndents();
code_stream_ << "DLTensor* arg" << i << " = "
<< "static_cast<DLTensor*>(value[" << i << "].v_handle);\n";
}
// Generate the call.
// Generate the internal call.
PrintIndents();
code_stream_ << func_name << "_(";
for (int i = 0; i < arg_cnt - 1; i++) {
code_stream_ << "static_cast<float*>(arg" << i << "->data), ";
code_stream_ << "static_cast<float*>(arg" << i << "->data),\n";
PrintIndents();
}
if (arg_cnt > 0) {
code_stream_ << "static_cast<float*>(arg" << arg_cnt - 1 << "->data)";
}
code_stream_ << ");\n\n";
code_stream_ << ");\n";
PrintIndents();
code_stream_ << "return 0;\n";
ExitScope();
code_stream_ << "}";
code_stream_ << "}\n\n";
// Generate the macro
code_stream_ << "TVM_DLL_EXPORT_TYPED_FUNC(" << func_name << ", "
<< func_name << "_wrapper_);\n\n";
}
/*!
......
src/relay/backend/contrib/dnnl/codegen.cc
......@@ -260,6 +260,7 @@ class DNNLModuleCodegen : public CSourceModuleCodegenBase {
code_stream_ << "#include <cstdlib>\n";
code_stream_ << "#include <cstring>\n";
code_stream_ << "#include <tvm/runtime/c_runtime_api.h>\n";
code_stream_ << "#include <tvm/runtime/packed_func.h>\n";
code_stream_ << "#include <dlpack/dlpack.h>\n";
// dnnl_kernel file is saved under src/runtime/contrib/dnnl so that we don't
// expose it to ordinary users. To make export_library use it, users need to
......
tests/python/relay/test_external_runtime.py
......@@ -33,6 +33,7 @@ def generate_csource_module():
code = r'''
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/packed_func.h>
#include <dlpack/dlpack.h>
#include <cstdint>
#include <cstring>
......@@ -69,22 +70,17 @@ def generate_csource_module():
free(buf_1);
}
extern "C" int json_rt_1(TVMValue* value, int* type_code, int nargs) {
if (nargs != 5) {
printf("Expect 5 args, but get %d", nargs);
return 1;
}
DLTensor* arg0 = static_cast<DLTensor*>(value[0].v_handle);
DLTensor* arg1 = static_cast<DLTensor*>(value[1].v_handle);
DLTensor* arg2 = static_cast<DLTensor*>(value[2].v_handle);
DLTensor* arg3 = static_cast<DLTensor*>(value[3].v_handle);
DLTensor* out = static_cast<DLTensor*>(value[4].v_handle);
gcc_1_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),
static_cast<float*>(arg2->data), static_cast<float*>(arg3->data),
static_cast<float*>(out->data));
return 0;
extern "C" int ccompiler_wrapper_1_(DLTensor* arg0, DLTensor* arg1,
DLTensor* arg2, DLTensor* arg3,
DLTensor* out) {
gcc_1_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),
static_cast<float*>(arg2->data), static_cast<float*>(arg3->data),
static_cast<float*>(out->data));
return 0;
}
TVM_DLL_EXPORT_TYPED_FUNC(json_rt_1, ccompiler_wrapper_1_);
GCC_BINARY_OP_2D(gcc_0_0, *, 10, 10);
GCC_BINARY_OP_2D(gcc_0_1, -, 10, 10);
GCC_BINARY_OP_2D(gcc_0_2, +, 10, 10);
......@@ -100,21 +96,17 @@ def generate_csource_module():
free(buf_1);
}
extern "C" int json_rt_0(TVMValue* value, int* type_code, int nargs) {
if (nargs != 5) {
printf("Expect 5 args, but get %d", nargs);
return 1;
}
DLTensor* arg0 = static_cast<DLTensor*>(value[0].v_handle);
DLTensor* arg1 = static_cast<DLTensor*>(value[1].v_handle);
DLTensor* arg2 = static_cast<DLTensor*>(value[2].v_handle);
DLTensor* arg3 = static_cast<DLTensor*>(value[3].v_handle);
DLTensor* out = static_cast<DLTensor*>(value[4].v_handle);
gcc_0_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),
static_cast<float*>(arg2->data), static_cast<float*>(arg3->data),
static_cast<float*>(out->data));
return 0;
extern "C" int ccompiler_wrapper_0_(DLTensor* arg0, DLTensor* arg1,
DLTensor* arg2, DLTensor* arg3,
DLTensor* out) {
gcc_0_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),
static_cast<float*>(arg2->data), static_cast<float*>(arg3->data),
static_cast<float*>(out->data));
return 0;
}
TVM_DLL_EXPORT_TYPED_FUNC(json_rt_0, ccompiler_wrapper_0_);
'''
csource_module = _tvm_module.csource_module_create(code, "cc")
return csource_module
......@@ -128,11 +120,12 @@ def generate_engine_module():
code = r'''
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/packed_func.h>
#include <dlpack/dlpack.h>
#include "gcc_engine.h"
#include "json_engine.h"
extern "C" void gcc_1_(float* gcc_input4, float* gcc_input5,
float* gcc_input6, float* gcc_input7, float* out) {
extern "C" void json_1_(float* json_input4, float* json_input5,
float* json_input6, float* json_input7, float* out) {
std::string graph =
"add_2d,10,10\n"
......@@ -140,28 +133,22 @@ def generate_engine_module():
"mul_2d,10,10\n";
Engine engine;
engine.run(graph, {gcc_input4, gcc_input5, gcc_input6, gcc_input7}, out);
engine.run(graph, {json_input4, json_input5, json_input6, json_input7}, out);
}
extern "C" int json_rt_1(TVMValue* value, int* type_code, int nargs) {
if (nargs != 5) {
printf("Expect 5 args, but get %d", nargs);
return 1;
}
DLTensor* arg0 = static_cast<DLTensor*>(value[0].v_handle);
DLTensor* arg1 = static_cast<DLTensor*>(value[1].v_handle);
DLTensor* arg2 = static_cast<DLTensor*>(value[2].v_handle);
DLTensor* arg3 = static_cast<DLTensor*>(value[3].v_handle);
DLTensor* out = static_cast<DLTensor*>(value[4].v_handle);
gcc_1_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),
extern "C" int json_wrapper_1_(DLTensor* arg0, DLTensor* arg1,
DLTensor* arg2, DLTensor* arg3,
DLTensor* out) {
json_1_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),
static_cast<float*>(arg2->data), static_cast<float*>(arg3->data),
static_cast<float*>(out->data));
return 0;
}
extern "C" void gcc_0_(float* gcc_input0, float* gcc_input1,
float* gcc_input2, float* gcc_input3, float* out) {
TVM_DLL_EXPORT_TYPED_FUNC(json_rt_1, json_wrapper_1_);
extern "C" void json_0_(float* json_input0, float* json_input1,
float* json_input2, float* json_input3, float* out) {
std::string graph =
"add_2d,10,10\n"
......@@ -169,40 +156,36 @@ def generate_engine_module():
"mul_2d,10,10\n";
Engine engine;
engine.run(graph, {gcc_input0, gcc_input1, gcc_input2, gcc_input3}, out);
engine.run(graph, {json_input0, json_input1, json_input2, json_input3}, out);
}
extern "C" int json_rt_0(TVMValue* value, int* type_code, int nargs) {
if (nargs != 5) {
printf("Expect 5 args, but get %d", nargs);
return 1;
}
DLTensor* arg0 = static_cast<DLTensor*>(value[0].v_handle);
DLTensor* arg1 = static_cast<DLTensor*>(value[1].v_handle);
DLTensor* arg2 = static_cast<DLTensor*>(value[2].v_handle);
DLTensor* arg3 = static_cast<DLTensor*>(value[3].v_handle);
DLTensor* out = static_cast<DLTensor*>(value[4].v_handle);
gcc_0_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),
extern "C" int json_wrapper_0_(DLTensor* arg0, DLTensor* arg1,
DLTensor* arg2, DLTensor* arg3,
DLTensor* out) {
json_0_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),
static_cast<float*>(arg2->data), static_cast<float*>(arg3->data),
static_cast<float*>(out->data));
return 0;
}
TVM_DLL_EXPORT_TYPED_FUNC(json_rt_0, json_wrapper_0_);
'''
gen_gcc_engine()
gen_json_engine()
csource_module = _tvm_module.csource_module_create(code, "cc")
return csource_module
def gen_gcc_engine():
def gen_json_engine():
"""An example of external backend runtime engine. This is supposed to be provided
by third-party vendors and included when building the generated external kernel code.
"""
code = r'''
#ifndef _GCC_ENGINE_H_
#define _GCC_ENGINE_H_
#ifndef _JSON_ENGINE_H_
#define _JSON_ENGINE_H_
#include <cstdint>
#include <string>
#include <sstream>
......@@ -298,9 +281,9 @@ def gen_gcc_engine():
std::vector<float*> buffers;
};
#endif
#endif // _JSON_ENGINE_H_
'''
header_file = tmp_path.relpath("gcc_engine.h")
header_file = tmp_path.relpath("json_engine.h")
with open(header_file, 'w') as f:
f.write(code)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment