[P1] FFT kernel issue

[n0thingNoob]

FFT generated ASM appears to update stage-level loop recurrences on the inner butterfly iteration

Summary

The generated FFT instruction YAML/ASM does not match the CPU result from fft_int.c. The issue appears to be in the generated loop-control / PHI wiring for the nested FFT loops.

In the C source, groupsPerStage, buttersPerGroup, and coef_base are stage-level variables. They should update only after all j/k iterations for the current stage complete.

for (i = 0; i < NSTAGES; ++i) {
  for (j = 0; j < groupsPerStage; ++j) {
    for (k = 0; k < buttersPerGroup; ++k) {
      ... butterfly ...
    }
  }
  groupsPerStage = groupsPerStage * 2;
  buttersPerGroup = buttersPerGroup / 2;
  coef_base = (coef_base << 1) + 1;
}

Suspect generated ASM

In tmp-generated-instructions.asm, PE(1,2) contains a PHI_START id=36 followed by DIV / 2. Dynamic trace shows this value evolves as:

128 -> 64 -> 32 -> 16 -> 8 -> 4 -> 2 -> 1

This looks like buttersPerGroup is being updated every dynamic butterfly iteration, rather than at the stage boundary.

In PE(2,2), PHI_START id=37 followed by SHL << 1 evolves as:

1 -> 2 -> 4 -> 8 -> 16 -> 32 -> 64 -> 128

This looks like groupsPerStage is also being updated every dynamic butterfly iteration.

ICMP_EQ id=64 then reaches 8 and triggers RETURN_VOID id=111 after only a few butterfly iterations:

Time 297: ICMP_EQ id=64 Src1=8 Src2=8 Result=1(true)
Time 299: RETURN_VOID id=111 Pred=true

Address stream symptom

The data load addresses become a linear sweep:

real lower: 0, 1, 2, ...
real upper: 128, 129, 130, ...
imag lower: 256, 257, 258, ...
imag upper: 384, 385, 386, ...

The CPU nested loop should switch address patterns between FFT stages, so this linear stream indicates the stage-level recurrence/predicate is likely connected to the wrong loop level.

Expected compiler behavior

The recurrence/predicate for:

groupsPerStage *= 2;
buttersPerGroup /= 2;
coef_base = (coef_base << 1) + 1;

should be guarded by the stage-boundary predicate, after the j/k loops finish. It should not be driven by the inner butterfly iteration.

Attached screenshots

• 01_c_source_loop.png: source loop and stage-level updates
• 02_asm_pe12_butters_recurrence.png: suspect buttersPerGroup recurrence
• 03_asm_pe22_groups_return.png: suspect groupsPerStage / return predicate
• 04_asm_address_generation_symptom.png: address-generation symptom
• 05_dfg_loop_control_subset.png: local DFG relation

[tancheng]

Hi @ShangkunLi, can you PTAL?

[tancheng]

How should we debug this @ShangkunLi? Do you need anything from @n0thingNoob?

[ShangkunLi]

How should we debug this @ShangkunLi? Do you need anything from @n0thingNoob?

Sry for the delay, I will try to reproduce this result tonight and debug it.

[ShangkunLi]

Hi~ @n0thingNoob, the problem is because we bypass some values in the --canonicalize-live-in pass through the DirectDominatingLiveIn pattern.

neura/lib/NeuraDialect/Transforms/CanonicalizeLiveInPass.cpp

Line 25 in 2fea5bf

