diff options
Diffstat (limited to 'src/optimize.ml')
| -rw-r--r-- | src/optimize.ml | 58 |
1 files changed, 36 insertions, 22 deletions
diff --git a/src/optimize.ml b/src/optimize.ml index 1fc2fbe8..b0d05bef 100644 --- a/src/optimize.ml +++ b/src/optimize.ml @@ -52,43 +52,57 @@ open Ast open Ast_util open Rewriter +let rec split_at_function' id defs acc = + match defs with + | [] -> None + | ([def], env) :: defs when is_fundef id def -> Some (acc, (def, env), defs) + | (def, env) :: defs -> split_at_function' id defs ((def, env) :: acc) + +let split_at_function id defs = + match split_at_function' id defs [] with + | None -> None + | Some (pre_defs, def, post_defs) -> + Some (List.rev pre_defs, def, post_defs) + let recheck (Defs defs) = let defs = Type_check.check_with_envs Type_check.initial_env defs in let rec find_optimizations = function - | ([DEF_pragma ("optimize", pragma, p_l)], env) - :: ([DEF_spec vs as def1], _) - :: ([DEF_fundef fdef as def2], _) - :: defs -> + | ([DEF_pragma ("optimize", pragma, p_l)], env) :: ([DEF_spec vs as def1], _) :: defs -> let id = id_of_val_spec vs in let args = Str.split (Str.regexp " +") (String.trim pragma) in begin match args with | ["unroll"; n]-> let n = int_of_string n in + begin match split_at_function id defs with + | Some (intervening_defs, ((DEF_fundef fdef as def2, _)), defs) -> + let rw_app subst (fn, args) = + if Id.compare id fn = 0 then E_app (subst, args) else E_app (fn, args) + in + let rw_exp subst = { id_exp_alg with e_app = rw_app subst } in + let rw_defs subst = { rewriters_base with rewrite_exp = (fun _ -> fold_exp (rw_exp subst)) } in - let rw_app subst (fn, args) = - if Id.compare id fn = 0 then E_app (subst, args) else E_app (fn, args) - in - let rw_exp subst = { id_exp_alg with e_app = rw_app subst } in - let rw_defs subst = { rewriters_base with rewrite_exp = (fun _ -> fold_exp (rw_exp subst)) } in - - let specs = ref [def1] in - let bodies = ref [rewrite_def (rw_defs (append_id id "_unroll_1")) def2] in + let specs = ref [def1] in + let bodies = ref [rewrite_def (rw_defs (append_id id "_unroll_1")) def2] in - for i = 1 to n do - let current_id = append_id id ("_unroll_" ^ string_of_int i) in - let next_id = if i = n then current_id else append_id id ("_unroll_" ^ string_of_int (i + 1)) in - (* Create a valspec for the new unrolled function *) - specs := !specs @ [DEF_spec (rename_valspec current_id vs)]; - (* Then duplicate it's function body and make it call the next unrolled function *) - bodies := !bodies @ [rewrite_def (rw_defs next_id) (DEF_fundef (rename_fundef current_id fdef))] - done; + for i = 1 to n do + let current_id = append_id id ("_unroll_" ^ string_of_int i) in + let next_id = if i = n then current_id else append_id id ("_unroll_" ^ string_of_int (i + 1)) in + (* Create a valspec for the new unrolled function *) + specs := !specs @ [DEF_spec (rename_valspec current_id vs)]; + (* Then duplicate it's function body and make it call the next unrolled function *) + bodies := !bodies @ [rewrite_def (rw_defs next_id) (DEF_fundef (rename_fundef current_id fdef))] + done; - !specs @ !bodies @ find_optimizations defs + !specs @ List.concat (List.map fst intervening_defs) @ !bodies @ find_optimizations defs + | _ -> + Reporting.warn "Could not find function body for unroll pragma at " p_l ""; + def1 :: find_optimizations defs + end | _ -> Reporting.warn "Unrecognised optimize pragma at" p_l ""; - def1 :: def2 :: find_optimizations defs + def1 :: find_optimizations defs end | (defs, _) :: defs' -> |