Reduce from an identity held in a register

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<mlir::vectorchain::VectorChainElementWiseCompareOperator> {
+  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<mlir::linalg::GenericOp> {
@@ -160,6 +185,12 @@ struct LowerLinalgGenericPattern
                     op, op.getLhs(), op.getRhs(), rewriter, identity_map,
                     mlir::vectorchain::VectorChainBinaryOperator::min, compute);
               })
+              .Case<mlir::arith::CmpFOp>([&](mlir::arith::CmpFOp op) {
+                return lowerCompareFOp(op, rewriter, compute);
+              })
+              .Case<mlir::arith::SelectOp>([&](mlir::arith::SelectOp op) {
+                return lowerSelectOp(op, rewriter, compute);
+              })
               .Case<mlir::arith::MaxNumFOp>([&](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>([&](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>([&](mlir::arith::CmpFOp op) {
+                return lowerCompareFOp(op, rewriter, compute);
+              })
+              .Case<mlir::arith::SelectOp>([&](mlir::arith::SelectOp op) {
+                return lowerSelectOp(op, rewriter, compute);
+              })
               .Case<mlir::memref::StoreOp>([&](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>([&](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<mlir::ShapedType>(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<mlir::ShapedType>(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:

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<mlir::ktdf::ReadFromFifoOp>())
-      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<mlir::ktdf::ReadFromFifoOp>()) {
+      const auto buffer = llvm::dyn_cast<mlir::MemRefType>(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();
   }

lib/Conversion/frontend/KTIRToScheduleIR/KTIRLegalityCheck.cpp

@@ -55,7 +55,7 @@ namespace {
   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::math::AbsFOp,
+                mlir::arith::CmpFOp, mlir::arith::SelectOp, mlir::math::AbsFOp,
                 mlir::linalg::YieldOp>(op)) {
     return true;
   }

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<MemRefType>(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();
   }

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<mlir::ModuleOp>().nest<mlir::func::FuncOp>();
+    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

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<mlir::PDLValue> values)
+    -> mlir::LogicalResult {
+  assert(values.size() == 2);
+  auto written = llvm::dyn_cast_if_present<mlir::DestinationStyleOpInterface>(
+      values[0].cast<mlir::Operation*>());
+  if (!written || written.getDpsInits().empty()) return mlir::failure();
+  auto params = llvm::dyn_cast_if_present<mlir::DictionaryAttr>(
+      values[1].cast<mlir::Attribute>());
+  if (!params) return mlir::failure();
+
+  auto accumulator =
+      llvm::dyn_cast<mlir::ShapedType>(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)

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<element_wise_compare_operator compare_lt>
+// 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<reduction_loop>}
+    }
+    return
+  }
+}