Merge pull request #2137 from j2kun:poly-approx-subpipeilne

PiperOrigin-RevId: 799153545

Author: copybara-github

lib/Pipelines/ArithmeticPipelineRegistration.cpp

@@ -35,10 +35,8 @@
 #include "lib/Transforms/LayoutOptimization/LayoutOptimization.h"
 #include "lib/Transforms/LayoutPropagation/LayoutPropagation.h"
 #include "lib/Transforms/LinalgCanonicalizations/LinalgCanonicalizations.h"
-#include "lib/Transforms/LowerPolynomialEval/LowerPolynomialEval.h"
 #include "lib/Transforms/OperationBalancer/OperationBalancer.h"
 #include "lib/Transforms/OptimizeRelinearization/OptimizeRelinearization.h"
-#include "lib/Transforms/PolynomialApproximation/PolynomialApproximation.h"
 #include "lib/Transforms/PopulateScale/PopulateScale.h"
 #include "lib/Transforms/PropagateAnnotation/PropagateAnnotation.h"
 #include "lib/Transforms/SecretInsertMgmt/Passes.h"
@@ -127,10 +125,6 @@ void mlirToSecretArithmeticPipelineBuilder(
   convertToDataObliviousPipelineBuilder(pm);
   pm.addPass(createSelectRewrite());
   pm.addPass(createCompareToSignRewrite());
-  pm.addPass(createPolynomialApproximation());
-  pm.addPass(createLowerPolynomialEval());
-  pm.addPass(createCanonicalizerPass());
-  pm.addPass(createCSEPass());
 
   // Simplify linalg ops for kernel lowering
   // Linalg canonicalization
@@ -154,6 +148,8 @@ void mlirToSecretArithmeticPipelineBuilder(
   pm.addPass(
       createConvertToCiphertextSemantics(convertToCiphertextSemanticsOptions));
 
+  mathToPolynomialApproximationBuilder(pm);
+
   // Balance Operations
   pm.addPass(createOperationBalancer());
 

lib/Pipelines/BUILD

@@ -18,7 +18,6 @@ cc_library(
         "@heir//lib/Dialect/Polynomial/Conversions/PolynomialToModArith",
         "@heir//lib/Dialect/Secret/Conversions/SecretToCGGI",
         "@heir//lib/Dialect/TOSA/Conversions/TosaToSecretArith",
-        "@heir//lib/Transforms/CompareToSignRewrite",
         "@heir//lib/Transforms/ConvertIfToSelect",
         "@heir//lib/Transforms/ConvertSecretExtractToStaticExtract",
         "@heir//lib/Transforms/ConvertSecretForToStaticFor",
@@ -29,6 +28,7 @@ cc_library(
         "@heir//lib/Transforms/MemrefToArith:ExpandCopy",
         "@heir//lib/Transforms/MemrefToArith:MemrefToArithRegistration",
         "@heir//lib/Transforms/PolynomialApproximation",
+        "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineToStandard",
         "@llvm-project//mlir:AffineTransforms",
         "@llvm-project//mlir:ArithTransforms",
@@ -131,11 +131,9 @@ cc_library(
         "@heir//lib/Transforms/LayoutOptimization",
         "@heir//lib/Transforms/LayoutPropagation",
         "@heir//lib/Transforms/LinalgCanonicalizations",
-        "@heir//lib/Transforms/LowerPolynomialEval",
         "@heir//lib/Transforms/MemrefToArith:MemrefToArithRegistration",
         "@heir//lib/Transforms/OperationBalancer",
         "@heir//lib/Transforms/OptimizeRelinearization",
-        "@heir//lib/Transforms/PolynomialApproximation",
         "@heir//lib/Transforms/PopulateScale",
         "@heir//lib/Transforms/PropagateAnnotation",
         "@heir//lib/Transforms/SecretInsertMgmt",

lib/Pipelines/PipelineRegistration.cpp

@@ -8,7 +8,9 @@
 #include "lib/Transforms/ConvertSecretInsertToStaticInsert/ConvertSecretInsertToStaticInsert.h"
 #include "lib/Transforms/ConvertSecretWhileToStaticFor/ConvertSecretWhileToStaticFor.h"
 #include "lib/Transforms/ElementwiseToAffine/ElementwiseToAffine.h"
+#include "lib/Transforms/LowerPolynomialEval/LowerPolynomialEval.h"
 #include "lib/Transforms/MemrefToArith/MemrefToArith.h"
+#include "lib/Transforms/PolynomialApproximation/PolynomialApproximation.h"
 #include "mlir/include/mlir/Conversion/AffineToStandard/AffineToStandard.h"  // from @llvm-project
 #include "mlir/include/mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h"  // from @llvm-project
 #include "mlir/include/mlir/Conversion/ConvertToLLVM/ToLLVMPass.h"  // from @llvm-project
@@ -91,6 +93,13 @@ void tosaPipelineBuilder(OpPassManager& manager, bool unroll) {
   manager.addPass(createSymbolDCEPass());
 }
 
+void mathToPolynomialApproximationBuilder(OpPassManager& pm) {
+  pm.addPass(createPolynomialApproximation());
+  pm.addPass(createLowerPolynomialEval());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createCSEPass());
+}
+
 void polynomialToLLVMPipelineBuilder(OpPassManager& manager) {
   // Poly
   manager.addPass(createElementwiseToAffine());

lib/Pipelines/PipelineRegistration.h

@@ -1,16 +1,19 @@
 #ifndef LIB_PIPELINES_PIPELINEREGISTRATION_H_
 #define LIB_PIPELINES_PIPELINEREGISTRATION_H_
 
-#include "mlir/include/mlir/Pass/PassManager.h"   // from @llvm-project
-#include "mlir/include/mlir/Pass/PassOptions.h"   // from @llvm-project
-#include "mlir/include/mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "llvm/include/llvm/Support/CommandLine.h"  // from @llvm-project
+#include "mlir/include/mlir/Pass/PassManager.h"     // from @llvm-project
+#include "mlir/include/mlir/Pass/PassOptions.h"     // from @llvm-project
+#include "mlir/include/mlir/Pass/PassRegistry.h"    // from @llvm-project
 
 namespace mlir::heir {
 
 void tosaToLinalg(OpPassManager& manager);
 
 void oneShotBufferize(OpPassManager& manager);
 
+void mathToPolynomialApproximationBuilder(OpPassManager& pm);
+
 struct TosaToArithOptions : public PassPipelineOptions<TosaToArithOptions> {
   PassOptions::Option<bool> unroll{*this, "full-unroll",
                                    llvm::cl::desc("Full unroll all loops."),

tests/Pipelines/math_to_polynomial_approximation/BUILD

@@ -0,0 +1,10 @@
+load("//bazel:lit.bzl", "glob_lit_tests")
+
+package(default_applicable_licenses = ["@heir//:license"])
+
+glob_lit_tests(
+    name = "all_tests",
+    data = ["@heir//tests:test_utilities"],
+    driver = "@heir//tests:run_lit.sh",
+    test_file_exts = ["mlir"],
+)

tests/Pipelines/math_to_polynomial_approximation/polynomial_approximation.mlir

@@ -0,0 +1,26 @@
+// RUN: heir-opt --math-to-polynomial-approximation %s | FileCheck %s
+
+// CHECK: @test_maximumf
+func.func @test_maximumf(%x: tensor<10xf32>) -> tensor<10xf32> {
+  // CHECK-NOT: arith.maximumf
+  // CHECK-NOT: polynomial.eval
+
+  // CHECK-DAG: arith.constant dense<0.03831
+  // CHECK-DAG: arith.constant dense<5.0{{0*}}e-01> : tensor<10xf32>
+  // CHECK-DAG: arith.constant dense<0.9370
+  // CHECK-DAG: arith.constant dense<-0.5062
+  // CHECK: arith.mulf
+  // CHECK: arith.addf
+  // CHECK: arith.mulf
+  // CHECK: arith.addf
+  // CHECK: arith.mulf
+  // CHECK: arith.addf
+  // CHECK: arith.mulf
+  // CHECK: arith.addf
+  // CHECK: arith.mulf
+  // CHECK: arith.addf
+  // CHECK: return
+  %c0 = arith.constant dense<0.0> : tensor<10xf32>
+  %0 = arith.maximumf %x, %c0 : tensor<10xf32>
+  return %0 : tensor<10xf32>
+}

tools/heir-opt.cpp

@@ -457,6 +457,12 @@ int main(int argc, char** argv) {
       "Transforms a native program to data-oblivious program",
       convertToDataObliviousPipelineBuilder);
 
+  PassPipelineRegistration<>(
+      "math-to-polynomial-approximation",
+      "Approximate math operations that cannot be expressed in FHE using "
+      "polynomial approximations.",
+      mathToPolynomialApproximationBuilder);
+
   return asMainReturnCode(
       MlirOptMain(argc, argv, "HEIR Pass Driver", registry));
 }