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/expand-integrate.sml
ViewVC logotype

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

Parent Directory Parent Directory | Revision Log Revision Log


Revision 2603 - (download) (annotate)
Thu Apr 24 03:58:37 2014 UTC (7 years, 3 months ago) by cchiw
File size: 11577 byte(s)
Added Lift Constructor
(* examples.sml
 *
 * COPYRIGHT (c) 2012 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *)


(*
A couple of different approaches.
One approach is to find all the Probe(Conv). Gerenerate exp for it
Then use Subst function to sub in. That takes care for index matching and

*)

(*This approach creates probe expanded terms, and adds params to the end. *)


structure Expand = struct

    local
   
    structure E = Ein
    structure mk= mkOperators


structure SrcIL = HighIL
structure SrcTy = HighILTypes
structure SrcOp = HighOps
structure SrcSV = SrcIL.StateVar
structure VTbl = SrcIL.Var.Tbl
structure DstIL = MidIL
structure DstTy = MidILTypes
structure DstOp = MidOps
structure DstV = DstIL.Var
structure SrcV = SrcIL.Var
structure P=Printer
structure shiftHtM=shiftHtM
structure split=splitHtM
structure F=Filter

datatype peanut=    O of  DstOp.rator | E of Ein.ein|C of DstTy.ty|S of int
datatype peanut2=    O2 of  SrcOp.rator | E2 of Ein.ein|C2 of SrcTy.ty|S2 of int
    in


fun assign (x, rator, args) = (x, DstIL.OP(rator, args))
fun assignEin (x, rator, args) = ((x, DstIL.EINAPP(rator, args)))






