[TIR][Schedule] Add FuseReductionEpilogue primitive to fuse epilogue … (#18418)

Currently it is not possible to fuse an epilogue operation (e.g., bias
addition) into a reduction block's initialization statement. This
limitation prevents leveraging hardware-specific instructions that
support bias accumulation in vector ISAs, such as MACC
(multiply-accumulate with bias) instructions.

This commit implements a new schedule primitive
'fuse_reduction_epilogue' that addresses the problem described in:

https://discuss.tvm.apache.org/t/tir-problem-inlining-addition-into-matmul-block/18066

The primitive transforms the following pattern:

  Before:
    for i, j, k in T.grid(M, N, K):
        with T.block("matmul"):
            with T.init():
                temp[vi, vj] = 0
            temp[vi, vj] = temp[vi, vj] + A[vi, vk] * B[vj, vk]

    for i, j in T.grid(M, N):
        with T.block("bias_add"):
            D[vi, vj] = temp[vi, vj] + C[vi, vj]

  After:
    for i, j, k in T.grid(M, N, K):
        with T.block("matmul"):
            T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
            T.writes(D[vi, vj])
            with T.init():
                D[vi, vj] = C[vi, vj]  # Fused epilogue into init
            D[vi, vj] = D[vi, vj] + A[vi, vk] * B[vj, vk]

The transformation removes the intermediate temp buffer and the separate
epilogue block, enabling better tensorization opportunities for hardware
with bias accumulation support.

Implementation:
- ReductionEpilogueFuser class for pattern validation and IR
transformation
  - BodyPatternAllowFusion: Validates epilogue can be fused
  - AnalyzeEpiloguePattern: Detects addition pattern (D = temp + C)
  - ExtractEpilogueInfo: Extracts buffer and region information
- CreateFusedReductionBlock: Creates single block with modified T.init()
- SingleBlockFusionReplacer: Replaces blocks and removes temp buffer
- Variable mapping between epilogue and reduction block iter vars
- Proper buffer and region updates with correct read/write ordering
- FFI bindings and Python API following TVM conventions

Changes:
- src/tir/schedule/primitive/compute_inline.cc: Core implementation
(~430 lines)
- src/tir/schedule/primitive.h: Function declaration
- include/tvm/tir/schedule/schedule.h: Virtual method in ScheduleNode
- src/tir/schedule/concrete_schedule.{h,cc}: ConcreteScheduleNode
implementation
- src/tir/schedule/traced_schedule.{h,cc}: TracedScheduleNode
implementation
- src/tir/schedule/schedule.cc: FFI binding registration
- python/tvm/tir/schedule/schedule.py: Python API with documentation
-
tests/python/tir-schedule/test_tir_schedule_fuse_reduction_epilogue.py:
Comprehensive tests including basic fusion, float32 variant, numerical
correctness verification, and trace roundtrip validation

Run tests with:
pytest
tests/python/tir-schedule/test_tir_schedule_fuse_reduction_epilogue.py
-v
  
 And, Could you please also take a look at #18240? Thx :)

---------

Co-authored-by: hyun gyu kim <[email protected]>

Author: kimm240

include/tvm/tir/schedule/schedule.h

@@ -608,6 +608,13 @@ class ScheduleNode : public runtime::Object {
    * \param block The block to be inlined to its producer
    */
   virtual void ReverseComputeInline(const BlockRV& block) = 0;
+  /*!
+   * \brief Fuse an epilogue block into a reduction block
+   * \param reduction_block The reduction block (e.g., matmul)
+   * \param epilogue_block The epilogue block to be fused (e.g., bias add)
+   */
+  virtual void FuseReductionEpilogue(const BlockRV& reduction_block,
+                                     const BlockRV& epilogue_block) = 0;
   /******** Schedule: Reduction ********/
   /*!
    * \brief Decompose a reduction block into two separate blocks.

python/tvm/tir/schedule/schedule.py

@@ -2345,6 +2345,33 @@ def after_inline(a: T.handle, c: T.handle) -> None:
         # pylint: disable-next=no-member
         _ffi_api.ScheduleReverseComputeInline(self, block)  # type: ignore
 
+    @type_checked
+    def fuse_reduction_epilogue(
+        self,
+        reduction_block: Union[BlockRV, str],
+        epilogue_block: Union[BlockRV, str],
+    ) -> None:
+        """Fuse an epilogue block into a reduction block.
+
+        It requires:
+        1) The reduction block is a complete reduction block
+        2) The epilogue block only reads from the reduction block's output
+        3) The epilogue performs a simple addition: output = reduction_result + bias
+
+        Parameters
+        ----------
+        reduction_block : Union[BlockRV, str]
+            The reduction block (e.g., matmul)
+        epilogue_block : Union[BlockRV, str]
+            The epilogue block to be fused (e.g., bias add)
+        """
+        reduction_block = self._normalize_block_arg(reduction_block)
+        epilogue_block = self._normalize_block_arg(epilogue_block)
+        # pylint: disable-next=no-member
+        _ffi_api.ScheduleFuseReductionEpilogue(
+            self, reduction_block, epilogue_block
+        )  # type: ignore
+
     ########## Schedule: Reduction ##########
 
     @type_checked

src/tir/schedule/concrete_schedule.cc

@@ -832,6 +832,15 @@ void ConcreteScheduleNode::ReverseComputeInline(const BlockRV& block_rv) {
   this->state_->DebugVerify();
 }
 
+void ConcreteScheduleNode::FuseReductionEpilogue(const BlockRV& reduction_block_rv,
+                                                 const BlockRV& epilogue_block_rv) {
+  TVM_TIR_SCHEDULE_BEGIN();
+  tir::FuseReductionEpilogue(state_, this->GetSRef(reduction_block_rv),
+                             this->GetSRef(epilogue_block_rv));
+  TVM_TIR_SCHEDULE_END("fuse-reduction-epilogue", this->error_render_level_);
+  this->state_->DebugVerify();
+}
+
 /******** Schedule: Block Annotation ********/
 
 void ConcreteScheduleNode::StorageAlign(const BlockRV& block_rv, int buffer_index, int axis,

src/tir/schedule/concrete_schedule.h

@@ -147,6 +147,8 @@ class ConcreteScheduleNode : public ScheduleNode {
                         int index = -1) override;
   void ComputeInline(const BlockRV& block) override;
   void ReverseComputeInline(const BlockRV& block) override;
+  void FuseReductionEpilogue(const BlockRV& reduction_block,
+                             const BlockRV& epilogue_block) override;
   /******** Schedule: Reduction ********/
   BlockRV RFactor(const LoopRV& loop_rv, int factor_axis) override;
   BlockRV DecomposeReduction(const BlockRV& block_rv, const LoopRV& loop_rv) override;

src/tir/schedule/primitive.h

@@ -509,6 +509,14 @@ TVM_DLL void ComputeInline(ScheduleState self, const StmtSRef& block_sref);
  * \param block_sref The sref to the block to be inlined to its producer
  */
 TVM_DLL void ReverseComputeInline(ScheduleState self, const StmtSRef& block_sref);
+/*!
+ * \brief Fuse an epilogue block into a reduction block
+ * \param self The state of the schedule
+ * \param reduction_block_sref The sref to the reduction block
+ * \param epilogue_block_sref The sref to the epilogue block to be fused
+ */
+TVM_DLL void FuseReductionEpilogue(ScheduleState self, const StmtSRef& reduction_block_sref,
+                                   const StmtSRef& epilogue_block_sref);
 /******** Schedule: Reduction ********/
 /*!
  * \brief Decompose a reduction block into two separate blocks.