Merge branch 'zdh/staged-eval' of github.com:Generative-Program-Analysis/GenSym into zdh/staged-eval

Author: butterunderflow

src/main/scala/wasm/StagedMiniWasm.scala

@@ -7,7 +7,7 @@ import lms.core.virtualize
 import lms.macros.SourceContext
 import lms.core.stub.{Base, ScalaGenBase, CGenBase}
 import lms.core.Backend._
-import lms.core.Backend.{Block => LMSBlock}
+import lms.core.Backend.{Block => LMSBlock, Const => LMSConst}
 import lms.core.Graph
 
 import gensym.wasm.ast._
@@ -28,6 +28,14 @@ trait StagedWasmEvaluator extends SAIOps {
   // a cache storing the compiled code for each function, to reduce re-compilation
   val compileCache = new HashMap[Int, Rep[(Cont[Unit]) => Unit]]
 
+  def funHere[A:Manifest,B:Manifest](f: Rep[A] => Rep[B], dummy: Rep[Unit] = "dummy".reflectCtrlWith[Unit]()): Rep[A => B] = {
+    // to avoid LMS lifting a function, we create a dummy node and read it inside function
+    fun((x: Rep[A]) => {
+      "dummy-op".reflectCtrlWith[Unit](dummy)
+      f(x)
+    })
+  }
+
   // NOTE: We don't support Ans type polymorphism yet
   def eval(insts: List[Instr],
            kont: Rep[Cont[Unit]],
@@ -88,37 +96,31 @@ trait StagedWasmEvaluator extends SAIOps {
         // no need to modify the stack when entering a block
         // the type system guarantees that we will never take more than the input size from the stack
         val funcTy = ty.funcType
-        val dummy = "dummy".reflectCtrlWith[Unit]()
         // TODO: somehow the type of exitSize in residual program is nothing
-        def restK: Rep[Cont[Unit]] = fun((_: Rep[Unit]) => {
+        def restK: Rep[Cont[Unit]] = funHere((_: Rep[Unit]) => {
           info(s"Exiting the block, stackSize =", Stack.size)
-          "dummy-op".reflectCtrlWith[Unit](dummy)
           eval(rest, kont, trail)
         })
         eval(inner, restK, restK :: trail)
       case Loop(ty, inner) =>
         val funcTy = ty.funcType
         val exitSize = Stack.size - funcTy.inps.size + funcTy.out.size
-        val dummy = "dummy".reflectCtrlWith[Unit]()
-        def restK = fun((_: Rep[Unit]) => {
-          "dummy-op".reflectCtrlWith[Unit](dummy)
+        def restK = funHere((_: Rep[Unit]) => {
           info(s"Exiting the loop, stackSize =", Stack.size)
           eval(rest, kont, trail)
         })
-        def loop : Rep[Unit => Unit] = fun((_u: Rep[Unit]) => {
-          "dummy-op".reflectCtrlWith[Unit](dummy)
+        val dummy = "dummy".reflectCtrlWith[Unit]()
+        def loop : Rep[Unit => Unit] = funHere((_u: Rep[Unit]) => {
           info(s"Entered the loop, stackSize =", Stack.size)
           eval(inner, restK, loop :: trail)
-        })
+        }, dummy) // <-- if we don't pass this dummy argument, lots of code will be generated
         loop(())
       case If(ty, thn, els) =>
         val funcTy = ty.funcType
         val exitSize = Stack.size - funcTy.inps.size + funcTy.out.size
         val cond = Stack.pop()
-        val dummy = "dummy".reflectCtrlWith[Unit]()
         // TODO: can we avoid code duplication here?
-        def restK = fun((_: Rep[Unit]) => {
-          "dummy-op".reflectCtrlWith[Unit](dummy)
+        def restK = funHere((_: Rep[Unit]) => {
           info(s"Exiting the if, stackSize =", Stack.size)
           eval(rest, kont, trail)
         })
@@ -694,6 +696,11 @@ object WasmToScalaCompiler {
 }
 
 trait StagedWasmCppGen extends CGenBase with CppSAICodeGenBase {
+  override def mayInline(n: Node): Boolean = n match {
+    case Node(s, "stack-pop", _, _) => false
+    case _ => super.mayInline(n)
+  }
+
   override def remap(m: Manifest[_]): String = {
     if (m.toString.endsWith("Num")) "Num"
     else if (m.toString.endsWith("Slice")) "Slice"
@@ -727,8 +734,11 @@ trait StagedWasmCppGen extends CGenBase with CppSAICodeGenBase {
       emit("Frames.set("); shallow(i); emit(", "); shallow(value); emit(");\n")
     case Node(_, "global-set", List(i, value), _) =>
       emit("Global.globalSet("); shallow(i); emit(", "); shallow(value); emit(");\n")
+    // Note: The following code is copied from the traverse of CppBackend.scala, try to avoid duplicated code
+    case n @ Node(f, "λ", (b: LMSBlock)::LMSConst(0)::rest, _) =>
+      // TODO: Is a leading block followed by 0 a hint for top function?
+      super.traverse(n)
     case n @ Node(f, "λ", (b: LMSBlock)::rest, _) =>
-      // Node: This code is copied from the traverse of CppSAICodeGenBase.scala, try to avoid code duplication
       val retType = remap(typeBlockRes(b.res))
       val argTypes = b.in.map(a => remap(typeMap(a))).mkString(", ")
       emitln(s"std::function<$retType(${argTypes})> ${quote(f)};")

src/test/scala/genwasym/TestStagedEval.scala

@@ -28,14 +28,14 @@ class TestStagedEval extends FunSuite {
     val moduleInst = ModuleInstance(Parser.parseFile(filename))
     val code = WasmToCppCompiler.compile(moduleInst, main, true)
     if (printRes) {
-      val writer = new java.io.PrintWriter(new java.io.File(s"$filename.cpp"))
-      try {
-        writer.write(code)
-      } finally {
-        writer.close()
-      }
+      println(code)
+    }
+    val writer = new java.io.PrintWriter(new java.io.File(s"$filename.cpp"))
+    try {
+      writer.write(code)
+    } finally {
+      writer.close()
     }
-    println(code)
   }
 
   test("ack-cpp") { testFileToCpp("./benchmarks/wasm/ack.wat", Some("real_main"), printRes = true) }