Remove autograd from FIR

Author: NicolasHug

src/libtorchaudio/lfilter.cpp

@@ -182,73 +182,60 @@ class DifferentiableIIR : public torch::autograd::Function<DifferentiableIIR> {
   }
 };
 
-class DifferentiableFIR : public torch::autograd::Function<DifferentiableFIR> {
- public:
-  static torch::Tensor forward(
-      torch::autograd::AutogradContext* ctx,
-      const torch::Tensor& waveform,
-      const torch::Tensor& b_coeffs) {
-    int64_t n_order = b_coeffs.size(1);
-    int64_t n_channel = b_coeffs.size(0);
-
-    namespace F = torch::nn::functional;
-    auto b_coeff_flipped = b_coeffs.flip(1).contiguous();
-    auto padded_waveform =
-        F::pad(waveform, F::PadFuncOptions({n_order - 1, 0}));
-
-    auto output = F::conv1d(
-        padded_waveform,
-        b_coeff_flipped.unsqueeze(1),
-        F::Conv1dFuncOptions().groups(n_channel));
-
-    ctx->save_for_backward({waveform, b_coeffs, output});
-    return output;
-  }
-
-  static torch::autograd::tensor_list backward(
-      torch::autograd::AutogradContext* ctx,
-      torch::autograd::tensor_list grad_outputs) {
-    auto saved = ctx->get_saved_variables();
-    auto x = saved[0];
-    auto b_coeffs = saved[1];
-    auto y = saved[2];
-
-    int64_t n_batch = x.size(0);
-    int64_t n_channel = x.size(1);
-    int64_t n_order = b_coeffs.size(1);
+// FIR filter forward and backward functions (no autograd inheritance)
+torch::Tensor fir_forward(
+    const torch::Tensor& waveform,
+    const torch::Tensor& b_coeffs) {
+  int64_t n_order = b_coeffs.size(1);
+  int64_t n_channel = b_coeffs.size(0);
 
-    auto dx = torch::Tensor();
-    auto db = torch::Tensor();
-    auto dy = grad_outputs[0];
+  namespace F = torch::nn::functional;
+  auto b_coeff_flipped = b_coeffs.flip(1).contiguous();
+  auto padded_waveform =
+      F::pad(waveform, F::PadFuncOptions({n_order - 1, 0}));
 
-    namespace F = torch::nn::functional;
+  auto output = F::conv1d(
+      padded_waveform,
+      b_coeff_flipped.unsqueeze(1),
+      F::Conv1dFuncOptions().groups(n_channel));
 
-    if (b_coeffs.requires_grad()) {
-      db = F::conv1d(
-               F::pad(x, F::PadFuncOptions({n_order - 1, 0}))
-                   .view({1, n_batch * n_channel, -1}),
-               dy.view({n_batch * n_channel, 1, -1}),
-               F::Conv1dFuncOptions().groups(n_batch * n_channel))
-               .view({n_batch, n_channel, -1})
-               .sum(0)
-               .flip(1);
-    }
+  return output;
+}
 
-    if (x.requires_grad()) {
-      dx = F::conv1d(
-          F::pad(dy, F::PadFuncOptions({0, n_order - 1})),
-          b_coeffs.unsqueeze(1),
-          F::Conv1dFuncOptions().groups(n_channel));
-    }
+std::tuple<torch::Tensor, torch::Tensor> fir_backward(
+    const torch::Tensor& grad_output,
+    const torch::Tensor& waveform,
+    const torch::Tensor& b_coeffs) {
+  int64_t n_batch = waveform.size(0);
+  int64_t n_channel = waveform.size(1);
+  int64_t n_order = b_coeffs.size(1);
+
+  auto dx = torch::Tensor();
+  auto db = torch::Tensor();
+
+  namespace F = torch::nn::functional;
+
+  // Compute gradient w.r.t. b_coeffs
+  if (b_coeffs.requires_grad()) {
+    db = F::conv1d(
+             F::pad(waveform, F::PadFuncOptions({n_order - 1, 0}))
+                 .view({1, n_batch * n_channel, -1}),
+             grad_output.view({n_batch * n_channel, 1, -1}),
+             F::Conv1dFuncOptions().groups(n_batch * n_channel))
+             .view({n_batch, n_channel, -1})
+             .sum(0)
+             .flip(1);
+  }
 
-    return {dx, db};
+  // Compute gradient w.r.t. waveform
+  if (waveform.requires_grad()) {
+    dx = F::conv1d(
+        F::pad(grad_output, F::PadFuncOptions({0, n_order - 1})),
+        b_coeffs.unsqueeze(1),
+        F::Conv1dFuncOptions().groups(n_channel));
   }
-};
 
-torch::Tensor differentiable_fir_apply(
-    const torch::Tensor& waveform,
-    const torch::Tensor& b_coeffs) {
-  return DifferentiableFIR::apply(waveform, b_coeffs);
+  return std::make_tuple(dx, db);
 }
 
 torch::Tensor differentiable_iir_apply(
@@ -267,14 +254,15 @@ TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
 
 TORCH_LIBRARY(torchaudio, m) {
   m.def(
-      "torchaudio::_lfilter(Tensor waveform, Tensor a_coeffs, Tensor b_coeffs) -> Tensor");
+      "torchaudio::_differentiable_iir_apply(Tensor waveform, Tensor a_coeffs_normalized) -> Tensor");
   m.def(
-      "torchaudio::_differentiable_fir_apply(Tensor waveform, Tensor b_coeffs) -> Tensor");
+      "torchaudio::_fir_forward(Tensor waveform, Tensor b_coeffs) -> Tensor");
   m.def(
-      "torchaudio::_differentiable_iir_apply(Tensor waveform, Tensor a_coeffs_normalized) -> Tensor");
+      "torchaudio::_fir_backward(Tensor grad_output, Tensor waveform, Tensor b_coeffs) -> (Tensor, Tensor)");
 }
 
 TORCH_LIBRARY_IMPL(torchaudio, CompositeImplicitAutograd, m) {
-  m.impl("torchaudio::_differentiable_fir_apply", differentiable_fir_apply);
   m.impl("torchaudio::_differentiable_iir_apply", differentiable_iir_apply);
+  m.impl("torchaudio::_fir_forward", fir_forward);
+  m.impl("torchaudio::_fir_backward", fir_backward);
 }

