From 4b81b227196973c3df6e153c3b2f58d3c742b9b6 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Sun, 20 Jan 2019 19:52:13 -0800 Subject: [PATCH 01/31] a preliminary version is done? --- .gitignore | 1 + python/tvm/hybrid/__init__.py | 2 +- python/tvm/hybrid/api.py | 4 + src/api/api_codegen.cc | 15 + src/codegen/build_module.cc | 2 + src/codegen/codegen_hybrid.cc | 508 ++++++++++++++++++++ src/codegen/codegen_hybrid.h | 171 +++++++ tests/python/unittest/test_hybrid_script.py | 2 + 8 files changed, 704 insertions(+), 1 deletion(-) create mode 100644 src/codegen/codegen_hybrid.cc create mode 100644 src/codegen/codegen_hybrid.h diff --git a/.gitignore b/.gitignore index 04dad2039860..ffca5ec16b0f 100644 --- a/.gitignore +++ b/.gitignore @@ -167,6 +167,7 @@ cscope* # vim temporary files *.swp *.swo +.ycm_extra_conf.py # TVM generated code perf diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index 6c137490c38e..ac65f4b8e2c4 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -7,5 +7,5 @@ 2. Developers can build HalideIR by writing Python code. """ -from .api import script +from .api import script, dump from .parser import parse_python diff --git a/python/tvm/hybrid/api.py b/python/tvm/hybrid/api.py index d43217ca5dfc..8b01cb8f359d 100644 --- a/python/tvm/hybrid/api.py +++ b/python/tvm/hybrid/api.py @@ -3,6 +3,7 @@ from .._ffi.base import decorate from .. import _api_internal as _tvm_internal +from .. import codegen as _codegen from ..tensor import Tensor from .parser import parse_python @@ -41,3 +42,6 @@ def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring return value return decorate(pyfunc, wrapped_func) + +def dump(ir): + return _codegen._HybridDump(ir) diff --git a/src/api/api_codegen.cc b/src/api/api_codegen.cc index 372cd0e262b7..ad91e577e158 100644 --- a/src/api/api_codegen.cc +++ b/src/api/api_codegen.cc @@ -9,6 +9,8 @@ #include #include +#include "../codegen/codegen_hybrid.h" + namespace tvm { namespace codegen { @@ -25,5 +27,18 @@ TVM_REGISTER_API("module._PackImportsToC") .set_body([](TVMArgs args, TVMRetValue *ret) { *ret = PackImportsToC(args[0], args[1]); }); + +TVM_REGISTER_API("codegen._HybridDump") +.set_body([](TVMArgs args, TVMRetValue *ret) { + Stmt stmt; + if (args[0].IsNodeType()) { + stmt = args[0]; + } else if (args[0].IsNodeType()) { + stmt = Evaluate::make(args[0]); + } + CodeGenHybrid generator; + generator.PrintStmt(stmt); + *ret = generator.Finish(); + }); } // namespace codegen } // namespace tvm diff --git a/src/codegen/build_module.cc b/src/codegen/build_module.cc index 32bb5f9d6617..f9190123a0a9 100644 --- a/src/codegen/build_module.cc +++ b/src/codegen/build_module.cc @@ -103,6 +103,8 @@ Target CreateTarget(const std::string& target_name, t->device_type = kDLCPU; } else if (target_name == "ext_dev") { t->device_type = kDLExtDev; + } else if (target_name == "hybrid") { + t->device_type = kDLCPU; } else { LOG(ERROR) << "Unknown target name " << target_name; return target::stackvm(); diff --git a/src/codegen/codegen_hybrid.cc b/src/codegen/codegen_hybrid.cc new file mode 100644 index 000000000000..4fde9618dc28 --- /dev/null +++ b/src/codegen/codegen_hybrid.cc @@ -0,0 +1,508 @@ +/*! Copyright (c) 2019 by Contributors + * \file codegen_hybrid.cc + */ +#include +#include +#include "codegen_hybrid.h" +#include "../pass/ir_util.h" +#include "../arithmetic/compute_expr.h" + +namespace tvm { +namespace codegen { + +using namespace ir; + +void CodeGenHybrid::InitFuncState(LoweredFunc f) { + alloc_storage_scope_.clear(); + handle_data_type_.clear(); + CodeGenSourceBase::ClearFuncState(); +} + +std::string CodeGenHybrid::GetVarID(const Variable* v) { + auto it = var_idmap_.find(v); + if (!simple_mode_) { + CHECK(it != var_idmap_.end()) << "Find undefined Variable " << v->name_hint; + } else { + if (it == var_idmap_.end()) + return AllocVarID(v); + } + return it->second; +} + +void CodeGenHybrid::ReserveKeywordsAsUnique() { + // skip the first underscore, so SSA variable starts from _1 + GetUniqueName("_"); + GetUniqueName("def"); + GetUniqueName("for"); + GetUniqueName("in"); + GetUniqueName("range"); + GetUniqueName("unroll"); + GetUniqueName("vectorize"); + GetUniqueName("parallel"); + GetUniqueName("if"); + GetUniqueName("else"); + GetUniqueName("and"); + GetUniqueName("or"); + GetUniqueName("not"); +} + +void CodeGenHybrid::AddFunction(LoweredFunc f) { + // clear previous generated state. + InitFuncState(f); + // reserve keywords + ReserveKeywordsAsUnique(); + // add to alloc buffer type. + for (const auto & kv : f->handle_data_type) { + RegisterHandleType(kv.first.get(), kv.second.type()); + } + + stream << "def " << f->name << "("; + for (size_t i = 0; i < f->args.size(); ++i) { + Var v = f->args[i]; + stream << ' ' << v->name_hint; + } + stream << "):\n"; + int func_scope = BeginScope(); + PrintStmt(f->body); + EndScope(func_scope); +} + +std::string CodeGenHybrid::Finish() { + return decl_stream.str() + stream.str(); +} + +void CodeGenHybrid::PrintExpr(const Expr& n, std::ostream& os) { // NOLINT(*) + VisitExpr(n, os); +} + +void CodeGenHybrid::PrintSSAAssign(const std::string& target, const std::string& src, Type t) { + LOG(FATAL) << "Python backend does not support SSA format."; +} + +// Print a reference expression to a buffer. +std::string CodeGenHybrid::GetBufferRef( + Type t, const Variable* buffer, Expr index) { + std::ostringstream os; + std::string vid = GetVarID(buffer); + os << vid << "["; + PrintExpr(index, os); + os << "]"; + return os.str(); +} + +// Print a reference expression to a buffer. +std::string CodeGenHybrid::GetStructRef( + Type t, const Expr& buffer, const Expr& index, int kind) { + if (kind < intrinsic::kArrKindBound_) { + std::ostringstream os; + os << "(((TVMArray*)"; + this->PrintExpr(buffer, os); + os << ")"; + if (kind == intrinsic::kArrAddr) { + os << " + "; + this->PrintExpr(index, os); + os << ")"; + return os.str(); + } + os << '['; + this->PrintExpr(index, os); + os << "]."; + // other case: get fields. + switch (kind) { + case intrinsic::kArrData: os << "data"; break; + case intrinsic::kArrShape: os << "shape"; break; + case intrinsic::kArrStrides: os << "strides"; break; + case intrinsic::kArrNDim: os << "ndim"; break; + case intrinsic::kArrTypeCode: os << "dtype.code"; break; + case intrinsic::kArrTypeBits: os << "dtype.bits"; break; + case intrinsic::kArrByteOffset: os << "byte_offset"; break; + case intrinsic::kArrTypeLanes: os << "dtype.lanes"; break; + case intrinsic::kArrDeviceId: os << "ctx.device_id"; break; + case intrinsic::kArrDeviceType: os << "ctx.device_type"; break; + default: LOG(FATAL) << "unknown field code"; + } + os << ')'; + return os.str(); + } else { + CHECK_LT(kind, intrinsic::kTVMValueKindBound_); + std::ostringstream os; + os << "(((TVMValue*)"; + this->PrintExpr(buffer, os); + os << ")[" << index << "]."; + if (t.is_handle()) { + os << "v_handle"; + } else if (t.is_float()) { + os << "v_float64"; + } else if (t.is_int()) { + os << "v_int64"; + } else { + LOG(FATAL) << "Do not know how to handle type" << t; + } + os << ")"; + return os.str(); + } +} + + +bool CodeGenHybrid::HandleTypeMatch(const Variable* buf_var, Type t) const { + auto it = handle_data_type_.find(buf_var); + if (it == handle_data_type_.end()) return false; + return it->second == t; +} + +void CodeGenHybrid::RegisterHandleType(const Variable* buf_var, Type t) { + auto it = handle_data_type_.find(buf_var); + if (it == handle_data_type_.end()) { + handle_data_type_[buf_var] = t; + } else { + CHECK(it->second == t) << "conflicting buf var type"; + } +} + +std::string CodeGenHybrid::CastFromTo(std::string value, Type from, Type target) { + if (from == target) return value; + std::ostringstream os; + this->PrintType(target, os); + os << "(" << value << ")"; + return os.str(); +} + +void CodeGenHybrid::BindThreadIndex(const IterVar& iv) { + LOG(FATAL) << "to be implemented"; +} + +void CodeGenHybrid::PrintStorageSync(const Call* op) { // NOLINT(*) + LOG(FATAL) << "to be implemented"; +} + +void CodeGenHybrid::PrintStorageScope(const std::string& scope, std::ostream& os) { // NOLINT(*) + CHECK_EQ(scope, "global"); +} + +void CodeGenHybrid::PrintType(Type t, std::ostream& os) { // NOLINT(*) + CHECK_EQ(t.lanes(), 1) << "do not yet support vector types"; + CHECK(!t.is_handle()) << "Buffer type cannot be a handle!"; + if (t.is_float()) { + CHECK(t.bits() == 32 || t.bits() == 64); + os << "float" << t.bits(); + } else if (t.is_uint() || t.is_int()) { + switch (t.bits()) { + case 8: case 16: case 32: case 64: { + os << "int" << t.bits(); return; + } + case 1: os << "int"; return; + } + } + LOG(FATAL) << "Cannot convert type " << t << " to Python type"; +} + +void CodeGenHybrid::VisitExpr_(const IntImm *op, std::ostream& os) { // NOLINT(*) + os << op->value; +} +void CodeGenHybrid::VisitExpr_(const UIntImm *op, std::ostream& os) { // NOLINT(*) + os << op->value; +} +void CodeGenHybrid::VisitExpr_(const FloatImm *op, std::ostream& os) { // NOLINT(*) + os << std::scientific << op->value; +} +void CodeGenHybrid::VisitExpr_(const StringImm *op, std::ostream& os) { // NOLINT(*) + os << "\"" << op->value << "\""; +} + +template +inline void PrintBinaryExpr(const T* op, + const char *opstr, + std::ostream& os, // NOLINT(*) + CodeGenHybrid* p) { + if (op->type.lanes() == 1) { + if (isalpha(opstr[0])) { + os << opstr << '('; + p->PrintExpr(op->a, os); + os << ", "; + p->PrintExpr(op->b, os); + os << ')'; + } else { + os << '('; + p->PrintExpr(op->a, os); + os << ' ' << opstr << ' '; + p->PrintExpr(op->b, os); + os << ')'; + } + } else { + LOG(FATAL) << "vec bin op to be implemented"; + } +} + +inline void PrintBinaryIntrinsitc(const Call* op, + const char *opstr, + std::ostream& os, // NOLINT(*) + CodeGenHybrid* p) { + if (op->type.lanes() == 1) { + CHECK_EQ(op->args.size(), 2U); + os << '('; + p->PrintExpr(op->args[0], os); + os << opstr; + p->PrintExpr(op->args[1], os); + os << ')'; + } else { + LOG(FATAL) << "vec bin intrin to be implemented"; + } +} +void CodeGenHybrid::VisitExpr_(const Cast *op, std::ostream& os) { // NOLINT(*) + std::stringstream value; + PrintExpr(op->value, value); + os << CastFromTo(value.str(), op->value.type(), op->type); +} +void CodeGenHybrid::VisitExpr_(const Variable *op, std::ostream& os) { // NOLINT(*) + os << GetVarID(op); +} +void CodeGenHybrid::VisitExpr_(const Add *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "+", os, this); +} +void CodeGenHybrid::VisitExpr_(const Sub *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "-", os, this); +} +void CodeGenHybrid::VisitExpr_(const Mul *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "*", os, this); +} +void CodeGenHybrid::VisitExpr_(const Div *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "/", os, this); +} +void CodeGenHybrid::VisitExpr_(const Mod *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "%", os, this); +} +void CodeGenHybrid::VisitExpr_(const Min *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "min", os, this); +} +void CodeGenHybrid::VisitExpr_(const Max *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "max", os, this); +} +void CodeGenHybrid::VisitExpr_(const EQ *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "==", os, this); +} +void CodeGenHybrid::VisitExpr_(const NE *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "!=", os, this); +} +void CodeGenHybrid::VisitExpr_(const LT *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "<", os, this); +} +void CodeGenHybrid::VisitExpr_(const LE *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "<=", os, this); +} +void CodeGenHybrid::VisitExpr_(const GT *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, ">", os, this); +} +void CodeGenHybrid::VisitExpr_(const GE *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, ">=", os, this); +} +void CodeGenHybrid::VisitExpr_(const And *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "&&", os, this); +} +void CodeGenHybrid::VisitExpr_(const Or *op, std::ostream& os) { // NOLINT(*) + PrintBinaryExpr(op, "||", os, this); +} +void CodeGenHybrid::VisitExpr_(const Not *op, std::ostream& os) { // NOLINT(*) + os << '!'; + PrintExpr(op->a, os); +} + +void CodeGenHybrid::VisitExpr_(const Call *op, std::ostream& os) { // NOLINT(*) + if (op->call_type == Call::Extern || + op->call_type == Call::PureExtern) { + os << op->name << "("; + for (size_t i = 0; i < op->args.size(); i++) { + this->PrintExpr(op->args[i], os); + if (i < op->args.size() - 1) { + os << ", "; + } + } + os << ")"; + } else if (op->is_intrinsic(Call::bitwise_and)) { + PrintBinaryIntrinsitc(op, " & ", os, this); + } else if (op->is_intrinsic(Call::bitwise_xor)) { + PrintBinaryIntrinsitc(op, " ^ ", os, this); + } else if (op->is_intrinsic(Call::bitwise_or)) { + PrintBinaryIntrinsitc(op, " | ", os, this); + } else if (op->is_intrinsic(Call::bitwise_not)) { + CHECK_EQ(op->args.size(), 1U); + os << "(~"; + this->PrintExpr(op->args[0], os); + os << ')'; + } else if (op->is_intrinsic(Call::shift_left)) { + PrintBinaryIntrinsitc(op, " << ", os, this); + } else if (op->is_intrinsic(Call::shift_right)) { + PrintBinaryIntrinsitc(op, " >> ", os, this); + } /*else if (op->is_intrinsic(intrinsic::tvm_if_then_else)) { + os << "("; + PrintExpr(op->args[0], os); + os << " ? "; + PrintExpr(op->args[1], os); + os << " : "; + PrintExpr(op->args[2], os); + os << ")"; + } else if (op->is_intrinsic(intrinsic::tvm_address_of)) { + const Load *l = op->args[0].as(); + CHECK(op->args.size() == 1 && l); + os << "(("; + this->PrintType(l->type.element_of(), os); + os << " *)" << this->GetVarID(l->buffer_var.get()) + << " + "; + this->PrintExpr(l->index, os); + os << ')'; + } else if (op->is_intrinsic(intrinsic::tvm_struct_get)) { + CHECK_EQ(op->args.size(), 3U); + os << GetStructRef( + op->type, op->args[0], op->args[1], + op->args[2].as()->value); + } else if (op->is_intrinsic(intrinsic::tvm_handle_is_null)) { + CHECK_EQ(op->args.size(), 1U); + os << "("; + this->PrintExpr(op->args[0], os); + os << " == NULL)"; + } else { + if (op->call_type == Call::Intrinsic || + op->call_type == Call::PureIntrinsic) { + LOG(FATAL) << "Unresolved intrinsic " << op->name + << " with return type " << op->type; + } else { + LOG(FATAL) << "Unresolved call type " << op->call_type; + } + }*/ +} + +void CodeGenHybrid::VisitExpr_(const Load* op, std::ostream& os) { // NOLINT(*) + // int lanes = op->type.lanes(); + // delcare type. + if (op->type.lanes() == 1) { + std::string ref = GetBufferRef(op->type, op->buffer_var.get(), op->index); + os << ref; + } else { + LOG(FATAL) << "vec load to be supported"; + } +} + +void CodeGenHybrid::VisitStmt_(const Store* op) { + Type t = op->value.type(); + if (t.lanes() == 1) { + std::string value = this->PrintExpr(op->value); + std::string ref = this->GetBufferRef(t, op->buffer_var.get(), op->index); + this->PrintIndent(); + stream << ref << " = " << value << "\n"; + } else { + LOG(FATAL) << "Vectorized store is not supported yet..."; + } +} + +void CodeGenHybrid::VisitExpr_(const Let* op, std::ostream& os) { // NOLINT(*) + std::string value = PrintExpr(op->value); + CHECK(!var_idmap_.count(op->var.get())); + var_idmap_[op->var.get()] = value; + os << PrintExpr(op->body); +} + +void CodeGenHybrid::VisitExpr_(const Ramp* op, std::ostream& os) { // NOLINT(*) + // TODO(@were): Support vectorization access in both frontend and backend + LOG(FATAL) << "ramp to be supported yet"; +} + +void CodeGenHybrid::VisitExpr_(const Broadcast* op, std::ostream& os) { // NOLINT(*) + LOG(FATAL) << "Broadcast: not supported "; +} + +void CodeGenHybrid::VisitExpr_(const Select* op, std::ostream& os) { // NOLINT(*) + PrintExpr(op->true_value, os); + os << " if "; + PrintExpr(op->condition, os); + os << " else "; + PrintExpr(op->false_value, os); + os << "\n"; +} + +void CodeGenHybrid::VisitStmt_(const LetStmt* op) { + std::string value = PrintExpr(op->value); + stream << AllocVarID(op->var.get()) + << " = " << value << ";\n"; + PrintStmt(op->body); +} + +void CodeGenHybrid::VisitStmt_(const Allocate* op) { + CHECK(!is_zero(op->condition)); + std::string vid = AllocVarID(op->buffer_var.get()); + PrintIndent(); + stream << vid << " = allocate(("; + for (size_t i = 0; i < op->extents.size(); ++i) { + if (!i) stream << ", "; + stream << PrintExpr(op->extents[i]); + } + stream << "), \"" << op-> type << "\")\n"; + RegisterHandleType(op->buffer_var.get(), op->type); + this->PrintStmt(op->body); +} + +void CodeGenHybrid::VisitStmt_(const AttrStmt* op) { + if (op->attr_key == ir::attr::thread_extent) { + LOG(FATAL) << "Thread binding support yet!\n"; + } else if (op->attr_key == ir::attr::storage_scope) { + const Variable* v = op->node.as(); + CHECK(v); + alloc_storage_scope_[v] = op->value.as()->value; + } else if (op->attr_key == ir::attr::volatile_scope) { + const Variable* v = op->node.as(); + CHECK(v); + volatile_buf_.insert(v); + } + PrintStmt(op->body); +} + +void CodeGenHybrid::VisitStmt_(const AssertStmt* op) { + //TODO(@were): Support AssertStmt in both hybrid parser and here + LOG(FATAL) << "assert to be supported yet!\n"; + PrintStmt(op->body); +} + +void CodeGenHybrid::VisitStmt_(const For* op) { + std::string extent = PrintExpr(op->extent); + PrintIndent(); + std::string vid = AllocVarID(op->loop_var.get()); + stream << "for " << vid << " in " << "range(" << extent << "):\n"; + int for_scope = BeginScope(); + PrintStmt(op->body); + EndScope(for_scope); + PrintIndent(); +} + +void CodeGenHybrid::VisitStmt_(const IfThenElse* op) { + std::string cond = PrintExpr(op->condition); + PrintIndent(); + stream << "if " << cond << ":\n"; + int then_scope = BeginScope(); + PrintStmt(op->then_case); + EndScope(then_scope); + + if (op->else_case.defined()) { + PrintIndent(); + stream << "else:\n"; + int else_scope = BeginScope(); + PrintStmt(op->else_case); + EndScope(else_scope); + } +} + +void CodeGenHybrid::VisitStmt_(const Block *op) { + PrintStmt(op->first); + if (op->rest.defined()) PrintStmt(op->rest); +} + +void CodeGenHybrid::VisitStmt_(const Evaluate *op) { + if (is_const(op->value)) return; + std::string str = PrintExpr(op->value); + if (!str.empty()) + stream << str << "\n"; +} + +void CodeGenHybrid::VisitStmt_(const ProducerConsumer *op) { + PrintStmt(op->body); +} + +} // namespace codegen +} // namespace tvm diff --git a/src/codegen/codegen_hybrid.h b/src/codegen/codegen_hybrid.h new file mode 100644 index 000000000000..017d1e3ecab4 --- /dev/null +++ b/src/codegen/codegen_hybrid.h @@ -0,0 +1,171 @@ +/*! + * Copyright (c) 2019 by Contributors + * \file codegen_hybrid.h + * \brief Common utilities to generated C style code. + */ +#ifndef TVM_CODEGEN_CODEGEN_HYBRID_H_ +#define TVM_CODEGEN_CODEGEN_HYBRID_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include "codegen_source_base.h" + +namespace tvm { +namespace codegen { + +using namespace ir; +/*! + * \brief A base class to generate Python script. + * + * Unlike C-family generators, CodeGenHybrid does not generate SSA formed Python script. + * Only normal form is supported. + * + * **NOTE** CodeGenHybrid does not aim at generating Python scripts consumed by Python2/3. + * For runtime support, please refer the decorator in ``tvm/python/hybrid/api.py``. + */ +class CodeGenHybrid : + public ExprFunctor, + public StmtFunctor, + public CodeGenSourceBase { + public: + /*! + * \brief Add the function to the generated module. + * \param f The function to be compiled. + */ + void AddFunction(LoweredFunc f); + /*! + * \brief Finalize the compilation and return the code. + * \return The code. + */ + std::string Finish(); + /*! + * \brief Print the Stmt n to CodeGenHybrid->stream + * \param n The statement to be printed. + */ + void PrintStmt(const Stmt& n) { + VisitStmt(n); + } + /*! + * \brief Print the expression n(or its ssa id if in ssa mode) into os + * \param n The expression to be printed. + * \param os The output stream + */ + void PrintExpr(const Expr& n, std::ostream& os); + /*! + * \brief Same as PrintExpr, but simply returns result string + * \param n The expression to be printed. + */ + std::string PrintExpr(const Expr& n) { + std::ostringstream os; + PrintExpr(n, os); + return os.str(); + } + /*! + * \brief Initialize codegen state for generating f. + * \param f The function to be compiled. + */ + virtual void InitFuncState(LoweredFunc f); + // expression + void VisitExpr_(const Variable* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Load* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Let* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Call* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Add* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Sub* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Mul* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Div* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Mod* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Min* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Max* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const EQ* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const NE* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const LT* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const LE* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const GT* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const GE* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const And* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Or* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Cast* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Not* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Select* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Ramp* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const Broadcast* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const IntImm* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const UIntImm* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const FloatImm* op, std::ostream& os) override; // NOLINT(*) + void VisitExpr_(const StringImm* op, std::ostream& os) override; // NOLINT(*) + // statment + void VisitStmt_(const LetStmt* op) override; + void VisitStmt_(const Store* op) override; + void VisitStmt_(const For* op) override; + void VisitStmt_(const IfThenElse* op) override; + void VisitStmt_(const Allocate* op) override; + void VisitStmt_(const AttrStmt* op) override; + void VisitStmt_(const AssertStmt* op) override; + void VisitStmt_(const Evaluate* op) override; + void VisitStmt_(const Block* op) override; + void VisitStmt_(const ProducerConsumer* op) override; + /*! + * Print Type represetnation of type t. + * \param t The type representation. + * \param os The stream to print the ctype into + */ + virtual void PrintType(Type t, std::ostream& os); // NOLINT(*) + /*! + * \brief Print expr representing the thread tag + * \param IterVar iv The thread index to be binded; + */ + virtual void BindThreadIndex(const IterVar& iv); // NOLINT(*) + virtual void PrintStorageScope(const std::string& scope, std::ostream& os); // NOLINT(*) + virtual void PrintStorageSync(const Call* op); // NOLINT(*) + // Get a cast type from to + virtual std::string CastFromTo(std::string value, Type from, Type target); + + protected: + std::string GetVarID(const Variable* v); + // Print reference to struct location + std::string GetStructRef( + Type t, const Expr& buffer, const Expr& index, int kind); + // print reference to a buffer as type t in index. + virtual std::string GetBufferRef( + Type t, const Variable* buffer, Expr index); + /*! + * \brief If buffer is allocated as type t. + * \param buf_var The buffer variable. + * \param t The type to be checked. + */ + bool HandleTypeMatch(const Variable* buf_var, Type t) const; + /*! + * \brief Register the data type of buf_var + * \param buf_var The buffer variable. + * \param t The type to be checked. + */ + void RegisterHandleType(const Variable* buf_var, Type t); + // override + void PrintSSAAssign( + const std::string& target, const std::string& src, Type t) final; + /*! \brief restrict keyword */ + std::string restrict_keyword_{""}; + /*! \brief the storage scope of allocation */ + std::unordered_map alloc_storage_scope_; + /*! \brief the data type of allocated buffers */ + std::unordered_map handle_data_type_; + /*! \brief reserves common C keywords */ + void ReserveKeywordsAsUnique(); + + private: + /*! \brief whether print a simple form */ + bool simple_mode_{true}; + /*! \brief set of volatile buf access */ + std::unordered_set volatile_buf_; +}; + +} // namespace codegen +} // namespace tvm +#endif // TVM_CODEGEN_CODEGEN_HYBRID_H_ diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index a54fec3a7bf7..4fae4b7eb252 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -408,6 +408,8 @@ def blur2d(a): return b a = tvm.placeholder((32, 32), 'float32', 'a') + b = blur2d(a) + sch = tvm.create_schedule(b.op) run_and_check(blur2d, [a]) if tvm.gpu().exist: From b2850686c78cafd665454acc693e4c0f9dcd3dd8 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Sun, 20 Jan 2019 20:13:17 -0800 Subject: [PATCH 02/31] we no longer need the redundant hybrid/api.py --- python/tvm/build_module.py | 2 +- python/tvm/hybrid/__init__.py | 50 ++++++++++++++++++++++++++++++++++- python/tvm/hybrid/api.py | 47 -------------------------------- src/api/api_codegen.cc | 15 ----------- src/codegen/codegen_hybrid.cc | 16 +++++++++++ src/codegen/codegen_hybrid.h | 1 + 6 files changed, 67 insertions(+), 64 deletions(-) delete mode 100644 python/tvm/hybrid/api.py diff --git a/python/tvm/build_module.py b/python/tvm/build_module.py index 2e270bc3b217..9c2fea860a9f 100755 --- a/python/tvm/build_module.py +++ b/python/tvm/build_module.py @@ -47,7 +47,7 @@ def dump(*args, **kwargs): if not isinstance(retv, (_stmt.Stmt, container.LoweredFunc, container.Array)): return retv fname = func.func_name if hasattr(func, 'func_name') else func.__name__ - pname = str(self._pass_id) + "_" + fname + "_ir.cc" + pname = str(self._pass_id) + "_" + fname + "_ir.py" with open(pname, "a") as f: out = retv.body if isinstance(retv, container.LoweredFunc) else retv f.write(str(out)) diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index ac65f4b8e2c4..1dcdace6031b 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -7,5 +7,53 @@ 2. Developers can build HalideIR by writing Python code. """ -from .api import script, dump +from __future__ import absolute_import as _abs + +from .._ffi.base import decorate +from .._ffi.function import _init_api +from .. import _api_internal as _tvm_internal +from ..tensor import Tensor + from .parser import parse_python +from .util import _pruned_source + + +def script(pyfunc): + """Decorate a python function function as hybrid script. + + The hybrid function support emulation mode and parsing to + the internal language IR. + + Returns + ------- + hybrid_func : function + A decorated hybrid script function. + """ + def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring + from .util import _enter_hybrid_runtime, _restore_runtime, _is_tvm_arg_types + if _is_tvm_arg_types(args): + src = _pruned_source(func) + parser = parse_python(src, func.__globals__, args) + + input_tensors = [] + for i in args: + if isinstance(i, Tensor): + input_tensors.append(i) + op = _tvm_internal._HybridOp(parser.func_name, "HybridOp", None, input_tensors, + parser.outputs, parser.parsed_body) + res = [op.output(i) for i in range(len(parser.outputs))] + return res[0] if len(res) == 1 else res + + intersect = _enter_hybrid_runtime(func) + value = func(*args, **kwargs) + _restore_runtime(func, intersect) + return value + + return decorate(pyfunc, wrapped_func) + + +_init_api("tvm.hybrid") + + +def dump(ir): + return _HybridDump(ir) diff --git a/python/tvm/hybrid/api.py b/python/tvm/hybrid/api.py deleted file mode 100644 index 8b01cb8f359d..000000000000 --- a/python/tvm/hybrid/api.py +++ /dev/null @@ -1,47 +0,0 @@ -"""APIs of lowering the Python subset to HalideIR""" -from __future__ import absolute_import as _abs - -from .._ffi.base import decorate -from .. import _api_internal as _tvm_internal -from .. import codegen as _codegen -from ..tensor import Tensor - -from .parser import parse_python -from .util import _pruned_source - - -def script(pyfunc): - """Decorate a python function function as hybrid script. - - The hybrid function support emulation mode and parsing to - the internal language IR. - - Returns - ------- - hybrid_func : function - A decorated hybrid script function. - """ - def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring - from .util import _enter_hybrid_runtime, _restore_runtime, _is_tvm_arg_types - if _is_tvm_arg_types(args): - src = _pruned_source(func) - parser = parse_python(src, func.__globals__, args) - - input_tensors = [] - for i in args: - if isinstance(i, Tensor): - input_tensors.append(i) - op = _tvm_internal._HybridOp(parser.func_name, "HybridOp", None, input_tensors, - parser.outputs, parser.parsed_body) - res = [op.output(i) for i in range(len(parser.outputs))] - return res[0] if len(res) == 1 else res - - intersect = _enter_hybrid_runtime(func) - value = func(*args, **kwargs) - _restore_runtime(func, intersect) - return value - - return decorate(pyfunc, wrapped_func) - -def dump(ir): - return _codegen._HybridDump(ir) diff --git a/src/api/api_codegen.cc b/src/api/api_codegen.cc index ad91e577e158..372cd0e262b7 100644 --- a/src/api/api_codegen.cc +++ b/src/api/api_codegen.cc @@ -9,8 +9,6 @@ #include #include -#include "../codegen/codegen_hybrid.h" - namespace tvm { namespace codegen { @@ -27,18 +25,5 @@ TVM_REGISTER_API("module._PackImportsToC") .set_body([](TVMArgs args, TVMRetValue *ret) { *ret = PackImportsToC(args[0], args[1]); }); - -TVM_REGISTER_API("codegen._HybridDump") -.set_body([](TVMArgs args, TVMRetValue *ret) { - Stmt stmt; - if (args[0].IsNodeType()) { - stmt = args[0]; - } else if (args[0].IsNodeType()) { - stmt = Evaluate::make(args[0]); - } - CodeGenHybrid generator; - generator.PrintStmt(stmt); - *ret = generator.Finish(); - }); } // namespace codegen } // namespace tvm diff --git a/src/codegen/codegen_hybrid.cc b/src/codegen/codegen_hybrid.cc index 4fde9618dc28..3c67fbd5c10d 100644 --- a/src/codegen/codegen_hybrid.cc +++ b/src/codegen/codegen_hybrid.cc @@ -12,6 +12,10 @@ namespace codegen { using namespace ir; +void CodeGenHybrid::Init(bool simple_mode) { + simple_mode_ = simple_mode; +} + void CodeGenHybrid::InitFuncState(LoweredFunc f) { alloc_storage_scope_.clear(); handle_data_type_.clear(); @@ -504,5 +508,17 @@ void CodeGenHybrid::VisitStmt_(const ProducerConsumer *op) { PrintStmt(op->body); } +TVM_REGISTER_API("hybrid._HybridDump") +.set_body([](TVMArgs args, TVMRetValue *ret) { + Stmt stmt; + if (args[0].IsNodeType()) { + stmt = args[0]; + } else if (args[0].IsNodeType()) { + stmt = Evaluate::make(args[0]); + } + CodeGenHybrid generator; + generator.PrintStmt(stmt); + *ret = generator.Finish(); + }); } // namespace codegen } // namespace tvm diff --git a/src/codegen/codegen_hybrid.h b/src/codegen/codegen_hybrid.h index 017d1e3ecab4..10ba9177d188 100644 --- a/src/codegen/codegen_hybrid.h +++ b/src/codegen/codegen_hybrid.h @@ -34,6 +34,7 @@ class CodeGenHybrid : public StmtFunctor, public CodeGenSourceBase { public: + void Init(bool simple_mode); /*! * \brief Add the function to the generated module. * \param f The function to be compiled. From de3ad83603318215d589b268e8635c6ac3d0dd19 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 21 Jan 2019 11:57:51 -0800 Subject: [PATCH 03/31] support assert stmt --- python/tvm/hybrid/parser.py | 20 ++++++++++++++++---- src/codegen/codegen_hybrid.cc | 8 ++++---- src/codegen/codegen_hybrid.h | 2 +- tests/python/unittest/test_hybrid_script.py | 5 +++++ 4 files changed, 26 insertions(+), 9 deletions(-) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index 1e3fe3301191..55ce882f26cf 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -18,17 +18,23 @@ from ..container import Array from ..tensor import Tensor, Operation from .. import expr as _expr +from .. import stmt as _stmt from .. import make as _make from .. import api as _api from .. import ir_pass as _ir_pass def pack_list_to_block(lst): - if len(lst) == 1: + n = len(lst) + if n == 1: return lst[0] - body = lst[0] - for i in lst[1:]: - body = _make.Block(body, i) + body = lst[n - 1] + for i in range(1, n): + stmt = lst[n - 1 - i] + if isinstance(stmt, _stmt.AssertStmt): + body = _make.AssertStmt(stmt.condition, stmt.message, body) + else: + body = _make.Block(stmt, body) return body @@ -496,6 +502,12 @@ def visit_Str(self, node): return node.s + def visit_Assert(self, node): + test = self.visit(node.test) + mesg = _api.convert(self.visit(node.msg)) + return _make.AssertStmt(test, mesg, util.make_nop()) + + def parse_python(src, symbols, args): """The helper function of calling the AST visitor diff --git a/src/codegen/codegen_hybrid.cc b/src/codegen/codegen_hybrid.cc index 3c67fbd5c10d..6b45ae2df767 100644 --- a/src/codegen/codegen_hybrid.cc +++ b/src/codegen/codegen_hybrid.cc @@ -315,7 +315,7 @@ void CodeGenHybrid::VisitExpr_(const Call *op, std::ostream& os) { // NOLINT(*) op->call_type == Call::PureExtern) { os << op->name << "("; for (size_t i = 0; i < op->args.size(); i++) { - this->PrintExpr(op->args[i], os); + PrintExpr(op->args[i], os); if (i < op->args.size() - 1) { os << ", "; } @@ -330,7 +330,7 @@ void CodeGenHybrid::VisitExpr_(const Call *op, std::ostream& os) { // NOLINT(*) } else if (op->is_intrinsic(Call::bitwise_not)) { CHECK_EQ(op->args.size(), 1U); os << "(~"; - this->PrintExpr(op->args[0], os); + PrintExpr(op->args[0], os); os << ')'; } else if (op->is_intrinsic(Call::shift_left)) { PrintBinaryIntrinsitc(op, " << ", os, this); @@ -424,8 +424,7 @@ void CodeGenHybrid::VisitExpr_(const Select* op, std::ostream& os) { // NOLINT( void CodeGenHybrid::VisitStmt_(const LetStmt* op) { std::string value = PrintExpr(op->value); - stream << AllocVarID(op->var.get()) - << " = " << value << ";\n"; + stream << AllocVarID(op->var.get()) << " = " << value << ";\n"; PrintStmt(op->body); } @@ -517,6 +516,7 @@ TVM_REGISTER_API("hybrid._HybridDump") stmt = Evaluate::make(args[0]); } CodeGenHybrid generator; + generator.Init(true); generator.PrintStmt(stmt); *ret = generator.Finish(); }); diff --git a/src/codegen/codegen_hybrid.h b/src/codegen/codegen_hybrid.h index 10ba9177d188..a1e900e1e24f 100644 --- a/src/codegen/codegen_hybrid.h +++ b/src/codegen/codegen_hybrid.h @@ -162,7 +162,7 @@ class CodeGenHybrid : private: /*! \brief whether print a simple form */ - bool simple_mode_{true}; + bool simple_mode_{false}; /*! \brief set of volatile buf access */ std::unordered_set volatile_buf_; }; diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 4fae4b7eb252..234ff140389b 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -69,6 +69,7 @@ def outer_product(n, m, a, b): c = output_tensor((n, m), a.dtype) for i in range(n): for j in range(m): + assert i < n and j < m, "index out of range!" c[i, j] = a[i] * b[j] return c @@ -100,6 +101,10 @@ def test_outer_product(): assert ibody.extent.name == 'm' #Check loop body jbody = ibody.body + assert isinstance(jbody, tvm.stmt.AssertStmt) + assert isinstance(jbody.message, tvm.expr.StringImm) + assert jbody.message.value == "index out of range!" + jbody = jbody.body assert isinstance(jbody, tvm.stmt.Provide) assert jbody.func.name == 'c' assert len(jbody.args) == 2 From 790396c7aaf5e8788db792d0c82f2e1ca25efbe1 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 21 Jan 2019 12:17:43 -0800 Subject: [PATCH 04/31] cast supported --- python/tvm/hybrid/calls.py | 8 ++++++++ python/tvm/hybrid/intrin.py | 6 ++++++ tests/python/unittest/test_hybrid_script.py | 8 ++++---- 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 3fd472c57afc..94d481a9b5e2 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -104,3 +104,11 @@ def len(func_id, args): except: #pylint: disable=bare-except _internal_assert(args[0].shape.__len__() == 1, "Only one-dimension array can get len") return _api.convert(args[0].shape[0]) + + +def _cast(func_id, args): + _internal_assert(args.__len__() == 1 and isinstance(args[0], _expr.Expr), \ + "Only one expression can be cast") + return _make.Cast(func_id, args[0]) + +int16 = int32 = int64 = float16 = float32 = float64 = _cast diff --git a/python/tvm/hybrid/intrin.py b/python/tvm/hybrid/intrin.py index cb6d0fdb74b8..e4c268779e7c 100644 --- a/python/tvm/hybrid/intrin.py +++ b/python/tvm/hybrid/intrin.py @@ -86,6 +86,12 @@ def sigmoid(x): 'tanh' : numpy.tanh, 'power' : numpy.power, 'exp' : numpy.exp, + 'int16' : numpy.int16, + 'int32' : numpy.int32, + 'int64' : numpy.int64, + 'float16' : numpy.float16, + 'float32' : numpy.float32, + 'float64' : numpy.float64, 'sigmoid' : sigmoid, 'popcount' : popcount, } diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 234ff140389b..88319756c005 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -589,19 +589,19 @@ def test_const_range(): @tvm.hybrid.script def foo(a, b): c = output_tensor(a.shape, a.dtype) - d = output_tensor(a.shape, a.dtype) + d = output_tensor(a.shape, 'int32') for i in const_range(2): for j in const_range(5): - c[i, j] = a[i, j] + b[i, j] + c[i, j] = float32(int32(a[i, j]) + b[i, j]) for i in const_range(len(b)): for j in const_range(len(b[0])): - d[i, j] = a[i, j] + b[i, j] + d[i, j] = int32(a[i, j] + b[i, j]) return c, d - a = tvm.placeholder((2, 5), name='a', dtype='int32') + a = tvm.placeholder((2, 5), name='a', dtype='float32') b = [[1, 2, 3, 4, 5], [5, 4, 3, 2, 1]] run_and_check(foo, [a, b]) From 4be2d5f3cf545c40c7180f1c132c35afe3e60241 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 21 Jan 2019 12:34:10 -0800 Subject: [PATCH 05/31] intrin -> runtime; util is mainly in charge of compilation time --- python/tvm/hybrid/__init__.py | 3 +- python/tvm/hybrid/calls.py | 3 +- python/tvm/hybrid/{intrin.py => runtime.py} | 25 +++++++++ python/tvm/hybrid/util.py | 60 +++++++-------------- python/tvm/hybrid/var_decl.py | 2 +- tests/python/unittest/test_hybrid_script.py | 2 +- 6 files changed, 51 insertions(+), 44 deletions(-) rename python/tvm/hybrid/{intrin.py => runtime.py} (73%) diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index 1dcdace6031b..0a9548f9a72b 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -30,7 +30,8 @@ def script(pyfunc): A decorated hybrid script function. """ def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring - from .util import _enter_hybrid_runtime, _restore_runtime, _is_tvm_arg_types + from .runtime import _enter_hybrid_runtime, _restore_runtime + from .util import _is_tvm_arg_types if _is_tvm_arg_types(args): src = _pruned_source(func) parser = parse_python(src, func.__globals__, args) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 94d481a9b5e2..d87eaf9db29e 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -111,4 +111,5 @@ def _cast(func_id, args): "Only one expression can be cast") return _make.Cast(func_id, args[0]) -int16 = int32 = int64 = float16 = float32 = float64 = _cast +uint8 = uint16 = uint32 = uint64 = int8 = int16 = int32 = int64 = _cast +float16 = float32 = float64 = _cast diff --git a/python/tvm/hybrid/intrin.py b/python/tvm/hybrid/runtime.py similarity index 73% rename from python/tvm/hybrid/intrin.py rename to python/tvm/hybrid/runtime.py index e4c268779e7c..2bc92c290eb3 100644 --- a/python/tvm/hybrid/intrin.py +++ b/python/tvm/hybrid/runtime.py @@ -86,6 +86,11 @@ def sigmoid(x): 'tanh' : numpy.tanh, 'power' : numpy.power, 'exp' : numpy.exp, + 'uint8' : numpy.uint8, + 'uint16' : numpy.uint16, + 'uint32' : numpy.uint32, + 'uint64' : numpy.uint64, + 'int8' : numpy.int8, 'int16' : numpy.int16, 'int32' : numpy.int32, 'int64' : numpy.int64, @@ -95,3 +100,23 @@ def sigmoid(x): 'sigmoid' : sigmoid, 'popcount' : popcount, } + + +def _enter_hybrid_runtime(func): + """Put hybrid runtime variables into the global scope""" + _globals = func.__globals__ + intersect = [] + for elem in list(HYBRID_GLOBALS.keys()): + if elem in _globals.keys(): + intersect.append((elem, _globals[elem])) + _globals[elem] = HYBRID_GLOBALS[elem] + return intersect + + +def _restore_runtime(func, intersect): + """Rollback the modification caused by hybrid runtime""" + _globals = func.__globals__ + for elem in list(HYBRID_GLOBALS.keys()): + _globals.pop(elem) + for k, v in intersect: + _globals[k] = v diff --git a/python/tvm/hybrid/util.py b/python/tvm/hybrid/util.py index 44222d2d80f7..56190a82765e 100644 --- a/python/tvm/hybrid/util.py +++ b/python/tvm/hybrid/util.py @@ -5,14 +5,13 @@ import logging import sys import numpy -from .intrin import HYBRID_GLOBALS -from .._ffi.base import numeric_types from .. import api as _api from .. import make as _make from .. import expr as _expr from .. import stmt as _stmt -from ..container import Array +from .._ffi.base import numeric_types from ..tensor import Tensor +from ..container import Array #pylint: disable=invalid-name @@ -20,6 +19,7 @@ tvm_arg_types = (Tensor, Array, _expr.Var, _expr.ConstExpr) halide_imm_types = (_expr.IntImm, _expr.FloatImm, _expr.UIntImm) + def _internal_assert(cond, err): """Simplify the code segment like if not XXX then raise an error""" if not cond: @@ -52,6 +52,23 @@ def _pruned_source(func): raise err +def replace_io(body, rmap): + """Replacing tensors usage according to the dict given""" + from .. import ir_pass + + def replace(op): + if isinstance(op, _stmt.Provide) and op.func in rmap.keys(): + buf = rmap[op.func] + return _make.Provide(buf.op, op.value_index, op.value, op.args) + elif isinstance(op, _expr.Call) and op.func in rmap.keys(): + buf = rmap[op.func] + return _make.Call(buf.dtype, buf.name, op.args, \ + _expr.Call.Halide, buf.op, buf.value_index) + return None + + return ir_pass.IRTransform(body, None, replace, ['Provide', 'Call']) + + def _is_tvm_arg_types(args): """Determine a list of element is either a list of tvm arguments of a list of numpy arguments. If neither is true, raise a value error.""" @@ -68,40 +85,3 @@ def _is_tvm_arg_types(args): _internal_assert(isinstance(elem, np_arg_types), \ "Expect a numpy type but %s get!" % str(type(elem))) return False - - -def _enter_hybrid_runtime(func): - """Put hybrid runtime variables into the global scope""" - _globals = func.__globals__ - intersect = [] - for elem in list(HYBRID_GLOBALS.keys()): - if elem in _globals.keys(): - intersect.append((elem, _globals[elem])) - _globals[elem] = HYBRID_GLOBALS[elem] - return intersect - - -def _restore_runtime(func, intersect): - """Rollback the modification caused by hybrid runtime""" - _globals = func.__globals__ - for elem in list(HYBRID_GLOBALS.keys()): - _globals.pop(elem) - for k, v in intersect: - _globals[k] = v - - -def replace_io(body, rmap): - """Replacing tensors usage according to the dict given""" - from .. import ir_pass - - def replace(op): - if isinstance(op, _stmt.Provide) and op.func in rmap.keys(): - buf = rmap[op.func] - return _make.Provide(buf.op, op.value_index, op.value, op.args) - elif isinstance(op, _expr.Call) and op.func in rmap.keys(): - buf = rmap[op.func] - return _make.Call(buf.dtype, buf.name, op.args, \ - _expr.Call.Halide, buf.op, buf.value_index) - return None - - return ir_pass.IRTransform(body, None, replace, ['Provide', 'Call']) diff --git a/python/tvm/hybrid/var_decl.py b/python/tvm/hybrid/var_decl.py index eb893a7f22a1..47201c44003d 100644 --- a/python/tvm/hybrid/var_decl.py +++ b/python/tvm/hybrid/var_decl.py @@ -2,7 +2,7 @@ import ast import sys -from .intrin import HYBRID_GLOBALS +from .runtime import HYBRID_GLOBALS from .util import _internal_assert diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 88319756c005..6aafb551f13b 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -1,6 +1,6 @@ import tvm, inspect, sys, traceback, numpy, nose, types from tvm.hybrid import script -from tvm.hybrid.intrin import HYBRID_GLOBALS +from tvm.hybrid.runtime import HYBRID_GLOBALS @nose.tools.nottest def run_and_check(func, args, var_dict={}, target='llvm', sch=None, outs=None): From b0981b04f56abf29e4caa1e28497a987b235c7e4 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 21 Jan 2019 12:59:44 -0800 Subject: [PATCH 06/31] assert statement --- src/codegen/codegen_hybrid.cc | 49 ++++++++++++++--------------------- 1 file changed, 20 insertions(+), 29 deletions(-) diff --git a/src/codegen/codegen_hybrid.cc b/src/codegen/codegen_hybrid.cc index 6b45ae2df767..07a92252fbca 100644 --- a/src/codegen/codegen_hybrid.cc +++ b/src/codegen/codegen_hybrid.cc @@ -34,8 +34,6 @@ std::string CodeGenHybrid::GetVarID(const Variable* v) { } void CodeGenHybrid::ReserveKeywordsAsUnique() { - // skip the first underscore, so SSA variable starts from _1 - GetUniqueName("_"); GetUniqueName("def"); GetUniqueName("for"); GetUniqueName("in"); @@ -336,34 +334,17 @@ void CodeGenHybrid::VisitExpr_(const Call *op, std::ostream& os) { // NOLINT(*) PrintBinaryIntrinsitc(op, " << ", os, this); } else if (op->is_intrinsic(Call::shift_right)) { PrintBinaryIntrinsitc(op, " >> ", os, this); - } /*else if (op->is_intrinsic(intrinsic::tvm_if_then_else)) { - os << "("; - PrintExpr(op->args[0], os); - os << " ? "; + } else if (op->is_intrinsic(intrinsic::tvm_if_then_else)) { PrintExpr(op->args[1], os); - os << " : "; + os << " if "; + PrintExpr(op->args[0], os); + os << " else "; PrintExpr(op->args[2], os); - os << ")"; - } else if (op->is_intrinsic(intrinsic::tvm_address_of)) { - const Load *l = op->args[0].as(); - CHECK(op->args.size() == 1 && l); - os << "(("; - this->PrintType(l->type.element_of(), os); - os << " *)" << this->GetVarID(l->buffer_var.get()) - << " + "; - this->PrintExpr(l->index, os); - os << ')'; - } else if (op->is_intrinsic(intrinsic::tvm_struct_get)) { - CHECK_EQ(op->args.size(), 3U); - os << GetStructRef( - op->type, op->args[0], op->args[1], - op->args[2].as()->value); - } else if (op->is_intrinsic(intrinsic::tvm_handle_is_null)) { - CHECK_EQ(op->args.size(), 1U); - os << "("; - this->PrintExpr(op->args[0], os); - os << " == NULL)"; } else { + // TODO(@were): Support tvm runtime intrinsics: + // intrinsic::tvm_address_of + // intrinsic::tvm_struct_get + // intrinsic::tvm_handle_is_null if (op->call_type == Call::Intrinsic || op->call_type == Call::PureIntrinsic) { LOG(FATAL) << "Unresolved intrinsic " << op->name @@ -371,7 +352,7 @@ void CodeGenHybrid::VisitExpr_(const Call *op, std::ostream& os) { // NOLINT(*) } else { LOG(FATAL) << "Unresolved call type " << op->call_type; } - }*/ + } } void CodeGenHybrid::VisitExpr_(const Load* op, std::ostream& os) { // NOLINT(*) @@ -443,6 +424,7 @@ void CodeGenHybrid::VisitStmt_(const Allocate* op) { } void CodeGenHybrid::VisitStmt_(const AttrStmt* op) { + // TODO(@were): Support thread and buffer binding if (op->attr_key == ir::attr::thread_extent) { LOG(FATAL) << "Thread binding support yet!\n"; } else if (op->attr_key == ir::attr::storage_scope) { @@ -459,7 +441,11 @@ void CodeGenHybrid::VisitStmt_(const AttrStmt* op) { void CodeGenHybrid::VisitStmt_(const AssertStmt* op) { //TODO(@were): Support AssertStmt in both hybrid parser and here - LOG(FATAL) << "assert to be supported yet!\n"; + stream << "assert "; + PrintExpr(op->condition, stream); + stream << ", "; + PrintExpr(op->message, stream); + stream << "\n"; PrintStmt(op->body); } @@ -504,11 +490,16 @@ void CodeGenHybrid::VisitStmt_(const Evaluate *op) { } void CodeGenHybrid::VisitStmt_(const ProducerConsumer *op) { + PrintIndent(); + stream << "# producing " << op->func->func_name() << "\n"; PrintStmt(op->body); + PrintIndent(); + stream << "# produced " << op->func->func_name() << "\n"; } TVM_REGISTER_API("hybrid._HybridDump") .set_body([](TVMArgs args, TVMRetValue *ret) { + // If the entrance is Python directly, we dump it as simple_mode. Stmt stmt; if (args[0].IsNodeType()) { stmt = args[0]; From 6c0d984a65f2fdee6f45e6b5e5b0bdded39c7a84 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 21 Jan 2019 13:13:39 -0800 Subject: [PATCH 07/31] fix python lint --- python/tvm/hybrid/__init__.py | 7 +++++-- python/tvm/hybrid/calls.py | 5 +++-- python/tvm/hybrid/parser.py | 9 +++++---- 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index 0a9548f9a72b..6712e17946f9 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -56,5 +56,8 @@ def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring _init_api("tvm.hybrid") -def dump(ir): - return _HybridDump(ir) +def dump(stmt): + """Dump a HalideIR node to hybrid script format. + To replace the default IRPrint when a mature timing. + """ + return _HybridDump(stmt) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index d87eaf9db29e..97883c65a99c 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -111,5 +111,6 @@ def _cast(func_id, args): "Only one expression can be cast") return _make.Cast(func_id, args[0]) -uint8 = uint16 = uint32 = uint64 = int8 = int16 = int32 = int64 = _cast -float16 = float32 = float64 = _cast +float16 = float32 = float64 = _cast #pylint: disable=invalid-name +int8 = int16 = int32 = int64 = _cast #pylint: disable=invalid-name +uint8 = uint16 = uint32 = uint64 = _cast #pylint: disable=invalid-name diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index 55ce882f26cf..9a191d2fbb5e 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -24,7 +24,8 @@ from .. import ir_pass as _ir_pass -def pack_list_to_block(lst): +def concat_list_to_block(lst): + """Concatenate a list of Python IR nodes to HalideIR Block""" n = len(lst) if n == 1: return lst[0] @@ -39,12 +40,12 @@ def pack_list_to_block(lst): def visit_list_to_block(visit, lst): - """Convert a list of Python IR nodes to HalideIR Block""" + """Visit and concatenate a list of Python IR nodes to HalideIR Block""" lst = [visit(stmt) for stmt in lst if not util.is_docstring(stmt)] lst = [stmt for stmt in lst if not _ir_pass.Equal(stmt, util.make_nop())] if not lst: return util.make_nop() - return pack_list_to_block(lst) + return concat_list_to_block(lst) class Symbol(Enum): @@ -447,7 +448,7 @@ def visit_For(self, node): body = visit_list_to_block(self.visit, node.body) body = self.wrap_up_realize(node, body) bodies.append(body) - return pack_list_to_block(bodies) + return concat_list_to_block(bodies) elif iter_var is None: _internal_assert(for_type is not None, "The loop bind function parse error!") From 557472be817660e05499bdd38aba2abf4cce7379 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 21 Jan 2019 13:14:28 -0800 Subject: [PATCH 08/31] fix cpp lint --- src/codegen/codegen_hybrid.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/codegen/codegen_hybrid.cc b/src/codegen/codegen_hybrid.cc index 07a92252fbca..a474ffb61e55 100644 --- a/src/codegen/codegen_hybrid.cc +++ b/src/codegen/codegen_hybrid.cc @@ -440,7 +440,7 @@ void CodeGenHybrid::VisitStmt_(const AttrStmt* op) { } void CodeGenHybrid::VisitStmt_(const AssertStmt* op) { - //TODO(@were): Support AssertStmt in both hybrid parser and here + // TODO(@were): Support AssertStmt in both hybrid parser and here stream << "assert "; PrintExpr(op->condition, stream); stream << ", "; From c1f4464c1a480c3720356791837d29499c78ecf0 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Tue, 22 Jan 2019 10:29:03 -0800 Subject: [PATCH 09/31] on the way to module --- docs/langref/hybrid_script.rst | 14 ++++++++ src/codegen/codegen_hybrid.cc | 59 +++++++++++++++++----------------- 2 files changed, 44 insertions(+), 29 deletions(-) diff --git a/docs/langref/hybrid_script.rst b/docs/langref/hybrid_script.rst index 7043281fcafb..9e80d5ba72ff 100644 --- a/docs/langref/hybrid_script.rst +++ b/docs/langref/hybrid_script.rst @@ -197,6 +197,20 @@ You can also do loop-thread bind by writing code like this: a[tx] = b[tx] +Assert Statement +~~~~~~~~~~~~~~~~ + +Assert statement is supported, you can simply use it as it is in standard Python. + +.. code-block:: python + + assert cond, mesg + +.. note:: + + ``Assert`` is NOT a function call. Users are encouraged to use assert in the way + presented above --- condition followed by message. It fits both Python AST and HalideIR. + Keywords ~~~~~~~~ - For keywords: ``serial``, ``range``, ``unroll``, ``parallel``, ``vectorize``, ``bind``, ``const_expr`` diff --git a/src/codegen/codegen_hybrid.cc b/src/codegen/codegen_hybrid.cc index a474ffb61e55..3b4cc0698655 100644 --- a/src/codegen/codegen_hybrid.cc +++ b/src/codegen/codegen_hybrid.cc @@ -98,16 +98,16 @@ std::string CodeGenHybrid::GetStructRef( if (kind < intrinsic::kArrKindBound_) { std::ostringstream os; os << "(((TVMArray*)"; - this->PrintExpr(buffer, os); + PrintExpr(buffer, os); os << ")"; if (kind == intrinsic::kArrAddr) { os << " + "; - this->PrintExpr(index, os); + PrintExpr(index, os); os << ")"; return os.str(); } os << '['; - this->PrintExpr(index, os); + PrintExpr(index, os); os << "]."; // other case: get fields. switch (kind) { @@ -129,7 +129,7 @@ std::string CodeGenHybrid::GetStructRef( CHECK_LT(kind, intrinsic::kTVMValueKindBound_); std::ostringstream os; os << "(((TVMValue*)"; - this->PrintExpr(buffer, os); + PrintExpr(buffer, os); os << ")[" << index << "]."; if (t.is_handle()) { os << "v_handle"; @@ -161,10 +161,25 @@ void CodeGenHybrid::RegisterHandleType(const Variable* buf_var, Type t) { } } +void CodeGenHybrid::PrintType(Type t, std::ostream &os) { + if (t.is_float()) { + os << "float"; + CHECK(t.bits() == 16 || t.bits() == 32 || t.bits() == 64); + } else if (t.is_int()) { + os << "int"; + CHECK(t.bits() == 8 || t.bits() == 16 || t.bits() == 32 || t.bits() == 64); + } else { + CHECK(t.is_uint()) << "Unsupported type " << t; + os << "uint"; + CHECK(t.bits() == 8 || t.bits() == 16 || t.bits() == 32 || t.bits() == 64); + } + os << t.bits(); +} + std::string CodeGenHybrid::CastFromTo(std::string value, Type from, Type target) { if (from == target) return value; std::ostringstream os; - this->PrintType(target, os); + PrintType(target, os); os << "(" << value << ")"; return os.str(); } @@ -181,31 +196,17 @@ void CodeGenHybrid::PrintStorageScope(const std::string& scope, std::ostream& os CHECK_EQ(scope, "global"); } -void CodeGenHybrid::PrintType(Type t, std::ostream& os) { // NOLINT(*) - CHECK_EQ(t.lanes(), 1) << "do not yet support vector types"; - CHECK(!t.is_handle()) << "Buffer type cannot be a handle!"; - if (t.is_float()) { - CHECK(t.bits() == 32 || t.bits() == 64); - os << "float" << t.bits(); - } else if (t.is_uint() || t.is_int()) { - switch (t.bits()) { - case 8: case 16: case 32: case 64: { - os << "int" << t.bits(); return; - } - case 1: os << "int"; return; - } - } - LOG(FATAL) << "Cannot convert type " << t << " to Python type"; -} - void CodeGenHybrid::VisitExpr_(const IntImm *op, std::ostream& os) { // NOLINT(*) - os << op->value; + PrintType(op->type, os); + os << "(" << op->value << ")"; } void CodeGenHybrid::VisitExpr_(const UIntImm *op, std::ostream& os) { // NOLINT(*) - os << op->value; + PrintType(op->type, os); + os << "(" << op->value << ")"; } void CodeGenHybrid::VisitExpr_(const FloatImm *op, std::ostream& os) { // NOLINT(*) - os << std::scientific << op->value; + PrintType(op->type, os); + os << "(" << op->value << ")"; } void CodeGenHybrid::VisitExpr_(const StringImm *op, std::ostream& os) { // NOLINT(*) os << "\"" << op->value << "\""; @@ -369,9 +370,9 @@ void CodeGenHybrid::VisitExpr_(const Load* op, std::ostream& os) { // NOLINT(*) void CodeGenHybrid::VisitStmt_(const Store* op) { Type t = op->value.type(); if (t.lanes() == 1) { - std::string value = this->PrintExpr(op->value); - std::string ref = this->GetBufferRef(t, op->buffer_var.get(), op->index); - this->PrintIndent(); + std::string value = PrintExpr(op->value); + std::string ref = GetBufferRef(t, op->buffer_var.get(), op->index); + PrintIndent(); stream << ref << " = " << value << "\n"; } else { LOG(FATAL) << "Vectorized store is not supported yet..."; @@ -420,7 +421,7 @@ void CodeGenHybrid::VisitStmt_(const Allocate* op) { } stream << "), \"" << op-> type << "\")\n"; RegisterHandleType(op->buffer_var.get(), op->type); - this->PrintStmt(op->body); + PrintStmt(op->body); } void CodeGenHybrid::VisitStmt_(const AttrStmt* op) { From f1d636ae384a8bf43a6d481812f3676dd7c351ce Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Tue, 22 Jan 2019 10:35:38 -0800 Subject: [PATCH 10/31] rollback .cc --- python/tvm/build_module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/build_module.py b/python/tvm/build_module.py index 9c2fea860a9f..2e270bc3b217 100755 --- a/python/tvm/build_module.py +++ b/python/tvm/build_module.py @@ -47,7 +47,7 @@ def dump(*args, **kwargs): if not isinstance(retv, (_stmt.Stmt, container.LoweredFunc, container.Array)): return retv fname = func.func_name if hasattr(func, 'func_name') else func.__name__ - pname = str(self._pass_id) + "_" + fname + "_ir.py" + pname = str(self._pass_id) + "_" + fname + "_ir.cc" with open(pname, "a") as f: out = retv.body if isinstance(retv, container.LoweredFunc) else retv f.write(str(out)) From 942d4b69e4b52595aa8144e0671d203d16dfb8e3 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Tue, 22 Jan 2019 10:37:40 -0800 Subject: [PATCH 11/31] fix typo, no direct expose then --- python/tvm/hybrid/__init__.py | 12 +----------- src/codegen/codegen_hybrid.cc | 15 --------------- 2 files changed, 1 insertion(+), 26 deletions(-) diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index 6712e17946f9..671d38d78ae2 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -19,7 +19,7 @@ def script(pyfunc): - """Decorate a python function function as hybrid script. + """Decorate a python function function as hybrid script. The hybrid function support emulation mode and parsing to the internal language IR. @@ -51,13 +51,3 @@ def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring return value return decorate(pyfunc, wrapped_func) - - -_init_api("tvm.hybrid") - - -def dump(stmt): - """Dump a HalideIR node to hybrid script format. - To replace the default IRPrint when a mature timing. - """ - return _HybridDump(stmt) diff --git a/src/codegen/codegen_hybrid.cc b/src/codegen/codegen_hybrid.cc index 3b4cc0698655..0bc0922ffff9 100644 --- a/src/codegen/codegen_hybrid.cc +++ b/src/codegen/codegen_hybrid.cc @@ -497,20 +497,5 @@ void CodeGenHybrid::VisitStmt_(const ProducerConsumer *op) { PrintIndent(); stream << "# produced " << op->func->func_name() << "\n"; } - -TVM_REGISTER_API("hybrid._HybridDump") -.set_body([](TVMArgs args, TVMRetValue *ret) { - // If the entrance is Python directly, we dump it as simple_mode. - Stmt stmt; - if (args[0].IsNodeType()) { - stmt = args[0]; - } else if (args[0].IsNodeType()) { - stmt = Evaluate::make(args[0]); - } - CodeGenHybrid generator; - generator.Init(true); - generator.PrintStmt(stmt); - *ret = generator.Finish(); - }); } // namespace codegen } // namespace tvm From 4f2e842c773a7e72c13701edb53db41d529bee51 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 23 Jan 2019 08:50:10 -0800 Subject: [PATCH 12/31] @vinx13 ceil is added i guess? --- python/tvm/build_module.py | 26 +++++++++++++++++++++----- python/tvm/hybrid/calls.py | 8 ++++++++ python/tvm/hybrid/runtime.py | 5 +++-- 3 files changed, 32 insertions(+), 7 deletions(-) diff --git a/python/tvm/build_module.py b/python/tvm/build_module.py index 2e270bc3b217..d3b40feba408 100755 --- a/python/tvm/build_module.py +++ b/python/tvm/build_module.py @@ -292,6 +292,24 @@ def get_binds(args, binds=None): return binds, arg_list +def form_body(sch): + """According to the given schedule, form the raw body + Parameters + ---------- + sch : tvm.schedule.Schedule + The given scheduler to form the raw body + + Returns + ------- + The body formed according to the given schedule + """ + # normalize schedule first + sch = sch.normalize() + bounds = schedule.InferBound(sch) + stmt = schedule.ScheduleOps(sch, bounds) + stmt = ir_pass.InjectPrefetch(stmt) + + def lower(sch, args, name="default_function", @@ -337,11 +355,7 @@ def lower(sch, # Phase 0 if isinstance(sch, schedule.Schedule): - # normalize schedule first - sch = sch.normalize() - bounds = schedule.InferBound(sch) - stmt = schedule.ScheduleOps(sch, bounds) - stmt = ir_pass.InjectPrefetch(stmt) + stmt = form_body(sch) for f in lower_phase0: stmt = f(stmt) @@ -533,6 +547,8 @@ def build(inputs, if isinstance(inputs, schedule.Schedule): if args is None: raise ValueError("args must be given for build from schedule") + if target == 'hybrid': + return form_body(sch) flist = lower(inputs, args, name=name, binds=binds) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 97883c65a99c..5749d50e45dd 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -114,3 +114,11 @@ def _cast(func_id, args): float16 = float32 = float64 = _cast #pylint: disable=invalid-name int8 = int16 = int32 = int64 = _cast #pylint: disable=invalid-name uint8 = uint16 = uint32 = uint64 = _cast #pylint: disable=invalid-name + + +def ceil_div(func_id, args): + _internal_assert(args.__len__() == 2, \ + "Only one expression can be cast") + for i in range(2): + _internal_assert(isinstance(args[0], _expr.Expr), "Only expressions can div") + return (a + b - 1) / b diff --git a/python/tvm/hybrid/runtime.py b/python/tvm/hybrid/runtime.py index 2bc92c290eb3..5df470e09085 100644 --- a/python/tvm/hybrid/runtime.py +++ b/python/tvm/hybrid/runtime.py @@ -86,6 +86,8 @@ def sigmoid(x): 'tanh' : numpy.tanh, 'power' : numpy.power, 'exp' : numpy.exp, + 'sigmoid' : sigmoid, + 'popcount' : popcount, 'uint8' : numpy.uint8, 'uint16' : numpy.uint16, 'uint32' : numpy.uint32, @@ -97,8 +99,7 @@ def sigmoid(x): 'float16' : numpy.float16, 'float32' : numpy.float32, 'float64' : numpy.float64, - 'sigmoid' : sigmoid, - 'popcount' : popcount, + 'ceil_div' : lambda a, b: (a + b - 1) / b } From 151f0fe73a8a14f6aef6d4ba6390df63f15d4bb9 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 30 Jan 2019 16:22:09 -0800 Subject: [PATCH 13/31] wip... --- CMakeLists.txt | 1 + cmake/config.cmake | 3 + cmake/modules/contrib/HybridDump.cmake | 5 + python/tvm/build_module.py | 3 +- python/tvm/hybrid/__init__.py | 9 +- python/tvm/hybrid/calls.py | 1 + python/tvm/hybrid/dump.py | 25 ++ python/tvm/hybrid/runtime.py | 1 - .../hybrid}/codegen_hybrid.cc | 353 ++++++------------ .../hybrid}/codegen_hybrid.h | 95 ++--- src/op/hybrid_op.cc | 29 +- 11 files changed, 203 insertions(+), 322 deletions(-) create mode 100644 cmake/modules/contrib/HybridDump.cmake create mode 100644 python/tvm/hybrid/dump.py rename src/{codegen => contrib/hybrid}/codegen_hybrid.cc (53%) rename src/{codegen => contrib/hybrid}/codegen_hybrid.h (65%) diff --git a/CMakeLists.txt b/CMakeLists.txt index cb9b2df2f284..494afbdff792 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -190,6 +190,7 @@ include(cmake/modules/contrib/BLAS.cmake) include(cmake/modules/contrib/Random.cmake) include(cmake/modules/contrib/Sort.cmake) include(cmake/modules/contrib/NNPack.cmake) +include(cmake/modules/contrib/HybridDump.cmake) add_library(tvm SHARED ${COMPILER_SRCS} ${RUNTIME_SRCS}) add_library(tvm_topi SHARED ${TOPI_SRCS}) diff --git a/cmake/config.cmake b/cmake/config.cmake index a97def410ddd..831dfa2af9e8 100644 --- a/cmake/config.cmake +++ b/cmake/config.cmake @@ -131,3 +131,6 @@ set(USE_SORT OFF) # Build ANTLR parser for Relay text format set(USE_ANTLR OFF) + +# Build ANTLR parser for Relay text format +set(USE_HYBRID_DUMP ON) diff --git a/cmake/modules/contrib/HybridDump.cmake b/cmake/modules/contrib/HybridDump.cmake new file mode 100644 index 000000000000..a4c8f626a64a --- /dev/null +++ b/cmake/modules/contrib/HybridDump.cmake @@ -0,0 +1,5 @@ +if(USE_HYBRID_DUMP) + message(STATUS "Build with contrib.hybriddump") + file(GLOB HYBRID_CONTRIB_SRC src/contrib/hybrid/*.cc) + list(APPEND COMPILER_SRCS ${HYBRID_CONTRIB_SRC}) +endif(USE_HYBRID_DUMP) diff --git a/python/tvm/build_module.py b/python/tvm/build_module.py index d3b40feba408..139f13364271 100755 --- a/python/tvm/build_module.py +++ b/python/tvm/build_module.py @@ -308,6 +308,7 @@ def form_body(sch): bounds = schedule.InferBound(sch) stmt = schedule.ScheduleOps(sch, bounds) stmt = ir_pass.InjectPrefetch(stmt) + return stmt def lower(sch, @@ -547,8 +548,6 @@ def build(inputs, if isinstance(inputs, schedule.Schedule): if args is None: raise ValueError("args must be given for build from schedule") - if target == 'hybrid': - return form_body(sch) flist = lower(inputs, args, name=name, binds=binds) diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index 671d38d78ae2..6cebf463e608 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -4,9 +4,15 @@ 1. Users can write some preliminary versions of the computation patterns have not been supported yet and verify it across the real execution and python semantic emulation. -2. Developers can build HalideIR by writing Python code. +2. So far, it is a text format dedicated to HalideIR Phase 0. Refer tvm.lower +for more details. A larger ambition of this module is to support all levels of +HalideIR. """ +# TODO(@were): Make this module more complete. +# 1. Support HalideIR dumping to Hybrid Script +# 2. Support multi-level HalideIR + from __future__ import absolute_import as _abs from .._ffi.base import decorate @@ -16,6 +22,7 @@ from .parser import parse_python from .util import _pruned_source +from .dump import dump def script(pyfunc): diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 5749d50e45dd..9a3229479ae2 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -121,4 +121,5 @@ def ceil_div(func_id, args): "Only one expression can be cast") for i in range(2): _internal_assert(isinstance(args[0], _expr.Expr), "Only expressions can div") + a, b = args[0], args[1] return (a + b - 1) / b diff --git a/python/tvm/hybrid/dump.py b/python/tvm/hybrid/dump.py new file mode 100644 index 000000000000..1eb2b538dfe2 --- /dev/null +++ b/python/tvm/hybrid/dump.py @@ -0,0 +1,25 @@ +"""Methods and data structures to support dumping HalideIR to Hybrid Script. +This allows users to do quick hack to generated HalideIR and cast it back to +TVM modules. +""" + +from .. import build_module +from .. import ir_pass +from .. import schedule +from .. import stmt as _stmt +from .. import expr as _expr + +class HybridModule(object): + """The usage of Hybrid Module is very similar to conventional TVM module, + but conventional TVM module requires a function body which is already fully + lowered. This contradicts to the fact that Hybrid Module is originally a text + format for Phase 0 HalideIR. Thus, a totally separated module is defined.""" + + def __init__(self): + pass + + def __call__(self): + pass + + def get_source(self): + pass diff --git a/python/tvm/hybrid/runtime.py b/python/tvm/hybrid/runtime.py index 5df470e09085..5013e367c540 100644 --- a/python/tvm/hybrid/runtime.py +++ b/python/tvm/hybrid/runtime.py @@ -73,7 +73,6 @@ def sigmoid(x): HYBRID_GLOBALS = { - 'len' : len, 'unroll' : range, 'vectorize' : range, 'parallel' : range, diff --git a/src/codegen/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc similarity index 53% rename from src/codegen/codegen_hybrid.cc rename to src/contrib/hybrid/codegen_hybrid.cc index 0bc0922ffff9..7a9d57829b24 100644 --- a/src/codegen/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -4,161 +4,27 @@ #include #include #include "codegen_hybrid.h" -#include "../pass/ir_util.h" -#include "../arithmetic/compute_expr.h" namespace tvm { -namespace codegen { +namespace contrib { using namespace ir; -void CodeGenHybrid::Init(bool simple_mode) { - simple_mode_ = simple_mode; -} - -void CodeGenHybrid::InitFuncState(LoweredFunc f) { - alloc_storage_scope_.clear(); - handle_data_type_.clear(); - CodeGenSourceBase::ClearFuncState(); -} - -std::string CodeGenHybrid::GetVarID(const Variable* v) { - auto it = var_idmap_.find(v); - if (!simple_mode_) { - CHECK(it != var_idmap_.end()) << "Find undefined Variable " << v->name_hint; - } else { - if (it == var_idmap_.end()) - return AllocVarID(v); - } - return it->second; -} - -void CodeGenHybrid::ReserveKeywordsAsUnique() { - GetUniqueName("def"); - GetUniqueName("for"); - GetUniqueName("in"); - GetUniqueName("range"); - GetUniqueName("unroll"); - GetUniqueName("vectorize"); - GetUniqueName("parallel"); - GetUniqueName("if"); - GetUniqueName("else"); - GetUniqueName("and"); - GetUniqueName("or"); - GetUniqueName("not"); -} - -void CodeGenHybrid::AddFunction(LoweredFunc f) { - // clear previous generated state. - InitFuncState(f); - // reserve keywords - ReserveKeywordsAsUnique(); - // add to alloc buffer type. - for (const auto & kv : f->handle_data_type) { - RegisterHandleType(kv.first.get(), kv.second.type()); +std::string CodeGenHybrid::GetUniqueName(std::string name) { + for (size_t i = 0; i < name.size(); ++i) { + if (name[i] == '.') + name[i] = '_'; } - - stream << "def " << f->name << "("; - for (size_t i = 0; i < f->args.size(); ++i) { - Var v = f->args[i]; - stream << ' ' << v->name_hint; + auto iter = ids_allocated_.find(name); + if (iter == ids_allocated_.end()) { + ids_allocated_[name] = 1; + return name; } - stream << "):\n"; - int func_scope = BeginScope(); - PrintStmt(f->body); - EndScope(func_scope); + return name + std::to_string(ids_allocated_[name]++); } std::string CodeGenHybrid::Finish() { - return decl_stream.str() + stream.str(); -} - -void CodeGenHybrid::PrintExpr(const Expr& n, std::ostream& os) { // NOLINT(*) - VisitExpr(n, os); -} - -void CodeGenHybrid::PrintSSAAssign(const std::string& target, const std::string& src, Type t) { - LOG(FATAL) << "Python backend does not support SSA format."; -} - -// Print a reference expression to a buffer. -std::string CodeGenHybrid::GetBufferRef( - Type t, const Variable* buffer, Expr index) { - std::ostringstream os; - std::string vid = GetVarID(buffer); - os << vid << "["; - PrintExpr(index, os); - os << "]"; - return os.str(); -} - -// Print a reference expression to a buffer. -std::string CodeGenHybrid::GetStructRef( - Type t, const Expr& buffer, const Expr& index, int kind) { - if (kind < intrinsic::kArrKindBound_) { - std::ostringstream os; - os << "(((TVMArray*)"; - PrintExpr(buffer, os); - os << ")"; - if (kind == intrinsic::kArrAddr) { - os << " + "; - PrintExpr(index, os); - os << ")"; - return os.str(); - } - os << '['; - PrintExpr(index, os); - os << "]."; - // other case: get fields. - switch (kind) { - case intrinsic::kArrData: os << "data"; break; - case intrinsic::kArrShape: os << "shape"; break; - case intrinsic::kArrStrides: os << "strides"; break; - case intrinsic::kArrNDim: os << "ndim"; break; - case intrinsic::kArrTypeCode: os << "dtype.code"; break; - case intrinsic::kArrTypeBits: os << "dtype.bits"; break; - case intrinsic::kArrByteOffset: os << "byte_offset"; break; - case intrinsic::kArrTypeLanes: os << "dtype.lanes"; break; - case intrinsic::kArrDeviceId: os << "ctx.device_id"; break; - case intrinsic::kArrDeviceType: os << "ctx.device_type"; break; - default: LOG(FATAL) << "unknown field code"; - } - os << ')'; - return os.str(); - } else { - CHECK_LT(kind, intrinsic::kTVMValueKindBound_); - std::ostringstream os; - os << "(((TVMValue*)"; - PrintExpr(buffer, os); - os << ")[" << index << "]."; - if (t.is_handle()) { - os << "v_handle"; - } else if (t.is_float()) { - os << "v_float64"; - } else if (t.is_int()) { - os << "v_int64"; - } else { - LOG(FATAL) << "Do not know how to handle type" << t; - } - os << ")"; - return os.str(); - } -} - - -bool CodeGenHybrid::HandleTypeMatch(const Variable* buf_var, Type t) const { - auto it = handle_data_type_.find(buf_var); - if (it == handle_data_type_.end()) return false; - return it->second == t; -} - -void CodeGenHybrid::RegisterHandleType(const Variable* buf_var, Type t) { - auto it = handle_data_type_.find(buf_var); - if (it == handle_data_type_.end()) { - handle_data_type_[buf_var] = t; - } else { - CHECK(it->second == t) << "conflicting buf var type"; - } + return stream.str(); } void CodeGenHybrid::PrintType(Type t, std::ostream &os) { @@ -176,26 +42,6 @@ void CodeGenHybrid::PrintType(Type t, std::ostream &os) { os << t.bits(); } -std::string CodeGenHybrid::CastFromTo(std::string value, Type from, Type target) { - if (from == target) return value; - std::ostringstream os; - PrintType(target, os); - os << "(" << value << ")"; - return os.str(); -} - -void CodeGenHybrid::BindThreadIndex(const IterVar& iv) { - LOG(FATAL) << "to be implemented"; -} - -void CodeGenHybrid::PrintStorageSync(const Call* op) { // NOLINT(*) - LOG(FATAL) << "to be implemented"; -} - -void CodeGenHybrid::PrintStorageScope(const std::string& scope, std::ostream& os) { // NOLINT(*) - CHECK_EQ(scope, "global"); -} - void CodeGenHybrid::VisitExpr_(const IntImm *op, std::ostream& os) { // NOLINT(*) PrintType(op->type, os); os << "(" << op->value << ")"; @@ -251,11 +97,19 @@ inline void PrintBinaryIntrinsitc(const Call* op, LOG(FATAL) << "vec bin intrin to be implemented"; } } + void CodeGenHybrid::VisitExpr_(const Cast *op, std::ostream& os) { // NOLINT(*) - std::stringstream value; - PrintExpr(op->value, value); - os << CastFromTo(value.str(), op->value.type(), op->type); + if (op->type == op->value.type()) { + PrintExpr(op->value, stream); + } else { + PrintType(op->type, stream); + stream << op->type; + os << "("; + PrintExpr(op->value, stream); + os << ")"; + } } + void CodeGenHybrid::VisitExpr_(const Variable *op, std::ostream& os) { // NOLINT(*) os << GetVarID(op); } @@ -310,16 +164,18 @@ void CodeGenHybrid::VisitExpr_(const Not *op, std::ostream& os) { // NOLINT(*) } void CodeGenHybrid::VisitExpr_(const Call *op, std::ostream& os) { // NOLINT(*) - if (op->call_type == Call::Extern || - op->call_type == Call::PureExtern) { - os << op->name << "("; - for (size_t i = 0; i < op->args.size(); i++) { - PrintExpr(op->args[i], os); - if (i < op->args.size() - 1) { - os << ", "; - } + if (op->call_type == Call::Halide) { + os << GetTensorID(op->func, op->value_index); + os << "["; + for (size_t i = 0; i < op->args.size(); ++i) { + if (i) os << ", "; + std::stringstream idx; + PrintExpr(op->args[i], idx); + os << idx.str(); } - os << ")"; + os << "]"; + } if (op->call_type == Call::Extern || + op->call_type == Call::PureExtern) { } else if (op->is_intrinsic(Call::bitwise_and)) { PrintBinaryIntrinsitc(op, " & ", os, this); } else if (op->is_intrinsic(Call::bitwise_xor)) { @@ -342,53 +198,35 @@ void CodeGenHybrid::VisitExpr_(const Call *op, std::ostream& os) { // NOLINT(*) os << " else "; PrintExpr(op->args[2], os); } else { - // TODO(@were): Support tvm runtime intrinsics: - // intrinsic::tvm_address_of - // intrinsic::tvm_struct_get - // intrinsic::tvm_handle_is_null - if (op->call_type == Call::Intrinsic || - op->call_type == Call::PureIntrinsic) { - LOG(FATAL) << "Unresolved intrinsic " << op->name - << " with return type " << op->type; - } else { - LOG(FATAL) << "Unresolved call type " << op->call_type; + os << op->name << "("; + for (size_t i = 0; i < op->args.size(); i++) { + PrintExpr(op->args[i], os); + if (i < op->args.size() - 1) { + os << ", "; + } } + os << ")"; } } void CodeGenHybrid::VisitExpr_(const Load* op, std::ostream& os) { // NOLINT(*) - // int lanes = op->type.lanes(); - // delcare type. - if (op->type.lanes() == 1) { - std::string ref = GetBufferRef(op->type, op->buffer_var.get(), op->index); - os << ref; - } else { - LOG(FATAL) << "vec load to be supported"; - } + LOG(FATAL) << "Phase 0 has no Load(s)!"; } void CodeGenHybrid::VisitStmt_(const Store* op) { - Type t = op->value.type(); - if (t.lanes() == 1) { - std::string value = PrintExpr(op->value); - std::string ref = GetBufferRef(t, op->buffer_var.get(), op->index); - PrintIndent(); - stream << ref << " = " << value << "\n"; - } else { - LOG(FATAL) << "Vectorized store is not supported yet..."; - } + LOG(FATAL) << "Phase 0 has no Store(s)!"; } void CodeGenHybrid::VisitExpr_(const Let* op, std::ostream& os) { // NOLINT(*) - std::string value = PrintExpr(op->value); - CHECK(!var_idmap_.count(op->var.get())); - var_idmap_[op->var.get()] = value; - os << PrintExpr(op->body); + LOG(FATAL) << "Phase 0 has no Let(s)!"; +} + +void CodeGenHybrid::VisitStmt_(const Allocate* op) { + LOG(FATAL) << "Phase 0 has no Allocate(s)!"; } void CodeGenHybrid::VisitExpr_(const Ramp* op, std::ostream& os) { // NOLINT(*) - // TODO(@were): Support vectorization access in both frontend and backend - LOG(FATAL) << "ramp to be supported yet"; + LOG(FATAL) << "Ramp to be supported yet"; } void CodeGenHybrid::VisitExpr_(const Broadcast* op, std::ostream& os) { // NOLINT(*) @@ -406,21 +244,7 @@ void CodeGenHybrid::VisitExpr_(const Select* op, std::ostream& os) { // NOLINT( void CodeGenHybrid::VisitStmt_(const LetStmt* op) { std::string value = PrintExpr(op->value); - stream << AllocVarID(op->var.get()) << " = " << value << ";\n"; - PrintStmt(op->body); -} - -void CodeGenHybrid::VisitStmt_(const Allocate* op) { - CHECK(!is_zero(op->condition)); - std::string vid = AllocVarID(op->buffer_var.get()); - PrintIndent(); - stream << vid << " = allocate(("; - for (size_t i = 0; i < op->extents.size(); ++i) { - if (!i) stream << ", "; - stream << PrintExpr(op->extents[i]); - } - stream << "), \"" << op-> type << "\")\n"; - RegisterHandleType(op->buffer_var.get(), op->type); + stream << GetVarID(op->var.get()) << " = " << value << ";\n"; PrintStmt(op->body); } @@ -428,20 +252,29 @@ void CodeGenHybrid::VisitStmt_(const AttrStmt* op) { // TODO(@were): Support thread and buffer binding if (op->attr_key == ir::attr::thread_extent) { LOG(FATAL) << "Thread binding support yet!\n"; - } else if (op->attr_key == ir::attr::storage_scope) { - const Variable* v = op->node.as(); - CHECK(v); + } else if (op->attr_key == ir::attr::realize_scope) { + auto v = FunctionRef(op->node.node_); alloc_storage_scope_[v] = op->value.as()->value; - } else if (op->attr_key == ir::attr::volatile_scope) { - const Variable* v = op->node.as(); - CHECK(v); - volatile_buf_.insert(v); } PrintStmt(op->body); } +void CodeGenHybrid::VisitStmt_(const Realize *op) { + PrintIndent(); + stream << GetTensorID(op->func, op->value_index) << " = allocate(("; + for (size_t i = 0; i < op->bounds.size(); ++i) { + if (i) stream << ", "; + stream << PrintExpr(op->bounds[i]->extent); + } + stream << "), \""; + PrintType(op->type, stream); + stream << "\", "; + CHECK(alloc_storage_scope_.count(op->func)); + stream << alloc_storage_scope_[op->func] << ")\n"; +} + void CodeGenHybrid::VisitStmt_(const AssertStmt* op) { - // TODO(@were): Support AssertStmt in both hybrid parser and here + PrintIndent(); stream << "assert "; PrintExpr(op->condition, stream); stream << ", "; @@ -450,31 +283,43 @@ void CodeGenHybrid::VisitStmt_(const AssertStmt* op) { PrintStmt(op->body); } +void CodeGenHybrid::VisitStmt_(const Provide* op) { + PrintIndent(); + stream << GetTensorID(op->func, op->value_index); + stream << "["; + for (size_t i = 0; i < op->args.size(); ++i) { + if (i) stream << ", "; + PrintExpr(op->args[i], stream); + } + stream << "]"; + PrintExpr(op->value, stream); + stream << "\n"; +} + void CodeGenHybrid::VisitStmt_(const For* op) { std::string extent = PrintExpr(op->extent); PrintIndent(); - std::string vid = AllocVarID(op->loop_var.get()); + std::string vid = GetVarID(op->loop_var.get()); stream << "for " << vid << " in " << "range(" << extent << "):\n"; - int for_scope = BeginScope(); + indent_ += tab_; PrintStmt(op->body); - EndScope(for_scope); - PrintIndent(); + indent_ -= tab_; } void CodeGenHybrid::VisitStmt_(const IfThenElse* op) { std::string cond = PrintExpr(op->condition); PrintIndent(); stream << "if " << cond << ":\n"; - int then_scope = BeginScope(); + indent_ += tab_; PrintStmt(op->then_case); - EndScope(then_scope); + indent_ -= tab_; if (op->else_case.defined()) { PrintIndent(); stream << "else:\n"; - int else_scope = BeginScope(); + indent_ += tab_; PrintStmt(op->else_case); - EndScope(else_scope); + indent_ -= tab_; } } @@ -497,5 +342,29 @@ void CodeGenHybrid::VisitStmt_(const ProducerConsumer *op) { PrintIndent(); stream << "# produced " << op->func->func_name() << "\n"; } -} // namespace codegen + +void CodeGenHybrid::PrintIndent() { + stream << std::string(indent_, ' '); +} + +std::string CodeGenHybrid::GetVarID(const Variable *v) { + auto node = v->GetNodePtr().get(); + if (id_map_.count(node)) { + return id_map_[node]; + } + return id_map_[node] = GetUniqueName(v->name_hint); +} + +std::string CodeGenHybrid::GetTensorID(const FunctionRef &func, int value_index) { + auto node = func.get(); + if (id_map_.count(node)) { + return id_map_[node]; + } + std::string name_hint = func->func_name(); + if (func->num_outputs() != 0) { + name_hint += ".v" + std::to_string(value_index); + } + return id_map_[node] = GetUniqueName(name_hint); +} +} // namespace contrib } // namespace tvm diff --git a/src/codegen/codegen_hybrid.h b/src/contrib/hybrid/codegen_hybrid.h similarity index 65% rename from src/codegen/codegen_hybrid.h rename to src/contrib/hybrid/codegen_hybrid.h index a1e900e1e24f..9eaede94cb64 100644 --- a/src/codegen/codegen_hybrid.h +++ b/src/contrib/hybrid/codegen_hybrid.h @@ -10,36 +10,31 @@ #include #include #include +#include #include #include #include #include -#include "codegen_source_base.h" namespace tvm { -namespace codegen { +namespace contrib { using namespace ir; /*! - * \brief A base class to generate Python script. - * - * Unlike C-family generators, CodeGenHybrid does not generate SSA formed Python script. - * Only normal form is supported. + * \brief A base class to generate Hybrid Script. * * **NOTE** CodeGenHybrid does not aim at generating Python scripts consumed by Python2/3. * For runtime support, please refer the decorator in ``tvm/python/hybrid/api.py``. */ class CodeGenHybrid : public ExprFunctor, - public StmtFunctor, - public CodeGenSourceBase { + public StmtFunctor { public: - void Init(bool simple_mode); /*! - * \brief Add the function to the generated module. - * \param f The function to be compiled. + * \brief Dump the given schedule to hybrid script. + * \param sch The schedule to be dumped to hybrid script. */ - void AddFunction(LoweredFunc f); + void DumpSchedule(const Schedule &sch); /*! * \brief Finalize the compilation and return the code. * \return The code. @@ -49,29 +44,26 @@ class CodeGenHybrid : * \brief Print the Stmt n to CodeGenHybrid->stream * \param n The statement to be printed. */ - void PrintStmt(const Stmt& n) { - VisitStmt(n); + void PrintStmt(const Stmt &n) { + this->VisitStmt(n); } /*! * \brief Print the expression n(or its ssa id if in ssa mode) into os * \param n The expression to be printed. * \param os The output stream */ - void PrintExpr(const Expr& n, std::ostream& os); + void PrintExpr(const Expr &n, std::ostream &os) { + this->VisitExpr(n, os); + } /*! * \brief Same as PrintExpr, but simply returns result string * \param n The expression to be printed. */ - std::string PrintExpr(const Expr& n) { + std::string PrintExpr(const Expr &n) { std::ostringstream os; PrintExpr(n, os); return os.str(); } - /*! - * \brief Initialize codegen state for generating f. - * \param f The function to be compiled. - */ - virtual void InitFuncState(LoweredFunc f); // expression void VisitExpr_(const Variable* op, std::ostream& os) override; // NOLINT(*) void VisitExpr_(const Load* op, std::ostream& os) override; // NOLINT(*) @@ -104,9 +96,11 @@ class CodeGenHybrid : // statment void VisitStmt_(const LetStmt* op) override; void VisitStmt_(const Store* op) override; + void VisitStmt_(const Provide* op) override; void VisitStmt_(const For* op) override; void VisitStmt_(const IfThenElse* op) override; void VisitStmt_(const Allocate* op) override; + void VisitStmt_(const Realize* op) override; void VisitStmt_(const AttrStmt* op) override; void VisitStmt_(const AssertStmt* op) override; void VisitStmt_(const Evaluate* op) override; @@ -118,55 +112,30 @@ class CodeGenHybrid : * \param os The stream to print the ctype into */ virtual void PrintType(Type t, std::ostream& os); // NOLINT(*) - /*! - * \brief Print expr representing the thread tag - * \param IterVar iv The thread index to be binded; - */ - virtual void BindThreadIndex(const IterVar& iv); // NOLINT(*) - virtual void PrintStorageScope(const std::string& scope, std::ostream& os); // NOLINT(*) - virtual void PrintStorageSync(const Call* op); // NOLINT(*) // Get a cast type from to virtual std::string CastFromTo(std::string value, Type from, Type target); - protected: - std::string GetVarID(const Variable* v); - // Print reference to struct location - std::string GetStructRef( - Type t, const Expr& buffer, const Expr& index, int kind); - // print reference to a buffer as type t in index. - virtual std::string GetBufferRef( - Type t, const Variable* buffer, Expr index); - /*! - * \brief If buffer is allocated as type t. - * \param buf_var The buffer variable. - * \param t The type to be checked. - */ - bool HandleTypeMatch(const Variable* buf_var, Type t) const; - /*! - * \brief Register the data type of buf_var - * \param buf_var The buffer variable. - * \param t The type to be checked. - */ - void RegisterHandleType(const Variable* buf_var, Type t); - // override - void PrintSSAAssign( - const std::string& target, const std::string& src, Type t) final; - /*! \brief restrict keyword */ - std::string restrict_keyword_{""}; + private: + // + int indent_{0}; + const int tab_{2}; + inline void PrintIndent(); + // + std::map ids_allocated_; + std::map id_map_; + // + std::string GetUniqueName(std::string s); + // + std::stringstream stream; + // + std::string GetVarID(const Variable *v); + std::string GetTensorID(const FunctionRef &func, int value_index); /*! \brief the storage scope of allocation */ - std::unordered_map alloc_storage_scope_; + std::map alloc_storage_scope_; /*! \brief the data type of allocated buffers */ std::unordered_map handle_data_type_; - /*! \brief reserves common C keywords */ - void ReserveKeywordsAsUnique(); - - private: - /*! \brief whether print a simple form */ - bool simple_mode_{false}; - /*! \brief set of volatile buf access */ - std::unordered_set volatile_buf_; }; -} // namespace codegen +} // namespace contrib } // namespace tvm #endif // TVM_CODEGEN_CODEGEN_HYBRID_H_ diff --git a/src/op/hybrid_op.cc b/src/op/hybrid_op.cc index 26daefa76d7f..0268498c7db2 100644 --- a/src/op/hybrid_op.cc +++ b/src/op/hybrid_op.cc @@ -173,25 +173,28 @@ Stmt HybridOpNode::BuildProvide( rmap[outputs[i]] = stage->op.output(i); } auto n = make_node(*this); - /* - * These two lines of codes replace tensors' reads & writes. + /* This is a story little bit complicated. + * The following two lines of codes replace output tensors' usage. * This is the simplest way I (@were) can come up with to glue - * hybrid scripts to the structure of TVM op. - * NAMING CONFLICT: In hybrid script all the tensors have their own - * names specified by the users. However, In TVM op, all the output - * tensors' names are the same as the op's name. I cannot change the - * name to the op's name in the function body after the op node is - * formed, because: - * 1. Output tensors all point to the corresponding op node. - * 2. Once OpNode is wrapped up by an Operation node, it can - * no longer be changed. + * hybrid operation node to TVM op system. + * In hybrid script all the tensors, especially the output tensors, + * have their own names defined by the users. However, In TVM + * conventional ops: + * 1. Output tensors refer the corresponding op node so that the output + * tensors have the same names as the operation produces them. + * 2. Once OpNode is wrapped up by an Operation node, it is finalized. + * Later access will be from a const OpNode*. * This is a chiken-egg paradox. It is impossible to put the output * tensors into the function body without forming the op node. The * function body is immutable after the node is formed. * * Finally, I decided to resolve this issue "lazily". During the - * pipeline of compilation, these tensors will be replaced when - * forming the function body and passing to next stage of compilation. + * pipeline of compilation, this stage is a very preliminary stage. + * Technically, it is before Phase 0. The actual tensors will be replaced + * here. + * Thus, the operation body is slightly different from the Phase 0 body. + * This is a major difference that HybridOpNode is NOT the same as + * ExternOpNode. * */ ret = op::ReplaceTensor(ret, rmap); ret = op::ReplaceProvideTensor(ret, rmap); From baaaa6a2fdb8c9e2dc17c1b0ce8a782c5ca6b3f5 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 13:03:20 -0800 Subject: [PATCH 14/31] temp commit --- src/contrib/hybrid/codegen_hybrid.cc | 93 ++++++++++++++-------------- src/contrib/hybrid/codegen_hybrid.h | 29 ++++++--- 2 files changed, 69 insertions(+), 53 deletions(-) diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index 7a9d57829b24..7c641934db54 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -10,17 +10,24 @@ namespace contrib { using namespace ir; -std::string CodeGenHybrid::GetUniqueName(std::string name) { - for (size_t i = 0; i < name.size(); ++i) { - if (name[i] == '.') - name[i] = '_'; +std::string CodeGenHybrid::GetUniqueName(std::string prefix) { + for (size_t i = 0; i < prefix.size(); ++i) { + if (prefix[i] == '.') prefix[i] = '_'; } - auto iter = ids_allocated_.find(name); - if (iter == ids_allocated_.end()) { - ids_allocated_[name] = 1; - return name; + auto it = ids_allocated_.find(prefix); + if (it != ids_allocated_.end()) { + while (true) { + std::ostringstream os; + os << prefix << (++it->second); + std::string name = os.str(); + if (ids_allocated_.count(name) == 0) { + prefix = name; + break; + } + } } - return name + std::to_string(ids_allocated_[name]++); + ids_allocated_[prefix] = 0; + return prefix; } std::string CodeGenHybrid::Finish() { @@ -52,7 +59,7 @@ void CodeGenHybrid::VisitExpr_(const UIntImm *op, std::ostream& os) { // NOLINT } void CodeGenHybrid::VisitExpr_(const FloatImm *op, std::ostream& os) { // NOLINT(*) PrintType(op->type, os); - os << "(" << op->value << ")"; + os << "(" << std::setprecision(20) << op->value << ")"; } void CodeGenHybrid::VisitExpr_(const StringImm *op, std::ostream& os) { // NOLINT(*) os << "\"" << op->value << "\""; @@ -63,22 +70,19 @@ inline void PrintBinaryExpr(const T* op, const char *opstr, std::ostream& os, // NOLINT(*) CodeGenHybrid* p) { - if (op->type.lanes() == 1) { - if (isalpha(opstr[0])) { - os << opstr << '('; - p->PrintExpr(op->a, os); - os << ", "; - p->PrintExpr(op->b, os); - os << ')'; - } else { - os << '('; - p->PrintExpr(op->a, os); - os << ' ' << opstr << ' '; - p->PrintExpr(op->b, os); - os << ')'; - } + CHECK(op->type.lanes() == 1) << "vec bin op not implemented"; + if (isalpha(opstr[0])) { + os << opstr << '('; + p->PrintExpr(op->a, os); + os << ", "; + p->PrintExpr(op->b, os); + os << ')'; } else { - LOG(FATAL) << "vec bin op to be implemented"; + os << '('; + p->PrintExpr(op->a, os); + os << ' ' << opstr << ' '; + p->PrintExpr(op->b, os); + os << ')'; } } @@ -86,16 +90,13 @@ inline void PrintBinaryIntrinsitc(const Call* op, const char *opstr, std::ostream& os, // NOLINT(*) CodeGenHybrid* p) { - if (op->type.lanes() == 1) { - CHECK_EQ(op->args.size(), 2U); - os << '('; - p->PrintExpr(op->args[0], os); - os << opstr; - p->PrintExpr(op->args[1], os); - os << ')'; - } else { - LOG(FATAL) << "vec bin intrin to be implemented"; - } + CHECK(op->type.lanes() == 1) << "vec bin intrin not implemented"; + CHECK_EQ(op->args.size(), 2U); + os << '('; + p->PrintExpr(op->args[0], os); + os << opstr; + p->PrintExpr(op->args[1], os); + os << ')'; } void CodeGenHybrid::VisitExpr_(const Cast *op, std::ostream& os) { // NOLINT(*) @@ -174,23 +175,21 @@ void CodeGenHybrid::VisitExpr_(const Call *op, std::ostream& os) { // NOLINT(*) os << idx.str(); } os << "]"; - } if (op->call_type == Call::Extern || - op->call_type == Call::PureExtern) { } else if (op->is_intrinsic(Call::bitwise_and)) { - PrintBinaryIntrinsitc(op, " & ", os, this); + PrintBinaryIntrinsitc(op, "&", os, this); } else if (op->is_intrinsic(Call::bitwise_xor)) { - PrintBinaryIntrinsitc(op, " ^ ", os, this); + PrintBinaryIntrinsitc(op, "^", os, this); } else if (op->is_intrinsic(Call::bitwise_or)) { - PrintBinaryIntrinsitc(op, " | ", os, this); + PrintBinaryIntrinsitc(op, "|", os, this); + } else if (op->is_intrinsic(Call::shift_left)) { + PrintBinaryIntrinsitc(op, "<<", os, this); + } else if (op->is_intrinsic(Call::shift_right)) { + PrintBinaryIntrinsitc(op, ">>", os, this); } else if (op->is_intrinsic(Call::bitwise_not)) { CHECK_EQ(op->args.size(), 1U); os << "(~"; PrintExpr(op->args[0], os); os << ')'; - } else if (op->is_intrinsic(Call::shift_left)) { - PrintBinaryIntrinsitc(op, " << ", os, this); - } else if (op->is_intrinsic(Call::shift_right)) { - PrintBinaryIntrinsitc(op, " >> ", os, this); } else if (op->is_intrinsic(intrinsic::tvm_if_then_else)) { PrintExpr(op->args[1], os); os << " if "; @@ -362,9 +361,13 @@ std::string CodeGenHybrid::GetTensorID(const FunctionRef &func, int value_index) } std::string name_hint = func->func_name(); if (func->num_outputs() != 0) { - name_hint += ".v" + std::to_string(value_index); + name_hint += "_v" + std::to_string(value_index); } return id_map_[node] = GetUniqueName(name_hint); } + +void CodeGenHybrid::DumpSchedule(const Schedule &sch) { + sch->outputs; +} } // namespace contrib } // namespace tvm diff --git a/src/contrib/hybrid/codegen_hybrid.h b/src/contrib/hybrid/codegen_hybrid.h index 9eaede94cb64..28874839c59c 100644 --- a/src/contrib/hybrid/codegen_hybrid.h +++ b/src/contrib/hybrid/codegen_hybrid.h @@ -107,28 +107,41 @@ class CodeGenHybrid : void VisitStmt_(const Block* op) override; void VisitStmt_(const ProducerConsumer* op) override; /*! - * Print Type represetnation of type t. + * \brief Print Type represetnation of type t. * \param t The type representation. * \param os The stream to print the ctype into */ virtual void PrintType(Type t, std::ostream& os); // NOLINT(*) - // Get a cast type from to virtual std::string CastFromTo(std::string value, Type from, Type target); private: - // + /*! \brief The current indent of the code dump. */ int indent_{0}; + /*! \brief The tab size of code indent. */ const int tab_{2}; + /*! \brief Print the current indent spaces. */ inline void PrintIndent(); - // + /*! \brief Keys are ids allocated, and values are the suffix to prevent double-name. */ std::map ids_allocated_; + /*! \brief Keys are either tensors or variables. Values are the corresponding IDs.*/ std::map id_map_; - // - std::string GetUniqueName(std::string s); - // + /*! + * \brief Find an unallocated name for the given prefix. + * \param prefix The given prefix. + */ + std::string GetUniqueName(std::string prefix); + /*! \brief The output code string builder. */ std::stringstream stream; - // + /*! + * \brief Get or allocate the ID for the given variable. + * \param v The given variable. + */ std::string GetVarID(const Variable *v); + /*! + * \brief Get or allocate the ID for the given tensor. + * \param func The tensor to allocate a name. + * \param value_index The value index of the given tensor. + */ std::string GetTensorID(const FunctionRef &func, int value_index); /*! \brief the storage scope of allocation */ std::map alloc_storage_scope_; From b64cd4e9fa1748311d6e77f65bb7c5b65495bacb Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 13:49:03 -0800 Subject: [PATCH 15/31] fix import --- python/tvm/hybrid/__init__.py | 4 +++- src/contrib/hybrid/codegen_hybrid.cc | 36 +++++++++++++++++++++++++--- src/contrib/hybrid/codegen_hybrid.h | 10 ++++---- 3 files changed, 42 insertions(+), 8 deletions(-) diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index 6cebf463e608..a317335e7937 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -22,7 +22,6 @@ from .parser import parse_python from .util import _pruned_source -from .dump import dump def script(pyfunc): @@ -58,3 +57,6 @@ def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring return value return decorate(pyfunc, wrapped_func) + + +_init_api("tvm.hybrid") diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index 7c641934db54..f2f5ec9d803b 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -366,8 +366,38 @@ std::string CodeGenHybrid::GetTensorID(const FunctionRef &func, int value_index) return id_map_[node] = GetUniqueName(name_hint); } -void CodeGenHybrid::DumpSchedule(const Schedule &sch) { - sch->outputs; -} +void CodeGenHybrid::DumpSchedule(const Stmt &stmt, + const Array &inputs, + const Array &outputs, + const std::string &name) { + stream << "def " << name << "("; + for (size_t i = 0; i < inputs.size(); ++i) { + if (i) stream << ", "; + stream << GetTensorID(inputs[i]->op, inputs[i]->value_index); + } + stream << "):\n"; + indent_ += tab_; + for (size_t i = 0; i < outputs.size(); ++i) { + PrintIndent(); + stream << GetTensorID(outputs[i]->op, outputs[i]->value_index) + << " = output_tensor(("; + for (size_t j = 0; j < outputs[i]->shape.size(); ++j) { + if (j) stream << ", "; + PrintExpr(outputs[i]->shape[j], stream); + stream << "), '" << outputs[i]->dtype << "')\n"; + } + } + PrintStmt(stmt); +} + +TVM_REGISTER_GLOBAL("hybrid.dump") +.set_body([](TVMArgs args, TVMRetValue* rv) { + CodeGenHybrid codegen; + if (args.size() == 4) + codegen.DumpSchedule(args[0], args[1], args[2], args[3]); + else + codegen.DumpSchedule(args[0], args[1], args[2]); + *rv = codegen.Finish(); + }); } // namespace contrib } // namespace tvm diff --git a/src/contrib/hybrid/codegen_hybrid.h b/src/contrib/hybrid/codegen_hybrid.h index 28874839c59c..2cec5a5d43ef 100644 --- a/src/contrib/hybrid/codegen_hybrid.h +++ b/src/contrib/hybrid/codegen_hybrid.h @@ -31,10 +31,13 @@ class CodeGenHybrid : public StmtFunctor { public: /*! - * \brief Dump the given schedule to hybrid script. - * \param sch The schedule to be dumped to hybrid script. + * \brief Dump the given function body to hybrid script. + * \param stmt The function body to be dumped to hybrid script. + * \param inputs Input tensors of this schedule. + * \param outputs Output tensors of this schedule. + * \param name The name of the function. */ - void DumpSchedule(const Schedule &sch); + void DumpSchedule(const Stmt &stmt, const Array &inputs, const Array &outputs, const std::string &name = "hybrid_func"); /*! * \brief Finalize the compilation and return the code. * \return The code. @@ -112,7 +115,6 @@ class CodeGenHybrid : * \param os The stream to print the ctype into */ virtual void PrintType(Type t, std::ostream& os); // NOLINT(*) - virtual std::string CastFromTo(std::string value, Type from, Type target); private: /*! \brief The current indent of the code dump. */ From 9847ef9883a3e0dd45830c91824a22eddd029181 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 15:10:58 -0800 Subject: [PATCH 16/31] i preliminary version is done? --- python/tvm/hybrid/calls.py | 7 +++++ python/tvm/hybrid/runtime.py | 1 + src/contrib/hybrid/codegen_hybrid.cc | 45 +++++++++++++++++----------- 3 files changed, 36 insertions(+), 17 deletions(-) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 9a3229479ae2..d854ef9a9763 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -8,6 +8,7 @@ from .. import ir_pass from ..stmt import For from .util import _internal_assert +from ..intrin import call_pure_intrin #pylint: disable=redefined-builtin @@ -123,3 +124,9 @@ def ceil_div(func_id, args): _internal_assert(isinstance(args[0], _expr.Expr), "Only expressions can div") a, b = args[0], args[1] return (a + b - 1) / b + + +def likely(func_id, args): + _internal_assert(args.__len__() == 1, \ + "Only one expression can be likely") + return call_pure_intrin(args[0].dtype, 'likely', args) diff --git a/python/tvm/hybrid/runtime.py b/python/tvm/hybrid/runtime.py index 5013e367c540..293e069c24ea 100644 --- a/python/tvm/hybrid/runtime.py +++ b/python/tvm/hybrid/runtime.py @@ -87,6 +87,7 @@ def sigmoid(x): 'exp' : numpy.exp, 'sigmoid' : sigmoid, 'popcount' : popcount, + 'likely' : lambda cond: cond, 'uint8' : numpy.uint8, 'uint16' : numpy.uint16, 'uint32' : numpy.uint32, diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index f2f5ec9d803b..3a82b04627b9 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -50,15 +50,13 @@ void CodeGenHybrid::PrintType(Type t, std::ostream &os) { } void CodeGenHybrid::VisitExpr_(const IntImm *op, std::ostream& os) { // NOLINT(*) - PrintType(op->type, os); - os << "(" << op->value << ")"; + os << op->value; } void CodeGenHybrid::VisitExpr_(const UIntImm *op, std::ostream& os) { // NOLINT(*) PrintType(op->type, os); - os << "(" << op->value << ")"; + os << op->value; } void CodeGenHybrid::VisitExpr_(const FloatImm *op, std::ostream& os) { // NOLINT(*) - PrintType(op->type, os); os << "(" << std::setprecision(20) << op->value << ")"; } void CodeGenHybrid::VisitExpr_(const StringImm *op, std::ostream& os) { // NOLINT(*) @@ -259,17 +257,21 @@ void CodeGenHybrid::VisitStmt_(const AttrStmt* op) { } void CodeGenHybrid::VisitStmt_(const Realize *op) { - PrintIndent(); - stream << GetTensorID(op->func, op->value_index) << " = allocate(("; - for (size_t i = 0; i < op->bounds.size(); ++i) { - if (i) stream << ", "; - stream << PrintExpr(op->bounds[i]->extent); - } - stream << "), \""; - PrintType(op->type, stream); - stream << "\", "; CHECK(alloc_storage_scope_.count(op->func)); - stream << alloc_storage_scope_[op->func] << ")\n"; + if (!alloc_storage_scope_[op->func].empty()) { + PrintIndent(); + stream << GetTensorID(op->func, op->value_index) << " = allocate(("; + for (size_t i = 0; i < op->bounds.size(); ++i) { + if (i) stream << ", "; + stream << PrintExpr(op->bounds[i]->extent); + } + if (op->bounds.size() == 1) stream << ", "; + stream << "), \""; + PrintType(op->type, stream); + stream << "\", '"; + stream << alloc_storage_scope_[op->func] << "')\n"; + } + PrintStmt(op->body); } void CodeGenHybrid::VisitStmt_(const AssertStmt* op) { @@ -290,7 +292,7 @@ void CodeGenHybrid::VisitStmt_(const Provide* op) { if (i) stream << ", "; PrintExpr(op->args[i], stream); } - stream << "]"; + stream << "] = "; PrintExpr(op->value, stream); stream << "\n"; } @@ -360,7 +362,7 @@ std::string CodeGenHybrid::GetTensorID(const FunctionRef &func, int value_index) return id_map_[node]; } std::string name_hint = func->func_name(); - if (func->num_outputs() != 0) { + if (func->num_outputs() > 1) { name_hint += "_v" + std::to_string(value_index); } return id_map_[node] = GetUniqueName(name_hint); @@ -384,10 +386,19 @@ void CodeGenHybrid::DumpSchedule(const Stmt &stmt, for (size_t j = 0; j < outputs[i]->shape.size(); ++j) { if (j) stream << ", "; PrintExpr(outputs[i]->shape[j], stream); - stream << "), '" << outputs[i]->dtype << "')\n"; } + if (outputs[i]->shape.size() == 1) + stream << ", "; + stream << "), '" << outputs[i]->dtype << "')\n"; } PrintStmt(stmt); + PrintIndent(); + stream << "return "; + for (size_t i = 0; i < outputs.size(); ++i) { + if (i) stream << ", "; + stream << GetTensorID(outputs[i]->op, outputs[i]->value_index); + } + stream << "\n"; } TVM_REGISTER_GLOBAL("hybrid.dump") From 6debdf1ce972b77334a1a201b67010445a253b80 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 16:43:30 -0800 Subject: [PATCH 17/31] on the way to build hybrid module --- python/tvm/hybrid/{dump.py => module.py} | 6 ++-- python/tvm/hybrid/var_decl.py | 2 +- src/contrib/hybrid/codegen_hybrid.cc | 32 ++++++++++++--------- src/contrib/hybrid/codegen_hybrid.h | 4 +-- tests/python/unittest/test_hybrid_script.py | 5 ++++ 5 files changed, 30 insertions(+), 19 deletions(-) rename python/tvm/hybrid/{dump.py => module.py} (90%) diff --git a/python/tvm/hybrid/dump.py b/python/tvm/hybrid/module.py similarity index 90% rename from python/tvm/hybrid/dump.py rename to python/tvm/hybrid/module.py index 1eb2b538dfe2..2e50a5a8ffbc 100644 --- a/python/tvm/hybrid/dump.py +++ b/python/tvm/hybrid/module.py @@ -15,11 +15,11 @@ class HybridModule(object): lowered. This contradicts to the fact that Hybrid Module is originally a text format for Phase 0 HalideIR. Thus, a totally separated module is defined.""" - def __init__(self): - pass + def __init__(self, src): + self.src_ = src def __call__(self): pass def get_source(self): - pass + return self.src_ diff --git a/python/tvm/hybrid/var_decl.py b/python/tvm/hybrid/var_decl.py index 47201c44003d..50b610567c74 100644 --- a/python/tvm/hybrid/var_decl.py +++ b/python/tvm/hybrid/var_decl.py @@ -45,7 +45,7 @@ def visit_Call(self, node): _internal_assert(isinstance(node.func, ast.Name), "Function call should be an id") func_id = node.func.id _internal_assert(func_id in list(HYBRID_GLOBALS.keys()) + \ - ['range', 'max', 'min'] + \ + ['range', 'max', 'min', 'len'] + \ list(self.symbols.keys()), \ "Function call id not in intrinsics' list") for elem in node.args: diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index 3a82b04627b9..021a044e4204 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -54,10 +54,10 @@ void CodeGenHybrid::VisitExpr_(const IntImm *op, std::ostream& os) { // NOLINT( } void CodeGenHybrid::VisitExpr_(const UIntImm *op, std::ostream& os) { // NOLINT(*) PrintType(op->type, os); - os << op->value; + os << "(" << op->value << ")"; } void CodeGenHybrid::VisitExpr_(const FloatImm *op, std::ostream& os) { // NOLINT(*) - os << "(" << std::setprecision(20) << op->value << ")"; + os << std::setprecision(20) << op->value; } void CodeGenHybrid::VisitExpr_(const StringImm *op, std::ostream& os) { // NOLINT(*) os << "\"" << op->value << "\""; @@ -101,10 +101,9 @@ void CodeGenHybrid::VisitExpr_(const Cast *op, std::ostream& os) { // NOLINT(*) if (op->type == op->value.type()) { PrintExpr(op->value, stream); } else { - PrintType(op->type, stream); - stream << op->type; + PrintType(op->type, os); os << "("; - PrintExpr(op->value, stream); + PrintExpr(op->value, os); os << ")"; } } @@ -368,14 +367,21 @@ std::string CodeGenHybrid::GetTensorID(const FunctionRef &func, int value_index) return id_map_[node] = GetUniqueName(name_hint); } -void CodeGenHybrid::DumpSchedule(const Stmt &stmt, - const Array &inputs, - const Array &outputs, - const std::string &name) { - stream << "def " << name << "("; +void CodeGenHybrid::DumpStmt(const Stmt &stmt, + const Array &inputs, + const Array &outputs, + const std::string &name) { + stream << "@tvm.hybrid.script\n" + << "def " << name << "("; for (size_t i = 0; i < inputs.size(); ++i) { if (i) stream << ", "; - stream << GetTensorID(inputs[i]->op, inputs[i]->value_index); + if (auto tensor = inputs[i].as()) { + stream << GetTensorID(tensor->op, tensor->value_index); + } else { + auto var = inputs[i].as(); + CHECK(var) << "Input should either be a tensor or a variable!"; + stream << GetVarID(var); + } } stream << "):\n"; indent_ += tab_; @@ -405,9 +411,9 @@ TVM_REGISTER_GLOBAL("hybrid.dump") .set_body([](TVMArgs args, TVMRetValue* rv) { CodeGenHybrid codegen; if (args.size() == 4) - codegen.DumpSchedule(args[0], args[1], args[2], args[3]); + codegen.DumpStmt(args[0], args[1], args[2], args[3]); else - codegen.DumpSchedule(args[0], args[1], args[2]); + codegen.DumpStmt(args[0], args[1], args[2]); *rv = codegen.Finish(); }); } // namespace contrib diff --git a/src/contrib/hybrid/codegen_hybrid.h b/src/contrib/hybrid/codegen_hybrid.h index 2cec5a5d43ef..3e75c670bf6e 100644 --- a/src/contrib/hybrid/codegen_hybrid.h +++ b/src/contrib/hybrid/codegen_hybrid.h @@ -37,7 +37,7 @@ class CodeGenHybrid : * \param outputs Output tensors of this schedule. * \param name The name of the function. */ - void DumpSchedule(const Stmt &stmt, const Array &inputs, const Array &outputs, const std::string &name = "hybrid_func"); + void DumpStmt(const Stmt &stmt, const Array &inputs, const Array &outputs, const std::string &name = "hybrid_func"); /*! * \brief Finalize the compilation and return the code. * \return The code. @@ -120,7 +120,7 @@ class CodeGenHybrid : /*! \brief The current indent of the code dump. */ int indent_{0}; /*! \brief The tab size of code indent. */ - const int tab_{2}; + const int tab_{4}; /*! \brief Print the current indent spaces. */ inline void PrintIndent(); /*! \brief Keys are ids allocated, and values are the suffix to prevent double-name. */ diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 6aafb551f13b..1a01b9b85090 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -43,6 +43,11 @@ def tvm_val_2_py_val(val): target=target) assert module + stmt = tvm.build_module.form_body(sch) + true_args = [i for i in args if isinstance(i, (tvm.tensor.Tensor, tvm.expr.Var))] + true_outs = [outs] if isinstance(outs, tvm.tensor.Tensor) else outs + print(tvm.hybrid.dump(stmt, true_args, true_outs)) + out_tensors = [] for i in range(op.num_outputs): output = op.output(i) From 5d614d40a0d6baf3e48f86ff1338c0c5b319d9c6 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 17:12:09 -0800 Subject: [PATCH 18/31] nearly fixed... --- python/tvm/hybrid/__init__.py | 1 + python/tvm/hybrid/calls.py | 3 ++- python/tvm/hybrid/module.py | 28 ++++++++++++++------- src/contrib/hybrid/codegen_hybrid.cc | 9 ++++--- src/contrib/hybrid/codegen_hybrid.h | 11 ++++---- tests/python/unittest/test_hybrid_script.py | 3 ++- 6 files changed, 35 insertions(+), 20 deletions(-) diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index a317335e7937..943c8145d750 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -22,6 +22,7 @@ from .parser import parse_python from .util import _pruned_source +from .module import HybridModule as Module def script(pyfunc): diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index d854ef9a9763..af796e9acce5 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -121,7 +121,7 @@ def ceil_div(func_id, args): _internal_assert(args.__len__() == 2, \ "Only one expression can be cast") for i in range(2): - _internal_assert(isinstance(args[0], _expr.Expr), "Only expressions can div") + _internal_assert(isinstance(args[i], _expr.Expr), "Only expressions can div") a, b = args[0], args[1] return (a + b - 1) / b @@ -129,4 +129,5 @@ def ceil_div(func_id, args): def likely(func_id, args): _internal_assert(args.__len__() == 1, \ "Only one expression can be likely") + _internal_assert(func_id == "likely", "This function cannot be directly invoked!") return call_pure_intrin(args[0].dtype, 'likely', args) diff --git a/python/tvm/hybrid/module.py b/python/tvm/hybrid/module.py index 2e50a5a8ffbc..eccdbfc81e33 100644 --- a/python/tvm/hybrid/module.py +++ b/python/tvm/hybrid/module.py @@ -3,11 +3,10 @@ TVM modules. """ -from .. import build_module -from .. import ir_pass -from .. import schedule -from .. import stmt as _stmt -from .. import expr as _expr +import imp +from ..contrib import util +from .util import _internal_assert + class HybridModule(object): """The usage of Hybrid Module is very similar to conventional TVM module, @@ -15,11 +14,22 @@ class HybridModule(object): lowered. This contradicts to the fact that Hybrid Module is originally a text format for Phase 0 HalideIR. Thus, a totally separated module is defined.""" - def __init__(self, src): - self.src_ = src - def __call__(self): - pass + def __init__(self, src, name): + temp = util.tempdir() + self.name_ = name + self.dst_ = temp.relpath(name) + self.src_ = 'import tvm\ntvm.hybrid.script\n%s' % src + with open(self.dst_, 'w') as f: + f.write(self.src_) + self.py_module_ = imp.load_source(name, self.dst_) + _internal_assert(hasattr(self.py_module_, name), \ + "The loaded source has no given function!") + + + def __call__(self, *args): + return getattr(self.py_module_, self.name_)(*args) + def get_source(self): return self.src_ diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index 021a044e4204..cbf2f48ebbd6 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -78,6 +78,8 @@ inline void PrintBinaryExpr(const T* op, } else { os << '('; p->PrintExpr(op->a, os); + if (!strcmp(opstr, "&&")) opstr = "and"; + if (!strcmp(opstr, "||")) opstr = "or"; os << ' ' << opstr << ' '; p->PrintExpr(op->b, os); os << ')'; @@ -157,7 +159,7 @@ void CodeGenHybrid::VisitExpr_(const Or *op, std::ostream& os) { // NOLINT(*) PrintBinaryExpr(op, "||", os, this); } void CodeGenHybrid::VisitExpr_(const Not *op, std::ostream& os) { // NOLINT(*) - os << '!'; + os << "not "; PrintExpr(op->a, os); } @@ -266,7 +268,7 @@ void CodeGenHybrid::VisitStmt_(const Realize *op) { } if (op->bounds.size() == 1) stream << ", "; stream << "), \""; - PrintType(op->type, stream); + PrintType(op->type, stream); stream << "\", '"; stream << alloc_storage_scope_[op->func] << "')\n"; } @@ -371,8 +373,7 @@ void CodeGenHybrid::DumpStmt(const Stmt &stmt, const Array &inputs, const Array &outputs, const std::string &name) { - stream << "@tvm.hybrid.script\n" - << "def " << name << "("; + stream << "def " << name << "("; for (size_t i = 0; i < inputs.size(); ++i) { if (i) stream << ", "; if (auto tensor = inputs[i].as()) { diff --git a/src/contrib/hybrid/codegen_hybrid.h b/src/contrib/hybrid/codegen_hybrid.h index 3e75c670bf6e..9e46a1085b07 100644 --- a/src/contrib/hybrid/codegen_hybrid.h +++ b/src/contrib/hybrid/codegen_hybrid.h @@ -3,8 +3,8 @@ * \file codegen_hybrid.h * \brief Common utilities to generated C style code. */ -#ifndef TVM_CODEGEN_CODEGEN_HYBRID_H_ -#define TVM_CODEGEN_CODEGEN_HYBRID_H_ +#ifndef TVM_CONTRIB_HYBRID_CODEGEN_HYBRID_H_ +#define TVM_CONTRIB_HYBRID_CODEGEN_HYBRID_H_ #include #include @@ -13,8 +13,8 @@ #include #include #include +#include #include -#include namespace tvm { namespace contrib { @@ -37,7 +37,8 @@ class CodeGenHybrid : * \param outputs Output tensors of this schedule. * \param name The name of the function. */ - void DumpStmt(const Stmt &stmt, const Array &inputs, const Array &outputs, const std::string &name = "hybrid_func"); + void DumpStmt(const Stmt &stmt, const Array &inputs, const Array &outputs, + const std::string &name = "hybrid_func"); /*! * \brief Finalize the compilation and return the code. * \return The code. @@ -153,4 +154,4 @@ class CodeGenHybrid : } // namespace contrib } // namespace tvm -#endif // TVM_CODEGEN_CODEGEN_HYBRID_H_ +#endif // TVM_CONTRIB_HYBRID_CODEGEN_HYBRID_H_ diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 1a01b9b85090..ca058d223848 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -46,7 +46,8 @@ def tvm_val_2_py_val(val): stmt = tvm.build_module.form_body(sch) true_args = [i for i in args if isinstance(i, (tvm.tensor.Tensor, tvm.expr.Var))] true_outs = [outs] if isinstance(outs, tvm.tensor.Tensor) else outs - print(tvm.hybrid.dump(stmt, true_args, true_outs)) + src = tvm.hybrid.dump(stmt, true_args, true_outs) + hmodule = tvm.hybrid.Module(src, "hybrid_func") out_tensors = [] for i in range(op.num_outputs): From 39048da204b768b7c39244c4c82cceb1e696c653 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 17:49:00 -0800 Subject: [PATCH 19/31] dumped python are equiv as original python --- python/tvm/hybrid/__init__.py | 28 +++++++++++++++++++++ python/tvm/hybrid/module.py | 22 ++++++++++------ src/contrib/hybrid/codegen_hybrid.cc | 25 +++++++++--------- src/contrib/hybrid/codegen_hybrid.h | 7 +++--- tests/python/unittest/test_hybrid_script.py | 14 ++++++++--- 5 files changed, 70 insertions(+), 26 deletions(-) diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index 943c8145d750..7e0d864a2e8a 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -19,6 +19,7 @@ from .._ffi.function import _init_api from .. import _api_internal as _tvm_internal from ..tensor import Tensor +from ..build_module import form_body from .parser import parse_python from .util import _pruned_source @@ -61,3 +62,30 @@ def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring _init_api("tvm.hybrid") + + +def build(sch, inputs, outputs, name="hybrid_func"): + """Dump the corrent schedule to hybrid module + + Parameters + ---------- + sch: Schedule + The schedule to be dumped + + inputs: An array of Tensors or Vars + The inputs of the function body + + outputs: An array of Tensors + The outputs of the function body + + Returns + ------- + module: HybridModule + The built results is wrapped in a HybridModule. + The usage of HybridModule is roughly the same as normal TVM-built modules. + """ + + stmt = form_body(sch) + src = dump(stmt, inputs, outputs, name) + + return Module(src, name) diff --git a/python/tvm/hybrid/module.py b/python/tvm/hybrid/module.py index eccdbfc81e33..1969656bf8ec 100644 --- a/python/tvm/hybrid/module.py +++ b/python/tvm/hybrid/module.py @@ -17,19 +17,27 @@ class HybridModule(object): def __init__(self, src, name): temp = util.tempdir() - self.name_ = name - self.dst_ = temp.relpath(name) - self.src_ = 'import tvm\ntvm.hybrid.script\n%s' % src - with open(self.dst_, 'w') as f: + dst = temp.relpath(name) + self.src_ = 'import tvm\n@tvm.hybrid.script\n%s' % src + with open(dst, 'w') as f: f.write(self.src_) - self.py_module_ = imp.load_source(name, self.dst_) - _internal_assert(hasattr(self.py_module_, name), \ + py_module = imp.load_source(name, dst) + _internal_assert(hasattr(py_module, name), \ "The loaded source has no given function!") + self.func_ = getattr(py_module, name) + _internal_assert(callable(self.func_), "This should be a function! At least callable!") def __call__(self, *args): - return getattr(self.py_module_, self.name_)(*args) + return self.func_(*args) def get_source(self): return self.src_ + + + def save(self, path): + if not path.endswith('.py'): + path = path + '.py' + with open(path, 'w') as f: + f.write(self.src_) diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index cbf2f48ebbd6..aed5bbe2cfed 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -123,7 +123,10 @@ void CodeGenHybrid::VisitExpr_(const Mul *op, std::ostream& os) { // NOLINT(*) PrintBinaryExpr(op, "*", os, this); } void CodeGenHybrid::VisitExpr_(const Div *op, std::ostream& os) { // NOLINT(*) - PrintBinaryExpr(op, "/", os, this); + if (op->type.is_int()) + PrintBinaryExpr(op, "//", os, this); + else + PrintBinaryExpr(op, "/", os, this); } void CodeGenHybrid::VisitExpr_(const Mod *op, std::ostream& os) { // NOLINT(*) PrintBinaryExpr(op, "%", os, this); @@ -338,11 +341,7 @@ void CodeGenHybrid::VisitStmt_(const Evaluate *op) { } void CodeGenHybrid::VisitStmt_(const ProducerConsumer *op) { - PrintIndent(); - stream << "# producing " << op->func->func_name() << "\n"; PrintStmt(op->body); - PrintIndent(); - stream << "# produced " << op->func->func_name() << "\n"; } void CodeGenHybrid::PrintIndent() { @@ -350,23 +349,23 @@ void CodeGenHybrid::PrintIndent() { } std::string CodeGenHybrid::GetVarID(const Variable *v) { - auto node = v->GetNodePtr().get(); - if (id_map_.count(node)) { - return id_map_[node]; + auto key = std::make_pair(v->GetNodePtr().get(), 0); + if (id_map_.count(key)) { + return id_map_[key]; } - return id_map_[node] = GetUniqueName(v->name_hint); + return id_map_[key] = GetUniqueName(v->name_hint); } std::string CodeGenHybrid::GetTensorID(const FunctionRef &func, int value_index) { - auto node = func.get(); - if (id_map_.count(node)) { - return id_map_[node]; + auto key = std::make_pair(func.get(), value_index); + if (id_map_.count(key)) { + return id_map_[key]; } std::string name_hint = func->func_name(); if (func->num_outputs() > 1) { name_hint += "_v" + std::to_string(value_index); } - return id_map_[node] = GetUniqueName(name_hint); + return id_map_[key] = GetUniqueName(name_hint); } void CodeGenHybrid::DumpStmt(const Stmt &stmt, diff --git a/src/contrib/hybrid/codegen_hybrid.h b/src/contrib/hybrid/codegen_hybrid.h index 9e46a1085b07..d45e744755c1 100644 --- a/src/contrib/hybrid/codegen_hybrid.h +++ b/src/contrib/hybrid/codegen_hybrid.h @@ -11,10 +11,11 @@ #include #include #include -#include -#include #include +#include #include +#include +#include namespace tvm { namespace contrib { @@ -127,7 +128,7 @@ class CodeGenHybrid : /*! \brief Keys are ids allocated, and values are the suffix to prevent double-name. */ std::map ids_allocated_; /*! \brief Keys are either tensors or variables. Values are the corresponding IDs.*/ - std::map id_map_; + std::map, std::string> id_map_; /*! * \brief Find an unallocated name for the given prefix. * \param prefix The given prefix. diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index ca058d223848..dfb1c6ff454e 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -22,15 +22,18 @@ def tvm_val_2_py_val(val): assert isinstance(outs, list) op = outs[0].op + boot_args = [] emu_args = [] nd_args = [] for i in args: if isinstance(i, tvm.tensor.Tensor): shape = [tvm_val_2_py_val(j) for j in i.shape] emu_args.append(numpy.random.randn(*shape).astype(i.dtype)) + boot_args.append(emu_args[-1]) nd_args.append(tvm.nd.array(emu_args[-1], ctx)) elif isinstance(i, tvm.expr.Var): emu_args.append(tvm_val_2_py_val(i)) + boot_args.append(emu_args[-1]) nd_args.append(emu_args[-1]) else: assert isinstance(i, list) @@ -43,11 +46,9 @@ def tvm_val_2_py_val(val): target=target) assert module - stmt = tvm.build_module.form_body(sch) true_args = [i for i in args if isinstance(i, (tvm.tensor.Tensor, tvm.expr.Var))] true_outs = [outs] if isinstance(outs, tvm.tensor.Tensor) else outs - src = tvm.hybrid.dump(stmt, true_args, true_outs) - hmodule = tvm.hybrid.Module(src, "hybrid_func") + h_module = tvm.hybrid.build(sch, true_args, true_outs) out_tensors = [] for i in range(op.num_outputs): @@ -60,11 +61,18 @@ def tvm_val_2_py_val(val): if isinstance(ref_data, numpy.ndarray): ref_data = [ref_data] + boot_res = h_module(*boot_args) + if isinstance(boot_res, numpy.ndarray): + boot_res = [boot_res] + module(*nd_args) for nd, np in zip(out_tensors, ref_data): tvm.testing.assert_allclose(nd.asnumpy(), np, rtol=1e-5, atol=1e-5) + for np0, np1 in zip(boot_res, ref_data): + tvm.testing.assert_allclose(np0, np1, rtol=1e-5, atol=1e-5) + @script def outer_product(n, m, a, b): From 11d2077418dcd1c22e2f86d8c964be3722f78894 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 18:28:27 -0800 Subject: [PATCH 20/31] on the way to bootstrap --- python/tvm/hybrid/__init__.py | 19 +---- python/tvm/hybrid/module.py | 8 +- python/tvm/hybrid/parser.py | 36 ++++++++- src/contrib/hybrid/codegen_hybrid.cc | 3 +- tests/python/unittest/test_hybrid_script.py | 83 ++++++++++++--------- 5 files changed, 95 insertions(+), 54 deletions(-) diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index 7e0d864a2e8a..831edbd6bef1 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -17,13 +17,11 @@ from .._ffi.base import decorate from .._ffi.function import _init_api -from .. import _api_internal as _tvm_internal -from ..tensor import Tensor from ..build_module import form_body -from .parser import parse_python -from .util import _pruned_source from .module import HybridModule as Module +from .parser import source_to_op +from .util import _pruned_source def script(pyfunc): @@ -38,21 +36,12 @@ def script(pyfunc): A decorated hybrid script function. """ def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring - from .runtime import _enter_hybrid_runtime, _restore_runtime from .util import _is_tvm_arg_types if _is_tvm_arg_types(args): src = _pruned_source(func) - parser = parse_python(src, func.__globals__, args) - - input_tensors = [] - for i in args: - if isinstance(i, Tensor): - input_tensors.append(i) - op = _tvm_internal._HybridOp(parser.func_name, "HybridOp", None, input_tensors, - parser.outputs, parser.parsed_body) - res = [op.output(i) for i in range(len(parser.outputs))] - return res[0] if len(res) == 1 else res + return source_to_op(src, func.__globals__, args) + from .runtime import _enter_hybrid_runtime, _restore_runtime intersect = _enter_hybrid_runtime(func) value = func(*args, **kwargs) _restore_runtime(func, intersect) diff --git a/python/tvm/hybrid/module.py b/python/tvm/hybrid/module.py index 1969656bf8ec..926ad79e3da5 100644 --- a/python/tvm/hybrid/module.py +++ b/python/tvm/hybrid/module.py @@ -6,6 +6,8 @@ import imp from ..contrib import util from .util import _internal_assert +from .util import _is_tvm_arg_types +from .parser import source_to_op class HybridModule(object): @@ -18,9 +20,9 @@ class HybridModule(object): def __init__(self, src, name): temp = util.tempdir() dst = temp.relpath(name) - self.src_ = 'import tvm\n@tvm.hybrid.script\n%s' % src + self.src_ = src with open(dst, 'w') as f: - f.write(self.src_) + f.write("import tvm\n@tvm.hybrid.script\n%s" % src) py_module = imp.load_source(name, dst) _internal_assert(hasattr(py_module, name), \ "The loaded source has no given function!") @@ -29,6 +31,8 @@ def __init__(self, src, name): def __call__(self, *args): + if _is_tvm_arg_types(args): + return source_to_op(self.src_, globals(), args) return self.func_(*args) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index 9a191d2fbb5e..159bac952cd0 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -17,6 +17,7 @@ from ..api import any as _any from ..container import Array from ..tensor import Tensor, Operation +from .. import _api_internal as _tvm_internal from .. import expr as _expr from .. import stmt as _stmt from .. import make as _make @@ -517,7 +518,7 @@ def parse_python(src, symbols, args): src : str The source code of the function to be parsed. - src : str + symbols : str The symbol list of the global context of the function. args : list of Tensors or Vars @@ -536,3 +537,36 @@ def parse_python(src, symbols, args): parser.parsed_body = parser.visit(root) _internal_assert(parser.returned, 'No valid return found in the function body!') return parser + + +def source_to_op(src, symbols, args): + """Another level of wrapper + + Parameters + ---------- + src : str + The source code of the function to be parsed. + + symbols : str + The symbol list of the global context of the function. + + args : list of Tensors or Vars + The argument lists to the function. + It is NOT encouraged to write a function without arguments. + It is NOT encouraged to write a function with side effect. + + Returns + ------- + res : list of output tensors + The result of output tensors of the formed OpNode. + """ + parser = parse_python(src, symbols, args) + + input_tensors = [] + for i in args: + if isinstance(i, Tensor): + input_tensors.append(i) + op = _tvm_internal._HybridOp(parser.func_name, "HybridOp", None, input_tensors, + parser.outputs, parser.parsed_body) + res = [op.output(i) for i in range(len(parser.outputs))] + return res[0] if len(res) == 1 else res diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index aed5bbe2cfed..70a96294fc06 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -57,7 +57,8 @@ void CodeGenHybrid::VisitExpr_(const UIntImm *op, std::ostream& os) { // NOLINT os << "(" << op->value << ")"; } void CodeGenHybrid::VisitExpr_(const FloatImm *op, std::ostream& os) { // NOLINT(*) - os << std::setprecision(20) << op->value; + PrintType(op->type, os); + os << "(" << std::setprecision(20) << op->value << ")"; } void CodeGenHybrid::VisitExpr_(const StringImm *op, std::ostream& os) { // NOLINT(*) os << "\"" << op->value << "\""; diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index dfb1c6ff454e..07743995e731 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -22,18 +22,15 @@ def tvm_val_2_py_val(val): assert isinstance(outs, list) op = outs[0].op - boot_args = [] emu_args = [] nd_args = [] for i in args: if isinstance(i, tvm.tensor.Tensor): shape = [tvm_val_2_py_val(j) for j in i.shape] emu_args.append(numpy.random.randn(*shape).astype(i.dtype)) - boot_args.append(emu_args[-1]) nd_args.append(tvm.nd.array(emu_args[-1], ctx)) elif isinstance(i, tvm.expr.Var): emu_args.append(tvm_val_2_py_val(i)) - boot_args.append(emu_args[-1]) nd_args.append(emu_args[-1]) else: assert isinstance(i, list) @@ -46,10 +43,6 @@ def tvm_val_2_py_val(val): target=target) assert module - true_args = [i for i in args if isinstance(i, (tvm.tensor.Tensor, tvm.expr.Var))] - true_outs = [outs] if isinstance(outs, tvm.tensor.Tensor) else outs - h_module = tvm.hybrid.build(sch, true_args, true_outs) - out_tensors = [] for i in range(op.num_outputs): output = op.output(i) @@ -61,18 +54,15 @@ def tvm_val_2_py_val(val): if isinstance(ref_data, numpy.ndarray): ref_data = [ref_data] - boot_res = h_module(*boot_args) - if isinstance(boot_res, numpy.ndarray): - boot_res = [boot_res] - module(*nd_args) for nd, np in zip(out_tensors, ref_data): tvm.testing.assert_allclose(nd.asnumpy(), np, rtol=1e-5, atol=1e-5) - for np0, np1 in zip(boot_res, ref_data): - tvm.testing.assert_allclose(np0, np1, rtol=1e-5, atol=1e-5) - + module_args = [i for i in args if isinstance(i, (tvm.tensor.Tensor, tvm.expr.Var))] + module_outs = [outs] if isinstance(outs, tvm.tensor.Tensor) else outs + h_module = tvm.hybrid.build(sch, module_args, module_outs) + return h_module, module_args, module_outs @script def outer_product(n, m, a, b): @@ -130,8 +120,8 @@ def test_outer_product(): assert mul.a.name == 'a' assert mul.b.name == 'b' - - run_and_check(outer_product, [n, m, a, b], {n: 99, m: 101}) + func, ins, outs = run_and_check(outer_product, [n, m, a, b], {n: 99, m: 101}) + run_and_check(func, ins, {n: 99, m: 101}, outs=outs) for key, _ in HYBRID_GLOBALS.items(): assert key not in globals().keys() @@ -216,7 +206,11 @@ def fanout(n, a): assert len(write.value.args) == 1 assert write.value.args[0].value == 0 - run_and_check(fanout, [n, a], {n: 10}) + func, ins, outs = run_and_check(fanout, [n, a], {n: 10}) + print(func.get_source()) + print(ins) + print(outs) + run_and_check(func, ins, {n: 10}, outs=outs) def test_looptype(): @@ -248,7 +242,8 @@ def looptype(a, b, c): assert jloop.for_type == tvm.stmt.For.Vectorized assert kloop.for_type == tvm.stmt.For.Unrolled - run_and_check(looptype, [a, b, c]) + func, ins, outs = run_and_check(looptype, [a, b, c]) + run_and_check(func, ins, outs=outs) def test_if(): @@ -267,7 +262,8 @@ def if_then_else(a): a = tvm.placeholder((10, ), dtype='int32', name='a') - run_and_check(if_then_else, [a]) + func, ins, outs = run_and_check(if_then_else, [a]) + run_and_check(func, ins, outs=outs) @script def if_triple_condition(a): @@ -279,7 +275,8 @@ def if_triple_condition(a): b[i] = a[i] + 1 return b - run_and_check(if_triple_condition, [a]) + func, ins, outs = run_and_check(if_triple_condition, [a]) + run_and_check(func, ins, outs=outs) @script def if_and(a): @@ -291,7 +288,8 @@ def if_and(a): b[i] = a[i] + 1 return b - run_and_check(if_and, [a]) + func, ins, outs = run_and_check(if_and, [a]) + run_and_check(func, ins, outs=outs) def test_bind(): @@ -307,7 +305,8 @@ def vec_add(a, b): a = tvm.placeholder((1000, ), dtype='float32', name='a') b = tvm.placeholder((1000, ), dtype='float32', name='b') - run_and_check(vec_add, [a, b], target='cuda') + func, ins, outs = run_and_check(vec_add, [a, b], target='cuda') + run_and_check(func, ins, outs=outs, target='cuda') @script def raw(a, b): @@ -320,7 +319,8 @@ def raw(a, b): sch = tvm.create_schedule(c.op) x = tvm.thread_axis('threadIdx.x') sch[c].bind(c.op.axis[0], x) - run_and_check(raw, [a, b], sch=sch, outs=[c], target='cuda') + func, ins, outs = run_and_check(raw, [a, b], sch=sch, outs=[c], target='cuda') + run_and_check(func, ins, outs=outs, target='cuda') # Test loop binds @tvm.hybrid.script @@ -337,7 +337,8 @@ def goo(a, b): b = [1, 2, 3, 4, 5] c = goo(a, tvm.convert(b)) sch = tvm.create_schedule(c.op) - run_and_check(goo, [a, b], sch=sch, outs=[c]) + func, ins, outs = run_and_check(goo, [a, b], sch=sch, outs=[c]) + run_and_check(func, ins, outs=outs) def test_math_intrin(): @script @@ -398,7 +399,8 @@ def blur(a): return b a = tvm.placeholder((32, 32), 'float32', 'a') - run_and_check(blur, [a]) + func, ins, outs = run_and_check(blur, [a]) + run_and_check(func, ins, outs=outs) @tvm.hybrid.script def triangle(a, b): @@ -411,7 +413,8 @@ def triangle(a, b): a = tvm.placeholder((10, ), dtype='float32', name='a') b = tvm.placeholder((10, ), dtype='float32', name='b') - run_and_check(triangle, [a, b]) + func, ins, outs = run_and_check(triangle, [a, b]) + run_and_check(func, ins, outs=outs) def test_allocate(): @tvm.hybrid.script @@ -429,7 +432,8 @@ def blur2d(a): a = tvm.placeholder((32, 32), 'float32', 'a') b = blur2d(a) sch = tvm.create_schedule(b.op) - run_and_check(blur2d, [a]) + func, ins, outs = run_and_check(blur2d, [a]) + run_and_check(func, ins, outs=outs) if tvm.gpu().exist: @tvm.hybrid.script @@ -447,7 +451,8 @@ def share_vec_add(a, b): a = tvm.placeholder((256, ), dtype='float32', name='a') b = tvm.placeholder((256, ), dtype='float32', name='b') - run_and_check(share_vec_add, [a, b], target='cuda') + func, ins, outs = run_and_check(share_vec_add, [a, b], target='cuda') + run_and_check(func, ins, outs=outs) else: print('[Warning] No GPU found! Skip shared mem test!') @@ -583,7 +588,8 @@ def foo(a, b): a = tvm.placeholder((10, ), name='a') b = tvm.placeholder((10, ), name='b') - run_and_check(foo, [a, b]) + func, ins, outs = run_and_check(foo, [a, b]) + run_and_check(func, ins, outs=outs) def test_bool(): @tvm.hybrid.script @@ -597,7 +603,8 @@ def foo(a): b[i] = 0.0 return b a = tvm.placeholder((10, ), name='a') - run_and_check(foo, [a]) + func, ins, outs = run_and_check(foo, [a]) + run_and_check(func, ins, outs=outs) def test_const_range(): @tvm.hybrid.script @@ -617,7 +624,8 @@ def foo(a, b): a = tvm.placeholder((2, 5), name='a', dtype='float32') b = [[1, 2, 3, 4, 5], [5, 4, 3, 2, 1]] - run_and_check(foo, [a, b]) + func, ins, outs = run_and_check(foo, [a, b]) + run_and_check(func, ins, outs=outs) @tvm.hybrid.script def goo(a, b): @@ -633,7 +641,8 @@ def goo(a, b): b = [1, 2, 3, 4, 5] c = goo(a, tvm.convert(b)) sch = tvm.create_schedule(c.op) - run_and_check(goo, [a, b]) + func, ins, outs = run_and_check(goo, [a, b]) + run_and_check(func, ins, outs=outs) @tvm.hybrid.script def hoo(a, b): @@ -647,7 +656,8 @@ def hoo(a, b): return c a = tvm.placeholder((5, ), name='a', dtype='int32') b = [1, 2, 3, 4, 5] - run_and_check(hoo, [a, b]) + func, ins, outs = run_and_check(hoo, [a, b]) + run_and_check(func, ins, outs=outs) def test_schedule(): @script @@ -689,7 +699,8 @@ def outer_product(a, b): assert isinstance(ir, tvm.stmt.For) assert ir.loop_var.name == 'j.outer.inner' ir = ir.body - run_and_check(outer_product, [a, b], sch=sch, outs=[c]) + func, ins, outs = run_and_check(outer_product, [a, b], sch=sch, outs=[c]) + run_and_check(func, ins, outs=outs) # Test fuse sch = tvm.create_schedule(c.op) @@ -701,13 +712,15 @@ def outer_product(a, b): ir = ir.body assert isinstance(ir, tvm.stmt.For) assert ir.loop_var.name == 'i.j.fused' - run_and_check(outer_product, [a, b], sch=sch, outs=[c]) + func, ins, outs = run_and_check(outer_product, [a, b], sch=sch, outs=[c]) + run_and_check(func, ins, outs=outs) # Test imperfect loop split sch = tvm.create_schedule(c.op) sch[c].split(c.op.axis[0], 3) ir = tvm.lower(sch, [a, b, c], simple_mode=True) run_and_check(outer_product, [a, b], sch=sch, outs=[c]) + run_and_check(func, ins, outs=outs) # Test loop binds From 3866e1c86b7918927cbf74bb3d6099b824bbd95c Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 19:16:14 -0800 Subject: [PATCH 21/31] cpu bootstrap done --- python/tvm/hybrid/calls.py | 2 +- src/contrib/hybrid/codegen_hybrid.cc | 10 +++++++++- tests/python/unittest/test_hybrid_script.py | 4 +--- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index af796e9acce5..17de1070219f 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -130,4 +130,4 @@ def likely(func_id, args): _internal_assert(args.__len__() == 1, \ "Only one expression can be likely") _internal_assert(func_id == "likely", "This function cannot be directly invoked!") - return call_pure_intrin(args[0].dtype, 'likely', args) + return call_pure_intrin(args[0].dtype, 'likely', *args) diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index 70a96294fc06..83d2c7fdb4ef 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -312,6 +312,14 @@ void CodeGenHybrid::VisitStmt_(const For* op) { indent_ -= tab_; } +bool is_noop(const Stmt &stmt) { + if (!stmt.defined()) + return true; + if (auto eval = stmt.as()) + return is_const(eval->value); + return false; +} + void CodeGenHybrid::VisitStmt_(const IfThenElse* op) { std::string cond = PrintExpr(op->condition); PrintIndent(); @@ -320,7 +328,7 @@ void CodeGenHybrid::VisitStmt_(const IfThenElse* op) { PrintStmt(op->then_case); indent_ -= tab_; - if (op->else_case.defined()) { + if (!is_noop(op->else_case)) { PrintIndent(); stream << "else:\n"; indent_ += tab_; diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 07743995e731..5178a7b9de29 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -62,6 +62,7 @@ def tvm_val_2_py_val(val): module_args = [i for i in args if isinstance(i, (tvm.tensor.Tensor, tvm.expr.Var))] module_outs = [outs] if isinstance(outs, tvm.tensor.Tensor) else outs h_module = tvm.hybrid.build(sch, module_args, module_outs) + return h_module, module_args, module_outs @script @@ -207,9 +208,6 @@ def fanout(n, a): assert write.value.args[0].value == 0 func, ins, outs = run_and_check(fanout, [n, a], {n: 10}) - print(func.get_source()) - print(ins) - print(outs) run_and_check(func, ins, {n: 10}, outs=outs) From 5375aba76b2bd50d86217728b481a11c1b1040da Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 20:57:54 -0800 Subject: [PATCH 22/31] bootstrap! --- cmake/config.cmake | 2 +- docker/install/install_tvm_cpu.sh | 1 + docker/install/install_tvm_gpu.sh | 1 + python/tvm/hybrid/__init__.py | 10 +-- python/tvm/hybrid/module.py | 1 + src/contrib/hybrid/codegen_hybrid.cc | 79 ++++++++++++++++++--- src/contrib/hybrid/codegen_hybrid.h | 8 ++- tests/python/unittest/test_hybrid_script.py | 2 +- 8 files changed, 85 insertions(+), 19 deletions(-) diff --git a/cmake/config.cmake b/cmake/config.cmake index 831dfa2af9e8..c02f569c8dfe 100644 --- a/cmake/config.cmake +++ b/cmake/config.cmake @@ -132,5 +132,5 @@ set(USE_SORT OFF) # Build ANTLR parser for Relay text format set(USE_ANTLR OFF) -# Build ANTLR parser for Relay text format +# Build with hybrid dump set(USE_HYBRID_DUMP ON) diff --git a/docker/install/install_tvm_cpu.sh b/docker/install/install_tvm_cpu.sh index 461ad244d37c..13e8d42a6b97 100644 --- a/docker/install/install_tvm_cpu.sh +++ b/docker/install/install_tvm_cpu.sh @@ -5,6 +5,7 @@ echo set\(USE_LLVM llvm-config-6.0\) >> config.cmake echo set\(USE_RPC ON\) >> config.cmake echo set\(USE_SORT ON\) >> config.cmake echo set\(USE_GRAPH_RUNTIME ON\) >> config.cmake +echo set\(USE_HYBRID_DUMP ON\) >> config.cmake echo set\(USE_BLAS openblas\) >> config.cmake echo set\(USE_SGX /opt/sgxsdk\) >> config.cmake echo set\(RUST_SGX_SDK /opt/rust-sgx-sdk\) >> config.cmake diff --git a/docker/install/install_tvm_gpu.sh b/docker/install/install_tvm_gpu.sh index 8a1324646fd5..8b697a943a92 100644 --- a/docker/install/install_tvm_gpu.sh +++ b/docker/install/install_tvm_gpu.sh @@ -7,6 +7,7 @@ echo set\(USE_CUDNN ON\) >> config.cmake echo set\(USE_RPC ON\) >> config.cmake echo set\(USE_SORT ON\) >> config.cmake echo set\(USE_GRAPH_RUNTIME ON\) >> config.cmake +echo set\(USE_HYBRID_DUMP ON\) >> config.cmake echo set\(USE_BLAS openblas\) >> config.cmake mkdir -p build cd build diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index 831edbd6bef1..bd0534d51024 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -19,7 +19,7 @@ from .._ffi.function import _init_api from ..build_module import form_body -from .module import HybridModule as Module +from . import module from .parser import source_to_op from .util import _pruned_source @@ -50,9 +50,6 @@ def wrapped_func(func, *args, **kwargs): #pylint: disable=missing-docstring return decorate(pyfunc, wrapped_func) -_init_api("tvm.hybrid") - - def build(sch, inputs, outputs, name="hybrid_func"): """Dump the corrent schedule to hybrid module @@ -77,4 +74,7 @@ def build(sch, inputs, outputs, name="hybrid_func"): stmt = form_body(sch) src = dump(stmt, inputs, outputs, name) - return Module(src, name) + return module.HybridModule(src, name) + + +_init_api("tvm.hybrid") diff --git a/python/tvm/hybrid/module.py b/python/tvm/hybrid/module.py index 926ad79e3da5..3847aba4d53a 100644 --- a/python/tvm/hybrid/module.py +++ b/python/tvm/hybrid/module.py @@ -4,6 +4,7 @@ """ import imp + from ..contrib import util from .util import _internal_assert from .util import _is_tvm_arg_types diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index 83d2c7fdb4ef..8bb7849a13b0 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -10,10 +10,14 @@ namespace contrib { using namespace ir; +std::string dot_to_underscore(std::string s) { + for (auto &ch : s) + if (ch == '.') ch = '_'; + return s; +} + std::string CodeGenHybrid::GetUniqueName(std::string prefix) { - for (size_t i = 0; i < prefix.size(); ++i) { - if (prefix[i] == '.') prefix[i] = '_'; - } + prefix = dot_to_underscore(prefix); auto it = ids_allocated_.find(prefix); if (it != ids_allocated_.end()) { while (true) { @@ -61,7 +65,7 @@ void CodeGenHybrid::VisitExpr_(const FloatImm *op, std::ostream& os) { // NOLINT os << "(" << std::setprecision(20) << op->value << ")"; } void CodeGenHybrid::VisitExpr_(const StringImm *op, std::ostream& os) { // NOLINT(*) - os << "\"" << op->value << "\""; + os << "'" << op->value << "'"; } template @@ -251,14 +255,26 @@ void CodeGenHybrid::VisitStmt_(const LetStmt* op) { } void CodeGenHybrid::VisitStmt_(const AttrStmt* op) { - // TODO(@were): Support thread and buffer binding if (op->attr_key == ir::attr::thread_extent) { - LOG(FATAL) << "Thread binding support yet!\n"; + auto iter_var = op->node.as(); + CHECK(iter_var); + binds_[iter_var->var.get()] = dot_to_underscore(iter_var->var->name_hint); + PrintIndent(); + stream << "for " << binds_[iter_var->var.get()] << " in bind('" + << iter_var->var->name_hint << "', "; + PrintExpr(op->value, stream); + stream << "):\n"; + indent_ += tab_; + PrintStmt(op->body); + indent_ -= tab_; } else if (op->attr_key == ir::attr::realize_scope) { auto v = FunctionRef(op->node.node_); alloc_storage_scope_[v] = op->value.as()->value; + PrintStmt(op->body); + } else { + // For now we ignore the unsupported AttrStmt + PrintStmt(op->body); } - PrintStmt(op->body); } void CodeGenHybrid::VisitStmt_(const Realize *op) { @@ -271,9 +287,9 @@ void CodeGenHybrid::VisitStmt_(const Realize *op) { stream << PrintExpr(op->bounds[i]->extent); } if (op->bounds.size() == 1) stream << ", "; - stream << "), \""; + stream << "), '"; PrintType(op->type, stream); - stream << "\", '"; + stream << "', '"; stream << alloc_storage_scope_[op->func] << "')\n"; } PrintStmt(op->body); @@ -358,6 +374,8 @@ void CodeGenHybrid::PrintIndent() { } std::string CodeGenHybrid::GetVarID(const Variable *v) { + if (binds_.count(v)) + return binds_[v]; auto key = std::make_pair(v->GetNodePtr().get(), 0); if (id_map_.count(key)) { return id_map_[key]; @@ -377,10 +395,53 @@ std::string CodeGenHybrid::GetTensorID(const FunctionRef &func, int value_index) return id_map_[key] = GetUniqueName(name_hint); } +void CodeGenHybrid::ReserveKeywords() { + GetUniqueName("def"); + GetUniqueName("for"); + GetUniqueName("in"); + GetUniqueName("range"); + GetUniqueName("unroll"); + GetUniqueName("const_range"); + GetUniqueName("parallel"); + GetUniqueName("vectorize"); + GetUniqueName("bind"); + GetUniqueName("threadIdx.x"); + GetUniqueName("threadIdx.y"); + GetUniqueName("threadIdx.z"); + GetUniqueName("blockIdx.x"); + GetUniqueName("blockIdx.y"); + GetUniqueName("blockIdx.z"); + GetUniqueName("allocate"); + GetUniqueName("output_tensor"); + GetUniqueName("sqrt"); + GetUniqueName("log"); + GetUniqueName("tanh"); + GetUniqueName("power"); + GetUniqueName("exp"); + GetUniqueName("sigmoid"); + GetUniqueName("popcount"); + GetUniqueName("likely"); + GetUniqueName("int8"); + GetUniqueName("int16"); + GetUniqueName("int32"); + GetUniqueName("int64"); + GetUniqueName("uint8"); + GetUniqueName("uint16"); + GetUniqueName("uint32"); + GetUniqueName("uint64"); + GetUniqueName("float16"); + GetUniqueName("float32"); + GetUniqueName("float64"); + GetUniqueName("ceil_div"); +} + void CodeGenHybrid::DumpStmt(const Stmt &stmt, const Array &inputs, const Array &outputs, const std::string &name) { + ReserveKeywords(); + GetUniqueName(name); + stream << "def " << name << "("; for (size_t i = 0; i < inputs.size(); ++i) { if (i) stream << ", "; diff --git a/src/contrib/hybrid/codegen_hybrid.h b/src/contrib/hybrid/codegen_hybrid.h index d45e744755c1..b45b7c43be09 100644 --- a/src/contrib/hybrid/codegen_hybrid.h +++ b/src/contrib/hybrid/codegen_hybrid.h @@ -45,6 +45,8 @@ class CodeGenHybrid : * \return The code. */ std::string Finish(); + /*! \brief Reserve keywords in avoid of name conflict. */ + void ReserveKeywords(); /*! * \brief Print the Stmt n to CodeGenHybrid->stream * \param n The statement to be printed. @@ -127,8 +129,10 @@ class CodeGenHybrid : inline void PrintIndent(); /*! \brief Keys are ids allocated, and values are the suffix to prevent double-name. */ std::map ids_allocated_; - /*! \brief Keys are either tensors or variables. Values are the corresponding IDs.*/ + /*! \brief Keys are either (tensors, value_index) or (variables, 0). Values are the corresponding IDs.*/ std::map, std::string> id_map_; + /*! \brief Variables (keys) binded to the threads (values). */ + std::map binds_; /*! * \brief Find an unallocated name for the given prefix. * \param prefix The given prefix. @@ -149,8 +153,6 @@ class CodeGenHybrid : std::string GetTensorID(const FunctionRef &func, int value_index); /*! \brief the storage scope of allocation */ std::map alloc_storage_scope_; - /*! \brief the data type of allocated buffers */ - std::unordered_map handle_data_type_; }; } // namespace contrib diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 5178a7b9de29..82c27e8948b5 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -450,7 +450,7 @@ def share_vec_add(a, b): a = tvm.placeholder((256, ), dtype='float32', name='a') b = tvm.placeholder((256, ), dtype='float32', name='b') func, ins, outs = run_and_check(share_vec_add, [a, b], target='cuda') - run_and_check(func, ins, outs=outs) + run_and_check(func, ins, outs=outs, target='cuda') else: print('[Warning] No GPU found! Skip shared mem test!') From 363c4a9f3657103580d2c8cfc6a68e7bf667b34b Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 21:08:26 -0800 Subject: [PATCH 23/31] fix lint --- python/tvm/hybrid/calls.py | 5 +++-- python/tvm/hybrid/module.py | 2 +- src/contrib/hybrid/codegen_hybrid.h | 4 +++- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index 17de1070219f..a7c0630cdc85 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -118,10 +118,11 @@ def _cast(func_id, args): def ceil_div(func_id, args): + _internal_assert(func_id == "ceil_div", "This function cannot be directly invoked!") _internal_assert(args.__len__() == 2, \ "Only one expression can be cast") - for i in range(2): - _internal_assert(isinstance(args[i], _expr.Expr), "Only expressions can div") + _internal_assert(isinstance(args[0], _expr.Expr), "Only expressions can div") + _internal_assert(isinstance(args[1], _expr.Expr), "Only expressions can div") a, b = args[0], args[1] return (a + b - 1) / b diff --git a/python/tvm/hybrid/module.py b/python/tvm/hybrid/module.py index 3847aba4d53a..0497e2910879 100644 --- a/python/tvm/hybrid/module.py +++ b/python/tvm/hybrid/module.py @@ -39,7 +39,7 @@ def __call__(self, *args): def get_source(self): return self.src_ - + def save(self, path): if not path.endswith('.py'): diff --git a/src/contrib/hybrid/codegen_hybrid.h b/src/contrib/hybrid/codegen_hybrid.h index b45b7c43be09..cdd6b85b9f9e 100644 --- a/src/contrib/hybrid/codegen_hybrid.h +++ b/src/contrib/hybrid/codegen_hybrid.h @@ -129,7 +129,9 @@ class CodeGenHybrid : inline void PrintIndent(); /*! \brief Keys are ids allocated, and values are the suffix to prevent double-name. */ std::map ids_allocated_; - /*! \brief Keys are either (tensors, value_index) or (variables, 0). Values are the corresponding IDs.*/ + /*! + * \brief Keys are either (tensors, value_index) or (variables, 0). + * Values are the corresponding IDs.*/ std::map, std::string> id_map_; /*! \brief Variables (keys) binded to the threads (values). */ std::map binds_; From d129948aff572c977e414a6120d7bab6a50dfc25 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Mon, 11 Feb 2019 21:15:41 -0800 Subject: [PATCH 24/31] fix doc --- python/tvm/hybrid/module.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/tvm/hybrid/module.py b/python/tvm/hybrid/module.py index 0497e2910879..f31e367f8764 100644 --- a/python/tvm/hybrid/module.py +++ b/python/tvm/hybrid/module.py @@ -1,6 +1,8 @@ """Methods and data structures to support dumping HalideIR to Hybrid Script. This allows users to do quick hack to generated HalideIR and cast it back to TVM modules. + +To enable this feature, you need to build with -DUSE_HYBRID_DUMP=ON. """ import imp From 0063637117e79e0683bfc3b9a40c6aaabb14b1bd Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Tue, 12 Feb 2019 10:10:26 -0800 Subject: [PATCH 25/31] resolve some review concerns --- .gitignore | 1 - python/tvm/hybrid/__init__.py | 2 +- python/tvm/hybrid/calls.py | 3 +-- src/contrib/hybrid/codegen_hybrid.cc | 2 +- 4 files changed, 3 insertions(+), 5 deletions(-) diff --git a/.gitignore b/.gitignore index ffca5ec16b0f..04dad2039860 100644 --- a/.gitignore +++ b/.gitignore @@ -167,7 +167,6 @@ cscope* # vim temporary files *.swp *.swo -.ycm_extra_conf.py # TVM generated code perf diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index bd0534d51024..ea8dbadaa416 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -72,7 +72,7 @@ def build(sch, inputs, outputs, name="hybrid_func"): """ stmt = form_body(sch) - src = dump(stmt, inputs, outputs, name) + src = _Dump(stmt, inputs, outputs, name) return module.HybridModule(src, name) diff --git a/python/tvm/hybrid/calls.py b/python/tvm/hybrid/calls.py index a7c0630cdc85..84ae537d49ab 100644 --- a/python/tvm/hybrid/calls.py +++ b/python/tvm/hybrid/calls.py @@ -119,8 +119,7 @@ def _cast(func_id, args): def ceil_div(func_id, args): _internal_assert(func_id == "ceil_div", "This function cannot be directly invoked!") - _internal_assert(args.__len__() == 2, \ - "Only one expression can be cast") + _internal_assert(args.__len__() == 2, "2 arguments expected for division!") _internal_assert(isinstance(args[0], _expr.Expr), "Only expressions can div") _internal_assert(isinstance(args[1], _expr.Expr), "Only expressions can div") a, b = args[0], args[1] diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index 8bb7849a13b0..f01195783ac0 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -477,7 +477,7 @@ void CodeGenHybrid::DumpStmt(const Stmt &stmt, stream << "\n"; } -TVM_REGISTER_GLOBAL("hybrid.dump") +TVM_REGISTER_GLOBAL("hybrid._Dump") .set_body([](TVMArgs args, TVMRetValue* rv) { CodeGenHybrid codegen; if (args.size() == 4) From 46df5636a2ecf1a7fc20946770bf03f1075bc7ee Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Tue, 12 Feb 2019 10:54:49 -0800 Subject: [PATCH 26/31] support load/save --- python/tvm/hybrid/__init__.py | 4 +- python/tvm/hybrid/module.py | 55 ++++++++++++++++----- python/tvm/hybrid/parser.py | 13 +++-- tests/python/unittest/test_hybrid_script.py | 9 +++- 4 files changed, 60 insertions(+), 21 deletions(-) diff --git a/python/tvm/hybrid/__init__.py b/python/tvm/hybrid/__init__.py index ea8dbadaa416..645ef992833f 100644 --- a/python/tvm/hybrid/__init__.py +++ b/python/tvm/hybrid/__init__.py @@ -19,7 +19,7 @@ from .._ffi.function import _init_api from ..build_module import form_body -from . import module +from .module import HybridModule from .parser import source_to_op from .util import _pruned_source @@ -74,7 +74,7 @@ def build(sch, inputs, outputs, name="hybrid_func"): stmt = form_body(sch) src = _Dump(stmt, inputs, outputs, name) - return module.HybridModule(src, name) + return HybridModule(src, name) _init_api("tvm.hybrid") diff --git a/python/tvm/hybrid/module.py b/python/tvm/hybrid/module.py index f31e367f8764..4199476f99f7 100644 --- a/python/tvm/hybrid/module.py +++ b/python/tvm/hybrid/module.py @@ -5,6 +5,7 @@ To enable this feature, you need to build with -DUSE_HYBRID_DUMP=ON. """ +import ast import imp from ..contrib import util @@ -20,22 +21,22 @@ class HybridModule(object): format for Phase 0 HalideIR. Thus, a totally separated module is defined.""" - def __init__(self, src, name): - temp = util.tempdir() - dst = temp.relpath(name) - self.src_ = src - with open(dst, 'w') as f: - f.write("import tvm\n@tvm.hybrid.script\n%s" % src) - py_module = imp.load_source(name, dst) - _internal_assert(hasattr(py_module, name), \ - "The loaded source has no given function!") - self.func_ = getattr(py_module, name) - _internal_assert(callable(self.func_), "This should be a function! At least callable!") + def __init__(self, src=None, name=None): + self.src_ = self.name = self.func_ = self.root_ = None + if src is not None: + temp = util.tempdir() + dst = temp.relpath("script.py") + with open(dst, 'w') as f: + f.write("import tvm\n@tvm.hybrid.script\n%s" % src) + + if name is not None: + self.name = name + self.load(dst) def __call__(self, *args): if _is_tvm_arg_types(args): - return source_to_op(self.src_, globals(), args) + return source_to_op(self.root_, globals(), args) return self.func_(*args) @@ -48,3 +49,33 @@ def save(self, path): path = path + '.py' with open(path, 'w') as f: f.write(self.src_) + + + def load(self, path): + with open(path, 'r') as f: + self.src_ = f.read() + + src = self.src_ + + class FindFunc(ast.NodeVisitor): + def __init__(self): + self.name = None + self.root = None + + + def visit_FunctionDef(self, node): + _internal_assert(self.name is None, "For now, only one function supported!") + self.name = node.name + _internal_assert(self.root is None, "For now, only one function supported!") + self.root = node + + root = ast.parse(src) + finder = FindFunc() + finder.visit(root) + _internal_assert(finder.name is not None and finder.root is not None, \ + "No function found!") + if self.name is None: + self.name = finder.name + self.root_ = finder.root + py_module = imp.load_source(self.name, path) + self.func_ = getattr(py_module, self.name) diff --git a/python/tvm/hybrid/parser.py b/python/tvm/hybrid/parser.py index 159bac952cd0..b9d64866b305 100644 --- a/python/tvm/hybrid/parser.py +++ b/python/tvm/hybrid/parser.py @@ -515,8 +515,9 @@ def parse_python(src, symbols, args): Parameters ---------- - src : str - The source code of the function to be parsed. + src : ast.node or str + If an ast.node, then directly lower it. + If a str, then parse it to ast and lower it. symbols : str The symbol list of the global context of the function. @@ -531,7 +532,8 @@ def parse_python(src, symbols, args): root : Stmt The result Halide IR and the parser class instance. """ - root = ast.parse(src) + root = ast.parse(src) if isinstance(src, str) else src + _internal_assert(root, ast.AST) var_usage = determine_variable_usage(root, args, symbols) parser = HybridParser(args, var_usage, symbols) parser.parsed_body = parser.visit(root) @@ -544,8 +546,9 @@ def source_to_op(src, symbols, args): Parameters ---------- - src : str - The source code of the function to be parsed. + src : ast.node or str + If an ast.node, then directly lower it. + If a str, then parse it to ast and lower it. symbols : str The symbol list of the global context of the function. diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index 82c27e8948b5..fc5d33d59182 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -1,4 +1,5 @@ -import tvm, inspect, sys, traceback, numpy, nose, types +import tvm, inspect, sys, traceback, numpy, nose, types, os +from tvm.contrib import util from tvm.hybrid import script from tvm.hybrid.runtime import HYBRID_GLOBALS @@ -122,6 +123,10 @@ def test_outer_product(): assert mul.b.name == 'b' func, ins, outs = run_and_check(outer_product, [n, m, a, b], {n: 99, m: 101}) + temp = util.tempdir() + path = temp.relpath('%s.py' % func.name) + func.save(path) + func_ = tvm.hybrid.HybridModule().load(path) run_and_check(func, ins, {n: 99, m: 101}, outs=outs) for key, _ in HYBRID_GLOBALS.items(): @@ -717,7 +722,7 @@ def outer_product(a, b): sch = tvm.create_schedule(c.op) sch[c].split(c.op.axis[0], 3) ir = tvm.lower(sch, [a, b, c], simple_mode=True) - run_and_check(outer_product, [a, b], sch=sch, outs=[c]) + func, ins, outs = run_and_check(outer_product, [a, b], sch=sch, outs=[c]) run_and_check(func, ins, outs=outs) # Test loop binds From 3d495f3feb7cdb817c75af6af9a6f147bfd2fd36 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Tue, 12 Feb 2019 11:06:41 -0800 Subject: [PATCH 27/31] fix lint --- python/tvm/hybrid/module.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/python/tvm/hybrid/module.py b/python/tvm/hybrid/module.py index 4199476f99f7..01557ba8b179 100644 --- a/python/tvm/hybrid/module.py +++ b/python/tvm/hybrid/module.py @@ -22,13 +22,23 @@ class HybridModule(object): def __init__(self, src=None, name=None): + """The constructor of this a hybrid module + + Parameters + ---------- + src : str + The source code of this module + + name : str + The name of this module + """ self.src_ = self.name = self.func_ = self.root_ = None if src is not None: temp = util.tempdir() dst = temp.relpath("script.py") with open(dst, 'w') as f: f.write("import tvm\n@tvm.hybrid.script\n%s" % src) - + if name is not None: self.name = name self.load(dst) @@ -52,12 +62,21 @@ def save(self, path): def load(self, path): + """Load the module from a python file + + Parameters + ---------- + path : str + Path to the given python file + """ with open(path, 'r') as f: self.src_ = f.read() src = self.src_ class FindFunc(ast.NodeVisitor): + """ Find the function in module to be loaded module. """ + #pylint: disable=invalid-name def __init__(self): self.name = None self.root = None From 3d086cbf04805a4e64211f2f40c0c4c6e480b23d Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Tue, 12 Feb 2019 17:54:50 -0800 Subject: [PATCH 28/31] thanks to xqdan fixed my typo --- tests/python/unittest/test_hybrid_script.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/python/unittest/test_hybrid_script.py b/tests/python/unittest/test_hybrid_script.py index fc5d33d59182..405577b05b3b 100644 --- a/tests/python/unittest/test_hybrid_script.py +++ b/tests/python/unittest/test_hybrid_script.py @@ -126,8 +126,9 @@ def test_outer_product(): temp = util.tempdir() path = temp.relpath('%s.py' % func.name) func.save(path) - func_ = tvm.hybrid.HybridModule().load(path) - run_and_check(func, ins, {n: 99, m: 101}, outs=outs) + func_ = tvm.hybrid.HybridModule() + func_.load(path) + run_and_check(func_, ins, {n: 99, m: 101}, outs=outs) for key, _ in HYBRID_GLOBALS.items(): assert key not in globals().keys() From fcc761d4bcca0365d5b32a19b2999a0747f1e83c Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 13 Feb 2019 10:56:11 -0800 Subject: [PATCH 29/31] fix build, make dump non-optional --- cmake/config.cmake | 3 --- cmake/modules/contrib/HybridDump.cmake | 8 +++----- docker/install/install_tvm_cpu.sh | 1 - docker/install/install_tvm_gpu.sh | 1 - 4 files changed, 3 insertions(+), 10 deletions(-) diff --git a/cmake/config.cmake b/cmake/config.cmake index c02f569c8dfe..a97def410ddd 100644 --- a/cmake/config.cmake +++ b/cmake/config.cmake @@ -131,6 +131,3 @@ set(USE_SORT OFF) # Build ANTLR parser for Relay text format set(USE_ANTLR OFF) - -# Build with hybrid dump -set(USE_HYBRID_DUMP ON) diff --git a/cmake/modules/contrib/HybridDump.cmake b/cmake/modules/contrib/HybridDump.cmake index a4c8f626a64a..c8d6d6e07756 100644 --- a/cmake/modules/contrib/HybridDump.cmake +++ b/cmake/modules/contrib/HybridDump.cmake @@ -1,5 +1,3 @@ -if(USE_HYBRID_DUMP) - message(STATUS "Build with contrib.hybriddump") - file(GLOB HYBRID_CONTRIB_SRC src/contrib/hybrid/*.cc) - list(APPEND COMPILER_SRCS ${HYBRID_CONTRIB_SRC}) -endif(USE_HYBRID_DUMP) +message(STATUS "Build with contrib.hybriddump") +file(GLOB HYBRID_CONTRIB_SRC src/contrib/hybrid/*.cc) +list(APPEND COMPILER_SRCS ${HYBRID_CONTRIB_SRC}) diff --git a/docker/install/install_tvm_cpu.sh b/docker/install/install_tvm_cpu.sh index 13e8d42a6b97..461ad244d37c 100644 --- a/docker/install/install_tvm_cpu.sh +++ b/docker/install/install_tvm_cpu.sh @@ -5,7 +5,6 @@ echo set\(USE_LLVM llvm-config-6.0\) >> config.cmake echo set\(USE_RPC ON\) >> config.cmake echo set\(USE_SORT ON\) >> config.cmake echo set\(USE_GRAPH_RUNTIME ON\) >> config.cmake -echo set\(USE_HYBRID_DUMP ON\) >> config.cmake echo set\(USE_BLAS openblas\) >> config.cmake echo set\(USE_SGX /opt/sgxsdk\) >> config.cmake echo set\(RUST_SGX_SDK /opt/rust-sgx-sdk\) >> config.cmake diff --git a/docker/install/install_tvm_gpu.sh b/docker/install/install_tvm_gpu.sh index 8b697a943a92..8a1324646fd5 100644 --- a/docker/install/install_tvm_gpu.sh +++ b/docker/install/install_tvm_gpu.sh @@ -7,7 +7,6 @@ echo set\(USE_CUDNN ON\) >> config.cmake echo set\(USE_RPC ON\) >> config.cmake echo set\(USE_SORT ON\) >> config.cmake echo set\(USE_GRAPH_RUNTIME ON\) >> config.cmake -echo set\(USE_HYBRID_DUMP ON\) >> config.cmake echo set\(USE_BLAS openblas\) >> config.cmake mkdir -p build cd build From 17c4a857af36efe5e7aad7ffb4ec2234c8d45c3e Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 13 Feb 2019 10:56:58 -0800 Subject: [PATCH 30/31] add vthread --- .ycm_extra_conf.py | 211 +++++++++++++++++++++++++++ src/contrib/hybrid/codegen_hybrid.cc | 1 + 2 files changed, 212 insertions(+) create mode 100644 .ycm_extra_conf.py diff --git a/.ycm_extra_conf.py b/.ycm_extra_conf.py new file mode 100644 index 000000000000..a1bf19e7d616 --- /dev/null +++ b/.ycm_extra_conf.py @@ -0,0 +1,211 @@ +# This file is NOT licensed under the GPLv3, which is the license for the rest +# of YouCompleteMe. +# +# Here's the license text for this file: +# +# This is free and unencumbered software released into the public domain. +# +# Anyone is free to copy, modify, publish, use, compile, sell, or +# distribute this software, either in source code form or as a compiled +# binary, for any purpose, commercial or non-commercial, and by any +# means. +# +# In jurisdictions that recognize copyright laws, the author or authors +# of this software dedicate any and all copyright interest in the +# software to the public domain. We make this dedication for the benefit +# of the public at large and to the detriment of our heirs and +# successors. We intend this dedication to be an overt act of +# relinquishment in perpetuity of all present and future rights to this +# software under copyright law. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +# IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR +# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, +# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR +# OTHER DEALINGS IN THE SOFTWARE. +# +# For more information, please refer to + +from distutils.sysconfig import get_python_inc +import platform +import os.path as p +import subprocess +import ycm_core + +DIR_OF_THIS_SCRIPT = p.abspath( p.dirname( __file__ ) ) +DIR_OF_THIRD_PARTY = p.join( DIR_OF_THIS_SCRIPT, 'third_party' ) +SOURCE_EXTENSIONS = [ '.cpp', '.cxx', '.cc', '.c', '.m', '.mm' ] + +# These are the compilation flags that will be used in case there's no +# compilation database set (by default, one is not set). +# CHANGE THIS LIST OF FLAGS. YES, THIS IS THE DROID YOU HAVE BEEN LOOKING FOR. +flags = [ +'-Wall', +'-Wextra', +'-Werror', +'-Wno-long-long', +'-Wno-variadic-macros', +'-fexceptions', +'-DNDEBUG', +# You 100% do NOT need -DUSE_CLANG_COMPLETER and/or -DYCM_EXPORT in your flags; +# only the YCM source code needs it. +'-DUSE_CLANG_COMPLETER', +'-DYCM_EXPORT=', +# THIS IS IMPORTANT! Without the '-x' flag, Clang won't know which language to +# use when compiling headers. So it will guess. Badly. So C++ headers will be +# compiled as C headers. You don't want that so ALWAYS specify the '-x' flag. +# For a C project, you would set this to 'c' instead of 'c++'. +'-x', +'c++', +'-isystem', +'cpp/pybind11', +'-isystem', +'cpp/BoostParts', +'-isystem', +get_python_inc(), +'-isystem', +'cpp/llvm/include', +'-isystem', +'cpp/llvm/tools/clang/include', +'-I', +'cpp/ycm', +'-I', +'cpp/ycm/ClangCompleter', +'-isystem', +'cpp/ycm/tests/gmock/gtest', +'-isystem', +'cpp/ycm/tests/gmock/gtest/include', +'-isystem', +'cpp/ycm/tests/gmock', +'-isystem', +'cpp/ycm/tests/gmock/include', +'-isystem', +'cpp/ycm/benchmarks/benchmark/include', +] + +# Clang automatically sets the '-std=' flag to 'c++14' for MSVC 2015 or later, +# which is required for compiling the standard library, and to 'c++11' for older +# versions. +if platform.system() != 'Windows': + flags.append( '-std=c++11' ) + + +# Set this to the absolute path to the folder (NOT the file!) containing the +# compile_commands.json file to use that instead of 'flags'. See here for +# more details: http://clang.llvm.org/docs/JSONCompilationDatabase.html +# +# You can get CMake to generate this file for you by adding: +# set( CMAKE_EXPORT_COMPILE_COMMANDS 1 ) +# to your CMakeLists.txt file. +# +# Most projects will NOT need to set this to anything; you can just change the +# 'flags' list of compilation flags. Notice that YCM itself uses that approach. +compilation_database_folder = '' + +database = ycm_core.CompilationDatabase( DIR_OF_THIS_SCRIPT + "/build/" ) + + +def IsHeaderFile( filename ): + extension = p.splitext( filename )[ 1 ] + return extension in [ '.h', '.hxx', '.hpp', '.hh' ] + + +def FindCorrespondingSourceFile( filename ): + if IsHeaderFile( filename ): + basename = p.splitext( filename )[ 0 ] + for extension in SOURCE_EXTENSIONS: + replacement_file = basename + extension + if p.exists( replacement_file ): + return replacement_file + return filename + + +def PathToPythonUsedDuringBuild(): + try: + filepath = p.join( DIR_OF_THIS_SCRIPT, 'PYTHON_USED_DURING_BUILDING' ) + with open( filepath ) as f: + return f.read().strip() + # We need to check for IOError for Python 2 and OSError for Python 3. + except ( IOError, OSError ): + return None + + +def Settings( **kwargs ): + language = kwargs[ 'language' ] + + if language == 'cfamily': + # If the file is a header, try to find the corresponding source file and + # retrieve its flags from the compilation database if using one. This is + # necessary since compilation databases don't have entries for header files. + # In addition, use this source file as the translation unit. This makes it + # possible to jump from a declaration in the header file to its definition + # in the corresponding source file. + filename = FindCorrespondingSourceFile( kwargs[ 'filename' ] ) + + if not database: + return { + 'flags': flags, + 'include_paths_relative_to_dir': DIR_OF_THIS_SCRIPT, + 'override_filename': filename + } + + compilation_info = database.GetCompilationInfoForFile( filename ) + if not compilation_info.compiler_flags_: + return {} + + # Bear in mind that compilation_info.compiler_flags_ does NOT return a + # python list, but a "list-like" StringVec object. + final_flags = list( compilation_info.compiler_flags_ ) + + # NOTE: This is just for YouCompleteMe; it's highly likely that your project + # does NOT need to remove the stdlib flag. DO NOT USE THIS IN YOUR + # ycm_extra_conf IF YOU'RE NOT 100% SURE YOU NEED IT. + try: + final_flags.remove( '-stdlib=libc++' ) + except ValueError: + pass + + return { + 'flags': final_flags, + 'include_paths_relative_to_dir': compilation_info.compiler_working_dir_, + 'override_filename': filename + } + + if language == 'python': + return { + 'interpreter_path': PathToPythonUsedDuringBuild() + } + + return {} + + +def GetStandardLibraryIndexInSysPath( sys_path ): + for index, path in enumerate( sys_path ): + if p.isfile( p.join( path, 'os.py' ) ): + return index + raise RuntimeError( 'Could not find standard library path in Python path.' ) + + +def PythonSysPath( **kwargs ): + sys_path = kwargs[ 'sys_path' ] + + interpreter_path = kwargs[ 'interpreter_path' ] + major_version = subprocess.check_output( [ + interpreter_path, '-c', 'import sys; print( sys.version_info[ 0 ] )' ] + ).rstrip().decode( 'utf8' ) + + sys_path.insert( GetStandardLibraryIndexInSysPath( sys_path ) + 1, + p.join( DIR_OF_THIRD_PARTY, 'python-future', 'src' ) ) + sys_path[ 0:0 ] = [ p.join( DIR_OF_THIS_SCRIPT ), + p.join( DIR_OF_THIRD_PARTY, 'bottle' ), + p.join( DIR_OF_THIRD_PARTY, 'cregex', + 'regex_{}'.format( major_version ) ), + p.join( DIR_OF_THIRD_PARTY, 'frozendict' ), + p.join( DIR_OF_THIRD_PARTY, 'jedi' ), + p.join( DIR_OF_THIRD_PARTY, 'parso' ), + p.join( DIR_OF_THIRD_PARTY, 'requests' ), + p.join( DIR_OF_THIRD_PARTY, 'waitress' ) ] + + return sys_path diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc index f01195783ac0..2117d471eeee 100644 --- a/src/contrib/hybrid/codegen_hybrid.cc +++ b/src/contrib/hybrid/codegen_hybrid.cc @@ -411,6 +411,7 @@ void CodeGenHybrid::ReserveKeywords() { GetUniqueName("blockIdx.x"); GetUniqueName("blockIdx.y"); GetUniqueName("blockIdx.z"); + GetUniqueName("vthread"); GetUniqueName("allocate"); GetUniqueName("output_tensor"); GetUniqueName("sqrt"); From f2b00f1611bcfecae375d0e69279493030a58866 Mon Sep 17 00:00:00 2001 From: Jian Weng Date: Wed, 13 Feb 2019 11:00:37 -0800 Subject: [PATCH 31/31] jesus why i added this --- .ycm_extra_conf.py | 211 --------------------------------------------- 1 file changed, 211 deletions(-) delete mode 100644 .ycm_extra_conf.py diff --git a/.ycm_extra_conf.py b/.ycm_extra_conf.py deleted file mode 100644 index a1bf19e7d616..000000000000 --- a/.ycm_extra_conf.py +++ /dev/null @@ -1,211 +0,0 @@ -# This file is NOT licensed under the GPLv3, which is the license for the rest -# of YouCompleteMe. -# -# Here's the license text for this file: -# -# This is free and unencumbered software released into the public domain. -# -# Anyone is free to copy, modify, publish, use, compile, sell, or -# distribute this software, either in source code form or as a compiled -# binary, for any purpose, commercial or non-commercial, and by any -# means. -# -# In jurisdictions that recognize copyright laws, the author or authors -# of this software dedicate any and all copyright interest in the -# software to the public domain. We make this dedication for the benefit -# of the public at large and to the detriment of our heirs and -# successors. We intend this dedication to be an overt act of -# relinquishment in perpetuity of all present and future rights to this -# software under copyright law. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, -# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF -# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. -# IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR -# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, -# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR -# OTHER DEALINGS IN THE SOFTWARE. -# -# For more information, please refer to - -from distutils.sysconfig import get_python_inc -import platform -import os.path as p -import subprocess -import ycm_core - -DIR_OF_THIS_SCRIPT = p.abspath( p.dirname( __file__ ) ) -DIR_OF_THIRD_PARTY = p.join( DIR_OF_THIS_SCRIPT, 'third_party' ) -SOURCE_EXTENSIONS = [ '.cpp', '.cxx', '.cc', '.c', '.m', '.mm' ] - -# These are the compilation flags that will be used in case there's no -# compilation database set (by default, one is not set). -# CHANGE THIS LIST OF FLAGS. YES, THIS IS THE DROID YOU HAVE BEEN LOOKING FOR. -flags = [ -'-Wall', -'-Wextra', -'-Werror', -'-Wno-long-long', -'-Wno-variadic-macros', -'-fexceptions', -'-DNDEBUG', -# You 100% do NOT need -DUSE_CLANG_COMPLETER and/or -DYCM_EXPORT in your flags; -# only the YCM source code needs it. -'-DUSE_CLANG_COMPLETER', -'-DYCM_EXPORT=', -# THIS IS IMPORTANT! Without the '-x' flag, Clang won't know which language to -# use when compiling headers. So it will guess. Badly. So C++ headers will be -# compiled as C headers. You don't want that so ALWAYS specify the '-x' flag. -# For a C project, you would set this to 'c' instead of 'c++'. -'-x', -'c++', -'-isystem', -'cpp/pybind11', -'-isystem', -'cpp/BoostParts', -'-isystem', -get_python_inc(), -'-isystem', -'cpp/llvm/include', -'-isystem', -'cpp/llvm/tools/clang/include', -'-I', -'cpp/ycm', -'-I', -'cpp/ycm/ClangCompleter', -'-isystem', -'cpp/ycm/tests/gmock/gtest', -'-isystem', -'cpp/ycm/tests/gmock/gtest/include', -'-isystem', -'cpp/ycm/tests/gmock', -'-isystem', -'cpp/ycm/tests/gmock/include', -'-isystem', -'cpp/ycm/benchmarks/benchmark/include', -] - -# Clang automatically sets the '-std=' flag to 'c++14' for MSVC 2015 or later, -# which is required for compiling the standard library, and to 'c++11' for older -# versions. -if platform.system() != 'Windows': - flags.append( '-std=c++11' ) - - -# Set this to the absolute path to the folder (NOT the file!) containing the -# compile_commands.json file to use that instead of 'flags'. See here for -# more details: http://clang.llvm.org/docs/JSONCompilationDatabase.html -# -# You can get CMake to generate this file for you by adding: -# set( CMAKE_EXPORT_COMPILE_COMMANDS 1 ) -# to your CMakeLists.txt file. -# -# Most projects will NOT need to set this to anything; you can just change the -# 'flags' list of compilation flags. Notice that YCM itself uses that approach. -compilation_database_folder = '' - -database = ycm_core.CompilationDatabase( DIR_OF_THIS_SCRIPT + "/build/" ) - - -def IsHeaderFile( filename ): - extension = p.splitext( filename )[ 1 ] - return extension in [ '.h', '.hxx', '.hpp', '.hh' ] - - -def FindCorrespondingSourceFile( filename ): - if IsHeaderFile( filename ): - basename = p.splitext( filename )[ 0 ] - for extension in SOURCE_EXTENSIONS: - replacement_file = basename + extension - if p.exists( replacement_file ): - return replacement_file - return filename - - -def PathToPythonUsedDuringBuild(): - try: - filepath = p.join( DIR_OF_THIS_SCRIPT, 'PYTHON_USED_DURING_BUILDING' ) - with open( filepath ) as f: - return f.read().strip() - # We need to check for IOError for Python 2 and OSError for Python 3. - except ( IOError, OSError ): - return None - - -def Settings( **kwargs ): - language = kwargs[ 'language' ] - - if language == 'cfamily': - # If the file is a header, try to find the corresponding source file and - # retrieve its flags from the compilation database if using one. This is - # necessary since compilation databases don't have entries for header files. - # In addition, use this source file as the translation unit. This makes it - # possible to jump from a declaration in the header file to its definition - # in the corresponding source file. - filename = FindCorrespondingSourceFile( kwargs[ 'filename' ] ) - - if not database: - return { - 'flags': flags, - 'include_paths_relative_to_dir': DIR_OF_THIS_SCRIPT, - 'override_filename': filename - } - - compilation_info = database.GetCompilationInfoForFile( filename ) - if not compilation_info.compiler_flags_: - return {} - - # Bear in mind that compilation_info.compiler_flags_ does NOT return a - # python list, but a "list-like" StringVec object. - final_flags = list( compilation_info.compiler_flags_ ) - - # NOTE: This is just for YouCompleteMe; it's highly likely that your project - # does NOT need to remove the stdlib flag. DO NOT USE THIS IN YOUR - # ycm_extra_conf IF YOU'RE NOT 100% SURE YOU NEED IT. - try: - final_flags.remove( '-stdlib=libc++' ) - except ValueError: - pass - - return { - 'flags': final_flags, - 'include_paths_relative_to_dir': compilation_info.compiler_working_dir_, - 'override_filename': filename - } - - if language == 'python': - return { - 'interpreter_path': PathToPythonUsedDuringBuild() - } - - return {} - - -def GetStandardLibraryIndexInSysPath( sys_path ): - for index, path in enumerate( sys_path ): - if p.isfile( p.join( path, 'os.py' ) ): - return index - raise RuntimeError( 'Could not find standard library path in Python path.' ) - - -def PythonSysPath( **kwargs ): - sys_path = kwargs[ 'sys_path' ] - - interpreter_path = kwargs[ 'interpreter_path' ] - major_version = subprocess.check_output( [ - interpreter_path, '-c', 'import sys; print( sys.version_info[ 0 ] )' ] - ).rstrip().decode( 'utf8' ) - - sys_path.insert( GetStandardLibraryIndexInSysPath( sys_path ) + 1, - p.join( DIR_OF_THIRD_PARTY, 'python-future', 'src' ) ) - sys_path[ 0:0 ] = [ p.join( DIR_OF_THIS_SCRIPT ), - p.join( DIR_OF_THIRD_PARTY, 'bottle' ), - p.join( DIR_OF_THIRD_PARTY, 'cregex', - 'regex_{}'.format( major_version ) ), - p.join( DIR_OF_THIRD_PARTY, 'frozendict' ), - p.join( DIR_OF_THIRD_PARTY, 'jedi' ), - p.join( DIR_OF_THIRD_PARTY, 'parso' ), - p.join( DIR_OF_THIRD_PARTY, 'requests' ), - p.join( DIR_OF_THIRD_PARTY, 'waitress' ) ] - - return sys_path