tvm_ffi_register_and_call_global_packed_functions.py

1import tvm_ffi
2import numpy as np
3
4# 1. Define a function that follows the TVM PackedCFunc signature
5# The function takes a list of arguments and returns a value
6def my_add(args):
7    a, b = args
8    return a + b
9
10# 2. Register the function to the TVM global registry
11tvm_ffi.register_func("my_add_func", my_add)
12
13# 3. Retrieve the function back from the registry
14f = tvm_ffi.get_global_func("my_add_func")
15
16# 4. Call the function and print the result
17# The FFI handles the conversion of Python types to TVM's internal representation
18result = f(10, 20)
19print(f"Result of 10 + 20: {result}")
20
21# 5. Example with more complex types (like DLPack/Arrays)
22def array_sum(args):
23    arr = args[0]
24    # In a real scenario, you would use tvm_ffi's array handling
25    return np.sum(arr)
26
27tvm_ffi.register_func("array_sum", array_sum)
28f_sum = tvm_ffi.get_global_func("array_sum")
29data = np.array([1, 2, 3, 4, 5], dtype="float32")
30print(f"Sum of array: {f_sum(data)}")