"""Memory-management tests for rocshmem4py. Covers the symmetric-heap allocators and the user-buffer registration APIs that previously had no Python coverage: * ``rocshmem_calloc`` -- zero-initialized collective allocation * ``rocshmem_align`` -- aligned collective allocation * ``rocshmem_buffer_register`` -- register a non-symmetric user buffer * ``rocshmem_buffer_unregister`` -- deregister a single user buffer * ``rocshmem_buffer_unregister_all`` -- deregister every user buffer These tests are torch-free: they drive HIP directly through ``ctypes`` so they run in any CI image. ``calloc`` / ``align`` are collective (each wraps an internal ``rocshmem_barrier_all``), so every PE must call them the same number of times -- pytest runs each test on every rank, which preserves that. The buffer register/unregister APIs are local backend operations and need real device memory (the GDA backend registers it with the NIC via ``ibv_reg_mr``), hence the ``hipMalloc`` shim below. """ import ctypes import pytest import rocshmem4py # noqa: E402 # --------------------------------------------------------------------------- # Minimal HIP ctypes shim (alloc/free + a device->host copy for verification) # --------------------------------------------------------------------------- try: _hip = ctypes.CDLL("libamdhip64.so") except OSError as e: pytest.skip(f"libamdhip64.so not loadable: {e}", allow_module_level=True) HIP_MEMCPY_DEVICE_TO_HOST = 2 _hip.hipMalloc.restype = ctypes.c_int _hip.hipMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t] _hip.hipFree.restype = ctypes.c_int _hip.hipFree.argtypes = [ctypes.c_void_p] _hip.hipMemcpy.restype = ctypes.c_int _hip.hipMemcpy.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ] _hip.hipDeviceSynchronize.restype = ctypes.c_int def _hip_check(rc: int, fn: str) -> None: if rc != 0: raise RuntimeError(f"{fn} failed with HIP error {rc}") def hip_malloc(nbytes: int) -> int: p = ctypes.c_void_p() _hip_check(_hip.hipMalloc(ctypes.byref(p), nbytes), "hipMalloc") return p.value or 0 def hip_free(device_ptr: int) -> None: _hip_check(_hip.hipFree(ctypes.c_void_p(device_ptr)), "hipFree") def d2h(host_buf, device_ptr: int) -> None: _hip_check( _hip.hipMemcpy( ctypes.cast(host_buf, ctypes.c_void_p), ctypes.c_void_p(device_ptr), ctypes.sizeof(host_buf), HIP_MEMCPY_DEVICE_TO_HOST, ), "hipMemcpy D2H", ) _hip_check(_hip.hipDeviceSynchronize(), "hipDeviceSynchronize") # --------------------------------------------------------------------------- # rocshmem_calloc # --------------------------------------------------------------------------- def test_calloc_zero_initialized(): """calloc returns a valid pointer to count*size zero-initialized bytes.""" count, size = 64, ctypes.sizeof(ctypes.c_int32) ptr = rocshmem4py.rocshmem_calloc(count, size) try: assert ptr > 0 rx = (ctypes.c_int32 * count)(*([0xABCD] * count)) d2h(rx, ptr) assert all(v == 0 for v in rx), f"calloc memory not zeroed: {list(rx)[:8]}..." finally: rocshmem4py.rocshmem_free(ptr) def test_calloc_zero_args_raises(): """count==0 or size==0 yields NULL in the C API, which the binding maps to RuntimeError.""" with pytest.raises(RuntimeError): rocshmem4py.rocshmem_calloc(0, 4) # --------------------------------------------------------------------------- # rocshmem_align # --------------------------------------------------------------------------- @pytest.mark.parametrize("alignment", [8, 64, 256, 4096]) def test_align_returns_aligned_pointer(alignment): ptr = rocshmem4py.rocshmem_align(alignment, 1024) try: assert ptr > 0 assert ptr % alignment == 0, ( f"pointer {hex(ptr)} not aligned to {alignment}" ) finally: rocshmem4py.rocshmem_free(ptr) @pytest.mark.parametrize("alignment", [0, 3, 24]) def test_align_invalid_alignment_raises(alignment): """alignment must be a non-zero power of two and a multiple of sizeof(void*); otherwise the C API returns NULL -> RuntimeError. Note: rocshmem_align is collective even on the invalid path (it still executes the internal barrier), so calling it on every rank is safe. """ with pytest.raises(RuntimeError): rocshmem4py.rocshmem_align(alignment, 1024) # --------------------------------------------------------------------------- # rocshmem_buffer_register / rocshmem_buffer_unregister # --------------------------------------------------------------------------- def test_buffer_register_unregister(): """A real (non-symmetric) HIP device buffer can be registered and then deregistered, both returning ROCSHMEM_SUCCESS.""" nbytes = 4096 addr = hip_malloc(nbytes) try: rc = rocshmem4py.rocshmem_buffer_register(addr, nbytes) assert rc == rocshmem4py.ROCSHMEM_SUCCESS, ( f"buffer_register returned {rc}" ) rc = rocshmem4py.rocshmem_buffer_unregister(addr) assert rc == rocshmem4py.ROCSHMEM_SUCCESS, ( f"buffer_unregister returned {rc}" ) finally: hip_free(addr) def test_buffer_register_invalid_args(): """Null address and zero length are rejected (non-zero status).""" assert rocshmem4py.rocshmem_buffer_register(0, 4096) != \ rocshmem4py.ROCSHMEM_SUCCESS addr = hip_malloc(4096) try: assert rocshmem4py.rocshmem_buffer_register(addr, 0) != \ rocshmem4py.ROCSHMEM_SUCCESS finally: hip_free(addr) def test_buffer_register_overlap_rejected(): """Registering an already-registered region must fail; the original registration is still valid and must be cleaned up.""" nbytes = 4096 addr = hip_malloc(nbytes) try: assert rocshmem4py.rocshmem_buffer_register(addr, nbytes) == \ rocshmem4py.ROCSHMEM_SUCCESS # Second registration of the same range overlaps -> rejected. assert rocshmem4py.rocshmem_buffer_register(addr, nbytes) != \ rocshmem4py.ROCSHMEM_SUCCESS finally: rocshmem4py.rocshmem_buffer_unregister(addr) hip_free(addr) def test_buffer_unregister_unknown_fails(): """Deregistering an address that was never registered returns non-success.""" addr = hip_malloc(4096) try: assert rocshmem4py.rocshmem_buffer_unregister(addr) != \ rocshmem4py.ROCSHMEM_SUCCESS finally: hip_free(addr) # --------------------------------------------------------------------------- # rocshmem_buffer_unregister_all # --------------------------------------------------------------------------- def test_buffer_unregister_all(): """unregister_all clears every registration; a subsequent single unregister of those buffers then fails (already cleared).""" nbytes = 4096 addrs = [hip_malloc(nbytes) for _ in range(3)] try: for a in addrs: assert rocshmem4py.rocshmem_buffer_register(a, nbytes) == \ rocshmem4py.ROCSHMEM_SUCCESS # Returns None (void); must not raise. assert rocshmem4py.rocshmem_buffer_unregister_all() is None # Everything has been deregistered already. for a in addrs: assert rocshmem4py.rocshmem_buffer_unregister(a) != \ rocshmem4py.ROCSHMEM_SUCCESS finally: for a in addrs: hip_free(a) if __name__ == "__main__": pytest.main([__file__, "-v"])