Repo sync

.bazelrc

@@ -29,6 +29,12 @@ build:avx --host_copt=-mavx
 build:avx --copt=-DCHECK_AVX
 build:avx --host_copt=-DCHECK_AVX
 
+# default off CUDA build
+build --@rules_cuda//cuda:enable=false
+
+# Only on when asked
+build:gpu --@rules_cuda//cuda:archs=compute_80:compute_80
+build:gpu --@rules_cuda//cuda:enable=true
 
 # Binary safety flags
 build --copt=-fPIC

CHANGELOG.md

@@ -10,6 +10,12 @@
 >
 > please add your unreleased change here.
 
+- [Bugfix] Fix compatibility with latest Jax
+- [Feature] Improve memory efficiency during encode/decode data
+- [Feature] Add radix sort support for SEMI2K
+- [Feature] Experimental: ABY3 matmul CUDA support
+- [Feature] Experimental: Private support under colocated mode
+
 ## 20230906
 
 - [SPU] 0.5.0 release
@@ -168,7 +174,7 @@
 ## 20220325
 
 - [SPU] 0.0.5.1 release
-- [Bugfix] Fix SEMI2K divivsion wrong answer
+- [Bugfix] Fix SEMI2K division wrong answer
 
 ## 20220324
 
@@ -198,7 +204,7 @@
 - [API] merge (config.proto, executable.proto, types.proto) into single spu.proto.
 - [API] change RuntimeConfig.enable_protocol_trace to enable_action_trace.
 - [API] change RuntimeConfig.fxp_recirptocal_goldschmdit_iters to fxp_reciprocal_goldschmdit_iters.
-- [API] add RuntimConfig.reveal_secret_condition to allow reveal secret control flow condition.
+- [API] add RuntimeConfig.reveal_secret_condition to allow reveal secret control flow condition.
 - [Bugfix] Fixed SEGV when reconstruct from an ABY3 scalar secret
 - [Feature] Left/right shift now properly supports non-scalar inputs
 

INSTALLATION.md

@@ -41,3 +41,9 @@ pip install spu
 python setup.py bdist_wheel
 pip install dist/*.whl --force-reinstall
 ```
+
+- Once GCC/bazel/python/Xcode version or other environment settings have changed, please run the following command to ensure a clean build
+
+```bash
+bazel clean --expunge
+```

bazel/nvidia_cutlass.BUILD

@@ -0,0 +1,19 @@
+load("@spulib//bazel:spu.bzl", "spu_cc_library")
+
+package(default_visibility = ["//visibility:public"])
+
+filegroup(
+    name = "all",
+    srcs = glob(["**"]),
+)
+
+spu_cc_library(
+    name = "cutlass",
+    srcs = [],
+    hdrs = glob([
+        "include/**/*.h",
+        "include/**/*.hpp",
+    ]),
+    strip_include_prefix = "include",
+    visibility = ["//visibility:public"],
+)

bazel/repositories.bzl

@@ -42,6 +42,7 @@ def spu_deps():
     _com_github_microsoft_gsl()
     _com_github_microsoft_kuku()
     _com_google_flatbuffers()
+    _com_github_nvidia_cutlass()
 
     maybe(
         git_repository,
@@ -76,9 +77,9 @@ def _rules_proto_grpc():
 def _rules_cuda():
     http_archive(
         name = "rules_cuda",
-        sha256 = "fa1462c4c3104de44489800a1da055f55afa57795789539c835e069818786f71",
-        strip_prefix = "rules_cuda-cab1fa2dd0e1f8489f566c91a5025856cf5ae572",
-        urls = ["https://github.com/bazel-contrib/rules_cuda/archive/cab1fa2dd0e1f8489f566c91a5025856cf5ae572.tar.gz"],
+        sha256 = "2f8c8c8c85f727bec4423efecec12d3b751cb0a98bda99f0f9d351608a23b858",
+        strip_prefix = "rules_cuda-v0.2.1",
+        urls = ["https://github.com/bazel-contrib/rules_cuda/releases/download/v0.2.1/rules_cuda-v0.2.1.tar.gz"],
     )
 
 def _bazel_platform():
@@ -357,3 +358,15 @@ def _com_google_flatbuffers():
             "https://github.com/google/flatbuffers/archive/refs/tags/v23.3.3.tar.gz",
         ],
     )
