[BUFFER/REFACTOR] Buffer byte_offset-> elem_offset, add buffer_bind_s…

…cope (apache#209)

include/tvm/buffer.h

@@ -16,10 +16,11 @@ namespace tvm {
 
 // Internal node container Buffer
 class BufferNode;
+
 /*!
  * \brief Buffer is a symbolic n-darray structure.
  *  It is a composition of primitive symbolic types,
- *  used to specify input/output strcuture of the program.
+ *  used to specify the memory layout of the Tensor used in program input.
  */
 class Buffer : public NodeRef {
  public:
@@ -38,6 +39,21 @@ class Buffer : public NodeRef {
    * \return The load expression.
    */
   Stmt MakeStore(Array<Expr> index, Expr value) const;
+  /*!
+   * \brief Return a new buffer that is equivalent with current one
+   *  but always add stride field.
+   * \return The strided version of the buffer.
+   */
+  Buffer MakeStrideView() const;
+  /*!
+   * \brief Make a new symbolic buffer representing a slice of the buffer.
+   * \param begins The beginning position of each dimension.
+   * \param extents The extent of each dimension.
+   * \note This function will make target buffer as compact as possible.
+   *  If stride is not needed in the slice, it won't be presented
+   * \return the result buffer.
+   */
+  Buffer MakeSlice(Array<Expr> begins, Array<Expr> extents) const;
   /*!
    * \brief access the internal node container
    * \return the pointer to the internal node container
@@ -63,17 +79,14 @@ class BufferNode : public Node {
    *  This can be an empty array, indicating array is contiguous
    */
   Array<Expr> strides;
-  /*!
-   * \brief The offset in bytes to the beginning pointer to data
-   *  Can be undefined, indicating this must be zero.
-   */
-  Expr byte_offset;
+  /*! \brief The offset in terms of number of dtype elements (including lanes) */
+  Expr elem_offset;
   // Meta data
   /*! \brief optional name of the buffer */
   std::string name;
   /*! \brief storage scope of the buffer, if other than global */
   std::string scope;
-  /*! \brief Alignment bytes size of byte_offset */
+  /*! \brief Alignment multiple in terms of dtype elements (including lanes) */
   int offset_alignment;
   /*! \brief constructor */
   BufferNode() {}
@@ -83,7 +96,7 @@ class BufferNode : public Node {
     v->Visit("dtype", &dtype);
     v->Visit("shape", &shape);
     v->Visit("strides", &strides);
-    v->Visit("byte_offset", &byte_offset);
+    v->Visit("elem_offset", &elem_offset);
     v->Visit("name", &name);
     v->Visit("scope", &scope);
     v->Visit("offset_alignment", &offset_alignment);

include/tvm/expr.h

@@ -61,6 +61,14 @@ inline TVMType Type2TVMType(Type t) {
   return ret;
 }
 
+// Get number of bytes considering vector type.
+inline int GetVectorBytes(Type dtype) {
+  int data_bits = dtype.bits() * dtype.lanes();
+  CHECK_EQ(data_bits % 8, 0U)
+      << "Need to load/store by multiple of bytes";
+  return data_bits / 8;
+}
+
 /*! \brief a named variable in TVM */
 class Var : public Halide::VarExpr {
  public:

include/tvm/ir.h

@@ -167,8 +167,16 @@ constexpr const char* prefetch_scope = "prefetch_scope";
 constexpr const char* scan_update_scope = "scan_update_scope";
 /*! \brief Mark of scan init scope */
 constexpr const char* scan_init_scope = "scan_init_scope";
-/*! \brief extern operator scope */
-constexpr const char* extern_op_scope = "extern_op_scope";
+/*!
+ * \brief Bind the buffer specification to the region of the op
+ *  When this scope occurs, the stmt.node is a Array<NodeRef> = [buffer, tensor]
+ *  stmt.value is a tvm_tuple(min0, extent0, min1, extent1, ...).
+ *  The scope represents that we need to bind the storage region of tensor to buffer.
+ *  This will affect replacement of some variables inside the scope that
+ *  corresponds to field of buffer to be the actual expressions of tensor during
+ *  storage flattening phase.
+ */
+constexpr const char* buffer_bind_scope = "buffer_bind_scope";
 // Pipeline related attributes
 /*! \brief channel read scope */
 constexpr const char* channel_read_scope = "channel_read_scope";
@@ -194,6 +202,14 @@ namespace intrinsic {
  *  }
  */
 constexpr const char* tvm_address_of = "tvm_address_of";
+/*!
+ * \brief tvm_tuple is not an actual function and cannot codegen.
+ *  It is used to represent tuple structure in value field of AttrStmt,
+ *  for the sake of giving hint to optimization.
+ *
+ *  Handle tvm_tuple(value0, value1, ..., value_n);
+ */
+constexpr const char* tvm_tuple = "tvm_tuple";
 /*!
  * \brief See pesudo code
  *
@@ -250,14 +266,14 @@ constexpr const char* tvm_stack_make_shape = "tvm_stack_make_shape";
  *                            Expr strides,
  *                            Expr ndim,
  *                            Expr dtype,
- *                            Expr byte_offset) {
+ *                            Expr elem_offset) {
  *     ret = alloca stack DLTensor();
  *     ret->data = data;
  *     ret->shape = shape;
  *     ret->strides = strides != 0 ? strides : nullptr;
  *     ret->ndim = ndim;
  *     ret->dtype = dtype.type();
- *     ret->byte_offset = byte_offset;
+ *     ret->byte_offset = elem_offset * sizeof(dtype);
  *     return ret;
  *  }
  */

include/tvm/operation.h

@@ -62,7 +62,7 @@ class OperationNode : public FunctionBaseNode {
   virtual Array<Expr> output_shape(size_t i) const = 0;
   /*!
    * \brief List all the input Tensors.
-   * \return List if input tensors.
+   * \return List of input tensors.
    */
   virtual Array<Tensor> InputTensors() const = 0;
   /*!

python/tvm/api.py

@@ -287,6 +287,7 @@ def scan(init, update, state_placeholder, inputs=None, name="scan"):
     res = [op.output(i) for i in range(len(update))]
     return res[0] if len(res) == 1 else res
 
+
 def extern(shape, inputs, fcompute,
            name="extern", dtype=None):
     """Compute several tensor via extern function.
@@ -374,7 +375,7 @@ def decl_buffer(shape,
                 name="buffer",
                 data=None,
                 strides=None,
-                byte_offset=None,
+                elem_offset=None,
                 scope="",
                 offset_alignment=0):
     """Decleare a new symbolic buffer.
@@ -401,8 +402,9 @@ def decl_buffer(shape,
     strides: array of Expr
         The stride of the buffer.
 
-    byte_offset: Expr, optional
-        The offset in bytes to data pointer.
+    elem_offset: Expr, optional
+        The beginning offset of the array to data.
+        In terms of number of elements of dtype.
 
     scope: str, optional
         The storage scope of the buffer, if not global.
@@ -423,7 +425,7 @@ def decl_buffer(shape,
     to create function that only handles specific case of data structure
     and make compiled function benefit from it.
 
-    If user pass strides and byte_offset is passed as None
+    If user pass strides and elem_offset is passed as None
     when constructing the function, then the function will be specialized
     for the DLTensor that is compact and aligned.
     If user pass a fully generic symbolic array to the strides,
@@ -436,7 +438,7 @@ def decl_buffer(shape,
         data = var(name, "handle")
 
     return _api_internal._Buffer(
-        data, dtype, shape, strides, byte_offset, name, scope, offset_alignment)
+        data, dtype, shape, strides, elem_offset, name, scope, offset_alignment)
 
 
 def _IterVar(dom, name, iter_type, thread_tag=''):
@@ -464,11 +466,11 @@ def _IterVar(dom, name, iter_type, thread_tag=''):
     if dom is not None:
         if isinstance(dom, (list, tuple)):
             if len(dom) != 2:
-                raise ValueError("need to list of ranges")
+                raise TypeError("need to be list of ranges")
             dom = Range(dom[0], dom[1])
 
         if not isinstance(dom, _collections.Range):
-            raise ValueError("dom need to be Range")
+            raise TypeError("dom need to be Range")
     name = name if name else 'iter'
     v = var(name)
     return _api_internal._IterVar(dom, v, iter_type, thread_tag)

python/tvm/collections.py

@@ -26,8 +26,6 @@ def __getitem__(self, i):
     def __len__(self):
         return _api_internal._ArraySize(self)
 
-    def __repr__(self):
-        return '[' + (','.join(str(x) for x in self)) + ']'
 
 @register_node
 class Map(NodeBase):
@@ -52,9 +50,6 @@ def items(self):
     def __len__(self):
         return _api_internal._MapSize(self)
 
-    def __repr__(self):
-        return '{' + (", ".join(str(x[0]) + ": " +str(x[1]) for x in self.items())) + '}'
-
 
 @register_node
 class Range(NodeBase):

python/tvm/expr.py

@@ -236,6 +236,10 @@ class Ramp(Expr):
 class Broadcast(Expr):
     pass
 
+@register_node
+class Shuffle(Expr):
+    pass
+
 @register_node
 class Call(Expr):
     Extern = 0

python/tvm/intrin.py

@@ -1,4 +1,4 @@
-"""Intrinsics and math functions in TVM."""
+"""Expression Intrinsics and math functions in TVM."""
 from __future__ import absolute_import as _abs
 
 from ._ffi.function import register_func as _register_func
@@ -20,7 +20,7 @@ def _pack_buffer(buf):
                  strides,
                  len(buf.shape),
                  const(0, dtype=buf.dtype),
-                 buf.byte_offset]
+                 buf.elem_offset]
     return _make.Call("handle", "tvm_stack_make_array",
                       pack_args, _Call.Intrinsic, None, 0)
 

python/tvm/stmt.py

@@ -73,3 +73,7 @@ class IfThenElse(Stmt):
 @register_node
 class Evaluate(Stmt):
     pass
+
+@register_node
+class Prefetch(Stmt):
+    pass

python/tvm/tensor.py

@@ -118,6 +118,12 @@ def num_outputs(self):
         """Number of outputs of this op."""
         return _api_internal._OpNumOutputs(self)
 
+    @property
+    def input_tensors(self):
+        """List of input tensors to this op."""
+        return _api_internal._OpInputTensors(self)
+
+
 @register_node
 class PlaceholderOp(Operation):
     """Placeholder operation."""

src/api/api_lang.cc

@@ -218,6 +218,11 @@ TVM_REGISTER_API("_OpNumOutputs")
     *ret = args[0].operator Operation()->num_outputs();
   });
 
+TVM_REGISTER_API("_OpInputTensors")
+.set_body([](TVMArgs args,  TVMRetValue* ret) {
+    *ret = args[0].operator Operation()->InputTensors();
+  });
+
 TVM_REGISTER_API("_IterVar")
 .set_body([](TVMArgs args,  TVMRetValue* ret) {
     *ret = IterVarNode::make(

src/contrib/cblas/cblas.cc

@@ -40,10 +40,13 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cblas.matmul")
                 transa ? A->shape[1] : A->shape[0],
                 transa ? B->shape[1] : B->shape[0],
                 1.0f,
-                static_cast<float*>(B->data), B->shape[1],
-                static_cast<float*>(A->data), A->shape[1],
+                reinterpret_cast<float*>(static_cast<char*>(B->data) + B->byte_offset),
+                B->shape[1],
+                reinterpret_cast<float*>(static_cast<char*>(A->data) + A->byte_offset),
+                A->shape[1],
                 0.0f,
-                static_cast<float*>(C->data), C->shape[1]);
+                reinterpret_cast<float*>(static_cast<char*>(C->data) + C->byte_offset),
+                C->shape[1]);
   });
 }  // namespace contrib
 }  // namespace tvm