4.3 按 Stream 顺序的内存分配器（Stream-Ordered Memory Allocator）

void *ptr;
size_t size = 512;
cudaMallocAsync(&ptr, size, cudaStreamPerThread);
// 使用该分配执行工作
kernel<<<..., cudaStreamPerThread>>>(ptr, ...);
// 可以指定异步释放，而不必同步 CPU 和 GPU
cudaFreeAsync(ptr, cudaStreamPerThread);

cudaMallocAsync(&ptr, size, stream1);
cudaEventRecord(event1, stream1);
// stream2 在访问前必须等待分配就绪
cudaStreamWaitEvent(stream2, event1);
kernel<<<..., stream2>>>(ptr, ...);
cudaEventRecord(event2, stream2);
// stream3 在释放分配前必须等待 stream2 完成访问
cudaStreamWaitEvent(stream3, event2);
cudaFreeAsync(ptr, stream3);

cudaMalloc(&ptr, size);
kernel<<<..., stream>>>(ptr, ...);
cudaFreeAsync(ptr, stream);

cudaMallocAsync(&ptr, size,stream);
kernel<<<..., stream>>>(ptr, ...);
// 需要同步，以免过早释放内存
cudaStreamSynchronize(stream);
cudaFree(ptr);

// 创建一个类似设备 0 隐式池的内存池
int device = 0;
cudaMemPoolProps poolProps = { };
poolProps.allocType = cudaMemAllocationTypePinned;
poolProps.location.id = device;
poolProps.location.type = cudaMemLocationTypeDevice;

cudaMemPoolCreate(&memPool, &poolProps));

// 创建驻留于支持 IPC 共享（通过文件描述符）的 CPU NUMA 节点上的池
int cpu_numa_id = 0;
cudaMemPoolProps poolProps = { };
poolProps.allocType = cudaMemAllocationTypePinned;
poolProps.location.id = cpu_numa_id;
poolProps.location.type = cudaMemLocationTypeHostNuma;
poolProps.handleType = cudaMemHandleTypePosixFileDescriptor;

cudaMemPoolCreate(&ipcMemPool, &poolProps));

// 展示 cudaMemPoolSetAccess 的用法
cudaError_t setAccessOnDevice(cudaMemPool_t memPool, int residentDevice,
              int accessingDevice) {
    cudaMemAccessDesc accessDesc = {};
    accessDesc.location.type = cudaMemLocationTypeDevice;
    accessDesc.location.id = accessingDevice;
    accessDesc.flags = cudaMemAccessFlagsProtReadWrite;

    int canAccess = 0;
    cudaError_t error = cudaDeviceCanAccessPeer(&canAccess, accessingDevice,
              residentDevice);
    if (error != cudaSuccess) {
        return error;
    } else if (canAccess == 0) {
        return cudaErrorPeerAccessUnsupported;
    }

    // 使地址可访问
    return cudaMemPoolSetAccess(memPool, &accessDesc, 1);
}

// 导出进程
// 在设备 0 上创建可导出的、具有 IPC 能力的池
cudaMemPoolProps poolProps = { };
poolProps.allocType = cudaMemAllocationTypePinned;
poolProps.location.id = 0;
poolProps.location.type = cudaMemLocationTypeDevice;

// 将 handleTypes 设置为非零值会使池可导出（具备 IPC 能力）
poolProps.handleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;

cudaMemPoolCreate(&memPool, &poolProps));

// 基于 FD 的句柄是整数类型
int fdHandle = 0;

// 获取池的操作系统原生句柄。
// 注意这里传入的是句柄内存的指针。
cudaMemPoolExportToShareableHandle(&fdHandle,
             memPool,
             CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR,
             0);

// 必须使用适当的操作系统专用 API 将句柄发送给导入进程。

// 导入进程
int fdHandle;
// 需要使用适当的操作系统专用 API 从导出进程取得句柄。
// 根据可共享句柄创建导入池。
// 注意这里按值传入句柄。
cudaMemPoolImportFromShareableHandle(&importedMemPool,
          (void*)fdHandle,
          CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR,
          0);

// 在导出进程中准备分配
cudaMemPoolPtrExportData exportData;
cudaEvent_t readyIpcEvent;
cudaIpcEventHandle_t readyIpcEventHandle;

// 用于进程间协调的 IPC event
// cudaEventInterprocess 标志使 event 成为 IPC event
// 出于性能原因设置 cudaEventDisableTiming
cudaEventCreate(&readyIpcEvent, cudaEventDisableTiming | cudaEventInterprocess)

// 从导出内存池分配
cudaMallocAsync(&ptr, size,exportMemPool, stream);

// 用于共享分配就绪状态的 event
cudaEventRecord(readyIpcEvent, stream);
cudaMemPoolExportPointer(&exportData, ptr);
cudaIpcGetEventHandle(&readyIpcEventHandle, readyIpcEvent);

// 使用任意机制将 IPC event 和指针导出数据共享给导入进程。
// 这里将数据复制到共享内存
shmem->ptrData = exportData;
shmem->readyIpcEventHandle = readyIpcEventHandle;
// 通知消费者数据已经就绪

// 导入分配
cudaMemPoolPtrExportData *importData = &shmem->prtData;
cudaEvent_t readyIpcEvent;
cudaIpcEventHandle_t *readyIpcEventHandle = &shmem->readyIpcEventHandle;

// 需要使用任意机制从导出进程取得 IPC event 句柄和导出数据。
// 这里使用共享内存，只需要同步以确保共享内存已经填充。
cudaIpcOpenEventHandle(&readyIpcEvent, readyIpcEventHandle);