+
+def _com_github_nvidia_cutlass():
+    maybe(
+        http_archive,
+        name = "com_github_nvidia_cutlass",
+        strip_prefix = "cutlass-3.2.0",
+        urls = [
+            "https://github.com/NVIDIA/cutlass/archive/refs/tags/v3.2.0.tar.gz",
+        ],
+        sha256 = "9637961560a9d63a6bb3f407faf457c7dbc4246d3afb54ac7dc1e014dd7f172f",
+        build_file = "@spulib//bazel:nvidia_cutlass.BUILD",
+    )

examples/cpp/pir/generate_pir_data.cc

@@ -13,11 +13,8 @@
 // limitations under the License.
 
 // clang-format off
-// build generate_pir_data
-// > bazel build //examples/cpp/pir:generate_pir_data -c opt
-//
 // To run the example, start two terminals:
-// > ./generate_pir_data -data_count 10000 -label_len 32 -server_out_path pir_server.csv -client_out_path pir_client.csv
+// > bazel run //examples/cpp/pir:generate_pir_data -c opt -- -data_count 10000 -label_len 32 -server_out_path pir_server.csv -client_out_path pir_client.csv
 // clang-format on
 
 #include <fstream>
@@ -100,4 +97,4 @@ int main(int argc, char **argv) {
   psi2_out_file.close();
 
   return 0;
-}
+}

examples/cpp/pir/keyword_pir_client.cc

@@ -13,11 +13,8 @@
 // limitations under the License.
 
 // clang-format off
-// build keyword_pir_client
-// > bazel build //examples/cpp/pir:keyword_pir_client -c opt
-//
 // To run the example, start terminals:
-// > ./keyword_pir_client -rank 1 -in_path ../../data/psi_client_data.csv.csv
+// > bazel run //examples/cpp/pir:keyword_pir_client -c opt -- -rank 1 -in_path ../../data/psi_client_data.csv.csv
 // >       -key_columns id -out_path pir_out.csv
 // clang-format on
 

examples/cpp/pir/keyword_pir_mem_server.cc

@@ -13,11 +13,8 @@
 // limitations under the License.
 
 // clang-format off
-// build keyword_pir_server
-// > bazel build //examples/cpp/pir:keyword_pir_server -c opt
-//
 // To run the example, start terminals:
-// > ./keyword_pir_server -rank 0 -setup_path pir_setup_dir
+// > bazel run //examples/cpp/pir:keyword_pir_server -c opt -- -rank 0 -setup_path pir_setup_dir
 // >        -oprf_key_path secret_key.bin
 // clang-format on
 

examples/cpp/pir/keyword_pir_server.cc

@@ -13,11 +13,8 @@
 // limitations under the License.
 
 // clang-format off
-// build keyword_pir_server
-// > bazel build //examples/cpp/pir:keyword_pir_server -c opt
-//
 // To run the example, start terminals:
-// > ./keyword_pir_server -rank 0 -setup_path pir_setup_dir
+// > bazel run //examples/cpp/pir:keyword_pir_server -c opt -- -rank 0 -setup_path pir_setup_dir
 // >        -oprf_key_path secret_key.bin
 // clang-format on
 

examples/cpp/pir/keyword_pir_setup.cc

@@ -13,14 +13,10 @@
 // limitations under the License.
 
 // clang-format off
