address comment

Author: Siyuan Feng

include/tvm/ir.h

@@ -1329,7 +1329,6 @@ inline bool IsPragmaKey(const std::string& attr_key) {
   return attr_key.compare(0, 7, "pragma_") == 0;
 }
 
-
 }  // namespace attr
 
 /*! \brief namespace of TVM Intrinsic functions */
@@ -1564,11 +1563,52 @@ constexpr const char* tvm_global_barrier_kinit = "tvm_global_barrier_kinit";
  */
 constexpr const char* tvm_thread_allreduce = "tvm_thread_allreduce";
 /*!
- * \brief tvm intrinsic for tensor core opeartors.
+ * \brief tvm intrinsic for tensor core load operators.
+ *
+ *  void tvm_load_matrix_sync(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
+ *                            Expr index, Expr buffer_ptr, Expr stride,
+ *                            StringImm layout) {
+ *    // m, n, k are the shape of wmma fragment.
+ *    // Determine fragment layout(column-major or row major) by layout.
+ *    // fragments must be in 'wmma.matrix_a' or 'wmma.matrix_b' scope.
+ *    nvcuda::wmma::load_matrix_sync(fragment[index], buffer_ptr, stride);
+ *  }
  */
 constexpr const char* tvm_load_matrix_sync = "tvm_load_matrix_sync";
+/*!
+ * \brief tvm intrinsic for tensor core mma_sync operators.
+ *
+ *  void tvm_mma_sync(Var fragment_d, Expr index_d,
+ *                    Var fragment_a, Expr index_a,
+ *                    Var fragment_b, Expr index_b,
+ *                    Var fragment_c, Expr index_c) {
+ *    nvcuda::wmma::mma_sync(fragment_d[index_d], fragment_a[index_a],
+ *                           fragment_b[index_b], fragment_c[index_c]);
+ *  }
+ */
 constexpr const char* tvm_mma_sync = "tvm_mma_sync";
+/*!
+ * \brief tvm intrinsic for tensor core fill_fragment operators.
+ *
+ *  void tvm_fill_fragment(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
+ *                         Expr index, Expr value) {
+ *    // m, n, k are the shape of wmma fragment
+ *    // fragments must be in 'wmma.accumulator' scope.
+ *    nvcuda::wmma::fill_fragment(fragment[index], value);
+ *  }
+ */
 constexpr const char* tvm_fill_fragment = "tvm_fill_fragment";
+/*!
+ * \brief tvm intrinsic for tensor core store operators.
+ *
+ *  void tvm_store_matrix_sync(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
+ *                             Expr index, Expr buffer_ptr, Expr stride,
+ *                             StringImm layout) {
+ *    // m, n, k are the shape of wmma fragment
+ *    // fragments must be in 'wmma.accumulator' scope.
+ *    nvcuda::wmma::store_matrix_sync(fragment[index], buffer_ptr, stride, layout);
+ *  }
+ */
 constexpr const char* tvm_store_matrix_sync = "tvm_store_matrix_sync";
 
 }   // namespace intrinsic

include/tvm/ir_pass.h

@@ -359,6 +359,15 @@ Stmt LiftAttrScope(Stmt stmt, std::string attr_key);
  */
 Stmt RewriteUnsafeSelect(Stmt stmt);
 
+/*!
+ * \brief Lower attached storage access information.
+ * Do this pass after all storage access analysis finish.
+ *
+ * \param stmt The stmt to be transformed
+ * \return Transformed stmt.
+ */
+Stmt LowerStorageAccessInfo(Stmt stmt);
+
 /*!
  * \brief Decorate the stmt with a device scope, this is helpful for
  * hardware accelerator without thread blocks.
@@ -505,13 +514,13 @@ LoweredFunc CombineContextCall(LoweredFunc f);
 LoweredFunc PointerValueTypeRewrite(LoweredFunc f);
 
 /*!
- * \brief Lower attached storage access information.
+ * \brief Lower attached storage access information on device.
  * Do this pass after all storage access analysis finish.
  *
  * \param func The device function to be lowered.
  * \return Transformed function.
  */
