Function run_prim_func

tests/python/tirx/codegen/test_codegen_nvshmem.py:36–46

Compile, export, load, and run a PrimFunc in the shared disco session.

(sess, prim_func, *args)

def run_prim_func(sess, prim_func, *args):
    """Compile, export, load, and run a PrimFunc in the shared disco session."""
    target = tvm.target.Target("cuda")
    with tempfile.TemporaryDirectory() as tmpdir:
        path = f"{tmpdir}/test.so"
        mod = tvm.compile(prim_func, target=target, tir_pipeline="tirx")
        print(mod.mod.imports[0].inspect_source())
        mod.export_library(path)
        rt_mod = sess.load_vm_module(path)
        rt_mod["main"](*args)
        sess._sync_all()


def create_nvshmem_array(sess, shape, dtype, init_data_fn=None, zero_out=True):