-// build keyword_pir_setup
-// > bazel build //examples/cpp/pir:keyword_pir_setup -c opt
-//
 // To generate ecc oprf secret key, start terminals:
 // > dd if=/dev/urandom of=secret_key.bin bs=32 count=1
-//
 // To run the example, start terminals:
-// > ./keyword_pir_setup -in_path ../../data/psi_server_data.csv -oprf_key_path secret_key.bin
+// > bazel run //examples/cpp/pir:keyword_pir_setup -c opt -- -in_path ../../data/psi_server_data.csv -oprf_key_path secret_key.bin
 // >     -key_columns id -label_columns label -data_per_query 1 -label_max_len 40
 // >     -setup_path pir_setup_dir
 // clang-format on

examples/cpp/simple_dp_psi.cc

@@ -13,12 +13,10 @@
 // limitations under the License.
 
 // clang-format off
-// build simple_dp_psi
-// > bazel build //examples/cpp:simple_dp_psi -c opt
-//
 // To run the example, start two terminals:
-// > ./simple_dp_psi -rank 0 -in_path examples/data/psi_1.csv -field_names id
-// > ./simple_dp_psi -rank 1 -in_path examples/data/psi_2.csv -field_names id -out_path /tmp/p2.out
+// > bazel run //examples/cpp/simple_dp_psi -c opt -- -rank 0 -in_path examples/data/psi_1.csv -field_names id
+// > bazel run //examples/cpp/simple_dp_psi -c opt -- -rank 1 -in_path examples/data/psi_2.csv -field_names id -out_path /tmp/p2.out
+// To run with non-default IP config, add -parties IP:port,IP:port to above commands
 // clang-format on
 
 #include <fstream>

examples/cpp/simple_in_memory_psi.cc

@@ -14,8 +14,9 @@
 
 // clang-format off
 // To run the example, start two terminals:
-// > bazel run //examples/cpp:simple_in_memory_psi -- --rank=0
-// > bazel run //examples/cpp:simple_in_memory_psi -- --rank=1
+// > bazel run //examples/cpp:simple_in_memory_psi -c opt -- -rank=0
+// > bazel run //examples/cpp:simple_in_memory_psi -c opt -- -rank=1
+// To run with non-default IP config, add -parties IP:port,IP:port to above commands
 // clang-format on
 
 #include <cassert>

examples/cpp/simple_lr.cc

@@ -14,8 +14,9 @@
 
 // clang-format off
 // To run the example, start two terminals:
-// > bazel run //examples/cpp:simple_lr -- --dataset=examples/data/perfect_logit_a.csv --has_label=true
-// > bazel run //examples/cpp:simple_lr -- --dataset=examples/data/perfect_logit_b.csv --rank=1
+// > bazel run //examples/cpp:simple_lr -c opt -- -dataset=examples/data/perfect_logit_a.csv -has_label=true
+// > bazel run //examples/cpp:simple_lr -c opt -- -dataset=examples/data/perfect_logit_b.csv -rank=1
+// To run with non-default IP config, add -parties IP:port,IP:port to above commands
 // clang-format on
 
 #include <fstream>

examples/cpp/simple_pphlo.cc

@@ -14,8 +14,9 @@
 
 // clang-format off
 // To run the example, start two terminals:
-// > bazel run //examples/cpp:simple_pphlo -- --rank=0
-// > bazel run //examples/cpp:simple_pphlo -- --rank=1
+// > bazel run //examples/cpp:simple_pphlo -c opt -- -rank=0
+// > bazel run //examples/cpp:simple_pphlo -c opt -- -rank=1
+// To run with non-default IP config, add -parties IP:port,IP:port to above commands
 // clang-format on
 
 #include "examples/cpp/utils.h"

examples/cpp/simple_psi.cc

@@ -14,8 +14,9 @@
 
 // clang-format off
 // To run the example, start two terminals:
