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

SCM Repository

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

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

Parent Directory Parent Directory | Revision Log Revision Log

Revision 2414 - (download) (annotate)
Mon Aug 19 05:02:14 2013 UTC (6 years, 1 month ago) by cchiw
File size: 10721 byte(s)
(* 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 *)
	  } -> 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 E=EinOp
    structure EinOp= Operators
    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 assignEin (x, rator, args) = (x, DstIL.EINAPP(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)))
            (*added copy from translate-basis*)
   fun pruneShape sv = (case TU.pruneShape(MV.toShape sv)
           of Ty.Shape dd => DstTy.tensorTy(List.map pruneDim dd)
            | shp => raise Fail("unresolved shape " ^ TU.shapeToString shp)
          (* end case *))
    fun shapeVarToTensor sv = pruneShape sv

    fun createVec d= S.transform(EinOp.tensorOp,[[d]],[])
    fun createdot(d)= S.transform(EinOp.innerProduct, [(0,[]),(0,[]),(0,[d])],[])


  (* 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 = createVec(2*s) (* vectors of coefficients cover support of kernel *)
	  val vecDimTy = createVec 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 *))
		  mkVar (0, l, [])


    fun doVoxelSample (result, v, k, s, deltas, {h, n, f, img}, offset) = let
	  val stride = ImageInfo.stride
	  val dim = ImageInfo.dim v
	  val vecsTy = createVec(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)
		  IT.create (dim-1, 2*s, fn _ => (), f, g, ([], []))
	  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)
				    intLit(t1, offset) ::
				    assign(t2, DstOp.Index(DstTy.iVecTy dim, i), [n]) ::
				    assign(index, DstOp.IAdd, [t1, t2]) ::
				    intLit(t1, offset) ::
				    assign(index, DstOp.IAdd,[t1, n]) ::
			    val indices = index::indices
			      (indices, code)
		      val (indices, indicesCode) = computeIndices (0, offsets)
		      val a = DstV.new ("a", DstTy.AddrTy v)
			indicesCode @ [
			    assign(a, DstOp.VoxelAddress(v, offset), img::indices),
			    assign(vox, DstOp.LoadVoxels(v, 2*s), [a])
			  ] @ code
		  IT.foldr genCode [] voxIter
	(* generate code to evaluate and construct the result tensor *)
	  val probeElem = probeElem {dim = dim, h = h, s = s, n = n, f = f, voxIter = voxIter}

		(*Used for result, *)

    (*Node, next to leaf*)

      fun genProbe (result, IT.ND(_, kids as (IT.LF _)::_), code) = let
		  fun genProbeCode (IT.LF arg, code) = probeElem arg @ code
		  fun getProbeVar (IT.LF(t, _)) = t
            List.foldr genProbeCode (cons (result, List.map getProbeVar kids) :: code) kids

    (*Node next to node*)

	    | genProbe (result, IT.ND(ty, kids), code) = let

            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))
		          lp (tmps, kids, code)
	    | genProbe (result, IT.LF(t, pdOp), code) = (* for scalar fields *)
		     probeElem (result, pdOp) @ code

(*Analyze *)

                    (**********     trasform to world space***********)
	  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 alpha= List.take(dim, length(dim)-1)
            val ilist=  [alpha,  tl(tensor), hd(tensor)]
            val TensorToWorldSpace=S.tranform(EinOp.innerProduct, ilist,[])
            val xform = assignEin(result, TensorToWorldSpace, [img, tensor])
		    genProbe (tensor, [xform])
		else genProbe (result, [])
	    (*loadCode @ probeCode*)


(*    from EinOps     of (kern, deltas), x)
          arguments   of fld-> (img->(v), h->(h,k))

  (* generate code for probing the field (D^k (v * h)) at pos *)
    fun expandEin {result, img, v, h, k, pos,deltas,index} = let
	  val dim = ImageInfo.dim v
	  val s = Kernel.support h
	  val vecsTy =createVec(2*s)
	  val vecDimTy = createVec(dim)
      val translate=DstOp.Translate v
      val transform=DstOp.Transform v
	(* 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 M = DstV.new ("M", transform)
      val T = DstV.new ("T", translate)
      val sub= S.transform(EinOp.subTensor,[dim],[])

       (* M_ij x_i*)
      val MXop=S.transform(EinOp.innerProduct,[[dim],[],[dim]],[])
      val MX = DstV.new ("MX", MXop)
      val PosToImgSpace=S.transform(EinOp.addTensor,[[dim]],[])
      val toImgSpaceCode = [
          assignEin(MX, Mxop, [M,pos]), (*Just added*)
          assignEin(x, PosToImgSpace,[MX,T])
		  assign(nd, DstOp.Floor dim, [x]),
		  assignEin(f, sub,[x,nd]),
		  assign(n, DstOp.RealToInt dim, [nd])


        This return an einExip, kernels are represented as scalar tensor in the epression, and
        the arguments for kernels are midOps-that have the kernel, direction and differentiation.

        Conv(h, delta)
            so look at deltas

      (* dim of kernel dimK*)

        (*maybe this should be somewhere else*)
    fun evalDelta(_,[])= 0
      | evalDelta(i, E.Delta(_,j)::dels)=
            if (i==j) then 1 + evalDelta(i,dels)
            else evalDelta(i,dels)

    fun evalPosition(position)=let
        fun generate(0, body, param, args)=
                val k=evalDelta(0 ,deltas)
                val x= DstOp.EvalKernel(2*s, h, 0, k)
                val arg= assign(v, x, position[0])
            in args@ [arg] end

        (* (body@[E.Tensor(id,[])],  param@[E.TEN], args@arg) end*)

        | generate(i, body, param, args)=
                val k=evalDelta(i ,deltas)
                val x= DstOp.EvalKernel(2*s, h, i, k)

                val arg= assign(v, x, position[i])
            in   generate(i-1, body@[E.Tensor(id,[])],  param@[E.TEN], args@[arg]) end
        val (body, param, args)=generate(dimk, [],[],[])

        val einexp=E.EIN{
                params=param @ [E.FLD, E.TEN]
                body= E.Prod body @[E.Field(id, [dim]), E.Tensor(id,[])]

        val r=S.transform(EinOp.Probe,[dimk])
        val t=R.App(r, [img, position])
        val t'=R.App(einexp, args@ [t])
        in  t' end

    call evalPosition on each position
    val positionE=evalPosition(position)
   val e= S.transform(EinOp.AddTensor,dim)
    in R.App(e, positionE) end 

    (*using sum hmm*)
    fun generate( paramk, indexk, bodyk, argk,0)=
           val einexp= E.EIN{
                params = paramk @ [E.FLD(dim), E.TEN],
                index = indexk,
                body = E.SUM(dimk, E.Prod(bodyk@[E.Probe(E.Field(dim), E.Tensor(dimk+1,[]))]))
            val args=argk @ [img, pos]
        in R.APP(einexp, args) end

      | generate( paramk, indexk, bodyk, argk,direction)=
            val deriv=evalDelta(deltas,direction)
            val arg=DStOp.EvalKernel(2*s, h, direction, deriv),[pos]
            val next=direction-1
            generate(paramk@[E.TEN], indexk@[E.SX(s-1)], bodyk@[E.Tensor(next, [])], argk@[arg],next)

    (*Add each tensors*)

    generate([],[],[], [],dimk)


ViewVC Help
Powered by ViewVC 1.0.0