Repo sync (#451)

libspu/compiler/tests/BUILD.bazel

@@ -12,26 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-load("//bazel:glob_lit_test.bzl", "glob_lit_tests")
+load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
+load("@llvm-project//llvm:lit_test.bzl", "lit_test", "package_path")
 load("//bazel:spu.bzl", "spu_cc_test")
 
-glob_lit_tests(
-    data = [":test_utilities"],
-    driver = "@llvm-project//mlir:run_lit.sh",
-    test_file_exts = ["mlir"],
-)
-
-# Bundle together all of the test utilities that are used by tests.
-filegroup(
-    name = "test_utilities",
-    testonly = True,
-    data = [
-        "lit.cfg",
-        "//libspu/compiler/tools:mlir-pphlo-opt",
-        "@llvm-project//llvm:FileCheck",
-    ],
-)
-
 spu_cc_test(
     name = "enum_conversion_test",
     srcs = ["enum_conversion_test.cc"],
@@ -40,3 +24,39 @@ spu_cc_test(
         "//libspu/dialect:pphlo_dialect",
     ],
 )
+
+# Equivalent of configure_lit_site_cfg from CMakeLists.txt.
+expand_template(
+    name = "lit_site_cfg_py_gen",
+    testonly = True,
+    out = "lit.site.cfg.py",
+    substitutions = {
+        "@LIT_SITE_CFG_IN_HEADER@": "# Autogenerated, do not edit.",
+        "@LLVM_TOOLS_DIR@": package_path("@llvm-project//llvm:BUILD"),
+        "\"@PPHLO_TOOLS_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], 'spulib', 'libspu', 'compiler', 'tools')",
+        "\"@PPHLO_SOURCE_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], 'spulib')",
+    },
+    template = "lit.site.cfg.py.in",
+)
+
+# Equivalent of add_lit_testsuite from CMakeLists.txt.
+[
+    lit_test(
+        name = "%s.test" % src,
+        size = "small",
+        srcs = [src],
+        data = [
+            "lit.cfg.py",
+            "lit.site.cfg.py",
+            "//libspu/compiler/tools:mlir-pphlo-opt",
+            "@llvm-project//llvm:FileCheck",
+        ] + glob(["%s.bc" % src]),
+        tags = ["pphlo_tests"],
+    )
+    for src in glob(["**/**/*.mlir"])
+]
+
+test_suite(
+    name = "pphlo_tests",
+    tags = ["pphlo_tests"],
+)

libspu/compiler/tests/lit.cfg.py

@@ -0,0 +1,36 @@
+"""Lit configuration to drive test in this repo."""
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2022 The StableHLO Authors.
+# 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.
+# -*- Python -*-
+# pylint: disable=undefined-variable
+import os
+import lit.formats
+from lit.llvm import llvm_config
+
+# Populate Lit configuration with the minimal required metadata.
+# Some metadata is populated in lit.site.cfg.py.in.
+config.name = 'PPHLO_TESTS_SUITE'
+config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
+config.suffixes = ['.mlir']
+config.test_source_root = os.path.dirname(__file__)
+# Make LLVM and StableHLO tools available in RUN directives
+tools = [
+    'FileCheck',
+    'mlir-pphlo-opt',
+]
+tool_dirs = [
+    config.llvm_tools_dir,
+    config.pphlo_tools_dir,
+]
+llvm_config.add_tool_substitutions(tools, tool_dirs)

libspu/compiler/tests/lit.site.cfg.py.in

@@ -0,0 +1,20 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2022 The StableHLO Authors.
+# Copyright 2023 SecretFlow Authors.
+# 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.
+@LIT_SITE_CFG_IN_HEADER@
+import lit.llvm
+lit.llvm.initialize(lit_config, config)
+config.llvm_tools_dir = "@LLVM_TOOLS_DIR@"
+config.pphlo_tools_dir = "@PPHLO_TOOLS_DIR@"
+lit_config.load_config(config, "@PPHLO_SOURCE_DIR@" + "/libspu/compiler/tests/lit.cfg.py")

libspu/device/io.cc

@@ -178,13 +178,8 @@ ColocatedIo::ColocatedIo(SPUContext *sctx) : sctx_(sctx) {}
 
 void ColocatedIo::hostSetVar(const std::string &name, const PtBufferView &bv,
                              Visibility vtype) {
-  if (vtype == VIS_PRIVATE) {
-    // handle SECRET/PRIVATE compiler/runtime trick.
-    unsynced_[name] = {convertToNdArray(bv), VIS_SECRET,
-                       static_cast<int>(sctx_->lctx()->Rank())};
-  } else {
-    unsynced_[name] = {convertToNdArray(bv), vtype};
-  }
+  unsynced_[name] = {convertToNdArray(bv), vtype,
+                     static_cast<int>(sctx_->lctx()->Rank())};
 }
 
 NdArrayRef ColocatedIo::hostGetVar(const std::string &name) const {

libspu/device/io_test.cc

@@ -143,13 +143,14 @@ TEST(ColocatedIoTest, PrivateWorks) {
     SPUContext sctx(hconf, lctx);
     ColocatedIo cio(&sctx);
 
-    // WHEN
+    // when experimental_enable_colocated_optimization is on,
+    // set secret with colocated io gets a prvate result
     if (lctx->Rank() == 0) {
       cio.hostSetVar("x", xt::xarray<int>{{1, -2, 3, 0}},
-                     Visibility::VIS_PRIVATE);
+                     Visibility::VIS_SECRET);
     } else if (lctx->Rank() == 1) {
       cio.hostSetVar("y", xt::xarray<float>{{1, -2, 3, 0}},
-                     Visibility::VIS_PRIVATE);
+                     Visibility::VIS_SECRET);
     }
     cio.sync();
 

libspu/mpc/aby3/arithmetic.cc

@@ -644,7 +644,7 @@ NdArrayRef MatMulAA::proc(KernelEvalContext* ctx, const NdArrayRef& x,
 
 #ifdef CUDA_ENABLED
   // FIXME: better heuristic?
-  if (spu::cuda::hasGPUDevice() && M * N <= 20000 && field == FM64) {
+  if (!spu::cuda::hasGPUDevice() || M * N <= 20000 || field != FM64) {
 #endif
     auto x1 = getFirstShare(x);
     auto x2 = getSecondShare(x);

spu/utils/frontend.py

@@ -88,7 +88,8 @@ def _argnames_partial_except(fn, static_argnames, kwargs):
     return functools.partial(fn, **static_kwargs), kwargs
 
 
-@cached(cache=LRUCache(maxsize=128), key=_jax_compilation_key)
+# FIXME: Figure out a proper way to hash lambda functions
+# @cached(cache=LRUCache(maxsize=128), key=_jax_compilation_key)
 def _jax_compilation(
     fn: Callable, static_argnums, static_argnames, args: List, kwargs: Dict
 ):