Multiple matmul fusion example

csrc/mma_type.cpp

@@ -142,4 +142,19 @@ size_t hash(const MatMulTileOptions& opts) {
       (hash(opts.cta_tile) << 2);
 }
 
+std::string toString(const MatmulDimRole role) {
+  switch (role) {
+    case MatmulDimRole::Batch:
+      return "Batch";
+    case MatmulDimRole::M:
+      return "M";
+    case MatmulDimRole::N:
+      return "N";
+    case MatmulDimRole::K:
+      return "K";
+  }
+  // Unreachable
+  return "Unrecognized role";
+}
+
 } // namespace nvfuser

csrc/mma_type.h

@@ -28,6 +28,8 @@ constexpr std::string_view MATMUL_LOG_PREFIX = "[MATMUL DEBUG] ";
 //! Named descriptors of domains in matmul
 enum class MatmulDimRole { M = 0, N, K, Batch };
 
+std::string toString(MatmulDimRole role);
+
 //! Named descriptors of TensorView roles in fusion
 //!  OPERAND_A - an input to the fusion that is a producer of a matmul "A" input
 //!  OPERAND_B - an input to the fusion that is a producer of a matmul "B" input

csrc/scheduler/mma_utils.cpp

@@ -281,6 +281,143 @@ void scheduleWarpTileWithReduction(TensorView* tv, MatMulTileOptions tile) {
   // [Kwo Mwo Nwo Mw Nw Mi Ni Ki]
 }
 
