Fat JAR executable that generates C code from Rise programs

.gitignore

@@ -18,4 +18,5 @@ modules.xml
 *.gz
 *.sc
 
+float-safe-optimizer.jar
 *.o

README.md

@@ -26,3 +26,20 @@ The source code for the compiler is organised into sub-packages of the `shine` p
 
 ### Setup and Documentation
 Please have a look at: https://rise-lang.org/doc/
+
+### Float Safe Optimizer
+
+This repository contains an optimizer executable that preserves floating-point semantics.
+To build a Fat JAR executable:
+```sh
+sbt float_safe_optimizer/assembly
+```
+To optimize a Rise program and generate code:
+```sh
+java -Xss20m -Xms512m -Xmx4G -jar float-safe-optimizer.jar $function_name $rise_source_path $output_path
+```
+For example:
+```sh
+java -Xss20m -Xms512m -Xmx4G -jar float-safe-optimizer.jar add3 float-safe-optimizer/examples/add3Seq.rise float-safe-optimizer/examples/add3Seq.c
+java -Xss20m -Xms512m -Xmx4G -jar float-safe-optimizer.jar add3 float-safe-optimizer/examples/add3.rise float-safe-optimizer/examples/add3.c
+```

build.sbt

@@ -104,5 +104,25 @@ clap := {
   "echo y" #| (baseDirectory.value + "/lib/clap/buildClap.sh") !
 }
 
-
+lazy val float_safe_optimizer = (project in file("float-safe-optimizer"))
+  .dependsOn(riseAndShine)
+  .enablePlugins(AssemblyPlugin)
+  .settings(
+    excludeDependencies ++= Seq(
+      ExclusionRule("org.scala-lang.modules", s"scala-xml_${scalaBinaryVersion.value}"),
+      ExclusionRule("junit", "junit"),
+      ExclusionRule("com.novocode", "junit-interface"),
+      ExclusionRule("org.scalacheck", "scalacheck"),
+      ExclusionRule("org.scalatest", "scalatest"),
+      ExclusionRule("com.lihaoyi", s"os-lib_${scalaBinaryVersion.value}"),
+      ExclusionRule("com.typesafe.play", s"play-json_${scalaBinaryVersion.value}"),
+      ExclusionRule("org.rise-lang", s"opencl-executor_${scalaBinaryVersion.value}"),
+      ExclusionRule("org.rise-lang", "CUexecutor"),
+      ExclusionRule("org.elevate-lang", s"cuda-executor_${scalaBinaryVersion.value}"),
+      ExclusionRule("org.elevate-lang", s"meta_${scalaBinaryVersion.value}"),
+    ),
+    name := "float-safe-optimizer",
+    javaOptions ++= Seq("-Xss20m", "-Xms512m", "-Xmx4G"),
+    assemblyOutputPath in assembly := file("float-safe-optimizer.jar"),
+  )
 

float-safe-optimizer/Main.scala

@@ -0,0 +1,43 @@
+package float_safe_optimizer
+
+import util.gen
+import rise.core.Expr
+import rise.core.DSL.ToBeTyped
+
+object Main {
+  def main(args: Array[String]): Unit = {
+    val name = args(0)
+    val exprSourcePath = args(1)
+    val outputPath = args(2)
+
+    val exprSource = util.readFile(exprSourcePath)
+    val untypedExpr = parseExpr(prefixImports(exprSource))
+    val typedExpr = untypedExpr.toExpr
+    println("typedExpr", typedExpr)
+    val optimizedExpr = Optimize(typedExpr)
+    println("optimizedExpr", optimizedExpr)
+    val code = gen.openmp.function.asStringFromExpr(optimizedExpr)
+    util.writeToPath(outputPath, code)
+  }
+
+  def prefixImports(source: String): String =
+    s"""
+       |import rise.core.DSL._
+       |import rise.core.DSL.Type._
+       |import rise.core.DSL.HighLevelConstructs._
+       |import rise.core.primitives._
+       |import rise.core.types._
+       |import rise.core.types.DataType._
+       |import rise.openmp.DSL._
+       |import rise.openmp.primitives._
+       |$source
+       |""".stripMargin
+
+  def parseExpr(source: String): ToBeTyped[Expr] = {
+    import scala.reflect.runtime.universe
+    import scala.tools.reflect.ToolBox
+
+    val toolbox = universe.runtimeMirror(getClass.getClassLoader).mkToolBox()
+    toolbox.eval(toolbox.parse(source)).asInstanceOf[ToBeTyped[Expr]]
+  }
+}

