generate_twiddles.py

"""************************************************************************
 *
 *  Copyright (C) Codeplay Software Ltd.
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *  Codeplay's portFFT
 *
 ************************************************************************"""
import math
MAX_SIZE = 64
DST = "./src/common/twiddle.hpp"

template = """
/***************************************************************************
 *
 *  Generated by scripts/generate_twiddles.py. Do not edit!
 *
 **************************************************************************/

#ifndef PORTFFT_COMMON_TWIDDLE_HPP
#define PORTFFT_COMMON_TWIDDLE_HPP

#pragma clang diagnostic push
// The twiddle precision can be lower than the constants used here when not using double precision.
#pragma clang diagnostic ignored "-Wimplicit-float-conversion"

namespace portfft::detail {{

template <typename T>
struct twiddle {{
  // twiddle_re[N][K] and twiddle_im[N][K] contain real and imaginary components of a twiddle factor K out of N.

  // We only have twiddles up to size 64 here. 64 is likely the largest size we will be able to handle within one
  // workitem on current GPUs

  // clang-format off
static constexpr T Re[{size}][{size}] = {{ {real_forward} }};

static constexpr T Im[{size}][{size}] = {{ {imag_forward} }};
  // clang-format on
}};
}}  // namespace portfft::detail

#pragma clang diagnostic pop

#endif
"""

def generate(max_size):
    res_real = []
    res_imag = []
    for size in range(max_size+1):
        tmp_real = []
        tmp_imag = []
        for i in range(size):
            # make sure zeros are exact to let compiler do further optimizations
            if i == 0:
                tmp_real.append(1.0)
                tmp_imag.append(0.0)
            elif 2*i == size:
                tmp_real.append(-1.0)
                tmp_imag.append(0.0)
            elif 4*i == size:
                tmp_real.append(0.0)
                tmp_imag.append(-1.0)
            elif 4*i == size*3:
                tmp_real.append(0.0)
                tmp_imag.append(1.0)
            else:
                theta = -2. * math.pi * i / size
                tmp_real.append(math.cos(theta))
                tmp_imag.append(math.sin(theta))

        # pad with zeros to max_size
        tmp_real += [0] * (max_size - size + 1)
        tmp_imag += [0] * (max_size - size + 1)

        res_real.append(tmp_real)
        res_imag.append(tmp_imag)
    
    return [res_real, res_imag]


def write(path, size):
    real, imag  = generate(size)
    with open(path, "w") as fil:
        realstr = ",\n".join("{" + ", ".join(str(j)
                                                     for j in i) + "}" for i in real)
        imagstr = ",\n".join("{" + ", ".join(str(j)
                                                     for j in i) + "}" for i in imag)
        
        content = template.format(size=size+1, real_forward=realstr, imag_forward=imagstr)
        fil.write(content)


if __name__ == "__main__":
    write(DST, MAX_SIZE)