diff --git a/include/tvm/buffer.h b/include/tvm/buffer.h index 2e4d7debcf42..7f8f5c0cc131 100644 --- a/include/tvm/buffer.h +++ b/include/tvm/buffer.h @@ -61,7 +61,7 @@ class BufferNode : public Node { /*! \brief optional name of the buffer */ std::string name; /*! \brief The pointer to the head of the data */ - Var ptr; + Var data; /*! \brief The shape of the buffer */ Array shape; /*! @@ -77,7 +77,7 @@ class BufferNode : public Node { void VisitAttrs(AttrVisitor* v) final { v->Visit("name", &name); - v->Visit("ptr", &ptr); + v->Visit("data", &data); v->Visit("shape", &shape); v->Visit("strides", &strides); v->Visit("dtype", &dtype); diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py index a1e2d4ff483b..c676e5cfeb67 100644 --- a/python/tvm/__init__.py +++ b/python/tvm/__init__.py @@ -17,3 +17,4 @@ from ._base import TVMError from .api import * +from .build import build diff --git a/python/tvm/api.py b/python/tvm/api.py index 7c7a5b33c4f9..85009186646a 100644 --- a/python/tvm/api.py +++ b/python/tvm/api.py @@ -145,7 +145,7 @@ def Buffer(shape, dtype=None, name="buffer", ptr=None, strides=None): - """Create a new buffer + """Create a new symbolic buffer Parameters ---------- diff --git a/python/tvm/build.py b/python/tvm/build.py new file mode 100644 index 000000000000..407a1ba146aa --- /dev/null +++ b/python/tvm/build.py @@ -0,0 +1,83 @@ +"""The build pipeline in python. + +Eventually some of these pipelines will be moved to C++. +But the first pipeline will be kept in python for ease of change and evolving. +""" +# pylint: disable=invalid-name, no-member, too-many-locals, too-many-arguments + +from . import api +from . import tensor +from . import schedule +from . import expr +from . import ir_pass +from . import codegen + +def build(sch, + args, + target, + name="default_function", + binds=None, + record_codes=None): + """Build a function with arguments as signiture. + + Parameters + ---------- + sch : tvm.Schedule + The schedule to be builded + + args : list of Buffer or Tensor or Var + The argument lists to the function. + + target : str + The target of the compilation. + + name : str + The name of result function. + + binds : dict, optional + Dictionary that maps the binding of symbolic buffer to Tensor. + By default, a new buffer is created for each tensor in the argument. + + Returns + ------- + f : Function, or pair of functions + The result function. + If the function requires host space allocation, + a pair of functions will be returned. + """ + binds = {} if binds is None else binds.copy() + arg_list = [] + for x in args: + if isinstance(x, tensor.Tensor): + buf = api.Buffer(x.shape, dtype=x.dtype, name=x.op.name) + assert x not in binds + binds[x] = buf + arg_list.append(buf) + elif isinstance(x, schedule.Buffer): + arg_list.append(x) + elif isinstance(x, expr.Var): + arg_list.append(x) + else: + raise ValueError("args must be Tensor, Buffer or Var") + + # lowering + bounds = schedule.InferBound(sch) + stmt = ir_pass.ScheduleOps(sch, bounds) + stmt = ir_pass.StorageFlatten(stmt, binds) + stmt = ir_pass.Simplify(stmt) + fapi = codegen.MakeAPI(stmt, name, arg_list, len(arg_list)) + fsplits = codegen.SplitHostDevice(fapi) + + if record_codes is not None: + output_ssa = False + for i, f in enumerate(fsplits): + t = target if i >= 1 else "c" + record_codes.append(codegen.CompileToC(f, output_ssa, t)) + + if target == "cuda": + ret = codegen.BuildNVRTC(fsplits, "stackvm") + elif target == "opencl": + ret = codegen.BuildOpenCL(fsplits, "stackvm") + else: + raise ValueError("Unknown target %s" % target) + return ret diff --git a/python/tvm/collections.py b/python/tvm/collections.py index c24b1d81b58f..f92275f06b05 100644 --- a/python/tvm/collections.py +++ b/python/tvm/collections.py @@ -58,12 +58,6 @@ class IterVar(NodeBase, _expr.ExprOp): pass -@register_node -class Buffer(NodeBase): - """Represent a Buffer in TVM.""" - pass - - @register_node class LoweredFunc(NodeBase): """Represent a LoweredFunc in TVM.""" diff --git a/python/tvm/schedule.py b/python/tvm/schedule.py index da5413e625f9..41a6afded977 100644 --- a/python/tvm/schedule.py +++ b/python/tvm/schedule.py @@ -5,6 +5,11 @@ from . import _api_internal from . import tensor as _tensor +@register_node +class Buffer(NodeBase): + """Represent a Buffer in TVM.""" + pass + @register_node class Split(NodeBase): """Split operation on axis.""" diff --git a/src/codegen/make_api.cc b/src/codegen/make_api.cc index 227faf37f410..3c1324a9aa6f 100644 --- a/src/codegen/make_api.cc +++ b/src/codegen/make_api.cc @@ -138,9 +138,9 @@ LoweredFunc MakeAPI(Stmt body, UIntImm::make(UInt(16), dtype.lanes())); seq_init.emplace_back(AssertStmt::make(cond, type_err_msg.str())); // Data Field - if (f_push(buf->ptr, TVMArrayGet(Handle(), v_arg, intrinsic::kData), + if (f_push(buf->data, TVMArrayGet(Handle(), v_arg, intrinsic::kData), v_arg->name_hint + ".data")) { - Var vptr(buf->ptr); + Var vptr(buf->data); handle_data_type.Set(vptr, make_const(buf->dtype, 0)); } // shape field diff --git a/src/lang/buffer.cc b/src/lang/buffer.cc index b44bca783834..8bbff8693921 100644 --- a/src/lang/buffer.cc +++ b/src/lang/buffer.cc @@ -45,23 +45,23 @@ inline Expr BufferOffset(const BufferNode* n, Array index) { Expr Buffer::MakeLoad(Array index) const { const BufferNode* n = operator->(); - return ir::Load::make(n->dtype, n->ptr, BufferOffset(n, index)); + return ir::Load::make(n->dtype, n->data, BufferOffset(n, index)); } Stmt Buffer::MakeStore(Array index, Expr value) const { const BufferNode* n = operator->(); CHECK_EQ(value.type(), n->dtype); - return ir::Store::make(n->ptr, value, BufferOffset(n, index)); + return ir::Store::make(n->data, value, BufferOffset(n, index)); } Buffer BufferNode::make(std::string name, - Var ptr, + Var data, Array shape, Array strides, Type dtype) { auto n = std::make_shared(); n->name = name; - n->ptr = ptr; + n->data = data; n->shape = shape; n->strides = strides; n->dtype = dtype; diff --git a/src/pass/storage_flatten.cc b/src/pass/storage_flatten.cc index 6058b6907fe7..ab344fcc0f3c 100644 --- a/src/pass/storage_flatten.cc +++ b/src/pass/storage_flatten.cc @@ -138,7 +138,7 @@ class StorageFlattener : public IRMutator { buf_map_[key].released = true; return Allocate::make( - e.buffer->ptr, e.buffer->dtype, e.buffer->shape, + e.buffer->data, e.buffer->dtype, e.buffer->shape, make_const(Bool(e.buffer->dtype.lanes()), true), body); } } diff --git a/tests/python/integration/test_ewise.py b/tests/python/integration/test_ewise.py new file mode 100644 index 000000000000..0395d633b5ed --- /dev/null +++ b/tests/python/integration/test_ewise.py @@ -0,0 +1,44 @@ +import tvm +import numpy as np + +def test_add(): + # graph + n = tvm.Var('n') + A = tvm.placeholder((n,), name='A') + B = tvm.placeholder((n,), name='B') + C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C') + # schedule + s = tvm.Schedule(C.op) + # create iter var and assign them tags. + num_thread = 256 + block_x = tvm.IterVar(thread_tag="blockIdx.x") + thread_x = tvm.IterVar((0, num_thread), thread_tag="threadIdx.x") + _, x = s[C].split(C.op.axis[0], factor=num_thread, outer=block_x) + _, x = s[C].split(x, outer=thread_x) + + # one line to build the function. + codes = [] + fadd = tvm.build(s, args=[A, B, C], + target="cuda", name="myadd", + record_codes=codes) + for c in codes: + print(c) + + # call the function + num_device = 1 + for i in range(num_device): + ctx = tvm.gpu(i) + if not ctx.enabled: + continue + # launch the kernel. + n = 1027 + a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx) + b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx) + c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx) + fadd(a, b, c) + np.testing.assert_allclose( + c.asnumpy(), a.asnumpy() + b.asnumpy()) + + +if __name__ == "__main__": + test_add() diff --git a/tests/python/test_codegen_device.py b/tests/python/unittest/test_codegen_device.py similarity index 100% rename from tests/python/test_codegen_device.py rename to tests/python/unittest/test_codegen_device.py diff --git a/tests/python/test_codegen_makeapi.py b/tests/python/unittest/test_codegen_makeapi.py similarity index 93% rename from tests/python/test_codegen_makeapi.py rename to tests/python/unittest/test_codegen_makeapi.py index ebe6f4e63da5..fd6522a2d50c 100644 --- a/tests/python/test_codegen_makeapi.py +++ b/tests/python/unittest/test_codegen_makeapi.py @@ -18,7 +18,7 @@ def test_makeapi(): stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, B:Bb, C:Cb}) num_packed_args = 2 f = tvm.codegen.MakeAPI(stmt, "myadd", [n, Ab, Bb, Cb], num_packed_args) - assert(f.handle_data_type[Ab.ptr].dtype == Ab.dtype) + assert(f.handle_data_type[Ab.data].dtype == Ab.dtype) assert(len(f.args) == 5) output_ssa = False diff --git a/tests/python/test_lang_basic.py b/tests/python/unittest/test_lang_basic.py similarity index 100% rename from tests/python/test_lang_basic.py rename to tests/python/unittest/test_lang_basic.py diff --git a/tests/python/test_lang_buffer.py b/tests/python/unittest/test_lang_buffer.py similarity index 85% rename from tests/python/test_lang_buffer.py rename to tests/python/unittest/test_lang_buffer.py index b1b6b7d9b8fb..c2a0affc5fcc 100644 --- a/tests/python/test_lang_buffer.py +++ b/tests/python/unittest/test_lang_buffer.py @@ -7,7 +7,7 @@ def test_buffer(): Ab = tvm.Buffer((m, n), tvm.float32) Bb = tvm.Buffer((n, l), tvm.float32) - assert isinstance(Ab, tvm.collections.Buffer) + assert isinstance(Ab, tvm.schedule.Buffer) assert Ab.dtype == tvm.float32 assert tuple(Ab.shape) == (m, n) diff --git a/tests/python/test_lang_container.py b/tests/python/unittest/test_lang_container.py similarity index 100% rename from tests/python/test_lang_container.py rename to tests/python/unittest/test_lang_container.py diff --git a/tests/python/test_lang_schedule.py b/tests/python/unittest/test_lang_schedule.py similarity index 100% rename from tests/python/test_lang_schedule.py rename to tests/python/unittest/test_lang_schedule.py diff --git a/tests/python/test_lang_tensor.py b/tests/python/unittest/test_lang_tensor.py similarity index 100% rename from tests/python/test_lang_tensor.py rename to tests/python/unittest/test_lang_tensor.py diff --git a/tests/python/test_pass_basic.py b/tests/python/unittest/test_pass_basic.py similarity index 100% rename from tests/python/test_pass_basic.py rename to tests/python/unittest/test_pass_basic.py diff --git a/tests/python/test_pass_inline.py b/tests/python/unittest/test_pass_inline.py similarity index 100% rename from tests/python/test_pass_inline.py rename to tests/python/unittest/test_pass_inline.py diff --git a/tests/python/test_pass_schedule_ops.py b/tests/python/unittest/test_pass_schedule_ops.py similarity index 100% rename from tests/python/test_pass_schedule_ops.py rename to tests/python/unittest/test_pass_schedule_ops.py diff --git a/tests/python/test_pass_storage_flatten.py b/tests/python/unittest/test_pass_storage_flatten.py similarity index 100% rename from tests/python/test_pass_storage_flatten.py rename to tests/python/unittest/test_pass_storage_flatten.py diff --git a/tests/python/test_runtime_ndarray.py b/tests/python/unittest/test_runtime_ndarray.py similarity index 100% rename from tests/python/test_runtime_ndarray.py rename to tests/python/unittest/test_runtime_ndarray.py diff --git a/tests/python/test_runtime_packed_func.py b/tests/python/unittest/test_runtime_packed_func.py similarity index 100% rename from tests/python/test_runtime_packed_func.py rename to tests/python/unittest/test_runtime_packed_func.py diff --git a/tests/python/test_runtime_stack_vm.py b/tests/python/unittest/test_runtime_stack_vm.py similarity index 85% rename from tests/python/test_runtime_stack_vm.py rename to tests/python/unittest/test_runtime_stack_vm.py index 435df5faadf2..363473661a3a 100644 --- a/tests/python/test_runtime_stack_vm.py +++ b/tests/python/unittest/test_runtime_stack_vm.py @@ -37,8 +37,8 @@ def test_stack_vm_loop(): stmt = tvm.make.For( i, 0, n - 1, 0, 0, tvm.make.Block( - tvm.make.Store(Ab.ptr, - tvm.make.Load(dtype, Ab.ptr, i) + 1, + tvm.make.Store(Ab.data, + tvm.make.Load(dtype, Ab.data, i) + 1, i + 1), tvm.make.Evaluate(tvm_call_global("tvm_stack_vm_print", i)))) print(stmt) @@ -59,10 +59,10 @@ def test_stack_vm_cond(): i, 0, n - 1, 0, 0, tvm.make.IfThenElse( tvm.make.EQ(i, 4), - tvm.make.Store(Ab.ptr, - tvm.make.Load(dtype, Ab.ptr, i) + 1, i + 1), - tvm.make.Store(Ab.ptr, - tvm.make.Load(dtype, Ab.ptr, i) + 2, i + 1))) + tvm.make.Store(Ab.data, + tvm.make.Load(dtype, Ab.data, i) + 1, i + 1), + tvm.make.Store(Ab.data, + tvm.make.Load(dtype, Ab.data, i) + 2, i + 1))) print(stmt) fapi = tvm.codegen.MakeAPI(stmt, "test", [Ab], 1) f = tvm.codegen.BuildStackVM(fapi) diff --git a/tests/python/test_schedule_bound_inference.py b/tests/python/unittest/test_schedule_bound_inference.py similarity index 100% rename from tests/python/test_schedule_bound_inference.py rename to tests/python/unittest/test_schedule_bound_inference.py diff --git a/tests/travis/run_test.sh b/tests/travis/run_test.sh index 5c332a5553c3..a64a39ea50fa 100755 --- a/tests/travis/run_test.sh +++ b/tests/travis/run_test.sh @@ -38,10 +38,10 @@ fi if [ ${TASK} == "python_test" ] || [ ${TASK} == "all_test" ]; then make all || exit -1 if [ ${TRAVIS_OS_NAME} == "osx" ]; then - python -m nose -v tests/python/ || exit -1 - python3 -m nose -v tests/python/ || exit -1 + python -m nose -v tests/python/unittest || exit -1 + python3 -m nose -v tests/python/unittest || exit -1 else - nosetests -v tests/python/ || exit -1 - nosetests3 -v tests/python/ || exit -1 + nosetests -v tests/python/unittest || exit -1 + nosetests3 -v tests/python/unittest || exit -1 fi fi