-LoweredFunc LowerStorageAccessInfo(LoweredFunc func);
+LoweredFunc LowerDeviceStorageAccessInfo(LoweredFunc func);
 
 /*!
  * \brief Lower intrinsic function calls.

python/tvm/build_module.py

@@ -494,8 +494,8 @@ def _build_for_device(flist, target, target_host):
         assert not fdevice
 
     target_host = _target.create(target_host)
-    fdevice = [ir_pass.LowerStorageAccessInfo(x) for x in fdevice]
-    fhost = [ir_pass.LowerStorageAccessInfo(x) for x in fhost]
+    fdevice = [ir_pass.LowerDeviceStorageAccessInfo(x) for x in fdevice]
+    fhost = [ir_pass.LowerDeviceStorageAccessInfo(x) for x in fhost]
     fdevice = [ir_pass.LowerIntrin(x, target.target_name) for x in fdevice]
     fhost = [ir_pass.LowerIntrin(x, target_host.target_name) for x in fhost]
     fhost = [ir_pass.CombineContextCall(x) for x in fhost]

src/api/api_pass.cc

@@ -118,6 +118,14 @@ TVM_REGISTER_API("ir_pass.PostOrderVisit")
       });
   });
 
+TVM_REGISTER_API("ir_pass.LowerStorageAccess")
+.set_body([](TVMArgs args, TVMRetValue *ret) {
+  LoweredFunc f = args[0];
+  auto n = make_node<LoweredFuncNode>(*f.operator->());
+  n->body = LowerStorageAccessInfo(f->body);
+  *ret = LoweredFunc(n);
+});
+
 // make from two arguments
 #define REGISTER_PASS(PassName)                                   \
   TVM_REGISTER_API("ir_pass."#PassName)                           \
@@ -140,6 +148,7 @@ REGISTER_PASS(SplitHostDevice);
 REGISTER_PASS(StorageRewrite);
 REGISTER_PASS(CoProcSync);
 REGISTER_PASS(LowerStorageAccessInfo);
+REGISTER_PASS(LowerDeviceStorageAccessInfo)
 REGISTER_PASS(InjectVirtualThread);
 REGISTER_PASS(InjectPrefetch);
 REGISTER_PASS(InjectDoubleBuffer);

src/codegen/build_module.cc

@@ -517,15 +517,15 @@ Array<Array<LoweredFunc> > split_dev_host_funcs(const Array<LoweredFunc>& funcs,
   for (size_t i = 0; i < fhost.size(); ++i) {
     auto func = fhost[i];
     func = ir::BindDeviceType(func, target->device_type);
-    func = ir::LowerStorageAccessInfo(func);
+    func = ir::LowerDeviceStorageAccessInfo(func);
     func = ir::LowerTVMBuiltin(func);
     fhost.Set(i, func);
   }
 
   for (size_t i = 0; i < fhost.size(); ++i) {
     auto func = fhost[i];
     func = ir::LowerIntrin(func, target_host->target_name);
-    func = ir::LowerStorageAccessInfo(func);
+    func = ir::LowerDeviceStorageAccessInfo(func);
     func = ir::CombineContextCall(func);
     fhost.Set(i, func);
   }

src/pass/infer_fragment.cc

@@ -18,7 +18,8 @@
  */
 
 /*!
- *  Copyright (c) 2019 by Contributors
+ * Copyright (c) 2019 by Contributors
+ * \brief Infer TensorCore metadata from tensor intrinsic.
  * \file tensorcore_fragment.cc
  */
 #include <tvm/ir.h>
