path: root/scheduling/RTLpathScheduleraux.ml

open DebugPrint
open Machine
open RTLpathLivegenaux
open RTLpath
open RTLpathCommon
open RTL
open Maps
open Registers
open ExpansionOracle

let config = Machine.config

let print_superblock (sb: superblock) code =
  let insts = sb.instructions in
  let li = sb.liveins in
  let outs = sb.s_output_regs in
  begin
    debug "{ instructions = "; print_instructions (Array.to_list insts) code; debug "\n";
    debug "  liveins = "; print_ptree_regset li; debug "\n";
    debug "  output_regs = "; print_regset outs; debug "}"
  end

let print_superblocks lsb code =
  let rec f = function
    | [] -> ()
    | sb :: lsb -> (print_superblock sb code; debug ",\n"; f lsb)
  in begin
    debug "[\n";
    f lsb;
    debug "]"
  end

let get_superblocks code entry pm typing =
  let visited = ref (PTree.map (fun n i -> false) code) in
  let rec get_superblocks_rec pc =
    let liveins = ref (PTree.empty) in
    let rec follow pc n =
      let inst = get_some @@ PTree.get pc code in
      if (n == 0) then begin
        (match (non_predicted_successors inst) with
          | [pcout] ->
              let live = (get_some @@ PTree.get pcout pm).input_regs in
                liveins := PTree.set pc live !liveins
          | _ -> ());
        ([pc], successors_inst inst)
      end else
        let nexts_from_exit = match (non_predicted_successors inst) with
          | [pcout] -> 
              let live = (get_some @@ PTree.get pcout pm).input_regs in begin
                liveins := PTree.set pc live !liveins;
                [pcout]
              end
          | [] -> []
          | _ -> failwith "Having more than one non_predicted_successor is not handled"
        in match (predicted_successor inst) with
          | None -> failwith "Incorrect path"
          | Some succ ->
              let (insts, nexts) = follow succ (n-1) in (pc :: insts, nexts_from_exit @ nexts)
    in if (get_some @@ PTree.get pc !visited) then []
    else begin
      visited := PTree.set pc true !visited;
      let pi = get_some @@ PTree.get pc pm in
      let (insts, nexts) = follow pc (Camlcoq.Nat.to_int pi.psize) in
      let superblock = { instructions = Array.of_list insts; liveins = !liveins;
        s_output_regs = pi.output_regs; typing = typing } in
      superblock :: (List.concat @@ List.map get_superblocks_rec nexts)
    end
  in let lsb = get_superblocks_rec entry in begin
    (* debug_flag := true; *)
    debug "Superblocks identified:"; print_superblocks lsb code; debug "\n";
    (* debug_flag := false; *)
    lsb
end

