Repo sync (secretflow#743)

.licenserc.yaml

@@ -40,6 +40,7 @@ header: # <1>
     - '**/*.mlir'
     - '**/*.csv'
     - '**/*.tmpl'
+    - 'libspu/compiler/tests/interpret/template/**.template'
     - 'LICENSE'
     - 'NOTICE'
     - '.bazelversion'

.vscode/settings.json

@@ -41,5 +41,5 @@
     "[python]": {
         "editor.defaultFormatter": "ms-python.black-formatter"
     },
-    "mlir.server_path": "bazel-bin/libspu/compiler/tools/pphlo-lsp"
+    "mlir.server_path": "bazel-bin/libspu/compiler/tools/spu-lsp"
 }

CHANGELOG.md

@@ -12,6 +12,12 @@
 
 ## TBD
 
+- [Feature] Support jax.numpy.bitwise_count
+- [Bugfix] Fix jax.numpy.signbit wrong answer with very large input
+
+## 20240621
+
+- [SPU] 0.9.1b0 release
 - [Feature] Add ORAM based dynamic_slice for ABY3
 - [Feature] Add Atan2Op support
 - [API] Add beaver cache support for semi2k (**experimental**)

bazel/repositories.bzl

@@ -136,8 +136,8 @@ def _bazel_skylib():
     )
 
 def _com_github_openxla_xla():
-    OPENXLA_COMMIT = "d9d0e780ff6a37c4d501c8e0e4f4a9fdca30cbd4"
-    OPENXLA_SHA256 = "77ef83491f409afbe549a2bd695d710a70fdf7f04db35eeb1fba3e97ef767113"
+    OPENXLA_COMMIT = "9b0dd58c9b625a2e958f4fc7787a1ff5c95dbb40"
+    OPENXLA_SHA256 = "f150c5b49e4d4497aae2c79232f1efe2baccaa72223b21dc8715be73eab74417"
 
     # We need openxla to handle xla/mhlo/stablehlo
     maybe(

docs/development/pipeline.rst

@@ -39,17 +39,17 @@ The horizontal part depicts the data pipeline, from left to right.
 1. Data providers use :ref:`SPU io <reference/py_api:Runtime IO>` module to encrypt input data.
 
    * For SPU MPC backend, *encrypt* means to split plaintext data into shares.
-   * For floating point data, encoding to fixed-point may be also required.
+   * For floating-point data, encoding to fixed-point may be also required.
 
-2. The encrypted data is send to :ref:`SPU runtime <reference/py_api:Runtime Setup>`.
-3. The output data is fetched *result owner*, and decrypted by the :ref:`SPU io <reference/py_api:Runtime IO>` module.
+2. The encrypted data is sent to :ref:`SPU runtime <reference/py_api:Runtime Setup>`.
+3. The output data is fetched by *result owner*, and decrypted by the :ref:`SPU io <reference/py_api:Runtime IO>` module.
 
 
 Just in time
 ------------
 
-Jit is short for `Just-in-time compilation <https://en.wikipedia.org/wiki/Just-in-time_compilation>`_, with this approach, the compiler can get more information, such as input shapes, than in `AOT mode <https://en.wikipedia.org/wiki/Ahead-of-time_compilation>`_. Jit may introduce more evaluation overhead, but it's really trivial in secure computation setting.
+JIT is short for `Just-in-time compilation <https://en.wikipedia.org/wiki/Just-in-time_compilation>`_, with this approach, the compiler can get more information, such as input shapes, than in `AOT mode <https://en.wikipedia.org/wiki/Ahead-of-time_compilation>`_. JIT may introduce more evaluation overhead, but it's really trivial in secure computation setting.
 
-In SPU, jit has more benefits since the backend engine may be orders of magnitude faster if it knows the *visibility* of data. For example, when multiplying two secrets, the backend MPC engine may involve expensive *beaver triple* progress, but when one of the inputs (of multiply) is public known to all parties, the operation will be much faster. So we should *mark* as much data as possible to be *public* (if it doesn't need to be protected), and tell the compiler these information.
+In SPU, JIT has more benefits since the backend engine may be orders of magnitude faster if it knows the *visibility* of data. For example, when multiplying two secrets, the backend MPC engine may involve expensive *beaver triple* progress, but when one of the inputs (of multiply) is public known to all parties, the operation will be much faster. So we should *mark* as much data as possible to be *public* (if it doesn't need to be protected), and tell the compiler these information.
 
 So, SPU compilation normally happens after all data infeed is done, and `just in time` before the real evaluation.

docs/reference/pphlo_doc.rst

@@ -3,7 +3,7 @@ PPHlo API reference
 
 PPHlo is short for (SPU High level ops), it's the assembly language of SPU.
 
-PPHlo is built on `MLIR <https://mlir.llvm.org/>`_ infrastructure, the concrete ops definition could be found :spu_code_host:`here <spu/blob/main/libspu/dialect/pphlo/ops.td>`.
+PPHlo is built on `MLIR <https://mlir.llvm.org/>`_ infrastructure, the concrete ops definition could be found :spu_code_host:`here <spu/blob/main/libspu/dialect/pphlo/IR/ops.td>`.
 
 Op List
 ~~~~~~~

docs/reference/update_pphlo_doc.sh

@@ -18,7 +18,7 @@
 SCRIPT=`realpath $0`
 SCRIPTPATH=`dirname $SCRIPT`
 
-bazel build //libspu/dialect/pphlo:op_doc
+bazel build //libspu/dialect/pphlo/IR:op_doc
 
-cp `bazel info workspace`/bazel-bin/libspu/dialect/pphlo/op_doc.md $SCRIPTPATH/pphlo_op_doc.md
+cp `bazel info workspace`/bazel-bin/libspu/dialect/pphlo/IR/op_doc.md $SCRIPTPATH/pphlo_op_doc.md
 

docs/tutorials/quick_start.ipynb

@@ -645,7 +645,7 @@
     }
    ],
    "source": [
-    "print(sigmoid.dump_pphlo(np.random.rand(3, 3)))"
+    "print(sigmoid.dump_ir(np.random.rand(3, 3)))"
    ]
   },
   {
@@ -698,7 +698,7 @@
    "source": [
     "X = ppd.device(\"P1\")(make_rand)()\n",
     "\n",
-    "print(sigmoid.dump_pphlo(X))"
+    "print(sigmoid.dump_ir(X))"
    ]
   },
   {

examples/cpp/simple_pphlo.cc

@@ -38,7 +38,7 @@ func.func @main() -> () {
     %0 = pphlo.constant dense<1> : tensor<i32>
     %1 = pphlo.constant dense<2> : tensor<i32>
     %2 = pphlo.add %0, %1 : tensor<i32>
-    pphlo.custom_call @dbg_print (%2) {has_side_effect = true} : (tensor<i32>)->()
+    pphlo.custom_call @spu.dbg_print (%2) {has_side_effect = true} : (tensor<i32>)->()
     return
 })";
 
