cannam@95: (*
cannam@95:  * Copyright (c) 1997-1999 Massachusetts Institute of Technology
cannam@95:  * Copyright (c) 2003, 2007-11 Matteo Frigo
cannam@95:  * Copyright (c) 2003, 2007-11 Massachusetts Institute of Technology
cannam@95:  *
cannam@95:  * This program is free software; you can redistribute it and/or modify
cannam@95:  * it under the terms of the GNU General Public License as published by
cannam@95:  * the Free Software Foundation; either version 2 of the License, or
cannam@95:  * (at your option) any later version.
cannam@95:  *
cannam@95:  * This program is distributed in the hope that it will be useful,
cannam@95:  * but WITHOUT ANY WARRANTY; without even the implied warranty of
cannam@95:  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
cannam@95:  * GNU General Public License for more details.
cannam@95:  *
cannam@95:  * You should have received a copy of the GNU General Public License
cannam@95:  * along with this program; if not, write to the Free Software
cannam@95:  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
cannam@95:  *
cannam@95:  *)
cannam@95: 
cannam@95: (* This file contains the instruction scheduler, which finds an
cannam@95:    efficient ordering for a given list of instructions.
cannam@95: 
cannam@95:    The scheduler analyzes the DAG (directed acyclic graph) formed by
cannam@95:    the instruction dependencies, and recursively partitions it.  The
cannam@95:    resulting schedule data structure expresses a "good" ordering
cannam@95:    and structure for the computation.
cannam@95: 
cannam@95:    The scheduler makes use of utilties in Dag and other packages to
cannam@95:    manipulate the Dag and the instruction list. *)
cannam@95: 
cannam@95: open Dag
cannam@95: (*************************************************
cannam@95:  *               Dag scheduler
cannam@95:  *************************************************)
cannam@95: let to_assignment node = (Expr.Assign (node.assigned, node.expression))
cannam@95: let makedag l = Dag.makedag 
cannam@95:     (List.map (function Expr.Assign (v, x) -> (v, x)) l)
cannam@95: 
cannam@95: let return x = x
cannam@95: let has_color c n = (n.color = c)
cannam@95: let set_color c n = (n.color <- c)
cannam@95: let has_either_color c1 c2 n = (n.color = c1 || n.color = c2)
cannam@95: 
cannam@95: let infinity = 100000 
cannam@95: 
cannam@95: let cc dag inputs =
cannam@95:   begin
cannam@95:     Dag.for_all dag (fun node -> 
cannam@95:       node.label <- infinity);
cannam@95:     
cannam@95:     (match inputs with 
cannam@95:       a :: _ -> bfs dag a 0
cannam@95:     | _ -> failwith "connected");
cannam@95: 
cannam@95:     return
cannam@95:       ((List.map to_assignment (List.filter (fun n -> n.label < infinity)
cannam@95: 				  (Dag.to_list dag))),
cannam@95:        (List.map to_assignment (List.filter (fun n -> n.label == infinity) 
cannam@95: 				  (Dag.to_list dag))))
cannam@95:   end
cannam@95: 
cannam@95: let rec connected_components alist =
cannam@95:   let dag = makedag alist in
cannam@95:   let inputs = 
cannam@95:     List.filter (fun node -> Util.null node.predecessors) 
cannam@95:       (Dag.to_list dag) in
cannam@95:   match cc dag inputs with
cannam@95:     (a, []) -> [a]
cannam@95:   | (a, b) -> a :: connected_components b
cannam@95: 
cannam@95: let single_load node =
cannam@95:   match (node.input_variables, node.predecessors) with
cannam@95:     ([x], []) -> 
cannam@95:       Variable.is_constant x ||
cannam@95:       (!Magic.locations_are_special && Variable.is_locative x)
cannam@95:   | _ -> false
cannam@95: 
cannam@95: let loads_locative node =
cannam@95:   match (node.input_variables, node.predecessors) with
cannam@95:     | ([x], []) -> Variable.is_locative x
cannam@95:     | _ -> false
cannam@95: 
cannam@95: let partition alist =
cannam@95:   let dag = makedag alist in
cannam@95:   let dag' = Dag.to_list dag in
cannam@95:   let inputs = 
cannam@95:     List.filter (fun node -> Util.null node.predecessors) dag'
cannam@95:   and outputs = 
cannam@95:     List.filter (fun node -> Util.null node.successors) dag'
cannam@95:   and special_inputs =  List.filter single_load dag' in
cannam@95:   begin
cannam@95:     
cannam@95:     let c = match !Magic.schedule_type with
cannam@95: 	| 1 -> RED; (* all nodes in the input partition *)
cannam@95: 	| -1 -> BLUE; (* all nodes in the output partition *)
cannam@95: 	| _ -> BLACK; (* node color determined by bisection algorithm *)
cannam@95:     in Dag.for_all dag (fun node -> node.color <- c);
cannam@95: 
cannam@95:     Util.for_list inputs (set_color RED);
cannam@95: 
cannam@95:     (*
cannam@95:        The special inputs are those input nodes that load a single
cannam@95:        location or twiddle factor.  Special inputs can end up either
cannam@95:        in the blue or in the red part.  These inputs are special
cannam@95:        because they inherit a color from their neighbors: If a red
cannam@95:        node needs a special input, the special input becomes red, but
cannam@95:        if all successors of a special input are blue, the special
cannam@95:        input becomes blue.  Outputs are always blue, whether they be
cannam@95:        special or not.
cannam@95: 
cannam@95:        Because of the processing of special inputs, however, the final
cannam@95:        partition might end up being composed only of blue nodes (which
cannam@95:        is incorrect).  In this case we manually reset all inputs
cannam@95:        (whether special or not) to be red.
cannam@95:     *)
cannam@95: 
cannam@95:     Util.for_list special_inputs (set_color YELLOW);
cannam@95: 
cannam@95:     Util.for_list outputs (set_color BLUE);
cannam@95: 
cannam@95:     let rec loopi donep = 
cannam@95:       match (List.filter
cannam@95: 	       (fun node -> (has_color BLACK node) &&
cannam@95: 		 List.for_all (has_either_color RED YELLOW) node.predecessors)
cannam@95: 	       dag') with
cannam@95: 	[] -> if (donep) then () else loopo true
cannam@95:       |	i -> 
cannam@95: 	  begin
cannam@95: 	    Util.for_list i (fun node -> 
cannam@95: 	      begin
cannam@95:       		set_color RED node;
cannam@95: 		Util.for_list node.predecessors (set_color RED);
cannam@95: 	      end);
cannam@95: 	    loopo false; 
cannam@95: 	  end
cannam@95: 
cannam@95:     and loopo donep =
cannam@95:       match (List.filter
cannam@95: 	       (fun node -> (has_either_color BLACK YELLOW node) &&
cannam@95: 		 List.for_all (has_color BLUE) node.successors)
cannam@95: 	       dag') with
cannam@95: 	[] -> if (donep) then () else loopi true
cannam@95:       |	o ->
cannam@95: 	  begin
cannam@95: 	    Util.for_list o (set_color BLUE);
cannam@95: 	    loopi false; 
cannam@95: 	  end
cannam@95: 
cannam@95:     in loopi false;
cannam@95: 
cannam@95:     (* fix the partition if it is incorrect *)
cannam@95:     if not (List.exists (has_color RED) dag') then 
cannam@95: 	Util.for_list inputs (set_color RED);
cannam@95:     
cannam@95:     return
cannam@95:       ((List.map to_assignment (List.filter (has_color RED) dag')),
cannam@95:        (List.map to_assignment (List.filter (has_color BLUE) dag')))
cannam@95:   end
cannam@95: 
cannam@95: type schedule = 
cannam@95:     Done
cannam@95:   | Instr of Expr.assignment
cannam@95:   | Seq of (schedule * schedule)
cannam@95:   | Par of schedule list
cannam@95: 
cannam@95: 
cannam@95: 
cannam@95: (* produce a sequential schedule determined by the user *)
cannam@95: let rec sequentially = function
cannam@95:     [] -> Done
cannam@95:   | a :: b -> Seq (Instr a, sequentially b)
cannam@95: 
cannam@95: let schedule =
cannam@95:   let rec schedule_alist = function
cannam@95:     | [] -> Done
cannam@95:     | [a] -> Instr a
cannam@95:     | alist -> match connected_components alist with
cannam@95: 	| ([a]) -> schedule_connected a
cannam@95: 	| l -> Par (List.map schedule_alist l)
cannam@95: 
cannam@95:   and schedule_connected alist = 
cannam@95:     match partition alist with
cannam@95:     | (a, b) -> Seq (schedule_alist a, schedule_alist b)
cannam@95: 
cannam@95:   in fun x ->
cannam@95:     let () = Util.info "begin schedule" in
cannam@95:     let res = schedule_alist x in
cannam@95:     let () = Util.info "end schedule" in
cannam@95:     res
cannam@95: 
cannam@95: 
cannam@95: (* partition a dag into two parts:
cannam@95: 
cannam@95:    1) the set of loads from locatives and their successors,
cannam@95:    2) all other nodes
cannam@95: 
cannam@95:    This step separates the ``body'' of the dag, which computes the
cannam@95:    actual fft, from the ``precomputations'' part, which computes e.g.
cannam@95:    twiddle factors.
cannam@95: *)
cannam@95: let partition_precomputations alist =
cannam@95:   let dag = makedag alist in
cannam@95:   let dag' = Dag.to_list dag in
cannam@95:   let loads =  List.filter loads_locative dag' in
cannam@95:     begin
cannam@95:       
cannam@95:       Dag.for_all dag (set_color BLUE);
cannam@95:       Util.for_list loads (set_color RED);
cannam@95: 
cannam@95:       let rec loop () = 
cannam@95: 	match (List.filter
cannam@95: 		 (fun node -> (has_color RED node) &&
cannam@95: 		    List.exists (has_color BLUE) node.successors)
cannam@95: 		 dag') with
cannam@95: 	    [] -> ()
cannam@95: 	  |	i -> 
cannam@95: 		  begin
cannam@95: 		    Util.for_list i 
cannam@95: 		      (fun node -> 
cannam@95: 			 Util.for_list node.successors (set_color RED));
cannam@95: 		    loop ()
cannam@95: 		  end
cannam@95: 
cannam@95:       in loop ();
cannam@95: 
cannam@95: 	return
cannam@95: 	  ((List.map to_assignment (List.filter (has_color BLUE) dag')),
cannam@95: 	   (List.map to_assignment (List.filter (has_color RED) dag')))
cannam@95:     end
cannam@95: 
cannam@95: let isolate_precomputations_and_schedule alist =
cannam@95:   let (a, b) = partition_precomputations alist in
cannam@95:     Seq (schedule a, schedule b)
cannam@95: