[CODEGEN/EXEC] CUDA, NVRTC pipeline complete

.travis.yml

@@ -4,7 +4,7 @@ language: cpp
 
 os:
  - linux
- # - osx
+ - osx
 
 env:
  # code analysis

Makefile

@@ -15,7 +15,7 @@ all: lib/libtvm.a lib/libtvm.so
 
 LIB_HALIDE_IR = HalideIR/lib/libHalideIR.a
 
-SRC = $(wildcard src/*.cc src/*/*.cc)
+SRC = $(wildcard src/*.cc src/*/*.cc src/*/*/*.cc)
 ALL_OBJ = $(patsubst src/%.cc, build/%.o, $(SRC))
 ALL_DEP = $(ALL_OBJ) $(LIB_HALIDE_IR)
 
@@ -39,7 +39,7 @@ endif
 
 ifeq ($(USE_CUDA), 1)
  CFLAGS += -DTVM_CUDA_RUNTIME=1
- LDFLAGS += -lcuda -lcudart
+ LDFLAGS += -lcuda -lcudart -lnvrtc
 else
  CFLAGS += -DTVM_CUDA_RUNTIME=0
 endif
@@ -92,3 +92,4 @@ clean:
 
 -include build/*.d
 -include build/*/*.d
+-include build/*/*/*.d

include/tvm/base.h

@@ -12,25 +12,9 @@
 #include <string>
 #include <memory>
 #include <functional>
