clean up LowerWarpmemory

masahi · May 17, 2022 · 754c83e · 754c83e
1 parent 178c3dc
commit 754c83e
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 42 deletions.
diff --git a/3rdparty/dmlc-core b/3rdparty/dmlc-core
diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc
@@ -27,7 +27,6 @@
 #include <tvm/runtime/registry.h>
 #include <tvm/tir/stmt_functor.h>
 #include <tvm/tir/index_map.h>
-#include <tvm/arith/iter_affine_map.h>
 
 #include <algorithm>
 #include <cmath>

diff --git a/src/tir/transforms/lower_warp_memory.cc b/src/tir/transforms/lower_warp_memory.cc
@@ -117,9 +117,10 @@ class WarpStoreCoeffFinder : private StmtExprVisitor {
     if (op->op.same_as(builtin::ptx_ldmatrix()) && op->args[3].as<VarNode>() == buffer_) {
       UpdatePattern(op->args[4]);
     } else if (op->op.same_as(builtin::mma_fill()) && op->args[1].as<VarNode>() == buffer_) {
-      auto* ptr = op->args[0].as<IntImmNode>();
-      CHECK(ptr);
-      warp_coeff_ = ptr->value;;
+      auto* local_size = op->args[0].as<IntImmNode>();
+      ICHECK(local_size) << "Integer expected for the first argument of mma_fill";
+      warp_coeff_ = local_size->value;
+      ;
     }
 
     StmtExprVisitor::VisitExpr_(op);
@@ -257,54 +258,32 @@ class WarpAccessRewriter : protected StmtExprMutator {
   }
 
  protected:
+  PrimExpr RewriteIndicesAt(const CallNode* op, const std::vector<int>& indices) {
+    Array<PrimExpr> new_args = op->args;
+    for (int i : indices) {
+      if (op->args[i].get() == buffer_) {
+        PrimExpr local_index = SplitIndexByGroup(op->args[i + 1]).first;
+        new_args.Set(i + 1, local_index);
+      }
+    }
+    return Call(op->dtype, op->op, new_args);
+  }
+
   PrimExpr VisitExpr_(const CallNode* op) override {
     if (op->op.same_as(builtin::ptx_mma())) {
-      Array<PrimExpr> new_args = op->args;
-      PrimExpr local_index, group;
-      bool changed = false;
-      int A_warp_arg_ind = 6;
-      for (int i = A_warp_arg_ind; i < A_warp_arg_ind + 6; i += 2) {
-        if (op->args[i].get() == buffer_) {
-          std::tie(local_index, group) = SplitIndexByGroup(op->args[i + 1]);
-          new_args.Set(i + 1, local_index);
-          changed = true;
-        }
-      }
-      if (!changed) return GetRef<PrimExpr>(op);
-      return Call(op->dtype, op->op, new_args);
+      return RewriteIndicesAt(op, {6, 8, 10});
     }
 
     if (op->op.same_as(builtin::ptx_ldmatrix())) {
-      Array<PrimExpr> new_args = op->args;
-      PrimExpr local_index, group;
-      if (op->args[3].get() == buffer_) {
-        std::tie(local_index, group) = SplitIndexByGroup(op->args[4]);
-        new_args.Set(4, local_index);
-        return Call(op->dtype, op->op, new_args);
-      }
-      return GetRef<PrimExpr>(op);
+      return RewriteIndicesAt(op, {3});
     }
 
     if (op->op.same_as(builtin::mma_store())) {
-      Array<PrimExpr> new_args = op->args;
-      PrimExpr local_offset, group;
-      if (op->args[3].get() == buffer_) {
-        std::tie(local_offset, group) = SplitIndexByGroup(op->args[4]);
-        new_args.Set(4, local_offset);
-        return Call(op->dtype, op->op, new_args);
-      }
-      return GetRef<PrimExpr>(op);
+      return RewriteIndicesAt(op, {3});
     }
 
     if (op->op.same_as(builtin::mma_fill())) {
-      Array<PrimExpr> new_args = op->args;
-      PrimExpr local_offset, group;
-      if (op->args[1].get() == buffer_) {
-        std::tie(local_offset, group) = SplitIndexByGroup(op->args[2]);
-        new_args.Set(2, local_offset);
-        return Call(op->dtype, op->op, new_args);
-      }
-      return GetRef<PrimExpr>(op);
+      return RewriteIndicesAt(op, {1});
     }
 
     return StmtExprMutator::VisitExpr_(op);
+3 −14		.github/workflows/githubci.yml
+1 −17		CMakeLists.txt
+0 −2		cmake/build_config.h.in
+8 −1		doc/Doxyfile
+8 −10		include/dmlc/base.h
+3 −9		include/dmlc/build_config_default.h
+1 −1		include/dmlc/json.h
+15 −160		include/dmlc/optional.h
+2 −2		include/dmlc/parameter.h
+0 −11		scripts/conda_env.yml
+6 −7		scripts/lint.py
+11 −15		scripts/test_script.sh
+0 −25		src/data.cc
+0 −183		src/data/parquet_parser.h
+1 −1		src/data/text_parser.h
+1 −2		test/.gitignore
+2 −0		test/dmlc_test.mk
+13 −0		test/unittest/dmlc_unittest.mk
+0 −52		test/unittest/unittest_optional.cc
+0 −120		test/unittest/unittest_parquet_parser.cc
+2 −3		test/unittest/unittest_threaditer.cc
+0 −6		test/unittest/unittest_threaditer.h
+9 −2		test/unittest/unittest_threaditer_exc_handling.cc
+1 −1		tracker/dmlc_tracker/sge.py