From 96225e800e0df7efd2b505aaaa3575d19c4545c9 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 6 Dec 2023 13:26:58 +0000
Subject: [PATCH 01/66] Start lowering serial grid reduction

---
 csrc/device_lower/lower2device.cpp |   4 +
 csrc/device_lower/pass/index.cpp   | 126 +++++++++++++++++++++++++++++
 csrc/device_lower/pass/index.h     |   3 +
 3 files changed, 133 insertions(+)

diff --git a/csrc/device_lower/lower2device.cpp b/csrc/device_lower/lower2device.cpp
index 89b79c60d46..fe86df5f209 100644
--- a/csrc/device_lower/lower2device.cpp
+++ b/csrc/device_lower/lower2device.cpp
@@ -265,13 +265,17 @@ GpuLower::GpuLower(Fusion* fusion, const CompileParams& cparams)
            {"loadStoreOpInserter", loadStoreOpInserter},
            {"insertAllocations", insertAllocations},
            {"insertRawThreadSynchronization", insertRawThreadSynchronization},
+           // NOTE: the smem/register reuse pass occurs here
            {"reuseMemoryAllocations", reuseMemoryAllocations},
            {"insertWarThreadSynchronization", insertWarThreadSynchronization},
            {"DoubleBufferPass", DoubleBufferPass::run},
            {"rotateLoops", rotateLoops},
            {"UnrollPass", UnrollPass::runPass},
            {"processMisalignedVectorization", processMisalignedVectorization},
+           // NOTE: serial GridReduction introduced here. New syncs can be
+           // introduced here which could impact smem reuse
            {"IndexLowering", IndexLowering::getIndexedExprs},
+           // NOTE: global reuse can be analyzed after this point
            {"fuseWarpReduce", fuseWarpReduce},
            {"generateConditionalFromPredicate",
             generateConditionalFromPredicate},
diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index 685feb23f3a..7650a507ad0 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -611,6 +611,128 @@ void IndexLowering::handleBlockReduction(
   GpuLower::current()->propagateExprInfo(rop, back());
 }
 
