feat: use parallel scan in certain conditions

DiffAPF · Jul 10, 2024 · 58e5334 · 58e5334
1 parent d9c3e1e
commit 58e5334
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/torchlpc/__init__.py b/torchlpc/__init__.py
@@ -2,6 +2,8 @@
 from typing import Optional
 
 from .core import LPC
+from .parallel_scan import WARPSIZE
+from .recurrence import RecurrenceCUDA
 
 __all__ = ["sample_wise_lpc"]
 
@@ -35,4 +37,7 @@ def sample_wise_lpc(
     else:
         assert zi.shape == (B, order)
 
+    if order == 1 and x.is_cuda and B * WARPSIZE < T:
+        return RecurrenceCUDA.apply(-a.squeeze(2), x, zi.squeeze(1))
+
     return LPC.apply(x, a, zi)