Usage of __restrict__ in CUDA

username: LukasKoestler
cooked: <p>Dear PyTorch Team,</p>
<p>thank you very much for the library, which is very pleasant to use and has great documentation!</p>
<p>In the CUDA example, you use <code>const torch::PackedTensorAccessor32&lt;scalar_t,3,torch::RestrictPtrTraits&gt;</code> when one would use <code>__restrict__ scalar_t*</code> without the accessor. This seems like a very nice solution that I would also want to use within some CUDA code. However, if I read <a href="https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/core/TensorAccessor.h" rel="noopener nofollow ugc">TensorAccessor.h</a> correctly, the restricted pointer is the member of a <code>class</code> and according to <a href="https://forums.developer.nvidia.com/t/restrict-seems-to-be-ignored-for-base-pointers-in-structs-having-base-pointers-with-restrict-as-kernel-arguments-directly-works-as-expected/154020" rel="noopener nofollow ugc">this discussion</a> in the Nvidia forums that is ignored. To validate this, I made this little <a href="https://godbolt.org/z/9c3z8GfT8" rel="noopener nofollow ugc">example</a> in the compiler explorer, which would show that the restrict is ignored. Could you maybe comment on this, as most likely I am missing something here?</p>
<p>Thank you<br>
Lukas</p>
<p><strong>Compiler Explorer Code</strong> (This looks much nicer on the compiler explorer website)</p>
<pre><code class="lang-cpp">// Some code copied from https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/core/TensorAccessor.h

// The PtrTraits argument to the TensorAccessor/GenericPackedTensorAccessor
// is used to enable the __restrict__ keyword/modifier for the data
// passed to cuda.
template &lt;typename T&gt;
struct DefaultPtrTraits {
  typedef T* PtrType;
};

template &lt;typename T&gt;
struct RestrictPtrTraits {
  typedef T* __restrict__ PtrType;
};

class wrapper{
public:
    RestrictPtrTraits&lt;double&gt;::PtrType ptr;
};

__global__ void assign(DefaultPtrTraits&lt;double&gt;::PtrType input, DefaultPtrTraits&lt;double&gt;::PtrType output) {
    int tid = blockIdx.x;
    {
        double la = input[tid];
        output[tid] = la;
        la = input[tid];
        output[tid] = la;
    }
        
}

__global__ void assign_restrict_direct(RestrictPtrTraits&lt;double&gt;::PtrType input, RestrictPtrTraits&lt;double&gt;::PtrType output) {
    int tid = blockIdx.x;
    {
        double la = input[tid];
        output[tid] = la;
        la = input[tid];
        output[tid] = la;
    }
        
}


__global__ void assign_restrict_struct(wrapper input, wrapper output) {
    int tid = blockIdx.x;
    {
        double la = input.ptr[tid];
        output.ptr[tid] = la;
        la = input.ptr[tid];
        output.ptr[tid] = la;
    }
        
}
</code></pre>
<p><strong>Compiler Explorer Assembly</strong> (NVCC 11.5.0 sm_52)</p>
<p>Note that <code>_Z22assign_restrict_struct7wrapperS_</code> uses two loads and stores and hence doesn’t exploit the <code>__restrict__</code> like <code>_Z22assign_restrict_directPdS_</code>.</p>
<pre data-code-wrap="assembly"><code class="lang-nohighlight">.visible .entry _Z6assignPdS_(
        .param .u64 _Z6assignPdS__param_0,
        .param .u64 _Z6assignPdS__param_1
)
{

        ld.param.u64    %rd1, [_Z6assignPdS__param_0];
        ld.param.u64    %rd2, [_Z6assignPdS__param_1];
        cvta.to.global.u64      %rd3, %rd2;
        cvta.to.global.u64      %rd4, %rd1;
        mov.u32         %r1, %ctaid.x;
        mul.wide.s32    %rd5, %r1, 8;
        add.s64         %rd6, %rd4, %rd5;
        ld.global.f64   %fd1, [%rd6];
        add.s64         %rd7, %rd3, %rd5;
        st.global.f64   [%rd7], %fd1;
        ld.global.f64   %fd2, [%rd6];
        st.global.f64   [%rd7], %fd2;
        ret;

}
.visible .entry _Z22assign_restrict_directPdS_(
        .param .u64 _Z22assign_restrict_directPdS__param_0,
        .param .u64 _Z22assign_restrict_directPdS__param_1
)
{

        ld.param.u64    %rd1, [_Z22assign_restrict_directPdS__param_0];
        ld.param.u64    %rd2, [_Z22assign_restrict_directPdS__param_1];
        cvta.to.global.u64      %rd3, %rd2;
        cvta.to.global.u64      %rd4, %rd1;
        mov.u32         %r1, %ctaid.x;
        mul.wide.s32    %rd5, %r1, 8;
        add.s64         %rd6, %rd4, %rd5;
        ld.global.nc.f64        %fd1, [%rd6];
        add.s64         %rd7, %rd3, %rd5;
        st.global.f64   [%rd7], %fd1;
        ret;

}
.visible .entry _Z22assign_restrict_struct7wrapperS_(
        .param .align 8 .b8 _Z22assign_restrict_struct7wrapperS__param_0[8],
        .param .align 8 .b8 _Z22assign_restrict_struct7wrapperS__param_1[8]
)
{

        ld.param.u64    %rd1, [_Z22assign_restrict_struct7wrapperS__param_1];
        ld.param.u64    %rd2, [_Z22assign_restrict_struct7wrapperS__param_0];
        cvta.to.global.u64      %rd3, %rd2;
        cvta.to.global.u64      %rd4, %rd1;
        mov.u32         %r1, %ctaid.x;
        mul.wide.s32    %rd5, %r1, 8;
        add.s64         %rd6, %rd3, %rd5;
        ld.global.f64   %fd1, [%rd6];
        add.s64         %rd7, %rd4, %rd5;
        st.global.f64   [%rd7], %fd1;
        ld.global.f64   %fd2, [%rd6];
        st.global.f64   [%rd7], %fd2;
        ret;

}
</code></pre>