Add `aten.stft.center` and decomposition #3880

giacs-epic · 2024-11-18T17:01:44Z

The choice to work with aten.stft.center instead of aten.stft is because the latter doesn't match the signature that gets exposed (see https://pytorch.org/docs/stable/generated/torch.stft.html).

This reverts commit ae145aa.

giacs-epic · 2025-01-09T09:23:50Z

CI errors (undefined symbols referenced by LazyNativeFunctions.cpp) are unrelated to the PR content.

zjgarvey · 2025-02-17T23:19:41Z

Might need to add tosa xfails if you haven't done so already. IIRC Tosa got added to the CI in your last sync.

zjgarvey

A general comment before reviewing further:

Is it absolutely necessary to use a loop? The only times I would ever consider using a loop is if there is a necessary loop-carried dependency, but I don't think that is the case here.

Even if there is a nice algorithm for rfft, I don't think decomposing to many rffts in a loop would be more efficient than converting this to, say, a convolution with something like window*exp (if that is possible with the configurations you are trying to support).

zjgarvey · 2025-02-18T17:03:45Z

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

+    // init_freq_tensor = aten.empty.memory_format([batch_dim?, n_freqs,
+    // n_frames],
+    //                                self.dtype, None, None, None, None)
+    // final_freq_tensor = prim.loop
+    // n_frames, %true, init(init_freq_tensor)
+    // {
+    //  ^bb0(frame, freq_tensor):
+    //    begin = frame * hop_length
+    //    end = begin + n_fft
+    //    narrow_length = min(end, signal_len) - begin
+    //    missing = n_fft - narrow_length
+    //    sliced = torch.narrow(self, axis_signal, begin, narrow_length) :
+    //            !torch.vtensor<[batch_dim?,?],f32>
+    //    padded_sliced = aten.pad(sliced, [0, missing], "constant", 0.0) :
+    //        !torch.vtensor<[batch_dim?,?],f32>
+    //    padded_sliced = tensor_static_info_cast(padded_sliced) :
+    //        !torch.vtensor<[batch_dim?,n_fft],f32>
+    //    weighted = aten.mul.Tensor(padded_sliced, window) :
+    //        !torch.vtensor<[batch_dim?,n_fft],f32>
+    //    f = onesidedBool ? aten.fft_rfft : aten.fft_fft
+    //    freq_slice_sq = f(weighted, None, axis_signal) :
+    //        !torch.vtensor<[batch_dim?,n_freqs],f32>
+    //    freq_slice = aten.unsqueeze(freq_slice_sq, axis_frames) :
+    //        !torch.vtensor<[batch_dim?,n_freqs, 1],f32>
+    //    new_freq_tensor = aten.slice_scatter(
+    //                      freq_tensor, freq_slice,
+    //                      dim=axis_frames, start=frame,
+    //                      end=None, step=1
+    //                      )
+    //    torch.prim.Loop.condition %true, iter(%new_freq_tensor)
+    // }
+    // return final_freq_tensor


Pseudo-IR isn't very helpful as a comment, since you include lit tests for this decomposition.

Removed in last commit.

zjgarvey · 2025-02-18T17:30:02Z

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

+    if (isa<Torch::NoneType>(hopLength.getType())) {
+      hopLength = rewriter.create<AtenFloordivIntOp>(
+          loc, n_fft,
+          rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(4)));


nit:

There is a builder for Torch::ConstantIntOp which allows passing an int directly, which is a bit easier to read.

Suggested change

rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(4)));

rewriter.create<ConstantIntOp>(loc, 4));

Changed in last commit.

zjgarvey · 2025-02-18T17:47:32Z

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

+    Value center = op.getCenter();
+    bool centerBool;
+    // TODO: add support for non-constant center and center=True
+    if (!matchPattern(center, m_TorchConstantBool(&centerBool)))
+      return rewriter.notifyMatchFailure(op,
+                                         "Unsupported: non-constant center");
+    if (centerBool)
+      return rewriter.notifyMatchFailure(op, "Unsupported: center=True");
+
+    Value normalized = op.getNormalized();
+    bool normalizedBool;
+    // TODO: add support for non-constant normalized and normalized=True
+    if (!matchPattern(normalized, m_TorchConstantBool(&normalizedBool)))
+      return rewriter.notifyMatchFailure(
+          op, "Unsupported: non-constant normalized");
+    if (normalizedBool)
+      return rewriter.notifyMatchFailure(op, "Unsupported: normalized=True");
+
+    bool onesidedBool;
+    // Default: True for real input and window, False otherwise.
+    // TODO: add support for non-constant onesided
+    if (isa<Torch::NoneType>(op.getOnesided().getType())) {
+      Type dtype = selfType.getDtype();
+      onesidedBool = !isa<mlir::ComplexType>(dtype);
+    } else if (!matchPattern(op.getOnesided(),
+                             m_TorchConstantBool(&onesidedBool)))
+      return rewriter.notifyMatchFailure(op,
+                                         "Unsupported: non-constant onesided");
+
+    Value returnComplex = op.getReturnComplex();
+    bool returnComplexBool;
+    // TODO: add support for non-constant return_complex and return_complex=True
+    if (!matchPattern(returnComplex, m_TorchConstantBool(&returnComplexBool)))
+      return rewriter.notifyMatchFailure(
+          op, "Unsupported: non-constant return_complex");
+    if (!returnComplex)
+      return rewriter.notifyMatchFailure(op,
+                                         "Unsupported: return_complex=False");