src/torchaudio/functional/filtering.py

@@ -933,8 +933,9 @@ def _lfilter_core_generic_loop(input_signal_windows: Tensor, a_coeffs_flipped: T
 
 if _IS_TORCHAUDIO_EXT_AVAILABLE:
     _lfilter_core_cpu_loop = torch.ops.torchaudio._lfilter_core_loop
-    _differentiable_fir_apply = torch.ops.torchaudio._differentiable_fir_apply
     _differentiable_iir_apply = torch.ops.torchaudio._differentiable_iir_apply
+    _fir_forward = torch.ops.torchaudio._fir_forward
+    _fir_backward = torch.ops.torchaudio._fir_backward
 else:
     _lfilter_core_cpu_loop = _lfilter_core_generic_loop
 
@@ -993,8 +994,26 @@ def _lfilter_core(
     return output
 
 
-# TODO find a better name for this, possibly renaming the existing `_lfilter_core`
-def _lfilter_core_in_python_calling_into_cpp_FIR_and_IIR(
+class _DifferentiableFIRFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, waveform, b_coeffs):
+        ctx.save_for_backward(waveform, b_coeffs)
+        
+        output = _fir_forward(waveform, b_coeffs)
+        return output
+    
+    @staticmethod
+    def backward(ctx, grad_output):
+        # Retrieve saved inputs
+        waveform, b_coeffs = ctx.saved_tensors
+        
+        # Call C++ backward function
+        grad_waveform, grad_b_coeffs = _fir_backward(grad_output, waveform, b_coeffs)
+        
+        return grad_waveform, grad_b_coeffs
+
+
+def _lfilter_core_python(
     waveform: Tensor,
     a_coeffs: Tensor,
     b_coeffs: Tensor,
@@ -1006,13 +1025,18 @@ def _lfilter_core_in_python_calling_into_cpp_FIR_and_IIR(
     a0 = a_coeffs[:, 0:1]  # Keep dimension for broadcasting
     b_coeffs_normalized = b_coeffs / a0
     a_coeffs_normalized = a_coeffs / a0
-
-    filtered_waveform = _differentiable_fir_apply(waveform, b_coeffs_normalized)
-    return _differentiable_iir_apply(filtered_waveform, a_coeffs_normalized)
+    
+    # Apply FIR filter using Python autograd function
+    filtered_waveform = _DifferentiableFIRFunction.apply(waveform, b_coeffs_normalized)
+    
+    # Apply IIR filter (still using C++ autograd)
+    output = _differentiable_iir_apply(filtered_waveform, a_coeffs_normalized)
+    
+    return output
 
 
 if _IS_TORCHAUDIO_EXT_AVAILABLE:
-    _lfilter = _lfilter_core_in_python_calling_into_cpp_FIR_and_IIR
+    _lfilter = _lfilter_core_python
 else:
     _lfilter = _lfilter_core