unroll.ml 2.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. open Types
  2. open Util
  3. (* Only unroll if the resulting number of statements is at most 25 *)
  4. let may_be_unrolled i_values body =
  5. List.length i_values * List.length body <= 25
  6. let rec range i j step =
  7. if i >= j then [] else i :: (range (Int32.add i step) j step)
  8. let rec assigns name = function
  9. | VarLet (dec, _, _, _) -> nameof dec = name
  10. | _ -> false
  11. let rec replace_var name replacement = function
  12. | VarUse (VarDec (_, var, _, _), None, _) when var = name -> replacement
  13. | node -> traverse_unit (replace_var name replacement) node
  14. let rec get_body_step i rest = function
  15. | [] -> None
  16. | [VarLet (
  17. VarDec (Int, assigned, None, _), None,
  18. Binop (Add,
  19. VarUse (VarDec (Int, added, None, _), None, _),
  20. Const (IntVal step, _),
  21. _),
  22. _)] when assigned = added -> Some (step, List.rev rest)
  23. | hd :: tl -> get_body_step i (hd :: rest) tl
  24. let rec unroll_body counters = function
  25. | [] -> []
  26. (*
  27. * Look for the following pattern:
  28. * i = 0;
  29. * while (i < stop) {
  30. * <body>;
  31. * i = i + step;
  32. * }
  33. * where start, stop, step are integer constants and i is a generated variable
  34. *)
  35. | (VarLet (VarDec (Int, i, None, _), None, Const (IntVal start, _), _) as init) ::
  36. (While (
  37. Binop (Lt,
  38. VarUse (VarDec (Int, comp, None, _), None, _),
  39. Const (IntVal stop, _),
  40. _),
  41. Block body,
  42. _) as loop) :: tl
  43. when is_generated_id i && comp = i ->
  44. begin
  45. match get_body_step i [] body with
  46. | Some (step, rest) ->
  47. (* First unroll inner loops in body before the may_be_unrolled check *)
  48. let rest = flatten_blocks (unroll_body counters rest) in
  49. let i_values = range start stop step in
  50. if may_be_unrolled i_values rest then begin
  51. Hashtbl.add counters i true;
  52. let dup_body value =
  53. replace_var i (Const (IntVal value, [Type Int])) (Block rest)
  54. in
  55. Block (List.map dup_body i_values) :: (unroll_body counters tl)
  56. end else
  57. init :: (unroll counters loop) :: (unroll_body counters tl)
  58. | None -> init :: (unroll counters loop) :: (unroll_body counters tl)
  59. end
  60. | hd :: tl -> (unroll counters hd) :: (unroll_body counters tl)
  61. and unroll counters = function
  62. | Block stats -> Block (unroll_body counters stats)
  63. | node -> traverse_unit (unroll counters) node
  64. let rec prune_vardecs counters = function
  65. | VarDec (_, name, _, _) when Hashtbl.mem counters name -> DummyNode
  66. | node -> traverse_unit (prune_vardecs counters) node
  67. let phase = function
  68. | Ast node ->
  69. let counters = Hashtbl.create 10 in
  70. Ast (unroll counters node |> prune_vardecs counters)
  71. | _ -> raise InvalidInput