ipc_event.cpp

#include "shmem.hpp"

namespace {

void print_ipc_handle(hipIpcEventHandle_t x, const char *message = nullptr) {
    int size = sizeof(x);
    if (message != nullptr) {
        printf("%s : ", message);
    }
    const uint8_t *ptr = reinterpret_cast<const uint8_t*>(&x);
    for(int i=0; i<size; i++) {
        printf("%d ", (int)ptr[i]);
    }
    printf("\n");
}

constexpr size_t SIZE_PER_PROCESS = 8 * 1024 * 1024; // 8MB

}

int main() {
    // <-- create MAX_NUM_PROCESS processes using fork()
    int rank = 0;
    for(int i=1; i<MAX_NUM_PROCESS; i++) {
        if (fork() == 0) {
            rank = i;
            break;
        }
    }
    printf("after fork, pid : %d, rank : %d\n", static_cast<int>(getpid()), rank);
    // -->

    CHECK(hipSetDevice(rank) == 0, "");

    SharedMemoryInfo *info = create_shared_memory_info("ipc_event", rank, MAX_NUM_PROCESS);
    SharedStruct *shm  = static_cast<SharedStruct *>(info->addr);
    
    // <-- hip variables
    void *local_dev_buffer = nullptr;
    CHECK(hipMalloc(&local_dev_buffer, SIZE_PER_PROCESS) == 0, "");
    hipIpcMemHandle_t mem_handle;
    CHECK(hipIpcGetMemHandle(&mem_handle, local_dev_buffer) == 0, "");
    memcpy(&shm->buffer_mem_handles[rank], &mem_handle, sizeof(hipIpcMemHandle_t));
    reusable_barrier(shm->barrier, shm->sense, MAX_NUM_PROCESS);

    void* buffer_mem_ptrs[MAX_NUM_PROCESS]{nullptr};
    for (size_t i = 0; i < MAX_NUM_PROCESS; i++) {
        if (i != rank) {
            CHECK(hipIpcOpenMemHandle(&buffer_mem_ptrs[i], shm->buffer_mem_handles[i],
                                                        hipIpcMemLazyEnablePeerAccess) == 0, "");
        } else {
            buffer_mem_ptrs[i] = local_dev_buffer;
        }
    }

    int leastPriority, greatestPriority;
    CHECK(hipDeviceGetStreamPriorityRange(&leastPriority, &greatestPriority) == 0, "");

    hipStream_t default_stream {nullptr};
    CHECK(hipStreamCreateWithFlags(&default_stream, hipStreamNonBlocking) == 0, "");

    hipStream_t copy_streams[MAX_NUM_PROCESS];
    hipEvent_t copy_events[MAX_NUM_PROCESS];
    for(size_t i=0; i<MAX_NUM_PROCESS; i++) {
        CHECK(hipStreamCreateWithPriority(&copy_streams[i], hipStreamNonBlocking, greatestPriority) == 0, "");
        CHECK(hipEventCreateWithFlags(&copy_events[i], hipEventDisableTiming) == 0, "");
    }

    hipEvent_t local_sync_event {nullptr};
    CHECK(hipEventCreateWithFlags(
                &local_sync_event, hipEventDisableTiming | hipEventInterprocess) == 0, "");
    hipIpcEventHandle_t local_sync_event_handle;
    CHECK(hipIpcGetEventHandle(&local_sync_event_handle, local_sync_event) == 0, "");
    memcpy(&shm->sync_event_handles[rank], &local_sync_event_handle, sizeof(hipIpcEventHandle_t));
    
    hipEvent_t local_exit_event {nullptr};
    CHECK(hipEventCreateWithFlags(
                &local_exit_event, hipEventDisableTiming | hipEventInterprocess) == 0, "");
    hipIpcEventHandle_t local_exit_event_handle;
    CHECK(hipIpcGetEventHandle(&local_exit_event_handle, local_exit_event) == 0, "");
    memcpy(&shm->exit_event_handles[rank], &local_exit_event_handle, sizeof(hipIpcEventHandle_t));
    
    reusable_barrier(shm->barrier, shm->sense, MAX_NUM_PROCESS);

    hipEvent_t exit_events[MAX_NUM_PROCESS];
    hipEvent_t sync_events[MAX_NUM_PROCESS];
    for(size_t i=0; i<MAX_NUM_PROCESS; i++) {
        if (i == rank) {
            sync_events[i] = local_sync_event;
            exit_events[i] = local_exit_event;
        }
        else {
            CHECK(hipIpcOpenEventHandle(&sync_events[i], shm->sync_event_handles[i]) == 0, "");
            CHECK(hipIpcOpenEventHandle(&exit_events[i], shm->exit_event_handles[i]) == 0, "");
        }
    }

    hipEvent_t entry_event {nullptr};
    CHECK(hipEventCreateWithFlags(&entry_event, hipEventDisableTiming) == 0, "");

    void *input {nullptr};
    CHECK(hipMalloc(&input, SIZE_PER_PROCESS) == 0, "");
    CHECK(hipMemsetAsync(input, 0, SIZE_PER_PROCESS, default_stream) == 0, "");
    void *output {nullptr};
    CHECK(hipMalloc(&output, SIZE_PER_PROCESS * MAX_NUM_PROCESS) == 0, "");
    // hip variables -->

    // CHECK(hipDeviceSynchronize() == 0, "");
    // <-- copy
    for(int iter = 0; iter < 60; iter++) {
        CHECK(hipEventRecord(entry_event, default_stream) == 0, "");
        CHECK(hipStreamWaitEvent(copy_streams[rank], entry_event, 0) == 0, "");

        CHECK(hipMemcpyAsync(buffer_mem_ptrs[rank], input, SIZE_PER_PROCESS,
            hipMemcpyDeviceToDeviceNoCU, copy_streams[rank]) == 0, "");
        CHECK(hipEventRecord(sync_events[rank], copy_streams[rank]) == 0, "");
        reusable_barrier(shm->barrier, shm->sense, MAX_NUM_PROCESS);

        for(size_t i=0; i<MAX_NUM_PROCESS; i++) {
            if (i != rank) {
                CHECK(hipStreamWaitEvent(copy_streams[i], sync_events[i], 0) == 0, "");
            }
        }

        for(size_t i=0; i<MAX_NUM_PROCESS; i++) {
            const void *src = buffer_mem_ptrs[i];
            void *dst = static_cast<uint8_t *>(output) + SIZE_PER_PROCESS * i;
            CHECK(hipMemcpyAsync(dst, src, SIZE_PER_PROCESS, hipMemcpyDeviceToDeviceNoCU, copy_streams[i]) == 0, "");

            CHECK(hipEventRecord(copy_events[i], copy_streams[i]) == 0, "");
        }
        
        for(size_t i=0; i<MAX_NUM_PROCESS; i++) {
            CHECK(hipStreamWaitEvent(default_stream, copy_events[i], 0) == 0, "");
        }
        CHECK(hipEventRecord(exit_events[rank], default_stream) == 0, "");
        reusable_barrier(shm->barrier, shm->sense, MAX_NUM_PROCESS);
        for(size_t i=0; i<MAX_NUM_PROCESS; i++) {
            if (i != rank) {
                CHECK(hipStreamWaitEvent(default_stream, exit_events[i], 0) == 0, "");
            }
        }

        CHECK(hipStreamSynchronize(default_stream) == 0, "");
    }
    // copy -->
    return 0;
}