Weekly sync (#43)

* repo-sync-2022-12-02T10:50:14+0800

* Fix merge

* Fix merge

* Fix yacl commit id.

CHANGELOG.md

@@ -8,8 +8,10 @@
 
 ## staging
 > please add your unreleased change here.
-
 - [Feature] ECDH-PSI supports white box interconnection mode
+- [Feature] Various performance improvements
+- [3p] Build with Tensorflow 2.11.0
+- [bugfix] Fix various crashes
 
 ## 20221116
 - [SPU] 0.3.0 release

examples/cpp/BUILD.bazel

@@ -49,8 +49,9 @@ spu_cc_binary(
     srcs = ["simple_pphlo.cc"],
     deps = [
         ":utils",
+        "//spu/device:api",
         "//spu/device:io",
-        "//spu/device/pphlo:executor",
+        "//spu/device/pphlo:pphlo_executor",
         "@llvm-project//llvm:Support",
     ],
 )

examples/cpp/simple_pphlo.cc

@@ -21,36 +21,37 @@
 #include "examples/cpp/utils.h"
 #include "spdlog/spdlog.h"
 
+#include "spu/device/api.h"
 #include "spu/device/io.h"
-#include "spu/device/pphlo/executor.h"
+#include "spu/device/pphlo/pphlo_executor.h"
 
 // This example demostrates the basic compute functionality of spu vm.
-void constant_add(spu::device::Executor* executor) {
+void constant_add(spu::HalContext* hctx) {
   // Write the assembly, this code simple add two numbers.
   // - `%1` is a constant public integer, with dtype int32 and value 1.
   // - `%2` is a constant public integer, with dtype int32 and value 2.
   // - `%3` is the sum of two integers.
   // - `dbg_print` print the value of `%3`
-  constexpr auto code = R"PPHlo(
-func @main() -> () {
+  constexpr auto code = R"(
+func.func @main() -> () {
     %0 = "pphlo.constant"() {value = dense<1> : tensor<i32>} : () -> tensor<!pphlo.pub<i32>>
     %1 = "pphlo.constant"() {value = dense<2> : tensor<i32>} : () -> tensor<!pphlo.pub<i32>>
     %2 = "pphlo.add"(%0, %1) : (tensor<!pphlo.pub<i32>>, tensor<!pphlo.pub<i32>>) -> tensor<!pphlo.pub<i32>>
     "pphlo.dbg_print"(%2) : (tensor<!pphlo.pub<i32>>) -> ()
     return
-}
-)PPHlo";
+})";
 
   // Run it, with no input and output, (since the program does not contain IO)
   spu::device::SymbolTable env;
-  executor->runWithEnv(code, {}, {}, &env);
+  spu::device::pphlo::PPHloExecutor executor;
+  spu::device::execute(&executor, hctx, code, {}, {}, &env);
 }
 
 // This example demostrates how to pass parameters.
