How does pytorch guarantee resources were held when asynchronously launch kernels

HfCloud: <p>I’m working with C++ extension for pytorch, and I have such a question: pytorch asynchronously launch CUDA kernels, suppose we have these codes:</p>
<pre><code class="lang-cpp">__global__ XXXKernel(float *a, float *b) {
         //some computational workload with a, b
}

at::Tensor XXX(at::Tensor a, at::Tensor b) {
        XXXKernel&lt;&lt;&lt;...&gt;&gt;&gt;(a.data_ptr&lt;float&gt;(), b.data_ptr&lt;float&gt;());
        //No cudaStreamSynchronize or cudaDeviceSynchronize here
}

//use pybind11 to bind XXX to a python function ext.XXX
</code></pre>
<p>python code:</p>
<pre><code class="lang-auto">def test() -&gt; torch.Tensor:
    a = torch.randn(10000)
    b = torch.randn(10000)
    return ext.XXX(a, b)

test()
</code></pre>
<p>Because of asynchronously launched kernel, test function may return before XXXKernel finished, is it right? Further question is it possible the variable a, b may be released by GC before XXXKernel finished? The internels  buffers of a, b may be freed before XXXKernel finished?  So, how to guarantee resources were held?</p>

ptrblck: <aside class="quote no-group" data-username="HfCloud" data-post="1" data-topic="177596">
<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/hfcloud/48/59526_2.png" class="avatar"> HfCloud:</div>
<blockquote>
<p>Because of asynchronously launched kernel</p>
</blockquote>
</aside>
<p>The kernel launches aren’t asynchronous but the kernel execution. The launches themselves are executed by the host and in stream-order.</p>
<aside class="quote no-group" data-username="HfCloud" data-post="1" data-topic="177596">
<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/hfcloud/48/59526_2.png" class="avatar"> HfCloud:</div>
<blockquote>
<p>test function may return before XXXKernel finished, is it right?</p>
</blockquote>
</aside>
<p>Yes, the host can return and continue its execution before the GPU work is done.</p>
<aside class="quote no-group" data-username="HfCloud" data-post="1" data-topic="177596">
<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/hfcloud/48/59526_2.png" class="avatar"> HfCloud:</div>
<blockquote>
<p>Further question is it possible the variable a, b may be released by GC before XXXKernel finished?</p>
</blockquote>
</aside>
<p>No, this shouldn’t be possible as the reference count should be increased while the object is used.</p>

HfCloud: <p>OK, thanks. So I should keep reference counts of a, b until the kernel execution finished.</p>
<p>But for the 3rd question I still don’t understand,  When C++ XXX function returns, a, b destruct and decrease reference counts. If python function test returned and python variable a, b were just deleted by GC(I’ m not sure, I think that a, b may be collected by GC after test function returned, or this hypothesis is impossible?) , at this time, the two tensors should have 0 references and free theirs buffers. meanwhile XXX kernel is not finished, this could cause Warp illegal address?</p>
<p>If I want to keep reference counts to tensor a,b until the kernel execution finished, Should I copy at::Tensor variables a,b internally in C++ code and keep them util kernel execution finished? like this:<br>
(BTW: I don’t know how to implement ‘isNotFinished’)</p>
<pre data-code-wrap="C++"><code class="lang-plaintext">

at::Tensor XXX(at::Tensor a, at::Tensor b) {
        XXXKernel&lt;&lt;&lt;...&gt;&gt;&gt;(a.data_ptr&lt;float&gt;(), b.data_ptr&lt;float&gt;());
        //No cudaStreamSynchronize or cudaDeviceSynchronize here

        std::thread([guard_a=a, guard_b=b]() {
                while(isNotFinished(XXXKernel)) 
                        std::this_thread::yield();
                guard_a.~Tensor(); 
                guard_b.~Tensor();
        }).detach();
}
</code></pre>
<p>Really appreciate your help!</p>

HfCloud: <p>Oh, I just learned how to implement ‘isNotFinished’, I can use cudaLaunchHostFunc to add a callback to release guard_a, guard_b after XXXKernel finish execution.</p>