+bool IndexLowering::handleSerialGridReduction(
+    const ReductionOp* rop,
+    Val* out,
+    Val* in) {
+  // We inspect the loop nest up to the point after which all outer loops are
+  // parallelized. This corresponds to the outer-most loop in the generated
+  // kernel code.
+  //
+  // Conditions for using serial grid reduction:
+  //  - All reduction dimensions are either unparallelized or parallelized as
+  //    BID, not TID or serial. Block and warp reductions could be allowed in
+  //    the future, but the current focus is on cases where all threads are
+  //    doing separate reductions simultaneously.
+  //  - rop is not an allreduce. Note that we could implement serial allreduce
+  //    but it would require inserting a separate grid sync after this outer
+  //    loop.
+  //  - There are no global loads (i.e. producer TVs in global memory) anywhere
+  //    in the loop nest. The reason for this restriction is that loads
+  //    introduce data dependencies which wind up serializing sequences of
+  //    loads/stores that can be disadvantageous. Consider that a non-serial
+  //    (default) gridReduce that is inlined with some producer global read
+  //    will enable all blocks to participate in the read in parallel, since
+  //    only the last block of the reduction segment is serialized, and only
+  //    for the reduction itself. In a serial reduction the entire loop is
+  //    serialized, meaning we will not allow all blocks to perform those reads
+  //    in parallel, which can be inefficient.
+  //    Global stores, on the other hand, are fine as they do not cause this
+  //    excessive serialization.
+  //    Note that in addition to global loads, we should also avoid inlining
+  //    expensive computation into these blocks, since as stated above the
+  //    whole block will be serialized. That should be left to the automatic
+  //    schedulers in cases where serial reductions will be used.
+  //  - There are no other reductions in this loop nest that are TID or BID
+  //    parallelized, unless they also satisfy the conditions above and their
+  //    reduction pattern matches this one. Otherwise our syncs will be
+  //    mismatched, and there is no good way to handle that yet.
+  // TODO: implement the conditions above
+  bool is_serial_reduce_candidate = std::none_of(
+      out_domain->leaf().begin(), out_domain->leaf().end(), [](IterDomain* id) {
+        return id->isReduction() &&
+            (id->isThreadDim() || !id->isParallelized());
+      });
+  if (is_serial_reduce_candidate) {
+    std::cout << "WARNING: Found serial grid reduce candidate: "
+              << rop->toString();
+  }
+  return false;
+
+  // If we do a grid reduction we can't have a reduction axis that is not bound
+  // to a grid or block dim.
+  NVF_ERROR(
+      std::none_of(
+          out_domain->leaf().begin(),
+          out_domain->leaf().end(),
+          [](IterDomain* id) {
+            return !id->isThread() && id->isReduction() &&
+                !id->extent()->isOneInt();
+          }),
+      "Found a reduction stage that has both a non-parallelized ",
+      "reduction and a grid reduction. This is not supported, ",
+      "please use rfactor to do the serialized reduction first, ",
+      "then the grid reduction. ",
+      rop->toString());
+
+  const bool is_persistent = rop->isAllreduce();
+  if (is_persistent) {
+    // TODO: enable persistent serial grid reduce yet
+    return false;
+  }
+
+  // TODO: allocate global work buffer TensorIndex
+  // The output index should look just like the reduction output, and be
+  // scheduled the same, except it should have no reduction axes.
+  // TODO: it might be simpler to just include reduction axes instead of
+  // stripping them here.
+
+  // auto work_buffer_index = IrBuilder::create<TensorIndex>(
+  //   work_buffer,
+  //   lowerSrcIndex(out, work_buffer));
+
+  auto sync_buffer_size =
+      getGridSyncBufferSize(out_domain, for_loops_, is_persistent);
+  auto sync_buffer = allocateUniqueBuffer(
+      sync_buffer_size, DataType::Int, true, out_tv, sync_buffer_map_);
+
+  // The thread predicate for GridReduction needs to be set
+  // separately from the main predicate. Do not combine them like
+  // other expressions.
+  const auto& thread_pred =
+      GpuLower::current()->threadPredMap().getPredicatedParallelTypes(out_tv);
+
+  auto serial_grid_reduction = IrBuilder::create<kir::GridReduction>(
+      rop->getReductionOpType(),
+      rop->init(),
+      out,
+      in,
+      work_buffer,
+      sync_buffer,
+      entrance_ind,
+      n_entrances,
+      is_persistent);
+
+  serial_grid_reduction =
+      serial_grid_reduction->withThreadPredicate(thread_pred);
+
+  if (rop->predicate()) {
+    serial_grid_reduction =
+        serial_grid_reduction->withPredicate(rop->predicate())
+            ->as<kir::GridReduction>();
+  }
+  if (rop->writePredicate()) {
+    serial_grid_reduction =
+        serial_grid_reduction->withWritePredicate(rop->writePredicate())
+            ->as<kir::GridReduction>();
+  }
+
+  pushBack(serial_grid_reduction);
+  GpuLower::current()->propagateExprInfo(rop, back());
+
+  // TODO: insert syncs here? or need to do in a separate pass afterward?
+}
+
 void IndexLowering::handleGridReduction(
     const ReductionOp* rop,
     Val* out,
@@ -620,6 +742,10 @@ void IndexLowering::handleGridReduction(
 
   NVF_ERROR(out_domain->hasGridReduction());
 
+  if (handleSerialGridReduction(rop, out, in)) {
+    return;
+  }
+
   // If we do a grid reduction we can't have a reduction axis that is not bound
   // to a grid or block dim.
   NVF_ERROR(
diff --git a/csrc/device_lower/pass/index.h b/csrc/device_lower/pass/index.h
index f8e5b0fddbc..e0ab6a24c0f 100644
--- a/csrc/device_lower/pass/index.h
+++ b/csrc/device_lower/pass/index.h
@@ -120,6 +120,9 @@ class IndexLowering : private OptOutConstDispatch {
 
   void handleBlockReduction(const ReductionOp* rop, Val* out, Val* in);
   void handleGridReduction(const ReductionOp* rop, Val* out, Val* in);
+  //! Called by handleGridReduction, this returns true if rop is lowered as a
+  //! serial grid reduction.
+  bool handleSerialGridReduction(const ReductionOp* rop, Val* out, Val* in);
 
   void handleBlockReduction(
       const GroupedReductionOp* rop,

From f9d2d01f915cdaec0b4afb79b23a19aa69f21046 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 8 Dec 2023 16:56:33 +0000
Subject: [PATCH 02/66] Add grid_serialization.{cpp,h}

---
 CMakeLists.txt                                |  1 +
 csrc/device_lower/pass/grid_serialization.cpp | 89 +++++++++++++++++++
 csrc/device_lower/pass/grid_serialization.h   | 27 ++++++
 3 files changed, 117 insertions(+)
 create mode 100644 csrc/device_lower/pass/grid_serialization.cpp
 create mode 100644 csrc/device_lower/pass/grid_serialization.h

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 577fa5abbe6..d60441b6b5c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -114,6 +114,7 @@ list(APPEND NVFUSER_SRCS
   ${NVFUSER_SRCS_DIR}/device_lower/pass/double_buffer.cpp
   ${NVFUSER_SRCS_DIR}/device_lower/pass/expr_sort.cpp
   ${NVFUSER_SRCS_DIR}/device_lower/pass/fusion_simplifier.cpp
+  ${NVFUSER_SRCS_DIR}/device_lower/pass/grid_serialization.cpp
   ${NVFUSER_SRCS_DIR}/device_lower/pass/index.cpp
   ${NVFUSER_SRCS_DIR}/device_lower/pass/scalar_hoist.cpp
   ${NVFUSER_SRCS_DIR}/device_lower/pass/insert_syncs.cpp
diff --git a/csrc/device_lower/pass/grid_serialization.cpp b/csrc/device_lower/pass/grid_serialization.cpp
new file mode 100644
index 00000000000..71fbc65fa61
--- /dev/null
+++ b/csrc/device_lower/pass/grid_serialization.cpp
@@ -0,0 +1,89 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/lower2device.h>
+#include <device_lower/pass/insert_syncs.h>
+#include <device_lower/utils.h>
+#include <dispatch.h>
+#include <instrumentation.h>
+#include <ir/utils.h>
+#include <kernel_ir.h>
+#include <kernel_ir_dispatch.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+//! Insert needed syncs in
+class GridSerializationSyncInserter : kir::ExprMutator {
+ public:
+  GridSerializationSyncInserter(const std::vector<Expr*>& exprs) {
+    kir::ExprMutator::traverseAndInsert(exprs);
+  }
+
+  static std::vector<Expr*> insert(const std::vector<Expr*>& exprs) {
+    GridSerializationSyncInserter inserter(exprs);
+    return inserter.exprs_;
+  }
+
+ private:
+  using kir::ExprMutator::handle;
+
+  void handle(ReductionOp* rop) final {
+    if (rop->serialGridReductionRequested()) {
+    }
+  }
+
+  //! Determine whether a ForLoop is a top-level (in generated kernel) loop
+  //! before handling.
+  void handle(kir::ForLoop* loop) final {
+  }
+
+ private:
+  //! Which Expr* is the current top-level containing the current Expr in the
+  //! generated kernel. When serial reductions are encountered, this expression
+  //! determines where we will place syncs: they will be placed before and after
+  //! this expression.
+  //!
+  //! For example, if we have
+  //!
+  //!   FOR iBlockIdx.x
+  //!     FOR iS{32}
+  //!       y = neg(x);
+  //!     ENDFOR iS{32}
+  //!     FOR iThreadIdx.x
+  //!       z = add(y, x);
+  //!     ENDFOR iThreadIdx.x
+  //!   ENDFOR iBlockIdx.x
+  //!
+  //! then when we are processing the `neg` Expr, cur_top_level_expr_ will be
+  //! the FOR iS{32} loop. However, when processing the `add` expression,
+  //! cur_top_level_expr_ will be nullptr since that expression itself will
+  //! appear in the main scope of the generated kernel.
+  //!
+  //! IfThenElse are treated similar to unparallelized ForLoops; if an
+  //! IfThenElse is at the top level, or is contained in a fully parallelized
+  //! loop nest, it is treated as a top level expr here. Note that this pass
+  //! will likely be run before any IfThenElse are placed in the kernel anyway.
+  Expr* cur_top_level_expr_ = nullptr;
+
+  //! If a serial grid reduction is found for the current expr, 
+  std::optional<std::tuple<bool, bool, bool>> current_expr_sync_pattern_ =
+      std::nullopt;
+};
+
+} // namespace
+
+std::vector<Expr*> insertGridSerializationSyncs(
+    const std::vector<Expr*>& exprs) {
+  FUSER_PERF_SCOPE("GpuLower::Lower::insertGridSerializationSyncs");
+  return GridSerializationSyncInserter::insert(exprs);
+}
+
+} // namespace nvfuser
diff --git a/csrc/device_lower/pass/grid_serialization.h b/csrc/device_lower/pass/grid_serialization.h
new file mode 100644
index 00000000000..9a30b0dd0b5
--- /dev/null
+++ b/csrc/device_lower/pass/grid_serialization.h
@@ -0,0 +1,27 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <c10/macros/Export.h>
+#include <exceptions.h>
+
+#include <ir/all_nodes.h>
+#include <kernel_ir.h>
+
+#include <vector>
+
+namespace nvfuser {
+
+//! Detect ReductionOps that have serialGridReductionRequested() == true. When
+//! found, confirm that no conflicting operations exist, then place sync nodes
+//! before and after outer-most non-parallelized loop. Also allocate a new
+//! global buffer to serve as the work buffer.
+std::vector<Expr*> insertGridSerializationSyncs(
+    const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser

From 6454b062c18cfdc5652e8db666298d5a962d6243 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 8 Dec 2023 16:57:32 +0000
Subject: [PATCH 03/66] Add requestSerialGridReduction

---
 csrc/ir/internal_nodes.h | 16 ++++++++++++++++
 csrc/ir/nodes.cpp        |  1 +
 2 files changed, 17 insertions(+)

diff --git a/csrc/ir/internal_nodes.h b/csrc/ir/internal_nodes.h
index 089b2bf41f9..ae724ce4ae9 100644
--- a/csrc/ir/internal_nodes.h
+++ b/csrc/ir/internal_nodes.h
@@ -913,6 +913,22 @@ class ReductionOp : public Expr {
   bool isAllreduce() const {
     return attribute<bool>(2);
   }
+
+  //! Scheduling method to request that this reduction be performed as a
+  //! serial grid reduction. Note that it is an error to use this method on a
+  //! reduction whose output has any of its reduction axes parallelized with a
+  //! threadIdx, even if that parallelization occurs after this method call.
+  //!
+  //! Also note that this operation should not be inlined with other reductions
+  //! unless they use the same parallelization pattern and they are also serial
+  //! gridreductions.
+  void requestSerialGridReduction(bool value = true) {
+    attribute<bool>(3) = value;
+  }
+
+  bool serialGridReductionRequested() const {
+    return attribute<bool>(3);
+  }
 };
 
 //! Grouped reduction operation for horizontal fusions. It works like
diff --git a/csrc/ir/nodes.cpp b/csrc/ir/nodes.cpp
index 9c5b10a9810..26bf2b3d626 100644
--- a/csrc/ir/nodes.cpp
+++ b/csrc/ir/nodes.cpp
@@ -1429,6 +1429,7 @@ ReductionOp::ReductionOp(
   addAttribute(init);
   addDataAttribute(reduction_op_type);
   addDataAttribute(is_allreduce);
+  addDataAttribute(false); // serial reduction
 }
 
 std::string ReductionOp::toString(int indent_size) const {

From 1142e4402611c241dd963f636e677fabb9032af6 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 8 Dec 2023 16:58:01 +0000
Subject: [PATCH 04/66] Disable previous changes to indexing pass.

Will revisit once sync pass is done, when we have a TensorIndex
---
 csrc/device_lower/pass/index.cpp | 23 ++++++++++++++++-------
 csrc/device_lower/pass/index.h   |  2 +-
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index 7650a507ad0..88115e3fa50 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -611,7 +611,7 @@ void IndexLowering::handleBlockReduction(
   GpuLower::current()->propagateExprInfo(rop, back());
 }
 
-bool IndexLowering::handleSerialGridReduction(
+void IndexLowering::handleSerialGridReduction(
     const ReductionOp* rop,
     Val* out,
     Val* in) {
@@ -648,6 +648,14 @@ bool IndexLowering::handleSerialGridReduction(
   //    reduction pattern matches this one. Otherwise our syncs will be
   //    mismatched, and there is no good way to handle that yet.
   // TODO: implement the conditions above
+  // TODO: Move this comment above to the lowering pass where it's implemented
+
+   TODO: uncomment once this function works again by associating a temporary T
+  // nsorIndex
+  /*
+  const auto out_tv = out->as<kir::TensorIndex>()->view();
+  const auto out_domain = out_tv->domain();
+
   bool is_serial_reduce_candidate = std::none_of(
       out_domain->leaf().begin(), out_domain->leaf().end(), [](IterDomain* id) {
         return id->isReduction() &&
@@ -729,23 +737,23 @@ bool IndexLowering::handleSerialGridReduction(
 
   pushBack(serial_grid_reduction);
   GpuLower::current()->propagateExprInfo(rop, back());
-
-  // TODO: insert syncs here? or need to do in a separate pass afterward?
+  */
 }
 
 void IndexLowering::handleGridReduction(
     const ReductionOp* rop,
     Val* out,
     Val* in) {
+  if (rop->serialGridReductionRequested()) {
+    handleSerialGridReduction(rop, out, in);
+    return;
+  }
+
   const auto out_tv = out->as<kir::TensorIndex>()->view();
   const auto out_domain = out_tv->domain();
 
   NVF_ERROR(out_domain->hasGridReduction());
 
-  if (handleSerialGridReduction(rop, out, in)) {
-    return;
-  }
-
   // If we do a grid reduction we can't have a reduction axis that is not bound
   // to a grid or block dim.
   NVF_ERROR(
@@ -1927,3 +1935,4 @@ void IndexLowering::handle(const CatOp* cat) {
 }
 
 } // namespace nvfuser
+
\ No newline at end of file
diff --git a/csrc/device_lower/pass/index.h b/csrc/device_lower/pass/index.h
index e0ab6a24c0f..ea875c3e4e3 100644
--- a/csrc/device_lower/pass/index.h
+++ b/csrc/device_lower/pass/index.h
@@ -122,7 +122,7 @@ class IndexLowering : private OptOutConstDispatch {
   void handleGridReduction(const ReductionOp* rop, Val* out, Val* in);
   //! Called by handleGridReduction, this returns true if rop is lowered as a
   //! serial grid reduction.
-  bool handleSerialGridReduction(const ReductionOp* rop, Val* out, Val* in);
+  void handleSerialGridReduction(const ReductionOp* rop, Val* out, Val* in);
 
   void handleBlockReduction(
       const GroupedReductionOp* rop,

From 1368ba80b86a2e521cb789c55c336f9680a668ad Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 8 Dec 2023 16:58:33 +0000
Subject: [PATCH 05/66] Call insertGridSerializationSyncs pass

---
 csrc/device_lower/lower2device.cpp            |  3 ++-
 csrc/device_lower/pass/grid_serialization.cpp |  5 ++--
 csrc/device_lower/pass/index.cpp              |  5 ++--
 .../pass/serial_reduction_syncs.h             | 24 +++++++++++++++++++
 4 files changed, 31 insertions(+), 6 deletions(-)
 create mode 100644 csrc/device_lower/pass/serial_reduction_syncs.h

diff --git a/csrc/device_lower/lower2device.cpp b/csrc/device_lower/lower2device.cpp
index fe86df5f209..330655cb017 100644
--- a/csrc/device_lower/lower2device.cpp
+++ b/csrc/device_lower/lower2device.cpp
@@ -16,6 +16,7 @@
 #include <device_lower/pass/double_buffer.h>
 #include <device_lower/pass/expr_sort.h>
 #include <device_lower/pass/fusion_simplifier.h>
+#include <device_lower/pass/grid_serialization.h>
 #include <device_lower/pass/index.h>
 #include <device_lower/pass/inline_ptx.h>
 #include <device_lower/pass/insert_syncs.h>
@@ -263,9 +264,9 @@ GpuLower::GpuLower(Fusion* fusion, const CompileParams& cparams)
           // const std::vector<Expr*>& and return a std::vector<Expr*>.
           {{"LoopNestGenerator", LoopNestGenerator::loweredExprs},
            {"loadStoreOpInserter", loadStoreOpInserter},
+           {"insertGridSerializationSyncs", insertGridSerializationSyncs},
            {"insertAllocations", insertAllocations},
            {"insertRawThreadSynchronization", insertRawThreadSynchronization},
-           // NOTE: the smem/register reuse pass occurs here
            {"reuseMemoryAllocations", reuseMemoryAllocations},
            {"insertWarThreadSynchronization", insertWarThreadSynchronization},
            {"DoubleBufferPass", DoubleBufferPass::run},
diff --git a/csrc/device_lower/pass/grid_serialization.cpp b/csrc/device_lower/pass/grid_serialization.cpp
index 71fbc65fa61..97b2253256e 100644
--- a/csrc/device_lower/pass/grid_serialization.cpp
+++ b/csrc/device_lower/pass/grid_serialization.cpp
@@ -42,8 +42,7 @@ class GridSerializationSyncInserter : kir::ExprMutator {
 
   //! Determine whether a ForLoop is a top-level (in generated kernel) loop
   //! before handling.
-  void handle(kir::ForLoop* loop) final {
-  }
+  void handle(kir::ForLoop* loop) final {}
 
  private:
   //! Which Expr* is the current top-level containing the current Expr in the
@@ -73,7 +72,7 @@ class GridSerializationSyncInserter : kir::ExprMutator {
   //! will likely be run before any IfThenElse are placed in the kernel anyway.
   Expr* cur_top_level_expr_ = nullptr;
 
-  //! If a serial grid reduction is found for the current expr, 
+  //! If a serial grid reduction is found for the current expr,
   std::optional<std::tuple<bool, bool, bool>> current_expr_sync_pattern_ =
       std::nullopt;
 };
diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index 88115e3fa50..a6bdb22b172 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -650,7 +650,8 @@ void IndexLowering::handleSerialGridReduction(
   // TODO: implement the conditions above
   // TODO: Move this comment above to the lowering pass where it's implemented
 
-   TODO: uncomment once this function works again by associating a temporary T
+TODO:
+  uncomment once this function works again by associating a temporary T
   // nsorIndex
   /*
   const auto out_tv = out->as<kir::TensorIndex>()->view();
@@ -1935,4 +1936,4 @@ void IndexLowering::handle(const CatOp* cat) {
 }
 
 } // namespace nvfuser
-
\ No newline at end of file
+
diff --git a/csrc/device_lower/pass/serial_reduction_syncs.h b/csrc/device_lower/pass/serial_reduction_syncs.h
new file mode 100644
index 00000000000..060e0d4b55f
--- /dev/null
+++ b/csrc/device_lower/pass/serial_reduction_syncs.h
@@ -0,0 +1,24 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <c10/macros/Export.h>
+#include <exceptions.h>
+
+#include <ir/all_nodes.h>
+#include <kernel_ir.h>
+
+#include <vector>
+
+namespace nvfuser {
+
+//! Insert syncs
+std::vector<Expr*> insertSerialGridReductionSyncs(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser
+

From b97db45da277037cd66bf765d475e3a334af3105 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 8 Dec 2023 16:59:18 +0000
Subject: [PATCH 06/66] Remove file added by mistake

---
 .../pass/serial_reduction_syncs.h             | 24 -------------------
 1 file changed, 24 deletions(-)
 delete mode 100644 csrc/device_lower/pass/serial_reduction_syncs.h

diff --git a/csrc/device_lower/pass/serial_reduction_syncs.h b/csrc/device_lower/pass/serial_reduction_syncs.h
deleted file mode 100644
index 060e0d4b55f..00000000000
--- a/csrc/device_lower/pass/serial_reduction_syncs.h
+++ /dev/null
@@ -1,24 +0,0 @@
-// clang-format off
-/*
- * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
- * All rights reserved.
- * SPDX-License-Identifier: BSD-3-Clause
- */
-// clang-format on
-#pragma once
-
-#include <c10/macros/Export.h>
-#include <exceptions.h>
-
-#include <ir/all_nodes.h>
-#include <kernel_ir.h>
-
-#include <vector>
-
-namespace nvfuser {
-
-//! Insert syncs
-std::vector<Expr*> insertSerialGridReductionSyncs(const std::vector<Expr*>& exprs);
-
-} // namespace nvfuser
-

From ac2da9d8689c4014950098187c477387d767949f Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 8 Dec 2023 17:11:56 +0000
Subject: [PATCH 07/66] Fix formatting lintrunner messed up

---
 csrc/device_lower/pass/index.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index a6bdb22b172..5995cf367ab 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -650,9 +650,8 @@ void IndexLowering::handleSerialGridReduction(
   // TODO: implement the conditions above
   // TODO: Move this comment above to the lowering pass where it's implemented
 
-TODO:
-  uncomment once this function works again by associating a temporary T
-  // nsorIndex
+  // TODO: uncomment once this function works again by associating a temporary
+  // TensorIndex
   /*
   const auto out_tv = out->as<kir::TensorIndex>()->view();
   const auto out_domain = out_tv->domain();
@@ -1936,4 +1935,3 @@ void IndexLowering::handle(const CatOp* cat) {
 }
 
 } // namespace nvfuser
-

From ebef797c8ef2e48153b140b34c66285fc69cd33f Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 8 Dec 2023 17:22:11 +0000
Subject: [PATCH 08/66] Bump num_reduction_op_attr

---
 csrc/kernel_ir.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/csrc/kernel_ir.h b/csrc/kernel_ir.h
index 893d32dace5..b34f9ef1c74 100644
--- a/csrc/kernel_ir.h
+++ b/csrc/kernel_ir.h
@@ -948,7 +948,7 @@ class IfThenElse final : public Expr {
 //! This node provides FusionExecutor the information it needs to allocate the
 //! reduction and sync buffers.
 class GridReduction final : public ReductionOp {
-  static constexpr int num_reduction_op_attr = 3;
+  static constexpr int num_reduction_op_attr = 4;
 
  public:
   using ReductionOp::ReductionOp;

From dcd860676fce0cfe84e5ff55a8fa77b5ac17ced4 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 8 Dec 2023 18:30:20 +0000
Subject: [PATCH 09/66] Add test

---
 test/test_serial_gridreduce.cpp | 93 +++++++++++++++++++++++++++++++++
 1 file changed, 93 insertions(+)

diff --git a/test/test_serial_gridreduce.cpp b/test/test_serial_gridreduce.cpp
index e7fc291ebbe..43675299b68 100644
--- a/test/test_serial_gridreduce.cpp
+++ b/test/test_serial_gridreduce.cpp
@@ -256,4 +256,97 @@ TEST_F(SerialGridReductionTest, CodegenNodes) {
   }
 }
 
+TEST_F(SerialGridReductionTest, Scheduling) {
+  for (bool serial : {true, false}) {
+    for (int64_t num_warps : {4, 8}) {
+      // B is size of inner serial loop. Outer loop is hardcoded at A=4
+      // Here we set B to a small value of 8 instead of 32 (i.e. 128 elements
+      // per thread), so that the non-serial compilation does not take too
+      // long.
+      for (int64_t B : {8}) {
+        std::unique_ptr<Fusion> fusion_ptr = std::make_unique<Fusion>();
+        auto fusion = fusion_ptr.get();
+        FusionGuard fg(fusion);
+
+        int64_t blocks_x = 8;
+        int64_t blocks_y = 8;
+        int64_t blocks_z = 5;
+        int64_t A = 4; // Size of outer serial loop
+        int64_t H = blocks_z;
+        int64_t W = A * B * blocks_x * blocks_y * num_warps * 32;
+
+        // Unreduced dimensions should be concrete. Reduced dimension could be
+        // symbolic, but is concrete here so that we can read tv0 to registers
+        TensorView* tv0 = TensorViewBuilder()
+                              .shape({H, W})
+                              .dtype(DataType::Float)
+                              .contiguity(true)
+                              .build();
+        fusion->addInput(tv0);
+
+        auto tv1 = sum(tv0, {0});
+        fusion->addOutput(tv1);
+
+        // Start with
+        //   [ rS{H}, iS{W} ]
+        // We are grid reducing the H dimension and we want to coalesce
+        // accesses in the W dimension. So we first reorder to
+        //   [ iS{W}, rS{H} ]
+        // then schedule as
+        //   [ iBIDx{blocks_x}, iBIDy{blocks_y}, iS{A}, iS{B}, iTIDy{num_warps},
+        //   iTIDx{32}, rBIDz{blocks_z} ]
+        auto tv2 = tv0->cacheAfter();
+        auto tv3 = tv1->cacheBefore();
+
+        tv3->reorder({{1, 0}, {0, 1}}); // blocks_x*blocks_y*A*B*num_warps*32, H
+        tv3->split(0, 32); // blocks_x*blocks_y*A*B*num_warps, 32, H
+        tv3->split(0, num_warps); // blocks_x*blocks_y*A*B, num_warps, 32, H
+        tv3->split(0, B); // blocks_x*blocks_y*A, B, num_warps, 32, H
+        tv3->split(0, A); // blocks_x*blocks_y, A, B, num_warps, 32, H
+        tv3->split(0, blocks_y); // blocks_x, blocks_y, A, B, num_warps, 32, H
+        tv3->axis(0)->parallelize(ParallelType::BIDx);
+        tv3->axis(1)->parallelize(ParallelType::BIDy);
+        tv3->axis(4)->parallelize(ParallelType::TIDy);
+        tv3->axis(5)->parallelize(ParallelType::TIDx);
+        tv3->axis(6)->parallelize(ParallelType::BIDz);
+        // Reorder to put parallel dims first for better inlining
+        tv3->reorder({
+            {4, 2},
+            {5, 3},
+            {2, 4},
+            {3, 5},
+        });
+
+        TransformPropagator propagator(tv3);
+        MaxRootDomainInfoSpanningTree(tv3).traverse(&propagator);
+        scheduler_utils::parallelizeAllLike(tv3);
+
+        // Here we just transpose A and B in tv2, so that it will be partially
+        // inlined with tv3, resulting in a separate loop to load tv0 into
+        // registers (tv2).
+        tv2->reorder({
+            {-2, -3},
+            {-3, -2},
+        });
+
+        inlineMost();
+
+        FusionExecutor fe;
+        if (serial) {
+          tv3->definition()->as<ReductionOp>()->requestSerialGridReduction();
+        }
+        fe.compileFusion(fusion);
+
+        auto input = at::randn(
+            {H, W}, at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0));
+        auto outputs = fe.runFusion({input});
+
+        if (serial) {
+          testValidate(fusion, outputs, {input}, __LINE__, __FILE__);
+        }
+      }
+    }
+  }
+}
+
 } // namespace nvfuser

From e196fa75dc6e61d5fe89f4d83e205d110a041123 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 11 Dec 2023 17:42:08 +0000
Subject: [PATCH 10/66] Fix sync insertion in lowering pass.

Still missing allocation/indexing of work buffer
---
 csrc/device_lower/pass/grid_serialization.cpp | 85 +++++++++++++++++--
 csrc/device_lower/pass/insert_syncs.cpp       | 30 +------
 csrc/device_lower/utils.cpp                   | 22 +++++
 csrc/device_lower/utils.h                     |  8 +-
 4 files changed, 109 insertions(+), 36 deletions(-)

diff --git a/csrc/device_lower/pass/grid_serialization.cpp b/csrc/device_lower/pass/grid_serialization.cpp
index 97b2253256e..e32d3c9b898 100644
--- a/csrc/device_lower/pass/grid_serialization.cpp
+++ b/csrc/device_lower/pass/grid_serialization.cpp
@@ -33,16 +33,86 @@ class GridSerializationSyncInserter : kir::ExprMutator {
   }
 
  private:
+  using kir::ExprMutator::dispatch;
   using kir::ExprMutator::handle;
 
-  void handle(ReductionOp* rop) final {
+  //! Record cur_expr_sync_pattern_ if this is a serial grid reduction
+  void handle(ReductionOp* rop) override {
     if (rop->serialGridReductionRequested()) {
+      ParallelTypeBitmap sync_pattern;
+      auto out = rop->out()->as<TensorView>();
+      NVF_ERROR(out != nullptr);
+      for (int i : c10::irange((int)out->nDims())) {
+        IterDomain* ax = out->axis(i);
+        if (!ax->isReduction()) {
+          continue;
+        }
+        NVF_ERROR(
+            !ax->isThreadDim(),
+            "Serial grid reduction cannot be applied with block reductions: ",
+            rop->toString());
+        if (ax->isBlockDim()) {
+          sync_pattern.set(ax->getParallelType());
+        }
+      }
+
+      if (!sync_pattern.hasBID()) {
+        // Don't set cur_expr_sync_pattern_ since this is not actually a grid
+        // reduction
+        return;
+      }
+
+      if (cur_expr_sync_pattern_.has_value()) {
+        NVF_ERROR(
+            cur_expr_sync_pattern_.value() == sync_pattern,
+            "Reduction op ",
+            rop->toString(),
+            " has requested serial grid reduction, but pattern ",
+            sync_pattern.toString(),
+            " conflicts with previous pattern: ",
+            cur_expr_sync_pattern_.value().toString());
+      } else {
+        cur_expr_sync_pattern_ = sync_pattern;
+      }
     }
   }
 
-  //! Determine whether a ForLoop is a top-level (in generated kernel) loop
-  //! before handling.
-  void handle(kir::ForLoop* loop) final {}
+  void dispatch(Expr* expr) override {
+    if (auto loop = dynamic_cast<kir::ForLoop*>(expr);
+        cur_top_level_expr_ || (loop && loop->isTrivial())) {
+      // Never sync around trivial loops. Also avoid redefining
+      // cur_top_level_expr_ if it is already set
+      kir::ExprMutator::dispatch(expr);
+      return;
+    }
+    // Any other expr, i.e. non-trivial loops or regular Exprs, can be synced if
+    // it is top-level and either is or contains a serial grid reduction
+    cur_top_level_expr_ = expr;
+    kir::ExprMutator::dispatch(expr);
+    // If a serial grid reduction was found in expr, then cur_expr_sync_pattern_
+    // will be set
+    if (cur_expr_sync_pattern_.has_value()) {
+      insertSyncs();
+    }
+    // reset state variables
+    cur_top_level_expr_ = nullptr;
+    cur_expr_sync_pattern_ = std::nullopt;
+  }
+
+  void insertSyncs() {
+    NVF_ERROR(cur_top_level_expr_ != nullptr);
+    NVF_ERROR(cur_expr_sync_pattern_.has_value());
+    kir::Allocate* alloc = lower_utils::allocGlobalBufferForGridComm(
+        lower_utils::getGridSyncBufferSize(cur_expr_sync_pattern_.value()),
+        DataType::Int,
+        true);
+    auto wait = IrBuilder::create<kir::BlockSerializeWait>(
+        cur_expr_sync_pattern_.value(), alloc->buffer());
+    registerInsertBefore(cur_top_level_expr_, wait);
+    auto release = IrBuilder::create<kir::BlockSerializeRelease>(
+        cur_expr_sync_pattern_.value(), alloc->buffer());
+    registerInsertAfter(cur_top_level_expr_, release);
+  }
 
  private:
   //! Which Expr* is the current top-level containing the current Expr in the
@@ -72,9 +142,10 @@ class GridSerializationSyncInserter : kir::ExprMutator {
   //! will likely be run before any IfThenElse are placed in the kernel anyway.
   Expr* cur_top_level_expr_ = nullptr;
 
-  //! If a serial grid reduction is found for the current expr,
-  std::optional<std::tuple<bool, bool, bool>> current_expr_sync_pattern_ =
-      std::nullopt;
+  //! If a serial grid reduction is found for the current expr, this indicates
+  //! parallel axes that are mapped to reduction domains in the serial
+  //! reduction.
+  std::optional<ParallelTypeBitmap> cur_expr_sync_pattern_ = std::nullopt;
 };
 
 } // namespace
diff --git a/csrc/device_lower/pass/insert_syncs.cpp b/csrc/device_lower/pass/insert_syncs.cpp
index 5dc28e41ef5..866b73d8771 100644
--- a/csrc/device_lower/pass/insert_syncs.cpp
+++ b/csrc/device_lower/pass/insert_syncs.cpp
@@ -369,32 +369,6 @@ class ValidatePlacementAfterWrites : private kir::IrVisitor {
   const std::unordered_set<Expr*>& writes_;
 };
 
-namespace {
-
-Val* getGridSyncBufferSize(const ParallelTypeBitmap& ptb) {
-  // See the comment above for getGridCommWorkBufferSize.
-  NVF_ERROR(
-      ptb.hasBID(),
-      "Detected  needing a grid sync but no grid bits set in bitmap.");
-  Val* buffer_size = GpuLower::current()->kernel()->oneVal();
-  for (auto pt : kParallelTypeBIDs) {
-    // Synchronized within pt, so all blocks of this PT use the same
-    // sync buffer location, and thus no need to expand the sync
-    // buffer size.
-    if (ptb.get(pt)) {
-      continue;
-    }
-    auto pt_dim = GpuLower::current()->parallelDimensionMap().get(pt);
-    if (pt_dim == nullptr || pt_dim->isOneInt()) {
-      continue;
-    }
-    buffer_size = IrBuilder::mulExpr(buffer_size, pt_dim);
-  }
-  return buffer_size;
-}
-
-} // namespace
-
 class ReadAfterWriteSyncs : public kir::ExprMutator {
  private:
   using kir::ExprMutator::handle;
@@ -489,7 +463,9 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
       kir::Allocate* maybe_alloc = nullptr;
       if (sync_bitmap.hasBID()) {
         maybe_alloc = lower_utils::allocGlobalBufferForGridComm(
-            getGridSyncBufferSize(sync_bitmap), DataType::Int, true);
+            lower_utils::getGridSyncBufferSize(sync_bitmap),
+            DataType::Int,
+            true);
         sync_expr = IrBuilder::create<kir::GridSync>(
             sync_bitmap, maybe_alloc->buffer());
       } else {
diff --git a/csrc/device_lower/utils.cpp b/csrc/device_lower/utils.cpp
index d07944ef212..e33c2268a2c 100644
--- a/csrc/device_lower/utils.cpp
+++ b/csrc/device_lower/utils.cpp
@@ -830,6 +830,28 @@ Val* u32IndexScalarSmemTv(TensorView* smem_tv) {
   return u32addr;
 }
 
+Val* getGridSyncBufferSize(const ParallelTypeBitmap& ptb) {
+  // See the comment above for getGridCommWorkBufferSize.
+  NVF_ERROR(
+      ptb.hasBID(),
+      "Detected  needing a grid sync but no grid bits set in bitmap.");
+  Val* buffer_size = GpuLower::current()->kernel()->oneVal();
+  for (auto pt : kParallelTypeBIDs) {
+    // Synchronized within pt, so all blocks of this PT use the same
+    // sync buffer location, and thus no need to expand the sync
+    // buffer size.
+    if (ptb.get(pt)) {
+      continue;
+    }
+    auto pt_dim = GpuLower::current()->parallelDimensionMap().get(pt);
+    if (pt_dim == nullptr || pt_dim->isOneInt()) {
+      continue;
+    }
+    buffer_size = SimplifyingIrBuilder::mulExpr(buffer_size, pt_dim);
+  }
+  return buffer_size;
+}
+
 } // namespace lower_utils
 
 } // namespace nvfuser
diff --git a/csrc/device_lower/utils.h b/csrc/device_lower/utils.h
index 39596fb77d2..749a6fbd4c5 100644
--- a/csrc/device_lower/utils.h
+++ b/csrc/device_lower/utils.h
@@ -298,10 +298,14 @@ bool isScalarExpr(Expr* expr);
 //!  IterDomain object.
 bool isExtentEqualToMaxParallelTypeExtent(const IterDomain* id);
 
-// Get the uint32_t index of a scalar TensorView. This is usually used for
-// indexing special items in shared memory, like mbarrier.
+//! Get the uint32_t index of a scalar TensorView. This is usually used for
+//! indexing special items in shared memory, like mbarrier.
 Val* u32IndexScalarSmemTv(TensorView* tv);
 
+//! Get the size of a global sync buffer needed to perform a grid reduction for
+//! each axis in bitmap.
+Val* getGridSyncBufferSize(const ParallelTypeBitmap& bitmap);
+
 } // namespace lower_utils
 
 } // namespace nvfuser

From 8be320af168e19a21cb0d5684cadf90cd520916a Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 11 Dec 2023 17:52:27 +0000
Subject: [PATCH 11/66] Fix missing allocation of sync flag buffer

---
 csrc/device_lower/pass/grid_serialization.cpp |  1 +
 csrc/device_lower/pass/index.cpp              | 10 ++++++++++
 csrc/device_lower/pass/index.h                |  2 ++
 3 files changed, 13 insertions(+)

diff --git a/csrc/device_lower/pass/grid_serialization.cpp b/csrc/device_lower/pass/grid_serialization.cpp
index e32d3c9b898..1bc24a30c9b 100644
--- a/csrc/device_lower/pass/grid_serialization.cpp
+++ b/csrc/device_lower/pass/grid_serialization.cpp
@@ -108,6 +108,7 @@ class GridSerializationSyncInserter : kir::ExprMutator {
         true);
     auto wait = IrBuilder::create<kir::BlockSerializeWait>(
         cur_expr_sync_pattern_.value(), alloc->buffer());
+    registerInsertBefore(cur_top_level_expr_, alloc);
     registerInsertBefore(cur_top_level_expr_, wait);
     auto release = IrBuilder::create<kir::BlockSerializeRelease>(
         cur_expr_sync_pattern_.value(), alloc->buffer());
diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index 5995cf367ab..27e7c6ea989 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -1728,6 +1728,16 @@ void IndexLowering::handle(const kir::AsyncCommit* commit) {
   pushBack(const_cast<kir::AsyncCommit*>(commit)); // NOLINT
 }
 
+void IndexLowering::handle(const kir::BlockSerializeWait* sync) {
+  // TODO(kir): remove the need for const_cast
+  pushBack(const_cast<kir::BlockSerializeWait*>(sync)); // NOLINT
+}
+
+void IndexLowering::handle(const kir::BlockSerializeRelease* sync) {
+  // TODO(kir): remove the need for const_cast
+  pushBack(const_cast<kir::BlockSerializeRelease*>(sync)); // NOLINT
+}
+
 void IndexLowering::generate(const std::vector<Expr*>& exprs) {
   for (auto expr : exprs) {
     OptOutConstDispatch::dispatch(expr);
diff --git a/csrc/device_lower/pass/index.h b/csrc/device_lower/pass/index.h
index ea875c3e4e3..86f18bf64b7 100644
--- a/csrc/device_lower/pass/index.h
+++ b/csrc/device_lower/pass/index.h
@@ -82,6 +82,8 @@ class IndexLowering : private OptOutConstDispatch {
   void handle(const kir::MBarrierInvalidate*) final;
   void handle(const kir::AsyncWait*) final;
   void handle(const kir::AsyncCommit*) final;
+  void handle(const kir::BlockSerializeWait*) final;
+  void handle(const kir::BlockSerializeRelease*) final;
 
   void generate(const std::vector<Expr*>& exprs);
 

From 41b125fee8a1104f8ccca5ca2ba6430a93d41c83 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Tue, 12 Dec 2023 14:55:07 +0000
Subject: [PATCH 12/66] Allocate global work buffer. Index is zero for now

I need to replay leaf transforms, then get index.
---
 csrc/codegen.cpp                   | 10 ++---
 csrc/device_lower/lower2device.cpp |  3 --
 csrc/device_lower/pass/index.cpp   | 65 ++++++++++++++----------------
 3 files changed, 35 insertions(+), 43 deletions(-)

diff --git a/csrc/codegen.cpp b/csrc/codegen.cpp
index 33735352754..94cfd217732 100644
--- a/csrc/codegen.cpp
+++ b/csrc/codegen.cpp
@@ -1576,17 +1576,17 @@ class CudaKernelGenerator : private kir::ConstIrVisitor {
     const auto data_type = grop->out()->dtype();
     const auto op_type = grop->getReductionOpType();
 
+    if (grop->isSerial()) {
+      generateSerialGridReduction(grop);
+      return;
+    }
+
     NVF_ERROR(grop->reduction_buffer()->buffer()->isA<TensorView>());
     NVF_ERROR(grop->sync_buffer()->buffer()->isA<TensorView>());
     const auto work_buffer =
         grop->reduction_buffer()->buffer()->as<TensorView>();
     const auto sync_buffer = grop->sync_buffer()->buffer()->as<TensorView>();
 
-    if (grop->isSerial()) {
-      generateSerialGridReduction(grop);
-      return;
-    }
-
     if (grop->isAllreduce()) {
       generateGridAllreduce(grop);
       return;
diff --git a/csrc/device_lower/lower2device.cpp b/csrc/device_lower/lower2device.cpp
index 330655cb017..51e8065dd5d 100644
--- a/csrc/device_lower/lower2device.cpp
+++ b/csrc/device_lower/lower2device.cpp
@@ -273,10 +273,7 @@ GpuLower::GpuLower(Fusion* fusion, const CompileParams& cparams)
            {"rotateLoops", rotateLoops},
            {"UnrollPass", UnrollPass::runPass},
            {"processMisalignedVectorization", processMisalignedVectorization},
-           // NOTE: serial GridReduction introduced here. New syncs can be
-           // introduced here which could impact smem reuse
            {"IndexLowering", IndexLowering::getIndexedExprs},
-           // NOTE: global reuse can be analyzed after this point
            {"fuseWarpReduce", fuseWarpReduce},
            {"generateConditionalFromPredicate",
             generateConditionalFromPredicate},
diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index 27e7c6ea989..7dee39aef22 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -652,21 +652,9 @@ void IndexLowering::handleSerialGridReduction(
 
   // TODO: uncomment once this function works again by associating a temporary
   // TensorIndex
-  /*
   const auto out_tv = out->as<kir::TensorIndex>()->view();
   const auto out_domain = out_tv->domain();
 
-  bool is_serial_reduce_candidate = std::none_of(
-      out_domain->leaf().begin(), out_domain->leaf().end(), [](IterDomain* id) {
-        return id->isReduction() &&
-            (id->isThreadDim() || !id->isParallelized());
-      });
-  if (is_serial_reduce_candidate) {
-    std::cout << "WARNING: Found serial grid reduce candidate: "
-              << rop->toString();
-  }
-  return false;
-
   // If we do a grid reduction we can't have a reduction axis that is not bound
   // to a grid or block dim.
   NVF_ERROR(
@@ -683,26 +671,31 @@ void IndexLowering::handleSerialGridReduction(
       "then the grid reduction. ",
       rop->toString());
 
-  const bool is_persistent = rop->isAllreduce();
-  if (is_persistent) {
-    // TODO: enable persistent serial grid reduce yet
-    return false;
-  }
+  NVF_ERROR(!rop->isAllreduce(), "Serial grid allReduce is not implemented");
 
-  // TODO: allocate global work buffer TensorIndex
+  // Allocate global work buffer TensorIndex.
   // The output index should look just like the reduction output, and be
   // scheduled the same, except it should have no reduction axes.
-  // TODO: it might be simpler to just include reduction axes instead of
-  // stripping them here.
-
-  // auto work_buffer_index = IrBuilder::create<TensorIndex>(
-  //   work_buffer,
-  //   lowerSrcIndex(out, work_buffer));
-
-  auto sync_buffer_size =
-      getGridSyncBufferSize(out_domain, for_loops_, is_persistent);
-  auto sync_buffer = allocateUniqueBuffer(
-      sync_buffer_size, DataType::Int, true, out_tv, sync_buffer_map_);
+  std::vector<IterDomain*> work_buffer_ids;
+  for (auto id : TensorDomain::noReductions(out_tv->getRootDomain())) {
+    // Clone output root without reductions as work buffer root domain
+    work_buffer_ids.push_back(IterDomainBuilder(id).build());
+  }
+  const auto work_buffer_root_domain =
+      IrBuilder::create<TensorDomain>(work_buffer_ids);
+  const auto work_buffer_tv = IrBuilder::create<TensorView>(
+      work_buffer_root_domain, out_tv->dtype(), MemoryType::Global);
+  // TODO: replay leaf transformations from out_tv on work_buffer_tv
+  auto work_alloc = IrBuilder::create<kir::Allocate>(
+      work_buffer_tv,
+      work_buffer_tv->getMemoryType(),
+      nullptr,
+      /*zero_init=*/false);
+  auto work_buffer_idx = IrBuilder::create<kir::TensorIndex>(
+      work_buffer_tv, out->fusion()->zeroVal(DataType::Index)
+      // lowerSrcIndex(out, work_buffer_tv)
+  );
+  pushBack(work_alloc);
 
   // The thread predicate for GridReduction needs to be set
   // separately from the main predicate. Do not combine them like
@@ -715,11 +708,14 @@ void IndexLowering::handleSerialGridReduction(
       rop->init(),
       out,
       in,
-      work_buffer,
-      sync_buffer,
-      entrance_ind,
-      n_entrances,
-      is_persistent);
+      // skip work_buffer, sync_buffer, entrance_ind, n_entrances for serial
+      // reduction node
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr,
+      false,
+      work_buffer_idx);
 
   serial_grid_reduction =
       serial_grid_reduction->withThreadPredicate(thread_pred);
@@ -737,7 +733,6 @@ void IndexLowering::handleSerialGridReduction(
 
   pushBack(serial_grid_reduction);
   GpuLower::current()->propagateExprInfo(rop, back());
-  */
 }
 
 void IndexLowering::handleGridReduction(

From 34d623cbfdd8bc9a91b6701eb086f3986d9e7db2 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 13 Dec 2023 12:46:00 +0000
Subject: [PATCH 13/66] Taking a stab at replay/indexing of intermediate

---
 csrc/device_lower/pass/index.cpp | 27 +++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index 7dee39aef22..6c4f9e921f8 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -15,6 +15,7 @@
 #include <options.h>
 #include <predicate_compute.h>
 #include <transform_iter.h>
+#include <transform_replay.h>
 
 #include <device_lower/pass/index.h>
 
@@ -674,8 +675,10 @@ void IndexLowering::handleSerialGridReduction(
   NVF_ERROR(!rop->isAllreduce(), "Serial grid allReduce is not implemented");
 
   // Allocate global work buffer TensorIndex.
-  // The output index should look just like the reduction output, and be
-  // scheduled the same, except it should have no reduction axes.
+  // The global work buffer will look just like the reduction output, and be
+  // scheduled the same, except it should have no reduction axes, it is not
+  // inlined, and it resides in global memory. Since it is not inlined it will
+  // use a global index
   std::vector<IterDomain*> work_buffer_ids;
   for (auto id : TensorDomain::noReductions(out_tv->getRootDomain())) {
     // Clone output root without reductions as work buffer root domain
@@ -683,18 +686,18 @@ void IndexLowering::handleSerialGridReduction(
   }
   const auto work_buffer_root_domain =
       IrBuilder::create<TensorDomain>(work_buffer_ids);
-  const auto work_buffer_tv = IrBuilder::create<TensorView>(
+  auto work_buffer_tv = IrBuilder::create<TensorView>(
       work_buffer_root_domain, out_tv->dtype(), MemoryType::Global);
-  // TODO: replay leaf transformations from out_tv on work_buffer_tv
+  // replay leaf transformations from out_tv on work_buffer_tv
+  const auto new_domain =
+      TransformReplay::replayCasP(out_tv, work_buffer_tv, 0).first;
+  work_buffer_tv = IrBuilder::create<TensorView>(
+      new_domain, out_tv->dtype(), MemoryType::Global);
+  auto work_buffer_idx =
+      Index::getConsumerIndex(work_buffer_tv, for_loops_, {});
+
   auto work_alloc = IrBuilder::create<kir::Allocate>(
-      work_buffer_tv,
-      work_buffer_tv->getMemoryType(),
-      nullptr,
-      /*zero_init=*/false);
-  auto work_buffer_idx = IrBuilder::create<kir::TensorIndex>(
-      work_buffer_tv, out->fusion()->zeroVal(DataType::Index)
-      // lowerSrcIndex(out, work_buffer_tv)
-  );
+      work_buffer_tv, work_buffer_tv->getMemoryType());
   pushBack(work_alloc);
 
   // The thread predicate for GridReduction needs to be set

From 910ff09862387b3783af1837c82d5d600dc8b893 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 13 Dec 2023 13:17:59 +0000
Subject: [PATCH 14/66] Use fullSelfReplay and getGlobalConsumerStridedIndices

Codegen is now like
```c++
  // Allocate global tensor T5
  reduction::serialReductionStep(
    T3[0LL],
    T2[(i14 + i18)],
    0.000000000e+00f,
    T5[((((((((((((nvfuser_index_t)blockIdx.x) * 8LL) + ((nvfuser_index_t)blockIdx.y)) * 4LL) + i13) * 8LL) + (i18 + nvfuser_zero)) * 4LL) + ((nvfuser_index_t)threadIdx.y)) * 32LL) + ((nvfuser_index_t)threadIdx.x))],
    [](float &a, float b) { a = a + b; },
    index_utils::maskedOffset<false, false, true>(blockIdx, gridDim) == 0,
    index_utils::maskedOffset<false, false, true>(blockIdx, gridDim) == index_utils::maskedSize<false, false, true>(gridDim) - 1,
    true,
    true);
```
This looks OK, although it will get a little better with hoisting. This
compiles, but I get an error in `runFusion`:
```
C++ exception with description "Expected T5_g[ iblockIdx.x59{( ceilDiv(( ceilDiv(( ceilDiv(( ceilDiv(( ceilDiv(262144, 32) ), 4) ), 8) ), 4) ), 8) )}, iblockIdx.y60{8}, ithreadIdx.y54{4}, ithreadIdx.x52{32}, iS58{4}, iS56{8}, rblockIdx.z49{5} ] to be bound to a tensor of rank 1, but got a tensor of rank 6
Exception raised from validateValWithConcreteValue at /opt/pytorch/nvfuser/csrc/expr_evaluator.cpp:38 (most recent call first):
```
This is happening when binding inputs I believe.
---
 csrc/device_lower/pass/index.cpp | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index 6c4f9e921f8..ba045098078 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -676,25 +676,33 @@ void IndexLowering::handleSerialGridReduction(
 
   // Allocate global work buffer TensorIndex.
   // The global work buffer will look just like the reduction output, and be
-  // scheduled the same, except it should have no reduction axes, it is not
+  // scheduled the same, including reduction axes, it is not
   // inlined, and it resides in global memory. Since it is not inlined it will
   // use a global index
   std::vector<IterDomain*> work_buffer_ids;
-  for (auto id : TensorDomain::noReductions(out_tv->getRootDomain())) {
+  for (auto id : out_tv->getRootDomain()) {
     // Clone output root without reductions as work buffer root domain
     work_buffer_ids.push_back(IterDomainBuilder(id).build());
   }
-  const auto work_buffer_root_domain =
+  auto work_buffer_root_domain =
       IrBuilder::create<TensorDomain>(work_buffer_ids);
-  auto work_buffer_tv = IrBuilder::create<TensorView>(
-      work_buffer_root_domain, out_tv->dtype(), MemoryType::Global);
   // replay leaf transformations from out_tv on work_buffer_tv
-  const auto new_domain =
-      TransformReplay::replayCasP(out_tv, work_buffer_tv, 0).first;
-  work_buffer_tv = IrBuilder::create<TensorView>(
+  const auto new_domain = TransformReplay::fullSelfReplay(
+      work_buffer_root_domain, out_tv->domain());
+  auto work_buffer_tv = IrBuilder::create<TensorView>(
       new_domain, out_tv->dtype(), MemoryType::Global);
+  // TODO: expose this in ir_utils instead of hidden in index_compute.cpp
+  const auto sumVals = [](const std::vector<Val*> vals) -> Val* {
+    Val* result_index = GpuLower::current()->kernel()->zeroVal();
+    for (auto v : vals) {
+      result_index = SimplifyingIrBuilder::addExpr(result_index, v);
+    }
+    return result_index;
+  };
+  Val* work_buffer_idx_val =
+      sumVals(Index::getGlobalConsumerStridedIndices(out_tv, for_loops_, {}));
   auto work_buffer_idx =
-      Index::getConsumerIndex(work_buffer_tv, for_loops_, {});
+      IrBuilder::create<kir::TensorIndex>(work_buffer_tv, work_buffer_idx_val);
 
   auto work_alloc = IrBuilder::create<kir::Allocate>(
       work_buffer_tv, work_buffer_tv->getMemoryType());

From 507cf47c6e1788adae4b649f56c74979f072b9d7 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 13 Dec 2023 13:46:46 +0000
Subject: [PATCH 15/66] Infer shape using allocation domain instead of root

Fixes execution error. Test passes!
---
 csrc/kernel_ir.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/csrc/kernel_ir.cpp b/csrc/kernel_ir.cpp
index d317c988760..dddd5b4d9be 100644
--- a/csrc/kernel_ir.cpp
+++ b/csrc/kernel_ir.cpp
@@ -149,7 +149,7 @@ Allocate::Allocate(
     NVF_ERROR(buffer->isA<TensorView>());
     NVF_ERROR(buffer->as<TensorView>()->getMemoryType() == memory_type);
     const auto domain = buffer->as<TensorView>()->domain();
-    for (auto axis : domain->noReductions()) {
+    for (auto axis : TensorDomain::noReductions(domain->maybeAllocation())) {
       shape.push_back(axis->extent());
     }
   }

From e44ef7ec6d40e3a748c5dc84bb165ee26d296f92 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 13 Dec 2023 14:09:58 +0000
Subject: [PATCH 16/66] Update comments

---
 csrc/device_lower/pass/grid_serialization.cpp | 19 +++++++++-
 csrc/device_lower/pass/index.cpp              | 38 +------------------
 2 files changed, 19 insertions(+), 38 deletions(-)

diff --git a/csrc/device_lower/pass/grid_serialization.cpp b/csrc/device_lower/pass/grid_serialization.cpp
index 1bc24a30c9b..4cd22873f16 100644
--- a/csrc/device_lower/pass/grid_serialization.cpp
+++ b/csrc/device_lower/pass/grid_serialization.cpp
@@ -20,7 +20,24 @@ namespace nvfuser {
 
 namespace {
 
-//! Insert needed syncs in
+//! Insert needed syncs in order to serialize blocks for serial grid reduction.
+//!
+//! We inspect the loop nest up to the point after which all outer loops are
+//! trivial. This corresponds to the outer-most loop in the generated
+//! kernel code.
+//!
+//! Conditions for using serial grid reduction:
+//!  - All reduction dimensions are either unparallelized or parallelized as
+//!    BID, not TID. Block and warp reductions could be allowed in the future,
+//!    but the current focus is on cases where all threads are doing separate
+//!    reductions simultaneously.
+//!  - rop is not an allreduce. Note that we could implement serial allreduce
+//!    but it would require inserting a separate grid sync after this outer
+//!    loop.
+//!  - There are no other reductions in this loop nest that are TID or BID
+//!    parallelized, unless they also satisfy the conditions above and their
+//!    reduction pattern matches this one. Otherwise our syncs will be
+//!    mismatched, and there is no good way to handle that yet.
 class GridSerializationSyncInserter : kir::ExprMutator {
  public:
   GridSerializationSyncInserter(const std::vector<Expr*>& exprs) {
diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index ba045098078..f1bc98519f1 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -616,43 +616,6 @@ void IndexLowering::handleSerialGridReduction(
     const ReductionOp* rop,
     Val* out,
     Val* in) {
-  // We inspect the loop nest up to the point after which all outer loops are
-  // parallelized. This corresponds to the outer-most loop in the generated
-  // kernel code.
-  //
-  // Conditions for using serial grid reduction:
-  //  - All reduction dimensions are either unparallelized or parallelized as
-  //    BID, not TID or serial. Block and warp reductions could be allowed in
-  //    the future, but the current focus is on cases where all threads are
-  //    doing separate reductions simultaneously.
-  //  - rop is not an allreduce. Note that we could implement serial allreduce
-  //    but it would require inserting a separate grid sync after this outer
-  //    loop.
-  //  - There are no global loads (i.e. producer TVs in global memory) anywhere
-  //    in the loop nest. The reason for this restriction is that loads
-  //    introduce data dependencies which wind up serializing sequences of
-  //    loads/stores that can be disadvantageous. Consider that a non-serial
-  //    (default) gridReduce that is inlined with some producer global read
-  //    will enable all blocks to participate in the read in parallel, since
-  //    only the last block of the reduction segment is serialized, and only
-  //    for the reduction itself. In a serial reduction the entire loop is
-  //    serialized, meaning we will not allow all blocks to perform those reads
-  //    in parallel, which can be inefficient.
-  //    Global stores, on the other hand, are fine as they do not cause this
-  //    excessive serialization.
-  //    Note that in addition to global loads, we should also avoid inlining
-  //    expensive computation into these blocks, since as stated above the
-  //    whole block will be serialized. That should be left to the automatic
-  //    schedulers in cases where serial reductions will be used.
-  //  - There are no other reductions in this loop nest that are TID or BID
-  //    parallelized, unless they also satisfy the conditions above and their
-  //    reduction pattern matches this one. Otherwise our syncs will be
-  //    mismatched, and there is no good way to handle that yet.
-  // TODO: implement the conditions above
-  // TODO: Move this comment above to the lowering pass where it's implemented
-
-  // TODO: uncomment once this function works again by associating a temporary
-  // TensorIndex
   const auto out_tv = out->as<kir::TensorIndex>()->view();
   const auto out_domain = out_tv->domain();
 
@@ -691,6 +654,7 @@ void IndexLowering::handleSerialGridReduction(
       work_buffer_root_domain, out_tv->domain());
   auto work_buffer_tv = IrBuilder::create<TensorView>(
       new_domain, out_tv->dtype(), MemoryType::Global);
+
   // TODO: expose this in ir_utils instead of hidden in index_compute.cpp
   const auto sumVals = [](const std::vector<Val*> vals) -> Val* {
     Val* result_index = GpuLower::current()->kernel()->zeroVal();

From 8a3134ea6eb20b7559a2efbe4dd659a768d4fb4f Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 13 Dec 2023 14:10:04 +0000
Subject: [PATCH 17/66] Hoist index scalar.

Generated kernel now looks like
```c++
  // Allocate global tensor T4
  grid_sync::blockSerializeWait<false, false, true>(&T4[index_utils::maskedOffset<true, true, false>(blockIdx, gridDim)]);
  #pragma unroll
  for(nvfuser_index_t i13 = 0; i13 < 4LL; ++i13) {
    nvfuser_index_t i14;
    i14 = 8LL * i13;
    nvfuser_index_t i15;
    i15 = 2048LL * i13;
    nvfuser_index_t i16;
    i16 = i4 + i15;
    nvfuser_index_t i17;
    i17 = -i15;
    #pragma unroll
    for(nvfuser_index_t i18 = 0; i18 < 8LL; ++i18) {
      nvfuser_index_t i19;
      i19 = 256LL * (i18 + nvfuser_zero);
      nvfuser_index_t i20;
      i20 = i16 + i19;
      float T3[1LL];
      T3[0LL] = 0.000000000e+00f;
      // Allocate global tensor T5
      reduction::serialReductionStep(
        T3[0LL],
        T2[(i14 + i18)],
        0.000000000e+00f,
        T5[i20],
        [](float &a, float b) { a = a + b; },
        index_utils::maskedOffset<false, false, true>(blockIdx, gridDim) == 0,
        index_utils::maskedOffset<false, false, true>(blockIdx, gridDim) == index_utils::maskedSize<false, false, true>(gridDim) - 1,
        true,
        true);
      if ((b6 && (i5 < (i17 - i19)))) {
        T1[i20]
           = T3[0LL];
      }
    }
  }
  NVFUSER_UPDATE_MAGIC_ZERO;
  grid_sync::blockSerializeRelease<false, false, true>(&T4[index_utils::maskedOffset<true, true, false>(blockIdx, gridDim)]);
```
Note that the index `i20` matches the output `T1`. This is what we need
to reclaim `T1` in a later PR; it will still be a challenge in that work
to exact map between `T5` and `T3` in order to get `T1` and `T5` exact
mapped...
---
 csrc/device_lower/pass/index.cpp | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index f1bc98519f1..df16cbd9c66 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -663,10 +663,14 @@ void IndexLowering::handleSerialGridReduction(
     }
     return result_index;
   };
+
   Val* work_buffer_idx_val =
       sumVals(Index::getGlobalConsumerStridedIndices(out_tv, for_loops_, {}));
-  auto work_buffer_idx =
-      IrBuilder::create<kir::TensorIndex>(work_buffer_tv, work_buffer_idx_val);
+
+  auto work_buffer_idx = IrBuilder::create<kir::TensorIndex>(
+      work_buffer_tv,
+      GpuLower::current()->commonScalarMap().hoistScalar(
+          work_buffer_idx_val, for_loops_));
 
   auto work_alloc = IrBuilder::create<kir::Allocate>(
       work_buffer_tv, work_buffer_tv->getMemoryType());

From d27a675b94c9fff3e14d172161b55ca5bf36afb8 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 13 Dec 2023 14:22:30 +0000
Subject: [PATCH 18/66] Clean up comments.

---
 csrc/device_lower/pass/grid_serialization.h | 3 +--
 csrc/device_lower/pass/index.cpp            | 6 +++---
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/csrc/device_lower/pass/grid_serialization.h b/csrc/device_lower/pass/grid_serialization.h
index 9a30b0dd0b5..725ce969848 100644
--- a/csrc/device_lower/pass/grid_serialization.h
+++ b/csrc/device_lower/pass/grid_serialization.h
@@ -19,8 +19,7 @@ namespace nvfuser {
 
 //! Detect ReductionOps that have serialGridReductionRequested() == true. When
 //! found, confirm that no conflicting operations exist, then place sync nodes
-//! before and after outer-most non-parallelized loop. Also allocate a new
-//! global buffer to serve as the work buffer.
+//! before and after outer-most non-parallelized loop.
 std::vector<Expr*> insertGridSerializationSyncs(
     const std::vector<Expr*>& exprs);
 
diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index df16cbd9c66..2e174a7c77f 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -638,10 +638,10 @@ void IndexLowering::handleSerialGridReduction(
   NVF_ERROR(!rop->isAllreduce(), "Serial grid allReduce is not implemented");
 
   // Allocate global work buffer TensorIndex.
+  //
   // The global work buffer will look just like the reduction output, and be
-  // scheduled the same, including reduction axes, it is not
-  // inlined, and it resides in global memory. Since it is not inlined it will
-  // use a global index
+  // scheduled the same way, including reduction axes. However, it is not
+  // inlined and it resides in global memory.
   std::vector<IterDomain*> work_buffer_ids;
   for (auto id : out_tv->getRootDomain()) {
     // Clone output root without reductions as work buffer root domain

From 46c70b64264f0b038578f8b5f3e382467f66f303 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 13 Dec 2023 20:03:13 +0000
Subject: [PATCH 19/66] Update NVFuserTest.Pipeline_CUDA

Also sort expected output by line to give clearer error messages.
---
 test/test_pipeline.cpp | 90 +++++++++++++++++++++++++-----------------
 1 file changed, 54 insertions(+), 36 deletions(-)

diff --git a/test/test_pipeline.cpp b/test/test_pipeline.cpp
index f81203cc666..58823dc4a76 100644
--- a/test/test_pipeline.cpp
+++ b/test/test_pipeline.cpp
@@ -21,6 +21,26 @@ namespace nvfuser {
 
 using namespace at::indexing;
 
+std::string sortByLine(const std::string& input) {
+  auto ss = std::stringstream(input);
+  std::vector<std::string> lines;
+  std::string line;
+  while (std::getline(ss, line, '\n')) {
+    lines.push_back(line);
+  }
+  std::sort(lines.begin(), lines.end());
+  std::stringstream output;
+  bool first = true;
+  for (auto line : lines) {
+    if (!first) {
+      output << std::endl;
+    }
+    first = false;
+    output << line;
+  }
+  return output.str();
+}
+
 TEST_F(NVFuserTest, Pipeline_CUDA) {
   // Fusion definition
   Fusion fusion;
@@ -81,7 +101,7 @@ TEST_F(NVFuserTest, Pipeline_CUDA) {
       " PipelineVal representing Val T0_g[ iS0{i0}, iS1{i2} ] on stage " +
       std::to_string(stage0.unique_id) +
       "\n"
-      " PipelineVal representing Val T4_g[ iS6{i13}, iS7{i14}, iS8{i15} ] on stage " +
+      " PipelineVal representing Val T4_g[ iS6{i15}, iS7{i16}, iS8{i17} ] on stage " +
       std::to_string(stage2.unique_id) +
       "\n"
       "}\n"
@@ -105,98 +125,96 @@ TEST_F(NVFuserTest, Pipeline_CUDA) {
       ".Inputs={T2_l[ iS4{i2} ], }. Outputs={T3_g[ rS5{i2} ], }.\n"
       "  PipelineStage representing Stage " +
       std::to_string(stage2.unique_id) +
-      ".Inputs={T4_g[ iS6{i13}, iS7{i14}, iS8{i15} ], }. Outputs={T5_l[ rS9{i13}, iS10{i14}, iS11{i15} ], }.\n"
-      "  PipelineVal representing Val T5_l[ rS9{i13}, iS10{i14}, iS11{i15} ] on stage " +
+      ".Inputs={T4_g[ iS6{i15}, iS7{i16}, iS8{i17} ], }. Outputs={T5_l[ rS9{i15}, iS10{i16}, iS11{i17} ], }.\n"
+      "  PipelineVal representing Val T5_l[ rS9{i15}, iS10{i16}, iS11{i17} ] on stage " +
       std::to_string(stage2.unique_id) +
       "\n"
-      "  PipelineCommunication that transfers PipelineVal representing Val T5_l[ rS9{i13}, iS10{i14}, iS11{i15} ] on stage " +
+      "  PipelineCommunication that transfers PipelineVal representing Val T5_l[ rS9{i15}, iS10{i16}, iS11{i17} ] on stage " +
       std::to_string(stage2.unique_id) +
-      " to PipelineVal representing Val T6_l[ iS12{i14}, iS13{i15} ] on stage " +
+      " to PipelineVal representing Val T6_l[ iS12{i16}, iS13{i17} ] on stage " +
       std::to_string(stage3.unique_id) +
       "\n"
-      "  PipelineVal representing Val T6_l[ iS12{i14}, iS13{i15} ] on stage " +
+      "  PipelineVal representing Val T6_l[ iS12{i16}, iS13{i17} ] on stage " +
       std::to_string(stage3.unique_id) +
       "\n"
       "  PipelineStage representing Stage " +
       std::to_string(stage3.unique_id) +
-      ".Inputs={T6_l[ iS12{i14}, iS13{i15} ], }. Outputs={T7_l[ iS14{i14}, iS15{i15} ], T8_l[ rS16{i14}, iS17{i15} ], }.\n"
-      "  PipelineVal representing Val T7_l[ iS14{i14}, iS15{i15} ] on stage " +
+      ".Inputs={T6_l[ iS12{i16}, iS13{i17} ], }. Outputs={T7_l[ iS14{i16}, iS15{i17} ], T8_l[ rS16{i16}, iS17{i17} ], }.\n"
+      "  PipelineVal representing Val T7_l[ iS14{i16}, iS15{i17} ] on stage " +
       std::to_string(stage3.unique_id) +
       "\n"
-      "  PipelineCommunication that transfers PipelineVal representing Val T7_l[ iS14{i14}, iS15{i15} ] on stage " +
+      "  PipelineCommunication that transfers PipelineVal representing Val T7_l[ iS14{i16}, iS15{i17} ] on stage " +
       std::to_string(stage3.unique_id) +
-      " to PipelineVal representing Val T12_l[ iS24{i14}, iS25{i15} ] on stage " +
+      " to PipelineVal representing Val T12_l[ iS24{i16}, iS25{i17} ] on stage " +
       std::to_string(stage5.unique_id) +
       "\n"
-      "  PipelineVal representing Val T12_l[ iS24{i14}, iS25{i15} ] on stage " +
+      "  PipelineVal representing Val T12_l[ iS24{i16}, iS25{i17} ] on stage " +
       std::to_string(stage5.unique_id) +
       "\n"
       "  PipelineStage representing Stage " +
       std::to_string(stage5.unique_id) +
-      ".Inputs={T12_l[ iS24{i14}, iS25{i15} ], }. Outputs={T13_g[ rS26{i14}, iS27{i15} ], }.\n"
-      "  PipelineVal representing Val T8_l[ rS16{i14}, iS17{i15} ] on stage " +
+      ".Inputs={T12_l[ iS24{i16}, iS25{i17} ], }. Outputs={T13_g[ rS26{i16}, iS27{i17} ], }.\n"
+      "  PipelineVal representing Val T8_l[ rS16{i16}, iS17{i17} ] on stage " +
       std::to_string(stage3.unique_id) +
       "\n"
-      "  PipelineCommunication that transfers PipelineVal representing Val T8_l[ rS16{i14}, iS17{i15} ] on stage " +
+      "  PipelineCommunication that transfers PipelineVal representing Val T8_l[ rS16{i16}, iS17{i17} ] on stage " +
       std::to_string(stage3.unique_id) +
-      " to PipelineVal representing Val T14_l[ iS28{i15} ] on stage " +
+      " to PipelineVal representing Val T14_l[ iS28{i17} ] on stage " +
       std::to_string(stage6.unique_id) +
       "\n"
-      "  PipelineVal representing Val T14_l[ iS28{i15} ] on stage " +
+      "  PipelineVal representing Val T14_l[ iS28{i17} ] on stage " +
       std::to_string(stage6.unique_id) +
       "\n"
-      "  PipelineCommunication that transfers PipelineVal representing Val T5_l[ rS9{i13}, iS10{i14}, iS11{i15} ] on stage " +
+      "  PipelineCommunication that transfers PipelineVal representing Val T5_l[ rS9{i15}, iS10{i16}, iS11{i17} ] on stage " +
       std::to_string(stage2.unique_id) +
-      " to PipelineVal representing Val T9_l[ iS18{i14}, iS19{i15} ] on stage " +
+      " to PipelineVal representing Val T9_l[ iS18{i16}, iS19{i17} ] on stage " +
       std::to_string(stage4.unique_id) +
       "\n"
-      "  PipelineVal representing Val T9_l[ iS18{i14}, iS19{i15} ] on stage " +
+      "  PipelineVal representing Val T9_l[ iS18{i16}, iS19{i17} ] on stage " +
       std::to_string(stage4.unique_id) +
       "\n"
       "  PipelineStage representing Stage " +
       std::to_string(stage4.unique_id) +
-      ".Inputs={T9_l[ iS18{i14}, iS19{i15} ], }. Outputs={T11_l[ rS22{i14}, iS23{i15} ], }.\n"
-      "  PipelineVal representing Val T11_l[ rS22{i14}, iS23{i15} ] on stage " +
+      ".Inputs={T9_l[ iS18{i16}, iS19{i17} ], }. Outputs={T11_l[ rS22{i16}, iS23{i17} ], }.\n"
+      "  PipelineVal representing Val T11_l[ rS22{i16}, iS23{i17} ] on stage " +
       std::to_string(stage4.unique_id) +
       "\n"
-      "  PipelineCommunication that transfers PipelineVal representing Val T11_l[ rS22{i14}, iS23{i15} ] on stage " +
+      "  PipelineCommunication that transfers PipelineVal representing Val T11_l[ rS22{i16}, iS23{i17} ] on stage " +
       std::to_string(stage4.unique_id) +
-      " to PipelineVal representing Val T15_l[ iS29{i15} ] on stage " +
+      " to PipelineVal representing Val T15_l[ iS29{i17} ] on stage " +
       std::to_string(stage6.unique_id) +
       "\n"
-      "  PipelineVal representing Val T15_l[ iS29{i15} ] on stage " +
+      "  PipelineVal representing Val T15_l[ iS29{i17} ] on stage " +
       std::to_string(stage6.unique_id) +
       "\n"
-      "  PipelineCommunication that transfers PipelineVal representing Val T13_g[ rS26{i14}, iS27{i15} ] on stage " +
+      "  PipelineCommunication that transfers PipelineVal representing Val T13_g[ rS26{i16}, iS27{i17} ] on stage " +
       std::to_string(stage5.unique_id) +
-      " to PipelineVal representing Val T16_l[ iS30{i15} ] on stage " +
+      " to PipelineVal representing Val T16_l[ iS30{i17} ] on stage " +
       std::to_string(stage6.unique_id) +
       "\n"
-      "  PipelineVal representing Val T16_l[ iS30{i15} ] on stage " +
+      "  PipelineVal representing Val T16_l[ iS30{i17} ] on stage " +
       std::to_string(stage6.unique_id) +
       "\n"
       "  PipelineStage representing Stage " +
       std::to_string(stage6.unique_id) +
-      ".Inputs={T14_l[ iS28{i15} ], T15_l[ iS29{i15} ], T16_l[ iS30{i15} ], }. Outputs={T19_g[ rS33{i15} ], }.\n"
+      ".Inputs={T14_l[ iS28{i17} ], T15_l[ iS29{i17} ], T16_l[ iS30{i17} ], }. Outputs={T19_g[ rS33{i17} ], }.\n"
       "}\n"
       "Pipeline's outputs:{\n"
       " PipelineVal representing Val T3_g[ rS5{i2} ] on stage " +
       std::to_string(stage1.unique_id) +
       "\n"
-      " PipelineVal representing Val T13_g[ rS26{i14}, iS27{i15} ] on stage " +
+      " PipelineVal representing Val T13_g[ rS26{i16}, iS27{i17} ] on stage " +
       std::to_string(stage5.unique_id) +
       "\n"
-      " PipelineVal representing Val T19_g[ rS33{i15} ] on stage " +
+      " PipelineVal representing Val T19_g[ rS33{i17} ] on stage " +
       std::to_string(stage6.unique_id) +
       "\n"
       "}"};
 
-  // We sort the string so it doesn't depend on the order of the Pipeline's DAG
-  // traversal
-
-  // TODO: we should sort on lines, not on characters
-  std::sort(obtained_string.begin(), obtained_string.end());
-  std::sort(ref_string.begin(), ref_string.end());
+  // We sort the string by line so it doesn't depend on the order of the
+  // Pipeline's DAG traversal
+  obtained_string = sortByLine(obtained_string);
+  ref_string = sortByLine(ref_string);
 
   EXPECT_EQ(obtained_string, ref_string);
 }

From 1149b43000c14c30ad141e6392629202f25b3ed2 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Thu, 4 Jan 2024 18:05:15 +0000
Subject: [PATCH 20/66] Clean up sum val computation

---
 csrc/device_lower/pass/index.cpp | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index d022559393f..cb9c191ba02 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -655,17 +655,11 @@ void IndexLowering::handleSerialGridReduction(
   auto work_buffer_tv = IrBuilder::create<TensorView>(
       new_domain, out_tv->dtype(), MemoryType::Global);
 
-  // TODO: expose this in ir_utils instead of hidden in index_compute.cpp
-  const auto sumVals = [](const std::vector<Val*> vals) -> Val* {
-    Val* result_index = GpuLower::current()->kernel()->zeroVal();
-    for (auto v : vals) {
-      result_index = SimplifyingIrBuilder::addExpr(result_index, v);
-    }
-    return result_index;
-  };
-
-  Val* work_buffer_idx_val =
-      sumVals(Index::getGlobalConsumerStridedIndices(out_tv, for_loops_, {}));
+  Val* work_buffer_idx_val = nullptr;
+  for (auto v :
+       Index::getGlobalConsumerStridedIndices(out_tv, for_loops_, {})) {
+    work_buffer_idx_val = SimplifyingIrBuilder::addExpr(work_buffer_idx_val, v);
+  }
 
   auto work_buffer_idx = IrBuilder::create<kir::TensorIndex>(
       work_buffer_tv,

From f2d7461a4b447d002bb507b7deedea21e4a028a9 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 10 Jan 2024 19:36:47 +0000
Subject: [PATCH 21/66] Clean up comments and reset sync pattern properly

---
 csrc/device_lower/pass/grid_serialization.cpp | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/csrc/device_lower/pass/grid_serialization.cpp b/csrc/device_lower/pass/grid_serialization.cpp
index 4cd22873f16..24a4ac76d31 100644
--- a/csrc/device_lower/pass/grid_serialization.cpp
+++ b/csrc/device_lower/pass/grid_serialization.cpp
@@ -95,19 +95,24 @@ class GridSerializationSyncInserter : kir::ExprMutator {
   }
 
   void dispatch(Expr* expr) override {
+    // We will detect top-level exprs here that require serialization and
+    // insert the required syncs before and after those exprs.
     if (auto loop = dynamic_cast<kir::ForLoop*>(expr);
-        cur_top_level_expr_ || (loop && loop->isTrivial())) {
-      // Never sync around trivial loops. Also avoid redefining
-      // cur_top_level_expr_ if it is already set
+        cur_top_level_expr_.has_value() || (loop && loop->isTrivial())) {
+      // Never sync around trivial loops since they do not appear in the
+      // generated CUDA code. Also avoid redefining cur_top_level_expr_ if it
+      // is already set, which indicates that this expression is contained in
+      // an outer non-trivial loop.
       kir::ExprMutator::dispatch(expr);
       return;
     }
     // Any other expr, i.e. non-trivial loops or regular Exprs, can be synced if
     // it is top-level and either is or contains a serial grid reduction
     cur_top_level_expr_ = expr;
+    // If a serial grid reduction was found when traversing expr, then
+    // cur_expr_sync_pattern_ will be set
+    cur_expr_sync_pattern_ = std::nullopt;
     kir::ExprMutator::dispatch(expr);
-    // If a serial grid reduction was found in expr, then cur_expr_sync_pattern_
-    // will be set
     if (cur_expr_sync_pattern_.has_value()) {
       insertSyncs();
     }

From a8641847e5c89fb4928a90f098a9c889184fda04 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Thu, 11 Jan 2024 17:36:31 +0000
Subject: [PATCH 22/66] Fix compile error

---
 csrc/device_lower/pass/grid_serialization.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/csrc/device_lower/pass/grid_serialization.cpp b/csrc/device_lower/pass/grid_serialization.cpp
index 24a4ac76d31..6508b030788 100644
--- a/csrc/device_lower/pass/grid_serialization.cpp
+++ b/csrc/device_lower/pass/grid_serialization.cpp
@@ -98,7 +98,7 @@ class GridSerializationSyncInserter : kir::ExprMutator {
     // We will detect top-level exprs here that require serialization and
     // insert the required syncs before and after those exprs.
     if (auto loop = dynamic_cast<kir::ForLoop*>(expr);
-        cur_top_level_expr_.has_value() || (loop && loop->isTrivial())) {
+        cur_top_level_expr_ != nullptr || (loop && loop->isTrivial())) {
       // Never sync around trivial loops since they do not appear in the
       // generated CUDA code. Also avoid redefining cur_top_level_expr_ if it
       // is already set, which indicates that this expression is contained in

From 819a6e0260940262dd875381bcfe2a1c33f89e03 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Tue, 16 Jan 2024 22:46:37 +0000
Subject: [PATCH 23/66] Re-use TensorDomain instead of replaying

---
 csrc/device_lower/pass/index.cpp | 17 +++--------------
 1 file changed, 3 insertions(+), 14 deletions(-)

diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index cb9c191ba02..348316d67b9 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -639,21 +639,10 @@ void IndexLowering::handleSerialGridReduction(
 
   // Allocate global work buffer TensorIndex.
   //
-  // The global work buffer will look just like the reduction output, and be
-  // scheduled the same way, including reduction axes. However, it is not
-  // inlined and it resides in global memory.
-  std::vector<IterDomain*> work_buffer_ids;
-  for (auto id : out_tv->getRootDomain()) {
-    // Clone output root without reductions as work buffer root domain
-    work_buffer_ids.push_back(IterDomainBuilder(id).build());
-  }
-  auto work_buffer_root_domain =
-      IrBuilder::create<TensorDomain>(work_buffer_ids);
-  // replay leaf transformations from out_tv on work_buffer_tv
-  const auto new_domain = TransformReplay::fullSelfReplay(
-      work_buffer_root_domain, out_tv->domain());
+  // The global work buffer will look just like the reduction output except
+  // that it resides in global memory.
   auto work_buffer_tv = IrBuilder::create<TensorView>(
-      new_domain, out_tv->dtype(), MemoryType::Global);
+      out_tv->domain(), out_tv->dtype(), MemoryType::Global);
 
   Val* work_buffer_idx_val = nullptr;
   for (auto v :

From 449a296faac1f78892130f49d7cce594fff0674b Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Thu, 18 Jan 2024 13:59:10 +0000
Subject: [PATCH 24/66] Copy domains to create new TensorDomain instead of
 reusing

---
 csrc/device_lower/pass/index.cpp | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index 348316d67b9..00066bdb3d9 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -641,8 +641,17 @@ void IndexLowering::handleSerialGridReduction(
   //
   // The global work buffer will look just like the reduction output except
   // that it resides in global memory.
+  //
+  // Note that we create a new TensorDomain, copying the domain vectors,
+  // instead of reusing out_tv->domain() directly. This prevents corruption of
+  // this new TensorDomain in case the original is modified later.
+  auto work_buffer_domain = IrBuilder::create<TensorDomain>(
+      out_tv->domain()->root(),
+      out_tv->domain()->rfactor(),
+      out_tv->domain()->allocation(),
+      out_tv->domain()->leaf());
   auto work_buffer_tv = IrBuilder::create<TensorView>(
-      out_tv->domain(), out_tv->dtype(), MemoryType::Global);
+      work_buffer_domain, out_tv->dtype(), MemoryType::Global);
 
   Val* work_buffer_idx_val = nullptr;
   for (auto v :

From 27e449a86bdf8cabe9d78111c5fe6f7e0e80d621 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 19 Jan 2024 12:18:46 +0000
Subject: [PATCH 25/66] Allocate work buffer like leaf of output

---
 csrc/device_lower/pass/index.cpp | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index 00066bdb3d9..cc5dcef5424 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -639,20 +639,18 @@ void IndexLowering::handleSerialGridReduction(
 
   // Allocate global work buffer TensorIndex.
   //
-  // The global work buffer will look just like the reduction output except
-  // that it resides in global memory.
-  //
-  // Note that we create a new TensorDomain, copying the domain vectors,
-  // instead of reusing out_tv->domain() directly. This prevents corruption of
-  // this new TensorDomain in case the original is modified later.
-  auto work_buffer_domain = IrBuilder::create<TensorDomain>(
-      out_tv->domain()->root(),
-      out_tv->domain()->rfactor(),
-      out_tv->domain()->allocation(),
-      out_tv->domain()->leaf());
+  // For convenience, the global work buffer is allocated like the leaf domain
+  // of the ReductionOp output. In the future, we may want the allocation
+  // domain to be different in order to enable re-use of global output buffers
+  // for in-place reduction.
+  std::vector<IterDomain*> work_buffer_root;
+  work_buffer_root.reserve(out_tv->nDims());
+  for (IterDomain* id : out_tv->getLeafDomain()) {
+    work_buffer_root.push_back(IterDomainBuilder(id).build());
+  }
+  auto work_buffer_domain = IrBuilder::create<TensorDomain>(work_buffer_root);
   auto work_buffer_tv = IrBuilder::create<TensorView>(
       work_buffer_domain, out_tv->dtype(), MemoryType::Global);
-
   Val* work_buffer_idx_val = nullptr;
   for (auto v :
        Index::getGlobalConsumerStridedIndices(out_tv, for_loops_, {})) {

From 2c3278b97f030755b657f867898725204c34cfe8 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 13 Dec 2023 15:06:19 +0000
Subject: [PATCH 26/66] Use serial grid reduction in split-K

---
 csrc/scheduler/matmul.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 366e0821b66..83b310073ff 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -899,6 +899,8 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     splitk_sum = mma_result;
     mma_result = splitk_sum->rFactor({-4, -1});
 
+    splitk_sum->definition()->as<ReductionOp>()->requestSerialGridReduction();
+
     num_splitk_dims = 1;
   }
 

From 4e40d68860573d559774e70c740ce5c3bd1b7254 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 13 Dec 2023 20:32:49 +0000
Subject: [PATCH 27/66] Set proper dtype for init in MmaOp

---
 csrc/ops/arith.cpp           | 2 +-
 csrc/scheduler/mma_utils.cpp | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/csrc/ops/arith.cpp b/csrc/ops/arith.cpp
index e75365969c0..f5001fb9819 100644
--- a/csrc/ops/arith.cpp
+++ b/csrc/ops/arith.cpp
@@ -2560,7 +2560,7 @@ TensorView* fusedMultiplySum(
     const std::vector<int>& axes,
     Val* init) {
   if (init == nullptr) {
-    init = IrBuilder::create<Val>(0.0);
+    init = IrBuilder::create<Val>(0.0, DataType::Float);
   }
 
   // TODO:
diff --git a/csrc/scheduler/mma_utils.cpp b/csrc/scheduler/mma_utils.cpp
index 67b44d78b5c..0c8f230bf4e 100644
--- a/csrc/scheduler/mma_utils.cpp
+++ b/csrc/scheduler/mma_utils.cpp
@@ -1288,10 +1288,10 @@ RolesMapOpt getTensorsRoles(Fusion* fusion) {
 namespace {
 
 void addMMAOp(Fusion* fusion_, std::vector<MulSumProperties>& props) {
-  auto* init = IrBuilder::create<Val>(0.0);
   for (auto prop : props) {
+    auto* init = IrBuilder::create<Val>(0.0, prop.out->getDataType().value());
     IrBuilder::create<MmaOp>(
-        prop.insouts.out, prop.insouts.a, prop.insouts.b, init);
+        prop.insouts.out, prop.insouts..a, prop.insouts..b, init);
   }
 }
 

From 7bfa7091a5e8f4cb9180fc0024308b0db5aa97ec Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 20 Dec 2023 12:59:58 +0000
Subject: [PATCH 28/66] Restore split-k benchmarks

These were disabled in #1545 because of slow compilation with gridReduce
---
 benchmark/matmul.cpp | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/benchmark/matmul.cpp b/benchmark/matmul.cpp
index 3a4fb8ae16a..4a1cc265f85 100644
--- a/benchmark/matmul.cpp
+++ b/benchmark/matmul.cpp
@@ -685,11 +685,7 @@ static void MatmulShapeWarpStageAutoSplitK(benchmark::internal::Benchmark* b) {
 
 ForAllLayouts(EagerModeBenchmark);
 ForAllLayouts(NvfuserMatmulBenchmark);
-// Disable split-K benchmarks due to slow compilation.
-// See https://github.com/NVIDIA/Fuser/issues/1389.
-// These benchmarks should be enabled again after merging
-// https://github.com/NVIDIA/Fuser/pull/1510
-// ForAllLayouts(AutoSplitKBenchmark);
+ForAllLayouts(AutoSplitKBenchmark);
 ForAllLayouts(AutoPartitionedKBenchmark);
 
 // Note: SplitK Reduction benchmarks are parametrized only by M, N. The splitk

From fc07a9a823f223466a48ac4f54077f025607719c Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 19 Jan 2024 12:30:46 +0000
Subject: [PATCH 29/66] Fix after rebase

---
 csrc/scheduler/mma_utils.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/csrc/scheduler/mma_utils.cpp b/csrc/scheduler/mma_utils.cpp
index 0c8f230bf4e..36f8da4a424 100644
--- a/csrc/scheduler/mma_utils.cpp
+++ b/csrc/scheduler/mma_utils.cpp
@@ -1289,9 +1289,10 @@ namespace {
 
 void addMMAOp(Fusion* fusion_, std::vector<MulSumProperties>& props) {
   for (auto prop : props) {
-    auto* init = IrBuilder::create<Val>(0.0, prop.out->getDataType().value());
+    auto* init =
+        IrBuilder::create<Val>(0.0, prop.insouts.out->getDataType().value());
     IrBuilder::create<MmaOp>(
-        prop.insouts.out, prop.insouts..a, prop.insouts..b, init);
+        prop.insouts.out, prop.insouts.a, prop.insouts.b, init);
   }
 }
 

From 49b7febafa11456b3f468277d8d54d6343158207 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 13 Dec 2023 20:32:49 +0000
Subject: [PATCH 30/66] Set proper dtype for init in MmaOp

---
 csrc/scheduler/mma_utils.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/csrc/scheduler/mma_utils.cpp b/csrc/scheduler/mma_utils.cpp
index 36f8da4a424..e21d209b6cf 100644
--- a/csrc/scheduler/mma_utils.cpp
+++ b/csrc/scheduler/mma_utils.cpp
@@ -1289,10 +1289,8 @@ namespace {
 
 void addMMAOp(Fusion* fusion_, std::vector<MulSumProperties>& props) {
   for (auto prop : props) {
-    auto* init =
-        IrBuilder::create<Val>(0.0, prop.insouts.out->getDataType().value());
-    IrBuilder::create<MmaOp>(
-        prop.insouts.out, prop.insouts.a, prop.insouts.b, init);
+    auto* init = IrBuilder::create<Val>(0.0, prop.insouts.out->getDataType().value());
+    IrBuilder::create<MmaOp>(prop.insouts.out, prop.insouts.a, prop.insouts.b, init);
   }
 }
 

From fe8fbf5de5828d195dcba6e60ad5f3b5581c0589 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 15 Dec 2023 13:51:40 +0000
Subject: [PATCH 31/66] Naive first step. Doesn't work yet

---
 csrc/scheduler/matmul.cpp | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 83b310073ff..08894d579c3 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -786,8 +786,11 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   auto mma_result = mma->out()->as<TensorView>();
 
   // Unswizzle mma result in shared memory
-  auto smem_epilogue =
-      params.use_smem_epilogue ? mma_result->cacheAfter() : mma_result;
+  // Note that if we are using split-K, we will set up this buffer after
+  // rfactoring the matmul, between the MmaOp and the ReductionOp, in order to
+  // take advantage of unswizzling during the grid reduction
+  auto smem_epilogue = (params.use_smem_epilogue && params.splitk_factor == 0)
+    ? mma_result->cacheAfter() : mma_result;
 
   // Clear MmaOp pointer, it's not needed from now on
   mma = nullptr;
@@ -899,6 +902,14 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     splitk_sum = mma_result;
     mma_result = splitk_sum->rFactor({-4, -1});
 
+    if (params.use_smem_epilogue) {
+      //   splitk_sum = sum(mma_result)
+      // becomes
+      //   smem_epilogue = set(mma_result)
+      //   splitk_sum = sum(smem_epilogue)
+      smem_epilogue = mma_result->cacheAfter();
+    }
+
     splitk_sum->definition()->as<ReductionOp>()->requestSerialGridReduction();
 
     num_splitk_dims = 1;

From ef4b3f350c9f779ee1156550ce7e08d09a7750cd Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Tue, 2 Jan 2024 16:49:13 +0000
Subject: [PATCH 32/66] Fix typo

---
 csrc/scheduler/matmul.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 08894d579c3..0c76d8274bc 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -789,7 +789,7 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   // Note that if we are using split-K, we will set up this buffer after
   // rfactoring the matmul, between the MmaOp and the ReductionOp, in order to
   // take advantage of unswizzling during the grid reduction
-  auto smem_epilogue = (params.use_smem_epilogue && params.splitk_factor == 0)
+  auto smem_epilogue = (params.use_smem_epilogue && params.splitk_factor == 1)
     ? mma_result->cacheAfter() : mma_result;
 
   // Clear MmaOp pointer, it's not needed from now on

From d5febec346e2a2bb89089d7aea3f92e1b988b046 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Tue, 2 Jan 2024 16:49:53 +0000
Subject: [PATCH 33/66] Add debug output

---
 csrc/scheduler/matmul.cpp | 39 ++++++++++++++++++++++++++++++++-------
 1 file changed, 32 insertions(+), 7 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 0c76d8274bc..b24506eee2f 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -863,6 +863,7 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
 
   // [... M,N,K]
   mma_utils::makeTile(mma_result, gemm_tile.cta_tile.toVector());
+  // [..., Mo, No, Ko, Mi, Ni, Ki]
 
   // Swizzle block tiles:
   if (params.grid_swizzle_factor != 1) {
@@ -885,15 +886,15 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     }
   }
 
-  // [..., Mo, No, Koo, Mi, Ni, Ki]
+  // [..., Mo, No, Ko, Mi, Ni, Ki]
   int num_splitk_dims = 0;
   TensorView* splitk_sum = nullptr;
   if (params.splitk_factor != 1) {
-    // Split Koo -> [Kf, Ko]
+    // Split Ko -> [rKf, rKg]
     mma_result->split(-4, params.splitk_factor, /*inner*/ false);
-    // After split [..., Mo, No, Kf, Ko, Mi, Ni, Ki]
+    // After split [..., Mo, No, Kf, Kg, Mi, Ni, Ki]
     // rFactor converts
-    //   mma_result = mma(A, B, {/*Kf*/-5, /*Ko*/-4, /*Ki*/-1});
+    //   mma_result = mma(A, B, {/*Kf*/-5, /*Kg*/-4, /*Ki*/-1});
     // to
     //   intermediate = mma(A, B, {-4, -1});
     //   final_sum = sum(intermediate, {/*Kf*/-3});
@@ -902,12 +903,17 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     splitk_sum = mma_result;
     mma_result = splitk_sum->rFactor({-4, -1});
 
+    // mma_result = [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
+    // splitk_sum = [..., iMo, iNo, rKf, iMi, iNi]
+
     if (params.use_smem_epilogue) {
       //   splitk_sum = sum(mma_result)
       // becomes
       //   smem_epilogue = set(mma_result)
       //   splitk_sum = sum(smem_epilogue)
       smem_epilogue = mma_result->cacheAfter();
+
+      // smem_epilogue = [..., iMo, iNo, iKf, iMi, iNi]
     }
 
     splitk_sum->definition()->as<ReductionOp>()->requestSerialGridReduction();
@@ -915,6 +921,25 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     num_splitk_dims = 1;
   }
 
+  // No split-K
+  //   mma_result      [..., iMo, iNo, rKo, iMi, iNi, rKi]
+  //   smem_epilogue   (unscheduled, same as original mma_result)
+  //   splitk_sum      (nullptr)
+  //
+  // With split-K & no smem epilogue
+  //   mma_result      [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
+  //   smem_epilogue   (mma_result)
+  //   splitk_sum      [..., iMo, iNo, rKf, iMi, iNi]
+  //
+  // With split-K & smem epilogue
+  //   mma_result      [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
+  //   smem_epilogue   [..., iMo, iNo, iKf, iMi, iNi]
+  //   splitk_sum      [..., iMo, iNo, rKf, iMi, iNi]
+  std::cout << "  After split-K split:\n";
+  std::cout << "    mma_result " << mma_result->toString() << std::endl;
+  std::cout << "    smem_epilogue " << smem_epilogue->toString() << std::endl;
+  std::cout << "    splitk_sum " << (splitk_sum != nullptr ? splitk_sum->toString() : "nullptr") << std::endl;
+
   // Propagate tiling globally
   scheduler_utils::transformPropagateToAllFrom(mma_result, -1);
 
@@ -1005,13 +1030,13 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   //  with splitk:
   // nbatch +   1   2   3   4    5    6    7    8     9    10    11   12
   //      -13 -12 -11 -10  -9   -8   -7   -6   -5    -4    -3    -2   -1
-  // [..., Mo, No, Kf, Ko, Kw, Mwo, Nwo, Mwi, Nwi, MNi1, MNi2, MNi3,  Ki]
+  // [..., Mo, No, Kf, Kg, Kw, Mwo, Nwo, Mwi, Nwi, MNi1, MNi2, MNi3,  Ki]
   //  (iS) iBx iBy rBz rS  rS  iTz  iTy   iS   iS  iMMA   iTx  iMMA rMMA
   //
   //  without splitk:
   // nbatch +   1       2   3    4    5    6    7     8     9    10   11
   //      -12 -11     -10  -9   -8   -7   -6   -5    -4    -3    -2   -1
-  // [..., Mo, No,     Ko, Kw, Mwo, Nwo, Mwi, Nwi, MNi1, MNi2, MNi3,  Ki]
+  // [..., Mo, No,     Kg, Kw, Mwo, Nwo, Mwi, Nwi, MNi1, MNi2, MNi3,  Ki]
   // (iBz) iBx iBy     rS  rS  iTz  iTy   iS   iS  iMMA   iTx  iMMA rMMA
 
   // When we have both batch dims and splitk, parallelize splitk only.
@@ -1159,7 +1184,7 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
 
   if (params.double_buffer_options.double_buffer_smem_read &&
       params.double_buffer_options.double_buffer_smem_write) {
-    // rotate Ko loop
+    // rotate Kg loop
     scheduler_utils::rotateLoop(
         mma_result, num_batch_dims + 2 + num_splitk_dims, {acr, bcr});
   }

From d358071b790e166c5cb89344af6cf3f2d623f00d Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Thu, 4 Jan 2024 20:11:08 +0000
Subject: [PATCH 34/66] TMP add more debug output, comments

---
 csrc/scheduler/matmul.cpp    | 69 +++++++++++++++++++++++++++++++++---
 test/test_gpu_tensorcore.cpp |  9 ++++-
 2 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index b24506eee2f..97eff3a9975 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -921,20 +921,36 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     num_splitk_dims = 1;
   }
 
-  // No split-K
+  // No split-K, smem_epilogue = false
   //   mma_result      [..., iMo, iNo, rKo, iMi, iNi, rKi]
   //   smem_epilogue   (unscheduled, same as original mma_result)
   //   splitk_sum      (nullptr)
+  //     mma_result T6_l[ iS27{( ceilDiv(i0, 128) )}, iS29{( ceilDiv(i4, 128) )}, rS31{( ceilDiv(i2, 32) )}, iS28{128}, iS30{128}, rS32{32} ]
+  //     smem_epilogue T6_l[ iS27{( ceilDiv(i0, 128) )}, iS29{( ceilDiv(i4, 128) )}, rS31{( ceilDiv(i2, 32) )}, iS28{128}, iS30{128}, rS32{32} ]
+  //
+  // No split-K, smem_epilogue = true
+  //   mma_result      [..., iMo, iNo, rKo, iMi, iNi, rKi]
+  //   smem_epilogue   (unscheduled, same as original mma_result)
+  //   splitk_sum      (nullptr)
+  //     mma_result T6_l[ iS29{( ceilDiv(i0, 128) )}, iS31{( ceilDiv(i4, 128) )}, rS33{( ceilDiv(i2, 32) )}, iS30{128}, iS32{128}, rS34{32} ]
+  //     smem_epilogue T7_l[ iS17{i0}, iS18{i4} ]
   //
   // With split-K & no smem epilogue
   //   mma_result      [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
   //   smem_epilogue   (mma_result)
   //   splitk_sum      [..., iMo, iNo, rKf, iMi, iNi]
-  //
+  //     mma_result T12_l[ iS38{( ceilDiv(i0, 128) )}, iS40{( ceilDiv(i4, 128) )}, iS44{2}rf, rS45{( ceilDiv(( ceilDiv(i2, 32) ), 2) )}rf, iS39{128}, iS41{128}, rS43{32}rf ]
+  //     smem_epilogue T6_l[ iS49{( ceilDiv(i0, 128) )}, iS51{( ceilDiv(i4, 128) )}, rS48{2}, iS50{128}, iS52{128} ]
+  //     splitk_sum T6_l[ iS49{( ceilDiv(i0, 128) )}, iS51{( ceilDiv(i4, 128) )}, rS48{2}, iS50{128}, iS52{128} ]
+
   // With split-K & smem epilogue
   //   mma_result      [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
   //   smem_epilogue   [..., iMo, iNo, iKf, iMi, iNi]
   //   splitk_sum      [..., iMo, iNo, rKf, iMi, iNi]
+  //     mma_result T12_l[ iS38{( ceilDiv(i0, 128) )}, iS40{( ceilDiv(i4, 128) )}, iS44{2}rf, rS45{( ceilDiv(( ceilDiv(i2, 32) ), 2) )}rf, iS39{128}, iS41{128}, rS43{32}rf ]
+  //     smem_epilogue T13_l[ iS56{( ceilDiv(i0, 128) )}, iS58{( ceilDiv(i4, 128) )}, iS55{2}, iS57{128}, iS59{128} ]
+  //     splitk_sum T6_l[ iS49{( ceilDiv(i0, 128) )}, iS51{( ceilDiv(i4, 128) )}, rS48{2}, iS50{128}, iS52{128} ]
+
   std::cout << "  After split-K split:\n";
   std::cout << "    mma_result " << mma_result->toString() << std::endl;
   std::cout << "    smem_epilogue " << smem_epilogue->toString() << std::endl;
@@ -952,18 +968,44 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
 
   // Schedule warp tile
   mma_utils::scheduleWarpTileWithReduction(mma_result, gemm_tile);
-  // [..., Mo, No, (Kf,) Ko, Kw, Mwo, Nwo, Mwi, Nwi, Mi, Ni, Ki]
+  // Single warp k:
+  //          -8   -7   -6  -5  -4  -3  -2  -1
+  //   [... rKwo iMwo iNwo iMw iNw iMi iNi rKi]
+  // If instead K is split across warps (unused):
+  //           -8  -7  -6  -5  -4  -3  -2  -1
+  //   [... iMNwo rKo iMw iNw rKw iMi iNi rKi]
+  // Note the positions of the reduction dimensions has changed
+  std::cout << "  mma_result after schedulWarpTileWithReduction:\n    " << mma_result->toString() << std::endl;
 
   // Propagate warp tile to main loop and epilog/output tvs
   scheduler_utils::BoundedDirectionalTransformPropagator::bothWays(
       mma_result, -1, {acw_smem, bcw_smem}, {smem_epilogue});
 
+  std::cout << "  After bothWays:\n";
+  std::cout << "    mma_result " << mma_result->toString() << std::endl;
+  std::cout << "    smem_epilogue " << smem_epilogue->toString() << std::endl;
+  std::cout << "    splitk_sum " << (splitk_sum != nullptr ? splitk_sum->toString() : "nullptr") << std::endl;
+
+  // No (cross-CTA) split-K
+  //   mma_result      [..., iMo iNo rKo rKwo iMwo iNwo iMw iNw iMi iNi rKi]
+  //   smem_epilogue   (unscheduled, same as original mma_result)
+  //   splitk_sum      (nullptr)
+  //
+  // With split-K & no smem epilogue
+  //   mma_result      [..., iMo iNo iKf rKg rKwo iMwo iNwo iMw iNw iMi iNi rKi]
+  //   smem_epilogue   (mma_result)
+  //   splitk_sum      [..., iMo iNo rKf iMwo iNwo iMw iNw iMi iNi]
+  //
+  // With split-K & smem epilogue
+  //   mma_result      [..., iMo iNo iKf rKg rKwo iMwo iNwo iMw iNw iMi iNi rKi]
+  //   smem_epilogue   [..., iMo iNo iKf iMwo iNwo iMw iNw iMi iNi]
+  //   splitk_sum      [..., iMo iNo rKf iMwo iNwo iMw iNw iMi iNi]
+
   // Schedule prolog:
   //   TODO: this section needs more configurability.
   // ------------------------------------------------------------------
   scheduleProlog(acw_smem, params);
   scheduleProlog(bcw_smem, params);
-  // [..., Mo, No, (Kf,) Ko, Kw, Mwo, Nwo, Mwi, Nwi, Mi, Ni, Ki]
 
   // Add mma swizzle:
   //   TODO: this section goes to a separate matmul util,
@@ -1027,6 +1069,25 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   //   - iMMA: unconracted axis in an MMA tensor core operation.
   //   - rMMA: contract in an MMA tensor core operation.
   //
+  // No (cross-CTA) split-K
+  //   mma_result      [..., iMo iNo rKo rKwo iMwo iNwo iMw iNw iMi iNi rKi]
+  //     nbatch +   1   2    3    4    5    6    7     8     9    10   11
+  //          -12 -11 -10   -9   -8   -7   -6   -5    -4    -3    -2   -1
+  //     [..., Mo, No, Kg, Kwo, Mwo, Nwo, Mwi, Nwi, MNi1, MNi2, MNi3,  Ki]
+  //     (iBz) iBx iBy rS   rS  iTz  iTy   iS   iS  iMMA   iTx  iMMA rMMA
+  //   smem_epilogue   (unscheduled, same as original mma_result)
+  //   splitk_sum      (nullptr)
+  //
+  // With split-K & no smem epilogue
+  //   mma_result      [..., iMo iNo iKf rKg rKwo iMwo iNwo iMw iNw iMi iNi rKi]
+  //   smem_epilogue   (mma_result)
+  //   splitk_sum      [..., iMo iNo rKf iMwo iNwo iMw iNw iMi iNi]
+  //
+  // With split-K & smem epilogue
+  //   mma_result      [..., iMo iNo iKf rKg rKwo iMwo iNwo iMw iNw iMi iNi rKi]
+  //   smem_epilogue   [..., iMo iNo iKf iMwo iNwo iMw iNw iMi iNi]
+  //   splitk_sum      [..., iMo iNo rKf iMwo iNwo iMw iNw iMi iNi]
+  //
   //  with splitk:
   // nbatch +   1   2   3   4    5    6    7    8     9    10    11   12
   //      -13 -12 -11 -10  -9   -8   -7   -6   -5    -4    -3    -2   -1
diff --git a/test/test_gpu_tensorcore.cpp b/test/test_gpu_tensorcore.cpp
index 07b1c29ee2d..d4ad99e1ec8 100644
--- a/test/test_gpu_tensorcore.cpp
+++ b/test/test_gpu_tensorcore.cpp
@@ -2393,6 +2393,8 @@ TEST_F(NVFuserTest, FusionAmpereMatmulSplitK_CUDA) {
   int M = 504, N = 136, K = 8096;
 
   for (auto layout : kAllSupportedMmaLayout) {
+  for (int splitk_factor : {1, 2}) {
+  for (int use_smem_epilogue : {false, true}) {
     Fusion fusion;
     FusionGuard fg(&fusion);
     auto tv0 = makeContigTensor(2, DataType::Half);
@@ -2413,7 +2415,10 @@ TEST_F(NVFuserTest, FusionAmpereMatmulSplitK_CUDA) {
     MatmulParams params;
     params.mma_macro = MmaMacro::Ampere_16_8_16;
     params.tile_sizes = gemm_tile;
-    params.splitk_factor = 2;
+    params.splitk_factor = splitk_factor;
+    params.use_smem_epilogue = use_smem_epilogue;
+
+    std::cout << "\n\n#####    layout=" << toString(layout) << " use_smem_epilogue=" << use_smem_epilogue << " splitk_factor=" << splitk_factor << std::endl;
     scheduleMatmul(&fusion, params);
 
     auto inputs = matmulAtInput(M, N, K, layout);
@@ -2429,6 +2434,8 @@ TEST_F(NVFuserTest, FusionAmpereMatmulSplitK_CUDA) {
     // Relax tolerance for larger sum due to large K
     NVF_CHECK(cg_outputs[0].allclose(tref, 1e-6 * K, 1e-6 * K));
   }
+  }
+  }
 }
 
 // Test splitk with bias epilogue

From 362bde9ef55e446f95d515e387a47df1f90a5a95 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 5 Jan 2024 14:33:26 +0000
Subject: [PATCH 35/66] Fix up to get kernel running, but still have bank
 conflicts

Bank conflict 0 8 at the caching set expression loading the swizzled
register T12_l to the smem epilogue buffer T13_s.
---
 csrc/scheduler/matmul.cpp | 44 +++++++++++++++------------------------
 1 file changed, 17 insertions(+), 27 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 97eff3a9975..99f2db66323 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -790,7 +790,8 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   // rfactoring the matmul, between the MmaOp and the ReductionOp, in order to
   // take advantage of unswizzling during the grid reduction
   auto smem_epilogue = (params.use_smem_epilogue && params.splitk_factor == 1)
-    ? mma_result->cacheAfter() : mma_result;
+      ? mma_result->cacheAfter()
+      : mma_result;
 
   // Clear MmaOp pointer, it's not needed from now on
   mma = nullptr;
@@ -925,36 +926,21 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   //   mma_result      [..., iMo, iNo, rKo, iMi, iNi, rKi]
   //   smem_epilogue   (unscheduled, same as original mma_result)
   //   splitk_sum      (nullptr)
-  //     mma_result T6_l[ iS27{( ceilDiv(i0, 128) )}, iS29{( ceilDiv(i4, 128) )}, rS31{( ceilDiv(i2, 32) )}, iS28{128}, iS30{128}, rS32{32} ]
-  //     smem_epilogue T6_l[ iS27{( ceilDiv(i0, 128) )}, iS29{( ceilDiv(i4, 128) )}, rS31{( ceilDiv(i2, 32) )}, iS28{128}, iS30{128}, rS32{32} ]
   //
   // No split-K, smem_epilogue = true
   //   mma_result      [..., iMo, iNo, rKo, iMi, iNi, rKi]
   //   smem_epilogue   (unscheduled, same as original mma_result)
   //   splitk_sum      (nullptr)
-  //     mma_result T6_l[ iS29{( ceilDiv(i0, 128) )}, iS31{( ceilDiv(i4, 128) )}, rS33{( ceilDiv(i2, 32) )}, iS30{128}, iS32{128}, rS34{32} ]
-  //     smem_epilogue T7_l[ iS17{i0}, iS18{i4} ]
   //
   // With split-K & no smem epilogue
   //   mma_result      [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
   //   smem_epilogue   (mma_result)
   //   splitk_sum      [..., iMo, iNo, rKf, iMi, iNi]
-  //     mma_result T12_l[ iS38{( ceilDiv(i0, 128) )}, iS40{( ceilDiv(i4, 128) )}, iS44{2}rf, rS45{( ceilDiv(( ceilDiv(i2, 32) ), 2) )}rf, iS39{128}, iS41{128}, rS43{32}rf ]
-  //     smem_epilogue T6_l[ iS49{( ceilDiv(i0, 128) )}, iS51{( ceilDiv(i4, 128) )}, rS48{2}, iS50{128}, iS52{128} ]
-  //     splitk_sum T6_l[ iS49{( ceilDiv(i0, 128) )}, iS51{( ceilDiv(i4, 128) )}, rS48{2}, iS50{128}, iS52{128} ]
 
   // With split-K & smem epilogue
   //   mma_result      [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
   //   smem_epilogue   [..., iMo, iNo, iKf, iMi, iNi]
   //   splitk_sum      [..., iMo, iNo, rKf, iMi, iNi]
-  //     mma_result T12_l[ iS38{( ceilDiv(i0, 128) )}, iS40{( ceilDiv(i4, 128) )}, iS44{2}rf, rS45{( ceilDiv(( ceilDiv(i2, 32) ), 2) )}rf, iS39{128}, iS41{128}, rS43{32}rf ]
-  //     smem_epilogue T13_l[ iS56{( ceilDiv(i0, 128) )}, iS58{( ceilDiv(i4, 128) )}, iS55{2}, iS57{128}, iS59{128} ]
-  //     splitk_sum T6_l[ iS49{( ceilDiv(i0, 128) )}, iS51{( ceilDiv(i4, 128) )}, rS48{2}, iS50{128}, iS52{128} ]
-
-  std::cout << "  After split-K split:\n";
-  std::cout << "    mma_result " << mma_result->toString() << std::endl;
-  std::cout << "    smem_epilogue " << smem_epilogue->toString() << std::endl;
-  std::cout << "    splitk_sum " << (splitk_sum != nullptr ? splitk_sum->toString() : "nullptr") << std::endl;
 
   // Propagate tiling globally
   scheduler_utils::transformPropagateToAllFrom(mma_result, -1);
@@ -975,25 +961,24 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   //           -8  -7  -6  -5  -4  -3  -2  -1
   //   [... iMNwo rKo iMw iNw rKw iMi iNi rKi]
   // Note the positions of the reduction dimensions has changed
-  std::cout << "  mma_result after schedulWarpTileWithReduction:\n    " << mma_result->toString() << std::endl;
 
   // Propagate warp tile to main loop and epilog/output tvs
+  std::vector<TensorView*> warp_tile_prop_tvs;
+  if (splitk_sum != nullptr) {
+    warp_tile_prop_tvs.push_back(splitk_sum);
+  } else if (smem_epilogue != nullptr) {
+    warp_tile_prop_tvs.push_back(smem_epilogue);
+  }
   scheduler_utils::BoundedDirectionalTransformPropagator::bothWays(
-      mma_result, -1, {acw_smem, bcw_smem}, {smem_epilogue});
-
-  std::cout << "  After bothWays:\n";
-  std::cout << "    mma_result " << mma_result->toString() << std::endl;
-  std::cout << "    smem_epilogue " << smem_epilogue->toString() << std::endl;
-  std::cout << "    splitk_sum " << (splitk_sum != nullptr ? splitk_sum->toString() : "nullptr") << std::endl;
+      mma_result, -1, {acw_smem, bcw_smem}, warp_tile_prop_tvs);
 
   // No (cross-CTA) split-K
   //   mma_result      [..., iMo iNo rKo rKwo iMwo iNwo iMw iNw iMi iNi rKi]
-  //   smem_epilogue   (unscheduled, same as original mma_result)
+  //   smem_epilogue   (unscheduled, same as original or current mma_result)
   //   splitk_sum      (nullptr)
   //
   // With split-K & no smem epilogue
   //   mma_result      [..., iMo iNo iKf rKg rKwo iMwo iNwo iMw iNw iMi iNi rKi]
-  //   smem_epilogue   (mma_result)
   //   splitk_sum      [..., iMo iNo rKf iMwo iNwo iMw iNw iMi iNi]
   //
   // With split-K & smem epilogue
@@ -1144,7 +1129,12 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
         scheduler_utils::BoundedDirectionalTransformPropagator::Options()
             .propagateParallelType()
             .propagateToBoundary());
-    smem_epilogue->axis(-1)->parallelize(ParallelType::Vectorize);
+    if (num_splitk_dims == 0) {
+      // TODO: Remove this once we are able to vectorize the splitk reduction
+      smem_epilogue->axis(-1)->parallelize(ParallelType::Vectorize);
+    } else {
+      splitk_sum->axis(-3)->parallelize(ParallelType::BIDz);
+    }
 
     for (auto [dc, d] : cached_and_forked_outputs) {
       // Schedule output tensor differently for better global memory access
@@ -1175,7 +1165,7 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     scheduleFusionInputsForEpilogue(roles_map, params.use_smem_epilogue);
   }
 
-  if (num_splitk_dims) {
+  if (num_splitk_dims && !params.use_smem_epilogue) {
     // Here we reorder splitk_sum so that the grid reduction in the z dimension
     // is placed last, ensuring that we can inline it with downstream tensors.
     //

From 4c71cee5e8961a03ca6a863871f32284fa8ba2e3 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 5 Jan 2024 14:34:28 +0000
Subject: [PATCH 36/66] Delint test, add printout of bank conflict

All the assertions about bank conflicts should actually be EXPECT
---
 test/test_gpu_tensorcore.cpp | 75 ++++++++++++++++++++----------------
 1 file changed, 42 insertions(+), 33 deletions(-)

diff --git a/test/test_gpu_tensorcore.cpp b/test/test_gpu_tensorcore.cpp
index d4ad99e1ec8..2cc65bee778 100644
--- a/test/test_gpu_tensorcore.cpp
+++ b/test/test_gpu_tensorcore.cpp
@@ -2393,48 +2393,57 @@ TEST_F(NVFuserTest, FusionAmpereMatmulSplitK_CUDA) {
   int M = 504, N = 136, K = 8096;
 
   for (auto layout : kAllSupportedMmaLayout) {
-  for (int splitk_factor : {1, 2}) {
-  for (int use_smem_epilogue : {false, true}) {
-    Fusion fusion;
-    FusionGuard fg(&fusion);
-    auto tv0 = makeContigTensor(2, DataType::Half);
-    auto tv1 = makeContigTensor(2, DataType::Half);
+    for (int splitk_factor : {1, 2}) {
+      for (int use_smem_epilogue : {false, true}) {
+        Fusion fusion;
+        FusionGuard fg(&fusion);
+        auto tv0 = makeContigTensor(2, DataType::Half);
+        auto tv1 = makeContigTensor(2, DataType::Half);
 
-    fusion.addInput(tv0);
-    fusion.addInput(tv1);
+        fusion.addInput(tv0);
+        fusion.addInput(tv1);
 
-    auto tv2 = matmul(tv0, tv1, layout, true);
+        auto tv2 = matmul(tv0, tv1, layout, true);
 
-    fusion.addOutput(tv2);
+        fusion.addOutput(tv2);
 
-    MatMulTileOptions gemm_tile;
-    gemm_tile.cta_tile = GemmTile(128, 128, 32);
-    gemm_tile.warp_tile = GemmTile(64, 64, 32);
-    gemm_tile.instruction_tile = GemmTile(16, 8, 16);
+        MatMulTileOptions gemm_tile;
+        gemm_tile.cta_tile = GemmTile(128, 128, 32);
+        gemm_tile.warp_tile = GemmTile(64, 64, 32);
+        gemm_tile.instruction_tile = GemmTile(16, 8, 16);
 
-    MatmulParams params;
-    params.mma_macro = MmaMacro::Ampere_16_8_16;
-    params.tile_sizes = gemm_tile;
-    params.splitk_factor = splitk_factor;
-    params.use_smem_epilogue = use_smem_epilogue;
+        MatmulParams params;
+        params.mma_macro = MmaMacro::Ampere_16_8_16;
+        params.tile_sizes = gemm_tile;
+        params.splitk_factor = splitk_factor;
+        params.use_smem_epilogue = use_smem_epilogue;
+        params.promote_prologue_smem_reuse = true;
 
-    std::cout << "\n\n#####    layout=" << toString(layout) << " use_smem_epilogue=" << use_smem_epilogue << " splitk_factor=" << splitk_factor << std::endl;
-    scheduleMatmul(&fusion, params);
+        std::cout << "\n\n#####    layout=" << toString(layout)
+                  << " use_smem_epilogue=" << use_smem_epilogue
+                  << " splitk_factor=" << splitk_factor << std::endl;
+        scheduleMatmul(&fusion, params);
 
-    auto inputs = matmulAtInput(M, N, K, layout);
+        auto inputs = matmulAtInput(M, N, K, layout);
 
-    FusionExecutor fe;
-    NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
-        7, 5, fe.compileFusion(&fusion, {inputs.first, inputs.second}));
-    ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
-    auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
-    auto tref = atMatmul(
-        inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
+        FusionExecutor fe;
+        NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
+            7, 5, fe.compileFusion(&fusion, {inputs.first, inputs.second}));
+        EXPECT_TRUE(getBankConflictInfo(fe.kernel()).empty());
+        std::cout << "BankConflictInfo:" << std::endl;
+        for (auto [expr, conflictways] : getBankConflictInfo(fe.kernel())) {
+          std::cout << "  " << expr->toString();
+          std::cout << "    " << conflictways.first << "  "
+                    << conflictways.second << std::endl;
+        }
+        auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
+        auto tref = atMatmul(
+            inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);
 
-    // Relax tolerance for larger sum due to large K
-    NVF_CHECK(cg_outputs[0].allclose(tref, 1e-6 * K, 1e-6 * K));
-  }
-  }
+        // Relax tolerance for larger sum due to large K
+        NVF_CHECK(cg_outputs[0].allclose(tref, 1e-6 * K, 1e-6 * K));
+      }
+    }
   }
 }
 

From 96fa361f5637ca783fdf1e17ced0afd1ebc3a345 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 5 Jan 2024 16:54:37 +0000
Subject: [PATCH 37/66] Undo special handling for bothWays

Without splitk, we want to propagate like

  ... <- mma_result -> smem_epilogue

This terminates at smem_epilogue if we provide `{smem_epilogue}` as the
forward boundary.

Previously, I was modifying this to be `splitk_sum` when we had both
splitk_sum and smem_epilogue. However, that meant we were propagating

  ... <- mma_result -> smem_epilogue -> splitk_sum

terminating before splitk_sum. But this was transforming smem_epilogue,
which was undesired and led to a lack of swizzling in smem_epilogue.
Now, we simply revert to the previous case. In fact, we could probably
just propagate backwards and not use smem_epilogue at all since it seems
we never want to propagate forward here...

Note that if use_smem_epilogue == false, then smem_epilogue ==
mma_result are the same pointer so there will be no forward propagation
in that case.
---
 csrc/scheduler/matmul.cpp | 13 +------------
 1 file changed, 1 insertion(+), 12 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 99f2db66323..58cd2d81f0f 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -954,23 +954,12 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
 
   // Schedule warp tile
   mma_utils::scheduleWarpTileWithReduction(mma_result, gemm_tile);
-  // Single warp k:
   //          -8   -7   -6  -5  -4  -3  -2  -1
   //   [... rKwo iMwo iNwo iMw iNw iMi iNi rKi]
-  // If instead K is split across warps (unused):
-  //           -8  -7  -6  -5  -4  -3  -2  -1
-  //   [... iMNwo rKo iMw iNw rKw iMi iNi rKi]
-  // Note the positions of the reduction dimensions has changed
 
   // Propagate warp tile to main loop and epilog/output tvs
-  std::vector<TensorView*> warp_tile_prop_tvs;
-  if (splitk_sum != nullptr) {
-    warp_tile_prop_tvs.push_back(splitk_sum);
-  } else if (smem_epilogue != nullptr) {
-    warp_tile_prop_tvs.push_back(smem_epilogue);
-  }
   scheduler_utils::BoundedDirectionalTransformPropagator::bothWays(
-      mma_result, -1, {acw_smem, bcw_smem}, warp_tile_prop_tvs);
+      mma_result, -1, {acw_smem, bcw_smem}, {smem_epilogue});
 
   // No (cross-CTA) split-K
   //   mma_result      [..., iMo iNo rKo rKwo iMwo iNwo iMw iNw iMi iNi rKi]

From 5ee942129ece37cbd089ffc44769b6e3823fc0dc Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 8 Jan 2024 18:36:51 +0000
Subject: [PATCH 38/66] Vectorize smem store

This resolves bank conflicts. Note that there is an inner-most iBIDz
dimension which was tripping up VectorizeValidator, so a patch is
included.
---
 csrc/device_lower/validation.cpp | 10 ++++++++++
 csrc/scheduler/matmul.cpp        |  6 ++----
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/csrc/device_lower/validation.cpp b/csrc/device_lower/validation.cpp
index 9cb4cc1200a..f28fe3cc51b 100644
--- a/csrc/device_lower/validation.cpp
+++ b/csrc/device_lower/validation.cpp
@@ -373,6 +373,16 @@ class VectorizeValidator : public OptInDispatch {
       if (r_id->isReduction() || r_id->isBroadcast()) {
         continue;
       }
+      if ((tv->getMemoryType() == MemoryType::Shared ||
+           tv->getMemoryType() == MemoryType::Local) &&
+          r_id->isBlockDim()) {
+        // Inner-most parallelized dimensions don't count in allocation of
+        // shared and local tensors.
+        continue;
+      }
+      if (tv->getMemoryType() == MemoryType::Local && r_id->isThreadDim()) {
+        continue;
+      }
       last_alloc_dim = r_id;
       last_alloc_dim_pos = i - 1;
       break;
diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 58cd2d81f0f..ad7131391e9 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -1118,10 +1118,8 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
         scheduler_utils::BoundedDirectionalTransformPropagator::Options()
             .propagateParallelType()
             .propagateToBoundary());
-    if (num_splitk_dims == 0) {
-      // TODO: Remove this once we are able to vectorize the splitk reduction
-      smem_epilogue->axis(-1)->parallelize(ParallelType::Vectorize);
-    } else {
+    smem_epilogue->axis(-1)->parallelize(ParallelType::Vectorize);
+    if (num_splitk_dims != 0) {
       splitk_sum->axis(-3)->parallelize(ParallelType::BIDz);
     }
 

From 21ece7cd51c215218f9d795dd46a12c47b1900a2 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 8 Jan 2024 18:42:57 +0000
Subject: [PATCH 39/66] Remove debug print

---
 test/test_gpu_tensorcore.cpp | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/test/test_gpu_tensorcore.cpp b/test/test_gpu_tensorcore.cpp
index 2cc65bee778..554d5a1005d 100644
--- a/test/test_gpu_tensorcore.cpp
+++ b/test/test_gpu_tensorcore.cpp
@@ -2419,9 +2419,6 @@ TEST_F(NVFuserTest, FusionAmpereMatmulSplitK_CUDA) {
         params.use_smem_epilogue = use_smem_epilogue;
         params.promote_prologue_smem_reuse = true;
 
-        std::cout << "\n\n#####    layout=" << toString(layout)
-                  << " use_smem_epilogue=" << use_smem_epilogue
-                  << " splitk_factor=" << splitk_factor << std::endl;
         scheduleMatmul(&fusion, params);
 
         auto inputs = matmulAtInput(M, N, K, layout);
@@ -2430,12 +2427,6 @@ TEST_F(NVFuserTest, FusionAmpereMatmulSplitK_CUDA) {
         NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
             7, 5, fe.compileFusion(&fusion, {inputs.first, inputs.second}));
         EXPECT_TRUE(getBankConflictInfo(fe.kernel()).empty());
-        std::cout << "BankConflictInfo:" << std::endl;
-        for (auto [expr, conflictways] : getBankConflictInfo(fe.kernel())) {
-          std::cout << "  " << expr->toString();
-          std::cout << "    " << conflictways.first << "  "
-                    << conflictways.second << std::endl;
-        }
         auto cg_outputs = fe.runFusion({inputs.first, inputs.second});
         auto tref = atMatmul(
             inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout);

From 8bc4d0e9a4211235c390fe5e97edbd95171eb558 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 8 Jan 2024 18:51:08 +0000
Subject: [PATCH 40/66] Use smem in all splitk tests (with/wo batch/bias)

---
 test/test_gpu_tensorcore.cpp | 226 +++++++++++++++++++----------------
 1 file changed, 124 insertions(+), 102 deletions(-)

diff --git a/test/test_gpu_tensorcore.cpp b/test/test_gpu_tensorcore.cpp
index 554d5a1005d..defe2973a37 100644
--- a/test/test_gpu_tensorcore.cpp
+++ b/test/test_gpu_tensorcore.cpp
@@ -2393,7 +2393,7 @@ TEST_F(NVFuserTest, FusionAmpereMatmulSplitK_CUDA) {
   int M = 504, N = 136, K = 8096;
 
   for (auto layout : kAllSupportedMmaLayout) {
-    for (int splitk_factor : {1, 2}) {
+    for (int splitk_factor : {2}) {
       for (int use_smem_epilogue : {false, true}) {
         Fusion fusion;
         FusionGuard fg(&fusion);
@@ -2449,48 +2449,55 @@ TEST_F(NVFuserTest, FusionAmpereMatmulSplitKBias_CUDA) {
   int M = 504, N = 136, K = 8096;
 
   for (auto layout : kAllSupportedMmaLayout) {
-    Fusion fusion;
-    FusionGuard fg(&fusion);
-    auto tv0 = makeContigTensor(2, DataType::Half);
-    auto tv1 = makeContigTensor(2, DataType::Half);
-    auto tv2 = makeContigTensor(1, DataType::Half);
+    for (int splitk_factor : {2}) {
+      for (int use_smem_epilogue : {false, true}) {
+        Fusion fusion;
+        FusionGuard fg(&fusion);
+        auto tv0 = makeContigTensor(2, DataType::Half);
+        auto tv1 = makeContigTensor(2, DataType::Half);
+        auto tv2 = makeContigTensor(1, DataType::Half);
 
-    fusion.addInput(tv0);
-    fusion.addInput(tv1);
-    fusion.addInput(tv2);
+        fusion.addInput(tv0);
+        fusion.addInput(tv1);
+        fusion.addInput(tv2);
 
-    auto tv3 = matmul(tv0, tv1, layout, true);
-    auto tv4 = broadcast(tv2, {false, true});
-    auto tv5 = add(tv3, tv4); // bias
+        auto tv3 = matmul(tv0, tv1, layout, true);
+        auto tv4 = broadcast(tv2, {false, true});
+        auto tv5 = add(tv3, tv4); // bias
 
-    fusion.addOutput(tv5);
+        fusion.addOutput(tv5);
 
-    MatMulTileOptions gemm_tile;
-    gemm_tile.cta_tile = GemmTile(128, 128, 32);
-    gemm_tile.warp_tile = GemmTile(64, 64, 32);
-    gemm_tile.instruction_tile = GemmTile(16, 8, 16);
+        MatMulTileOptions gemm_tile;
+        gemm_tile.cta_tile = GemmTile(128, 128, 32);
+        gemm_tile.warp_tile = GemmTile(64, 64, 32);
+        gemm_tile.instruction_tile = GemmTile(16, 8, 16);
 
-    MatmulParams params;
-    params.mma_macro = MmaMacro::Ampere_16_8_16;
-    params.tile_sizes = gemm_tile;
-    params.splitk_factor = 2;
-    scheduleMatmul(&fusion, params);
+        MatmulParams params;
+        params.mma_macro = MmaMacro::Ampere_16_8_16;
+        params.tile_sizes = gemm_tile;
+        params.splitk_factor = splitk_factor;
+        params.use_smem_epilogue = use_smem_epilogue;
+        params.promote_prologue_smem_reuse = true;
+
+        scheduleMatmul(&fusion, params);
 
-    auto [aten_a, aten_b] = matmulAtInput(M, N, K, layout);
-    at::Tensor aten_bias = at::randn({M}, aten_a.options());
-    std::vector<c10::IValue> inputs = {aten_a, aten_b, aten_bias};
+        auto [aten_a, aten_b] = matmulAtInput(M, N, K, layout);
+        at::Tensor aten_bias = at::randn({M}, aten_a.options());
+        std::vector<c10::IValue> inputs = {aten_a, aten_b, aten_bias};
 
-    FusionExecutor fe;
-    NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
-        7, 5, fe.compileFusion(&fusion, inputs));
-    ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
-    auto cg_outputs = fe.runFusion(inputs);
-    auto tref = atBiasEpilogue(
-        atMatmul(aten_a.to(at::kFloat), aten_b.to(at::kFloat), layout),
-        aten_bias);
+        FusionExecutor fe;
+        NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
+            7, 5, fe.compileFusion(&fusion, inputs));
+        ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
+        auto cg_outputs = fe.runFusion(inputs);
+        auto tref = atBiasEpilogue(
+            atMatmul(aten_a.to(at::kFloat), aten_b.to(at::kFloat), layout),
+            aten_bias);
 
-    // Relax tolerance for larger sum due to large K
-    NVF_CHECK(cg_outputs[0].allclose(tref, 1e-6 * K, 1e-6 * K));
+        // Relax tolerance for larger sum due to large K
+        NVF_CHECK(cg_outputs[0].allclose(tref, 1e-6 * K, 1e-6 * K));
+      }
+    }
   }
 }
 
@@ -2505,45 +2512,53 @@ TEST_F(NVFuserTest, FusionAmpereMatmulBatchSplitK_CUDA) {
   int B = 2, M = 504, N = 136, K = 2048;
 
   for (auto layout : kAllSupportedMmaLayout) {
-    Fusion fusion;
-    FusionGuard fg(&fusion);
-    auto tv0 = makeContigTensor(3, DataType::Half);
-    auto tv1 = makeContigTensor(3, DataType::Half);
+    for (int splitk_factor : {2}) {
+      for (int use_smem_epilogue : {false, true}) {
+        Fusion fusion;
+        FusionGuard fg(&fusion);
+        auto tv0 = makeContigTensor(3, DataType::Half);
+        auto tv1 = makeContigTensor(3, DataType::Half);
 
-    fusion.addInput(tv0);
-    fusion.addInput(tv1);
+        fusion.addInput(tv0);
+        fusion.addInput(tv1);
 
-    auto tv2 = matmul(tv0, tv1, layout, true);
+        auto tv2 = matmul(tv0, tv1, layout, true);
 
-    fusion.addOutput(tv2);
+        fusion.addOutput(tv2);
 
-    MatMulTileOptions gemm_tile;
-    gemm_tile.cta_tile = GemmTile(128, 128, 32);
-    gemm_tile.warp_tile = GemmTile(64, 64, 32);
-    gemm_tile.instruction_tile = GemmTile(16, 8, 16);
+        MatMulTileOptions gemm_tile;
+        gemm_tile.cta_tile = GemmTile(128, 128, 32);
+        gemm_tile.warp_tile = GemmTile(64, 64, 32);
+        gemm_tile.instruction_tile = GemmTile(16, 8, 16);
 
-    MatmulParams params;
-    params.mma_macro = MmaMacro::Ampere_16_8_16;
-    params.tile_sizes = gemm_tile;
-    params.splitk_factor = 2;
-    scheduleMatmul(&fusion, params);
+        MatmulParams params;
+        params.mma_macro = MmaMacro::Ampere_16_8_16;
+        params.tile_sizes = gemm_tile;
+        params.splitk_factor = splitk_factor;
+        params.use_smem_epilogue = use_smem_epilogue;
+        params.promote_prologue_smem_reuse = true;
 
-    at::Tensor aten_a =
-        matmulAtInput(layout, TensorMatmulPos::A, at::kHalf, M, N, K, B);
-    at::Tensor aten_b =
-        matmulAtInput(layout, TensorMatmulPos::B, at::kHalf, M, N, K, B);
+        scheduleMatmul(&fusion, params);
 
-    std::vector<c10::IValue> inputs = {aten_a, aten_b};
+        at::Tensor aten_a =
+            matmulAtInput(layout, TensorMatmulPos::A, at::kHalf, M, N, K, B);
+        at::Tensor aten_b =
+            matmulAtInput(layout, TensorMatmulPos::B, at::kHalf, M, N, K, B);
 
-    FusionExecutor fe;
-    NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
-        7, 5, fe.compileFusion(&fusion, inputs));
-    ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
-    auto cg_outputs = fe.runFusion(inputs);
-    auto tref = atMatmul(aten_a.to(at::kFloat), aten_b.to(at::kFloat), layout);
+        std::vector<c10::IValue> inputs = {aten_a, aten_b};
 
-    // Relax tolerance for larger sum due to large K
-    EXPECT_TRUE(cg_outputs[0].allclose(tref, 1e-6 * K, 1e-6 * K));
+        FusionExecutor fe;
+        NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
+            7, 5, fe.compileFusion(&fusion, inputs));
+        ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
+        auto cg_outputs = fe.runFusion(inputs);
+        auto tref =
+            atMatmul(aten_a.to(at::kFloat), aten_b.to(at::kFloat), layout);
+
+        // Relax tolerance for larger sum due to large K
+        EXPECT_TRUE(cg_outputs[0].allclose(tref, 1e-6 * K, 1e-6 * K));
+      }
+    }
   }
 }
 
@@ -2558,52 +2573,59 @@ TEST_F(NVFuserTest, FusionAmpereMatmulBatchSplitKBias_CUDA) {
   int B = 2, M = 504, N = 136, K = 2048;
 
   for (auto layout : kAllSupportedMmaLayout) {
-    Fusion fusion;
-    FusionGuard fg(&fusion);
-    auto tv0 = makeContigTensor(3, DataType::Half);
-    auto tv1 = makeContigTensor(3, DataType::Half);
-    auto tv2 = makeContigTensor(1, DataType::Half);
+    for (int splitk_factor : {2}) {
+      for (int use_smem_epilogue : {false, true}) {
+        Fusion fusion;
+        FusionGuard fg(&fusion);
+        auto tv0 = makeContigTensor(3, DataType::Half);
+        auto tv1 = makeContigTensor(3, DataType::Half);
+        auto tv2 = makeContigTensor(1, DataType::Half);
 
-    fusion.addInput(tv0);
-    fusion.addInput(tv1);
-    fusion.addInput(tv2);
+        fusion.addInput(tv0);
+        fusion.addInput(tv1);
+        fusion.addInput(tv2);
 
-    auto tv3 = matmul(tv0, tv1, layout, true);
-    auto tv4 = broadcast(tv2, {true, false, true});
-    auto tv5 = add(tv3, tv4);
+        auto tv3 = matmul(tv0, tv1, layout, true);
+        auto tv4 = broadcast(tv2, {true, false, true});
+        auto tv5 = add(tv3, tv4);
 
-    fusion.addOutput(tv5);
+        fusion.addOutput(tv5);
 
-    MatMulTileOptions gemm_tile;
-    gemm_tile.cta_tile = GemmTile(128, 128, 32);
-    gemm_tile.warp_tile = GemmTile(64, 64, 32);
-    gemm_tile.instruction_tile = GemmTile(16, 8, 16);
+        MatMulTileOptions gemm_tile;
+        gemm_tile.cta_tile = GemmTile(128, 128, 32);
+        gemm_tile.warp_tile = GemmTile(64, 64, 32);
+        gemm_tile.instruction_tile = GemmTile(16, 8, 16);
 
-    MatmulParams params;
-    params.mma_macro = MmaMacro::Ampere_16_8_16;
-    params.tile_sizes = gemm_tile;
-    params.splitk_factor = 2;
-    scheduleMatmul(&fusion, params);
+        MatmulParams params;
+        params.mma_macro = MmaMacro::Ampere_16_8_16;
+        params.tile_sizes = gemm_tile;
+        params.splitk_factor = splitk_factor;
+        params.use_smem_epilogue = use_smem_epilogue;
+        params.promote_prologue_smem_reuse = true;
 
-    at::Tensor aten_a =
-        matmulAtInput(layout, TensorMatmulPos::A, at::kHalf, M, N, K, B);
-    at::Tensor aten_b =
-        matmulAtInput(layout, TensorMatmulPos::B, at::kHalf, M, N, K, B);
-    at::Tensor aten_bias = at::randn({M}, aten_a.options());
+        scheduleMatmul(&fusion, params);
 
-    std::vector<c10::IValue> inputs = {aten_a, aten_b, aten_bias};
+        at::Tensor aten_a =
+            matmulAtInput(layout, TensorMatmulPos::A, at::kHalf, M, N, K, B);
+        at::Tensor aten_b =
+            matmulAtInput(layout, TensorMatmulPos::B, at::kHalf, M, N, K, B);
+        at::Tensor aten_bias = at::randn({M}, aten_a.options());
 
-    FusionExecutor fe;
-    NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
-        7, 5, fe.compileFusion(&fusion, inputs));
-    ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
-    auto cg_outputs = fe.runFusion(inputs);
-    auto tref = atBiasEpilogue(
-        atMatmul(aten_a.to(at::kFloat), aten_b.to(at::kFloat), layout),
-        aten_bias);
+        std::vector<c10::IValue> inputs = {aten_a, aten_b, aten_bias};
 
-    // Relax tolerance for larger sum due to large K
-    EXPECT_TRUE(cg_outputs[0].allclose(tref, 1e-6 * K, 1e-6 * K));
+        FusionExecutor fe;
+        NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK(
+            7, 5, fe.compileFusion(&fusion, inputs));
+        ASSERT_TRUE(getBankConflictInfo(fe.kernel()).empty());
+        auto cg_outputs = fe.runFusion(inputs);
+        auto tref = atBiasEpilogue(
+            atMatmul(aten_a.to(at::kFloat), aten_b.to(at::kFloat), layout),
+            aten_bias);
+
+        // Relax tolerance for larger sum due to large K
+        EXPECT_TRUE(cg_outputs[0].allclose(tref, 1e-6 * K, 1e-6 * K));
+      }
+    }
   }
 }
 

From 49edbfe83126b4b2176f24d523a3314e8f943caa Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Tue, 9 Jan 2024 10:27:07 -0500
Subject: [PATCH 41/66] [WIP] Add benchmarks for nanogpt bwd split-K sizes

---
 benchmark/matmul.cpp | 51 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 50 insertions(+), 1 deletion(-)

diff --git a/benchmark/matmul.cpp b/benchmark/matmul.cpp
index 4a1cc265f85..dfba864a2e7 100644
--- a/benchmark/matmul.cpp
+++ b/benchmark/matmul.cpp
@@ -250,6 +250,11 @@ MatmulParams getMatmulParams(
   params.double_buffer_options.double_buffer_smem_read = true;
   params.double_buffer_options.smem_double_buffer_stage = stage_number;
   params.splitk_factor = splitk_factor;
+  if (splitk_factor > 1) {
+    // always use smem_epilogue when doing splitk
+    params.promote_prologue_smem_reuse = true;
+    params.use_smem_epilogue = true;
+  }
 
   return params;
 }
@@ -496,7 +501,17 @@ static void NvFuserScheduler_MatmulSplitKReduction(
     {784, 72, 8},        \
     {784, 8, 72},        \
     /* {1, 1, 2048}, */  \
-    {1024, 1024, 1024}   \
+    {1024, 1024, 1024},  \
+    /* NanoGPT bwd sizes */  \
+    {1024, 2048, 4096},      \
+    {1024, 2048, 50304}     \
+  }
+
+#define SplitKSpecificShapes \
+  {                          \
+    /* NanoGPT bwd sizes */  \
+    {1024, 2048, 4096},      \
+    {1024, 2048, 50304}     \
   }
 // clang-format on
 
@@ -603,6 +618,22 @@ static void MatmulShapeWarpStageAutoSplitK(benchmark::internal::Benchmark* b) {
   }
 }
 
+// Use this for manual splitk.
+static void MatmulShapeWarpStageSpecificSplitK(
+    benchmark::internal::Benchmark* b) {
+  b->ArgNames({"M", "N", "K", "warps", "stages", "splitk_factor"});
+  for (long int num_warps : NumWarps) {
+    for (long int num_stages : NumStages) {
+      for (auto [m, n, k] :
+           std::vector<std::tuple<int, int, int>>(SplitKSpecificShapes)) {
+        for (auto splitk_factor : {2, 3, 4, 5, 6}) {
+          b->Args({m, n, k, num_warps, num_stages, splitk_factor});
+        }
+      }
+    }
+  }
+}
+
 #define EagerModeBenchmark(layout)                                         \
   BENCHMARK_CAPTURE(                                                       \
       Baseline_Matmul, eagermode_legacyshapes_##layout, MmaLayout::layout) \
@@ -683,9 +714,27 @@ static void MatmulShapeWarpStageAutoSplitK(benchmark::internal::Benchmark* b) {
       ->UseManualTime()                   \
       ->Apply(MatmulShapeWarpStageAutoSplitK);
 
+static void NvFuserScheduler_Matmul_Manual(
+    benchmark::State& benchmark_state,
+    MmaLayout layout) {
+  int splitk_factor = benchmark_state.range(5);
+  NvFuserScheduler_Matmul(
+      benchmark_state, layout, splitk_factor, /*partitionedk=*/false);
+}
+
+#define SpecificSplitKBenchmark(layout) \
+  BENCHMARK_CAPTURE(                    \
+      NvFuserScheduler_Matmul_Manual,   \
+      nvfuser_splitk_##layout,          \
+      MmaLayout::layout)                \
+      ->Unit(benchmark::kMicrosecond)   \
+      ->UseManualTime()                 \
+      ->Apply(MatmulShapeWarpStageSpecificSplitK);
+
 ForAllLayouts(EagerModeBenchmark);
 ForAllLayouts(NvfuserMatmulBenchmark);
 ForAllLayouts(AutoSplitKBenchmark);
+ForAllLayouts(SpecificSplitKBenchmark);
 ForAllLayouts(AutoPartitionedKBenchmark);
 
 // Note: SplitK Reduction benchmarks are parametrized only by M, N. The splitk

From bc50c8dfc9af84c7d98f0ebaca4328250ecb0cd4 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Thu, 14 Dec 2023 18:07:35 +0000
Subject: [PATCH 42/66] Vectorized serial grid reduction

---
 csrc/codegen.cpp                 | 13 +++++++----
 csrc/device_lower/validation.cpp |  8 +++++--
 runtime/grid_reduction.cu        | 26 +++++++++++++++-------
 test/test_serial_gridreduce.cpp  | 37 +++++++++++++++++++++-----------
 4 files changed, 58 insertions(+), 26 deletions(-)

diff --git a/csrc/codegen.cpp b/csrc/codegen.cpp
index 9b68cf4fdb5..6e11cf93d89 100644
--- a/csrc/codegen.cpp
+++ b/csrc/codegen.cpp
@@ -1693,11 +1693,16 @@ class CudaKernelGenerator : private kir::ConstIrVisitor {
         block_flags,
         ArgumentBuilder().arg("gridDim"));
 
+    int64_t vectorize_size = ir_utils::getVectorizeSize(out->view());
+
+    ArgumentBuilder template_args;
+    template_args.arg("/*vec_size=*/").append(std::to_string(vectorize_size));
+
     ArgumentBuilder func_args(block_nest_level_ + 1, kTab);
-    func_args.arg(gen(out));
-    func_args.arg(gen(grop->in()));
+    func_args.arg("&").append(gen(out));
+    func_args.arg("&").append(gen(grop->in()));
     func_args.arg(gen(grop->init()));
-    func_args.arg(gen(grop->serialReductionTensor()));
+    func_args.arg("&").append(gen(grop->serialReductionTensor()));
     func_args.arg(genReductionOp(op_type, out->dtype()));
 
     // Whether this is the first or last step
@@ -1720,7 +1725,7 @@ class CudaKernelGenerator : private kir::ConstIrVisitor {
       func_args.arg(read_pred);
     }
 
-    indent() << "reduction::serialReductionStep(\n";
+    indent() << "reduction::serialReductionStep<" << template_args << ">(\n";
     indent() << kTab << func_args << ");\n";
   }
 
diff --git a/csrc/device_lower/validation.cpp b/csrc/device_lower/validation.cpp
index f28fe3cc51b..8a3c47a81c9 100644
--- a/csrc/device_lower/validation.cpp
+++ b/csrc/device_lower/validation.cpp
@@ -283,6 +283,7 @@ class VectorizeValidator : public OptInDispatch {
 
   void handle(Split* s) final {
     if (s->outer() == vectorized_id_) {
+      std::cout << "Split is invalid: " << s->toString();
       is_valid = false;
     } else if (s->inner() == vectorized_id_) {
       vectorized_id_ = s->in();
@@ -314,6 +315,7 @@ class VectorizeValidator : public OptInDispatch {
     if (swizzle->outX() == vectorized_id_ || swizzle->inX() == vectorized_id_ ||
         swizzle->outY() == vectorized_id_ || swizzle->inY() == vectorized_id_) {
       // Do not (yet) allow vectorization across any swizzled id.
+      std::cout << "Swizzle is invalid: " << swizzle->toString();
       is_valid = false;
     }
   }
@@ -322,6 +324,7 @@ class VectorizeValidator : public OptInDispatch {
     if (swizzle->outX() == vectorized_id_ || swizzle->inX() == vectorized_id_ ||
         swizzle->outY() == vectorized_id_ || swizzle->inY() == vectorized_id_) {
       // Do not (yet) allow vectorization across any swizzled id.
+      std::cout << "Swizzle is invalid: " << swizzle->toString();
       is_valid = false;
     }
   }
@@ -589,9 +592,10 @@ void validateAndCollectVectorizeInfo(Fusion* fusion) {
       }
     }
     if (has_vectorize_dim) {
+      Expr* def = tv->definition();
       NVF_ERROR(
-          tv->definition() == nullptr || tv->definition()->isA<LoadStoreOp>() ||
-              tv->definition()->isA<SliceOp>(),
+          def == nullptr || def->isA<LoadStoreOp>() ||
+              def->isA<SliceOp>() || (def->isA<ReductionOp>()), // && def->as<ReductionOp>()->SerialGridReductionRequested(),
           "Vectorized accesses cannot be inline with computation, they are only supported with a Set operation.",
           "TensorView: ",
           tv);
diff --git a/runtime/grid_reduction.cu b/runtime/grid_reduction.cu
index 76afb0038f9..aba070a6118 100644
--- a/runtime/grid_reduction.cu
+++ b/runtime/grid_reduction.cu
@@ -625,19 +625,19 @@ __device__ void gridReduceGroup(
 
 // This performs a single reduction step, combining a single element "in" with
 // a previous value "work". For a serial grid reduction, "work" resides in
-// global memory.
+// global memory, while "in" and "out" are in registers.
 //
 // If the write predicate is false, this function returns early (noop). If the
 // read predicate is false, "init" is used in place of "in".
 //
 // If first_step is false, "work" will be read and reduction_op will be called.
 // The result will be written back to "work" unless last_step is true.
-template <typename T, typename Func>
+template <int64_t vec_size, typename T, typename Func>
 __device__ void serialReductionStep(
-    T& out,
-    T in,
+    T* out,
+    T* in,
     T init,
-    volatile T& work,
+    volatile T* work,
     Func reduction_op,
     bool first_step,
     bool last_step,
@@ -646,12 +646,22 @@ __device__ void serialReductionStep(
   if (!write_pred) {
     return;
   }
-  out = read_pred ? in : init;
+  if (read_pred) {
+    loadGeneric<T, vec_size>(out, in);
+  }
   if (!first_step) {
-    reduction_op(out, work);
+    T work_reg[vec_size];
+    loadGenericVolatile<T, vec_size, false, true>(work_reg, work);
+#pragma unroll
+    for (int i = 0; i < vec_size; ++i) {
+      if (!read_pred) {
+        out[i] = init;
+      }
+      reduction_op(out[i], work_reg[i]);
+    }
   }
   if (!last_step) {
-    work = out;
+    loadGenericVolatile<T, vec_size, true, false>(work, out);
   }
 }
 
diff --git a/test/test_serial_gridreduce.cpp b/test/test_serial_gridreduce.cpp
index 43675299b68..95deed2a7a6 100644
--- a/test/test_serial_gridreduce.cpp
+++ b/test/test_serial_gridreduce.cpp
@@ -40,6 +40,7 @@ using SerialGridReductionTest = NVFuserTest;
 TEST_F(SerialGridReductionTest, CodegenNodes) {
   for (bool serial : {true, false}) {
     for (int64_t num_warps : {4, 8}) {
+    for (int64_t vec_size : {4, 2, 1}) {
       // B is size of inner serial loop. Outer loop is hardcoded at A=4
       // Here we set B to a small value of 8 instead of 32 (i.e. 128 elements
       // per thread), so that the non-serial compilation does not take too
@@ -54,7 +55,7 @@ TEST_F(SerialGridReductionTest, CodegenNodes) {
         int64_t blocks_z = 5;
         int64_t A = 4; // Size of outer serial loop
         int64_t H = blocks_z;
-        int64_t W = A * B * blocks_x * blocks_y * num_warps * 32;
+        int64_t W = A * B * blocks_x * blocks_y * num_warps * 32 * vec_size;
 
         // Unreduced dimensions should be concrete. Reduced dimension could be
         // symbolic, but is concrete here so that we can read tv0 to registers
@@ -75,21 +76,23 @@ TEST_F(SerialGridReductionTest, CodegenNodes) {
         //   [ iS{W}, rS{H} ]
         // then schedule as
         //   [ iBIDx{blocks_x}, iBIDy{blocks_y}, iS{A}, iS{B}, iTIDy{num_warps},
-        //   iTIDx{32}, rBIDz{blocks_z} ]
+        //   iTIDx{32}, iV{vec_size} rBIDz{blocks_z} ]
         auto tv2 = tv0->cacheAfter();
         auto tv3 = tv1->cacheBefore();
 
-        tv3->reorder({{1, 0}, {0, 1}}); // blocks_x*blocks_y*A*B*num_warps*32, H
-        tv3->split(0, 32); // blocks_x*blocks_y*A*B*num_warps, 32, H
-        tv3->split(0, num_warps); // blocks_x*blocks_y*A*B, num_warps, 32, H
-        tv3->split(0, B); // blocks_x*blocks_y*A, B, num_warps, 32, H
-        tv3->split(0, A); // blocks_x*blocks_y, A, B, num_warps, 32, H
-        tv3->split(0, blocks_y); // blocks_x, blocks_y, A, B, num_warps, 32, H
+        tv3->reorder({{1, 0}, {0, 1}}); // blocks_x*blocks_y*A*B*num_warps*32*vec_size, H
+        tv3->split(0, vec_size); // blocks_x*blocks_y*A*B*num_warps, 32, vec_size, H
+        tv3->split(0, 32); // blocks_x*blocks_y*A*B*num_warps, 32, vec_size, H
+        tv3->split(0, num_warps); // blocks_x*blocks_y*A*B, num_warps, 32, vec_size, H
+        tv3->split(0, B); // blocks_x*blocks_y*A, B, num_warps, 32, vec_size, H
+        tv3->split(0, A); // blocks_x*blocks_y, A, B, num_warps, 32, vec_size, H
+        tv3->split(0, blocks_y); // blocks_x, blocks_y, A, B, num_warps, 32, vec_size, H
         tv3->axis(0)->parallelize(ParallelType::BIDx);
         tv3->axis(1)->parallelize(ParallelType::BIDy);
         tv3->axis(4)->parallelize(ParallelType::TIDy);
         tv3->axis(5)->parallelize(ParallelType::TIDx);
-        tv3->axis(6)->parallelize(ParallelType::BIDz);
+        tv3->axis(6)->parallelize(ParallelType::Vectorize);
+        tv3->axis(7)->parallelize(ParallelType::BIDz);
         // Reorder to put parallel dims first for better inlining
         tv3->reorder({
             {4, 2},
@@ -106,8 +109,8 @@ TEST_F(SerialGridReductionTest, CodegenNodes) {
         // inlined with tv3, resulting in a separate loop to load tv0 into
         // registers (tv2).
         tv2->reorder({
-            {-2, -3},
-            {-3, -2},
+            {-3, -4},
+            {-4, -3},
         });
 
         inlineMost();
@@ -164,6 +167,10 @@ TEST_F(SerialGridReductionTest, CodegenNodes) {
                                       ->body(); // B
             // We will need the store op output TensorIndex
             LoadStoreOp* output_store_expr = B_scope.exprs()
+                                                 .back()
+                                                 ->as<kir::ForLoop>()  // vectorize loop
+                                                 ->body()
+                                                 .exprs()
                                                  .back()
                                                  ->as<kir::IfThenElse>()
                                                  ->thenBody()
@@ -171,7 +178,12 @@ TEST_F(SerialGridReductionTest, CodegenNodes) {
                                                  ->as<LoadStoreOp>();
             // bidz_scope is the scope containing the GridReduction expression
             kir::Scope& bidz_scope =
-                B_scope.exprs().at(4)->as<kir::ForLoop>()->body(); // BIDz
+                B_scope.exprs().at(4)->as<kir::ForLoop>() // vectorize loop
+                ->body()
+                .exprs()
+                .back()
+                ->as<kir::ForLoop>()
+                ->body(); // BIDz
             auto old_grop = bidz_scope.at(0)->as<kir::GridReduction>();
             // Store the TensorIndex for the output tensor T1_g, so that we can
             // re-use its index
@@ -253,6 +265,7 @@ TEST_F(SerialGridReductionTest, CodegenNodes) {
         }
       }
     }
+    }
   }
 }
 

From f740f1256fdff297dc8bcd789cf0fb53ed1e8e07 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Thu, 14 Dec 2023 18:39:35 +0000
Subject: [PATCH 43/66] Remove debug prints

---
 csrc/device_lower/validation.cpp | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/csrc/device_lower/validation.cpp b/csrc/device_lower/validation.cpp
index 8a3c47a81c9..01db5efd379 100644
--- a/csrc/device_lower/validation.cpp
+++ b/csrc/device_lower/validation.cpp
@@ -283,7 +283,6 @@ class VectorizeValidator : public OptInDispatch {
 
   void handle(Split* s) final {
     if (s->outer() == vectorized_id_) {
-      std::cout << "Split is invalid: " << s->toString();
       is_valid = false;
     } else if (s->inner() == vectorized_id_) {
       vectorized_id_ = s->in();
@@ -315,7 +314,6 @@ class VectorizeValidator : public OptInDispatch {
     if (swizzle->outX() == vectorized_id_ || swizzle->inX() == vectorized_id_ ||
         swizzle->outY() == vectorized_id_ || swizzle->inY() == vectorized_id_) {
       // Do not (yet) allow vectorization across any swizzled id.
-      std::cout << "Swizzle is invalid: " << swizzle->toString();
       is_valid = false;
     }
   }
@@ -324,7 +322,6 @@ class VectorizeValidator : public OptInDispatch {
     if (swizzle->outX() == vectorized_id_ || swizzle->inX() == vectorized_id_ ||
         swizzle->outY() == vectorized_id_ || swizzle->inY() == vectorized_id_) {
       // Do not (yet) allow vectorization across any swizzled id.
-      std::cout << "Swizzle is invalid: " << swizzle->toString();
       is_valid = false;
     }
   }

From 71bae3a58f332c09c3dfa120189cc3b9e690b300 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 18 Dec 2023 10:42:11 -0500
Subject: [PATCH 44/66] Use loadGlobalToLocal instead of loadGenericVolatile

This actually means we will insert st.global.cg as inline Asm instead of
the reinterpret_cast-based copy that loadGenericVolatile uses. This
fixed the error I was seeing with vec_size=8, indicating that
loadGenericVolatile is not obeying the same ordering semantics as we
need for serial reduction.
---
 runtime/grid_reduction.cu | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/runtime/grid_reduction.cu b/runtime/grid_reduction.cu
index aba070a6118..8a75410d9de 100644
--- a/runtime/grid_reduction.cu
+++ b/runtime/grid_reduction.cu
@@ -648,20 +648,22 @@ __device__ void serialReductionStep(
   }
   if (read_pred) {
     loadGeneric<T, vec_size>(out, in);
+  } else {
+#pragma unroll
+    for (int i = 0; i < vec_size; ++i) {
+      out[i] = init;
+    }
   }
   if (!first_step) {
     T work_reg[vec_size];
-    loadGenericVolatile<T, vec_size, false, true>(work_reg, work);
+    loadGlobalToLocal<T, vec_size, true, CacheOp::Global>(work_reg, work);
 #pragma unroll
     for (int i = 0; i < vec_size; ++i) {
-      if (!read_pred) {
-        out[i] = init;
-      }
       reduction_op(out[i], work_reg[i]);
     }
   }
   if (!last_step) {
-    loadGenericVolatile<T, vec_size, true, false>(work, out);
+    loadLocalToGlobal<T, vec_size, true>(work, out);
   }
 }
 

From 747b466c4220829163594faef7f101002e53c5b1 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Tue, 9 Jan 2024 13:26:43 -0500
Subject: [PATCH 45/66] Use smem_epilogue when possible in benchmarks

---
 benchmark/matmul.cpp | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/benchmark/matmul.cpp b/benchmark/matmul.cpp
index dfba864a2e7..992057133b4 100644
--- a/benchmark/matmul.cpp
+++ b/benchmark/matmul.cpp
@@ -16,6 +16,7 @@
 #include <scheduler/all_schedulers.h>
 #include <scheduler/matmul.h>
 #include <scheduler/matmul_heuristic.h>
+#include <scheduler/mma_utils.h>
 #include <utils.h>
 
 #include <benchmark/benchmark.h>
@@ -250,11 +251,14 @@ MatmulParams getMatmulParams(
   params.double_buffer_options.double_buffer_smem_read = true;
   params.double_buffer_options.smem_double_buffer_stage = stage_number;
   params.splitk_factor = splitk_factor;
-  if (splitk_factor > 1) {
-    // always use smem_epilogue when doing splitk
-    params.promote_prologue_smem_reuse = true;
-    params.use_smem_epilogue = true;
-  }
+  const auto [use_smem_epilogue, promote_prologue_smem_reuse] =
+      mma_utils::generateSharedMemoryEpilogueHeuristics(
+          gemm_tile,
+          stage_number,
+          {DataType::Half, DataType::Half, DataType::Float},
+          false);
+  params.promote_prologue_smem_reuse = use_smem_epilogue;
+  params.use_smem_epilogue = promote_prologue_smem_reuse;
 
   return params;
 }

From 0b13a59caacf40116148f6dc8235bc6bb52cada5 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 19 Jan 2024 16:28:41 +0000
Subject: [PATCH 46/66] Vectorize innermost non-trivial dimension.

Pattern is now identical to output
---
 csrc/scheduler/matmul.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index ad7131391e9..c3c80f0493b 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -1133,6 +1133,12 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
       scheduler_utils::BoundedDirectionalTransformPropagator::backward(
           d, -1, {smem_epilogue});
     }
+    if (num_splitk_dims != 0) {
+      // Now that transforms are propagated backward to smem_epilogue, which is
+      // before splitk_sum, we can vectorize the inner-most non-trivial
+      // dimension of splitk_sum
+      splitk_sum->axis(-2)->parallelize(ParallelType::Vectorize);
+    }
   } else {
     for (auto [dc, d] : cached_and_forked_outputs) {
       scheduler_utils::BoundedDirectionalTransformPropagator::forward(

From 21203ea4d8329b2d5a022204bf647e87668a6981 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 19 Jan 2024 11:43:44 -0500
Subject: [PATCH 47/66] Clean up check for use_smem_epilogue in benchmark

---
 benchmark/matmul.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/benchmark/matmul.cpp b/benchmark/matmul.cpp
index 992057133b4..b5937502bff 100644
--- a/benchmark/matmul.cpp
+++ b/benchmark/matmul.cpp
@@ -251,14 +251,14 @@ MatmulParams getMatmulParams(
   params.double_buffer_options.double_buffer_smem_read = true;
   params.double_buffer_options.smem_double_buffer_stage = stage_number;
   params.splitk_factor = splitk_factor;
-  const auto [use_smem_epilogue, promote_prologue_smem_reuse] =
+  std::tie(params.use_smem_epilogue, params.promote_prologue_smem_reuse) =
       mma_utils::generateSharedMemoryEpilogueHeuristics(
           gemm_tile,
           stage_number,
           {DataType::Half, DataType::Half, DataType::Float},
+          true,
+          true,
           false);
-  params.promote_prologue_smem_reuse = use_smem_epilogue;
-  params.use_smem_epilogue = promote_prologue_smem_reuse;
 
   return params;
 }

From aef008dd5a6d1031d37722f80f3885def1e9a894 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 19 Jan 2024 13:56:45 -0500
Subject: [PATCH 48/66] Set manual seed before generating inputs in
 SingleMatmulBase

---
 benchmark/matmul.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/benchmark/matmul.cpp b/benchmark/matmul.cpp
index b5937502bff..40c87c80337 100644
--- a/benchmark/matmul.cpp
+++ b/benchmark/matmul.cpp
@@ -145,6 +145,9 @@ static void SingleMatmulBase(
   int64_t n = benchmark_state.range(1);
   int64_t k = benchmark_state.range(2);
 
+  // inputs
+  at::manual_seed(0);
+
   // Tensor inputs
   auto inputs = matmulAtInput(m, n, k, layout);
   auto expected_output = atMatmul(
@@ -164,9 +167,6 @@ static void SingleMatmulBase(
 
   optimization::OptimizationPass<optimization::PreSegmenter>::runPass(fusion);
 
-  // inputs
-  at::manual_seed(0);
-
   KernelArgumentHolder args = KernelArgumentHolder::createKernelArgumentHolder(
       {inputs.first, inputs.second});
 

From e36740db7d8cf67a0db6b340d108d2cdc5b21346 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Fri, 19 Jan 2024 13:57:09 -0500
Subject: [PATCH 49/66] Disable double buffering smem writes

---
 benchmark/matmul.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/benchmark/matmul.cpp b/benchmark/matmul.cpp
index 40c87c80337..e04d8ad43b5 100644
--- a/benchmark/matmul.cpp
+++ b/benchmark/matmul.cpp
@@ -247,7 +247,7 @@ MatmulParams getMatmulParams(
   params.mma_macro = MmaMacro::Ampere_16_16_16;
   params.tile_sizes = gemm_tile;
   params.async_gmem_load_operands = true;
-  params.double_buffer_options.double_buffer_smem_write = true;
+  params.double_buffer_options.double_buffer_smem_write = false;
   params.double_buffer_options.double_buffer_smem_read = true;
   params.double_buffer_options.smem_double_buffer_stage = stage_number;
   params.splitk_factor = splitk_factor;

From 907e87781a4bdcf23c69ee215937a5fc883f720f Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 7 Feb 2024 13:31:15 +0000
Subject: [PATCH 50/66] Add smem_epilogue:{0,1} cases in benchmarks

---
 benchmark/matmul.cpp | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/benchmark/matmul.cpp b/benchmark/matmul.cpp
index e04d8ad43b5..6090130cfdc 100644
--- a/benchmark/matmul.cpp
+++ b/benchmark/matmul.cpp
@@ -258,7 +258,7 @@ MatmulParams getMatmulParams(
           {DataType::Half, DataType::Half, DataType::Float},
           true,
           true,
-          false);
+          true);
 
   return params;
 }
@@ -376,7 +376,8 @@ static void NvFuserScheduler_Matmul(
     benchmark::State& benchmark_state,
     MmaLayout layout,
     int splitk_factor = 1,
-    bool partitionedk = false) {
+    bool partitionedk = false,
+    bool use_smem_epilogue = false) {
   int num_warps = benchmark_state.range(3);
   int number_of_stage = benchmark_state.range(4);
 
@@ -390,6 +391,14 @@ static void NvFuserScheduler_Matmul(
 
   auto params = getMatmulParams(
       cta_tile, number_of_stage, layout, partitionedk ? 1 : splitk_factor);
+  if (use_smem_epilogue) {
+    if (!params.use_smem_epilogue) {
+      benchmark_state.SkipWithError("Insufficient shared mem for smem epilogue");
+    }
+  } else {
+    params.use_smem_epilogue = false;
+    params.promote_prologue_smem_reuse = false;
+  }
 
   NVFUSER_BENCHMARK_ARCH_SMEM_GUARD(
       8, 0, getSmemSize(cta_tile, number_of_stage), benchmark_state);
@@ -625,13 +634,15 @@ static void MatmulShapeWarpStageAutoSplitK(benchmark::internal::Benchmark* b) {
 // Use this for manual splitk.
 static void MatmulShapeWarpStageSpecificSplitK(
     benchmark::internal::Benchmark* b) {
-  b->ArgNames({"M", "N", "K", "warps", "stages", "splitk_factor"});
+  b->ArgNames({"M", "N", "K", "warps", "stages", "splitk_factor", "smem_epilogue"});
   for (long int num_warps : NumWarps) {
     for (long int num_stages : NumStages) {
       for (auto [m, n, k] :
            std::vector<std::tuple<int, int, int>>(SplitKSpecificShapes)) {
         for (auto splitk_factor : {2, 3, 4, 5, 6}) {
-          b->Args({m, n, k, num_warps, num_stages, splitk_factor});
+          for (bool use_smem_epilogue: {false, true}) {
+            b->Args({m, n, k, num_warps, num_stages, splitk_factor, use_smem_epilogue});
+          }
         }
       }
     }
@@ -722,8 +733,9 @@ static void NvFuserScheduler_Matmul_Manual(
     benchmark::State& benchmark_state,
     MmaLayout layout) {
   int splitk_factor = benchmark_state.range(5);
+  bool use_smem_epilogue = benchmark_state.range(6);
   NvFuserScheduler_Matmul(
-      benchmark_state, layout, splitk_factor, /*partitionedk=*/false);
+      benchmark_state, layout, splitk_factor, /*partitionedk=*/false, use_smem_epilogue);
 }
 
 #define SpecificSplitKBenchmark(layout) \
@@ -737,9 +749,9 @@ static void NvFuserScheduler_Matmul_Manual(
 
 ForAllLayouts(EagerModeBenchmark);
 ForAllLayouts(NvfuserMatmulBenchmark);
-ForAllLayouts(AutoSplitKBenchmark);
+//ForAllLayouts(AutoSplitKBenchmark);
 ForAllLayouts(SpecificSplitKBenchmark);
-ForAllLayouts(AutoPartitionedKBenchmark);
+//ForAllLayouts(AutoPartitionedKBenchmark);
 
 // Note: SplitK Reduction benchmarks are parametrized only by M, N. The splitk
 // factor is deduced automatically from N

From b04a23f482b0134d5dc5a5fc7439770f486fb1b1 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 12 Feb 2024 19:39:20 +0000
Subject: [PATCH 51/66] Split inner axis if vec width is too large

For example, if we use vec_size 8 for a Half epilogue, but we are
reducing in Float, we need to convert to vec_size 4 for the reduction.
---
 csrc/scheduler/matmul.cpp | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 5303063abc4..e938b816fcd 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -1137,7 +1137,21 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
       // Now that transforms are propagated backward to smem_epilogue, which is
       // before splitk_sum, we can vectorize the inner-most non-trivial
       // dimension of splitk_sum
+      Val* vec_ext = splitk_sum->axis(-2)->extent();
+      NVF_ERROR(vec_ext->isConstInt());
+      int64_t vec_ext_int = vec_ext->evaluate().as<int64_t>();
+      if (vec_ext_int * dataTypeSize(splitk_sum->dtype()) > 16) {
+        // NOTE: We might encounter an illegal vectorization size if we are using
+        // Float for this reduction and Half for output. So here we first check
+        // whether the vectorize size is at most 16 bytes. If not , we split into
+        // an unrolled loop that will do multiple vectorized reads/writes instead.
+        // Note that we reorder such that the axes are in order UR TIDx V.
+        splitk_sum->split(-2, 16 / dataTypeSize(splitk_sum->dtype()), /*inner_split=*/true);
+        splitk_sum->axis(-3)->parallelize(ParallelType::Unroll);
+        splitk_sum->reorder({{-4, -3}});
+      }
       splitk_sum->axis(-2)->parallelize(ParallelType::Vectorize);
+      splitk_sum->axis(-3)->parallelize(ParallelType::TIDx);
     }
   } else {
     for (auto [dc, d] : cached_and_forked_outputs) {

From 3fcdfed472db2943c4c704c3c8ddfd090e401c8a Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 12 Feb 2024 19:53:44 +0000
Subject: [PATCH 52/66] Delint

---
 csrc/scheduler/matmul.cpp | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index e938b816fcd..48253c06692 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -1141,12 +1141,14 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
       NVF_ERROR(vec_ext->isConstInt());
       int64_t vec_ext_int = vec_ext->evaluate().as<int64_t>();
       if (vec_ext_int * dataTypeSize(splitk_sum->dtype()) > 16) {
-        // NOTE: We might encounter an illegal vectorization size if we are using
-        // Float for this reduction and Half for output. So here we first check
-        // whether the vectorize size is at most 16 bytes. If not , we split into
-        // an unrolled loop that will do multiple vectorized reads/writes instead.
-        // Note that we reorder such that the axes are in order UR TIDx V.
-        splitk_sum->split(-2, 16 / dataTypeSize(splitk_sum->dtype()), /*inner_split=*/true);
+        // NOTE: We might encounter an illegal vectorization size if we are
+        // using Float for this reduction and Half for output. So here we first
+        // check whether the vectorize size is at most 16 bytes. If not , we
+        // split into an unrolled loop that will do multiple vectorized
+        // reads/writes instead. Note that we reorder such that the axes are in
+        // order UR TIDx V.
+        splitk_sum->split(
+            -2, 16 / dataTypeSize(splitk_sum->dtype()), /*inner_split=*/true);
         splitk_sum->axis(-3)->parallelize(ParallelType::Unroll);
         splitk_sum->reorder({{-4, -3}});
       }

From 687c13abaa7cb59d8ace92737d47cbd63cf3a6d7 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 12 Feb 2024 20:15:47 +0000
Subject: [PATCH 53/66] Delint benchmark

---
 benchmark/matmul.cpp | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/benchmark/matmul.cpp b/benchmark/matmul.cpp
index 3bf1254a936..d05889c1302 100644
--- a/benchmark/matmul.cpp
+++ b/benchmark/matmul.cpp
@@ -403,7 +403,8 @@ static void NvFuserScheduler_Matmul(
       cta_tile, number_of_stage, layout, partitionedk ? 1 : splitk_factor);
   if (use_smem_epilogue) {
     if (!params.use_smem_epilogue) {
-      benchmark_state.SkipWithError("Insufficient shared mem for smem epilogue");
+      benchmark_state.SkipWithError(
+          "Insufficient shared mem for smem epilogue");
     }
   } else {
     params.use_smem_epilogue = false;
@@ -665,14 +666,22 @@ static void MatmulShapeWarpStageAutoSplitK(benchmark::internal::Benchmark* b) {
 // Use this for manual splitk.
 static void MatmulShapeWarpStageSpecificSplitK(
     benchmark::internal::Benchmark* b) {
-  b->ArgNames({"M", "N", "K", "warps", "stages", "splitk_factor", "smem_epilogue"});
+  b->ArgNames(
+      {"M", "N", "K", "warps", "stages", "splitk_factor", "smem_epilogue"});
   for (long int num_warps : NumWarps) {
     for (long int num_stages : NumStages) {
       for (auto [m, n, k] :
            std::vector<std::tuple<int, int, int>>(SplitKSpecificShapes)) {
         for (auto splitk_factor : {2, 3, 4, 5, 6}) {
-          for (bool use_smem_epilogue: {false, true}) {
-            b->Args({m, n, k, num_warps, num_stages, splitk_factor, use_smem_epilogue});
+          for (bool use_smem_epilogue : {false, true}) {
+            b->Args(
+                {m,
+                 n,
+                 k,
+                 num_warps,
+                 num_stages,
+                 splitk_factor,
+                 use_smem_epilogue});
           }
         }
       }
@@ -766,7 +775,11 @@ static void NvFuserScheduler_Matmul_Manual(
   int splitk_factor = benchmark_state.range(5);
   bool use_smem_epilogue = benchmark_state.range(6);
   NvFuserScheduler_Matmul(
-      benchmark_state, layout, splitk_factor, /*partitionedk=*/false, use_smem_epilogue);
+      benchmark_state,
+      layout,
+      splitk_factor,
+      /*partitionedk=*/false,
+      use_smem_epilogue);
 }
 
 #define SpecificSplitKBenchmark(layout) \

From 7f230825a3ad856e0fb01f7f09766c23f6ee7be3 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 12 Feb 2024 20:18:21 +0000
Subject: [PATCH 54/66] Guard against limited smem in test

---
 test/test_gpu_tensorcore.cpp | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/test/test_gpu_tensorcore.cpp b/test/test_gpu_tensorcore.cpp
index 5b686e53205..fb50029b06f 100644
--- a/test/test_gpu_tensorcore.cpp
+++ b/test/test_gpu_tensorcore.cpp
@@ -2520,8 +2520,24 @@ TEST_F(NVFuserTest, FusionAmpereMatmulSplitK_CUDA) {
         params.mma_macro = MmaMacro::Ampere_16_8_16;
         params.tile_sizes = gemm_tile;
         params.splitk_factor = splitk_factor;
-        params.use_smem_epilogue = use_smem_epilogue;
-        params.promote_prologue_smem_reuse = true;
+        if (use_smem_epilogue) {
+          std::tie(
+              params.use_smem_epilogue, params.promote_prologue_smem_reuse) =
+              mma_utils::generateSharedMemoryEpilogueHeuristics(
+                  gemm_tile,
+                  stage_number,
+                  {DataType::Half, DataType::Half, DataType::Float},
+                  true,
+                  true,
+                  true);
+          if (!params.use_smem_epilogue) {
+            std::cout
+                << "Skipping smem epilogue due to shared memory constraints on this device"
+                << std::endl;
+            continue;
+          }
+          params.promote_prologue_smem_reuse = true;
+        }
 
         scheduleMatmul(&fusion, params);
 

From b6b5287f0adb0bd7974bbdbf559da67c821dd439 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 12 Feb 2024 20:19:17 +0000
Subject: [PATCH 55/66] Fix missing stage number in test smem check

---
 test/test_gpu_tensorcore.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_gpu_tensorcore.cpp b/test/test_gpu_tensorcore.cpp
index fb50029b06f..134bb544d63 100644
--- a/test/test_gpu_tensorcore.cpp
+++ b/test/test_gpu_tensorcore.cpp
@@ -2525,7 +2525,7 @@ TEST_F(NVFuserTest, FusionAmpereMatmulSplitK_CUDA) {
               params.use_smem_epilogue, params.promote_prologue_smem_reuse) =
               mma_utils::generateSharedMemoryEpilogueHeuristics(
                   gemm_tile,
-                  stage_number,
+                  1,
                   {DataType::Half, DataType::Half, DataType::Float},
                   true,
                   true,

From af710d62ab7d241f0fe52bb22ddd5b11a322b7e0 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 11 Mar 2024 17:23:22 +0000
Subject: [PATCH 56/66] Fix batch split-K tests.

We should not expect to re-use prologue smem when we have batch &
split-K, because in that case the batch loop will exist in the kernel so
there will be no smem re-use in those cases. Note that this _should_ be
possible but requires a block sync at the end of the batch loop and we
currently do not support that pattern. See #1899
---
 test/test_gpu_tensorcore.cpp | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/test/test_gpu_tensorcore.cpp b/test/test_gpu_tensorcore.cpp
index c6a7a09e194..620526a02ae 100644
--- a/test/test_gpu_tensorcore.cpp
+++ b/test/test_gpu_tensorcore.cpp
@@ -2912,7 +2912,12 @@ TEST_F(NVFuserTest, FusionAmpereMatmulBatchSplitK_CUDA) {
         mma_utils::MmaDataTypes data_types = {
             DataType::Half, DataType::Half, DataType::Float};
         int64_t estimated_smem = mma_utils::computeExpectedSharedMemoryUsage(
-            params, data_types, true, true);
+            params,
+            data_types,
+            // NOTE: Batch split-K matmuls cannot currently re-use smem due to
+            // outer batch loop
+            /*smem_a_reuse_guaranteed=*/false,
+            /*smem_b_reuse_guaranteed=*/false);
         int64_t actual_smem = fe.lastLaunchParams().smem();
         EXPECT_EQ(estimated_smem, actual_smem);
       }
@@ -2989,7 +2994,12 @@ TEST_F(NVFuserTest, FusionAmpereMatmulBatchSplitKBias_CUDA) {
         mma_utils::MmaDataTypes data_types = {
             DataType::Half, DataType::Half, DataType::Float};
         int64_t estimated_smem = mma_utils::computeExpectedSharedMemoryUsage(
-            params, data_types, true, true);
+            params,
+            data_types,
+            // NOTE: Batch split-K matmuls cannot currently re-use smem due to
+            // outer batch loop
+            /*smem_a_reuse_guaranteed=*/false,
+            /*smem_b_reuse_guaranteed=*/false);
         int64_t actual_smem = fe.lastLaunchParams().smem();
         EXPECT_EQ(estimated_smem, actual_smem);
       }

From 27db63e24c200e20f647b0adacd3f9c12b4a7062 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 18 Mar 2024 16:48:53 +0000
Subject: [PATCH 57/66] Revert inadvertent change to benchmark

---
 benchmark/matmul.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/benchmark/matmul.cpp b/benchmark/matmul.cpp
index 67100983413..70234cd222a 100644
--- a/benchmark/matmul.cpp
+++ b/benchmark/matmul.cpp
@@ -247,7 +247,7 @@ MatmulParams getMatmulParams(
   params.mma_macro = MmaMacro::Ampere_16_16_16;
   params.tile_sizes = gemm_tile;
   params.async_gmem_load_operands = true;
-  params.double_buffer_options.double_buffer_smem_write = false;
+  params.double_buffer_options.double_buffer_smem_write = true;
   params.double_buffer_options.double_buffer_smem_read = true;
   params.double_buffer_options.smem_double_buffer_stage = stage_number;
   params.splitk_factor = splitk_factor;

From 3ebc7338c1b65310455a740474f9df9bdbcb495d Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Mon, 18 Mar 2024 16:50:25 +0000
Subject: [PATCH 58/66] Comment on bool options

---
 benchmark/matmul.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/benchmark/matmul.cpp b/benchmark/matmul.cpp
index 70234cd222a..7188a407315 100644
--- a/benchmark/matmul.cpp
+++ b/benchmark/matmul.cpp
@@ -256,9 +256,9 @@ MatmulParams getMatmulParams(
           gemm_tile,
           stage_number,
           {DataType::Half, DataType::Half, DataType::Float},
-          true,
-          true,
-          true);
+          /*smem_a_reuse_guaranteed=*/true,
+          /*smem_b_reuse_guaranteed=*/true,
+          /*ignore_occupancy_drop=*/true);
 
   return params;
 }

From c7cfcaf82d5ace79089543b453c5b403ffe9d834 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Tue, 19 Mar 2024 12:52:13 +0000
Subject: [PATCH 59/66] Move definition of smem_epilogue further down

---
 csrc/scheduler/matmul.cpp | 34 ++++++++++++++++------------------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index e03d496cafc..8e38abd4001 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -797,14 +797,6 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   // Setup accumulator register.
   auto mma_result = mma->out()->as<TensorView>();
 
-  // Unswizzle mma result in shared memory
-  // Note that if we are using split-K, we will set up this buffer after
-  // rfactoring the matmul, between the MmaOp and the ReductionOp, in order to
-  // take advantage of unswizzling during the grid reduction
-  auto smem_epilogue = (params.use_smem_epilogue && params.splitk_factor == 1)
-      ? mma_result->cacheAfter()
-      : mma_result;
-
   // TODO:
   //  Significant build out needed here
   //   for more flexibility and data type support.
@@ -873,6 +865,12 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   mma_utils::makeTile(mma_result, gemm_tile.cta_tile.toVector());
   // [..., Mo, No, Ko, Mi, Ni, Ki]
 
+  // Unswizzle mma result in shared memory
+  // Note that if we are using split-K, we will set up this buffer after
+  // rfactoring the matmul, between the MmaOp and the ReductionOp, in order to
+  // take advantage of unswizzling during the grid reduction
+  TensorView* smem_epilogue = mma_result;
+
   // Swizzle block tiles:
   if (params.grid_swizzle_factor != 1) {
     int factor = std::max(1, params.grid_swizzle_factor); // must be >=1
@@ -914,21 +912,21 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     // mma_result = [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
     // splitk_sum = [..., iMo, iNo, rKf, iMi, iNi]
 
-    if (params.use_smem_epilogue) {
-      //   splitk_sum = sum(mma_result)
-      // becomes
-      //   smem_epilogue = set(mma_result)
-      //   splitk_sum = sum(smem_epilogue)
-      smem_epilogue = mma_result->cacheAfter();
-
-      // smem_epilogue = [..., iMo, iNo, iKf, iMi, iNi]
-    }
-
     splitk_sum->definition()->as<ReductionOp>()->requestSerialGridReduction();
 
     num_splitk_dims = 1;
   }
 
+  if (params.use_smem_epilogue) {
+    // For split-K
+    //   splitk_sum = sum(mma_result)
+    // becomes
+    //   smem_epilogue = set(mma_result)
+    //   splitk_sum = sum(smem_epilogue)
+    smem_epilogue = mma_result->cacheAfter();
+    // smem_epilogue = [..., iMo, iNo, iKf, iMi, iNi]
+  }
+
   // No split-K, smem_epilogue = false
   //   mma_result      [..., iMo, iNo, rKo, iMi, iNi, rKi]
   //   smem_epilogue   (unscheduled, same as original mma_result)

From e6377f479f44a7346e624a51261fb455010e9e8b Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Tue, 19 Mar 2024 20:54:53 +0000
Subject: [PATCH 60/66] Refactor splitk_sum scheduling into separate function

I think this clarifies things quite a bit.
---
 csrc/scheduler/matmul.cpp | 206 ++++++++++++++++++--------------------
 1 file changed, 95 insertions(+), 111 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 8e38abd4001..81e32bacb41 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -687,6 +687,51 @@ void scheduleFusionInputsForEpilogue(
   }
 }
 
+void scheduleSplitKSum(
+    TensorView* splitk_sum,
+    const int num_batch_dims, // TODO: this should not be needed
+    bool use_smem_epilogue) {
+  if (splitk_sum == nullptr) {
+    // This indicates no split-K was used
+    return;
+  }
+
+  // Always use serial grid reduction for split-K sum
+  splitk_sum->definition()->as<ReductionOp>()->requestSerialGridReduction();
+
+  if (use_smem_epilogue) {
+    // Now that transforms are propagated backward to smem_epilogue, which is
+    // before splitk_sum, we can vectorize the inner-most non-trivial
+    // dimension of splitk_sum
+    Val* vec_ext = splitk_sum->axis(-2)->extent();
+    NVF_ERROR(vec_ext->isConstInt());
+    int64_t vec_ext_int = vec_ext->evaluate().as<int64_t>();
+    // Parallelize as [... Tx V Bz]
+    splitk_sum->axis(-1)->parallelize(ParallelType::BIDz);
+    splitk_sum->axis(-3)->parallelize(ParallelType::TIDx);
+    if (vec_ext_int * dataTypeSize(splitk_sum->dtype()) > 16) {
+      // NOTE: We might encounter an illegal vectorization size if we are
+      // using Float for this reduction and Half for output. So here we first
+      // check whether the vectorize size is at most 16 bytes. If not, then we
+      // split into an unrolled loop that will do multiple vectorized
+      // reads/writes instead. Note that we reorder such that the axes are in
+      // order UR TIDx V.
+      splitk_sum->split(
+          -2, 16 / dataTypeSize(splitk_sum->dtype()), /*inner_split=*/true);
+      splitk_sum->axis(-3)->parallelize(ParallelType::Unroll);
+      splitk_sum->axis(-2)->parallelize(ParallelType::Vectorize);
+      splitk_sum->reorder({{-4, -3}});
+      // In this case, we have [... UR Tx V Bz]
+    }
+    splitk_sum->axis(-2)->parallelize(ParallelType::Vectorize);
+  } else { // no smem epilogue
+    // Reorder to become [... rBz V]
+    splitk_sum->reorder({{-8, -2}});
+    // Vectorize inner-most dimension
+    splitk_sum->axis(-1)->parallelize(ParallelType::Vectorize);
+  }
+}
+
 } // namespace
 
 void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
@@ -912,8 +957,6 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     // mma_result = [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
     // splitk_sum = [..., iMo, iNo, rKf, iMi, iNi]
 
-    splitk_sum->definition()->as<ReductionOp>()->requestSerialGridReduction();
-
     num_splitk_dims = 1;
   }
 
@@ -927,12 +970,7 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     // smem_epilogue = [..., iMo, iNo, iKf, iMi, iNi]
   }
 
-  // No split-K, smem_epilogue = false
-  //   mma_result      [..., iMo, iNo, rKo, iMi, iNi, rKi]
-  //   smem_epilogue   (unscheduled, same as original mma_result)
-  //   splitk_sum      (nullptr)
-  //
-  // No split-K, smem_epilogue = true
+  // No split-K, smem_epilogue = false or true
   //   mma_result      [..., iMo, iNo, rKo, iMi, iNi, rKi]
   //   smem_epilogue   (unscheduled, same as original mma_result)
   //   splitk_sum      (nullptr)
@@ -958,27 +996,32 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   }
 
   // Schedule warp tile
+  // Incoming mma_result = [... iMo iNo (iKf) rKg iMi iNi rKi]
   mma_utils::scheduleWarpTileWithReduction(mma_result, gemm_tile);
-  //          -8   -7   -6  -5  -4  -3  -2  -1
-  //   [... rKwo iMwo iNwo iMw iNw iMi iNi rKi]
+  // After scheduling warp tile, the last three dimensions are split and
+  // rearranged:
+  //        -3 -2 -1
+  //   [...  M  N  K]
+  // maps to
+  //         -8  -7 -6  -5 -4 -3 -2 -1
+  //   [... Kwo Mwo Nwo Mw Nw Mi Ni Ki]
+  // so now
+  //                   -12 -11  -10   -9   -8   -7   -6  -5  -4   -3   -2   -1
+  // mma_result = [... iMo iNo (iKf) rKg rKwo iMwo iNwo iMw iNw iMin iNin rKin]
+  // splitk_sum = [... iMo iNo  rKf  iMi  iNi]
 
   // Propagate warp tile to main loop and epilog/output tvs
   scheduler_utils::BoundedDirectionalTransformPropagator::bothWays(
       mma_result, -1, {acw_smem, bcw_smem}, {smem_epilogue});
 
   // No (cross-CTA) split-K
-  //   mma_result      [..., iMo iNo rKo rKwo iMwo iNwo iMw iNw iMi iNi rKi]
+  //   mma_result      [..., iMo iNo rKo rKwo iMwo iNwo iMw iNw iMin iNin rKin]
   //   smem_epilogue   (unscheduled, same as original or current mma_result)
   //   splitk_sum      (nullptr)
   //
-  // With split-K & no smem epilogue
-  //   mma_result      [..., iMo iNo iKf rKg rKwo iMwo iNwo iMw iNw iMi iNi rKi]
-  //   splitk_sum      [..., iMo iNo rKf iMwo iNwo iMw iNw iMi iNi]
-  //
-  // With split-K & smem epilogue
-  //   mma_result      [..., iMo iNo iKf rKg rKwo iMwo iNwo iMw iNw iMi iNi rKi]
-  //   smem_epilogue   [..., iMo iNo iKf iMwo iNwo iMw iNw iMi iNi]
-  //   splitk_sum      [..., iMo iNo rKf iMwo iNwo iMw iNw iMi iNi]
+  // With split-K
+  //   mma_result      [... iMo iNo iKf  rKg rKwo iMwo iNwo iMw iNw iMin iNin
+  //   rKin] splitk_sum      [... iMo iNo rKf  iMi  iNi]
 
   // Schedule prolog:
   //   TODO: this section needs more configurability.
@@ -1026,6 +1069,14 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   };
   propagate_mma_input_schedule_to(acw_smem, bcw_smem);
 
+  // This does a split-reorder-merge swizzle of the last two M and N dimensions
+  // (and a possible final reduction dim).
+  // eg. [M64, N24, R]  -> [WarpGroup128, N3, M2, N2, Ro, R4, R2]
+  // Before
+  //   mma_result  [... iMo iNo (iKf) rKg rKwo iMwo iNwo iMw iNw iMin iNin rKin]
+  // After
+  //   mma_result  [... iMo iNo (iKf) rKg rKwo iMwo iNwo iMw iNw iMino iNino
+  //   iMin2 iNin2 rKino rKin4 rKin2]
   mma_result->applyMmaSwizzle(MmaOperand::Accumulator);
 
   // Set parallelization:
@@ -1061,39 +1112,35 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   //   - B: block
   //   - T: thread
   //   - S: serial. This will become a for loop in the generated kernel
-  //   - iMMA: unconracted axis in an MMA tensor core operation.
+  //   - iMMA: uncontracted axis in an MMA tensor core operation.
   //   - rMMA: contract in an MMA tensor core operation.
   //
-  // No (cross-CTA) split-K
-  //   mma_result      [..., iMo iNo rKo rKwo iMwo iNwo iMw iNw iMi iNi rKi]
-  //     nbatch +   1   2    3    4    5    6    7     8     9    10   11
-  //          -12 -11 -10   -9   -8   -7   -6   -5    -4    -3    -2   -1
-  //     [..., Mo, No, Kg, Kwo, Mwo, Nwo, Mwi, Nwi, MNi1, MNi2, MNi3,  Ki]
-  //     (iBz) iBx iBy rS   rS  iTz  iTy   iS   iS  iMMA   iTx  iMMA rMMA
+  // With split-K:
+  //   mma_result
+  //     nbatch +   1    2    3    4    5    6   7   8     9    10    11    12
+  //     13    14    15
+  //              -15  -14  -13  -12  -11  -10  -9  -8    -7    -6    -5    -4
+  //              -3    -2    -1
+  //     [... iMo iNo (iKf) rKg rKwo iMwo iNwo iMw iNw iMino iNino iMin2 iNin2
+  //     rKino rKin4 rKin2]
+  //          iBx iBy  iBz   rS   rS  iTz  iTy  iS  iS  iTx   iMMA  iMMA  iMMA
+  //          rMMA  rMMA  rMMA
   //   smem_epilogue   (unscheduled, same as original mma_result)
   //   splitk_sum      (nullptr)
   //
-  // With split-K & no smem epilogue
-  //   mma_result      [..., iMo iNo iKf rKg rKwo iMwo iNwo iMw iNw iMi iNi rKi]
-  //   smem_epilogue   (mma_result)
-  //   splitk_sum      [..., iMo iNo rKf iMwo iNwo iMw iNw iMi iNi]
-  //
-  // With split-K & smem epilogue
-  //   mma_result      [..., iMo iNo iKf rKg rKwo iMwo iNwo iMw iNw iMi iNi rKi]
-  //   smem_epilogue   [..., iMo iNo iKf iMwo iNwo iMw iNw iMi iNi]
-  //   splitk_sum      [..., iMo iNo rKf iMwo iNwo iMw iNw iMi iNi]
-  //
-  //  with splitk:
-  // nbatch +   1   2   3   4    5    6    7    8     9    10    11   12
-  //      -13 -12 -11 -10  -9   -8   -7   -6   -5    -4    -3    -2   -1
-  // [..., Mo, No, Kf, Kg, Kw, Mwo, Nwo, Mwi, Nwi, MNi1, MNi2, MNi3,  Ki]
-  //  (iS) iBx iBy rBz rS  rS  iTz  iTy   iS   iS  iMMA   iTx  iMMA rMMA
-  //
-  //  without splitk:
-  // nbatch +   1       2   3    4    5    6    7     8     9    10   11
-  //      -12 -11     -10  -9   -8   -7   -6   -5    -4    -3    -2   -1
-  // [..., Mo, No,     Kg, Kw, Mwo, Nwo, Mwi, Nwi, MNi1, MNi2, MNi3,  Ki]
-  // (iBz) iBx iBy     rS  rS  iTz  iTy   iS   iS  iMMA   iTx  iMMA rMMA
+  // Without split-K:
+  //   mma_result
+  //     nbatch +   1    2    3    4    5   6   7    8      9    10    11    12
+  //     13    14
+  //              -14  -13  -12  -11  -10  -9  -8   -7     -6    -5    -4    -3
+  //              -2    -1
+  //     [... iMo iNo  rKg rKwo iMwo iNwo iMw iNw iMino iNino iMin2 iNin2 rKino
+  //     rKin4 rKin2]
+  //    (iBz) iBx iBy   rS   rS  iTz  iTy  iS  iS  iTx   iMMA  iMMA  iMMA  rMMA
+  //    rMMA  rMMA
+  //   smem_epilogue   (unscheduled, same as original mma_result)
+  //   splitk_sum
+  //     [... iMo iNo rKf  iMi  iNi]
 
   // When we have both batch dims and splitk, parallelize splitk only.
   // If we only have batch dim, parallelize the batch dim.
@@ -1140,9 +1187,6 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
             .propagateParallelType()
             .propagateToBoundary());
     smem_epilogue->axis(-1)->parallelize(ParallelType::Vectorize);
-    if (num_splitk_dims != 0) {
-      splitk_sum->axis(-3)->parallelize(ParallelType::BIDz);
-    }
 
     for (auto [dc, d] : cached_outputs) {
       // Schedule output tensor differently for better global memory access
@@ -1154,28 +1198,6 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
       scheduler_utils::BoundedDirectionalTransformPropagator::backward(
           d, -1, {smem_epilogue});
     }
-    if (num_splitk_dims != 0) {
-      // Now that transforms are propagated backward to smem_epilogue, which is
-      // before splitk_sum, we can vectorize the inner-most non-trivial
-      // dimension of splitk_sum
-      Val* vec_ext = splitk_sum->axis(-2)->extent();
-      NVF_ERROR(vec_ext->isConstInt());
-      int64_t vec_ext_int = vec_ext->evaluate().as<int64_t>();
-      if (vec_ext_int * dataTypeSize(splitk_sum->dtype()) > 16) {
-        // NOTE: We might encounter an illegal vectorization size if we are
-        // using Float for this reduction and Half for output. So here we first
-        // check whether the vectorize size is at most 16 bytes. If not , we
-        // split into an unrolled loop that will do multiple vectorized
-        // reads/writes instead. Note that we reorder such that the axes are in
-        // order UR TIDx V.
-        splitk_sum->split(
-            -2, 16 / dataTypeSize(splitk_sum->dtype()), /*inner_split=*/true);
-        splitk_sum->axis(-3)->parallelize(ParallelType::Unroll);
-        splitk_sum->reorder({{-4, -3}});
-      }
-      splitk_sum->axis(-2)->parallelize(ParallelType::Vectorize);
-      splitk_sum->axis(-3)->parallelize(ParallelType::TIDx);
-    }
   } else {
     for (auto [dc, d] : cached_outputs) {
       scheduler_utils::BoundedDirectionalTransformPropagator::forward(
@@ -1195,45 +1217,7 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     scheduleFusionInputsForEpilogue(roles_map, params.use_smem_epilogue);
   }
 
-  if (num_splitk_dims && !params.use_smem_epilogue) {
-    // Here we reorder splitk_sum so that the grid reduction in the z dimension
-    // is placed last, ensuring that we can inline it with downstream tensors.
-    //
-    // Epilogue tensors:
-    // nbatch +   1    2   3     4    5     6     7     8
-    // [...  Mo  No  Mwo  Nwo  Mwi  Nwi  MNi1  MNi2  MNi3]
-    //  (iS) iBx iBy iTz  iTy   iS   iS    iS   iTx    iV
-    //
-    // Before reordering, splitk_sum is similar to mma_result but with the
-    // reduction axes removed. The Kf dimension causes inlining between nbatch
-    // + 1 and nbatch + 2.
-    // nbatch +   1    2    3    4    5    6     7     8     9
-    // [...  Mo  No   Kf  Mwo  Nwo  Mwi  Nwi  MNi1  MNi2  MNi3]
-    //  (iS) iBx iBy rBz  iTz  iTy   iS   iS    iS   iTx    iS
-    //
-    // splitk_sum (after the reordering below)
-    // nbatch +   1    2    3    4    5     6     7     8    9
-    // [...  Mo  No  Mwo  Nwo  Mwi  Nwi  MNi1  MNi2  MNi3   Kf]
-    //  (iS) iBx iBy iTz  iTy   iS   iS    iS   iTx    iS  rBz
-    //
-    // This reordering step lets us inline all but the last dim MNi3 (position
-    // nbatch + 7) which might be vectorized.
-    //
-    // NOTE: we need to do this reorder after the propagation above so that it
-    // doesn't get reset.
-    splitk_sum->reorder({
-        {num_batch_dims + 2, num_batch_dims + 9},
-        {num_batch_dims + 3, num_batch_dims + 2},
-        {num_batch_dims + 4, num_batch_dims + 3},
-        {num_batch_dims + 5, num_batch_dims + 4},
-        {num_batch_dims + 6, num_batch_dims + 5},
-        {num_batch_dims + 7, num_batch_dims + 6},
-        {num_batch_dims + 8, num_batch_dims + 7},
-        {num_batch_dims + 9, num_batch_dims + 8},
-    });
-    // Vectorize inner-most dimension
-    splitk_sum->axis(-1)->parallelize(ParallelType::Vectorize);
-  }
+  scheduleSplitKSum(splitk_sum, num_batch_dims, params.use_smem_epilogue);
 
   // auto inline for all tensors except register tensors
   inlineMost(ir_utils::allTvsExcept(fusion, {acr, bcr, ab, bb}));

From c7f8bac24d8332c9e5648be160eb3c5334045abe Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Tue, 19 Mar 2024 23:09:29 +0000
Subject: [PATCH 61/66] Clean up comments

---
 csrc/scheduler/matmul.cpp | 31 ++++++++++---------------------
 1 file changed, 10 insertions(+), 21 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 81e32bacb41..1b947bbef1f 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -937,13 +937,13 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     }
   }
 
-  // [..., Mo, No, Ko, Mi, Ni, Ki]
+  // [..., iMo, iNo, rKo, iMi, iNi, rKi]
   int num_splitk_dims = 0;
   TensorView* splitk_sum = nullptr;
   if (params.splitk_factor != 1) {
     // Split Ko -> [rKf, rKg]
     mma_result->split(-4, params.splitk_factor, /*inner*/ false);
-    // After split [..., Mo, No, Kf, Kg, Mi, Ni, Ki]
+    // After split [..., iMo, iNo, rKf, rKg, iMi, iNi, rKi]
     // rFactor converts
     //   mma_result = mma(A, B, {/*Kf*/-5, /*Kg*/-4, /*Ki*/-1});
     // to
@@ -954,14 +954,18 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     splitk_sum = mma_result;
     mma_result = splitk_sum->rFactor({-4, -1});
 
-    // mma_result = [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
-    // splitk_sum = [..., iMo, iNo, rKf, iMi, iNi]
-
     num_splitk_dims = 1;
   }
 
+  // At this point we have the following schedule:
+  //   No split-K
+  //     mma_result      [..., iMo, iNo, rKo, iMi, iNi, rKi]
+  //   Split-K
+  //     mma_result      [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
+  //     splitk_sum      [..., iMo, iNo, rKf, iMi, iNi]
+
   if (params.use_smem_epilogue) {
-    // For split-K
+    // Note that for split-K
     //   splitk_sum = sum(mma_result)
     // becomes
     //   smem_epilogue = set(mma_result)
@@ -970,21 +974,6 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
     // smem_epilogue = [..., iMo, iNo, iKf, iMi, iNi]
   }
 
-  // No split-K, smem_epilogue = false or true
-  //   mma_result      [..., iMo, iNo, rKo, iMi, iNi, rKi]
-  //   smem_epilogue   (unscheduled, same as original mma_result)
-  //   splitk_sum      (nullptr)
-  //
-  // With split-K & no smem epilogue
-  //   mma_result      [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
-  //   smem_epilogue   (mma_result)
-  //   splitk_sum      [..., iMo, iNo, rKf, iMi, iNi]
-
-  // With split-K & smem epilogue
-  //   mma_result      [..., iMo, iNo, iKf, rKg, iMi, iNi, rKi]
-  //   smem_epilogue   [..., iMo, iNo, iKf, iMi, iNi]
-  //   splitk_sum      [..., iMo, iNo, rKf, iMi, iNi]
-
   // Propagate tiling globally
   scheduler_utils::transformPropagateToAllFrom(mma_result, -1);
 

From b20cc734b9c48cdaaebeeedd40e7fd85cd20bccc Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 20 Mar 2024 00:04:08 +0000
Subject: [PATCH 62/66] Simplify scheduleSplitKSum and fix comment formatting

Also fixes print for vectorization validation error.
---
 csrc/device_lower/validation.cpp |  4 ++--
 csrc/scheduler/matmul.cpp        | 31 ++++++++++++++++---------------
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/csrc/device_lower/validation.cpp b/csrc/device_lower/validation.cpp
index 61e9eae6fce..89c6d7472fb 100644
--- a/csrc/device_lower/validation.cpp
+++ b/csrc/device_lower/validation.cpp
@@ -411,9 +411,9 @@ class VectorizeValidator : public OptInDispatch {
           ", allocation domain: ",
           ir_utils::toString(tv->getMaybeAllocationDomain()),
           ", vectorized id: ",
-          validator.vectorized_id_,
+          validator.vectorized_id_->toString(),
           ", innermost id: ",
-          last_alloc_dim,
+          last_alloc_dim->toString(),
           ", contiguity: ",
           contiguity.has_value() ? (*contiguity ? "t" : "f") : "n");
     }
diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 1b947bbef1f..45f55b884a7 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -703,6 +703,8 @@ void scheduleSplitKSum(
     // Now that transforms are propagated backward to smem_epilogue, which is
     // before splitk_sum, we can vectorize the inner-most non-trivial
     // dimension of splitk_sum
+    //
+    // Note that the split-K reduction is the inner-most dimension.
     Val* vec_ext = splitk_sum->axis(-2)->extent();
     NVF_ERROR(vec_ext->isConstInt());
     int64_t vec_ext_int = vec_ext->evaluate().as<int64_t>();
@@ -719,17 +721,16 @@ void scheduleSplitKSum(
       splitk_sum->split(
           -2, 16 / dataTypeSize(splitk_sum->dtype()), /*inner_split=*/true);
       splitk_sum->axis(-3)->parallelize(ParallelType::Unroll);
-      splitk_sum->axis(-2)->parallelize(ParallelType::Vectorize);
       splitk_sum->reorder({{-4, -3}});
-      // In this case, we have [... UR Tx V Bz]
+      // In this case, we have [... iUR iTx rBz iS]
     }
-    splitk_sum->axis(-2)->parallelize(ParallelType::Vectorize);
+    splitk_sum->reorder({{-2, -1}});
   } else { // no smem epilogue
-    // Reorder to become [... rBz V]
-    splitk_sum->reorder({{-8, -2}});
-    // Vectorize inner-most dimension
-    splitk_sum->axis(-1)->parallelize(ParallelType::Vectorize);
+    // Reorder to place the split-K reduction innermost [... rBz iS]
+    splitk_sum->reorder({{-9, -2}});
   }
+  // Vectorize inner-most dimension [... (iUR iTx) rBz iS]
+  splitk_sum->axis(-1)->parallelize(ParallelType::Vectorize);
 }
 
 } // namespace
@@ -1106,14 +1107,14 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   //
   // With split-K:
   //   mma_result
-  //     nbatch +   1    2    3    4    5    6   7   8     9    10    11    12
-  //     13    14    15
-  //              -15  -14  -13  -12  -11  -10  -9  -8    -7    -6    -5    -4
-  //              -3    -2    -1
-  //     [... iMo iNo (iKf) rKg rKwo iMwo iNwo iMw iNw iMino iNino iMin2 iNin2
-  //     rKino rKin4 rKin2]
-  //          iBx iBy  iBz   rS   rS  iTz  iTy  iS  iS  iTx   iMMA  iMMA  iMMA
-  //          rMMA  rMMA  rMMA
+  //     nbatch +   1    2    3    4    5    6   7   8
+  //              -15  -14  -13  -12  -11  -10  -9  -8
+  //     [... iMo iNo (iKf) rKg rKwo iMwo iNwo iMw iNw     ...
+  //          iBx iBy  iBz   rS   rS  iTz  iTy  iS  iS
+  //                              9    10    11    12    13    14    15
+  //                             -7    -6    -5    -4    -3    -2    -1
+  //                    ...   iMino iNino iMin2 iNin2 rKino rKin4 rKin2]
+  //                            iTx  iMMA  iMMA  iMMA  rMMA  rMMA  rMMA
   //   smem_epilogue   (unscheduled, same as original mma_result)
   //   splitk_sum      (nullptr)
   //

From 26dc559ff205f6032ac34e1d5f1544d2510bcb9d Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 20 Mar 2024 00:06:17 +0000
Subject: [PATCH 63/66] Fix comments

---
 csrc/scheduler/matmul.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 45f55b884a7..811aecdf586 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -708,7 +708,6 @@ void scheduleSplitKSum(
     Val* vec_ext = splitk_sum->axis(-2)->extent();
     NVF_ERROR(vec_ext->isConstInt());
     int64_t vec_ext_int = vec_ext->evaluate().as<int64_t>();
-    // Parallelize as [... Tx V Bz]
     splitk_sum->axis(-1)->parallelize(ParallelType::BIDz);
     splitk_sum->axis(-3)->parallelize(ParallelType::TIDx);
     if (vec_ext_int * dataTypeSize(splitk_sum->dtype()) > 16) {
@@ -729,7 +728,7 @@ void scheduleSplitKSum(
     // Reorder to place the split-K reduction innermost [... rBz iS]
     splitk_sum->reorder({{-9, -2}});
   }
-  // Vectorize inner-most dimension [... (iUR iTx) rBz iS]
+  // Vectorize inner-most dimension [... (iUR iTx) rBz iV]
   splitk_sum->axis(-1)->parallelize(ParallelType::Vectorize);
 }
 

From 83ec2102293fba07c71bb3daf1ea88c8dd11f02d Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Wed, 20 Mar 2024 11:17:42 +0000
Subject: [PATCH 64/66] Fix comment

---
 csrc/scheduler/matmul.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index 811aecdf586..aa74db16d1b 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -725,7 +725,7 @@ void scheduleSplitKSum(
     }
     splitk_sum->reorder({{-2, -1}});
   } else { // no smem epilogue
-    // Reorder to place the split-K reduction innermost [... rBz iS]
+    // Reorder to place the split-K reduction next to innermost [... rBz iS]
     splitk_sum->reorder({{-9, -2}});
   }
   // Vectorize inner-most dimension [... (iUR iTx) rBz iV]

From 992925517179f1ffc4ea4110abc6ea042ed34f69 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Thu, 21 Mar 2024 13:22:00 +0000
Subject: [PATCH 65/66] Fix comment formatting

---
 csrc/scheduler/matmul.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index aa74db16d1b..18875ca1f3b 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -1009,8 +1009,8 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   //   splitk_sum      (nullptr)
   //
   // With split-K
-  //   mma_result      [... iMo iNo iKf  rKg rKwo iMwo iNwo iMw iNw iMin iNin
-  //   rKin] splitk_sum      [... iMo iNo rKf  iMi  iNi]
+  //   mma_result   [... iMo iNo iKf  rKg rKwo iMwo iNwo iMw iNw iMin iNin rKin]
+  //   splitk_sum   [... iMo iNo rKf  iMi  iNi]
 
   // Schedule prolog:
   //   TODO: this section needs more configurability.

From 590b9a82e1e8d54f693f591499558079acff1881 Mon Sep 17 00:00:00 2001
From: Jacob Hinkle <jhinkle@nvidia.com>
Date: Thu, 21 Mar 2024 18:54:40 +0000
Subject: [PATCH 66/66] Fix another clang-format wrapped comment

---
 csrc/scheduler/matmul.cpp | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/csrc/scheduler/matmul.cpp b/csrc/scheduler/matmul.cpp
index a4084e49e69..f4917a5e663 100644
--- a/csrc/scheduler/matmul.cpp
+++ b/csrc/scheduler/matmul.cpp
@@ -1064,8 +1064,8 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   // Before
   //   mma_result  [... iMo iNo (iKf) rKg rKwo iMwo iNwo iMw iNw iMin iNin rKin]
   // After
-  //   mma_result  [... iMo iNo (iKf) rKg rKwo iMwo iNwo iMw iNw iMino iNino
-  //   iMin2 iNin2 rKino rKin4 rKin2]
+  //   mma_result  [... iMo iNo (iKf) rKg rKwo iMwo iNwo iMw
+  //                              iNw iMino iNino iMin2 iNin2 rKino rKin4 rKin2]
   mma_result->applyMmaSwizzle(MmaOperand::Accumulator);
 
   // Set parallelization:
@@ -1119,14 +1119,14 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   //
   // Without split-K:
   //   mma_result
-  //     nbatch +   1    2    3    4    5   6   7    8      9    10    11    12
-  //     13    14
-  //              -14  -13  -12  -11  -10  -9  -8   -7     -6    -5    -4    -3
-  //              -2    -1
-  //     [... iMo iNo  rKg rKwo iMwo iNwo iMw iNw iMino iNino iMin2 iNin2 rKino
-  //     rKin4 rKin2]
-  //    (iBz) iBx iBy   rS   rS  iTz  iTy  iS  iS  iTx   iMMA  iMMA  iMMA  rMMA
-  //    rMMA  rMMA
+  //     nbatch +   1   2    3    4    5   6   7    8
+  //              -14 -13  -12  -11  -10  -9  -8   -7
+  //     [... iMo iNo rKg rKwo iMwo iNwo iMw iNw iMino
+  //    (iBz) iBx iBy  rS   rS  iTz  iTy  iS  iS  iTx
+  //                                   9    10    11     12    13    14
+  //                                  -6    -5    -4     -3    -2    -1
+  //                               iNino iMin2 iNin2  rKino rKin4 rKin2]
+  //                                iMMA  iMMA  iMMA   rMMA  rMMA  rMMA
   //   smem_epilogue   (unscheduled, same as original mma_result)
   //   splitk_sum
   //     [... iMo iNo rKf  iMi  iNi]