-void parameters(spu::device::Executor* executor) {
+void parameters(spu::HalContext* hctx) {
   // In this example, data owner also participates the computation progress,
   // which is called "colocated mode" in spu system.
-  spu::device::ColocatedIo cio(executor->getContext());
+  spu::device::ColocatedIo cio(hctx);
 
   if (cio.getRank() == 0) {
     // rank-0, set a float variable 3.14 as 'x' to the device.
@@ -72,29 +73,29 @@ void parameters(spu::device::Executor* executor) {
   // - `%3` is the product of two values, it will do auto type promotion.
   // - `dbg_print` print the value of `%3`
   constexpr auto code = R"PPHlo(
-func @main(%arg0: tensor<!pphlo.sec<i32>>, %arg1: tensor<!pphlo.sec<i32>>) -> () {
-  %0 = "pphlo.multiply"(%arg0, %arg1) : (tensor<!pphlo.sec<i32>>, tensor<!pphlo.sec<i32>>) -> tensor<!pphlo.sec<i32>>
-  "pphlo.dbg_print"(%0) : (tensor<!pphlo.sec<i32>>) -> ()
+func.func @main(%arg0: tensor<!pphlo.sec<f32>>, %arg1: tensor<!pphlo.sec<i32>>) -> () {
+  %0 = "pphlo.multiply"(%arg0, %arg1) : (tensor<!pphlo.sec<f32>>, tensor<!pphlo.sec<i32>>) -> tensor<!pphlo.sec<f32>>
+  "pphlo.dbg_print"(%0) : (tensor<!pphlo.sec<f32>>) -> ()
   return
-}
-  )PPHlo";
+})PPHlo";
 
   // run the assembly, with
   // - "x" binding to the first parameter (position 0).
   // - "y" binding to the second parameter (position 1).
   // - there is no output bindings.
-  executor->runWithEnv(code, {"x", "y"}, {}, &cio.deviceSymbols());
+  spu::device::pphlo::PPHloExecutor executor;
+  spu::device::execute(&executor, hctx, code, {"x", "y"}, {},
+                       &cio.deviceSymbols());
 }
 
 int main(int argc, char** argv) {
   llvm::cl::ParseCommandLineOptions(argc, argv);
 
   auto hctx = MakeHalContext();
-  spu::device::pphlo::PPHloExecutor executor(hctx.get());
 
-  parameters(&executor);
+  parameters(hctx.get());
 
-  constant_add(&executor);
+  constant_add(hctx.get());
 
   return 0;
 }

examples/python/millionare.py

@@ -17,7 +17,7 @@
 # > bazel run -c opt //examples/python/utils:nodectl -- up
 #
 # Run this example script.
-# > bazel run //examples/python:millionare
+# > bazel run -c opt //examples/python:millionare
 
 
 import argparse

examples/python/ml/flax_mlp.py

@@ -16,7 +16,7 @@
 # > bazel run -c opt //examples/python/utils:nodectl -- up
 #
 # Run this example script.
-# > bazel run //examples/python/ml:flax_mlp
+# > bazel run -c opt //examples/python/ml:flax_mlp
 
 import argparse
 import json

examples/python/ml/jax_lr.py

@@ -16,7 +16,7 @@
 # > bazel run -c opt //examples/python/utils:nodectl -- up
 #
 # Run this example script.
-# > bazel run //examples/python/ml:jax_lr
+# > bazel run -c opt //examples/python/ml:jax_lr
 
 
 import argparse

examples/python/ml/ss_lr.py

@@ -16,7 +16,7 @@
 # > bazel run -c opt //examples/python/utils:nodectl -- up
 #
 # Run this example script.
-# > bazel run //examples/python/ml:ss_lr
+# > bazel run -c opt //examples/python/ml:ss_lr
 
 import argparse
 import json

examples/python/ml/ss_xgb.py

@@ -16,7 +16,7 @@
 # > bazel run -c opt //examples/python/utils:nodectl -- up
 #
 # Run this example script.
-# > bazel run //examples/python/ml:ss_xgb
+# > bazel run -c opt //examples/python/ml:ss_xgb
 
 import argparse
 import json

examples/python/ml/stax_mnist_classifier.py

@@ -23,7 +23,7 @@
 # > bazel run -c opt //examples/python/utils:nodectl -- up
 #
 # Run this example script.
-# > bazel run //examples/python/ml:stax_mnist_classifier
+# > bazel run -c opt //examples/python/ml:stax_mnist_classifier
 
 import time
 import itertools

examples/python/ml/stax_nn.py

@@ -50,12 +50,12 @@
 import argparse
 
 parser = argparse.ArgumentParser(description='distributed driver.')
-parser.add_argument("--model", default='network_a')
-parser.add_argument("-c", "--config", default="examples/python/conf/3pc.json")
-parser.add_argument("-l", "--learning_rate", default=0.01)
-parser.add_argument("-e", "--epoch", default=5)
-parser.add_argument("-b", "--batch_size", default=128)
-parser.add_argument("-o", "--optimizer", default="SGD")
+parser.add_argument("--model", default='network_a', type=str)
+parser.add_argument("-c", "--config", default="examples/python/conf/3pc.json", type=str)
+parser.add_argument("-l", "--learning_rate", default=0.01, type=float)
+parser.add_argument("-e", "--epoch", default=5, type=int)
+parser.add_argument("-b", "--batch_size", default=128, type=int)
+parser.add_argument("-o", "--optimizer", default="SGD", type=str)
 args = parser.parse_args()
 
 # Follows https://arxiv.org/pdf/2107.00501.pdf Appendix C.

examples/python/ml/tf_experiment.py

@@ -29,7 +29,7 @@
 # > bazel run -c opt //examples/python/utils:nodectl -- up
 #
 # Run this example script.
-# > bazel run //examples/python/ml:tf_experiment
+# > bazel run -c opt //examples/python/ml:tf_experiment
 # This example is tf counterpart to //examples/python/ml:jax_lr
 
 

examples/python/stats/corr.py

@@ -16,7 +16,7 @@
 # > bazel run -c opt //examples/python/utils:nodectl -- up
 #
 # Run this example script.
-# > bazel run //examples/python/stats:corr
+# > bazel run -c opt //examples/python/stats:corr
 
 import argparse
 import json

examples/python/stats/pvalue.py

@@ -16,7 +16,7 @@
 # > bazel run -c opt //examples/python/utils:nodectl -- up
 #
 # Run this example script.
-# > bazel run //examples/python/stats:pvalue
+# > bazel run -c opt //examples/python/stats:pvalue
 
 
 import argparse

examples/python/stats/woe.py

@@ -16,7 +16,7 @@
 # > bazel run -c opt //examples/python/utils:nodectl -- up
 #
 # Run this example script.
-# > bazel run //examples/python/stats:woe
+# > bazel run -c opt //examples/python/stats:woe
 
 
 import argparse

spu/BUILD.bazel

@@ -20,7 +20,7 @@ load("//bazel:spu.bzl", "spu_version_file")
 
 package(default_visibility = ["//visibility:public"])
 
-SPU_VERSION = "0.3.1b0"
+SPU_VERSION = "0.3.1b2"
 
 proto_library(
     name = "spu_proto",

spu/binding/BUILD.bazel

@@ -41,8 +41,9 @@ pybind_extension(
         ":version_script.lds",
         "//spu/compiler:compile",
         "//spu/compiler/common:compilation_context",
+        "//spu/device:api",
         "//spu/device:io",
-        "//spu/device/pphlo:executor",
+        "//spu/device/pphlo:pphlo_executor",
         "//spu/psi:bucket_psi",
         "//spu/psi:memory_psi",
         "@yacl//yacl/link",

spu/binding/_lib.cc

@@ -22,8 +22,9 @@
 #include "spu/compiler/common/compilation_context.h"
 #include "spu/compiler/compile.h"
 #include "spu/core/type_util.h"
+#include "spu/device/api.h"
 #include "spu/device/io.h"
-#include "spu/device/pphlo/executor.h"
+#include "spu/device/pphlo/pphlo_executor.h"
 #include "spu/kernel/context.h"
 #include "spu/kernel/value.h"
 #include "spu/psi/bucket_psi.h"
@@ -205,8 +206,7 @@ void BindLink(py::module& m) {
         });
 }
 
-// Wrap Processor, it's workaround for protobuf pybind11/protoc conflict.
-
+// Wrap Runtime, it's workaround for protobuf pybind11/protoc conflict.
 class RuntimeWrapper {
   std::unique_ptr<spu::HalContext> hctx_;
 
@@ -226,8 +226,8 @@ class RuntimeWrapper {
     spu::ExecutableProto exec;
     YACL_ENFORCE(exec.ParseFromString(exec_pb));
 
-    spu::device::pphlo::PPHloExecutor executor(hctx_.get());
-    executor.runWithEnv(exec, &env_);
+    spu::device::pphlo::PPHloExecutor executor;
+    spu::device::execute(&executor, hctx_.get(), exec, &env_);
   }
 
   void SetVar(const std::string& name, const py::bytes& value) {

spu/binding/util/frontend_test.py

@@ -37,14 +37,12 @@ def test_jax_compile(self):
         self.assertEqual(executable.name, "add")
         self.assertEqual(executable.input_names, ["in1", "in2"])
         self.assertEqual(executable.output_names, ["test-out0"])
-        self.assertMultiLineEqual(
-            executable.code.decode(),
-            "module @xla_computation_add {\n"
+        self.assertTrue(
             "  func.func @main(%arg0: tensor<2x!pphlo.pub<i32>>,"
             " %arg1: tensor<2x!pphlo.pub<i32>>) -> tensor<2x!pphlo.pub<i32>> {\n"
             "    %0 = \"pphlo.add\"(%arg0, %arg1) : (tensor<2x!pphlo.pub<i32>>,"
             " tensor<2x!pphlo.pub<i32>>) -> tensor<2x!pphlo.pub<i32>>\n"
-            "    return %0 : tensor<2x!pphlo.pub<i32>>\n  }\n}\n",
+            "    return %0 : tensor<2x!pphlo.pub<i32>>\n  }" in executable.code.decode()
         )
         self.assertEqual(output.shape, (2,))
         self.assertEqual(output.dtype, np.dtype("int32"))
@@ -62,14 +60,12 @@ def test_tf_compile(self):
         self.assertEqual(executable.name, "add")
         self.assertEqual(executable.input_names, ["in1", "in2"])
         self.assertEqual(executable.output_names, ["test-out0"])
-        self.assertMultiLineEqual(
-            executable.code.decode(),
-            "module @a_inference_add_9__.9 {\n"
+        self.assertTrue(
             "  func.func @main(%arg0: tensor<2x!pphlo.pub<i64>>,"
             " %arg1: tensor<2x!pphlo.pub<i64>>) -> tensor<2x!pphlo.pub<i64>> {\n"
             "    %0 = \"pphlo.add\"(%arg0, %arg1) : (tensor<2x!pphlo.pub<i64>>,"
             " tensor<2x!pphlo.pub<i64>>) -> tensor<2x!pphlo.pub<i64>>\n"
-            "    return %0 : tensor<2x!pphlo.pub<i64>>\n  }\n}\n",
+            "    return %0 : tensor<2x!pphlo.pub<i64>>\n  }" in executable.code.decode()
         )
         self.assertEqual(output.shape, (2,))
         self.assertEqual(output.dtype, np.dtype("int64"))

spu/compiler/core/core.cc

@@ -47,6 +47,9 @@ void Core::buildPipeline(mlir::PassManager *pm) {
   optPM.addPass(mlir::pphlo::createOptimizeMaxPoolingPass());
   optPM.addPass(mlir::pphlo::createDecomposeComparisonPass());
   optPM.addPass(mlir::pphlo::createDecomposeMinMaxPass());
+
+  optPM.addPass(mlir::createCSEPass());
+
   optPM.addPass(mlir::pphlo::createReduceTruncationPass());
   optPM.addPass(mlir::pphlo::createLowerMixedTypeOpPass());