@@ -76,7 +76,7 @@ void parameters(spu::SPUContext* sctx) {
   constexpr auto code = R"PPHlo(
 func.func @main(%arg0: tensor<!pphlo.secret<f32>>, %arg1: tensor<!pphlo.secret<i32>>) -> () {
   %0 = pphlo.multiply %arg0, %arg1 : (tensor<!pphlo.secret<f32>>, tensor<!pphlo.secret<i32>>) -> tensor<!pphlo.secret<f32>>
-  pphlo.custom_call @dbg_print (%0) {has_side_effect = true} : (tensor<!pphlo.secret<f32>>)->()
+  pphlo.custom_call @spu.dbg_print (%0) {has_side_effect = true} : (tensor<!pphlo.secret<f32>>)->()
   return
 })PPHlo";
 

libspu/compiler/codegen/BUILD.bazel

@@ -23,6 +23,6 @@ spu_cc_library(
     srcs = ["codegen.cc"],
     hdrs = ["codegen.h"],
     deps = [
-        "//libspu/dialect/pphlo:dialect",
+        "//libspu/dialect/pphlo/IR:dialect",
     ],
 )

libspu/compiler/core/BUILD.bazel

@@ -21,7 +21,7 @@ spu_cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "//libspu/compiler/common:compilation_context",
-        "//libspu/compiler/passes:all_passes",
+        "//libspu/dialect/pphlo/transforms:all_passes",
         "@llvm-project//mlir:Transforms",
     ],
 )

libspu/compiler/core/core.cc

@@ -21,8 +21,8 @@
 #include "mlir/Transforms/Passes.h"
 
 #include "libspu/compiler/common/compilation_context.h"
