Update FFT example.

Author: maleadt

examples/fft.jl

@@ -20,7 +20,8 @@ using FFTW
 # in columns. In Julia column-major, reshape (F1F2, F0) puts stride-F0 elements in rows.
 # We use right-multiply X @ W instead of W @ X to process rows instead of columns.
 #
-# Input/output layout: (D, BS, N2D) where D=2 for real/imag interleaving.
+# Input/output memory layout: (D, BS, N2D) where D=2 for real/imag interleaving.
+# Internally, BS is permuted to trailing position for batched matmul convention.
 function fft_kernel(
     x_packed_in::ct.TileArray{Float32, 3},   # Input (D, BS, N2D) - natural Julia complex layout
     y_packed_out::ct.TileArray{Float32, 3},  # Output (D, BS, N2D)
@@ -55,96 +56,94 @@ function fft_kernel(
     bid = ct.bid(1)
 
     # --- Load Input Data ---
-    # Input is (D, BS, N2D) where D=2 for real/imag. Load and reshape to (2, BS, N).
-    X_ri = reshape(ct.load(x_packed_in; index=(1, bid, 1), shape=(D, BS, N2D)), (2, BS, N))
+    # Input is (D, BS, N2D) where D=2 for real/imag. Load and permute BS to trailing.
+    X_ri_mem = reshape(ct.load(x_packed_in; index=(1, bid, 1), shape=(D, BS, N2D)), (2, BS, N))
+    X_ri = permutedims(X_ri_mem, (1, 3, 2))  # (2, N, BS) — trailing batch
 
     # Split real and imaginary parts (extract from first dimension)
-    X_r = reshape(ct.extract(X_ri, (1, 1, 1), (1, BS, N)), (BS, F1F2, F0))
-    X_i = reshape(ct.extract(X_ri, (2, 1, 1), (1, BS, N)), (BS, F1F2, F0))
+    X_r = reshape(ct.extract(X_ri, (1, 1, 1), (1, N, BS)), (F1F2, F0, BS))
+    X_i = reshape(ct.extract(X_ri, (2, 1, 1), (1, N, BS)), (F1F2, F0, BS))
 
     # --- Load DFT Matrices ---
-    # W0 (F0 x F0) - for right-multiply X @ W0
+    # W0 (F0 x F0) - for right-multiply X @ W0, batch dim trailing
     W0_ri = reshape(ct.load(W0; index=(1, 1, 1), shape=(F0, F0, 2)), (F0, F0, 2))
-    W0_r = ct.broadcast_to(reshape(ct.extract(W0_ri, (1, 1, 1), (F0, F0, 1)), (1, F0, F0)), (BS, F0, F0))
-    W0_i = ct.broadcast_to(reshape(ct.extract(W0_ri, (1, 1, 2), (F0, F0, 1)), (1, F0, F0)), (BS, F0, F0))
+    W0_r = ct.broadcast_to(reshape(ct.extract(W0_ri, (1, 1, 1), (F0, F0, 1)), (F0, F0, 1)), (F0, F0, BS))
+    W0_i = ct.broadcast_to(reshape(ct.extract(W0_ri, (1, 1, 2), (F0, F0, 1)), (F0, F0, 1)), (F0, F0, BS))
 
     # W1 (F1 x F1)
     W1_ri = reshape(ct.load(W1; index=(1, 1, 1), shape=(F1, F1, 2)), (F1, F1, 2))
-    W1_r = ct.broadcast_to(reshape(ct.extract(W1_ri, (1, 1, 1), (F1, F1, 1)), (1, F1, F1)), (BS, F1, F1))
-    W1_i = ct.broadcast_to(reshape(ct.extract(W1_ri, (1, 1, 2), (F1, F1, 1)), (1, F1, F1)), (BS, F1, F1))
+    W1_r = ct.broadcast_to(reshape(ct.extract(W1_ri, (1, 1, 1), (F1, F1, 1)), (F1, F1, 1)), (F1, F1, BS))
+    W1_i = ct.broadcast_to(reshape(ct.extract(W1_ri, (1, 1, 2), (F1, F1, 1)), (F1, F1, 1)), (F1, F1, BS))
 
     # W2 (F2 x F2)
     W2_ri = reshape(ct.load(W2; index=(1, 1, 1), shape=(F2, F2, 2)), (F2, F2, 2))
-    W2_r = ct.broadcast_to(reshape(ct.extract(W2_ri, (1, 1, 1), (F2, F2, 1)), (1, F2, F2)), (BS, F2, F2))
-    W2_i = ct.broadcast_to(reshape(ct.extract(W2_ri, (1, 1, 2), (F2, F2, 1)), (1, F2, F2)), (BS, F2, F2))
+    W2_r = ct.broadcast_to(reshape(ct.extract(W2_ri, (1, 1, 1), (F2, F2, 1)), (F2, F2, 1)), (F2, F2, BS))
+    W2_i = ct.broadcast_to(reshape(ct.extract(W2_ri, (1, 1, 2), (F2, F2, 1)), (F2, F2, 1)), (F2, F2, BS))
 
     # --- Load Twiddle Factors ---
     # T0 (F1F2, F0) - note swapped from Python's (F0, F1F2)
     T0_ri = reshape(ct.load(T0; index=(1, 1, 1), shape=(F1F2, F0, 2)), (F1F2, F0, 2))
-    T0_r = reshape(ct.extract(T0_ri, (1, 1, 1), (F1F2, F0, 1)), (1, N))
-    T0_i = reshape(ct.extract(T0_ri, (1, 1, 2), (F1F2, F0, 1)), (1, N))
+    T0_r = reshape(ct.extract(T0_ri, (1, 1, 1), (F1F2, F0, 1)), (N, 1))
+    T0_i = reshape(ct.extract(T0_ri, (1, 1, 2), (F1F2, F0, 1)), (N, 1))
 
     # T1 (F0F2, F1) - note swapped from Python's (F1, F2)
     T1_ri = reshape(ct.load(T1; index=(1, 1, 1), shape=(F0F2, F1, 2)), (F0F2, F1, 2))
-    T1_r = reshape(ct.extract(T1_ri, (1, 1, 1), (F0F2, F1, 1)), (1, F0F2 * F1))
-    T1_i = reshape(ct.extract(T1_ri, (1, 1, 2), (F0F2, F1, 1)), (1, F0F2 * F1))
+    T1_r = reshape(ct.extract(T1_ri, (1, 1, 1), (F0F2, F1, 1)), (F0F2 * F1, 1))
+    T1_i = reshape(ct.extract(T1_ri, (1, 1, 2), (F0F2, F1, 1)), (F0F2 * F1, 1))
 
     # --- Stage 0: F0-point DFT ---
-    # X is (BS, F1F2, F0), W0 is (BS, F0, F0)
+    # X is (F1F2, F0, BS), W0 is (F0, F0, BS) — trailing batch
     # Right-multiply: X @ W0 processes each row (F1F2 rows, each with F0 elements)
-    # Each row has elements at stride F1F2 in the original array - exactly what we need!
-    X_r_ = X_r * W0_r - X_i * W0_i  # (BS, F1F2, F0) @ (BS, F0, F0) → (BS, F1F2, F0)
+    X_r_ = X_r * W0_r - X_i * W0_i  # (F1F2, F0, BS) @ (F0, F0, BS) → (F1F2, F0, BS)
     X_i_ = X_r * W0_i + X_i * W0_r
 
     # --- Twiddle & Permute 0 ---
-    # Reshape to (BS, N) for element-wise twiddle multiply
-    X_r_flat = reshape(X_r_, (BS, N))
-    X_i_flat = reshape(X_i_, (BS, N))
+    # Reshape to (N, BS) for element-wise twiddle multiply
+    X_r_flat = reshape(X_r_, (N, BS))
+    X_i_flat = reshape(X_i_, (N, BS))
     X_r2 = T0_r .* X_r_flat .- T0_i .* X_i_flat
     X_i2 = T0_i .* X_r_flat .+ T0_r .* X_i_flat
 
     # Reshape and permute for stage 1
-    # Current logical layout after reshape (BS, F1F2, F0): data at (bs, f1*F2+f2, f0)
-    # Reshape to (BS, F2, F1, F0) then permute to (BS, F0F2, F1) for stage 1
-    X_r3 = reshape(X_r2, (BS, F2, F1, F0))
-    X_i3 = reshape(X_i2, (BS, F2, F1, F0))
-    X_r4 = permutedims(X_r3, (1, 2, 4, 3))  # (BS, F2, F0, F1)
-    X_i4 = permutedims(X_i3, (1, 2, 4, 3))
-    X_r5 = reshape(X_r4, (BS, F0F2, F1))
-    X_i5 = reshape(X_i4, (BS, F0F2, F1))
+    # Reshape to (F2, F1, F0, BS) then permute to (F0F2, F1, BS) for stage 1
+    X_r3 = reshape(X_r2, (F2, F1, F0, BS))
+    X_i3 = reshape(X_i2, (F2, F1, F0, BS))
+    X_r4 = permutedims(X_r3, (1, 3, 2, 4))  # (F2, F0, F1, BS)
+    X_i4 = permutedims(X_i3, (1, 3, 2, 4))
+    X_r5 = reshape(X_r4, (F0F2, F1, BS))
+    X_i5 = reshape(X_i4, (F0F2, F1, BS))
 
     # --- Stage 1: F1-point DFT ---
-    # X is (BS, F0F2, F1), W1 is (BS, F1, F1)
+    # X is (F0F2, F1, BS), W1 is (F1, F1, BS)
     X_r6 = X_r5 * W1_r - X_i5 * W1_i
     X_i6 = X_r5 * W1_i + X_i5 * W1_r
 
     # --- Twiddle & Permute 1 ---
-    X_r_flat2 = reshape(X_r6, (BS, N))
-    X_i_flat2 = reshape(X_i6, (BS, N))
+    X_r_flat2 = reshape(X_r6, (N, BS))
+    X_i_flat2 = reshape(X_i6, (N, BS))
     X_r7 = T1_r .* X_r_flat2 .- T1_i .* X_i_flat2
     X_i7 = T1_i .* X_r_flat2 .+ T1_r .* X_i_flat2
 
     # Reshape and permute for stage 2
-    X_r8 = reshape(X_r7, (BS, F2, F0, F1))
-    X_i8 = reshape(X_i7, (BS, F2, F0, F1))
-    X_r9 = permutedims(X_r8, (1, 3, 4, 2))  # (BS, F0, F1, F2)
-    X_i9 = permutedims(X_i8, (1, 3, 4, 2))
-    X_r10 = reshape(X_r9, (BS, F0F1, F2))
-    X_i10 = reshape(X_i9, (BS, F0F1, F2))
+    X_r8 = reshape(X_r7, (F2, F0, F1, BS))
+    X_i8 = reshape(X_i7, (F2, F0, F1, BS))
+    X_r9 = permutedims(X_r8, (2, 3, 1, 4))  # (F0, F1, F2, BS)
+    X_i9 = permutedims(X_i8, (2, 3, 1, 4))
+    X_r10 = reshape(X_r9, (F0F1, F2, BS))
+    X_i10 = reshape(X_i9, (F0F1, F2, BS))
 
     # --- Stage 2: F2-point DFT ---
-    # X is (BS, F0F1, F2), W2 is (BS, F2, F2)
+    # X is (F0F1, F2, BS), W2 is (F2, F2, BS)
     X_r11 = X_r10 * W2_r - X_i10 * W2_i
     X_i11 = X_r10 * W2_i + X_i10 * W2_r
 
     # --- Final Output ---
-    # After stage 2, data is in (BS, F0F1, F2) layout
-    # Reshape to (BS, F0, F1, F2) - output is already in frequency order
-    X_r_final = reshape(X_r11, (1, BS, N))
-    X_i_final = reshape(X_i11, (1, BS, N))
+    X_r_final = reshape(X_r11, (1, N, BS))
+    X_i_final = reshape(X_i11, (1, N, BS))
 
     # --- Concatenate and Store ---
-    Y_ri = reshape(ct.cat((X_r_final, X_i_final), 1), (D, BS, N2D))
+    # Permute BS back to middle for memory layout (D, BS, N2D)
+    Y_ri = permutedims(reshape(ct.cat((X_r_final, X_i_final), 1), (D, N2D, BS)), (1, 3, 2))
     ct.store(y_packed_out; index=(1, bid, 1), tile=Y_ri)
 
     return