()
| 130 | |
| 131 | |
| 132 | def main() -> int: |
| 133 | import argparse |
| 134 | |
| 135 | parser = argparse.ArgumentParser( |
| 136 | description="Share a GPU buffer between two processes via CUDA IPC" |
| 137 | ) |
| 138 | parser.add_argument( |
| 139 | "--elements", |
| 140 | type=int, |
| 141 | default=1024, |
| 142 | help="Number of float32 elements in the shared buffer (default: 1024)", |
| 143 | ) |
| 144 | parser.add_argument("--device", type=int, default=0, help="CUDA device id") |
| 145 | args = parser.parse_args() |
| 146 | |
| 147 | # CUDA is incompatible with the ``fork`` start method because forked |
| 148 | # children inherit a corrupt CUDA state. Always use ``spawn``. |
| 149 | mp.set_start_method("spawn", force=True) |
| 150 | |
| 151 | device = Device(args.device) |
| 152 | device.set_current() |
| 153 | print_gpu_info(device) |
| 154 | |
| 155 | if not check_ipc_support(device): |
| 156 | print("\nCUDA IPC memory pools are not supported on this platform.") |
| 157 | return 2 |
| 158 | |
| 159 | N = args.elements |
| 160 | nbytes = N * np.dtype(np.float32).itemsize |
| 161 | parent_seed = 100 |
| 162 | child_seed = -1.0 |
| 163 | |
| 164 | # Create an IPC-enabled memory pool. Buffers allocated from this MR |
| 165 | # are picklable and can be shared across processes. |
| 166 | mr = DeviceMemoryResource( |
| 167 | device, |
| 168 | options=DeviceMemoryResourceOptions( |
| 169 | max_size=max(nbytes * 4, 1 << 20), |
| 170 | ipc_enabled=True, |
| 171 | ), |
| 172 | ) |
| 173 | print( |
| 174 | "Created IPC-enabled DeviceMemoryResource " |
| 175 | f"(is_ipc_enabled={mr.is_ipc_enabled})" |
| 176 | ) |
| 177 | |
| 178 | buffer = mr.allocate(nbytes, stream=device.default_stream) |
| 179 | try: |
| 180 | # Fill the buffer with a known pattern from the parent side. |
| 181 | arr = cp.from_dlpack(buffer).view(dtype=cp.float32) |
| 182 | arr[:] = cp.arange(N, dtype=cp.float32) + float(parent_seed) |
| 183 | device.sync() |
| 184 | print(f"Parent wrote pattern (first 5 values): {arr[:5].get()}") |
| 185 | |
| 186 | # Launch the child process and hand the buffer over. |
| 187 | q_to_child = mp.Queue() |
| 188 | q_from_child = mp.Queue() |
| 189 | child = mp.Process( |
no test coverage detected