kernel fusion wip

src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala

@@ -0,0 +1,123 @@
+package lms.transformation.tensor
+
+import scala.annotation.implicitNotFound
+import scala.collection._
+
+import lms.core._
+import lms.core.stub._
+import lms.collection.mutable._
+import lms.macros.SourceContext
+import lms.thirdparty.array_computation.{ArrayCPUOps, CUDATypeLess, CudaOps}
+
+import Backend._
+
+object FusedTensorTypeLess {
+  import BaseTypeLess._
+  import PrimitiveTypeLess._
+  import ArrayTypeLess._
+  import CUDATypeLess._
+
+  type E = Backend.Exp
+  def C(a: Any) = Backend.Const(a)
+
+  /// typeless frontend
+  def TENSOR(size: Int, array: ARRAY)(f: Backend.Exp => Backend.Exp)(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectRead("tensor", C(size), array.x, Adapter.g.reify(xn => f(xn)))(array.x))).withSrcType(__pos, array.et)
+  }
+
+  def ZEROS(size: Int, array: ARRAY)(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectUnsafe("tensor_zeros", C(size), array.x))).withSrcType(__pos, array.et)
+  }
+
+  def ONES(size: Int, array: ARRAY)(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectUnsafe("tensor_ones", C(size), array.x))).withSrcType(__pos, array.et)
+  }
+
+  class TENSOR(override val x: Backend.Exp, override val useOldMetadata: Boolean = false) extends TOP(x) {
+    def withEleType(m: Manifest[_]): this.type = { Adapter.typeMap(x) = m; this }
+    override def withSrcType(pos: SourceContext, m: Manifest[_]): this.type =
+      withSource(pos).withEleType(m)
+
+    def et: Manifest[_] = {
+      if (useOldMetadata) Adapter.oldTypeMap(x) else Adapter.typeMap(x)
+    }
+
+    def size: Int = {
+      gc.get(x.asInstanceOf[Backend.Sym]) match {
+        case Some(Node(_, s, Backend.Const(size:Int)::_, _)) => size
+        case a => System.out.println(a); ???
+      }
+    }
+
+    def show(implicit __pos: SourceContext): UNIT = {
+      UNIT(Adapter.g.reflectEffect("show_tensor", x)()(Adapter.CTRL))
+    }
+
+    def + (y: TENSOR)(implicit __pos: SourceContext): TENSOR = {
+      (new TENSOR(Adapter.g.reflectUnsafe("tensor_add", x, y.x))).withSrcType(__pos, et)
+    }
+
+    def - (y: TENSOR)(implicit __pos: SourceContext): TENSOR = {
+      (new TENSOR(Adapter.g.reflect("tensor_minus", x, y.x))).withSrcType(__pos, et)
+    }
+
+    def apply(e: Backend.Exp)(implicit __pos: SourceContext): INT = {
+      // todo: change to correct effect
+      INT(Adapter.g.reflectEffect("tensor_apply", x, e)()(Adapter.CTRL)).withSrcType(__pos, et)
+    }
+  }
+}
+
+
+trait FusedTensorOps extends Dsl with ArrayOps with CudaOps {
+
+  import PrimitiveTypeLess._
+  import ArrayTypeLess._
+  import FusedTensorTypeLess._
+
+  // def NewArray[T:Manifest](x: Rep[Int])(implicit __pos: SourceContext): Rep[Array[T]] = {
+  //   Wrap[Array[T]](ARRAY(new INT(Unwrap(x)), manifest[T]).x)
+  // }
+
+  /// Typed Frontend
+  class Tensor[+T]
+  object Tensor {
+    def zeros[T:Manifest](size: Int, array: Rep[Array[T]])(implicit __pos: SourceContext): Rep[Tensor[T]] = {
+      val tensor = ZEROS(size, new ARRAY(Unwrap(array)))
+      Wrap[Tensor[T]](tensor.x)
+    }
+
+    def ones[T:Manifest](size: Int, array: Rep[Array[T]])(implicit __pos: SourceContext): Rep[Tensor[T]] = {
+      val tensor = ONES(size, new ARRAY(Unwrap(array)))
+      Wrap[Tensor[T]](tensor.x)
+    }
+
+    def apply[T:Numeric:Manifest](size: Int, array: Rep[Array[T]], f: Rep[Int] => Rep[Int])(implicit __pos: SourceContext): Rep[Tensor[T]] = {
+      Wrap[Tensor[T]](TENSOR(size, new ARRAY(Unwrap(array)))(unwrapFun[Int, Int](f)).x)
+    }
+  }
+
+  def tensor[T:Numeric:Manifest](x: Rep[Tensor[T]]): TENSOR = new TENSOR(Unwrap(x))
+
+  implicit class TensorOps[T:Numeric:Manifest](x: Rep[Tensor[T]]) {
+    val self = tensor(x)
+
+    def show(implicit __pos: SourceContext): Rep[Unit] = Wrap[Unit](self.show.x)
+
+    def + (y: Rep[Tensor[T]])(implicit  __pos: SourceContext): Rep[Tensor[T]] = {
+      val t = self + tensor(y)
+      Wrap[Tensor[T]](t.x)
+    }
+
+    def - (y: Rep[Tensor[T]])(implicit  __pos: SourceContext): Rep[Tensor[T]] = {
+      val t = self - tensor(y)
+      Wrap[Tensor[T]](t.x)
+    }
+
+    def apply(y: Rep[Int])(implicit  __pos: SourceContext): Rep[T] = {
+      val t = self.apply(Unwrap(y))
+      Wrap[T](t.x)
+    }
+
+  }
+}