float-safe-optimizer/Optimize.scala

@@ -0,0 +1,100 @@
+package float_safe_optimizer
+
+import rise.eqsat._
+
+object Optimize {
+  def apply(e: rise.core.Expr): rise.core.Expr = {
+    val expr = Expr.fromNamed(e)
+    val (body, annotation, wrapBody) = analyseTopLevel(expr)
+    println("eqsat expr", expr)
+
+    val rules = {
+      import rise.eqsat.rules._
+      Seq(
+        // implementation choices:
+        reduceSeq,
+        mapSeq,
+        // satisfying read/write annotations:
+        mapSeqArray,
+        // simplifications:
+        // mapFusion,
+        // reduceSeqMapFusion,
+        removeTransposePair,
+        fstReduction,
+        sndReduction,
+        reduceSeqOne,
+      /* maybe:
+        omp.mapPar --> need heuristic vs mapSeq
+        toMemAfterMapSeq / storeToMem
+        storeToMem
+        reduceSeqMapFusion
+        mapSeqUnroll/reduceSeqUnroll --> need heuristic
+        eliminateMapIdentity
+
+        is it worth the cost?:
+        betaExtract
+        betaNatExtract
+        eta
+
+        not generic enough, use Elevate passes or custom applier?:
+        idxReduction_i_n
+      */
+      )
+    }
+
+    LoweringSearch.init().run(BENF, Cost, Seq(body), rules, Some(annotation)) match {
+      case Some(resBody) =>
+        val res = wrapBody(resBody)
+        Expr.toNamed(res)
+      case None => throw new Exception("could not find valid low-level expression")
+    }
+  }
+
+  // TODO: this code might be avoidable by making DPIA+codegen rely on top-level type instead of top-level constructs
+  def analyseTopLevel(e: Expr)
+  : (Expr, (BeamExtractRW.TypeAnnotation, Map[Int, BeamExtractRW.TypeAnnotation]), Expr => Expr) = {
+    import rise.eqsat.RWAnnotationDSL._
+
+    // returns (body, argCount, wrapBody)
+    def rec(e: Expr): (Expr, Int, Expr => Expr) = {
+      e.node match {
+        case NatLambda(e2) =>
+          val (b, a, w) = rec(e2)
+          (b, a, b => Expr(NatLambda(w(b)), e.t))
+        case DataLambda(e2) =>
+          val (b, a, w) = rec(e2)
+          (b, a, b => Expr(DataLambda(w(b)), e.t))
+        case Lambda(e2) =>
+          e.t.node match {
+            case FunType(Type(_: DataTypeNode[_, _]), _) => ()
+            case _ => throw new Exception("top level higher-order functions are not supported")
+          }
+          val (b, a, w) = rec(e2)
+          (b, a + 1, b => Expr(Lambda(w(b)), e.t))
+        case _ =>
+          if (!e.t.node.isInstanceOf[DataTypeNode[_, _]]) {
+            throw new Exception("expected body with data type")
+          }
+          (e, 0, b => b)
+      }
+    }
+
+    val (b, a, w) = rec(e)
+    (b, (write, List.tabulate(a)(i => i -> read).toMap), w)
+  }
+
+  object Cost extends CostFunction[Int] {
+    val ordering = implicitly
+
+    override def cost(egraph: EGraph, enode: ENode, t: TypeId, costs: EClassId => Int): Int = {
+      import rise.core.primitives._
+
+      val nodeCost = enode match {
+        // prefer avoiding mapSeq
+        case Primitive(mapSeq()) => 5
+        case _ => 1
+      }
+      enode.children().foldLeft(nodeCost) { case (acc, eclass) => acc + costs(eclass) }
+    }
+  }
+}

