opencl/fourier_transform.cc

#include <atomic>
#include <cassert>
#include <chrono>
#include <cstdio>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <sstream>

#include <openclx/compiler>

#include <vtestbed/config/openmp.hh>
#include <vtestbed/opencl/fourier_transform.hh>
#include <vtestbed/opencl/pipeline.hh>

namespace {

    #if defined(VTB_DEBUG_CHIRP_Z)
    template <class T, int N>
    void
    dump(vtb::opencl::Buffer<T> d_x, const blitz::TinyVector<int,N>& shape,
         const char* name) {
        blitz::Array<T,N> x(shape);
        ppl.copy(d_x, x);
        ppl.wait();
        std::clog << name << '=' << x << std::endl;
    }
    #define VTB_DUMP(x, shape, name) ::dump(x,shape,name)
    #else
    #define VTB_DUMP(x, shape, name)
    #endif


    const char* base_kernels =
        "#ifndef M_PI\n"
        "#define M_PI 0x1.921fb54442d18p+1\n"
        "#endif\n"
        "#define complexMul(a,b) ((float2)(mad(-(a).y, (b).y, (a).x * (b).x), mad((a).y, (b).x, (a).x * (b).y)))\n"
        "#define conj(a) ((float2)((a).x, -(a).y))\n"
        "#define conjTransp(a) ((float2)(-(a).y, (a).x))\n"
        "\n"
        "#define fftKernel2(a,dir) \\\n"
        "{ \\\n"
        "    float2 c = (a)[0];    \\\n"
        "    (a)[0] = c + (a)[1];  \\\n"
        "    (a)[1] = c - (a)[1];  \\\n"
        "}\n"
        "\n"
        "#define fftKernel2S(d1,d2,dir) \\\n"
        "{ \\\n"
        "    float2 c = (d1);   \\\n"
        "    (d1) = c + (d2);   \\\n"
        "    (d2) = c - (d2);   \\\n"
        "}\n"
        "\n"
        "#define fftKernel4(a,dir) \\\n"
        "{ \\\n"
        "    fftKernel2S((a)[0], (a)[2], dir); \\\n"
        "    fftKernel2S((a)[1], (a)[3], dir); \\\n"
        "    fftKernel2S((a)[0], (a)[1], dir); \\\n"
        "    (a)[3] = (float2)(dir)*(conjTransp((a)[3])); \\\n"
        "    fftKernel2S((a)[2], (a)[3], dir); \\\n"
        "    float2 c = (a)[1]; \\\n"
        "    (a)[1] = (a)[2]; \\\n"
        "    (a)[2] = c; \\\n"
        "}\n"
        "\n"
        "#define fftKernel4s(a0,a1,a2,a3,dir) \\\n"
        "{ \\\n"
        "    fftKernel2S((a0), (a2), dir); \\\n"
        "    fftKernel2S((a1), (a3), dir); \\\n"
        "    fftKernel2S((a0), (a1), dir); \\\n"
        "    (a3) = (float2)(dir)*(conjTransp((a3))); \\\n"
        "    fftKernel2S((a2), (a3), dir); \\\n"
        "    float2 c = (a1); \\\n"
        "    (a1) = (a2); \\\n"
        "    (a2) = c; \\\n"
        "}\n"
        "\n"
        "#define bitreverse8(a) \\\n"
        "{ \\\n"
        "    float2 c; \\\n"
        "    c = (a)[1]; \\\n"
        "    (a)[1] = (a)[4]; \\\n"
        "    (a)[4] = c; \\\n"
        "    c = (a)[3]; \\\n"
        "    (a)[3] = (a)[6]; \\\n"
        "    (a)[6] = c; \\\n"
        "}\n"
        "\n"
        "#define fftKernel8(a,dir) \\\n"
        "{ \\\n"
        "   const float2 w1  = (float2)(0x1.6a09e6p-1f,  dir*0x1.6a09e6p-1f);  \\\n"
        "   const float2 w3  = (float2)(-0x1.6a09e6p-1f, dir*0x1.6a09e6p-1f);  \\\n"
        "   float2 c; \\\n"
        "   fftKernel2S((a)[0], (a)[4], dir); \\\n"
        "   fftKernel2S((a)[1], (a)[5], dir); \\\n"
        "   fftKernel2S((a)[2], (a)[6], dir); \\\n"
        "   fftKernel2S((a)[3], (a)[7], dir); \\\n"
        "   (a)[5] = complexMul(w1, (a)[5]); \\\n"
        "   (a)[6] = (float2)(dir)*(conjTransp((a)[6])); \\\n"
        "   (a)[7] = complexMul(w3, (a)[7]); \\\n"
        "   fftKernel2S((a)[0], (a)[2], dir); \\\n"
        "   fftKernel2S((a)[1], (a)[3], dir); \\\n"
        "   fftKernel2S((a)[4], (a)[6], dir); \\\n"
        "   fftKernel2S((a)[5], (a)[7], dir); \\\n"
        "   (a)[3] = (float2)(dir)*(conjTransp((a)[3])); \\\n"
        "   (a)[7] = (float2)(dir)*(conjTransp((a)[7])); \\\n"
        "   fftKernel2S((a)[0], (a)[1], dir); \\\n"
        "   fftKernel2S((a)[2], (a)[3], dir); \\\n"
        "   fftKernel2S((a)[4], (a)[5], dir); \\\n"
        "   fftKernel2S((a)[6], (a)[7], dir); \\\n"
        "   bitreverse8((a)); \\\n"
        "}\n"
        "\n"
        "#define bitreverse4x4(a) \\\n"
        "{ \\\n"
        "   float2 c; \\\n"
        "   c = (a)[1];  (a)[1]  = (a)[4];  (a)[4]  = c; \\\n"
        "   c = (a)[2];  (a)[2]  = (a)[8];  (a)[8]  = c; \\\n"
        "   c = (a)[3];  (a)[3]  = (a)[12]; (a)[12] = c; \\\n"
        "   c = (a)[6];  (a)[6]  = (a)[9];  (a)[9]  = c; \\\n"
        "   c = (a)[7];  (a)[7]  = (a)[13]; (a)[13] = c; \\\n"
        "   c = (a)[11]; (a)[11] = (a)[14]; (a)[14] = c; \\\n"
        "}\n"
        "\n"
        "#define fftKernel16(a,dir) \\\n"
        "{ \\\n"
        "    const float w0 = 0x1.d906bcp-1f; \\\n"
        "    const float w1 = 0x1.87de2ap-2f; \\\n"
        "    const float w2 = 0x1.6a09e6p-1f; \\\n"
        "    fftKernel4s((a)[0], (a)[4], (a)[8],  (a)[12], dir); \\\n"
        "    fftKernel4s((a)[1], (a)[5], (a)[9],  (a)[13], dir); \\\n"
        "    fftKernel4s((a)[2], (a)[6], (a)[10], (a)[14], dir); \\\n"
        "    fftKernel4s((a)[3], (a)[7], (a)[11], (a)[15], dir); \\\n"
        "    (a)[5]  = complexMul((a)[5], (float2)(w0, dir*w1)); \\\n"
        "    (a)[6]  = complexMul((a)[6], (float2)(w2, dir*w2)); \\\n"
        "    (a)[7]  = complexMul((a)[7], (float2)(w1, dir*w0)); \\\n"
        "    (a)[9]  = complexMul((a)[9], (float2)(w2, dir*w2)); \\\n"
        "    (a)[10] = (float2)(dir)*(conjTransp((a)[10])); \\\n"
        "    (a)[11] = complexMul((a)[11], (float2)(-w2, dir*w2)); \\\n"
        "    (a)[13] = complexMul((a)[13], (float2)(w1, dir*w0)); \\\n"
        "    (a)[14] = complexMul((a)[14], (float2)(-w2, dir*w2)); \\\n"
        "    (a)[15] = complexMul((a)[15], (float2)(-w0, dir*-w1)); \\\n"
        "    fftKernel4((a), dir); \\\n"
        "    fftKernel4((a) + 4, dir); \\\n"
        "    fftKernel4((a) + 8, dir); \\\n"
        "    fftKernel4((a) + 12, dir); \\\n"
        "    bitreverse4x4((a)); \\\n"
        "}\n"
        "\n"
        "#define bitreverse32(a) \\\n"
        "{ \\\n"
        "    float2 c1, c2; \\\n"
        "    c1 = (a)[2];   (a)[2] = (a)[1];   c2 = (a)[4];   (a)[4] = c1;   c1 = (a)[8];   (a)[8] = c2;    c2 = (a)[16];  (a)[16] = c1;   (a)[1] = c2; \\\n"
        "    c1 = (a)[6];   (a)[6] = (a)[3];   c2 = (a)[12];  (a)[12] = c1;  c1 = (a)[24];  (a)[24] = c2;   c2 = (a)[17];  (a)[17] = c1;   (a)[3] = c2; \\\n"
        "    c1 = (a)[10];  (a)[10] = (a)[5];  c2 = (a)[20];  (a)[20] = c1;  c1 = (a)[9];   (a)[9] = c2;    c2 = (a)[18];  (a)[18] = c1;   (a)[5] = c2; \\\n"
        "    c1 = (a)[14];  (a)[14] = (a)[7];  c2 = (a)[28];  (a)[28] = c1;  c1 = (a)[25];  (a)[25] = c2;   c2 = (a)[19];  (a)[19] = c1;   (a)[7] = c2; \\\n"
        "    c1 = (a)[22];  (a)[22] = (a)[11]; c2 = (a)[13];  (a)[13] = c1;  c1 = (a)[26];  (a)[26] = c2;   c2 = (a)[21];  (a)[21] = c1;   (a)[11] = c2; \\\n"
        "    c1 = (a)[30];  (a)[30] = (a)[15]; c2 = (a)[29];  (a)[29] = c1;  c1 = (a)[27];  (a)[27] = c2;   c2 = (a)[23];  (a)[23] = c1;   (a)[15] = c2; \\\n"
        "}\n"
        "\n"
        "#define fftKernel32(a,dir) \\\n"
        "{ \\\n"
        "    fftKernel2S((a)[0],  (a)[16], dir); \\\n"
        "    fftKernel2S((a)[1],  (a)[17], dir); \\\n"
        "    fftKernel2S((a)[2],  (a)[18], dir); \\\n"
        "    fftKernel2S((a)[3],  (a)[19], dir); \\\n"
        "    fftKernel2S((a)[4],  (a)[20], dir); \\\n"
        "    fftKernel2S((a)[5],  (a)[21], dir); \\\n"
        "    fftKernel2S((a)[6],  (a)[22], dir); \\\n"
        "    fftKernel2S((a)[7],  (a)[23], dir); \\\n"
        "    fftKernel2S((a)[8],  (a)[24], dir); \\\n"
        "    fftKernel2S((a)[9],  (a)[25], dir); \\\n"
        "    fftKernel2S((a)[10], (a)[26], dir); \\\n"
        "    fftKernel2S((a)[11], (a)[27], dir); \\\n"
        "    fftKernel2S((a)[12], (a)[28], dir); \\\n"
        "    fftKernel2S((a)[13], (a)[29], dir); \\\n"
        "    fftKernel2S((a)[14], (a)[30], dir); \\\n"
        "    fftKernel2S((a)[15], (a)[31], dir); \\\n"
        "    (a)[17] = complexMul((a)[17], (float2)(0x1.f6297cp-1f, dir*0x1.8f8b84p-3f)); \\\n"
        "    (a)[18] = complexMul((a)[18], (float2)(0x1.d906bcp-1f, dir*0x1.87de2ap-2f)); \\\n"
        "    (a)[19] = complexMul((a)[19], (float2)(0x1.a9b662p-1f, dir*0x1.1c73b4p-1f)); \\\n"
        "    (a)[20] = complexMul((a)[20], (float2)(0x1.6a09e6p-1f, dir*0x1.6a09e6p-1f)); \\\n"
        "    (a)[21] = complexMul((a)[21], (float2)(0x1.1c73b4p-1f, dir*0x1.a9b662p-1f)); \\\n"
        "    (a)[22] = complexMul((a)[22], (float2)(0x1.87de2ap-2f, dir*0x1.d906bcp-1f)); \\\n"
        "    (a)[23] = complexMul((a)[23], (float2)(0x1.8f8b84p-3f, dir*0x1.f6297cp-1f)); \\\n"
        "    (a)[24] = complexMul((a)[24], (float2)(0x0p+0f, dir*0x1p+0f)); \\\n"
        "    (a)[25] = complexMul((a)[25], (float2)(-0x1.8f8b84p-3f, dir*0x1.f6297cp-1f)); \\\n"
        "    (a)[26] = complexMul((a)[26], (float2)(-0x1.87de2ap-2f, dir*0x1.d906bcp-1f)); \\\n"
        "    (a)[27] = complexMul((a)[27], (float2)(-0x1.1c73b4p-1f, dir*0x1.a9b662p-1f)); \\\n"
        "    (a)[28] = complexMul((a)[28], (float2)(-0x1.6a09e6p-1f, dir*0x1.6a09e6p-1f)); \\\n"
        "    (a)[29] = complexMul((a)[29], (float2)(-0x1.a9b662p-1f, dir*0x1.1c73b4p-1f)); \\\n"
        "    (a)[30] = complexMul((a)[30], (float2)(-0x1.d906bcp-1f, dir*0x1.87de2ap-2f)); \\\n"
        "    (a)[31] = complexMul((a)[31], (float2)(-0x1.f6297cp-1f, dir*0x1.8f8b84p-3f)); \\\n"
        "    fftKernel16((a), dir); \\\n"
        "    fftKernel16((a) + 16, dir); \\\n"
        "    bitreverse32((a)); \\\n"
        "}\n\n";

