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,15 +260,21 @@ 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 the outermost reduction dimension 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).
+    pipeline with one scf.for loop over the outermost reduction dimension's
+    chunk count, each iteration containing one ktdf.pipeline (Load / Compute /
+    Store stages).
+
+    Only the outermost reduction dimension (smallest loop-dim index) 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.  Inner reduction dims are left for other passes.
+
     First-vs-rest accumulation behaviour is selected at runtime via
-    `%condition = AND(iv_j == 0 for all active chunk loop IVs)`:
+    `%condition = (chunk_iv == 0)`:
 
       Load stage   : transfers the current input chunk slice to fifo_in.
                      When !condition, also transfers the partial accumulator
@@ -294,27 +301,23 @@ def ReductionDimChunkingPass : Pass<"reduction-dim-chunking", "mlir::ModuleOp">
     - The inner Compute stage body contains exactly one linalg.generic with
       at least one `reduction` iterator type.
 
-    Option `num-chunks` sets the static chunk count per reduction dim (one
-    value per 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 omitted (the default), the pass infers the
-    chunk count automatically from the input tensor size and the
+    Option `num-chunks` sets the number of sequential chunks for the outermost
+    reduction dimension.  When omitted or set to 0 (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.
+    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
     ignored when `num-chunks` is explicitly set.
   }];
   let constructor = "::mlir::ktdf::createReductionDimChunkingPass()";
   let options = [
-    ListOption<"numChunks", "num-chunks", "unsigned",
-               "Number of sequential chunks per reduction dimension (one value "
-               "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.">,
+    Option<"numChunks", "num-chunks", "unsigned",
+           "0",
+           "Number of sequential chunks for the outermost reduction dimension. "
+           "0 (the default) means auto-infer from chunk-size-threshold.">,
     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))) {