# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Basic tests for a Disco nvshmem support""" # pylint: disable=missing-docstring import tempfile import numpy as np import pytest import tvm import tvm.testing from tvm.runtime import ShapeTuple from tvm.runtime import disco as di from tvm.script import tirx as T from tvm.support.popen_pool import PopenWorker from tvm.testing import env NUM_WORKERS = 4 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): """Create and optionally initialize an nvshmem-accessible DNDArray.""" nvshmem_empty = sess.get_global_func("runtime.disco.nvshmem.empty") arr = nvshmem_empty(ShapeTuple(shape), dtype, None) if init_data_fn: for i in range(NUM_WORKERS): arr.debug_copy_from(i, init_data_fn(i, shape, dtype)) elif zero_out: zero_data = np.zeros(shape, dtype=dtype) for i in range(NUM_WORKERS): arr.debug_copy_from(i, zero_data) return arr @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") @pytest.mark.skip(reason="nvshmem doesn't work with pytest") def test_codegen_nvshmem(): def _test_func(): ############ setup ############ 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() def test_thread_info(sess): @T.prim_func def main(res: T.Buffer((2,), "int32")): T.device_entry() cta_id = T.cta_id([1]) tid = T.thread_id([nwarps * 32]) res[0] = T.nvshmem.my_pe() res[1] = T.nvshmem.n_pes() res_array = sess.empty((2,), "int32") run_prim_func(sess, main, res_array) def test_transfer(sess, scope, shape, nwarps, nelems, op_name): """Tests data transfer operations (get/put) at thread, warp, and block scopes.""" dtype = "float32" is_get = "get" in op_name op_func = getattr(T.nvshmem, op_name) if scope != "thread": op_func = getattr(op_func, scope) # fmt: off @T.prim_func def main(A: T.Buffer(shape, dtype), B: T.Buffer(shape, dtype)): T.device_entry() cta_id = T.cta_id([1]) warp_id = T.warp_id([nwarps]) lane_id = T.lane_id([32]) tid = T.thread_id([nwarps * 32]) my_pe = T.nvshmem.my_pe() n_pes = T.nvshmem.n_pes() offset = T.if_then_else( scope == "block", 0, T.if_then_else(scope == "thread", tid, warp_id * 32) ) op_func(dst=B.ptr_to([offset]), src=A.ptr_to([offset]), nelems=nelems, pe=(my_pe + 1) % n_pes) # noqa: E501 T.nvshmem.quiet() # fmt: on def init_fn(i, s, d): return np.arange(s[0], dtype=d) + i * 100 A_array = create_nvshmem_array(sess, shape, dtype, init_fn) B_array = create_nvshmem_array(sess, shape, dtype) sess.sync_worker_0() run_prim_func(sess, main, A_array, B_array) for i in range(NUM_WORKERS): if is_get: expected_B = A_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy() actual_B = B_array.debug_get_from_remote(i).numpy() else: # put expected_B = A_array.debug_get_from_remote(i).numpy() actual_B = B_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy() np.testing.assert_equal(actual_B, expected_B) def test_signal_op(sess, sig_op): """Tests signal_op and wait_until to implement a barrier-like pattern.""" cmp_value = 1 if sig_op == "set" else 2 # fmt: off @T.prim_func def main(res: T.Buffer((1,), "uint64")): T.device_entry() cta_id = T.cta_id([1]) tid = T.thread_id([nwarps * 32]) my_pe = T.nvshmem.my_pe() n_pes = T.nvshmem.n_pes() dst_pe = (my_pe + 1) % n_pes if sig_op == "add": res[0] = 1 T.nvshmem.barrier_all() T.nvshmem.signal_op(sig_addr=res.ptr_to([0]), signal=1, sig_op=sig_op, pe=dst_pe) T.nvshmem.wait_until(ivar=res.ptr_to([0]), cmp="eq", cmp_value=cmp_value) # fmt: on res_array = create_nvshmem_array(sess, (1,), "uint64") sess.sync_worker_0() run_prim_func(sess, main, res_array) for i in range(NUM_WORKERS): res = res_array.debug_get_from_remote(i).numpy() if sig_op == "set": np.testing.assert_equal(res[0], 1) elif sig_op == "add": np.testing.assert_equal(res[0], 2) def test_put_signal(sess, scope, shape, nwarps, nelems, cmp_value): """Tests combined data transfer and signal operations at thread/warp/block scopes.""" dtype = "float32" op_func = getattr(T.nvshmem, "putmem_signal_nbi") if scope != "thread": op_func = getattr(op_func, scope) @T.prim_func def main( A: T.Buffer(shape, dtype), B: T.Buffer(shape, dtype), signal_array: T.Buffer((1,), "uint64"), ): T.device_entry() cta_id = T.cta_id([1]) warp_id = T.warp_id([nwarps]) lane_id = T.lane_id([32]) tid = T.thread_id([nwarps * 32]) my_pe = T.nvshmem.my_pe() n_pes = T.nvshmem.n_pes() dst_pe = (my_pe + 1) % n_pes offset = T.if_then_else( scope == "block", 0, T.if_then_else(scope == "thread", tid, warp_id * 32), ) op_func( dst=B.access_ptr("w", offset=offset), src=A.access_ptr("r", offset=offset), nelems=nelems, sig_addr=signal_array.access_ptr("w", offset=0), signal=1, sig_op="set", pe=dst_pe, ) T.nvshmem.wait_until( ivar=signal_array.access_ptr("r", offset=0), cmp="eq", cmp_value=cmp_value, ) def init_A(i, s, d): return np.arange(s[0], dtype=d) + i * 100 A_array = create_nvshmem_array(sess, shape, dtype, init_A) B_array = create_nvshmem_array(sess, shape, dtype) signal_array = create_nvshmem_array(sess, (1,), "uint64") sess.sync_worker_0() run_prim_func(sess, main, A_array, B_array, signal_array) for i in range(NUM_WORKERS): expected = A_array.debug_get_from_remote(i).numpy() actual = B_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy() signal_np = signal_array.debug_get_from_remote(i).numpy() np.testing.assert_equal(actual, expected) np.testing.assert_equal(signal_np[0], cmp_value) def test_fence_barrier(sess): shape = (64,) dtype = "float32" # fmt: off @T.prim_func def main(A: T.Buffer(shape, dtype), B: T.Buffer(shape, dtype), res: T.Buffer((1,), "uint64")): # noqa: E501 T.device_entry() cta_id = T.cta_id([1]) warp_id = T.warp_id([nwarps]) lane_id = T.lane_id([32]) tid = T.thread_id([2 * 32]) my_pe = T.nvshmem.my_pe() n_pes = T.nvshmem.n_pes() dst_pe = (my_pe + 1) % n_pes T.nvshmem.barrier_all() T.nvshmem.putmem_nbi.block(dst=B.ptr_to([0]), src=A.ptr_to([0]), nelems=4 * 64, pe=(my_pe + 1) % n_pes) # noqa: E501 T.nvshmem.fence() if tid == 0: T.nvshmem.signal_op(sig_addr=res.ptr_to([0]), signal=1, sig_op="set", pe=dst_pe) T.nvshmem.wait_until(ivar=res.ptr_to([0]), cmp="eq", cmp_value=1) # fmt: on def init_fn(i, s, d): return np.arange(s[0], dtype=d) + i * 100 A_array = create_nvshmem_array(sess, shape, dtype, init_fn) B_array = create_nvshmem_array(sess, shape, dtype) res_array = create_nvshmem_array(sess, (1,), "uint64") run_prim_func(sess, main, A_array, B_array, res_array) for i in range(NUM_WORKERS): expected_B = A_array.debug_get_from_remote(i).numpy() actual_B = B_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy() np.testing.assert_equal(actual_B, expected_B) # test thread info test_thread_info(sess) print("\n\ntest_thread_info done\n\n") # test transfer for scope, shape, nwarps, nelems, op_name in [ ("thread", (32,), 1, 4, "getmem_nbi"), ("thread", (32,), 1, 4, "putmem_nbi"), ("warp", (64,), 2, 4 * 32, "getmem_nbi"), ("warp", (64,), 2, 4 * 32, "putmem_nbi"), ("block", (64,), 2, 4 * 64, "getmem_nbi"), ("block", (64,), 2, 4 * 64, "putmem_nbi"), ]: test_transfer(sess, scope, shape, nwarps, nelems, op_name) print(f"\n\ntest_transfer done for {scope}, {shape}, {nwarps}, {nelems}, {op_name}\n\n") # test signal op for sig_op in ["set", "add"]: test_signal_op(sess, sig_op) print(f"\n\ntest_signal_op done for {sig_op}\n\n") # test put signal for scope, shape, nwarps, nelems, cmp_value in [ ("thread", (32,), 1, 4, 32), ("warp", (64,), 2, 4 * 32, 2), ("block", (64,), 2, 4 * 64, 1), ]: test_put_signal(sess, scope, shape, nwarps, nelems, cmp_value) print( f"\n\ntest_put_signal done for {scope}, {shape}, {nwarps}, {nelems}, {cmp_value}\n\n" # noqa: E501 ) # test fence barrier test_fence_barrier(sess) print("\n\ntest_fence_barrier done\n\n") ############ cleanup ############ finalize_dfunc = sess.get_global_func("runtime.disco.nvshmem.finalize_nvshmem") finalize_dfunc() sess.sync_worker_0() sess.shutdown() return True def run_and_check(): worker = PopenWorker() try: worker.send(_test_func) assert worker.recv() finally: worker.kill() tvm.testing.run_with_gpu_lock(run_and_check) if __name__ == "__main__": tvm.testing.main()