comb-to-cggi: refactor comb-to-cggi pass in prep for bool gate lowerings

PiperOrigin-RevId: 618461130
google · Mar 23, 2024 · 1125bfd · 1125bfd
1 parent a119175
commit 1125bfd
Show file tree

Hide file tree

Showing 7 changed files with 262 additions and 236 deletions.
diff --git a/lib/Conversion/CombToCGGI/CombToCGGI.cpp b/lib/Conversion/CombToCGGI/CombToCGGI.cpp
diff --git a/tests/comb_to_cggi/add_one.mlir b/tests/comb_to_cggi/add_one.mlir
@@ -1,4 +1,4 @@
-// RUN: heir-opt --comb-to-cggi -cse %s | FileCheck %s
+// RUN: heir-opt --secret-distribute-generic --comb-to-cggi -cse %s | FileCheck %s
 
 // This test was produced by running
 //   heir-opt --yosys-optimizer --canonicalize tests/yosys_optimizer/add_one.mlir

diff --git a/tests/comb_to_cggi/memref_ops.mlir b/tests/comb_to_cggi/memref_ops.mlir
@@ -1,4 +1,4 @@
-// RUN: heir-opt --comb-to-cggi %s | FileCheck %s
+// RUN: heir-opt --secret-distribute-generic --comb-to-cggi %s | FileCheck %s
 
 // CHECK: module
 module attributes {tf_saved_model.semantics} {

diff --git a/tests/comb_to_cggi/secret_cast.mlir b/tests/comb_to_cggi/secret_cast.mlir
@@ -1,6 +1,6 @@
 // This test ensures that secret casting before and after generics lowers to CGGI properly.
 
-// RUN: heir-opt --comb-to-cggi -cse %s | FileCheck %s
+// RUN: heir-opt --secret-distribute-generic --comb-to-cggi -cse %s | FileCheck %s
 
 // CHECK: module
 module attributes {tf_saved_model.semantics} {

diff --git a/tests/comb_to_cggi/secret_generic.mlir b/tests/comb_to_cggi/secret_generic.mlir
@@ -1,4 +1,4 @@
-// RUN: heir-opt --comb-to-cggi %s | FileCheck %s
+// RUN: heir-opt --secret-distribute-generic --comb-to-cggi %s | FileCheck %s
 
 module {
   // CHECK-NOT: secret

diff --git a/tests/comb_to_cggi/truth_table.mlir b/tests/comb_to_cggi/truth_table.mlir
@@ -1,4 +1,4 @@
-// RUN: heir-opt --split-input-file --comb-to-cggi --cse %s | FileCheck %s
+// RUN: heir-opt --secret-distribute-generic --split-input-file --comb-to-cggi --cse %s | FileCheck %s
 
 // CHECK-NOT: secret
 // CHECK: @truth_table_all_secret([[ARG:%.*]]: [[LWET:!lwe.lwe_ciphertext<.*>]]) -> [[LWET]]

diff --git a/tools/heir-opt.cpp b/tools/heir-opt.cpp
@@ -238,6 +238,7 @@ void tosaToBooleanTfhePipeline(const std::string &yosysFilesPath,
 
         // Lower combinational circuit to CGGI
         pm.addPass(mlir::createCSEPass());
+        pm.addPass(secret::createSecretDistributeGeneric());
         pm.addPass(comb::createCombToCGGI());
 
         // CGGI to Tfhe-Rust exit dialect