    std::vector<int>
    radix_array(int n, int max) {
        std::vector<int> result;
        max = std::min(n, max);
        while (n > max) {
            result.push_back(max);
            n /= max;
        }
        result.push_back(n);
        return result;
    }

    std::vector<int>
    radix_array(int n) {
        std::vector<int> result;
        switch (n) {
            case 2: result = {2}; break;
            case 4: result = {4}; break;
            case 8: result = {8}; break;
            case 16: result = {8,2}; break;
            case 32: result = {8,4}; break;
            case 64: result = {8,8}; break;
            case 128: result = {8,4,4}; break;
            case 256: result = {4,4,4,4}; break;
            case 512: result = {8,8,8}; break;
            case 1024: result = {16,16,4}; break;
            case 2048: result = {8,8,8,4}; break;
            default: throw std::runtime_error{"unable to generate radix array"};
        }
        return result;
    }

    void
    formattedLoad(
        std::ostream& out,
        int aIndex,
        int gIndex,
        vtb::opencl::Fourier_transform_format dataFormat
    ) {
        using vtb::opencl::Fourier_transform_format;
        if (dataFormat == Fourier_transform_format::Interleaved_complex) {
            out << "        a[" << (aIndex) << "] = in[" << (gIndex) << "];\n";
        } else {
            out << "        a[" << (aIndex) << "].x = in_real[" << (gIndex) << "];\n";
            out << "        a[" << (aIndex) << "].y = in_imag[" << (gIndex) << "];\n";
        }
    }

    void
    formattedStore(
        std::ostream& out,
        int aIndex,
        int gIndex,
        vtb::opencl::Fourier_transform_format dataFormat
    ) {
        using vtb::opencl::Fourier_transform_format;
        if (dataFormat == Fourier_transform_format::Interleaved_complex) {
            out << "        out[" << (gIndex) << "] = a[" << (aIndex) << "];\n";
        } else {
            out << "        out_real[" << (gIndex) << "] = a[" << (aIndex) << "].x;\n";
            out << "        out_imag[" << (gIndex) << "] = a[" << (aIndex) << "].y;\n";
        }
    }

    void
    insertHeader(
        std::ostream& out,
        std::string kernelName,
        vtb::opencl::Fourier_transform_format dataFormat
    ) {
        using vtb::opencl::Fourier_transform_format;
        if (dataFormat == Fourier_transform_format::Split_complex) {
            out << "__kernel void " + kernelName
                << "(__global float *in_real, __global float *in_imag, __global float *out_real, __global float *out_imag, int dir, int S)\n";
        } else {
            out << "__kernel void " + kernelName
                << "(__global float2 *in, __global float2 *out, int dir, int S)\n";
        }
    }

    void
    insertVariables(std::ostream& out, int maxRadix) {
        out << "    int i, j, r, indexIn, indexOut, index, tid, bNum, xNum, k, l;\n";
        out << "    int s, ii, jj, offset;\n";
        out << "    float2 w;\n";
        out << "    float ang, angf, ang1;\n";
        out << "    __local float *lMemStore, *lMemLoad;\n";
        out << "    float2 a[" << maxRadix << "];\n";
        out << "    int lId = get_local_id( 0 );\n";
        out << "    int groupId = get_group_id( 0 );\n";
    }

    void
    insertfftKernel(std::ostream& out, int Nr, int numIter) {
        for (int i=0; i<numIter; ++i) {
            out << "    fftKernel" << (Nr) << "(a+" << (i*Nr) << ", dir);\n";
        }
    }

