|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 3 | + * or more contributor license agreements. See the NOTICE file |
| 4 | + * distributed with this work for additional information |
| 5 | + * regarding copyright ownership. The ASF licenses this file |
| 6 | + * to you under the Apache License, Version 2.0 (the |
| 7 | + * "License"); you may not use this file except in compliance |
| 8 | + * with the License. You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, |
| 13 | + * software distributed under the License is distributed on an |
| 14 | + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 15 | + * KIND, either express or implied. See the License for the |
| 16 | + * specific language governing permissions and limitations |
| 17 | + * under the License. |
| 18 | + */ |
| 19 | +/*! |
| 20 | + * \file example.cc |
| 21 | + * \brief Example of a tvm-ffi based library that registers various functions. |
| 22 | + * |
| 23 | + * It is a simple example that demonstrates how to package a tvm-ffi library into a python wheel. |
| 24 | + * The library is written in C++ and can be compiled into a shared library. |
| 25 | + * The shared library can then be loaded into python and used to call the functions. |
| 26 | + */ |
| 27 | +#include <tvm/ffi/dtype.h> |
| 28 | +#include <tvm/ffi/error.h> |
| 29 | +#include <tvm/ffi/function.h> |
| 30 | +#include <tvm/ffi/reflection/registry.h> |
| 31 | + |
| 32 | +namespace tvm_ffi_extension { |
| 33 | + |
| 34 | +namespace ffi = tvm::ffi; |
| 35 | + |
| 36 | +/*! |
| 37 | + * \brief Raises a runtime error |
| 38 | + * |
| 39 | + * This is an example function to show how to raise and propagate |
| 40 | + * an error across the language boundary. |
| 41 | + * |
| 42 | + * \param msg The message to raise the error with |
| 43 | + */ |
| 44 | +void RaiseError(ffi::String msg) { TVM_FFI_THROW(RuntimeError) << msg; } |
| 45 | + |
| 46 | +void AddOne(DLTensor* x, DLTensor* y) { |
| 47 | + // implementation of a library function |
| 48 | + TVM_FFI_ICHECK(x->ndim == 1) << "x must be a 1D tensor"; |
| 49 | + DLDataType f32_dtype{kDLFloat, 32, 1}; |
| 50 | + TVM_FFI_ICHECK(x->dtype == f32_dtype) << "x must be a float tensor"; |
| 51 | + TVM_FFI_ICHECK(y->ndim == 1) << "y must be a 1D tensor"; |
| 52 | + TVM_FFI_ICHECK(y->dtype == f32_dtype) << "y must be a float tensor"; |
| 53 | + TVM_FFI_ICHECK(x->shape[0] == y->shape[0]) << "x and y must have the same shape"; |
| 54 | + for (int i = 0; i < x->shape[0]; ++i) { |
| 55 | + static_cast<float*>(y->data)[i] = static_cast<float*>(x->data)[i] + 1; |
| 56 | + } |
| 57 | +} |
| 58 | + |
| 59 | +// expose global symbol add_one |
| 60 | +TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, tvm_ffi_extension::AddOne); |
| 61 | + |
| 62 | +// The static initialization block is |
| 63 | +// called once when the library is loaded. |
| 64 | +TVM_FFI_STATIC_INIT_BLOCK({ |
| 65 | + namespace refl = tvm::ffi::reflection; |
| 66 | + // In this particular example, we use the reflection mechanisms to |
| 67 | + // register the functions directly into the global function table. |
| 68 | + // |
| 69 | + // This is an alternative approach to TVM_FFI_DLL_EXPORT_TYPED_FUNC |
| 70 | + // that exports the function directly as C symbol that follows tvm-ffi abi. |
| 71 | + // |
| 72 | + // - For functions that are expected to be static part of tvm_ffi_example project, |
| 73 | + // one can use reflection mechanisms to register the globa function. |
| 74 | + // - For functions that are compiled and dynamically loaded at runtime, consider |
| 75 | + // using the normal export mechanism so they won't be exposed to the global function table. |
| 76 | + // |
| 77 | + // Make sure to have a unique name across all registered functions, |
| 78 | + // always prefix with a package namespace name to avoid name collision. |
| 79 | + // |
| 80 | + // The function can then be found via tvm_ffi.get_global_func(name) |
| 81 | + // If the function is expected to stay throughout the lifetime of the program/ |
| 82 | + // |
| 83 | + // When registering via reflection mechanisms, the library do not need to be loaded via |
| 84 | + // tvm::ffi::Module::LoadFromFile, instead, just load the dll or simply bundle into the |
| 85 | + // final project |
| 86 | + refl::GlobalDef().def("tvm_ffi_extension.raise_error", RaiseError); |
| 87 | +}); |
| 88 | +} // namespace tvm_ffi_extension |
0 commit comments