Skip to content

Commit 9b5930d

Browse files
authored
[FFI][DOCS] Add missing files in packaging example (#18261)
This PR adds the missing files in packaging example also renames get_started to quick_start
1 parent 46eac56 commit 9b5930d

File tree

9 files changed

+128
-0
lines changed

9 files changed

+128
-0
lines changed
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations.
16+
# Base logic to load library for extension package
17+
import torch
18+
import sys
19+
import tvm_ffi_extension
20+
21+
22+
def run_add_one():
23+
x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32)
24+
y = torch.empty_like(x)
25+
tvm_ffi_extension.add_one(x, y)
26+
print(y)
27+
28+
29+
def run_raise_error():
30+
tvm_ffi_extension.raise_error("This is an error")
31+
32+
33+
if __name__ == "__main__":
34+
if len(sys.argv) > 1:
35+
if sys.argv[1] == "add_one":
36+
run_add_one()
37+
elif sys.argv[1] == "raise_error":
38+
run_raise_error()
39+
else:
40+
print("Usage: python run_example.py <add_one|raise_error>")
Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
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
File renamed without changes.

0 commit comments

Comments
 (0)