(* simplify-vars.sml
*
* This module analyses the use of variables in the Simple AST and rationalizes
* their use in the following ways:
*
* -- for any strand parameter that is used in a method, we create a shadow state
* variable
*
* -- strand-invariant state variables (except outputs) and expressions are
* lifted to global scope (TODO)
*
* We assume that contraction has already been run and that unused variables have
* been eliminated.
*
* This code is part of the Diderot Project (http://diderot-language.cs.uchicago.edu)
*
* COPYRIGHT (c) 2015 The University of Chicago
* All rights reserved.
*)
structure SimplifyVars : sig
val transform : Simple.program -> Simple.program
end = struct
structure S = Simple
structure SV = SimpleVar
structure VMap = SimpleVar.Map
(* track if a strand parameter is referenced in a strand method *)
local
val {setFn, getFn} = SV.newFlag()
in
fun markUsedInMethod x = (case SV.kindOf x
of SV.StrandParam => setFn(x, true)
| _ => ()
(* end case *))
fun isUsedInMethod x = getFn x
fun clrUsedInMethodMark x = setFn(x, false)
end (* local *)
(* analyze a method for uses of parameters *)
fun analyzeMethod blk = let
fun analyzeBlk (S.Block stms) = List.app analyzeStm stms
and analyzeStm stm = (case stm
of S.S_Var(x, NONE) => ()
| S.S_Var(x, SOME e) => analyzeExp e
| S.S_Assign(x, e) => analyzeExp e
| S.S_IfThenElse(x, b1, b2) => (markUsedInMethod x; analyzeBlk b1; analyzeBlk b2)
| S.S_Foreach(x, xs, blk) => analyzeBlk blk
| S.S_New(strnd, xs) => List.app markUsedInMethod xs
| S.S_Continue => ()
| S.S_Die => ()
| S.S_Stabilize => ()
| S.S_Return x => markUsedInMethod x
| S.S_Print xs => List.app markUsedInMethod xs
| S.S_MapReduce{args, ...} => raise Fail "unexpected MapReduce in method"
(* end case *))
and analyzeExp exp = (case exp
of S.E_Var x => markUsedInMethod x
| S.E_Lit _ => ()
| S.E_Select(x, fld) => markUsedInMethod x
| S.E_Apply(f, xs, _) => List.app markUsedInMethod xs
| S.E_Prim(_, _, xs, _) => List.app markUsedInMethod xs
| S.E_Tensor(xs, _) => List.app markUsedInMethod xs
| S.E_Seq(xs, _) => List.app markUsedInMethod xs
| S.E_Slice(x, indices, _) => (
markUsedInMethod x;
List.app (Option.app markUsedInMethod) indices)
| S.E_Coerce{x, ...} => markUsedInMethod x
| S.E_LoadSeq _ => ()
| S.E_LoadImage _ => ()
(* end case *))
in
analyzeBlk blk
end
(* rename the free variables in a block according to the given mapping. Variables that are
* not in the domain of the map are unchanged.
*)
fun renameBlock env = let
fun rename x = (case VMap.find(env, x)
of SOME x' => x'
| NONE => x
(* end case *))
val renameList = List.map rename
fun renameBlk (S.Block stms) = S.Block(List.map renameStm stms)
and renameStm stm = (case stm
of S.S_Var(x, NONE) => stm
| S.S_Var(x, SOME e) => S.S_Var(x, SOME(renameExp e))
| S.S_Assign(x, e) => S.S_Assign(rename x, renameExp e)
| S.S_IfThenElse(x, b1, b2) =>
S.S_IfThenElse(rename x, renameBlk b1, renameBlk b2)
| S.S_Foreach(x, xs, blk) =>
S.S_Foreach(rename x, rename xs, renameBlk blk)
| S.S_New(strnd, xs) => S.S_New(strnd, renameList xs)
| S.S_Continue => stm
| S.S_Die => stm
| S.S_Stabilize => stm
| S.S_Return x => S.S_Return(rename x)
| S.S_Print xs => S.S_Print(renameList xs)
| S.S_MapReduce{results, reductions, body, args, source} => S.S_MapReduce{
results = results,
reductions = reductions,
body = body,
args = renameList args,
source = source
}
(* end case *))
and renameExp exp = (case exp
of S.E_Var x => S.E_Var(rename x)
| S.E_Select(x, fld) => S.E_Select(rename x, fld)
| S.E_Lit _ => exp
| S.E_Apply(f, xs, ty) => S.E_Apply(f, renameList xs, ty)
| S.E_Prim(f, tys, xs, ty) =>
S.E_Prim(f, tys, renameList xs, ty)
| S.E_Tensor(xs, ty) => S.E_Tensor(renameList xs, ty)
| S.E_Seq(xs, ty) => S.E_Seq(renameList xs, ty)
| S.E_Slice(x, xs, ty) =>
S.E_Slice(rename x, List.map (Option.map rename) xs, ty)
| S.E_Coerce{srcTy, dstTy, x} =>
S.E_Coerce{srcTy=srcTy, dstTy=dstTy, x=rename x}
| S.E_LoadSeq _ => exp
| S.E_LoadImage _ => exp
(* end case *))
in
renameBlk
end
(* transform a strand definition by introducing shadow state variables for
* parameters.
*)
fun doStrand (S.Strand{name, params, state, stateInit, initM, updateM, stabilizeM}) = let
(* analyze the methods *)
val () = (
Option.app analyzeMethod initM;
analyzeMethod updateM;
Option.app analyzeMethod stabilizeM)
(* identify parameters that need to be shadowed *)
val shadowParams = List.filter (not o isUsedInMethod) params
in
case List.filter isUsedInMethod params
of [] => NONE
| used => let
val (shadowParams, initStms, env) = let
fun f (x, (xs, stms, env)) = let
val x' = SV.copy(x, SV.StrandStateVar)
val stm = S.S_Assign(x', S.E_Var x)
in
(x'::xs, stm::stms, VMap.insert(env, x, x'))
end
in
List.foldr f ([], [], VMap.empty) used
end
val rename = renameBlock env
in
SOME(S.Strand{
name = name, params = params,
state = state @ shadowParams,
stateInit = let
val S.Block stms = stateInit
in
S.Block(stms @ initStms)
end,
initM = Option.map rename initM,
updateM = rename updateM,
stabilizeM = Option.map rename stabilizeM
})
end
(* end case *)
end
fun transform prog = let
val S.Program{
props, consts, inputs, constInit, globals, funcs, init, strand, create, update
} = prog
in
case doStrand strand
of NONE => prog
| SOME strand => S.Program{
props = props, consts = consts, inputs = inputs, constInit = constInit,
globals = globals, funcs = funcs, init = init, strand = strand,
create = create, update = update
}
(* end case *)
end
end