struct DirectDominatingLiveIn {

Specifically,

^bb1(%4: i32, %5: i32, %6: i32, %7: i32, %8: i64):
  %9 = "neura.icmp"(%5) <{cmpType = "sgt"}> {rhs_value = 0 : i32} : (i32) -> i1
  neura.cond_br %9 : i1 then %4, %5, %6, %8 : i32, i32, i32, i64 to ^bb2 else %8 : i64 to ^bb6

...

^bb6(%79: i64):  // stage update block
  %80 = "neura.shl"(%6) {rhs_value = 1 : i32} : (i32) -> i32
  %81 = "neura.div"(%5) {rhs_value = 2 : i32} : (i32) -> i32
  %82 = "neura.shl"(%4) {rhs_value = 1 : i32} : (i32) -> i32
  %83 = "neura.or"(%82) {rhs_value = 1 : i32} : (i32) -> i32
  %84 = "neura.add"(%7) {rhs_value = 1 : i32} : (i32) -> i32
  %85 = "neura.icmp"(%84) <{cmpType = "eq"}> {rhs_value = 8 : i32} : (i32) -> i1
  neura.cond_br %85 : i1 then %85 : i1 to ^bb7 else %83, %81, %80, %84, %79 : i32, i32, i32, i32, i64 to ^bb1

The problem is that ^bb6 only takes %79 as a block argument, but it directly uses %4, %5, %6, and %7 from ^bb1.

These values are FFT stage-level variables:

%4 = coef_base
%5 = buttersPerGroup
%6 = groupsPerStage
%7 = stage index i

They should only be updated once per FFT stage, after the inner butterfly loops finish. But because they are treated as direct dominating live-ins, the later control-to-dataflow lowering does not predicate them correctly. After CFG flattening, these update operations can fire at the inner-loop/butterfly rate.

The fix is to reject the direct-live-in optimization when there is a backedge on the path from the defining block to the using block.

I have attached the generated files after fixing the bug locally. fft-codegen.zip Could you plz run it and see if it is correct now? If so, I can make a pr.

[tancheng]

Hi @ShangkunLi, you can directly make a PR, and provide simple test to demonstrate the issue there. Thanks~!

[n0thingNoob]

Hi~ @n0thingNoob, the problem is because we bypass some values in the --canonicalize-live-in pass through the DirectDominatingLiveIn pattern.

neura/lib/NeuraDialect/Transforms/CanonicalizeLiveInPass.cpp

Line 25 in 2fea5bf

struct DirectDominatingLiveIn {
Specifically,

^bb1(%4: i32, %5: i32, %6: i32, %7: i32, %8: i64):
%9 = "neura.icmp"(%5) <{cmpType = "sgt"}> {rhs_value = 0 : i32} : (i32) -> i1
neura.cond_br %9 : i1 then %4, %5, %6, %8 : i32, i32, i32, i64 to ^bb2 else %8 : i64 to ^bb6

...

^bb6(%79: i64): // stage update block
%80 = "neura.shl"(%6) {rhs_value = 1 : i32} : (i32) -> i32
%81 = "neura.div"(%5) {rhs_value = 2 : i32} : (i32) -> i32
%82 = "neura.shl"(%4) {rhs_value = 1 : i32} : (i32) -> i32
%83 = "neura.or"(%82) {rhs_value = 1 : i32} : (i32) -> i32
%84 = "neura.add"(%7) {rhs_value = 1 : i32} : (i32) -> i32
%85 = "neura.icmp"(%84) <{cmpType = "eq"}> {rhs_value = 8 : i32} : (i32) -> i1
neura.cond_br %85 : i1 then %85 : i1 to ^bb7 else %83, %81, %80, %84, %79 : i32, i32, i32, i32, i64 to ^bb1
The problem is that ^bb6 only takes %79 as a block argument, but it directly uses %4, %5, %6, and %7 from ^bb1.

These values are FFT stage-level variables:

%4 = coef_base
%5 = buttersPerGroup
%6 = groupsPerStage
%7 = stage index i
They should only be updated once per FFT stage, after the inner butterfly loops finish. But because they are treated as direct dominating live-ins, the later control-to-dataflow lowering does not predicate them correctly. After CFG flattening, these update operations can fire at the inner-loop/butterfly rate.

The fix is to reject the direct-live-in optimization when there is a backedge on the path from the defining block to the using block.

I have attached the generated files after fixing the bug locally. fft-codegen.zip Could you plz run it and see if it is correct now? If so, I can make a pr.

ok, let me see.

[n0thingNoob]

Looks like it works. The FFT test has been passed on Zeonica. Nice work!