(sess)
| 220 | np.testing.assert_equal(signal_np[0], cmp_value) |
| 221 | |
| 222 | def test_fence_barrier(sess): |
| 223 | shape = (64,) |
| 224 | dtype = "float32" |
| 225 | |
| 226 | # fmt: off |
| 227 | @T.prim_func |
| 228 | def main(A: T.Buffer(shape, dtype), B: T.Buffer(shape, dtype), res: T.Buffer((1,), "uint64")): # noqa: E501 |
| 229 | T.device_entry() |
| 230 | cta_id = T.cta_id([1]) |
| 231 | warp_id = T.warp_id([nwarps]) |
| 232 | lane_id = T.lane_id([32]) |
| 233 | tid = T.thread_id([2 * 32]) |
| 234 | my_pe = T.nvshmem.my_pe() |
| 235 | n_pes = T.nvshmem.n_pes() |
| 236 | dst_pe = (my_pe + 1) % n_pes |
| 237 | T.nvshmem.barrier_all() |
| 238 | 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 |
| 239 | T.nvshmem.fence() |
| 240 | if tid == 0: |
| 241 | T.nvshmem.signal_op(sig_addr=res.ptr_to([0]), signal=1, sig_op="set", pe=dst_pe) |
| 242 | T.nvshmem.wait_until(ivar=res.ptr_to([0]), cmp="eq", cmp_value=1) |
| 243 | # fmt: on |
| 244 | def init_fn(i, s, d): |
| 245 | return np.arange(s[0], dtype=d) + i * 100 |
| 246 | |
| 247 | A_array = create_nvshmem_array(sess, shape, dtype, init_fn) |
| 248 | B_array = create_nvshmem_array(sess, shape, dtype) |
| 249 | res_array = create_nvshmem_array(sess, (1,), "uint64") |
| 250 | run_prim_func(sess, main, A_array, B_array, res_array) |
| 251 | |
| 252 | for i in range(NUM_WORKERS): |
| 253 | expected_B = A_array.debug_get_from_remote(i).numpy() |
| 254 | actual_B = B_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy() |
| 255 | np.testing.assert_equal(actual_B, expected_B) |
| 256 | |
| 257 | # test thread info |
| 258 | test_thread_info(sess) |
no test coverage detected
searching dependent graphs…