furnace/extern/fftw/genfft/schedule.ml

237 lines
7.1 KiB
OCaml

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