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

SCM Repository

[diderot] View of /branches/vis15/src/compiler/target-cpu/gen-outputs.sml
ViewVC logotype

View of /branches/vis15/src/compiler/target-cpu/gen-outputs.sml

Parent Directory Parent Directory | Revision Log Revision Log


Revision 3996 - (download) (annotate)
Sun Jun 19 15:21:12 2016 UTC (3 years ago) by jhr
File size: 17273 byte(s)
working on merge: code generation
(* gen-outputs.sml
 *
 * This code is part of the Diderot Project (http://diderot-language.cs.uchicago.edu)
 *
 * COPYRIGHT (c) 2016 The University of Chicago
 * All rights reserved.
 *
 * Generate strand output functions.  The output formats always have a single axis for the
 * data elements followed by one, or more, axes for the output structure.  There are four
 * cases that we handle:
 *
 *      grid, fixed-size elements:
 *              nrrd has object axis followed by grid axes
 *
 *      collection, fixed-size elements
 *              nrrd has object axis followed by a single axis
 *
 *      grid, dynamic-size elements
 *              nLengths nrrd has size 2 for objects (offset, length) followed by grid axes
 *              nData nrrd has object axis followed by a single axis
 *
 *      collection, dynamic-size elements
 *              nLengths nrrd has size 2 for objects (offset, length) followed by a single axis
 *              nData nrrd has object axis followed by a single axis
 *
 * The object axis kind depends on the output type, but it will either be one of the tensor types
 * that Teem knows about or else nrrdKindList.  In any case, the data elements are written as a
 * flat vector following the in-memory layout.  The other axes in the file will have nrrdKindSpace
 * as their kind.
 *
 * TODO: some of this code will be common across all targets (e.g., writing outputs to files), so
 * we will want to refactor it.
 *
 * TODO: for sequences of tensors (e.g., tensor[3][2]), we should use a separate axis for the
 * sequence dimension with kind nrrdKindList.
 *
 * TODO: since the runtime tracks numbers of strands in various states, we should be
 * able to use that information directly from the world without having to recompute it!
 *)

structure GenOutputs : sig

  (* gen (props, nAxes, outputs)
   *    returns a list of function declarations for getting the output/snapshot nrrds from
   *    the program state.  The arguments are:
   *        props       - the target information
   *        nAxes       - the number of axes in the grid of strands (NONE for a collection)
   *        outputs     - the list of output state variables paired with their API types
   *)
    val gen : CodeGenEnv.t * int option * OutputUtil.output_info list -> CLang.decl list

  end = struct

    structure IR = TreeIR
    structure V = TreeVar
    structure Ty = APITypes
    structure CL = CLang
    structure Nrrd = NrrdEnums
    structure U = GenOutputsUtil
    structure RN = CxxNames
    structure Env = CodeGenEnv

    fun mapi f l = let
          fun mapf (i, [], l) = List.rev l
            | mapf (i, x::xs, l) = mapf (i+1, xs, f(i, x)::l)
          in
            mapf (0, l, [])
          end

    val nrrdPtrTy = CL.T_Ptr(CL.T_Named "Nrrd")
    val sizeTy = CL.T_Named "size_t"
    val wrldPtr = RN.worldPtrTy
    fun mkInt i = CL.mkInt(IntInf.fromInt i)

  (* variables in the generated code *)
    val wrldV = CL.mkVar "wrld"
    val sizesV = CL.mkVar "sizes"
    val iV = CL.mkVar "i"
    val nV = CL.mkVar "n"
    val cpV = CL.mkVar "cp"
    val ipV = CL.mkVar "ip"
    val msgV = CL.mkVar "msg"
    val offsetV = CL.mkVar "offset"
    val nDataV = CL.mkVar "nData"
    val nLengthsV = CL.mkVar "nLengths"
    val numStableV = CL.mkVar "numStable"
    val numElemsV = CL.mkVar "numElems"
    val outSV = CL.mkVar "outS"
    val DIDEROT_DEAD = CL.mkVar "diderot::kDead"
    val DIDEROT_STABLE = CL.mkVar "diderot::kStable"
    val NRRD = CL.mkVar "NRRD"

  (* dymanic sequence operations *)
    fun seqLength arg = CL.mkApply("diderot::DynSeqLength", [arg])
    fun seqCopy (elemSz, dst, seq) = CL.mkApply("diderot::DynSeqCopy", [elemSz, dst, seq])

  (* utility functions for initializing the sizes array *)
    fun sizes i = CL.mkSubscript(sizesV, mkInt i)
    fun setSizes (i, v) = CL.mkAssign(sizes i, v)

  (* code to access state variable
        wrld->outState[i]->name
   * or
        wrld->state[i].name
   *)
    fun stateVar spec name = if TargetSpec.dualState spec
          then CL.mkIndirect(CL.mkSubscript(CL.mkIndirect(wrldV, "_outState"), iV), "sv_"^name)
          else CL.mkSelect(CL.mkSubscript(CL.mkIndirect(wrldV, "_state"), iV), "sv_"^name)

  (* code fragment to loop over strands
        for (unsigned int i = 0;  i < wrld->_nstrands;  i++) ...
   *)
    fun forStrands stm = CL.mkFor(
          CL.uint32, [("i", mkInt 0)],
          CL.mkBinOp(iV, CL.#<, CL.mkIndirect(wrldV, "_nstrands")),
          [CL.mkPostOp(iV, CL.^++)],
          stm)

  (* code fragment to test for stable strands in a loop
        if (wrld->_status[i] == diderot::kStable)
            ...
   *)
    fun ifStable stm = CL.mkIfThen(
          CL.mkBinOp(CL.mkSubscript(CL.mkIndirect(wrldV, "_status"), iV), CL.#==, DIDEROT_STABLE),
          stm)

  (* code fragment to test for active strands in a loop; note that NEW strands are considered active.
        if (wrld->_status[i] != DIDEROT_DIE)
            ...
   *)
    fun ifActive stm = CL.mkIfThen(
          CL.mkBinOp(CL.mkSubscript(CL.mkIndirect(wrldV, "_status"), iV), CL.#!=, DIDEROT_DEAD),
          stm)

  (* code fragment to initialize the axes kinds; the data axis (axis[0]) is given, but we skip it
   * (by convention) if it is scalar. The other axes are the specified domAxisKind.
   *)
    fun initAxisKinds (nrrd, dataAxisKind, nAxes, domAxisKind) = let
        (* nData->axis[0].kind *)
          fun axisKind i = CL.mkSelect(CL.mkSubscript(CL.mkIndirect(nrrd, "axis"), mkInt i), "kind")
          fun init (i, k) = CL.mkAssign (axisKind i, CL.mkVar(Nrrd.kindToEnum k))
          val (firstSpace, dataAxis) = (case dataAxisKind
                 of Nrrd.KindScalar => (0, [])
                  | _ => (1, [init(0, dataAxisKind)])
                (* end case *))
          in
            dataAxis @ List.tabulate(nAxes, fn i => init(i+firstSpace, domAxisKind))
          end

  (* create the body of an output function for dynamic-size outputs.  The structure of the
   * function body is:
   *
   *    declarations
   *    compute sizes array for nLengths
   *    allocate nrrd for nLengths
   *    compute sizes array for nData
   *    allocate nrrd for nData
   *    copy data from strands to nrrd
   *)
    fun genDynOutput (env, snapshot, nAxes, ty, name) = let
	  val spec = Env.target env
          val (elemCTy, nrrdType, axisKind, nElems) = OutputUtil.infoOf (env, ty)
          val stateVar = stateVar spec
          val (nAxes, domAxisKind) = (case nAxes
                 of NONE => (1, Nrrd.KindList)
                  | SOME n => (n, Nrrd.KindSpace)
                (* end case *))
        (* declarations *)
          val sizesDecl = CL.mkDecl(CL.T_Array(sizeTy, SOME(nAxes+1)), "sizes", NONE)
        (* count number of elements (and stable strands) *)
          val countElems = let
                val nElemsInit = CL.mkDeclInit(CL.uint32, "numElems", CL.mkInt 0)
                val cntElems = CL.S_Exp(CL.mkAssignOp(numElemsV, CL.+=, seqLength(stateVar name)))
                in
                  if #isGrid spec
                    then [
                        CL.mkComment["count number of elements"],
                        nElemsInit, forStrands cntElems
                      ]
                    else let
                      val cntBlk = CL.mkBlock[cntElems, CL.S_Exp(CL.mkPostOp(numStableV, CL.^++))]
                      val lpBody = if snapshot
                            then ifActive cntBlk
                            else ifStable cntBlk
                      in [
                        CL.mkComment["count number of output elements and stable strands"],
                        CL.mkDeclInit(CL.uint32, "numStable", CL.mkInt 0),
                        nElemsInit,
                        forStrands lpBody
                      ] end
                end
	(* code to check for zero outputs, which happens for collections with no active strands *)
	  val checkForNoStrands = if #isGrid spec
		then []
		else [
		    CL.mkComment["check for no output"],
		    CL.mkIfThen(
		      CL.mkBinOp(mkInt 0, CL.#==, numStableV),
		      CL.mkBlock[
			  CL.mkCall("nrrdEmpty", [nLengthsV]),
			  CL.mkCall("nrrdEmpty", [nDataV]),
			  CL.mkReturn(SOME(CL.mkVar "false"))
			])
		  ]
        (* generate code to allocate the nLengths nrrd *)
          val lengthsNrrd = let
                val dimSizes = setSizes(0, CL.mkInt 2)  (* nLengths is 2-element vector *)
                in
                  CL.mkComment["allocate nLengths nrrd"] ::
                  (if #isGrid spec
                    then dimSizes ::
                      List.tabulate (nAxes, fn i =>
                        setSizes(i+1, CL.mkSubscript(CL.mkIndirect(wrldV, "_size"), mkInt(nAxes-i-1)))) @
                      [U.maybeAlloc (env, nLengthsV, Nrrd.tyToEnum Nrrd.TypeInt, nAxes+1)]
                    else [
                        dimSizes, setSizes(1, numStableV),
                        U.maybeAlloc (env, nLengthsV, Nrrd.tyToEnum Nrrd.TypeInt, 2)
                      ])
                end
	(* code to check for no data to output (i.e., all of the output sequences are empty) *)
	  val checkForEmpty = [
		  CL.mkComment["check for empty output"],
		  CL.mkIfThen(
		    CL.mkBinOp(mkInt 0, CL.#==, numElemsV),
		    CL.mkBlock[
			CL.mkCall("nrrdEmpty", [nDataV]),
			CL.mkReturn(SOME(CL.mkVar "false"))
		      ])
		]
        (* generate code to allocate the data nrrd *)
          val dataNrrd = if (axisKind = Nrrd.KindScalar)
                then [ (* drop data axis for scalar data by convention *)
                    CL.mkComment["allocate nData nrrd"],
                    setSizes(0, numElemsV),
                    U.maybeAlloc (env, nDataV, Nrrd.tyToEnum nrrdType, 1)
                  ]
                else [
                    CL.mkComment["allocate nData nrrd"],
                    setSizes(0, mkInt nElems),
                    setSizes(1, numElemsV),
                    U.maybeAlloc (env, nDataV, Nrrd.tyToEnum nrrdType, 2)
                  ]
        (* generate the nLengths copy code *)
          val copyLengths = let
                val pInit = CL.mkDeclInit(CL.T_Ptr CL.uint32, "ip",
                      CL.mkReinterpretCast(CL.T_Ptr(CL.uint32), CL.mkIndirect(nLengthsV, "data")))
                val offsetDecl = CL.mkDeclInit(CL.uint32, "offset", CL.mkInt 0)
                val copyBlk = CL.mkBlock[
                        CL.mkDeclInit(CL.uint32, "n", seqLength(stateVar name)),
                        CL.mkAssign(CL.mkUnOp(CL.%*, CL.mkPostOp(ipV, CL.^++)), offsetV),
                        CL.mkAssign(CL.mkUnOp(CL.%*, CL.mkPostOp(ipV, CL.^++)), nV),
                        CL.S_Exp(CL.mkAssignOp(offsetV, CL.+=, nV))
                      ]
                val copyStm = if #isGrid spec
                      then copyBlk
                      else if #snapshot spec
                        then ifActive copyBlk
                        else ifStable copyBlk
                in
                  CL.mkComment["initialize nLengths nrrd"] ::
                  pInit ::
                  offsetDecl ::
                  forStrands copyStm ::
                  initAxisKinds (nLengthsV, Nrrd.Kind2Vector, nAxes, domAxisKind)
                end
        (* generate the nData copy code *)
          val copyData = let
                val pInit = CL.mkDeclInit(CL.charPtr, "cp",
                      CL.mkReinterpretCast(CL.charPtr, CL.mkIndirect(nDataV, "data")))
                val copyStm = CL.mkAssign(cpV, seqCopy(
                      CL.mkBinOp(mkInt nElems, CL.#*, CL.mkSizeof(elemCTy)), cpV, stateVar name))
                val copyStm = if #isGrid spec
                       then copyStm
                      else if #snapshot spec
                        then ifActive copyStm
                        else ifStable copyStm
                in
                  CL.mkComment["initialize nLengths nrrd"] ::
                  pInit ::
                  forStrands copyStm ::
                  initAxisKinds (nDataV, axisKind, 1, Nrrd.KindList)
                end
        (* the function body *)
          val stms =
                sizesDecl ::
                countElems @
		checkForNoStrands @
                lengthsNrrd @
		checkForEmpty @
                dataNrrd @
                copyLengths @
                copyData @
                [CL.mkReturn(SOME(CL.mkVar "false"))]
          in
            ([CL.PARAM([], nrrdPtrTy, "nLengths"), CL.PARAM([], nrrdPtrTy, "nData")], CL.mkBlock stms)
          end

  (* create the body of an output function for fixed-size outputs.  The structure of the
   * function body is:
   *
   *    declare and compute sizes array
   *    allocate nrrd nData
   *    copy data from strands to nrrd
   *)
    fun genFixedOutput (env, snapshot, nAxes, ty, name) = let
	  val spec = Env.target env
          val (elemCTy, nrrdType, axisKind, nElems) = OutputUtil.infoOf (env, ty)
          val stateVar = stateVar spec
          val (nAxes, domAxisKind) = (case nAxes
                 of NONE => (1, Nrrd.KindList)
                  | SOME n => (n, Nrrd.KindSpace)
                (* end case *))
          val nDataAxes = if (axisKind = Nrrd.KindScalar) then 0 else 1
        (* generate the sizes initialization code *)
          val initSizes = let
                val dimSizes = let
                      val dcl = CL.mkDecl(CL.T_Array(sizeTy, SOME(nAxes+nDataAxes)), "sizes", NONE)
                      in
                        if (axisKind = Nrrd.KindScalar)
                          then [dcl]
                          else [dcl, setSizes(0, mkInt nElems)]
                      end
                in
                  if #isGrid spec
                    then dimSizes @
                      List.tabulate (nAxes, fn i =>
                        setSizes(i+nDataAxes, CL.mkSubscript(CL.mkIndirect(wrldV, "_size"), mkInt(nAxes-i-1))))
                    else let
                      val cntStm = CL.S_Exp(CL.mkPostOp(numStableV, CL.^++))
                      val lpBody = if snapshot
                            then ifActive cntStm
                            else ifStable cntStm
                      in
                        CL.mkDeclInit(sizeTy, "numStable", mkInt 0) ::
                        forStrands lpBody ::
                        dimSizes @ [setSizes(nDataAxes, numStableV)]
                      end
                end
	(* code to check for no data to output (i.e., no active strands) *)
	  val checkForEmpty = if (#isGrid spec)
		then []
		else [
		    CL.mkComment["check for empty output"],
		    CL.mkIfThen(
		      CL.mkBinOp(mkInt 0, CL.#==, numStableV),
		      CL.mkBlock[
			  CL.mkCall("nrrdEmpty", [nDataV]),
			  CL.mkReturn(SOME(CL.mkVar "false"))
			])
		  ]
        (* generate the copy code *)
          val copyCode = let
                val pDecl = CL.mkDeclInit(CL.charPtr, "cp",
                      CL.mkReinterpretCast(CL.charPtr, CL.mkIndirect(nDataV, "data")))
                val copyBlk = CL.mkBlock[
                        CL.mkCall("memcpy", [
                            cpV,
                            CL.mkUnOp(CL.%&, stateVar name),
                            CL.mkBinOp(mkInt nElems, CL.#*, CL.mkSizeof elemCTy)
                          ]),
                        CL.mkExpStm(CL.mkAssignOp(cpV, CL.+=,
                          CL.mkBinOp(mkInt nElems, CL.#*, CL.mkSizeof elemCTy)))
                      ]
                val copyStm = if #isGrid spec
                       then copyBlk
                      else if snapshot
                        then ifActive copyBlk
                        else ifStable copyBlk
                in
                  pDecl :: forStrands copyStm :: initAxisKinds (nDataV, axisKind, nAxes, domAxisKind)
                end
        (* the function body *)
          val stms =
                CL.mkComment["Compute sizes of nrrd file"] ::
                initSizes @
		checkForEmpty @
                CL.mkComment["Allocate nData nrrd"] ::
                U.maybeAlloc (env, nDataV, Nrrd.tyToEnum  nrrdType, nAxes+nDataAxes) ::
                CL.mkComment["copy data to output nrrd"] ::
                copyCode @
                [CL.mkReturn(SOME(CL.mkVar "false"))]
          in
            ([CL.PARAM([], nrrdPtrTy, "nData")], CL.mkBlock stms)
          end

    fun gen (env, nAxes, outputs) = let
	  val spec = Env.target env
          fun getFn snapshot {name, ty, isGlobal} = let
                val funcName = if snapshot
                      then GenLibraryInterface.snapshotGet(spec, name)
                      else GenLibraryInterface.outputGet(spec, name)
                fun mkFunc (params, body) =
                      CL.D_Func(
			[], CL.boolTy, funcName, CL.PARAM([], RN.worldPtrTy, "wrld")::params, body)
                in
                  case ty
                   of Ty.SeqTy(ty', NONE) => mkFunc (genDynOutput(env, snapshot, nAxes, ty', name))
                    | _ => mkFunc (genFixedOutput(env, snapshot, nAxes, ty, name))
                  (* end case *)
                end
	  val getFns = List.map (getFn false) outputs
	  in
	    if (#exec spec)
	      then getFns @ U.genOutput(env, outputs)
	    else if (#snapshot spec)
	      then List.map (getFn true) outputs @ getFns
	      else getFns
	  end

  end

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