Add composite relu

lib/Pipelines/ArithmeticPipelineRegistration.cpp

@@ -201,7 +201,7 @@ void mlirToSecretArithmeticPipelineBuilder(
   // Vectorize and optimize rotations
   // TODO(#2320): figure out where this fits in the new pipeline
   hecoSIMDVectorizerPipelineBuilder(pm, options.experimentalDisableLoopUnroll);
-  mathToPolynomialApproximationBuilder(pm);
+  mathToPolynomialApproximationBuilder(pm, options.useCompositeRelu);
 
   // Layout assignment and optimization
   LayoutPropagationOptions layoutPropagationOptions;
@@ -715,6 +715,7 @@ void torchLinalgToCkksBuilder(OpPassManager& manager,
   suboptions.enableArithmetization = true;
   suboptions.minSlotCount = options.minSlotCount;
   suboptions.greedyBootstrapWaterline = options.greedyBootstrapWaterline;
+  suboptions.useCompositeRelu = options.useCompositeRelu;
   suboptions.scalingModBits = options.scalingModBits;
   suboptions.firstModBits = options.firstModBits;
   suboptions.enableSplitPreprocessing = options.enableSplitPreprocessing;

lib/Pipelines/ArithmeticPipelineRegistration.h

@@ -77,6 +77,13 @@ struct MlirToRLWEPipelineOptions : public LoopOptions {
       *this, "bfv-mod-bits",
       llvm::cl::desc("The number of bits for all moduli for B/FV"),
       llvm::cl::init(60)};
+  PassOptions::Option<bool> useCompositeRelu{
+      *this, "use-composite-relu",
+      llvm::cl::desc("Approximate ReLU with the composite-sign method "
+                     "(x*step(x/B), 3 chained minimax Chebyshev polys) "
+                     "instead of a single-polynomial max(x,0) fit. More "
+                     "accurate for deep nets; needs more depth/bootstrapping."),
+      llvm::cl::init(false)};
   PassOptions::Option<bool> debug{
       *this, "debug",
       llvm::cl::desc("Insert debug ports after every secret operation."),

lib/Pipelines/PipelineRegistration.cpp

@@ -66,8 +66,11 @@ void oneShotBufferize(OpPassManager& manager, bool includeDeallocation) {
   manager.addPass(createCanonicalizerPass());
 }
 
-void mathToPolynomialApproximationBuilder(OpPassManager& pm) {
-  pm.addPass(createPolynomialApproximation());
+void mathToPolynomialApproximationBuilder(OpPassManager& pm,
+                                          bool useCompositeRelu) {
+  PolynomialApproximationOptions polyApproxOptions;
+  polyApproxOptions.useCompositeRelu = useCompositeRelu;
+  pm.addPass(createPolynomialApproximation(polyApproxOptions));
   pm.addPass(createLowerPolynomialEval());
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createCSEPass());

lib/Pipelines/PipelineRegistration.h

@@ -14,7 +14,8 @@ void prepareForBufferize(OpPassManager& manager);
 
 void oneShotBufferize(OpPassManager& manager, bool includeDeallocation = true);
 
-void mathToPolynomialApproximationBuilder(OpPassManager& pm);
+void mathToPolynomialApproximationBuilder(OpPassManager& pm,
+                                          bool useCompositeRelu = false);
 
 void polynomialToLLVMPipelineBuilder(OpPassManager& manager);
 

lib/Transforms/ActivationCanonicalizations/ActivationCanonicalizations.cpp

@@ -4,6 +4,8 @@
 
 #include "lib/Dialect/MathExt/IR/MathExtOps.h"
 #include "llvm/include/llvm/ADT/APFloat.h"               // from @llvm-project
+#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
+#include "mlir/include/mlir/Dialect/Linalg/IR/Linalg.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/Attributes.h"             // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
@@ -40,6 +42,93 @@ static bool IsOne(mlir::Attribute attr) {
 // populateWithGenerated, which can conflict with other generated patterns.
 #include "lib/Transforms/ActivationCanonicalizations/Rewrites.cpp.inc"
 
+// select(a > c, a, c) = max(a, c) for floats. This replaces the DRR
+// `SelectGreaterThanEqualFloat` pattern and folds in
+// two attr-forwarding behaviors so the polynomial-approximation domain survives
+// regardless of where torch-mlir attached it:
+//   (a) copy discardable attrs off the select itself (the old DRR behavior),
+//   (b) if the domain is still missing, copy `domain_lower`/`domain_upper` from
+//       an enclosing `linalg.generic` (a torch ReLU imports as a generic
+//       carrying those attrs, with `cmpf+select` in its body) The
+//       generic's copy is dropped afterwards by stripForwardedDomains().
+// Either way the domain lands on the `arith.maximumf` that
+// PolynomialApproximation / ReluViaCompositeSign read; without it they fall
+// back to [-1, 1].
+struct SelectGreaterThanEqualFloatPattern
+    : public OpRewritePattern<arith::SelectOp> {
+  using OpRewritePattern<arith::SelectOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(arith::SelectOp op,
+                                PatternRewriter& rewriter) const override {
+    auto cmpOp = op.getCondition().getDefiningOp<arith::CmpFOp>();
+    if (!cmpOp)
+      return rewriter.notifyMatchFailure(op, "condition is not arith.cmpf");
+
+    auto pred = cmpOp.getPredicate();
+    if (pred != arith::CmpFPredicate::UGT && pred != arith::CmpFPredicate::UGE)
+      return rewriter.notifyMatchFailure(op, "predicate is not ugt/uge");
+
+    // Must be the ReLU/max shape: select(a >? c, a, c).
+    if (cmpOp.getLhs() != op.getTrueValue() ||
+        cmpOp.getRhs() != op.getFalseValue())
+      return rewriter.notifyMatchFailure(op,
+                                         "operands are not select(a>c,a,c)");
+
+    auto maxOp =
+        arith::MaximumFOp::create(rewriter, op.getLoc(), op.getTrueValue(),
+                                  op.getFalseValue(), cmpOp.getFastmathAttr());
+
+    // (a) Forward any discardable attrs annotated on the select op itself onto
+    // the maximumf (the old DRR `SelectGreaterThanEqualFloat` behavior). Covers
+    // IR where the domain is attached directly to the select.
+    for (auto attr : op->getDiscardableAttrs())
+      maxOp->setAttr(attr.getName(), attr.getValue());
+
+    // (b) If the domain still isn't on the maximumf, it lives on an enclosing
+    // linalg.generic instead (where torch-mlir's importer attaches the ReLU
+    // domain). Copy it down onto the maximumf — the op PolynomialApproximation
+    // actually reads. We only COPY here; the generic's own bounds are dropped
+    // afterwards by stripForwardedDomains(). A single generic can hold several
+    // ReLUs, so stripping as soon as the first one is rewritten would starve
+    // the rest and silently leave them on the default [-1, 1] domain.
+    if (auto generic = dyn_cast<linalg::GenericOp>(op->getParentOp())) {
+      Attribute lo = generic->getAttr("domain_lower");
+      Attribute hi = generic->getAttr("domain_upper");
+      if (lo && !maxOp->hasAttr("domain_lower"))
+        maxOp->setAttr("domain_lower", lo);
+      if (hi && !maxOp->hasAttr("domain_upper"))
+        maxOp->setAttr("domain_upper", hi);
+    }
+
+    rewriter.replaceOp(op, maxOp.getResult());
+    return success();
+  }
+};
+
+// The domain bounds must end up on exactly one op. Once the patterns above have
+// copied an enclosing generic's bounds down onto every ReLU in its body, the
+// generic's own copy is redundant, so drop it: leaving the bounds on BOTH makes
+// a later activation-lifting pass merge two `domain_lower` entries into one
+// dictionary, tripping DictionaryAttr's uniqueness assertion.
+static void stripForwardedDomains(Operation* root) {
+  root->walk([](linalg::GenericOp generic) {
+    if (!generic->hasAttr("domain_lower") && !generic->hasAttr("domain_upper"))
+      return;
+    // Only strip if the bounds actually made it onto an op inside the body;
+    // otherwise there was no ReLU to forward them to and they are still the
+    // only record of the domain.
+    bool forwarded = false;
+    generic->getRegion(0).walk([&](Operation* inner) {
+      if (inner->hasAttr("domain_lower") || inner->hasAttr("domain_upper"))
+        forwarded = true;
+    });
+    if (forwarded) {
+      generic->removeAttr("domain_lower");
+      generic->removeAttr("domain_upper");
+    }
+  });
+}
+
 struct ActivationCanonicalizations
     : impl::ActivationCanonicalizationsBase<ActivationCanonicalizations> {
   using ActivationCanonicalizationsBase::ActivationCanonicalizationsBase;
@@ -48,8 +137,11 @@ struct ActivationCanonicalizations
     MLIRContext* context = &getContext();
     RewritePatternSet patterns(context);
     populateWithGenerated(patterns);
+    patterns.add<SelectGreaterThanEqualFloatPattern>(context);
 
     (void)walkAndApplyPatterns(getOperation(), std::move(patterns));
+
+    stripForwardedDomains(getOperation());
   }
 };
 

lib/Transforms/ActivationCanonicalizations/BUILD

@@ -31,6 +31,7 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LinalgDialect",
         "@llvm-project//mlir:MathDialect",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",

lib/Transforms/ActivationCanonicalizations/Rewrites.td

@@ -9,22 +9,8 @@ include "lib/Dialect/MathExt/IR/MathExtOps.td"
 
 def IsOne : Constraint<CPred<"IsOne($0)">>;
 
-def GetDefiningOp : NativeCodeCall<"$0.getDefiningOp()">;
-
-// These are attributes that are not inherent to the definition of the op, e.g.
-// domain_upper and domain_lower for activation functions.
-def CopyDiscardableAttrs : NativeCodeCallVoid<
-  "for (auto attr : $0->getDiscardableAttrs()) $1->setAttr(attr.getName(), attr.getValue())">;
-
-// select(a >= c, a, c) = max(a, c)
-def SelectGreaterThanEqualFloat : Pat<
-  (SelectOp:$src
-    (Arith_CmpFOp $pred, $a, $c, $flags), $a, $c),
-  (Arith_MaximumFOp:$dest $a, $c, $flags),
-  [(Constraint<
-            CPred<"$0.getValue() == arith::CmpFPredicate::UGT || "
-                  "$0.getValue() == arith::CmpFPredicate::UGE">> $pred)],
-  [(CopyDiscardableAttrs (GetDefiningOp $src), (GetDefiningOp $dest))]>;
+// select(a >= c, a, c) = max(a, c) for floats is handled by the C++
+// `SelectGreaterThanEqualFloatPattern` in ActivationCanonicalizations.cpp.
 
 // select(a >= c, a, c) = max(a, c)
 def SelectGreaterThanEqualUnsigned : Pat<

lib/Transforms/PolynomialApproximation/BUILD

@@ -15,6 +15,7 @@ cc_library(
         "@heir//lib/Analysis/SecretnessAnalysis",
         "@heir//lib/Dialect/MathExt/IR:Dialect",
         "@heir//lib/Dialect/Polynomial/IR:Dialect",
+        "@heir//lib/Utils",
         "@heir//lib/Utils/Approximation:CaratheodoryFejer",
         "@heir//lib/Utils/Polynomial",
         "@llvm-project//llvm:Support",