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/load-voxels.sml
ViewVC logotype

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

Parent Directory Parent Directory | Revision Log Revision Log


Revision 3802 - (download) (annotate)
Tue May 3 17:13:30 2016 UTC (3 years, 4 months ago) by jhr
File size: 7708 byte(s)
  working on merge
(* load-voxels.sml
 *
 * This module handles the expansion of the Mir IR operators LoadVoxels and
 * LoadVoxelsWithCtl into Low IR code.
 *
 * 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 LoadVoxels : sig

  (* expand the assignment
   *
   *	lhs = LoadVoxels<info,s>(img, idx)
   *
   * into Low IR code.  If the image has type `image(d)[dd]`, then the resulting lhs tensor
   * has type `tensor[dd,s^d]`, where `s^d` means `d` occurrences of `s`.
   *)
    val expand : {
	    lhs : LowIR.var,
	    info : ImageInfo.info,
	    s : int,
	    img : LowIR.var,
	    idx : LowIR.var
	  } -> LowIR.assign list

  (* expand the assignment
   *
   *	lhs = LoadVoxelsWithCtl<info,s,ctl>(img, idx)
   *
   * into Low IR code.
   *)
    val expandWithCtl : {
	    lhs : LowIR.var,
	    info : ImageInfo.info,
	    s : int,
	    ctl : IndexCtl.ctl,
	    img : LowIR.var,
	    idx : LowIR.var
	  } -> LowIR.assign list

  end = struct

    structure IR = LowIR
    structure V = IR.Var
    structure Ty = LowTypes
    structure Op = LowOps

    fun assignOp (avail, pre, ty, opss, args) =
	  AvailRHS.addAssign(avail, pre, ty, IR.OP(opss, args))
    fun assignCons (_, _, []) = raise Fail "empty cons"
      | assignCons (avail, pre, args as x::_) = let
	  val ty = Ty.TensorTy(List.length args :: Ty.tensorShape(V.ty x))
	  in
	    AvailRHS.addAssign(avail, pre, ty, IR.CONS(args, ty))
	  end

    fun imul' (avail, pre, a, b) = assignOp (avail, pre, Ty.IntTy, Op.IMul, [a, b])
    fun iadd' (avail, pre, a, b) = assignOp (avail, pre, Ty.IntTy, Op.IAdd, [a, b])
    fun ilit' (avail, pre, n) =
	  AvailRHS.addAssign (avail, pre, Ty.IntTy, IR.LIT(Literal.Int(IntLit.fromInt n)))
    fun imul (avail, a, b) = imul' (avail, "mulRes", a, b)
    fun iadd (avail, a, b) = iadd' (avail, "addRes", a, b)
    fun ilit (avail, n) = ilit' (avail, "ilit", n)

  (* load the indices into individual variables; return the number of indices and
   * the list of variables.
   *)
    fun loadIndices (avail, idx) = (case V.ty idx
	   of Ty.IntTy => (1, [idx])
	    | ty as Ty.SeqTy(Ty.IntTy, SOME n) => let
		fun loadIdx i =
		      assignOp (avail, "idx", Ty.IntTy, Op.Subscript ty, [idx, ilit (avail, i)])
		in
		  (n, List.tabulate (n, loadIdx))
		end
	    | ty => raise Fail(concat["loadIndices: bogus type " ^ Ty.toString ty])
	  (* end case *))

  (* load the sizes of the axes into individual variables (except the slowest).
   * The result will be a list of variables holding the axis sizes ordered from
   * fastest to slowest.
   *)
    fun loadSizes (avail, info, img) = (case ImageInfo.sizes info
	   of [] => let (* no proxy, so we need load sizes from image object *)
		fun load i = assignOp (avail, "sz", Ty.IntTy, Op.ImageDim(info, i), [img])
		in
		  List.tabulate (ImageInfo.dim info - 1, load)
		end
	    | sz::szs => let
		fun load (_, [], xs) = List.rev xs (* drop last size, since it is the slowest *)
		  | load (sz, sz'::szs, xs) = load (sz', szs, ilit (avail, sz) :: xs)
		in
		  load (sz, szs, [])
		end
	  (* end case *))

  (* add code to compute the total offset from the image base address, which is
   *
   *	offp = stride * ix + offset
   *)
    fun adjustForStrideAndOffset (avail, 1, 0, ix) = ix
      | adjustForStrideAndOffset (avail, stride, 0, ix) =
	  imul' (avail, "offp", ilit' (avail, "stride", stride), ix)
      | adjustForStrideAndOffset (avail, 1, offset, ix) =
	  iadd' (avail, "offp", ix, ilit' (avail, "offset", offset))
      | adjustForStrideAndOffset (avail, stride, offset, ix) =
	  iadd' (avail, "offp",
	    imul (avail, ilit' (avail, "stride", stride), ix),
	    ilit' (avail, "offset", offset))

  (* compute the voxel address for a sample.  For the call
   *
   *    voxelAddress (avail, info, img) (offset, [i_1, ..., i_d])
   *
   * the address is given by
   *
   *    base + stride * (i_1 + N_1 * (i_2 + N_2 * (... + N_{d-1} * i_d) ...)) + offset
   *
   * where
   *    base    -- base address of the image data
   *    stride  -- number of samples per voxel
   *    offset  -- offset of sample being addressed
   *    N_i     -- size of ith axis in elements
   *
   * Note that we are following the Nrrd convention that the axes are ordered
   * in fastest to slowest order.  We are also assuming the C semantics of address
   * arithmetic, where the offset will be automatically scaled by the size of the
   * elements.
   *)
    fun voxelAddress (avail, info, img) = let
          val stride = ImageInfo.stride info
        (* get N_1 ... N_{d-1} *)
	  val sizes = loadSizes (avail, info, img)
	  fun voxelAddr (offset, []) = raise Fail "empty index list"
	    | voxelAddr (offset, [ix]) = let
		val offp = adjustForStrideAndOffset (avail, stride, offset, ilit (avail, ix))
		in
		  assignOp (avail, "baseAddr", Ty.AddrTy info, Op.VoxelAddress info, [img, offp])
		end
	    | voxelAddr (offset, ixs) = let
		fun gen ([sz], [ix]) = imul (avail, sz, ilit (avail, ix))
		  | gen (sz::szs, ix::ixs) =
		      imul (avail, sz, iadd (avail, ilit (avail, ix), gen (szs, ixs)))
		  | gen _ = raise Fail "voxelAddress: sizes/indices arity mismatch"
                val t = gen (sizes, ixs)
		val offp = adjustForStrideAndOffset (avail, stride, offset, t)
		in
		  assignOp (avail, "baseAddr", Ty.AddrTy info, Op.VoxelAddress info, [img, offp])
		end
	  in
	    voxelAddr
	  end

  (* iterate over an N-dimensional grid of s^n indices.  For each location,
   * the function f gets called with a list of integers [idx1, ..., idxn],
   * where idx1 varies the fastest and idxn varies the slowest.
   *)
    fun gridIterate (avail, n, s) (f : int list -> IR.var) = let
	  fun iter (0, idxs) = f idxs
	    | iter (i, idxs) = let
		fun lp (j, xs) = if (j < s)
		      then lp (j+1, iter(i-1, j::idxs) :: xs)
		      else assignCons (avail, "samp", List.rev xs)
		in
		  lp (0, [])
		end
	  in
	    iter (n, [])
	  end

  (* iterate over the shape of a voxel.  For each iteration, the sample offset
   * is passed into the function f, which returns a variable that has been
   * assigned to the corresponding tensor of samples.
   *)
    fun shapeIterate avail shp (f : int -> IR.var) = let
	  fun iter ([], offset) = (f offset, offset+1)
	    | iter (d::dd, offset) = let
		fun lp (i, offset, xs) = if (i < d)
		      then let
			val (x, offset) = iter (dd, offset)
			in
			  lp (i+1, offset, x::xs)
			end
		      else (assignCons (avail, "vox", List.rev xs), offset)
		in
		  lp (0, offset, [])
		end
	  val (y, _) = iter (shp, 0)
	  in
	    y
	  end

    fun expand {lhs, info, s, img, idx} = let
	  val avail = AvailRHS.new ()
	  val (dim, idxs) = loadIndices (avail, idx)
	  val voxelAddr = voxelAddress (avail, info, img)
	  fun load offset idxs = assignOp (
		avail, "vec", Ty.TensorTy[s],
		Op.LoadVoxels(info, s),
		[voxelAddr (offset, idxs)])
	  val result = shapeIterate avail (ImageInfo.voxelShape info)
		(fn offset => gridIterate (avail, dim-1, s) (load offset))
	  in
	    AvailRHS.addAssignToList (avail, (lhs, IR.VAR result));
	    List.rev (AvailRHS.getAssignments avail)
	  end

    fun expandWithCtl {lhs, info, s, ctl, img, idx} = let
	  val avail = AvailRHS.new ()
	  val (dim, idxs) = loadIndices (avail, idx)
	  val voxelAddr = voxelAddress (avail, info, img)
	  fun load offset idxs = assignOp (
		avail, "val", Ty.realTy,
		Op.LoadVoxels(info, 1),
		[voxelAddr (offset, idxs)])
	  val result = shapeIterate avail (ImageInfo.voxelShape info)
		(fn offset => gridIterate (avail, dim, s) (load offset))
	  in
	    AvailRHS.addAssignToList (avail, (lhs, IR.VAR result));
	    List.rev (AvailRHS.getAssignments avail)
	  end

  end

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