kernel fusion wip

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

@@ -0,0 +1,125 @@
+package lms.transformation.tensor
+
+import scala.annotation.implicitNotFound
+import scala.collection._
+import scala.collection.mutable.HashMap
+import scala.collection.immutable.Set
+
+import lms.core._
+import lms.core.stub._
+import lms.collection.mutable._
+import lms.macros.SourceContext
+import lms.thirdparty.array_computation.{ArrayCPUTypeLess, CUDATypeLess, CUBLASTypeLess}
+import lms.thirdparty.{CLibTypeLess}
+
+import Backend._
+
+abstract class FusedTensorToCuda extends Transformer {
+
+  override val name = "FusedTensorToCuda"
+
+  import BaseTypeLess._
+  import PrimitiveTypeLess._
+  import ArrayTypeLess._
+  import FusedTensorTypeLess._
+  import PrimitiveTypeLess._
+  import CUDATypeLess._
+  import CLibTypeLess._
+
+  def gpu_array(size: Int, m: Manifest[_], device: INT)(implicit __pos: SourceContext): ARRAY = {
+    CUDA_SET_DEVICE(device)
+    CUDA_MALLOC(size, m)
+  }
+
+  def ScanFile(scan: ARRAY, count: INT, filenameFormat: Rep[String], filenameArgs: Rep[Any]*)(implicit pos: SourceContext) = {
+    val function = scan.et match {
+      case m if m == manifest[Float] => "scan_float_array"
+      case m if m == manifest[Int] => "scan_int_array"
+      case m => throw new Exception(s"not yet supporting manifest ${m}")
+    }
+    LIB_FUNCTION(manifest[Unit], function, scan.x::count.x::Unwrap(filenameFormat)::filenameArgs.map(Unwrap).toList:_*)(Seq[Int](), Seq[Int](0,1), Set[Int]())
+  }
+
+    // map concrete (input) tensors to CUDA arrays
+  val tensor2arr = new mutable.HashMap[Backend.Sym, Backend.Exp]
+
+  override def transform(n: Node): Backend.Exp = n match {
+    case Node(s, "tensor_input", Backend.Const(inputs:Seq[View])::_, _) =>
+      implicit val __pos = Adapter.oldSourceMap(s)
+      val t = new TENSOR(s, useOldMetadata = true)
+      val size = t.size
+      val m = manifest[Int]
+      val cpuArray = ARRAY(size, m)
+      ScanFile(cpuArray, size, unit("input"))
+      CUDA_SET_DEVICE(INT(0))
+      val gpuArray = CUDA_MALLOC(size, m) // allocate CUDA array for input tensors
+      CUDA_MEMCPY(gpuArray, cpuArray, size, HOST2DEVICE, m)
+
+      tensor2arr(s) = gpuArray.x
+      gpuArray.x
+
+    case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) if tensor2arr.contains(a) =>
+      implicit val __pos = Adapter.oldSourceMap(s)
+      val arr = new ARRAY(tensor2arr(a))
+      (arr.apply(INT(transform(b)))).x // change tensor apply to array apply
+
+    case Node(s, "tensor", Backend.Const(inputs:Seq[View])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _) =>
+      implicit val __pos = Adapter.oldSourceMap(s)
+      val t = new TENSOR(s, useOldMetadata = true)
+      val sz1 = t.size
+
+      // input array. for now, assume only one input
+      val in_arr = new ARRAY(tensor2arr(inputs.head.t))
+      // System.out.println("in_arr:" + in_arr)
+      // allocate output array. assume only one output
+      val out_arr = CUDA_MALLOC(sz1, manifest[Int])
+      // System.out.println("in_arr:" + in_arr)
+      // System.out.println("out_arr:" + out_arr)
+
+      val kernel = CUDA_KERNEL3({ xn: List[Backend.Exp] => 
+        val in_array = (new ARRAY(xn(0))).withSrcType(__pos, manifest[Int])
+        val out_array = (new ARRAY(xn(1))).withSrcType(__pos, manifest[Int])
+        val size  = (new INT(xn(2))).withSrcType(__pos, manifest[Int])
+
+        val stride = gridDimX * blockDimX
+        val tid = threadIdxX + blockIdxX * blockDimX
+
+        val i = var_new(Wrap[Int](tid.x))
+
+        __whileDo(ordering_lt(readVar(i), Wrap[Int](size.x)), {
+          // replace input to function argument, tensor lambda to loop index
+          try {
+            subst(inputs.head.t) = in_array.x
+            subst(arg) = Unwrap(readVar(i))
+            traverse(f)
+          } finally {
+            subst -= inputs.head.t
+            subst -= arg
+          }
+          out_array(INT(Unwrap(readVar(i)))) = INT(transform(r))
+          i += Wrap[Int](stride.x)
+        })
+
+        Backend.Const(())
+      }, manifest[Array[Int]], manifest[Array[Int]], manifest[Int])
+
+      (kernel(in_arr, out_arr, INT(sz1), DIM3(0), DIM3(0))).x
+
+
+    case Node(s, "tensor_show", (x:Backend.Sym)::_, _) =>
+      implicit val pos = Adapter.oldSourceMap(s)
+      Backend.Const(())
+    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/fusedTensor.scala

@@ -0,0 +1,232 @@
+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)
+
+  case class View(t: Backend.Sym, from: Int, to: Int)
+
+  /// typeless frontend
+
+  /*
+  def TENSOR(size: Seq[Int], inputs: Seq[Backend.Sym])(f: Backend.Exp => Backend.Exp)(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectUnsafe("tensor", C(size), C(inputs), Adapter.g.reify(xn => f(xn))))).withSrcType(__pos, manifest[Int])
+  }*/
+
+  def TENSOR(inputs: Seq[View])(f: Backend.Exp => Backend.Exp)(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectUnsafe("tensor", C(inputs), Adapter.g.reify(xn => f(xn))))).withSrcType(__pos, manifest[Int])
+  }
+
+  def ZEROS(size: Int)(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectUnsafe("tensor_zeros", C(size)))).withSrcType(__pos, manifest[Int])
+  }
+
+  def ONES(size: Int)(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectUnsafe("tensor_ones", C(size)))).withSrcType(__pos, manifest[Int])
+  }
+
+  def CONSTS(size: Int, num: Int)(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectUnsafe("tensor_consts", C(size), C(num)))).withSrcType(__pos, manifest[Int])
+  }
+
+  def INPUT(size: Int)(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectUnsafe("tensor_input", C(size)))).withSrcType(__pos, manifest[Int])
+  }
+
+  // used to track input by itself
+  def INPUT1(inputs: Seq[View])(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectUnsafe("tensor_input", C(inputs)))).withSrcType(__pos, manifest[Int])
+  }
+
+  def TENSORS(inputs: Seq[Backend.Exp])(implicit __pos: SourceContext): TENSOR = {
+    (new TENSOR(Adapter.g.reflectUnsafe("tensors", C(inputs)))).withSrcType(__pos, manifest[Int])
+  }
+
+  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:Seq[Int])::_, _)) => size
+        case Some(Node(_, s, Backend.Const(_)::Backend.Const(size:Seq[Int])::_, _)) => size*/
+        case Some(Node(_, s, Backend.Const(inputs:Seq[View])::_, _)) =>
+          val sz = inputs(0).to - inputs(0).from
+          inputs foreach {
+            case View(_, from, to) => assert(to - from == sz, "operation shape mismatch")
+          }; sz
+        case a => System.out.println(a); ???
+      }
+    }
+
+    def inputs: Seq[View] = {
+      gc.get(x.asInstanceOf[Backend.Sym]) match {
+        // case Some(Node(_, op, _::Backend.Const(ins:Seq[Backend.Sym])::_, _)) => ins
+        case Some(Node(_, s, Backend.Const(inputs:Seq[View])::_, _)) => inputs
+        case a => Seq()
+      }
+    }
+
+    def body: Backend.Block = {
+      gc.get(x.asInstanceOf[Backend.Sym]) match {
+        // case Some(Node(_, "tensor", _::_::(f:Backend.Block)::_, _)) => f
+        case Some(Node(_, "tensor", _::(f:Backend.Block)::_, _)) => f
+        case a => ???
+      }
+    }
+
+    def show(implicit __pos: SourceContext): UNIT = {
+      UNIT(Adapter.g.reflectWrite("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 tanh(implicit __pos: SourceContext): TENSOR = {
+      (new TENSOR(Adapter.g.reflect("tensor_tanh", x))).withSrcType(__pos, et)
+    }
+
+    def relu(implicit __pos: SourceContext): TENSOR = {
+      (new TENSOR(Adapter.g.reflect("tensor_relu", x))).withSrcType(__pos, et)
+    }
+
+    def apply(e: Backend.Exp)(implicit __pos: SourceContext): INT = {
+      // INT(Adapter.g.reflectEffect("tensor_apply", x, e)()(Adapter.CTRL)).withSrcType(__pos, et)
+      // read effect?
+      INT(Adapter.g.reflect("tensor_apply", x, e)).withSrcType(__pos, et)
+    }
+
+    def split(sh: Seq[Int])(implicit __pos: SourceContext): TENSOR = {
+      (new TENSOR(Adapter.g.reflect("tensor_split", x, C(sh)))).withSrcType(__pos, et)
+    }
+
+    def result(i: Int)(implicit __pos: SourceContext): TENSOR = {
+      (new TENSOR(Adapter.g.reflect("tensor_result", x, C(i)))).withSrcType(__pos, et)
+    }
+
+    /*def concat(y: TENSOR)(implicit __pos: SourceContext): TENSOR = {
+      (new TENSOR(Adapter.g.reflectUnsafe("tensor_concat", x, y.x))).withSrcType(__pos, et)
+    }*/
+
+    def concat(y: Seq[TENSOR])(implicit __pos: SourceContext): TENSOR = {
+      val tmp = x +: (y map { _.x })
+      (new TENSOR(Adapter.g.reflectUnsafe("tensor_concat", tmp:_*))).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)(implicit __pos: SourceContext): Rep[Tensor[T]] = {
+      val tensor = ZEROS(size)
+      Wrap[Tensor[T]](tensor.x)
+    }
+
+    def ones[T:Manifest](size: Int)(implicit __pos: SourceContext): Rep[Tensor[T]] = {
+      val tensor = ONES(size)
+      Wrap[Tensor[T]](tensor.x)
+    }
+
+    def consts[T:Manifest](size: Int, num: Int)(implicit __pos: SourceContext): Rep[Tensor[T]] = {
+      val tensor = CONSTS(size, num)
+      Wrap[Tensor[T]](tensor.x)
+    }
+
+    def input[T:Manifest](size: Int)(implicit  __pos: SourceContext): Rep[Tensor[T]] = {
+      val tensor = INPUT(size)
+      Wrap[Tensor[T]](tensor.x)
+    }
+
+    def apply[T:Numeric:Manifest](size: Int, f: Rep[Int] => Rep[Int])(implicit __pos: SourceContext): Rep[Tensor[T]] = {
+      Wrap[Tensor[T]](TENSOR(Seq(View(null, 0, size)))(unwrapFun[Int, Int](f)).x) // is the input correct?
+    }
+  }
+
+  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 tanh(implicit  __pos: SourceContext): Rep[Tensor[T]] = {
+      val t = self.tanh
+      Wrap[Tensor[T]](t.x)
+    }
+
+    def relu(implicit  __pos: SourceContext): Rep[Tensor[T]] = {
+      val t = self.relu
+      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)
+    }
+
+    def split(sh: Seq[Int])(implicit  __pos: SourceContext): Rep[Tensor[T]] = {
+      val t = self.split(sh)
+      Wrap[Tensor[T]](t.x)
+    }
+
+    def result(i: Int)(implicit  __pos: SourceContext): Rep[Tensor[T]] = {
+      val t = self.result(i)
+      Wrap[Tensor[T]](t.x)
+    }
+
+    def concat(ys: Seq[Rep[Tensor[T]]])(implicit  __pos: SourceContext): Rep[Tensor[T]] = {
+      val y = ys map { tensor(_) }
+      val t = self.concat(y)
+      Wrap[Tensor[T]](t.x)
+    }
+  }
+}