[TensorIR] Fix tensorir after rebase

Author: merrymercy

python/tvm/tensorir/tensor_intrin.py

@@ -0,0 +1,15 @@
+
+from .._ffi.node import NodeBase
+from .. import make as _make, api as _api, _api_internal
+from .tree_node import register_tensorir_node
+
+@register_tensorir_node
+class TensorIntrinsic(NodeBase):
+    def __init__(self, op, intrin_func, name):
+        self.__init_handle_by_constructor__(_make._TensorIntrinsic, op, intrin_func, name)
+
+    def __call__(self, inputs, outputs):
+        return _api_internal._TensorIntrinsic_Instantiate(self, inputs, outputs)
+
+def decl_tensor_intrin(op, intrin_func, name):
+    return _make._TensorIntrinsic(op, intrin_func, name)

python/tvm/tensorir/tensor_region.py

@@ -0,0 +1,53 @@
+from .._ffi.node import NodeBase
+from .. import make as _make, api as _api, intrin as _intrin, expr as _expr, ir_pass as _ir_pass, \
+    _api_internal
+from ..build_module import current_build_config
+
+from .tree_node import register_tensorir_node
+
+@register_tensorir_node
+class TensorRegion(NodeBase):
+    def __init__(self, tensor_slice):
+        ranges = []
+        for x in tensor_slice.indices:
+            assert x.step is None or x.step == 1, "Only support step = 1"
+            ranges.append(_api.Range(x.start, x.stop))
+        self.__init_handle_by_constructor__(_make.TensorRegion,
+                                            tensor_slice.tensor,
+                                            ranges)
+
+    def __getitem__(self, item):
+        assert len(item) == self.data.ndim, "The dimension of index is wrong"
+
+        mins = []
+        extents = []
+        for x in item:
+            mins.append(x.start)
+            extents.append((x.stop - x.start))
+            assert x.step is None or x.step == 1, "Only support step == 1"
+
+        return _api_internal._TensorRegion_MakeView(self, mins, extents)
+
+    def emit_buffer_bind(self, ib, **kwargs):
+        """Emit buffer_bind_scope Attr Stmt to an IRBuilder"""
+        data = self.data
+
+        # skip ones   todo(lmzheng) : fix this to match inputs placeholder
+        shape = [_ir_pass.Simplify(x.extent) for x in self.ranges]
+        while isinstance(shape[0], _expr.IntImm) and shape[0].value == 1 and len(shape) > 1:
+            shape = shape[1:]
+
+        shape_dtype = shape[0].dtype if hasattr(shape[0], "dtype") else "int32"
+        tmp = _api.const(1, shape_dtype)
+
+        cfg = current_build_config()
+        B = _api.decl_buffer(shape, dtype=data.dtype, name="B" + data.name,
+                             elem_offset=_api.var(data.name + "_offset", dtype=shape_dtype),
+                             **kwargs)
+        ranges = []
+        for x in self.ranges:
+            ranges.append(x.min)
+            ranges.append(x.extent)
+        ib.scope_attr([B, data], "buffer_bind_scope", _intrin.call_intrin("handle", "tvm_tuple", *ranges))
+
+        return B

python/tvm/tensorir/tree_node.py

@@ -0,0 +1,48 @@
+from .. import make as _make, expr as _expr
+from ..api import _init_api
+from .._ffi.node import NodeBase, register_node
+
+def register_tensorir_node(type_key=None):
+    """Register a Relay node type.
+
+    Parameters
+    ----------
+    type_key : str or cls
+        The type key of the node.
+    """
+    if not isinstance(type_key, str):
+        return register_node(
+            "tensorir." + type_key.__name__)(type_key)
+    return register_node(type_key)
+
+
+@register_tensorir_node
+class ScheduleTreeNode(NodeBase):
+    def __str__(self):
+        return PrintTreeNode(self)
+
+
+@register_tensorir_node
+class AxisTreeNode(ScheduleTreeNode):
+    def __init__(self, loop_var, min, extent, axis_type, children):
+        self.__init_handle_by_constructor__(_make.AxisTreeNode,
+                                            loop_var,
+                                            min,
+                                            extent,
+                                            axis_type,
+                                            children)
+
+
+@register_tensorir_node
+class BlockTreeNode(ScheduleTreeNode):
+    def __init__(self, args, vars, inputs, outputs, stmt, children):
+        stmt = _make.Evaluate(stmt) if isinstance(stmt, _expr.Expr) else stmt
+        self.__init_handle_by_constructor__(_make.BlockTreeNode,
+                                            args,
+                                            vars,
+                                            inputs,
+                                            outputs,
+                                            stmt,
+                                            children)
+
+_init_api('tvm.tensorir.tree_node')

