diff --git a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp index ba1e7d33..6875d474 100644 --- a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp +++ b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp @@ -69,6 +69,31 @@ static bool matchAbsMaxOperands(mlir::arith::MaxNumFOp maxnum_op, return true; } +/// Gets the compare operator standing for \p predicate, or nothing where the +/// unit has none. Only the ordered predicates map: an unordered one asks about +/// NaN, which the compare does not answer. +[[nodiscard]] auto compareOperatorFor(mlir::arith::CmpFPredicate predicate) + -> std::optional { + using Predicate = mlir::arith::CmpFPredicate; + using Operator = mlir::vectorchain::VectorChainElementWiseCompareOperator; + switch (predicate) { + case Predicate::OEQ: + return Operator::compare_eq; + case Predicate::ONE: + return Operator::compare_neq; + case Predicate::OLT: + return Operator::compare_lt; + case Predicate::OLE: + return Operator::compare_le; + case Predicate::OGT: + return Operator::compare_gt; + case Predicate::OGE: + return Operator::compare_ge; + default: + return std::nullopt; + } +} + /// Pattern to lower linalg.generic compute operations struct LowerLinalgGenericPattern : public mlir::OpRewritePattern { @@ -160,6 +185,12 @@ struct LowerLinalgGenericPattern op, op.getLhs(), op.getRhs(), rewriter, identity_map, mlir::vectorchain::VectorChainBinaryOperator::min, compute); }) + .Case([&](mlir::arith::CmpFOp op) { + return lowerCompareFOp(op, rewriter, compute); + }) + .Case([&](mlir::arith::SelectOp op) { + return lowerSelectOp(op, rewriter, compute); + }) .Case([&](mlir::arith::MaxNumFOp op) { mlir::Value lhs, rhs; if (matchAbsMaxOperands(op, lhs, rhs)) { @@ -179,9 +210,14 @@ struct LowerLinalgGenericPattern } return res; } - return lowerBinaryFOp( - op, op.getLhs(), op.getRhs(), rewriter, identity_map, - mlir::vectorchain::VectorChainBinaryOperator::max, compute); + // Only the abs-of-both shape lowers. maxnumf and minnumf + // are the withdrawn 754-2008 operations, whose handling of NaN + // and of signed zero is left to the implementation, and what + // this unit does is not written down -- so mapping them to a + // plain max or min would be a guess. minimumf and maximumf say + // what they mean and are lowered instead. + return rewriter.notifyMatchFailure( + op, "maxnumf outside the abs-max shape is not lowered"); }) .Case([&](mlir::math::AbsFOp op) -> mlir::LogicalResult { @@ -306,6 +342,12 @@ struct LowerLinalgGenericPattern op, op.getLhs(), op.getRhs(), rewriter, identity_map, mlir::vectorchain::VectorChainBinaryOperator::min, compute); }) + .Case([&](mlir::arith::CmpFOp op) { + return lowerCompareFOp(op, rewriter, compute); + }) + .Case([&](mlir::arith::SelectOp op) { + return lowerSelectOp(op, rewriter, compute); + }) .Case([&](mlir::memref::StoreOp op) { return lowerMemRefStore(op, rewriter); }) @@ -327,9 +369,14 @@ struct LowerLinalgGenericPattern } return res; } - return lowerBinaryFOp( - op, op.getLhs(), op.getRhs(), rewriter, identity_map, - mlir::vectorchain::VectorChainBinaryOperator::max, compute); + // Only the abs-of-both shape lowers. maxnumf and minnumf + // are the withdrawn 754-2008 operations, whose handling of NaN + // and of signed zero is left to the implementation, and what + // this unit does is not written down -- so mapping them to a + // plain max or min would be a guess. minimumf and maximumf say + // what they mean and are lowered instead. + return rewriter.notifyMatchFailure( + op, "maxnumf outside the abs-max shape is not lowered"); }) .Case([&](mlir::math::AbsFOp op) -> mlir::LogicalResult { @@ -640,6 +687,49 @@ struct LowerLinalgGenericPattern rewriter.replaceOp(op, binary_op.getData()); return mlir::success(); } + + /// Lowers \p op to an element-wise compare. + /// + /// The result carries the operands' type rather than a boolean: it is what a + /// selection takes as its condition, and the unit keeps it in a lane of the + /// same width. The i1 form of this op is the separate mask operand. + mlir::LogicalResult lowerCompareFOp( + mlir::arith::CmpFOp op, mlir::PatternRewriter& rewriter, + mlir::ktdf_arch::ExecutionUnitOp compute) const { + const auto compare_kind = compareOperatorFor(op.getPredicate()); + if (!compare_kind) return mlir::failure(); + + const auto lhs_ty = llvm::dyn_cast(op.getLhs().getType()); + if (!lhs_ty) return mlir::failure(); + auto operands = getFlattenedVectorType(lhs_ty, compute); + if (!operands) return mlir::failure(); + + auto compare_op = mlir::vectorchain::ElementWiseCompareOp::create( + rewriter, op->getLoc(), operands, op.getLhs(), op.getRhs(), + /*mask=*/nullptr, /*dbgName=*/nullptr, *compare_kind); + + rewriter.replaceOp(op, compare_op.getData()); + return mlir::success(); + } + + /// Lowers \p op to an element-wise selection, taking a lane from one side or + /// the other by the mask a compare left. + mlir::LogicalResult lowerSelectOp( + mlir::arith::SelectOp op, mlir::PatternRewriter& rewriter, + mlir::ktdf_arch::ExecutionUnitOp compute) const { + const auto picked = + llvm::dyn_cast(op.getTrueValue().getType()); + if (!picked) return mlir::failure(); + auto result = getFlattenedVectorType(picked, compute); + if (!result) return mlir::failure(); + + auto selection_op = mlir::vectorchain::ElementWiseSelectionOp::create( + rewriter, op->getLoc(), result, op.getCondition(), op.getTrueValue(), + op.getFalseValue(), /*mask=*/nullptr, /*dbgName=*/nullptr); + + rewriter.replaceOp(op, selection_op.getData()); + return mlir::success(); + } }; /// Pattern to lower linalg.fill into: diff --git a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/OperationLowerings.cpp b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/OperationLowerings.cpp index eb98ccbf..32cd0368 100644 --- a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/OperationLowerings.cpp +++ b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/OperationLowerings.cpp @@ -581,12 +581,17 @@ struct LowerSignalPattern } }; -/// Lower a memref.copy whose source is a ktdf.read_from_fifo (memref form). +/// Lower a memref.copy that MapReductionPartials emitted. /// -/// MapReductionPartials emits: +/// Two sources reach here. A fifo read is already a value, so it is stored +/// straight into the destination: /// %r = ktdf.read_from_fifo ... -> memref<...> /// memref.copy %r, %dest /// +/// A buffer is not, so it is loaded first -- which is how a reduction's +/// accumulator is copied from the buffer holding its identity: +/// memref.copy %identity, %accumulator +/// /// Runs at higher benefit (2) than FoldEmptyCopy (1, a canonicalization /// pattern) so it fires first, preventing FoldEmptyCopy from crashing on /// opaque memory-space attributes. @@ -601,12 +606,27 @@ struct LowerMemRefCopyFromFifoPattern mlir::LogicalResult matchAndRewrite( mlir::memref::CopyOp copy_op, mlir::PatternRewriter& rewriter) const override { - if (!copy_op.getSource().getDefiningOp()) - return mlir::failure(); - rewriter.setInsertionPoint(copy_op); - emitVectorStore(rewriter, copy_op.getLoc(), copy_op.getSource(), - copy_op.getTarget()); + const mlir::Location loc = copy_op.getLoc(); + + mlir::Value stored = copy_op.getSource(); + if (!stored.getDefiningOp()) { + const auto buffer = llvm::dyn_cast(stored.getType()); + if (!buffer) { + return rewriter.notifyMatchFailure( + copy_op, "source is neither a fifo read nor a memref"); + } + if (!buffer.hasStaticShape()) { + return rewriter.notifyMatchFailure( + copy_op, "source has no static shape to load as one vector"); + } + stored = emitVectorLoad(rewriter, loc, + mlir::VectorType::get({buffer.getNumElements()}, + buffer.getElementType()), + stored); + } + + emitVectorStore(rewriter, loc, stored, copy_op.getTarget()); rewriter.eraseOp(copy_op); return mlir::success(); } diff --git a/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp b/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp index bb7242ee..8e1d7faa 100644 --- a/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp +++ b/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp @@ -55,7 +55,7 @@ namespace { if (mlir::isa(op)) { return true; } diff --git a/lib/Dialect/KTDF/Transforms/MapReductionPartials.cpp b/lib/Dialect/KTDF/Transforms/MapReductionPartials.cpp index 94592610..9bc9bd57 100644 --- a/lib/Dialect/KTDF/Transforms/MapReductionPartials.cpp +++ b/lib/Dialect/KTDF/Transforms/MapReductionPartials.cpp @@ -258,8 +258,19 @@ static LogicalResult lowerIterArgInitializer(Value init_val, Value alloc_val, Location loc = empty_op.getLoc(); Value neutral_val = arith::ConstantOp::create(builder, loc, neutral.value()); + + // The identity goes to a buffer of its own, which the accumulator is then + // copied from. Filling the accumulator directly would put the value in the + // copy instruction's immediate, and that field is too narrow for a 32-bit + // one. A buffer written once and never read back is invariant, so what + // hoists such a write can lift it to where a register is initialised + // instead -- which carries the whole value. + auto identity = memref::AllocOp::create( + builder, loc, cast(alloc_val.getType())); linalg::FillOp::create(builder, loc, ValueRange{neutral_val}, - ValueRange{alloc_val}); + ValueRange{identity.getResult()}); + memref::CopyOp::create(builder, loc, identity.getResult(), alloc_val); + empty_op->erase(); return success(); } diff --git a/lib/Pipeline.cpp b/lib/Pipeline.cpp index 59c85d4e..0a4c39b2 100644 --- a/lib/Pipeline.cpp +++ b/lib/Pipeline.cpp @@ -88,6 +88,15 @@ void scheduler::buildSchedulerOptimizationPipeline( pm.addPass(mlir::ktdf::createSplitReductionInnerOuterDimPass()); pm.addPass(mlir::ktdf::createReductionLoopExposurePass()); pm.addPass(mlir::ktdf::createMapReductionPartialsPass()); + // The pass above writes a reduction's identity into a buffer of its own and + // copies the accumulator from it, and can only insert where the reduction is. + // Hoisting lifts that write clear of the loop, which is what lets it become a + // register's initial value rather than an immediate. + { + auto& nested = pm.nest().nest(); + nested.addPass(createHoistInvariantsPass()); + nested.addPass(createHoistConstantStoragePass()); + } pm.addPass(mlir::ktdf::createBroadcastPromotionPass()); pm.addPass(createDoubleBufferingPass(scheduler_ctx)); // Parallelizing before tile selection is beneficial because the tile size diff --git a/lib/Transforms/ApplyDevicePatterns.cpp b/lib/Transforms/ApplyDevicePatterns.cpp index 36543cbd..49ce44f3 100644 --- a/lib/Transforms/ApplyDevicePatterns.cpp +++ b/lib/Transforms/ApplyDevicePatterns.cpp @@ -128,6 +128,48 @@ auto reductionKindOf(mlir::Value yielded) -> llvm::StringRef { return "absmax"; } +// Rewrite helper: ktdf.with_precision(op, params) → DictionaryAttr +// +// Expected `values` entries, two of them in this order: +// [0] mlir::Operation* — the op whose accumulator gives the element type +// [1] mlir::Attribute — the DictionaryAttr to add the precision to +// +// Adds to those parameters the precision a template's `mode` field takes, read +// off the element type the op accumulates in, and leaves the rest of them +// alone. Fails for a type no template names, so the compute is left as it is +// rather than lowered at the wrong width. +auto ktdfWithPrecision(mlir::PatternRewriter& rewriter, + mlir::PDLResultList& results, + llvm::ArrayRef values) + -> mlir::LogicalResult { + assert(values.size() == 2); + auto written = llvm::dyn_cast_if_present( + values[0].cast()); + if (!written || written.getDpsInits().empty()) return mlir::failure(); + auto params = llvm::dyn_cast_if_present( + values[1].cast()); + if (!params) return mlir::failure(); + + auto accumulator = + llvm::dyn_cast(written.getDpsInits().front().getType()); + if (!accumulator) return mlir::failure(); + + const mlir::Type element = accumulator.getElementType(); + llvm::StringRef precision; + if (element.isF16()) { + precision = "fp16"; + } else if (element.isF32()) { + precision = "fp32"; + } else { + return mlir::failure(); + } + + mlir::NamedAttrList named(params); + named.set("prec", rewriter.getStringAttr(precision)); + results.push_back(named.getDictionary(rewriter.getContext())); + return mlir::success(); +} + // Constraint: ktdf.is_reduction_kind(op, expected_kind) // // Expected `values` entries (in order): @@ -173,6 +215,7 @@ class PatternCache : public mlir::ktdf_arch::PatternCache { patterns.registerConstraintFunction("ktdf.is_reduction_kind", ktdfIsReductionKind); patterns.registerRewriteFunction("ktdf.subview_source", ktdfSubviewSource); + patterns.registerRewriteFunction("ktdf.with_precision", ktdfWithPrecision); } MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PatternCache) diff --git a/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-cmpf-select-lowering.mlir b/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-cmpf-select-lowering.mlir new file mode 100644 index 00000000..453249ff --- /dev/null +++ b/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-cmpf-select-lowering.mlir @@ -0,0 +1,64 @@ +// RUN: dataflow-scheduler-opt -pass-pipeline="builtin.module(ktdflowering-to-dfir)" %s | FileCheck %s + +// Tests that a compare and a selection lower to their vectorchain counterparts. +// Neither had a lowering, so the four comparisons and the select of the bundle +// spec could not be expressed at all. +// +// The compare's result carries the operands' type rather than a boolean: it is +// what a selection takes as its condition. The i1 form of that op is its separate +// mask operand. + +// CHECK-LABEL: func.func @cmpf_select +// CHECK: vectorchain.element_wise_compare +// CHECK-SAME: compare_op = #vectorchain +// CHECK: vectorchain.element_wise_selection + +#map_in = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#map_out = affine_map<(d0, d1, d2) -> (d0, d2)> + +module { + ktdf_arch.device @sample_device attributes {} import("../../../../Dialect/KTDFArch/sample_device.mlir") + + // ── a compare feeding a selection ────────────────────────────────────────── + func.func @cmpf_select() attributes {grid = [2]} { + %l1lu0 = dataflow.get_unit {core = 0 : i32, name = "C0-L1LU", type = "L1LU"} : index + %l1lu1 = dataflow.get_unit {core = 1 : i32, name = "C1-L1LU", type = "L1LU"} : index + %sfu0 = dataflow.get_unit {core = 0 : i32, name = "C0-SFU", type = "SFU"} : index + %sfu1 = dataflow.get_unit {core = 1 : i32, name = "C1-SFU", type = "SFU"} : index + %tile_id = ktdp.get_compute_tile_id : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c256 = arith.constant 256 : index + %map_l1lu = uniform.def_immutable_mapping([%c0 -> %l1lu0], [%c1 -> %l1lu1]) : index + %u_l1lu = uniform.query_map(map:%map_l1lu, key:%tile_id) : index + %map_sfu = uniform.def_immutable_mapping([%c0 -> %sfu0], [%c1 -> %sfu1]) : index + %u_sfu = uniform.query_map(map:%map_sfu, key:%tile_id) : index + + %alloc_l1 = memref.alloc() : memref<1x256x64xf16, "L1"> + %fifo = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> + + ktdf_lowering.execute_on %u_l1lu { + scf.for %i = %c0 to %c256 step %c1 { + ktdf.data_transfer from %alloc_l1[%c0, %i, %c0] size [1, 1, 64] to %fifo size [64] : memref<1x256x64xf16, "L1">, !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> + } + } + ktdf_lowering.execute_on %u_sfu { + %alloc = memref.alloc() : memref<1x64xf16, "SFU_REG"> + %neg_inf = arith.constant 0xFF80 : f16 + linalg.fill ins(%neg_inf : f16) outs(%alloc : memref<1x64xf16, "SFU_REG">) + scf.for %i = %c0 to %c256 step %c1 { + %input = ktdf.read_from_fifo %fifo : !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> -> memref<1x1x64xf16> + linalg.generic { + indexing_maps = [#map_in, #map_out], + iterator_types = ["parallel", "reduction", "parallel"] + } ins(%input : memref<1x1x64xf16>) outs(%alloc : memref<1x64xf16, "SFU_REG">) { + ^bb0(%in: f16, %out: f16): + %cmp = arith.cmpf olt, %in, %out : f16 + %result = arith.select %cmp, %in, %out : f16 + linalg.yield %result : f16 + } + } {loop_type = #ktdf.loop_type} + } + return + } +} diff --git a/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-maxnumf-lowering.mlir b/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-maxnumf-lowering.mlir index 8e658ba8..c745cb62 100644 --- a/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-maxnumf-lowering.mlir +++ b/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-maxnumf-lowering.mlir @@ -1,13 +1,10 @@ // RUN: dataflow-scheduler-opt -pass-pipeline="builtin.module(ktdflowering-to-dfir)" %s | FileCheck %s -// Tests that arith.maxnumf lowers to the correct vectorchain binary operator: +// Tests that absf(x) maxnumf absf(y) lowers to the abs_max binary operator. // -// maxnumf alone → binary_operator max -// absf(x) maxnumf absf(y) → binary_operator abs_max - -// CHECK-LABEL: func.func @maxnumf_plain -// CHECK: vectorchain.binary -// CHECK-SAME: binary_op = #vectorchain +// That shape is the only one maxnumf lowers in: on its own it is the withdrawn +// 754-2008 operation, whose NaN and signed-zero handling is left to the +// implementation, so linalg-maxnumf-plain-rejected covers it being turned away. // CHECK-LABEL: func.func @maxnumf_abs_max // CHECK: vectorchain.binary @@ -23,48 +20,6 @@ module { ktdf_arch.device @sample_device attributes {} import("../../../../Dialect/KTDFArch/sample_device.mlir") - // ── plain maxnumf: should lower to max, not abs_max ──────────────────────── - func.func @maxnumf_plain() attributes {grid = [2]} { - %l1lu0 = dataflow.get_unit {core = 0 : i32, name = "C0-L1LU", type = "L1LU"} : index - %l1lu1 = dataflow.get_unit {core = 1 : i32, name = "C1-L1LU", type = "L1LU"} : index - %sfu0 = dataflow.get_unit {core = 0 : i32, name = "C0-SFU", type = "SFU"} : index - %sfu1 = dataflow.get_unit {core = 1 : i32, name = "C1-SFU", type = "SFU"} : index - %tile_id = ktdp.get_compute_tile_id : index - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c256 = arith.constant 256 : index - %map_l1lu = uniform.def_immutable_mapping([%c0 -> %l1lu0], [%c1 -> %l1lu1]) : index - %u_l1lu = uniform.query_map(map:%map_l1lu, key:%tile_id) : index - %map_sfu = uniform.def_immutable_mapping([%c0 -> %sfu0], [%c1 -> %sfu1]) : index - %u_sfu = uniform.query_map(map:%map_sfu, key:%tile_id) : index - - %alloc_l1 = memref.alloc() : memref<1x256x64xf16, "L1"> - %fifo = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> - - ktdf_lowering.execute_on %u_l1lu { - scf.for %i = %c0 to %c256 step %c1 { - ktdf.data_transfer from %alloc_l1[%c0, %i, %c0] size [1, 1, 64] to %fifo size [64] : memref<1x256x64xf16, "L1">, !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> - } - } - ktdf_lowering.execute_on %u_sfu { - %alloc = memref.alloc() : memref<1x64xf16, "SFU_REG"> - %neg_inf = arith.constant 0xFF80 : f16 - linalg.fill ins(%neg_inf : f16) outs(%alloc : memref<1x64xf16, "SFU_REG">) - scf.for %i = %c0 to %c256 step %c1 { - %input = ktdf.read_from_fifo %fifo : !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> -> memref<1x1x64xf16> - linalg.generic { - indexing_maps = [#map_in, #map_out], - iterator_types = ["parallel", "reduction", "parallel"] - } ins(%input : memref<1x1x64xf16>) outs(%alloc : memref<1x64xf16, "SFU_REG">) { - ^bb0(%in: f16, %out: f16): - %result = arith.maxnumf %in, %out : f16 - linalg.yield %result : f16 - } - } {loop_type = #ktdf.loop_type} - } - return - } - // ── abs_max fusion: absf(x) maxnumf absf(y) → abs_max ───────────────────── func.func @maxnumf_abs_max() attributes {grid = [2]} { %l1lu0 = dataflow.get_unit {core = 0 : i32, name = "C0-L1LU", type = "L1LU"} : index diff --git a/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-maxnumf-plain-rejected.mlir b/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-maxnumf-plain-rejected.mlir new file mode 100644 index 00000000..a283242e --- /dev/null +++ b/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-maxnumf-plain-rejected.mlir @@ -0,0 +1,62 @@ +// RUN: not dataflow-scheduler-opt -pass-pipeline="builtin.module(ktdflowering-to-dfir)" %s 2>&1 | FileCheck %s + +// maxnumf on its own is not lowered, and neither is minnumf: they are the +// withdrawn 754-2008 operations, which leave NaN and signed-zero handling to the +// implementation, and what this unit does with either is not written down. Mapping +// them to a plain max or min would be a guess, so they are turned away and +// minimumf and maximumf -- which say what they mean -- are lowered instead. +// +// absf(x) maxnumf absf(y) still lowers, to abs_max; linalg-maxnumf-lowering covers +// it. + +// CHECK: error: failed to run operation lowerings for maxnumf_plain + +#map_in = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#map_out = affine_map<(d0, d1, d2) -> (d0, d2)> + +module { + ktdf_arch.device @sample_device attributes {} import("../../../../Dialect/KTDFArch/sample_device.mlir") + + // ── plain maxnumf: turned away ───────────────────────────────────────────── + func.func @maxnumf_plain() attributes {grid = [2]} { + %l1lu0 = dataflow.get_unit {core = 0 : i32, name = "C0-L1LU", type = "L1LU"} : index + %l1lu1 = dataflow.get_unit {core = 1 : i32, name = "C1-L1LU", type = "L1LU"} : index + %sfu0 = dataflow.get_unit {core = 0 : i32, name = "C0-SFU", type = "SFU"} : index + %sfu1 = dataflow.get_unit {core = 1 : i32, name = "C1-SFU", type = "SFU"} : index + %tile_id = ktdp.get_compute_tile_id : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c256 = arith.constant 256 : index + %map_l1lu = uniform.def_immutable_mapping([%c0 -> %l1lu0], [%c1 -> %l1lu1]) : index + %u_l1lu = uniform.query_map(map:%map_l1lu, key:%tile_id) : index + %map_sfu = uniform.def_immutable_mapping([%c0 -> %sfu0], [%c1 -> %sfu1]) : index + %u_sfu = uniform.query_map(map:%map_sfu, key:%tile_id) : index + + %alloc_l1 = memref.alloc() : memref<1x256x64xf16, "L1"> + %fifo = ktdf.fifo.allocate() -> !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> + + ktdf_lowering.execute_on %u_l1lu { + scf.for %i = %c0 to %c256 step %c1 { + ktdf.data_transfer from %alloc_l1[%c0, %i, %c0] size [1, 1, 64] to %fifo size [64] : memref<1x256x64xf16, "L1">, !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> + } + } + ktdf_lowering.execute_on %u_sfu { + %alloc = memref.alloc() : memref<1x64xf16, "SFU_REG"> + %neg_inf = arith.constant 0xFF80 : f16 + linalg.fill ins(%neg_inf : f16) outs(%alloc : memref<1x64xf16, "SFU_REG">) + scf.for %i = %c0 to %c256 step %c1 { + %input = ktdf.read_from_fifo %fifo : !ktdf.fifo.slot<"L1LU" -> "SFU", 64xf16> -> memref<1x1x64xf16> + linalg.generic { + indexing_maps = [#map_in, #map_out], + iterator_types = ["parallel", "reduction", "parallel"] + } ins(%input : memref<1x1x64xf16>) outs(%alloc : memref<1x64xf16, "SFU_REG">) { + ^bb0(%in: f16, %out: f16): + %result = arith.maxnumf %in, %out : f16 + linalg.yield %result : f16 + } + } {loop_type = #ktdf.loop_type} + } + return + } + +} diff --git a/test/Transforms/MapReductionPartials/basic.mlir b/test/Transforms/MapReductionPartials/basic.mlir index 818841ec..244e9f73 100644 --- a/test/Transforms/MapReductionPartials/basic.mlir +++ b/test/Transforms/MapReductionPartials/basic.mlir @@ -57,7 +57,9 @@ // CHECK-NEXT: ktdf.stage depends_in(%[[VAL_6:.*]]#2) depends_out(%[[VAL_6]]#3) { // CHECK-NEXT: %[[ALLOC_2:.*]] = memref.alloc() : memref<1x64xf16, "SFU_REG"> // CHECK-NEXT: %[[CONSTANT_8:.*]] = arith.constant 0.000000e+00 : f16 -// CHECK-NEXT: linalg.fill ins(%[[CONSTANT_8]] : f16) outs(%[[ALLOC_2]] : memref<1x64xf16, "SFU_REG">) +// CHECK-NEXT: %[[IDENTITY:.*]] = memref.alloc() : memref<1x64xf16, "SFU_REG"> +// CHECK-NEXT: linalg.fill ins(%[[CONSTANT_8]] : f16) outs(%[[IDENTITY]] : memref<1x64xf16, "SFU_REG">) +// CHECK-NEXT: memref.copy %[[IDENTITY]], %[[ALLOC_2]] : memref<1x64xf16, "SFU_REG"> to memref<1x64xf16, "SFU_REG"> // CHECK-NEXT: scf.for %[[VAL_7:.*]] = %[[CONSTANT_4]] to %[[CONSTANT_6]] step %[[CONSTANT_5]] { // CHECK-NEXT: %[[READ_FROM_FIFO_0:.*]] = ktdf.read_from_fifo %[[VAL_6]]#0 : <"L1LU" -> "SFU", 64xf16> -> memref<1x1x64xf16> // CHECK-NEXT: linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]]], iterator_types = ["parallel", "reduction", "parallel"]} ins(%[[READ_FROM_FIFO_0]] : memref<1x1x64xf16>) outs(%[[ALLOC_2]] : memref<1x64xf16, "SFU_REG">) { diff --git a/test/Transforms/MapReductionPartials/inner_outer_dim_reduction.mlir b/test/Transforms/MapReductionPartials/inner_outer_dim_reduction.mlir index e0436e2c..2533a853 100644 --- a/test/Transforms/MapReductionPartials/inner_outer_dim_reduction.mlir +++ b/test/Transforms/MapReductionPartials/inner_outer_dim_reduction.mlir @@ -68,7 +68,9 @@ // CHECK-NEXT: ktdf.stage depends_in(%[[VAL_6:.*]]#2) depends_out(%[[VAL_6]]#3) { // CHECK-NEXT: %[[ALLOC_2:.*]] = memref.alloc() : memref<1x64xf16, "SFU_REG"> // CHECK-NEXT: %[[CONSTANT_8:.*]] = arith.constant 0.000000e+00 : f16 -// CHECK-NEXT: linalg.fill ins(%[[CONSTANT_8]] : f16) outs(%[[ALLOC_2]] : memref<1x64xf16, "SFU_REG">) +// CHECK-NEXT: %[[IDENTITY:.*]] = memref.alloc() : memref<1x64xf16, "SFU_REG"> +// CHECK-NEXT: linalg.fill ins(%[[CONSTANT_8]] : f16) outs(%[[IDENTITY]] : memref<1x64xf16, "SFU_REG">) +// CHECK-NEXT: memref.copy %[[IDENTITY]], %[[ALLOC_2]] : memref<1x64xf16, "SFU_REG"> to memref<1x64xf16, "SFU_REG"> // CHECK-NEXT: %[[CONSTANT_9:.*]] = arith.constant 2 : index // CHECK-NEXT: scf.for %[[VAL_7:.*]] = %[[CONSTANT_4]] to %[[CONSTANT_9]] step %[[CONSTANT_5]] { // CHECK-NEXT: %[[READ_FROM_FIFO_0:.*]] = ktdf.read_from_fifo %[[VAL_6]]#0 : <"L1LU" -> "SFU", 64xf16> -> memref<1x1x64xf16> diff --git a/test/Transforms/MapReductionPartials/multi_outer_dim_reduction.mlir b/test/Transforms/MapReductionPartials/multi_outer_dim_reduction.mlir index 07b338b1..725d7fa1 100644 --- a/test/Transforms/MapReductionPartials/multi_outer_dim_reduction.mlir +++ b/test/Transforms/MapReductionPartials/multi_outer_dim_reduction.mlir @@ -57,7 +57,9 @@ // CHECK-NEXT: ktdf.stage depends_in(%[[VAL_7:.*]]#2) depends_out(%[[VAL_7]]#3) { // CHECK-NEXT: %[[ALLOC_2:.*]] = memref.alloc() : memref<64xf16, "SFU_REG"> // CHECK-NEXT: %[[CONSTANT_10:.*]] = arith.constant 0.000000e+00 : f16 -// CHECK-NEXT: linalg.fill ins(%[[CONSTANT_10]] : f16) outs(%[[ALLOC_2]] : memref<64xf16, "SFU_REG">) +// CHECK-NEXT: %[[IDENTITY:.*]] = memref.alloc() : memref<64xf16, "SFU_REG"> +// CHECK-NEXT: linalg.fill ins(%[[CONSTANT_10]] : f16) outs(%[[IDENTITY]] : memref<64xf16, "SFU_REG">) +// CHECK-NEXT: memref.copy %[[IDENTITY]], %[[ALLOC_2]] : memref<64xf16, "SFU_REG"> to memref<64xf16, "SFU_REG"> // CHECK-NEXT: %[[CONSTANT_11:.*]] = arith.constant 2 : index // CHECK-NEXT: %[[CONSTANT_12:.*]] = arith.constant 256 : index // CHECK-NEXT: scf.for %[[VAL_8:.*]] = %[[CONSTANT_4]] to %[[CONSTANT_11]] step %[[CONSTANT_5]] { diff --git a/test/Transforms/MapReductionPartials/neutral_element.mlir b/test/Transforms/MapReductionPartials/neutral_element.mlir index a135b1fc..450f5534 100644 --- a/test/Transforms/MapReductionPartials/neutral_element.mlir +++ b/test/Transforms/MapReductionPartials/neutral_element.mlir @@ -5,26 +5,32 @@ // CHECK-LABEL: func.func @mulf_reduction // CHECK: arith.constant 1.000000e+00 : f16 +// CHECK-NEXT: memref.alloc // CHECK-NEXT: linalg.fill // CHECK-LABEL: func.func @maximumf_reduction // CHECK: arith.constant 0xFC00 : f16 +// CHECK-NEXT: memref.alloc // CHECK-NEXT: linalg.fill // CHECK-LABEL: func.func @minimumf_reduction // CHECK: arith.constant 0x7C00 : f16 +// CHECK-NEXT: memref.alloc // CHECK-NEXT: linalg.fill // CHECK-LABEL: func.func @subf_reduction // CHECK: arith.constant 0.000000e+00 : f16 +// CHECK-NEXT: memref.alloc // CHECK-NEXT: linalg.fill // CHECK-LABEL: func.func @absmax_reduction // CHECK: arith.constant 0.000000e+00 : f16 +// CHECK-NEXT: memref.alloc // CHECK-NEXT: linalg.fill // CHECK-LABEL: func.func @addi_reduction // CHECK: arith.constant 0 : i16 +// CHECK-NEXT: memref.alloc // CHECK-NEXT: linalg.fill #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> diff --git a/test/Transforms/MapReductionPartials/scf_if_init.mlir b/test/Transforms/MapReductionPartials/scf_if_init.mlir index 34673018..54b15aa4 100644 --- a/test/Transforms/MapReductionPartials/scf_if_init.mlir +++ b/test/Transforms/MapReductionPartials/scf_if_init.mlir @@ -59,7 +59,9 @@ // CHECK-NEXT: %[[ALLOC_2:.*]] = memref.alloc() : memref<1x64xf16, "SFU_REG"> // CHECK-NEXT: scf.if %[[CMPI_INIT]] { // CHECK-NEXT: %[[CONSTANT_8:.*]] = arith.constant 0.000000e+00 : f16 -// CHECK-NEXT: linalg.fill ins(%[[CONSTANT_8]] : f16) outs(%[[ALLOC_2]] : memref<1x64xf16, "SFU_REG">) +// CHECK-NEXT: %[[IDENTITY:.*]] = memref.alloc() : memref<1x64xf16, "SFU_REG"> +// CHECK-NEXT: linalg.fill ins(%[[CONSTANT_8]] : f16) outs(%[[IDENTITY]] : memref<1x64xf16, "SFU_REG">) +// CHECK-NEXT: memref.copy %[[IDENTITY]], %[[ALLOC_2]] : memref<1x64xf16, "SFU_REG"> to memref<1x64xf16, "SFU_REG"> // CHECK-NEXT: } else { // CHECK-NEXT: %[[READ_INIT:.*]] = ktdf.read_from_fifo %[[VAL_6]]#0 : <"L1LU" -> "SFU", 64xf16> -> memref<1x64xf16> // CHECK-NEXT: memref.copy %[[READ_INIT]], %[[ALLOC_2]] : memref<1x64xf16> to memref<1x64xf16, "SFU_REG"> diff --git a/test/dataflow-scheduler-opt/emit-dfir-pipeline.mlir b/test/dataflow-scheduler-opt/emit-dfir-pipeline.mlir index 4a674886..37809f41 100644 --- a/test/dataflow-scheduler-opt/emit-dfir-pipeline.mlir +++ b/test/dataflow-scheduler-opt/emit-dfir-pipeline.mlir @@ -45,6 +45,12 @@ // CHECK-NEXT: split-reduction-inner-outer-dim // CHECK-NEXT: reduction-loop-exposure // CHECK-NEXT: map-reduction-partials +// CHECK-NEXT: builtin.module( +// CHECK-NEXT: func.func( +// CHECK-NEXT: hoist-invariants +// CHECK-NEXT: hoist-constant-storage +// CHECK-NEXT: ) +// CHECK-NEXT: ) // CHECK-NEXT: broadcast-promotion // CHECK-NEXT: double-buffering // CHECK-NEXT: parallelize-loops-across-instances