Update lower-polynomial-eval to lower to kernel.eval_chebyshev if the backend supports it

lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.cpp

@@ -234,7 +234,7 @@ void annotateSecretness(Operation* top, DataFlowSolver* solver, bool verbose) {
   });
 }
 
-bool isSecret(Value value, DataFlowSolver* solver) {
+bool isSecret(Value value, const DataFlowSolver* solver) {
   auto* lattice = solver->lookupState<SecretnessLattice>(value);
   return isSecret(lattice);
 }
@@ -249,7 +249,7 @@ bool isSecret(const SecretnessLattice* lattice) {
   return lattice->getValue().getSecretness();
 }
 
-bool isSecret(ValueRange values, DataFlowSolver* solver) {
+bool isSecret(ValueRange values, const DataFlowSolver* solver) {
   if (values.empty()) {
     return false;
   }

lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h

@@ -244,11 +244,11 @@ void annotateSecretness(Operation* top, DataFlowSolver* solver, bool verbose);
 
 // this method is used when DataFlowSolver has finished running the secretness
 // analysis
-bool isSecret(Value value, DataFlowSolver* solver);
+bool isSecret(Value value, const DataFlowSolver* solver);
 
 bool isSecret(const SecretnessLattice* lattice);
 
-bool isSecret(ValueRange values, DataFlowSolver* solver);
+bool isSecret(ValueRange values, const DataFlowSolver* solver);
 
 void getSecretOperands(Operation* op,
                        SmallVectorImpl<OpOperand*>& secretOperands,

lib/Dialect/Polynomial/IR/PolynomialOps.cpp

@@ -1079,7 +1079,9 @@ OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
       return nullptr;
     }
 
-    RNSPolynomial resultPoly = lhsPoly.mul(rhsPoly);
+    std::optional<RNSPolynomial> resultPolyOpt = lhsPoly.mul(rhsPoly);
+    if (!resultPolyOpt) return nullptr;
+    RNSPolynomial resultPoly = *resultPolyOpt;
 
     auto resultType = getResult().getType();
     auto elementType = lhsAttr.getCoefficients().getElementType();
@@ -1102,7 +1104,9 @@ OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
     auto lhs = getSingleLimbRNSPolynomial(lhsIntAttr, lhsPoly);
     auto rhs = getSingleLimbRNSPolynomial(rhsIntAttr, rhsPoly);
     if (lhs && rhs) {
-      RNSPolynomial result = lhs->mul(*rhs);
+      std::optional<RNSPolynomial> resultOpt = lhs->mul(*rhs);
+      if (!resultOpt) return nullptr;
+      RNSPolynomial result = *resultOpt;
       return getTypedIntPolynomialAttr(getContext(), result.getData(),
                                        getResult().getType());
     }
@@ -1126,7 +1130,9 @@ OpFoldResult NTTOp::fold(FoldAdaptor adaptor) {
     if (!rnsRootAttr) return nullptr;
 
     RNSPolynomial poly = inputAttr.getPolynomial();
-    RNSPolynomial resultPoly = poly.toNtt(rnsRootAttr);
+    std::optional<RNSPolynomial> resultPolyOpt = poly.toNtt(rnsRootAttr);
+    if (!resultPolyOpt) return nullptr;
+    RNSPolynomial resultPoly = *resultPolyOpt;
 
     auto resultType = getResult().getType();
     auto elementType = inputAttr.getCoefficients().getElementType();
@@ -1151,7 +1157,9 @@ OpFoldResult NTTOp::fold(FoldAdaptor adaptor) {
   if (!poly) return nullptr;
   SmallVector<uint64_t> roots = {
       modArithRootAttr.getValue().getValue().getZExtValue()};
-  RNSPolynomial resultPoly = poly->toNtt(roots);
+  std::optional<RNSPolynomial> resultPolyOpt = poly->toNtt(roots);
+  if (!resultPolyOpt) return nullptr;
+  RNSPolynomial resultPoly = *resultPolyOpt;
   return getTypedIntPolynomialAttr(getContext(), resultPoly.getData(),
                                    getResult().getType());
 }
@@ -1166,7 +1174,10 @@ OpFoldResult INTTOp::fold(FoldAdaptor adaptor) {
     if (!rnsRootAttr) return nullptr;
 
     RNSPolynomial poly = inputAttr.getPolynomial();
-    RNSPolynomial resultPoly = poly.toCoefficient(rnsRootAttr);
+    std::optional<RNSPolynomial> resultPolyOpt =
+        poly.toCoefficient(rnsRootAttr);
+    if (!resultPolyOpt) return nullptr;
+    RNSPolynomial resultPoly = *resultPolyOpt;
 
     auto resultType = getResult().getType();
     auto elementType = inputAttr.getCoefficients().getElementType();
@@ -1191,7 +1202,9 @@ OpFoldResult INTTOp::fold(FoldAdaptor adaptor) {
   if (!poly) return nullptr;
   SmallVector<uint64_t> roots = {
       modArithRootAttr.getValue().getValue().getZExtValue()};
-  RNSPolynomial resultPoly = poly->toCoefficient(roots);
+  std::optional<RNSPolynomial> resultPolyOpt = poly->toCoefficient(roots);
+  if (!resultPolyOpt) return nullptr;
+  RNSPolynomial resultPoly = *resultPolyOpt;
   return getTypedIntPolynomialAttr(getContext(), resultPoly.getData(),
                                    getResult().getType());
 }

lib/Target/Lattigo/BUILD

@@ -31,6 +31,7 @@ cc_library(
         "@llvm-project//mlir:DialectUtils",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MathDialect",
         "@llvm-project//mlir:MemRefDialect",
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:Support",

lib/Target/Lattigo/LattigoEmitter.cpp

@@ -148,6 +148,7 @@ LogicalResult LattigoEmitter::translate(Operation& op) {
               [&](auto op) { return printOperation(op); })
           .Case<preprocessing::LoadResourceOp>(
               [&](auto op) { return printOperation(op); })
+          .Case<math::SqrtOp>([&](auto op) { return printOperation(op); })
 
           // Lattigo ops
           .Case<
@@ -1059,6 +1060,28 @@ LogicalResult LattigoEmitter::printOperation(arith::XOrIOp op) {
   return printBinaryOp(op, op.getLhs(), op.getRhs(), "^");
 }
 
+LogicalResult LattigoEmitter::printOperation(math::SqrtOp op) {
+  imports.insert("\"math\"");
+  Type type = op.getOperand().getType();
+  auto typeStringResult = convertType(type);
+  if (failed(typeStringResult)) return failure();
+  std::string typeString = typeStringResult.value();
+
+  std::string operandName = getName(op.getOperand());
+  std::string resultName = getName(op.getResult());
+
+  if (typeString == "float32") {
+    os << resultName << " := float32(math.Sqrt(float64(" << operandName
+       << ")))\n";
+  } else if (typeString == "float64") {
+    os << resultName << " := math.Sqrt(" << operandName << ")\n";
+  } else {
+    return op.emitOpError("Unsupported float type for math.sqrt: ")
+           << typeString;
+  }
+  return success();
+}
+
 LogicalResult LattigoEmitter::printOperation(arith::RemSIOp op) {
   return printBinaryOp(op, op.getLhs(), op.getRhs(), "%");
 }
@@ -2604,7 +2627,7 @@ void registerToLattigoTranslation() {
                     func::FuncDialect, tensor::TensorDialect,
                     tensor_ext::TensorExtDialect, lattigo::LattigoDialect,
                     memref::MemRefDialect, mgmt::MgmtDialect, scf::SCFDialect,
-                    preprocessing::PreprocessingDialect>();
+                    preprocessing::PreprocessingDialect, math::MathDialect>();
       });
 }
 
@@ -2625,7 +2648,7 @@ void registerToLattigoPreprocessingTranslation() {
                     func::FuncDialect, tensor::TensorDialect,
                     tensor_ext::TensorExtDialect, lattigo::LattigoDialect,
                     memref::MemRefDialect, mgmt::MgmtDialect, scf::SCFDialect,
-                    preprocessing::PreprocessingDialect>();
+                    preprocessing::PreprocessingDialect, math::MathDialect>();
       });
 }
 
@@ -2646,7 +2669,7 @@ void registerToLattigoPreprocessedTranslation() {
                     func::FuncDialect, tensor::TensorDialect,
                     tensor_ext::TensorExtDialect, lattigo::LattigoDialect,
                     memref::MemRefDialect, mgmt::MgmtDialect, scf::SCFDialect,
-                    preprocessing::PreprocessingDialect>();
+                    preprocessing::PreprocessingDialect, math::MathDialect>();
       });
 }
 

