[PASS] Memory barrier detection, storage access lower. (#317)

apache · Aug 13, 2017 · 79e482b · 79e482b
1 parent afa2086
commit 79e482b
Show file tree

Hide file tree

Showing 10 changed files with 570 additions and 283 deletions.
diff --git a/include/tvm/ir_pass.h b/include/tvm/ir_pass.h
@@ -266,6 +266,15 @@ Stmt CoProcSync(Stmt stmt);
  */
 Stmt LiftAttrScope(Stmt stmt, std::string attr_key);
 
+/*!
+ * \brief Lower attached storage access information.
+ * Do this pass after all storage access analysis finish.
+ *
+ * \param stmt The stmt to be trasnformed
+ * \return Transformed stmt.
+ */
+Stmt LowerStorageAccessInfo(Stmt stmt);
+
 /*!
  * \brief Make an user callable API LoweredFunc.
  *

diff --git a/include/tvm/target_info.h b/include/tvm/target_info.h
@@ -23,11 +23,17 @@ struct MemoryInfoNode : public Node {
   int max_num_bits;
   /*! \brief maximum number of bits to be used in simd op */
   int max_simd_bits;
+  /*!
+   * \brief head address of the buffer, if visible to CPU
+   *  This address can be None.
+   */
+  Expr head_address;
 
   void VisitAttrs(AttrVisitor* v) final {
     v->Visit("unit_bits", &unit_bits);
     v->Visit("max_num_bits", &max_num_bits);
     v->Visit("max_simd_bits", &max_simd_bits);
+    v->Visit("head_address", &head_address);
   }
 
   static constexpr const char* _type_key = "MemoryInfo";

diff --git a/python/tvm/build_module.py b/python/tvm/build_module.py
@@ -197,7 +197,6 @@ def lower(sch,
     stmt = ir_pass.VectorizeLoop(stmt)
     stmt = ir_pass.InjectVirtualThread(stmt)
     stmt = ir_pass.StorageRewrite(stmt)
-    stmt = ir_pass.CoProcSync(stmt)
     cfg = BuildConfig.current
     stmt = ir_pass.UnrollLoop(
         stmt,
@@ -210,6 +209,7 @@ def lower(sch,
     stmt = ir_pass.Simplify(stmt)
     if simple_mode:
         return stmt
+    stmt = ir_pass.LowerStorageAccessInfo(stmt)
     return ir_pass.MakeAPI(stmt, name, arg_list, 0, cfg.restricted_func)
 
 

diff --git a/src/api/api_pass.cc b/src/api/api_pass.cc
@@ -95,6 +95,7 @@ REGISTER_PASS2(BindDeviceType);
 REGISTER_PASS1(SplitHostDevice);
 REGISTER_PASS1(StorageRewrite);
 REGISTER_PASS1(CoProcSync);
+REGISTER_PASS1(LowerStorageAccessInfo);
 REGISTER_PASS1(InjectVirtualThread);
 REGISTER_PASS1(InjectPrefetch);
 REGISTER_PASS1(LoopPartition);