derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
"""Team-scoped collective binding tests."""
|
|
|
|
from conftest import requires_multi_pe, requires_torch
|
|
|
|
import rocshmem4py
|
|
from rocshmem4py.interop import torch as rshmem_torch
|
|
|
|
pytestmark = [requires_torch, requires_multi_pe]
|
|
|
|
|
|
def test_world_alltoall_torch_wrapper():
|
|
import torch
|
|
|
|
me = rocshmem4py.rocshmem_my_pe()
|
|
n = rocshmem4py.rocshmem_n_pes()
|
|
chunk = 4
|
|
|
|
src = rshmem_torch.create_tensor((n, chunk), torch.int32)
|
|
dst = rshmem_torch.create_tensor((n, chunk), torch.int32)
|
|
try:
|
|
for peer in range(n):
|
|
src[peer].fill_(me * 100 + peer)
|
|
dst.fill_(-1)
|
|
torch.cuda.synchronize()
|
|
rshmem_torch.barrier_all()
|
|
|
|
rshmem_torch.alltoall(rocshmem4py.ROCSHMEM_TEAM_WORLD, dst, src)
|
|
rshmem_torch.barrier_all()
|
|
torch.cuda.synchronize()
|
|
|
|
expected = torch.empty_like(dst)
|
|
for peer in range(n):
|
|
expected[peer].fill_(peer * 100 + me)
|
|
torch.testing.assert_close(dst, expected)
|
|
finally:
|
|
rshmem_torch.free_tensor(src)
|
|
rshmem_torch.free_tensor(dst)
|
|
|
|
|
|
def test_world_broadcast_torch_wrapper():
|
|
import torch
|
|
|
|
me = rocshmem4py.rocshmem_my_pe()
|
|
root = 0
|
|
|
|
src = rshmem_torch.create_tensor((16,), torch.int32)
|
|
dst = rshmem_torch.create_tensor((16,), torch.int32)
|
|
try:
|
|
src.fill_(me)
|
|
dst.fill_(-1)
|
|
if me == root:
|
|
dst.copy_(src)
|
|
torch.cuda.synchronize()
|
|
rshmem_torch.barrier_all()
|
|
|
|
rshmem_torch.broadcast(rocshmem4py.ROCSHMEM_TEAM_WORLD, dst, src, root)
|
|
rshmem_torch.barrier_all()
|
|
torch.cuda.synchronize()
|
|
|
|
torch.testing.assert_close(dst, torch.full_like(dst, root))
|
|
finally:
|
|
rshmem_torch.free_tensor(src)
|
|
rshmem_torch.free_tensor(dst)
|
|
|
|
|
|
def test_team_collective_symbols_are_exported():
|
|
for name in (
|
|
"rocshmem_alltoallmem_on_stream",
|
|
"rocshmem_broadcastmem_on_stream",
|
|
"rocshmem_barrier",
|
|
"rocshmem_team_sync",
|
|
"rocshmem_barrier_on_stream",
|
|
"rocshmem_team_sync_on_stream",
|
|
"rocshmem_global_exit",
|
|
"rocshmem_dump_stats",
|
|
"rocshmem_reset_stats",
|
|
):
|
|
assert hasattr(rocshmem4py, name)
|