    void
    insertTwiddleKernel(
        std::ostream& out,
        int Nr,
        int numIter,
        int Nprev,
        int len,
        int numWorkItemsPerXForm
    ) {
        int logNPrev = (int)std::log2(Nprev);
        for (int z=0; z<numIter; ++z) {
            if (z == 0) {
                if (Nprev > 1) {
                    out << "    angf = (float) (ii >> " << (logNPrev) << ");\n";
                } else {
                    out << "    angf = (float) ii;\n";
                }
            } else {
                if (Nprev > 1) {
                    out << "    angf = (float) ((" << (z*numWorkItemsPerXForm) << " + ii) >>" << (logNPrev) << ");\n";
                } else {
                    out << "    angf = (float) (" << (z*numWorkItemsPerXForm) << " + ii);\n";
                }
            }
            for (int k=1; k<Nr; ++k) {
                int ind = z*Nr + k;
                //float fac =  (float) (2.0 * M_PI * (double) k / (double) len);
                out << "    ang = dir * ( 2.0f * M_PI * " << (k) << ".0f / " << (len) << ".0f )" << " * angf;\n";
                out << "    w = (float2)(native_cos(ang), native_sin(ang));\n";
                out << "    a[" << (ind) << "] = complexMul(a[" << (ind) << "], w);\n";
            }
        }
    }

