(*
 * Copyright (c) 1997-1999 Massachusetts Institute of Technology
 * Copyright (c) 2000-2002 Stefan Kral
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 *)

(* This file contains the entry point for the genfft program: it
   parses the command-line parameters and calls the rest of the
   program as needed to generate the requested code. *)

open List
open Util
open Expr
open Variable
open Number
open Fft
open Symmetry
open VScheduler
open VAnnotatedScheduler
open VFpBasics
open Printf
open VSimdBasics
open VSimdUnparsing

(****************************************************************************)

let fftwcomplex_data_type = "fftw_complex"
let simd2_data_type = "fftw_simd2"
let basic_integer_type = "int"

let iarray = "input"
let oarray = "output"
let istride = "istride"
let ostride = "ostride"
let twiddle_order = "twiddle_order"

let ioarray = "inout"
let iostride = "iostride"
let twarray = "W"

(****************************************************************************)

type codelet_type = TWIDDLE | NO_TWIDDLE | REAL2HC | HC2HC | HC2REAL
                    | REALEVEN | REALODD | REALEVEN2 | REALODD2
                    | REALEVEN_TWIDDLE | REALODD_TWIDDLE

let rec list_to_c = function
    [] -> ""
  | [a] -> (string_of_int a)
  | a :: b -> (string_of_int a) ^ ", " ^ (list_to_c b)

let codelet_description n dir ty name =
  let (ctype, itype) = match ty with
    | TWIDDLE    -> "FFTW_TWIDDLE", 0
    | NO_TWIDDLE -> "FFTW_NOTW", 1
    | _ -> failwith "codelet_description: unsupported"
  and (cdir, idir) = match dir with
    | FORWARD -> "FFTW_FORWARD", 0
    | BACKWARD -> "FFTW_BACKWARD", 1
  and (_, num_twiddle, tw_o) = Twiddle.twiddle_policy ()
  in let (declare_order, order, nt) =
    match ty with
      TWIDDLE -> ("static const int " ^ twiddle_order ^ "[] = {" ^
		  (list_to_c (tw_o n)) ^ "};\n\n"),
	twiddle_order,
	num_twiddle n
    | NO_TWIDDLE -> "", "(const int *) 0", 0
    | _ -> failwith "codelet_description: unsupported"

  (* this should be replaced by CRC/MD5 of the codelet *)
  and signature = 11 * (2 * n + idir) + itype 

  in "\n\n" ^ declare_order ^
  "fftw_codelet_desc " ^ name ^ "_desc = {\n" ^
  "\"" ^ name ^ "\",\n" ^
  "(void (*)()) " ^ name ^ ",\n" ^
  (string_of_int n) ^ ",\n" ^
  cdir ^ ",\n" ^
  ctype ^ ",\n" ^
  (string_of_int signature) ^ ",\n" ^
  (string_of_int nt) ^ ",\n" ^
  order ^ ",\n" ^
  "};\n"

(****************************************************************************)

let optimize varinfo n expr =
  let _ = info "simplifying..." in
  let code0 = Exprdag.simplify_to_alist expr in
  let _ = info "mapping assignments to vfpinstrs..." in
  let code1 = AssignmentsToVfpinstrs.assignmentsToVfpinstrs varinfo code0 in
  let _ = info "vectorizing..." in
  let code2 = VTwoWayVectorization.vectorize varinfo n code1 in
  let _ = info "optimizing..." in
  let code3 = VOptimization.apply_optrules code2 in
  let _ = info "vectorized scheduling..." in
  let code4 = schedule code3 in
  let _ = info "vectorized annotating..." in
  let code5 = annotate code4 in
  let _ = info "vectorized linearizing..." in
  let code6 = annotatedscheduleToVsimdinstrs code5 in
  let _ = info "improving..." in
     VImproveSchedule.improve_schedule code6

(****************************************************************************)

let emit_constant_declaration (n,m) =
  printf "DECLARE_SIMD_CONST2(%s%s,%s,%s);\n"
	 (unparse n) (unparse m) (to_string n) (to_string m)
and emit_tmpvar_declaration i =
  printf "DECLARE_SIMD_TMPVAR2(%s);\n" (vsimdregToString i)