lib/Target/Lattigo/LattigoEmitter.h

@@ -18,6 +18,7 @@
 #include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"   // from @llvm-project
+#include "mlir/include/mlir/Dialect/Math/IR/Math.h"      // from @llvm-project
 #include "mlir/include/mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/SCF/IR/SCF.h"        // from @llvm-project
 #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
@@ -136,6 +137,7 @@ class LattigoEmitter {
   LogicalResult printOperation(::mlir::arith::SubIOp op);
   LogicalResult printOperation(::mlir::arith::SubFOp op);
   LogicalResult printOperation(::mlir::arith::XOrIOp op);
+  LogicalResult printOperation(::mlir::math::SqrtOp op);
   LogicalResult printOperation(::mlir::func::CallOp op);
   LogicalResult printOperation(::mlir::func::FuncOp op);
   LogicalResult printOperation(::mlir::func::ReturnOp op);

lib/Transforms/LowerPolynomialEval/BUILD

@@ -13,7 +13,11 @@ cc_library(
     deps = [
         ":Patterns",
         ":pass_inc_gen",
+        "@heir//lib/Analysis/SecretnessAnalysis",
+        "@heir//lib/Dialect/Kernel/IR:Dialect",
         "@heir//lib/Dialect/Polynomial/IR:Dialect",
+        "@heir//lib/Target/CompilationTarget",
+        "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:TransformUtils",
@@ -26,6 +30,8 @@ cc_library(
     srcs = ["Patterns.cpp"],
     hdrs = ["Patterns.h"],
     deps = [
+        "@heir//lib/Analysis/SecretnessAnalysis",
+        "@heir//lib/Dialect/Kernel/IR:Dialect",
         "@heir//lib/Dialect/Polynomial/IR:Dialect",
         "@heir//lib/Kernel:AbstractValue",
         "@heir//lib/Kernel:ArithmeticDag",

lib/Transforms/LowerPolynomialEval/LowerPolynomialEval.cpp

@@ -2,9 +2,13 @@
 
 #include <utility>
 
+#include "lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h"
+#include "lib/Dialect/Kernel/IR/KernelDialect.h"
+#include "lib/Target/CompilationTarget/CompilationTarget.h"
 #include "lib/Transforms/LowerPolynomialEval/Patterns.h"
-#include "mlir/include/mlir/IR/MLIRContext.h"   // from @llvm-project
-#include "mlir/include/mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/include/mlir/Analysis/DataFlow/Utils.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/MLIRContext.h"           // from @llvm-project
+#include "mlir/include/mlir/IR/PatternMatch.h"          // from @llvm-project
 #include "mlir/include/mlir/Transforms/WalkPatternRewriteDriver.h"  // from @llvm-project
 
 // IWYU pragma: begin_keep
@@ -17,21 +21,62 @@ namespace heir {
 #define GEN_PASS_DEF_LOWERPOLYNOMIALEVAL
 #include "lib/Transforms/LowerPolynomialEval/LowerPolynomialEval.h.inc"
 
+static bool hasBackendAttribute(ModuleOp module) {
+  if (!module) return false;
+  for (NamedAttribute attr : module->getAttrs()) {
+    if (!isa<UnitAttr>(attr.getValue())) continue;
+    if (attr.getName().strref().starts_with("backend.")) {
+      return true;
+    }
+  }
+  return false;
+}
+
 struct LowerPolynomialEval
     : impl::LowerPolynomialEvalBase<LowerPolynomialEval> {
   using LowerPolynomialEvalBase::LowerPolynomialEvalBase;
 
   void runOnOperation() override {
     MLIRContext* context = &getContext();
+
+    ModuleOp module = dyn_cast<ModuleOp>(getOperation());
+    if (!module) {
+      module = getOperation()->getParentOfType<ModuleOp>();
+    }
+
+    bool hasKernelChebyshev = false;
+    if (module && hasBackendAttribute(module)) {
+      auto target = getTargetConfig(module);
+      if (succeeded(target)) {
+        hasKernelChebyshev = target->has_kernel_chebyshev;
+      }
+    }
+
     RewritePatternSet patterns(context);
 
+    DataFlowSolver solver;
+    dataflow::loadBaselineAnalyses(solver);
+    solver.load<SecretnessAnalysis>();
+    if (failed(solver.initializeAndRun(getOperation()))) {
+      getOperation()->emitOpError() << "Failed to run SecretnessAnalysis.\n";
+      return signalPassFailure();
+    }
+
     switch (method) {
       case PolynomialApproximationMethod::Automatic:
         patterns.add<LowerViaHorner, LowerViaPatersonStockmeyerMonomial>(
             context, /*force=*/false);
-        patterns.add<LowerViaPatersonStockmeyerChebyshev>(
-            context,
-            /*force=*/false, minCoefficientThreshold);
+        if (hasKernelChebyshev) {
+          patterns.add<LowerToKernelEvalChebyshev>(context, solver,
+                                                   /*force=*/false);
+          patterns.add<LowerViaPatersonStockmeyerChebyshev>(
+              context,
+              /*force=*/false, minCoefficientThreshold);
+        } else {
+          patterns.add<LowerViaPatersonStockmeyerChebyshev>(
+              context,
+              /*force=*/false, minCoefficientThreshold);
+        }
         break;
       case PolynomialApproximationMethod::Horner:
         patterns.add<LowerViaHorner>(context, /*force=*/true);
@@ -41,9 +86,17 @@ struct LowerPolynomialEval
                                                          /*force=*/true);
         break;
       case PolynomialApproximationMethod::PatersonStockmeyerChebyshev:
-        patterns.add<LowerViaPatersonStockmeyerChebyshev>(
-            context,
-            /*force=*/true, minCoefficientThreshold);
+        if (hasKernelChebyshev) {
+          patterns.add<LowerToKernelEvalChebyshev>(context, solver,
+                                                   /*force=*/true);
+          patterns.add<LowerViaPatersonStockmeyerChebyshev>(
+              context,
+              /*force=*/true, minCoefficientThreshold);
+        } else {
+          patterns.add<LowerViaPatersonStockmeyerChebyshev>(
+              context,
+              /*force=*/true, minCoefficientThreshold);
+        }
         break;
       default:
         getOperation()->emitError() << "Unknown lowering method: " << method;

lib/Transforms/LowerPolynomialEval/LowerPolynomialEval.td

@@ -29,6 +29,7 @@ def LowerPolynomialEval : Pass<"lower-polynomial-eval"> {
   }];
   let dependentDialects = [
     "::mlir::heir::polynomial::PolynomialDialect",
+    "::mlir::heir::kernel::KernelDialect",
   ];
   let options = [
     Option<"method", "method", "mlir::heir::PolynomialApproximationMethod",

lib/Transforms/LowerPolynomialEval/Patterns.cpp

@@ -4,6 +4,8 @@
 #include <cstdint>
 #include <map>
 
+#include "lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h"
+#include "lib/Dialect/Kernel/IR/KernelOps.h"
 #include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
 #include "lib/Dialect/Polynomial/IR/PolynomialOps.h"
 #include "lib/Kernel/AbstractValue.h"
@@ -226,5 +228,51 @@ LogicalResult LowerViaPatersonStockmeyerChebyshev::matchAndRewrite(
   return success();
 }
 
+LogicalResult LowerToKernelEvalChebyshev::matchAndRewrite(
+    EvalOp op, PatternRewriter& rewriter) const {
+  if (!mlir::heir::isSecret(op.getValue(), &solver)) {
+    return rewriter.notifyMatchFailure(op, "operand is not secret");
+  }
+  auto attr = dyn_cast<polynomial::TypedChebyshevPolynomialAttr>(
+      op.getPolynomialAttr());
+  if (!attr) return failure();
+
+  auto lowerAttr = op->getAttrOfType<FloatAttr>("domain_lower");
+  auto upperAttr = op->getAttrOfType<FloatAttr>("domain_upper");
+  if (!lowerAttr || !upperAttr) return failure();
+
+  double lower = lowerAttr.getValue().convertToDouble();
+  double upper = upperAttr.getValue().convertToDouble();
+
+  ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+  Value xInput = op.getValue();
+
+  if (std::abs(lower - -1.0) > 1e-9 || std::abs(upper - 1.0) > 1e-9) {
+    APFloat rescale = APFloat(2.0 / (upper - lower));
+    APFloat shift = APFloat(-(upper + lower) / (upper - lower));
+
+    Type inputType = xInput.getType();
+
+    if (!rescale.isExactlyValue(1.0)) {
+      xInput = arith::MulFOp::create(
+                   b, xInput,
+                   arith::ConstantOp::create(
+                       b, inputType, getScalarOrDenseAttr(inputType, rescale)))
+                   .getResult();
+    }
+    if (!shift.isZero()) {
+      xInput = arith::AddFOp::create(
+                   b, xInput,
+                   arith::ConstantOp::create(
+                       b, inputType, getScalarOrDenseAttr(inputType, shift)))
+                   .getResult();
+    }
+  }
+
+  rewriter.replaceOpWithNewOp<kernel::EvalChebyshevOp>(
+      op, op.getType(), xInput, attr.getValue().getCoefficients());
+  return success();
+}
+
 }  // namespace heir
 }  // namespace mlir