diff --git a/bazel/glob_lit_test.bzl b/bazel/glob_lit_test.bzl deleted file mode 100644 index e64cdc64..00000000 --- a/bazel/glob_lit_test.bzl +++ /dev/null @@ -1,138 +0,0 @@ -# Copyright 2022 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. - -# Test definitions for Lit, the LLVM test runner. -# -# This is reusing the LLVM Lit test runner in the interim until the new build -# rules are upstreamed. -# TODO(b/136126535): remove this custom rule. -"""Lit runner globbing test -""" - -load("@bazel_skylib//lib:paths.bzl", "paths") - -# Default values used by the test runner. -_default_test_file_exts = ["mlir", ".pbtxt", ".td"] -_default_driver = "@llvm-project//mlir:run_lit.sh" -_default_size = "small" -_default_tags = [] - -# These are patterns which we should never match, for tests, subdirectories, or -# test input data files. -_ALWAYS_EXCLUDE = [ - "**/LICENSE.txt", - "**/README.txt", - "**/lit.local.cfg", - # Exclude input files that have spaces in their names, since bazel - # cannot cope with such "targets" in the srcs list. - "**/* *", - "**/* */**", -] - -def _run_lit_test(name, data, size, tags, driver, features, exec_properties): - # Remove the default_driver from the data: it does not exist as a file and is - # just a placeholder from the copybara rewrite. - data = [d for d in data if d != _default_driver] - - # Disable tests on windows for now, to enable testing rest of all xla and mlir. - native.py_test( - name = name, - srcs = ["@llvm-project//llvm:lit"], - tags = tags + ["no_pip", "no_windows"], - args = [ - "libspu/compiler/tests/" + paths.basename(data[-1]) + " -v", - ] + features, - data = data + [ - "@llvm-project//llvm:FileCheck", - "@llvm-project//llvm:count", - "@llvm-project//llvm:not", - ], - size = size, - main = "lit.py", - exec_properties = exec_properties, - ) - -def glob_lit_tests( - exclude = [], - test_file_exts = _default_test_file_exts, - default_size = _default_size, - size_override = {}, - data = [], - per_test_extra_data = {}, - default_tags = _default_tags, - tags_override = {}, - driver = _default_driver, - features = [], - exec_properties = {}): - """Creates all plausible Lit tests (and their inputs) under this directory. - - Args: - exclude: [str], paths to exclude (for tests and inputs). - test_file_exts: [str], extensions for files that are tests. - default_size: str, the test size for targets not in "size_override". - size_override: {str: str}, sizes to use for specific tests. - data: [str], additional input data to the test. - per_test_extra_data: {str: [str]}, extra data to attach to a given file. - default_tags: [str], additional tags to attach to the test. - tags_override: {str: str}, tags to add to specific tests. - driver: str, label of the driver shell script. - Note: use of a custom driver is not currently supported - and specifying a default driver will abort the tests. - features: [str], list of extra features to enable. - exec_properties: a dictionary of properties to pass on. - """ - - # Ignore some patterns by default for tests and input data. - exclude = _ALWAYS_EXCLUDE + exclude - - tests = native.glob( - ["*." + ext for ext in test_file_exts], - exclude = exclude, - ) - - # Run tests individually such that errors can be attributed to a specific - # failure. - for curr_test in tests: - # Instantiate this test with updated parameters. - lit_test( - name = curr_test, - data = data + per_test_extra_data.get(curr_test, []), - size = size_override.get(curr_test, default_size), - tags = default_tags + tags_override.get(curr_test, []), - driver = driver, - features = features, - exec_properties = exec_properties, - ) - -def lit_test( - name, - data = [], - size = _default_size, - tags = _default_tags, - driver = _default_driver, - features = [], - exec_properties = {}): - """Runs test files under lit. - - Args: - name: str, the name of the test. - data: [str], labels that should be provided as data inputs. - size: str, the size of the test. - tags: [str], tags to attach to the test. - driver: str, label of the driver shell script. - Note: use of a custom driver is not currently supported - and specifying a default driver will abort the tests. - features: [str], list of extra features to enable. - """ - _run_lit_test(name + ".test", data + [name], size, tags, driver, features, exec_properties) diff --git a/libspu/compiler/tests/BUILD.bazel b/libspu/compiler/tests/BUILD.bazel index 1649fc05..e9436aee 100644 --- a/libspu/compiler/tests/BUILD.bazel +++ b/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"], +) diff --git a/libspu/compiler/tests/lit.cfg b/libspu/compiler/tests/lit.cfg deleted file mode 100644 index 5b10873f..00000000 --- a/libspu/compiler/tests/lit.cfg +++ /dev/null @@ -1,73 +0,0 @@ -"""Lit configuration to drive test in this repo.""" - -# -*- Python -*- -# pylint: disable=undefined-variable - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import os - -import lit.formats -import lit.llvm -from lit.llvm.subst import ToolSubst -import lit.util - -# Configuration file for the 'lit' test runner. -real_test_srcdir = os.environ['TEST_SRCDIR'] -external_srcdir = real_test_srcdir - -# name: The name of this test suite. -config.name = 'MLIR_PPHLO_OPT' -config.llvm_tools_dir = os.path.join(external_srcdir, 'llvm-project', 'llvm') -config.mlir_obj_root = os.path.join(real_test_srcdir) -config.mlir_tools_dir = os.path.join(external_srcdir, 'llvm-project', 'mlir') -# test_source_root: The root path where tests are located. -config.test_source_root = os.path.dirname(__file__) -# test_exec_root: The root path where tests should be run. -config.test_exec_root = os.path.join( - real_test_srcdir, 'spulib', 'libspu', 'compiler', 'tests' -) -config.mlir_pphlo_tools_dir = os.path.join( - real_test_srcdir, 'spulib', 'libspu', 'compiler', 'tools' -) - -# suffixes: A list of file extensions to treat as test files. -config.suffixes = ['.mlir'] - -test_dir = os.environ['TEST_TARGET'] -test_dir = test_dir.strip('/').rsplit(':', 1)[0] -config.mlir_test_dir = os.path.join( - real_test_srcdir, os.environ['TEST_WORKSPACE'], test_dir -) - -lit.llvm.initialize(lit_config, config) - -from lit.llvm import llvm_config - -config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell) - -config.substitutions.append(('%PATH%', config.environment['PATH'])) - -llvm_config.with_system_environment(['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP']) - -llvm_config.use_default_substitutions() - -# excludes: A list of directories to exclude from the testsuite. The 'Inputs' -# subdirectories contain auxiliary inputs for various tests in their parent -# directories. -config.excludes = ['Inputs', 'Examples', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt'] - -# Tweak the PATH to include the tools dir. -llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True) - -tool_dirs = [ - config.mlir_pphlo_tools_dir, - config.llvm_tools_dir, -] -tools = [ - 'mlir-pphlo-opt', -] - -llvm_config.add_tool_substitutions(tools, tool_dirs) diff --git a/libspu/compiler/tests/lit.cfg.py b/libspu/compiler/tests/lit.cfg.py new file mode 100644 index 00000000..a5d000df --- /dev/null +++ b/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) diff --git a/libspu/compiler/tests/lit.site.cfg.py.in b/libspu/compiler/tests/lit.site.cfg.py.in new file mode 100644 index 00000000..f9073fd2 --- /dev/null +++ b/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") diff --git a/libspu/device/io.cc b/libspu/device/io.cc index 69d7674e..44bc4508 100644 --- a/libspu/device/io.cc +++ b/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(sctx_->lctx()->Rank())}; - } else { - unsynced_[name] = {convertToNdArray(bv), vtype}; - } + unsynced_[name] = {convertToNdArray(bv), vtype, + static_cast(sctx_->lctx()->Rank())}; } NdArrayRef ColocatedIo::hostGetVar(const std::string &name) const { diff --git a/libspu/device/io_test.cc b/libspu/device/io_test.cc index 12fb2f64..b67642e4 100644 --- a/libspu/device/io_test.cc +++ b/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{{1, -2, 3, 0}}, - Visibility::VIS_PRIVATE); + Visibility::VIS_SECRET); } else if (lctx->Rank() == 1) { cio.hostSetVar("y", xt::xarray{{1, -2, 3, 0}}, - Visibility::VIS_PRIVATE); + Visibility::VIS_SECRET); } cio.sync(); diff --git a/libspu/mpc/aby3/arithmetic.cc b/libspu/mpc/aby3/arithmetic.cc index 2be30c79..a23749ff 100644 --- a/libspu/mpc/aby3/arithmetic.cc +++ b/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); diff --git a/spu/utils/frontend.py b/spu/utils/frontend.py index 097c3a4e..de7e04c8 100644 --- a/spu/utils/frontend.py +++ b/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 ):