src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala

@@ -0,0 +1,66 @@
+package lms.transformation.tensor
+
+import scala.annotation.implicitNotFound
+import scala.collection._
+
+import lms.core._
+import lms.core.stub._
+import lms.collection.mutable._
+import lms.macros.SourceContext
+import lms.thirdparty.array_computation.{ArrayCPUTypeLess, CUDATypeLess, CUBLASTypeLess}
+
+import Backend._
+
+abstract class FusedTensorLowering extends Transformer {
+
+  import BaseTypeLess._
+  import PrimitiveTypeLess._
+  import ArrayTypeLess._
+  import ArrayCPUTypeLess._
+  import FusedTensorTypeLess._
+
+  val tensors = new mutable.HashMap[Backend.Sym, (Node, List[Backend.Sym], Seq[Node])]
+
+  override def transform(n: Node): Backend.Exp = n match {
+    case Node(s, "tensor_zeros", (Backend.Const(sz:Int))::(arr:Backend.Exp)::_, _) =>
+      System.out.println(n)
+      implicit val pos = Adapter.oldSourceMap(s)
+      val array = new ARRAY(arr)
+      val t = TENSOR(sz, array)(i => INT(0).x)
+      t.x
+    case Node(s, "tensor_ones", (Backend.Const(sz:Int))::(arr:Backend.Exp)::_, _) =>
+      System.out.println(n)
+      implicit val pos = Adapter.oldSourceMap(s)
+      val array = new ARRAY(arr)
+      val t = TENSOR(sz, array)(i => INT(1).x)
+      t.x
+    case Node(s, "tensor_add", (x:Backend.Sym)::(y:Backend.Sym)::_, _) =>
+      System.out.println(n)
+      implicit val pos = Adapter.oldSourceMap(s)
+      val a = new TENSOR(x, useOldMetadata = true)
+      val b = new TENSOR(y, useOldMetadata = true)
+      val sz = a.size
+      val array = ARRAY(10, manifest[Int])
+      // val t = TENSOR(sz, array)(i => INT(a.apply(i)+ b.apply(i)).x)
+      // val t = TENSOR(sz, array)(i => INT(-1).x)
+      val x1 = a.apply(INT(0).x)
+      val x2 = b.apply(INT(0).x)
+      val r = x1 + x2
+      System.out.println(r)
+      val t = TENSOR(sz, array)(i => r.x)
+      System.out.println(t)
+      t.x
+    case _ => super.transform(n)
+  }
+
+  override def transform(graph: Graph): Graph = {
+    assert (g == null)
+    g = new GraphBuilderOpt()
+    Adapter.g = g
+    try {
+      super.transform(graph)
+    } finally {
+      g = null; Adapter.g = null
+    }
+  }
+}

src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala

