From 1f7b5512881c50132b4a36e13e948c44d2e8a0ac Mon Sep 17 00:00:00 2001 From: Swagath Venkataramani Date: Fri, 4 Sep 2026 13:34:32 -0400 Subject: [PATCH 1/8] [ApplyDevicePatterns] Let a pattern take a template's precision from the IR A template whose instructions are written for either width needs the caller to say which one, and the callers that say it hardcode a precision per pattern -- so a pattern covering both widths has to be written twice. ktdf.with_precision(op, params) reads the element type the op accumulates in and adds the precision to the pattern's own parameter dictionary, leaving whatever else that dictionary carries alone. It fails for a type no template names, so the compute is left as it is rather than lowered at the wrong width. Signed-off-by: Swagath Venkataramani --- lib/Transforms/ApplyDevicePatterns.cpp | 38 ++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/lib/Transforms/ApplyDevicePatterns.cpp b/lib/Transforms/ApplyDevicePatterns.cpp index 36543cbd..c672140d 100644 --- a/lib/Transforms/ApplyDevicePatterns.cpp +++ b/lib/Transforms/ApplyDevicePatterns.cpp @@ -128,6 +128,43 @@ auto reductionKindOf(mlir::Value yielded) -> llvm::StringRef { return "absmax"; } +// Rewrite helper: ktdf.with_precision(op, params) → DictionaryAttr +// +// Adds to \p params the precision a template's `mode` field takes, read off the +// element type \p op accumulates in. Fails for a type no template names, so the +// compute is left alone 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 +210,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) From f1ec6ac51e8aa4ad3df9ac25971d9671706ac9fa Mon Sep 17 00:00:00 2001 From: Swagath Venkataramani Date: Fri, 4 Sep 2026 13:34:32 -0400 Subject: [PATCH 2/8] [KTDFLowToDFIR] Load a buffer before storing it in a memref.copy The copy lowering stored its source straight into the destination, which is right for a fifo read -- that is already a value -- and wrong for a buffer, where the store is handed a memref: error: 'agen.vector_store' op operand #0 must be vector of any type values or Dataflow custom vector type, but got 'memref<1x32xf32>' A buffer source is loaded first now. Nothing emits that form yet; the pass in the commit below does. Signed-off-by: Swagath Venkataramani --- .../KTDFLowToDFIR/OperationLowerings.cpp | 27 ++++++++++++++----- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/OperationLowerings.cpp b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/OperationLowerings.cpp index eb98ccbf..afd5191a 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,20 @@ 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 || !buffer.hasStaticShape()) return mlir::failure(); + 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(); } From 29d8123566293259caba3e92b2228e7520e920eb Mon Sep 17 00:00:00 2001 From: Swagath Venkataramani Date: Fri, 4 Sep 2026 13:34:32 -0400 Subject: [PATCH 3/8] [MapReductionPartials] Reduce from an identity held in a register A reduction's accumulator was filled with its identity in place. Where a device turns such a fill into an instruction's immediate, the width of that field bounds what the identity may be: enough for an f16 negative infinity and not for an f32 one, which failed late and said only that an immediate was out of range. Giving a register its value up front has no such bound, and reaching it needs the write to be invariant and clear of the loop. So the identity goes to a buffer of its own and the accumulator is copied from it. That buffer is written once and never read back, where the accumulator never was invariant -- the reduction writes it too, and the combiner reads a subview of it. Hoisting after this pass then lifts the fill out of the loop, which is what lets the value be given to the register rather than carried in an instruction. This applies at every width, so an f16 reduction now spends a register on an identity that would have fitted an immediate. Making it conditional on the value fitting would avoid that, at the cost of two ways of doing one thing. Signed-off-by: Swagath Venkataramani --- .../KTDF/Transforms/MapReductionPartials.cpp | 13 ++++++++++++- lib/Pipeline.cpp | 9 +++++++++ test/Transforms/MapReductionPartials/basic.mlir | 4 +++- .../inner_outer_dim_reduction.mlir | 4 +++- .../multi_outer_dim_reduction.mlir | 4 +++- .../MapReductionPartials/neutral_element.mlir | 6 ++++++ .../MapReductionPartials/scf_if_init.mlir | 4 +++- test/dataflow-scheduler-opt/emit-dfir-pipeline.mlir | 6 ++++++ 8 files changed, 45 insertions(+), 5 deletions(-) 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/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 From 533d4ad1eceab1ce13bee6312c1f4b0b96f0f56d Mon Sep 17 00:00:00 2001 From: Swagath Venkataramani Date: Tue, 8 Sep 2026 22:03:26 -0400 Subject: [PATCH 4/8] [KTDFLowToDFIR] Lower arith.minnumf to a min The frontend legality check accepts minnumf alongside minimumf, maxnumf and maximumf, but only minnumf had no lowering. A body using it got through the check and died further down saying only error: unsupported operation type in linalg.generic body It lowers to a plain min. There is no absolute-min operator, so unlike maxnumf it has no second shape to select for an abs-of-both pattern. Signed-off-by: Swagath Venkataramani --- .../KTDFLowToDFIR/LinalgLowering.cpp | 14 +++++ .../linalg-minnumf-lowering.mlir | 62 +++++++++++++++++++ 2 files changed, 76 insertions(+) create mode 100644 test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-minnumf-lowering.mlir diff --git a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp index ba1e7d33..1df333f3 100644 --- a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp +++ b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp @@ -160,6 +160,13 @@ struct LowerLinalgGenericPattern op, op.getLhs(), op.getRhs(), rewriter, identity_map, mlir::vectorchain::VectorChainBinaryOperator::min, compute); }) + .Case([&](mlir::arith::MinNumFOp op) { + // A plain min: there is no absolute-min operator for the + // abs-of-both shape that maxnumf has. + return lowerBinaryFOp( + op, op.getLhs(), op.getRhs(), rewriter, identity_map, + mlir::vectorchain::VectorChainBinaryOperator::min); + }) .Case([&](mlir::arith::MaxNumFOp op) { mlir::Value lhs, rhs; if (matchAbsMaxOperands(op, lhs, rhs)) { @@ -306,6 +313,13 @@ struct LowerLinalgGenericPattern op, op.getLhs(), op.getRhs(), rewriter, identity_map, mlir::vectorchain::VectorChainBinaryOperator::min, compute); }) + .Case([&](mlir::arith::MinNumFOp op) { + // A plain min: there is no absolute-min operator for the + // abs-of-both shape that maxnumf has. + return lowerBinaryFOp( + op, op.getLhs(), op.getRhs(), rewriter, identity_map, + mlir::vectorchain::VectorChainBinaryOperator::min); + }) .Case([&](mlir::memref::StoreOp op) { return lowerMemRefStore(op, rewriter); }) diff --git a/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-minnumf-lowering.mlir b/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-minnumf-lowering.mlir new file mode 100644 index 00000000..0b532636 --- /dev/null +++ b/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-minnumf-lowering.mlir @@ -0,0 +1,62 @@ +// RUN: dataflow-scheduler-opt -pass-pipeline="builtin.module(ktdflowering-to-dfir)" %s | FileCheck %s + +// Tests that arith.minnumf lowers to the min binary operator. It was accepted by +// the frontend legality check and had no lowering, so a body using it got as far +// as here and then said only "unsupported operation type in linalg.generic body". +// +// There is no absolute-min operator, so minnumf has no second shape to select the +// way maxnumf picks abs_max for the abs-of-both case. + +// CHECK-LABEL: func.func @minnumf_plain +// CHECK: vectorchain.binary +// CHECK-SAME: binary_op = #vectorchain + +#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") + + // ── minnumf: lowers to min ───────────────────────────────────────────────── + func.func @minnumf_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.minnumf %in, %out : f16 + linalg.yield %result : f16 + } + } {loop_type = #ktdf.loop_type} + } + return + } + +} From f92e951c5319fbbff3d0c4aac13370339e7bebdf Mon Sep 17 00:00:00 2001 From: Swagath Venkataramani Date: Wed, 9 Sep 2026 06:48:35 -0400 Subject: [PATCH 5/8] [KTDFLowToDFIR] Lower a compare and a selection Neither arith.cmpf nor arith.select had a lowering, and neither was in the frontend's allow-list, so a comparison could not be expressed at all. Both dialect ops they map to already existed, along with the lowerings below them. The compare's result carries the operands' type rather than a boolean: it is what a selection takes as its condition, and the i1 form of that op is its separate mask operand. Only the ordered predicates map -- an unordered one asks about NaN, which the compare does not answer. What this reaches depends on the shape. A select whose two sides are the values compared fuses with the compare into a single min or max below here, and works end to end. A select over anything else -- picking between two values the comparison did not name -- still leaves the compare with no reader further down, so it is expressible but not yet compilable. Signed-off-by: Swagath Venkataramani --- .../KTDFLowToDFIR/LinalgLowering.cpp | 75 +++++++++++++++++++ .../KTIRToScheduleIR/KTIRLegalityCheck.cpp | 11 +-- .../linalg-cmpf-select-lowering.mlir | 64 ++++++++++++++++ 3 files changed, 145 insertions(+), 5 deletions(-) create mode 100644 test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-cmpf-select-lowering.mlir diff --git a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp index 1df333f3..525a4900 100644 --- a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp +++ b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp @@ -70,6 +70,31 @@ static bool matchAbsMaxOperands(mlir::arith::MaxNumFOp maxnum_op, } /// Pattern to lower linalg.generic compute operations +/// 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; + } +} + struct LowerLinalgGenericPattern : public mlir::OpRewritePattern { LowerLinalgGenericPattern(mlir::MLIRContext* context, @@ -167,6 +192,12 @@ struct LowerLinalgGenericPattern op, op.getLhs(), op.getRhs(), rewriter, identity_map, mlir::vectorchain::VectorChainBinaryOperator::min); }) + .Case([&](mlir::arith::CmpFOp op) { + return lowerCompareFOp(op, rewriter); + }) + .Case([&](mlir::arith::SelectOp op) { + return lowerSelectOp(op, rewriter); + }) .Case([&](mlir::arith::MaxNumFOp op) { mlir::Value lhs, rhs; if (matchAbsMaxOperands(op, lhs, rhs)) { @@ -320,6 +351,12 @@ struct LowerLinalgGenericPattern op, op.getLhs(), op.getRhs(), rewriter, identity_map, mlir::vectorchain::VectorChainBinaryOperator::min); }) + .Case([&](mlir::arith::CmpFOp op) { + return lowerCompareFOp(op, rewriter); + }) + .Case([&](mlir::arith::SelectOp op) { + return lowerSelectOp(op, rewriter); + }) .Case([&](mlir::memref::StoreOp op) { return lowerMemRefStore(op, rewriter); }) @@ -654,6 +691,44 @@ 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) const { + const auto compare_kind = compareOperatorFor(op.getPredicate()); + if (!compare_kind) return mlir::failure(); + + auto operands = + getFlattenedVectorType(op.getLhs().getType(), resource_kinds_); + 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) const { + auto result = + getFlattenedVectorType(op.getTrueValue().getType(), resource_kinds_); + 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/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp b/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp index bb7242ee..1ee47d8f 100644 --- a/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp +++ b/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp @@ -52,11 +52,12 @@ namespace { /// Determines whether @p op is legal within the body of a 'linalg.generic'. [[nodiscard]] auto isLegalGenericBodyOp(mlir::Operation* op) -> bool { // Accept supported arith operations and the 'linalg.yield' terminator. - if (mlir::isa(op)) { + if (mlir::isa< + mlir::arith::AddFOp, mlir::arith::MulFOp, mlir::arith::SubFOp, + mlir::arith::AddIOp, mlir::arith::MaximumFOp, mlir::arith::MinimumFOp, + mlir::arith::MaxNumFOp, mlir::arith::MinNumFOp, mlir::arith::CmpFOp, + mlir::arith::SelectOp, mlir::math::AbsFOp, mlir::linalg::YieldOp>( + op)) { return true; } 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 + } +} From 066a373abe039ab8c69e4264fc4d1082ac18b08c Mon Sep 17 00:00:00 2001 From: Swagath Venkataramani Date: Wed, 9 Sep 2026 13:05:24 -0400 Subject: [PATCH 6/8] Address review on the compare, the copy and the precision helper - A copy that cannot be lowered says why through notifyMatchFailure rather than failing silently, so the reason shows up when debugging. - ktdf.with_precision names its two values entries in the header, as the constraints beside it do. - compareOperatorFor was inserted between LowerLinalgGenericPattern's own doc comment and the struct, leaving that comment on the wrong thing. - The minnumf case says what it assumes: the unit has one minimum and minnumf differs from minimumf only in which operand a NaN takes, so mapping both to it is exact only where neither is NaN. maxnumf and maximumf already map that way, so restricting one without the other would leave the two inconsistent. Signed-off-by: Swagath Venkataramani --- .../KTDFLowToDFIR/LinalgLowering.cpp | 14 +++++++++----- .../KTDFLowToDFIR/OperationLowerings.cpp | 9 ++++++++- lib/Transforms/ApplyDevicePatterns.cpp | 11 ++++++++--- 3 files changed, 25 insertions(+), 9 deletions(-) diff --git a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp index 525a4900..79d4bc10 100644 --- a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp +++ b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp @@ -69,7 +69,6 @@ static bool matchAbsMaxOperands(mlir::arith::MaxNumFOp maxnum_op, return true; } -/// Pattern to lower linalg.generic compute operations /// 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. @@ -95,6 +94,7 @@ static bool matchAbsMaxOperands(mlir::arith::MaxNumFOp maxnum_op, } } +/// Pattern to lower linalg.generic compute operations struct LowerLinalgGenericPattern : public mlir::OpRewritePattern { LowerLinalgGenericPattern(mlir::MLIRContext* context, @@ -186,8 +186,10 @@ struct LowerLinalgGenericPattern mlir::vectorchain::VectorChainBinaryOperator::min, compute); }) .Case([&](mlir::arith::MinNumFOp op) { - // A plain min: there is no absolute-min operator for the - // abs-of-both shape that maxnumf has. + // The unit has one minimum, and minnumf differs from minimumf + // only in which operand a NaN takes, so both map to it and + // neither is exact where one is NaN. maxnumf and maximumf + // already map this way. return lowerBinaryFOp( op, op.getLhs(), op.getRhs(), rewriter, identity_map, mlir::vectorchain::VectorChainBinaryOperator::min); @@ -345,8 +347,10 @@ struct LowerLinalgGenericPattern mlir::vectorchain::VectorChainBinaryOperator::min, compute); }) .Case([&](mlir::arith::MinNumFOp op) { - // A plain min: there is no absolute-min operator for the - // abs-of-both shape that maxnumf has. + // The unit has one minimum, and minnumf differs from minimumf + // only in which operand a NaN takes, so both map to it and + // neither is exact where one is NaN. maxnumf and maximumf + // already map this way. return lowerBinaryFOp( op, op.getLhs(), op.getRhs(), rewriter, identity_map, mlir::vectorchain::VectorChainBinaryOperator::min); diff --git a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/OperationLowerings.cpp b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/OperationLowerings.cpp index afd5191a..32cd0368 100644 --- a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/OperationLowerings.cpp +++ b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/OperationLowerings.cpp @@ -612,7 +612,14 @@ struct LowerMemRefCopyFromFifoPattern mlir::Value stored = copy_op.getSource(); if (!stored.getDefiningOp()) { const auto buffer = llvm::dyn_cast(stored.getType()); - if (!buffer || !buffer.hasStaticShape()) return mlir::failure(); + 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()), diff --git a/lib/Transforms/ApplyDevicePatterns.cpp b/lib/Transforms/ApplyDevicePatterns.cpp index c672140d..49ce44f3 100644 --- a/lib/Transforms/ApplyDevicePatterns.cpp +++ b/lib/Transforms/ApplyDevicePatterns.cpp @@ -130,9 +130,14 @@ auto reductionKindOf(mlir::Value yielded) -> llvm::StringRef { // Rewrite helper: ktdf.with_precision(op, params) → DictionaryAttr // -// Adds to \p params the precision a template's `mode` field takes, read off the -// element type \p op accumulates in. Fails for a type no template names, so the -// compute is left alone rather than lowered at the wrong width. +// 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) From f17cf0568a40baffc03b8a61c95e93e0eb9c4e45 Mon Sep 17 00:00:00 2001 From: Swagath Venkataramani Date: Thu, 10 Sep 2026 08:11:45 -0400 Subject: [PATCH 7/8] [KTDFLowToDFIR] Lower maxnumf only as an abs-max, and not minnumf maxnumf and minnumf are the 754-2008 operations, withdrawn in 2019. They leave NaN handling and signed zero to the implementation, and what this unit does with either is not written down -- so mapping them to a plain max or min was a guess. minimumf and maximumf say what they mean and are lowered instead. So minnumf goes back to being rejected, at the frontend rather than below it, and maxnumf lowers only in the abs-of-both shape that selects abs_max. The reduction path and MapReductionPartials already turned both away, so this is the elementwise path catching up rather than a new position. linalg-maxnumf-lowering keeps the two abs-max cases; the plain one moves to linalg-maxnumf-plain-rejected, which pins it being turned away. Signed-off-by: Swagath Venkataramani --- .../KTDFLowToDFIR/LinalgLowering.cpp | 40 ++++++-------- .../KTIRToScheduleIR/KTIRLegalityCheck.cpp | 11 ++-- .../linalg-maxnumf-lowering.mlir | 53 ++----------------- ...lir => linalg-maxnumf-plain-rejected.mlir} | 24 ++++----- 4 files changed, 37 insertions(+), 91 deletions(-) rename test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/{linalg-minnumf-lowering.mlir => linalg-maxnumf-plain-rejected.mlir} (70%) diff --git a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp index 79d4bc10..12c878dc 100644 --- a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp +++ b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp @@ -185,15 +185,6 @@ struct LowerLinalgGenericPattern op, op.getLhs(), op.getRhs(), rewriter, identity_map, mlir::vectorchain::VectorChainBinaryOperator::min, compute); }) - .Case([&](mlir::arith::MinNumFOp op) { - // The unit has one minimum, and minnumf differs from minimumf - // only in which operand a NaN takes, so both map to it and - // neither is exact where one is NaN. maxnumf and maximumf - // already map this way. - return lowerBinaryFOp( - op, op.getLhs(), op.getRhs(), rewriter, identity_map, - mlir::vectorchain::VectorChainBinaryOperator::min); - }) .Case([&](mlir::arith::CmpFOp op) { return lowerCompareFOp(op, rewriter); }) @@ -219,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 { @@ -346,15 +342,6 @@ struct LowerLinalgGenericPattern op, op.getLhs(), op.getRhs(), rewriter, identity_map, mlir::vectorchain::VectorChainBinaryOperator::min, compute); }) - .Case([&](mlir::arith::MinNumFOp op) { - // The unit has one minimum, and minnumf differs from minimumf - // only in which operand a NaN takes, so both map to it and - // neither is exact where one is NaN. maxnumf and maximumf - // already map this way. - return lowerBinaryFOp( - op, op.getLhs(), op.getRhs(), rewriter, identity_map, - mlir::vectorchain::VectorChainBinaryOperator::min); - }) .Case([&](mlir::arith::CmpFOp op) { return lowerCompareFOp(op, rewriter); }) @@ -382,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 { diff --git a/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp b/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp index 1ee47d8f..8e1d7faa 100644 --- a/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp +++ b/lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp @@ -52,12 +52,11 @@ namespace { /// Determines whether @p op is legal within the body of a 'linalg.generic'. [[nodiscard]] auto isLegalGenericBodyOp(mlir::Operation* op) -> bool { // Accept supported arith operations and the 'linalg.yield' terminator. - if (mlir::isa< - mlir::arith::AddFOp, mlir::arith::MulFOp, mlir::arith::SubFOp, - mlir::arith::AddIOp, mlir::arith::MaximumFOp, mlir::arith::MinimumFOp, - mlir::arith::MaxNumFOp, mlir::arith::MinNumFOp, mlir::arith::CmpFOp, - mlir::arith::SelectOp, mlir::math::AbsFOp, mlir::linalg::YieldOp>( - op)) { + if (mlir::isa(op)) { return true; } 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-minnumf-lowering.mlir b/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-maxnumf-plain-rejected.mlir similarity index 70% rename from test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-minnumf-lowering.mlir rename to test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-maxnumf-plain-rejected.mlir index 0b532636..a283242e 100644 --- a/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-minnumf-lowering.mlir +++ b/test/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/linalg-maxnumf-plain-rejected.mlir @@ -1,15 +1,15 @@ -// RUN: dataflow-scheduler-opt -pass-pipeline="builtin.module(ktdflowering-to-dfir)" %s | FileCheck %s +// RUN: not dataflow-scheduler-opt -pass-pipeline="builtin.module(ktdflowering-to-dfir)" %s 2>&1 | FileCheck %s -// Tests that arith.minnumf lowers to the min binary operator. It was accepted by -// the frontend legality check and had no lowering, so a body using it got as far -// as here and then said only "unsupported operation type in linalg.generic body". +// 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. // -// There is no absolute-min operator, so minnumf has no second shape to select the -// way maxnumf picks abs_max for the abs-of-both case. +// absf(x) maxnumf absf(y) still lowers, to abs_max; linalg-maxnumf-lowering covers +// it. -// CHECK-LABEL: func.func @minnumf_plain -// CHECK: vectorchain.binary -// CHECK-SAME: binary_op = #vectorchain +// 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)> @@ -17,8 +17,8 @@ module { ktdf_arch.device @sample_device attributes {} import("../../../../Dialect/KTDFArch/sample_device.mlir") - // ── minnumf: lowers to min ───────────────────────────────────────────────── - func.func @minnumf_plain() attributes {grid = [2]} { + // ── 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 @@ -51,7 +51,7 @@ module { iterator_types = ["parallel", "reduction", "parallel"] } ins(%input : memref<1x1x64xf16>) outs(%alloc : memref<1x64xf16, "SFU_REG">) { ^bb0(%in: f16, %out: f16): - %result = arith.minnumf %in, %out : f16 + %result = arith.maxnumf %in, %out : f16 linalg.yield %result : f16 } } {loop_type = #ktdf.loop_type} From aaec52d04d8fb31a2829fc6c7b13de839ce22e33 Mon Sep 17 00:00:00 2001 From: Swagath Venkataramani Date: Thu, 10 Sep 2026 08:23:47 -0400 Subject: [PATCH 8/8] Follow the compute parameter getFlattenedVectorType now takes Rebasing onto main brings #144, which gives getFlattenedVectorType a ShapedType and an execution unit in place of the resource kinds, and threads that unit through lowerBinaryFOp. The compare and the selection take it the same way. The dialects submodule moves with it: the pin main records carries the header the ResourceKinds view moved to, which the commit before it does not have. Signed-off-by: Swagath Venkataramani --- .../KTDFLowToDFIR/LinalgLowering.cpp | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp index 12c878dc..6875d474 100644 --- a/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp +++ b/lib/Conversion/backend/ScheduleIRToDFIR/KTDFLowToDFIR/LinalgLowering.cpp @@ -186,10 +186,10 @@ struct LowerLinalgGenericPattern mlir::vectorchain::VectorChainBinaryOperator::min, compute); }) .Case([&](mlir::arith::CmpFOp op) { - return lowerCompareFOp(op, rewriter); + return lowerCompareFOp(op, rewriter, compute); }) .Case([&](mlir::arith::SelectOp op) { - return lowerSelectOp(op, rewriter); + return lowerSelectOp(op, rewriter, compute); }) .Case([&](mlir::arith::MaxNumFOp op) { mlir::Value lhs, rhs; @@ -343,10 +343,10 @@ struct LowerLinalgGenericPattern mlir::vectorchain::VectorChainBinaryOperator::min, compute); }) .Case([&](mlir::arith::CmpFOp op) { - return lowerCompareFOp(op, rewriter); + return lowerCompareFOp(op, rewriter, compute); }) .Case([&](mlir::arith::SelectOp op) { - return lowerSelectOp(op, rewriter); + return lowerSelectOp(op, rewriter, compute); }) .Case([&](mlir::memref::StoreOp op) { return lowerMemRefStore(op, rewriter); @@ -693,13 +693,15 @@ struct LowerLinalgGenericPattern /// 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) const { + 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(); - auto operands = - getFlattenedVectorType(op.getLhs().getType(), resource_kinds_); + 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( @@ -712,10 +714,13 @@ struct LowerLinalgGenericPattern /// 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) const { - auto result = - getFlattenedVectorType(op.getTrueValue().getType(), resource_kinds_); + 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(