Using CUDA IPC memory handles in pytorch

username: abweiss
cooked: <p>I want to insert a trained pytorch model into the middle of a multi-process pipeline.  The input/output data for the model should never move off the GPU.  Device pointers to the data need to be passed back and forth between processes using CUDA IPC memory handles.</p>
<p>Basically, I need a way to access/create the IPC handles and to convert to/from torch.cuda.*Tensor objects.</p>
<p>What is the best way to implement this?  I know pycuda gives access to CUDA IPC handles (e.g. pycuda.driver.mem_get_ipc_handle), but from my experience pycuda does not play nicely with pytorch.  Are there any other simple solutions in the python realm?</p>

username: colesbury
cooked: <p>You can share CUDA tensors across processes using multiprocessing queues. (e.g. <code>multiprocessing.SimpleQueue</code>) The PyTorch code will create an IPC handle when the tensor is added to the queue and open that handle when the tensor is retrieved from the queue.</p>
<p>Beware that you need to keep the original CUDA tensor alive for at least as long as any view of it is accessible in another process.</p>

username: abweiss
cooked: <p>Thanks for the quick response <a class="mention" href="/u/colesbury">@colesbury</a>.</p>
<p>Just to clarify, the other processes in the pipeline are not python processes (they are C/C++/CUDA).  So it’s important that I can access/create IPC handles with device pointers to the raw underlying tensor data.  My confusion is how to work with these handles within the python/pytorch process.  Correct me if I’m wrong, but it seems that <code>multiprocessing.SimpleQueue</code> will only work the way you describe if both processes are using Pytorch.</p>
<p>So, just to be absolutely clear, the full plan is to use shared memory to pass IPC handles between processes.  For example, the shared memory file will include a 64byte cudaIpcMemHandle_t (containing a pointer to the raw data in GPU memory), plus additional bytes to specify the number of rows and columns in the tensor.</p>

username: colesbury
cooked: <p>It will be a bit tricky to do correctly because small PyTorch storages are packed into the same CUDA allocation block. You will have to rely on implementation details of PyTorch that may change in the future:</p>
<pre><code class="lang-auto">x = torch.randn(100, device='cuda')
storage = x.storage()
device, handle, size, offset, view_size = storage._share_cuda_()
</code></pre>
<p><code>device</code> is the index of the GPU (i.e. 0 for the first GPU)<br>
<code>handle</code> is the <code>cudaIpcMemHandle_t</code> as a Python byte string<br>
<code>size</code> is the size of the allocation (not the Storage!, in elements, not bytes!)<br>
<code>offset</code> is the offset in bytes of the storage data pointer from the CUDA allocation<br>
<code>view_size</code> is the size of the storage (in elements, not bytes!)</p>

username: abweiss
cooked: <p>Thanks again <a class="mention" href="/u/colesbury">@colesbury</a>.</p>
<p>So <code>_share_cuda_()</code> gives me access to the <code>cudaIpcMemHandle_t</code> of an existing <code>torch.cuda.Tensor</code>.  It’s unfortunate that the handle is not exposed through a regular function call, but it’s a good start.</p>
<p>Now, what about when I need to convert the other way around, from handle to tensor?  If I have a <code>cudaIpcMemHandle_t</code>, read in from shared memory and converted to a Python byte string, can I insert that into a <code>torch.cuda.Storage</code> and thereby produce a <code>torch.cuda.Tensor</code> which points to the appropriate data from GPU memory?</p>
<p>Also, can you explain the <code>offset</code> a bit more?  It sounds like multiple different <code>torch.cuda.Storage</code> objects share the same <code>cudaIpcMemHandle_t</code>, but with different offsets in memory.  Is that correct?  I don’t see that as a major problem.  I’ll just have to write the <code>offset</code> to shared memory as well.</p>
<p>Another idea altogether:  What about using Pytorch’s <code>extension-ffi</code> to access the <code>cudaIpcMemHandle_t</code> and storing the data into a <code>THCudaTensor</code>?  I’ve never played with the <code>extension-ffi</code> before, so I don’t really understand it’s capabilities.  I’ll need to make calls to  functions like <code>cudaIpcOpenMemHandle</code>, which are part of CUDAs runtime API.  Is this possible?</p>

username: colesbury
cooked: <p>If you want to go back and forth between C/C++ and Python you probably want to use an extension. You should prefer <a href="https://github.com/pytorch/extension-cpp">https://github.com/pytorch/extension-cpp</a> over <code>extension-ffi</code> as <code>TH/THC</code> is being slowly deprecated and moved into <code>ATen</code>.</p>
<p><code>ATen</code> provides a <code>Type::storageFromBlob</code> function which you can use after you open the IPC handle.</p>
<p>I don’t think there’s an equivalent function in Python. It would probably be good for us to add something like that.</p>

username: abweiss
cooked: <p><a class="mention" href="/u/colesbury">@colesbury</a> Thanks so much for all the help on this.  I think I’m almost there.</p>
<p>I’ve been playing around with <code>extension-cpp</code> and I’m running into a couple of issues.</p>
<p>As a reference point, I am mostly following the <code>extension-cpp</code> tutorial here:<br>
<a href="https://pytorch.org/tutorials/advanced/cpp_extension.html#writing-a-mixed-c-cuda-extension" class="onebox" target="_blank" rel="nofollow noopener">https://pytorch.org/tutorials/advanced/cpp_extension.html#writing-a-mixed-c-cuda-extension</a></p>
<p>So I have three files, a <code>.py</code>, a <code>.cpp</code>, and a <code>.cu</code>.  I am using the the JIT method for compiling my extension.</p>
<p>In the <code>.cu</code> file, I am using the CUDA runtime API to extract a float* device pointer from a <code>cudaIpcMemHandle</code>.  I am then using <code>tensorFromBlob</code> to fill an <code>at::Tensor</code> object.  Here is how I am using <code>tensorFromBlob</code>:</p>
<p><code>at::Tensor cuda_tensor_from_shm = at::CUDA(at::kFloat).tensorFromBlob(d_img, {rows,cols});</code></p>
<p>My first problem is that the above line of code takes about three seconds to execute.  Does it only take so long the first time I call the extension, or is it going to be slow every time?  Obviously the whole point of using shared memory and CUDA IPC handles was to make the cost of transferring data negligibly small; I was hoping for sub-millisecond times.</p>
<p>The second problem is that I get a segmentation fault happening at some point between the <code>.cpp</code> code and the <code>.py</code> code.  I haven’t precisely pinpointed it yet.  However, my guess is that after calling <code>tensorToBlob</code>, I need to copy the data to a new <code>at::Tensor</code> before I can use it in Pytorch.  Is that correct?  If so, is there a super-fast ATen device-to-device copy I can use?</p>

username: abweiss
cooked: <p>Everything works after modifying my <code>tensorFromBlob</code> code from:</p>
<p><code>at::Tensor cuda_tensor_from_shm = at::CUDA(at::kFloat).tensorFromBlob(d_img, {rows,cols});</code></p>
<p>to:<br>
<code>at::Tensor cuda_tensor_from_shm = torch::CUDA(at::kFloat).tensorFromBlob(d_img, {rows,cols});</code></p>
<p>I’ll need to dig into the code to understand why <code>torch::CUDA</code> is the correct scoping, but anyway it works.</p>

username: AmirFB
cooked: <p>How can I do the opposite in Python? I have transferred an IPC handle to a python code, but I do now know how to convert it into a torch tensor.</p>