-#include "libspu/compiler/passes/passes.h"
 #include "libspu/core/prelude.h"
+#include "libspu/dialect/pphlo/transforms/passes.h"
 
 namespace spu::compiler {
 
@@ -51,14 +51,17 @@ void Core::buildPipeline(mlir::PassManager *pm) {
     optPM.addPass(mlir::createCSEPass());
     optPM.addPass(mlir::spu::pphlo::createOptimizeMaxPoolingPass());
   }
-  optPM.addPass(mlir::spu::pphlo::createDecomposeComparisonPass());
-  optPM.addPass(mlir::spu::pphlo::createDecomposeMinMaxPass());
+  optPM.addPass(mlir::spu::pphlo::createDecomposeOps());
   optPM.addPass(mlir::spu::pphlo::createSortLowering());
 
   if (!options.disable_partial_sort_optimization()) {
     optPM.addPass(mlir::spu::pphlo::createPartialSortToTopK());
   }
 
+  optPM.addPass(mlir::spu::pphlo::createRewriteSignbitPatterns());
+
+  optPM.addPass(mlir::spu::pphlo::createInlineSecretControlFlow());
+
   if (!options.disable_sqrt_plus_epsilon_rewrite()) {
     optPM.addPass(mlir::spu::pphlo::createOptimizeSqrtPlusEps());
   }

libspu/compiler/core/core.h

@@ -34,6 +34,11 @@ class Core final {
 private:
   CompilationContext *ctx_;
 
+#ifdef EXPOSE_PIPELINE_BUILDER
+public:
+#else
+private: // NOLINT:
+#endif
   void buildPipeline(mlir::PassManager *pm);
 };
 

libspu/compiler/front_end/BUILD.bazel

@@ -71,8 +71,8 @@ spu_cc_library(
     deps = [
         ":hlo_importer",
         "//libspu/compiler/common:compilation_context",
-        "//libspu/compiler/passes:hlo_legalize_to_pphlo",
-        "//libspu/compiler/passes:lower_conversion_cast",
+        "//libspu/compiler/utils",
+        "//libspu/dialect/pphlo/transforms:all_passes",
         "@llvm-project//mlir:FuncExtensions",
         "@llvm-project//mlir:Parser",
         "@xla//xla/mlir_hlo:mhlo_passes",

libspu/compiler/front_end/fe.cc

@@ -29,21 +29,13 @@
 
 #include "libspu/compiler/common/compilation_context.h"
 #include "libspu/compiler/front_end/hlo_importer.h"
-#include "libspu/compiler/passes/passes.h"
+#include "libspu/compiler/utils/utils.h"
 #include "libspu/core/prelude.h"
-#include "libspu/dialect/pphlo/dialect.h"
+#include "libspu/dialect/pphlo/IR/dialect.h"
+#include "libspu/dialect/pphlo/transforms/passes.h"
 
 namespace spu::compiler {
 
-namespace {
-
-mlir::LogicalResult argparser_error_handler(const llvm::Twine &msg) {
-  SPDLOG_ERROR(msg.str());
-  return mlir::failure();
-}
-
-} // namespace
-
 FE::FE(CompilationContext *ctx) : ctx_(ctx) {
   ctx_->getMLIRContext()
       ->loadDialect<mlir::spu::pphlo::PPHloDialect, mlir::mhlo::MhloDialect,
@@ -133,8 +125,8 @@ void FE::buildFrontEndPipeline(mlir::PassManager *pm, const std::string &args) {
   {
     auto l = mlir::spu::pphlo::createLegalizeToPPHloPass();
     if (!args.empty()) {
-      SPU_ENFORCE(
-          l->initializeOptions(args, argparser_error_handler).succeeded());
+      SPU_ENFORCE(l->initializeOptions(args, mlir::spu::argparser_error_handler)
+                      .succeeded());
     }
     pm->addPass(std::move(l));
   }

libspu/compiler/front_end/hlo_importer.h

@@ -28,7 +28,7 @@ class CompilationContext;
 
 class HloImporter final {
 public:
-  explicit HloImporter(CompilationContext *context) : context_(context){};
+  explicit HloImporter(CompilationContext *context) : context_(context) {};
 
   /// Load a xla module and returns a mlir-hlo module
   mlir::OwningOpRef<mlir::ModuleOp>