Home My Page Projects Code Snippets Project Openings diderot
Summary Activity Tracker Tasks SCM

SCM Repository

[diderot] View of /branches/vis12/src/compiler/high-to-mid/probe.sml
ViewVC logotype

View of /branches/vis12/src/compiler/high-to-mid/probe.sml

Parent Directory Parent Directory | Revision Log Revision Log


Revision 3072 - (download) (annotate)
Sun Mar 15 16:33:00 2015 UTC (4 years, 3 months ago) by jhr
File size: 14719 byte(s)
  working on border control
(* probe.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Expansion of probe operations in the HighIL to MidIL translation.
 *)

structure Probe : sig

    val expand : {
            result : MidIL.var,         (* result variable for probe *)
            img : MidIL.var,            (* probe image argument *)
            v : ImageInfo.info,         (* summary info about image *)
            h : Kernel.kernel,          (* reconstruction kernel *)
            k : int,                    (* number of levels of differentiation *)
            pos : MidIL.var,            (* probe position argument *)
            ctl : MidIL.var BorderCtl.ctl (* border-control for image indexing *)
          } -> MidIL.assign list

  end = struct

    structure SrcIL = HighIL
    structure SrcOp = HighOps
    structure DstIL = MidIL
    structure DstTy = MidILTypes
    structure DstOp = MidOps
    structure DstV = DstIL.Var
    structure VMap = SrcIL.Var.Map
    structure BCtl = BorderCtl
    structure IT = Shape

  (* generate a new variable indexed by dimension *)
    fun newVar_dim (prefix, d, ty) =
          DstV.new (prefix ^ Partials.axisToString(Partials.axis d), ty)

    fun assign (x, rator, args) = (x, DstIL.OP(rator, args))
    fun cons (x, args) = (x, DstIL.CONS(DstV.ty x, args))
    fun realLit (x, i) = (x, DstIL.LIT(Literal.Float(FloatLit.fromInt(IntInf.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
          val vecsTy = DstTy.vecTy(2*s) (* vectors of coefficients cover support of kernel *)
          val vecDimTy = DstTy.vecTy dim
        (* generate the variables that hold the convolution coefficients.  The
         * resulting list is in slowest-to-fastest axes order.
         *)
          val convCoeffs = let
                val Partials.D l = pdOp
                fun mkVar (_, [], coeffs) = coeffs
                  | mkVar (i, d::dd, coeffs) = (case d
                       of 0 => mkVar(i+1, dd, newVar_dim("h", i, vecsTy) :: coeffs)
                        | 1 => mkVar(i+1, dd, newVar_dim("dh", i, vecsTy) :: coeffs)
                        | _ => mkVar(i+1, dd, newVar_dim(concat["d", Int.toString d, "h"], i, vecsTy) :: coeffs)
                      (* end case *))
                in
                  mkVar (0, l, [])
                end
        (* for each dimension in space, we evaluate the kernel at the coordinates for that axis.
         * the coefficients are
         *      h_{s-i} (f - i) for 1-s <= i <= s
         *)
          val coeffCode = let
                fun gen (x, k, (d, code)) = let
                    (* note that for 1D images, the f vector is a scalar *)
                      val fd = if (dim > 1)
                            then newVar_dim ("f", d, DstTy.realTy)
                            else f
                      val a = DstV.new ("a", vecsTy)
                    (* note that we reverse the order of the list since the convolution
                     * space is flipped from the image space and we want the voxel vector
                     * to be in increasing address order.
                     *)
                      val tmps = List.rev(List.tabulate(2*s,
                            fn i => (DstV.new("t"^Int.toString i, DstTy.realTy), i - s)))
                      fun mkArg ((t, 0), code) = (t, DstIL.VAR fd) :: code
                        | mkArg ((t, n), code) = let
                            val (rator, n) = if (n < 0) then (DstOp.Sub, ~n) else (DstOp.Add, n)
                            val t' = DstV.new ("r", DstTy.realTy)
                            in
                              realLit (t', n) ::
                              assign (t, rator DstTy.realTy, [fd, t']) ::
                              code
                            end
                      val code =
                            cons(a, List.map #1 tmps) ::                    
                            assign(x, DstOp.EvalKernel(2*s, h, k), [a]) ::
                              code
                      val code = List.foldr mkArg code tmps
                      val code = if (dim > 1)
                            then assign(fd, DstOp.Index(DstTy.vecTy dim, d), [f]) :: code
                            else code
                      in
                        (d+1, code)
                      end
                val Partials.D l = pdOp
                in
                (* we iterate from fastest to slowest axis *)
                  #2 (ListPair.foldr gen (0, []) (convCoeffs, List.rev l))
                end
        (* generate the reduction code in reverse order *)
          fun genReduce (result, [hh], IT.LF{vox, offsets}, code) =
                assign (result, DstOp.Dot(2*s), [vox, hh]) :: code
            | genReduce (result, hh::r, IT.ND(_, kids), code) = let
                val tv = DstV.new ("tv", vecsTy)
                val tmps = List.tabulate(2*s, fn i => DstV.new("t"^Int.toString i, DstTy.realTy))
                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
            | genReduce _ = raise Fail "genReduce"
          val reduceCode = genReduce (result, convCoeffs, voxIter, [])
          in
            coeffCode @ reduceCode
          end

    fun loadVoxelVec (vox, v, offset, s, img, indices, BCtl.None) = let
	  val a = DstV.new ("a", DstTy.AddrTy v)
	  in [
	    assign(a, DstOp.VoxelAddress(v, offset), img::indices),
	    assign(vox, DstOp.LoadVoxels(v, 2*s), [a])
	  ] end
      | loadVoxelVec (vox, v, offset, s, img, indices, BCtl.Default x) = raise Fail "Border control default"
      | loadVoxelVec (vox, v, offset, s, img, indices, BCtl.Index ictl) = let
	  val numVoxels = 2*s
	(* split the indices into a reverse-order prefix and the last index *)
	  val (revIndices, baseIdx) = let
		fun split ([i], prefix) = (prefix, i)
		  | split (idx::idxs, prefix) = split (idxs, idx::prefix)
		in
		  split (indices, [])
		end
	  fun mk (i, voxs, stms) = if (i < numVoxels)
		then let
		  val (indices, stms) = if (i = 0)
			then (indices, stms)
			else let
			  val idx = DstV.new ("idx", DstTy.IntTy)
			  val n = DstV.new ("n", DstTy.IntTy)
			  val stms = assign(idx, DstOp.Add DstTy.IntTy, [baseIdx, n]) ::
				intLit(n, i) :: stms
			  in
			    (List.revAppend(revIndices, [idx]), stms)
			  end
		  val a = DstV.new ("a", DstTy.AddrTy v)
		  val addrStm = assign(a, DstOp.VoxelAddressWithCtl(v, offset, ictl), img::indices)
		  val vox = DstV.new ("vox", DstTy.realTy)
		  val loadStm = assign(vox, DstOp.LoadVoxels(v, 1), [a])
		  in
		    mk (i+1, vox::voxs, loadStm :: addrStm :: stms)
		  end
		else List.rev(cons(vox, List.rev voxs) :: stms)
	  in
	    mk (0, [], [])
	  end

    fun doVoxelSample (result, v, k, s, diffIter, {h, n, f, img}, offset, ctl) = let
          val stride = ImageInfo.stride
          val dim = ImageInfo.dim v
          val vecsTy = DstTy.vecTy(2*s) (* vectors of coefficients cover support of kernel *)
        (* generate code to load the voxel data; since we use a vector load operation to load the
         * fastest dimension, the height of the tree is one less than the dimension of space.
         *)
          val voxIter = let
                fun f (i, (offsets, id)) = (i - (s - 1) :: offsets, i::id)
                fun g (offsets, id) = {
                        offsets = ~(s-1) :: offsets,
                        vox = DstV.new(String.concat("v" :: List.map Int.toString id), vecsTy)
                      }
                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, DstTy.intTy)
                            val t1 = DstV.new ("t1", DstTy.intTy)
                            val t2 = DstV.new ("t2", DstTy.intTy)
                            val (indices, code) = computeIndices (i+1, offsets)
                            val code = if (dim > 1)
                                  then
                                    intLit(t1, offset) ::
                                    assign(t2, DstOp.Index(DstTy.iVecTy dim, i), [n]) ::
                                    assign(index, DstOp.Add(DstTy.intTy), [t1, t2]) ::
                                    code
                                  else
                                    intLit(t1, offset) ::
                                    assign(index, DstOp.Add(DstTy.intTy), [t1, n]) ::
                                    code
                            val indices = index::indices
                            in
                              (indices, code)
                            end
                      val (indices, indicesCode) = computeIndices (0, offsets)
                      val a = DstV.new ("a", DstTy.AddrTy v)
		      val loadStms = loadVoxelVec (vox, v, offset, s, img, indices, ctl)
                      in
                        indicesCode @ loadStms @ 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
                  List.foldr genProbeCode (cons (result, List.map getProbeVar kids) :: code) kids
                end
            | genProbe (result, IT.ND(ty, kids), code) = let
(* FIXME: the type of the tmps depends on the types of the kids *)
                val tmps = List.tabulate(dim, fn i => DstV.new("t"^Int.toString i, ty))
                val code = cons(result, tmps) :: code
                fun lp ([], [], code) = code
                  | lp (t::ts, kid::kids, code) = genProbe(t, kid, lp(ts, kids, code))
                in
                  lp (tmps, kids, code)
                end
            | genProbe (result, IT.LF(t, pdOp), code) = (* for scalar fields *)
                probeElem (result, pdOp) @ code
          val probeCode = if (k > 0)
                then let
                (* for gradients, etc. we have to transform back to world space *)
                  val ty = DstV.ty result
                  val tensor = DstV.new("tensor", ty)
                  val xform = assign(result, DstOp.TensorToWorldSpace(v, ty), [img, tensor])
                  in
                    genProbe (tensor, diffIter, [xform])
                  end
                else genProbe (result, diffIter, [])
          in
(* FIXME: for dim > 1 and k > 1, we need to transform the result back into world space *)
            loadCode @ probeCode
          end

  (* generate code for probing the field (D^k (v * h)) at pos *)
    fun expand {result, img, v, h, k, pos, ctl} = let
          val dim = ImageInfo.dim v
          val s = Kernel.support h
          val vecsTy = DstTy.vecTy(2*s) (* vectors of coefficients to cover support of kernel *)
          val vecDimTy = DstTy.vecTy dim
        (* generate the transform code *)
          val x = DstV.new ("x", vecDimTy)      (* image-space position *)
          val f = DstV.new ("f", vecDimTy)
          val nd = DstV.new ("nd", vecDimTy)
          val n = DstV.new ("n", DstTy.iVecTy dim)
          val toImgSpaceCode = [
                  assign(x, DstOp.PosToImgSpace v, [img, pos]),
                  assign(nd, DstOp.Floor dim, [x]),
                  assign(f, DstOp.Sub vecDimTy, [x, nd]),
                  assign(n, DstOp.RealToInt dim, [nd])
                ]
        (* generate the shape of the differentiation tensor with variables representing
         * the elements
         *)
          fun diffIter () = let
                val partial = Partials.partial dim
                fun f (i, (_::dd, axes)) = (dd, Partials.axis i :: axes)
                fun labelNd (_::dd, _) = DstTy.tensorTy dd
                fun labelLf (_, axes) = let
                        val r = DstV.new(
                                String.concat("r" :: List.map Partials.axisToString axes),
                                DstTy.realTy)
                        in
                          (r, partial axes)
                        end
                in
                  IT.create (k, dim, labelNd, f, labelLf, (List.tabulate(k, fn _ => dim), []))
                end
          val vars = {h=h, n=n, f=f, img=img}
          in
            case ImageInfo.voxelShape v
             of [] => toImgSpaceCode @ doVoxelSample (result, v, k, s, diffIter(), vars, 0, ctl)
              | shape => let
                (* the result will be a d1 x d2 matrix of k-order tensors; each element of the matrix
                 * has the sampleTy = tensor[dim,...,dim].
                 *)
                  val sampeShape = List.tabulate(k, fn _ => dim)
                  fun gen (y, d::dd, offset, code) = let
                        val prefix = DstV.name y
                        val xs = List.tabulate(d,
                              fn i => DstV.new(
                                concat[prefix, "_", Int.toString i],
                                DstTy.TensorTy(dd @ sampeShape)))
                        val code = cons(y, xs) :: code
                        in
                          List.foldr (fn (x, (off, cd)) => gen(x, dd, off, cd)) (offset, code) xs
                        end
                    | gen (y, [], offset, code) =
                        (offset-1, doVoxelSample (y, v, k, s, diffIter(), vars, offset, ctl) @ code)
                  val (_, code) = gen (result, shape, ImageInfo.stride v - 1, [])
                  in
                    toImgSpaceCode @ code
                  end
            (* end case *)
          end

  end

root@smlnj-gforge.cs.uchicago.edu
ViewVC Help
Powered by ViewVC 1.0.0