-#include <typeinfo>
-#include <type_traits>
 
 namespace tvm {
 
-/*!
- *\brief whether to use CUDA runtime
- */
-#ifndef TVM_CUDA_RUNTIME
-#define TVM_CUDA_RUNTIME 1
-#endif
-
-/*!
- *\brief whether to use opencl runtime
- */
-#ifndef TVM_OPENCL_RUNTIME
-#define TVM_OPENCL_RUNTIME 0
-#endif
-
 using ::tvm::Node;
 using ::tvm::NodeRef;
 using ::tvm::AttrVisitor;

include/tvm/codegen.h

@@ -9,13 +9,18 @@
 #include <string>
 #include "./base.h"
 #include "./expr.h"
-#include "./module.h"
+#include "./lowered_func.h"
 #include "./runtime/packed_func.h"
 
 
 namespace tvm {
 /*! \brief namespace for lowlevel IR pass and codegen */
 namespace codegen {
+// use packed function from runtime.
+using runtime::PackedFunc;
+using runtime::TVMArgs;
+using runtime::TVMRetValue;
+
 /*!
  * \brief Make an user callable API LoweredFunc.
  *
@@ -64,8 +69,35 @@ Array<Var> UndefinedVars(const LoweredFunc& f);
  */
 Array<LoweredFunc> SplitHostDevice(LoweredFunc func);
 
+/*!
+ * \brief Build a stack VM function.
+ * \param func The LoweredFunc to be build
+ * \param device_funcs The additional device functions
+ * \return A packed function representing the func.
+ */
+PackedFunc BuildStackVM(
+ LoweredFunc func,
+ const std::unordered_map<LoweredFunc, PackedFunc>& device_funcs);
+
+/*!
+ * \brief Build a CUDA function with NVRTC
+ *
+ * \param fsplits The LoweredFuncs to be build (after SplitHostDevice)
+ * The first element is the host function, followed by device functions.
+ * \param host_mode The host side compilation mode:
+ * - "stackvm": use stack vm to interpret host side code.
+ */
+PackedFunc BuildNVRTC(Array<LoweredFunc> fsplits, std::string host_mode);
 
-runtime::PackedFunc BuildStackVM(LoweredFunc func);
+/*!
+ * \brief Build a OpenCL function.
+ *
+ * \param fsplits The LoweredFuncs to be build (after SplitHostDevice)
+ * The first element is the host function, followed by device functions.
+ * \param host_mode The host side compilation mode:
+ * - "stackvm": use stack vm to interpret host side code.
+ */
+PackedFunc BuildOpenCL(Array<LoweredFunc> fsplits, std::string host_mode);
 
 } // namespace codegen
 } // namespace tvm

include/tvm/expr.h

@@ -12,7 +12,7 @@
 #include <string>
 #include <algorithm>
 #include "./base.h"
-#include "./runtime/packed_func.h"
+#include "./runtime/c_runtime_api.h"
 
 namespace tvm {
 
@@ -33,6 +33,19 @@ using Halide::Internal::Variable;
 
 using Halide::Internal::make_const;
 
+
+inline Type TVMType2Type(TVMType t) {
+ return Type(static_cast<halide_type_code_t>(t.code), t.bits, t.lanes);
+}
+
+inline TVMType Type2TVMType(Type t) {
+ TVMType ret;
+ ret.code = static_cast<uint8_t>(t.code());
+ ret.bits = static_cast<uint8_t>(t.bits());
+ ret.lanes = static_cast<uint16_t>(t.lanes());
+ return ret;
+}
+
 /*! \brief a named variable in TVM */
 class Var : public Halide::VarExpr {
  public:

include/tvm/module.h → include/tvm/lowered_func.h

@@ -1,11 +1,11 @@
 /*!
- * Copyright (c) 2016 by Contributors
- * \file module.h
- * \brief Low level IR module,
- * Contains lowered function information.
+ * Copyright (c) 2017 by Contributors
+ * \file lowered_func.h
+ * \brief Information about a lowered TVM function.
+ * This data structure is final step toward codegen.
  */
-#ifndef TVM_MODULE_H_
-#define TVM_MODULE_H_
+#ifndef TVM_LOWERED_FUNC_H_
+#define TVM_LOWERED_FUNC_H_
 
 #include <tvm/container.h>
 #include <ir/FunctionBase.h>
@@ -102,4 +102,13 @@ inline const LoweredFuncNode* LoweredFunc::operator->() const {
 
 } // namespace tvm
 
-#endif // TVM_MODULE_H_
+namespace std {
+template <>
+struct hash<::tvm::LoweredFunc> {
+ std::size_t operator()(const ::tvm::LoweredFunc& k) const {
+ return k.hash();
+ }
+};
+}
+
+#endif // TVM_LOWERED_FUNC_H_

include/tvm/packed_func_ext.h

@@ -14,6 +14,7 @@
 
 #include "./base.h"
 #include "./expr.h"
+#include "./runtime/packed_func.h"
 
 namespace tvm {
 using runtime::TVMArgs;
@@ -162,19 +163,7 @@ inline void TVMArgsSetter::operator()(size_t i, NodeRef& other) const { // NOLI
  type_codes_[i] = kNodeHandle;
 }
 
-// Type related stuffs
-inline Type TVMType2Type(TVMType t) {
- return Type(static_cast<halide_type_code_t>(t.code), t.bits, t.lanes);
-}
-
-inline TVMType Type2TVMType(Type t) {
- TVMType ret;
- ret.code = static_cast<uint8_t>(t.code());
- ret.bits = static_cast<uint8_t>(t.bits());
- ret.lanes = static_cast<uint16_t>(t.lanes());
- return ret;
-}
-
+// type related stuffs
 inline TVMRetValue& TVMRetValue::operator=(const Halide::Type& t) {
  return this->operator=(Type2TVMType(t));
 }

include/tvm/runtime/config.h

@@ -0,0 +1,23 @@
+/*!
+ * Copyright (c) 2017 by Contributors
+ * \file config.h
+ * \brief Runtime library related configurations.
+ */
+#ifndef TVM_RUNTIME_CONFIG_H_
+#define TVM_RUNTIME_CONFIG_H_
+
+/*!
+ *\brief whether to use CUDA runtime
+ */
+#ifndef TVM_CUDA_RUNTIME
+#define TVM_CUDA_RUNTIME 1
+#endif
+
+/*!
+ *\brief whether to use opencl runtime
+ */
+#ifndef TVM_OPENCL_RUNTIME
+#define TVM_OPENCL_RUNTIME 0
+#endif
+
+#endif // TVM_RUNTIME_CONFIG_H_

include/tvm/runtime/packed_func.h

@@ -163,6 +163,13 @@ inline const char* TypeCode2Str(int type_code);
  */
 inline TVMType String2TVMType(std::string s);
 
+/*!
+ * \brief convert a TVM type to string.
+ * \param t The type to be converted.
+ * \return The corresponding tvm type in string.
+ */
+inline std::string TVMType2String(TVMType t);
+
 // macro to check type code.
 #define TVM_CHECK_TYPE_CODE(CODE, T) \
  CHECK_EQ(CODE, T) << " expected " \
@@ -258,6 +265,9 @@ class TVMArgValue : public TVMPODValue_ {
  using TVMPODValue_::operator TVMArray*;
  // conversion operator.
  operator std::string() const {
+ if (type_code_ == kTVMType) {
+ return TVMType2String(operator TVMType());
+ }
  TVM_CHECK_TYPE_CODE(type_code_, kStr);
  return std::string(value_.v_str);
  }
@@ -308,7 +318,6 @@ class TVMRetValue : public TVMPODValue_ {
  */
  TVMRetValue(TVMRetValue&& other)
  : TVMPODValue_(other.value_, other.type_code_) {
- other.type_code_ = kNull;
  }
  /*! \brief destructor */
  ~TVMRetValue() {
@@ -328,6 +337,9 @@ class TVMRetValue : public TVMPODValue_ {
  }
  // conversion operators
  operator std::string() const {
+ if (type_code_ == kTVMType) {
+ return TVMType2String(operator TVMType());
+ }
  TVM_CHECK_TYPE_CODE(type_code_, kStr);
  return *ptr<std::string>();
  }
@@ -418,6 +430,13 @@ class TVMRetValue : public TVMPODValue_ {
  *ret_type_code = type_code_;
  type_code_ = kNull;
  }
+ /*! \return The value field, if the data is POD */
+ const TVMValue& value() const {
+ CHECK(type_code_ != kNodeHandle &&
+ type_code_ != kFuncHandle &&
+ type_code_ != kStr) << "TVMRetValue.value can only be used for POD data";
+ return value_;
+ }
  // NodeRef related extenstions: in tvm/packed_func_ext.h
  inline TVMRetValue& operator=(const NodeRef& other);
  inline TVMRetValue& operator=(const std::shared_ptr<Node>& other);
@@ -488,7 +507,7 @@ inline const char* TypeCode2Str(int type_code) {
  case kInt: return "int";
  case kFloat: return "float";
  case kStr: return "str";
- case kHandle: return "Handle";
+ case kHandle: return "handle";
  case kNull: return "NULL";
  case kNodeHandle: return "NodeHandle";
  case kArrayHandle: return "ArrayHandle";
@@ -499,6 +518,21 @@ inline const char* TypeCode2Str(int type_code) {
  }
 }
 
+inline std::ostream& operator<<(std::ostream& os, TVMType t) { // NOLINT(*)
+ os << TypeCode2Str(t.code)
+ << static_cast<int>(t.bits);
+ if (t.lanes != 1) {
+ os << 'x' << static_cast<int>(t.lanes);
+ }
+ return os;
+}
+
+inline std::string TVMType2String(TVMType t) {
+ std::ostringstream os;
+ os << t;
+ return os.str();
+}
+
 inline TVMType String2TVMType(std::string s) {
  TVMType t;
  t.bits = 32; t.lanes = 1;

python/tvm/__init__.py

@@ -13,7 +13,7 @@
 from . import schedule
 
 from . import ndarray as nd
-from .ndarray import cpu, gpu, opencl, init_opencl
+from .ndarray import cpu, gpu, opencl, init_opencl, cl
 
 from ._base import TVMError
 from .api import *

python/tvm/_ctypes/_function.py

@@ -90,8 +90,8 @@ def _make_tvm_args(args, temp_args):
  values[i].v_float64 = arg
  type_codes[i] = TypeCode.FLOAT
  elif isinstance(arg, TVMType):
- values[i].v_type = arg
- type_codes[i] = TypeCode.TVM_TYPE
+ values[i].v_str = c_str(str(arg))
+ type_codes[i] = TypeCode.STR
  elif isinstance(arg, string_types):
  values[i].v_str = c_str(arg)
  type_codes[i] = TypeCode.STR

python/tvm/_ctypes/_types.py

@@ -86,8 +86,7 @@ class TVMValue(ctypes.Union):
  _fields_ = [("v_int64", ctypes.c_int64),
  ("v_float64", ctypes.c_double),
  ("v_handle", ctypes.c_void_p),
- ("v_str", ctypes.c_char_p),
- ("v_type", TVMType)]
+ ("v_str", ctypes.c_char_p)]
 
 
 TVMPackedCFunc = ctypes.CFUNCTYPE(
@@ -117,7 +116,6 @@ def _return_handle(x):
  TypeCode.FLOAT: lambda x: x.v_float64,
  TypeCode.HANDLE: _return_handle,
  TypeCode.NULL: lambda x: None,
- TypeCode.TVM_TYPE: lambda x: x.v_type,
  TypeCode.STR: lambda x: py_str(x.v_str)
 }
 
@@ -127,6 +125,5 @@ def _return_handle(x):
  TypeCode.FLOAT: lambda x: x.v_float64,
  TypeCode.HANDLE: _return_handle,
  TypeCode.NULL: lambda x: None,
- TypeCode.TVM_TYPE: lambda x: x.v_type,
  TypeCode.STR: lambda x: py_str(x.v_str)
 }

python/tvm/api.py

@@ -16,9 +16,9 @@
 from . import expr as _expr
 from . import collections as _collections
 
-int32 = TVMType("int32")
-float32 = TVMType("float32")
-handle = TVMType("handle")
+int32 = "int32"
+float32 = "float32"
+handle = "handle"
 
 def const(value, dtype=None):
  """construct a constant"""

python/tvm/collections.py

@@ -9,6 +9,12 @@
 class Array(NodeBase):
  """Array container of TVM"""
  def __getitem__(self, i):
+ if isinstance(i, slice):
+ start = i.start if i.start is not None else 0
+ stop = i.stop if i.stop is not None else len(self)
+ step = i.step if i.step is not None else 1
+ return [self[idx] for idx in range(start, stop, step)]
+
  if i >= len(self):
  raise IndexError("array index out ot range")
  return _api_internal._ArrayGetItem(self, i)