Repo sync (#234)

CHANGELOG.md

@@ -10,6 +10,8 @@
 >
 > please add your unreleased change here.
 
+- [API] Change IO interface, split share into chunk
+
 ## 20230705
 - [SPU] 0.4.1 release
 - [Improvement] Improve tanh performance

bazel/repositories.bzl

@@ -18,7 +18,7 @@ load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
 
 SECRETFLOW_GIT = "https://github.com/secretflow"
 
-YACL_COMMIT_ID = "fff880ff9e1e50d5260c5a881ff7abc051cc27d5"
+YACL_COMMIT_ID = "a85ab1f6cb642c8e36dc429714a108d5368df762"
 
 def spu_deps():
     _bazel_platform()

bazel/seal.BUILD

@@ -59,5 +59,5 @@ spu_cmake_external(
     out_static_libs = ["libseal-4.1.a"],
     deps = [
         "@com_github_facebook_zstd//:zstd",
-    ]
+    ],
 )

benchmark/setup_dockers_and_run.sh

@@ -42,7 +42,7 @@ sleep 10
 
 echo -e "${COLOR_GREEN}Run benchmark${COLOR_END}"
 docker run --rm --mount type=bind,source="$(pwd)",target=/home/admin/dev/ --network nn-benchmark spu-build:v1 \
-                sh -c "cd /home/admin/dev && bash benchmark/run_bench.sh $@" | tee benchmark_results.log;
+    sh -c "cd /home/admin/dev && bash benchmark/run_bench.sh $@" | tee benchmark_results.log;
 
 echo -e "${COLOR_GREEN}Shutdown docker compose${COLOR_END}"
 docker-compose -f .circleci/benchmark.yml down

build_wheel_entrypoint.sh

@@ -15,7 +15,6 @@
 # limitations under the License.
 #
 
-
 pip install numpy
 
 python setup.py bdist_wheel

docs/development/add_protocols.rst

@@ -167,7 +167,7 @@ Inside the **makeAby3Protocol** function, it does three things.
 - The third is to register the protocol kernels (functions). We can see that three types of kernels are registered. \
   The first type is the kernels implemented in the `pv2k.cc <https://github.com/secretflow/spu/blob/main/libspu/mpc/common/pv2k.cc>`_ \
   file, using **Pub2k** as the naming prefix of kernel classes. The second type is the kernels implemented in the \
-  `ab_api.cc <https://github.com/secretflow/spu/blob/main/libspu/mpc/common/ab_api.cc>`_ file, using **ABProt** as the \ 
+  `ab_api.cc <https://github.com/secretflow/spu/blob/main/libspu/mpc/ab_api.cc>`_ file, using **ABProt** as the \
   naming prefix of kernel classes. The third type is implemented in `arithmetic.cc <https://github.com/secretflow/spu/blob/main/libspu/mpc/aby3/arithmetic.cc>`_, \
   `boolean.cc <https://github.com/secretflow/spu/blob/main/libspu/mpc/aby3/boolean.cc>`_ and other files under the aby3 directory.
 

libspu/compiler/core/core.cc

@@ -83,6 +83,8 @@ void Core::buildPipeline(mlir::PassManager *pm) {
 
   optPM.addPass(mlir::createLoopInvariantCodeMotionPass());
   optPM.addPass(mlir::createCSEPass());
+
+  optPM.addPass(mlir::pphlo::createInsertDeallocationOp());
 }
 
 } // namespace spu::compiler

libspu/compiler/passes/BUILD.bazel

@@ -234,6 +234,18 @@ spu_cc_library(
     ],
 )
 