@@ -34,10 +35,14 @@
 namespace tvm {
 namespace ir {
 
+// Get fragment information from tensor intrinsics
 class FragmentGetter : public IRVisitor {
  public:
+  // fragment metadata
   struct FragmentInfo {
+    // fragment shape
     int m, n, k;
+    // fragment layout (row-major or column-major)
     std::string layout;
     FragmentInfo() = default;
     FragmentInfo(int _m, int _n, int _k, const std::string& _layout)
@@ -49,9 +54,11 @@ class FragmentGetter : public IRVisitor {
 
     if (op->is_intrinsic(intrinsic::tvm_load_matrix_sync) ||
         op->is_intrinsic(intrinsic::tvm_store_matrix_sync)) {
+      // Get shape and layout information from load and store intrinsic
       CHECK_EQ(op->args.size(), 8U);
       const Variable* buffer_var = op->args[0].as<Variable>();
       CHECK(buffer_var);
+      // Get shape
       const IntImm* m = op->args[1].as<IntImm>();
       const IntImm* n = op->args[2].as<IntImm>();
       const IntImm* k = op->args[3].as<IntImm>();
@@ -63,6 +70,7 @@ class FragmentGetter : public IRVisitor {
 
       std::string scope = scopes[buffer_var];
       if (fragments.count(buffer_var)) {
+        // check if the fragment has met before
         FragmentInfo info = fragments[buffer_var];
         CHECK_EQ(m->value, info.m);
         CHECK_EQ(n->value, info.n);
@@ -71,6 +79,7 @@ class FragmentGetter : public IRVisitor {
           CHECK_EQ(layout->value, info.layout);
         }
       } else {
+        // store metadata
         FragmentInfo info;
         if (scope == "wmma.matrix_a" || scope == "wmma.matrix_b") {
           info = FragmentInfo(m->value, n->value, k->value, layout->value);
@@ -80,9 +89,11 @@ class FragmentGetter : public IRVisitor {
         fragments[buffer_var] = info;
       }
     } else if (op->is_intrinsic(intrinsic::tvm_fill_fragment)) {
+      // Get shape information from fill intrinsic
       CHECK_EQ(op->args.size(), 6U);
       const Variable* buffer_var = op->args[0].as<Variable>();
       CHECK(buffer_var);
+      // Get shape
       const IntImm* m = op->args[1].as<IntImm>();
       const IntImm* n = op->args[2].as<IntImm>();
       const IntImm* k = op->args[3].as<IntImm>();
@@ -91,6 +102,7 @@ class FragmentGetter : public IRVisitor {
       CHECK(k);
 
       std::string scope = scopes[buffer_var];
+      // Only wmma.accumulator can use tvm_fill_fragment
       CHECK_EQ(scope, "wmma.accumulator");
       if (fragments.count(buffer_var)) {
         FragmentInfo info = fragments[buffer_var];
@@ -104,6 +116,7 @@ class FragmentGetter : public IRVisitor {
     }
   }
 
+  // Get memory scope
   void Visit_(const AttrStmt* op) final {
     if (op->attr_key == attr::storage_scope) {
       const Variable* buffer = op->node.as<Variable>();
@@ -113,15 +126,19 @@ class FragmentGetter : public IRVisitor {
     IRVisitor::Visit_(op);
   }
 
+  // Memory scope for allocations
   std::unordered_map<const Variable*, std::string> scopes;
+  // Fragment metadata for all fragments
   std::unordered_map<const Variable*, FragmentInfo> fragments;
 };
 
+// Check shape of fragment making sure it is a valid shape for tvm_mma_sync
 class FragmentChecker : public IRVisitor {
  public:
   explicit FragmentChecker(const FragmentGetter &getter) : fragment_getter(getter) {}
 
   void Visit_(const Call* op) final {
+    // Check shape when calling tvm_mma_sync
     if (op->is_intrinsic(intrinsic::tvm_mma_sync)) {
       CHECK_EQ(op->args.size(), 8U);
       const Variable* buffer_var_d = op->args[0].as<Variable>();
@@ -132,24 +149,28 @@ class FragmentChecker : public IRVisitor {
       CHECK(buffer_var_a);
       CHECK(buffer_var_b);
       CHECK(buffer_var_c);
+
+      // Check all fragment A, B, C and D have the same shape
       CHECK(CheckShape(buffer_var_d, buffer_var_a));
       CHECK(CheckShape(buffer_var_d, buffer_var_b));
       CHECK(CheckShape(buffer_var_d, buffer_var_c));
     }
   }
 
  private:
+  // A tool for checking shapes of two fragments
   bool CheckShape(const Variable* buffer1, const Variable* buffer2) {
     CHECK(fragment_getter.fragments.count(buffer1));
     CHECK(fragment_getter.fragments.count(buffer2));
     FragmentGetter::FragmentInfo info1 = fragment_getter.fragments.at(buffer1);
     FragmentGetter::FragmentInfo info2 = fragment_getter.fragments.at(buffer2);
     return info1.m == info2.m && info1.n == info2.n && info1.k == info2.k;
   }
-
+  // Fragment infomation
   const FragmentGetter &fragment_getter;
 };
 
+// Store the metadata into attributes
 class InferFragmenter : public IRMutator {
  public:
   explicit InferFragmenter(const FragmentGetter &getter) : fragment_getter(getter) {}
@@ -158,13 +179,17 @@ class InferFragmenter : public IRMutator {
     Stmt stmt = IRMutator::Mutate_(op, s);
     const Variable* buffer = op->buffer_var.get();
     if (fragment_getter.fragments.count(buffer)) {
+      // Add attribute to fragments allocation
       FragmentGetter::FragmentInfo info = fragment_getter.fragments.at(buffer);
+
+      // Add shape attribute to all fragments
       std::string shape = std::to_string(info.n) + ", " +
                           std::to_string(info.m) + ", " +
                           std::to_string(info.k);
       Expr shape_expr = StringImm::make(shape);
       Stmt shape_attr = AttrStmt::make(op->buffer_var, attr::fragment_shape, shape_expr, stmt);
       if (info.layout != "") {
+        // Add shape attribute to matrix_a and matrix_b
         Stmt layout_attr = AttrStmt::make(op->buffer_var, attr::fragment_layout,
                                           StringImm::make(info.layout), shape_attr);
         return layout_attr;
@@ -176,6 +201,7 @@ class InferFragmenter : public IRMutator {
   }
 
  private:
+  // Fragment infomation
   const FragmentGetter &fragment_getter;
 };
 

src/pass/storage_access.cc

@@ -341,7 +341,7 @@ Stmt LowerStorageAccessInfo(Stmt stmt) {
   return StorageAccessInfoLower().Mutate(stmt);
 }
 
-LoweredFunc LowerStorageAccessInfo(LoweredFunc f) {
+LoweredFunc LowerDeviceStorageAccessInfo(LoweredFunc f) {
   auto n = make_node<LoweredFuncNode>(*f.operator->());
   n->body = LowerStorageAccessInfo(f->body);
   return LoweredFunc(n);

tests/python/unittest/test_schedule_tensor_core.py

@@ -191,6 +191,7 @@ def test_tensor_core_batch_matmal():
     s[BF].tensorize(BF.op.axis[-2], intrin_wmma_load_matrix('wmma.matrix_b'))
     s[C].tensorize(kernel_i, intrin_wmma_store_matrix())
     s[CF].tensorize(_i, intrin_wmma_gemm())
+
     func = tvm.build(s, [A, B, C], 'cuda')
 
     ctx = tvm.gpu(0)

tests/scripts/task_lint.sh

@@ -30,19 +30,19 @@ trap cleanup 0
 echo "Check file types..."
 python3 tests/lint/check_file_type.py
 
-#echo "Check ASF license header..."
-#java -jar /bin/apache-rat.jar -E tests/lint/rat-excludes  -d . | (grep "== File" > /tmp/$$.apache-rat.txt || true)
-#if grep --quiet -E "File" /tmp/$$.apache-rat.txt; then
-#    echo "Need to add ASF header to the following files."
-#    echo "----------------File List----------------"
-#    cat /tmp/$$.apache-rat.txt
-#    echo "-----------------------------------------"
-#    echo "Use the following steps to add the headers:"
-#    echo "- Create file_list.txt in your text editor"
-#    echo "- Copy paste the above content in file-list into file_list.txt"
-#    echo "- python3 tests/lint/add_asf_header.py file_list.txt"
-#    exit 1
-#fi
+echo "Check ASF license header..."
+java -jar /bin/apache-rat.jar -E tests/lint/rat-excludes  -d . | (grep "== File" > /tmp/$$.apache-rat.txt || true)
+if grep --quiet -E "File" /tmp/$$.apache-rat.txt; then
+    echo "Need to add ASF header to the following files."
+    echo "----------------File List----------------"
+    cat /tmp/$$.apache-rat.txt
+    echo "-----------------------------------------"
+    echo "Use the following steps to add the headers:"
+    echo "- Create file_list.txt in your text editor"
+    echo "- Copy paste the above content in file-list into file_list.txt"
+    echo "- python3 tests/lint/add_asf_header.py file_list.txt"
+    exit 1
+fi
 
 echo "Check codestyle of c++ code..."
 make cpplint

tutorials/optimize/opt_conv_tensorcore.py

@@ -55,6 +55,12 @@
 import numpy as np
 from tvm.contrib import nvcc
 
+# Use nvcc compiler for better perf
+@tvm.register_func
+def tvm_callback_cuda_compile(code):
+    ptx =  nvcc.compile_cuda(code, target="ptx")
+    return ptx
+
 # The sizes of inputs and filters
 batch_size = 256
 height = 14
@@ -250,10 +256,10 @@ def intrin_func(ins, outs):
 
 
 # Define tiling sizes
-block_row_warps = 2
-block_col_warps = 4
-warp_row_tiles = 4
-warp_col_tiles = 2
+block_row_warps = 4
+block_col_warps = 2
+warp_row_tiles = 2
+warp_col_tiles = 4
 warp_size = 32
 chunk = 2
 
@@ -333,7 +339,8 @@ def intrin_func(ins, outs):
 
 ctx = tvm.gpu(0)
 if nvcc.have_tensorcore(ctx.compute_version):
-    func = tvm.build(s, [A, W, Conv], 'cuda')
+    with tvm.build_config(auto_unroll_max_step=16):
+        func = tvm.build(s, [A, W, Conv], 'cuda')
     a_np = np.random.uniform(size=data_shape).astype(A.dtype)
     w_np = np.random.uniform(size=kernel_shape).astype(W.dtype)
     a = tvm.nd.array(a_np, ctx)