-// > bazel run //examples/cpp:simple_psi -- -rank 0 -protocol 1 -in_path examples/data/psi_1.csv -field_names id -out_path /tmp/p1.out
-// > bazel run //examples/cpp:simple_psi -- -rank 1 -protocol 1 -in_path examples/data/psi_2.csv -field_names id -out_path /tmp/p2.out
+// > bazel run //examples/cpp:simple_psi -c opt -- -rank 0 -protocol 1 -in_path examples/data/psi_1.csv -field_names id -out_path /tmp/p1.out
+// > bazel run //examples/cpp:simple_psi -c opt -- -rank 1 -protocol 1 -in_path examples/data/psi_2.csv -field_names id -out_path /tmp/p2.out
+// To run with non-default IP config, add -parties IP:port,IP:port to above commands
 // clang-format on
 
 #include "absl/strings/str_split.h"

examples/python/conf/3pc_colocated.json

@@ -0,0 +1,52 @@
+{
+    "id": "colocated.3pc",
+    "nodes": {
+        "node:0": "127.0.0.1:9920",
+        "node:1": "127.0.0.1:9921",
+        "node:2": "127.0.0.1:9922"
+    },
+    "devices": {
+        "SPU": {
+            "kind": "SPU",
+            "config": {
+                "node_ids": [
+                    "node:0",
+                    "node:1",
+                    "node:2"
+                ],
+                "spu_internal_addrs": [
+                    "127.0.0.1:9930",
+                    "127.0.0.1:9931",
+                    "127.0.0.1:9932"
+                ],
+                "experimental_data_folder": [
+                    "/tmp/spu_data_0/",
+                    "/tmp/spu_data_1/",
+                    "/tmp/spu_data_2/"
+                ],
+                "runtime_config": {
+                    "protocol": "ABY3",
+                    "field": "FM64",
+                    "enable_pphlo_profile": true,
+                    "enable_hal_profile": true,
+                    "experimental_disable_mmul_split": true,
+                    "enable_pphlo_trace": false,
+                    "fxp_exp_mode": 1,
+                    "experimental_enable_colocated_optimization": true
+                }
+            }
+        },
+        "P1": {
+            "kind": "PYU",
+            "config": {
+                "node_id": "node:1"
+            }
+        },
+        "P2": {
+            "kind": "PYU",
+            "config": {
+                "node_id": "node:2"
+            }
+        }
+    }
+}

examples/python/conf/BUILD.bazel

@@ -19,6 +19,7 @@ filegroup(
     srcs = [
         "2pc.json",
         "3pc.json",
+        "3pc_colocated.json",
         "ds_breast_cancer_basic.json",
         "ds_mock_regression_basic.json",
     ],

examples/python/ml/flax_llama7b/README.md

@@ -3,6 +3,8 @@
 This example demonstrates how to use SPU to run secure inference on a pre-trained
 [Llama-7B](https://research.facebook.com/publications/llama-open-and-efficient-foundation-language-models/) model using [Puma](https://arxiv.org/abs/2307.12533).
 
+> **_NOTE:_**  To run LLaMA-7B with ABY3, each node requires at least 1TB of RAM
+
 1. Install huggingface transformers library
 
     ```sh

libspu/compiler/core/core.cc

@@ -70,6 +70,8 @@ void Core::buildPipeline(mlir::PassManager *pm) {
 
   optPM.addPass(mlir::createCSEPass());
 
+  optPM.addPass(mlir::pphlo::createConvertPushDownPass());
+
   if (!options.disable_reduce_truncation_optimization()) {
     optPM.addPass(mlir::pphlo::createReduceTruncationPass());
   }

libspu/compiler/passes/BUILD.bazel

@@ -258,10 +258,23 @@ spu_cc_library(
     ],
 )
 
+spu_cc_library(
+    name = "convert_push_down",
+    srcs = ["convert_push_down.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"],
     deps = [
+        ":convert_push_down",
         ":decompose_comparison",
         ":decompose_minmax",
         ":expand_secret_gather",