calculate mma store dst index using inverse affine map

Author: masahi

src/target/source/codegen_c.cc

@@ -460,6 +460,13 @@ void CodeGenC::VisitExpr_(const MulNode* op, std::ostream& os) {  // NOLINT(*)
 void CodeGenC::VisitExpr_(const DivNode* op, std::ostream& os) {  // NOLINT(*)
   PrintBinaryExpr(op, "/", os, this);
 }
+void CodeGenC::VisitExpr_(const FloorDivNode* op, std::ostream& os) {  // NOLINT(*)
+  PrintBinaryExpr(op, "/", os, this);
+}
+void CodeGenC::VisitExpr_(const FloorModNode* op, std::ostream& os) {  // NOLINT(*)
+  PrintBinaryExpr(op, "%", os, this);
+}
+
 void CodeGenC::VisitExpr_(const ModNode* op, std::ostream& os) {  // NOLINT(*)
   if (op->dtype.is_int() || op->dtype.is_uint()) {
     PrintBinaryExpr(op, "%", os, this);

src/target/source/codegen_c.h

@@ -134,6 +134,8 @@ class CodeGenC : public ExprFunctor<void(const PrimExpr&, std::ostream&)>,
   void VisitExpr_(const MulNode* op, std::ostream& os) override;         // NOLINT(*)
   void VisitExpr_(const DivNode* op, std::ostream& os) override;         // NOLINT(*)
   void VisitExpr_(const ModNode* op, std::ostream& os) override;         // NOLINT(*)
+  void VisitExpr_(const FloorModNode* op, std::ostream& os) override;    // NOLINT(*)
+  void VisitExpr_(const FloorDivNode* op, std::ostream& os) override;    // NOLINT(*)
   void VisitExpr_(const MinNode* op, std::ostream& os) override;         // NOLINT(*)
   void VisitExpr_(const MaxNode* op, std::ostream& os) override;         // NOLINT(*)
   void VisitExpr_(const EQNode* op, std::ostream& os) override;          // NOLINT(*)

src/target/source/codegen_cuda.cc

@@ -34,6 +34,7 @@
 
 #include "literal/cuda_half_t.h"
 #include "ptx.h"
+#include "tvm/arith/iter_affine_map.h"
 
 namespace tvm {
 namespace codegen {
@@ -838,21 +839,38 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) {
     std::string dst = this->PrintExpr(op->args[2]);
     std::string src = this->PrintExpr(op->args[3]);
     std::string src_offset = this->PrintExpr(op->args[4]);
-    std::string stride = this->PrintExpr(op->args[5]);
 
     if (m == 16 && n == 8) {
+      std::string stride = this->PrintExpr(op->args[5]);
       os << "for (int i = 0; i < 4; ++i) {\n";
       os << dst << "[(i / 2 * 8 + threadIdx.x / 4) * " << stride
          << " + (threadIdx.x % 4) * 2 + i % 2]"
          << " = " << src << "[" << src_offset << " + i];\n";
       os << "}\n";
     } else if (m == 16 && n == 16) {
-      os << "for (int outer = 0; outer < 2; ++outer) {\n";
-      os << "for (int i = 0; i < 4; ++i) {\n";
-      os << dst << "[(i / 2 * 8 + threadIdx.x / 4) * " << stride
-         << " + outer * 8 + (threadIdx.x % 4) * 2 + i % 2]"
-         << " = " << src << "[" << src_offset << " + outer * 4 + i];\n";
-      os << "}\n";
+      const auto* index_map =
+          runtime::Registry::Get("tir.index_map.shared_16x16_to_ldmatrix_32x8_layout");
+      ICHECK(index_map);
+
+      Var var_i("i");
+      Var var_j("j");
+      Array<PrimExpr> forward_map = (*index_map)(var_i, var_j);
+
+      arith::Analyzer ana;
+      auto iter_map = arith::DetectIterMap(
+          forward_map, {{var_i, Range(0, 16)}, {var_j, Range(0, 16)}}, true, true, &ana, true);
+
+      Var thread_id("threadIdx.x");
+      Var local_id("local_id");
+      auto inverse_map = arith::InverseAffineIterMap(iter_map, {thread_id, local_id});
+      PrimExpr stride = op->args[5];
+      auto dst_idx = inverse_map[var_i] * stride + inverse_map[var_j];
+
+      var_idmap_[thread_id.get()] = "threadIdx.x";
+      var_idmap_[local_id.get()] = "local_id";
+      os << "for (int local_id = 0; local_id < 8; ++local_id) {\n";
+      os << dst << "[" + this->PrintExpr(dst_idx) + "]"
+         << " = " << src << "[" << src_offset << " + local_id];\n";
       os << "}\n";
     }
   } else if (op->op.same_as(builtin::mma_fill())) {

tests/python/unittest/test_mma_16x8x16_4k_tune.py

@@ -13,6 +13,12 @@ def shared_16x16_to_ldmatrix_32x8_layout(i, j):
     return thread_id, 4 * (j // 8) + (i // 8) * 2 + (j % 2)
 
 
+@tvm._ffi.register_func("tir.index_map.shared_16x16_to_ldmatrix_32x8_layout")
+def index_map_shared_16x16_to_ldmatrix_32x8_layout(i, j):
+    thread_id, local_id = shared_16x16_to_ldmatrix_32x8_layout(i, j)
+    return tvm.runtime.convert([thread_id, local_id])
+
+
 @T.prim_func
 def ldmatrix_a_desc(a: T.handle, c: T.handle) -> None:
     A_shared = T.match_buffer(a, (16, 16), "float16", align=128, offset_factor=16, scope="shared")