Torchaudio.functional.lfilter runs very slow on gpu, but fast on cpu

username: knitvoger
cooked: <p>As the title said, lfilter on cpu cost about 1 ms to run. But on the gpu, it is 1000x slower!  I posted my code and test result below. Could anyone tell me why and how to accelerate the speed on gpu?</p>
<pre><code class="lang-auto">device="cuda"
waveform = torch.rand([16000]).to(device)
a = torch.tensor([1.0, 0.0]).to(device)
b = torch.tensor([1.0, -0.97]).to(device)
for i in range(5):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    waveform = torchaudio.functional.lfilter(waveform, a, b)
    end.record() 
    torch.cuda.synchronize()
    print("run on %s cost %.3f ms" % (device, start.elapsed_time(end)))
</code></pre>
<p>Test result</p>
<pre><code class="lang-auto">run on cpu cost 1.473 ms
run on cpu cost 1.531 ms
run on cpu cost 0.905 ms
run on cpu cost 0.774 ms
run on cpu cost 1.007 ms
run on cuda cost 961.567 ms
run on cuda cost 955.971 ms
run on cuda cost 962.749 ms
run on cuda cost 957.605 ms
run on cuda cost 965.437 ms
</code></pre>
<p>My env is v100 gpu, cuda 11, torch ‘1.10.1+cu113’ and torchaudio ‘0.10.1+cu113’</p>
<p>Thank you very much!</p>

username: ptrblck
cooked: <p>If I’m not mistaken, <a href="https://github.com/pytorch/audio/blob/a525abbcbe7858ec22902666e90a5bbde5fefdce/torchaudio/csrc/lfilter.cpp#L72-L97">this code</a> will be used, which will launch a lot of GPU kernels (<code>nsys</code> shows 160045 <code>cudaLaunchKernel</code> calls) which should create a lot of CPU overhead.<br>
To reduce the overhead caused by the kernel launches, you could try to use CUDA graphs as described <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-graphs">here</a> and see if it would match your use case.</p>
<p>Example:</p>
<pre><code class="lang-python">import torch
import torchaudio
import time

device = 'cuda'
nb_iters = 100

# Placeholder input used for capture
waveform = torch.rand([16000]).to(device)
a = torch.tensor([1.0, 0.0]).to(device)
b = torch.tensor([1.0, -0.97]).to(device)

# warmup
for _ in range(10):
    waveform = torchaudio.functional.lfilter(waveform, a, b)

# profile
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(nb_iters):
    waveform = torchaudio.functional.lfilter(waveform, a, b)
torch.cuda.synchronize()
t1 = time.perf_counter()

print('Eager, {}s/iter'.format((t1 - t0)/nb_iters))

# CUDA graphs

g = torch.cuda.CUDAGraph()

# Warmup before capture
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for _ in range(3):
        waveform = torchaudio.functional.lfilter(waveform, a, b)
torch.cuda.current_stream().wait_stream(s)

# Captures the graph
# To allow capture, automatically sets a side stream as the current stream in the context
with torch.cuda.graph(g):
    waveform = torchaudio.functional.lfilter(waveform, a, b)

torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(nb_iters):
    g.replay()
torch.cuda.synchronize()
t1 = time.perf_counter()

print('CUDA graphs, {}s/iter'.format((t1 - t0)/nb_iters))
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">Eager, 0.6360926309000934s/iter
CUDA graphs, 0.07975950390042272s/iter
</code></pre>