How to jit compile with `cupy.cuda.compile_with_cache`

tommy19970714: <p>I’m trying to compile <a href="https://github.com/geyuying/PF-AFN.git" rel="noopener nofollow ugc">PF-AFN</a> model to jit.<br>
I got the following error when I compile the model with <code>cupy.cuda.compile_with_cache</code> to jit.</p>
<pre><code class="lang-auto">NVRTCError: NVRTC_ERROR_COMPILATION (6)

During handling of the above exception, another exception occurred:

CompileException                          Traceback (most recent call last)

cupy/util.pyx in cupy.util.memoize.decorator.ret()

/usr/local/lib/python3.7/dist-packages/cupy/cuda/compiler.py in compile(self, options)
    440         except nvrtc.NVRTCError:
    441             log = nvrtc.getProgramLog(self.ptr)
--&gt; 442             raise CompileException(log, self.src, self.name, options, 'nvrtc')
    443 
    444 

CompileException: /tmp/tmpan1ut480/3b7c153ce98d06488f1cbac8793f6dff_2.cubin.cu(16): error: identifier "tensor" is undefined

1 error detected in the compilation of "/tmp/tmpan1ut480/3b7c153ce98d06488f1cbac8793f6dff_2.cubin.cu".
</code></pre>
<h2>
<a name="to-reproduce-1" class="anchor" href="#to-reproduce-1"></a>To Reproduce</h2>
<p>This is a colab to reproduce the error.</p><aside class="onebox allowlistedgeneric">
  <header class="source">
      <img src="https://colab.research.google.com/img/favicon.ico?vrz=colab-20210517-060154-RC00_374160454" class="site-icon" width="16" height="16">
      <a href="https://colab.research.google.com/drive/1WDRCN6wPIAsl5tBFKfne0ABN49estM9P?usp=sharing" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>
  <article class="onebox-body">
    <img src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar" width="60" height="60">

<h3><a href="https://colab.research.google.com/drive/1WDRCN6wPIAsl5tBFKfne0ABN49estM9P?usp=sharing" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



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

<p>This is a minimum code.</p>
<pre><code class="lang-auto">@cupy.util.memoize(for_each_device=True)
def cupy_launch(strFunction, strKernel):
	return cupy.cuda.compile_with_cache(strKernel).get_function(strFunction)

kernel_Correlation_rearrange = " .... "

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

    def forward(self, x_warp_after, x_cond):
        cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
          'intStride': 1,
          'input': x_warp_after,
          'output': x_cond
        }))(
        )
        return x_warp_after, x_cond

net = Net().cuda()
input1 = torch.randn([1, 256, 8, 6]).cuda()
input2 = torch.randn([1, 256, 8, 6]).cuda()
trace_model = torch.jit.trace(net, [input1, input2])
</code></pre>
<h2>
<a name="expected-behavior-2" class="anchor" href="#expected-behavior-2"></a>Expected behavior</h2>
<p>I think the above error occurs when I use cupy.cuda.compile_with_cache.</p>
<h2>
<a name="environment-3" class="anchor" href="#environment-3"></a>Environment</h2>
<ul>
<li>PyTorch Version (e.g., 1.0): 1.8.1+cu101</li>
<li>OS (e.g., Linux): Ubuntu 18.04.5 LTS (x86_64)</li>
<li>How you installed PyTorch (<code>conda</code>, <code>pip</code>, source): pip</li>
<li>Build command you used (if compiling from source): no</li>
<li>Python version: 3.7 (64-bit runtime)</li>
<li>CUDA/cuDNN version: 11.0.221</li>
<li>GPU models and configuration: GPU 0: Tesla T4</li>
<li>Any other relevant information:</li>
</ul>

ptrblck: <p>Based on the error message it seems that <code>cupy</code> is unable to compile the PyTorch methods and I’m unsure if this would even be supported.<br>
Do you have any resources claiming that this should work and some examples demonstrating it?</p>

tommy19970714: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thank you for your reply!<br>
You can check this example from the following colab.</p><aside class="onebox allowlistedgeneric">
  <header class="source">
      <img src="https://colab.research.google.com/img/favicon.ico?vrz=colab-20210518-060121-RC00_374369754" class="site-icon" width="16" height="16">
      <a href="https://colab.research.google.com/drive/1WDRCN6wPIAsl5tBFKfne0ABN49estM9P?usp=sharing" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>
  <article class="onebox-body">
    <img src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar" width="60" height="60">

<h3><a href="https://colab.research.google.com/drive/1WDRCN6wPIAsl5tBFKfne0ABN49estM9P?usp=sharing" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



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


ptrblck: <p>Thanks for the code! It shows the error, which might be helpful for debugging, but my previous question was regarding the expectations that this would be supported. Do you have any working example, demos, blog posts etc., which explain how <code>cupy</code> can be used, as I’m unfamiliar with it?</p>