codelets.cc

// Code generation

#define CODELETS 1

#include "arith.h"
#include "complex.h"
#include "exp.h"
#include "debug.h"
#include "expansion.h"
#include "format.h"
#include "join.h"
#include "nearest.h"
#include "noncopyable.h"
#include "print.h"
#include "series.h"
#include "sig.h"
#include <deque>
#include <fstream>
#include <memory>
#include <random>
#include <tuple>
#include <type_traits>
#include <unordered_map>
#include <unordered_set>
#include <vector>
namespace mandelbrot {

using std::decay_t;
using std::deque;
using std::endl;
using std::make_pair;
using std::make_tuple;
using std::ofstream;
using std::ostream;
using std::remove_cvref_t;
using std::shared_ptr;
using std::string_view;
using std::tie;
using std::tuple;
using std::unique_ptr;
using std::unordered_map;
using std::unordered_set;
using std::vector;
typedef vector<Exp> Exps;
typedef Complex<Exp> CExp;
typedef vector<CExp> CExps;

// Container utilities
template<class C> auto pop_front(C& xs) { auto x = xs.front(); xs.pop_front(); return x; }
template<class C> auto pop_back(C& xs) { auto x = xs.back(); xs.pop_back(); return x; }
template<class C> bool any(const C& xs) { return std::any_of(xs.begin(), xs.end(), [](auto&& x) { return bool(x); }); }
template<class T,class C> void extend(vector<T>& xs, const C& ys) { xs.insert(xs.end(), ys.begin(), ys.end()); }

// Indented printing
unique_ptr<ostream> out;
int indent = 0;
const bool debug = false;

struct Indent : public Noncopyable {
  const int outer;
  Indent() : outer(indent) { indent += 2; }
  ~Indent() { indent = outer; }
};

void line() {
  slow_assert(out);
  if (debug) print();
  *out << endl;
}
template<class T> void line(T&& x) {
  slow_assert(out);
  const auto s = tfm::format("%*s%s", indent, "", x);
  if (debug) print(s);
  *out << s << endl;
}
template<class... Args> void line(const Args&... args) { line(tfm::format(args...)); }

struct Scope : public Noncopyable {
  const string close;
  unique_ptr<Indent> i;
  template<class... Args> Scope(const string& close, const Args&... args)
    : close(close) {
    line(tfm::format(args...));
    i.reset(new Indent);
  }
  ~Scope() { i.reset(); line(close); }
};

struct Blank : public Noncopyable { ~Blank() { line(); }};

struct Header : public Noncopyable {
  Header(const string& path, const string& comment, const vector<string>& includes) {
    slow_assert(!out);
    out.reset(new ofstream(path));
    line("// " + comment);
    line("// Autogenerated by codelets.cc");
    line("#pragma once");
    line();
    for (const auto& i : includes)
      line("#include \"%s\"", i);
    line("namespace mandelbrot {");
    line();
  }
  ~Header() {
    line("}  // namespace mandelbrot");
    out.reset();
  }
};

// A basic block
struct Block : public Noncopyable {
private:
  unordered_set<string> names;  // Names used so far
  unordered_map<Sig,int,SigHash> counts;  // Number of times an expression is used (doesn't count above 2)
  unordered_map<Sig,Exp,SigHash> lets;  // Expressions we've already assigned as variables
  Stats stats;

  string fresh(const string& prefix) {
    for (int n = -1;; n++) {
      const auto s = n < 0 ? prefix : tfm::format("%s%d", prefix, n);
      if (names.insert(s).second)
        return s;
    }
  }
public:
  ~Block() { line("// %s", stats.show()); }

  template<class C> void inputs(const C& xs) { for (const auto& x : xs) input(x); }

  template<class E=Exp> E input(const string& name) {
    slow_assert(names.insert(name).second, name);
    const auto x = Exp(name, random_sig());
    if constexpr (is_same_v<E,CExp>) return split(x);
    else { static_assert(is_same_v<E,Exp>); return x; }
  }

  Exps inputs(const string& prefix, const int n) {
    Exps xs;
    for (int i = 0; i < n; i++)
      xs.push_back(input(tfm::format("%s%d", prefix, i)));
    return xs;
  }

