Tests data transfer operations (get/put) at thread, warp, and block scopes.
(sess, scope, shape, nwarps, nelems, op_name)
| 88 | run_prim_func(sess, main, res_array) |
| 89 | |
| 90 | def test_transfer(sess, scope, shape, nwarps, nelems, op_name): |
| 91 | """Tests data transfer operations (get/put) at thread, warp, and block scopes.""" |
| 92 | dtype = "float32" |
| 93 | is_get = "get" in op_name |
| 94 | op_func = getattr(T.nvshmem, op_name) |
| 95 | if scope != "thread": |
| 96 | op_func = getattr(op_func, scope) |
| 97 | |
| 98 | # fmt: off |
| 99 | @T.prim_func |
| 100 | def main(A: T.Buffer(shape, dtype), B: T.Buffer(shape, dtype)): |
| 101 | T.device_entry() |
| 102 | cta_id = T.cta_id([1]) |
| 103 | warp_id = T.warp_id([nwarps]) |
| 104 | lane_id = T.lane_id([32]) |
| 105 | tid = T.thread_id([nwarps * 32]) |
| 106 | |
| 107 | my_pe = T.nvshmem.my_pe() |
| 108 | n_pes = T.nvshmem.n_pes() |
| 109 | offset = T.if_then_else( |
| 110 | scope == "block", 0, T.if_then_else(scope == "thread", tid, warp_id * 32) |
| 111 | ) |
| 112 | op_func(dst=B.ptr_to([offset]), src=A.ptr_to([offset]), nelems=nelems, pe=(my_pe + 1) % n_pes) # noqa: E501 |
| 113 | T.nvshmem.quiet() |
| 114 | # fmt: on |
| 115 | |
| 116 | def init_fn(i, s, d): |
| 117 | return np.arange(s[0], dtype=d) + i * 100 |
| 118 | |
| 119 | A_array = create_nvshmem_array(sess, shape, dtype, init_fn) |
| 120 | B_array = create_nvshmem_array(sess, shape, dtype) |
| 121 | sess.sync_worker_0() |
| 122 | run_prim_func(sess, main, A_array, B_array) |
| 123 | |
| 124 | for i in range(NUM_WORKERS): |
| 125 | if is_get: |
| 126 | expected_B = A_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy() |
| 127 | actual_B = B_array.debug_get_from_remote(i).numpy() |
| 128 | else: # put |
| 129 | expected_B = A_array.debug_get_from_remote(i).numpy() |
| 130 | actual_B = B_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy() |
| 131 | np.testing.assert_equal(actual_B, expected_B) |
| 132 | |
| 133 | def test_signal_op(sess, sig_op): |
| 134 | """Tests signal_op and wait_until to implement a barrier-like pattern.""" |
no test coverage detected
searching dependent graphs…