(*shape, dtype=torch.float32, device=FIRST_CUDA_DEVICE)
| 89 | |
| 90 | |
| 91 | def get_paged(*shape, dtype=torch.float32, device=FIRST_CUDA_DEVICE): |
| 92 | num_bytes = dtype.itemsize * prod(shape) |
| 93 | managed_ptr = lib.cget_managed_ptr(ct.c_size_t(num_bytes)) |
| 94 | c_ptr = ct.cast(managed_ptr, ct.POINTER(ct.c_int)) |
| 95 | new_array = np.ctypeslib.as_array(c_ptr, shape=shape) |
| 96 | out = torch.frombuffer(new_array, dtype=dtype, count=prod(shape)).view(shape) |
| 97 | out.is_paged = True |
| 98 | out.page_deviceid = device.index |
| 99 | return out |
| 100 | |
| 101 | |
| 102 | def prefetch_tensor(A: torch.Tensor, to_cpu=False): |