[TIR][Schedule] FuseReductionEpilogue: Add Clipping pattern support

Currently, the FuseReductionEpilogue primitive only supports Bias
(addition) and BiasReLU (addition + ReLU) epilogue patterns. However,
clipping operations (min(max(x, lower), upper)) are commonly used in
deep learning models and would benefit from the same fusion optimization.

This commit extends FuseReductionEpilogue to support Clipping patterns
by:

1. Adding EpilogueType::Clipping to the enum to distinguish clipping
   patterns from other epilogue types.

2. Adding clipping_lower_ and clipping_upper_ members to
   ReductionEpilogueFuser to store clipping bounds extracted from the
   epilogue pattern.

3. Extending AnalyzeEpiloguePattern to detect clipping patterns:
   - min(max(temp, lower), upper)
   - max(min(temp, upper), lower)
   - All commutative variants of min/max at each level

4. Updating BiasReLU pattern matching to handle max(0, x) form in
   addition to max(x, 0) for better commutativity support.

5. Modifying CreateFusedReductionBlock to apply clipping to the init
   value: init = min(max(0, lower), upper)

6. Updating BufferReplacer to apply clipping per-iteration:
   value = min(max(value, lower), upper)

7. Adding validation in BodyPatternAllowFusion to ensure temp appears
   exactly once in clipping patterns.

8. Creating comprehensive test coverage with 8 test cases:
   - Basic fusion test
   - Numerical correctness verification
   - Multiple epilogue blocks test
   - 5 commutative variant tests

This implementation follows the same per-iteration semantics as BiasReLU,
where clipping is applied at each reduction step rather than
post-reduction. This semantic change is documented in the docstring with
a warning about potential numerical differences.

The test suite verifies that all commutative forms of clipping patterns
are correctly recognized and that the fused implementation produces
numerically identical results to the per-iteration reference
implementation.

Author: hyun gyu kim

3rdparty/tvm-ffi

@@ -0,0 +1 @@
+Subproject commit 3ea601bb413074c49a77c4ce3218bc08f8c4703c

ffi/3rdparty/dlpack

@@ -0,0 +1 @@
+Subproject commit 793921876c981ce49759114d7bb89bb89b2d3a2d

ffi/3rdparty/libbacktrace

@@ -2356,14 +2356,41 @@ def fuse_reduction_epilogue(
         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
+        3) The epilogue matches one of the supported patterns:
+           - Bias: ``output = reduction_result + bias``
+           - BiasReLU: ``output = max(reduction_result + bias, 0)``
+           - Clipping: ``output = min(max(reduction_result, lower), upper)``
+           or their commutative variants
+
+        .. warning::
+
+            **Semantic Change for Non-Linear Epilogues (BiasReLU, Clipping):**
+
+            For non-linear epilogues (BiasReLU and Clipping), fusion changes the
+            computation semantics from post-reduction application to per-iteration
+            application. This can lead to different numerical results.
+
+            **Example with Clipping to [-5, 5] and inputs [6, -2]:**
+
+            - **Post-reduction clipping** (original): ``clip(sum([6, -2])) = clip(4) = 4``
+            - **Per-iteration clipping** (fused): ``acc=0 → clip(0+6)=5 → clip(5+(-2))=3``
+
+            The fused version applies clipping at each reduction iteration, which
+            may be an intended optimization for some models but can cause unexpected
+            correctness issues if users are not aware of this behavior.
+
+            For linear epilogues (Bias), fusion preserves exact numerical equivalence.
 
         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)
