Compute sin and cos simultaneously on the same input

vadimkantorov: <p>A pattern found in the wild at: <a href="https://github.com/chenhaoxing/DiffusionInst/blob/main/diffusioninst/head.py#L70" class="inline-onebox" rel="noopener nofollow ugc">DiffusionInst/head.py at main · chenhaoxing/DiffusionInst · GitHub</a><br>
<code>embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)</code>.</p>
<p>This is often found in positional encoding computation / maybe some random fourier features.</p>
<p>Is there a nice fused fast way of doing this in PyTorch?</p>
<p>This is equivalent to taking complex exponential (but could be a feature request for a separate alias as <code>torch.sincos(x)</code>), but in NumPy at least it’s a bit slower than computing sin and cos separately (despite that some sincos_ps SSE recipe exists) <a href="https://stackoverflow.com/questions/32397347/is-there-a-fast-way-to-return-sin-and-cos-of-the-same-value-in-python" class="inline-onebox" rel="noopener nofollow ugc">numpy - Is there a fast Way to return Sin and Cos of the same value in Python? - Stack Overflow</a> / <a href="http://gruntthepeon.free.fr/ssemath/sse_mathfun.h" rel="noopener nofollow ugc">http://gruntthepeon.free.fr/ssemath/sse_mathfun.h</a></p>
<p>Ideally, this pattern can also be fused in the future.</p>

ptrblck: <p>That’s an interesting use case and yes, it should be possible to write a fused version of this operation.<br>
One approach would be to use the nvFuser Python frontend, which is available in the latest nightly releases.<br>
Here is a complete code snippet showing the fused kernel with the limitation of the <code>torch.cat</code> operation. This op is currently not supported and still WIP, so you would still need to execute an additional eager kernel.</p>
<pre><code class="lang-python">import torch
from torch.profiler import profile, record_function, ProfilerActivity
from torch._C._nvfuser import Fusion, FusionCache, FusionDefinition, DataType
from typing import List


def sincos_eager(embeddings):
    return torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)


def nvFuser_sincos(fd: FusionDefinition):
    embeddings = fd.define_tensor(symbolic_sizes=[-1, -1], contiguous=[True, True], dtype=DataType.Float)
    output1 = fd.ops.sin(embeddings)
    output2 = fd.ops.cos(embeddings)
    #output = fd.ops.cat((output1, output2), dim=dim)
    fd.add_output(output1)
    fd.add_output(output2)


sincos = Fusion()
with FusionDefinition(sincos) as fd:
    nvFuser_sincos(fd)


def python_frontend_sincos(
    fs: Fusion, input: torch.Tensor
) -&gt; List[torch.Tensor]:
    out = fs.execute([input])
    return out

# setup
embeddings = torch.randn(1024, 1024, device='cuda')

# compare
out_pynvf = python_frontend_sincos(sincos, embeddings)
out_pynvf = torch.cat(out_pynvf, -1)
out_eager = sincos_eager(embeddings)

print((out_pynvf - out_eager).abs().max())
# tensor(0., device='cuda:0')

# profile
with profile(activities=[ProfilerActivity.CUDA], record_shapes=True) as prof:
    with record_function("embedding_eager"):
        sincos_eager(embeddings)
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                                                    Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                                        cudaLaunchKernel        88.10%     348.000us        88.10%     348.000us     116.000us       0.000us         0.00%       0.000us       0.000us             3  
#                                   cudaDeviceSynchronize        11.90%      47.000us        11.90%      47.000us      47.000us       0.000us         0.00%       0.000us       0.000us             1  
# void at::native::vectorized_elementwise_kernel&lt;4, at...         0.00%       0.000us         0.00%       0.000us       0.000us      12.000us        24.49%      12.000us      12.000us             1  
# void at::native::vectorized_elementwise_kernel&lt;4, at...         0.00%       0.000us         0.00%       0.000us       0.000us      12.000us        24.49%      12.000us      12.000us             1  
# void at::native::(anonymous namespace)::CatArrayBatc...         0.00%       0.000us         0.00%       0.000us       0.000us      25.000us        51.02%      25.000us      25.000us             1  
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  

with profile(activities=[ProfilerActivity.CUDA], record_shapes=True) as prof:
    with record_function("embedding_nvFuser"):
        python_frontend_sincos(sincos, embeddings)
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                                                    Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                                   cudaDeviceSynchronize       100.00%      18.000us       100.00%      18.000us      18.000us       0.000us         0.00%       0.000us       0.000us             1  
# CudaCodeGen::kernel9(CudaCodeGen::Tensor&lt;float, 2&gt;, ...         0.00%       0.000us         0.00%       0.000us       0.000us      17.000us       100.00%      17.000us      17.000us             1  
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
</code></pre>

vadimkantorov: <p>Created a follow-up discussion at <a href="https://github.com/pytorch/pytorch/issues/90559" class="inline-onebox" rel="noopener nofollow ugc">[feature request, idea] Fused torch.sincos(x) or cossin(x) - somewhat complex exponential · Issue #90559 · pytorch/pytorch · GitHub</a></p>