src/tensorir/intrinsic.cc

@@ -0,0 +1,89 @@
+/*!
+ *  Copyright (c) 2019 by Contributors
+ *  \brief Tensor intrinsics
+ */
+
+#include <tvm/packed_func_ext.h>
+#include "tree_node.h"
+#include "tree_builder.h"
+#include "intrinsic.h"
+#include "util.h"
+
+namespace tvm {
+namespace tensorir {
+
+// maker
+TensorIntrinsic TensorIntrinsicNode::make(Operation op, PackedFunc intrin_func, std::string name) {
+  NodePtr<TensorIntrinsicNode> node = make_node<TensorIntrinsicNode>();
+
+  node->op = std::move(op);
+  node->intrin_func = std::move(intrin_func);
+  node->name = std::move(name);
+
+  // todo (lmzheng): build BlockTreeNode `from`  for checking and untensorize
+
+  return TensorIntrinsic(node);
+}
+
+const TensorIntrinsicNode* TensorIntrinsic::operator->() const {
+  return static_cast<const TensorIntrinsicNode*>(node_.get());
+}
+
+ScheduleTreeNode TensorIntrinsic::Instantiate(Array<TensorRegion> inputs,
+                                              Array<TensorRegion> outputs) const {
+  NodeRef ret = operator->()->intrin_func(inputs, outputs);
+
+  if (ret->derived_from<ScheduleTreeNodeNode>()) {
+    return Downcast<ScheduleTreeNode>(ret);
+  } else if (ret->derived_from<StmtNode>()) {
+    Stmt stmt = Downcast<Stmt>(ret);
+    Array<Expr> args;
+    Array<Var> vars;
+    Set<Var> used_vars;
+    Map<Var, Expr> var_map;
+    arith::Analyzer analyzer;
+
+    // gather vars
+    for (const auto& x : inputs) {
+      for (const auto& ran: x->ranges) {
+        used_vars.insert(GatherVars(ran->min));
+        used_vars.insert(GatherVars(ran->extent));
+      }
+    }
+    for (const auto& x : outputs) {
+      for (const auto& ran: x->ranges) {
+        used_vars.insert(GatherVars(ran->min));
+        used_vars.insert(GatherVars(ran->extent));
+      }
+    }
+
+    // canonicalize outputs
+    Array<TensorRegion> new_outputs;
+    std::tie(args, vars, new_outputs, var_map) = CreateOutputRegions(
+        outputs, used_vars, &analyzer);
+
+    // replace inputs
+    Array<TensorRegion> new_inputs;
+    for (const auto& x : inputs) {
+      Array<Range> ranges;
+      for (const auto& ran: x->ranges) {
+        ranges.push_back(Range::make_by_min_extent(
+            SubstituteAndEquationSimplify(ran->min, var_map, &analyzer),
+            SubstituteAndEquationSimplify(ran->extent, var_map, &analyzer)));
+      }
+      new_inputs.push_back(TensorRegionNode::make(x->data, ranges));
+    }
+
+    // replace stmt
+    stmt = SubstituteAndEquationSimplify(stmt, var_map, &analyzer);
+
+    return BlockTreeNodeNode::make(args, vars, new_inputs, new_outputs,
+                                   stmt, Array<ScheduleTreeNode>{});
+  } else {
+    LOG(FATAL) << "The intrin func returns invalid value";
+  }
+  return ScheduleTreeNode(nullptr);
+}
+
+} // namespace tensorir
+} // namespace tvm

src/tensorir/intrinsic.h

@@ -0,0 +1,53 @@
+/*!
+ *  Copyright (c) 2019 by Contributors
+ *  \brief Tensor intrinsics
+ */
+
+#ifndef TVM_TENSORIR_INTRINSIC_H_
+#define TVM_TENSORIR_INTRINSIC_H_
+
+#include "tree_node.h"
+
+namespace tvm {
+namespace tensorir {
+
+using runtime::PackedFunc;
+
+// A tensor intrinsic replaces a block to another block
+class TensorIntrinsic;
+class TensorIntrinsicNode : public Node {
+ public:
+  Operation op;             // semantic form
+  PackedFunc intrin_func;   // (Array<TensorRegion>, Array<TensorRegion>) -> ScheduleTreeNode or Stmt,
+                            // todo(lmzheng): use TypedPackedFunc？
+  std::string name;
+
+  void VisitAttrs(AttrVisitor *v) final {
+    v->Visit("op", &op);
+    //v->Visit("intrin_func", &intrin_func);  // todo(lmzheng): fix AttrVisitor
+    v->Visit("name", &name);
+  }
+
+  TVM_DLL static TensorIntrinsic make(Operation op, PackedFunc intrin_func, std::string name);
+
+  static constexpr const char *_type_key = "tensorir.TensorIntrinsic";
+  TVM_DECLARE_NODE_TYPE_INFO(TensorIntrinsicNode, Node);
+};
+
+class TensorIntrinsic : public NodeRef {
+ public:
+  TensorIntrinsic() {}
+  explicit TensorIntrinsic(NodePtr<Node> n): NodeRef(n) {}
+
+  const TensorIntrinsicNode* operator->() const;
+  ScheduleTreeNode Instantiate(Array<TensorRegion> inputs, Array<TensorRegion> outputs) const;
+
+  using ContainerType = TensorIntrinsicNode;
+};
+
+
+} // namespace tensorir
+} // namespace tvm
+
+
+#endif // TVM_TENSORIR_INTRINSIC_H_

