repo-sync-2023-08-08T15:56:58+0800

.circleci/testpypi-config.yml

@@ -0,0 +1,95 @@
+# 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.
+
+# Use the latest 2.1 version of CircleCI pipeline process engine.
+# See: https://circleci.com/docs/2.0/configuration-reference
+version: 2.1
+
+parameters:
+  GHA_Action:
+    type: string
+    default: ""
+  GHA_Actor:
+    type: string
+    default: ""
+  GHA_Event:
+    type: string
+    default: ""
+  GHA_Meta:
+    type: string
+    default: ""
+
+
+# Define a job to be invoked later in a workflow.
+# See: https://circleci.com/docs/2.0/configuration-reference/#jobs
+jobs:
+  macOS_arm64_publish:
+    macos:
+      xcode: 14.2
+    environment:
+      HOMEBREW_NO_AUTO_UPDATE: 1
+    resource_class: macos.m1.large.gen1
+    steps:
+      - checkout
+      - run:
+          name: "Install homebrew dependencies"
+          command: |
+            brew install bazel cmake ninja libomp wget go
+            (cd /opt/homebrew/Cellar/bazel/*.*.*/libexec/bin && curl -fLO https://github.com/bazelbuild/bazel/releases/download/6.2.1/bazel-6.2.1-darwin-arm64 && chmod +x bazel-6.2.1-darwin-arm64)
+      - run:
+          name: "Install Miniconda"
+          command: |
+            wget https://repo.anaconda.com/miniconda/Miniconda3-py38_23.1.0-1-MacOSX-arm64.sh -O ~/miniconda.sh
+            bash ~/miniconda.sh -b -p $HOME/miniconda
+            source $HOME/miniconda/bin/activate
+            conda init zsh bash
+      - run:
+          name: "build package and publish"
+          command: |
+            set +e
+            conda create -n build python=3.8 -y
+            conda activate build
+
+            sh ./build_wheel_entrypoint.sh
+            python3 -m pip install twine
+            ls dist/*.whl
+
+            python3 -m twine upload -r testpypi -u __token__ -p ${TEST_PYPI_TWINE_TOKEN} dist/*.whl
+  linux_publish:
+    docker:
+      - image: secretflow/release-ci:latest
+    resource_class: 2xlarge
+    shell: /bin/bash --login -eo pipefail
+    steps:
+      - checkout
+      - run:
+          name: "build package and publish"
+          command: |
+            conda update -n base -c defaults conda -y
+            conda create -n build python=3.8 -y
+            conda activate build
+
+            sh ./build_wheel_entrypoint.sh
+            python3 -m pip install twine
+            ls dist/*.whl
+
+            python3 -m twine upload -r testpypi -u __token__ -p ${TEST_PYPI_TWINE_TOKEN} dist/*.whl
+
+# Invoke jobs via workflows
+# See: https://circleci.com/docs/2.0/configuration-reference/#workflows
+workflows:
+  publish:
+    jobs:
+      - linux_publish
+      - macOS_arm64_publish

.github/workflows/test-pypi.yml

@@ -0,0 +1,30 @@
+# 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.
+
+name: "Publish TestPypi Package"
+
+on:
+  workflow_dispatch:
+
+jobs:
+  trigger-circleci:
+    runs-on: ubuntu-latest
+    steps:
+      - name: secretflow-spu-testpypi
+        id: secretflow-spu-testpypi
+        uses: CircleCI-Public/trigger-circleci-pipeline-action@v1.1.0
+        with:
+          GHA_Meta: "test_pypi"
+        env:
+          CCI_TOKEN: ${{ secrets.CCI_TOKEN }}

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 = "6d07cc784f9d7425a68fb251e646b59f69727d07"
+YACL_COMMIT_ID = "5d6112505d52cfa7a27cd10e364e1c9893e50c8f"
 
 def spu_deps():
     _bazel_platform()
@@ -65,22 +65,22 @@ def _bazel_platform():
     http_archive(
         name = "platforms",
         urls = [
-            "https://mirror.bazel.build/github.com/bazelbuild/platforms/releases/download/0.0.6/platforms-0.0.6.tar.gz",
-            "https://github.com/bazelbuild/platforms/releases/download/0.0.6/platforms-0.0.6.tar.gz",
+            "https://mirror.bazel.build/github.com/bazelbuild/platforms/releases/download/0.0.7/platforms-0.0.7.tar.gz",
+            "https://github.com/bazelbuild/platforms/releases/download/0.0.7/platforms-0.0.7.tar.gz",
         ],
-        sha256 = "5308fc1d8865406a49427ba24a9ab53087f17f5266a7aabbfc28823f3916e1ca",
+        sha256 = "3a561c99e7bdbe9173aa653fd579fe849f1d8d67395780ab4770b1f381431d51",
     )
 
 def _com_github_facebook_zstd():
     maybe(
         http_archive,
         name = "com_github_facebook_zstd",
         build_file = "@spulib//bazel:zstd.BUILD",
-        strip_prefix = "zstd-1.5.0",
-        sha256 = "5194fbfa781fcf45b98c5e849651aa7b3b0a008c6b72d4a0db760f3002291e94",
+        strip_prefix = "zstd-1.5.5",
+        sha256 = "98e9c3d949d1b924e28e01eccb7deed865eefebf25c2f21c702e5cd5b63b85e1",
         type = ".tar.gz",
         urls = [
-            "https://github.com/facebook/zstd/releases/download/v1.5.0/zstd-1.5.0.tar.gz",
+            "https://github.com/facebook/zstd/archive/refs/tags/v1.5.5.tar.gz",
         ],
     )
 
@@ -169,10 +169,10 @@ def _com_github_pybind11_bazel():
     maybe(
         http_archive,
         name = "pybind11_bazel",
-        sha256 = "6426567481ee345eb48661e7db86adc053881cb4dd39fbf527c8986316b682b9",
-        strip_prefix = "pybind11_bazel-fc56ce8a8b51e3dd941139d329b63ccfea1d304b",
+        sha256 = "2d3316d89b581966fc11eab9aa9320276baee95c8233c7a8efc7158623a48de0",
+        strip_prefix = "pybind11_bazel-ff261d2e9190955d0830040b20ea59ab9dbe66c8",
         urls = [
-            "https://github.com/pybind/pybind11_bazel/archive/fc56ce8a8b51e3dd941139d329b63ccfea1d304b.zip",
+            "https://github.com/pybind/pybind11_bazel/archive/ff261d2e9190955d0830040b20ea59ab9dbe66c8.zip",
         ],
     )
 
@@ -181,10 +181,10 @@ def _com_github_pybind11():
         http_archive,
         name = "pybind11",
         build_file = "@pybind11_bazel//:pybind11.BUILD",
-        sha256 = "5d8c4c5dda428d3a944ba3d2a5212cb988c2fae4670d58075a5a49075a6ca315",
-        strip_prefix = "pybind11-2.10.3",
+        sha256 = "d475978da0cdc2d43b73f30910786759d593a9d8ee05b1b6846d1eb16c6d2e0c",
+        strip_prefix = "pybind11-2.11.1",
         urls = [
-            "https://github.com/pybind/pybind11/archive/refs/tags/v2.10.3.tar.gz",
+            "https://github.com/pybind/pybind11/archive/refs/tags/v2.11.1.tar.gz",
         ],
     )
 

docs/reference/gen_benchmark_report.py

@@ -139,9 +139,9 @@ def format_time_with_last_col(self, df: pd.DataFrame):
                 df[df[time_col_name] == time_col_value], 'ns'
             )
             unit = get_time_unit(weight, 'ns')
-            formated = self.format_time(df[df[time_col_name] == time_col_value], unit)
-            for i in formated.index:
-                df.loc[i, :] = formated.loc[i, :]
+            formatted = self.format_time(df[df[time_col_name] == time_col_value], unit)
+            for i in formatted.index:
+                df.loc[i, :] = formatted.loc[i, :]
             df[time_col_name].replace(
                 time_col_value,
                 time_col_value + '/' + unit,

examples/cpp/simple_lr.cc

@@ -24,10 +24,7 @@
 
 #include "examples/cpp/utils.h"
 #include "spdlog/spdlog.h"
-#include "xtensor/xarray.hpp"
 #include "xtensor/xcsv.hpp"
-#include "xtensor/xio.hpp"
-#include "xtensor/xview.hpp"
 
 #include "libspu/device/io.h"
 #include "libspu/kernel/hal/hal.h"

libspu/compiler/common/ir_printer_config.cc

@@ -70,7 +70,7 @@ void IRPrinterConfig::printAfterIfEnabled(Pass *pass, Operation *operation,
 }
 
 std::string IRPrinterConfig::genFileName(StringRef pass_name, StringRef stage) {
-  return fmt::format("{}-{}-{}.mlir", pp_cnt++, pass_name, stage);
+  return fmt::format("{}-{}-{}.mlir", pp_cnt++, pass_name.str(), stage.str());
 }
 
 } // namespace mlir::pphlo

libspu/compiler/front_end/BUILD.bazel

@@ -34,23 +34,31 @@ spu_cc_library(
         "@xla//xla/service:cholesky_expander",
         "@xla//xla/service:conditional_simplifier",
         "@xla//xla/service:conditional_to_select",
+        "@xla//xla/service:convolution_4d_expander",
         "@xla//xla/service:convolution_group_converter",
         "@xla//xla/service:dot_decomposer",
-        "@xla//xla/service:dot_merger",
+        "@xla//xla/service:eigh_expander",
         "@xla//xla/service:float_normalization",
+        "@xla//xla/service:gather_simplifier",
         "@xla//xla/service:hlo_constant_folding",
+        "@xla//xla/service:hlo_cse",
         "@xla//xla/service:hlo_dce",
         "@xla//xla/service:hlo_pass_pipeline",
         "@xla//xla/service:map_inliner",
         "@xla//xla/service:operand_upcaster",
         "@xla//xla/service:qr_expander",
+        "@xla//xla/service:real_imag_expander",
         "@xla//xla/service:reshape_mover",
         "@xla//xla/service:result_caster",
         "@xla//xla/service:scatter_expander",
+        "@xla//xla/service:slice_sinker",
         "@xla//xla/service:sort_simplifier",
+        "@xla//xla/service:stable_sort_expander",
         "@xla//xla/service:triangular_solve_expander",
+        "@xla//xla/service:while_loop_constant_sinking",
         "@xla//xla/service:while_loop_simplifier",
         "@xla//xla/service:zero_sized_hlo_elimination",
+        "@xla//xla/service/gpu:dot_dimension_sorter",
         "@xla//xla/translate/hlo_to_mhlo:hlo_module_importer",
     ],
 )

libspu/compiler/front_end/hlo_importer.cc

@@ -23,12 +23,16 @@
 #include "xla/service/cholesky_expander.h"
 #include "xla/service/conditional_simplifier.h"
 #include "xla/service/conditional_to_select.h"
+#include "xla/service/convolution_4d_expander.h"
 #include "xla/service/convolution_group_converter.h"
 #include "xla/service/dot_decomposer.h"
-#include "xla/service/dot_merger.h"
+#include "xla/service/eigh_expander.h"
 #include "xla/service/float_normalization.h"
 #include "xla/service/float_support.h"
+#include "xla/service/gather_simplifier.h"
+#include "xla/service/gpu/dot_dimension_sorter.h"
 #include "xla/service/hlo_constant_folding.h"
+#include "xla/service/hlo_cse.h"
 #include "xla/service/hlo_dce.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_pass_fix.h"
@@ -37,12 +41,16 @@
 #include "xla/service/map_inliner.h"
 #include "xla/service/operand_upcaster.h"
 #include "xla/service/qr_expander.h"
+#include "xla/service/real_imag_expander.h"
 #include "xla/service/reshape_mover.h"
 #include "xla/service/result_caster.h"
 #include "xla/service/scatter_expander.h"
+#include "xla/service/slice_sinker.h"
 #include "xla/service/sort_simplifier.h"
+#include "xla/service/stable_sort_expander.h"
 #include "xla/service/triangular_solve_expander.h"
 #include "xla/service/tuple_simplifier.h"
+#include "xla/service/while_loop_constant_sinking.h"
 #include "xla/service/while_loop_simplifier.h"
 #include "xla/service/zero_sized_hlo_elimination.h"
 #include "xla/translate/hlo_to_mhlo/hlo_module_importer.h"
@@ -54,7 +62,17 @@
 
 namespace xla {
 void runHloPasses(xla::HloModule *module) {
-  HloPassPipeline pipeline("HLO passes");
+
+  // Simplifier options
+  AlgebraicSimplifierOptions options;
+  // For MPC, dot is way faster than reduce
+  options.set_enable_dot_strength_reduction(false);
+  // We do not handle nan, so just use faster minmax
+  options.set_minmax_propagate_nan(false);
+  // Transpose and reshape is cheep for us
+  options.set_unconditionally_simplify_reduce_of_transpose_or_reshape(true);
+
+  HloPassPipeline pipeline("optimization");
   pipeline.AddInvariantChecker<HloVerifier>(/*layout_sensitive=*/false,
                                             /*allow_mixed_precision=*/false);
 