    int
    insertGlobalLoadsAndTranspose(
        std::ostream& out,
        int N,
        int numWorkItemsPerXForm,
        int numXFormsPerWG,
        int R0,
        int mem_coalesce_width,
        vtb::opencl::Fourier_transform_format dataFormat
    ) {
        using vtb::opencl::Fourier_transform_format;
        int log2NumWorkItemsPerXForm = (int) log2(numWorkItemsPerXForm);
        int groupSize = numWorkItemsPerXForm * numXFormsPerWG;
        int lMemSize = 0;
        if (numXFormsPerWG > 1) {
            out << "        s = S & " << (numXFormsPerWG-1) << ";\n";
        }
        if (numWorkItemsPerXForm >= mem_coalesce_width) {
            if (numXFormsPerWG > 1) {
                out << "    ii = lId & " << (numWorkItemsPerXForm-1) << ";\n";
                out << "    jj = lId >> " << log2NumWorkItemsPerXForm << ";\n";
                out << "    if( !s || (groupId < get_num_groups(0)-1) || (jj < s) ) {\n";
                out << "        offset = mad24( mad24(groupId, "
                    << numXFormsPerWG
                    << ", jj), " << N
                    << ", ii );\n";
                if (dataFormat == Fourier_transform_format::Interleaved_complex) {
                    out << "        in += offset;\n";
                    out << "        out += offset;\n";
                } else {
                    out << "        in_real += offset;\n";
                    out << "        in_imag += offset;\n";
                    out << "        out_real += offset;\n";
                    out << "        out_imag += offset;\n";
                }
                for (int i=0; i<R0; ++i) {
                    formattedLoad(out, i, i*numWorkItemsPerXForm, dataFormat);
                }
                out << "    }\n";
            } else {
                out << "    ii = lId;\n";
                out << "    jj = 0;\n";
                out << "    offset =  mad24(groupId, " << N << ", ii);\n";
                if (dataFormat == Fourier_transform_format::Interleaved_complex) {
                    out << "        in += offset;\n";
                    out << "        out += offset;\n";
                } else {
                    out << "        in_real += offset;\n";
                    out << "        in_imag += offset;\n";
                    out << "        out_real += offset;\n";
                    out << "        out_imag += offset;\n";
                }
                for (int i=0; i<R0; ++i) {
                    formattedLoad(out, i, i*numWorkItemsPerXForm, dataFormat);
                }
            }
        } else if (N >= mem_coalesce_width) {
            int numInnerIter = N / mem_coalesce_width;
            int numOuterIter = numXFormsPerWG / (groupSize / mem_coalesce_width);

            out << "    ii = lId & " << (mem_coalesce_width - 1) << ";\n";
            out << "    jj = lId >> " << ((int)log2(mem_coalesce_width)) << ";\n";
            out << "    lMemStore = sMem + mad24( jj, " << (N + numWorkItemsPerXForm) << ", ii );\n";
            out << "    offset = mad24( groupId, " << (numXFormsPerWG) << ", jj);\n";
            out << "    offset = mad24( offset, " << (N) << ", ii );\n";
            if (dataFormat == Fourier_transform_format::Interleaved_complex) {
                out << "        in += offset;\n";
                out << "        out += offset;\n";
            } else {
                out << "        in_real += offset;\n";
                out << "        in_imag += offset;\n";
                out << "        out_real += offset;\n";
                out << "        out_imag += offset;\n";
            }
            out << "if((groupId == get_num_groups(0)-1) && s) {\n";
            for(int i=0; i<numOuterIter; ++i) {
                out << "    if( jj < s ) {\n";
                for (int j=0; j<numInnerIter; ++j) {
                    formattedLoad(
                        out,
                        i * numInnerIter + j,
                        j * mem_coalesce_width + i * ( groupSize / mem_coalesce_width ) * N,
                        dataFormat
                    );
                }
                out << "    }\n";
                if (i != numOuterIter-1) {
                    out << "    jj += " << (groupSize / mem_coalesce_width) << ";\n";
                }
            }
            out << "}\n ";
            out << "else {\n";
            for (int i = 0; i < numOuterIter; i++ ) {
                for (int j = 0; j < numInnerIter; j++ ) {
                    formattedLoad(
                        out,
                        i * numInnerIter + j,
                        j * mem_coalesce_width + i * ( groupSize / mem_coalesce_width ) * N,
                        dataFormat
                    );
                }
            }
            out << "}\n";
            out << "    ii = lId & " << (numWorkItemsPerXForm - 1) << ";\n";
            out << "    jj = lId >> " << (log2NumWorkItemsPerXForm) << ";\n";
            out << "    lMemLoad  = sMem + mad24( jj, " << (N + numWorkItemsPerXForm) << ", ii);\n";
            for (int i=0; i<numOuterIter; ++i) {
                for (int j=0; j<numInnerIter; ++j) {
                    out << "    lMemStore["
                        << (j * mem_coalesce_width + i * ( groupSize / mem_coalesce_width ) * (N + numWorkItemsPerXForm ))
                        << "] = a["
                        << (i * numInnerIter + j)
                        << "].x;\n";
                }
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for (int i=0; i<R0; ++i) {
                out << "    a["
                    << (i)
                    << "].x = lMemLoad["
                    << (i * numWorkItemsPerXForm)
                    << "];\n";
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for (int i=0; i<numOuterIter; ++i) {
                for (int j=0; j<numInnerIter; ++j) {
                    out << "    lMemStore["
                        << (j * mem_coalesce_width + i * ( groupSize / mem_coalesce_width ) * (N + numWorkItemsPerXForm ))
                        << "] = a["
                        << (i * numInnerIter + j)
                        << "].y;\n";
                }
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for (int i=0; i<R0; ++i) {
                out << "    a["
                    << (i)
                    << "].y = lMemLoad["
                    << (i * numWorkItemsPerXForm) << "];\n";
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            lMemSize = (N + numWorkItemsPerXForm) * numXFormsPerWG;
        } else {
            out << "    offset = mad24( groupId,  " << (N * numXFormsPerWG) << ", lId );\n";
            if (dataFormat == Fourier_transform_format::Interleaved_complex) {
                out << "        in += offset;\n";
                out << "        out += offset;\n";
            } else {
                out << "        in_real += offset;\n";
                out << "        in_imag += offset;\n";
                out << "        out_real += offset;\n";
                out << "        out_imag += offset;\n";
            }
            out << "    ii = lId & " << (N-1) << ";\n";
            out << "    jj = lId >> " << ((int)log2(N)) << ";\n";
            out << "    lMemStore = sMem + mad24( jj, " << (N + numWorkItemsPerXForm) << ", ii );\n";
            out << "if((groupId == get_num_groups(0)-1) && s) {\n";
            for (int i=0; i<R0; ++i) {
                out << "    if(jj < s )\n";
                formattedLoad(out, i, i*groupSize, dataFormat);
                if (i != R0-1) {
                    out << "    jj += " << (groupSize / N) << ";\n";
                }
            }
            out << "}\n";
            out << "else {\n";
            for (int i=0; i<R0; ++i) {
                formattedLoad(out, i, i*groupSize, dataFormat);
            }
            out << "}\n";
            if (numWorkItemsPerXForm > 1) {
                out << "    ii = lId & " << (numWorkItemsPerXForm - 1) << ";\n";
                out << "    jj = lId >> " << (log2NumWorkItemsPerXForm) << ";\n";
                out << "    lMemLoad = sMem + mad24( jj, " << (N + numWorkItemsPerXForm) << ", ii );\n";
            } else {
                out << "    ii = 0;\n";
                out << "    jj = lId;\n";
                out << "    lMemLoad = sMem + mul24( jj, " << (N + numWorkItemsPerXForm) << ");\n";
            }
            for (int i=0; i<R0; ++i) {
                out << "    lMemStore[" << (i * ( groupSize / N ) * ( N + numWorkItemsPerXForm )) << "] = a[" << (i) << "].x;\n";
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";

            for (int i=0; i<R0; ++i) {
                out << "    a[" << (i) << "].x = lMemLoad[" << (i * numWorkItemsPerXForm) << "];\n";
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for (int i=0; i<R0; ++i) {
                out << "    lMemStore[" << (i * ( groupSize / N ) * ( N + numWorkItemsPerXForm )) << "] = a[" << (i) << "].y;\n";
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for (int i=0; i<R0; ++i) {
                out << "    a[" << (i) << "].y = lMemLoad[" << (i * numWorkItemsPerXForm) << "];\n";
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            lMemSize = (N + numWorkItemsPerXForm) * numXFormsPerWG;
        }
        return lMemSize;
    }

    int
    insertGlobalStoresAndTranspose(
        std::ostream& out,
        int N,
        int maxRadix,
        int Nr,
        int numWorkItemsPerXForm,
        int numXFormsPerWG,
        int mem_coalesce_width,
        vtb::opencl::Fourier_transform_format dataFormat
    ) {
        int groupSize = numWorkItemsPerXForm * numXFormsPerWG;
        int i, j, k, ind;
        int lMemSize = 0;
        int numIter = maxRadix / Nr;
        if( numWorkItemsPerXForm >= mem_coalesce_width )
        {
            if(numXFormsPerWG > 1)
            {
                out << "    if( !s || (groupId < get_num_groups(0)-1) || (jj < s) ) {\n";
            }
            for(i = 0; i < maxRadix; i++)
            {
                j = i % numIter;
                k = i / numIter;
                ind = j * Nr + k;
                formattedStore(out, ind, i*numWorkItemsPerXForm, dataFormat);
            }
            if(numXFormsPerWG > 1)
                out << "    }\n";
        }
        else if( N >= mem_coalesce_width )
        {
            int numInnerIter = N / mem_coalesce_width;
            int numOuterIter = numXFormsPerWG / ( groupSize / mem_coalesce_width );
            out << "    lMemLoad  = sMem + mad24( jj, " << (N + numWorkItemsPerXForm) << ", ii );\n";
            out << "    ii = lId & " << (mem_coalesce_width - 1) << ";\n";
            out << "    jj = lId >> " << ((int)log2(mem_coalesce_width)) << ";\n";
            out << "    lMemStore = sMem + mad24( jj," << (N + numWorkItemsPerXForm) << ", ii );\n";
            for( i = 0; i < maxRadix; i++ )
            {
                j = i % numIter;
                k = i / numIter;
                ind = j * Nr + k;
                out << "    lMemLoad[" << (i*numWorkItemsPerXForm) << "] = a[" << (ind) << "].x;\n";
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for( i = 0; i < numOuterIter; i++ )
                for( j = 0; j < numInnerIter; j++ )
                    out << "    a[" << (i*numInnerIter + j) << "].x = lMemStore[" << (j*mem_coalesce_width + i*( groupSize / mem_coalesce_width )*(N + numWorkItemsPerXForm)) << "];\n";
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for( i = 0; i < maxRadix; i++ )
            {
                j = i % numIter;
                k = i / numIter;
                ind = j * Nr + k;
                out << "    lMemLoad[" << (i*numWorkItemsPerXForm) << "] = a[" << (ind) << "].y;\n";
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for( i = 0; i < numOuterIter; i++ )
                for( j = 0; j < numInnerIter; j++ )
                    out << "    a[" << (i*numInnerIter + j) << "].y = lMemStore[" << (j*mem_coalesce_width + i*( groupSize / mem_coalesce_width )*(N + numWorkItemsPerXForm)) << "];\n";
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            out << "if((groupId == get_num_groups(0)-1) && s) {\n";
            for(i = 0; i < numOuterIter; i++ )
            {
                out << "    if( jj < s ) {\n";
                for (int j = 0; j < numInnerIter; j++ )
                    formattedStore(out, i*numInnerIter + j, j*mem_coalesce_width + i*(groupSize/mem_coalesce_width)*N, dataFormat);
                out << "    }\n";
                if(i != numOuterIter - 1)
                    out << "    jj += " << (groupSize / mem_coalesce_width) << ";\n";
            }
            out << "}\n";
            out << "else {\n";
            for(i = 0; i < numOuterIter; i++ )
            {
                for (int j = 0; j < numInnerIter; j++ )
                    formattedStore(out, i*numInnerIter + j, j*mem_coalesce_width + i*(groupSize/mem_coalesce_width)*N, dataFormat);
            }
            out << "}\n";
            lMemSize = (N + numWorkItemsPerXForm) * numXFormsPerWG;
        }
        else
        {
            out << "    lMemLoad  = sMem + mad24( jj," << (N + numWorkItemsPerXForm) << ", ii );\n";
            out << "    ii = lId & " << (N - 1) << ";\n";
            out << "    jj = lId >> " << ((int) log2(N)) << ";\n";
            out << "    lMemStore = sMem + mad24( jj," << (N + numWorkItemsPerXForm) << ", ii );\n";
            for( i = 0; i < maxRadix; i++ )
            {
                j = i % numIter;
                k = i / numIter;
                ind = j * Nr + k;
                out << "    lMemLoad[" << (i*numWorkItemsPerXForm) << "] = a[" << (ind) << "].x;\n";
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for( i = 0; i < maxRadix; i++ )
                out << "    a[" << (i) << "].x = lMemStore[" << (i*( groupSize / N )*( N + numWorkItemsPerXForm )) << "];\n";
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for( i = 0; i < maxRadix; i++ )
            {
                j = i % numIter;
                k = i / numIter;
                ind = j * Nr + k;
                out << "    lMemLoad[" << (i*numWorkItemsPerXForm) << "] = a[" << (ind) << "].y;\n";
            }
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            for( i = 0; i < maxRadix; i++ )
                out << "    a[" << (i) << "].y = lMemStore[" << (i*( groupSize / N )*( N + numWorkItemsPerXForm )) << "];\n";
            out << "    barrier( CLK_LOCAL_MEM_FENCE );\n";
            out << "if((groupId == get_num_groups(0)-1) && s) {\n";
            for( i = 0; i < maxRadix; i++ )
            {
                out << "    if(jj < s ) {\n";
                formattedStore(out, i, i*groupSize, dataFormat);
                out << "    }\n";
                if( i != maxRadix - 1)
                    out << "    jj +=" << (groupSize / N) << ";\n";
            }
            out << "}\n";
            out << "else {\n";
            for( i = 0; i < maxRadix; i++ )
            {
                formattedStore(out, i, i*groupSize, dataFormat);
            }
            out << "}\n";
            lMemSize = (N + numWorkItemsPerXForm) * numXFormsPerWG;
        }
        return lMemSize;
    }

    static void
    insertLocalLoadIndexArithmatic(
        std::ostream& out,
        int Nprev,
        int Nr,
        int numWorkItemsReq,
        int numWorkItemsPerXForm,
        int numXFormsPerWG,
        int offset,
        int midPad
    )
    {
        int Ncurr = Nprev * Nr;
        int logNcurr = (int)std::log2(Ncurr);
        int logNprev = (int)std::log2(Nprev);
        int incr = (numWorkItemsReq + offset) * Nr + midPad;
        if (Ncurr < numWorkItemsPerXForm) {
            if (Nprev == 1) {
                out << "    j = ii & " << (Ncurr - 1) << ";\n";
            } else {
                out << "    j = (ii & " << (Ncurr - 1) << ") >> " << logNprev << ";\n";
            }
            if (Nprev == 1) {
                out << "    i = ii >> " << logNcurr << ";\n";
            } else {
                out << "    i = mad24(ii >> " << logNcurr << ", "
                    << Nprev << ", ii & " << (Nprev-1) << ");\n";
            }
        } else {
            if (Nprev == 1) {
                out << "    j = ii;\n";
            } else {
                out << "    j = ii >> " << logNprev << ";\n";
            }
            if (Nprev == 1) {
                out << "    i = 0;\n";
            } else {
                out << "    i = ii & " << (Nprev-1) << ";\n";
            }
        }
        if (numXFormsPerWG > 1) {
            out << "    i = mad24(jj, " << incr << ", i);\n";
        }
        out << "    lMemLoad = sMem + mad24(j, "
            << (numWorkItemsReq + offset)
            << ", i);\n";
    }

    void
    insertLocalStoreIndexArithmatic(
        std::ostream& out,
        int numWorkItemsReq,
        int numXFormsPerWG,
        int Nr,
        int offset,
        int midPad
    ) {
        if (numXFormsPerWG == 1) {
            out << "    lMemStore = sMem + ii;\n";
        } else {
            out << "    lMemStore = sMem + mad24(jj, "
                << ((numWorkItemsReq + offset)*Nr + midPad)
                <<", ii);\n";
        }
    }

    void
    insertLocalStores(
        std::ostream& out,
        int numIter,
        int Nr,
        int numWorkItemsPerXForm,
        int numWorkItemsReq,
        int offset,
        const char* comp
    ) {
        for (int z=0; z<numIter; ++z) {
            for (int k=0; k<Nr; ++k) {
                int index = k*(numWorkItemsReq + offset) + z*numWorkItemsPerXForm;
                out << "    lMemStore[" << (index)
                    << "] = a[" << (z*Nr + k) << "]." << comp << ";\n";
            }
        }
        out << "    barrier(CLK_LOCAL_MEM_FENCE);\n";
    }

    void
    insertLocalLoads(
        std::ostream& out,
        int n,
        int Nr,
        int Nrn,
        int Nprev,
        int Ncurr,
        int numWorkItemsPerXForm,
        int numWorkItemsReq,
        int offset,
        const char* comp
    ) {
        int numWorkItemsReqN = n / Nrn;
        int interBlockHNum = std::max( Nprev / numWorkItemsPerXForm, 1 );
        int interBlockHStride = numWorkItemsPerXForm;
        int vertWidth = std::max(numWorkItemsPerXForm / Nprev, 1);
        vertWidth = std::min( vertWidth, Nr);
        int vertNum = Nr / vertWidth;
        int vertStride = ( n / Nr + offset ) * vertWidth;
        int iter = std::max( numWorkItemsReqN / numWorkItemsPerXForm, 1);
        int intraBlockHStride = (numWorkItemsPerXForm / (Nprev*Nr)) > 1 ? (numWorkItemsPerXForm / (Nprev*Nr)) : 1;
        intraBlockHStride *= Nprev;
        int stride = numWorkItemsReq / Nrn;
        int i;
        for(i = 0; i < iter; i++) {
            int ii = i / (interBlockHNum * vertNum);
            int zz = i % (interBlockHNum * vertNum);
            int jj = zz % interBlockHNum;
            int kk = zz / interBlockHNum;
            int z;
            for(z = 0; z < Nrn; z++) {
                int st = kk * vertStride + jj * interBlockHStride + ii * intraBlockHStride + z * stride;
                out << "    a[" << (i*Nrn + z) << "]."
                    << comp << " = lMemLoad[" << (st) << "];\n";
            }
        }
        out << "    barrier(CLK_LOCAL_MEM_FENCE);\n";
    }

    int
    getPadding(
        int numWorkItemsPerXForm,
        int Nprev,
        int numWorkItemsReq,
        int numXFormsPerWG,
        int Nr,
        int numBanks,
        int* offset,
        int* midPad
    ) {
        if((numWorkItemsPerXForm <= Nprev) || (Nprev >= numBanks))
            *offset = 0;
        else {
            int numRowsReq = ((numWorkItemsPerXForm < numBanks) ? numWorkItemsPerXForm : numBanks) / Nprev;
            int numColsReq = 1;
            if(numRowsReq > Nr)
                numColsReq = numRowsReq / Nr;
            numColsReq = Nprev * numColsReq;
            *offset = numColsReq;
        }
        if(numWorkItemsPerXForm >= numBanks || numXFormsPerWG == 1)
            *midPad = 0;
        else {
            int bankNum = ( (numWorkItemsReq + *offset) * Nr ) & (numBanks - 1);
            if( bankNum >= numWorkItemsPerXForm )
                *midPad = 0;
            else
                *midPad = numWorkItemsPerXForm - bankNum;
        }
        int lMemSize = ( numWorkItemsReq + *offset) * Nr * numXFormsPerWG + *midPad * (numXFormsPerWG - 1);
        return lMemSize;
    }

    void
    getGlobalRadixInfo(
        int n,
        int *radix,
        int *R1,
        int *R2,
        int *numRadices
    ) {
        int baseRadix = std::min(n, 128);
        int numR = 0;
        int N = n;
        while (N > baseRadix) {
            N /= baseRadix;
            numR++;
        }
        for (int i = 0; i < numR; i++) {
            radix[i] = baseRadix;
        }
        radix[numR] = N;
        numR++;
        *numRadices = numR;
        for (int i = 0; i < numR; i++) {
            int B = radix[i];
            if (B <= 8) {
                R1[i] = B;
                R2[i] = 1;
                continue;
            }
            int r1 = 2;
            int r2 = B / r1;
            while (r2 > r1) {
               r1 *= 2;
               r2 = B / r1;
            }
            R1[i] = r1;
            R2[i] = r2;
        }
    }

    inline void
    trim_right(std::string& rhs) {
        while (!rhs.empty() && rhs.back() <= ' ') { rhs.pop_back(); }
    }

    inline std::string
    trim(std::string rhs) {
        trim_right(rhs);
        return rhs;
    }

    template <class T, int N>
    inline size_t
    num_bytes(const blitz::TinyVector<int,N>& n) {
        return blitz::product(n)*sizeof(T);
    }

}

std::string
vtb::opencl::Fourier_transform_base::kernel_name(const char* prefix) {
    const char sep = '_';
    std::stringstream name;
    name << prefix << sep
        << this->_shape(0) << sep
        << this->_shape(1) << sep
        << this->_shape(2) << sep
        << ++this->_kindex;
    return name.str();
}

void
vtb::opencl::Fourier_transform_base::generate_source_code() {
    this->_src = base_kernels;
    this->_kindex = 0;
    this->_kernels.clear();
    for (int i=0; i<3; ++i) {
        this->generate_fft(i);
    }
    for (const auto& kernel : this->_kernels) {
        if (!kernel.in_place_possible) {
            this->temp_buffer_needed = true;
            break;
        }
    }
}

void
vtb::opencl::Fourier_transform_base::generate_fft(int axis) {
    if (axis == 0) {
        int nx = this->_shape(0);
        if (nx > this->max_localmem_fft_size) {
            generate_fft_global(nx, 1, axis, 1);
        } else if (nx > 1) {
            std::vector<int> radices{radix_array(nx)};
            if (nx/radices[0] <= this->_maxworkgroupsize) {
                generate_fft_local();
            } else {
                radices = radix_array(nx, this->_maxradix);
                if (nx/radices[0] <= this->_maxworkgroupsize) {
                    generate_fft_local();
                } else {
                    generate_fft_global(nx, 1, axis, 1);
                }
            }
        }
    }
    if (axis == 1) {
        int ny = this->_shape(1);
        if (ny > 1) {
            int stride = this->_shape(0);
            generate_fft_global(ny, stride, axis, 1);
        }
    }
    if (axis == 2) {
        int nz = this->_shape(2);
        if (nz > 1) {
            int stride = _shape(0)*_shape(1);
            generate_fft_global(nz, stride, axis, 1);
        }
    }
}

void
vtb::opencl::Fourier_transform_base::generate_fft_local() {
    int n = this->_shape(0);
    if (n > this->_maxworkgroupsize*this->_maxradix) {
        throw std::invalid_argument{"signal length too big for local mem fft"};
    }
    std::vector<int> radices{radix_array(n)};
    if (n/radices[0] > this->_maxworkgroupsize) {
        radices = radix_array(n, this->_maxradix);
    }
    if (radices.front() > this->_maxradix) {
        throw std::invalid_argument{"bad radix array"};
    }
    if (n/radices.front() > this->_maxworkgroupsize) {
        throw std::invalid_argument{
            "required work items per xform greater than "
            "maximum work items allowed per work group for local mem fft"
        };
    }
    int numRadix = radices.size();
    {
        int prod = 1;
        for (int i=0; i<numRadix; ++i) {
            prod *= radices[i];
        }
        if (prod != n) {
            throw std::invalid_argument{"bad radices"};
        }
    }
    int offset, midPad;
    std::stringstream out;
    Kernel_info kernel{};
    kernel.name = this->kernel_name("fft_local");
    kernel.axis = 0;
    kernel.in_place_possible = true;
    int numWorkItemsPerXForm = n / radices[0];
    int numWorkItemsPerWG = numWorkItemsPerXForm <= 64 ? 64 : numWorkItemsPerXForm;
    assert(numWorkItemsPerWG <= this->_maxworkgroupsize);
    int numXFormsPerWG = numWorkItemsPerWG / numWorkItemsPerXForm;
    kernel.num_workgroups = 1;
    kernel.num_xforms_per_workgroup = numXFormsPerWG;
    kernel.num_workitems_per_workgroup = numWorkItemsPerWG;
    int maxRadix = radices[0];
    int lMemSize = 0;
    insertVariables(out, maxRadix);
    lMemSize = insertGlobalLoadsAndTranspose(
        out,
        n,
        numWorkItemsPerXForm,
        numXFormsPerWG,
        maxRadix,
        this->min_mem_coalesce_width,
        this->_format
    );
    kernel.lmem_size = (lMemSize > kernel.lmem_size) ? lMemSize : kernel.lmem_size;
    int Nprev = 1;
    int len = n;
    for(int r = 0; r<numRadix; ++r) {
        int numIter = radices[0] / radices[r];
        int numWorkItemsReq = n / radices[r];
        int Ncurr = Nprev * radices[r];
        insertfftKernel(out, radices[r], numIter);
        if (r < (numRadix-1)) {
            insertTwiddleKernel(
                out,
                radices[r],
                numIter,
                Nprev,
                len,
                numWorkItemsPerXForm
            );
            lMemSize = getPadding(
                numWorkItemsPerXForm,
                Nprev,
                numWorkItemsReq,
                numXFormsPerWG,
                radices[r],
                this->num_local_mem_banks,
                &offset,
                &midPad
            );
            kernel.lmem_size = (lMemSize > kernel.lmem_size)
                ? lMemSize
                : kernel.lmem_size;
            insertLocalStoreIndexArithmatic(
                out,
                numWorkItemsReq,
                numXFormsPerWG,
                radices[r],
                offset,
                midPad
            );
            insertLocalLoadIndexArithmatic(
                out,
                Nprev,
                radices[r],
                numWorkItemsReq,
                numWorkItemsPerXForm,
                numXFormsPerWG,
                offset,
                midPad
            );
            insertLocalStores(
                out,
                numIter,
                radices[r],
                numWorkItemsPerXForm,
                numWorkItemsReq,
                offset,
                "x"
            );
            insertLocalLoads(
                out,
                n,
                radices[r],
                radices[r+1],
                Nprev,
                Ncurr,
                numWorkItemsPerXForm,
                numWorkItemsReq,
                offset,
                "x"
            );
            insertLocalStores(
                out,
                numIter,
                radices[r],
                numWorkItemsPerXForm,
                numWorkItemsReq,
                offset,
                "y"
            );
            insertLocalLoads(
                out,
                n,
                radices[r],
                radices[r+1],
                Nprev,
                Ncurr,
                numWorkItemsPerXForm,
                numWorkItemsReq,
                offset,
                "y"
            );
            Nprev = Ncurr;
            len = len / radices[r];
        }
    }
    lMemSize = insertGlobalStoresAndTranspose(
        out,
        n,
        maxRadix,
        radices[numRadix - 1],
        numWorkItemsPerXForm,
        numXFormsPerWG,
        this->min_mem_coalesce_width,
        this->_format
    );
    kernel.lmem_size = (lMemSize > kernel.lmem_size) ? lMemSize : kernel.lmem_size;
    std::stringstream result;
    result << this->_src;
    insertHeader(result, kernel.name, this->_format);
    result << "{\n";
    if (kernel.lmem_size) {
        result << "    __local float sMem[" << kernel.lmem_size << "];\n";
    }
    result << out.str();
    result << "}\n";
    this->_src += result.str();
    this->_kernels.emplace_back(kernel);
}

void
vtb::opencl::Fourier_transform_base::generate_fft_global(
    int n,
    int BS,
    int axis,
    int vertBS
) {
    int k, t;
    int radixArr[10] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
    int R1Arr[10] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
    int R2Arr[10] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
    int radix, R1, R2;
    int numRadices;
    int maxThreadsPerBlock = this->_maxworkgroupsize;
    int batchSize = this->min_mem_coalesce_width;
    int vertical = (axis == 0) ? 0 : 1;
    getGlobalRadixInfo(n, radixArr, R1Arr, R2Arr, &numRadices);
    int numPasses = numRadices;
    int N = n;
    int m = (int)log2(n);
    int Rinit = vertical ? BS : 1;
    batchSize = vertical ? std::min(BS, batchSize) : batchSize;
    std::stringstream out;
    for (int passNum=0; passNum<numPasses; ++passNum) {
        out.str("");
        radix = radixArr[passNum];
        R1 = R1Arr[passNum];
        R2 = R2Arr[passNum];
        int strideI = Rinit;
        for (int i=0; i<numPasses; ++i) {
            if (i != passNum){
                strideI *= radixArr[i];
            }
        }
        int strideO = Rinit;
        for (int i=0; i<passNum; ++i) {
            strideO *= radixArr[i];
        }
        int threadsPerXForm = R2;
        batchSize = R2 == 1 ? this->_maxworkgroupsize : batchSize;
        batchSize = std::min(batchSize, strideI);
        int threadsPerBlock = batchSize * threadsPerXForm;
        threadsPerBlock = std::min(threadsPerBlock, maxThreadsPerBlock);
        batchSize = threadsPerBlock / threadsPerXForm;
        assert(R2 <= R1);
        assert(R1*R2 == radix);
        assert(R1 <= this->_maxradix);
        assert(threadsPerBlock <= maxThreadsPerBlock);
        int numIter = R1 / R2;
        int gInInc = threadsPerBlock / batchSize;
        int lgStrideO = (int)log2(strideO);
        int numBlocksPerXForm = strideI / batchSize;
        int numBlocks = numBlocksPerXForm;
        if (!vertical) {
            numBlocks *= BS;
        } else {
            numBlocks *= vertBS;
        }
        Kernel_info kernel{};
        kernel.name = this->kernel_name("fft_global");
        if (R2 == 1) {
            kernel.lmem_size = 0;
        } else {
            if (strideO == 1) {
                kernel.lmem_size = (radix + 1)*batchSize;
            } else {
                kernel.lmem_size = threadsPerBlock*R1;
            }
        }
        kernel.num_workgroups = numBlocks;
        kernel.num_xforms_per_workgroup = 1;
        kernel.num_workitems_per_workgroup = threadsPerBlock;
        kernel.axis = axis;
        if((passNum == (numPasses - 1)) && (numPasses & 1)) {
            kernel.in_place_possible = true;
        } else {
            kernel.in_place_possible = false;
        }
        insertVariables(out, R1);
        if (vertical) {
            out << "xNum = groupId >> " << ((int)log2(numBlocksPerXForm)) << ";\n";
            out << "groupId = groupId & " << (numBlocksPerXForm - 1) << ";\n";
            out << "indexIn = mad24(groupId, " << (batchSize) << ", xNum << " << ((int)log2(n*BS)) << ");\n";
            out << "tid = mul24(groupId, " << (batchSize) << ");\n";
            out << "i = tid >> " << (lgStrideO) << ";\n";
            out << "j = tid & " << (strideO - 1) << ";\n";
            int stride = radix*Rinit;
            for (int i=0; i<passNum; ++i) {
                stride *= radixArr[i];
            }
            out << "indexOut = mad24(i, " << (stride) << ", j + " << "(xNum << " << ((int) log2(n*BS)) << "));\n";
            out << "bNum = groupId;\n";
        } else {
            int lgNumBlocksPerXForm = (int)log2(numBlocksPerXForm);
            out << "bNum = groupId & " << (numBlocksPerXForm - 1) << ";\n";
            out << "xNum = groupId >> " << (lgNumBlocksPerXForm) << ";\n";
            out << "indexIn = mul24(bNum, " << (batchSize) << ");\n";
            out << "tid = indexIn;\n";
            out << "i = tid >> " << (lgStrideO) << ";\n";
            out << "j = tid & " << (strideO - 1) << ";\n";
            int stride = radix*Rinit;
            for (int i=0; i<passNum; ++i) {
                stride *= radixArr[i];
            }
            out << "indexOut = mad24(i, " << (stride) << ", j);\n";
            out << "indexIn += (xNum << " << (m) << ");\n";
            out << "indexOut += (xNum << " << (m) << ");\n";
        }
        // Load Data
        int lgBatchSize = (int)log2(batchSize);
        out << "tid = lId;\n";
        out << "i = tid & " << (batchSize - 1) << ";\n";
        out << "j = tid >> " << (lgBatchSize) << ";\n";
        out << "indexIn += mad24(j, " << (strideI) << ", i);\n";
        if (this->_format == Fourier_transform_format::Split_complex) {
            out << "in_real += indexIn;\n";
            out << "in_imag += indexIn;\n";
            for (int j=0; j<R1; ++j)
                out << "a[" << (j) << "].x = in_real[" << (j*gInInc*strideI) << "];\n";
            for (int j=0; j<R1; ++j)
                out << "a[" << (j) << "].y = in_imag[" << (j*gInInc*strideI) << "];\n";
        } else {
            out << "in += indexIn;\n";
            for (int j=0; j<R1; ++j) {
                out << "a[" << (j) << "] = in[" << (j*gInInc*strideI) << "];\n";
            }
        }
        out << "fftKernel" << (R1) << "(a, dir);\n";
        if (R2 > 1) {
            // twiddle
            for (int k = 1; k < R1; k++) {
                out << "ang = dir*(2.0f*M_PI*" << (k) << "/" << (radix) << ")*j;\n";
                out << "w = (float2)(native_cos(ang), native_sin(ang));\n";
                out << "a[" << (k) << "] = complexMul(a[" << (k) << "], w);\n";
            }
            // shuffle
            numIter = R1 / R2;
            out << "indexIn = mad24(j, " << (threadsPerBlock*numIter) << ", i);\n";
            out << "lMemStore = sMem + tid;\n";
            out << "lMemLoad = sMem + indexIn;\n";
            for (int k = 0; k < R1; k++) {
                out << "lMemStore[" << (k*threadsPerBlock) << "] = a[" << (k) << "].x;\n";
            }
            out << "barrier(CLK_LOCAL_MEM_FENCE);\n";
            for(k = 0; k < numIter; k++)
                for(t = 0; t < R2; t++)
                    out << "a[" << (k*R2+t) << "].x = lMemLoad[" << (t*batchSize + k*threadsPerBlock) << "];\n";
            out << "barrier(CLK_LOCAL_MEM_FENCE);\n";
            for(k = 0; k < R1; k++)
                out << "lMemStore[" << (k*threadsPerBlock) << "] = a[" << (k) << "].y;\n";
            out << "barrier(CLK_LOCAL_MEM_FENCE);\n";
            for(k = 0; k < numIter; k++)
                for(t = 0; t < R2; t++)
                    out << "a[" << (k*R2+t) << "].y = lMemLoad[" << (t*batchSize + k*threadsPerBlock) << "];\n";
            out << "barrier(CLK_LOCAL_MEM_FENCE);\n";
            for(int j = 0; j < numIter; j++)
                out << "fftKernel" << (R2) << "(a + " << (j*R2) << ", dir);\n";
        }
        // twiddle
        if (passNum < (numPasses - 1)) {
            out << "l = ((bNum << " << (lgBatchSize) << ") + i) >> " << (lgStrideO) << ";\n";
            out << "k = j << " << ((int)log2(R1/R2)) << ";\n";
            out << "ang1 = dir*(2.0f*M_PI/" << (N) << ")*l;\n";
            for(t = 0; t < R1; t++)
            {
                out << "ang = ang1*(k + " << ((t%R2)*R1 + (t/R2)) << ");\n";
                out << "w = (float2)(native_cos(ang), native_sin(ang));\n";
                out << "a[" << (t) << "] = complexMul(a[" << (t) << "], w);\n";
            }
        }
        // Store Data
        if(strideO == 1) {
            out << "lMemStore = sMem + mad24(i, " << (radix + 1) << ", j << " << ((int)log2(R1/R2)) << ");\n";
            out << "lMemLoad = sMem + mad24(tid >> " << ((int)log2(radix)) << ", " << (radix+1) << ", tid & " << (radix-1) << ");\n";
            for(int i = 0; i < R1/R2; i++)
                for(int j = 0; j < R2; j++)
                    out << "lMemStore[ " << (i + j*R1) << "] = a[" << (i*R2+j) << "].x;\n";
            out << "barrier(CLK_LOCAL_MEM_FENCE);\n";
            if(threadsPerBlock >= radix)
            {
                for(int i = 0; i < R1; i++)
                    out << "a[" << (i) << "].x = lMemLoad[" << (i*(radix+1)*(threadsPerBlock/radix)) << "];\n";
            }
            else
            {
                int innerIter = radix/threadsPerBlock;
                int outerIter = R1/innerIter;
                for(int i = 0; i < outerIter; i++)
                    for(int j = 0; j < innerIter; j++)
                        out << "a[" << (i*innerIter+j) << "].x = lMemLoad[" << (j*threadsPerBlock + i*(radix+1)) << "];\n";
            }
            out << "barrier(CLK_LOCAL_MEM_FENCE);\n";
            for (int i = 0; i < R1/R2; i++)
                for(int j = 0; j < R2; j++)
                    out << "lMemStore[ " << (i + j*R1) << "] = a[" << (i*R2+j) << "].y;\n";
            out << "barrier(CLK_LOCAL_MEM_FENCE);\n";
            if (threadsPerBlock >= radix) {
                for(int i = 0; i < R1; i++)
                    out << "a[" << (i) << "].y = lMemLoad[" << (i*(radix+1)*(threadsPerBlock/radix)) << "];\n";
            }
            else
            {
                int innerIter = radix/threadsPerBlock;
                int outerIter = R1/innerIter;
                for(int i = 0; i < outerIter; i++)
                    for (int j = 0; j < innerIter; j++)
                        out << "a[" << (i*innerIter+j) << "].y = lMemLoad[" << (j*threadsPerBlock + i*(radix+1)) << "];\n";
            }
            out << "barrier(CLK_LOCAL_MEM_FENCE);\n";
            out << "indexOut += tid;\n";
            if(this->_format == Fourier_transform_format::Split_complex) {
                out << "out_real += indexOut;\n";
                out << "out_imag += indexOut;\n";
                for(k = 0; k < R1; k++)
                    out << "out_real[" << (k*threadsPerBlock) << "] = a[" << (k) << "].x;\n";
                for(k = 0; k < R1; k++)
                    out << "out_imag[" << (k*threadsPerBlock) << "] = a[" << (k) << "].y;\n";
            }
            else {
                out << "out += indexOut;\n";
                for(k = 0; k < R1; k++)
                    out << "out[" << (k*threadsPerBlock) << "] = a[" << (k) << "];\n";
            }
        } else {
            out << "indexOut += mad24(j, " << (numIter*strideO) << ", i);\n";
            if(this->_format == Fourier_transform_format::Split_complex) {
                out << "out_real += indexOut;\n";
                out << "out_imag += indexOut;\n";
                for(k = 0; k < R1; k++)
                    out << "out_real[" << (((k%R2)*R1 + (k/R2))*strideO) << "] = a[" << (k) << "].x;\n";
                for(k = 0; k < R1; k++)
                    out << "out_imag[" << (((k%R2)*R1 + (k/R2))*strideO) << "] = a[" << (k) << "].y;\n";
            }
            else {
                out << "out += indexOut;\n";
                for(k = 0; k < R1; k++)
                    out << "out[" << (((k%R2)*R1 + (k/R2))*strideO) << "] = a[" << (k) << "];\n";
            }
        }
        std::stringstream result;
        insertHeader(result, kernel.name, this->_format);
        result << "{\n";
        if (kernel.lmem_size) {
            result << "    __local float sMem[" << (kernel.lmem_size) << "];\n";
        }
        result << out.str();
        result << "}\n";
        this->_src += result.str();
        N /= radix;
        this->_kernels.emplace_back(kernel);
    }
}

void
vtb::opencl::Fourier_transform_base::getKernelWorkDimensions(
    const Kernel_info& kernel,
    int* batchSize,
    size_t* gWorkItems,
    size_t* lWorkItems
) {
    *lWorkItems = kernel.num_workitems_per_workgroup;
    int numWorkGroups = kernel.num_workgroups;
    int numXFormsPerWG = kernel.num_xforms_per_workgroup;
    int ny = this->_shape(1);
    int nz = this->_shape(2);
    if (kernel.axis == 0) {
        *batchSize *= ny*nz;
        numWorkGroups = (*batchSize % numXFormsPerWG)
            ? (*batchSize/numXFormsPerWG + 1)
            : (*batchSize/numXFormsPerWG);
        numWorkGroups *= kernel.num_workgroups;
    }
    if (kernel.axis == 1) {
        *batchSize *= nz;
        numWorkGroups *= *batchSize;
    }
    if (kernel.axis == 2) {
        numWorkGroups *= *batchSize;
    }
    *gWorkItems = numWorkGroups * *lWorkItems;
}

void
vtb::opencl::Fourier_transform_base::precompile(const int3& max_power, Context* context) {
    using clock = std::chrono::system_clock;
    using std::chrono::seconds;
    int ni = max_power(0);
    int nj = max_power(1);
    std::atomic<int> count{};
    auto t0 = clock::now();
    int max_count = product(max_power);
    std::atomic<bool> slow{false};
    #if defined(VTB_WITH_OPENMP)
    #pragma omp parallel for collapse(2) schedule(dynamic,1)
    #endif
    for (int i=1; i<=ni; ++i) {
        for (int j=1; j<=nj; ++j) {
            Fourier_transform_base fft;
            fft.context(context);
            fft.shape({1<<i, 1<<j, 1});
            if (clock::now()-t0 > seconds(1)) { slow = true; }
            auto cnt = ++count;
            if (slow && cnt%10 == 0 && cnt >= 10) {
                std::fprintf(stderr, "%5d/%-5d compile fft\n", cnt, max_count);
            }
        }
    }
}

void
vtb::opencl::Fourier_transform_base::dump(std::ostream& out) {
    size_t global = 0, local = 0;
    for (const auto& kernel : this->_kernels) {
        cl_int batch_size = 1;
        getKernelWorkDimensions(kernel, &batch_size, &global, &local);
        out << std::setw(20) << kernel.name
            << std::setw(20) << global
            << std::setw(20) << local
            << std::endl;
    }
    out << this->_src;
}

vtb::opencl::Fourier_transform_base::Fourier_transform_base(const int3& shape):
_shape{shape} {
    this->init();
}

void
vtb::opencl::Fourier_transform_base::allocate_temporary_buffer(int batch_size) {
    if (this->last_batch_size != batch_size) {
        this->last_batch_size = batch_size;
        size_t n = this->buffer_size(batch_size);
        this->_workarea = context()->context().buffer(clx::memory_flags::read_write, n);
    }
}

void
vtb::opencl::Fourier_transform_base::enqueue(
    clx::buffer x,
    int direction,
    int batch_size
) {
    auto& ppl = context()->pipeline();
    clx::buffer buffer_in = x, buffer_out(nullptr);
    Kernel_info* first = this->_kernels.data();
    Kernel_info* last = this->_kernels.data() + this->_kernels.size();
    // compute in-place transforms
    while (first != last && first->in_place_possible) {
        auto& k = *first;
        int new_batch_size = batch_size;
        size_t local = 0, global = 0;
        this->getKernelWorkDimensions(k, &new_batch_size, &global, &local);
        k.kernel.argument(0, buffer_in);
        k.kernel.argument(1, buffer_in);
        k.kernel.argument(2, direction);
        k.kernel.argument(3, new_batch_size);
        ppl.step();
        ppl.kernel(k.kernel, clx::range{global}, clx::range{local});
        ++first;
    }
    // allocate temporary buffer to compute the remaining out-of-place transforms
    if (first != last) {
        this->allocate_temporary_buffer(batch_size);
        buffer_out = this->_workarea;
        auto remaining = last - first;
        if (remaining%2 != 0) {
            // TODO this is not tested
            ppl.step();
            ppl.copy(buffer_in, buffer_out);
            std::swap(buffer_in, buffer_out);
        }
    }
    // compute out-of-place transforms by swapping between actual and
    // temporary buffer
    while (first != last) {
        auto& k = *first;
        int new_batch_size = batch_size;
        size_t local = 0, global = 0;
        this->getKernelWorkDimensions(k, &new_batch_size, &global, &local);
        k.kernel.argument(0, buffer_in);
        k.kernel.argument(1, buffer_out);
        k.kernel.argument(2, direction);
        k.kernel.argument(3, new_batch_size);
        ppl.step();
        ppl.kernel(k.kernel, clx::range{global}, clx::range{local});
        std::swap(buffer_in, buffer_out);
        ++first;
    }
}

void
vtb::opencl::Fourier_transform_base::init() {
    if (!blitz::is_power_of_two(blitz::product(this->_shape))) {
        throw std::invalid_argument{"bad shape"};
    }
    int dim = 0;
    while (dim < 3 && this->_shape(dim) != 1) {
        ++dim;
    }
    if (dim < 0 || dim > 2) {
        throw std::invalid_argument("OpenCL FFT supports 1,2,3 dimensions only");
    }
    this->_ndimensions = dim;
    bool success = false;
    clx::compiler cc = context()->compiler_copy();
    cc.options(cc.options() + " -cl-mad-enable");
    auto device = cc.devices().front();
    this->_maxworkgroupsize = device.max_work_group_size();
    while (!success) {
        this->generate_source_code();
        clx::program prog = cc.compile("fft.cl", this->_src);
        auto all_kernels = prog.kernels();
        for (auto& kernel : all_kernels) {
            auto name = kernel.name();
            for (auto& k : this->_kernels) {
                if (name == k.name) { k.kernel = kernel; break; }
            }
        }
        success = true;
        size_t min_wg_size = std::numeric_limits<size_t>::max();
        for (const auto& k : this->_kernels) {
            auto wg = k.kernel.work_group(device);
            if (wg.size < size_t(k.num_workitems_per_workgroup)) { success = false; }
            if (wg.size < min_wg_size) { min_wg_size = wg.size; }
        }
        this->_maxworkgroupsize = min_wg_size;
    }
}

void
vtb::opencl::Chirp_Z_transform_base::context(Context* rhs) {
    this->_fft.context(rhs);
    auto prog = context()->compiler().compile("chirp_z_transform.cl");
    _makechirp = prog.kernel("make_chirp");
    _reciprocal_chirp = prog.kernel("reciprocal_chirp");
    _mult1 = prog.kernel("multiply_1");
    _mult2 = prog.kernel("multiply_2");
    _mult3 = prog.kernel("multiply_3");
    _zero_init = prog.kernel("zero_init");
}

void
vtb::opencl::Chirp_Z_transform_base::make_chirp(
    const int3& shape,
    const int3& fft_shape
) {
    auto& ppl = context()->pipeline();
    _shape = shape;
    int3 chirp_shape{shape*2-1};
    ppl.allocate(product(chirp_shape), this->_chirp);
    ppl.allocate(product(fft_shape), this->_xp);
    ppl.allocate(product(fft_shape), this->_ichirp);
    auto& kernel = this->_makechirp;
    kernel.arguments(this->_chirp, shape(0), shape(1), shape(2));
    ppl.kernel(kernel, chirp_shape);
    ppl.step();
    VTB_DUMP(this->_chirp, chirp_shape, "chirp");
}

void
vtb::opencl::Chirp_Z_transform_base::enqueue(
    clx::buffer x,
    int direction,
    int batch_size
) {
    using blitz::product;
    if (batch_size != 1) {
        throw std::runtime_error{"batch size > 1 not supported"};
    }
    auto& ppl = this->_fft.context()->pipeline();
    // TODO Step is deprecated. Replace it with Stack.
    // TODO zero_init is not copied, but this is not a problem until
    // we run two steps in parallel, which modern gpus cannot do
//  Step st;
    {
//      Step st1;
        clx::kernel zero = _zero_init;
        zero.argument(0, _xp);
        _mult1.argument(0, x);
        _mult1.argument(1, _chirp);
        _mult1.argument(2, direction);
        _mult1.argument(3, _fft.shape()(0));
        _mult1.argument(4, _fft.shape()(1));
        _mult1.argument(5, _fft.shape()(2));
        _mult1.argument(6, _xp);
        ppl.kernel(zero, _fft.shape());
        ppl.step();
        ppl.kernel(_mult1, _shape);
        ppl.step();
        VTB_DUMP(this->_xp, this->_fft.shape(), "xp");
        _fft.enqueue(_xp, direction, batch_size);
        VTB_DUMP(this->_xp, this->_fft.shape(), "fft(xp)");
    }
    int3 chirp_shape{_shape*2-1};
    {
//      Step st1;
        clx::kernel zero = _zero_init;
        zero.argument(0, _ichirp);
        _reciprocal_chirp.argument(0, _chirp);
        _reciprocal_chirp.argument(1, direction);
        _reciprocal_chirp.argument(2, _fft.shape()(0));
        _reciprocal_chirp.argument(3, _fft.shape()(1));
        _reciprocal_chirp.argument(4, _fft.shape()(2));
        _reciprocal_chirp.argument(5, _ichirp);
        ppl.kernel(zero, _fft.shape());
        ppl.step();
        ppl.kernel(_reciprocal_chirp, chirp_shape);
        ppl.step();
        VTB_DUMP(this->_ichirp, this->_fft.shape(), "ichirp");
        _fft.enqueue(_ichirp, direction, batch_size);
        VTB_DUMP(this->_ichirp, this->_fft.shape(), "fft(ichirp)");
    }
    _mult2.argument(0, _ichirp);
    _mult2.argument(1, _xp);
    ppl.step();
    ppl.kernel(_mult2, _fft.shape());
    VTB_DUMP(this->_ichirp, this->_fft.shape(), "mult2");
    ppl.step();
    _fft.enqueue(_ichirp, -direction, batch_size);
    VTB_DUMP(this->_ichirp, this->_fft.shape(), "ifft");
    ppl.step();
    _mult3.argument(0, _ichirp);
    _mult3.argument(1, _chirp);
    _mult3.argument(2, x);
    _mult3.argument(3, _fft.shape()(0));
    _mult3.argument(4, _fft.shape()(1));
    _mult3.argument(5, _fft.shape()(2));
    _mult3.argument(6, 1.f / product(_fft.shape()));
    ppl.kernel(_mult3, _shape);
}