From d4af7ad6ab6fe0aa75dfae899ec65a3d4e78fbd7 Mon Sep 17 00:00:00 2001
From: Tianqi Chen <tqchen@users.noreply.github.com>
Date: Tue, 31 Jan 2017 14:40:47 -0800
Subject: [PATCH] [TEST/PYTHON] Add unittest folder, add a build pipeline.
Rename Buffer.ptr to Buffer.data to be consistent with Array. (#29)
---
include/tvm/buffer.h | 4 +-
python/tvm/__init__.py | 1 +
python/tvm/api.py | 2 +-
python/tvm/build.py | 83 +++++++++++++++++++
python/tvm/collections.py | 6 --
python/tvm/schedule.py | 5 ++
src/codegen/make_api.cc | 4 +-
src/lang/buffer.cc | 8 +-
src/pass/storage_flatten.cc | 2 +-
tests/python/integration/test_ewise.py | 44 ++++++++++
.../{ => unittest}/test_codegen_device.py | 0
.../{ => unittest}/test_codegen_makeapi.py | 2 +-
.../python/{ => unittest}/test_lang_basic.py | 0
.../python/{ => unittest}/test_lang_buffer.py | 2 +-
.../{ => unittest}/test_lang_container.py | 0
.../{ => unittest}/test_lang_schedule.py | 0
.../python/{ => unittest}/test_lang_tensor.py | 0
.../python/{ => unittest}/test_pass_basic.py | 0
.../python/{ => unittest}/test_pass_inline.py | 0
.../{ => unittest}/test_pass_schedule_ops.py | 0
.../test_pass_storage_flatten.py | 0
.../{ => unittest}/test_runtime_ndarray.py | 0
.../test_runtime_packed_func.py | 0
.../{ => unittest}/test_runtime_stack_vm.py | 12 +--
.../test_schedule_bound_inference.py | 0
tests/travis/run_test.sh | 8 +-
26 files changed, 155 insertions(+), 28 deletions(-)
create mode 100644 python/tvm/build.py
create mode 100644 tests/python/integration/test_ewise.py
rename tests/python/{ => unittest}/test_codegen_device.py (100%)
rename tests/python/{ => unittest}/test_codegen_makeapi.py (93%)
rename tests/python/{ => unittest}/test_lang_basic.py (100%)
rename tests/python/{ => unittest}/test_lang_buffer.py (85%)
rename tests/python/{ => unittest}/test_lang_container.py (100%)
rename tests/python/{ => unittest}/test_lang_schedule.py (100%)
rename tests/python/{ => unittest}/test_lang_tensor.py (100%)
rename tests/python/{ => unittest}/test_pass_basic.py (100%)
rename tests/python/{ => unittest}/test_pass_inline.py (100%)
rename tests/python/{ => unittest}/test_pass_schedule_ops.py (100%)
rename tests/python/{ => unittest}/test_pass_storage_flatten.py (100%)
rename tests/python/{ => unittest}/test_runtime_ndarray.py (100%)
rename tests/python/{ => unittest}/test_runtime_packed_func.py (100%)
rename tests/python/{ => unittest}/test_runtime_stack_vm.py (85%)
rename tests/python/{ => unittest}/test_schedule_bound_inference.py (100%)
diff --git a/include/tvm/buffer.h b/include/tvm/buffer.h
index 2e4d7debc..7f8f5c0cc 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<Expr> 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 a1e2d4ff4..c676e5cfe 100644
--- a/python/tvm/__init__.py
+++ b/python/tvm/__init__.py
@@ -17,3 +17,4 @@ from .ndarray import cpu, gpu, opencl, init_opencl, cl
from ._base import TVMError
from .api import *
+from .build import build
diff --git a/python/tvm/api.py b/python/tvm/api.py
index 7c7a5b33c..850091866 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 000000000..407a1ba14
--- /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 c24b1d81b..f92275f06 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 da5413e62..41a6afded 100644
--- a/python/tvm/schedule.py
+++ b/python/tvm/schedule.py
@@ -5,6 +5,11 @@ from ._ctypes._node import NodeBase, register_node
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 227faf37f..3c1324a9a 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 b44bca783..8bbff8693 100644
--- a/src/lang/buffer.cc
+++ b/src/lang/buffer.cc
@@ -45,23 +45,23 @@ inline Expr BufferOffset(const BufferNode* n, Array<Expr> index) {
Expr Buffer::MakeLoad(Array<Expr> 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<Expr> 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<Expr> shape,
Array<Expr> strides,
Type dtype) {
auto n = std::make_shared<BufferNode>();
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 6058b6907..ab344fcc0 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 000000000..0395d633b
--- /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 ebe6f4e63..fd6522a2d 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 b1b6b7d9b..c2a0affc5 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 435df5faa..363473661 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 5c332a555..a64a39ea5 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
--
GitLab