@@ -70,17 +88,25 @@ void runHloPasses(xla::HloModule *module) {
 
   pipeline.AddPass<CholeskyExpander>(); // Eliminate chol
   pipeline.AddPass<QrExpander>();       // Eliminate qr
+  pipeline.AddPass<EighExpander>();
   pipeline.AddPass<TriangularSolveExpander>();
 
-  // Inline computations with a single call site.
-  pipeline.AddPass<CallInliner>(/*single_call_site=*/true);
-  pipeline.AddPass<BatchDotSimplification>();
-  pipeline.AddPass<DotDecomposer>(); // Simplify dot
   // Convert BF16 operations to F32 operations so that the SPU backend can
   // support BF16 operations without directly implementing a BF16 lowering for
   // most ops.
   FloatSupport bf16_support(BF16);
   pipeline.AddPass<FloatNormalization>(&bf16_support);
+
+  // Inline computations with a single call site.
+  pipeline.AddPass<CallInliner>(/*single_call_site=*/true);
+  pipeline.AddPass<gpu::DotDimensionSorter>();
+  pipeline.AddPass<BatchDotSimplification>();
+  pipeline.AddPass<DotDecomposer>(); // Simplify dot
+
+  pipeline.AddPass<Convolution4DExpander>();
+
+  pipeline.AddPass<StableSortExpander>();
+
   // After canonicalization, there may be more batch dots that can be
   // simplified.
   pipeline.AddPass<BatchDotSimplification>();
@@ -95,34 +121,33 @@ void runHloPasses(xla::HloModule *module) {
       /*rewrite_training_op=*/true,
       /*rewrite_inference_op=*/true,
       /*rewrite_grad_op=*/true);
-  pipeline.AddPass<ScatterExpander>(ScatterExpander::kEliminateAllScatters);
 
   // Run the following passes to a fixed point.
-  [&pipeline =
-       pipeline.AddPass<HloPassFix<HloPassPipeline>>("simplification")] {
+  [&, &pipeline =
+          pipeline.AddPass<HloPassFix<HloPassPipeline>>("simplification")] {
     pipeline.AddInvariantCheckerDebug<HloVerifier>(
         /*layout_sensitive=*/false,
         /*allow_mixed_precision=*/false);
 
-    AlgebraicSimplifierOptions options;
-    options.set_enable_dot_strength_reduction(false);
+    pipeline.AddPass<GatherSimplifier>();
+    pipeline.AddPass<ScatterExpander>(ScatterExpander::kEliminateAllScatters);
     pipeline.AddPass<AlgebraicSimplifier>(options);
-    pipeline.AddPass<SortSimplifier>();
-    pipeline.AddPass<HloDCE>();
-
-    // BatchNormExpander can create zero-sized ops, so zero-sized HLO
-    // elimination has to come after that pass.
-    pipeline.AddPass<ZeroSizedHloElimination>();
+    pipeline.AddPass<BitcastDtypesExpander>();
+    // AlgebraicSimplifier may add contracting dimensions to a dot.
+    pipeline.AddPass<gpu::DotDimensionSorter>();
+    pipeline.AddPass<DotDecomposer>();
 
+    pipeline.AddPass<SortSimplifier>();
     pipeline.AddPass<TupleSimplifier>();
     pipeline.AddPass<WhileLoopSimplifier>();
-
-    pipeline.AddPass<HloDCE>();
+    pipeline.AddPass<SliceSinker>();
     pipeline.AddPass<ReshapeMover>();
     pipeline.AddPass<HloConstantFolding>();
     pipeline.AddPass<ConditionalSimplifier>();
+    pipeline.AddPass<RealImagExpander>();
+    pipeline.AddPass<HloCSE>(/*is_layout_sensitive=*/false);
+    pipeline.AddPass<HloDCE>();
   }();
-  pipeline.AddPass<BitcastDtypesExpander>();
 
   auto status = pipeline.Run(module).status();
 
@@ -194,4 +219,4 @@ HloImporter::parseXlaModuleFromString(const std::string &content) {
   return mlir_hlo;
 }
 
-} // namespace spu::compiler
+} // namespace spu::compiler