repo-sync-2025-02-06T20:02:27+0800 (#988)

.vscode/settings.json

@@ -41,5 +41,11 @@
     "[python]": {
         "editor.defaultFormatter": "ms-python.black-formatter"
     },
-    "mlir.server_path": "bazel-bin/libspu/compiler/tools/spu-lsp"
+    "mlir.server_path": "bazel-bin/libspu/compiler/tools/spu-lsp",
+    "files.exclude": {
+        // "**/bazel-*/**": true,
+        "external":true,
+        ".cache":true,
+        "**/__pycache__":true
+    }
 }

MODULE.bazel

@@ -21,7 +21,7 @@
 
 module(
     name = "spu",
-    version = "0.9.4.dev20250123",
+    version = "0.9.4.dev20250209",
     compatibility_level = 1,
 )
 
@@ -32,13 +32,7 @@ local_path_override(
     path = "src",
 )
 
-bazel_dep(name = "psi")
-git_override(
-    module_name = "psi",
-    commit = "8ead92f1bb10329c7e7e56d541fecb3dcd47ee03",
-    remote = "https://github.com/secretflow/psi.git",
-)
-
+bazel_dep(name = "psi", version = "0.6.0.dev250123")
 bazel_dep(name = "yacl", version = "20241212.0-871832a")
 bazel_dep(name = "grpc", version = "1.66.0.bcr.3")
 single_version_override(

examples/cpp/utils.cc

@@ -19,8 +19,6 @@
 
 #include "libspu/core/config.h"
 
-#include "libspu/spu.pb.h"
-
 llvm::cl::opt<std::string> Parties(
     "parties", llvm::cl::init("127.0.0.1:61530,127.0.0.1:61531"),
     llvm::cl::desc("server list, format: host1:port1[,host2:port2, ...]"));
@@ -52,13 +50,13 @@ std::unique_ptr<spu::SPUContext> MakeSPUContext() {
   auto lctx = MakeLink(Parties.getValue(), Rank.getValue());
 
   spu::RuntimeConfig config;
-  config.set_protocol(static_cast<spu::ProtocolKind>(ProtocolKind.getValue()));
-  config.set_field(static_cast<spu::FieldType>(Field.getValue()));
+  config.protocol = static_cast<spu::ProtocolKind>(ProtocolKind.getValue());
+  config.field = static_cast<spu::FieldType>(Field.getValue());
 
   populateRuntimeConfig(config);
 
-  config.set_enable_action_trace(EngineTrace.getValue());
-  config.set_enable_type_checker(EngineTrace.getValue());
+  config.enable_action_trace = EngineTrace.getValue();
+  config.enable_type_checker = EngineTrace.getValue();
 
   return std::make_unique<spu::SPUContext>(config, lctx);
 }

examples/python/ir_dump/ir_dump.py

@@ -28,7 +28,7 @@
 import jax.numpy as jnp
 import numpy as np
 
-import spu.spu_pb2 as spu_pb2
+import spu.libspu as libspu
 import spu.utils.distributed as ppd
 
 logging.basicConfig(level=logging.INFO)
@@ -48,7 +48,7 @@
 dump_path = os.path.join(os.path.expanduser("~"), args.dir)
 logging.info(f"Dump path: {dump_path}")
 # refer to spu.proto for more detailed configuration
-copts = spu_pb2.CompilerOptions()
+copts = libspu.CompilerOptions()
 copts.enable_pretty_print = True
 copts.pretty_print_dump_dir = dump_path
 copts.xla_pp_kind = 2

examples/python/ml/flax_llama7b/flax_llama7b.py

@@ -31,7 +31,7 @@
 from transformers import LlamaTokenizer
 
 import spu.intrinsic as intrinsic
-import spu.spu_pb2 as spu_pb2
+import spu.libspu as libspu
 import spu.utils.distributed as ppd
 
 parser = argparse.ArgumentParser(description='distributed driver.')
@@ -46,7 +46,7 @@
 
 ppd.init(conf["nodes"], conf["devices"])
 
-copts = spu_pb2.CompilerOptions()
+copts = libspu.CompilerOptions()
 copts.enable_pretty_print = False
 copts.xla_pp_kind = 2
 # enable x / broadcast(y) -> x * broadcast(1/y)

examples/python/ml/flax_llama7b_split/flax_llama7b_split.py

@@ -39,7 +39,7 @@
 from flax.linen.linear import Array
 from transformers import LlamaTokenizer
 
-import spu.spu_pb2 as spu_pb2
+import spu.libspu as libspu
 import spu.utils.distributed as ppd
 
 parser = argparse.ArgumentParser(description='distributed driver.')
@@ -59,7 +59,7 @@
 
 ppd.init(conf["nodes"], conf["devices"])
 
-copts = spu_pb2.CompilerOptions()
+copts = libspu.CompilerOptions()
 copts.enable_pretty_print = False
 copts.xla_pp_kind = 2
 # enable x / broadcast(y) -> x * broadcast(1/y)

examples/python/ml/flax_whisper/flax_whisper.py

@@ -27,7 +27,7 @@
 from transformers import FlaxWhisperForConditionalGeneration, WhisperProcessor
 
 import spu.utils.distributed as ppd
-from spu import spu_pb2
+from spu import libspu
 
 parser = argparse.ArgumentParser(description='distributed driver.')
 parser.add_argument(
@@ -68,7 +68,7 @@ def run_on_spu():
     inputs_ids = processor(ds[0]["audio"]["array"], return_tensors="np")
 
     # Enable rewrite for better performance
-    copts = spu_pb2.CompilerOptions()
+    copts = libspu.CompilerOptions()
     copts.enable_optimize_denominator_with_broadcast = True
 
     input_ids = ppd.device("P1")(lambda x: x)(inputs_ids.input_features)

experimental/squirrel/objectives.cc

@@ -291,11 +291,11 @@ namespace {
     spu::SPUContext* ctx, const spu::Value& _x, float threshold,
     spu::FieldType working_ft) {
   namespace sk = spu::kernel;
-  auto src_field = ctx->config().field();
+  auto src_field = ctx->config().field;
 
   spu::Value x(CastRing(_x.data(), working_ft), _x.dtype());
   // FIXME(lwj): dirty hack
-  const_cast<spu::RuntimeConfig*>(&ctx->config())->set_field(working_ft);
+  const_cast<spu::RuntimeConfig*>(&ctx->config())->field = working_ft;
   ctx->getState<spu::mpc::Z2kState>()->setField(working_ft);
 
   const auto ONE = sk::hal::_constant(ctx, 1, x.shape());
@@ -314,7 +314,7 @@ namespace {
   auto is_too_large = sk::hal::_xor(
       ctx, True, sk::hal::_or(ctx, is_neg, is_inside_range));  // x > t
   // FIXME(lwj): dirty hack
-  const_cast<spu::RuntimeConfig*>(&ctx->config())->set_field(src_field);
+  const_cast<spu::RuntimeConfig*>(&ctx->config())->field = src_field;
   ctx->getState<spu::mpc::Z2kState>()->setField(src_field);
 
   is_neg = spu::Value(CastRing(is_neg.data(), src_field), spu::DT_I1);

experimental/squirrel/objectives_test.cc

@@ -110,9 +110,9 @@ TEST_P(ObjectivesTest, MaxGain) {
 
   spu::mpc::utils::simulate(2, [&](std::shared_ptr<yacl::link::Context> lctx) {
     spu::RuntimeConfig rt_config;
-    rt_config.set_protocol(ProtocolKind::REF2K);
-    rt_config.set_field(field);
-    rt_config.set_fxp_fraction_bits(16);
+    rt_config.protocol = ProtocolKind::REF2K;
+    rt_config.field = field;
+    rt_config.fxp_fraction_bits = 16;
 
     auto _ctx = std::make_unique<spu::SPUContext>(rt_config, lctx);
     auto ctx = _ctx.get();
@@ -170,13 +170,12 @@ TEST_P(ObjectivesTest, Logistic) {
 
   spu::mpc::utils::simulate(2, [&](std::shared_ptr<yacl::link::Context> lctx) {
     spu::RuntimeConfig rt_config;
-    rt_config.set_protocol(ProtocolKind::CHEETAH);
-    rt_config.mutable_cheetah_2pc_config()->set_ot_kind(
-        CheetahOtKind::YACL_Softspoken);
-    rt_config.set_field(field);
-    rt_config.set_fxp_fraction_bits(17);
-    rt_config.set_enable_hal_profile(true);
-    rt_config.set_enable_pphlo_profile(true);
+    rt_config.protocol = ProtocolKind::CHEETAH;
+    rt_config.cheetah_2pc_config.ot_kind = CheetahOtKind::YACL_Softspoken;
+    rt_config.field = field;
+    rt_config.fxp_fraction_bits = 17;
+    rt_config.enable_hal_profile = true;
+    rt_config.enable_pphlo_profile = true;
 
     auto _ctx = std::make_unique<spu::SPUContext>(rt_config, lctx);
     auto ctx = _ctx.get();
@@ -198,7 +197,7 @@ TEST_P(ObjectivesTest, Logistic) {
       return;
     }
 
-    double fxp = std::pow(2., rt_config.fxp_fraction_bits());
+    double fxp = std::pow(2., rt_config.fxp_fraction_bits);
     double max_err = 0.;
 
     for (int64_t i = 0; i < logistic.numel(); ++i) {
@@ -241,13 +240,12 @@ TEST_P(ObjectivesTest, Sigmoid) {
 
   spu::mpc::utils::simulate(2, [&](std::shared_ptr<yacl::link::Context> lctx) {
     spu::RuntimeConfig rt_config;
-    rt_config.set_protocol(ProtocolKind::CHEETAH);
-    rt_config.mutable_cheetah_2pc_config()->set_ot_kind(
-        CheetahOtKind::YACL_Softspoken);
-    rt_config.set_field(field);
-    rt_config.set_fxp_fraction_bits(17);
-    rt_config.set_enable_hal_profile(true);
-    rt_config.set_enable_pphlo_profile(true);
+    rt_config.protocol = ProtocolKind::CHEETAH;
+    rt_config.cheetah_2pc_config.ot_kind = CheetahOtKind::YACL_Softspoken;
+    rt_config.field = field;
+    rt_config.fxp_fraction_bits = 17;
+    rt_config.enable_hal_profile = true;
+    rt_config.enable_pphlo_profile = true;
 
     auto _ctx = std::make_unique<spu::SPUContext>(rt_config, lctx);
     auto ctx = _ctx.get();
@@ -269,7 +267,7 @@ TEST_P(ObjectivesTest, Sigmoid) {
       return;
     }
 
-    double fxp = std::pow(2., rt_config.fxp_fraction_bits());
+    double fxp = std::pow(2., rt_config.fxp_fraction_bits);
     double max_err = 0.;
 
     for (int64_t i = 0; i < logistic.numel(); ++i) {

experimental/squirrel/squirrel_demo_main.cc

@@ -118,11 +118,11 @@ std::unique_ptr<spu::SPUContext> MakeSPUContext() {
   auto lctx = MakeLink(Parties.getValue(), Rank.getValue());
 
   spu::RuntimeConfig config;
-  config.set_protocol(spu::ProtocolKind::CHEETAH);
-  config.set_field(static_cast<spu::FieldType>(Field.getValue()));
-  config.set_fxp_fraction_bits(18);
-  config.set_fxp_div_goldschmidt_iters(1);
-  config.set_enable_hal_profile(EngineTrace.getValue());
+  config.protocol = spu::ProtocolKind::CHEETAH;
+  config.field = static_cast<spu::FieldType>(Field.getValue());
+  config.fxp_fraction_bits = 18;
+  config.fxp_div_goldschmidt_iters = 1;
+  config.enable_hal_profile = EngineTrace.getValue();
   auto hctx = std::make_unique<spu::SPUContext>(config, lctx);
   spu::mpc::Factory::RegisterProtocol(hctx.get(), lctx);
   return hctx;
@@ -218,7 +218,7 @@ void RunTest(spu::SPUContext* hctx, squirrel::XGBTreeBuilder& builder,
   }
 
   const int64_t nsamples = dframe.shape(0);
-  const double fxp = std::pow(2., hctx->config().fxp_fraction_bits());
+  const double fxp = std::pow(2., hctx->config().fxp_fraction_bits);
 
   SPDLOG_DEBUG("Computing inference on testing set ...");
 
@@ -293,7 +293,7 @@ int main(int argc, char** argv) {
       bucket_size, nfeatures, peer_nfeatures);
 
   worker->BuildMap(dframe);
-  worker->Setup(8 * spu::SizeOf(hctx->config().field()), hctx->lctx());
+  worker->Setup(8 * spu::SizeOf(hctx->config().field), hctx->lctx());
 
   std::string act = Activation.getValue();
   SPU_ENFORCE(act == "log" or act == "sig", "invalid activation type={}", act);
@@ -335,7 +335,7 @@ int main(int argc, char** argv) {
   }
 
   // Test on train set
-  double fxp = std::pow(2., hctx->config().fxp_fraction_bits());
+  double fxp = std::pow(2., hctx->config().fxp_fraction_bits);
   int32_t correct = 0;
   SPDLOG_DEBUG("Computing inference on training set ...");
   for (int64_t i = 0; i < (int64_t)nsamples; ++i) {

experimental/squirrel/tree_builder.cc

@@ -589,7 +589,7 @@ double XGBTreeBuilder::DEBUG_OpenObjects(
   Gsum = hal::reveal(ctx, Gsum);
   Hsum = hal::reveal(ctx, Hsum);
   weights = hal::reveal(ctx, weights);
-  const double fxp = std::pow(2., ctx->config().fxp_fraction_bits());
+  const double fxp = std::pow(2., ctx->config().fxp_fraction_bits);
   double object = 0.0;
   for (int64_t i = 0; i < weights.numel(); ++i) {
     double G = Gsum.data().at<int64_t>(i) / fxp;
@@ -606,7 +606,7 @@ double XGBTreeBuilder::DEBUG_OpenLoss(spu::SPUContext* ctx,
   using namespace spu::kernel;
   auto _pred = hal::reveal(ctx, pred);
   auto _label = hal::reveal(ctx, label);
-  double fxp = std::pow(2., ctx->config().fxp_fraction_bits());
+  double fxp = std::pow(2., ctx->config().fxp_fraction_bits);
   double loss = 0.;
   for (int64_t i = 0; i < _pred.numel(); ++i) {
     double y = _label.data().at<int64_t>(i) / fxp;

experimental/squirrel/utils.cc

@@ -127,7 +127,7 @@ spu::Value ArgMax(spu::SPUContext* ctx, const spu::Value& x, int axis,
 
 spu::Value MulArithShareWithPrivateBoolean(spu::SPUContext* ctx,
                                            const spu::Value& ashr) {
-  SPU_ENFORCE(ctx->config().protocol() == spu::ProtocolKind::CHEETAH);
+  SPU_ENFORCE(ctx->config().protocol == spu::ProtocolKind::CHEETAH);
   SPU_ENFORCE(ashr.isSecret());
 
   spu::KernelEvalContext kctx(ctx);
@@ -143,7 +143,7 @@ spu::Value MulArithShareWithPrivateBoolean(spu::SPUContext* ctx,
 spu::Value MulArithShareWithPrivateBoolean(
     spu::SPUContext* ctx, const spu::Value& ashr,
     absl::Span<const uint8_t> prv_boolean) {
-  SPU_ENFORCE(ctx->config().protocol() == spu::ProtocolKind::CHEETAH);
+  SPU_ENFORCE(ctx->config().protocol == spu::ProtocolKind::CHEETAH);
   SPU_ENFORCE(ashr.isSecret());
   SPU_ENFORCE_EQ(ashr.numel(), (int64_t)prv_boolean.size());
 
@@ -160,9 +160,9 @@ spu::Value MulArithShareWithPrivateBoolean(
 spu::Value MulPrivateArithWithPrivateBoolean(spu::SPUContext* ctx,
                                              const spu::Value& arith) {
   using namespace spu;
-  SPU_ENFORCE(ctx->config().protocol() == spu::ProtocolKind::CHEETAH);
+  SPU_ENFORCE(ctx->config().protocol == spu::ProtocolKind::CHEETAH);
   spu::KernelEvalContext kctx(ctx);
-  auto ft = ctx->config().field();
+  auto ft = ctx->config().field;
   auto out = mpc::cheetah::TiledDispatchOTFunc(
       &kctx, arith.data(),
       [&](const NdArrayRef& input,
@@ -184,11 +184,11 @@ spu::Value MulPrivateArithWithPrivateBoolean(spu::SPUContext* ctx,
                                              const spu::DataType dtype,
                                              const spu::Shape& shape) {
   using namespace spu;
-  SPU_ENFORCE(ctx->config().protocol() == spu::ProtocolKind::CHEETAH);
+  SPU_ENFORCE(ctx->config().protocol == spu::ProtocolKind::CHEETAH);
   SPU_ENFORCE_EQ(boolean.size(), (size_t)shape.numel());
 
   spu::KernelEvalContext kctx(ctx);
-  auto ft = ctx->config().field();
+  auto ft = ctx->config().field;
   auto out = mpc::cheetah::TiledDispatchOTFunc(
       &kctx, boolean,
       [&](absl::Span<const uint8_t> input,
@@ -210,10 +210,10 @@ spu::Value MulArithShareWithANDBoolShare(spu::SPUContext* ctx,
   SPU_ENFORCE(ashr.isSecret());
   SPU_ENFORCE_EQ(ashr.numel(), (int64_t)bshr.size());
 
-  SPU_ENFORCE(ctx->config().protocol() == spu::ProtocolKind::CHEETAH);
+  SPU_ENFORCE(ctx->config().protocol == spu::ProtocolKind::CHEETAH);
 
   spu::KernelEvalContext kctx(ctx);
-  auto ft = ctx->config().field();
+  auto ft = ctx->config().field;
   int rank = ctx->lctx()->Rank();
 
   auto out = mpc::cheetah::TiledDispatchOTFunc(