Skip to content

Instantly share code, notes, and snippets.

@kayush2O6
Last active November 12, 2022 15:37
Show Gist options
  • Save kayush2O6/202875ef816025653980bd1d6e5a0071 to your computer and use it in GitHub Desktop.
Save kayush2O6/202875ef816025653980bd1d6e5a0071 to your computer and use it in GitHub Desktop.
convert numba cuda array to pytorch tensor
from numba import cuda
import ctypes
import numpy as np
import torch
def devndarray2torch(dev_arr):
t = torch.empty(size=dev_arr.shape, dtype=dtyp).cuda()
ctx = cuda.cudadrv.driver.driver.get_context()
# constant value of #bytes in case of float32 = 4
mp = cuda.cudadrv.driver.MemoryPointer(ctx, ctypes.c_ulong(t.data_ptr()), t.numel()*4)
tmp_arr = cuda.cudadrv.devicearray.DeviceNDArray(t.size(), [i*4 for i in t.stride()], np.dtype('float32'),
gpu_data=mp, stream=torch.cuda.current_stream().cuda_stream)
# To verify whether the data pointer is same or not.
# print(tmp_arr.__cuda_array_interface__)
# print(dev_arr.__cuda_array_interface__)
tmp_arr.copy_to_device(dev_arr)
return t
d_arr = cuda.to_device(np.array([[10,20,30],[40,50,60.0]], dtype=np.float32))
tensor = devndarray2tensor(d_arr)
print(tensor)
@kayush2O6
Copy link
Author

your code is running as it is, in colab. here is the running notebook: https://colab.research.google.com/drive/1R9V8qNo2qj-yUbuNnq8IaEkoB1pw-oUP

you are running on windows, I haven't tested on windows machines.

@kayush2O6
Copy link
Author

what is the output of these lines??

print(tmp_arr.__cuda_array_interface__)
print(dev_arr.__cuda_array_interface__)

@BarakChamo
Copy link

Huh interesting, I'm running on Windows with PyTorch 1.2.0, numba 0.45.1, Python 3.7.4

The output is:

{'shape': (2, 3), 'strides': (12, 4), 'data': (37748736, False), 'typestr': '<f4', 'version': 1}
{'shape': (2, 3), 'strides': (12, 4), 'data': (21510488064, False), 'typestr': '<f4', 'version': 1}

@mantouRobot
Copy link

I also encountered the same problem running on Windows.

@mantouRobot
Copy link

@AK-ayush, @BarakChamo
any ideas?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment