Fix chunking non outermost dims and insert new generic for cross-chunk accumulation to avoid fp numerical mismatch

include/dataflow-scheduler/Dialect/KTDF/Transforms/Passes.td

@@ -93,8 +93,9 @@ def ReductionLoopExposurePass : Pass<"reduction-loop-exposure", "mlir::ModuleOp"
   let description = [{
     Walks the module for any ktdf.stage that contains a linalg.generic with a
     reduction iterator (the Compute stage).  For each one found it rewrites the
-    parent ktdf.pipeline so that the outer reduction becomes one nested scf.for loop per reduction dimension, with the accumulator tensor
-    threaded as iter_args.
+    parent ktdf.pipeline so that the outer reduction becomes one nested scf.for
+    loop per reduction dimension, with the accumulator tensor threaded as
+    iter_args.
 
     The innermost reduction dimension — identified as the reduction loop dim
     that maps to the rightmost input tensor dimension — is excluded from loop
@@ -259,13 +260,23 @@ def MapReductionPartialsPass : Pass<"map-reduction-partials", "mlir::ModuleOp">
 }
 
 def ReductionDimChunkingPass : Pass<"reduction-dim-chunking", "mlir::ModuleOp"> {
-  let summary = "Split the reduction dimension of inner linalg.generic ops into sequential chunks";
+  let summary = "Split reduction dimensions of inner linalg.generic ops into sequential chunks";
   let description = [{
     This pass targets the inner Load→Compute→Store pipeline that contains a
     linalg.generic with a reduction iterator.  It replaces that single
     pipeline with nested scf.for loops — one per reduction dimension that
     has more than one chunk — each innermost iteration containing one
     ktdf.pipeline (Load / Compute / Store stages).
+
+    When the chunk count is inferred from the threshold (num-chunks not
+    specified), only the outermost reduction dimension is chunked.  If that dimension is also the inner dim (the
+    rightmost input tensor dimension, belonging to the hardware SIMD path),
+    the pass skips chunking entirely.
+
+    When num-chunks is specified explicitly, multi-dim chunking is applied:
+    one value per reduction dimension (in iterator-type order), one nested
+    scf.for per dim (dims with chunk count 1 produce no loop).
+
     First-vs-rest accumulation behaviour is selected at runtime via
     `%condition = AND(iv_j == 0 for all active chunk loop IVs)`:
 
@@ -299,9 +310,10 @@ def ReductionDimChunkingPass : Pass<"reduction-dim-chunking", "mlir::ModuleOp">
     the pass emits one nested scf.for per dim (dims with chunk count 1
     produce no loop).  When omitted (the default), the pass infers the
     chunk count automatically from the input tensor size and the
-    `chunk-size-threshold` option: it picks the smallest integer N ≥ 1 such
-    that (total_input_bytes / N) ≤ chunk-size-threshold AND N evenly divides
-    every reduction-dimension size.
+    `chunk-size-threshold` option (outermost dim only): it picks the
+    smallest integer N ≥ 1 such that (total_input_bytes / N) ≤
+    chunk-size-threshold AND N evenly divides the outermost reduction
+    dimension size.
 
     Option `chunk-size-threshold` (default 1 MiB = 1048576 bytes) is the
     per-chunk byte budget used when `num-chunks` is not provided.  It is
@@ -314,7 +326,8 @@ def ReductionDimChunkingPass : Pass<"reduction-dim-chunking", "mlir::ModuleOp">
                "per reduction dim, in iterator-type order).  Values may differ "
                "across dims; the pass emits one nested scf.for per dim (dims "
                "with chunk count 1 produce no loop).  When empty (the default) "
-               "the pass infers the chunk count from chunk-size-threshold.">,
+               "the pass infers the chunk count from chunk-size-threshold "
+               "(outermost dim only).">,
     Option<"chunkSizeThreshold", "chunk-size-threshold", "int64_t",
            "1048576",
            "Per-chunk size budget in bytes used for automatic numChunks "

lib/Dialect/KTDF/Transforms/MapReductionPartials.cpp

@@ -559,6 +559,107 @@ static LogicalResult rewriteGeneric(
   return success();
 }
 
+// ---------------------------------------------------------------------------
+// Lower the "combine-with-partial" scf.if generated by ReductionLoopExposure
+// to a result-less scf.if that writes directly into `alloc_val`.
+//
+// After rewriteGeneric has run, the combine scf.if has this shape:
+//
+//   %result = scf.if %is_first -> (tensor<DxElem>) {
+//     scf.yield %alloc_val        // then: first chunk — alloc already filled
+//   } else {
+//     %r = ktdf.read_from_fifo %slot -> tensor<DxElem>
+//     %g = linalg.generic(parallel) ins(%r) outs(%alloc_val) -> tensor<DxElem>
+//     scf.yield %g
+//   }
+//
+// Both branches now reference `alloc_val` (a memref), so we lower to:
+//
+//   scf.if %is_first {
+//     // nothing — alloc already filled by the outer-dim rewrite
+//   } else {
+//     %r = ktdf.read_from_fifo %slot -> memref<DxElem>
+//     linalg.generic(parallel) ins(%r) outs(%alloc_val)   // buffer form
+//   }
+//
+// The original scf.if is erased.  The caller must update the downstream
+// generic's ins[0] to `alloc_val` directly.
+// ---------------------------------------------------------------------------
+static LogicalResult lowerCombineIfToMemref(scf::IfOp if_op, Value alloc_val) {
+  // ── then-branch: must be "scf.yield %alloc_val" ─────────────────────────
+  auto then_yield =
+      cast<scf::YieldOp>(if_op.getThenRegion().front().getTerminator());
+  assert(then_yield.getNumOperands() == 1 &&
+         then_yield.getOperand(0) == alloc_val &&
+         "lowerCombineIfToMemref: unexpected then-branch shape");
+
+  // ── else-branch: read_from_fifo + parallel linalg.generic + yield ────────
+  Block& else_block = if_op.getElseRegion().front();
+  auto else_yield = cast<scf::YieldOp>(else_block.getTerminator());
+  assert(else_yield.getNumOperands() == 1 &&
+         "lowerCombineIfToMemref: expected one yield operand in else-branch");
+
+  // The yield operand is the result of the parallel combine generic.
+  auto combine_generic =
+      else_yield.getOperand(0).getDefiningOp<linalg::GenericOp>();
+  if (!combine_generic)
+    return if_op.emitError(
+        "lowerCombineIfToMemref: else-branch yield is not a linalg.generic "
+        "result");
+
+  // The generic's single input is the tensor read from the partial FIFO slot.
+  if (combine_generic.getInputs().size() != 1)
+    return if_op.emitError(
+        "lowerCombineIfToMemref: combine generic must have exactly one input");
+  Value tensor_read = combine_generic.getInputs()[0];
+  auto read_op = tensor_read.getDefiningOp<ktdf::ReadFromFifoOp>();
+  if (!read_op)
+    return if_op.emitError(
+        "lowerCombineIfToMemref: combine generic input is not a "
+        "ktdf.read_from_fifo");
+
+  // ── Emit the lowered result-less scf.if ──────────────────────────────────
+  OpBuilder builder(if_op);
+  Location loc = if_op.getLoc();
+  auto new_if =
+      scf::IfOp::create(builder, loc, /*resultTypes=*/TypeRange{},
+                        if_op.getCondition(), /*withElseRegion=*/true);
+
+  // then-branch: empty — alloc_val already holds the correct data.
+  // (The default result-less scf.yield terminator is already present.)
+
+  // else-branch: memref read + buffer-semantics parallel generic into
+  // alloc_val.
+  {
+    Block& new_else = new_if.getElseRegion().front();
+    OpBuilder else_b = OpBuilder::atBlockBegin(&new_else);
+
+    auto tensor_type = cast<RankedTensorType>(tensor_read.getType());
+    auto memref_type =
+        MemRefType::get(tensor_type.getShape(), tensor_type.getElementType());
+    Value new_read = ktdf::ReadFromFifoOp::create(else_b, loc, memref_type,
+                                                  read_op.getFifoSlot())
+                         .getResult();
+
+    // Buffer-semantics generic: ins = memref read, outs = alloc_val, no result.
+    auto buf_generic = linalg::GenericOp::create(
+        else_b, loc, /*resultTensorTypes=*/TypeRange{},
+        /*inputs=*/ValueRange{new_read}, /*outputs=*/ValueRange{alloc_val},
+        combine_generic.getIndexingMapsAttr(),
+        combine_generic.getIteratorTypesAttr(),
+        /*doc=*/StringAttr{}, /*library_call=*/StringAttr{});
+    IRMapping mapping;
+    combine_generic.getRegion().cloneInto(&buf_generic.getRegion(), mapping);
+    Block& placeholder = buf_generic.getRegion().front();
+    if (&placeholder != &buf_generic.getRegion().back()) placeholder.erase();
+    // Terminator (result-less scf.yield) is already present in new_else.
+  }
+
+  // NOTE: the caller must reassign all uses of if_op's result before
+  // erasing it — we do not erase here so the caller can rewire ins[0] first.
+  return success();
+}
+
 // ---------------------------------------------------------------------------
 // Set `new_type` on a PrivateOp result and its corresponding inner value
 // (the private_yield operand at the same index).
@@ -899,6 +1000,25 @@ struct MapReductionPartialsPass
         stale_tensor_read = generic_op.getInputs()[0].getDefiningOp();
         Value new_read = convertInputToMemref(builder, generic_op);
         generic_op.getInputsMutable().assign(new_read);
+      } else if (auto if_op =
+                     generic_op.getInputs()[0].getDefiningOp<scf::IfOp>()) {
+        // ins[0] is the "combine-with-partial" scf.if emitted by
+        // ReductionLoopExposure.  After rewriteGeneric ran on the outer-dim
+        // accumulator, the if's then-branch yields the outer-dim alloc
+        // (a memref) and the else-branch still has a tensor read + parallel
+        // generic.  Lower to a result-less scf.if writing into that alloc,
+        // then wire the alloc directly as ins[0] for the inner-dim generic.
+        Value alloc_val =
+            if_op.getThenRegion().front().getTerminator()->getOperand(0);
+        if (failed(lowerCombineIfToMemref(if_op, alloc_val))) {
+          signalPassFailure();
+          return;
+        }
+        // RAUW replaces all uses of the scf.if result (including ins[0]) with
+        // alloc_val, leaving the old if with no uses so it can be erased
+        // safely.
+        if_op.getResult(0).replaceAllUsesWith(alloc_val);
+        if_op.erase();
       }
       // Transform inner-dim generic into memref-typed generic.
       if (failed(rewriteInnerDimGeneric(generic_op, group_local_mem))) {

lib/Dialect/KTDF/Transforms/ReductionDimChunking.cpp

@@ -16,9 +16,17 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// ReductionDimChunking: split the reduction dimension of the inner
-// linalg.generic into sequential chunks so that each chunk fits in the
-// hardware FIFO path.
+// ReductionDimChunking: split reduction dimensions of the inner linalg.generic
+// into sequential chunks so that each chunk fits in the hardware FIFO path.
+//
+// When the chunk count is inferred from the threshold (--num-chunks not given),
+// only the outermost reduction dimension is chunked.
+// If that dimension is also the inner dim (the rightmost input tensor
+// dimension, belonging to the hardware SIMD path), chunking is skipped.
+//
+// When --num-chunks is given explicitly, multi-dim chunking is applied using
+// one value per reduction dimension (in iterator-type order); one nested
+// scf.for is emitted per dim; dims with chunk count 1 produce no loop.
 //
 // The pass operates on the shape produced by StageCoarsening.  It expects a
 // top-level ktdf.pipeline with three sibling stages:
@@ -37,18 +45,17 @@
 // the entire dimension).  Each innermost iteration contains one ktdf.pipeline
 // with three stages built by ktdf::StageFactory (see ReductionUtils.h).
 // First-vs-rest accumulation behaviour is selected at runtime via
-// %condition = (all active loop IVs == 0):
+// %condition = AND(iv_j == 0 for all active chunk loop IVs):
 //
-//   Load stage   : transfers each input chunk slice (memref → fifo_in[i]).
+//   Load stage   : transfers the current input chunk slice (memref → fifo_in).
 //                  When !condition, also transfers the partial accumulator
-//                  (local memory output buffer → fifo_partial[i]) so the
+//                  (local memory output buffer → fifo_partial) so the
 //                  Compute stage can read it back.
 //   Compute stage: when condition, initialises the output tensor with
 //                  tensor.empty; otherwise reads the partial result from
 //                  fifo_partial.  The linalg.generic and write_to_fifo are
 //                  unconditional.
-//   Store stage  : unconditionally writes each fifo_out[i] back to the local
-//   memory
+//   Store stage  : unconditionally writes fifo_out back to the local memory
 //                  output buffer.
 //
 // The existing local memory output buffer (discovered via the original Store
@@ -194,16 +201,47 @@ struct ReductionDimChunkingPass
     // ------------------------------------------------------------------
     // Determine reduction dims, num_chunks, and per-dim chunk sizes.
     //
-    // When --num-chunks was provided by the user, validate and use it.
-    // Otherwise delegate to ReductionChunkAnalysis which picks the
-    // smallest N that fits within chunkSizeThreshold bytes per chunk.
+    // Threshold path (numChunks empty): only the outermost reduction dim is
+    // chunked.  If it is also the inner (rightmost input tensor) dim it belongs
+    // to the hardware SIMD path — skip chunking entirely.
+    //
+    // Explicit path (numChunks non-empty): multi-dim chunking — one value per
+    // reduction dim, nested scf.for loops, dims with count 1 produce no loop.
     // ------------------------------------------------------------------
     SmallVector<int64_t> reduction_dims;
     SmallVector<int64_t> chunk_sizes;
     unsigned loop_num_chunks = 0;
 
     if (numChunks.empty()) {
-      // Auto-infer via analysis.
+      // Auto-infer via analysis — outermost reduction dim only.
+
+      // Locate the outermost reduction dim (first iterator typed `reduction`).
+      int64_t outermost_red_dim = -1;
+      {
+        auto iter_types = generic_op.getIteratorTypesArray();
+        for (int64_t i = 0; i < static_cast<int64_t>(iter_types.size()); ++i) {
+          if (iter_types[i] == utils::IteratorType::reduction) {
+            outermost_red_dim = i;
+            break;
+          }
+        }
+      }
+      if (outermost_red_dim < 0) {
+        LDBG(1) << PASS_NAME
+            ": could not find a reduction dimension — skipping";
+        return success();
+      }
+
+      // If the outermost reduction dim is also the inner (rightmost input
+      // tensor) dim it belongs to the hardware SIMD path — skip entirely.
+      std::optional<unsigned> inner_dim = findInnerDimLoopDim(generic_op);
+      if (inner_dim && static_cast<int64_t>(*inner_dim) == outermost_red_dim) {
+        LDBG(1) << PASS_NAME << ": outermost reduction dim "
+                << outermost_red_dim
+                << " is the inner (rightmost input) dim — skipping chunking";
+        return success();
+      }
+
       auto result = analyzeReductionChunks(generic_op, chunkSizeThreshold);
       if (!result) {
         inner_pipeline.emitError(PASS_NAME
@@ -212,8 +250,9 @@ struct ReductionDimChunkingPass
         return failure();
       }
       loop_num_chunks = result->num_chunks;
-      chunk_sizes = std::move(result->chunk_sizes);
-      reduction_dims = std::move(result->reduction_dims);
+      // Restrict to outermost dim only: take just the first entry.
+      reduction_dims = {outermost_red_dim};
+      chunk_sizes = {result->chunk_sizes[0]};
     } else {
       // User-supplied --num-chunks path: collect reduction dims and validate.
       auto iter_types = generic_op.getIteratorTypesArray();
@@ -261,18 +300,13 @@ struct ReductionDimChunkingPass
       }
     }
 
-    if (reduction_dims.empty()) {
-      LDBG(1) << PASS_NAME ": could not find reduction dimension — skipping";
-      return success();
-    }
-
     LDBG(1) << PASS_NAME ": num_reduction_dims=" << reduction_dims.size()
             << " total_chunks=" << loop_num_chunks;
 
     // Collect per-dim chunk counts in reduction-dim order.
     SmallVector<int64_t> per_dim_num_chunks;
     if (numChunks.empty()) {
-      // Auto-inferred path: all dims use the same num_chunks.
+      // Auto-inferred path: only the outermost dim was kept, assign its count.
       per_dim_num_chunks.assign(reduction_dims.size(),
                                 static_cast<int64_t>(loop_num_chunks));
     } else {