fix: lower chained secret Chebyshev evaluations natively

SecretnessAnalysis runs before walk-based rewriting. Replacing an earlier eval creates a new operand with no lattice entry, while each later eval result remains analyzed. Query the result so chained secret evaluations consistently use native lowering.

Author: AlexanderViand

lib/Transforms/LowerPolynomialEval/Patterns.cpp

@@ -230,8 +230,8 @@ LogicalResult LowerViaPatersonStockmeyerChebyshev::matchAndRewrite(
 
 LogicalResult LowerToKernelEvalChebyshev::matchAndRewrite(
     EvalOp op, PatternRewriter& rewriter) const {
-  if (!mlir::heir::isSecret(op.getValue(), &solver)) {
-    return rewriter.notifyMatchFailure(op, "operand is not secret");
+  if (!mlir::heir::isSecret(op.getResult(), &solver)) {
+    return rewriter.notifyMatchFailure(op, "result is not secret");
   }
   auto attr = dyn_cast<polynomial::TypedChebyshevPolynomialAttr>(
       op.getPolynomialAttr());

tests/Examples/lattigo/ckks/relu_composite/relu_composite_test.go

@@ -7,7 +7,7 @@ import (
 )
 
 func TestReluComposite(t *testing.T) {
-	evaluator, params, ecd, enc, dec := Relu_composite__configure()
+	bootstrapEvaluator, evaluator, params, ecd, enc, dec := Relu_composite__configure()
 
 	const n = 16
 	arg0 := make([]float32, n)
@@ -23,7 +23,7 @@ func TestReluComposite(t *testing.T) {
 	ct0 := Relu_composite__encrypt__arg0(evaluator, params, ecd, enc, arg0)
 
 	start := time.Now()
-	resultCt := Relu_composite(evaluator, params, ecd, ct0)
+	resultCt := Relu_composite(bootstrapEvaluator, evaluator, params, ecd, ct0)
 	t.Logf("composite-sign ReLU took %s", time.Since(start))
 
 	result := Relu_composite__decrypt__result0(evaluator, params, ecd, dec, resultCt)

tests/Transforms/lower_polynomial_eval/conditional_lower.mlir

@@ -18,6 +18,21 @@ module attributes {
     return %0 : !secret.secret<f64>
   }
 
+  // CHECK: @test_secret_kernel_chain
+  func.func @test_secret_kernel_chain(%x: !secret.secret<f64>) -> !secret.secret<f64> {
+    // CHECK: %[[EVAL_0:.*]] = kernel.eval_chebyshev
+    // CHECK: %[[EVAL_1:.*]] = kernel.eval_chebyshev %[[EVAL_0]]
+    // CHECK: kernel.eval_chebyshev %[[EVAL_1]]
+    %0 = secret.generic(%x : !secret.secret<f64>) {
+    ^body(%x_val: f64):
+      %1 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.0, 2.0]> : !polynomial.polynomial<ring=<coefficientType=f64>>, %x_val {domain_lower = -1.0 : f64, domain_upper = 1.0 : f64} : f64
+      %2 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.0, 2.0]> : !polynomial.polynomial<ring=<coefficientType=f64>>, %1 {domain_lower = -1.0 : f64, domain_upper = 1.0 : f64} : f64
+      %3 = polynomial.eval #polynomial.typed_chebyshev_polynomial<[1.0, 2.0]> : !polynomial.polynomial<ring=<coefficientType=f64>>, %2 {domain_lower = -1.0 : f64, domain_upper = 1.0 : f64} : f64
+      secret.yield %3 : f64
+    } -> (!secret.secret<f64>)
+    return %0 : !secret.secret<f64>
+  }
+
   // CHECK: @test_public_kernel
   func.func @test_public_kernel(%x: f64) -> f64 {
     // CHECK-NOT: kernel.eval_chebyshev