Provide external functions as TVMScript?

Questions
1

I have a simple TVMScript program:

import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T

@T.prim_func
def main(A: T.Buffer[(128,), "float32"], 
            B: T.Buffer[(128,), "float32"], 
            C: T.Buffer[(128,), "float32"]):
    T.func_attr({"global_symbol": "main", "tir.noalias": True})
    for i in range(128):
        with T.block("C"):
            vi = T.axis.spatial(128, i)
            C[vi] = A[vi] + B[vi]

mod = tvm.IRModule.from_expr(main)
print(mod["main"].script())

Can I provide an external definition fo the kernel:

import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T

def foo(A, B, C):
    # Convert A,B,C to numpy.ndarray or something?
    __ my_random_kernel(A, B, C)

@T.prim_func
def main(A: T.Buffer[(128,), "float32"], 
            B: T.Buffer[(128,), "float32"], 
            C: T.Buffer[(128,), "float32"]):
    T.func_attr({"global_symbol": "main", "tir.noalias": True})
    foo(A, B, C)

It seems like external definitions can only be provided at the Tensor Expression level External Tensor Functions — tvm 0.9.0 documentation.

Is there a way to provide external function definitions at the TVMScript level?

2

The easies way to do this is via ffi call or Packed Func call, where you can call from both relax and tir.

see Customize Optimization — tvm 0.21.dev0 documentation for relax dispatch. For TIR functions you can use intrinsic call_packed to achieve the goal