(* probe.sml
*
* COPYRIGHT (c) 2010 The Diderot Project (http://diderot.cs.uchicago.edu)
* All rights reserved.
*
* Expansion of probe operations in the HighIL to MidIL translation.
*)
structure Probe : sig
val expand : MidIL.var * FieldDef.field_def * MidIL.var -> MidIL.assign list
end = struct
structure SrcIL = HighIL
structure SrcOp = HighOps
structure DstIL = MidIL
structure DstOp = MidOps
structure DstV = DstIL.Var
structure VMap = SrcIL.Var.Map
structure IT = Shape
(* generate a new variable indexed by dimension *)
fun newVar_dim (prefix, d) =
DstV.new (prefix ^ Partials.axisToString(Partials.axis d))
fun assign (x, rator, args) = (x, DstIL.OP(rator, args))
fun cons (x, args) = (x, DstIL.CONS args)
fun realLit (x, i) = (x, DstIL.LIT(Literal.Float(FloatLit.fromInt i)))
fun intLit (x, i) = (x, DstIL.LIT(Literal.Int(IntInf.fromInt i)))
(* generate code for a evaluating a single element of a probe operation *)
fun probeElem {
dim, (* dimension of space *)
h, s, (* kernel h with support s *)
n, f, (* Dst vars for integer and fractional components of position *)
voxIter (* iterator over voxels *)
} (result, pdOp) = let
(* generate the variables that hold the convolution coefficients *)
val convCoeffs = let
val Partials.D l = pdOp
fun mkVar 0 = newVar_dim("h", dim)
| mkVar 1 = newVar_dim("dh", dim)
| mkVar i = newVar_dim(concat["d", Int.toString i, "h"], dim)
in
List.map mkVar l
end
(* for each dimension, we evaluate the kernel at the coordinates for that axis *)
val coeffCode = let
fun gen (x, k, (d, code)) = let
val d = d-1
val fd = newVar_dim ("f", d)
val a = DstV.new "a"
val tmps = List.tabulate(2*s, fn i => (DstV.new("t"^Int.toString i), s - (i+1)))
fun mkArg ((t, n), code) = let
val t' = DstV.new "r"
in
realLit (t', n) ::
assign (t, DstOp.Add DstOp.realTy, [fd, t']) ::
code
end
(* code in reverse order *)
val code =
assign(x, DstOp.EvalKernel(2*s, h, k), [a]) ::
assign(fd, DstOp.Select(dim, d), [f]) ::
cons(a, List.map #1 tmps) ::
(List.foldl mkArg [] tmps)
in
(d, List.rev code)
end
val Partials.D l = pdOp
in
#2 (ListPair.foldr gen (dim, []) (convCoeffs, l))
end
(* generate the reduction code *)
fun genReduce (result, [hh], IT.ND(_, kids), code) = let
(* the kids will all be leaves *)
val vv = DstV.new "vv"
fun getVox (IT.LF{vox, offsets}) = vox
in
cons (vv, List.map getVox kids) ::
assign (result, DstOp.Dot(2*s), [hh, vv]) :: code
end
| genReduce (result, hh::r, IT.ND(_, kids), code) = let
val tv = DstV.new "tv"
val tmps = List.tabulate(2*s, fn i => DstV.new("t"^Int.toString i))
fun lp ([], [], code) = code
| lp (t::ts, kid::kids, code) = genReduce(t, r, kid, lp(ts, kids, code))
val code = cons(tv, tmps) :: assign(result, DstOp.Dot(2*s), [hh, tv]) :: code
in
lp (tmps, kids, code)
end
val reduceCode = genReduce (result, convCoeffs, voxIter, [])
in
coeffCode @ reduceCode
end
(* generate code for probing the field (D^k (v * h)) at pos *)
fun probe (result, (k, v, h), pos) = let
val ImageInfo.ImgInfo{dim, ty=([], ty), ...} = v
val dimTy = DstOp.VecTy dim
val s = Kernel.support h
(* generate the transform code *)
val x = DstV.new "x" (* image-space position *)
val f = DstV.new "f"
val nd = DstV.new "nd"
val n = DstV.new "n"
val transformCode = [
assign(x, DstOp.Transform v, [pos]),
assign(nd, DstOp.Floor dim, [x]),
assign(f, DstOp.Sub dimTy, [x, nd]),
assign(n, DstOp.TruncToInt dim, [nd])
]
(* generate the shape of the differentiation tensor with variables representing
* the elements
*)
val diffIter = let
val partial = Partials.partial dim
fun f (i, axes) = Partials.axis i :: axes
fun g axes =
(DstV.new(String.concat("y" :: List.map Partials.axisToString axes)), partial axes)
in
IT.create (k-1, dim, fn _ => (), f, g, [])
end
(* generate code to load the voxel data *)
val voxIter = let
fun f (i, (offsets, id)) = (i - (s - 1) :: offsets, i::id)
fun g (offsets, id) = {
offsets = offsets,
vox = DstV.new(String.concat("v" :: List.map Int.toString id))
}
in
IT.create (dim-1, 2*s, fn _ => (), f, g, ([], []))
end
val loadCode = let
fun genCode ({offsets, vox}, code) = let
fun computeIndices (_, []) = ([], [])
| computeIndices (i, offset::offsets) = let
val index = newVar_dim("i", i)
val t1 = DstV.new "t1"
val t2 = DstV.new "t2"
val (indices, code) = computeIndices (i+1, offsets)
val code =
intLit(t1, offset) ::
assign(t2, DstOp.Select(2*s, i), [n]) ::
assign(index, DstOp.Add(DstOp.IntTy), [t1, t2]) ::
code
val indices = index::indices
in
(indices, code)
end
val (indices, indicesCode) = computeIndices (0, ~(s-1) :: offsets)
val a = DstV.new "a"
in
indicesCode @ [
assign(a, DstOp.VoxelAddress v, indices),
assign(vox, DstOp.LoadVoxels(ty, 2*s), [a])
] @ code
end
in
IT.foldr genCode [] voxIter
end
(* generate code to evaluate and construct the result tensor *)
val probeElem = probeElem {dim = dim, h = h, s = s, n = n, f = f, voxIter = voxIter}
fun genProbe (result, IT.ND(_, kids as (IT.LF _)::_), code) = let
(* the kids will all be leaves *)
fun genProbeCode (IT.LF arg, code) = probeElem arg @ code
fun getProbeVar (IT.LF(t, _)) = t
in
cons (result, List.map getProbeVar kids) :: List.foldr genProbeCode code kids
end
| genProbe (result, IT.ND(_, kids), code) = let
val tmps = List.tabulate(dim, fn i => DstV.new("t"^Int.toString i))
fun lp ([], [], code) = code
| lp (t::ts, kid::kids, code) = genProbe(t, kid, lp(ts, kids, code))
val code = cons(result, tmps) :: code
in
lp (tmps, kids, code)
end
val probeCode = genProbe (result, diffIter, [])
in
transformCode @ loadCode @ probeCode
end
fun expand (result, fld, pos) = let
fun expand' (result, FieldDef.CONV(k, v, h)) = let
val x = DstV.new "x"
val xformStm = (x, DstIL.OP(DstOp.Transform v, [pos]))
in
probe (result, (k, v, h), x) @ [xformStm]
end
(* should push negation down to probe operation
| expand' (result, FieldDef.NEG fld) = let
val r = DstV.new "value"
val stms = expand' (r, fld)
val ty = ??
in
expand' (r, fld) @ [assign(r, DstOp.Neg ty, [r])]
end
*)
| expand' (result, FieldDef.SUM(fld1, dlf2)) = raise Fail "expandInside: SUM"
in
List.rev (expand' (result, fld))
end
end