+void scheduleWarpTile(
+    TensorView* tv,
+    MatMulTileOptions tile,
+    const std::vector<MatmulDimRole>& merged_roles) {
+  // Tiles are always given in [M, N, K]
+  auto cta_tile = tile.cta_tile;
+  auto warp_tile = tile.warp_tile;
+  auto instruction_tile = tile.instruction_tile;
+
+  // Do not split K dimension of CTA tile into multiple warp tiles
+  NVF_CHECK(
+      cta_tile.k == warp_tile.k,
+      "CTA tile and warp tile must have same K dimension");
+
+  // Find _inner_ tile dimensions present in tv
+  std::vector<int64_t> inner_dims;
+  std::vector<int64_t> required_inner_dim_sizes;
+  int64_t m_dim = -1, n_dim = -1, k_dim = -1;
+  int64_t min_inner_pos = tv->nDims(); // position of outermost inner dim
+  for (size_t i : c10::irange(merged_roles.size())) {
+    int64_t i_tv = tv->nDims() - 1 - (int64_t)i;
+    int64_t i_roles = (int64_t)merged_roles.size() - 1 - (int64_t)i;
+    if (i_tv < 0 || i_roles < 0) {
+      break;
+    }
+    inner_dims.push_back(i_tv);
+    MatmulDimRole role = merged_roles.at(i_roles);
+    switch (role) {
+      case MatmulDimRole::M:
+        m_dim = i_tv;
+        required_inner_dim_sizes.push_back(cta_tile.m);
+        break;
+      case MatmulDimRole::N:
+        n_dim = i_tv;
+        required_inner_dim_sizes.push_back(cta_tile.n);
+        break;
+      case MatmulDimRole::K:
+        k_dim = i_tv;
+        required_inner_dim_sizes.push_back(cta_tile.k);
+        break;
+      default:
+        continue;
+        break;
+    }
+    min_inner_pos = std::min(min_inner_pos, i_tv);
+  }
+
+  NVF_ERROR(
+      k_dim != -1,
+      "scheduleWarpTile should only be called on an operand with K dimension but found ",
+      tv->toString());
+
+  mma_utils::checkDimSize(tv, inner_dims, required_inner_dim_sizes);
+
+  //     m_dim  n_dim k_dim
+  //[...     M,     N,    K]
+  // Distribute warp tile:
+  if (m_dim != -1) {
+    tv->split(m_dim, warp_tile.m);
+    if (n_dim > m_dim) {
+      n_dim++;
+    }
+    if (k_dim > m_dim) {
+      k_dim++;
+    }
+  }
+  if (n_dim != -1) {
+    tv->split(n_dim, warp_tile.n);
+    if (m_dim > n_dim) {
+      m_dim++;
+    }
+    if (k_dim > m_dim) {
+      k_dim++;
+    }
+  }
+
+  //  m_dim  m_dim+1  n_dim  n_dim+1  k_dim
+  // [  Mwo       Mw    Nwo       Nw      K]
+  if (m_dim != -1) {
+    tv->split(m_dim + 1, instruction_tile.m);
+    if (n_dim > m_dim) {
+      n_dim++;
+    }
+    if (k_dim > m_dim) {
+      k_dim++;
+    }
+  }
+  if (n_dim != -1) {
+    tv->split(n_dim + 1, instruction_tile.n);
+    if (m_dim > n_dim) {
+      m_dim++;
+    }
+    if (k_dim > m_dim) {
+      k_dim++;
+    }
+  }
+  if (k_dim != -1) {
+    tv->split(k_dim, instruction_tile.k);
+    if (m_dim > k_dim) {
+      m_dim++;
+    }
+    if (n_dim > k_dim) {
+      n_dim++;
+    }
+  }
+
+  //  m_dim  m_dim+1  m_dim+2  n_dim  n_dim+1  n_dim+2  k_dim  k_dim+1
+  // [  Mwo       Mw       Mi    Nwo       Nw       Ni    Kwo       Ki]
+
+  // Nowwe reorder. The new_order will be
+  // [Kwo Mwo Nwo Mw Nw Mi Ni Ki]
+
+  std::vector<IterDomain*> new_loop;
+  new_loop.reserve(tv->nDims());
+  // Add all the outer dimensions as-is
+  for (size_t i : c10::irange(min_inner_pos)) {
+    new_loop.push_back(tv->axis((int64_t)i));
+  }
+  if (k_dim != -1) {
+    new_loop.push_back(tv->axis(k_dim));
+  }
+  for (size_t j : c10::irange(3)) {
+    if (m_dim != -1) {
+      new_loop.push_back(tv->axis(m_dim + (int64_t)j));
+    }
+    if (n_dim != -1) {
+      new_loop.push_back(tv->axis(n_dim + (int64_t)j));
+    }
+  }
+  if (k_dim != -1) {
+    new_loop.push_back(tv->axis(k_dim + 1));
+  }
+  NVF_ERROR(new_loop.size() == tv->nDims());
+
+  tv->setLoopDomain(new_loop);
+}
+
 void scheduleWarpTileWithNoReduction(TensorView* tv, MatMulTileOptions tile) {
   // Assumes
   // [M, N, K]
@@ -434,7 +571,7 @@ void makeTile(
   abten.reorder(reorder_map_old_to_new);
 }
 
-void makeTile(TensorView* tv, std::vector<int64_t> tile_sizes) {
+void makeTile(TensorView* tv, const std::vector<int64_t>& tile_sizes) {
   // We will create an AbstractMatmulTensor so that we can use the abstract
   // makeTile implementation above.
 
@@ -2145,4 +2282,22 @@ std::optional<std::pair<DimRolesMap, TensorRolesMap>> allPatternRoles(
 
 } // namespace mma_utils
 
+std::string toString(const mma_utils::AbstractMatmulTensor& abten) {
+  std::ostringstream ss;
+  ss << "AbstractMatmulTensor (" << abten.size() << "):" << std::endl;
+  for (size_t i : c10::irange(abten.size())) {
+    const AbstractId& abs_id = abten[(int64_t)i];
+    const std::optional<MatmulDimRole> role = abten.getTag((int64_t)i).value();
+    ss << "  " << (role.has_value() ? toString(role.value()) : "no role");
+    if (abs_id.is<ValGroupAndItsGraph>()) {
+      const ValGroup& g = abs_id.as<ValGroupAndItsGraph>().group;
+      for (Val* v : g->vector()) {
+        ss << " " << v->toString();
+      }
+    }
+    ss << std::endl;
+  }
+  return ss.str();
+}
+
 } // namespace nvfuser

csrc/scheduler/mma_utils.h

@@ -49,6 +49,15 @@ NVF_API void scheduleWarpTileWithReduction(
     TensorView* tv,
     MatMulTileOptions tile);
 
+//! Same as above, but takes a vector of dim roles that existed in tv before
+//! splitting/reordering to create block tiles.
+//! Note that if the input has a dimension ordering other than M/N K, it will
+//! be reordered to have M/N K inner dimensions in this utility.
+NVF_API void scheduleWarpTile(
+    TensorView* tv,
+    MatMulTileOptions tile,
+    const std::vector<MatmulDimRole>& merged_dim_roles);
+
 //! Schedule utility for mma output in matmul main loop:
 //!  Realize the hierarchical tiling based on the given tiling options
 //! on consumers of mma ops in epilog.
@@ -61,10 +70,10 @@ NVF_API void scheduleWarpTileWithNoReduction(
 //! Eg.
 //!  A[B,I0,I1,I2] -> makeTile({1,2,3})
 //! Gives A[B, I0o, I1o, I2o, I0i(1), I1i(2), I2i(3)]
-void makeTile(TensorView* tv, std::vector<int64_t> tile_sizes);
+void makeTile(TensorView* tv, const std::vector<int64_t>& tile_sizes);
 
-//! The above call assumes the axes in TV are [(B), M, N, K]. In this version,
-//! we provide the dimension roles that are present for this tensor.
+//! The above call assumes the axes are [(B), M, N, K]. In this version, we
+//! provide the dimension roles that are present for this tensor.
 void makeTile(
     TensorView* tv,
     const GemmTile& tile_sizes,
@@ -465,4 +474,6 @@ std::optional<std::pair<DimRolesMap, TensorRolesMap>> allPatternRoles(
 
 } // namespace mma_utils
 
+std::string toString(const mma_utils::AbstractMatmulTensor& abten);
+
 } // namespace nvfuser