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 2417 - (download) (annotate)
Mon Aug 19 14:15:03 2013 UTC (5 years, 11 months ago) by cchiw
File size: 11176 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 *))
		in
		  mkVar (0, l, [])
		end
 

		



					(********************************)
    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)
		      }
		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.IAdd, [t1, t2]) ::
				    code
				  else
				    intLit(t1, offset) ::
				    assign(index, DstOp.IAdd,[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)
		      in
			indicesCode @ [
			    assign(a, DstOp.VoxelAddress(v, offset), img::indices),
			    assign(vox, DstOp.LoadVoxels(v, 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}

		(*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
          in
            List.foldr genProbeCode (cons (result, List.map getProbeVar kids) :: code) kids
          end

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


(*Analyze *)

                            
                    (**********     trasform to world space***********)
	  val probeCode = if (*(k > 0)*)
            length(deltas)>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])
		   in
		    genProbe (tensor, [xform])
		  end
		else genProbe (result, [])
	  in
        probeCode
	    (*loadCode @ probeCode*)
	  end

(****************************************************************************)

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


  (* 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, deltas',body, param, args)=
            let
                val k=evalDelta(0 ,deltas')
                val x= DstOp.EvalKernel(2*s, h, 0, k)
                
                val arg= assign(v, x, position[0])
            in
                (body@[E.Tensor(id,[])],  param@[E.TEN], args@[arg])
            end

        | generate(i, deltas',body, param, args)=
            let
                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,deltas', body@[E.Tensor(id,[])],  param@[E.TEN], args@[arg]) end
    



        fun onedirection(0,code)=code
        | onedirection(direction,code) =let
                    (*deltas-> deltas' need to be rewritten to reflect direction
                        Delta(i, j)=> i, is kenel,j is direction of result.
                    *)
                val (body, param, args)=generate(dimk, deltas',[],[],[])

                val einexp=E.EIN{
                    params=param @ [E.FLD, E.TEN]
                    index=[],
                    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  onedirection(direction-1,code@[t']) end

    in
        ondirection(dimk, [])
    end

        
    (*
    on
        returns code, where each element is a direction in result
    val positionE=evalPosition(position)
   val e= S.transform(EinOp.AddTensor,dim)
    in R.App(e, positionE) end 
*)




    (*using sum hmmmm
    fun generate( paramk, indexk, bodyk, argk,0)=
        let
           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)=
        let
            val deriv=evalDelta(deltas,direction)
            val arg=DStOp.EvalKernel(2*s, h, direction, deriv),[pos]
            val next=direction-1
        in
            generate(paramk@[E.TEN], indexk@[E.SX(s-1)], bodyk@[E.Tensor(next, [])], argk@[arg],next)
        end



    in


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

end
  end

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