float-safe-optimizer/examples/.gitignore

@@ -0,0 +1 @@
+*.c

float-safe-optimizer/examples/add3.rise

@@ -0,0 +1,3 @@
+depFun((n: Nat) => fun((n`.`i32) ->: (n`.`i32))(in =>
+  map(add(li32(3)))(in)
+))

float-safe-optimizer/examples/add3Seq.rise

@@ -0,0 +1,3 @@
+depFun((n: Nat) => fun((n`.`i32) ->: (n`.`i32))(in =>
+  mapSeq(add(li32(3)))(in)
+))

float-safe-optimizer/examples/add3TypeError.rise

@@ -0,0 +1,3 @@
+depFun((n: Nat) => fun((n`.`i32) ->: (n`.`f32))(in =>
+  map(add(li32(3)))(in)
+))

float-safe-optimizer/examples/addv.rise

@@ -0,0 +1,4 @@
+depFun((n: Nat) => depFun((m: Nat) => depFun((o: Nat) =>
+  fun(((n+o)`.`i32) ->: ((m+o)`.`i32) ->: (o`.`i32))((a, b) =>
+  zip(take(o)(a))(take(o)(b)) |> map(fun(x => fst(x) + snd(x)))
+))))

float-safe-optimizer/examples/cos.rise

@@ -0,0 +1,3 @@
+depFun((n: Nat) => fun(n`.`f64)(x =>
+  x |> mapSeq(foreignFun("cos", f64 ->: f64))
+))

float-safe-optimizer/examples/generate.rise

@@ -0,0 +1,3 @@
+depFun((n: Nat) => fun(i32 ->: i32 ->: 2`.`i32)((x, y) =>
+  generate(fun(i => select(i =:= lidx(0, 2))(x)(y))) |> mapSeq(fun(x => x))
+))

float-safe-optimizer/examples/larr.rise

@@ -0,0 +1,7 @@
+import rise.core.semantics._
+
+def larrToMem(s: Seq[Data]) = larr(s) |> mapSeqUnroll(fun(x => x)) |> toMem
+
+depFun((n: Nat) =>
+  larrToMem(Seq(IntData(1), IntData(2), IntData(3)))
+)

float-safe-optimizer/examples/length-rw.rise

@@ -0,0 +1,5 @@
+def length() = reduceSeq(fun(acc => fun(x => acc + li32(1))))(li32(0))
+
+depFun((n: Nat) => fun(n`.`i32)(x =>
+    x |> mapSeq(fun(y => y + length()(x)))
+))

float-safe-optimizer/examples/length.rise

@@ -0,0 +1,5 @@
+def length() = impl { n: Nat => impl{ dt: DataType => fun(n`.`dt)(x => l(n)) } }
+
+depFun((n: Nat) => fun(n`.`i32)(x =>
+    x |> mapSeq(fun(y => y + (cast(length()(x)) :: i32)))
+))

float-safe-optimizer/examples/makeArray.rise

@@ -0,0 +1,3 @@
+depFun((n: Nat) => fun(i32 ->: i32 ->: 2`.`i32)((x, y) =>
+  makeArray(2)(x)(y) |> mapSeqUnroll(fun(x => x))
+))

float-safe-optimizer/examples/scalarOut.rise

@@ -0,0 +1,3 @@
+depFun((n: Nat) => fun(n`.`f64)(x =>
+  x |> reduceSeq(add)(lf64(0.0))
+))

project/plugins.sbt

@@ -1,3 +1,5 @@
 addSbtPlugin("ch.epfl.scala" % "sbt-bloop"    % "1.4.0-RC1")
 addSbtPlugin("org.scalameta" % "sbt-scalafmt" % "2.3.0")
 addSbtPlugin("org.scalameta" % "sbt-mdoc" % "2.2.17")
+addSbtPlugin("com.eed3si9n" % "sbt-assembly" % "0.15.0")
+addDependencyTreePlugin