Move all of these match failures before the generation of runtime assert ops in the if (hasWindow) block.

Changed in last commit.

giacs-epic · 2025-03-28T15:19:50Z

A general comment before reviewing further:

Is it absolutely necessary to use a loop? The only times I would ever consider using a loop is if there is a necessary loop-carried dependency, but I don't think that is the case here.

Even if there is a nice algorithm for rfft, I don't think decomposing to many rffts in a loop would be more efficient than converting this to, say, a convolution with something like window*exp (if that is possible with the configurations you are trying to support).

The two approaches have different time complexities: assuming N is both the number of frequency samples and the window length (as per default) and T is the number of frames, STFT via Conv would be O(T N^2), while STFT via many FFTs (Cooley-Tukey algorithm) would be O(T N log2(N)).
Our reference model (Parakeet) for example would have the following settings for STFT when inputting a single signal of 10 seconds:

signal = float32[1,160512]
frame_step = 160
window = float32[512]
frame_length = 512

For a single core execution the many FFTs would give a ~56x speedup. In a parallel execution setting the Conv-based algorithm would have 1 more axis for parallelization (dimension N=512), so which algorithm wins would be highly dependent on the number of cores available and the signal length (and also the batches dimension).

Regarding the use of a loop, the many FFTs requires it as the signal length dimension could be unknown. If I remember correctly the loop gets unrolled completely in case the signal length dimension is known.

giacs-epic added 21 commits October 18, 2024 07:58

Emit aten.stft.center

e5f20fd

Add WIP code to decompose STFT into fft or rfft

7becc57

Further decompose stack into unsqueeze+cat

34af833

Add simplifications to unroll+fold list generation

ae145aa

Fix unsqueeze axis

d2a5d52

Fix output sizes

476573b

Fix onesided

794df5a

Add e2e test case

abe7dcf

Fix result type check

d00a51e

Fix frame padding constant

412c5af

Add aten.stft.center shape and dtype functions

5621908

Add more tests

39ad96d

Merge branch 'main' into decompose_STFT

bc33b19

Switch to implementation that doesn't depend on simplification

da353c4

Merge branch 'main' into decompose_STFT

15cf656

Regenerate AbstractInterpLibrary.cpp

80d3eb2

Fix empty/dynamic-sized window and dynamic-sized signal handling

fd1563a

Add unit tests

bac4d77

Revert "Add simplifications to unroll+fold list generation"

4bfc92c

This reverts commit ae145aa.

Merge branch 'main' into decompose_STFT

93a49a9

Fix torch onnx to torch tests

421420b

giacs-epic marked this pull request as ready for review December 2, 2024 13:50

giacs-epic added 3 commits December 2, 2024 14:53

Update xfail_sets

410d2bf

Add AtenStftOp canonicalizer (required due to latest versions of torch)

cf13b48

Merge branch 'main' into decompose_STFT

ae4db67

Merge branch 'main' into decompose_STFT

44ac929

Add Stft tests to TOSA xfail sets

815a989

zjgarvey requested changes Feb 18, 2025

View reviewed changes

giacs-epic added 3 commits March 28, 2025 15:43

Review feedback

e4388fe

Merge branch 'main' into decompose_STFT

8c4c72b

Merge branch 'main' into decompose_STFT

089a108

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `aten.stft.center` and decomposition #3880

Add `aten.stft.center` and decomposition #3880

giacs-epic commented Nov 18, 2024 •

edited

Loading

giacs-epic commented Jan 9, 2025

zjgarvey commented Feb 17, 2025

zjgarvey left a comment

zjgarvey Feb 18, 2025

giacs-epic Mar 28, 2025

zjgarvey Feb 18, 2025

giacs-epic Mar 28, 2025

zjgarvey Feb 18, 2025

giacs-epic Mar 28, 2025

giacs-epic commented Mar 28, 2025

	rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(4)));
	rewriter.create<ConstantIntOp>(loc, 4));

Add aten.stft.center and decomposition #3880

Are you sure you want to change the base?

Add aten.stft.center and decomposition #3880

Conversation

giacs-epic commented Nov 18, 2024 • edited Loading

giacs-epic commented Jan 9, 2025

zjgarvey commented Feb 17, 2025

zjgarvey left a comment

Choose a reason for hiding this comment

zjgarvey Feb 18, 2025

Choose a reason for hiding this comment

giacs-epic Mar 28, 2025

Choose a reason for hiding this comment

zjgarvey Feb 18, 2025

Choose a reason for hiding this comment

giacs-epic Mar 28, 2025

Choose a reason for hiding this comment

zjgarvey Feb 18, 2025

Choose a reason for hiding this comment

giacs-epic Mar 28, 2025

Choose a reason for hiding this comment

giacs-epic commented Mar 28, 2025

Add `aten.stft.center` and decomposition #3880

Add `aten.stft.center` and decomposition #3880

giacs-epic commented Nov 18, 2024 •

edited

Loading