(* TODO David *)
let schedule_superblock sb code =
  if not !Clflags.option_fprepass
  then sb.instructions
  else
    (* let old_flag = !debug_flag in
    debug_flag := true;
    print_endline "ORIGINAL SUPERBLOCK";
    print_superblock sb code;
    debug_flag := old_flag; *)
    let nr_instr = Array.length sb.instructions in
    let trailer_length =
      match PTree.get (sb.instructions.(nr_instr-1)) code with
      | None -> 0
      | Some ii ->
         match predicted_successor ii with
         | Some _ -> 0
         | None -> 1 in
    match PrepassSchedulingOracle.schedule_sequence
            (Array.map (fun i ->
                 (match PTree.get i code with
                 | Some ii -> ii
                 | None -> failwith "RTLpathScheduleraux.schedule_superblock"),
                 (match PTree.get i sb.liveins with
                 | Some s -> s
                 | None -> Regset.empty))
               (Array.sub sb.instructions 0 (nr_instr-trailer_length))) with
    | None -> sb.instructions
    | Some order ->
       let ins' =
         Array.append 
           (Array.map (fun i -> sb.instructions.(i)) order)
           (Array.sub sb.instructions (nr_instr-trailer_length) trailer_length) in
       (* Printf.printf "REORDERED SUPERBLOCK %d\n" (Array.length ins');
       debug_flag := true;
       print_instructions (Array.to_list ins') code;
       debug_flag := old_flag;
       flush stdout; *)
       assert ((Array.length sb.instructions) = (Array.length ins'));
       (*sb.instructions; *)
       ins';;

  (* stub2: reverse function *)
  (*
  let reversed = Array.of_list @@ List.rev @@ Array.to_list (sb.instructions) in
  let tmp = reversed.(0) in
  let last_index = Array.length reversed - 1 in
  begin
    reversed.(0) <- reversed.(last_index);
    reversed.(last_index) <- tmp;
    reversed
  end *)
  (* stub: identity function *)

(**
 * Perform basic checks on the new order :
 * - must have the same length as the old order
 * - non basic instructions (call, tailcall, return, jumptable, non predicted CB) must not move
 *)
let check_order code old_order new_order = begin
  assert ((Array.length old_order) == (Array.length new_order));
  let length = Array.length new_order in
  if length > 0 then
    let last_inst = Array.get old_order (length - 1) in
    let instr = get_some @@ PTree.get last_inst code in
    match predicted_successor instr with
    | None ->
        if (last_inst != Array.get new_order (length - 1)) then
          failwith "The last instruction of the superblock is not basic, but was moved"
    | _ -> ()
end

type sinst =
  (* Each middle instruction has a direct successor *)
  (* A Smid can be the last instruction of a superblock, but a Send cannot be moved *)
  | Smid of RTL.instruction * node
  | Send of RTL.instruction

let rinst_to_sinst inst =
  match inst with
  | Inop n -> Smid(inst, n)
  | Iop (_,_,_,n) -> Smid(inst, n)
  | Iload (_,_,_,_,_,n) -> Smid(inst, n)
  | Istore (_,_,_,_,n) -> Smid(inst, n)
  | Icond (_,_,n1,n2,p) -> (
      match p with
      | Some true -> Smid(inst, n1)
      | Some false -> Smid(inst, n2)
      | None -> Send(inst)
    )
  | Icall _ | Ibuiltin _ | Ijumptable _ | Itailcall _ | Ireturn _ -> Send(inst)

let change_predicted_successor s = function
  | Smid(i, n) -> Smid(i, s)
  | Send _ -> failwith "Called change_predicted_successor on Send. Are you trying to move a non-basic instruction in the middle of the block?"

(* Forwards the successor changes into an RTL instruction *)
let sinst_to_rinst = function
  | Smid(inst, s) -> (
      match inst with
      | Inop n -> Inop s
      | Iop (a,b,c,n) -> Iop (a,b,c,s)
      | Iload (a,b,c,d,e,n) -> Iload (a,b,c,d,e,s)
      | Istore (a,b,c,d,n) -> Istore (a,b,c,d,s)
      | Icond (a,b,n1,n2,p) -> (
        match p with
        | Some true -> Icond(a, b, s, n2, p)
        | Some false -> Icond(a, b, n1, s, p)
        | None -> failwith "Non predicted Icond as a middle instruction!"
        )
      | _ -> failwith "That instruction shouldn't be a middle instruction"
      )
  | Send i -> i

let is_a_cb = function Icond _ -> true | _ -> false
let is_a_load = function Iload _ -> true | _ -> false

let find_array arr n =
  let index = ref None in
  begin
    Array.iteri (fun i n' ->
      if n = n' then
        match !index with
        | Some _ -> failwith "More than one element present"
        | None -> index := Some i
    ) arr;
    !index
  end

let rec hashedset_from_list = function
  | [] -> HashedSet.PSet.empty
  | n::ln -> HashedSet.PSet.add n (hashedset_from_list ln)

let hashedset_map f hs = hashedset_from_list @@ List.map f @@ HashedSet.PSet.elements hs

let apply_schedule code sb new_order =
  let tc = ref code in
  let old_order = sb.instructions in
  let count_cbs order code =
    let current_cbs = ref HashedSet.PSet.empty in
    let cbs_above = ref PTree.empty in
    Array.iter (fun n ->
      let inst = get_some @@ PTree.get n code in
      if is_a_cb inst then current_cbs := HashedSet.PSet.add n !current_cbs
      else if is_a_load inst then cbs_above := PTree.set n !current_cbs !cbs_above
    ) order;
    !cbs_above
  in let fmap n =
    let index = get_some @@ find_array new_order n in
    old_order.(index)
  in begin
    check_order code old_order new_order;
    (* First pass - modify the positions, nothing else *)
    Array.iteri (fun i n' ->
      let inst' = get_some @@ PTree.get n' code in
      let iend = Array.length old_order - 1 in
      let new_inst =
        if (i == iend) then
          let final_inst_node = Array.get old_order iend in
          let sinst' = rinst_to_sinst inst' in
          match sinst' with
          (* The below assert fails if a Send is in the middle of the original superblock *)
          | Send i -> (assert (final_inst_node == n'); i)
          | Smid _ ->
              let final_inst = get_some @@ PTree.get final_inst_node code in
              match rinst_to_sinst final_inst with
              | Smid (_, s') -> sinst_to_rinst @@ change_predicted_successor s' sinst'
              | Send _ -> assert(false) (* should have failed earlier *)
        else
          sinst_to_rinst
          (* this will fail if the moved instruction is a Send *)
          @@ change_predicted_successor (Array.get old_order (i+1))
          @@ rinst_to_sinst inst'
      in tc := PTree.set (Array.get old_order i) new_inst !tc
    ) new_order;
    (* Second pass - turn the loads back into trapping when it was not needed *)
    (* 1) We remember which CBs are "above" a given load *)
    let cbs_above = count_cbs old_order code in
    (* 2) We do the same for new_order *)
    let cbs_above' = count_cbs (Array.map fmap new_order) !tc in
    (* 3) We examine each load, turn it back into trapping if cbs_above is included in cbs_above' *)
    Array.iter (fun n ->
      let n' = fmap n in
      let inst' = get_some @@ PTree.get n' !tc in
      match inst' with
      | Iload (t,a,b,c,d,s) ->
          let pset = hashedset_map fmap @@ get_some @@ PTree.get n cbs_above in
          let pset' = get_some @@ PTree.get n' cbs_above' in
          if HashedSet.PSet.is_subset pset pset' then tc := PTree.set n' (Iload (AST.TRAP,a,b,c,d,s)) !tc
          else assert !config.has_non_trapping_loads
      | _ -> ()
    ) old_order;
    !tc
  end

let turn_all_loads_nontrap sb code =
  if not !config.has_non_trapping_loads then code
  else begin
    let code' = ref code in
    Array.iter (fun n ->
      let inst = get_some @@ PTree.get n code in
      match inst with
      | Iload (t,a,b,c,d,s) -> code' := PTree.set n (Iload (AST.NOTRAP,a,b,c,d,s)) !code'
      | _ -> ()
    ) sb.instructions;
    !code'
  end

let rec do_schedule code pm = function
  | [] -> (code, pm)
  | sb :: lsb ->
      (*debug_flag := true;*)
      let (code_exp, pm) =
        if !Clflags.option_fexpanse_rtlcond then (expanse sb code pm)
        else (code, pm) in
      (*debug_flag := false;*)
      (* Trick: instead of turning loads into non trap as needed..
       * First, we turn them all into non-trap.
       * Then, we turn back those who didn't need to be turned, into TRAP again
       * This is because the scheduler (rightfully) refuses to schedule ahead of a branch
       * operations that might trap *)
      let code' = turn_all_loads_nontrap sb code_exp in
      let schedule = schedule_superblock sb code' in
      let new_code = apply_schedule code' sb schedule in
      begin
        (*debug_flag := true;*)
        if code != code_exp then (
        debug "Old Code: "; print_code code;
        debug "Exp Code: "; print_code code_exp);
        (*debug "\nSchedule to apply: "; print_arrayp schedule;*)
        (*debug "\nNew Code: "; print_code new_code;*)
        debug "\n";
        do_schedule new_code pm lsb
      end

let get_ok r = match r with Errors.OK x -> x | _ -> failwith "Did not get OK"

let scheduler f =
  let code = f.fn_RTL.fn_code in
  let id_ptree = PTree.map (fun n i -> n) (f.fn_path) in
  let entry = f.fn_RTL.fn_entrypoint in
  let pm = f.fn_path in
  let typing = get_ok @@ RTLtyping.type_function f.fn_RTL in
  let lsb = get_superblocks code entry pm typing in
  begin
    (* debug_flag := true; *)
    debug "Pathmap:\n"; debug "\n";
    print_path_map pm;
    debug "Superblocks:\n";
    (*debug_flag := true; *)
    (*print_code code; flush stdout; flush stderr;*)
    (*debug_flag := false;*)
    (*print_superblocks lsb code; debug "\n";*)
    find_last_node_reg (PTree.elements code);
    (*node := !node - 1;*)
    let (tc, pm) = do_schedule code pm lsb in
    (((tc, entry), pm), id_ptree)
  end