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

SCM Repository

[diderot] View of /branches/vis15/src/compiler/mid-to-low/field-to-low.sml
ViewVC logotype

View of /branches/vis15/src/compiler/mid-to-low/field-to-low.sml

Parent Directory Parent Directory | Revision Log Revision Log


Revision 3796 - (download) (annotate)
Fri Apr 29 21:40:04 2016 UTC (3 years, 3 months ago) by cchiw
File size: 5760 byte(s)
use ProjectLast to index image tensor
(* field-to-low.sml
 *
 * NOTE: this code will need to be changed if we ever want to support different kernels
 * for different axes
 *
 * This code is part of the Diderot Project (http://diderot-language.cs.uchicago.edu)
 *
 * COPYRIGHT (c) 2016 The University of Chicago
 * All rights reserved.
 *)

structure FieldToLow : sig

  (* expand a MidIR probe to LowIR code.  The arguments are:
   *
   *    avail   -- available LowIR assignments
   *    mapp    -- mapping from iteration indices to deBruijn indices
   *    sx      -- summation bounds
   *    prod    -- body of summation, which is a product of Ein expressions
   *    args    -- the actual arguments of the enclosing Ein expression
   *)
    val expand : {
            avail : AvailRHS.t,
            mapp : int IntRedBlackMap.map,
            sx : Ein.sumrange list,
            prod : Ein.ein_exp list,
            args : LowIR.var list
          }  -> LowIR.var

  end = struct

    structure IR = LowIR
    structure Ty = LowTypes
    structure Op = LowOps
    structure Var = LowIR.Var
    structure E = Ein
    structure Mk = MkLowIR
    structure IMap = IntRedBlackMap

  (* Index cons at piece *)
    fun getHolder (args, id, piece) = let
          val t = List.nth(args, id)
          in
            case IR.Var.getDef t
             of IR.CONS(eargs, _) => List.nth(eargs, piece)
              | IR.OP(_, _) => t
              | rhs => raise Fail(String.concat[
                    "getHolder found ", Var.name t, "=", IR.RHS.toString rhs,
                    " at ", Int.toString id
                  ])
            (* end case *)
          end

(* FIXME: I think this function can be reimplemented in a more straightforward way *)
  (* evaluate image expression *)
    fun imgToArgs (avail, mapp, sx, E.Img(Vid, alpha, vs, s, _), args) = let
          val vI = List.nth(args, Vid)
          val range1 = 2*s
          val range = List.tabulate (range1, fn e => e)
          val beta = List.map (fn id => Mk.lookupMu(mapp, id)) alpha
          val dim = length(vs)

        (* Index tensor with image shape and position v_beta[n0,n1,n2]*)
        (* change here depending on data layout created by load voxel*)

        fun getIX idxs =
            AvailRHS.addAssign (
                avail, "proj", Ty.TensorTy [range1],
                IR.OP(Op.ProjectLast(IR.Var.ty (vI), beta @ idxs), [vI]))

        (* 2-d case*)
          fun iter2 ([], rest) = rest
            | iter2 (i::es, rest) = iter2(es, getIX [i]::rest)
        (* 3-d case*)
          fun iter3 ([], [], rest) = rest
            | iter3 ([_], [], rest) = rest
            | iter3 (i::es, [], rest) = iter3(es, range, rest)
            | iter3 (i::es, j::js, rest) = iter3(i::es, js, getIX [i, j]::rest)
          in
            case vs
             of [_]       => [getIX []]                 (* 1-d case *)
              | [_,_]     => iter2 (range, [])          (* 2-d case*)
              | [_, _, _] => iter3 (range, range, [])   (* 3-d case*)
	      | _ => raise Fail "unsupported dimension"
           (* end case *)
          end

  (* Convolution product of Image and Kernel *)
    fun prodImgKrn (avail, imgArg, krnArg, s) = let
        (* Number of arguments for Cons *)
          val range1 = 2*s
          val range0 = range1-1
          fun ConsInt args = let
                val ty = Ty.TensorTy [range1]
                val rhs = IR.CONS (args, ty)
                in
                  AvailRHS.addAssign (avail,"cons"^"_", ty, rhs)
                end
          fun mkDotVec (a,b) = Mk.vecDot(avail, range1, a, b)
          fun mul2d ([], rest, hy) = ConsInt (List.rev rest)
            | mul2d (e::es, rest, hy) = let
                val vA = mkDotVec (e, hy)
                in
                  mul2d (es, vA::rest, hy)
                end
          fun mul3d ([], _ , _, rest, hz) = rest
            | mul3d (e1::es, rest, 0, consrest, hz) = let
                val vA = mkDotVec (hz, e1)
                val vD = ConsInt (rest@[vA])
                in
                  mul3d (es, [], range0, consrest@[vD], hz)
                end
            | mul3d (e1::es, rest, n, consrest, hz) = let
                val vA = mkDotVec (hz, e1)
                in
                  mul3d (es, rest@[vA], n-1, consrest, hz)
                end
        (*Create Product by doing case analysis of the dimension*)
          in
            case (krnArg, imgArg)
             of ([h0], [i]) => mkDotVec (i, h0)   (*1-D case*)
              | ([h0, h1], _) => let
                  val vA = mul2d (imgArg, [], h0)
                  in
                    mkDotVec (vA, h1)
                  end
              | ([h0, h1, h2], _) => let
                  val restZ = mul3d (imgArg, [], range0, [], h0)
                  val restY = mul2d (restZ, [], h1)
                  in
                    mkDotVec (h2, restY)
                  end
              | _ =>  raise Fail "Kernel dimensions not between 1-3"
           (* end case *)
          end

  (* expand a MidIR probe to LowIR code. *)
    fun expand {avail, mapp, sx, prod as E.Img e1::krnexps, args} = let
          val imgArgs = imgToArgs(avail, mapp, sx, E.Img e1, args)
        (* get piece for each kernel *)
          fun getf (E.Krn(id, dels, _)) = let
              (* evaluate dels to integer *)
                val delta = List.foldl (fn((i, j), y) => Mk.evalDelta(mapp, i, j) + y) 0 dels
                in
                  getHolder (args, id, delta)  (* selects variable in holder/cons list *)
                end
        (* evaluate kernel expression(s) *)
          val krnArgs = List.map getf krnexps  (* doesn't create code *)
          val rtn = prodImgKrn (avail, imgArgs, krnArgs, #4 e1)
          in
            rtn
          end

  end

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