[LANG] Generalize compute to tensor region

include/tvm/operation.h

@@ -186,12 +186,14 @@ class TensorComputeOpNode : public OperationNode {
  public:
   Array<IterVar> axis;
 
-  Array<IterVar> out_axis;
+  // Array<IterVar> out_axis;
 
-  Array<IterVar> tensor_axis;
+  // Array<IterVar> tensor_axis;
 
   Array<IterVar> reduce_axis;
 
+  int sch_ndim;
+
   Array<Tensor> inputs;
 
   Array<Region> input_regions;
@@ -231,23 +233,18 @@ class TensorComputeOpNode : public OperationNode {
     v->Visit("name", &name);
     v->Visit("tag", &tag);
     v->Visit("axis", &axis);
-    v->Visit("out_axis", &out_axis);
-    v->Visit("tensor_axis", &tensor_axis);
     v->Visit("reduce_axis", &reduce_axis);
+    v->Visit("sch_ndim", &sch_ndim);
     v->Visit("inputs", &inputs);
+    v->Visit("input_regions", &input_regions);
+    v->Visit("intrin", &intrin);
   }
 
   static Operation make(std::string name,
                         std::string tag,
-                        Array<IterVar> out_axis,
-                        Array<IterVar> tensor_axis,
-                        TensorIntrinCall intrin_call);
-
-  static Operation make(std::string name,
-                        std::string tag,
-                        Array<IterVar> out_axis,
-                        Array<IterVar> tensor_axis,
+                        Array<IterVar> axis,
                         Array<IterVar> reduce_axis,
+                        int sch_ndim,
                         Array<Tensor> tensors,
                         Array<Region> regions,
                         TensorIntrin intrin);

python/tvm/api.py

@@ -262,15 +262,17 @@ def compute(shape, fcompute, name="compute", tag="", attrs=None):
     body = fcompute(*[v.var for v in dim_var])
 
     if isinstance(body, _tensor.TensorIntrinCall):
-        tensor_var = []
         for i, s in enumerate(shape[out_ndim:]):
             var_name = "ax" + str(i)
-            tensor_var.append(_IterVar((0, s), var_name, 4))
+            dim_var.append(_IterVar((0, s), var_name, 4))
         op_node = _api_internal._TensorComputeOp(name,
                                                  tag,
                                                  dim_var,
-                                                 tensor_var,
-                                                 body)
+                                                 body.reduce_axis,
+                                                 out_ndim,
+                                                 body.tensors,
+                                                 body.regions,
+                                                 body.intrin)
     else:
         if not isinstance(body, (list, tuple)):
             body = [body]

src/api/api_lang.cc

@@ -292,7 +292,10 @@ TVM_REGISTER_API("_TensorComputeOp")
                                      args[1],
                                      args[2],
                                      args[3],
-                                     args[4]);
+                                     args[4],
+                                     args[5],
+                                     args[6],
+                                     args[7]);
   });
 
 TVM_REGISTER_API("_ExternOp")

src/op/tensor_compute_op.cc

