[API] Prefetch schedule supported (#258)

* prefetch interface added

* prefetch python comments modified. prefetch info data structure maintained.

* start injecting prefetches. first step (domain touch) implemented.

* domain touch tested.

* Prefetch ir_mutator and ir_visitor dispatch registered.

* modify domain touched from passing a func_ref to passing a tensor

* modify domain touched from passing a func_ref to passing a tensor

* modify Tensor copy to Tensor ref

* temp commit for rebase

* debug info removed, typo fixed, ready to rebase

* prefetch flatten test add!

* roll back builtin functions to side effect functions

* lint error fixed!

* add cache line size to storage flatten argument

* forgot modifications add

* change code style to dmlc-like; get rid of can_prove, use manually compute instead

* python lint error fixed

* modify instrinsic name to pass tests

* [TEST] get rid of str(), replace them by accessing attributes

* change map to list comprehension

* redundant numpy import removed

include/tvm/arithmetic.h

@@ -12,6 +12,9 @@
 #include "./expr.h"
 
 namespace tvm {
+
+class Tensor;
+
 /*! \brief namespace of arithmetic */
 namespace arith {
 /*!
@@ -255,6 +258,16 @@ IntSet DeduceBound(Expr v, Expr cond,
                    const std::unordered_map<const Variable*, IntSet>& relax_map);
 
 /*!
+ * \brief Infer a regular domain that covers all the calls or provides within the given statement.
+ * \param body The given statement.
+ * \param tensor The name of the calls or provides.
+ * \param consider_calls If calls (read) are considered.
+ * \param consider_provides If provides (write) are considered.
+ * \return The domain that covers all the calls or provides within the given statement.
+ */
+Domain DomainTouched(Stmt body, const Tensor &tensor, bool consider_calls, bool consider_provides);
+
+/*!
  * \brief Evaluate the expression with modular analysis
  * \param e The expression to be evaluated.
  * \param mod_map Map of modular statistics of known variables.

include/tvm/ir_pass.h

@@ -169,10 +169,12 @@ Stmt Inline(Stmt stmt,
  * \param stmt The stmt to be trasnformed.
  * \param extern_buffer Map specifies external
  *    buffer assignment of input and outputs.
+ * \param cache_line_size The size of CPU cache line.
  * \return Transformed stmt.
  */
 Stmt StorageFlatten(Stmt stmt,
-                    Map<Tensor, Buffer> extern_buffer);
+                    Map<Tensor, Buffer> extern_buffer,
+                    int cache_line_size);
 
 /*!
  * \brief Remove No Op from the Stmt.
@@ -223,6 +225,13 @@ Stmt VectorizeLoop(Stmt stmt);
 Stmt InjectVirtualThread(Stmt stmt);
 
 /*!
+ * \brief Inject prefetch instructions into stmt.
+ * \param stmt The statment to be transformed.
+ * \return Transformed stmt.
+ */
+Stmt InjectPrefetch(Stmt stmt);
+
+/*!
  * \brief Rewrite storage allocation pattern.
  *  Moves the allocation to outer most possible scope.
  *  Trying to share space between allocations to make

include/tvm/schedule.h

@@ -181,6 +181,14 @@ class Stage : public NodeRef {
    */
   Stage& parallel(IterVar var);   // NOLINT(*)
   /*!
+   * \brief Fetch data in advance.
+   * \param domain the tensor to be prefetched
+   * \param var the iteration point at which to apply prefetching
+   * \param offset the number of iterations be to fetched in advance
+   * \return reference to self
+   */
+  Stage& prefetch(const Tensor &domain, IterVar var, Expr offset); //NOLINT(*)
+  /*!
    * \brief whether the stage has been scheduled.
    * \return whether the stage has been scheduled.
    */

python/tvm/build.py

@@ -185,7 +185,8 @@ def lower(sch,
     sch = sch.normalize()
     bounds = schedule.InferBound(sch)
     stmt = schedule.ScheduleOps(sch, bounds)
-    stmt = ir_pass.StorageFlatten(stmt, binds)
+    stmt = ir_pass.InjectPrefetch(stmt)
+    stmt = ir_pass.StorageFlatten(stmt, binds, 64)
     stmt = ir_pass.CanonicalSimplify(stmt)
     if not simple_mode:
         stmt = ir_pass.LoopPartition(stmt)

python/tvm/schedule.py

@@ -422,4 +422,18 @@ class Stage(NodeBase):
         """
         _api_internal._StageParallel(self, var)
 
+    def prefetch(self, tensor, var, offset):
+        """Prefetch the specified variable
+
+        Parameters
+        ----------
+        tensor : Tensor
+            The tensor to be prefetched
+        var : IterVar
+            The loop point at which the prefetching is applied
+        offset : Expr
+            The number of iterations to be prefetched before actual execution
+        """
+        _api_internal._StagePrefetch(self, tensor, var, offset)
+
 _init_api("tvm.schedule")

src/api/api_arith.cc

@@ -6,7 +6,7 @@
 #include <tvm/expr.h>
 #include <tvm/ir.h>
 #include <tvm/api_registry.h>
-#include <tvm/arithmetic.h>
+#include <tvm/tensor.h>
 
 namespace tvm {
 namespace arith {
@@ -38,6 +38,13 @@ TVM_REGISTER_API("arith.DeduceBound")
         args[3].operator Map<Var, IntSet>());
   });
 
+
+TVM_REGISTER_API("arith.DomainTouched")
+.set_body([](TVMArgs args, TVMRetValue *ret) {
+    *ret = DomainTouched(args[0], args[1], args[2], args[3]);
+  });
+
+
 TVM_REGISTER_API("_IntervalSetGetMin")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     *ret = args[0].operator IntSet().min();

src/api/api_lang.cc

@@ -358,6 +358,12 @@ TVM_REGISTER_API("_StageParallel")
         .parallel(args[1]);
   });
 
+TVM_REGISTER_API("_StagePrefetch")
+  .set_body([](TVMArgs args, TVMRetValue *ret) {
+    args[0].operator Stage()
+      .prefetch(args[1], args[2], args[3]);
+  });
+
 TVM_REGISTER_API("_ScheduleNormalize")
 .set_body([](TVMArgs args, TVMRetValue* ret) {
     *ret = args[0].operator Schedule()

src/api/api_pass.cc

@@ -86,7 +86,7 @@ TVM_REGISTER_API("ir_pass.PostOrderVisit")
 REGISTER_PASS1(ConvertSSA);
 REGISTER_PASS1(VerifySSA);
 REGISTER_PASS4(Inline);
-REGISTER_PASS2(StorageFlatten);
+REGISTER_PASS3(StorageFlatten);
 REGISTER_PASS1(VectorizeLoop);
 REGISTER_PASS4(UnrollLoop);
 REGISTER_PASS2(StorageSync);
@@ -95,6 +95,7 @@ REGISTER_PASS2(BindDeviceType);
 REGISTER_PASS1(SplitHostDevice);
 REGISTER_PASS1(StorageRewrite);
 REGISTER_PASS1(InjectVirtualThread);
+REGISTER_PASS1(InjectPrefetch);
 REGISTER_PASS1(LoopPartition);
 REGISTER_PASS1(RemoveNoOp);
 REGISTER_PASS2(SplitPipeline);

src/arithmetic/domain_touched.cc

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file bound_deducer.cc
+ * \brief Utility to deduce bound of expression
+ */
+#include <tvm/expr.h>
+#include <tvm/ir_pass.h>
+#include <tvm/ir_visitor.h>
+#include <tvm/tensor.h>
+#include <tvm/api_registry.h>
+
+#include <unordered_set>
+#include <unordered_map>
+
+namespace tvm {
+namespace arith {
+
+using namespace ir;
+
+// Find Read region of the tensor in the stmt.
+class FuncTouchedDomain final : public IRVisitor {
+ public:
+  FuncTouchedDomain(const Tensor &tensor, bool consider_calls, bool consider_provides)
+    : tensor_(tensor), consider_calls_(consider_calls), consider_provides_(consider_provides)  {}
+
+  Domain Find(const Stmt& stmt) {
+    this->Visit(stmt);
+    Domain ret;
+    Range none;
+    for (size_t i = 0; i < bounds_.size(); ++i) {
+      ret.push_back(arith::Union(bounds_[i]).cover_range(none));
+    }
+    return ret;
+  }
+
+  void Visit_(const For *op) final {
+    const Variable* var = op->loop_var.get();
+    dom_map_[var] = IntSet::range(
+        Range::make_by_min_extent(op->min, op->extent));
+    IRVisitor::Visit_(op);
+    dom_map_.erase(var);
+  }
+
+  void Visit_(const LetStmt* op) final {
+    dom_map_[op->var.get()] =
+        arith::EvalSet(op->value, dom_map_);
+    IRVisitor::Visit_(op);
+    dom_map_.erase(op->var.get());
+  }
+
+  /* TODO: Thread extent unitest not generated.*/
+  void Visit_(const AttrStmt* op) final {
+    if (op->attr_key == attr::thread_extent) {
+      const IterVarNode* thread_axis = op->node.as<IterVarNode>();
+      CHECK(thread_axis);
+      const Variable* var = thread_axis->var.get();
+      dom_map_[var] = IntSet::range(Range(make_zero(op->value.type()), op->value));
+      IRVisitor::Visit_(op);
+      dom_map_.erase(var);
+    } else {
+      IRVisitor::Visit_(op);
+    }
+  }
+
+  void Visit_(const Call* op) final {
+    if (consider_calls_ && tensor_->op.same_as(op->func)
+        && tensor_->value_index == op->value_index) {
+      Touch(op->args);
+    }
+    IRVisitor::Visit_(op);
+  }
+
+  void Visit_(const Provide* op) final {
+    if (consider_provides_ && tensor_->op.same_as(op->func)
+        && tensor_->value_index == op->value_index) {
+      Touch(op->args);
+    }
+    IRVisitor::Visit_(op);
+  }
+
+ private:
+  void Touch(const Array<Expr>& args) {
+    if (args.size() > bounds_.size()) {
+      bounds_.resize(args.size());
+    }
+    for (size_t i = 0; i < args.size(); ++i) {
+      bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
+    }
+  }
+
+  const Tensor &tensor_;
+  bool consider_calls_, consider_provides_;
+  std::vector<std::vector<IntSet> > bounds_;
+  std::unordered_map<const Variable*, IntSet> dom_map_;
+};
+
+Domain DomainTouched(Stmt stmt, const Tensor &tensor, bool consider_calls, bool consider_provides) {
+  return FuncTouchedDomain(tensor, consider_calls, consider_provides).Find(stmt);
+}
+
+}  // namespace arith
+}  // namespace tvm

src/codegen/llvm/codegen_llvm.cc

@@ -807,7 +807,7 @@ llvm::Value* CodeGenLLVM::CreateIntrinstic(const Call* op) {
     llvm::Function* f = llvm::Intrinsic::getDeclaration(
         module_.get(), id, arg_types);
     return builder_->CreateCall(f, arg_values);
-  } else if (op->is_intrinsic("llvm_buildin")) {
+  } else if (op->is_intrinsic("llvm_builtin")) {
     std::vector<llvm::Value*> arg_values;
     for (size_t i = 1; i < op->args.size(); ++i) {
       llvm::Value* v = MakeValue(op->args[i]);

src/codegen/llvm/intrin_rule_llvm.cc

@@ -28,9 +28,12 @@ inline void DispatchLLVMBuildin(const TVMArgs& targs, TVMRetValue* rv) {
     cargs.push_back(arg);
   }
   *rv = Call::make(
-      call->type, "llvm_buildin", cargs, Call::Intrinsic);
+      call->type, "llvm_builtin", cargs, Call::Intrinsic);
 }
 
+TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
+.set_body(DispatchLLVMBuildin<::llvm::Intrinsic::prefetch>);
+
 template<unsigned id>
 inline void DispatchLLVMPureIntrin(const TVMArgs& targs, TVMRetValue* rv) {
   Expr e = targs[0];
@@ -46,8 +49,20 @@ inline void DispatchLLVMPureIntrin(const TVMArgs& targs, TVMRetValue* rv) {
       call->type, "llvm_intrin", cargs, Call::PureIntrinsic);
 }
 
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.__buildin_prefetch")
-.set_body(DispatchLLVMBuildin<::llvm::Intrinsic::prefetch>);
+template<unsigned id>
+inline void DispatchLLVMIntrin(const TVMArgs& targs, TVMRetValue* rv) {
+  Expr e = targs[0];
+  const Call* call = e.as<Call>();
+  CHECK(call != nullptr);
+  Array<Expr> cargs;
+  // intrin id.
+  cargs.push_back(UIntImm::make(UInt(32), id));
+  for (Expr arg : call->args) {
+    cargs.push_back(arg);
+  }
+  *rv = Call::make(
+      call->type, "llvm_intrin", cargs, Call::Intrinsic);
+}
 
 TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
 .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp>);

src/op/op_util.cc

@@ -95,11 +95,11 @@ MakeLoopNest(const Stage& stage,
             << "Cannot prefetch on trivial loop with extent=1";
         CHECK_EQ(it_attr->prefetch_data.size(),
                  it_attr->prefetch_offset.size());
-        for (size_t i = 0; i < it_attr->prefetch_data.size(); ++i) {
+        for (size_t j = 0; j < it_attr->prefetch_data.size(); ++j) {
           nest[i + 1].emplace_back(
-              AttrStmt::make(it_attr->prefetch_data[i],
+              AttrStmt::make(it_attr->prefetch_data[j],
                              ir::attr::prefetch_scope,
-                             it_attr->prefetch_offset[i], no_op));
+                             it_attr->prefetch_offset[j], no_op));
         }
       }
     } else if (bind_iv->thread_tag == "vthread") {

src/pass/inject_prefetch.cc

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file inject_prefetch.cc
+ */
+#include <tvm/ir.h>
+#include <tvm/ir_mutator.h>
+#include <tvm/ir_visitor.h>
+#include <tvm/ir_pass.h>
+#include <tvm/arithmetic.h>
+#include <unordered_set>
+
+namespace tvm {
+namespace ir {
+
+using arith::IntSet;
+using arith::DomainTouched;
+using Halide::Internal::Region;
+
+class PrefetchInjector : public IRMutator {
+ public:
+  Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
+    Stmt ret = IRMutator::Mutate_(op, s);
+    op = ret.as<AttrStmt>();
+    if (op && op->attr_key == attr::prefetch_scope) {
+      Tensor ts(op->node.node_);
+      CHECK_NE(loop_nest_.size(), 0U);
+      Domain domain = DomainTouched(op->body, ts, true, false);
+      Region region;
+
+      auto iter_var = loop_nest_.back().get();
+      vectorized_[iter_var] = IntSet::single_point(loop_nest_.back() + op->value);
+
+      for (Range r : domain) {
+        if (!r.defined()) {
+          LOG(WARNING) << "Cannot decide prefetch region for " << ts;
+          return op->body;
+        }
+        Range res(EvalSet(r, vectorized_).cover_range(none));
+        region.push_back(Range::make_by_min_extent(res->min, res->extent));
+      }
+
+      vectorized_.erase(iter_var);
+
+      Stmt prefetch = Prefetch::make(ts->op, ts->value_index, ts->dtype, region);
+      return Block::make(prefetch, op->body);
+    }
+    return ret;
+  }
+
+  Stmt Mutate_(const For* op, const Stmt& s) final {
+    auto &var = op->loop_var;
+    loop_nest_.push_back(var);
+    if (op->for_type == ForType::Vectorized) {
+      vectorized_[var.get()] = IntSet::interval(op->min, (op->min + op->extent) - 1);
+    }
+    Stmt ret = IRMutator::Mutate_(op, s);
+    if (op->for_type == ForType::Vectorized) {
+      vectorized_.erase(var.get());
+    }
+    loop_nest_.pop_back();
+    return ret;
+  }
+
+ private:
+  std::vector<VarExpr> loop_nest_;
+  std::unordered_map<const Variable *, IntSet> vectorized_;
+  static const Range none;
+};
+
+const Range PrefetchInjector::none;
+
+Stmt InjectPrefetch(Stmt stmt) {
+  return PrefetchInjector().Mutate(stmt);
+}
+
+}  // namespace ir
+}  // namespace tvm

src/pass/ir_mutator.cc

@@ -265,7 +265,8 @@ TVM_STATIC_IR_FUNCTOR(IRMutator, vtable_stmt)
 .DISPATCH_TO_MUTATE_STMT(Provide)
 .DISPATCH_TO_MUTATE_STMT(Realize)
 .DISPATCH_TO_MUTATE_STMT(Block)
-.DISPATCH_TO_MUTATE_STMT(Evaluate);
+.DISPATCH_TO_MUTATE_STMT(Evaluate)
+.DISPATCH_TO_MUTATE_STMT(Prefetch);
 
 
 // Mutate Expr

src/pass/ir_visitor.cc

@@ -255,7 +255,8 @@ TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
 .DISPATCH_TO_VISIT(IntImm)
 .DISPATCH_TO_VISIT(UIntImm)
 .DISPATCH_TO_VISIT(FloatImm)
-.DISPATCH_TO_VISIT(StringImm);
+.DISPATCH_TO_VISIT(StringImm)
+.DISPATCH_TO_VISIT(Prefetch);
 
 }  // namespace ir
 }  // namespace tvm

src/pass/storage_flatten.cc

@@ -5,6 +5,7 @@
 #include <tvm/ir.h>
 #include <tvm/expr.h>
 #include <tvm/ir_mutator.h>
+#include <tvm/ir_operator.h>
 #include <tvm/ir_pass.h>
 #include <tvm/buffer.h>
 #include <tvm/runtime/device_api.h>
@@ -20,16 +21,18 @@ namespace ir {
 using Halide::Internal::Region;
 using runtime::StorageScope;
 using runtime::ThreadScope;
+using intrinsic::tvm_address_of;
 
 class StorageFlattener : public IRMutator {
  public:
-  explicit StorageFlattener(Map<Tensor, Buffer> extern_buffer) {
+  explicit StorageFlattener(Map<Tensor, Buffer> extern_buffer, int cache_line_size) {
     for (auto kv : extern_buffer) {
       BufferEntry e;
       e.buffer = kv.second;
       e.external = true;
       buf_map_[TensorKey{kv.first->op, kv.first->value_index}] = e;
     }
+    cache_line_size_ = cache_line_size;
   }
   Stmt Mutate_(const Store* op, const Stmt& s) final {
     Stmt stmt = IRMutator::Mutate_(op, s);
@@ -169,6 +172,62 @@ class StorageFlattener : public IRMutator {
     }
   }
 
+  Stmt Mutate_(const Prefetch *op, const Stmt &s) final {
+    Stmt stmt = IRMutator::Mutate_(op, s);
+    op = stmt.as<Prefetch>();
+    CHECK(op != nullptr);
+    TensorKey key{op->func, op->value_index};
+    auto it = buf_map_.find(key);
+    CHECK(it != buf_map_.end())
+        << "Cannot find allocated buffer for " << key.f;
+    const BufferEntry& e = it->second;
+
+    CHECK(!e.released)
+        << "Read a buffer that is already out of scope";
+    CHECK_EQ(e.buffer->shape.size(), op->bounds.size())
+      << "Prefetch dim should be the same as buffer dim";
+
+    int block_size = 1,
+        elem_cnt = cache_line_size_ / e.buffer->dtype.bytes(),
+        shape = 0;
+
+    int starts = op->bounds.size() - 1;
+    while (starts > 0 && arith::GetConstInt(e.buffer->shape[starts], &shape)
+        && elem_cnt >= block_size * shape) {
+      block_size *= shape;
+      starts--;
+    }
+    Expr stride(elem_cnt / block_size);
+
+    Array<Expr> args;
+    std::vector<VarExpr> vars;
+
+    for (int i = op->bounds.size() - 1; i > starts; --i) {
+      args.push_back(op->bounds[i]->min);
+    }
+    auto &func_name = op->func->func_name();
+    vars.push_back(VarExpr("prefetch." + func_name + "." + std::to_string(starts), Int(32)));
+    args.push_back(op->bounds[starts]->min + stride * vars.back());
+    for (int i = starts - 1; i >= 0; --i) {
+      vars.push_back(VarExpr("prefetch." + func_name + "." + std::to_string(i), Int(32)));
+      args.push_back(vars.back() + op->bounds[i]->min);
+    }
+    for (int i = starts; i >= 0; --i) {
+      if (i < starts) {
+        stmt = For::make(
+            vars[i], 0, op->bounds[i]->extent, ForType::Serial, DeviceAPI::Host, stmt);
+      } else {
+        Expr load = e.buffer.MakeLoad(e.RelIndex(args));
+        Expr address = Call::make(Handle(), tvm_address_of, {load}, Call::PureIntrinsic);
+        Expr prefetch = Call::make(op->type, Call::prefetch, {address, 0, 3, 1}, Call::Intrinsic);
+        stmt = Evaluate::make(prefetch);
+        Expr extent = (op->bounds[i]->extent - 1) / stride + 1;
+        stmt = For::make(vars[i], 0, extent, ForType::Serial, DeviceAPI::Host, stmt);
+      }
+    }
+    return stmt;
+  }
+
  private:
   // Start bind
   Stmt HandleBufferBindScope(const AttrStmt* op) {
@@ -252,11 +311,14 @@ class StorageFlattener : public IRMutator {
   std::unordered_map<const Node*, std::string> storage_scope_;
   // The current thread scope.
   std::vector<ThreadScope> curr_thread_scope_;
+  // The size of cacheline
+  int cache_line_size_;
 };
 
 Stmt StorageFlatten(Stmt stmt,
-                    Map<Tensor, Buffer> extern_buffer) {
-  stmt = StorageFlattener(extern_buffer).Mutate(stmt);
+                    Map<Tensor, Buffer> extern_buffer,
+                    int cache_line_size) {
+  stmt = StorageFlattener(extern_buffer, cache_line_size).Mutate(stmt);
   return stmt;
 }
 

src/schedule/schedule_lang.cc

@@ -337,6 +337,24 @@ Stage& Stage::parallel(IterVar var) {   // NOLINT(*)
   return *this;
 }
 
+Stage& Stage::prefetch(const Tensor &tensor, IterVar var, Expr offset) {
+  StageNode *self = operator->();
+  ArrayNode* all_vars = self->all_iter_vars.CopyOnWrite();
+  ArrayNode* leaf_vars = self->leaf_iter_vars.CopyOnWrite();
+  FindLeafVar(all_vars, leaf_vars, var);
+  auto it = self->iter_var_attrs.find(var);
+  std::shared_ptr<IterVarAttrNode> n;
+  if (it != self->iter_var_attrs.end()) {
+    n = std::make_shared<IterVarAttrNode>(*(*it).second.operator->());
+  } else {
+    n = std::make_shared<IterVarAttrNode>();
+  }
+  n->prefetch_data.push_back(tensor);
+  n->prefetch_offset.push_back(offset);
+  self->iter_var_attrs.Set(var, IterVarAttr(n));
+  return *this;
+}
+
 Stage CopyStage(const Stage& s) {
   std::shared_ptr<StageNode> n =
       std::make_shared<StageNode>(*s.operator->());

tests/python/integration/test_dot.py

@@ -13,7 +13,7 @@ def lower(s, args, name="mydot"):
     s = s.normalize()
     bounds = tvm.schedule.InferBound(s)
     stmt = tvm.schedule.ScheduleOps(s, bounds)
-    stmt = tvm.ir_pass.StorageFlatten(stmt, binds)
+    stmt = tvm.ir_pass.StorageFlatten(stmt, binds, 16)
     stmt = tvm.ir_pass.CanonicalSimplify(stmt)
     stmt = tvm.ir_pass.Simplify(stmt)
     fapi = tvm.ir_pass.MakeAPI(stmt, name, arg_list, 0, True)

tests/python/unittest/test_arith_domain_touched.py

+import tvm
+
+def test_domain_touched():
+    i = tvm.var('i')
+    j = tvm.var('j')
+    n = tvm.convert(100)
+    m = tvm.var('m')
+    a = tvm.placeholder((n, m), name = 'a')
+    b = tvm.placeholder((n, m), name = 'b')
+    ir = tvm.make.For(
+            i, 0, n, 0, 0,
+            tvm.make.For(j, 0, m, 0, 0,
+                tvm.make.Provide(
+                    a.op,
+                    0,
+                    tvm.make.Call(b.dtype, 'b', [i - 1, j + 1], 3, b.op, 0) + 
+                    tvm.make.Call(a.dtype, 'a', [i - 1, j - 1], 3, a.op, 0),
+                    [i, j]
+                )
+            )
+    )
+    a_domain_r = tvm.arith.DomainTouched(ir, a, True, False)
+    assert a_domain_r[0].min.value == -1
+    assert a_domain_r[0].extent.value == 100
+    assert a_domain_r[1].min.value == -1
+    assert a_domain_r[1].extent.name == 'm'
+
+    a_domain_w = tvm.arith.DomainTouched(ir, a, False, True)
+    assert a_domain_w[0].min.value == 0
+    assert a_domain_w[0].extent.value == 100
+    assert a_domain_w[1].min.value == 0
+    assert a_domain_w[1].extent.name == 'm'
+
+    a_domain_rw= tvm.arith.DomainTouched(ir, a, True, True)
+    assert a_domain_rw[0].min.value == -1
+    assert a_domain_rw[0].extent.value == 101
+    assert a_domain_rw[1].min.value == -1
+    assert isinstance(a_domain_rw[1].extent, tvm.expr.Add)
+    assert a_domain_rw[1].extent.a.name == 'm'
+    assert a_domain_rw[1].extent.b.value == 1
+
+    b_domain_r = tvm.arith.DomainTouched(ir, b, True, False)
+    assert b_domain_r
+    assert b_domain_r[0].min.value == -1
+    assert b_domain_r[0].extent.value == 100
+    assert b_domain_r[1].min.value == 1
+    assert b_domain_r[1].extent.name == 'm'
+
+    b_domain_w = tvm.arith.DomainTouched(ir, b, False, True)
+    assert isinstance(b_domain_w, tvm.container.Array)
+    assert len(b_domain_w) == 0
+
+if __name__ == "__main__":
+    test_domain_touched()
+

tests/python/unittest/test_codegen_device.py

@@ -28,7 +28,7 @@ def test_add_pipeline():
     Bb = tvm.decl_buffer(B.shape, B.dtype, name='B')
     Cb = tvm.decl_buffer(C.shape, C.dtype, name='C')
     stmt = tvm.ir_pass.LoopPartition(stmt)
-    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, B:Bb, C:Cb})
+    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, B:Bb, C:Cb}, 64)
     stmt = tvm.ir_pass.Simplify(stmt)
     fapi = tvm.ir_pass.MakeAPI(stmt, "myadd", [Ab, Bb, Cb], 0, True)
     fsplits = [x for x in tvm.ir_pass.SplitHostDevice(fapi)]

tests/python/unittest/test_codegen_llvm.py

@@ -12,7 +12,7 @@ def test_llvm_intrin():
     ]
     ib.emit(tvm.make.Evaluate(
         tvm.make.Call(
-            "int32", "__buildin_prefetch", args, tvm.expr.Call.Intrinsic, None, 0)))
+            "int32", "prefetch", args, tvm.expr.Call.Intrinsic, None, 0)))
     body = ib.get()
     func = tvm.ir_pass.MakeAPI(body, "prefetch", [A], 0, True)
     fcode = tvm.build(func, None, "llvm")

tests/python/unittest/test_pass_loop_partition.py

@@ -20,7 +20,7 @@ def lower(sch, args):
     bounds = tvm.schedule.InferBound(sch)
     stmt = tvm.schedule.ScheduleOps(sch, bounds)
     stmt = tvm.ir_pass.LoopPartition(stmt)
-    stmt = tvm.ir_pass.StorageFlatten(stmt, binds)
+    stmt = tvm.ir_pass.StorageFlatten(stmt, binds, 64)
     stmt = tvm.ir_pass.CanonicalSimplify(stmt)
     stmt = tvm.ir_pass.VectorizeLoop(stmt)
     stmt = tvm.ir_pass.Simplify(stmt)

tests/python/unittest/test_pass_makeapi.py

@@ -15,7 +15,7 @@ def test_makeapi():
     Ab = tvm.decl_buffer(A.shape, A.dtype, name='A')
     Bb = tvm.decl_buffer(B.shape, B.dtype, name='B')
     Cb = tvm.decl_buffer(C.shape, C.dtype, name='C')
-    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, B:Bb, C:Cb})
+    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, B:Bb, C:Cb}, 64)
 
     num_unpacked_args = 2
     f = tvm.ir_pass.MakeAPI(

tests/python/unittest/test_pass_split_pipeline.py

@@ -12,7 +12,7 @@ def lower(s, args):
     s.normalize()
     bounds = tvm.schedule.InferBound(s)
     stmt = tvm.schedule.ScheduleOps(s, bounds)
-    stmt = tvm.ir_pass.StorageFlatten(stmt, binds)
+    stmt = tvm.ir_pass.StorageFlatten(stmt, binds, 64)
     stmt = tvm.ir_pass.CanonicalSimplify(stmt)
     stmt = tvm.ir_pass.Simplify(stmt)
     return stmt

tests/python/unittest/test_pass_storage_flatten.py

@@ -16,8 +16,22 @@ def test_flatten2():
 
     Ab = tvm.decl_buffer(A.shape, A.dtype, name='A')
     A2b = tvm.decl_buffer(A2.shape, A2.dtype, name='A2')
-    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, A2: A2b})
+    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, A2: A2b}, 64)
     stmt = tvm.ir_pass.Simplify(stmt)
 
+def test_flatten_prefetch():
+    A = tvm.placeholder((25, 100, 4), name = 'A')
+    _A= tvm.decl_buffer(A.shape, A.dtype, name = 'A');
+    i = tvm.var('i')
+    j = tvm.var('j')
+    region = [tvm.make.range_by_min_extent(i[0], i[1]) for i in [(i, 2), (j, 8), (0, 4)]]
+    stmt = tvm.make.Prefetch(A.op, 0, A.dtype, region)
+    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: _A}, 64)
+    stmt = tvm.ir_pass.Simplify(stmt)
+    assert stmt.extent.value == 2
+    assert isinstance(stmt.body, tvm.stmt.For)
+    assert stmt.body.extent.value == 2
+
 if __name__ == "__main__":
     test_flatten2()
+    test_flatten_prefetch()

tests/python/unittest/test_pass_storage_rewrite.py

@@ -15,7 +15,7 @@ def test_storage_share():
     stmt = tvm.schedule.ScheduleOps(s, bounds)
     Ab = tvm.decl_buffer(A.shape, A.dtype, name='A')
     Bb = tvm.decl_buffer(B.shape, B.dtype, name='B')
-    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, B: Bb})
+    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, B: Bb}, 64)
     stmt = tvm.ir_pass.CanonicalSimplify(stmt)
     stmt = tvm.ir_pass.Simplify(stmt)
     stmt = tvm.ir_pass.StorageRewrite(stmt)
@@ -51,7 +51,7 @@ def test_storage_share_gpu():
     stmt = tvm.schedule.ScheduleOps(s, bounds)
     Ab = tvm.decl_buffer(A[0].shape, A[0].dtype, name='A')
     Bb = tvm.decl_buffer(A[0].shape, A[0].dtype, name='B')
-    stmt = tvm.ir_pass.StorageFlatten(stmt, {A[0]: Ab, A[-1]: Bb})
+    stmt = tvm.ir_pass.StorageFlatten(stmt, {A[0]: Ab, A[-1]: Bb}, 64)
     stmt = tvm.ir_pass.CanonicalSimplify(stmt)
     stmt = tvm.ir_pass.Simplify(stmt)
     stmt = tvm.ir_pass.StorageRewrite(stmt)

tests/python/unittest/test_pass_storage_sync.py

@@ -19,7 +19,7 @@ def test_storage_sync():
     stmt = tvm.schedule.ScheduleOps(s, bounds)
     Ab = tvm.decl_buffer(A.shape, A.dtype, name='A')
     A2b = tvm.decl_buffer(A2.shape, A2.dtype, name='A2')
-    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, A2: A2b})
+    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, A2: A2b}, 64)
     f = tvm.ir_pass.MakeAPI(stmt, "test", [Ab, A2b], 0, True)
     flist = tvm.ir_pass.SplitHostDevice(f)
     f = flist[1]

tests/python/unittest/test_pass_virtual_thread.py

@@ -19,7 +19,7 @@ def test_virtual_thread():
 
     Ab = tvm.decl_buffer(A.shape, A.dtype, name='A')
     A2b = tvm.decl_buffer(A2.shape, A2.dtype, name='A2')
-    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, A2: A2b})
+    stmt = tvm.ir_pass.StorageFlatten(stmt, {A: Ab, A2: A2b}, 64)
     stmt = tvm.ir_pass.Simplify(stmt)
     stmt = tvm.ir_pass.InjectVirtualThread(stmt)
     print(stmt)

tests/verilog/integration/test_codegen_verilog.py

@@ -14,7 +14,7 @@ def lower(s, args, name):
     s = s.normalize()
     bounds = tvm.schedule.InferBound(s)
     stmt = tvm.schedule.ScheduleOps(s, bounds)
-    stmt = tvm.ir_pass.StorageFlatten(stmt, binds)
+    stmt = tvm.ir_pass.StorageFlatten(stmt, binds, 64)
     stmt = tvm.ir_pass.CanonicalSimplify(stmt)
     stmt = tvm.ir_pass.Simplify(stmt)
     stmt = tvm.ir_pass.SplitPipeline(stmt, True)