diff --git a/include/dataflow-scheduler/Dialect/KTDF/Transforms/Passes.td b/include/dataflow-scheduler/Dialect/KTDF/Transforms/Passes.td index 303efb56..40a837e1 100644 --- a/include/dataflow-scheduler/Dialect/KTDF/Transforms/Passes.td +++ b/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 " diff --git a/lib/Dialect/KTDF/Transforms/MapReductionPartials.cpp b/lib/Dialect/KTDF/Transforms/MapReductionPartials.cpp index 94592610..7e2bd6a7 100644 --- a/lib/Dialect/KTDF/Transforms/MapReductionPartials.cpp +++ b/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) { +// scf.yield %alloc_val // then: first chunk — alloc already filled +// } else { +// %r = ktdf.read_from_fifo %slot -> tensor +// %g = linalg.generic(parallel) ins(%r) outs(%alloc_val) -> tensor +// 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 +// 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(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(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(); + 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(); + 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(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()) { + // 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))) { diff --git a/lib/Dialect/KTDF/Transforms/ReductionDimChunking.cpp b/lib/Dialect/KTDF/Transforms/ReductionDimChunking.cpp index 9fe59b40..f1a89ae7 100644 --- a/lib/Dialect/KTDF/Transforms/ReductionDimChunking.cpp +++ b/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 reduction_dims; SmallVector 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(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 inner_dim = findInnerDimLoopDim(generic_op); + if (inner_dim && static_cast(*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 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(loop_num_chunks)); } else { @@ -298,7 +332,7 @@ struct ReductionDimChunkingPass // ----------------------------------------------------------------------- // Replace inner_pipeline with nested scf.for loops — one per reduction - // dimension whose num_chunks > 1. Dimensions with num_chunks == 1 need no + // dimension whose num_chunks > 1. Dimensions with num_chunks == 1 need no // loop; their IV is treated as the constant 0 for offset and condition // computation. // diff --git a/lib/Dialect/KTDF/Transforms/ReductionLoopExposure.cpp b/lib/Dialect/KTDF/Transforms/ReductionLoopExposure.cpp index 4724ab36..682e7f29 100644 --- a/lib/Dialect/KTDF/Transforms/ReductionLoopExposure.cpp +++ b/lib/Dialect/KTDF/Transforms/ReductionLoopExposure.cpp @@ -878,16 +878,12 @@ struct ReductionLoopExposurePass // Rewrite the compute stage with N nested scf.for loops (one per // reduction dim), each carrying the accumulator tensor as iter_arg. // - // The accumulator seed is determined before the outermost loop: - // - On the first chunk (is_first_chunk == true): tensor.empty (zero init). - // - On subsequent chunks: read the previous partial result from - // fifo_in_partial. + // The accumulator is always seeded with tensor.empty before the outermost + // loop. When a partial FIFO path exists (cross-chunk accumulation), a + // combine scf.if is emitted *after* the loop to merge the previous chunk's + // partial result into the loop output: // - // %seed = scf.if %is_first -> tensor<...> { - // %e = tensor.empty(); scf.yield %e - // } else { - // %p = ktdf.read_from_fifo fifo_in_partial; scf.yield %p - // } + // %seed = tensor.empty() // scf.for %r0 = 0 to D0 iter_args(%a0 = %seed) {loop_type = reduction} // ... // %slice = ktdf.read_from_fifo fifo_in @@ -896,6 +892,14 @@ struct ReductionLoopExposurePass // scf.yield %updated // ... // scf.yield %r0_result + // // (when fifo_in_partial): + // %combined = scf.if %is_first -> tensor<...> { + // scf.yield %r0_result // first chunk: loop result is final + // } else { + // %p = ktdf.read_from_fifo fifo_in_partial + // %g = linalg.generic(parallel, addf) ins(%p) outs(%r0_result) + // scf.yield %g // subsequent chunks: add prior partial + // } // ------------------------------------------------------------------------- LogicalResult rewriteComputeStage(IRRewriter& rewriter, Location loc, MLIRContext* ctx, ktdf::StageOp stage, @@ -917,49 +921,15 @@ struct ReductionLoopExposurePass // already in front of the loops, and what reads its result is behind them. rewriter.setInsertionPoint(generic_op); - // Accumulator seed: when a partial FIFO path exists, on the first chunk - // zero-init via tensor.empty; on subsequent chunks read the previous - // partial result from fifo_in_partial. When there is no partial path - // (pipeline has no accumulator feedback), always use tensor.empty. + // Accumulator seed: always initialize with tensor.empty. const unsigned results = static_cast(generic_op.getNumResults()); SmallVector seeds; - if (fifo_in_partial && is_first_chunk) { - if (results != 1) { - return generic_op.emitError( - PASS_NAME - ": a compute with more than one accumulator has no partial fifo " - "per accumulator to read the previous chunk from"); - } - // Build the seed scf.if with an else region. The regions start empty, - // so we use OpBuilder::atBlockBegin (not getTerminator()) to populate - // them before inserting the scf.yield terminator. - auto seed_if = - scf::IfOp::create(rewriter, loc, TypeRange{output_tensor_type}, - is_first_chunk, /*withElseRegion=*/true); - { - Block& then_block = seed_if.getThenRegion().front(); - OpBuilder then_b = OpBuilder::atBlockBegin(&then_block); - auto empty = - tensor::EmptyOp::create(then_b, loc, output_tensor_type.getShape(), - output_tensor_type.getElementType()); - scf::YieldOp::create(then_b, loc, ValueRange{empty.getResult()}); - } - { - Block& else_block = seed_if.getElseRegion().front(); - OpBuilder else_b = OpBuilder::atBlockBegin(&else_block); - auto partial_read = ktdf::ReadFromFifoOp::create( - else_b, loc, output_tensor_type, fifo_in_partial); - scf::YieldOp::create(else_b, loc, ValueRange{partial_read.getResult()}); - } - seeds.push_back(seed_if.getResult(0)); - } else { - for (unsigned r = 0; r < results; ++r) { - seeds.push_back(tensor::EmptyOp::create( - rewriter, loc, output_tensor_type.getShape(), - output_tensor_type.getElementType()) - .getResult()); - } + for (unsigned r = 0; r < results; ++r) { + seeds.push_back( + tensor::EmptyOp::create(rewriter, loc, output_tensor_type.getShape(), + output_tensor_type.getElementType()) + .getResult()); } // Build the nested loops, threading each accumulator through every level. @@ -989,25 +959,89 @@ struct ReductionLoopExposurePass body_builder.clone(*generic_op.getOperation(), mapping)); ValueRange updated = new_generic.getResults(); + // Replace the placeholder yield in the innermost loop with the real one. + for (unsigned r = 0; r < results; ++r) { + inner_yield->setOperand(r, updated[r]); + } + + // Accumulation combining after loop: if there's a partial FIFO input, + // emit an scf.if after the loop where the else branch adds the partial + // FIFO value to the loop reduction result via a parallel linalg.generic. + rewriter.setInsertionPointAfter(nested.outermost_loop); + + SmallVector final_results; + if (fifo_in_partial && is_first_chunk) { + if (results != 1) { + return generic_op.emitError( + PASS_NAME + ": a compute with more than one accumulator has no partial fifo " + "per accumulator to read the previous chunk from"); + } + auto if_op = + scf::IfOp::create(rewriter, loc, TypeRange{output_tensor_type}, + is_first_chunk, /*withElseRegion=*/true); + { + Block& then_block = if_op.getThenRegion().front(); + OpBuilder then_b = OpBuilder::atBlockBegin(&then_block); + scf::YieldOp::create(then_b, loc, + ValueRange{nested.outermost_loop.getResult(0)}); + } + { + Block& else_block = if_op.getElseRegion().front(); + OpBuilder else_b = OpBuilder::atBlockBegin(&else_block); + auto partial_read = ktdf::ReadFromFifoOp::create( + else_b, loc, output_tensor_type, fifo_in_partial); + + int64_t rank = output_tensor_type.getRank(); + AffineMap id_map = else_b.getMultiDimIdentityMap(rank); + SmallVector indexing_maps = {id_map, id_map}; + SmallVector iter_types( + rank, utils::IteratorType::parallel); + + auto add_generic = linalg::GenericOp::create( + else_b, loc, TypeRange{output_tensor_type}, + /*inputs=*/ValueRange{partial_read.getResult()}, + /*outputs=*/ValueRange{nested.outermost_loop.getResult(0)}, + indexing_maps, iter_types, + [&](OpBuilder& b, Location b_loc, ValueRange args) { + Value in_val = args[0]; + Value out_val = args[1]; + Value sum_val; + Type elem_type = output_tensor_type.getElementType(); + if (isa(elem_type)) { + sum_val = arith::AddFOp::create(b, b_loc, in_val, out_val); + } else { + sum_val = arith::AddIOp::create(b, b_loc, in_val, out_val); + } + linalg::YieldOp::create(b, b_loc, sum_val); + }); + + scf::YieldOp::create(else_b, loc, ValueRange{add_generic.getResult(0)}); + } + final_results.push_back(if_op.getResult(0)); + } else { + for (unsigned r = 0; r < results; ++r) { + final_results.push_back(nested.outermost_loop.getResult(r)); + } + } + if (fifo_outs.front()) { - // Each result feeds a write_to_fifo directly. Emit the guarded writes on - // the last iteration and drop the original write ops. + // Emit the guarded writes on the last iteration inside the innermost + // loop. Value is_last = buildAllLast(body_builder, loc, nested.ivs, last_vals); - auto if_op = scf::IfOp::create(body_builder, loc, TypeRange{}, is_last, - /*withElseRegion=*/false); - OpBuilder then_builder(if_op.getThenRegion().front().getTerminator()); + auto last_if = scf::IfOp::create(body_builder, loc, TypeRange{}, is_last, + /*withElseRegion=*/false); + OpBuilder then_builder(last_if.getThenRegion().front().getTerminator()); for (unsigned r = 0; r < results; ++r) { ktdf::WriteToFifoOp::create(then_builder, loc, updated[r], fifo_outs[r]); } } else { // The generic's result feeds other ops (e.g. a downstream inner-dim - // reduction). Replace all uses of the original generic with the - // outermost loop result so those ops pick up the fully-accumulated - // tensor after the loop completes. + // reduction). Replace all uses of the original generic with the + // final results. for (unsigned r = 0; r < results; ++r) { - generic_op.getResult(r).replaceAllUsesWith( - nested.outermost_loop.getResult(r)); + generic_op.getResult(r).replaceAllUsesWith(final_results[r]); } // write_to_fifo has no results so use_empty() is always true; exclude // it from the erase list so it is preserved together with the ops that @@ -1016,11 +1050,6 @@ struct ReductionLoopExposurePass to_erase, [](Operation* op) { return isa(op); }); } - // Replace the placeholder yield in the innermost loop with the real one. - for (unsigned r = 0; r < results; ++r) { - inner_yield->setOperand(r, updated[r]); - } - // Erase original body ops (reverse order, only if unused). for (auto* op : llvm::reverse(to_erase)) if (op->use_empty()) rewriter.eraseOp(op); diff --git a/test/Transforms/MapReductionPartials/outer_inner_dim_with_partial.mlir b/test/Transforms/MapReductionPartials/outer_inner_dim_with_partial.mlir new file mode 100644 index 00000000..79dd1b1c --- /dev/null +++ b/test/Transforms/MapReductionPartials/outer_inner_dim_with_partial.mlir @@ -0,0 +1,217 @@ +// RUN: dataflow-scheduler-opt --map-reduction-partials %s | FileCheck %s + +// Tests the case where ReductionLoopExposure has produced: +// G1 (outer-dim, loop-exposed): absmax reduction tensor<1x1x64xf16> -> tensor<1x64xf16> +// carried as iter_arg through a scf.for reduction loop. +// combine scf.if: on the first chunk yields the loop result directly; on +// subsequent chunks reads the previous partial from a FIFO and adds it. +// G2 (inner-dim): absmax reduction tensor<1x64xf16> -> tensor<1x64xf16> +// whose input is the combine scf.if result. +// +// MapReductionPartials must: +// - Lower G1: alloc + linalg.fill(0.0) + buffer scf.for (no iter_arg). +// - Lower the combine scf.if to a result-less scf.if: +// then: empty (alloc already holds the loop result) +// else: memref read_from_fifo + buffer linalg.generic addf into alloc +// - Lower G2 using the G1 alloc as ins, a rank-reducing subview as outs, +// and write the full alloc to the widened FIFO. + +// CHECK: #[[$MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)> +// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)> +// CHECK: #[[$MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> + +// CHECK-LABEL: module @local_schedule_0 + +// Outer pipeline private: unchanged +// CHECK: %[[PRIV0:.*]]:4 = ktdf.private -> (memref<2x1x256x64xf16 +// CHECK: memref.alloc() : memref<2x1x256x64xf16 +// CHECK: memref.alloc() : memref<2x1x64xf16 + +// Inner pipeline private: slots unchanged (G2 widens slot #2 below) +// CHECK: %[[PRIV1:.*]]:5 = ktdf.private +// CHECK-SAME: !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> +// CHECK-SAME: !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> +// CHECK-SAME: !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16> + +// L1LU load stage: unchanged +// CHECK: ktdf.stage depends_in(none) +// CHECK: scf.if + +// SFU compute stage — key transformations: +// CHECK: ktdf.stage depends_in(%[[PRIV1]]#3) +// +// G1: alloc + fill(0.0) hoisted to top of stage +// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc() : memref<1x64xf16, "SFU_REG"> +// CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f16 +// CHECK-NEXT: linalg.fill ins(%[[ZERO]] : f16) outs(%[[ALLOC]] : memref<1x64xf16, "SFU_REG">) +// +// G1: buffer scf.for reduction loop (no iter_arg) +// CHECK: scf.for %[[IV:.*]] = {{.*}} { +// CHECK-NEXT: %[[RD0:.*]] = ktdf.read_from_fifo %[[PRIV1]]#0 : <"L1LU" -> "SFU", 64xf16> -> memref<1x1x64xf16> +// CHECK-NEXT: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "reduction", "parallel"]} ins(%[[RD0]] : memref<1x1x64xf16>) outs(%[[ALLOC]] : memref<1x64xf16, "SFU_REG">) +// CHECK: } {loop_type = #ktdf.loop_type} +// +// combine scf.if: result-less; then empty, else read_from_fifo + addf into alloc +// CHECK: scf.if %[[IS_FIRST:.*]] { +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[RD1:.*]] = ktdf.read_from_fifo %[[PRIV1]]#1 : <"L1LU" -> "SFU", 64xf16> -> memref<1x64xf16> +// CHECK-NEXT: linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%[[RD1]] : memref<1x64xf16>) outs(%[[ALLOC]] : memref<1x64xf16, "SFU_REG">) +// CHECK: } +// +// G2: rank-reducing subview of alloc (reduction dim d1 collapsed to size 1) +// CHECK: %[[SV:.*]] = memref.subview %[[ALLOC]][0, 0] [1, 1] [1, 1] +// CHECK-SAME: memref<1x64xf16, "SFU_REG"> to memref<1x1xf16, strided<[64, 1]>, "SFU_REG"> +// CHECK-NEXT: linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP1]]], iterator_types = ["parallel", "reduction", "parallel"]} ins(%[[ALLOC]] : memref<1x64xf16, "SFU_REG">) outs(%[[SV]] : memref<1x1xf16 +// +// write_to_fifo sends full alloc (not subview) +// CHECK: ktdf.write_to_fifo %[[ALLOC]], %[[PRIV1]]#2 + + +#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d0, d2)> +#map2 = affine_map<(d0, d1) -> (d0, d1)> +#map3 = affine_map<(d0, d1, d2) -> (d0, d1)> +#set = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 255 >= 0, d2 >= 0, -d2 + 63 >= 0)> +#set1 = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 63 >= 0)> + +module { + module { + func.func @absmax_onstick_1core() attributes {grid = [1]} { + call @local_schedule_0() : () -> () + return + } + func.func private @local_schedule_0() + } + ktdf_arch.device @sample_device import("../../Dialect/KTDFArch/sample_device.mlir") + module @local_schedule_0 { + func.func @local_schedule_0() attributes {grid = [1]} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8589934592 = arith.constant 8589934592 : index + %c2 = arith.constant 2 : index + %0 = ktdp.construct_memory_view %c0, sizes: [2, 256, 64], strides: [16384, 64, 1] {coordinate_set = #set, memory_space = #ktdp.memory_space} : memref<2x256x64xf16> + %1 = ktdp.construct_memory_view %c8589934592, sizes: [2, 64], strides: [64, 1] {coordinate_set = #set1, memory_space = #ktdp.memory_space} : memref<2x64xf16> + %memspacecast = memref.memory_space_cast %0 : memref<2x256x64xf16> to memref<2x256x64xf16, #ktdp.memory_space> + %reinterpret_cast = memref.reinterpret_cast %memspacecast to offset: [0], sizes: [2, 256, 64], strides: [16384, 64, 1] : memref<2x256x64xf16, #ktdp.memory_space> to memref<2x256x64xf16, strided<[16384, 64, 1]>, #ktdp.memory_space> + %cast = memref.cast %reinterpret_cast : memref<2x256x64xf16, strided<[16384, 64, 1]>, #ktdp.memory_space> to memref<2x256x64xf16, strided<[16384, 64, 1], offset: ?>, #ktdp.memory_space> + %memspacecast_0 = memref.memory_space_cast %1 : memref<2x64xf16> to memref<2x64xf16, #ktdp.memory_space> + %reinterpret_cast_1 = memref.reinterpret_cast %memspacecast_0 to offset: [0], sizes: [2, 64], strides: [64, 1] : memref<2x64xf16, #ktdp.memory_space> to memref<2x64xf16, strided<[64, 1]>, #ktdp.memory_space> + %cast_2 = memref.cast %reinterpret_cast_1 : memref<2x64xf16, strided<[64, 1]>, #ktdp.memory_space> to memref<2x64xf16, strided<[64, 1], offset: ?>, #ktdp.memory_space> + ktdf.pipeline { + %3:4 = ktdf.private -> (memref<2x1x256x64xf16, #ktdp.memory_space>, memref<2x1x64xf16, #ktdp.memory_space>, !ktdf.token, !ktdf.token) { + %alloc = memref.alloc() : memref<2x1x256x64xf16, #ktdp.memory_space> + %alloc_3 = memref.alloc() : memref<2x1x64xf16, #ktdp.memory_space> + %4 = ktdf.create_token : !ktdf.token + %5 = ktdf.create_token : !ktdf.token + ktdf.private_yield %alloc, %alloc_3, %4, %5 : memref<2x1x256x64xf16, #ktdp.memory_space>, memref<2x1x64xf16, #ktdp.memory_space>, !ktdf.token, !ktdf.token + } + ktdf.stage depends_in(none) depends_out(%3#2) { + scf.for %arg0 = %c0 to %c2 step %c1 { + %4 = arith.subi %arg0, %c0 : index + %5 = arith.divsi %4, %c1 : index + ktdf.data_transfer from %cast[%arg0, 0, 0] size [1, 256, 64] to %3#0[%5, 0, 0, 0] size [1, 1, 256, 64] : memref<2x256x64xf16, strided<[16384, 64, 1], offset: ?>, #ktdp.memory_space>, memref<2x1x256x64xf16, #ktdp.memory_space> + } {loop_type = #ktdf.loop_type} + } {applicable_units = ["MNILU"]} + ktdf.stage depends_in(%3#2) depends_out(%3#3) { + scf.for %arg0 = %c0 to %c2 step %c1 { + %c4 = arith.constant 4 : index + %c0_3 = arith.constant 0 : index + %c1_4 = arith.constant 1 : index + scf.for %arg1 = %c0_3 to %c4 step %c1_4 { + %4 = arith.cmpi eq, %arg1, %c0_3 : index + %c0_5 = arith.constant 0 : index + %c1_6 = arith.constant 1 : index + %c63 = arith.constant 63 : index + ktdf.pipeline { + %5:5 = ktdf.private -> (!ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16>, !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16>, !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, !ktdf.token, !ktdf.token) { + %6 = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> + %7 = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> + %8 = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16> + %9 = ktdf.create_token : !ktdf.token + %10 = ktdf.create_token : !ktdf.token + ktdf.private_yield %6, %7, %8, %9, %10 : !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16>, !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16>, !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, !ktdf.token, !ktdf.token + } + ktdf.stage depends_in(none) depends_out(%5#3) { + %c0_7 = arith.constant 0 : index + %c1_8 = arith.constant 1 : index + %6 = arith.subi %arg0, %c0_7 : index + %7 = arith.divsi %6, %c1_8 : index + scf.if %4 { + } else { + ktdf.data_transfer from %3#1[%7, %c0_7, %c0_7] size [1, 1, 64] to %5#1 size [64] : memref<2x1x64xf16, #ktdp.memory_space>, !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> + } + %c64 = arith.constant 64 : index + scf.for %arg2 = %c0_5 to %c64 step %c1_6 { + ktdf.data_transfer from %3#0[%7, %c0_7, %arg1 * 64 + %arg2, %c0_7] size [1, 1, 1, 64] to %5#0 size [64] : memref<2x1x256x64xf16, #ktdp.memory_space>, !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> + } {loop_type = #ktdf.loop_type} + } {applicable_units = ["L1LU"]} + ktdf.stage depends_in(%5#3) depends_out(%5#4) { + // G1: outer-dim absmax reduction loop, iter_arg initialized + // to tensor.empty. + %6 = tensor.empty() : tensor<1x64xf16> + %c64 = arith.constant 64 : index + %7 = scf.for %arg2 = %c0_5 to %c64 step %c1_6 iter_args(%arg3 = %6) -> (tensor<1x64xf16>) { + %11 = ktdf.read_from_fifo %5#0 : <"L1LU" -> "SFU", 64xf16> -> tensor<1x1x64xf16> + %12 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "reduction", "parallel"]} ins(%11 : tensor<1x1x64xf16>) outs(%arg3 : tensor<1x64xf16>) { + ^bb0(%in: f16, %out: f16): + %13 = math.absf %in : f16 + %14 = math.absf %out : f16 + %15 = arith.maxnumf %13, %14 : f16 + linalg.yield %15 : f16 + } -> tensor<1x64xf16> + scf.yield %12 : tensor<1x64xf16> + } {loop_type = #ktdf.loop_type} + // combine scf.if: first chunk -> pass loop result through; + // subsequent chunks -> add previous partial from FIFO. + %8 = scf.if %4 -> (tensor<1x64xf16>) { + scf.yield %7 : tensor<1x64xf16> + } else { + %11 = ktdf.read_from_fifo %5#1 : <"L1LU" -> "SFU", 64xf16> -> tensor<1x64xf16> + %12 = linalg.generic {indexing_maps = [#map2, #map2], iterator_types = ["parallel", "parallel"]} ins(%11 : tensor<1x64xf16>) outs(%7 : tensor<1x64xf16>) { + ^bb0(%in: f16, %out: f16): + %13 = arith.addf %in, %out : f16 + linalg.yield %13 : f16 + } -> tensor<1x64xf16> + scf.yield %12 : tensor<1x64xf16> + } + // G2: inner-dim absmax reduction over the combined result. + %9 = tensor.empty() : tensor<1x64xf16> + %10 = linalg.generic {indexing_maps = [#map3, #map1], iterator_types = ["parallel", "reduction", "parallel"]} ins(%8 : tensor<1x64xf16>) outs(%9 : tensor<1x64xf16>) { + ^bb0(%in: f16, %out: f16): + %11 = math.absf %in : f16 + %12 = math.absf %out : f16 + %13 = arith.maxnumf %11, %12 : f16 + linalg.yield %13 : f16 + } -> tensor<1x64xf16> + ktdf.write_to_fifo %10, %5#2 : tensor<1x64xf16>, <"SFU" -> "L1SU", 64xf16> + } {applicable_units = ["SFU"]} + ktdf.stage depends_in(%5#4) depends_out(none) { + %c64 = arith.constant 64 : index + scf.for %arg2 = %c0_5 to %c64 step %c1_6 { + %c0_7 = arith.constant 0 : index + %c1_8 = arith.constant 1 : index + %6 = arith.subi %arg0, %c0_7 : index + %7 = arith.divsi %6, %c1_8 : index + %8 = arith.cmpi eq, %arg2, %c63 : index + scf.if %8 { + ktdf.data_transfer from %5#2 size [64] to %3#1[%7, %c0_7, %c0_7] size [1, 1, 64] : !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, memref<2x1x64xf16, #ktdp.memory_space> + } + } {loop_type = #ktdf.loop_type} + } {applicable_units = ["L1SU"]} + } + } + } {loop_type = #ktdf.loop_type} + } {applicable_units = ["L1LU", "SFU", "L1SU"]} + ktdf.stage depends_in(%3#3) depends_out(none) { + scf.for %arg0 = %c0 to %c2 step %c1 { + %4 = arith.subi %arg0, %c0 : index + %5 = arith.divsi %4, %c1 : index + ktdf.data_transfer from %3#1[%5, 0, 0] size [1, 1, 64] to %cast_2[%arg0, %c0 * 64] size [1, 64] : memref<2x1x64xf16, #ktdp.memory_space>, memref<2x64xf16, strided<[64, 1], offset: ?>, #ktdp.memory_space> + } {loop_type = #ktdf.loop_type} + } {applicable_units = ["MNISU"]} + } + return + } + } +} diff --git a/test/Transforms/ReductionDimChunking/outermost_dim_only.mlir b/test/Transforms/ReductionDimChunking/outermost_dim_only.mlir new file mode 100644 index 00000000..40fc36af --- /dev/null +++ b/test/Transforms/ReductionDimChunking/outermost_dim_only.mlir @@ -0,0 +1,198 @@ +// RUN: dataflow-scheduler-opt --reduction-dim-chunking="chunk-size-threshold=32768" %s | FileCheck %s + +// Test: when using the threshold path (no --num-chunks), only the outermost +// reduction dim is chunked even if there are multiple reduction dims. +// +// Input linalg.generic: iterator_types = ["reduction", "reduction", "parallel"] +// over tensor<2x256x64xf16> → tensor<64xf16>. +// +// Total input bytes = 2*256*64*2 = 65536. threshold=32768 forces 2 chunks. +// Outermost dim-0 (size=2) → 2 chunks of 1. +// dim-1 (size=256) is left intact as a reduction inside the linalg.generic. +// +// One scf.for over dim-0 (bound=2). One ktdf.pipeline per iteration. +// is_first = (iv_0 == 0). +// Chunk tensor shape: tensor<1x256x64xf16> (only dim-0 halved). + +// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1, d2) -> (d2)> +// CHECK: #[[$ATTR_2:.+]] = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 255 >= 0, d2 >= 0, -d2 + 63 >= 0)> +// CHECK: #[[$ATTR_3:.+]] = affine_set<(d0) : (d0 >= 0, -d0 + 63 >= 0)> +// CHECK-LABEL: module { +// CHECK: func.func @sum_1core() attributes {grid = [1]} { +// CHECK: call @local_schedule_0() : () -> () +// CHECK: return +// CHECK: } +// CHECK: func.func private @local_schedule_0() +// CHECK: } +// CHECK: ktdf_arch.device @spyre_single_corelet import("../../Dialect/KTDFArch/sample_device.mlir") + +// CHECK-LABEL: module @local_schedule_0 { +// CHECK-NEXT: func.func @local_schedule_0() attributes {grid = [1]} { +// CHECK-NEXT: %[[CONSTANT_0:.*]] = arith.constant 0 : index +// CHECK-NEXT: %[[CONSTANT_1:.*]] = arith.constant 1 : index +// CHECK-NEXT: %[[CONSTANT_2:.*]] = arith.constant 8589934592 : index +// CHECK-NEXT: %[[CONSTRUCT_MEMORY_VIEW_0:.*]] = ktdp.construct_memory_view %[[CONSTANT_0]], sizes: [2, 256, 64], strides: [16384, 64, 1] {coordinate_set = #[[$ATTR_2]], memory_space = #ktdp.memory_space} : memref<2x256x64xf16> +// CHECK-NEXT: %[[CONSTRUCT_MEMORY_VIEW_1:.*]] = ktdp.construct_memory_view %[[CONSTANT_2]], sizes: [64], strides: [1] {coordinate_set = #[[$ATTR_3]], memory_space = #ktdp.memory_space} : memref<64xf16> +// CHECK-NEXT: %[[MEMORY_SPACE_CAST_0:.*]] = memref.memory_space_cast %[[CONSTRUCT_MEMORY_VIEW_0]] : memref<2x256x64xf16> to memref<2x256x64xf16, "DDR"> +// CHECK-NEXT: %[[REINTERPRET_CAST_0:.*]] = memref.reinterpret_cast %[[MEMORY_SPACE_CAST_0]] to offset: [0], sizes: [2, 256, 64], strides: [16384, 64, 1] : memref<2x256x64xf16, "DDR"> to memref<2x256x64xf16, strided<[16384, 64, 1]>, "DDR"> +// CHECK-NEXT: %[[CAST_0:.*]] = memref.cast %[[REINTERPRET_CAST_0]] : memref<2x256x64xf16, strided<[16384, 64, 1]>, "DDR"> to memref<2x256x64xf16, strided<[16384, 64, 1], offset: ?>, "DDR"> +// CHECK-NEXT: %[[MEMORY_SPACE_CAST_1:.*]] = memref.memory_space_cast %[[CONSTRUCT_MEMORY_VIEW_1]] : memref<64xf16> to memref<64xf16, "DDR"> +// CHECK-NEXT: %[[REINTERPRET_CAST_1:.*]] = memref.reinterpret_cast %[[MEMORY_SPACE_CAST_1]] to offset: [0], sizes: [64], strides: [1] : memref<64xf16, "DDR"> to memref<64xf16, strided<[1]>, "DDR"> +// CHECK-NEXT: %[[CAST_1:.*]] = memref.cast %[[REINTERPRET_CAST_1]] : memref<64xf16, strided<[1]>, "DDR"> to memref<64xf16, strided<[1], offset: ?>, "DDR"> +// CHECK-NEXT: ktdf.pipeline { +// CHECK-NEXT: %[[PRIVATE_0:.*]]:4 = ktdf.private -> (memref<1x2x256x64xf16, "L1">, memref<1x64xf16, "L1">, !ktdf.token, !ktdf.token) { +// CHECK-NEXT: %[[ALLOC_0:.*]] = memref.alloc() : memref<1x2x256x64xf16, "L1"> +// CHECK-NEXT: %[[ALLOC_1:.*]] = memref.alloc() : memref<1x64xf16, "L1"> +// CHECK-NEXT: %[[CREATE_TOKEN_0:.*]] = ktdf.create_token : !ktdf.token +// CHECK-NEXT: %[[CREATE_TOKEN_1:.*]] = ktdf.create_token : !ktdf.token +// CHECK-NEXT: ktdf.private_yield %[[ALLOC_0]], %[[ALLOC_1]], %[[CREATE_TOKEN_0]], %[[CREATE_TOKEN_1]] : memref<1x2x256x64xf16, "L1">, memref<1x64xf16, "L1">, !ktdf.token, !ktdf.token +// CHECK-NEXT: } +// CHECK-NEXT: ktdf.stage depends_in(none) depends_out(%[[VAL_0:.*]]#2) { +// CHECK-NEXT: scf.for %[[VAL_1:.*]] = %[[CONSTANT_0]] to %[[CONSTANT_1]] step %[[CONSTANT_1]] { +// CHECK-NEXT: ktdf.data_transfer from %[[CAST_0]]{{\[}}%[[CONSTANT_0]], %[[CONSTANT_0]], %[[CONSTANT_0]]] size [2, 256, 64] to %[[VAL_0]]#0{{\[}}%[[VAL_1]], 0, 0, 0] size [1, 2, 256, 64] : memref<2x256x64xf16, strided<[16384, 64, 1], offset: ?>, "DDR">, memref<1x2x256x64xf16, "L1"> +// CHECK-NEXT: } {loop_type = #ktdf.loop_type} +// CHECK-NEXT: } {applicable_units = ["MNILU"]} +// CHECK-NEXT: ktdf.stage depends_in(%[[VAL_2:.*]]#2) depends_out(%[[VAL_2]]#3) { +// CHECK-NEXT: scf.for %[[VAL_3:.*]] = %[[CONSTANT_0]] to %[[CONSTANT_1]] step %[[CONSTANT_1]] { +// CHECK-NEXT: %[[CONSTANT_3:.*]] = arith.constant 2 : index +// CHECK-NEXT: %[[CONSTANT_4:.*]] = arith.constant 0 : index +// CHECK-NEXT: %[[CONSTANT_5:.*]] = arith.constant 1 : index +// CHECK-NEXT: scf.for %[[VAL_4:.*]] = %[[CONSTANT_4]] to %[[CONSTANT_3]] step %[[CONSTANT_5]] { +// CHECK-NEXT: %[[CMPI_0:.*]] = arith.cmpi eq, %[[VAL_4]], %[[CONSTANT_4]] : index +// CHECK-NEXT: ktdf.pipeline { +// CHECK-NEXT: %[[PRIVATE_1:.*]]:5 = ktdf.private -> (!ktdf.fifo.slot<"L1LU" -> "SFU", 16384xf16>, !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16>, !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, !ktdf.token, !ktdf.token) { +// CHECK-NEXT: %[[FIFO_0:.*]] = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"L1LU" -> "SFU", 16384xf16> +// CHECK-NEXT: %[[FIFO_1:.*]] = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> +// CHECK-NEXT: %[[FIFO_2:.*]] = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16> +// CHECK-NEXT: %[[CREATE_TOKEN_2:.*]] = ktdf.create_token : !ktdf.token +// CHECK-NEXT: %[[CREATE_TOKEN_3:.*]] = ktdf.create_token : !ktdf.token +// CHECK-NEXT: ktdf.private_yield %[[FIFO_0]], %[[FIFO_1]], %[[FIFO_2]], %[[CREATE_TOKEN_2]], %[[CREATE_TOKEN_3]] : !ktdf.fifo.slot<"L1LU" -> "SFU", 16384xf16>, !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16>, !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, !ktdf.token, !ktdf.token +// CHECK-NEXT: } +// CHECK-NEXT: ktdf.stage depends_in(none) depends_out(%[[VAL_5:.*]]#3) { +// CHECK-NEXT: %[[CONSTANT_6:.*]] = arith.constant 0 : index +// CHECK-NEXT: %[[CONSTANT_7:.*]] = arith.constant 1 : index +// CHECK-NEXT: %[[SUBI_0:.*]] = arith.subi %[[VAL_3]], %[[CONSTANT_6]] : index +// CHECK-NEXT: %[[DIVSI_0:.*]] = arith.divsi %[[SUBI_0]], %[[CONSTANT_7]] : index +// CHECK-NEXT: ktdf.data_transfer from %[[VAL_2]]#0{{\[}}%[[DIVSI_0]], %[[VAL_4]], %[[CONSTANT_6]], %[[CONSTANT_6]]] size [1, 1, 256, 64] to %[[VAL_5]]#0 size [16384] : memref<1x2x256x64xf16, "L1">, !ktdf.fifo.slot<"L1LU" -> "SFU", 16384xf16> +// CHECK-NEXT: scf.if %[[CMPI_0]] { +// CHECK-NEXT: } else { +// CHECK-NEXT: ktdf.data_transfer from %[[VAL_2]]#1{{\[}}%[[DIVSI_0]], %[[CONSTANT_6]]] size [1, 64] to %[[VAL_5]]#1 size [64] : memref<1x64xf16, "L1">, !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> +// CHECK-NEXT: } +// CHECK-NEXT: } {applicable_units = ["L1LU"]} +// CHECK-NEXT: ktdf.stage depends_in(%[[VAL_6:.*]]#3) depends_out(%[[VAL_6]]#4) { +// CHECK-NEXT: %[[READ_FROM_FIFO_0:.*]] = ktdf.read_from_fifo %[[VAL_6]]#0 : <"L1LU" -> "SFU", 16384xf16> -> tensor<1x256x64xf16> +// CHECK-NEXT: %[[IF_0:.*]] = scf.if %[[CMPI_0]] -> (tensor<64xf16>) { +// CHECK-NEXT: %[[EMPTY_0:.*]] = tensor.empty() : tensor<64xf16> +// CHECK-NEXT: scf.yield %[[EMPTY_0]] : tensor<64xf16> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[READ_FROM_FIFO_1:.*]] = ktdf.read_from_fifo %[[VAL_6]]#1 : <"L1LU" -> "SFU", 64xf16> -> tensor<64xf16> +// CHECK-NEXT: scf.yield %[[READ_FROM_FIFO_1]] : tensor<64xf16> +// CHECK-NEXT: } +// CHECK-NEXT: %[[GENERIC_0:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]]], iterator_types = ["reduction", "reduction", "parallel"]} ins(%[[READ_FROM_FIFO_0]] : tensor<1x256x64xf16>) outs(%[[IF_0]] : tensor<64xf16>) { +// CHECK-NEXT: ^bb0(%[[VAL_7:.*]]: f16, %[[VAL_8:.*]]: f16): +// CHECK-NEXT: %[[ADDF_0:.*]] = arith.addf %[[VAL_7]], %[[VAL_8]] : f16 +// CHECK-NEXT: linalg.yield %[[ADDF_0]] : f16 +// CHECK-NEXT: } -> tensor<64xf16> +// CHECK-NEXT: ktdf.write_to_fifo %[[GENERIC_0]], %[[VAL_6]]#2 : tensor<64xf16>, <"SFU" -> "L1SU", 64xf16> +// CHECK-NEXT: } {applicable_units = ["SFU"]} +// CHECK-NEXT: ktdf.stage depends_in(%[[VAL_9:.*]]#4) depends_out(none) { +// CHECK-NEXT: %[[CONSTANT_8:.*]] = arith.constant 0 : index +// CHECK-NEXT: %[[CONSTANT_9:.*]] = arith.constant 1 : index +// CHECK-NEXT: %[[SUBI_1:.*]] = arith.subi %[[VAL_3]], %[[CONSTANT_8]] : index +// CHECK-NEXT: %[[DIVSI_1:.*]] = arith.divsi %[[SUBI_1]], %[[CONSTANT_9]] : index +// CHECK-NEXT: ktdf.data_transfer from %[[VAL_9]]#2 size [64] to %[[VAL_2]]#1{{\[}}%[[DIVSI_1]], %[[CONSTANT_8]]] size [1, 64] : !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, memref<1x64xf16, "L1"> +// CHECK-NEXT: } {applicable_units = ["L1SU"]} +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } {loop_type = #ktdf.loop_type} +// CHECK-NEXT: } {applicable_units = ["L1LU", "SFU", "L1SU"]} +// CHECK-NEXT: ktdf.stage depends_in(%[[VAL_10:.*]]#3) depends_out(none) { +// CHECK-NEXT: scf.for %[[VAL_11:.*]] = %[[CONSTANT_0]] to %[[CONSTANT_1]] step %[[CONSTANT_1]] { +// CHECK-NEXT: ktdf.data_transfer from %[[VAL_10]]#1{{\[}}%[[VAL_11]], 0] size [1, 64] to %[[CAST_1]]{{\[}}%[[CONSTANT_0]]] size [64] : memref<1x64xf16, "L1">, memref<64xf16, strided<[1], offset: ?>, "DDR"> +// CHECK-NEXT: } {loop_type = #ktdf.loop_type} +// CHECK-NEXT: } {applicable_units = ["MNISU"]} +// CHECK-NEXT: } +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-NEXT: } + + +#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#map1 = affine_map<(d0, d1, d2) -> (d2)> +#set = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 255 >= 0, d2 >= 0, -d2 + 63 >= 0)> +#set1 = affine_set<(d0) : (d0 >= 0, -d0 + 63 >= 0)> +module { + module { + func.func @sum_1core() attributes {grid = [1]} { + call @local_schedule_0() : () -> () + return + } + func.func private @local_schedule_0() + } + ktdf_arch.device @spyre_single_corelet import("../../Dialect/KTDFArch/sample_device.mlir") + module @local_schedule_0 { + func.func @local_schedule_0() attributes {grid = [1]} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8589934592 = arith.constant 8589934592 : index + %0 = ktdp.construct_memory_view %c0, sizes: [2, 256, 64], strides: [16384, 64, 1] {coordinate_set = #set, memory_space = #ktdp.memory_space} : memref<2x256x64xf16> + %1 = ktdp.construct_memory_view %c8589934592, sizes: [64], strides: [1] {coordinate_set = #set1, memory_space = #ktdp.memory_space} : memref<64xf16> + %memspacecast = memref.memory_space_cast %0 : memref<2x256x64xf16> to memref<2x256x64xf16, "DDR"> + %reinterpret_cast = memref.reinterpret_cast %memspacecast to offset: [0], sizes: [2, 256, 64], strides: [16384, 64, 1] : memref<2x256x64xf16, "DDR"> to memref<2x256x64xf16, strided<[16384, 64, 1]>, "DDR"> + %cast = memref.cast %reinterpret_cast : memref<2x256x64xf16, strided<[16384, 64, 1]>, "DDR"> to memref<2x256x64xf16, strided<[16384, 64, 1], offset: ?>, "DDR"> + %memspacecast_0 = memref.memory_space_cast %1 : memref<64xf16> to memref<64xf16, "DDR"> + %reinterpret_cast_1 = memref.reinterpret_cast %memspacecast_0 to offset: [0], sizes: [64], strides: [1] : memref<64xf16, "DDR"> to memref<64xf16, strided<[1]>, "DDR"> + %cast_2 = memref.cast %reinterpret_cast_1 : memref<64xf16, strided<[1]>, "DDR"> to memref<64xf16, strided<[1], offset: ?>, "DDR"> + ktdf.pipeline { + %2:4 = ktdf.private -> (memref<1x2x256x64xf16, "L1">, memref<1x64xf16, "L1">, !ktdf.token, !ktdf.token) { + %alloc = memref.alloc() : memref<1x2x256x64xf16, "L1"> + %alloc_3 = memref.alloc() : memref<1x64xf16, "L1"> + %3 = ktdf.create_token : !ktdf.token + %4 = ktdf.create_token : !ktdf.token + ktdf.private_yield %alloc, %alloc_3, %3, %4 : memref<1x2x256x64xf16, "L1">, memref<1x64xf16, "L1">, !ktdf.token, !ktdf.token + } + ktdf.stage depends_in(none) depends_out(%2#2) { + scf.for %arg0 = %c0 to %c1 step %c1 { + ktdf.data_transfer from %cast[%c0, %c0, %c0] size [2, 256, 64] to %2#0[%arg0, 0, 0, 0] size [1, 2, 256, 64] : memref<2x256x64xf16, strided<[16384, 64, 1], offset: ?>, "DDR">, memref<1x2x256x64xf16, "L1"> + } {loop_type = #ktdf.loop_type} + } {applicable_units = ["MNILU"]} + ktdf.stage depends_in(%2#2) depends_out(%2#3) { + scf.for %arg0 = %c0 to %c1 step %c1 { + ktdf.pipeline { + %3:4 = ktdf.private -> (!ktdf.fifo.slot<"L1LU" -> "SFU", 32768xf16>, !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, !ktdf.token, !ktdf.token) { + %4 = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"L1LU" -> "SFU", 32768xf16> + %5 = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16> + %6 = ktdf.create_token : !ktdf.token + %7 = ktdf.create_token : !ktdf.token + ktdf.private_yield %4, %5, %6, %7 : !ktdf.fifo.slot<"L1LU" -> "SFU", 32768xf16>, !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, !ktdf.token, !ktdf.token + } + ktdf.stage depends_in(none) depends_out(%3#2) { + ktdf.data_transfer from %2#0[%arg0, 0, 0, 0] size [1, 2, 256, 64] to %3#0 size [32768] : memref<1x2x256x64xf16, "L1">, !ktdf.fifo.slot<"L1LU" -> "SFU", 32768xf16> + } {applicable_units = ["L1LU"]} + ktdf.stage depends_in(%3#2) depends_out(%3#3) { + %4 = ktdf.read_from_fifo %3#0 : <"L1LU" -> "SFU", 32768xf16> -> tensor<2x256x64xf16> + %5 = tensor.empty() : tensor<64xf16> + %6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["reduction", "reduction", "parallel"]} ins(%4 : tensor<2x256x64xf16>) outs(%5 : tensor<64xf16>) { + ^bb0(%in: f16, %out: f16): + %7 = arith.addf %in, %out : f16 + linalg.yield %7 : f16 + } -> tensor<64xf16> + ktdf.write_to_fifo %6, %3#1 : tensor<64xf16>, <"SFU" -> "L1SU", 64xf16> + } {applicable_units = ["SFU"]} + ktdf.stage depends_in(%3#3) depends_out(none) { + ktdf.data_transfer from %3#1 size [64] to %2#1[%arg0, 0] size [1, 64] : !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, memref<1x64xf16, "L1"> + } {applicable_units = ["L1SU"]} + } + } {loop_type = #ktdf.loop_type} + } {applicable_units = ["L1LU", "SFU", "L1SU"]} + ktdf.stage depends_in(%2#3) depends_out(none) { + scf.for %arg0 = %c0 to %c1 step %c1 { + ktdf.data_transfer from %2#1[%arg0, 0] size [1, 64] to %cast_2[%c0] size [64] : memref<1x64xf16, "L1">, memref<64xf16, strided<[1], offset: ?>, "DDR"> + } {loop_type = #ktdf.loop_type} + } {applicable_units = ["MNISU"]} + } + return + } + } +} diff --git a/test/Transforms/ReductionLoopExposure/muil_fifo_dest.mlir b/test/Transforms/ReductionLoopExposure/muil_fifo_dest.mlir index 2053eb74..9573aa36 100644 --- a/test/Transforms/ReductionLoopExposure/muil_fifo_dest.mlir +++ b/test/Transforms/ReductionLoopExposure/muil_fifo_dest.mlir @@ -3,8 +3,9 @@ // CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1, d2) -> (d2)> -// CHECK: #[[$ATTR_2:.+]] = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 255 >= 0, d2 >= 0, -d2 + 63 >= 0)> -// CHECK: #[[$ATTR_3:.+]] = affine_set<(d0) : (d0 >= 0, -d0 + 63 >= 0)> +// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0) -> (d0)> +// CHECK: #[[$ATTR_3:.+]] = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 255 >= 0, d2 >= 0, -d2 + 63 >= 0)> +// CHECK: #[[$ATTR_4:.+]] = affine_set<(d0) : (d0 >= 0, -d0 + 63 >= 0)> // CHECK-LABEL: module { // CHECK: func.func @sum_1core() attributes {grid = [1]} { // CHECK: call @local_schedule_0() : () -> () @@ -19,8 +20,8 @@ // CHECK-NEXT: %[[CONSTANT_0:.*]] = arith.constant 0 : index // CHECK-NEXT: %[[CONSTANT_1:.*]] = arith.constant 1 : index // CHECK-NEXT: %[[CONSTANT_2:.*]] = arith.constant 8589934592 : index -// CHECK-NEXT: %[[CONSTRUCT_MEMORY_VIEW_0:.*]] = ktdp.construct_memory_view %[[CONSTANT_0]], sizes: [2, 256, 64], strides: [16384, 64, 1] {coordinate_set = #[[$ATTR_2]], memory_space = #ktdp.memory_space} : memref<2x256x64xf16> -// CHECK-NEXT: %[[CONSTRUCT_MEMORY_VIEW_1:.*]] = ktdp.construct_memory_view %[[CONSTANT_2]], sizes: [64], strides: [1] {coordinate_set = #[[$ATTR_3]], memory_space = #ktdp.memory_space} : memref<64xf16> +// CHECK-NEXT: %[[CONSTRUCT_MEMORY_VIEW_0:.*]] = ktdp.construct_memory_view %[[CONSTANT_0]], sizes: [2, 256, 64], strides: [16384, 64, 1] {coordinate_set = #[[$ATTR_3]], memory_space = #ktdp.memory_space} : memref<2x256x64xf16> +// CHECK-NEXT: %[[CONSTRUCT_MEMORY_VIEW_1:.*]] = ktdp.construct_memory_view %[[CONSTANT_2]], sizes: [64], strides: [1] {coordinate_set = #[[$ATTR_4]], memory_space = #ktdp.memory_space} : memref<64xf16> // CHECK-NEXT: %[[MEMORY_SPACE_CAST_0:.*]] = memref.memory_space_cast %[[CONSTRUCT_MEMORY_VIEW_0]] : memref<2x256x64xf16> to memref<2x256x64xf16, "DDR"> // CHECK-NEXT: %[[REINTERPRET_CAST_0:.*]] = memref.reinterpret_cast %[[MEMORY_SPACE_CAST_0]] to offset: [0], sizes: [2, 256, 64], strides: [16384, 64, 1] : memref<2x256x64xf16, "DDR"> to memref<2x256x64xf16, strided<[16384, 64, 1]>, "DDR"> // CHECK-NEXT: %[[CAST_0:.*]] = memref.cast %[[REINTERPRET_CAST_0]] : memref<2x256x64xf16, strided<[16384, 64, 1]>, "DDR"> to memref<2x256x64xf16, strided<[16384, 64, 1], offset: ?>, "DDR"> @@ -76,17 +77,11 @@ // CHECK-NEXT: } {loop_type = #ktdf.loop_type} // CHECK-NEXT: } {applicable_units = ["L1LU"]} // CHECK-NEXT: ktdf.stage depends_in(%[[VAL_8:.*]]#3) depends_out(%[[VAL_8]]#4) { -// CHECK-NEXT: %[[IF_0:.*]] = scf.if %[[ANDI_0]] -> (tensor<64xf16>) { -// CHECK-NEXT: %[[EMPTY_0:.*]] = tensor.empty() : tensor<64xf16> -// CHECK-NEXT: scf.yield %[[EMPTY_0]] : tensor<64xf16> -// CHECK-NEXT: } else { -// CHECK-NEXT: %[[READ_FROM_FIFO_0:.*]] = ktdf.read_from_fifo %[[VAL_8]]#1 : <"L1LU" -> "SFU", 64xf16> -> tensor<64xf16> -// CHECK-NEXT: scf.yield %[[READ_FROM_FIFO_0]] : tensor<64xf16> -// CHECK-NEXT: } +// CHECK-NEXT: %[[EMPTY_0:.*]] = tensor.empty() : tensor<64xf16> // CHECK-NEXT: %[[CONSTANT_11:.*]] = arith.constant 64 : index -// CHECK-NEXT: %[[FOR_0:.*]] = scf.for %[[VAL_9:.*]] = %[[CONSTANT_7]] to %[[CONSTANT_11]] step %[[CONSTANT_8]] iter_args(%[[VAL_10:.*]] = %[[IF_0]]) -> (tensor<64xf16>) { -// CHECK-NEXT: %[[READ_FROM_FIFO_1:.*]] = ktdf.read_from_fifo %[[VAL_8]]#0 : <"L1LU" -> "SFU", 64xf16> -> tensor<1x1x64xf16> -// CHECK-NEXT: %[[GENERIC_0:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]]], iterator_types = ["reduction", "reduction", "parallel"]} ins(%[[READ_FROM_FIFO_1]] : tensor<1x1x64xf16>) outs(%[[VAL_10]] : tensor<64xf16>) { +// CHECK-NEXT: %[[FOR_0:.*]] = scf.for %[[VAL_9:.*]] = %[[CONSTANT_7]] to %[[CONSTANT_11]] step %[[CONSTANT_8]] iter_args(%[[VAL_10:.*]] = %[[EMPTY_0]]) -> (tensor<64xf16>) { +// CHECK-NEXT: %[[READ_FROM_FIFO_0:.*]] = ktdf.read_from_fifo %[[VAL_8]]#0 : <"L1LU" -> "SFU", 64xf16> -> tensor<1x1x64xf16> +// CHECK-NEXT: %[[GENERIC_0:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]]], iterator_types = ["reduction", "reduction", "parallel"]} ins(%[[READ_FROM_FIFO_0]] : tensor<1x1x64xf16>) outs(%[[VAL_10]] : tensor<64xf16>) { // CHECK-NEXT: ^bb0(%[[VAL_11:.*]]: f16, %[[VAL_12:.*]]: f16): // CHECK-NEXT: %[[ADDF_0:.*]] = arith.addf %[[VAL_11]], %[[VAL_12]] : f16 // CHECK-NEXT: linalg.yield %[[ADDF_0]] : f16 @@ -97,15 +92,26 @@ // CHECK-NEXT: } // CHECK-NEXT: scf.yield %[[GENERIC_0]] : tensor<64xf16> // CHECK-NEXT: } {loop_type = #ktdf.loop_type} +// CHECK-NEXT: %[[IF_0:.*]] = scf.if %[[ANDI_0]] -> (tensor<64xf16>) { +// CHECK-NEXT: scf.yield %[[FOR_0]] : tensor<64xf16> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[READ_FROM_FIFO_1:.*]] = ktdf.read_from_fifo %[[VAL_8]]#1 : <"L1LU" -> "SFU", 64xf16> -> tensor<64xf16> +// CHECK-NEXT: %[[GENERIC_1:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_2]], #[[$ATTR_2]]], iterator_types = ["parallel"]} ins(%[[READ_FROM_FIFO_1]] : tensor<64xf16>) outs(%[[FOR_0]] : tensor<64xf16>) { +// CHECK-NEXT: ^bb0(%[[VAL_13:.*]]: f16, %[[VAL_14:.*]]: f16): +// CHECK-NEXT: %[[ADDF_1:.*]] = arith.addf %[[VAL_13]], %[[VAL_14]] : f16 +// CHECK-NEXT: linalg.yield %[[ADDF_1]] : f16 +// CHECK-NEXT: } -> tensor<64xf16> +// CHECK-NEXT: scf.yield %[[GENERIC_1]] : tensor<64xf16> +// CHECK-NEXT: } // CHECK-NEXT: } {applicable_units = ["SFU"]} -// CHECK-NEXT: ktdf.stage depends_in(%[[VAL_13:.*]]#4) depends_out(none) { +// CHECK-NEXT: ktdf.stage depends_in(%[[VAL_15:.*]]#4) depends_out(none) { // CHECK-NEXT: %[[CONSTANT_12:.*]] = arith.constant 64 : index -// CHECK-NEXT: scf.for %[[VAL_14:.*]] = %[[CONSTANT_7]] to %[[CONSTANT_12]] step %[[CONSTANT_8]] { +// CHECK-NEXT: scf.for %[[VAL_16:.*]] = %[[CONSTANT_7]] to %[[CONSTANT_12]] step %[[CONSTANT_8]] { // CHECK-NEXT: %[[SUBI_1:.*]] = arith.subi %[[VAL_3]], %[[CONSTANT_3]] : index // CHECK-NEXT: %[[DIVSI_1:.*]] = arith.divsi %[[SUBI_1]], %[[CONSTANT_4]] : index -// CHECK-NEXT: %[[CMPI_3:.*]] = arith.cmpi eq, %[[VAL_14]], %[[CONSTANT_9]] : index +// CHECK-NEXT: %[[CMPI_3:.*]] = arith.cmpi eq, %[[VAL_16]], %[[CONSTANT_9]] : index // CHECK-NEXT: scf.if %[[CMPI_3]] { -// CHECK-NEXT: ktdf.data_transfer from %[[VAL_13]]#2 size [64] to %[[VAL_2]]#1{{\[}}%[[DIVSI_1]], %[[CONSTANT_3]]] size [1, 64] : !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, memref<1x64xf16, "L1"> +// CHECK-NEXT: ktdf.data_transfer from %[[VAL_15]]#2 size [64] to %[[VAL_2]]#1{{\[}}%[[DIVSI_1]], %[[CONSTANT_3]]] size [1, 64] : !ktdf.fifo.slot<"SFU" -> "L1SU", 64xf16>, memref<1x64xf16, "L1"> // CHECK-NEXT: } // CHECK-NEXT: } {loop_type = #ktdf.loop_type} // CHECK-NEXT: } {applicable_units = ["L1SU"]} @@ -114,9 +120,9 @@ // CHECK-NEXT: } // CHECK-NEXT: } {loop_type = #ktdf.loop_type} // CHECK-NEXT: } {applicable_units = ["L1LU", "SFU", "L1SU"]} -// CHECK-NEXT: ktdf.stage depends_in(%[[VAL_15:.*]]#3) depends_out(none) { -// CHECK-NEXT: scf.for %[[VAL_16:.*]] = %[[CONSTANT_0]] to %[[CONSTANT_1]] step %[[CONSTANT_1]] { -// CHECK-NEXT: ktdf.data_transfer from %[[VAL_15]]#1{{\[}}%[[VAL_16]], 0] size [1, 64] to %[[CAST_1]]{{\[}}%[[CONSTANT_0]]] size [64] : memref<1x64xf16, "L1">, memref<64xf16, strided<[1], offset: ?>, "DDR"> +// CHECK-NEXT: ktdf.stage depends_in(%[[VAL_17:.*]]#3) depends_out(none) { +// CHECK-NEXT: scf.for %[[VAL_18:.*]] = %[[CONSTANT_0]] to %[[CONSTANT_1]] step %[[CONSTANT_1]] { +// CHECK-NEXT: ktdf.data_transfer from %[[VAL_17]]#1{{\[}}%[[VAL_18]], 0] size [1, 64] to %[[CAST_1]]{{\[}}%[[CONSTANT_0]]] size [64] : memref<1x64xf16, "L1">, memref<64xf16, strided<[1], offset: ?>, "DDR"> // CHECK-NEXT: } {loop_type = #ktdf.loop_type} // CHECK-NEXT: } {applicable_units = ["MNISU"]} // CHECK-NEXT: }