  // Compute an expression now, without analyzing it
  Exp now(const string& prefix, const string& exp, const Sig sig) {
    const auto name = fresh(prefix);
    line("const auto %s = %s;", name, exp);
    return Exp(name, sig);
  }
  Exp now(const string& prefix, const Exp& exp) { return now(prefix, str(exp), exp.sig()); }

  // Register that e is used, increasing the counts of newly used expressions
  void use(const Exp& x) {
    int& n = counts[x.sig()];
    if (!n)
      for (const auto& y : x.args())
        use(y);
    n++;
  }
  void use(const Complex<Exp>& z) { use(z.r); use(z.i); }
  void use(const Exps& xs) { for (const auto& x : xs) use(x); }

  // Compute an expression, caching intermediate expressions as variables
  Exp compute(const Exp& x) {
    if (const auto it = lets.find(x.sig()); it != lets.end()) return it->second;
    const auto it = counts.find(x.sig());
    slow_assert(it != counts.end(), "No count for x = %s", x);
    const int count = it->second;

    // Compute it
    const auto y = x.map_args([this](const Exp& y) { return compute(y); });
    stats.add(y);
    if (x.fast() || count == 1) return y;

    // Expression is used more than once.  Cache it in a variable.
    const auto v = now("t", y);
    lets.insert(make_pair(x.sig(), v));
    return v;
  }
  Exps compute(const Exps& xs) {
    Exps ys;
    for (const auto& x : xs)
      ys.push_back(compute(x));
    return ys;
  }
};

// Expansion arithmetic
// For each i, ulp(x[i]) >= |x[i+1]|

// Turn a + b with no overlap properties into an expansion.
// This is Theorem 7 in Shewchuck.
tuple<Exp,Exp> two_sum(const Exp& a, const Exp& b) {
  if (a.zero()) return {b, 0};
  if (b.zero()) return {a, 0};
  const auto x = a + b;
  const auto v = x - a;
  const auto y = (a - (x - v)) + (b - v);
  return {x, y};
}

// a * b as an expansion, using fused multiply-add
tuple<Exp,Exp> two_prod(const Exp& a, const Exp& b) {
  const auto x = a * b;
  const auto y = fma(a, b, -x);
  return {x, y};
}

// Sum two expansions, producing an expansion of the same size.
// This is Figure 2 of Collange et al.
Exps collange_add(const Exps& x, const Exps& y) {
  const int n = int(x.size());
  slow_assert(n == int(y.size()));
  deque<Exp> rest;
  for (int i = 0; i < n; i++) {
    rest.push_back(x[i]);
    rest.push_back(y[i]);
  }
  vector<Exp> state;
  while (rest.size()) {
    auto t = pop_front(rest);
    for (int i = 0; i < int(state.size()); i++)
      tie(state[i], t) = two_sum(t, state[i]);
    state.push_back(t);
  }
  slow_assert(int(state.size()) == 2*n);
  return {state.begin(), state.begin() + n};
}

// Sum via negation + symbolic add
Exps collange_sub(const Exps& x, const Exps& y) {
  return collange_add(x, -y);
}

// Multiply two expansions, producing an expansion of the same size
// This is Algorithm 2 of Collange et al.
Exps collange_mul(const Exps& x, const Exps& y) {
  const int n = int(x.size());
  slow_assert(n == int(y.size()));
  vector<Exp> pi;
  deque<Exp> s(n);
  for (int i = 0; i < n; i++) {
    deque<Exp> e(n), ep(n);
    for (int j = 0; j < n; j++) {
      const auto [p, ej] = two_prod(x[j], y[i]);
      e[j] = ej;
      tie(s[j], ep[j]) = two_sum(s[j], p);
    }
    pi.push_back(pop_front(s));
    s.emplace_back();
    while (any(e)) {
      for (int j = 0; j < n; j++)
        tie(s[j], e[j]) = two_sum(s[j], e[j]);
      e.pop_back();
      e.emplace_front();
    }
    while (any(ep)) {
      for (int j = 0; j < n; j++)
        tie(s[j], ep[j]) = two_sum(s[j], ep[j]);
      ep.pop_back();
      ep.emplace_front();
    }
  }
  slow_assert(n == int(pi.size()));
  return pi;
}

void expansion_arithmetic(const string& path) {
  Header h(path, "Expansion arithmetic codelets", {});
  for (const int n : {2, 3, 4}) {
    // Negation
    {
      Blank b;
      Scope fun("}", "__host__ __device__ static inline Expansion<%d> operator-(const Expansion<%d> x) {", n, n);
      vector<string> nx;
      for (int i = 0; i < n; i++)
        nx.push_back(tfm::format("-x.x[%d]", i));
      line("return Expansion<%d>(%s, nonoverlap);", n, join(nx));
    }

    // Add, subtract, multiply
    const auto binary = [n](const string& op, const auto& body) {
      Blank b;
      Scope fun("}", "__host__ __device__ static inline Expansion<%d>\n"
                "operator%s(const Expansion<%d> x, const Expansion<%d> y) {",
                n, op, n, n);
      line("#ifdef __clang__");
      line("#pragma clang fp reassociate(off)");
      line("#endif  // __clang__");
      CSE cse(false);  // No Schwartz-Zippel for expansion arithmetic!
      Block B;
      B.input("x");
      B.input("y");
      const auto x = B.inputs("x", n), y = B.inputs("y", n);
      line("const auto [%s] = x.x;", join(x));
      line("const auto [%s] = y.x;", join(y));
      const auto s = body(x, y);
      B.use(s);
      line("return Expansion<%d>(%s, nonoverlap);", n, join(B.compute(s)));
    };
    binary("+", [](const auto& x, const auto& y) { return collange_add(x, y); });
    binary("-", [](const auto& x, const auto& y) { return collange_sub(x, y); });
    binary("*", [](const auto& x, const auto& y) { return collange_mul(x, y); });
  }
}

// Base cases for series multiplication and squaring
void mul_bases(const string& path) {
  Header h(path, "Multiplication and squaring base cases", {"loops.h"});
  const int n = 8;
  line("static constexpr int mul_base_n = %d;", n);
  line("static constexpr int sqr_base_n = %d;\n", n);

  // Zero extension
  const auto extend = [](Block& B, const string& x, const int i) {
    const auto v = tfm::format("%s%d", x, i);
    line("const auto %s = %d < n%s ? %s[%d] : S(0);", v, i, x, x, i);
    return B.input(v);
  };

  // Early exit
  const auto exit = [](const string& n, const int i) { line("if (%s <= %d) return;", n, i); };

  // Multiplication
  {
    Blank b;
    Scope f(")", "DEF_SERIAL(mul_base, (S* z, const int nz, const S* x, const int nx, const S* y, const int ny),");
    CSE cse(true);  // Assume exact arithmetic for CSE
    Block B;
    exit("nz", 0);
    // Cache inputs so that aliasing works
    Exps xi, yi;
    for (int i = 0; i < n; i++) xi.push_back(extend(B, "x", i));
    for (int i = 0; i < n; i++) yi.push_back(extend(B, "y", i));
    line();
    for (int i = 0; i < n; i++) {
      if (i) exit("nz", i);
      vector<Exp> ts;
      for (int j = 0; j <= i; j++)
        ts.push_back(xi[j] * yi[i - j]);
      const auto t = sum(ts);
      B.use(t);
      line("z[%d] = %s;", i, B.compute(t));
    }
  };

  // Squaring
  {
    Blank b;
    Scope f(")", "DEF_SERIAL(sqr_base, (S* y, const int ny, const S* x, const int nx),");
    CSE cse(true);  // Assume exact arithmetic for CSE
    Block B;
    exit("ny", 0);
    Exps xi;
    for (int i = 0; i < n; i++) xi.push_back(extend(B, "x", i)); // Cache inputs so that aliasing works
    line();
    for (int i = 0; i < n; i++) {
      if (i) exit("ny", i);
      vector<Exp> ts;
      for (int j = 0; j < i-j; j++)
        ts.push_back(xi[j] * xi[i - j]);
      auto t = 2 * sum(ts);
      if (i % 2 == 0)
        t = t + xi[i/2] * xi[i/2];
      B.use(t);
      line("y[%d] = %s;", i, B.compute(t));
    }
  };
}

// Override series arithmetic
void add_scalar(Series<Exp>& x, const Exp a) {
  slow_assert(x.nonzero());
  x[0] = x[0] + a;
}
Exp ldexp(const Exp x, const int b) {
  slow_assert(!b);
  return x;
}
void high_addsub_ldexp(Series<Exp>& y, const int sign, const int b, const int64_t s, SeriesView<const Exp> x) {
  const auto ynz = y.nonzero(), xnz = x.nonzero();
  const auto nk = min(y.known(), x.known() + s);
  const auto nz = min(nk, max(ynz, xnz ? xnz + s : 0));
  slow_assert(abs(sign) == 1 && nz <= y.limit());
  const auto x_ = x.copy(xnz);  // Watch for aliasing
  y.set_counts(nk, nz);
  for (int i = 0; i < nz; i++) {
    const auto yi = i < ynz ? y[i] : Exp(0);
    auto xi = ldexp(uint32_t(i-s) < uint32_t(xnz) ? x_[i-s] : Exp(0), b);
    if (sign < 0) xi = -xi;
    y[i] = yi + xi;
  }
}
void fft_mul(span<Exp> z, span<const Exp> x, span<const Exp> y) {
  // Actually, we use naive multiplication
  const int nz = z.size(), nx = x.size(), ny = y.size();
  for (int i = 0; i < nz; i++) {
    vector<Exp> ts;
    for (int j = 0; j <= i; j++)
      if (j < nx && i - j < ny)
        ts.push_back(x[j] * y[i - j]);
    z[i] = sum(ts);
  }
}

// Base cases for series functions
void series_bases(const string& path) {
  Header h(path, "Series function base cases", {"loops.h"});
  const int n = 5;

  // Unary
  const auto unary = [n](const string& name, const int leading, auto&& set) {
    {
      Blank b;
      Scope f(")", "DEF_SERIAL(%s_base_serial, (S* y, const int ny, const S* x, const int nx),", name);
      CSE cse(true);  // Assume exact arithmetic for CSE
      Block B;
      B.inputs(vector<string>{"y", "ny", "x", "nx"});
      for (int i = 0; i < leading; i++) B.input(tfm::format("x%d", i));  // Make names nicer
      Series<Exp> x(n), y(n);
      x.set_counts(n, n);
      for (int i = 0; i < n; i++) {
        if (i < leading)
          x[i] = 0;
        else {
          x[i] = B.input(tfm::format("x%d", i));
          line("const auto %s = %d < nx ? x[%d] : S(0);", x[i], i, i);
        }
      }
      set(y, x.view());
      for (int i = 0; i < n; i++)
        B.use(y[i]);
      for (int i = 0; i < n; i++) {
        const auto yi = B.compute(y[i]);
        line("y[%d] = %s;", i, B.compute(y[i]));
        if (i+1 < n) {
          line();
          line("if (ny == %d) return;", i+1);
        }
      }
    } {
      Blank b;
      Scope f("}", "template<class T> void %s_base(Series<T>& y, type_identity_t<SeriesView<const T>> x) {", name);
      line("const int ny = min(%d, int(x.known()));", n);
      line("y.set_counts(ny, ny);");
      line("if (ny) %s_base_serial(y.data(), ny, x.data(), x.nonzero());", name);
    }
  };
  unary("inv", 0, [=](auto& y, const auto x) {
    Series<Exp> r(n), rx(n), inv_rx(n);
    r.set_scalar(n, inv(x[0]));
    rx = mul(r, x);
    inv_rx = inv(rx);
    y = mul(r, inv_rx);
  });
  unary("exp", 1, [=](auto& y, const auto x) { y = exp(x); });
}

// Read/write locations
template<class T> struct Loc { typedef T E; function<T()> get; function<void(T)> set; };
template<class T> using Locs = vector<Loc<T>>;

template<class E> Loc<E> loc(Block& B, const string& name, const Exp& j, const bool add = false);

template<> Loc<Exp> loc(Block& B, const string& name, const Exp& j, const bool add) {
  return Loc<Exp>{
    [&B,name,j]() { line("const auto %s = %s < xn ? x[%s] : 0;", name, j, j); return B.input(name); },
    [&B,j,add](Exp x) {
      string op = add ? "+=" : "=";
      if (add) {
        if (auto nx = x.unneg()) {
          op = "-=";
          x = *nx;
        }
      }
      line("if (%s < xn) x[%s] %s %s;", j, j, op, B.compute(x)); },
  };
};

template<> Loc<CExp> loc(Block& B, const string& name, const Exp& j, const bool add) {
  slow_assert(!add, "We only use add for output into real arrays");
  return Loc<CExp>{
    [&B,name,j] { line("const auto %s = y[%s];", name, j); return split(B.input(name)); },
    [&B,j](const CExp& z) { line("y[%s] = Complex<S>(%s, %s);", j, B.compute(z.r), B.compute(z.i)); },
  };
}

template<class E> Locs<E> locs(Block& B, const int n, const Exp& j, const Exp& dj, const bool add = false) {
  Locs<E> zs;
  for (int i = 0; i < n; i++)
    zs.push_back(loc<E>(B, tfm::format("%c%d", is_same_v<E,Exp> ? 'x' : 'y', i), j + i*dj, add));
  return zs;
}

// Symbolic gradient of the circuit from xs to ys, given differentials on the outputs
Exps gradient(span<const Exp> xs, span<const Exp> ys, span<const Exp> dys) {
  slow_assert(ys.size() == dys.size());
  struct State {
    int need;  // Number of output contributions we need
    Exps terms;  // Contributions to the gradient
  };
  unordered_map<Sig,State,SigHash> state;

  // Sweep through the graph from ys to xs, computing need counts
  {
    Exps work(ys.begin(), ys.end());
    while (work.size()) {
      const auto z = pop_back(work);
      for (const auto& x : z.args()) {
        auto& n = state[x.sig()].need;
        if (!n) work.push_back(x);
        n++;
      }
    }
  }

  // Compute gradients in topological sort order
  for (int i = 0; i < int(ys.size()); i++)
    state[ys[i].sig()].terms.push_back(dys[i]);
  Exps work(ys.begin(), ys.end());
  while (work.size()) {
    const auto z = pop_back(work);
    auto& sz = state[z.sig()];
    const auto dz = sum(sz.terms);
    sz.terms.resize(1);
    sz.terms[0] = dz;
    const auto args = z.args();
    const auto grad = z.grad();
    slow_assert(args.size() == grad.size(), "%s: %d %d", z, args.size(), grad.size());
    for (int i = 0; i < int(args.size()); i++) {
      auto& s = state[args[i].sig()];
      s.terms.push_back(grad[i] * dz);
      s.need--;
      if (!s.need) work.push_back(args[i]);
    }
  }

  // Read off input gradients
  Exps grads;
  for (const auto& x : xs) {
    auto& s = state[x.sig()];
    slow_assert(!s.need && s.terms.size() == 1);
    grads.push_back(s.terms[0]);
  }
  return grads;
}

// Convert between lists of Exp and CExp
span<const Exp> flatten(span<const Exp> xs) { return xs; }
span<const Exp> flatten(span<const CExp> zs) {
  return span<const Exp>(reinterpret_cast<const Exp*>(zs.data()), 2*zs.size());
}
template<class Ts> Ts thicken(const Exps& xs);
template<> Exps thicken(const Exps& xs) { return xs; }
template<> CExps thicken(const Exps& xs) {
  const int n = exact_div(int(xs.size()), 2);
  const CExp* p = reinterpret_cast<const CExp*>(xs.data());
  return CExps(p, p + n);
}

// Gradients with complex inputs or outputs
template<class Xs,class Ys> Xs gradient(const Xs xs, const Ys ys, const Ys dys) {
  return thicken<Xs>(gradient(flatten(xs), flatten(ys), flatten(dys)));
}

Exp constant(const Expansion<2> x, const Sig sig) {
  const auto s = tfm::format("%.17g", x.span());
  return other(tfm::format("constant<S>(%s)", s.substr(1, s.size()-2)), 2, sig);
}

Exp constant_sqrt_half() {
  typedef Expansion<2> S;
  return constant(nearest_sqrt<S>(1, 2), inv(sqrt(Sig(2))));
}

CExp constant_twiddle(const int j, const int s) {
  typedef Expansion<2> S;
  slow_assert(s == 1 || s == 2 || s == 3);
  const auto sh = inv(sqrt(Sig(2)));
  const auto base = s == 1 ? Complex<Sig>(-1, 0)
                  : s == 2 ? Complex<Sig>(0, 1)
                           : Complex<Sig>(sh, sh);
  const auto sig = pow(base, j);
  const auto t = nearest_twiddle<S>(j, 2<<s);
  return CExp(constant(t.r, sig.r), constant(t.i, sig.i));
}

typedef function<Complex<Exp>(const Exp&)> Twiddle;

Complex<Exp> base_twiddle(const Exp& j) {
  return split(call("twiddle", j, arbitrary("twiddle", j.sig())));
}

// Forward compute from inputs to outputs
template<class I,class O> void forward(Block& B, const vector<I>& inputs, const vector<O>& outputs, const auto& compute) {
  vector<typename I::E> xs;
  for (const auto& i : inputs) xs.push_back(i.get());
  const auto ys = compute(xs);
  slow_assert(ys.size() == outputs.size());
  for (const auto& y : ys) B.use(y);
  for (int i = 0; i < int(outputs.size()); i++) outputs[i].set(ys[i]);
}

// Either forward, or backward via symbolic transpose
template<class I, class O> void
way(Block& B, const bool fwd, const vector<I>& inputs, const vector<O>& outputs, const auto& compute) {
  if (fwd) return forward(B, inputs, outputs, compute);
  // Take the symbolic transpose
  vector<typename I::E> pxs;
  for (int i = 0; i < int(inputs.size()); i++) pxs.push_back(B.input<typename I::E>(tfm::format("px%d", i)));
  const auto pys = compute(pxs);
  forward(B, outputs, inputs, [pxs, pys](const auto& ys) { return gradient(pxs, pys, ys); });
}

enum Mode { NoX, GetX, SetX };

void loop(const string& name, const bool serial, const Mode mode, const string& n, const vector<string>& args,
          const auto& compute) {
  // Collect all arguments
  vector<string> full = {"Complex<S>* y"};
  if (mode != NoX) {
    full.push_back(mode == GetX ? "const S* x" : "S* x");
    full.push_back("const int xn");
  }
  full.insert(full.end(), args.begin(), args.end());

  // Write code
  Blank b;
  CSE cse(true);  // Assume exact arithmetic for CSE
  const auto inputs = [](Block& B) {
    B.input("x");
    B.input("xn");
    B.input("y");
  };
  if (serial) {
    Scope f(")", "DEF_SERIAL(%s, (%s),", name, join(full));
    Block B;
    inputs(B);
    compute(B, 1, 0);
  } else {
    Scope f(")", "DEF_LOOP(%s, %s, j, (%s),", name, n, join(full));
    Block B;
    inputs(B);
    compute(B, B.input(n), B.input("j"));
  }
}

// Forward and backward FFTs
void both(const bool allow_add, const function<void(bool,bool,string,string)>& one) {
  for (const bool fwd : {true, false}) {
    for (const auto add : {false, true}) {
      if (add && (!allow_add || fwd)) continue;
      const string i = fwd ? "" : "i";
      const string arrow = fwd ? "->" : "<-";
      one(fwd, add, i, arrow);
    }
  }
}

// Decimination-in-frequency shifted real-to-complex FFT, using the commutator notation:
//   t j2 j1 j0 -> t k2/2 j1 j0
//              -> k2 t k1/2 j0
//              -> k2 k1 t k0/2
//              == k2 k1 k0

// First srfft butterfly, copying real to complex, without twiddle factors:
//   t j2 j1 j0 -> t k2/2 j1 j0
CExp butterfly_0(const Exp x0, const Exp x1) {
  return CExp(x0, -x1);
}

// Second ssrfft butterfly, in place, shifted twiddling on input:
//   t k(p-1)/2 j(p-2) ...j... -> k(p-1) t k(p-2)/2 ...j...
tuple<CExp,CExp> butterfly_1(const Exp j, const Twiddle& twiddle, const CExp y0, const CExp y1) {
  const auto z1 = diag<-1>(constant_sqrt_half(), y1);
  const auto u0 = conj(twiddle(j)) * (y0 + z1);
  const auto u1 = conj(twiddle(3*j) * (y0 - z1));
  return make_tuple(u0, u1);
}

// Radix-2 shifted butterfly
tuple<CExp,CExp> butterfly_s1(const Exp j, const Twiddle& twiddle, const CExp y0, const CExp y1) {
  const auto u0 = y0 + y1;
  const auto u1 = conj(twiddle(j) * (y0 - y1));
  return make_tuple(u0, u1);
}

// Radix-4 shifted butterfly
CExps butterfly_s2(const Exp j, const Twiddle& twiddle, const CExps y) {
  slow_assert(y.size() == 4);
  const auto w1 = twiddle(j);
  const auto w2 = twiddle(2*j);

  const auto u0 = y[0] + y[2];
  const auto u1 = y[1] + y[3];
  const auto t0 = y[0] - y[2];
  const auto t1 = y[1] - y[3];

  const auto c0 = u0 + u1;
  const auto c1 = conj(w2 * (u0 - u1));
  const auto c2 = conj(w1 * (t0 + left(t1)));
  const auto c3 = conj(w1) * (t0 - left(t1));
  return CExps{c0, c1, c2, c3};
}

// Radix-8 shifted butterfly
CExps butterfly_s3(const Exp j, const Twiddle& twiddle, const CExps y) {
  slow_assert(y.size() == 8);
  const auto w1 = twiddle(j);
  const auto w2 = twiddle(2*j);
  const auto w3 = twiddle(3*j);
  const auto w4 = twiddle(4*j);
  const auto sqrt_half = constant_sqrt_half();

  const auto u0 = y[0] + y[4];
  const auto u1 = y[1] + y[5];
  const auto u2 = y[2] + y[6];
  const auto u3 = y[3] + y[7];
  const auto t0 = y[0] - y[4];
  const auto t1 = y[1] - y[5];
  const auto t2 = y[2] - y[6];
  const auto t3 = y[3] - y[7];

  const auto b0 = u0 + u2;
  const auto b1 = u1 + u3;
  const auto b2 = t0 + left(t2);
  const auto b3 = diag<1>(sqrt_half, t1 + left(t3));
  const auto c0 = u0 - u2;
  const auto c1 = u1 - u3;
  const auto c2 = t0 - left(t2);
  const auto c3 = diag<-1>(sqrt_half, t1 - left(t3));

  const auto d0 = b0 + b1;
  const auto d1 = conj(w4 * (b0 - b1));
  const auto d2 = conj(w2 * (c0 + left(c1)));
  const auto d3 = conj(w2) * (c0 - left(c1));
  const auto d4 = conj(w1 * (b2 + b3));
  const auto d5 = conj(w3) * (b2 - b3);
  const auto d6 = conj(w1) * (c2 + c3);
  const auto d7 = conj(w3 * (c2 - c3));
  return CExps{d0, d1, d2, d3, d4, d5, d6, d7};
}

// Higher radix fft codelets
void butterflies(const string& path) {
  Header h(path, "Fourier transform base cases", {"loops.h"});

  line("namespace {");
  line("template<class S> struct TwiddleSlice;");
  line("template<class S> struct BigTwiddleSlice;");
  line("}  // namespace");
  line();

  // The first r butterfly levels in a single loop
  const vector<string> twiddle_args = {"BigTwiddleSlice<S> twiddle", "const int p"};
  for (int r = 1; r <= 3; r++) {
    string suffix;
    for (int i = 0; i < r; i++) suffix += tfm::format("%d", i);
    const auto nr = tfm::format("n%d", 1 << r);
    both(true, [r,nr,suffix,twiddle_args](const bool fwd, const bool add, const string& i, const string& arrow) {
      const auto name = tfm::format("%s%ssrfft_butterfly_%s", add ? "add_" : "", i, suffix);
      loop(name, r < 3, fwd ? GetX : SetX, nr, r == 3 ? twiddle_args : vector<string>(),
        [fwd,r,add](Block& B, const Exp nr, const Exp j) {
          const auto p = B.input("p");
          const auto m = r > 2 ? B.now("m", tfm::format("1 << (%s)", p-r), random_sig()) : 0;
          const int R = 1 << r;
          const auto x = locs<Exp>(B, R, j, nr, add);
          const auto y = locs<CExp>(B, R/2, j, nr);
          const auto tw = [r](const Exp& j) {
            if (r == 3) return base_twiddle(j);
            const auto jj = j.unint();
            slow_assert(jj);
            return constant_twiddle(*jj, r);
          };
          way(B, fwd, x, y, [&B,m,r,nr,R,p,j,tw](const Exps x) {
            CExps y;
            for (int i = 0; i < R/2; i++)
              y.push_back(butterfly_0(x[i], x[i+R/2]));
            if (r > 1) {
              for (int i = 0; i < R/4; i++)
                tie(y[i], y[i+R/4]) = butterfly_1(j+i*nr, tw, y[i], y[i+R/4]);
            }
            for (int a = 3; a <= r; a++) {
              const auto ma = B.now(tfm::format("m%d", a), tfm::format("m << %d", r-a), m.sig() << (r-a));
              const auto j1 = B.now("j1", tfm::format("%s & (%s-1)", j, ma), random_sig());
              const int R0 = 1 << (a-2);
              const int R1 = R >> a;
              const auto twa = [tw,a](const Exp& j) { return tw((1 << a) * j); };
              for (int i0 = 0; i0 < R0; i0++) {
                for (int i1 = 0; i1 < R1; i1++) {
                  auto& y0 = y[2*i0 + i1];
                  auto& y1 = y[2*i0 + i1 + R1];
                  tie(y0, y1) = butterfly_s1(j1+i1*nr, twa, y0, y1);
                }
              }
            }
            return y;
          });
        }
      );
    });
  }

  // Remaining butterflies, using varying radix
  for (int r = 1; r <= 3; r++) {
    const auto nr = tfm::format("n%d", 1 << (r+1));
    both(false, [r,nr](const bool fwd, const bool add, const string& i, const string& arrow) {
      slow_assert(!add);
      line("// Radix-%d %ssrfft butterfly, in place", r, i);
      loop(tfm::format("%ssrfft_butterfly_s%d", i, r), false, NoX, nr, {"TwiddleSlice<S> twiddle", "const int s"},
        [fwd,r](Block& B, const Exp nr, const Exp j) {
          const auto s = B.input("s");
          const auto m = B.now("m", "1 << s", random_sig());
          const auto j1 = B.now("j1", tfm::format("%s & (m-1)", j), random_sig());
          const auto j0 = B.now("j0", tfm::format("(%s - j1) << %d", j, r), random_sig());
          const auto y = locs<CExp>(B, 1 << r, j0 + j1, m);
          way(B, fwd, y, y, [r,s,m,j1](CExps y) {
            const auto tw = [](const Exp& j) { return base_twiddle(j); };
            if (r == 1)
              tie(y[0], y[1]) = butterfly_s1(j1, tw, y[0], y[1]);
            else if (r == 2)
              y = butterfly_s2(j1, tw, y);
            else if (r == 3)
              y = butterfly_s3(j1, tw, y);
            return y;
          });
        }
      );
    });
  }
}

}  // namespace mandelbrot
using namespace mandelbrot;

int main(int argc, char** argv) {
  try {
    const vector<string> paths(argv + 1, argv + argc);
    for (const auto& path : paths) {
      if (path.ends_with("gen-expansion.h")) expansion_arithmetic(path);
      else if (path.ends_with("gen-mul-bases.h")) mul_bases(path);
      else if (path.ends_with("gen-series-bases.h")) series_bases(path);
      else if (path.ends_with("gen-butterflies.h")) butterflies(path);
      else die("Unmatched path '%s'", path);
    }
  } catch (const std::exception& e) {
    die("%s (%s)", e.what(), typeid(e).name());
  }
}