+spu_cc_library(
+    name = "insert_deallocation",
+    srcs = ["insert_deallocation.cc"],
+    hdrs = ["passes.h"],
+    deps = [
+        ":pass_details",
+        "//libspu/dialect:pphlo_dialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
 spu_cc_library(
     name = "all_passes",
     hdrs = ["register_passes.h"],
@@ -242,6 +254,7 @@ spu_cc_library(
         ":decompose_minmax",
         ":expand_secret_gather",
         ":hlo_legalize_to_pphlo",
+        ":insert_deallocation",
         ":lower_conversion_cast",
         ":lower_mixed_type_op",
         ":optimize_denominator_with_broadcast",

libspu/compiler/passes/insert_deallocation.cc

@@ -0,0 +1,127 @@
+// Copyright 2023 Ant Group Co., Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <memory>
+
+#include "mlir/Analysis/Liveness.h"
+#include "mlir/Pass/Pass.h"
+
+#include "libspu/compiler/passes/pass_details.h"
+#include "libspu/dialect/pphlo_ops.h"
+
+#ifdef ENABLE_LIVENESS_DEBUG
+
+#include "spdlog/spdlog.h"
+
+void printLiveness(mlir::Liveness *liveness) {
+  std::string buf;
+  llvm::raw_string_ostream os(buf);
+
+  liveness->print(os);
+
+  SPDLOG_INFO("liveness = {}", os.str());
+}
+
+#endif
+
+namespace mlir::pphlo {
+
+namespace {
+
+struct Deallocator {
+private:
+  std::unique_ptr<Liveness> top_liveness_;
+
+public:
+  LogicalResult transformOp(Operation *op,
+                            const LivenessBlockInfo *block_liveness) {
+    for (const auto &operand : op->getOperands()) {
+      if (block_liveness->isLiveOut(operand) ||
+          mlir::isa<BlockArgument>(operand)) {
+        // skip live out values and block args
+        continue;
+      }
+
+      if (operand.getDefiningOp()->getParentRegion() != op->getParentRegion()) {
+        // This value is captured by current region, right now we do not handle
+        // cross region ownership.. skip
+        continue;
+      }
+
+      if (top_liveness_->isDeadAfter(operand, op)) {
+        OpBuilder builder(op->getContext());
+        builder.setInsertionPointAfter(op);
+        builder.create<FreeOp>(op->getLoc(), operand);
+      }
+    }
+
+    for (int64_t idx = 0; idx < op->getNumRegions(); ++idx) {
+      if (failed(transformRegion(op->getRegion(idx)))) {
+        return failure();
+      }
+    }
+
+    return success();
+  }
+
+  LogicalResult transformBlock(Block &block) {
+    const auto *block_liveness = top_liveness_->getLiveness(&block);
+    for (auto &op : llvm::make_early_inc_range(block.without_terminator())) {
+      auto opResult = transformOp(&op, block_liveness);
+      if (failed(opResult)) {
+        return failure();
+      }
+    }
+    return success();
+  }
+
+  LogicalResult transformRegion(Region &r) {
+    for (auto &b : r.getBlocks()) {
+      if (failed(transformBlock(b))) {
+        return failure();
+      }
+    }
+    return success();
+  }
+
+  LogicalResult transformFuncOp(func::FuncOp op) {
+    if (op->getNumRegions() == 0) {
+      return success();
+    }
+
+    top_liveness_ = std::make_unique<Liveness>(op);
+
+    // Transform function body.
+    if (failed(transformRegion(op.getBody()))) {
+      return failure();
+    }
+
+    return success();
+  }
+};
+
+struct InsertDeallocation : public InsertDeallocationBase<InsertDeallocation> {
+  void runOnOperation() override {
+    if (failed(Deallocator().transformFuncOp(getOperation()))) {
+      signalPassFailure();
+    }
+  }
+};
+} // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createInsertDeallocationOp() {
+  return std::make_unique<InsertDeallocation>();
+}
+
+} // namespace mlir::pphlo

libspu/compiler/passes/passes.h

@@ -71,6 +71,8 @@ std::unique_ptr<OperationPass<func::FuncOp>> createRewriteDivSqrtPatterns();
 std::unique_ptr<OperationPass<func::FuncOp>>
 createOptimizeDenominatorWithBroadcast();
 
+std::unique_ptr<OperationPass<func::FuncOp>> createInsertDeallocationOp();
+
 } // namespace pphlo
 
 } // namespace mlir

libspu/compiler/passes/passes.td

@@ -92,3 +92,9 @@ def OptimizeDenominatorWithBcast: Pass<"optimize-denominator-with-broadcast", "f
   let constructor = "createOptimizeDenominatorWithBroadcast()";
   let dependentDialects = ["pphlo::PPHloDialect"];
 }
