def input_transform(ymm_d): assert isinstance(ymm_d, list) and len(ymm_d) == 8 and all(isinstance(ymm, YMMRegister) for ymm in ymm_d) ymm_wd = [YMMRegister() for _ in range(8)] # wd0 = (d0 - d6) + 5.25 * (d4 - d2) # wd1 = (d6 + d2 - 4.25 * d4) + (d1 + d5 - 4.25 * d3) # wd2 = (d6 + d2 - 4.25 * d4) - (d1 + d5 - 4.25 * d3) # wd3 = (d6 + 0.25 * d2 - 1.25 * d4) + 2.0 * (d5 + 0.25 * d1 - 1.25 * d3) # wd4 = (d6 + 0.25 * d2 - 1.25 * d4) - 2.0 * (d5 + 0.25 * d1 - 1.25 * d3) # wd5 = (d6 - 5.0 * d4 + 4.0 * d2) + 2.0 * (d1 + 0.25 * d5 - 1.25 * d3) # wd6 = (d6 - 5.0 * d4 + 4.0 * d2) - 2.0 * (d1 + 0.25 * d5 - 1.25 * d3) # wd7 = (d7 - d1) + 5.25 * (d3 - d5) ymm_0_25 = YMMRegister() VMOVAPS(ymm_0_25, Constant.float32x8(0.25)) # Compute wd0 := d0 - d6 VSUBPS(ymm_wd[0], ymm_d[0], ymm_d[6]) ymm_d4_sub_d2 = YMMRegister() VSUBPS(ymm_d4_sub_d2, ymm_d[4], ymm_d[2]) # Compute wd7 := d7 - d1 VSUBPS(ymm_wd[7], ymm_d[7], ymm_d[1]) ymm_d3_sub_d5 = YMMRegister() VSUBPS(ymm_d3_sub_d5, ymm_d[3], ymm_d[5]) # Compute wd1 := d2 + d6 VADDPS(ymm_wd[1], ymm_d[2], ymm_d[6]) # Compute wd2 := d1 + d5 VADDPS(ymm_wd[2], ymm_d[1], ymm_d[5]) # Compute wd4 := d5 + 0.25 * d1 VMOVAPS(ymm_wd[4], ymm_d[5]) VFMADD231PS(ymm_wd[4], ymm_d[1], ymm_0_25) # Compute wd5 := d6 - 5.0 * d4 VMOVAPS(ymm_wd[5], Constant.float32x8(5.0)) VFNMADD132PS(ymm_wd[5], ymm_d[6], ymm_d[4]) # Compute wd3 := d6 + 0.25 * d2 VFMADD231PS(ymm_d[6], ymm_d[2], ymm_0_25) SWAP.REGISTERS(ymm_wd[3], ymm_d[6]) # Compute wd6 := d1 + 0.25 * d5 VFMADD231PS(ymm_d[1], ymm_d[5], ymm_0_25) SWAP.REGISTERS(ymm_wd[6], ymm_d[1]) ymm_5_25 = YMMRegister() VMOVAPS(ymm_5_25, Constant.float32x8(5.25)) # Compute wd0 := (d0 - d6) + 5.25 * (d4 - d2) VFMADD231PS(ymm_wd[0], ymm_d4_sub_d2, ymm_5_25) # Compute wd7 := (d7 - d1) + 5.25 * (d3 - d5) VFMADD231PS(ymm_wd[7], ymm_d3_sub_d5, ymm_5_25) ymm_4_25 = YMMRegister() VMOVAPS(ymm_4_25, Constant.float32x8(4.25)) # Compute # wd1 := (d6 + d2) - 4.25 * d4 # wd2 := (d1 + d5) - 4.25 * d3 VFNMADD231PS(ymm_wd[1], ymm_d[4], ymm_4_25) VFNMADD231PS(ymm_wd[2], ymm_d[3], ymm_4_25) ymm_1_25 = YMMRegister() VMOVAPS(ymm_1_25, Constant.float32x8(1.25)) # Compute # wd3 := (d6 + 0.25 * d2) - 1.25 * d4 # wd4 := (d5 + 0.25 * d1) - 1.25 * d3 # wd6 := (d1 + 0.25 * d5) - 1.25 * d3 # wd5 := (d6 - 5.0 * d4) + 4.0 * d2 VFNMADD231PS(ymm_wd[3], ymm_d[4], ymm_1_25) VFNMADD231PS(ymm_wd[4], ymm_d[3], ymm_1_25) VFMADD231PS(ymm_wd[5], ymm_d[2], Constant.float32x8(4.0)) VFNMADD231PS(ymm_wd[6], ymm_d[3], ymm_1_25) ymm_2 = YMMRegister() VMOVAPS(ymm_2, Constant.float32x8(2.0)) butterfly(ymm_wd[1], ymm_wd[2]) butterfly(ymm_wd[3], ymm_wd[4], scale_b=ymm_2) butterfly(ymm_wd[5], ymm_wd[6], scale_b=ymm_2) return ymm_wd
def ifft8_across_rows(ymm_data, bias=None): assert isinstance(ymm_data, list) and len(ymm_data) == 8 ymm_real = ymm_data[0::2] ymm_imag = ymm_data[1::2] if bias is None: # Do 1/N scaling before IFFT ymm_one_eighth = YMMRegister() VMOVAPS(ymm_one_eighth, Constant.float32x8(0.125)) for ymm_row in ymm_data: if ymm_row is ymm_real[2]: VMULPS(ymm_row, ymm_row, Constant.float32x8(0.25)) elif ymm_row is ymm_imag[2]: VMULPS(ymm_row, ymm_row, Constant.float32x8(-0.25)) else: VMULPS(ymm_row, ymm_row, ymm_one_eighth) else: # Do 1/N scaling after FFT (merge with bias addition) VMULPS(ymm_real[2], ymm_real[2], Constant.float32x8(2.0)) VMULPS(ymm_imag[2], ymm_imag[2], Constant.float32x8(-2.0)) butterfly(ymm_real[0], ymm_imag[0]) # H1.real, H1.imag = W1.real - W3.real, W1.imag + W3.imag ymm_h1_real, ymm_h1_imag = YMMRegister(), YMMRegister() VSUBPS(ymm_h1_real, ymm_real[1], ymm_real[3]) VADDPS(ymm_h1_imag, ymm_imag[1], ymm_imag[3]) # G1.real, G1.imag = W1.real + W3.real, W1.imag - W3.imag ymm_g1_real, ymm_g1_imag = YMMRegister(), YMMRegister() VADDPS(ymm_g1_real, ymm_real[1], ymm_real[3]) VSUBPS(ymm_g1_imag, ymm_imag[1], ymm_imag[3]) # H1+, H1- = H1.real + H1.imag, H1.real - H1.imag ymm_h1_plus, ymm_h1_minus = YMMRegister(), YMMRegister() VADDPS(ymm_h1_plus, ymm_h1_real, ymm_h1_imag) VSUBPS(ymm_h1_minus, ymm_h1_real, ymm_h1_imag) ymm_sqrt2_over_2 = YMMRegister() VMOVAPS(ymm_sqrt2_over_2, Constant.float32x8(sqrt2_over_2)) # w1.real = G1.real - SQRT2_OVER_2 * H1.plus; # w3.real = G1.real + SQRT2_OVER_2 * H1.plus; VMOVAPS(ymm_real[1], ymm_g1_real) VFNMADD231PS(ymm_real[1], ymm_h1_plus, ymm_sqrt2_over_2) VFMADD231PS(ymm_g1_real, ymm_h1_plus, ymm_sqrt2_over_2) SWAP.REGISTERS(ymm_real[3], ymm_g1_real) # w1.imag = G1.imag + SQRT2_OVER_2 * H1.minus; # w3.imag = -G1.imag + SQRT2_OVER_2 * H1.minus; VMOVAPS(ymm_imag[1], ymm_g1_imag) VFMADD231PS(ymm_imag[1], ymm_h1_minus, ymm_sqrt2_over_2) VFMSUB231PS(ymm_g1_imag, ymm_h1_minus, ymm_sqrt2_over_2) SWAP.REGISTERS(ymm_imag[3], ymm_g1_imag) fft.complex_soa.fft4_across_rows(ymm_real, ymm_imag, transformation="inverse") if bias is not None: ymm_bias = bias if not isinstance(bias, YMMRegister): ymm_bias = YMMRegister() VMOVAPS(ymm_bias, bias) ymm_one_eighth = YMMRegister() VMOVAPS(ymm_one_eighth, Constant.float32x8(0.125)) # 1/N scaling for ymm_row in ymm_data: VFMADD132PS(ymm_row, ymm_bias, ymm_one_eighth)
def input_transform(ymm_d): assert isinstance(ymm_d, list) and len(ymm_d) == 8 and all( isinstance(ymm, YMMRegister) for ymm in ymm_d) ymm_wd = [YMMRegister() for _ in range(8)] # wd0 = (d0 - d6) + 5.25 * (d4 - d2) # wd1 = (d6 + d2 - 4.25 * d4) + (d1 + d5 - 4.25 * d3) # wd2 = (d6 + d2 - 4.25 * d4) - (d1 + d5 - 4.25 * d3) # wd3 = (d6 + 0.25 * d2 - 1.25 * d4) + 2.0 * (d5 + 0.25 * d1 - 1.25 * d3) # wd4 = (d6 + 0.25 * d2 - 1.25 * d4) - 2.0 * (d5 + 0.25 * d1 - 1.25 * d3) # wd5 = (d6 - 5.0 * d4 + 4.0 * d2) + 2.0 * (d1 + 0.25 * d5 - 1.25 * d3) # wd6 = (d6 - 5.0 * d4 + 4.0 * d2) - 2.0 * (d1 + 0.25 * d5 - 1.25 * d3) # wd7 = (d7 - d1) + 5.25 * (d3 - d5) ymm_0_25 = YMMRegister() VMOVAPS(ymm_0_25, Constant.float32x8(0.25)) # Compute wd0 := d0 - d6 VSUBPS(ymm_wd[0], ymm_d[0], ymm_d[6]) ymm_d4_sub_d2 = YMMRegister() VSUBPS(ymm_d4_sub_d2, ymm_d[4], ymm_d[2]) # Compute wd7 := d7 - d1 VSUBPS(ymm_wd[7], ymm_d[7], ymm_d[1]) ymm_d3_sub_d5 = YMMRegister() VSUBPS(ymm_d3_sub_d5, ymm_d[3], ymm_d[5]) # Compute wd1 := d2 + d6 VADDPS(ymm_wd[1], ymm_d[2], ymm_d[6]) # Compute wd2 := d1 + d5 VADDPS(ymm_wd[2], ymm_d[1], ymm_d[5]) # Compute wd4 := d5 + 0.25 * d1 VMOVAPS(ymm_wd[4], ymm_d[5]) VFMADD231PS(ymm_wd[4], ymm_d[1], ymm_0_25) # Compute wd5 := d6 - 5.0 * d4 VMOVAPS(ymm_wd[5], Constant.float32x8(5.0)) VFNMADD132PS(ymm_wd[5], ymm_d[6], ymm_d[4]) # Compute wd3 := d6 + 0.25 * d2 VFMADD231PS(ymm_d[6], ymm_d[2], ymm_0_25) SWAP.REGISTERS(ymm_wd[3], ymm_d[6]) # Compute wd6 := d1 + 0.25 * d5 VFMADD231PS(ymm_d[1], ymm_d[5], ymm_0_25) SWAP.REGISTERS(ymm_wd[6], ymm_d[1]) ymm_5_25 = YMMRegister() VMOVAPS(ymm_5_25, Constant.float32x8(5.25)) # Compute wd0 := (d0 - d6) + 5.25 * (d4 - d2) VFMADD231PS(ymm_wd[0], ymm_d4_sub_d2, ymm_5_25) # Compute wd7 := (d7 - d1) + 5.25 * (d3 - d5) VFMADD231PS(ymm_wd[7], ymm_d3_sub_d5, ymm_5_25) ymm_4_25 = YMMRegister() VMOVAPS(ymm_4_25, Constant.float32x8(4.25)) # Compute # wd1 := (d6 + d2) - 4.25 * d4 # wd2 := (d1 + d5) - 4.25 * d3 VFNMADD231PS(ymm_wd[1], ymm_d[4], ymm_4_25) VFNMADD231PS(ymm_wd[2], ymm_d[3], ymm_4_25) ymm_1_25 = YMMRegister() VMOVAPS(ymm_1_25, Constant.float32x8(1.25)) # Compute # wd3 := (d6 + 0.25 * d2) - 1.25 * d4 # wd4 := (d5 + 0.25 * d1) - 1.25 * d3 # wd6 := (d1 + 0.25 * d5) - 1.25 * d3 # wd5 := (d6 - 5.0 * d4) + 4.0 * d2 VFNMADD231PS(ymm_wd[3], ymm_d[4], ymm_1_25) VFNMADD231PS(ymm_wd[4], ymm_d[3], ymm_1_25) VFMADD231PS(ymm_wd[5], ymm_d[2], Constant.float32x8(4.0)) VFNMADD231PS(ymm_wd[6], ymm_d[3], ymm_1_25) ymm_2 = YMMRegister() VMOVAPS(ymm_2, Constant.float32x8(2.0)) butterfly(ymm_wd[1], ymm_wd[2]) butterfly(ymm_wd[3], ymm_wd[4], scale_b=ymm_2) butterfly(ymm_wd[5], ymm_wd[6], scale_b=ymm_2) return ymm_wd
def fft4_across_rows(ymm_real, ymm_imag, transformation="forward"): assert isinstance(ymm_real, list) and len(ymm_real) == 4 assert isinstance(ymm_imag, list) and len(ymm_imag) == 4 assert transformation in {"forward", "inverse"} ymm_data = sum(zip(ymm_real, ymm_imag), ()) # FFT-4 Butterfly for i in range(4): butterfly(ymm_data[i], ymm_data[i + 4]) # Multiply by FFT-4 twiddle factors SWAP.REGISTERS(ymm_real[3], ymm_imag[3]) # 2x FFT-2 Butterfly butterfly(ymm_data[0], ymm_data[2]) butterfly(ymm_data[1], ymm_data[3]) if transformation == "forward": butterfly(ymm_data[4], ymm_data[6]) butterfly(ymm_data[5], ymm_data[7], negate_b=True) else: butterfly(ymm_data[4], ymm_data[6], negate_b=True) butterfly(ymm_data[5], ymm_data[7]) # Bit reversal: not needed SWAP.REGISTERS(ymm_real[1], ymm_real[2]) SWAP.REGISTERS(ymm_imag[1], ymm_imag[2])
def ifft8_across_rows(ymm_data, bias=None): assert isinstance(ymm_data, list) and len(ymm_data) == 8 ymm_real = ymm_data[0::2] ymm_imag = ymm_data[1::2] if bias is None: # Do 1/N scaling before IFFT ymm_one_eighth = YMMRegister() VMOVAPS(ymm_one_eighth, Constant.float32x8(0.125)) for ymm_row in ymm_data: if ymm_row is ymm_real[2]: VMULPS(ymm_row, ymm_row, Constant.float32x8(0.25)) elif ymm_row is ymm_imag[2]: VMULPS(ymm_row, ymm_row, Constant.float32x8(-0.25)) else: VMULPS(ymm_row, ymm_row, ymm_one_eighth) else: # Do 1/N scaling after FFT (merge with bias addition) VMULPS(ymm_real[2], ymm_real[2], Constant.float32x8(2.0)) VMULPS(ymm_imag[2], ymm_imag[2], Constant.float32x8(-2.0)) butterfly(ymm_real[0], ymm_imag[0]) # H1.real, H1.imag = W1.real - W3.real, W1.imag + W3.imag ymm_h1_real, ymm_h1_imag = YMMRegister(), YMMRegister() VSUBPS(ymm_h1_real, ymm_real[1], ymm_real[3]) VADDPS(ymm_h1_imag, ymm_imag[1], ymm_imag[3]) # G1.real, G1.imag = W1.real + W3.real, W1.imag - W3.imag ymm_g1_real, ymm_g1_imag = YMMRegister(), YMMRegister() VADDPS(ymm_g1_real, ymm_real[1], ymm_real[3]) VSUBPS(ymm_g1_imag, ymm_imag[1], ymm_imag[3]) # H1+, H1- = H1.real + H1.imag, H1.real - H1.imag ymm_h1_plus, ymm_h1_minus = YMMRegister(), YMMRegister() VADDPS(ymm_h1_plus, ymm_h1_real, ymm_h1_imag) VSUBPS(ymm_h1_minus, ymm_h1_real, ymm_h1_imag) ymm_sqrt2_over_2 = YMMRegister() VMOVAPS(ymm_sqrt2_over_2, Constant.float32x8(sqrt2_over_2)) # w1.real = G1.real - SQRT2_OVER_2 * H1.plus; # w3.real = G1.real + SQRT2_OVER_2 * H1.plus; VMOVAPS(ymm_real[1], ymm_g1_real) VFNMADD231PS(ymm_real[1], ymm_h1_plus, ymm_sqrt2_over_2) VFMADD231PS(ymm_g1_real, ymm_h1_plus, ymm_sqrt2_over_2) SWAP.REGISTERS(ymm_real[3], ymm_g1_real) # w1.imag = G1.imag + SQRT2_OVER_2 * H1.minus; # w3.imag = -G1.imag + SQRT2_OVER_2 * H1.minus; VMOVAPS(ymm_imag[1], ymm_g1_imag) VFMADD231PS(ymm_imag[1], ymm_h1_minus, ymm_sqrt2_over_2) VFMSUB231PS(ymm_g1_imag, ymm_h1_minus, ymm_sqrt2_over_2) SWAP.REGISTERS(ymm_imag[3], ymm_g1_imag) complex_soa.fft4_across_rows(ymm_real, ymm_imag, transformation="inverse") if bias is not None: ymm_bias = bias if not isinstance(bias, YMMRegister): ymm_bias = YMMRegister() VMOVAPS(ymm_bias, bias) ymm_one_eighth = YMMRegister() VMOVAPS(ymm_one_eighth, Constant.float32x8(0.125)) # 1/N scaling for ymm_row in ymm_data: VFMADD132PS(ymm_row, ymm_bias, ymm_one_eighth)
def fft16_within_rows(ymm_real_rows, ymm_imag_rows, bit_reversal=True): if isinstance(ymm_real_rows, tuple) and isinstance(ymm_imag_rows, tuple): return fft16_within_rows([ymm_real_rows], [ymm_imag_rows]) assert isinstance(ymm_real_rows, list) and all( isinstance(ymm_real, tuple) and all( isinstance(ymm, YMMRegister) for ymm in ymm_real) for ymm_real in ymm_real_rows) assert isinstance(ymm_imag_rows, list) and all( isinstance(ymm_imag, tuple) and all( isinstance(ymm, YMMRegister) for ymm in ymm_imag) for ymm_imag in ymm_imag_rows) # FFT16: Butterfly for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1]) butterfly(ymm_imag[0], ymm_imag[1]) # FFT16: Multiplication by twiddle factors ymm_fft16_cos_twiddle_factor, ymm_fft16_sin_twiddle_factor = YMMRegister( ), YMMRegister() VMOVAPS(ymm_fft16_cos_twiddle_factor, Constant.float32x8(*cos_npi_over_8)) VMOVAPS(ymm_fft16_sin_twiddle_factor, Constant.float32x8(*sin_npi_over_8)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): ymm_new_real1, ymm_new_imag1 = YMMRegister(), YMMRegister() VMULPS(ymm_new_real1, ymm_real[1], ymm_fft16_cos_twiddle_factor) VMULPS(ymm_new_imag1, ymm_imag[1], ymm_fft16_cos_twiddle_factor) VFMADD231PS(ymm_new_real1, ymm_imag[1], ymm_fft16_sin_twiddle_factor) VFNMADD231PS(ymm_new_imag1, ymm_real[1], ymm_fft16_sin_twiddle_factor) SWAP.REGISTERS(ymm_real[1], ymm_new_real1) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag1) # 2x FFT8: Butterfly for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): transpose2x2x128(ymm_real[0], ymm_real[1]) transpose2x2x128(ymm_imag[0], ymm_imag[1]) # w[0] = x0 x1 x2 x3 x8 x9 x10 x11 # w[1] = x4 x5 x6 x7 x12 x13 x14 x15 for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1]) butterfly(ymm_imag[0], ymm_imag[1]) # 2x FFT8: Multiplication by twiddle factors ymm_fft8_cos_twiddle_factor, ymm_fft8_sin_twiddle_factor = YMMRegister( ), YMMRegister() VMOVAPS(ymm_fft8_cos_twiddle_factor, Constant.float32x8(*(cos_npi_over_4 * 2))) VMOVAPS(ymm_fft8_sin_twiddle_factor, Constant.float32x8(*(sin_npi_over_4 * 2))) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): ymm_new_real1, ymm_new_imag1 = YMMRegister(), YMMRegister() VMULPS(ymm_new_real1, ymm_real[1], ymm_fft8_cos_twiddle_factor) VMULPS(ymm_new_imag1, ymm_imag[1], ymm_fft8_cos_twiddle_factor) VFMADD231PS(ymm_new_real1, ymm_imag[1], ymm_fft8_sin_twiddle_factor) VFNMADD231PS(ymm_new_imag1, ymm_real[1], ymm_fft8_sin_twiddle_factor) SWAP.REGISTERS(ymm_real[1], ymm_new_real1) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag1) # 4x FFT4: Butterfly for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): transpose2x2x2x64(ymm_real[0], ymm_real[1]) transpose2x2x2x64(ymm_imag[0], ymm_imag[1]) # w[0] = x0 x1 x4 x5 x8 x9 x12 x13 # w[1] = x2 x3 x6 x7 x10 x11 x14 x15 for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1]) butterfly(ymm_imag[0], ymm_imag[1]) # 4x FFT4: Multiplication by twiddle factors and 8x FFT2: Butterfly ymm_fft4_twiddle_factor = YMMRegister() VMOVAPS(ymm_fft4_twiddle_factor, Constant.float32x8(+1.0, +1.0, -1.0, -1.0, +1.0, +1.0, -1.0, -1.0)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): ymm_new_real = YMMRegister(), YMMRegister() VSHUFPS(ymm_new_real[0], ymm_real[0], ymm_real[1], _MM_SHUFFLE(2, 0, 2, 0)) VSHUFPS(ymm_new_real[1], ymm_real[0], ymm_imag[1], _MM_SHUFFLE(3, 1, 3, 1)) butterfly(ymm_new_real[0], ymm_new_real[1]) ymm_new_imag = YMMRegister(), YMMRegister() VSHUFPS(ymm_new_imag[0], ymm_imag[0], ymm_imag[1], _MM_SHUFFLE(2, 0, 2, 0)) VSHUFPS(ymm_new_imag[1], ymm_imag[0], ymm_real[1], _MM_SHUFFLE(3, 1, 3, 1)) butterfly(ymm_new_imag[0], ymm_new_imag[1], scale_b=ymm_fft4_twiddle_factor) SWAP.REGISTERS(ymm_real[0], ymm_new_real[0]) SWAP.REGISTERS(ymm_real[1], ymm_new_real[1]) SWAP.REGISTERS(ymm_imag[0], ymm_new_imag[0]) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag[1]) # w[0] = x0 x4 x2 x6 x8 x12 x10 x14 # w[1] = x1 x5 x3 x7 x9 x11 x13 x15 if bit_reversal: # Bit reversal ymm_bit_reversal_mask = YMMRegister() VMOVDQA(ymm_bit_reversal_mask, Constant.uint32x8(0, 4, 1, 5, 2, 6, 3, 7)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): for i in range(2): VPERMPS(ymm_real[i], ymm_bit_reversal_mask, ymm_real[i]) VPERMPS(ymm_imag[i], ymm_bit_reversal_mask, ymm_imag[i])
def ifft16_within_rows(ymm_real_rows, ymm_imag_rows, bit_reversal=True): if isinstance(ymm_real_rows, tuple) and isinstance(ymm_imag_rows, tuple): return ifft16_within_rows([ymm_real_rows], [ymm_imag_rows]) assert isinstance(ymm_real_rows, list) and all( isinstance(ymm_real, tuple) and all( isinstance(ymm, YMMRegister) for ymm in ymm_real) for ymm_real in ymm_real_rows) assert isinstance(ymm_imag_rows, list) and all( isinstance(ymm_imag, tuple) and all( isinstance(ymm, YMMRegister) for ymm in ymm_imag) for ymm_imag in ymm_imag_rows) if bit_reversal: # Bit reversal # w[0] = x0 x8 x4 x12 x2 x10 x6 x14 # w[1] = x1 x9 x5 x13 x3 x11 x7 x15 ymm_bit_reversal_mask = YMMRegister() VMOVDQA(ymm_bit_reversal_mask, Constant.uint32x8(0, 2, 4, 6, 1, 3, 5, 7)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): for i in range(2): VPERMPS(ymm_real[i], ymm_bit_reversal_mask, ymm_real[i]) VPERMPS(ymm_imag[i], ymm_bit_reversal_mask, ymm_imag[i]) # 8x FFT2: Butterfly # w[0] = x0 x4 x2 x6 x8 x12 x10 x14 # w[1] = x1 x5 x3 x7 x9 x13 x11 x15 for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1]) butterfly(ymm_imag[0], ymm_imag[1]) ymm_new_real = YMMRegister(), YMMRegister() VUNPCKLPS(ymm_new_real[0], ymm_real[0], ymm_real[1]) VUNPCKHPS(ymm_new_real[1], ymm_real[0], ymm_imag[1]) ymm_new_imag = YMMRegister(), YMMRegister() VUNPCKLPS(ymm_new_imag[0], ymm_imag[0], ymm_imag[1]) VUNPCKHPS(ymm_new_imag[1], ymm_imag[0], ymm_real[1]) SWAP.REGISTERS(ymm_imag[0], ymm_new_imag[0]) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag[1]) SWAP.REGISTERS(ymm_real[0], ymm_new_real[0]) SWAP.REGISTERS(ymm_real[1], ymm_new_real[1]) # w[0] = x0 x1 x4 x5 x8 x9 x12 x13 # w[1] = x2 x3 x6 x7 x10 x11 x14 x15 # 4x FFT4: Butterfly and multiplication by twiddle factors ymm_fft4_twiddle_factor = YMMRegister() VMOVAPS(ymm_fft4_twiddle_factor, Constant.float32x8(+1.0, -1.0, +1.0, -1.0, +1.0, -1.0, +1.0, -1.0)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1], scale_b=ymm_fft4_twiddle_factor) butterfly(ymm_imag[0], ymm_imag[1]) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): transpose2x2x2x64(ymm_real[0], ymm_real[1]) transpose2x2x2x64(ymm_imag[0], ymm_imag[1]) # w[0] = x0 x1 x2 x3 x8 x9 x10 x11 # w[1] = x4 x5 x6 x7 x12 x13 x14 x15 # 2x FFT8: Multiplication by twiddle factors ymm_fft8_cos_twiddle_factor, ymm_fft8_sin_twiddle_factor = YMMRegister( ), YMMRegister() VMOVAPS(ymm_fft8_cos_twiddle_factor, Constant.float32x8(*(cos_npi_over_4 * 2))) VMOVAPS(ymm_fft8_sin_twiddle_factor, Constant.float32x8(*(sin_npi_over_4 * 2))) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): ymm_new_real1, ymm_new_imag1 = YMMRegister(), YMMRegister() VMULPS(ymm_new_real1, ymm_real[1], ymm_fft8_cos_twiddle_factor) VMULPS(ymm_new_imag1, ymm_imag[1], ymm_fft8_cos_twiddle_factor) VFNMADD231PS(ymm_new_real1, ymm_imag[1], ymm_fft8_sin_twiddle_factor) VFMADD231PS(ymm_new_imag1, ymm_real[1], ymm_fft8_sin_twiddle_factor) SWAP.REGISTERS(ymm_real[1], ymm_new_real1) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag1) # 2x FFT8: Butterfly for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1]) butterfly(ymm_imag[0], ymm_imag[1]) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): transpose2x2x128(ymm_real[0], ymm_real[1]) transpose2x2x128(ymm_imag[0], ymm_imag[1]) # w[0] = x0 x1 x2 x3 x4 x5 x6 x7 # w[1] = x8 x9 x10 x11 x12 x13 x14 x15 # FFT16: Multiplication by twiddle factors and scale scale_factor = 0.0625 ymm_fft16_cos_scale_twiddle_factor, ymm_fft16_sin_scale_twiddle_factor = YMMRegister( ), YMMRegister() VMOVAPS( ymm_fft16_cos_scale_twiddle_factor, Constant.float32x8(*[cos * scale_factor for cos in cos_npi_over_8])) VMOVAPS( ymm_fft16_sin_scale_twiddle_factor, Constant.float32x8(*[sin * scale_factor for sin in sin_npi_over_8])) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): ymm_new_real1, ymm_new_imag1 = YMMRegister(), YMMRegister() VMULPS(ymm_new_real1, ymm_real[1], ymm_fft16_cos_scale_twiddle_factor) VMULPS(ymm_new_imag1, ymm_imag[1], ymm_fft16_cos_scale_twiddle_factor) VFNMADD231PS(ymm_new_real1, ymm_imag[1], ymm_fft16_sin_scale_twiddle_factor) VFMADD231PS(ymm_new_imag1, ymm_real[1], ymm_fft16_sin_scale_twiddle_factor) SWAP.REGISTERS(ymm_real[1], ymm_new_real1) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag1) # FFT16: Butterfly and scale ymm_scale_factor = YMMRegister() VMOVAPS(ymm_scale_factor, Constant.float32x8(scale_factor)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1], scale_a=ymm_scale_factor) butterfly(ymm_imag[0], ymm_imag[1], scale_a=ymm_scale_factor)
def inverse_vfft(reg_t0, reg_t8, reg_t_stride, data_in, reg_row_start=None, reg_row_end=None, store_mask=None): assert isinstance(reg_t0, GeneralPurposeRegister64) assert isinstance(reg_t8, GeneralPurposeRegister64) assert isinstance(reg_t_stride, GeneralPurposeRegister64) assert isinstance(data_in, list) and len(data_in) == 16 assert reg_row_end is None or isinstance(reg_row_end, GeneralPurposeRegister32) assert store_mask is None or isinstance(store_mask, LocalVariable) and store_mask.size == YMMRegister.size in_real, in_imag = data_in[0::2], data_in[1::2] ymm_scale_factor = YMMRegister() VMOVAPS(ymm_scale_factor, Constant.float32x8(0.0625)) ymm_W1_real, ymm_W1_imag = YMMRegister(), YMMRegister() VMULPS(ymm_W1_real, ymm_scale_factor, in_real[1]) VMULPS(ymm_W1_imag, ymm_scale_factor, in_imag[1]) ymm_W2_real, ymm_W2_imag = YMMRegister(), YMMRegister() VMULPS(ymm_W2_real, ymm_scale_factor, in_real[2]) VMULPS(ymm_W2_imag, ymm_scale_factor, in_imag[2]) ymm_W3_real, ymm_W3_imag = YMMRegister(), YMMRegister() VMULPS(ymm_W3_real, ymm_scale_factor, in_real[3]) VMULPS(ymm_W3_imag, ymm_scale_factor, in_imag[3]) # G[n].real, H[n].real = W[n].real + W[8-n].real, W[n].real - W[8-n].real # G[n].imag, H[n].imag = W[n].imag - W[8-n].imag, W[n].imag + W[8-n].imag ymm_W7_real, ymm_W7_imag = YMMRegister(), YMMRegister() VMOVUPS(ymm_W7_real, in_real[7]) ymm_G1_real, ymm_H1_real = butterfly(ymm_W1_real, ymm_W7_real, scale_b=ymm_scale_factor) VMOVUPS(ymm_W7_imag, in_imag[7]) ymm_G1_imag, ymm_H1_imag = butterfly(ymm_W1_imag, ymm_W7_imag, scale_b=ymm_scale_factor, negate_b=True) ymm_W6_real, ymm_W6_imag = YMMRegister(), YMMRegister() VMOVUPS(ymm_W6_real, in_real[6]) ymm_G2_real, ymm_H2_real = butterfly(ymm_W2_real, ymm_W6_real, scale_b=ymm_scale_factor) VMOVUPS(ymm_W6_imag, in_imag[6]) ymm_G2_imag, ymm_H2_imag = butterfly(ymm_W2_imag, ymm_W6_imag, scale_b=ymm_scale_factor, negate_b=True) ymm_W5_real, ymm_W5_imag = YMMRegister(), YMMRegister() VMOVUPS(ymm_W5_real, in_real[5]) ymm_G3_real, ymm_H3_real = butterfly(ymm_W3_real, ymm_W5_real, scale_b=ymm_scale_factor) VMOVUPS(ymm_W5_imag, in_imag[5]) ymm_G3_imag, ymm_H3_imag = butterfly(ymm_W3_imag, ymm_W5_imag, scale_b=ymm_scale_factor, negate_b=True) # H[2]+, H[2]- = H[2].real + H[2].imag, H[2].real - H[2].imag ymm_H2_add, ymm_H2_sub = butterfly(ymm_H2_real, ymm_H2_imag) # w[ n].real = G[ n].real - H[ n].real * cos((N-n)*pi/2N) - H[ n].imag * cos(n*pi/2N) # w[2N-n].real = G[ n].real + H[ n].real * cos((N-n)*pi/2N) + H[ n].imag * cos(n*pi/2N) # w[ n].imag = G[ n].imag + H[ n].real * cos(n*pi/2N) - H[ n].imag * cos((N-n)*pi/2N) # w[2N-n].imag = -G[ n].imag + H[ n].real * cos(n*pi/2N) - H[ n].imag * cos((N-n)*pi/2N) # w[ N-n].real = G[N-n].real - H[N-n].real * cos(n*pi/2N) - H[N-n].imag * cos((N-n)*pi/2N) # w[ N+n].real = G[N-n].real + H[N-n].real * cos(n*pi/2N) + H[N-n].imag * cos((N-n)*pi/2N) # w[ N-n].imag = G[N-n].imag + H[N-n].real * cos((N-n)*pi/2N) - H[N-n].imag * cos(n*pi/2N) # w[ N+n].imag = -G[N-n].imag + H[N-n].real * cos((N-n)*pi/2N) - H[N-n].imag * cos(n*pi/2N) ymm_cos_1pi_over_8, ymm_cos_3pi_over_8 = YMMRegister(), YMMRegister() VMOVAPS(ymm_cos_3pi_over_8, Constant.float32x8(cos_npi_over_8[3])) VMOVAPS(ymm_cos_1pi_over_8, Constant.float32x8(cos_npi_over_8[1])) ymm_w1_real, ymm_w7_real = YMMRegister(), ymm_G1_real VMOVAPS(ymm_w1_real, ymm_G1_real) VFNMADD231PS(ymm_w1_real, ymm_H1_real, ymm_cos_3pi_over_8) VFMADD231PS(ymm_w7_real, ymm_H1_real, ymm_cos_3pi_over_8) ymm_w1_imag, ymm_w7_imag = YMMRegister(), ymm_G1_imag VMOVAPS(ymm_w1_imag, ymm_G1_imag) VFMADD231PS(ymm_w1_imag, ymm_H1_real, ymm_cos_1pi_over_8) VFMSUB231PS(ymm_w7_imag, ymm_H1_real, ymm_cos_1pi_over_8) ymm_w3_real, ymm_w5_real = YMMRegister(), ymm_G3_real VMOVAPS(ymm_w3_real, ymm_G3_real) VFNMADD231PS(ymm_w3_real, ymm_H3_real, ymm_cos_1pi_over_8) VFMADD231PS(ymm_w5_real, ymm_H3_real, ymm_cos_1pi_over_8) ymm_w3_imag, ymm_w5_imag = YMMRegister(), ymm_G3_imag VMOVAPS(ymm_w3_imag, ymm_G3_imag) VFMADD231PS(ymm_w3_imag, ymm_H3_real, ymm_cos_3pi_over_8) VFMSUB231PS(ymm_w5_imag, ymm_H3_real, ymm_cos_3pi_over_8) ymm_sqrt2_over_2 = YMMRegister() VMOVAPS(ymm_sqrt2_over_2, Constant.float32x8(sqrt2_over_2)) # w[ N/2].real = G[N/2].real - H[N/2]+ * sqrt(2)/2 # w[ N/2].imag = G[N/2].imag + H[N/2]- * sqrt(2)/2 # w[3N/2].real = G[N/2].real + H[N/2]+ * sqrt(2)/2 # w[3N/2].imag = -G[N/2].imag + H[N/2]- * sqrt(2)/2 ymm_w2_real, ymm_w6_real = YMMRegister(), ymm_G2_real VMOVAPS(ymm_w2_real, ymm_G2_real) VFNMADD231PS(ymm_w2_real, ymm_H2_add, ymm_sqrt2_over_2) VFMADD231PS(ymm_w6_real, ymm_H2_add, ymm_sqrt2_over_2) ymm_w2_imag, ymm_w6_imag = YMMRegister(), ymm_G2_imag VMOVAPS(ymm_w2_imag, ymm_G2_imag) VFMADD231PS(ymm_w2_imag, ymm_H2_sub, ymm_sqrt2_over_2) VFMSUB231PS(ymm_w6_imag, ymm_H2_sub, ymm_sqrt2_over_2) # w[ n].real = G[ n].real - H[ n].real * cos((N-n)*pi/2N) - H[ n].imag * cos(n*pi/2N) # w[2N-n].real = G[ n].real + H[ n].real * cos((N-n)*pi/2N) + H[ n].imag * cos(n*pi/2N) # w[ n].imag = G[ n].imag + H[ n].real * cos(n*pi/2N) - H[ n].imag * cos((N-n)*pi/2N) # w[2N-n].imag = -G[ n].imag + H[ n].real * cos(n*pi/2N) - H[ n].imag * cos((N-n)*pi/2N) # w[ N-n].real = G[N-n].real - H[N-n].real * cos(n*pi/2N) - H[N-n].imag * cos((N-n)*pi/2N) # w[ N+n].real = G[N-n].real + H[N-n].real * cos(n*pi/2N) + H[N-n].imag * cos((N-n)*pi/2N) # w[ N-n].imag = G[N-n].imag + H[N-n].real * cos((N-n)*pi/2N) - H[N-n].imag * cos(n*pi/2N) # w[ N+n].imag = -G[N-n].imag + H[N-n].real * cos((N-n)*pi/2N) - H[N-n].imag * cos(n*pi/2N) ymm_cos_1pi_over_8, ymm_cos_3pi_over_8 = YMMRegister(), YMMRegister() VMOVAPS(ymm_cos_1pi_over_8, Constant.float32x8(cos_npi_over_8[1])) VMOVAPS(ymm_cos_3pi_over_8, Constant.float32x8(cos_npi_over_8[3])) VFNMADD231PS(ymm_w1_real, ymm_H1_imag, ymm_cos_1pi_over_8) VFMADD231PS(ymm_w7_real, ymm_H1_imag, ymm_cos_1pi_over_8) VFNMADD231PS(ymm_w1_imag, ymm_H1_imag, ymm_cos_3pi_over_8) VFNMADD231PS(ymm_w7_imag, ymm_H1_imag, ymm_cos_3pi_over_8) VFNMADD231PS(ymm_w3_real, ymm_H3_imag, ymm_cos_3pi_over_8) VFMADD231PS(ymm_w5_real, ymm_H3_imag, ymm_cos_3pi_over_8) VFNMADD231PS(ymm_w3_imag, ymm_H3_imag, ymm_cos_1pi_over_8) VFNMADD231PS(ymm_w5_imag, ymm_H3_imag, ymm_cos_1pi_over_8) data = [ LocalVariable(YMMRegister.size), YMMRegister(), ymm_w1_real, ymm_w1_imag, ymm_w2_real, ymm_w2_imag, ymm_w3_real, ymm_w3_imag, LocalVariable(YMMRegister.size), LocalVariable(YMMRegister.size), ymm_w5_real, ymm_w5_imag, ymm_w6_real, ymm_w6_imag, ymm_w7_real, ymm_w7_imag ] real, imag = data[0::2], data[1::2] # TODO: optimize ymm_w0_real, ymm_w0_imag = YMMRegister(), imag[0] VMOVUPS(ymm_w0_real, in_real[0]) VMOVUPS(ymm_w0_imag, in_imag[0]) VMULPS(ymm_w0_real, ymm_w0_real, Constant.float32x8(0.0625)) butterfly(ymm_w0_real, ymm_w0_imag, scale_b=Constant.float32x8(0.0625)) VMOVAPS(real[0], ymm_w0_real) # TODO: optimize ymm_w4_real, ymm_w4_imag = YMMRegister(), YMMRegister() VMOVUPS(ymm_w4_real, in_real[4]) VMOVUPS(ymm_w4_imag, in_imag[4]) VMULPS(ymm_w4_real, ymm_w4_real, Constant.float32x8(0.125)) VMULPS(ymm_w4_imag, ymm_w4_imag, Constant.float32x8(-0.125)) VMOVAPS(real[4], ymm_w4_real) VMOVAPS(imag[4], ymm_w4_imag) # Bit reversal for i in range(8): new_i = fft8_bitreverse(i) if new_i > i: real[i], real[new_i] = real[new_i], real[i] imag[i], imag[new_i] = imag[new_i], imag[i] data = interleave(real, imag) # 4x FFT2: butterfly for i, (data_lo, data_hi) in enumerate(zip(data[0:2] + data[4:6] + data[8:10] + data[12:14], data[2:4] + data[6:8] + data[10:12] + data[14:16])): butterfly(data_lo, data_hi) # 2x FFT4: multiplication by twiddle factors fft4_scale_b, fft4_negate_b = {}, {} fft8_scale_b, fft8_negate_b = {}, {} # w3.re, w3.im = -w3.im, w3.re # w7.re, w7.im = -w7.im, w7.re SWAP.REGISTERS(real[3], imag[3]) fft4_negate_b[id(real[3])] = True SWAP.REGISTERS(real[7], imag[7]) fft4_negate_b[id(real[7])] = True # 2x FFT4: butterfly for data_lo, data_hi in zip(data[0:4] + data[8:12], data[4:8] + data[12:16]): butterfly(data_lo, data_hi, negate_b=fft4_negate_b.get(id(data_hi), False)) # FFT8: multiplication by twiddle factors # w6.re, w6.im = -w6.im, w6.re SWAP.REGISTERS(real[6], imag[6]) fft8_negate_b[id(real[6])] = True # w5.re, w5.im = SQRT2_OVER_2 * (w5.re - w5.im), SQRT2_OVER_2 * (w5.re + w5.im) butterfly(real[5], imag[5], negate_b=True) fft8_scale_b[id(real[5])] = Constant.float32x8(sqrt2_over_2) fft8_scale_b[id(imag[5])] = Constant.float32x8(sqrt2_over_2) # w7.re, w7.im = -SQRT2_OVER_2 * (w7.re + w7.im), SQRT2_OVER_2 * (w7.re - w7.im) butterfly(real[7], imag[7]) fft8_scale_b[id(real[7])] = Constant.float32x8(sqrt2_over_2) fft8_negate_b[id(real[7])] = True fft8_scale_b[id(imag[7])] = Constant.float32x8(sqrt2_over_2) ymm_store_mask = YMMRegister() if store_mask: VMOVAPS(ymm_store_mask, store_mask) # FFT8: butterfly with Block() as store_data: for i, (data_lo, data_hi) in enumerate(zip(data[0:8], data[8:16])): row_lo = i row_hi = row_lo + 8 ymm_data_lo, ymm_data_hi = \ butterfly(data_lo, data_hi, scale_b=fft8_scale_b.get(id(data_hi)), negate_b=fft8_negate_b.get(id(data_hi), False), writeback=False) with Block() as store_data_lo: if reg_row_start: CMP(reg_row_start, row_lo) JA(store_data_lo.end) if reg_row_end: CMP(reg_row_end, row_lo) JBE(store_data_lo.end) elif reg_row_end: CMP(reg_row_end, row_lo) JBE(store_data.end) if store_mask: VMASKMOVPS([reg_t0], ymm_store_mask, ymm_data_lo) else: VMOVUPS([reg_t0], ymm_data_lo) if i + 1 != 8: ADD(reg_t0, reg_t_stride) with Block() as store_data_hi: if reg_row_start: CMP(reg_row_start, row_hi) JA(store_data_hi.end) if reg_row_end: CMP(reg_row_end, row_hi) JBE(store_data_hi.end) if store_mask: VMASKMOVPS([reg_t8], ymm_store_mask, ymm_data_hi) else: VMOVUPS([reg_t8], ymm_data_hi) if i + 1 != 8: ADD(reg_t8, reg_t_stride)
def forward_vfft(reg_t0, reg_t8, reg_t_stride, data_out, reg_row_start=None, reg_row_end=None, ymm_load_mask=None): assert isinstance(reg_t0, GeneralPurposeRegister64) assert isinstance(reg_t8, GeneralPurposeRegister64) assert isinstance(reg_t_stride, GeneralPurposeRegister64) assert isinstance(data_out, list) and len(data_out) == 16 assert ymm_load_mask is None or isinstance(ymm_load_mask, YMMRegister) out_real, out_imag = data_out[0::2], data_out[1::2] real, imag = [YMMRegister() for _ in range(8)], [YMMRegister() for _ in range(8)] imag[0] = LocalVariable(YMMRegister.size) imag[4] = LocalVariable(YMMRegister.size) data = interleave(real, imag) for i, (data_lo, data_hi) in enumerate(zip(data[0:8], data[8:16])): row_lo = i row_hi = row_lo + 8 ymm_data_lo, ymm_data_hi = data_lo, data_hi if isinstance(data_lo, LocalVariable): ymm_data_lo = YMMRegister() if isinstance(data_hi, LocalVariable): ymm_data_hi = YMMRegister() VXORPS(ymm_data_lo, ymm_data_lo, ymm_data_lo) skip_data_lo = Label() if reg_row_start: CMP(reg_row_start, row_lo) JA(skip_data_lo) if reg_row_end: CMP(reg_row_end, row_lo) JBE(skip_data_lo) if ymm_load_mask is None: VMOVUPS(ymm_data_lo, [reg_t0]) else: VMASKMOVPS(ymm_data_lo, ymm_load_mask, [reg_t0]) if i + 1 != 8: ADD(reg_t0, reg_t_stride) LABEL(skip_data_lo) VMOVAPS(ymm_data_hi, ymm_data_lo) skip_data_hi = Label() if reg_row_start: CMP(reg_row_start, row_hi) JA(skip_data_hi) if reg_row_end: CMP(reg_row_end, row_hi) JBE(skip_data_hi) if ymm_load_mask is None: VMOVUPS(ymm_data_hi, [reg_t8]) butterfly(ymm_data_lo, ymm_data_hi) else: ymm_temp_hi = YMMRegister() VMASKMOVPS(ymm_temp_hi, ymm_load_mask, [reg_t8]) VSUBPS(ymm_data_hi, ymm_data_lo, ymm_temp_hi) VADDPS(ymm_data_lo, ymm_data_lo, ymm_temp_hi) if i + 1 != 8: ADD(reg_t8, reg_t_stride) LABEL(skip_data_hi) if isinstance(data_lo, LocalVariable): VMOVAPS(data_lo, ymm_data_lo) if isinstance(data_hi, LocalVariable): VMOVAPS(data_hi, ymm_data_hi) # FFT8: multiplication by twiddle factors fft4_scale_b, fft4_negate_b = {}, {} fft2_scale_b, fft2_negate_b = {}, {} # w6.re, w6.im = w6.im, -w6.re SWAP.REGISTERS(real[6], imag[6]) fft4_negate_b[id(imag[6])] = True # w5.re, w5.im = SQRT2_OVER_2 * (w5.re + w5.im), SQRT2_OVER_2 * (w5.im - w5.re) butterfly(imag[5], real[5]) SWAP.REGISTERS(real[5], imag[5]) # w7.re, w7.im = -SQRT2_OVER_2 * (w7.re - w7.im), -SQRT2_OVER_2 * (w7.re + w7.im) butterfly(real[7], imag[7], negate_b=True) fft4_negate_b[id(real[7])] = True fft4_negate_b[id(imag[7])] = True # Propogate multiplication by sqrt2_over_2 until the last butterfly in FFT2 ymm_sqrt2_over_2 = YMMRegister() fft2_scale_b[id(real[5])] = ymm_sqrt2_over_2 fft2_scale_b[id(imag[5])] = ymm_sqrt2_over_2 fft2_scale_b[id(real[7])] = ymm_sqrt2_over_2 fft2_scale_b[id(imag[7])] = ymm_sqrt2_over_2 # 2x FFT4: butterfly for data_lo, data_hi in zip(data[0:4] + data[8:12], data[4:8] + data[12:16]): butterfly(data_lo, data_hi, negate_b=fft4_negate_b.get(id(data_hi), False), scale_b=fft4_scale_b.get(id(data_hi))) # 2x FFT4: multiplication by twiddle factors # w3.re, w3.im = w3.im, -w3.re # w7.re, w7.im = w7.im, -w7.re SWAP.REGISTERS(real[3], imag[3]) SWAP.REGISTERS(real[7], imag[7]) fft2_negate_b[id(imag[3])] = True fft2_negate_b[id(imag[7])] = True # 4x FFT2: butterfly # Process the first two elements separately ymm_real0, ymm_real1 = butterfly(real[0], real[1], writeback=False) store_ymm_result(out_real[4], ymm_real1) # bit-reversal: 1->4 ymm_imag0, ymm_imag1 = butterfly(imag[0], imag[1], negate_out_b=True, writeback=False) store_ymm_result(out_imag[4], ymm_imag1) # bit-reversal: 1->4 VMOVAPS(ymm_sqrt2_over_2, Constant.float32x8(sqrt2_over_2)) for i, (data_lo, data_hi) in enumerate(zip(data[4:6] + data[8:10] + data[12:14], data[6:8] + data[10:12] + data[14:16])): butterfly(data_lo, data_hi, negate_b=fft2_negate_b.get(id(data_hi), False), scale_b=fft2_scale_b.get(id(data_hi))) butterfly(ymm_real0, ymm_imag0) store_ymm_result(out_real[0], ymm_real0) store_ymm_result(out_imag[0], ymm_imag0) # Bit reversal for i in range(8): new_i = fft8_bitreverse(i) if new_i > i: real[i], real[new_i] = real[new_i], real[i] imag[i], imag[new_i] = imag[new_i], imag[i] data = interleave(real, imag) ymm_two_g2_real, ymm_two_g2_imag = YMMRegister(), YMMRegister() ymm_two_h2_real, ymm_two_h2_imag = YMMRegister(), YMMRegister() VADDPS(ymm_two_g2_real, real[2], real[6]) VSUBPS(ymm_two_h2_imag, real[6], real[2]) VSUBPS(ymm_two_g2_imag, imag[2], imag[6]) VADDPS(ymm_two_h2_real, imag[2], imag[6]) ymm_two_g1_real, ymm_two_g1_imag = YMMRegister(), YMMRegister() ymm_two_h1_real, ymm_two_h1_imag = YMMRegister(), YMMRegister() ymm_real1 = load_ymm_variable(real[1]) VADDPS(ymm_two_g1_real, ymm_real1, real[7]) VSUBPS(ymm_two_h1_imag, real[7], ymm_real1) ymm_imag1 = load_ymm_variable(imag[1]) VSUBPS(ymm_two_g1_imag, ymm_imag1, imag[7]) VADDPS(ymm_two_h1_real, ymm_imag1, imag[7]) ymm_two_h2_add, ymm_two_h2_sub = YMMRegister(), YMMRegister() VADDPS(ymm_two_h2_add, ymm_two_h2_real, ymm_two_h2_imag) VSUBPS(ymm_two_h2_sub, ymm_two_h2_imag, ymm_two_h2_real) ymm_two_g3_real, ymm_two_g3_imag = YMMRegister(), YMMRegister() ymm_two_h3_real, ymm_two_h3_imag = YMMRegister(), YMMRegister() VADDPS(ymm_two_g3_real, real[3], real[5]) VSUBPS(ymm_two_h3_imag, real[5], real[3]) VSUBPS(ymm_two_g3_imag, imag[3], imag[5]) VADDPS(ymm_two_h3_real, imag[3], imag[5]) # const float two_w2_real = two_g2_real + SQRT2_OVER_2 * (two_h2_real + two_h2_imag); # const float two_w2_imag = two_g2_imag + SQRT2_OVER_2 * (two_h2_imag - two_h2_real); # const float two_w6_real = two_g2_real - SQRT2_OVER_2 * (two_h2_real + two_h2_imag); # const float two_w6_imag = -two_g2_imag + SQRT2_OVER_2 * (two_h2_imag - two_h2_real); ymm_sqrt2_over_2 = YMMRegister() VMOVAPS(ymm_sqrt2_over_2, Constant.float32x8(sqrt2_over_2)) ymm_two_w2_real, ymm_two_w6_real = YMMRegister(), ymm_two_g2_real VMOVAPS(ymm_two_w2_real, ymm_two_g2_real) VFMADD231PS(ymm_two_w2_real, ymm_two_h2_add, ymm_sqrt2_over_2) VFNMADD231PS(ymm_two_w6_real, ymm_two_h2_add, ymm_sqrt2_over_2) ymm_two_w2_imag, ymm_two_w6_imag = YMMRegister(), ymm_two_g2_imag VMOVAPS(ymm_two_w2_imag, ymm_two_g2_imag) VFMADD231PS(ymm_two_w2_imag, ymm_two_h2_sub, ymm_sqrt2_over_2) VFMSUB231PS(ymm_two_w6_imag, ymm_two_h2_sub, ymm_sqrt2_over_2) ymm_half = YMMRegister() VMOVAPS(ymm_half, Constant.float32x8(0.5)) VMULPS(ymm_two_w2_real, ymm_two_w2_real, ymm_half) store_ymm_result(out_real[2], ymm_two_w2_real) VMULPS(ymm_two_w6_real, ymm_two_w6_real, ymm_half) store_ymm_result(out_real[6], ymm_two_w6_real) VMULPS(ymm_two_w2_imag, ymm_two_w2_imag, ymm_half) store_ymm_result(out_imag[2], ymm_two_w2_imag) VMULPS(ymm_two_w6_imag, ymm_two_w6_imag, ymm_half) store_ymm_result(out_imag[6], ymm_two_w6_imag) # const float two_w1_real = two_g1_real + two_h1_real * COS_1PI_OVER_8 + two_h1_imag * COS_3PI_OVER_8; # const float two_w1_imag = two_g1_imag + two_h1_imag * COS_1PI_OVER_8 - two_h1_real * COS_3PI_OVER_8; # const float two_w7_real = two_g1_real - two_h1_real * COS_1PI_OVER_8 - two_h1_imag * COS_3PI_OVER_8; # const float two_w7_imag = -two_g1_imag + two_h1_imag * COS_1PI_OVER_8 - two_h1_real * COS_3PI_OVER_8; # const float two_w3_real = two_g3_real + two_h3_real * COS_3PI_OVER_8 + two_h3_imag * COS_1PI_OVER_8; # const float two_w3_imag = two_g3_imag + two_h3_imag * COS_3PI_OVER_8 - two_h3_real * COS_1PI_OVER_8; # const float two_w5_real = two_g3_real - two_h3_real * COS_3PI_OVER_8 - two_h3_imag * COS_1PI_OVER_8; # const float two_w5_imag = -two_g3_imag + two_h3_imag * COS_3PI_OVER_8 - two_h3_real * COS_1PI_OVER_8; ymm_cos_1pi_over_8 = YMMRegister() VMOVAPS(ymm_cos_1pi_over_8, Constant.float32x8(cos_npi_over_8[1])) ymm_two_w1_real, ymm_two_w7_real = YMMRegister(), ymm_two_g1_real VMOVAPS(ymm_two_w1_real, ymm_two_g1_real) VFMADD231PS(ymm_two_w1_real, ymm_two_h1_real, ymm_cos_1pi_over_8) VFNMADD231PS(ymm_two_w7_real, ymm_two_h1_real, ymm_cos_1pi_over_8) ymm_two_w1_imag, ymm_two_w7_imag = YMMRegister(), ymm_two_g1_imag VMOVAPS(ymm_two_w1_imag, ymm_two_g1_imag) VFMADD231PS(ymm_two_w1_imag, ymm_two_h1_imag, ymm_cos_1pi_over_8) VFMSUB231PS(ymm_two_w7_imag, ymm_two_h1_imag, ymm_cos_1pi_over_8) ymm_two_w3_real, ymm_two_w5_real = YMMRegister(), ymm_two_g3_real VMOVAPS(ymm_two_w3_real, ymm_two_g3_real) VFMADD231PS(ymm_two_w3_real, ymm_two_h3_imag, ymm_cos_1pi_over_8) VFNMADD231PS(ymm_two_w5_real, ymm_two_h3_imag, ymm_cos_1pi_over_8) ymm_two_w3_imag, ymm_two_w5_imag = YMMRegister(), ymm_two_g3_imag VMOVAPS(ymm_two_w3_imag, ymm_two_g3_imag) VFNMADD231PS(ymm_two_w3_imag, ymm_two_h3_real, ymm_cos_1pi_over_8) VFNMSUB231PS(ymm_two_w5_imag, ymm_two_h3_real, ymm_cos_1pi_over_8) ymm_cos_3pi_over_8 = YMMRegister() VMOVAPS(ymm_cos_3pi_over_8, Constant.float32x8(cos_npi_over_8[3])) VFMADD231PS(ymm_two_w1_real, ymm_two_h1_imag, ymm_cos_3pi_over_8) VFNMADD231PS(ymm_two_w7_real, ymm_two_h1_imag, ymm_cos_3pi_over_8) VFNMADD231PS(ymm_two_w1_imag, ymm_two_h1_real, ymm_cos_3pi_over_8) VFNMADD231PS(ymm_two_w7_imag, ymm_two_h1_real, ymm_cos_3pi_over_8) VFMADD231PS(ymm_two_w3_real, ymm_two_h3_real, ymm_cos_3pi_over_8) VFNMADD231PS(ymm_two_w5_real, ymm_two_h3_real, ymm_cos_3pi_over_8) VFMADD231PS(ymm_two_w3_imag, ymm_two_h3_imag, ymm_cos_3pi_over_8) VFMADD231PS(ymm_two_w5_imag, ymm_two_h3_imag, ymm_cos_3pi_over_8) ymm_half = YMMRegister() VMOVAPS(ymm_half, Constant.float32x8(0.5)) VMULPS(ymm_two_w1_real, ymm_two_w1_real, ymm_half) store_ymm_result(out_real[1], ymm_two_w1_real) VMULPS(ymm_two_w7_real, ymm_two_w7_real, ymm_half) store_ymm_result(out_real[7], ymm_two_w7_real) VMULPS(ymm_two_w1_imag, ymm_two_w1_imag, ymm_half) store_ymm_result(out_imag[1], ymm_two_w1_imag) VMULPS(ymm_two_w7_imag, ymm_two_w7_imag, ymm_half) store_ymm_result(out_imag[7], ymm_two_w7_imag) VMULPS(ymm_two_w3_real, ymm_two_w3_real, ymm_half) store_ymm_result(out_real[3], ymm_two_w3_real) VMULPS(ymm_two_w5_real, ymm_two_w5_real, ymm_half) store_ymm_result(out_real[5], ymm_two_w5_real) VMULPS(ymm_two_w3_imag, ymm_two_w3_imag, ymm_half) store_ymm_result(out_imag[3], ymm_two_w3_imag) VMULPS(ymm_two_w5_imag, ymm_two_w5_imag, ymm_half) store_ymm_result(out_imag[5], ymm_two_w5_imag)
def ifft16_within_rows(ymm_real_rows, ymm_imag_rows, bit_reversal=True): if isinstance(ymm_real_rows, tuple) and isinstance(ymm_imag_rows, tuple): return ifft16_within_rows([ymm_real_rows], [ymm_imag_rows]) assert isinstance(ymm_real_rows, list) and all(isinstance(ymm_real, tuple) and all(isinstance(ymm, YMMRegister) for ymm in ymm_real) for ymm_real in ymm_real_rows) assert isinstance(ymm_imag_rows, list) and all(isinstance(ymm_imag, tuple) and all(isinstance(ymm, YMMRegister) for ymm in ymm_imag) for ymm_imag in ymm_imag_rows) if bit_reversal: # Bit reversal # w[0] = x0 x8 x4 x12 x2 x10 x6 x14 # w[1] = x1 x9 x5 x13 x3 x11 x7 x15 ymm_bit_reversal_mask = YMMRegister() VMOVDQA(ymm_bit_reversal_mask, Constant.uint32x8(0, 2, 4, 6, 1, 3, 5, 7)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): for i in range(2): VPERMPS(ymm_real[i], ymm_bit_reversal_mask, ymm_real[i]) VPERMPS(ymm_imag[i], ymm_bit_reversal_mask, ymm_imag[i]) # 8x FFT2: Butterfly # w[0] = x0 x4 x2 x6 x8 x12 x10 x14 # w[1] = x1 x5 x3 x7 x9 x13 x11 x15 for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1]) butterfly(ymm_imag[0], ymm_imag[1]) ymm_new_real = YMMRegister(), YMMRegister() VUNPCKLPS(ymm_new_real[0], ymm_real[0], ymm_real[1]) VUNPCKHPS(ymm_new_real[1], ymm_real[0], ymm_imag[1]) ymm_new_imag = YMMRegister(), YMMRegister() VUNPCKLPS(ymm_new_imag[0], ymm_imag[0], ymm_imag[1]) VUNPCKHPS(ymm_new_imag[1], ymm_imag[0], ymm_real[1]) SWAP.REGISTERS(ymm_imag[0], ymm_new_imag[0]) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag[1]) SWAP.REGISTERS(ymm_real[0], ymm_new_real[0]) SWAP.REGISTERS(ymm_real[1], ymm_new_real[1]) # w[0] = x0 x1 x4 x5 x8 x9 x12 x13 # w[1] = x2 x3 x6 x7 x10 x11 x14 x15 # 4x FFT4: Butterfly and multiplication by twiddle factors ymm_fft4_twiddle_factor = YMMRegister() VMOVAPS(ymm_fft4_twiddle_factor, Constant.float32x8(+1.0, -1.0, +1.0, -1.0, +1.0, -1.0, +1.0, -1.0)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1], scale_b=ymm_fft4_twiddle_factor) butterfly(ymm_imag[0], ymm_imag[1]) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): transpose2x2x2x64(ymm_real[0], ymm_real[1]) transpose2x2x2x64(ymm_imag[0], ymm_imag[1]) # w[0] = x0 x1 x2 x3 x8 x9 x10 x11 # w[1] = x4 x5 x6 x7 x12 x13 x14 x15 # 2x FFT8: Multiplication by twiddle factors ymm_fft8_cos_twiddle_factor, ymm_fft8_sin_twiddle_factor = YMMRegister(), YMMRegister() VMOVAPS(ymm_fft8_cos_twiddle_factor, Constant.float32x8(*(cos_npi_over_4 * 2))) VMOVAPS(ymm_fft8_sin_twiddle_factor, Constant.float32x8(*(sin_npi_over_4 * 2))) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): ymm_new_real1, ymm_new_imag1 = YMMRegister(), YMMRegister() VMULPS(ymm_new_real1, ymm_real[1], ymm_fft8_cos_twiddle_factor) VMULPS(ymm_new_imag1, ymm_imag[1], ymm_fft8_cos_twiddle_factor) VFNMADD231PS(ymm_new_real1, ymm_imag[1], ymm_fft8_sin_twiddle_factor) VFMADD231PS(ymm_new_imag1, ymm_real[1], ymm_fft8_sin_twiddle_factor) SWAP.REGISTERS(ymm_real[1], ymm_new_real1) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag1) # 2x FFT8: Butterfly for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1]) butterfly(ymm_imag[0], ymm_imag[1]) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): transpose2x2x128(ymm_real[0], ymm_real[1]) transpose2x2x128(ymm_imag[0], ymm_imag[1]) # w[0] = x0 x1 x2 x3 x4 x5 x6 x7 # w[1] = x8 x9 x10 x11 x12 x13 x14 x15 # FFT16: Multiplication by twiddle factors and scale scale_factor = 0.0625 ymm_fft16_cos_scale_twiddle_factor, ymm_fft16_sin_scale_twiddle_factor = YMMRegister(), YMMRegister() VMOVAPS(ymm_fft16_cos_scale_twiddle_factor, Constant.float32x8(*[cos * scale_factor for cos in cos_npi_over_8])) VMOVAPS(ymm_fft16_sin_scale_twiddle_factor, Constant.float32x8(*[sin * scale_factor for sin in sin_npi_over_8])) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): ymm_new_real1, ymm_new_imag1 = YMMRegister(), YMMRegister() VMULPS(ymm_new_real1, ymm_real[1], ymm_fft16_cos_scale_twiddle_factor) VMULPS(ymm_new_imag1, ymm_imag[1], ymm_fft16_cos_scale_twiddle_factor) VFNMADD231PS(ymm_new_real1, ymm_imag[1], ymm_fft16_sin_scale_twiddle_factor) VFMADD231PS(ymm_new_imag1, ymm_real[1], ymm_fft16_sin_scale_twiddle_factor) SWAP.REGISTERS(ymm_real[1], ymm_new_real1) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag1) # FFT16: Butterfly and scale ymm_scale_factor = YMMRegister() VMOVAPS(ymm_scale_factor, Constant.float32x8(scale_factor)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1], scale_a=ymm_scale_factor) butterfly(ymm_imag[0], ymm_imag[1], scale_a=ymm_scale_factor)
def fft4_across_rows(ymm_real, ymm_imag, transformation="forward"): assert isinstance(ymm_real, list) and len(ymm_real) == 4 assert isinstance(ymm_imag, list) and len(ymm_imag) == 4 assert transformation in {"forward", "inverse"} ymm_data = sum(zip(ymm_real, ymm_imag), ()) # FFT-4 Butterfly for i in range(4): butterfly(ymm_data[i], ymm_data[i + 4]) # Multiply by FFT-4 twiddle factors SWAP.REGISTERS(ymm_real[3], ymm_imag[3]) # 2x FFT-2 Butterfly butterfly(ymm_data[0], ymm_data[2]) butterfly(ymm_data[1], ymm_data[3]) if transformation == "forward": butterfly(ymm_data[4], ymm_data[6]) butterfly(ymm_data[5], ymm_data[7], negate_b=True) else: butterfly(ymm_data[4], ymm_data[6], negate_b=True) butterfly(ymm_data[5], ymm_data[7]) # Bit reversal: not needed SWAP.REGISTERS(ymm_real[1], ymm_real[2]) SWAP.REGISTERS(ymm_imag[1], ymm_imag[2])
def fft16_within_rows(ymm_real_rows, ymm_imag_rows, bit_reversal=True): if isinstance(ymm_real_rows, tuple) and isinstance(ymm_imag_rows, tuple): return fft16_within_rows([ymm_real_rows], [ymm_imag_rows]) assert isinstance(ymm_real_rows, list) and all(isinstance(ymm_real, tuple) and all(isinstance(ymm, YMMRegister) for ymm in ymm_real) for ymm_real in ymm_real_rows) assert isinstance(ymm_imag_rows, list) and all(isinstance(ymm_imag, tuple) and all(isinstance(ymm, YMMRegister) for ymm in ymm_imag) for ymm_imag in ymm_imag_rows) # FFT16: Butterfly for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1]) butterfly(ymm_imag[0], ymm_imag[1]) # FFT16: Multiplication by twiddle factors ymm_fft16_cos_twiddle_factor, ymm_fft16_sin_twiddle_factor = YMMRegister(), YMMRegister() VMOVAPS(ymm_fft16_cos_twiddle_factor, Constant.float32x8(*cos_npi_over_8)) VMOVAPS(ymm_fft16_sin_twiddle_factor, Constant.float32x8(*sin_npi_over_8)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): ymm_new_real1, ymm_new_imag1 = YMMRegister(), YMMRegister() VMULPS(ymm_new_real1, ymm_real[1], ymm_fft16_cos_twiddle_factor) VMULPS(ymm_new_imag1, ymm_imag[1], ymm_fft16_cos_twiddle_factor) VFMADD231PS(ymm_new_real1, ymm_imag[1], ymm_fft16_sin_twiddle_factor) VFNMADD231PS(ymm_new_imag1, ymm_real[1], ymm_fft16_sin_twiddle_factor) SWAP.REGISTERS(ymm_real[1], ymm_new_real1) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag1) # 2x FFT8: Butterfly for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): transpose2x2x128(ymm_real[0], ymm_real[1]) transpose2x2x128(ymm_imag[0], ymm_imag[1]) # w[0] = x0 x1 x2 x3 x8 x9 x10 x11 # w[1] = x4 x5 x6 x7 x12 x13 x14 x15 for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1]) butterfly(ymm_imag[0], ymm_imag[1]) # 2x FFT8: Multiplication by twiddle factors ymm_fft8_cos_twiddle_factor, ymm_fft8_sin_twiddle_factor = YMMRegister(), YMMRegister() VMOVAPS(ymm_fft8_cos_twiddle_factor, Constant.float32x8(*(cos_npi_over_4 * 2))) VMOVAPS(ymm_fft8_sin_twiddle_factor, Constant.float32x8(*(sin_npi_over_4 * 2))) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): ymm_new_real1, ymm_new_imag1 = YMMRegister(), YMMRegister() VMULPS(ymm_new_real1, ymm_real[1], ymm_fft8_cos_twiddle_factor) VMULPS(ymm_new_imag1, ymm_imag[1], ymm_fft8_cos_twiddle_factor) VFMADD231PS(ymm_new_real1, ymm_imag[1], ymm_fft8_sin_twiddle_factor) VFNMADD231PS(ymm_new_imag1, ymm_real[1], ymm_fft8_sin_twiddle_factor) SWAP.REGISTERS(ymm_real[1], ymm_new_real1) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag1) # 4x FFT4: Butterfly for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): transpose2x2x2x64(ymm_real[0], ymm_real[1]) transpose2x2x2x64(ymm_imag[0], ymm_imag[1]) # w[0] = x0 x1 x4 x5 x8 x9 x12 x13 # w[1] = x2 x3 x6 x7 x10 x11 x14 x15 for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): butterfly(ymm_real[0], ymm_real[1]) butterfly(ymm_imag[0], ymm_imag[1]) # 4x FFT4: Multiplication by twiddle factors and 8x FFT2: Butterfly ymm_fft4_twiddle_factor = YMMRegister() VMOVAPS(ymm_fft4_twiddle_factor, Constant.float32x8(+1.0, +1.0, -1.0, -1.0, +1.0, +1.0, -1.0, -1.0)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): ymm_new_real = YMMRegister(), YMMRegister() VSHUFPS(ymm_new_real[0], ymm_real[0], ymm_real[1], _MM_SHUFFLE(2, 0, 2, 0)) VSHUFPS(ymm_new_real[1], ymm_real[0], ymm_imag[1], _MM_SHUFFLE(3, 1, 3, 1)) butterfly(ymm_new_real[0], ymm_new_real[1]) ymm_new_imag = YMMRegister(), YMMRegister() VSHUFPS(ymm_new_imag[0], ymm_imag[0], ymm_imag[1], _MM_SHUFFLE(2, 0, 2, 0)) VSHUFPS(ymm_new_imag[1], ymm_imag[0], ymm_real[1], _MM_SHUFFLE(3, 1, 3, 1)) butterfly(ymm_new_imag[0], ymm_new_imag[1], scale_b=ymm_fft4_twiddle_factor) SWAP.REGISTERS(ymm_real[0], ymm_new_real[0]) SWAP.REGISTERS(ymm_real[1], ymm_new_real[1]) SWAP.REGISTERS(ymm_imag[0], ymm_new_imag[0]) SWAP.REGISTERS(ymm_imag[1], ymm_new_imag[1]) # w[0] = x0 x4 x2 x6 x8 x12 x10 x14 # w[1] = x1 x5 x3 x7 x9 x11 x13 x15 if bit_reversal: # Bit reversal ymm_bit_reversal_mask = YMMRegister() VMOVDQA(ymm_bit_reversal_mask, Constant.uint32x8(0, 4, 1, 5, 2, 6, 3, 7)) for ymm_real, ymm_imag in zip(ymm_real_rows, ymm_imag_rows): for i in range(2): VPERMPS(ymm_real[i], ymm_bit_reversal_mask, ymm_real[i]) VPERMPS(ymm_imag[i], ymm_bit_reversal_mask, ymm_imag[i])
data = [YMMRegister() for _ in range(16)] data[0] = LocalVariable(data[0]) data[8] = LocalVariable(data[8]) real, imag = data[0::2], data[1::2] for i, (data_lo, data_hi) in enumerate(zip(data[0:8], data[8:16])): ymm_data_lo, ymm_data_hi = data_lo, data_hi if isinstance(data_lo, LocalVariable): ymm_data_lo = YMMRegister() if isinstance(data_hi, LocalVariable): ymm_data_hi = YMMRegister() VMOVUPS(ymm_data_lo, [reg_t + i * YMMRegister.size]) VMOVUPS(ymm_data_hi, [reg_t + (i + 8) * YMMRegister.size]) butterfly(ymm_data_lo, ymm_data_hi) if isinstance(data_lo, LocalVariable): VMOVAPS(data_lo, ymm_data_lo) if isinstance(data_hi, LocalVariable): VMOVAPS(data_hi, ymm_data_hi) # FFT8: multiplication by twiddle factors fft4_scale_b, fft4_negate_b = {}, {} fft2_scale_b, fft2_negate_b = {}, {} # w6.re, w6.im = w6.im, -w6.re SWAP.REGISTERS(real[6], imag[6]) fft4_negate_b[id(imag[6])] = True # w5.re, w5.im = SQRT2_OVER_2 * (w5.re + w5.im), SQRT2_OVER_2 * (w5.im - w5.re)
def fft8_across_rows(ymm_data): assert isinstance(ymm_data, list) and len(ymm_data) == 8 ymm_real = ymm_data[0::2] ymm_imag = ymm_data[1::2] fft.complex_soa.fft4_across_rows(ymm_real, ymm_imag) butterfly(ymm_real[0], ymm_imag[0]) # const float two_gdata1_real = crealf(data1) + crealf(data3); # const float two_gdata1_imag = cimagf(data1) - cimagf(data3); ymm_two_gdata1_real, ymm_two_gdata1_imag = YMMRegister(), YMMRegister() VADDPS(ymm_two_gdata1_real, ymm_real[1], ymm_real[3]) VSUBPS(ymm_two_gdata1_imag, ymm_imag[1], ymm_imag[3]) # const float two_hdata1_real = cimagf(data1) + cimagf(data3); # const float two_hdata1_imag = crealf(data3) - crealf(data1); ymm_two_hdata1_real, ymm_two_hdata1_imag = YMMRegister(), YMMRegister() VADDPS(ymm_two_hdata1_real, ymm_imag[1], ymm_imag[3]) VSUBPS(ymm_two_hdata1_imag, ymm_real[3], ymm_real[1]) # const float two_hdata1_real_plus_imag = two_hdata1_real + two_hdata1_imag; # const float two_hdata1_real_minus_imag = two_hdata1_real - two_hdata1_imag; ymm_two_hdata1_plus, ymm_two_hdata1_minus = YMMRegister(), YMMRegister() VADDPS(ymm_two_hdata1_plus, ymm_two_hdata1_real, ymm_two_hdata1_imag) VSUBPS(ymm_two_hdata1_minus, ymm_two_hdata1_real, ymm_two_hdata1_imag) ymm_sqrt2_over_2 = YMMRegister() VMOVAPS(ymm_sqrt2_over_2, Constant.float32x8(sqrt2_over_2)) # const float two_data1_real = two_gdata1_real + SQRT2_OVER_2 * two_hdata1_real_plus_imag; # const float two_data1_imag = two_gdata1_imag - SQRT2_OVER_2 * two_hdata1_real_minus_imag; # const float two_data3_real = two_gdata1_real - SQRT2_OVER_2 * two_hdata1_real_plus_imag; # const float two_data3_imag = -two_gdata1_imag - SQRT2_OVER_2 * two_hdata1_real_minus_imag; ymm_two_data1_real, ymm_two_data1_imag = YMMRegister(), YMMRegister() ymm_two_data3_real, ymm_two_data3_imag = YMMRegister(), YMMRegister() VMOVAPS(ymm_two_data3_real, ymm_two_gdata1_real) VMOVAPS(ymm_two_data3_imag, ymm_two_gdata1_imag) VFMADD231PS(ymm_two_gdata1_real, ymm_two_hdata1_plus, ymm_sqrt2_over_2) VFNMADD231PS(ymm_two_gdata1_imag, ymm_two_hdata1_minus, ymm_sqrt2_over_2) SWAP.REGISTERS(ymm_two_data1_real, ymm_two_gdata1_real) SWAP.REGISTERS(ymm_two_data1_imag, ymm_two_gdata1_imag) VFNMADD231PS(ymm_two_data3_real, ymm_two_hdata1_plus, ymm_sqrt2_over_2) VFNMSUB231PS(ymm_two_data3_imag, ymm_two_hdata1_minus, ymm_sqrt2_over_2) # /* Store outputs */ # fdata[0] = crealf(data0) + cimagf(data0); # fdata[1] = crealf(data0) - cimagf(data0); # fdata[2] = 0.5f * two_data1_real; # fdata[3] = 0.5f * two_data1_imag; # fdata[4] = crealf(data2); # fdata[5] = -cimagf(data2); # fdata[6] = 0.5f * two_data3_real; # fdata[7] = 0.5f * two_data3_imag; ymm_half = YMMRegister() VMOVAPS(ymm_half, Constant.float32x8(0.5)) VMULPS(ymm_real[1], ymm_two_data1_real, ymm_half) VMULPS(ymm_imag[1], ymm_two_data1_imag, ymm_half) VXORPS(ymm_imag[2], ymm_imag[2], Constant.float32x8(-0.0)) VMULPS(ymm_real[3], ymm_two_data3_real, ymm_half) VMULPS(ymm_imag[3], ymm_two_data3_imag, ymm_half)
data = [YMMRegister() for _ in range(16)] data[0] = LocalVariable(data[0]) data[8] = LocalVariable(data[8]) real, imag = data[0::2], data[1::2] for i, (data_lo, data_hi) in enumerate(zip(data[0:8], data[8:16])): ymm_data_lo, ymm_data_hi = data_lo, data_hi if isinstance(data_lo, LocalVariable): ymm_data_lo = YMMRegister() if isinstance(data_hi, LocalVariable): ymm_data_hi = YMMRegister() VMOVUPS(ymm_data_lo, [reg_t + i * YMMRegister.size]) VMOVUPS(ymm_data_hi, [reg_t + (i + 8) * YMMRegister.size]) butterfly(ymm_data_lo, ymm_data_hi) if isinstance(data_lo, LocalVariable): VMOVAPS(data_lo, ymm_data_lo) if isinstance(data_hi, LocalVariable): VMOVAPS(data_hi, ymm_data_hi) # FFT8: multiplication by twiddle factors fft4_scale_b, fft4_negate_b = {}, {} fft2_scale_b, fft2_negate_b = {}, {} # w6.re, w6.im = w6.im, -w6.re SWAP.REGISTERS(real[6], imag[6]) fft4_negate_b[id(imag[6])] = True # w5.re, w5.im = SQRT2_OVER_2 * (w5.re + w5.im), SQRT2_OVER_2 * (w5.im - w5.re)
def fft8_across_rows(ymm_data): assert isinstance(ymm_data, list) and len(ymm_data) == 8 ymm_real = ymm_data[0::2] ymm_imag = ymm_data[1::2] complex_soa.fft4_across_rows(ymm_real, ymm_imag) butterfly(ymm_real[0], ymm_imag[0]) # const float two_gdata1_real = crealf(data1) + crealf(data3); # const float two_gdata1_imag = cimagf(data1) - cimagf(data3); ymm_two_gdata1_real, ymm_two_gdata1_imag = YMMRegister(), YMMRegister() VADDPS(ymm_two_gdata1_real, ymm_real[1], ymm_real[3]) VSUBPS(ymm_two_gdata1_imag, ymm_imag[1], ymm_imag[3]) # const float two_hdata1_real = cimagf(data1) + cimagf(data3); # const float two_hdata1_imag = crealf(data3) - crealf(data1); ymm_two_hdata1_real, ymm_two_hdata1_imag = YMMRegister(), YMMRegister() VADDPS(ymm_two_hdata1_real, ymm_imag[1], ymm_imag[3]) VSUBPS(ymm_two_hdata1_imag, ymm_real[3], ymm_real[1]) # const float two_hdata1_real_plus_imag = two_hdata1_real + two_hdata1_imag; # const float two_hdata1_real_minus_imag = two_hdata1_real - two_hdata1_imag; ymm_two_hdata1_plus, ymm_two_hdata1_minus = YMMRegister(), YMMRegister() VADDPS(ymm_two_hdata1_plus, ymm_two_hdata1_real, ymm_two_hdata1_imag) VSUBPS(ymm_two_hdata1_minus, ymm_two_hdata1_real, ymm_two_hdata1_imag) ymm_sqrt2_over_2 = YMMRegister() VMOVAPS(ymm_sqrt2_over_2, Constant.float32x8(sqrt2_over_2)) # const float two_data1_real = two_gdata1_real + SQRT2_OVER_2 * two_hdata1_real_plus_imag; # const float two_data1_imag = two_gdata1_imag - SQRT2_OVER_2 * two_hdata1_real_minus_imag; # const float two_data3_real = two_gdata1_real - SQRT2_OVER_2 * two_hdata1_real_plus_imag; # const float two_data3_imag = -two_gdata1_imag - SQRT2_OVER_2 * two_hdata1_real_minus_imag; ymm_two_data1_real, ymm_two_data1_imag = YMMRegister(), YMMRegister() ymm_two_data3_real, ymm_two_data3_imag = YMMRegister(), YMMRegister() VMOVAPS(ymm_two_data3_real, ymm_two_gdata1_real) VMOVAPS(ymm_two_data3_imag, ymm_two_gdata1_imag) VFMADD231PS(ymm_two_gdata1_real, ymm_two_hdata1_plus, ymm_sqrt2_over_2) VFNMADD231PS(ymm_two_gdata1_imag, ymm_two_hdata1_minus, ymm_sqrt2_over_2) SWAP.REGISTERS(ymm_two_data1_real, ymm_two_gdata1_real) SWAP.REGISTERS(ymm_two_data1_imag, ymm_two_gdata1_imag) VFNMADD231PS(ymm_two_data3_real, ymm_two_hdata1_plus, ymm_sqrt2_over_2) VFNMSUB231PS(ymm_two_data3_imag, ymm_two_hdata1_minus, ymm_sqrt2_over_2) # /* Store outputs */ # fdata[0] = crealf(data0) + cimagf(data0); # fdata[1] = crealf(data0) - cimagf(data0); # fdata[2] = 0.5f * two_data1_real; # fdata[3] = 0.5f * two_data1_imag; # fdata[4] = crealf(data2); # fdata[5] = -cimagf(data2); # fdata[6] = 0.5f * two_data3_real; # fdata[7] = 0.5f * two_data3_imag; ymm_half = YMMRegister() VMOVAPS(ymm_half, Constant.float32x8(0.5)) VMULPS(ymm_real[1], ymm_two_data1_real, ymm_half) VMULPS(ymm_imag[1], ymm_two_data1_imag, ymm_half) VXORPS(ymm_imag[2], ymm_imag[2], Constant.float32x8(-0.0)) VMULPS(ymm_real[3], ymm_two_data3_real, ymm_half) VMULPS(ymm_imag[3], ymm_two_data3_imag, ymm_half)