+            The epilogue block to be fused (e.g., bias add, ReLU, clipping)
+
+        Examples
+        --------
+        See :py:func:`test_tir_schedule_fuse_reduction_epilogue` for examples.
         """
         reduction_block = self._normalize_block_arg(reduction_block)
         epilogue_block = self._normalize_block_arg(epilogue_block)

python/tvm/tir/schedule/schedule.py

@@ -992,6 +992,7 @@ void ReverseComputeInline(ScheduleState self, const StmtSRef& consumer_block_sre
 enum class EpilogueType {
   Bias,      // temp + C
   BiasReLU,  // max(temp + C, 0)
+  Clipping,  // min(max(temp, lower), upper)
 };
 
 class ReductionEpilogueFuser : public BaseInliner {
@@ -1058,6 +1059,8 @@ class ReductionEpilogueFuser : public BaseInliner {
   Buffer epilogue_addend_buffer_{nullptr};                 // Addend buffer C
   BufferRegion epilogue_addend_region_{nullptr};           // Read region of C
   EpilogueType epilogue_type_;                             // Type of epilogue operation
+  PrimExpr clipping_lower_{nullptr};                       // Lower bound for clipping
+  PrimExpr clipping_upper_{nullptr};                       // Upper bound for clipping
 };
 
 bool ReductionEpilogueFuser::BodyPatternAllowFusion(const BlockRealize& epilogue_block_realize) {
@@ -1080,19 +1083,28 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const BlockRealize& epilogue
     return false;
   }
 
-  // 4. Analyze epilogue pattern: D[i,j] = temp[i,j] + C[i,j]
+  // 4. Analyze epilogue pattern: D[i,j] = temp[i,j] + C[i,j] or
+  //    D[i,j] = min(max(temp[i,j], lower), upper)
   if (!AnalyzeEpiloguePattern(inlined_store_->value)) {
-    // Failure: epilogue is not a simple addition pattern
+    // Failure: epilogue is not a supported pattern (Bias, BiasReLU, or Clipping)
     return false;
   }
 
-  // 5. Check if producer is a reduction block
+  // 5. For Clipping pattern, verify temp appears exactly once
+  if (epilogue_type_ == EpilogueType::Clipping) {
+    if (loads.size() != 1) {
+      // Failure: temp must appear exactly once in clipping pattern
+      return false;
+    }
+  }
+
+  // 6. Check if producer is a reduction block
   if (!IsReductionBlock(reduction_block_)) {
     // Failure: producer is not a reduction block
     return false;
   }
 
-  // 6. Extract epilogue information (output buffer, indices, regions, etc.)
+  // 7. Extract epilogue information (output buffer, indices, regions, etc.)
   ExtractEpilogueInfo();
 
   return true;
@@ -1115,19 +1127,97 @@ bool ReductionEpilogueFuser::AnalyzeEpiloguePattern(const PrimExpr& value) {
     }
   }
 
-  // Pattern 2: max(temp[i,j] + C[i,j], 0) or max(C[i,j] + temp[i,j], 0) (BiasReLU)
+  // Pattern 2: min(max(temp[i,j], lower), upper) or max(min(temp[i,j], upper), lower) (Clipping)
+  // Handle all commutative variants of min/max at each level.
+
+  // Helper to check if an expression is a load from the reduction buffer, and
+  // return the other operand as `other` if so.
+  auto match_buffer_in_commutative_op = [this](const PrimExpr& a, const PrimExpr& b,
+                                               PrimExpr* other) -> bool {
+    if (const auto* load_a = a.as<BufferLoadNode>()) {
+      if (load_a->buffer.same_as(inlined_buffer_)) {
+        *other = b;
+        return true;
+      }
+    }
+    if (const auto* load_b = b.as<BufferLoadNode>()) {
+      if (load_b->buffer.same_as(inlined_buffer_)) {
+        *other = a;
+        return true;
+      }
+    }
+    return false;
+  };
+
+  // Check for min(max(temp, lower), upper) and commutative variants
+  if (const auto* min_node = value.as<MinNode>()) {
+    const MaxNode* max_node = nullptr;
+    PrimExpr upper;
+    // Try both (a, b) as possible positions of the inner max
+    if ((max_node = min_node->a.as<MaxNode>())) {
+      upper = min_node->b;
+    } else if ((max_node = min_node->b.as<MaxNode>())) {
+      upper = min_node->a;
+    }
+    if (max_node != nullptr) {
+      PrimExpr lower;
+      if (match_buffer_in_commutative_op(max_node->a, max_node->b, &lower)) {
+        clipping_lower_ = lower;
+        clipping_upper_ = upper;
+        epilogue_type_ = EpilogueType::Clipping;
+        return true;
+      }
+    }
+  }
+
+  // Check for max(min(temp[i,j], upper), lower) and commutative variants
   if (const auto* max_node = value.as<MaxNode>()) {
-    // Check if second operand is zero (ReLU: max(x, 0))
-    // Support both integer and float zero constants
+    const MinNode* min_node = nullptr;
+    PrimExpr lower;
+    // Try both (a, b) as possible positions of the inner min
+    if ((min_node = max_node->a.as<MinNode>())) {
+      lower = max_node->b;
+    } else if ((min_node = max_node->b.as<MinNode>())) {
+      lower = max_node->a;
+    }
+    if (min_node != nullptr) {
+      PrimExpr upper;
+      if (match_buffer_in_commutative_op(min_node->a, min_node->b, &upper)) {
+        clipping_lower_ = lower;
+        clipping_upper_ = upper;
+        epilogue_type_ = EpilogueType::Clipping;
+        return true;
+      }
+    }
+  }
+
+  // Pattern 3: max(temp[i,j] + C[i,j], 0) or max(C[i,j] + temp[i,j], 0) (BiasReLU)
+  // Also handle max(0, temp[i,j] + C[i,j]) or max(0, C[i,j] + temp[i,j])
+  if (const auto* max_node = value.as<MaxNode>()) {
+    // Check if either operand is zero (ReLU: max(x, 0) or max(0, x))
+    // Support both integer and float zero constants.
+    const PrimExpr* add_candidate = nullptr;
     bool is_zero_const = false;
-    if (tir::is_zero(max_node->b)) {
+    auto is_zero_expr = [](const PrimExpr& expr) -> bool {
+      if (tir::is_zero(expr)) {
+        return true;
+      }
+      if (const auto* float_imm = expr.as<FloatImmNode>()) {
+        return float_imm->value == 0.0;
+      }
+      return false;
+    };
+
+    if (is_zero_expr(max_node->a)) {
+      is_zero_const = true;
+      add_candidate = &max_node->b;
+    } else if (is_zero_expr(max_node->b)) {
       is_zero_const = true;
-    } else if (const auto* float_imm = max_node->b.as<FloatImmNode>()) {
-      is_zero_const = (float_imm->value == 0.0);
+      add_candidate = &max_node->a;
     }
-    if (is_zero_const) {
-      // Check if first operand is AddNode
-      if (const auto* add = max_node->a.as<AddNode>()) {
+
+    if (is_zero_const && add_candidate != nullptr) {
+      if (const auto* add = add_candidate->as<AddNode>()) {
         const auto* load_a = add->a.as<BufferLoadNode>();
         const auto* load_b = add->b.as<BufferLoadNode>();
 
@@ -1218,6 +1308,14 @@ Block ReductionEpilogueFuser::CreateFusedReductionBlock(const BlockNode* reducti
     PrimExpr zero = tir::make_zero(init_value.dtype());
     new_init_store = BufferStore(epilogue_output_buffer_, Max(init_value, zero),
                                  Substitute(epilogue_output_indices_, var_map));
+  } else if (epilogue_type_ == EpilogueType::Clipping) {
+    // For Clipping, init should be min(max(init_value, lower), upper)
+    // Since init is typically 0, this becomes min(max(0, lower), upper)
+    PrimExpr init_value = tir::make_zero(epilogue_output_buffer_->dtype);
+    PrimExpr clipped_init = Min(Max(init_value, Substitute(clipping_lower_, var_map)),
+                                Substitute(clipping_upper_, var_map));
+    new_init_store = BufferStore(epilogue_output_buffer_, clipped_init,
+                                 Substitute(epilogue_output_indices_, var_map));
   } else {
     // Bias: D[vi, vj] = C[vi, vj]
     new_init_store = BufferStore(epilogue_output_buffer_, Substitute(epilogue_addend_, var_map),
@@ -1228,11 +1326,14 @@ Block ReductionEpilogueFuser::CreateFusedReductionBlock(const BlockNode* reducti
   // 3. Replace output buffer from temp to D in body
   class BufferReplacer : public StmtExprMutator {
    public:
-    BufferReplacer(Buffer old_buf, Buffer new_buf, EpilogueType epilogue_type, DataType dtype)
+    BufferReplacer(Buffer old_buf, Buffer new_buf, EpilogueType epilogue_type, DataType dtype,
+                   PrimExpr clipping_lower = PrimExpr(), PrimExpr clipping_upper = PrimExpr())
         : old_buffer_(old_buf),
           new_buffer_(new_buf),
           epilogue_type_(epilogue_type),
-          dtype_(dtype) {}
+          dtype_(dtype),
+          clipping_lower_(clipping_lower),
+          clipping_upper_(clipping_upper) {}
 
     Stmt VisitStmt_(const BufferStoreNode* op) final {
       BufferStore store = Downcast<BufferStore>(StmtExprMutator::VisitStmt_(op));
@@ -1242,6 +1343,9 @@ Block ReductionEpilogueFuser::CreateFusedReductionBlock(const BlockNode* reducti
         if (epilogue_type_ == EpilogueType::BiasReLU) {
           PrimExpr zero = tir::make_zero(dtype_);
           new_value = Max(new_value, zero);
+        } else if (epilogue_type_ == EpilogueType::Clipping) {
+          // For Clipping, apply min(max(value, lower), upper) per iteration
+          new_value = Min(Max(new_value, clipping_lower_), clipping_upper_);
         }
         return BufferStore(new_buffer_, new_value, store->indices);
       }
@@ -1261,10 +1365,17 @@ Block ReductionEpilogueFuser::CreateFusedReductionBlock(const BlockNode* reducti
     Buffer new_buffer_;
     EpilogueType epilogue_type_;
     DataType dtype_;
+    PrimExpr clipping_lower_;
+    PrimExpr clipping_upper_;
   };
 
   DataType dtype = epilogue_output_buffer_->dtype;
-  BufferReplacer replacer(inlined_buffer_, epilogue_output_buffer_, epilogue_type_, dtype);
+  PrimExpr clipping_lower_subst =
+      epilogue_type_ == EpilogueType::Clipping ? Substitute(clipping_lower_, var_map) : PrimExpr();
+  PrimExpr clipping_upper_subst =
+      epilogue_type_ == EpilogueType::Clipping ? Substitute(clipping_upper_, var_map) : PrimExpr();
+  BufferReplacer replacer(inlined_buffer_, epilogue_output_buffer_, epilogue_type_, dtype,
+                          clipping_lower_subst, clipping_upper_subst);
   new_block->body = replacer(reduction_block->body);
 
   // 4. Update write regions