+
+def InsertDeallocation: Pass<"insert-deallocation", "func::FuncOp"> {
+  let summary = "Insert deallocation ops";
+  let constructor = "createInsertDeallocationOp()";
+  let dependentDialects = ["pphlo::PPHloDialect"];
+}

libspu/compiler/tests/pphlo_simple_dealloc.mlir

@@ -0,0 +1,13 @@
+// RUN: mlir-pphlo-opt --insert-deallocation --split-input-file %s | FileCheck %s
+
+func.func @main() -> (tensor<!pphlo.pub<f32>>) {
+    %0 = "pphlo.constant"() {value = dense<0xFF800000> : tensor<f32>} : () -> tensor<!pphlo.pub<f32>>
+    %1 = "pphlo.constant"() {value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<!pphlo.pub<f32>>
+    %2 = "pphlo.less"(%0, %1): (tensor<!pphlo.pub<f32>>, tensor<!pphlo.pub<f32>>) -> tensor<!pphlo.pub<i1>>
+    %3 = "pphlo.select"(%2, %0, %1): (tensor<!pphlo.pub<i1>>, tensor<!pphlo.pub<f32>>, tensor<!pphlo.pub<f32>>) -> tensor<!pphlo.pub<f32>>
+    //CHECK: "pphlo.free"(%1)
+    //CHECK: "pphlo.free"(%0)
+    //CHECK: "pphlo.free"(%2)
+    //CHECK-NOT: "pphlo.free"(%3)
+    return %3: tensor<!pphlo.pub<f32>>
+}

libspu/core/BUILD.bazel

@@ -38,6 +38,7 @@ spu_cc_library(
     deps = [
         "//libspu/core:prelude",
         "@com_google_absl//absl/types:span",
+        "@yacl//yacl/link",
         "@yacl//yacl/utils:scope_guard",
     ],
 )

libspu/core/context.h

@@ -33,7 +33,7 @@ namespace spu {
 class SPUContext final {
   RuntimeConfig config_;
 
-  // A dynamic object for polymophic(multi-stage) operations.
+  // A dynamic object for polymorphic(multi-stage) operations.
   std::unique_ptr<Object> prot_;
 
   // TODO(jint): do we really need a link here? how about a FHE context.
@@ -107,6 +107,10 @@ class KernelEvalContext final {
 
   SPUContext* sctx() { return sctx_; }
 
+  const std::shared_ptr<yacl::link::Context>& lctx() const {
+    return sctx_->lctx();
+  }
+
   const std::string& id() { return sctx_->id(); }
   const std::string& pid() { return sctx_->pid(); }
 

libspu/core/ndarray_ref.cc

@@ -250,6 +250,11 @@ NdArrayRef NdArrayRef::reshape(absl::Span<const int64_t> to_shape) const {
   SPU_ENFORCE(calcNumel(shape()) == calcNumel(to_shape),
               "reshape from {} to {} is changing numel", shape(), to_shape);
 
+  // Reshape empty is always a noop
+  if (calcNumel(to_shape) == 0) {
+    return NdArrayRef(buf(), eltype(), to_shape, strides(), offset());
+  }
+
   std::vector<int64_t> new_strides(to_shape.size(), 0);
   if (attempt_nocopy_reshape(*this, to_shape, new_strides)) {
     // No copy reshape
@@ -548,11 +553,16 @@ NdArrayRef unflatten(const ArrayRef& arr, absl::Span<const int64_t> shape) {
                       std::vector<int64_t>(shape.size(), 0), arr.offset());
   }
 
-  ArrayRef compact = arr.isCompact() ? arr : arr.clone();
-
   auto strides = makeCompactStrides(shape);
-  return NdArrayRef(compact.buf(), compact.eltype(), shape, std::move(strides),
-                    compact.offset());
+
+  if (arr.stride() != 1) {
+    for (auto& s : strides) {
+      s *= arr.stride();
+    }
+  }
+
+  return NdArrayRef(arr.buf(), arr.eltype(), shape, std::move(strides),
+                    arr.offset());
 }
 
 namespace {