fun getRHS2 x  = (case SrcIL.Var.binding x
    of SrcIL.VB_RHS(SrcIL.OP(rator, args)) => (O2 rator, args)
    | SrcIL.VB_RHS(SrcIL.VAR x') => getRHS2 x'
    | SrcIL.VB_RHS(SrcIL.EINAPP (e,args))=>(E2 e,args)
    | SrcIL.VB_RHS(SrcIL.CONS (ty,args))=>(C2 ty,args)
    | SrcIL.VB_NONE=>(S2 2,[])
    | vb => raise Fail(concat[
    "expected rhs operator for ", SrcIL.Var.toString x,
    "but found ", SrcIL.vbToString vb])
    (* end case *))





(*Create fractional, and integer position vectors*)
fun transformToImgSpace  (dim,v,posx)=let

    val translate=DstOp.Translate v
    val transform=DstOp.Transform v
    val M = DstV.new ("M", DstTy.tensorTy [dim,dim]) (*transform dim by dim?*)
    val T = DstV.new ("T", DstTy.tensorTy [dim,dim])          (*translate*)
    val x = DstV.new ("x", DstTy.vecTy dim)
    val f = DstV.new ("f", DstTy.vecTy dim)         (*fractional*)
    val nd = DstV.new ("nd", DstTy.vecTy dim)       (*real position*)
    val n = DstV.new ("n", DstTy.iVecTy dim)        (*interger position*)
 

    val PosToImgSpace=mk.transform(dim,dim)
    val code=[
        assign(M, transform, []),
        assign(T, translate, []),

        assignEin(x, PosToImgSpace,[M,posx,T]) ,  (* MX+T*)
        assign(nd, DstOp.Floor dim, [x]),   (*nd *)
        assignEin(f, mk.subTen([dim]),[x,nd]),           (*fractional*)
        assign(n, DstOp.RealToInt dim, [nd]) (*real to Int*)
    ]

    in ([n,f],code)
    end

fun replaceH(kvar, place,args)=let
    val l1=List.take(args, place)
    val l2=List.drop(args,place+1)
    in l1@[kvar]@l2 end

(*Get Img, and Kern Args*)
fun getArgs(hid,hArg,V,imgArg,args,lift)=
    case (getRHS2 hArg,getRHS2 imgArg)
        of ((O2(SrcOp.Kernel(h, i)),arg),(O2(SrcOp.LoadImage img),_))=> let
            val hvar=DstV.new ("KNL", DstTy.KernelTy)
            val imgvar=DstV.new ("IMG", DstTy.ImageTy img)
            val args2= (case lift
                of 0=> let
                    val args1=replaceH(hvar, hid,args)
                    in replaceH(imgvar, V,args1) end
                | _ => [imgvar, hvar]
                (* end case *))
            in 
                (Kernel.support h ,img, [assign (hvar, DstOp.Kernel(h, i), []), assign(imgvar,DstOp.LoadImage img,[])],args2)
            end
        | ((O2(SrcOp.Kernel(h, i)),arg),_)=> raise Fail "Not an img Argument"
        |  _ => raise Fail "Not a kernel argument"



fun handleArgs(E.Probe(E.Conv(V,shape,h, deltas),E.Tensor(t,alpha)),(params,args),origargs,lift)=let
    val E.IMG(dim)=List.nth(params,V)
    val kArg=List.nth(origargs,h)
    val imgArg=List.nth(origargs,V)
    val newposArg=List.nth(args, t)
    val (s,img,argcode,argsVH) =getArgs(h,kArg,V,imgArg,args,lift)
    val (argsT,code')=transformToImgSpace(dim,img,newposArg)
    in (dim,argsVH@argsT,argcode@code', s)
    end



(*createDels=> creates the kronecker deltas for each Kernel*)
fun createDels([],_)= []
    | createDels(d::ds,dim)= [( E.C dim,d)]@createDels(ds,dim)


(*Created new body for probe*)
fun createBody(dim, s,sx,shape,deltas,V, h, nid, fid)=let

    (*sumIndex creating summaiton Index for body*)
    fun sumIndex(0)=[]
    |sumIndex(dim)= sumIndex(dim-1)@[(E.V (dim+sx-1),1-s,s)]

    (*createKRN Image field and kernels *)
    fun createKRN(0,imgpos,rest)=E.Prod ([E.Img(V,shape,imgpos)] @rest)
    | createKRN(dim,imgpos,rest)=let

        val dim'=dim-1
        val sum=sx+dim'
        val dels=createDels(deltas,dim')
        val pos=[E.Add[E.Tensor(fid,[E.C dim']),E.Value(sum)]]
        val rest'= E.Krn(h,dels,E.Sub(E.Tensor(nid,[E.C dim']),E.Value(sum)))
        in
            createKRN(dim',pos@imgpos,[rest']@rest)
        end

    val exp=createKRN(dim, [],[])
    val esum=sumIndex (dim)
    in E.Sum(esum, exp)
    end

fun ShapeConv([],n)=[]
    | ShapeConv(E.C c::es, n)=ShapeConv(es, n)
    | ShapeConv(E.V v::es, n)=
        if(n>v) then [E.V v] @ ShapeConv(es, n)
        else ShapeConv(es,n)


(*Lift probe*)
fun liftProbe(b,(params,args),index, sumIndex,origargs)=let
    val _=print "\n Lift Probe\n"
    val E.Probe(E.Conv(_,alpha,_,dx),pos)=b

    val newId=length(params)
    val n=length(index)

    (*Create new tensor replacement*)
    val shape=ShapeConv(alpha@dx, n)
    val newB=E.Tensor(newId,shape)

    (* Create new Param*)
    val shape'=List.map (fn E.V v=>(List.nth(index, v))) shape
    val newP= E.TEN(1,shape')

    (*Create new Arg*)
    val newArg = DstV.new ("PC", DstTy.tensorTy shape')

    (*Expand Probe*)
    val ns=length sumIndex

    val (dim,args',code,s) = handleArgs(b,(params,args), origargs,1)
    
    val body' =(case ns
        of 0=>    createBody(dim, s,n,alpha,dx,0, 1, 3, 2)
        |_=>let  
            val (E.V v,_,_)=List.nth(sumIndex, ns-1)
            val body'=createBody(dim, s,v+1,alpha,dx,0, 1, 3, 2)
            in  E.Sum(sumIndex ,body')
            end
    (* end case *))

    val params'=[E.IMG(dim),E.KRN,E.TEN(3,[dim]),E.TEN(1,[dim])]
    val (p',i',b',a')=shiftHtM.clean(params', index,body', args')

    val newbie'=Ein.EIN{params=p', index=i', body=b'}
    val _ = print(String.concat["\n ", split.printA(newArg, newbie', a'),"\n"])
    val data=assignEin (newArg, newbie', a')
    in (newB, (params@[newP],args@[newArg]) ,code@[data])
    end
 
 
 (* Expand probe in place *)
 fun replaceProbe(b,(params,args),index, sumIndex,origargs)=let
    val _=print "\n Don't replace probe \n"
    val E.Probe(E.Conv(V,alpha,h,dx),pos)=b
 
    val fid=length(params)
    val n=length(index)
  
    (*Expand Probe*)
    val ns=length sumIndex

    val (dim,args',code,s) = handleArgs(b,(params,args), origargs,0)
    val nid=fid+1


    val body' =(case ns
        of 0=> createBody(dim, s,n,alpha,dx,V, h, nid, fid)
        |_=>let
            val (E.V v,_,_)=List.nth(sumIndex, ns-1)
            in createBody(dim, s,v+1,alpha,dx,V, h, nid, fid)
            end
    (* end case *))

    val params'=params@[E.TEN(3,[dim]),E.TEN(1,[dim])]
    val UUU=Ein.EIN{params=params', index=index, body=body'}
    val _ =print(String.concat["\n $$$ new sub-expression $$$ \n",P.printerE(UUU),"\n"])

    in (body',(params',args') ,code)
    end
 

fun flatten []=[]
    | flatten(e1::es)=e1@(flatten es)

 
 (* sx-[] then move out, otherwise keep in *)
fun expandEinOp ( Ein.EIN{params, index, body}, origargs,args) = let

    val dummy=E.Const 0
    val sumIndex=ref []

    

    (*b-current body, info-original ein op, data-new assigments*)
    fun rewriteBody(b,info)= let

        fun callfn(c1,body)=let
            val ref x=sumIndex
            val c'=[c1]@x
            val (bodyK,infoK,dataK)= (sumIndex:=c';rewriteBody(body ,info))
            val ref s=sumIndex
            val z=hd(s)
            in
                (sumIndex:=tl(s);(E.Sum(z,bodyK),infoK,dataK))
            end
    
        (*val t=print(String.concat["\n\n ***:",Printer.printbody(b),"$ \n"])*)
        in (case b
            of  E.Sum(c,  E.Probe(E.Conv v, E.Tensor t)) =>let
                val ref sx=sumIndex
                in (case sx 
                    of [] => liftProbe(E.Probe(E.Conv v, E.Tensor t ), info,index, c,origargs)
                    | _ =>  replaceProbe(E.Probe(E.Conv v, E.Tensor t ), info,index, (flatten sx)@c,origargs)
                (* end case*))
            end
        | E.Probe(E.Conv _, E.Tensor _) =>let
            val ref sx=sumIndex
            in (case sx
                of []=> liftProbe(b, info,index, [],origargs)
                | _=> replaceProbe(b, info,index, flatten sx,origargs)
             (* end case*))
            end
        | E.Probe _=> (dummy,info,[])
        | E.Conv _=>  (dummy,info,[])
        | E.Lift _=> (dummy,info,[])
        | E.Field _ => (dummy,info,[])
        | E.Apply _ => (dummy,info,[])
        | E.Neg e=> let
            val (body',info',data')=rewriteBody(e,info)
            in
                (E.Neg(body'),info',data')
            end
        | E.Sum([c1],E.Prod p1) => (case F.pushSum2(c1,p1)
            of ([],[])=>raise Fail"Empty summation product"
            | ([],[keep])=> callfn([c1],keep)
            | ([],keep)=> callfn([c1],E.Prod(keep))
            | ([s], [])=> rewriteBody(s ,info)
            | (s, [])=> rewriteBody(E.Prod s ,info)
            | (s,keep)=>let
                val (bodyS,infoS,dataS)= rewriteBody(E.Prod s ,info)
                val ref x=sumIndex
                val c'=[[c1]]@x
                val (bodyK,infoK,dataK)= (sumIndex:=c';rewriteBody(E.Prod keep ,infoS))
                val ref s=sumIndex
                val z=hd(s)
                in (sumIndex:=tl(s);(E.Prod[bodyS,E.Sum(z,bodyK)],infoK, dataS@dataK))
                end
            (*end case*))
        | E.Sum([(v, lb, ub)],e1)=>(case F.foundSx(v,e1)
            of NONE => rewriteBody(e1, info)
            | SOME _ => callfn([(v, lb, ub)],e1)
            (* end case *))
        | E.Sum (c,e)=> callfn(c,e)
        | E.Sub(a,b)=>let
            val (bodyA,infoA,dataA)= rewriteBody(a,info)
            val (bodyB, infoB, dataB)= rewriteBody(b,infoA)
            in   (E.Sub(bodyA, bodyB),infoB,dataA@dataB)
            end
        | E.Div(a,b)=>let
            val (bodyA,infoA,dataA)= rewriteBody(a,info)
            val (bodyB, infoB,dataB)= rewriteBody(b,infoA)
            in  (E.Div(bodyA, bodyB),infoB,dataA@dataB) end
        | E.Add es=> let
            fun filter([], done, info', data)= (E.Add done, info',data)
                | filter(e::es, done, info',data)= let
                    val (body', info'',data')= rewriteBody(e,info')
                    in filter(es, done@[body'], info'',data@data') end
            in filter(es, [],info,[]) end
        
        | E.Prod es=> let
            fun filter([], done, info',data)= (E.Prod done,info', data)
                | filter(e::es, done, info',data)= let
                    val (body', info'',data')= rewriteBody(e, info')
                    in filter(es, done@[body'], info'',data@data') end
                in filter(es, [],info,[]) end
        | _=>  (b,info,[])
        (* end case *))
        end

     val empty =fn key =>NONE
    val mm=print "\n ************************** \n Starting Exapnd"
    val (body',(params',args'),newbies)=rewriteBody(body,(params,args))
    val e'=Ein.EIN{params=params', index=index, body=body'}
    
    val rr=print (String.concat[P.printerE(e'),"\n DONE expand ************************** \n "])
    in ((e',args'),newbies) end

  end; (* local *)

end (* local *)

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