From e0d926bf9f67f2b2c98f9f8a16c5ecc1c600dc5c Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Fri, 19 Nov 2021 17:02:34 -0500 Subject: [PATCH 01/21] vertical fusion --- .../kernel_fusion/dot/dot.check.cu | 38 +++++++++ .../transformation/test_fused_tensor.scala | 82 +++++++++++++++++++ 2 files changed, 120 insertions(+) create mode 100644 src/out/transformer/kernel_fusion/dot/dot.check.cu create mode 100644 src/test/scala/lms/transformation/test_fused_tensor.scala diff --git a/src/out/transformer/kernel_fusion/dot/dot.check.cu b/src/out/transformer/kernel_fusion/dot/dot.check.cu new file mode 100644 index 00000000..c1b2b5c0 --- /dev/null +++ b/src/out/transformer/kernel_fusion/dot/dot.check.cu @@ -0,0 +1,38 @@ +/***************************************** +Emitting C Generated Code +*******************************************/ +#ifndef _GNU_SOURCE +#define _GNU_SOURCE +#endif +#include "cudnn_header.h" +#include "nccl_header.h" +#include +#include +#include +#include "cuda_header.h" +#include +#include +#include "cublas_header.h" +#include +#include "mpi_header.h" +#include "scanner_header.h" +/**************** Snippet ****************/ +void Snippet(int x0) { + module([&]() { + float x1 = tensor_zeros(<4d2x4d3xf32>, Split d0 at devices=[GPU(0), GPU(1)]); + float x2 = tensor_zeros(<4d2x4d3xf32>, Split d0 at devices=[GPU(0), GPU(1)]); + float x3 = tensor_add(<4d2x4d3xf32>, Split d0 at devices=[GPU(0), GPU(1)], x1, x2); + return x3; + })("loss"); +} +/***************************************** +End of C Generated Code +*******************************************/ +int main(int argc, char *argv[]) { + if (argc != 2) { + printf("usage: %s \n", argv[0]); + return 0; + } + Snippet(atoi(argv[1])); + return 0; +} diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala new file mode 100644 index 00000000..5a4d2957 --- /dev/null +++ b/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)) + } +} \ No newline at end of file From 6046496e2b156c7e1f7b76bd298ec7ba4385f406 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Fri, 19 Nov 2021 17:04:29 -0500 Subject: [PATCH 02/21] vertical fusion --- .../fused_tensor/fusedTensor.scala | 123 ++++++++++++++++++ .../fused_tensor/fusedTensorLowering.scala | 63 +++++++++ .../fused_tensor/fusedTensorVertical.scala | 52 ++++++++ 3 files changed, 238 insertions(+) create mode 100644 src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala create mode 100644 src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala create mode 100644 src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala new file mode 100644 index 00000000..144073c7 --- /dev/null +++ b/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) + } + + } +} diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala new file mode 100644 index 00000000..f7d6d69e --- /dev/null +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala @@ -0,0 +1,63 @@ +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) + 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) + 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) + 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 + } + } +} diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala new file mode 100644 index 00000000..c975f033 --- /dev/null +++ b/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 + } + } +} From 536e36976269b1f2e3fce057b165a245b2e5d9a7 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Fri, 19 Nov 2021 17:16:57 -0500 Subject: [PATCH 03/21] implicit pos --- .../lms/transformation/fused_tensor/fusedTensorLowering.scala | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala index f7d6d69e..72dc8199 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala @@ -24,16 +24,19 @@ abstract class FusedTensorLowering extends Transformer { 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 From dc238cd8417261500e8ddecd5c10a0379af3f1b3 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Fri, 19 Nov 2021 17:41:44 -0500 Subject: [PATCH 04/21] remove old test --- .../kernel_fusion/dot/dot.check.cu | 38 ------------------- 1 file changed, 38 deletions(-) delete mode 100644 src/out/transformer/kernel_fusion/dot/dot.check.cu diff --git a/src/out/transformer/kernel_fusion/dot/dot.check.cu b/src/out/transformer/kernel_fusion/dot/dot.check.cu deleted file mode 100644 index c1b2b5c0..00000000 --- a/src/out/transformer/kernel_fusion/dot/dot.check.cu +++ /dev/null @@ -1,38 +0,0 @@ -/***************************************** -Emitting C Generated Code -*******************************************/ -#ifndef _GNU_SOURCE -#define _GNU_SOURCE -#endif -#include "cudnn_header.h" -#include "nccl_header.h" -#include -#include -#include -#include "cuda_header.h" -#include -#include -#include "cublas_header.h" -#include -#include "mpi_header.h" -#include "scanner_header.h" -/**************** Snippet ****************/ -void Snippet(int x0) { - module([&]() { - float x1 = tensor_zeros(<4d2x4d3xf32>, Split d0 at devices=[GPU(0), GPU(1)]); - float x2 = tensor_zeros(<4d2x4d3xf32>, Split d0 at devices=[GPU(0), GPU(1)]); - float x3 = tensor_add(<4d2x4d3xf32>, Split d0 at devices=[GPU(0), GPU(1)], x1, x2); - return x3; - })("loss"); -} -/***************************************** -End of C Generated Code -*******************************************/ -int main(int argc, char *argv[]) { - if (argc != 2) { - printf("usage: %s \n", argv[0]); - return 0; - } - Snippet(atoi(argv[1])); - return 0; -} From 9b7bab095206bc2a8ee5a40db2859bdb1ca4a632 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Tue, 23 Nov 2021 13:51:38 -0500 Subject: [PATCH 05/21] fix tensor add issue by adding one more pass --- .../fused_tensor/fusedTensor.scala | 25 ++++++++ .../fused_tensor/fusedTensorLowering.scala | 19 +----- .../fused_tensor/fusedTensorSimplify.scala | 62 +++++++++++++++++++ .../fused_tensor/fusedTensorVertical.scala | 2 - .../fused_tensor/show/show.check.cu | 22 +++++++ .../transformation/test_fused_tensor.scala | 50 ++++++++++----- 6 files changed, 146 insertions(+), 34 deletions(-) create mode 100644 src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala create mode 100644 src/out/transformer/fused_tensor/show/show.check.cu diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index 144073c7..283fdd8a 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -49,6 +49,13 @@ object FusedTensorTypeLess { } } + def arr: Backend.Sym = { + gc.get(x.asInstanceOf[Backend.Sym]) match { + case Some(Node(_, s, Backend.Const(size:Int)::(arr:Backend.Sym)::_, _)) => arr + case a => System.out.println(a); ??? + } + } + def show(implicit __pos: SourceContext): UNIT = { UNIT(Adapter.g.reflectEffect("show_tensor", x)()(Adapter.CTRL)) } @@ -61,6 +68,14 @@ object FusedTensorTypeLess { (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 = { // todo: change to correct effect INT(Adapter.g.reflectEffect("tensor_apply", x, e)()(Adapter.CTRL)).withSrcType(__pos, et) @@ -114,6 +129,16 @@ trait FusedTensorOps extends Dsl with ArrayOps with CudaOps { 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) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala index 72dc8199..465b0249 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala @@ -18,38 +18,21 @@ abstract class FusedTensorLowering extends Transformer { import ArrayTypeLess._ import ArrayCPUTypeLess._ import FusedTensorTypeLess._ + import PrimitiveTypeLess._ 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) } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala new file mode 100644 index 00000000..16113a14 --- /dev/null +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala @@ -0,0 +1,62 @@ +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 FusedTensorSimplify extends Transformer { + + import BaseTypeLess._ + import PrimitiveTypeLess._ + import ArrayTypeLess._ + import ArrayCPUTypeLess._ + import FusedTensorTypeLess._ + import PrimitiveTypeLess._ + + 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_add", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val a = new TENSOR(transform(x), useOldMetadata = true) + val b = new TENSOR(transform(y), useOldMetadata = true) + val array = ARRAY(10, manifest[Int]) + val t = TENSOR(10, array)(i => ( + a.apply(INT(i).x) + b.apply(INT(i).x)).x) + t.x + case Node(s, "tensor_tanh", (x:Backend.Sym)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val t = new TENSOR(transform(x), useOldMetadata = true) + val array = ARRAY(10, manifest[Int]) + val res = TENSOR(10, array)(i => t.apply(INT(i).x).tanh().x) + res.x + case Node(s, "tensor_relu", (x:Backend.Sym)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val t = new TENSOR(x, useOldMetadata = true) + val array = new ARRAY(t.arr) + val res = TENSOR(t.size, array)(i => { + // IF(c: BOOL)(a: => TOP)(b: => TOP) + (IF(INT(i) < INT(0))(INT(0))(INT(i))).x + }) + res.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 + } + } +} diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala index c975f033..d01bc2fc 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala @@ -23,11 +23,9 @@ abstract class FusedTensorVertical extends Transformer { 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) diff --git a/src/out/transformer/fused_tensor/show/show.check.cu b/src/out/transformer/fused_tensor/show/show.check.cu new file mode 100644 index 00000000..d1436430 --- /dev/null +++ b/src/out/transformer/fused_tensor/show/show.check.cu @@ -0,0 +1,22 @@ +/***************************************** +Emitting C Generated Code +*******************************************/ +#include +#include +#include +#include +/**************** Snippet ****************/ +void Snippet(int x0) { + printf("%d", 0); +} +/***************************************** +End of C Generated Code +*******************************************/ +int main(int argc, char *argv[]) { + if (argc != 2) { + printf("usage: %s \n", argv[0]); + return 0; + } + Snippet(atoi(argv[1])); + return 0; +} diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index 5a4d2957..acbf0b92 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -13,7 +13,7 @@ import lms.thirdparty.array_computation.{CCodeGenCBLASOps, CCodeGenCudaOps} import Backend._ class FixedSizeFusedTensorTest extends TutorialFunSuite { - val under = "transformer/fused_tensor" + val under = "transformer/fused_tensor/" abstract class CompilerCFusedTensor[A: Manifest, B: Manifest] extends CompilerC[A,B] with FusedTensorOps { q => @@ -23,10 +23,11 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { override val passes = List( new FusedTensorLowering {}, - new FusedTensorVertical {} + new FusedTensorSimplify {}, + new FusedTensorVertical {}, + new Canonicalize {} ) - /* var log_path: String = "" def setLogPath(path: String) { log_path = path } @@ -48,11 +49,12 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { val stream = new java.io.PrintStream(source) stream.println("==================") for (node <- graph.nodes) - node.toString + // node.toString + stream.println(node) stream.println(graph.block) stream.println("==================") source.toString - }*/ + } } test("show") { @@ -62,21 +64,41 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { @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) + printf("%d", a(0)) + } + } + checkWithLogPath("show", driver.code, "cu", driver.setLogPath) + } + + test("add") { + val driver = new CompilerCFusedTensor[Int, Unit] { + import FusedTensorTypeLess._ + + @virtualize + def snippet(arg: Rep[Int]): Rep[Unit] = { + val array = NewArray[Int](10) 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) + } + } + checkWithLogPath("add", driver.code, "cu", driver.setLogPath) + } + + test("tanh") { + val driver = new CompilerCFusedTensor[Int, Unit] { + import FusedTensorTypeLess._ - // 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; () + @virtualize + def snippet(arg: Rep[Int]): Rep[Unit] = { + val array = NewArray[Int](10) + val a = Tensor.zeros[Int](10, array) + val c = a.tanh + printf("%d", c(0)) } } - System.out.println(indent(driver.code)) + checkWithLogPath("tanh", driver.code, "cu", driver.setLogPath) } } \ No newline at end of file From ee11eabdffeb2e7bcb7863e1c192ecdc89dd61de Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Tue, 23 Nov 2021 13:59:02 -0500 Subject: [PATCH 06/21] tensor const --- .../lms/transformation/fused_tensor/fusedTensor.scala | 9 +++++++++ .../fused_tensor/fusedTensorLowering.scala | 5 +++++ src/out/transformer/fused_tensor/show/show.check.cu | 1 + .../scala/lms/transformation/test_fused_tensor.scala | 2 ++ 4 files changed, 17 insertions(+) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index 283fdd8a..3e4e004c 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -33,6 +33,10 @@ object FusedTensorTypeLess { (new TENSOR(Adapter.g.reflectUnsafe("tensor_ones", C(size), array.x))).withSrcType(__pos, array.et) } + def CONSTS(size: Int, num: Int, array: ARRAY)(implicit __pos: SourceContext): TENSOR = { + (new TENSOR(Adapter.g.reflectUnsafe("tensor_consts", C(size), C(num), 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 = @@ -107,6 +111,11 @@ trait FusedTensorOps extends Dsl with ArrayOps with CudaOps { Wrap[Tensor[T]](tensor.x) } + def consts[T:Manifest](size: Int, num: Int, array: Rep[Array[T]])(implicit __pos: SourceContext): Rep[Tensor[T]] = { + val tensor = CONSTS(size, num, 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) } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala index 465b0249..aff437cc 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala @@ -33,6 +33,11 @@ abstract class FusedTensorLowering extends Transformer { val array = new ARRAY(arr) val t = TENSOR(sz, array)(i => INT(1).x) t.x + case Node(s, "tensor_consts", (Backend.Const(sz:Int))::(Backend.Const(n:Int))::(arr:Backend.Exp)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val array = new ARRAY(arr) + val t = TENSOR(sz, array)(i => INT(n).x) + t.x case _ => super.transform(n) } diff --git a/src/out/transformer/fused_tensor/show/show.check.cu b/src/out/transformer/fused_tensor/show/show.check.cu index d1436430..f6a63ebf 100644 --- a/src/out/transformer/fused_tensor/show/show.check.cu +++ b/src/out/transformer/fused_tensor/show/show.check.cu @@ -8,6 +8,7 @@ Emitting C Generated Code /**************** Snippet ****************/ void Snippet(int x0) { printf("%d", 0); + printf("%d", 5); } /***************************************** End of C Generated Code diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index acbf0b92..04d894a0 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -65,7 +65,9 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { def snippet(arg: Rep[Int]): Rep[Unit] = { val array = NewArray[Int](10) val a = Tensor.zeros[Int](10, array) + val b = Tensor.consts[Int](10, 5, array) printf("%d", a(0)) + printf("%d", b(0)) } } checkWithLogPath("show", driver.code, "cu", driver.setLogPath) From 88e4d60b8b38e3c461ff8cf9c9cddf377c73a080 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Tue, 23 Nov 2021 14:10:07 -0500 Subject: [PATCH 07/21] tensor minus --- .../transformation/fused_tensor/fusedTensorSimplify.scala | 8 ++++++++ src/test/scala/lms/transformation/test_fused_tensor.scala | 2 ++ 2 files changed, 10 insertions(+) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala index 16113a14..f5813190 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala @@ -31,6 +31,14 @@ abstract class FusedTensorSimplify extends Transformer { val t = TENSOR(10, array)(i => ( a.apply(INT(i).x) + b.apply(INT(i).x)).x) t.x + case Node(s, "tensor_minus", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val a = new TENSOR(transform(x), useOldMetadata = true) + val b = new TENSOR(transform(y), useOldMetadata = true) + val array = ARRAY(10, manifest[Int]) + val t = TENSOR(10, array)(i => ( + a.apply(INT(i).x) - b.apply(INT(i).x)).x) + t.x case Node(s, "tensor_tanh", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(transform(x), useOldMetadata = true) diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index 04d894a0..ae35cf40 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -83,7 +83,9 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { val a = Tensor.zeros[Int](10, array) val b = Tensor.ones[Int](10, array) val c = a + b + val d = a - b printf("%d", c(0)) + printf("%d", d(0)) } } checkWithLogPath("add", driver.code, "cu", driver.setLogPath) From 14b8e3d2abc090ac204e9c9de9c55a75e138f410 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Tue, 23 Nov 2021 14:17:05 -0500 Subject: [PATCH 08/21] tensor effect; relu --- .../fused_tensor/fusedTensor.scala | 4 ++-- .../transformer/fused_tensor/add/add.check.cu | 23 +++++++++++++++++++ .../fused_tensor/relu/relu.check.cu | 22 ++++++++++++++++++ .../fused_tensor/tanh/tanh.check.cu | 23 +++++++++++++++++++ .../transformation/test_fused_tensor.scala | 15 ++++++++++++ 5 files changed, 85 insertions(+), 2 deletions(-) create mode 100644 src/out/transformer/fused_tensor/add/add.check.cu create mode 100644 src/out/transformer/fused_tensor/relu/relu.check.cu create mode 100644 src/out/transformer/fused_tensor/tanh/tanh.check.cu diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index 3e4e004c..e3ee0e12 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -81,8 +81,8 @@ object FusedTensorTypeLess { } 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) + // INT(Adapter.g.reflectEffect("tensor_apply", x, e)()(Adapter.CTRL)).withSrcType(__pos, et) + INT(Adapter.g.reflect("tensor_apply", x, e)).withSrcType(__pos, et) } } } diff --git a/src/out/transformer/fused_tensor/add/add.check.cu b/src/out/transformer/fused_tensor/add/add.check.cu new file mode 100644 index 00000000..5e0ffca4 --- /dev/null +++ b/src/out/transformer/fused_tensor/add/add.check.cu @@ -0,0 +1,23 @@ +/***************************************** +Emitting C Generated Code +*******************************************/ +#include +#include +#include +#include +/**************** Snippet ****************/ +void Snippet(int x0) { + printf("%d", 1); + printf("%d", -1); +} +/***************************************** +End of C Generated Code +*******************************************/ +int main(int argc, char *argv[]) { + if (argc != 2) { + printf("usage: %s \n", argv[0]); + return 0; + } + Snippet(atoi(argv[1])); + return 0; +} diff --git a/src/out/transformer/fused_tensor/relu/relu.check.cu b/src/out/transformer/fused_tensor/relu/relu.check.cu new file mode 100644 index 00000000..d1436430 --- /dev/null +++ b/src/out/transformer/fused_tensor/relu/relu.check.cu @@ -0,0 +1,22 @@ +/***************************************** +Emitting C Generated Code +*******************************************/ +#include +#include +#include +#include +/**************** Snippet ****************/ +void Snippet(int x0) { + printf("%d", 0); +} +/***************************************** +End of C Generated Code +*******************************************/ +int main(int argc, char *argv[]) { + if (argc != 2) { + printf("usage: %s \n", argv[0]); + return 0; + } + Snippet(atoi(argv[1])); + return 0; +} diff --git a/src/out/transformer/fused_tensor/tanh/tanh.check.cu b/src/out/transformer/fused_tensor/tanh/tanh.check.cu new file mode 100644 index 00000000..5c2def63 --- /dev/null +++ b/src/out/transformer/fused_tensor/tanh/tanh.check.cu @@ -0,0 +1,23 @@ +/***************************************** +Emitting C Generated Code +*******************************************/ +#include +#include +#include +#include +#include +/**************** Snippet ****************/ +void Snippet(int x0) { + printf("%d", tanh(0)); +} +/***************************************** +End of C Generated Code +*******************************************/ +int main(int argc, char *argv[]) { + if (argc != 2) { + printf("usage: %s \n", argv[0]); + return 0; + } + Snippet(atoi(argv[1])); + return 0; +} diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index ae35cf40..7cc4a913 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -105,4 +105,19 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { } checkWithLogPath("tanh", driver.code, "cu", driver.setLogPath) } + + test("relu") { + val driver = new CompilerCFusedTensor[Int, Unit] { + import FusedTensorTypeLess._ + + @virtualize + def snippet(arg: Rep[Int]): Rep[Unit] = { + val array = NewArray[Int](10) + val a = Tensor.ones[Int](10, array) + val c = a.relu + printf("%d", c(0)) + } + } + checkWithLogPath("relu", driver.code, "cu", driver.setLogPath) + } } \ No newline at end of file From 7c824bd87f7e333819e45adbbea9ed808aabe3bd Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Tue, 23 Nov 2021 16:32:40 -0500 Subject: [PATCH 09/21] remove array allocation --- .../fused_tensor/fusedTensor.scala | 34 +++++++++---------- .../fused_tensor/fusedTensorLowering.scala | 15 ++++---- .../fused_tensor/fusedTensorSimplify.scala | 14 ++++---- .../fused_tensor/fusedTensorVertical.scala | 5 +-- .../transformation/test_fused_tensor.scala | 21 +++++------- 5 files changed, 41 insertions(+), 48 deletions(-) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index e3ee0e12..ea8d81a4 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -21,20 +21,20 @@ object FusedTensorTypeLess { 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 TENSOR(size: Int)(f: Backend.Exp => Backend.Exp)(implicit __pos: SourceContext): TENSOR = { + (new TENSOR(Adapter.g.reflectUnsafe("tensor", C(size), Adapter.g.reify(xn => f(xn))))).withSrcType(__pos, manifest[Int]) } - 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 ZEROS(size: Int)(implicit __pos: SourceContext): TENSOR = { + (new TENSOR(Adapter.g.reflectUnsafe("tensor_zeros", C(size)))).withSrcType(__pos, manifest[Int]) } - 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) + 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, array: ARRAY)(implicit __pos: SourceContext): TENSOR = { - (new TENSOR(Adapter.g.reflectUnsafe("tensor_consts", C(size), C(num), array.x))).withSrcType(__pos, array.et) + 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]) } class TENSOR(override val x: Backend.Exp, override val useOldMetadata: Boolean = false) extends TOP(x) { @@ -52,7 +52,7 @@ object FusedTensorTypeLess { case a => System.out.println(a); ??? } } - + def arr: Backend.Sym = { gc.get(x.asInstanceOf[Backend.Sym]) match { case Some(Node(_, s, Backend.Const(size:Int)::(arr:Backend.Sym)::_, _)) => arr @@ -101,23 +101,23 @@ trait FusedTensorOps extends Dsl with ArrayOps with CudaOps { /// 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))) + 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, array: Rep[Array[T]])(implicit __pos: SourceContext): Rep[Tensor[T]] = { - val tensor = ONES(size, new ARRAY(Unwrap(array))) + 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, array: Rep[Array[T]])(implicit __pos: SourceContext): Rep[Tensor[T]] = { - val tensor = CONSTS(size, num, new ARRAY(Unwrap(array))) + 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 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 apply[T:Numeric:Manifest](size: Int, f: Rep[Int] => Rep[Int])(implicit __pos: SourceContext): Rep[Tensor[T]] = { + Wrap[Tensor[T]](TENSOR(size)(unwrapFun[Int, Int](f)).x) } } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala index aff437cc..0bef0aa1 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala @@ -23,20 +23,17 @@ abstract class FusedTensorLowering extends Transformer { 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)::_, _) => + case Node(s, "tensor_zeros", (Backend.Const(sz:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val array = new ARRAY(arr) - val t = TENSOR(sz, array)(i => INT(0).x) + val t = TENSOR(sz)(i => INT(0).x) t.x - case Node(s, "tensor_ones", (Backend.Const(sz:Int))::(arr:Backend.Exp)::_, _) => + case Node(s, "tensor_ones", (Backend.Const(sz:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val array = new ARRAY(arr) - val t = TENSOR(sz, array)(i => INT(1).x) + val t = TENSOR(sz)(i => INT(1).x) t.x - case Node(s, "tensor_consts", (Backend.Const(sz:Int))::(Backend.Const(n:Int))::(arr:Backend.Exp)::_, _) => + case Node(s, "tensor_consts", (Backend.Const(sz:Int))::(Backend.Const(n:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val array = new ARRAY(arr) - val t = TENSOR(sz, array)(i => INT(n).x) + val t = TENSOR(sz)(i => INT(n).x) t.x case _ => super.transform(n) } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala index f5813190..5608e393 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala @@ -23,33 +23,31 @@ abstract class FusedTensorSimplify extends Transformer { 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_add", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x), useOldMetadata = true) val b = new TENSOR(transform(y), useOldMetadata = true) - val array = ARRAY(10, manifest[Int]) - val t = TENSOR(10, array)(i => ( + val t = TENSOR(10)(i => ( a.apply(INT(i).x) + b.apply(INT(i).x)).x) t.x case Node(s, "tensor_minus", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x), useOldMetadata = true) val b = new TENSOR(transform(y), useOldMetadata = true) - val array = ARRAY(10, manifest[Int]) - val t = TENSOR(10, array)(i => ( + val t = TENSOR(10)(i => ( a.apply(INT(i).x) - b.apply(INT(i).x)).x) t.x + case Node(s, "tensor_tanh", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(transform(x), useOldMetadata = true) - val array = ARRAY(10, manifest[Int]) - val res = TENSOR(10, array)(i => t.apply(INT(i).x).tanh().x) + val res = TENSOR(10)(i => t.apply(INT(i).x).tanh().x) res.x case Node(s, "tensor_relu", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(x, useOldMetadata = true) - val array = new ARRAY(t.arr) - val res = TENSOR(t.size, array)(i => { + val res = TENSOR(t.size)(i => { // IF(c: BOOL)(a: => TOP)(b: => TOP) (IF(INT(i) < INT(0))(INT(0))(INT(i))).x }) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala index d01bc2fc..fe76620f 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala @@ -22,11 +22,12 @@ abstract class FusedTensorVertical extends Transformer { 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)::_, _) => + case Node(s, "tensor", Backend.Const(size:Int)::_, _) => + System.out.println(n) tensors(s) = (n, path, inner) super.transform(n) case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) => - val (Node(_, _, Backend.Const(szy:Int)::(_)::(f@Backend.Block(arg::Nil, r, block, eff))::_, _), path0, inner0) = tensors(a) + 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) { diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index 7cc4a913..d5c49e95 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -63,25 +63,24 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { @virtualize def snippet(arg: Rep[Int]): Rep[Unit] = { - val array = NewArray[Int](10) - val a = Tensor.zeros[Int](10, array) - val b = Tensor.consts[Int](10, 5, array) + val a = Tensor.zeros[Int](10) + val b = Tensor.consts[Int](10, 5) printf("%d", a(0)) printf("%d", b(0)) } } checkWithLogPath("show", driver.code, "cu", driver.setLogPath) } - + test("add") { val driver = new CompilerCFusedTensor[Int, Unit] { import FusedTensorTypeLess._ @virtualize def snippet(arg: Rep[Int]): Rep[Unit] = { - val array = NewArray[Int](10) - val a = Tensor.zeros[Int](10, array) - val b = Tensor.ones[Int](10, array) + // val array = NewArray[Int](10) + val a = Tensor.zeros[Int](10) + val b = Tensor.ones[Int](10) val c = a + b val d = a - b printf("%d", c(0)) @@ -97,23 +96,21 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { @virtualize def snippet(arg: Rep[Int]): Rep[Unit] = { - val array = NewArray[Int](10) - val a = Tensor.zeros[Int](10, array) + val a = Tensor.zeros[Int](10) val c = a.tanh printf("%d", c(0)) } } checkWithLogPath("tanh", driver.code, "cu", driver.setLogPath) } - + test("relu") { val driver = new CompilerCFusedTensor[Int, Unit] { import FusedTensorTypeLess._ @virtualize def snippet(arg: Rep[Int]): Rep[Unit] = { - val array = NewArray[Int](10) - val a = Tensor.ones[Int](10, array) + val a = Tensor.ones[Int](10) val c = a.relu printf("%d", c(0)) } From 933b0d9cd604d70b6c09bcafa4240380223b90f7 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Tue, 23 Nov 2021 16:37:52 -0500 Subject: [PATCH 10/21] fix relu simpl. --- .../lms/transformation/fused_tensor/fusedTensorSimplify.scala | 2 +- src/out/transformer/fused_tensor/relu/relu.check.cu | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala index 5608e393..dbf84896 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala @@ -49,7 +49,7 @@ abstract class FusedTensorSimplify extends Transformer { val t = new TENSOR(x, useOldMetadata = true) val res = TENSOR(t.size)(i => { // IF(c: BOOL)(a: => TOP)(b: => TOP) - (IF(INT(i) < INT(0))(INT(0))(INT(i))).x + (IF(t.apply(INT(i).x) < INT(0))(INT(0))(t.apply(INT(i).x))).x }) res.x case _ => super.transform(n) diff --git a/src/out/transformer/fused_tensor/relu/relu.check.cu b/src/out/transformer/fused_tensor/relu/relu.check.cu index d1436430..ef6e8a95 100644 --- a/src/out/transformer/fused_tensor/relu/relu.check.cu +++ b/src/out/transformer/fused_tensor/relu/relu.check.cu @@ -7,7 +7,7 @@ Emitting C Generated Code #include /**************** Snippet ****************/ void Snippet(int x0) { - printf("%d", 0); + printf("%d", false ? 0 : 1); } /***************************************** End of C Generated Code From 9d1778ffa12a685a35754dcdfa0683b994477ea9 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Tue, 23 Nov 2021 16:42:33 -0500 Subject: [PATCH 11/21] remove arr --- .../lms/transformation/fused_tensor/fusedTensor.scala | 7 ------- .../transformation/fused_tensor/fusedTensorVertical.scala | 1 - 2 files changed, 8 deletions(-) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index ea8d81a4..4028fce1 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -52,13 +52,6 @@ object FusedTensorTypeLess { case a => System.out.println(a); ??? } } - - def arr: Backend.Sym = { - gc.get(x.asInstanceOf[Backend.Sym]) match { - case Some(Node(_, s, Backend.Const(size:Int)::(arr:Backend.Sym)::_, _)) => arr - case a => System.out.println(a); ??? - } - } def show(implicit __pos: SourceContext): UNIT = { UNIT(Adapter.g.reflectEffect("show_tensor", x)()(Adapter.CTRL)) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala index fe76620f..8b36fa7b 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala @@ -23,7 +23,6 @@ abstract class FusedTensorVertical extends Transformer { override def transform(n: Node): Backend.Exp = n match { case Node(s, "tensor", Backend.Const(size:Int)::_, _) => - System.out.println(n) tensors(s) = (n, path, inner) super.transform(n) case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) => From 00143ff199b0acd5799ce19ffd6996e304166870 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Tue, 23 Nov 2021 16:53:16 -0500 Subject: [PATCH 12/21] fix size --- .../transformation/fused_tensor/fusedTensorSimplify.scala | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala index dbf84896..c981e042 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala @@ -28,21 +28,21 @@ abstract class FusedTensorSimplify extends Transformer { implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x), useOldMetadata = true) val b = new TENSOR(transform(y), useOldMetadata = true) - val t = TENSOR(10)(i => ( + val t = TENSOR(a.size)(i => ( a.apply(INT(i).x) + b.apply(INT(i).x)).x) t.x case Node(s, "tensor_minus", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x), useOldMetadata = true) val b = new TENSOR(transform(y), useOldMetadata = true) - val t = TENSOR(10)(i => ( + val t = TENSOR(a.size)(i => ( a.apply(INT(i).x) - b.apply(INT(i).x)).x) t.x case Node(s, "tensor_tanh", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(transform(x), useOldMetadata = true) - val res = TENSOR(10)(i => t.apply(INT(i).x).tanh().x) + val res = TENSOR(t.size)(i => t.apply(INT(i).x).tanh().x) res.x case Node(s, "tensor_relu", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) From 1623596ad0e39087d135e5da16773fefd9b9d08d Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Mon, 29 Nov 2021 13:49:57 -0500 Subject: [PATCH 13/21] track tesor input; simple cuda backend --- .../fused_tensor/fusedTensor.scala | 9 ++- .../fused_tensor/fusedTensorKernel.scala | 69 +++++++++++++++++++ .../fused_tensor/fusedTensorLowering.scala | 6 +- .../fused_tensor/fusedTensorSimplify.scala | 9 ++- .../fused_tensor/fusedTensorVertical.scala | 4 +- .../transformer/fused_tensor/add/add.check.cu | 25 ++++++- .../fused_tensor/show/show.check.cu | 25 ++++++- .../fused_tensor/tanh/tanh.check.cu | 14 +++- .../transformation/test_fused_tensor.scala | 19 +++-- 9 files changed, 157 insertions(+), 23 deletions(-) create mode 100644 src/main/scala/lms/transformation/fused_tensor/fusedTensorKernel.scala diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index 4028fce1..e47673b5 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -21,8 +21,13 @@ object FusedTensorTypeLess { def C(a: Any) = Backend.Const(a) /// typeless frontend + /* def TENSOR(size: Int)(f: Backend.Exp => Backend.Exp)(implicit __pos: SourceContext): TENSOR = { (new TENSOR(Adapter.g.reflectUnsafe("tensor", C(size), Adapter.g.reify(xn => f(xn))))).withSrcType(__pos, manifest[Int]) + }*/ + + def TENSOR(size: 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 ZEROS(size: Int)(implicit __pos: SourceContext): TENSOR = { @@ -49,6 +54,7 @@ object FusedTensorTypeLess { def size: Int = { gc.get(x.asInstanceOf[Backend.Sym]) match { case Some(Node(_, s, Backend.Const(size:Int)::_, _)) => size + case Some(Node(_, s, Backend.Const(_)::Backend.Const(size:Int)::_, _)) => size case a => System.out.println(a); ??? } } @@ -75,6 +81,7 @@ object FusedTensorTypeLess { 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) } } @@ -110,7 +117,7 @@ trait FusedTensorOps extends Dsl with ArrayOps with CudaOps { } def apply[T:Numeric:Manifest](size: Int, f: Rep[Int] => Rep[Int])(implicit __pos: SourceContext): Rep[Tensor[T]] = { - Wrap[Tensor[T]](TENSOR(size)(unwrapFun[Int, Int](f)).x) + Wrap[Tensor[T]](TENSOR(size, Seq())(unwrapFun[Int, Int](f)).x) // is the input correct? } } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorKernel.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorKernel.scala new file mode 100644 index 00000000..c0dcabf2 --- /dev/null +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorKernel.scala @@ -0,0 +1,69 @@ +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 FusedTensorKernel extends Transformer { + + import BaseTypeLess._ + import PrimitiveTypeLess._ + import ArrayTypeLess._ + // import ArrayCPUTypeLess._ + import FusedTensorTypeLess._ + import PrimitiveTypeLess._ + import CUDATypeLess._ + + def gpu_array(size: Int, m: Manifest[_], device: INT)(implicit __pos: SourceContext): ARRAY = { + CUDA_SET_DEVICE(device) + CUDA_MALLOC(size, m) + } + + 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(sz:Int)::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _) => + implicit val __pos = Adapter.oldSourceMap(s) + val arr = CUDA_MALLOC(sz, manifest[Int]) + + val tmp = CUDA_KERNEL3({xn: List[Backend.Exp] => + val array = (new ARRAY(xn(0))).withSrcType(__pos, manifest[Int]) + val value = (new NUM(xn(1))).withSrcType(__pos, manifest[Int]) + val size = (new INT(xn(2))).withSrcType(__pos, manifest[Int]) + + val r1 = (new INT(r)).withSrcType(__pos, manifest[Int]) + + val stride = gridDimX * blockDimX + val tid = threadIdxX + blockIdxX * blockDimX + for (i <- range_until_step(Wrap[Int](tid.x), Wrap[Int](size.x), Wrap[Int](stride.x))) { + array(INT(Unwrap(i))) = r1; () + } + Backend.Const(()) + }, manifest[Array[Int]], manifest[Int], manifest[Int]) + + (tmp(arr, INT(0), INT(0), DIM3(0), DIM3(0))).x + + case Node(s, "tensor_show", Backend.Sym(x)::_, _) => + 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 + } + } +} diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala index 0bef0aa1..b0483b92 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala @@ -25,15 +25,15 @@ abstract class FusedTensorLowering extends Transformer { override def transform(n: Node): Backend.Exp = n match { case Node(s, "tensor_zeros", (Backend.Const(sz:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val t = TENSOR(sz)(i => INT(0).x) + val t = TENSOR(sz, Seq())(i => INT(0).x) t.x case Node(s, "tensor_ones", (Backend.Const(sz:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val t = TENSOR(sz)(i => INT(1).x) + val t = TENSOR(sz, Seq())(i => INT(1).x) t.x case Node(s, "tensor_consts", (Backend.Const(sz:Int))::(Backend.Const(n:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val t = TENSOR(sz)(i => INT(n).x) + val t = TENSOR(sz, Seq())(i => INT(n).x) t.x case _ => super.transform(n) } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala index c981e042..239bc63e 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala @@ -28,26 +28,25 @@ abstract class FusedTensorSimplify extends Transformer { implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x), useOldMetadata = true) val b = new TENSOR(transform(y), useOldMetadata = true) - val t = TENSOR(a.size)(i => ( + val t = TENSOR(a.size, Seq())(i => ( a.apply(INT(i).x) + b.apply(INT(i).x)).x) t.x case Node(s, "tensor_minus", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x), useOldMetadata = true) val b = new TENSOR(transform(y), useOldMetadata = true) - val t = TENSOR(a.size)(i => ( + val t = TENSOR(a.size, Seq(x, y))(i => ( a.apply(INT(i).x) - b.apply(INT(i).x)).x) t.x - case Node(s, "tensor_tanh", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(transform(x), useOldMetadata = true) - val res = TENSOR(t.size)(i => t.apply(INT(i).x).tanh().x) + val res = TENSOR(t.size, Seq(x))(i => t.apply(INT(i).x).tanh().x) // ad-hoc!!! res.x case Node(s, "tensor_relu", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(x, useOldMetadata = true) - val res = TENSOR(t.size)(i => { + val res = TENSOR(t.size, Seq(x))(i => { // IF(c: BOOL)(a: => TOP)(b: => TOP) (IF(t.apply(INT(i).x) < INT(0))(INT(0))(t.apply(INT(i).x))).x }) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala index 8b36fa7b..3012d829 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala @@ -22,11 +22,11 @@ abstract class FusedTensorVertical extends Transformer { 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)::_, _) => + case Node(s, "tensor", _, _) => tensors(s) = (n, path, inner) super.transform(n) case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) => - val (Node(_, _, Backend.Const(szy:Int)::(f@Backend.Block(arg::Nil, r, block, eff))::_, _), path0, inner0) = tensors(a) + val (Node(_, _, Backend.Const(szy:Int)::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _), path0, inner0) = tensors(a) try { subst(arg) = transform(b) withResetScope(path0, inner0) { diff --git a/src/out/transformer/fused_tensor/add/add.check.cu b/src/out/transformer/fused_tensor/add/add.check.cu index 5e0ffca4..938e384d 100644 --- a/src/out/transformer/fused_tensor/add/add.check.cu +++ b/src/out/transformer/fused_tensor/add/add.check.cu @@ -5,10 +5,31 @@ Emitting C Generated Code #include #include #include +/************* Functions **************/ +__global__ void x3(int x4, int x5, int x6) { + int x7 = gridDim.x * blockDim.x; + int x8 = threadIdx.x + blockIdx.x * blockDim.x; + while (x8 < x6) { + x4[x8] = 1; + x8 = x8 + x7; + } +} +__global__ void x9(int x10, int x11, int x12) { + int x13 = gridDim.x * blockDim.x; + int x14 = threadIdx.x + blockIdx.x * blockDim.x; + while (x14 < x12) { + x10[x14] = -1; + x14 = x14 + x13; + } +} /**************** Snippet ****************/ void Snippet(int x0) { - printf("%d", 1); - printf("%d", -1); + int* x1 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); + int* x2 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x2, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); + show_tensor(x3(x1, 0, 0, dim3(0, 1, 1), dim3(0, 1, 1))); + show_tensor(x9(x2, 0, 0, dim3(0, 1, 1), dim3(0, 1, 1))); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/show/show.check.cu b/src/out/transformer/fused_tensor/show/show.check.cu index f6a63ebf..fa606910 100644 --- a/src/out/transformer/fused_tensor/show/show.check.cu +++ b/src/out/transformer/fused_tensor/show/show.check.cu @@ -5,10 +5,31 @@ Emitting C Generated Code #include #include #include +/************* Functions **************/ +__global__ void x3(int x4, int x5, int x6) { + int x7 = gridDim.x * blockDim.x; + int x8 = threadIdx.x + blockIdx.x * blockDim.x; + while (x8 < x6) { + x4[x8] = 0; + x8 = x8 + x7; + } +} +__global__ void x9(int x10, int x11, int x12) { + int x13 = gridDim.x * blockDim.x; + int x14 = threadIdx.x + blockIdx.x * blockDim.x; + while (x14 < x12) { + x10[x14] = 5; + x14 = x14 + x13; + } +} /**************** Snippet ****************/ void Snippet(int x0) { - printf("%d", 0); - printf("%d", 5); + int* x1 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); + int* x2 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x2, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); + show_tensor(x3(x1, 0, 0, dim3(0, 1, 1), dim3(0, 1, 1))); + show_tensor(x9(x2, 0, 0, dim3(0, 1, 1), dim3(0, 1, 1))); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/tanh/tanh.check.cu b/src/out/transformer/fused_tensor/tanh/tanh.check.cu index 5c2def63..0b6f2e08 100644 --- a/src/out/transformer/fused_tensor/tanh/tanh.check.cu +++ b/src/out/transformer/fused_tensor/tanh/tanh.check.cu @@ -6,9 +6,21 @@ Emitting C Generated Code #include #include #include +/************* Functions **************/ +__global__ void x3(int x4, int x5, int x6) { + int x7 = gridDim.x * blockDim.x; + int x8 = threadIdx.x + blockIdx.x * blockDim.x; + while (x8 < x6) { + x4[x8] = x1; + x8 = x8 + x7; + } +} /**************** Snippet ****************/ void Snippet(int x0) { - printf("%d", tanh(0)); + int x1 = tanh(0); + int* x2 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x2, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); + show_tensor(x3(x2, 0, 0, dim3(0, 1, 1), dim3(0, 1, 1))); } /***************************************** End of C Generated Code diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index d5c49e95..0d248be5 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -8,7 +8,7 @@ import macros.SourceContext import lms.core._ import lms.core.stub._ import lms.thirdparty.{CCodeGenLibs} -import lms.thirdparty.array_computation.{CCodeGenCBLASOps, CCodeGenCudaOps} +// import lms.transformation.tensor.{CCodeGenCudaCustomOps} import Backend._ @@ -25,7 +25,8 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { new FusedTensorLowering {}, new FusedTensorSimplify {}, new FusedTensorVertical {}, - new Canonicalize {} + new Canonicalize {}, + new FusedTensorKernel {} ) var log_path: String = "" @@ -65,8 +66,10 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { def snippet(arg: Rep[Int]): Rep[Unit] = { val a = Tensor.zeros[Int](10) val b = Tensor.consts[Int](10, 5) - printf("%d", a(0)) - printf("%d", b(0)) + // printf("%d", a(0)) + // printf("%d", b(0)) + a.show + b.show } } checkWithLogPath("show", driver.code, "cu", driver.setLogPath) @@ -83,8 +86,10 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { val b = Tensor.ones[Int](10) val c = a + b val d = a - b - printf("%d", c(0)) - printf("%d", d(0)) + // printf("%d", c(0)) + // printf("%d", d(0)) + c.show + d.show } } checkWithLogPath("add", driver.code, "cu", driver.setLogPath) @@ -98,7 +103,7 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { def snippet(arg: Rep[Int]): Rep[Unit] = { val a = Tensor.zeros[Int](10) val c = a.tanh - printf("%d", c(0)) + c.show } } checkWithLogPath("tanh", driver.code, "cu", driver.setLogPath) From effba06c26a42727bd6f1ecab4f89d1b308362e1 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Thu, 2 Dec 2021 00:25:37 -0500 Subject: [PATCH 14/21] cuda backend --- ...orKernel.scala => FusedTensorToCuda.scala} | 43 ++++++++++---- .../fused_tensor/fusedTensor.scala | 21 +++++++ ...ring.scala => fusedTensorFunctional.scala} | 10 +++- .../fused_tensor/fusedTensorSimplify.scala | 17 +++--- .../fused_tensor/fusedTensorVertical.scala | 19 ++++++- .../transformer/fused_tensor/add/add.check.cu | 30 +++++----- .../show.check.cu => input/input.check.cu} | 25 +++------ .../fused_tensor/relu/relu.check.cu | 14 ++++- .../fused_tensor/tanh/tanh.check.cu | 19 +++---- .../transformation/test_fused_tensor.scala | 56 +++++++++---------- 10 files changed, 155 insertions(+), 99 deletions(-) rename src/main/scala/lms/transformation/fused_tensor/{fusedTensorKernel.scala => FusedTensorToCuda.scala} (61%) rename src/main/scala/lms/transformation/fused_tensor/{fusedTensorLowering.scala => fusedTensorFunctional.scala} (82%) rename src/out/transformer/fused_tensor/{show/show.check.cu => input/input.check.cu} (52%) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorKernel.scala b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala similarity index 61% rename from src/main/scala/lms/transformation/fused_tensor/fusedTensorKernel.scala rename to src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala index c0dcabf2..238d00c7 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorKernel.scala +++ b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala @@ -11,12 +11,13 @@ import lms.thirdparty.array_computation.{ArrayCPUTypeLess, CUDATypeLess, CUBLAST import Backend._ -abstract class FusedTensorKernel extends Transformer { +abstract class FusedTensorToCuda extends Transformer { + + override val name = "FusedTensorToCuda" import BaseTypeLess._ import PrimitiveTypeLess._ import ArrayTypeLess._ - // import ArrayCPUTypeLess._ import FusedTensorTypeLess._ import PrimitiveTypeLess._ import CUDATypeLess._ @@ -26,29 +27,47 @@ abstract class FusedTensorKernel extends Transformer { CUDA_MALLOC(size, m) } - 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(sz:Int)::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _) => implicit val __pos = Adapter.oldSourceMap(s) - val arr = CUDA_MALLOC(sz, manifest[Int]) - val tmp = CUDA_KERNEL3({xn: List[Backend.Exp] => + val arr = new ARRAY(inputs.head) // for now, assume only one input + // System.out.println("input: " + inputs.head) + // System.out.println("arr: " + arr) + // System.out.println("res:" + r) + + val kernel = CUDA_KERNEL3({ xn: List[Backend.Exp] => val array = (new ARRAY(xn(0))).withSrcType(__pos, manifest[Int]) - val value = (new NUM(xn(1))).withSrcType(__pos, manifest[Int]) + val value = (new NUM(xn(1))).withSrcType(__pos, manifest[Int]) // not used val size = (new INT(xn(2))).withSrcType(__pos, manifest[Int]) - val r1 = (new INT(r)).withSrcType(__pos, manifest[Int]) - val stride = gridDimX * blockDimX val tid = threadIdxX + blockIdxX * blockDimX - for (i <- range_until_step(Wrap[Int](tid.x), Wrap[Int](size.x), Wrap[Int](stride.x))) { - array(INT(Unwrap(i))) = r1; () - } + + val i = var_new(Wrap[Int](tid.x)) + + // System.out.println("i: " + UnwrapV(i)) + // PRINTF("%d", INT(Unwrap(readVar(i)))) + + __whileDo(ordering_lt(readVar(i), Wrap[Int](size.x)), { + // replace input to function argument, tensor lambda to loop index + try { + subst(inputs.head) = array.x + subst(arg) = UnwrapV(i) + traverse(f) + } finally { + subst -= inputs.head + subst -= arg + } + array(INT(Unwrap(readVar(i)))) = INT(transform(r)) + i += Wrap[Int](stride.x) + }) + Backend.Const(()) }, manifest[Array[Int]], manifest[Int], manifest[Int]) - (tmp(arr, INT(0), INT(0), DIM3(0), DIM3(0))).x + (kernel(arr, INT(0), INT(sz), DIM3(0), DIM3(0))).x case Node(s, "tensor_show", Backend.Sym(x)::_, _) => implicit val pos = Adapter.oldSourceMap(s) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index e47673b5..d5d8904e 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -42,6 +42,15 @@ object FusedTensorTypeLess { (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), C(Seq())))).withSrcType(__pos, manifest[Int]) + } + + // used to track input by itself + def INPUT1(size: Int, inputs: Seq[Backend.Sym])(implicit __pos: SourceContext): TENSOR = { + (new TENSOR(Adapter.g.reflectUnsafe("tensor_input", C(size), 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 = @@ -59,6 +68,13 @@ object FusedTensorTypeLess { } } + def inputs: Seq[Backend.Sym] = { + gc.get(x.asInstanceOf[Backend.Sym]) match { + case Some(Node(_, op, _::Backend.Const(ins:Seq[Backend.Sym])::_, _)) => ins + case a => Seq() + } + } + def show(implicit __pos: SourceContext): UNIT = { UNIT(Adapter.g.reflectEffect("show_tensor", x)()(Adapter.CTRL)) } @@ -116,6 +132,11 @@ trait FusedTensorOps extends Dsl with ArrayOps with CudaOps { 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(size, Seq())(unwrapFun[Int, Int](f)).x) // is the input correct? } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala similarity index 82% rename from src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala rename to src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala index b0483b92..469e5253 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorLowering.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala @@ -11,7 +11,9 @@ import lms.thirdparty.array_computation.{ArrayCPUTypeLess, CUDATypeLess, CUBLAST import Backend._ -abstract class FusedTensorLowering extends Transformer { +abstract class FusedTensorFunctional extends Transformer { + + override val name = "FusedTensorFunctional" import BaseTypeLess._ import PrimitiveTypeLess._ @@ -20,8 +22,6 @@ abstract class FusedTensorLowering extends Transformer { import FusedTensorTypeLess._ import PrimitiveTypeLess._ - 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))::_, _) => implicit val pos = Adapter.oldSourceMap(s) @@ -35,6 +35,10 @@ abstract class FusedTensorLowering extends Transformer { implicit val pos = Adapter.oldSourceMap(s) val t = TENSOR(sz, Seq())(i => INT(n).x) t.x + case Node(s, "tensor_input", (Backend.Const(sz:Int))::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val t1 = INPUT1(sz, Seq(s)) // necessary? + t1.x case _ => super.transform(n) } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala index 239bc63e..669814fc 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala @@ -13,6 +13,8 @@ import Backend._ abstract class FusedTensorSimplify extends Transformer { + override val name = "FusedTensorSimplify" + import BaseTypeLess._ import PrimitiveTypeLess._ import ArrayTypeLess._ @@ -28,28 +30,27 @@ abstract class FusedTensorSimplify extends Transformer { implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x), useOldMetadata = true) val b = new TENSOR(transform(y), useOldMetadata = true) - val t = TENSOR(a.size, Seq())(i => ( - a.apply(INT(i).x) + b.apply(INT(i).x)).x) + val t = TENSOR(a.size, a.inputs ++ b.inputs){ i => + (a.apply(INT(i).x) + b.apply(INT(i).x)).x } t.x case Node(s, "tensor_minus", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x), useOldMetadata = true) val b = new TENSOR(transform(y), useOldMetadata = true) - val t = TENSOR(a.size, Seq(x, y))(i => ( - a.apply(INT(i).x) - b.apply(INT(i).x)).x) + val t = TENSOR(a.size, a.inputs ++ b.inputs){ i => + (a.apply(INT(i).x) - b.apply(INT(i).x)).x } t.x case Node(s, "tensor_tanh", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(transform(x), useOldMetadata = true) - val res = TENSOR(t.size, Seq(x))(i => t.apply(INT(i).x).tanh().x) // ad-hoc!!! + val res = TENSOR(t.size, t.inputs){ i => t.apply(INT(i).x).tanh().x } // ad-hoc!!! res.x case Node(s, "tensor_relu", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(x, useOldMetadata = true) - val res = TENSOR(t.size, Seq(x))(i => { + val res = TENSOR(t.size, t.inputs){ i => // IF(c: BOOL)(a: => TOP)(b: => TOP) - (IF(t.apply(INT(i).x) < INT(0))(INT(0))(t.apply(INT(i).x))).x - }) + (IF(t.apply(INT(i).x) < INT(0))(INT(0))(t.apply(INT(i).x))).x } res.x case _ => super.transform(n) } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala index 3012d829..37a69d1b 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala @@ -13,19 +13,34 @@ import Backend._ abstract class FusedTensorVertical extends Transformer { + override val name = "FusedTensorVerticalFusion" + import BaseTypeLess._ import PrimitiveTypeLess._ import ArrayTypeLess._ - import ArrayCPUTypeLess._ import FusedTensorTypeLess._ + import CUDATypeLess._ + // map virtual tensors syms to node and context val tensors = new mutable.HashMap[Backend.Sym, (Node, List[Backend.Sym], Seq[Node])] + // 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(sz:Int)::_, _) => + implicit val __pos = Adapter.oldSourceMap(s) + val arr = CUDA_MALLOC(sz, manifest[Int]) // allocate CUDA array for input tensors + tensor2arr(s) = arr.x + arr.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", _, _) => tensors(s) = (n, path, inner) super.transform(n) - case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) => + case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) if tensors.contains(a) => val (Node(_, _, Backend.Const(szy:Int)::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _), path0, inner0) = tensors(a) try { subst(arg) = transform(b) diff --git a/src/out/transformer/fused_tensor/add/add.check.cu b/src/out/transformer/fused_tensor/add/add.check.cu index 938e384d..c200fc92 100644 --- a/src/out/transformer/fused_tensor/add/add.check.cu +++ b/src/out/transformer/fused_tensor/add/add.check.cu @@ -6,30 +6,28 @@ Emitting C Generated Code #include #include /************* Functions **************/ -__global__ void x3(int x4, int x5, int x6) { - int x7 = gridDim.x * blockDim.x; - int x8 = threadIdx.x + blockIdx.x * blockDim.x; - while (x8 < x6) { - x4[x8] = 1; - x8 = x8 + x7; +__global__ void x2(int x3, int x4, int x5) { + int x6 = gridDim.x * blockDim.x; + int x7 = threadIdx.x + blockIdx.x * blockDim.x; + while (x7 < x5) { + x3[x7] = x3[x7]; + x7 = x7 + x6; } } -__global__ void x9(int x10, int x11, int x12) { - int x13 = gridDim.x * blockDim.x; - int x14 = threadIdx.x + blockIdx.x * blockDim.x; - while (x14 < x12) { - x10[x14] = -1; - x14 = x14 + x13; +__global__ void x8(int x9, int x10, int x11) { + int x12 = gridDim.x * blockDim.x; + int x13 = threadIdx.x + blockIdx.x * blockDim.x; + while (x13 < x11) { + x9[x13] = x9[x13] - 1; + x13 = x13 + x12; } } /**************** Snippet ****************/ void Snippet(int x0) { int* x1 = (int*)malloc(0 * sizeof(int)); CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - int* x2 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x2, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - show_tensor(x3(x1, 0, 0, dim3(0, 1, 1), dim3(0, 1, 1))); - show_tensor(x9(x2, 0, 0, dim3(0, 1, 1), dim3(0, 1, 1))); + show_tensor(x2(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); + show_tensor(x8(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/show/show.check.cu b/src/out/transformer/fused_tensor/input/input.check.cu similarity index 52% rename from src/out/transformer/fused_tensor/show/show.check.cu rename to src/out/transformer/fused_tensor/input/input.check.cu index fa606910..dcf64aff 100644 --- a/src/out/transformer/fused_tensor/show/show.check.cu +++ b/src/out/transformer/fused_tensor/input/input.check.cu @@ -6,30 +6,19 @@ Emitting C Generated Code #include #include /************* Functions **************/ -__global__ void x3(int x4, int x5, int x6) { - int x7 = gridDim.x * blockDim.x; - int x8 = threadIdx.x + blockIdx.x * blockDim.x; - while (x8 < x6) { - x4[x8] = 0; - x8 = x8 + x7; - } -} -__global__ void x9(int x10, int x11, int x12) { - int x13 = gridDim.x * blockDim.x; - int x14 = threadIdx.x + blockIdx.x * blockDim.x; - while (x14 < x12) { - x10[x14] = 5; - x14 = x14 + x13; +__global__ void x2(int x3, int x4, int x5) { + int x6 = gridDim.x * blockDim.x; + int x7 = threadIdx.x + blockIdx.x * blockDim.x; + while (x7 < x5) { + x3[x7] = x3[x7] + 1; + x7 = x7 + x6; } } /**************** Snippet ****************/ void Snippet(int x0) { int* x1 = (int*)malloc(0 * sizeof(int)); CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - int* x2 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x2, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - show_tensor(x3(x1, 0, 0, dim3(0, 1, 1), dim3(0, 1, 1))); - show_tensor(x9(x2, 0, 0, dim3(0, 1, 1), dim3(0, 1, 1))); + show_tensor(x2(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/relu/relu.check.cu b/src/out/transformer/fused_tensor/relu/relu.check.cu index ef6e8a95..3dc850f7 100644 --- a/src/out/transformer/fused_tensor/relu/relu.check.cu +++ b/src/out/transformer/fused_tensor/relu/relu.check.cu @@ -5,9 +5,21 @@ Emitting C Generated Code #include #include #include +/************* Functions **************/ +__global__ void x2(int x3, int x4, int x5) { + int x6 = gridDim.x * blockDim.x; + int x7 = threadIdx.x + blockIdx.x * blockDim.x; + while (x7 < x5) { + int x8 = x3[x7]; + x3[x7] = x8 < 0 ? 0 : x8; + x7 = x7 + x6; + } +} /**************** Snippet ****************/ void Snippet(int x0) { - printf("%d", false ? 0 : 1); + int* x1 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); + show_tensor(x2(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/tanh/tanh.check.cu b/src/out/transformer/fused_tensor/tanh/tanh.check.cu index 0b6f2e08..06e36f22 100644 --- a/src/out/transformer/fused_tensor/tanh/tanh.check.cu +++ b/src/out/transformer/fused_tensor/tanh/tanh.check.cu @@ -7,20 +7,19 @@ Emitting C Generated Code #include #include /************* Functions **************/ -__global__ void x3(int x4, int x5, int x6) { - int x7 = gridDim.x * blockDim.x; - int x8 = threadIdx.x + blockIdx.x * blockDim.x; - while (x8 < x6) { - x4[x8] = x1; - x8 = x8 + x7; +__global__ void x2(int x3, int x4, int x5) { + int x6 = gridDim.x * blockDim.x; + int x7 = threadIdx.x + blockIdx.x * blockDim.x; + while (x7 < x5) { + x3[x7] = tanh(x3[x7]); + x7 = x7 + x6; } } /**************** Snippet ****************/ void Snippet(int x0) { - int x1 = tanh(0); - int* x2 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x2, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - show_tensor(x3(x2, 0, 0, dim3(0, 1, 1), dim3(0, 1, 1))); + int* x1 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); + show_tensor(x2(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); } /***************************************** End of C Generated Code diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index 0d248be5..3878fd78 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -22,11 +22,11 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { } override val passes = List( - new FusedTensorLowering {}, + new FusedTensorFunctional {}, new FusedTensorSimplify {}, new FusedTensorVertical {}, new Canonicalize {}, - new FusedTensorKernel {} + new FusedTensorToCuda {} ) var log_path: String = "" @@ -58,38 +58,21 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { } } - test("show") { - val driver = new CompilerCFusedTensor[Int, Unit] { - import FusedTensorTypeLess._ - - @virtualize - def snippet(arg: Rep[Int]): Rep[Unit] = { - val a = Tensor.zeros[Int](10) - val b = Tensor.consts[Int](10, 5) - // printf("%d", a(0)) - // printf("%d", b(0)) - a.show - b.show - } - } - checkWithLogPath("show", driver.code, "cu", driver.setLogPath) - } - test("add") { val driver = new CompilerCFusedTensor[Int, Unit] { import FusedTensorTypeLess._ @virtualize def snippet(arg: Rep[Int]): Rep[Unit] = { - // val array = NewArray[Int](10) - val a = Tensor.zeros[Int](10) - val b = Tensor.ones[Int](10) - val c = a + b - val d = a - b + val a = Tensor.input[Int](10) + val b = Tensor.zeros[Int](10) + val c = Tensor.ones[Int](10) + val d = a + b + val e = a - c // printf("%d", c(0)) // printf("%d", d(0)) - c.show d.show + e.show } } checkWithLogPath("add", driver.code, "cu", driver.setLogPath) @@ -101,7 +84,7 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { @virtualize def snippet(arg: Rep[Int]): Rep[Unit] = { - val a = Tensor.zeros[Int](10) + val a = Tensor.input[Int](10) val c = a.tanh c.show } @@ -115,11 +98,26 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { @virtualize def snippet(arg: Rep[Int]): Rep[Unit] = { - val a = Tensor.ones[Int](10) - val c = a.relu - printf("%d", c(0)) + val a = Tensor.input[Int](10) + val b = a.relu + b.show } } checkWithLogPath("relu", driver.code, "cu", driver.setLogPath) } + + test("input") { + val driver = new CompilerCFusedTensor[Int, Unit] { + import FusedTensorTypeLess._ + + @virtualize + def snippet(arg: Rep[Int]): Rep[Unit] = { + val a = Tensor.input[Int](10) + val b = Tensor.ones[Int](10) + val c = a + b + c.show + } + } + checkWithLogPath("input", driver.code, "cu", driver.setLogPath) + } } \ No newline at end of file From dad771c8f0327d6782aad79aaedb044edae65ed1 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Thu, 2 Dec 2021 16:57:26 -0500 Subject: [PATCH 15/21] split --- .../fused_tensor/fusedTensor.scala | 30 ++++++++++ .../fused_tensor/fusedTensorSimplify.scala | 2 +- .../fused_tensor/fusedTensorSplit.scala | 59 +++++++++++++++++++ .../fused_tensor/split/split.check.cu | 33 +++++++++++ .../transformation/test_fused_tensor.scala | 18 ++++++ 5 files changed, 141 insertions(+), 1 deletion(-) create mode 100644 src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala create mode 100644 src/out/transformer/fused_tensor/split/split.check.cu diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index d5d8904e..fea7e4d9 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -51,6 +51,10 @@ object FusedTensorTypeLess { (new TENSOR(Adapter.g.reflectUnsafe("tensor_input", C(size), 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 = @@ -100,6 +104,18 @@ object FusedTensorTypeLess { // 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) + } } } @@ -174,5 +190,19 @@ trait FusedTensorOps extends Dsl with ArrayOps with CudaOps { 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(y: Rep[Tensor[T]])(implicit __pos: SourceContext): Rep[Tensor[T]] = { + val t = self.concat(tensor(y)) + Wrap[Tensor[T]](t.x) + } } } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala index 669814fc..fc4d618c 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala @@ -22,7 +22,7 @@ abstract class FusedTensorSimplify extends Transformer { import FusedTensorTypeLess._ import PrimitiveTypeLess._ - val tensors = new mutable.HashMap[Backend.Sym, (Node, List[Backend.Sym], Seq[Node])] + // val tensors = new mutable.HashMap[Backend.Sym, (Node, List[Backend.Sym], Seq[Node])] override def transform(n: Node): Backend.Exp = n match { diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala new file mode 100644 index 00000000..f9f4de48 --- /dev/null +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala @@ -0,0 +1,59 @@ +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 FusedTensorSplit extends Transformer { + + override val name = "FusedTensorSplit" + + import BaseTypeLess._ + import PrimitiveTypeLess._ + import ArrayTypeLess._ + import ArrayCPUTypeLess._ + import FusedTensorTypeLess._ + import PrimitiveTypeLess._ + + val splits = new mutable.HashMap[(Backend.Sym, Int), TENSOR] + val results = new mutable.HashMap[Backend.Sym, TENSOR] + + override def transform(n: Node): Backend.Exp = n match { + case Node(s, "tensor_split", (x:Backend.Sym)::(Backend.Const(sz:Seq[Int]))::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val t = new TENSOR(x, useOldMetadata = true) + val t1 = TENSOR(t.size, t.inputs){ i => t.apply(INT(i).x).x } + val t2 = TENSOR(t.size, t.inputs){ i => t.apply(INT(i).x).x } + splits((s, 0)) = t1 + splits((s, 1)) = t2 + // System.out.println("t1: " + t1) + TENSORS(Seq(t1.x, t2.x)).x + + case Node(s, "tensor_result",(x:Backend.Sym)::(Backend.Const(i:Int))::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val t = splits((x, i)) + System.out.println("t:" + t) + System.out.println("s:" + s) + 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 + } + } +} \ No newline at end of file diff --git a/src/out/transformer/fused_tensor/split/split.check.cu b/src/out/transformer/fused_tensor/split/split.check.cu new file mode 100644 index 00000000..dcf64aff --- /dev/null +++ b/src/out/transformer/fused_tensor/split/split.check.cu @@ -0,0 +1,33 @@ +/***************************************** +Emitting C Generated Code +*******************************************/ +#include +#include +#include +#include +/************* Functions **************/ +__global__ void x2(int x3, int x4, int x5) { + int x6 = gridDim.x * blockDim.x; + int x7 = threadIdx.x + blockIdx.x * blockDim.x; + while (x7 < x5) { + x3[x7] = x3[x7] + 1; + x7 = x7 + x6; + } +} +/**************** Snippet ****************/ +void Snippet(int x0) { + int* x1 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); + show_tensor(x2(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); +} +/***************************************** +End of C Generated Code +*******************************************/ +int main(int argc, char *argv[]) { + if (argc != 2) { + printf("usage: %s \n", argv[0]); + return 0; + } + Snippet(atoi(argv[1])); + return 0; +} diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index 3878fd78..615d50e3 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -23,6 +23,7 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { override val passes = List( new FusedTensorFunctional {}, + new FusedTensorSplit {}, new FusedTensorSimplify {}, new FusedTensorVertical {}, new Canonicalize {}, @@ -58,6 +59,23 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { } } + test("split") { + val driver = new CompilerCFusedTensor[Int, Unit] { + import FusedTensorTypeLess._ + + @virtualize + def snippet(arg: Rep[Int]): Rep[Unit] = { + val a = Tensor.input[Int](10) + val b = Tensor.ones[Int](5) + val c = a.split(Seq(5, 5)) + val d = c.result(0) + val e = d + b + e.show + } + } + checkWithLogPath("split", driver.code, "cu", driver.setLogPath) + } + test("add") { val driver = new CompilerCFusedTensor[Int, Unit] { import FusedTensorTypeLess._ From 484e90560fb4073dba9518e2c81358f377be00f1 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Thu, 2 Dec 2021 17:15:04 -0500 Subject: [PATCH 16/21] track input range --- .../fused_tensor/FusedTensorToCuda.scala | 5 +++-- .../fused_tensor/fusedTensor.scala | 20 +++++++++---------- .../fused_tensor/fusedTensorFunctional.scala | 8 ++++---- .../fused_tensor/fusedTensorVertical.scala | 7 ++++--- 4 files changed, 21 insertions(+), 19 deletions(-) diff --git a/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala index 238d00c7..33646ab1 100644 --- a/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala +++ b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala @@ -29,8 +29,9 @@ abstract class FusedTensorToCuda extends Transformer { override def transform(n: Node): Backend.Exp = n match { - case Node(s, "tensor", Backend.Const(sz:Int)::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _) => + case Node(s, "tensor", Backend.Const(sz:Seq[Int])::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _) => implicit val __pos = Adapter.oldSourceMap(s) + val sz1 = sz.sum val arr = new ARRAY(inputs.head) // for now, assume only one input // System.out.println("input: " + inputs.head) @@ -67,7 +68,7 @@ abstract class FusedTensorToCuda extends Transformer { Backend.Const(()) }, manifest[Array[Int]], manifest[Int], manifest[Int]) - (kernel(arr, INT(0), INT(sz), DIM3(0), DIM3(0))).x + (kernel(arr, INT(0), INT(sz1), DIM3(0), DIM3(0))).x case Node(s, "tensor_show", Backend.Sym(x)::_, _) => implicit val pos = Adapter.oldSourceMap(s) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index fea7e4d9..1925ecc3 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -26,28 +26,28 @@ object FusedTensorTypeLess { (new TENSOR(Adapter.g.reflectUnsafe("tensor", C(size), Adapter.g.reify(xn => f(xn))))).withSrcType(__pos, manifest[Int]) }*/ - def TENSOR(size: Int, inputs: Seq[Backend.Sym])(f: Backend.Exp => Backend.Exp)(implicit __pos: SourceContext): TENSOR = { + 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 ZEROS(size: Int)(implicit __pos: SourceContext): TENSOR = { - (new TENSOR(Adapter.g.reflectUnsafe("tensor_zeros", C(size)))).withSrcType(__pos, manifest[Int]) + (new TENSOR(Adapter.g.reflectUnsafe("tensor_zeros", C(Seq(0, 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]) + (new TENSOR(Adapter.g.reflectUnsafe("tensor_ones", C(Seq(0, 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]) + (new TENSOR(Adapter.g.reflectUnsafe("tensor_consts", C(Seq(0, size)), C(num)))).withSrcType(__pos, manifest[Int]) } def INPUT(size: Int)(implicit __pos: SourceContext): TENSOR = { - (new TENSOR(Adapter.g.reflectUnsafe("tensor_input", C(size), C(Seq())))).withSrcType(__pos, manifest[Int]) + (new TENSOR(Adapter.g.reflectUnsafe("tensor_input", C(Seq(0, size)), C(Seq())))).withSrcType(__pos, manifest[Int]) } // used to track input by itself - def INPUT1(size: Int, inputs: Seq[Backend.Sym])(implicit __pos: SourceContext): TENSOR = { + def INPUT1(size: Seq[Int], inputs: Seq[Backend.Sym])(implicit __pos: SourceContext): TENSOR = { (new TENSOR(Adapter.g.reflectUnsafe("tensor_input", C(size), C(inputs)))).withSrcType(__pos, manifest[Int]) } @@ -64,10 +64,10 @@ object FusedTensorTypeLess { if (useOldMetadata) Adapter.oldTypeMap(x) else Adapter.typeMap(x) } - def size: Int = { + def size: Seq[Int] = { gc.get(x.asInstanceOf[Backend.Sym]) match { - case Some(Node(_, s, Backend.Const(size:Int)::_, _)) => size - case Some(Node(_, s, Backend.Const(_)::Backend.Const(size:Int)::_, _)) => size + case Some(Node(_, s, Backend.Const(size:Seq[Int])::_, _)) => size + case Some(Node(_, s, Backend.Const(_)::Backend.Const(size:Seq[Int])::_, _)) => size case a => System.out.println(a); ??? } } @@ -154,7 +154,7 @@ trait FusedTensorOps extends Dsl with ArrayOps with CudaOps { } def apply[T:Numeric:Manifest](size: Int, f: Rep[Int] => Rep[Int])(implicit __pos: SourceContext): Rep[Tensor[T]] = { - Wrap[Tensor[T]](TENSOR(size, Seq())(unwrapFun[Int, Int](f)).x) // is the input correct? + Wrap[Tensor[T]](TENSOR(Seq(size), Seq())(unwrapFun[Int, Int](f)).x) // is the input correct? } } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala index 469e5253..6a69c967 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala @@ -23,19 +23,19 @@ abstract class FusedTensorFunctional extends Transformer { import PrimitiveTypeLess._ override def transform(n: Node): Backend.Exp = n match { - case Node(s, "tensor_zeros", (Backend.Const(sz:Int))::_, _) => + case Node(s, "tensor_zeros", (Backend.Const(sz:Seq[Int]))::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = TENSOR(sz, Seq())(i => INT(0).x) t.x - case Node(s, "tensor_ones", (Backend.Const(sz:Int))::_, _) => + case Node(s, "tensor_ones", (Backend.Const(sz:Seq[Int]))::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = TENSOR(sz, Seq())(i => INT(1).x) t.x - case Node(s, "tensor_consts", (Backend.Const(sz:Int))::(Backend.Const(n:Int))::_, _) => + case Node(s, "tensor_consts", (Backend.Const(sz:Seq[Int]))::(Backend.Const(n:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = TENSOR(sz, Seq())(i => INT(n).x) t.x - case Node(s, "tensor_input", (Backend.Const(sz:Int))::_, _) => + case Node(s, "tensor_input", (Backend.Const(sz:Seq[Int]))::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t1 = INPUT1(sz, Seq(s)) // necessary? t1.x diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala index 37a69d1b..0c7bb45c 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala @@ -27,9 +27,10 @@ abstract class FusedTensorVertical extends Transformer { 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(sz:Int)::_, _) => + case Node(s, "tensor_input", Backend.Const(sz:Seq[Int])::_, _) => implicit val __pos = Adapter.oldSourceMap(s) - val arr = CUDA_MALLOC(sz, manifest[Int]) // allocate CUDA array for input tensors + val size = sz.sum + val arr = CUDA_MALLOC(size, manifest[Int]) // allocate CUDA array for input tensors tensor2arr(s) = arr.x arr.x case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) if tensor2arr.contains(a) => @@ -41,7 +42,7 @@ abstract class FusedTensorVertical extends Transformer { tensors(s) = (n, path, inner) super.transform(n) case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) if tensors.contains(a) => - val (Node(_, _, Backend.Const(szy:Int)::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _), path0, inner0) = tensors(a) + val (Node(_, _, _::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _), path0, inner0) = tensors(a) try { subst(arg) = transform(b) withResetScope(path0, inner0) { From a10236f070fadbee7d07f326318daec4d333e205 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Fri, 3 Dec 2021 03:40:34 -0500 Subject: [PATCH 17/21] split/concat; merge passes --- .../fused_tensor/FusedTensorToCuda.scala | 3 +- .../fused_tensor/fusedTensor.scala | 7 ++ .../fused_tensor/fusedTensorConcat.scala | 61 +++++++++++++++++ .../fused_tensor/fusedTensorFunctional.scala | 67 +++++++++++++++++++ .../fused_tensor/fusedTensorSplit.scala | 12 ++-- .../fused_tensor/split/split.check.cu | 5 +- .../transformation/test_fused_tensor.scala | 16 +++-- 7 files changed, 156 insertions(+), 15 deletions(-) create mode 100644 src/main/scala/lms/transformation/fused_tensor/fusedTensorConcat.scala diff --git a/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala index 33646ab1..e9ed19ae 100644 --- a/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala +++ b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala @@ -55,7 +55,8 @@ abstract class FusedTensorToCuda extends Transformer { // replace input to function argument, tensor lambda to loop index try { subst(inputs.head) = array.x - subst(arg) = UnwrapV(i) + // subst(arg) = UnwrapV(i) + subst(arg) = Unwrap(readVar(i)) traverse(f) } finally { subst -= inputs.head diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index 1925ecc3..a77899ea 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -79,6 +79,13 @@ object FusedTensorTypeLess { } } + def body: Backend.Block = { + gc.get(x.asInstanceOf[Backend.Sym]) match { + case Some(Node(_, "tensor", _::_::(f:Backend.Block)::_, _)) => f + case a => ??? + } + } + def show(implicit __pos: SourceContext): UNIT = { UNIT(Adapter.g.reflectEffect("show_tensor", x)()(Adapter.CTRL)) } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorConcat.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorConcat.scala new file mode 100644 index 00000000..4ab94c2b --- /dev/null +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorConcat.scala @@ -0,0 +1,61 @@ +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 FusedTensorConcat extends Transformer { + + override val name = "FusedTensorConcat" + + import BaseTypeLess._ + import PrimitiveTypeLess._ + import ArrayTypeLess._ + import ArrayCPUTypeLess._ + import FusedTensorTypeLess._ + import PrimitiveTypeLess._ + + val splits = new mutable.HashMap[(Backend.Sym, Int), TENSOR] + val results = new mutable.HashMap[Backend.Sym, TENSOR] + + override def transform(n: Node): Backend.Exp = n match { + case Node(s, "tensor_concat", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val a = new TENSOR(transform(x), useOldMetadata = true) + val b = new TENSOR(transform(y), useOldMetadata = true) + + require(a.size.last == b.size.head && a.inputs == b.inputs, "cannot concat") + + System.out.println("a: " + a.body) + System.out.println("b: " + b.body) + + val Backend.Block(a_arg::Nil, a_r, _, _) = a.body + val Backend.Block(b_arg::Nil, b_r, _, _) = b.body + + val sz = a.size.sum + b.size.sum + val res = TENSOR(Seq(0, sz), a.inputs){ i => + (IF(INT(i) < INT(a.size.sum))(a.apply(i))(b.apply(i))).x + } + res.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 + } + } +} diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala index 6a69c967..30d25184 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala @@ -22,6 +22,8 @@ abstract class FusedTensorFunctional extends Transformer { import FusedTensorTypeLess._ import PrimitiveTypeLess._ + val splits = new mutable.HashMap[(Backend.Sym, Int), TENSOR] // source sym, idx |-> TENSOR + override def transform(n: Node): Backend.Exp = n match { case Node(s, "tensor_zeros", (Backend.Const(sz:Seq[Int]))::_, _) => implicit val pos = Adapter.oldSourceMap(s) @@ -39,6 +41,71 @@ abstract class FusedTensorFunctional extends Transformer { implicit val pos = Adapter.oldSourceMap(s) val t1 = INPUT1(sz, Seq(s)) // necessary? t1.x + case Node(s, "tensor_add", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val a = new TENSOR(transform(x)) + val b = new TENSOR(transform(y)) + val t = TENSOR(a.size, a.inputs ++ b.inputs){ i => + (a.apply(INT(i).x) + b.apply(INT(i).x)).x } + t.x + case Node(s, "tensor_minus", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val a = new TENSOR(transform(x)) + val b = new TENSOR(transform(y)) + val t = TENSOR(a.size, a.inputs ++ b.inputs){ i => + (a.apply(INT(i).x) - b.apply(INT(i).x)).x } + t.x + case Node(s, "tensor_tanh", (x:Backend.Sym)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val t = new TENSOR(transform(x)) + val res = TENSOR(t.size, t.inputs){ i => t.apply(INT(i).x).tanh().x } // ad-hoc!!! + res.x + case Node(s, "tensor_relu", (x:Backend.Sym)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val t = new TENSOR(x) + val res = TENSOR(t.size, t.inputs){ i => + // IF(c: BOOL)(a: => TOP)(b: => TOP) + (IF(t.apply(INT(i).x) < INT(0))(INT(0))(t.apply(INT(i).x))).x } + res.x + case Node(s, "tensor_split", (x:Backend.Sym)::(Backend.Const(sz:Seq[Int]))::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val t = new TENSOR(x) // get tensor in this level + // val sz = t.size + require(sz.sum == t.size.sum, "invalid split pattern") + + val t1 = TENSOR(Seq(0, sz(0)), t.inputs){ i => t.apply(INT(i).x).x } // fixme: sizes are ad-hoc + val t2 = TENSOR(Seq(sz(0), t.size.sum), t.inputs){ i => t.apply(INT(i).x).x } + splits((s, 0)) = t1 + splits((s, 1)) = t2 + // System.out.println("t: " + t) + // System.out.println("t1: " + t1) + // System.out.println("t2: " + t2) + + //TENSORS(Seq(t1.x, t2.x)).x + Backend.Const(()) + + case Node(s, "tensor_result",(x:Backend.Sym)::(Backend.Const(i:Int))::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val t = splits((x, i)) + t.x + case Node(s, "tensor_concat", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => + implicit val pos = Adapter.oldSourceMap(s) + val a = new TENSOR(transform(x)) + val b = new TENSOR(transform(y)) + + require(a.size.last == b.size.head && a.inputs == b.inputs, "cannot concat") + + System.out.println("a: " + a.body) + System.out.println("b: " + b.body) + + val Backend.Block(a_arg::Nil, a_r, _, _) = a.body + val Backend.Block(b_arg::Nil, b_r, _, _) = b.body + + val sz = a.size.sum + b.size.sum + val res = TENSOR(Seq(0, sz), a.inputs){ i => + (IF(INT(i) < INT(a.size.sum))(a.apply(i))(b.apply(i))).x + } + res.x case _ => super.transform(n) } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala index f9f4de48..57f12663 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala @@ -29,20 +29,18 @@ abstract class FusedTensorSplit extends Transformer { case Node(s, "tensor_split", (x:Backend.Sym)::(Backend.Const(sz:Seq[Int]))::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(x, useOldMetadata = true) - val t1 = TENSOR(t.size, t.inputs){ i => t.apply(INT(i).x).x } - val t2 = TENSOR(t.size, t.inputs){ i => t.apply(INT(i).x).x } + require(sz.sum == t.size.sum, "invalid split pattern") + + val t1 = TENSOR(Seq(0, sz(0)), t.inputs){ i => t.apply(INT(i).x).x } // fixme: sizes are ad-hoc + val t2 = TENSOR(Seq(sz(0), t.size.sum), t.inputs){ i => t.apply(INT(i).x).x } splits((s, 0)) = t1 splits((s, 1)) = t2 - // System.out.println("t1: " + t1) TENSORS(Seq(t1.x, t2.x)).x case Node(s, "tensor_result",(x:Backend.Sym)::(Backend.Const(i:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = splits((x, i)) - System.out.println("t:" + t) - System.out.println("s:" + s) t.x - case _ => super.transform(n) } @@ -56,4 +54,4 @@ abstract class FusedTensorSplit extends Transformer { g = null; Adapter.g = null } } -} \ No newline at end of file +} diff --git a/src/out/transformer/fused_tensor/split/split.check.cu b/src/out/transformer/fused_tensor/split/split.check.cu index dcf64aff..b8690541 100644 --- a/src/out/transformer/fused_tensor/split/split.check.cu +++ b/src/out/transformer/fused_tensor/split/split.check.cu @@ -10,7 +10,8 @@ __global__ void x2(int x3, int x4, int x5) { int x6 = gridDim.x * blockDim.x; int x7 = threadIdx.x + blockIdx.x * blockDim.x; while (x7 < x5) { - x3[x7] = x3[x7] + 1; + int x8 = x7; + x3[x7] = x8 < 5 ? x3[x8] + 1 : x3[x8]; x7 = x7 + x6; } } @@ -18,7 +19,7 @@ __global__ void x2(int x3, int x4, int x5) { void Snippet(int x0) { int* x1 = (int*)malloc(0 * sizeof(int)); CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - show_tensor(x2(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); + show_tensor(x2(x1, 0, 20, dim3(0, 1, 1), dim3(0, 1, 1))); } /***************************************** End of C Generated Code diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index 615d50e3..24764a33 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -23,11 +23,13 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { override val passes = List( new FusedTensorFunctional {}, - new FusedTensorSplit {}, - new FusedTensorSimplify {}, + // new FusedTensorSplit {}, + // new FusedTensorSimplify {}, + // new FusedTensorConcat {}, + // new Canonicalize {}, new FusedTensorVertical {}, new Canonicalize {}, - new FusedTensorToCuda {} + new FusedTensorToCuda {}, ) var log_path: String = "" @@ -67,10 +69,14 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { def snippet(arg: Rep[Int]): Rep[Unit] = { val a = Tensor.input[Int](10) val b = Tensor.ones[Int](5) + val c = a.split(Seq(5, 5)) val d = c.result(0) - val e = d + b - e.show + val e = c.result(1) + + val f = d + b + val g = f.concat(e) + g.show } } checkWithLogPath("split", driver.code, "cu", driver.setLogPath) From 2988a08d501c9106482b8b30a1668095b7b2fa78 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Fri, 3 Dec 2021 20:07:35 -0500 Subject: [PATCH 18/21] fix cuda size_t codegen --- src/out/transformer/fused_tensor/add/add.check.cu | 8 +++++--- src/out/transformer/fused_tensor/input/input.check.cu | 6 ++++-- src/out/transformer/fused_tensor/relu/relu.check.cu | 6 ++++-- src/out/transformer/fused_tensor/split/split.check.cu | 6 ++++-- src/out/transformer/fused_tensor/tanh/tanh.check.cu | 6 ++++-- src/test/scala/lms/transformation/test_fused_tensor.scala | 4 ++-- 6 files changed, 23 insertions(+), 13 deletions(-) diff --git a/src/out/transformer/fused_tensor/add/add.check.cu b/src/out/transformer/fused_tensor/add/add.check.cu index c200fc92..cffaeaf4 100644 --- a/src/out/transformer/fused_tensor/add/add.check.cu +++ b/src/out/transformer/fused_tensor/add/add.check.cu @@ -1,7 +1,9 @@ /***************************************** Emitting C Generated Code *******************************************/ +#include #include +#include "cuda_header.h" #include #include #include @@ -25,9 +27,9 @@ __global__ void x8(int x9, int x10, int x11) { /**************** Snippet ****************/ void Snippet(int x0) { int* x1 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - show_tensor(x2(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); - show_tensor(x8(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); + CUDA_CALL(cudaMalloc(&x1, (size_t)(10 * sizeof(int)))); + show_tensor(x2<<>>(x1, 0, 10)); + show_tensor(x8<<>>(x1, 0, 10)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/input/input.check.cu b/src/out/transformer/fused_tensor/input/input.check.cu index dcf64aff..107dcf1e 100644 --- a/src/out/transformer/fused_tensor/input/input.check.cu +++ b/src/out/transformer/fused_tensor/input/input.check.cu @@ -1,7 +1,9 @@ /***************************************** Emitting C Generated Code *******************************************/ +#include #include +#include "cuda_header.h" #include #include #include @@ -17,8 +19,8 @@ __global__ void x2(int x3, int x4, int x5) { /**************** Snippet ****************/ void Snippet(int x0) { int* x1 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - show_tensor(x2(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); + CUDA_CALL(cudaMalloc(&x1, (size_t)(10 * sizeof(int)))); + show_tensor(x2<<>>(x1, 0, 10)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/relu/relu.check.cu b/src/out/transformer/fused_tensor/relu/relu.check.cu index 3dc850f7..592e6eba 100644 --- a/src/out/transformer/fused_tensor/relu/relu.check.cu +++ b/src/out/transformer/fused_tensor/relu/relu.check.cu @@ -1,7 +1,9 @@ /***************************************** Emitting C Generated Code *******************************************/ +#include #include +#include "cuda_header.h" #include #include #include @@ -18,8 +20,8 @@ __global__ void x2(int x3, int x4, int x5) { /**************** Snippet ****************/ void Snippet(int x0) { int* x1 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - show_tensor(x2(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); + CUDA_CALL(cudaMalloc(&x1, (size_t)(10 * sizeof(int)))); + show_tensor(x2<<>>(x1, 0, 10)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/split/split.check.cu b/src/out/transformer/fused_tensor/split/split.check.cu index b8690541..b8539e59 100644 --- a/src/out/transformer/fused_tensor/split/split.check.cu +++ b/src/out/transformer/fused_tensor/split/split.check.cu @@ -1,7 +1,9 @@ /***************************************** Emitting C Generated Code *******************************************/ +#include #include +#include "cuda_header.h" #include #include #include @@ -18,8 +20,8 @@ __global__ void x2(int x3, int x4, int x5) { /**************** Snippet ****************/ void Snippet(int x0) { int* x1 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - show_tensor(x2(x1, 0, 20, dim3(0, 1, 1), dim3(0, 1, 1))); + CUDA_CALL(cudaMalloc(&x1, (size_t)(10 * sizeof(int)))); + show_tensor(x2<<>>(x1, 0, 20)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/tanh/tanh.check.cu b/src/out/transformer/fused_tensor/tanh/tanh.check.cu index 06e36f22..564e3de3 100644 --- a/src/out/transformer/fused_tensor/tanh/tanh.check.cu +++ b/src/out/transformer/fused_tensor/tanh/tanh.check.cu @@ -1,7 +1,9 @@ /***************************************** Emitting C Generated Code *******************************************/ +#include #include +#include "cuda_header.h" #include #include #include @@ -18,8 +20,8 @@ __global__ void x2(int x3, int x4, int x5) { /**************** Snippet ****************/ void Snippet(int x0) { int* x1 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x1, (lms.thirdparty.size_ttypeless$sizet)(10 * sizeof(Int)))); - show_tensor(x2(x1, 0, 10, dim3(0, 1, 1), dim3(0, 1, 1))); + CUDA_CALL(cudaMalloc(&x1, (size_t)(10 * sizeof(int)))); + show_tensor(x2<<>>(x1, 0, 10)); } /***************************************** End of C Generated Code diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index 24764a33..9d51c487 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -8,7 +8,7 @@ import macros.SourceContext import lms.core._ import lms.core.stub._ import lms.thirdparty.{CCodeGenLibs} -// import lms.transformation.tensor.{CCodeGenCudaCustomOps} +import lms.thirdparty.array_computation.{CCodeGenCBLASOps, CCodeGenCudaOps, CCodeGenCuBLAS} import Backend._ @@ -17,7 +17,7 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { abstract class CompilerCFusedTensor[A: Manifest, B: Manifest] extends CompilerC[A,B] with FusedTensorOps { q => - override val codegen = new DslGenC with CCodeGenLibs { + override val codegen = new DslGenCPP with CCodeGenCudaOps with CCodeGenLibs { val IR: q.type = q } From ee30fe15df41a887674025e0b9d373f7da9fb386 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Fri, 3 Dec 2021 21:23:25 -0500 Subject: [PATCH 19/21] add scanning for input tensors --- .../fused_tensor/FusedTensorToCuda.scala | 26 ++++++------ .../fused_tensor/fusedTensor.scala | 2 +- .../fused_tensor/fusedTensorFunctional.scala | 4 +- .../fused_tensor/fusedTensorVertical.scala | 27 +++++++++++-- .../transformer/fused_tensor/add/add.check.cu | 40 +++++++++++-------- .../fused_tensor/input/input.check.cu | 24 ++++++----- .../fused_tensor/relu/relu.check.cu | 26 +++++++----- .../fused_tensor/split/split.check.cu | 26 +++++++----- .../fused_tensor/tanh/tanh.check.cu | 24 ++++++----- 9 files changed, 125 insertions(+), 74 deletions(-) diff --git a/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala index e9ed19ae..cbc83775 100644 --- a/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala +++ b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala @@ -27,34 +27,32 @@ abstract class FusedTensorToCuda extends Transformer { CUDA_MALLOC(size, m) } - override def transform(n: Node): Backend.Exp = n match { case Node(s, "tensor", Backend.Const(sz:Seq[Int])::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _) => implicit val __pos = Adapter.oldSourceMap(s) val sz1 = sz.sum - val arr = new ARRAY(inputs.head) // for now, assume only one input - // System.out.println("input: " + inputs.head) - // System.out.println("arr: " + arr) - // System.out.println("res:" + r) + // input array. for now, assume only one input + val in_arr = new ARRAY(inputs.head) + // 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 array = (new ARRAY(xn(0))).withSrcType(__pos, manifest[Int]) - val value = (new NUM(xn(1))).withSrcType(__pos, manifest[Int]) // not used + 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)) - - // System.out.println("i: " + UnwrapV(i)) - // PRINTF("%d", INT(Unwrap(readVar(i)))) __whileDo(ordering_lt(readVar(i), Wrap[Int](size.x)), { // replace input to function argument, tensor lambda to loop index try { - subst(inputs.head) = array.x + subst(inputs.head) = in_array.x // subst(arg) = UnwrapV(i) subst(arg) = Unwrap(readVar(i)) traverse(f) @@ -62,14 +60,14 @@ abstract class FusedTensorToCuda extends Transformer { subst -= inputs.head subst -= arg } - array(INT(Unwrap(readVar(i)))) = INT(transform(r)) + out_array(INT(Unwrap(readVar(i)))) = INT(transform(r)) i += Wrap[Int](stride.x) }) Backend.Const(()) - }, manifest[Array[Int]], manifest[Int], manifest[Int]) + }, manifest[Array[Int]], manifest[Array[Int]], manifest[Int]) - (kernel(arr, INT(0), INT(sz1), DIM3(0), DIM3(0))).x + (kernel(in_arr, out_arr, INT(sz1), DIM3(0), DIM3(0))).x case Node(s, "tensor_show", Backend.Sym(x)::_, _) => implicit val pos = Adapter.oldSourceMap(s) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index a77899ea..f9e5ad05 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -87,7 +87,7 @@ object FusedTensorTypeLess { } def show(implicit __pos: SourceContext): UNIT = { - UNIT(Adapter.g.reflectEffect("show_tensor", x)()(Adapter.CTRL)) + UNIT(Adapter.g.reflectEffect("show_tensor", x)(x)(Adapter.CTRL)) } def + (y: TENSOR)(implicit __pos: SourceContext): TENSOR = { diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala index 30d25184..1feb2233 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala @@ -95,8 +95,8 @@ abstract class FusedTensorFunctional extends Transformer { require(a.size.last == b.size.head && a.inputs == b.inputs, "cannot concat") - System.out.println("a: " + a.body) - System.out.println("b: " + b.body) + // System.out.println("a: " + a.body) + // System.out.println("b: " + b.body) val Backend.Block(a_arg::Nil, a_r, _, _) = a.body val Backend.Block(b_arg::Nil, b_r, _, _) = b.body diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala index 0c7bb45c..1a2c5155 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala @@ -2,12 +2,16 @@ 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._ @@ -20,6 +24,16 @@ abstract class FusedTensorVertical extends Transformer { import ArrayTypeLess._ import FusedTensorTypeLess._ import CUDATypeLess._ + import CLibTypeLess._ + + 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 virtual tensors syms to node and context val tensors = new mutable.HashMap[Backend.Sym, (Node, List[Backend.Sym], Seq[Node])] @@ -30,9 +44,16 @@ abstract class FusedTensorVertical extends Transformer { case Node(s, "tensor_input", Backend.Const(sz:Seq[Int])::_, _) => implicit val __pos = Adapter.oldSourceMap(s) val size = sz.sum - val arr = CUDA_MALLOC(size, manifest[Int]) // allocate CUDA array for input tensors - tensor2arr(s) = arr.x - arr.x + 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 + Backend.Const(()) 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)) diff --git a/src/out/transformer/fused_tensor/add/add.check.cu b/src/out/transformer/fused_tensor/add/add.check.cu index cffaeaf4..d69cfbad 100644 --- a/src/out/transformer/fused_tensor/add/add.check.cu +++ b/src/out/transformer/fused_tensor/add/add.check.cu @@ -8,28 +8,36 @@ Emitting C Generated Code #include #include /************* Functions **************/ -__global__ void x2(int x3, int x4, int x5) { - int x6 = gridDim.x * blockDim.x; - int x7 = threadIdx.x + blockIdx.x * blockDim.x; - while (x7 < x5) { - x3[x7] = x3[x7]; - x7 = x7 + x6; +__global__ void x4(int x5, int x6, int x7) { + int x8 = gridDim.x * blockDim.x; + int x9 = threadIdx.x + blockIdx.x * blockDim.x; + while (x9 < x7) { + x6[x9] = x2[x9]; + x9 = x9 + x8; } } -__global__ void x8(int x9, int x10, int x11) { - int x12 = gridDim.x * blockDim.x; - int x13 = threadIdx.x + blockIdx.x * blockDim.x; - while (x13 < x11) { - x9[x13] = x9[x13] - 1; - x13 = x13 + x12; +__global__ void x11(int x12, int x13, int x14) { + int x15 = gridDim.x * blockDim.x; + int x16 = threadIdx.x + blockIdx.x * blockDim.x; + while (x16 < x14) { + x13[x16] = x2[x16] - 1; + x16 = x16 + x15; } } /**************** Snippet ****************/ void Snippet(int x0) { - int* x1 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x1, (size_t)(10 * sizeof(int)))); - show_tensor(x2<<>>(x1, 0, 10)); - show_tensor(x8<<>>(x1, 0, 10)); + int* x1 = (int*)malloc(10 * sizeof(int)); + scan_int_array(x1, 10, "input"); + CUDA_CALL(cudaSetDevice(0)); + int* x2 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x2, (size_t)(10 * sizeof(int)))); + CUDA_CALL(cudaMemcpy(x2, x1, (size_t)(10 * sizeof(int)), cudaMemcpyHostToDevice)); + int* x3 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x3, (size_t)(10 * sizeof(int)))); + int* x10 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x10, (size_t)(10 * sizeof(int)))); + show_tensor(x4<<>>(x1, x3, 10)); + show_tensor(x11<<>>(x1, x10, 10)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/input/input.check.cu b/src/out/transformer/fused_tensor/input/input.check.cu index 107dcf1e..d9d7f708 100644 --- a/src/out/transformer/fused_tensor/input/input.check.cu +++ b/src/out/transformer/fused_tensor/input/input.check.cu @@ -8,19 +8,25 @@ Emitting C Generated Code #include #include /************* Functions **************/ -__global__ void x2(int x3, int x4, int x5) { - int x6 = gridDim.x * blockDim.x; - int x7 = threadIdx.x + blockIdx.x * blockDim.x; - while (x7 < x5) { - x3[x7] = x3[x7] + 1; - x7 = x7 + x6; +__global__ void x4(int x5, int x6, int x7) { + int x8 = gridDim.x * blockDim.x; + int x9 = threadIdx.x + blockIdx.x * blockDim.x; + while (x9 < x7) { + x6[x9] = x2[x9] + 1; + x9 = x9 + x8; } } /**************** Snippet ****************/ void Snippet(int x0) { - int* x1 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x1, (size_t)(10 * sizeof(int)))); - show_tensor(x2<<>>(x1, 0, 10)); + int* x1 = (int*)malloc(10 * sizeof(int)); + scan_int_array(x1, 10, "input"); + CUDA_CALL(cudaSetDevice(0)); + int* x2 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x2, (size_t)(10 * sizeof(int)))); + CUDA_CALL(cudaMemcpy(x2, x1, (size_t)(10 * sizeof(int)), cudaMemcpyHostToDevice)); + int* x3 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x3, (size_t)(10 * sizeof(int)))); + show_tensor(x4<<>>(x1, x3, 10)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/relu/relu.check.cu b/src/out/transformer/fused_tensor/relu/relu.check.cu index 592e6eba..b215ebfd 100644 --- a/src/out/transformer/fused_tensor/relu/relu.check.cu +++ b/src/out/transformer/fused_tensor/relu/relu.check.cu @@ -8,20 +8,26 @@ Emitting C Generated Code #include #include /************* Functions **************/ -__global__ void x2(int x3, int x4, int x5) { - int x6 = gridDim.x * blockDim.x; - int x7 = threadIdx.x + blockIdx.x * blockDim.x; - while (x7 < x5) { - int x8 = x3[x7]; - x3[x7] = x8 < 0 ? 0 : x8; - x7 = x7 + x6; +__global__ void x4(int x5, int x6, int x7) { + int x8 = gridDim.x * blockDim.x; + int x9 = threadIdx.x + blockIdx.x * blockDim.x; + while (x9 < x7) { + int x10 = x2[x9]; + x6[x9] = x10 < 0 ? 0 : x10; + x9 = x9 + x8; } } /**************** Snippet ****************/ void Snippet(int x0) { - int* x1 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x1, (size_t)(10 * sizeof(int)))); - show_tensor(x2<<>>(x1, 0, 10)); + int* x1 = (int*)malloc(10 * sizeof(int)); + scan_int_array(x1, 10, "input"); + CUDA_CALL(cudaSetDevice(0)); + int* x2 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x2, (size_t)(10 * sizeof(int)))); + CUDA_CALL(cudaMemcpy(x2, x1, (size_t)(10 * sizeof(int)), cudaMemcpyHostToDevice)); + int* x3 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x3, (size_t)(10 * sizeof(int)))); + show_tensor(x4<<>>(x1, x3, 10)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/split/split.check.cu b/src/out/transformer/fused_tensor/split/split.check.cu index b8539e59..391993b6 100644 --- a/src/out/transformer/fused_tensor/split/split.check.cu +++ b/src/out/transformer/fused_tensor/split/split.check.cu @@ -8,20 +8,26 @@ Emitting C Generated Code #include #include /************* Functions **************/ -__global__ void x2(int x3, int x4, int x5) { - int x6 = gridDim.x * blockDim.x; - int x7 = threadIdx.x + blockIdx.x * blockDim.x; - while (x7 < x5) { - int x8 = x7; - x3[x7] = x8 < 5 ? x3[x8] + 1 : x3[x8]; - x7 = x7 + x6; +__global__ void x4(int x5, int x6, int x7) { + int x8 = gridDim.x * blockDim.x; + int x9 = threadIdx.x + blockIdx.x * blockDim.x; + while (x9 < x7) { + int x10 = x9; + x6[x9] = x10 < 5 ? x2[x10] + 1 : x2[x10]; + x9 = x9 + x8; } } /**************** Snippet ****************/ void Snippet(int x0) { - int* x1 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x1, (size_t)(10 * sizeof(int)))); - show_tensor(x2<<>>(x1, 0, 20)); + int* x1 = (int*)malloc(10 * sizeof(int)); + scan_int_array(x1, 10, "input"); + CUDA_CALL(cudaSetDevice(0)); + int* x2 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x2, (size_t)(10 * sizeof(int)))); + CUDA_CALL(cudaMemcpy(x2, x1, (size_t)(10 * sizeof(int)), cudaMemcpyHostToDevice)); + int* x3 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x3, (size_t)(20 * sizeof(int)))); + show_tensor(x4<<>>(x1, x3, 20)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/tanh/tanh.check.cu b/src/out/transformer/fused_tensor/tanh/tanh.check.cu index 564e3de3..61a983fe 100644 --- a/src/out/transformer/fused_tensor/tanh/tanh.check.cu +++ b/src/out/transformer/fused_tensor/tanh/tanh.check.cu @@ -9,19 +9,25 @@ Emitting C Generated Code #include #include /************* Functions **************/ -__global__ void x2(int x3, int x4, int x5) { - int x6 = gridDim.x * blockDim.x; - int x7 = threadIdx.x + blockIdx.x * blockDim.x; - while (x7 < x5) { - x3[x7] = tanh(x3[x7]); - x7 = x7 + x6; +__global__ void x4(int x5, int x6, int x7) { + int x8 = gridDim.x * blockDim.x; + int x9 = threadIdx.x + blockIdx.x * blockDim.x; + while (x9 < x7) { + x6[x9] = tanh(x2[x9]); + x9 = x9 + x8; } } /**************** Snippet ****************/ void Snippet(int x0) { - int* x1 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x1, (size_t)(10 * sizeof(int)))); - show_tensor(x2<<>>(x1, 0, 10)); + int* x1 = (int*)malloc(10 * sizeof(int)); + scan_int_array(x1, 10, "input"); + CUDA_CALL(cudaSetDevice(0)); + int* x2 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x2, (size_t)(10 * sizeof(int)))); + CUDA_CALL(cudaMemcpy(x2, x1, (size_t)(10 * sizeof(int)), cudaMemcpyHostToDevice)); + int* x3 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x3, (size_t)(10 * sizeof(int)))); + show_tensor(x4<<>>(x1, x3, 10)); } /***************************************** End of C Generated Code From cc050b2541d88cf32001409e5de5184190ec5717 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Mon, 6 Dec 2021 02:10:59 -0500 Subject: [PATCH 20/21] use views; fix kernel call bug --- .../fused_tensor/FusedTensorToCuda.scala | 55 +++++++++++++-- .../fused_tensor/fusedTensor.scala | 45 +++++++----- .../fused_tensor/fusedTensorConcat.scala | 61 ----------------- .../fused_tensor/fusedTensorFunctional.scala | 51 +++++++------- .../fused_tensor/fusedTensorSimplify.scala | 68 ------------------- .../fused_tensor/fusedTensorSplit.scala | 57 ---------------- .../fused_tensor/fusedTensorVertical.scala | 39 ++--------- .../transformer/fused_tensor/add/add.check.cu | 4 +- .../fused_tensor/input/input.check.cu | 2 +- .../fused_tensor/relu/relu.check.cu | 2 +- .../fused_tensor/split/split.check.cu | 4 +- .../fused_tensor/tanh/tanh.check.cu | 2 +- 12 files changed, 112 insertions(+), 278 deletions(-) delete mode 100644 src/main/scala/lms/transformation/fused_tensor/fusedTensorConcat.scala delete mode 100644 src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala delete mode 100644 src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala diff --git a/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala index cbc83775..bc5fea11 100644 --- a/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala +++ b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala @@ -2,12 +2,15 @@ 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._ @@ -21,19 +24,57 @@ abstract class FusedTensorToCuda extends Transformer { 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", Backend.Const(sz:Seq[Int])::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _) => + 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 + System.out.println("xxx:" + 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 sz1 = sz.sum + 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(inputs.head) + System.out.println("hd:" + inputs) + val xxx = tensor2arr(inputs.head.t) + System.out.println("xxx:" + xxx) + val in_arr = new ARRAY(xxx) + 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) @@ -52,12 +93,11 @@ abstract class FusedTensorToCuda extends Transformer { __whileDo(ordering_lt(readVar(i), Wrap[Int](size.x)), { // replace input to function argument, tensor lambda to loop index try { - subst(inputs.head) = in_array.x - // subst(arg) = UnwrapV(i) + subst(inputs.head.t) = in_array.x subst(arg) = Unwrap(readVar(i)) traverse(f) } finally { - subst -= inputs.head + subst -= inputs.head.t subst -= arg } out_array(INT(Unwrap(readVar(i)))) = INT(transform(r)) @@ -69,7 +109,8 @@ abstract class FusedTensorToCuda extends Transformer { (kernel(in_arr, out_arr, INT(sz1), DIM3(0), DIM3(0))).x - case Node(s, "tensor_show", Backend.Sym(x)::_, _) => + + case Node(s, "tensor_show", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) Backend.Const(()) case _ => super.transform(n) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index f9e5ad05..43b988d6 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -20,35 +20,38 @@ object FusedTensorTypeLess { 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: Int)(f: Backend.Exp => Backend.Exp)(implicit __pos: SourceContext): TENSOR = { - (new TENSOR(Adapter.g.reflectUnsafe("tensor", C(size), Adapter.g.reify(xn => f(xn))))).withSrcType(__pos, manifest[Int]) - }*/ + /* 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(Seq(0, size))))).withSrcType(__pos, manifest[Int]) + (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(Seq(0, size))))).withSrcType(__pos, manifest[Int]) + (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(Seq(0, size)), C(num)))).withSrcType(__pos, manifest[Int]) + (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(Seq(0, size)), C(Seq())))).withSrcType(__pos, manifest[Int]) + (new TENSOR(Adapter.g.reflectUnsafe("tensor_input", C(size)))).withSrcType(__pos, manifest[Int]) } // used to track input by itself - def INPUT1(size: Seq[Int], inputs: Seq[Backend.Sym])(implicit __pos: SourceContext): TENSOR = { - (new TENSOR(Adapter.g.reflectUnsafe("tensor_input", C(size), C(inputs)))).withSrcType(__pos, manifest[Int]) + 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 = { @@ -64,30 +67,38 @@ object FusedTensorTypeLess { if (useOldMetadata) Adapter.oldTypeMap(x) else Adapter.typeMap(x) } - def size: Seq[Int] = { + 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(_)::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[Backend.Sym] = { + 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(_, 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 Some(Node(_, "tensor", _::(f:Backend.Block)::_, _)) => f case a => ??? } } def show(implicit __pos: SourceContext): UNIT = { - UNIT(Adapter.g.reflectEffect("show_tensor", x)(x)(Adapter.CTRL)) + UNIT(Adapter.g.reflectWrite("show_tensor", x)(Adapter.CTRL)) } def + (y: TENSOR)(implicit __pos: SourceContext): TENSOR = { @@ -161,7 +172,7 @@ trait FusedTensorOps extends Dsl with ArrayOps with CudaOps { } def apply[T:Numeric:Manifest](size: Int, f: Rep[Int] => Rep[Int])(implicit __pos: SourceContext): Rep[Tensor[T]] = { - Wrap[Tensor[T]](TENSOR(Seq(size), Seq())(unwrapFun[Int, Int](f)).x) // is the input correct? + Wrap[Tensor[T]](TENSOR(Seq(View(null, 0, size)))(unwrapFun[Int, Int](f)).x) // is the input correct? } } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorConcat.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorConcat.scala deleted file mode 100644 index 4ab94c2b..00000000 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorConcat.scala +++ /dev/null @@ -1,61 +0,0 @@ -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 FusedTensorConcat extends Transformer { - - override val name = "FusedTensorConcat" - - import BaseTypeLess._ - import PrimitiveTypeLess._ - import ArrayTypeLess._ - import ArrayCPUTypeLess._ - import FusedTensorTypeLess._ - import PrimitiveTypeLess._ - - val splits = new mutable.HashMap[(Backend.Sym, Int), TENSOR] - val results = new mutable.HashMap[Backend.Sym, TENSOR] - - override def transform(n: Node): Backend.Exp = n match { - case Node(s, "tensor_concat", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => - implicit val pos = Adapter.oldSourceMap(s) - val a = new TENSOR(transform(x), useOldMetadata = true) - val b = new TENSOR(transform(y), useOldMetadata = true) - - require(a.size.last == b.size.head && a.inputs == b.inputs, "cannot concat") - - System.out.println("a: " + a.body) - System.out.println("b: " + b.body) - - val Backend.Block(a_arg::Nil, a_r, _, _) = a.body - val Backend.Block(b_arg::Nil, b_r, _, _) = b.body - - val sz = a.size.sum + b.size.sum - val res = TENSOR(Seq(0, sz), a.inputs){ i => - (IF(INT(i) < INT(a.size.sum))(a.apply(i))(b.apply(i))).x - } - res.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 - } - } -} diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala index 1feb2233..bdd0a43f 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala @@ -25,63 +25,59 @@ abstract class FusedTensorFunctional extends Transformer { val splits = new mutable.HashMap[(Backend.Sym, Int), TENSOR] // source sym, idx |-> TENSOR override def transform(n: Node): Backend.Exp = n match { - case Node(s, "tensor_zeros", (Backend.Const(sz:Seq[Int]))::_, _) => + case Node(s, "tensor_zeros", (Backend.Const(sz:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val t = TENSOR(sz, Seq())(i => INT(0).x) + val t = TENSOR(Seq(View(null, 0, sz)))(i => INT(0).x) t.x - case Node(s, "tensor_ones", (Backend.Const(sz:Seq[Int]))::_, _) => + case Node(s, "tensor_ones", (Backend.Const(sz:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val t = TENSOR(sz, Seq())(i => INT(1).x) + val t = TENSOR(Seq(View(null, 0, sz)))(i => INT(1).x) t.x - case Node(s, "tensor_consts", (Backend.Const(sz:Seq[Int]))::(Backend.Const(n:Int))::_, _) => + case Node(s, "tensor_consts", (Backend.Const(sz:Int))::(Backend.Const(n:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val t = TENSOR(sz, Seq())(i => INT(n).x) + val t = TENSOR(Seq(View(null, 0, sz)))(i => INT(n).x) t.x - case Node(s, "tensor_input", (Backend.Const(sz:Seq[Int]))::_, _) => + case Node(s, "tensor_input", (Backend.Const(sz:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val t1 = INPUT1(sz, Seq(s)) // necessary? + val t1 = INPUT1(Seq(View(s, 0, sz))) // necessary? t1.x case Node(s, "tensor_add", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x)) val b = new TENSOR(transform(y)) - val t = TENSOR(a.size, a.inputs ++ b.inputs){ i => + val t = TENSOR(a.inputs ++ b.inputs){ i => (a.apply(INT(i).x) + b.apply(INT(i).x)).x } t.x case Node(s, "tensor_minus", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x)) val b = new TENSOR(transform(y)) - val t = TENSOR(a.size, a.inputs ++ b.inputs){ i => + val t = TENSOR(a.inputs ++ b.inputs){ i => (a.apply(INT(i).x) - b.apply(INT(i).x)).x } t.x case Node(s, "tensor_tanh", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(transform(x)) - val res = TENSOR(t.size, t.inputs){ i => t.apply(INT(i).x).tanh().x } // ad-hoc!!! + val res = TENSOR(t.inputs){ i => t.apply(INT(i).x).tanh().x } // ad-hoc!!! res.x case Node(s, "tensor_relu", (x:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) - val t = new TENSOR(x) - val res = TENSOR(t.size, t.inputs){ i => + val t = new TENSOR(transform(x)) + val res = TENSOR(t.inputs){ i => // IF(c: BOOL)(a: => TOP)(b: => TOP) (IF(t.apply(INT(i).x) < INT(0))(INT(0))(t.apply(INT(i).x))).x } res.x case Node(s, "tensor_split", (x:Backend.Sym)::(Backend.Const(sz:Seq[Int]))::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = new TENSOR(x) // get tensor in this level - // val sz = t.size - require(sz.sum == t.size.sum, "invalid split pattern") + require(t.inputs.size == 1, "tensor to be splitted must have a single input") + val input = t.inputs(0).t + require(sz.sum == t.size, "invalid split pattern") - val t1 = TENSOR(Seq(0, sz(0)), t.inputs){ i => t.apply(INT(i).x).x } // fixme: sizes are ad-hoc - val t2 = TENSOR(Seq(sz(0), t.size.sum), t.inputs){ i => t.apply(INT(i).x).x } + val t1 = TENSOR(Seq(View(input, 0, sz(0)))){ i => t.apply(INT(i).x).x } // fixme: sizes are ad-hoc + val t2 = TENSOR(Seq(View(input, sz(0), t.size))){ i => t.apply(INT(i).x).x } splits((s, 0)) = t1 splits((s, 1)) = t2 - // System.out.println("t: " + t) - // System.out.println("t1: " + t1) - // System.out.println("t2: " + t2) - - //TENSORS(Seq(t1.x, t2.x)).x Backend.Const(()) case Node(s, "tensor_result",(x:Backend.Sym)::(Backend.Const(i:Int))::_, _) => @@ -92,8 +88,9 @@ abstract class FusedTensorFunctional extends Transformer { implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x)) val b = new TENSOR(transform(y)) - - require(a.size.last == b.size.head && a.inputs == b.inputs, "cannot concat") + + require(a.inputs(0).t == b.inputs(0).t, "cannot concat from different sources") + require(a.inputs(0).to == b.inputs(0).from, "cannot concat unmatched shapes") // System.out.println("a: " + a.body) // System.out.println("b: " + b.body) @@ -101,9 +98,9 @@ abstract class FusedTensorFunctional extends Transformer { val Backend.Block(a_arg::Nil, a_r, _, _) = a.body val Backend.Block(b_arg::Nil, b_r, _, _) = b.body - val sz = a.size.sum + b.size.sum - val res = TENSOR(Seq(0, sz), a.inputs){ i => - (IF(INT(i) < INT(a.size.sum))(a.apply(i))(b.apply(i))).x + val sz = a.size + b.size + val res = TENSOR(Seq(View(a.inputs(0).t, 0, sz))){ i => + (IF(INT(i) < INT(a.size))(a.apply(i))(b.apply(i))).x } res.x case _ => super.transform(n) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala deleted file mode 100644 index fc4d618c..00000000 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSimplify.scala +++ /dev/null @@ -1,68 +0,0 @@ -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 FusedTensorSimplify extends Transformer { - - override val name = "FusedTensorSimplify" - - import BaseTypeLess._ - import PrimitiveTypeLess._ - import ArrayTypeLess._ - import ArrayCPUTypeLess._ - import FusedTensorTypeLess._ - import PrimitiveTypeLess._ - - // 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_add", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => - implicit val pos = Adapter.oldSourceMap(s) - val a = new TENSOR(transform(x), useOldMetadata = true) - val b = new TENSOR(transform(y), useOldMetadata = true) - val t = TENSOR(a.size, a.inputs ++ b.inputs){ i => - (a.apply(INT(i).x) + b.apply(INT(i).x)).x } - t.x - case Node(s, "tensor_minus", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => - implicit val pos = Adapter.oldSourceMap(s) - val a = new TENSOR(transform(x), useOldMetadata = true) - val b = new TENSOR(transform(y), useOldMetadata = true) - val t = TENSOR(a.size, a.inputs ++ b.inputs){ i => - (a.apply(INT(i).x) - b.apply(INT(i).x)).x } - t.x - case Node(s, "tensor_tanh", (x:Backend.Sym)::_, _) => - implicit val pos = Adapter.oldSourceMap(s) - val t = new TENSOR(transform(x), useOldMetadata = true) - val res = TENSOR(t.size, t.inputs){ i => t.apply(INT(i).x).tanh().x } // ad-hoc!!! - res.x - case Node(s, "tensor_relu", (x:Backend.Sym)::_, _) => - implicit val pos = Adapter.oldSourceMap(s) - val t = new TENSOR(x, useOldMetadata = true) - val res = TENSOR(t.size, t.inputs){ i => - // IF(c: BOOL)(a: => TOP)(b: => TOP) - (IF(t.apply(INT(i).x) < INT(0))(INT(0))(t.apply(INT(i).x))).x } - res.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 - } - } -} diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala deleted file mode 100644 index 57f12663..00000000 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorSplit.scala +++ /dev/null @@ -1,57 +0,0 @@ -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 FusedTensorSplit extends Transformer { - - override val name = "FusedTensorSplit" - - import BaseTypeLess._ - import PrimitiveTypeLess._ - import ArrayTypeLess._ - import ArrayCPUTypeLess._ - import FusedTensorTypeLess._ - import PrimitiveTypeLess._ - - val splits = new mutable.HashMap[(Backend.Sym, Int), TENSOR] - val results = new mutable.HashMap[Backend.Sym, TENSOR] - - override def transform(n: Node): Backend.Exp = n match { - case Node(s, "tensor_split", (x:Backend.Sym)::(Backend.Const(sz:Seq[Int]))::_, _) => - implicit val pos = Adapter.oldSourceMap(s) - val t = new TENSOR(x, useOldMetadata = true) - require(sz.sum == t.size.sum, "invalid split pattern") - - val t1 = TENSOR(Seq(0, sz(0)), t.inputs){ i => t.apply(INT(i).x).x } // fixme: sizes are ad-hoc - val t2 = TENSOR(Seq(sz(0), t.size.sum), t.inputs){ i => t.apply(INT(i).x).x } - splits((s, 0)) = t1 - splits((s, 1)) = t2 - TENSORS(Seq(t1.x, t2.x)).x - - case Node(s, "tensor_result",(x:Backend.Sym)::(Backend.Const(i:Int))::_, _) => - implicit val pos = Adapter.oldSourceMap(s) - val t = splits((x, i)) - 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 - } - } -} diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala index 1a2c5155..2e0dd034 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorVertical.scala @@ -2,8 +2,8 @@ package lms.transformation.tensor import scala.annotation.implicitNotFound import scala.collection._ -import scala.collection.mutable.HashMap -import scala.collection.immutable.Set +// import scala.collection.mutable.HashMap +// import scala.collection.immutable.Set import lms.core._ @@ -11,7 +11,6 @@ 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._ @@ -24,46 +23,18 @@ abstract class FusedTensorVertical extends Transformer { import ArrayTypeLess._ import FusedTensorTypeLess._ import CUDATypeLess._ - import CLibTypeLess._ + - 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 virtual tensors syms to node and context val tensors = new mutable.HashMap[Backend.Sym, (Node, List[Backend.Sym], Seq[Node])] - // 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(sz:Seq[Int])::_, _) => - implicit val __pos = Adapter.oldSourceMap(s) - val size = sz.sum - 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 - Backend.Const(()) - 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", _, _) => + case Node(s, "tensor", Backend.Const(inputs:Seq[View])::_, _) => tensors(s) = (n, path, inner) super.transform(n) case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) if tensors.contains(a) => - val (Node(_, _, _::Backend.Const(inputs:Seq[Backend.Sym])::(f@Backend.Block(arg::Nil, r, block, eff))::_, _), path0, inner0) = tensors(a) + val (Node(_, _, _::(f@Backend.Block(arg::Nil, r, block, eff))::_, _), path0, inner0) = tensors(a) try { subst(arg) = transform(b) withResetScope(path0, inner0) { diff --git a/src/out/transformer/fused_tensor/add/add.check.cu b/src/out/transformer/fused_tensor/add/add.check.cu index d69cfbad..f921cffc 100644 --- a/src/out/transformer/fused_tensor/add/add.check.cu +++ b/src/out/transformer/fused_tensor/add/add.check.cu @@ -36,8 +36,8 @@ void Snippet(int x0) { CUDA_CALL(cudaMalloc(&x3, (size_t)(10 * sizeof(int)))); int* x10 = (int*)malloc(0 * sizeof(int)); CUDA_CALL(cudaMalloc(&x10, (size_t)(10 * sizeof(int)))); - show_tensor(x4<<>>(x1, x3, 10)); - show_tensor(x11<<>>(x1, x10, 10)); + show_tensor(x4<<>>(x2, x3, 10)); + show_tensor(x11<<>>(x2, x10, 10)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/input/input.check.cu b/src/out/transformer/fused_tensor/input/input.check.cu index d9d7f708..eef5b491 100644 --- a/src/out/transformer/fused_tensor/input/input.check.cu +++ b/src/out/transformer/fused_tensor/input/input.check.cu @@ -26,7 +26,7 @@ void Snippet(int x0) { CUDA_CALL(cudaMemcpy(x2, x1, (size_t)(10 * sizeof(int)), cudaMemcpyHostToDevice)); int* x3 = (int*)malloc(0 * sizeof(int)); CUDA_CALL(cudaMalloc(&x3, (size_t)(10 * sizeof(int)))); - show_tensor(x4<<>>(x1, x3, 10)); + show_tensor(x4<<>>(x2, x3, 10)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/relu/relu.check.cu b/src/out/transformer/fused_tensor/relu/relu.check.cu index b215ebfd..4e2647df 100644 --- a/src/out/transformer/fused_tensor/relu/relu.check.cu +++ b/src/out/transformer/fused_tensor/relu/relu.check.cu @@ -27,7 +27,7 @@ void Snippet(int x0) { CUDA_CALL(cudaMemcpy(x2, x1, (size_t)(10 * sizeof(int)), cudaMemcpyHostToDevice)); int* x3 = (int*)malloc(0 * sizeof(int)); CUDA_CALL(cudaMalloc(&x3, (size_t)(10 * sizeof(int)))); - show_tensor(x4<<>>(x1, x3, 10)); + show_tensor(x4<<>>(x2, x3, 10)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/split/split.check.cu b/src/out/transformer/fused_tensor/split/split.check.cu index 391993b6..2fa6c321 100644 --- a/src/out/transformer/fused_tensor/split/split.check.cu +++ b/src/out/transformer/fused_tensor/split/split.check.cu @@ -26,8 +26,8 @@ void Snippet(int x0) { CUDA_CALL(cudaMalloc(&x2, (size_t)(10 * sizeof(int)))); CUDA_CALL(cudaMemcpy(x2, x1, (size_t)(10 * sizeof(int)), cudaMemcpyHostToDevice)); int* x3 = (int*)malloc(0 * sizeof(int)); - CUDA_CALL(cudaMalloc(&x3, (size_t)(20 * sizeof(int)))); - show_tensor(x4<<>>(x1, x3, 20)); + CUDA_CALL(cudaMalloc(&x3, (size_t)(10 * sizeof(int)))); + show_tensor(x4<<>>(x2, x3, 10)); } /***************************************** End of C Generated Code diff --git a/src/out/transformer/fused_tensor/tanh/tanh.check.cu b/src/out/transformer/fused_tensor/tanh/tanh.check.cu index 61a983fe..314ad2ae 100644 --- a/src/out/transformer/fused_tensor/tanh/tanh.check.cu +++ b/src/out/transformer/fused_tensor/tanh/tanh.check.cu @@ -27,7 +27,7 @@ void Snippet(int x0) { CUDA_CALL(cudaMemcpy(x2, x1, (size_t)(10 * sizeof(int)), cudaMemcpyHostToDevice)); int* x3 = (int*)malloc(0 * sizeof(int)); CUDA_CALL(cudaMalloc(&x3, (size_t)(10 * sizeof(int)))); - show_tensor(x4<<>>(x1, x3, 10)); + show_tensor(x4<<>>(x2, x3, 10)); } /***************************************** End of C Generated Code From eebd2eaf46abfcb132a50fd07f235df691361cf0 Mon Sep 17 00:00:00 2001 From: Luke Jiang <39086121+luke-jiang@users.noreply.github.com> Date: Mon, 6 Dec 2021 15:00:32 -0500 Subject: [PATCH 21/21] add support for multi-dim split/concat --- .../fused_tensor/FusedTensorToCuda.scala | 8 +-- .../fused_tensor/fusedTensor.scala | 12 +++- .../fused_tensor/fusedTensorFunctional.scala | 67 ++++++++++++++++--- .../fused_tensor/split3/split3.check.cu | 41 ++++++++++++ .../transformation/test_fused_tensor.scala | 22 +++++- 5 files changed, 132 insertions(+), 18 deletions(-) create mode 100644 src/out/transformer/fused_tensor/split3/split3.check.cu diff --git a/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala index bc5fea11..bcdcb7f4 100644 --- a/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala +++ b/src/main/scala/lms/transformation/fused_tensor/FusedTensorToCuda.scala @@ -56,7 +56,6 @@ abstract class FusedTensorToCuda extends Transformer { CUDA_MEMCPY(gpuArray, cpuArray, size, HOST2DEVICE, m) tensor2arr(s) = gpuArray.x - System.out.println("xxx:" + gpuArray.x) gpuArray.x case Node(s, "tensor_apply", (a:Backend.Sym)::(b:Backend.Exp)::_, _) if tensor2arr.contains(a) => @@ -70,11 +69,8 @@ abstract class FusedTensorToCuda extends Transformer { val sz1 = t.size // input array. for now, assume only one input - System.out.println("hd:" + inputs) - val xxx = tensor2arr(inputs.head.t) - System.out.println("xxx:" + xxx) - val in_arr = new ARRAY(xxx) - System.out.println("in_arr:" + in_arr) + 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) diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala index 43b988d6..b475db2b 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensor.scala @@ -131,8 +131,13 @@ object FusedTensorTypeLess { (new TENSOR(Adapter.g.reflect("tensor_result", x, C(i)))).withSrcType(__pos, et) } - def concat(y: TENSOR)(implicit __pos: SourceContext): TENSOR = { + /*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) } } } @@ -218,8 +223,9 @@ trait FusedTensorOps extends Dsl with ArrayOps with CudaOps { Wrap[Tensor[T]](t.x) } - def concat(y: Rep[Tensor[T]])(implicit __pos: SourceContext): Rep[Tensor[T]] = { - val t = self.concat(tensor(y)) + 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) } } diff --git a/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala index bdd0a43f..0907ed61 100644 --- a/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala +++ b/src/main/scala/lms/transformation/fused_tensor/fusedTensorFunctional.scala @@ -74,17 +74,21 @@ abstract class FusedTensorFunctional extends Transformer { val input = t.inputs(0).t require(sz.sum == t.size, "invalid split pattern") - val t1 = TENSOR(Seq(View(input, 0, sz(0)))){ i => t.apply(INT(i).x).x } // fixme: sizes are ad-hoc - val t2 = TENSOR(Seq(View(input, sz(0), t.size))){ i => t.apply(INT(i).x).x } - splits((s, 0)) = t1 - splits((s, 1)) = t2 + val froms = sz.init.scanLeft(0) { _ + _ } + val tos = sz.scanLeft(0) { _ + _ }.tail + val indices = Range(0, sz.length, 1).toList + + froms zip tos zip indices foreach { + case ((from, to), i) => + splits((s, i)) = TENSOR(Seq(View(input, from, to))){ i => t.apply(INT(i).x).x } + } Backend.Const(()) case Node(s, "tensor_result",(x:Backend.Sym)::(Backend.Const(i:Int))::_, _) => implicit val pos = Adapter.oldSourceMap(s) val t = splits((x, i)) t.x - case Node(s, "tensor_concat", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => + /*case Node(s, "tensor_concat", (x:Backend.Sym)::(y:Backend.Sym)::_, _) => implicit val pos = Adapter.oldSourceMap(s) val a = new TENSOR(transform(x)) val b = new TENSOR(transform(y)) @@ -92,9 +96,6 @@ abstract class FusedTensorFunctional extends Transformer { require(a.inputs(0).t == b.inputs(0).t, "cannot concat from different sources") require(a.inputs(0).to == b.inputs(0).from, "cannot concat unmatched shapes") - // System.out.println("a: " + a.body) - // System.out.println("b: " + b.body) - val Backend.Block(a_arg::Nil, a_r, _, _) = a.body val Backend.Block(b_arg::Nil, b_r, _, _) = b.body @@ -102,6 +103,56 @@ abstract class FusedTensorFunctional extends Transformer { val res = TENSOR(Seq(View(a.inputs(0).t, 0, sz))){ i => (IF(INT(i) < INT(a.size))(a.apply(i))(b.apply(i))).x } + res.x*/ + case Node(s, "tensor_concat", (xs:List[Backend.Sym]), _) => + implicit val pos = Adapter.oldSourceMap(s) + + val tensors = xs map { t => new TENSOR(transform(t)) } + val inputs = tensors map { t => + val real_inputs = t.inputs.filter { i => i.t != null } + require(real_inputs.size == 1, "need only 1 input for concat") + real_inputs(0) + } + val sorted = (tensors zip inputs).sortWith(_._2.from < _._2.from) + // System.out.println(sorted) + + val s_tensors = sorted map { _._1 } + val s_inputs = sorted map { _._2 } + + // System.out.println(s_tensors) + // System.out.println(s_inputs) + + val input = s_inputs(0).t + s_inputs foreach { v => require(v.t == input, "cannot concat from different sources") } + + // TODO: rewrite me + var end = s_inputs.head.to + for (v <- s_inputs.tail) { + if (v.from != end) { + require(false, "cannot merge range") + } else { + end = v.to + } + } + // System.out.println(end) + + // val gaps = sorted.init map { _._2.to } + val gaps = s_inputs.init map { _.to } + // System.out.println(gaps) + + val res = TENSOR(Seq(View(input, 0, end))){ i => + /* + val inner1 = sorted(2)._1.apply(i) + val inner = IF(INT(i) < INT(gaps(1)))(sorted(1)._1.apply(i))(inner1) + (IF(INT(i) < INT(gaps(0)))(sorted(0)._1.apply(i))(inner)).x*/ + var inner = s_tensors(gaps.size).apply(i) + var j = gaps.size - 1 + while (j >= 0) { + inner = INT(IF(INT(i) < INT(gaps(j)))(s_tensors(j).apply(i))(inner)) + j = j - 1 + } + inner.x + } res.x case _ => super.transform(n) } diff --git a/src/out/transformer/fused_tensor/split3/split3.check.cu b/src/out/transformer/fused_tensor/split3/split3.check.cu new file mode 100644 index 00000000..7100d43e --- /dev/null +++ b/src/out/transformer/fused_tensor/split3/split3.check.cu @@ -0,0 +1,41 @@ +/***************************************** +Emitting C Generated Code +*******************************************/ +#include +#include +#include "cuda_header.h" +#include +#include +#include +/************* Functions **************/ +__global__ void x4(int x5, int x6, int x7) { + int x8 = gridDim.x * blockDim.x; + int x9 = threadIdx.x + blockIdx.x * blockDim.x; + while (x9 < x7) { + x6[x9] = x2[x9]; + x9 = x9 + x8; + } +} +/**************** Snippet ****************/ +void Snippet(int x0) { + int* x1 = (int*)malloc(15 * sizeof(int)); + scan_int_array(x1, 15, "input"); + CUDA_CALL(cudaSetDevice(0)); + int* x2 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x2, (size_t)(15 * sizeof(int)))); + CUDA_CALL(cudaMemcpy(x2, x1, (size_t)(15 * sizeof(int)), cudaMemcpyHostToDevice)); + int* x3 = (int*)malloc(0 * sizeof(int)); + CUDA_CALL(cudaMalloc(&x3, (size_t)(15 * sizeof(int)))); + show_tensor(x4<<>>(x2, x3, 15)); +} +/***************************************** +End of C Generated Code +*******************************************/ +int main(int argc, char *argv[]) { + if (argc != 2) { + printf("usage: %s \n", argv[0]); + return 0; + } + Snippet(atoi(argv[1])); + return 0; +} diff --git a/src/test/scala/lms/transformation/test_fused_tensor.scala b/src/test/scala/lms/transformation/test_fused_tensor.scala index 9d51c487..aa5f289b 100644 --- a/src/test/scala/lms/transformation/test_fused_tensor.scala +++ b/src/test/scala/lms/transformation/test_fused_tensor.scala @@ -75,13 +75,33 @@ class FixedSizeFusedTensorTest extends TutorialFunSuite { val e = c.result(1) val f = d + b - val g = f.concat(e) + val g = f.concat(Seq(e)) g.show } } checkWithLogPath("split", driver.code, "cu", driver.setLogPath) } + test("split3") { + val driver = new CompilerCFusedTensor[Int, Unit] { + import FusedTensorTypeLess._ + + @virtualize + def snippet(arg: Rep[Int]): Rep[Unit] = { + val a = Tensor.input[Int](15) + + val b = a.split(Seq(5, 5, 5)) + val c = b.result(0) + val d = b.result(1) + val e = b.result(2) + + val g = c.concat(Seq(e, d)) + g.show + } + } + checkWithLogPath("split3", driver.code, "cu", driver.setLogPath) + } + test("add") { val driver = new CompilerCFusedTensor[Int, Unit] { import FusedTensorTypeLess._