@@ -0,0 +1,52 @@
+package lms.transformation.tensor
+
+import scala.annotation.implicitNotFound
+import scala.collection._
+
+import lms.core._
+import lms.core.stub._
+import lms.collection.mutable._
+import lms.macros.SourceContext
+import lms.thirdparty.array_computation.{ArrayCPUTypeLess, CUDATypeLess, CUBLASTypeLess}
+
+import Backend._
+
+abstract class FusedTensorVertical extends Transformer {
+
+  import BaseTypeLess._
+  import PrimitiveTypeLess._
+  import ArrayTypeLess._
+  import ArrayCPUTypeLess._
+  import FusedTensorTypeLess._
+
+  val tensors = new mutable.HashMap[Backend.Sym, (Node, List[Backend.Sym], Seq[Node])]
+
+  override def transform(n: Node): Backend.Exp = n match {
+    case Node(s, "tensor", Backend.Const(size:Int)::(x:Backend.Sym)::_, _) =>
+      System.out.println(n)
+      tensors(s) = (n, path, inner)
+      super.transform(n)
+    case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) =>
+      System.out.println(n)
+      val (Node(_, _, Backend.Const(szy:Int)::(_)::(f@Backend.Block(arg::Nil, r, block, eff))::_, _), path0, inner0) = tensors(a)
+      try {
+        subst(arg) = transform(b)
+        withResetScope(path0, inner0) {
+          traverse(f)
+        }
+        transform(r)
+      } finally subst -= arg
+    case _ => super.transform(n)
+  }
+
+  override def transform(graph: Graph): Graph = {
+    assert (g == null)
+    g = new GraphBuilderOpt()
+    Adapter.g = g
+    try {
+      super.transform(graph)
+    } finally {
+      g = null; Adapter.g = null
+    }
+  }
+}

src/test/scala/lms/transformation/test_fused_tensor.scala

@@ -0,0 +1,82 @@
+package lms
+package transformation.tensor
+
+import scala.annotation.implicitNotFound
+import lms.core.virtualize
+import macros.SourceContext
+
+import lms.core._
+import lms.core.stub._
+import lms.thirdparty.{CCodeGenLibs}
+import lms.thirdparty.array_computation.{CCodeGenCBLASOps, CCodeGenCudaOps}
+
+import Backend._
+
+class FixedSizeFusedTensorTest extends TutorialFunSuite {
+  val under = "transformer/fused_tensor"
+
+  abstract class CompilerCFusedTensor[A: Manifest, B: Manifest] extends CompilerC[A,B] with FusedTensorOps { q =>
+
+    override val codegen = new DslGenC with CCodeGenLibs {
+      val IR: q.type = q
+    }
+
+    override val passes = List(
+      new FusedTensorLowering {},
+      new FusedTensorVertical {}
+    )
+
+    /*
+    var log_path: String = ""
+    def setLogPath(path: String) { log_path = path }
+
+    override def transform(graph: Graph): Graph = {
+      logGraph(show_graph(graph), log_path)
+      super.transform(graph)
+    }
+
+    override def transformOnePass(pass: Transformer, index: Int, graph: Graph) = {
+      val new_graph = pass.transform(graph)
+      if (log_path == "") throw new Exception("should set log_path first")
+      logGraph(show_graph(new_graph), log_path, index, pass.name)
+      new_graph
+    }
+
+    def show_graph(graph: Graph): String = {
+      // return a string representation of the graph
+      val source = new java.io.ByteArrayOutputStream()
+      val stream = new java.io.PrintStream(source)
+      stream.println("==================")
+      for (node <- graph.nodes)
+        node.toString
+      stream.println(graph.block)
+      stream.println("==================")
+      source.toString
+    }*/
+  }
+
+  test("show") {
+    val driver = new CompilerCFusedTensor[Int, Unit] {
+      import FusedTensorTypeLess._
+
+      @virtualize
+      def snippet(arg: Rep[Int]): Rep[Unit] = {
+        val array = NewArray[Int](10)
+        // val a = Tensor(6, array, i => 0)
+        val a = Tensor.zeros[Int](10, array)
+        val b = Tensor.ones[Int](10, array)
+        val c = a + b
+        // c.show; ()
+        printf("%d", c(0))
+        // a.show; ()
+        // val b = Tensor(6, array, i => 1)
+
+        // val c = Tensor.zeros[Int](10, array)
+        // printf("%d", a(1) + b(1))
+        //val c = Tensor(6, array, i => a(i) + b(i))
+        //c.show; ()
+      }
+    }
+    System.out.println(indent(driver.code))
+  }
+}