src/tensorir/schedule.cc

@@ -13,6 +13,7 @@
 #include "schedule.h"
 #include "tree_builder.h"
 #include "util.h"
+#include "../arithmetic/int_set.h"
 
 namespace tvm {
 
@@ -388,26 +389,26 @@ BlockTreeNode Schedule::compute_at(BlockTreeNode block, AxisTreeNode axis) {
       new_args[i] = block->args[i];
     }
 
-    if (const arith::IntervalSet* set = iter_domain[i].as<arith::IntervalSet>()) {
+    if (const arith::IntervalSetNode* set = iter_domain[i].as<arith::IntervalSetNode>()) {
       node = AxisTreeNodeNode::make(iter_var,
-                                    set->i.min,
-                                    set->i.max - set->i.min + 1,
+                                    set->min_value,
+                                    set->max_value - set->min_value + 1,
                                     kOpaque, // todo(lmzheng): fill correct type to replace kOpaque x 3
                                     Array<ScheduleTreeNode>{last});
       new_args[i] = iter_var;
-    } else if (const arith::StrideSet* set = iter_domain[i].as<arith::StrideSet>()) {
+    } else if (const arith::StrideSetNode* set = iter_domain[i].as<arith::StrideSetNode>()) {
       CHECK(set->extents.size() == 1);
-      CHECK(set->base.is_single_point());
+      CHECK(is_one(set->base_extent));
       if (is_one(set->extents[0])) {
         node = AxisTreeNode(nullptr);
-        new_args[i] = set->base.min;
+        new_args[i] = set->base_min;
       } else {
         node = AxisTreeNodeNode::make(iter_var,
                                       0,
                                       set->extents[0],
                                       kOpaque,
                                       Array<ScheduleTreeNode>{last});
-        new_args[i] = iter_var * set->strides[0] + set->base.min;
+        new_args[i] = iter_var * set->strides[0] + set->base_min;
       }
     } else {
       LOG(FATAL) << "Cannot handle int set : " << iter_domain[i];
@@ -539,11 +540,11 @@ BlockTreeNode Schedule::blockize(AxisTreeNode axis) {
     for (size_t i = 0; i < iter.first.ndim(); ++i) {
       Array<IntSet> to_merge;
       for (const std::vector<IntSet>& y : iter.second) {
-        const arith::IntervalSet* set = y[i].as<arith::IntervalSet>();
+        const arith::IntervalSetNode* set = y[i].as<arith::IntervalSetNode>();
         CHECK(set != nullptr);
-        IntSet b = arith::IntervalSet::make(
-            SubstituteAndEquationSimplify(set->i.min, var_map, &analyzer),
-            SubstituteAndEquationSimplify(set->i.max, var_map, &analyzer));
+        IntSet b = arith::IntervalSet(
+            SubstituteAndEquationSimplify(set->min_value, var_map, &analyzer),
+            SubstituteAndEquationSimplify(set->max_value, var_map, &analyzer));
         to_merge.push_back(b);
       }
       IntSet merged = arith::Union(to_merge);
@@ -600,7 +601,7 @@ BlockTreeNode Schedule::tensorize(BlockTreeNode block, TensorIntrinsic intrin) {
                                         block->inputs, block->outputs,
                                         Stmt(NodePtr<Node>(nullptr)),
                                         Array<ScheduleTreeNode>{Downcast<ScheduleTreeNode>(ret)});
-  } else if (ret->derived_from<HalideIR::Internal::BaseStmtNode>()) {
+  } else if (ret->derived_from<StmtNode>()) {
     new_block = BlockTreeNodeNode::make(block->args, block->vars,
                                         block->inputs, block->outputs,
                                         Downcast<Stmt>(ret), Array<ScheduleTreeNode>{});