and emit_instruction var_unparser i =
  printf "%s;\n" (vsimdinstrToString' var_unparser i)

(***)

let variablestride_unparser stride = function
  | 0 -> "0"
  | 1 -> stride
  | n -> stride ^ "*" ^ intToString n

let no_twiddle_unparser arr i = match arr with
  | Output -> oarray ^ "[" ^ variablestride_unparser ostride i ^ "]"
  | Input  -> iarray ^ "[" ^ variablestride_unparser istride i ^ "]"
  | _ -> failwith "no_twiddle_unparser: unsupported"

let varinfoNOTW =
	[(Input,   ComplexArray(InterleavedFormat)); 
	 (Output,  ComplexArray(InterleavedFormat))]

let no_twiddle_genemit n dir =
  let _ = info "generating..." in
  let fn_name = match dir with
    | FORWARD  -> "fftw_no_twiddle_" ^ intToString n
    | BACKWARD -> "fftwi_no_twiddle_" ^ intToString n in
  let code = optimize varinfoNOTW n (no_twiddle_gen_expr n no_sym dir) in
  begin
    iter emit_constant_declaration (rev (vsimdinstrsToConsts code));
    print_newline ();
    printf "void %s(const %s *inC, %s *outC, %s %s, %s %s)\n" 
	   fn_name 
	   fftwcomplex_data_type 
	   fftwcomplex_data_type
	   basic_integer_type istride
	   basic_integer_type ostride;
    print_string "{\n";
      printf "const %s *%s = (const %s *) inC;\n" 
	     simd2_data_type iarray simd2_data_type;
      printf "%s *%s = (%s *) outC;\n" simd2_data_type oarray simd2_data_type;
      iter emit_tmpvar_declaration (rev (vsimdinstrsToDstregs code));
      print_newline ();
      iter (emit_instruction no_twiddle_unparser) code;
    print_string "}\n";
    print_string (codelet_description n dir NO_TWIDDLE fn_name)
  end

(***)

let twiddle_unparser arr i = match arr with
  | Input | Output -> ioarray ^ "[" ^ variablestride_unparser iostride i ^ "]"
  | Twiddle        -> twarray ^ "[" ^ (intToString i) ^ "]"

let varinfoTW = 
	[(Input,   ComplexArray(InterleavedFormat)); 
	 (Twiddle, ComplexArray(InterleavedFormat)); 
	 (Output,  ComplexArray(InterleavedFormat))]

let twiddle_gen n dir =
  let _ = info "generating..." in
  let fn_name = match dir with
    | FORWARD  -> "fftw_twiddle_" ^ intToString n
    | BACKWARD -> "fftwi_twiddle_" ^ intToString n in
  let code = optimize varinfoTW n (twiddle_dit_gen_expr n no_sym no_sym dir) in
  let (_, num_twiddle, _) = Twiddle.twiddle_policy () in
  begin
    iter emit_constant_declaration (rev (vsimdinstrsToConsts code));
    print_newline ();
    printf "void %s(%s *ioC, const %s *twC, %s %s, %s m, %s dist)\n{\n" 
	   fn_name
	   fftwcomplex_data_type
	   fftwcomplex_data_type
	   basic_integer_type iostride
	   basic_integer_type 
	   basic_integer_type;
      printf "%s *%s = (%s *) ioC;\n" simd2_data_type ioarray simd2_data_type;
      printf "const %s *%s = (const %s *) twC;\n" 
	     simd2_data_type twarray simd2_data_type;
      print_newline ();
      printf "for ( ; m>0; m--, %s += dist, %s += %d)\n" 
	     ioarray twarray (num_twiddle n);
      print_string "{\n";
        iter emit_tmpvar_declaration (rev (vsimdinstrsToDstregs code));
         print_newline ();
        iter (emit_instruction twiddle_unparser) code;
      print_string "}\n";
    print_string "}\n";
    print_string (codelet_description n dir TWIDDLE fn_name)
  end


type mode = 
  | GEN_NOTHING
  | GEN_NOTWID of int
  | GEN_NOTWIDI of int
  | GEN_TWIDDLE of int
  | GEN_TWIDDLEI of int
  | GEN_REAL2HC of int
  | GEN_HC2REAL of int
  | GEN_HC2HC_FORWARD of int
  | GEN_HC2HC_BACKWARD of int
  | GEN_REAL_EVEN of int
  | GEN_REAL_ODD of int
  | GEN_REAL_EVEN2 of int
  | GEN_REAL_ODD2 of int
  | GEN_REAL_EVEN_TWIDDLE of int
  | GEN_REAL_ODD_TWIDDLE of int


let usage = "Usage: genfft [ -notwiddle | -notwiddleinv | -twiddle | -twiddleinv ] <number>"

let mode = ref GEN_NOTHING

let undocumented = " Undocumented voodoo parameter"

let main () =
  Arg.parse [
  "-notwiddle", 
    Arg.Int(fun i -> mode := GEN_NOTWID i), 
    "<n> : Generate a no twiddle codelet of size <n>";
  "-notwiddleinv", 
    Arg.Int(fun i -> mode := GEN_NOTWIDI i), 
    "<n> : Generate a no twiddle inverse codelet of size <n>";

  "-twiddle", 
    Arg.Int(fun i -> mode := GEN_TWIDDLE i), 
    "<n> : Generate a twiddle codelet of size <n>";
  "-twiddleinv", 
    Arg.Int(fun i -> mode := GEN_TWIDDLEI i), 
    "<n> : Generate a twiddle inverse codelet of size <n>";

  "-verbose", 
    Arg.Unit(fun () -> Magic.verbose := true),
    " Enable verbose logging messages to stderr";

  "-rader-min", Arg.Int(fun i -> Magic.rader_min := i),
    "<n> : Use Rader's algorithm for prime sizes >= <n>";

  "-magic-alternate-convolution", 
    Arg.Int(fun i -> Magic.alternate_convolution := i),
    undocumented;

  "-magic-window", 
    Arg.Int(fun i -> Magic.window := i), 
    undocumented;
  "-magic-variables", 
    Arg.Int(fun i -> Magic.number_of_variables := i),
    undocumented;

  "-magic-loopo", 
    Arg.Unit(fun () -> Magic.loopo := true),
    undocumented;
  "-magic-loopi", 
    Arg.Unit(fun () -> Magic.loopo := false),
    undocumented;

  "-magic-times-3-3", 
    Arg.Unit(fun () -> Magic.times_3_3 := true),
    undocumented;
  "-magic-times-4-2", 
    Arg.Unit(fun () -> Magic.times_3_3 := false),
    undocumented;

  "-magic-inline-single", 
    Arg.Unit(fun () -> Magic.inline_single := true),
    undocumented;
  "-magic-no-inline-single", 
    Arg.Unit(fun () -> Magic.inline_single := false),
    undocumented;

  "-magic-inline-loads", 
    Arg.Unit(fun () -> Magic.inline_loads := true),
    undocumented;
  "-magic-no-inline-loads", 
    Arg.Unit(fun () -> Magic.inline_loads := false),
    undocumented;

  "-magic-enable-fma", 
    Arg.Unit(fun () -> Magic.enable_fma := true),
    undocumented;
  "-magic-disable-fma", 
    Arg.Unit(fun () -> Magic.enable_fma := false),
    undocumented;

  "-magic-enable-fma-expansion", 
    Arg.Unit(fun () -> Magic.enable_fma_expansion := true),
    undocumented;
  "-magic-disable-fma-expansion", 
    Arg.Unit(fun () -> Magic.enable_fma_expansion := false),
    undocumented;

  "-magic-collect-common-twiddle", 
    Arg.Unit(fun () -> Magic.collect_common_twiddle := true),
    undocumented;
  "-magic-no-collect-common-twiddle", 
    Arg.Unit(fun () -> Magic.collect_common_twiddle := false),
    undocumented;

  "-magic-collect-common-inputs", 
    Arg.Unit(fun () -> Magic.collect_common_inputs := true),
    undocumented;
  "-magic-no-collect-common-inputs", 
    Arg.Unit(fun () -> Magic.collect_common_inputs := false),
    undocumented;

  "-magic-use-wsquare", 
    Arg.Unit(fun () -> Magic.use_wsquare := true),
    undocumented;
  "-magic-no-wsquare", 
    Arg.Unit(fun () -> Magic.use_wsquare := false),
    undocumented;

  "-magic-vectsteps-limit", 
    Arg.Int(fun i -> Magic.vectsteps_limit := i), 
    undocumented;

  "-magic-twiddle-load-all", 
    Arg.Unit(fun () -> Magic.twiddle_policy := Magic.TWIDDLE_LOAD_ALL),
    undocumented;
  "-magic-twiddle-iter", 
    Arg.Unit(fun () -> Magic.twiddle_policy := Magic.TWIDDLE_ITER),
    undocumented;
  "-magic-twiddle-load-odd", 
    Arg.Unit(fun () -> Magic.twiddle_policy := Magic.TWIDDLE_LOAD_ODD),
    undocumented;
  "-magic-twiddle-square1", 
    Arg.Unit(fun () -> Magic.twiddle_policy := Magic.TWIDDLE_SQUARE1),
    undocumented;
  "-magic-twiddle-square2", 
    Arg.Unit(fun () -> Magic.twiddle_policy := Magic.TWIDDLE_SQUARE2),
    undocumented;
  "-magic-twiddle-square3", 
    Arg.Unit(fun () -> Magic.twiddle_policy := Magic.TWIDDLE_SQUARE3),
    undocumented;
] (fun _ -> failwith "too many arguments") usage;
  match !mode with
    | GEN_TWIDDLE i  -> twiddle_gen i Fft.FORWARD
    | GEN_TWIDDLEI i -> twiddle_gen i Fft.BACKWARD
    | GEN_NOTWID i   -> no_twiddle_genemit i Fft.FORWARD
    | GEN_NOTWIDI i  -> no_twiddle_genemit i Fft.BACKWARD
    | _ -> failwith "one of -notwiddle, -notwiddleinv, -twiddle, -twiddleinv must be specified"

let _ = main ()

