Accelerating layers with Numba

malmaud: <p>Hi all,<br>
Has anyone had success integrating functions decorated with <code>numba.cuda.jit</code> into a PyTorch workflow? <a href="http://pytorch.org/about/" rel="nofollow noopener">http://pytorch.org/about/</a> mentions using your ‘favorite libraries…such as Numba’ but I’m not clear on how to actually do that.<br>
Of course it’s no problem on the CPU to decorate a layer with <code>numba.jit</code> and convert inputs and outputs to and from numpy arrays, although even that doesn’t seem particularly elegant because of all the  numpy/Tensor conversions required (I’m aware they are zero-overhead, though.)</p>

tom: <pre><code class="lang-auto">from numba import cuda
os.environ['NUMBAPRO_LIBDEVICE']='/usr/lib/nvidia-cuda-toolkit/libdevice/'
os.environ['NUMBAPRO_NVVM']='/usr/lib/x86_64-linux-gnu/libnvvm.so.3.1.0'


@cuda.jit('(float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], int32, int32, int32)')
def cu_something(A, c, d, u, v, b, n, m):
   ...

def get_devicendarray(t):
    assert t.type() == 'torch.cuda.FloatTensor'
    ctx = cuda.cudadrv.driver.driver.get_context()
    mp = cuda.cudadrv.driver.MemoryPointer(ctx, ctypes.c_ulong(t.data_ptr()), t.numel()*4)
    return cuda.cudadrv.devicearray.DeviceNDArray(t.size(), [i*4 for i in t.stride()], numpy.dtype('float32'), 
                                                  gpu_data=mp, stream=torch.cuda.current_stream().cuda_stream)
</code></pre>
<p>and then later used it with <code>torch.cuda.FloatTensors</code> A, target, c,d,u,v.</p>
<pre><code class="lang-auto">Ad,targetd,cd,dd,ud,vd = (get_devicendarray(x) for x in (A,target,c,d,u,v))
cu_something[((b-1)//BLOCK+1,(n-1)//BLOCK+1),(BLOCK,BLOCK)](Ad,targetd,cd,dd,ud,vd,b,n,m)
</code></pre>
<p>It seemed to work but <strong>I was warned by people who know better than me that it was dangerous and potentially unstable</strong> (e.g. is the context above the right one).<br>
I have not done too much with it.</p>
<p>Best regards</p>
<p>Thomas</p>

Adam_Dziedzic: <p>Thank you <a class="mention" href="/u/tom">@tom</a> This works like a charm for me, even without the explicit os.environ. Suggestion: <a class="mention" href="/u/smth">@smth</a> It would be great if PyTorch supported officially the conversion of a torch tensor to the DeviceNDArray (Numba gpu resident numpy array). It could be named: torch.Tensor.numba()  and would not require some complicated, not supported code, and would be done without the additional data movement from gpu to cpu and back to gpu.</p>

tom: <p>There is a great PR open for numba integration:<br>
<aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/pull/11984" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
      <a href="https://github.com/asford" target="_blank" rel="nofollow noopener">
    <img alt="asford" src="https://avatars1.githubusercontent.com/u/282792?v=4" class="thumbnail onebox-avatar" width="90" height="90">
  </a>

<h4>
  <a href="https://github.com/pytorch/pytorch/pull/11984" target="_blank" rel="nofollow noopener">Implement Tensor.__cuda_array_interface__.</a>
</h4>

<div class="date">
  by <a href="https://github.com/asford" target="_blank" rel="nofollow noopener">asford</a>
  on <a href="https://github.com/pytorch/pytorch/pull/11984" target="_blank" rel="nofollow noopener">08:59PM - 22 Sep 18</a>
</div>

<div class="github-commit-stats">
  <strong>14 commits</strong>
  changed <strong>6 files</strong>
  with <strong>332 additions</strong>
  and <strong>1 deletions</strong>.
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>

Etienne_Perot: <aside class="quote no-group" data-username="tom" data-post="4" data-topic="13923">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>or numba inte</p>
</blockquote>
</aside>
<p>Hello, using the <strong>cuda_array_interface</strong>, how can i build the numba array?</p>

tom: <p><code>numba_cuda_array = numba.cuda.as_cuda_array(cuda_tensor)</code></p>
<p>The <a href="https://github.com/pytorch/pytorch/blob/master/test/test_numba_integration.py" rel="nofollow noopener">Numba integration test script</a> probably is a good source of information what works with PyTorch and Numba.</p>
<p>Best regards</p>
<p>Thomas</p>

Etienne_Perot: <p>thanks Thomas V! Just another one I don’t find in the test script, regarding cuda streams, can we build/ communicate a numba.cuda.stream to pytorch or the opposite?</p>

tom: <p>You could try to use the torch <code>Stream</code> object’s <code>cuda_stream</code> property and feed that to <code>numba.cuda.external_stream</code>, but I didn’t try.</p>

Peter_Featherstone: <p>Does this work with autograd ? I want to be able to insert a numba accelerated operation inside my model and have backpropagation work as normal. Is this possible? Cheers</p>

ptrblck: <p>No, since Autograd won’t be able to understand 3rd party library operations. You would thus need to implement a custom <code>autograd.Function</code> as described <a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html">here</a> including the <code>backward</code> pass.</p>