Replacing torch.zeros internals with cudaMemset instead of fill kernel

username: anguswang
cooked: <p>I’ve noticed that a fill kernel is called when constructing a tensor on the GPU with <code>torch.zeros</code>. Would it be possible and/or preferable to use <code>cudaMemset</code> instead? It seems like the more intuitive and efficient choice to me. If this is not suitable, why?</p>
<p><code>torch.zeros</code> first calls <code>cudaMalloc</code> to allocate memory on the device and then calls <a href="https://github.com/pytorch/pytorch/blob/4bae7ae3d9a204c3ea2c6c801e4aedde9331be6a/aten/src/ATen/native/cuda/FillKernel.cu#L22" rel="noopener nofollow ugc"><code>fill_kernel_cuda</code></a>. From my understanding, it should be possible to just replace <code>fill_kernel_cuda</code> with a suitable implementation calling <code>cudaMemset</code>.</p>

username: ptrblck
cooked: <p>Are you seeing any benefits in using <code>cudaMemset</code> instead of the <code>TensorIterator</code> approach we are using, i.e. did you compare both approaches?</p>

username: anguswang
cooked: <p>Sorry for the late response. I have done some basic profiling using the PyTorch Profiler. It looks like changing the implementation to using <code>cudaMemset</code> does improve the performance of <code>torch.zeros</code>.</p>
<p>I profiled the creation of a fp16 zeros tensor of shape (1024, 1024, 1024) on the GPU as shown below:</p>
<pre><code class="lang-auto">import torch
from torch.profiler import profile, record_function, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], 
             record_shapes=True, 
             profile_memory=True, 
             use_cuda=True) as prof:
    with record_function("torch_zeros"):
        zero_tensor = torch.zeros((1024, 1024, 1024),
                                  dtype=torch.float16,
                                  pin_memory=False,
                                  device='cuda')

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
</code></pre>
<p>I got the following results for the current implementation:</p>
<pre><code class="lang-auto">-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                            torch_zeros         6.35%      19.267ms        99.19%     300.836ms     300.836ms       0.000us         0.00%       2.798ms       2.798ms           0 b           0 b       2.00 Gb           0 b             1  
                                            aten::zeros         0.04%     130.736us        92.18%     279.572ms     279.572ms       0.000us         0.00%       2.798ms       2.798ms           0 b           0 b       2.00 Gb           0 b             1  
                                            aten::zero_         0.01%      38.260us         2.17%       6.594ms       6.594ms       0.000us         0.00%       2.798ms       2.798ms           0 b           0 b           0 b           0 b             1  
                                            aten::fill_         0.02%      57.585us         2.16%       6.556ms       6.556ms       2.798ms       100.00%       2.798ms       2.798ms           0 b           0 b           0 b           0 b             1  
void at::native::vectorized_elementwise_kernel&lt;4, at...         0.00%       0.000us         0.00%       0.000us       0.000us       2.798ms       100.00%       2.798ms       2.798ms           0 b           0 b           0 b           0 b             1  
                                            torch_zeros         0.00%       0.000us         0.00%       0.000us       0.000us       2.798ms       100.00%       2.798ms       2.798ms           0 b           0 b           0 b           0 b             1  
                                     cudaGetDeviceCount         0.00%       1.407us         0.00%       1.407us       0.703us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             2  
                             cudaGetDeviceProperties_v2         0.66%       1.996ms         0.66%       1.996ms       1.996ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             1  
                                            aten::empty         0.04%     108.655us        89.96%     272.847ms     272.847ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b       2.00 Gb       2.00 Gb             1  
                       cudaDeviceGetStreamPriorityRange        89.84%     272.472ms        89.84%     272.472ms     272.472ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             1  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 303.297ms
Self CUDA time total: 2.798ms
</code></pre>
<p>And the following results if I swap the fill kernel with a <code>cudaMemset</code> call:</p>
<pre><code class="lang-auto">------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg       CPU Mem  Self CPU Mem      CUDA Mem  Self CUDA Mem    # of Calls  
------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                         torch_zeros         6.61%      19.356ms        99.57%     291.549ms     291.549ms       0.000us         0.00%       1.579ms       1.579ms           0 b           0 b       2.00 Gb           0 b             1  
                         aten::zeros         0.05%     133.432us        92.29%     270.246ms     270.246ms       0.000us         0.00%       1.579ms       1.579ms           0 b           0 b       2.00 Gb           0 b             1  
                         aten::zero_         0.01%      38.235us         0.04%     102.532us     102.532us       0.000us         0.00%       1.579ms       1.579ms           0 b           0 b           0 b           0 b             1  
                         aten::fill_         0.01%      39.738us         0.02%      64.297us      64.297us       1.579ms       100.00%       1.579ms       1.579ms           0 b           0 b           0 b           0 b             1  
                     Memset (Device)         0.00%       0.000us         0.00%       0.000us       0.000us       1.579ms       100.00%       1.579ms       1.579ms           0 b           0 b           0 b           0 b             1  
                         torch_zeros         0.00%       0.000us         0.00%       0.000us       0.000us       1.579ms       100.00%       1.579ms       1.579ms           0 b           0 b           0 b           0 b             1  
                  cudaGetDeviceCount         0.00%       1.285us         0.00%       1.285us       0.643us       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             2  
          cudaGetDeviceProperties_v2         0.66%       1.946ms         0.66%       1.946ms       1.946ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             1  
                         aten::empty         0.04%     120.588us        92.21%     270.010ms     270.010ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b       2.00 Gb       2.00 Gb             1  
    cudaDeviceGetStreamPriorityRange        92.08%     269.620ms        92.08%     269.620ms     269.620ms       0.000us         0.00%       0.000us       0.000us           0 b           0 b           0 b           0 b             1  
------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 292.815ms
Self CUDA time total: 1.579ms
</code></pre>
<p>The time saved is not very significant, but it is still faster and in my opinion, using <code>cudaMemset</code> seems a lot more straightforward and intuitive here.</p>