Function _kernel_compile

tests/python/disco/test_nvshmem.py:282–363

Compile and run a kernel that calls NVSHMEM functions. Runs in a fresh process, so setting the env var is safe.

(compile_mode)

def _kernel_compile(compile_mode):
    """Compile and run a kernel that calls NVSHMEM functions.

    Runs in a fresh process, so setting the env var is safe.
    """
    os.environ["TVM_CUDA_COMPILE_MODE"] = compile_mode

    num_workers = 2
    sess = di.ProcessSession(num_workers=num_workers)

    f_init_nvshmem_uid = tvm.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid")
    uid = f_init_nvshmem_uid()
    init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem")
    init_dfunc(uid, num_workers, 0)
    sess.sync_worker_0()

    try:

        @I.ir_module(s_tir=True)
        class NvshmemQueryModule:
            @T.prim_func(s_tir=True)
            def query_pe(
                my_pe_out: T.Buffer((1,), "int32"),
                n_pes_out: T.Buffer((1,), "int32"),
            ):
                with T.sblock("root"):
                    T.reads()
                    T.writes(my_pe_out[0:1], n_pes_out[0:1])
                    T.call_kernel(
                        NVSHMEM_QUERY_KERNEL_SOURCE,
                        ((1,), (1,)),  # grid=(1,), block=(1,)
                        my_pe_out.data,
                        n_pes_out.data,
                        kernel_name="nvshmem_query_kernel",
                    )

            @R.function
            def main() -> R.Tuple(R.Tensor((1,), "int32"), R.Tensor((1,), "int32")):
                cls = NvshmemQueryModule
                with R.dataflow():
                    my_pe = R.call_tir(
                        cls.query_pe,
                        (),
                        out_sinfo=[
                            R.Tensor((1,), "int32"),
                            R.Tensor((1,), "int32"),
                        ],
                    )
                    R.output(my_pe)
                return my_pe

        tmpdir = tempfile.mkdtemp()
        try:
            path = tmpdir + "/test_nvshmem_kernel.so"

            target = tvm.target.Target("cuda")
            tvm.compile(NvshmemQueryModule, target=target).export_library(path)
            mod = sess.load_vm_module(path)