@@ -29,10 +29,7 @@ int TensorComputeOpNode::num_outputs() const {
 }
 
 Array<IterVar> TensorComputeOpNode::root_iter_vars() const {
-  Array<IterVar> ret = out_axis;
-  for (IterVar iv : tensor_axis) {
-    ret.push_back(iv);
-  }
+  Array<IterVar> ret = axis;
   for (IterVar iv : reduce_axis) {
     ret.push_back(iv);
   }
@@ -45,57 +42,45 @@ Type TensorComputeOpNode::output_dtype(size_t i) const {
 
 Array<Expr> TensorComputeOpNode::output_shape(size_t i) const {
   Array<Expr> shape;
-  for (const auto& ivar : this->out_axis) {
+  for (const auto& ivar : this->axis) {
     shape.push_back(ivar->dom->extent);
   }
-  size_t index = this->inputs.size() + i;
-  for (const auto& dim : this->intrin->buffers[index]->shape) {
-    shape.push_back(dim);
-  }
   return shape;
 }
 
 
-Operation TensorComputeOpNode::make(std::string name,
-                                    std::string tag,
-                                    Array<IterVar> out_axis,
-                                    Array<IterVar> tensor_axis,
-                                    TensorIntrinCall intrin_call) {
-  return TensorComputeOpNode::make(name,
-                                   tag,
-                                   out_axis,
-                                   tensor_axis,
-                                   intrin_call->reduce_axis,
-                                   intrin_call->tensors,
-                                   intrin_call->regions,
-                                   intrin_call->intrin);
-}
+// Operation TensorComputeOpNode::make(std::string name,
+//                                     std::string tag,
+//                                     Array<IterVar> out_axis,
+//                                     Array<IterVar> tensor_axis,
+//                                     TensorIntrinCall intrin_call) {
+//   return TensorComputeOpNode::make(name,
+//                                    tag,
+//                                    out_axis,
+//                                    tensor_axis,
+//                                    intrin_call->reduce_axis,
+//                                    intrin_call->tensors,
+//                                    intrin_call->regions,
+//                                    intrin_call->intrin);
+// }
 
 Operation TensorComputeOpNode::make(std::string name,
                                     std::string tag,
-                                    Array<IterVar> out_axis,
-                                    Array<IterVar> tensor_axis,
+                                    Array<IterVar> axis,
                                     Array<IterVar> reduce_axis,
+                                    int sch_ndim,
                                     Array<Tensor> tensors,
                                     Array<Region> regions,
                                     TensorIntrin intrin) {
   auto n = make_node<TensorComputeOpNode>();
-  n->name = name;
-  n->tag = tag;
-  Array<IterVar> axis;
-  for (auto iv : out_axis) {
-    axis.push_back(iv);
-  }
-  for (auto iv : tensor_axis) {
-    axis.push_back(iv);
-  }
-  n->axis = axis;
-  n->out_axis = out_axis;
-  n->tensor_axis = tensor_axis;
-  n->reduce_axis = reduce_axis;
-  n->inputs = tensors;
-  n->input_regions = regions;
-  n->intrin = intrin;
+  n->name = std::move(name);
+  n->tag = std::move(tag);
+  n->axis = std::move(axis);
+  n->reduce_axis = std::move(reduce_axis);
+  n->sch_ndim = sch_ndim;
+  n->inputs = std::move(tensors);
+  n->input_regions = std::move(regions);
+  n->intrin = std::move(intrin);
   return Operation(n);
 }
 
@@ -156,11 +141,12 @@ void TensorComputeOpNode::GatherBound(
     const std::unordered_map<Tensor, TensorDom>& tensor_dom,
     std::unordered_map<IterVar, Range>* out_dom_map) const {
   const TensorDom& tdom = tensor_dom.at(self.output(0));
-  for (size_t i = 0; i < this->out_axis.size(); ++i) {
-    Range r = arith::Union(tdom.data.at(i)).cover_range(this->out_axis[i]->dom);
-    CHECK(!out_dom_map->count(this->out_axis[i]));
-    (*out_dom_map)[this->out_axis[i]] = r;
+  for (size_t i = 0; i < this->axis.size(); ++i) {
+    Range r = arith::Union(tdom.data.at(i)).cover_range(this->axis[i]->dom);
+    CHECK(!out_dom_map->count(this->axis[i]));
+    (*out_dom_map)[this->axis[i]] = r;
   }
+  // should I add dom of tensor_vars
   for (size_t i = 0; i < this->reduce_axis.size(); ++i) {
     CHECK(!out_dom_map->count(this->reduce_axis[i]));
     (*out_dom_map)[this->reduce_axis[i]] = this->reduce_axis[i]->dom;
@@ -173,21 +159,17 @@ Stmt TensorComputeOpNode::BuildRealize(
     const Stmt& body) const {
   CHECK_EQ(stage->op.get(), this);
   HalideIR::Internal::Region bounds;
-  for (IterVar iv : this->out_axis) {
+  for (IterVar iv : this->axis) {
     bounds.push_back(realize_map.at(iv));
   }
-  size_t out_buff_idx = this->intrin->buffers.size();
-  for (const Expr extent : this->intrin->buffers[out_buff_idx - 1]->shape) {
-    bounds.push_back(Range(0, extent));
-  }
   Stmt realize = body;
   for (int i = this->num_outputs(); i > 0; --i) {
     Tensor t = stage->op.output(i-1);
     realize = ir::Realize::make(t->op, t->value_index,
       t->dtype, bounds, const_true(), realize);
     // alignment requirement, only useful for compute
-    for (size_t i = 0; i < this->out_axis.size(); ++i) {
-      auto it = stage->iter_var_attrs.find(this->out_axis[i]);
+    for (int i = 0; i < sch_ndim; ++i) {
+      auto it = stage->iter_var_attrs.find(this->axis[i]);
       if (it != stage->iter_var_attrs.end()) {
         IterVarAttr attr = (*it).second;
         if (attr->dim_align_factor != 0) {
@@ -232,8 +214,8 @@ ComputeLoopNest MakeLoopNest(
     for (IterVar iv : self->reduce_axis) {
       update_state[iv] = 2;
     }
-    for (IterVar iv : self->out_axis) {
-      update_state[iv] = 1;
+    for (int i = 0; i < self->sch_ndim; ++i) {
+      update_state[self->axis[i]] = 1;
     }
     // find which iter var is related to reduction and which is related to axis.
     schedule::PassDownBitMaskOr(stage, &update_state);
@@ -308,13 +290,16 @@ Stmt TensorComputeOpNode::BuildProvide(
     Array<NodeRef> bind_spec{buffer, tensor};
 
     Array<Expr> tuple;
-    for (const IterVar ivar : this->out_axis) {
-      tuple.push_back(ivar->var);
-      tuple.push_back(1);
-    }
-    for (const Expr extent : buffer->shape) {
-      tuple.push_back(0);
-      tuple.push_back(extent);
+    for (size_t i = 0; i < this->axis.size(); ++i) {
+      auto ivar = this->axis[i];
+      if (i < static_cast<size_t>(this->sch_ndim)) {
+        tuple.push_back(ivar->var);
+        tuple.push_back(1);
+      } else {
+        Range dom = ivar->dom;
+        tuple.push_back(dom->min);
+        tuple.push_back(dom->extent);
+      }
     }
 
     output_bind_nest.emplace_back(AttrStmt::make(

src/schedule/schedule_dataflow_rewrite.cc

@@ -340,6 +340,12 @@ Array<Tensor> CacheWriteWithReLayoutTensor(Schedule sch,
     &red_axis, &new_axis, &dom_map, &vsub, &vsub2newvar, &predicates);
 
 
+  for (int i = tensor_op->sch_ndim; i < static_cast<int>(tensor_op->axis.size()); ++i) {
+    IterVar iv = tensor_op->axis[i];
+    IterVar new_iv = IterVarNode::make(
+      iv->dom, iv->var.copy_with_suffix(".c"), iv->iter_type);
+    new_axis.push_back(new_iv);
+  }
   Array<Region> new_regions;
   for (Region old_region : tensor_op->input_regions) {
     Region region;
@@ -353,15 +359,15 @@ Array<Tensor> CacheWriteWithReLayoutTensor(Schedule sch,
 
   Operation cache_op = TensorComputeOpNode::make(
       tensor_op->name + "." + scope, tensor_op->tag, new_axis,
-      tensor_op->tensor_axis, tensor_op->reduce_axis,
+      tensor_op->reduce_axis, tensor_op->sch_ndim,
       tensor_op->inputs, new_regions, tensor_op->intrin);
 
   // axis will be used in generating compute op
-  Array<IterVar> compute_axis = tensor_op->out_axis;
-  for (IterVar iv : tensor_op->tensor_axis) {
-    // new tensor axis with kDataPar IterVar type
+  Array<IterVar> compute_axis = tensor_op->axis;
+  for (size_t i = tensor_op->sch_ndim; i < tensor_op->axis.size(); ++i) {
+    IterVar iv = tensor_op->axis[i];
     IterVar aiv = IterVarNode::make(iv->dom, iv->var, kDataPar);
-    compute_axis.push_back(aiv);
+    compute_axis.Set(i, aiv);
   }
 
   // The reader args
@@ -378,8 +384,8 @@ Array<Tensor> CacheWriteWithReLayoutTensor(Schedule sch,
       args.push_back(value_map.at(iv));
     }
     // tensorized region axis
-    for (size_t i = 0; i < tensor_op->tensor_axis.size(); ++i) {
-      IterVar iv = compute_axis[tensor_op->out_axis.size() + i];
+    for (size_t i = tensor_op->sch_ndim; i < tensor_op->axis.size(); ++i) {
+      IterVar iv = compute_axis[i];
       args.push_back(value_map.at(iv));
     }
   }

tests/python/unittest/test_lang_tensor.py

@@ -111,6 +111,7 @@ def intrin_func(ins, outs):
 
     s = tvm.create_schedule(C.op)
     stmt = tvm.lower(s, [A, B, C], simple_mode=True)
+    print(stmt)
     assert isinstance(stmt.body.body, tvm.stmt.Evaluate)
 
 def test_tensor_compute2():
@@ -154,6 +155,7 @@ def intrin_func(ins, outs):
 
     s = tvm.create_schedule(C.op)
     stmt = tvm.lower(s, [A, B, C], simple_mode=True)
+    print(stmt)
     assert isinstance(stmt.body.body.body.first, tvm.stmt.Evaluate)
     assert isinstance(stmt.body.body.body.rest.body, tvm.stmt.Evaluate)
 

tests/python/unittest/test_schedule_schedule_ops.py

@@ -387,6 +387,7 @@ def test_schedule_tensor_compute2():
 
 
 def test_schedule_tensor_compute3():
+    # compute_at
     M = 1024
     factor = 16
     dtype = 'float32'