// 导入分配。该操作不会阻塞等待分配就绪。
cudaMemPoolImportPointer(&ptr, importedMemPool, importData);

// 在导入进程使用分配之前，等待分配 stream 中的先前操作完成。
cudaStreamWaitEvent(stream, readyIpcEvent);
kernel<<<..., stream>>>(ptr, ...);

// 导入进程必须先于导出进程释放
kernel<<<..., stream>>>(ptr, ...);

// 导入进程中的最后一次访问
cudaFreeAsync(ptr, stream);

// 导入进程释放后，不允许再从导入进程访问
cudaIpcEventRecord(finishedIpcEvent, stream);

// 导出进程
// 导出进程需要将自己的释放与导入进程的 stream 顺序释放协调起来。
cudaStreamWaitEvent(stream, finishedIpcEvent);
kernel<<<..., stream>>>(ptrInExportingProcess, ...);

// 导入进程的释放不会阻止导出进程使用该分配。
cudFreeAsync(ptrInExportingProcess,stream);

int driverVersion = 0;
int deviceSupportsMemoryPools = 0;
int poolSupportedHandleTypes = 0;
cudaDriverGetVersion(&driverVersion);
if (driverVersion >= 11020) {
    cudaDeviceGetAttribute(&deviceSupportsMemoryPools,
                           cudaDevAttrMemoryPoolsSupported, device);
}
if (deviceSupportsMemoryPools != 0) {
    // `device` 支持按 Stream 顺序的内存分配器
}

if (driverVersion >= 11030) {
    cudaDeviceGetAttribute(&poolSupportedHandleTypes,
              cudaDevAttrMemoryPoolSupportedHandleTypes, device);
}
if (poolSupportedHandleTypes & cudaMemHandleTypePosixFileDescriptor) {
   // 指定设备上的池可以使用基于 POSIX 文件描述符的 IPC 创建
}

Cuuint64_t setVal = UINT64_MAX;
cudaMemPoolSetAttribute(memPool, cudaMemPoolAttrReleaseThreshold, &setVal);

Cuuint64_t setVal = UINT64_MAX;
cudaMemPoolSetAttribute(memPool, cudaMemPoolAttrReleaseThreshold, &setVal);

// 应用阶段：需要从按 Stream 顺序的分配器获得大量内存
for (i=0; i<10; i++) {
    for (j=0; j<10; j++) {
        cudaMallocAsync(&ptrs[j],size[j], stream);
    }
    kernel<<<...,stream>>>(ptrs,...);
    for (j=0; j<10; j++) {
        cudaFreeAsync(ptrs[j], stream);
    }
}

// 下一阶段不需要这么多内存。
// 同步，使 trim 操作知道这些分配已不再使用。
cudaStreamSynchronize(stream);
cudaMemPoolTrimTo(mempool, 0);

// 其他进程/分配机制现在可以使用 trim 操作释放的物理内存。

// 批量获取使用统计的辅助函数示例
struct usageStatistics {
    cuuint64_t reserved;
    cuuint64_t reservedHigh;
    cuuint64_t used;
    cuuint64_t usedHigh;
};

void getUsageStatistics(cudaMemoryPool_t memPool, struct usageStatistics *statistics)
{
    cudaMemPoolGetAttribute(memPool, cudaMemPoolAttrReservedMemCurrent, statistics->reserved);
    cudaMemPoolGetAttribute(memPool, cudaMemPoolAttrReservedMemHigh, statistics->reservedHigh);
    cudaMemPoolGetAttribute(memPool, cudaMemPoolAttrUsedMemCurrent, statistics->used);
    cudaMemPoolGetAttribute(memPool, cudaMemPoolAttrUsedMemHigh, statistics->usedHigh);
}

// 重置水位线后，它们会取当前值。
void resetStatistics(cudaMemoryPool_t memPool)
{
    cuuint64_t value = 0;
    cudaMemPoolSetAttribute(memPool, cudaMemPoolAttrReservedMemHigh, &value);
    cudaMemPoolSetAttribute(memPool, cudaMemPoolAttrUsedMemHigh, &value);
}

cudaMallocAsync(&ptr, size, originalStream);
kernel<<<..., originalStream>>>(ptr, ...);
cudaFreeAsync(ptr, originalStream);
cudaEventRecord(event,originalStream);

// 等待捕获了另一个 stream 中释放操作的 event，
// 可以让分配器在另一个 stream 的新分配请求中复用内存，
// 前提是启用了 cudaMemPoolReuseFollowEventDependencies。
cudaStreamWaitEvent(otherStream, event);
cudaMallocAsync(&ptr2, size, otherStream);

cudaMallocAsync(&ptr, size, originalStream);
kernel<<<..., originalStream>>>(ptr, ...);
cudaFreeAsync(ptr, originalStream);

// 经过一段时间后，kernel 完成运行
wait(10);

// 启用 cudaMemPoolReuseAllowOpportunistic 时，
// 分配器可以根据 originalStream 的进度使用之前的分配满足请求。
cudaMallocAsync(&ptr2, size, otherStream);

cudaMallocAsync(&ptr, size, originalStream);
kernel<<<..., originalStream>>>(ptr, ...);
cudaFreeAsync(ptr, originalStream);

// 启用 cudaMemPoolReuseAllowInternalDependencies，且驱动无法分配更多物理内存时，
// 驱动可能等效于在分配 stream 中执行 cudaStreamWaitEvent，
// 确保 futureStream 中的未来工作发生在 originalStream 中
// 仍可能访问原始分配的工作之后。
cudaMallocAsync(&ptr2, size, otherStream);