Feat/implement fusion for irfft (#4736)

* ignore my tags file in .gitignore

* add fusion support for rfft

* remove warnings for missing documentation

* update router runner, add rfft call

* fix typos in doc for rfft

* implement fusion for irfft

* update router runner, add irfft call

* add fusion support for rfft

* remove warnings for missing documentation

* implement fusion for irfft

* update router runner, add irfft call

Author: Sublime12

crates/burn-fusion/src/ops/module.rs

@@ -1593,10 +1593,36 @@ impl<B: FusionBackend> ModuleOps<Fusion<B>> for Fusion<B> {
     }
 
     fn irfft(
-        _spectrum_re: FloatTensor<Fusion<B>>,
-        _spectrum_im: FloatTensor<Fusion<B>>,
-        _dim: usize,
+        spectrum_re: FloatTensor<Fusion<B>>,
+        spectrum_im: FloatTensor<Fusion<B>>,
+        dim: usize,
     ) -> FloatTensor<Fusion<B>> {
-        todo!("irfft is not yet supported for fusion")
+        make_ops!(IRfftOps, IRfftOpIr, |desc: &IRfftOpIr,
+                                        handles: &mut HandleContainer<
+            B::Handle,
+        >| {
+            let input_re = handles.get_float_tensor::<B>(&desc.input_re);
+            let input_im = handles.get_float_tensor::<B>(&desc.input_im);
+
+            let signal = B::irfft(input_re, input_im, desc.dim);
+            handles.register_float_tensor::<B>(&desc.out_signal.id, signal);
+        });
+
+        let streams = OperationStreams::with_inputs([&spectrum_re, &spectrum_im]);
+        let client = spectrum_re.client.clone();
+
+        let desc = IRfftOpIr::create(spectrum_re.into_ir(), spectrum_im.into_ir(), dim, || {
+            client.create_empty_handle()
+        });
+
+        let mut outputs = client
+            .register(
+                streams,
+                OperationIr::Module(ModuleOperationIr::IRfft(desc.clone())),
+                IRfftOps::<B>::new(desc),
+            )
+            .into_iter();
+
+        outputs.next().unwrap()
     }
 }

crates/burn-fusion/src/stream/context.rs

@@ -491,6 +491,12 @@ impl RelativeOps for ModuleOperationIr {
                 out_re: desc.out_re.to_relative(converter),
                 out_im: desc.out_re.to_relative(converter),
             }),
+            ModuleOperationIr::IRfft(desc) => ModuleOperationIr::IRfft(IRfftOpIr {
+                input_re: desc.input_re.to_relative(converter),
+                input_im: desc.input_im.to_relative(converter),
+                dim: desc.dim,
+                out_signal: desc.out_signal.to_relative(converter),
+            }),
             ModuleOperationIr::Attention(desc) => ModuleOperationIr::Attention(AttentionOpIr {
                 query: desc.query.to_relative(converter),
                 key: desc.key.to_relative(converter),

crates/burn-ir/src/operation.rs

@@ -252,6 +252,8 @@ pub enum ModuleOperationIr {
     InterpolateBackward(InterpolateBackwardOpIr),
     /// Operation corresponding to [Rfft](burn_backend::ops::ModuleOps::rfft)
     Rfft(RfftOpIr),
+    /// Operation corresponding to [IRfft](burn_backend::ops::ModuleOps::irfft)
+    IRfft(IRfftOpIr),
     /// Operation corresponding to [attention](burn_backend::ops::ModuleOps::attention).
     Attention(AttentionOpIr),
 }
@@ -1599,6 +1601,15 @@ pub struct RfftOpIr {
     pub out_im: TensorIr,
 }
 
+#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
+#[allow(missing_docs)]
+pub struct IRfftOpIr {
+    pub input_re: TensorIr,
+    pub input_im: TensorIr,
+    pub dim: usize,
+    pub out_signal: TensorIr,
+}
+
 #[allow(missing_docs)]
 impl RfftOpIr {
     pub fn create<F>(signal: TensorIr, dim: usize, mut new_id: F) -> Self
@@ -1618,6 +1629,25 @@ impl RfftOpIr {
     }
 }
 
+#[allow(missing_docs)]
+impl IRfftOpIr {
+    pub fn create<F>(input_re: TensorIr, input_im: TensorIr, dim: usize, mut new_id: F) -> Self
+    where
+        F: FnMut() -> crate::TensorId,
+    {
+        let mut shape = input_re.shape.clone();
+        shape[dim] = (shape[dim] - 1) * 2;
+        let dtype = input_re.dtype;
+
+        Self {
+            input_re,
+            input_im,
+            dim,
+            out_signal: TensorIr::uninit(new_id(), shape, dtype),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
 #[allow(missing_docs)]
 pub struct AttentionOptionsIr {
@@ -2703,6 +2733,9 @@ impl ModuleOperationIr {
                 Box::new([&repr.x, &repr.grad].into_iter())
             }
             ModuleOperationIr::Rfft(repr) => Box::new([&repr.signal].into_iter()),
+            ModuleOperationIr::IRfft(repr) => {
+                Box::new([&repr.input_re, &repr.input_im].into_iter())
+            }
             ModuleOperationIr::Attention(repr) => {
                 if let Some(mask) = &repr.mask {
                     if let Some(attn_bias) = &repr.attn_bias {
@@ -2798,6 +2831,7 @@ impl ModuleOperationIr {
             ModuleOperationIr::Interpolate(repr) => Box::new([&repr.out].into_iter()),
             ModuleOperationIr::InterpolateBackward(repr) => Box::new([&repr.out].into_iter()),
             ModuleOperationIr::Rfft(repr) => Box::new([&repr.out_re, &repr.out_im].into_iter()),
+            ModuleOperationIr::IRfft(repr) => Box::new([&repr.out_signal].into_iter()),
             ModuleOperationIr::Attention(repr) => Box::new([&repr.out].into_iter()),
         }
     }
@@ -2998,6 +3032,10 @@ impl ModuleOperationIr {
             ModuleOperationIr::Rfft(repr) => {
                 repr.signal.mark_read_only(nodes, &mut output);
             }
+            ModuleOperationIr::IRfft(repr) => {
+                repr.input_re.mark_read_only(nodes, &mut output);
+                repr.input_im.mark_read_only(nodes, &mut output);
+            }
             ModuleOperationIr::Attention(repr) => {
                 repr.query.mark_read_only(nodes, &mut output);
                 repr.key.mark_read_only(nodes, &mut output);

crates/burn-router/src/runner.rs

@@ -1524,6 +1524,13 @@ impl<B: BackendIr> RunnerClient for Runner<B> {
                     handles.register_float_tensor::<B>(&desc.out_re.id, out_re);
                     handles.register_float_tensor::<B>(&desc.out_im.id, out_im);
                 }
+                ModuleOperationIr::IRfft(desc) => {
+                    let spectrum_re = handles.get_float_tensor::<B>(&desc.input_re);
+                    let spectrum_im = handles.get_float_tensor::<B>(&desc.input_im);
+                    let signal = B::irfft(spectrum_re, spectrum_im, desc.dim);
+
+                    handles.register_float_tensor::<B>(&desc.out_signal.id, signal);
+                }
                 ModuleOperationIr::Attention(desc) => {
                     let query = handles.get_float_tensor::<B>(&desc.query);
                     let key = handles.get_float_tensor::<B>(&desc.key);