Repo sync (#651)

secretflow · Apr 15, 2024 · ad1461f · ad1461f
1 parent 8bf3c97
commit ad1461f
Show file tree

Hide file tree

Showing 18 changed files with 1,374 additions and 1 deletion.
diff --git a/libspu/compiler/core/core.cc b/libspu/compiler/core/core.cc
@@ -63,6 +63,8 @@ void Core::buildPipeline(mlir::PassManager *pm) {
     optPM.addPass(mlir::spu::pphlo::createOptimizeSqrtPlusEps());
   }
 
+  optPM.addPass(mlir::spu::pphlo::createExpandSecretGatherPass());
+
   if (!options.disable_div_sqrt_rewrite()) {
     optPM.addPass(mlir::spu::pphlo::createRewriteDivSqrtPatterns());
   }

diff --git a/libspu/compiler/front_end/hlo_importer.cc b/libspu/compiler/front_end/hlo_importer.cc
@@ -127,7 +127,7 @@ void runHloPasses(xla::HloModule *module) {
         /*allow_mixed_precision=*/false);
 
     pipeline.AddPass<GatherSimplifier>();
-    pipeline.AddPass<GatherExpander>(GatherExpander::kEliminateAllGathers);
+    pipeline.AddPass<GatherExpander>(GatherExpander::kEliminateSimpleGathers);
     pipeline.AddPass<ScatterExpander>(ScatterExpander::kEliminateAllScatters);
     pipeline.AddPass<AlgebraicSimplifier>(options);
     pipeline.AddPass<BitcastDtypesExpander>();

diff --git a/libspu/compiler/passes/BUILD.bazel b/libspu/compiler/passes/BUILD.bazel
@@ -85,6 +85,18 @@ spu_cc_library(
     ],
 )
 
+spu_cc_library(
+    name = "expand_secret_gather",
+    srcs = ["expand_secret_gather.cc"],
+    hdrs = ["passes.h"],
+    deps = [
+        ":pass_details",
+        "//libspu/dialect/pphlo:dialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
 spu_cc_library(
     name = "hlo_legalize_to_pphlo",
     srcs = ["hlo_legalize_to_pphlo.cc"],
@@ -276,6 +288,7 @@ spu_cc_library(
         ":convert_push_down",
         ":decompose_comparison",
         ":decompose_minmax",
+        ":expand_secret_gather",
         ":hlo_legalize_to_pphlo",
         ":insert_deallocation",
         ":lower_conversion_cast",