From 263d7b97c5ade43786229c3220ffafd30529ae9b Mon Sep 17 00:00:00 2001 From: tran-go Date: Thu, 16 Jul 2026 19:43:57 +0800 Subject: [PATCH 1/2] Add block arithmetic operations --- openfhe_numpy/cpp/lib/numpy_bindings.cpp | 9 + openfhe_numpy/cpp/lib/numpy_enc_matrix.cpp | 302 ++++++++++--- .../cpp/lib/numpy_helper_functions.cpp | 3 +- openfhe_numpy/cpp/lib/numpy_utils.cpp | 107 +++-- openfhe_numpy/operations/__init__.py | 17 +- openfhe_numpy/operations/arithmetic_utils.py | 144 ++++++ openfhe_numpy/operations/block_arithmetic.py | 411 ++++++++++++++++++ .../operations/block_arithmetic_utils.py | 375 ++++++++++++++++ openfhe_numpy/operations/broadcast.py | 1 - openfhe_numpy/operations/crypto_helper.py | 115 ++++- openfhe_numpy/operations/matrix_api.py | 73 ++-- openfhe_numpy/operations/matrix_arithmetic.py | 98 +---- openfhe_numpy/tensor/block_ctarray.py | 2 + openfhe_numpy/tensor/block_tensor.py | 12 +- openfhe_numpy/tensor/ctarray.py | 32 +- openfhe_numpy/tensor/ptarray.py | 2 + openfhe_numpy/tensor/tensor.py | 51 +-- 17 files changed, 1472 insertions(+), 282 deletions(-) create mode 100644 openfhe_numpy/operations/arithmetic_utils.py create mode 100644 openfhe_numpy/operations/block_arithmetic.py create mode 100644 openfhe_numpy/operations/block_arithmetic_utils.py diff --git a/openfhe_numpy/cpp/lib/numpy_bindings.cpp b/openfhe_numpy/cpp/lib/numpy_bindings.cpp index 1fc4806..1f34220 100644 --- a/openfhe_numpy/cpp/lib/numpy_bindings.cpp +++ b/openfhe_numpy/cpp/lib/numpy_bindings.cpp @@ -191,6 +191,15 @@ void bind_matrix_funcs(py::module& m) { py::arg("numCols"), py::arg("numRepeats")); + m.def("EvalMatMulSquare", + [](const std::vector& matrixA, ConstCiphertext& matrixB, uint32_t numCols) + { + return EvalMatMulSquare(matrixA, matrixB, numCols); + }, + py::arg("matrixA"), + py::arg("matrixB"), + py::arg("numCols")); + m.def("EvalMatMulSquare", [](ConstCiphertext& matrixA, ConstCiphertext& matrixB, uint32_t numCols) { diff --git a/openfhe_numpy/cpp/lib/numpy_enc_matrix.cpp b/openfhe_numpy/cpp/lib/numpy_enc_matrix.cpp index cf6cf37..5ca45a6 100644 --- a/openfhe_numpy/cpp/lib/numpy_enc_matrix.cpp +++ b/openfhe_numpy/cpp/lib/numpy_enc_matrix.cpp @@ -28,8 +28,9 @@ // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. //============================================================================== -#include "numpy_enc_matrix.h" +#include +#include "numpy_enc_matrix.h" #include "numpy_utils.h" using namespace lbcrypto; @@ -46,23 +47,35 @@ namespace openfhe_numpy { * @param numRepeats Optional offset used by PHI and PSI types. * @return std::vector List of rotation indices to be used for EvalRotateKeyGen. **/ -static std::vector GenLinTransIndices(uint32_t numCols, LinTransType type, uint32_t numRepeats = 0) { - if (numCols < 0) { +static std::vector GenLinTransIndices( + uint32_t numCols, + LinTransType type, + uint32_t numRepeats = 0 +) { + if (numCols == 0) { OPENFHE_THROW("numCols must be positive"); } - if (numCols > std::numeric_limits::max() / 2 || // conservative upper bound - numRepeats < 0 || numRepeats > std::numeric_limits::max() / 2) { + if (numCols > static_cast(std::numeric_limits::max() / 2) || + numRepeats > static_cast(std::numeric_limits::max() / 2)) { OPENFHE_THROW("numCols or numRepeats too large"); } + const int32_t nCols = static_cast(numCols); + const int32_t repeats = static_cast(numRepeats); + + if (numCols > static_cast(std::numeric_limits::max() / numCols)) { + OPENFHE_THROW("numCols * numCols is too large"); + } + + const int32_t blockSize = nCols * nCols; + std::vector rotationIndices; - int32_t nCols = static_cast(numCols); switch (type) { case LinTransType::SIGMA: // Generate indices from -numCols to numCols - 1 - rotationIndices.reserve(2*nCols); + rotationIndices.reserve(2 * nCols); for (int32_t k = -nCols; k < nCols; ++k) { rotationIndices.push_back(k); } @@ -70,9 +83,12 @@ static std::vector GenLinTransIndices(uint32_t numCols, LinTransType ty case LinTransType::TAU: // Generate indices: 0, numCols, 2*numCols, ..., (numCols-1)*numCols - rotationIndices.reserve(nCols); + rotationIndices.reserve(2 * nCols - 1); for (int32_t k = 0; k < nCols; ++k) { rotationIndices.push_back(nCols * k); + if (k != 0) { + rotationIndices.push_back(nCols * k - blockSize); + } } break; @@ -80,18 +96,22 @@ static std::vector GenLinTransIndices(uint32_t numCols, LinTransType ty // Generate indices: numRepeats, numRepeats - numCols rotationIndices.reserve(2); for (int32_t k = 0; k < 2; ++k) { - rotationIndices.push_back(numRepeats - k * nCols); + rotationIndices.push_back(repeats - k * nCols); } break; case LinTransType::PSI: // Generate a single index based on offset - rotationIndices.push_back(nCols * numRepeats); + rotationIndices.reserve(2); + rotationIndices.push_back(nCols * repeats); + if (repeats != 0) { + rotationIndices.push_back(nCols * repeats - blockSize); + } break; case LinTransType::TRANSPOSE: // Generate indices for transposing a square matrix via diagonals - rotationIndices.reserve(2*nCols); + rotationIndices.reserve(2 * nCols); for (int32_t k = -nCols + 1; k < nCols; ++k) { rotationIndices.push_back((nCols - 1) * k); } @@ -99,7 +119,6 @@ static std::vector GenLinTransIndices(uint32_t numCols, LinTransType ty default: OPENFHE_THROW("Linear transformation is undefined"); - break; } return rotationIndices; @@ -189,7 +208,7 @@ Ciphertext EvalMultMatVec(std::shared_ptr& ctVector, ConstCiphertext& ctMatrix) { - if (numCols < 0) { + if (numCols <= 0) { OPENFHE_THROW("numCols must be positive"); } @@ -295,48 +314,192 @@ std::vector EvalLinTransSigma(const std::vector& vector, uint32_ * @param numCols The number of columns in the transformation matrix. * @return The resulting ciphertext after the transformation. */ + +namespace { + +void ValidateCompactSquareTransform(size_t slots, uint32_t numCols) { + if (numCols == 0) { + OPENFHE_THROW("numCols must be positive"); + } + + const size_t d = static_cast(numCols); + const size_t blockSize = d * d; + + if (slots < blockSize || slots % blockSize != 0) { + OPENFHE_THROW("slots must be a multiple of numCols * numCols"); + } +} + +// The ciphertext holds slots/blockSize replicated copies of the matrix +// (see EncodeMatrix), so the mask must repeat every blockSize entries -- +// matching the tiling already done by GenSigmaDiag/GenTauDiag/GenPhiDiag -- +// otherwise every replica after the first is left zeroed out. +std::vector GenCompactRowShiftDiag(size_t slots, uint32_t numCols, uint32_t k, bool wrapPart) { + ValidateCompactSquareTransform(slots, numCols); + + const size_t d = static_cast(numCols); + const size_t kk = static_cast(k); + const size_t n = d * d; + + if (kk >= d) { + OPENFHE_THROW("row-shift index k is out of range"); + } + + std::vector diag(slots, 0.0); + + for (size_t t = 0; t < slots / n; ++t) { + const size_t base = t * n; + for (size_t row = 0; row < d; ++row) { + const bool wraps = (row + kk >= d); + if (wraps != wrapPart) { + continue; + } + + for (size_t col = 0; col < d; ++col) { + diag[base + row * d + col] = 1.0; + } + } + } + + return diag; +} + +std::vector GenCompactTauDiag(size_t slots, uint32_t numCols, uint32_t col, bool wrapPart) { + ValidateCompactSquareTransform(slots, numCols); + + const size_t d = static_cast(numCols); + const size_t j = static_cast(col); + const size_t n = d * d; + + if (j >= d) { + OPENFHE_THROW("Tau column index is out of range"); + } + + std::vector diag(slots, 0.0); + + for (size_t t = 0; t < slots / n; ++t) { + const size_t base = t * n; + for (size_t row = 0; row < d; ++row) { + const bool wraps = (row + j >= d); + if (wraps == wrapPart) { + diag[base + row * d + j] = 1.0; + } + } + } + + return diag; +} + +std::vector EvalCompactRowShift(const std::vector& vector, uint32_t numCols, uint32_t k) { + const size_t slots = vector.size(); + ValidateCompactSquareTransform(slots, numCols); + + const int32_t d = static_cast(numCols); + const int32_t kk = static_cast(k); + const int32_t blockSize = d * d; + + std::vector result(slots, 0.0); + + auto addMaskedRotation = [&](int32_t shift, bool wrapPart) { + std::vector diag = GenCompactRowShiftDiag(slots, numCols, k, wrapPart); + const int32_t signedSlots = static_cast(slots); + const int32_t normalizedShift = ((shift % signedSlots) + signedSlots) % signedSlots; + + for (size_t i = 0; i < slots; ++i) { + result[i] += vector[(i + normalizedShift) % slots] * diag[i]; + } + }; + + addMaskedRotation(d * kk, false); + + if (kk != 0) { + addMaskedRotation(d * kk - blockSize, true); + } + + return result; +} + +} // namespace + Ciphertext EvalLinTransTau(ConstCiphertext& ctVector, uint32_t numCols) { - CryptoContext cc = ctVector->GetCryptoContext(); - bool flag = true; + CryptoContext cc = ctVector->GetCryptoContext(); + const size_t slots = cc->GetEncodingParams()->GetBatchSize(); + + ValidateCompactSquareTransform(slots, numCols); + + const int32_t nCols = static_cast(numCols); + const int32_t blockSize = nCols * nCols; + + bool flag = true; Ciphertext ctResult; - const size_t slots = cc->GetEncodingParams()->GetBatchSize(); + for (uint32_t k = 0; k < numCols; ++k) { + const int32_t kk = static_cast(k); - for (size_t k = 0; k < numCols; ++k) { - size_t shift = (static_cast(numCols) * k) % slots; - Ciphertext ctRotated = cc->EvalRotate(ctVector, shift); - Plaintext ptDiagonal = cc->MakeCKKSPackedPlaintext(GenTauDiag(slots, numCols, k)); - Ciphertext ctProd = cc->EvalMult(ctRotated, ptDiagonal); + Plaintext ptNoWrap = cc->MakeCKKSPackedPlaintext( + GenCompactTauDiag(slots, numCols, k, false) + ); + Ciphertext ctNoWrap = cc->EvalMult( + cc->EvalRotate(ctVector, nCols * kk), + ptNoWrap + ); if (flag) { - ctResult = ctProd; - flag = false; + ctResult = ctNoWrap; + flag = false; + } else { + cc->EvalAddInPlace(ctResult, ctNoWrap); } - else { - cc->EvalAddInPlace(ctResult, ctProd); + + if (k != 0) { + Plaintext ptWrap = cc->MakeCKKSPackedPlaintext( + GenCompactTauDiag(slots, numCols, k, true) + ); + Ciphertext ctWrap = cc->EvalMult( + cc->EvalRotate(ctVector, nCols * kk - blockSize), + ptWrap + ); + cc->EvalAddInPlace(ctResult, ctWrap); } } + return ctResult; } - /** * @brief LinTrans on a packed matrix */ std::vector EvalLinTransTau(const std::vector& vector, uint32_t numCols) { const size_t slots = vector.size(); - std::vector result(slots, 0); + ValidateCompactSquareTransform(slots, numCols); - for (size_t k = 0; k < numCols; ++k) { - size_t shift = (numCols * k) % slots; - std::vector diag = GenTauDiag(slots, numCols, k); - for (size_t i = 0; i < slots; ++i) { - result[i] += vector[(i + shift) % slots] * diag[i]; + const int32_t nCols = static_cast(numCols); + const int32_t blockSize = nCols * nCols; + + std::vector result(slots, 0.0); + + for (uint32_t k = 0; k < numCols; ++k) { + const int32_t kk = static_cast(k); + + auto addMaskedRotation = [&](int32_t shift, bool wrapPart) { + std::vector diag = GenCompactTauDiag(slots, numCols, k, wrapPart); + const int32_t signedSlots = static_cast(slots); + const int32_t normalizedShift = ((shift % signedSlots) + signedSlots) % signedSlots; + + for (size_t i = 0; i < slots; ++i) { + result[i] += vector[(i + normalizedShift) % slots] * diag[i]; + } + }; + + addMaskedRotation(nCols * kk, false); + + if (k != 0) { + addMaskedRotation(nCols * kk - blockSize, true); } } + return result; } - Ciphertext EvalLinTransTau(PrivateKey& secretKey, ConstCiphertext& ciphertext, uint32_t numCols) { @@ -355,26 +518,40 @@ Ciphertext EvalLinTransTau(PrivateKey& secretKey, * * @param numCols The number of padded cols in the encoded matrix */ -Ciphertext EvalLinTransPhi(ConstCiphertext& ctVector, uint32_t numCols, uint32_t numRepeats) { +Ciphertext EvalLinTransPhi( + ConstCiphertext& ctVector, + uint32_t numCols, + uint32_t numRepeats +) { + if (numCols == 0) { + OPENFHE_THROW("numCols must be positive"); + } - CryptoContext cc = ctVector->GetCryptoContext(); + CryptoContext cc = ctVector->GetCryptoContext(); const size_t slots = cc->GetEncodingParams()->GetBatchSize(); - bool flag = true; + const int32_t nCols = static_cast(numCols); + const int32_t repeats = static_cast(numRepeats); + + bool flag = true; Ciphertext ctResult; - for (size_t i = 0; i < 2; ++i) { - int32_t rotateIdx = numRepeats - i * numCols; + for (int32_t i = 0; i < 2; ++i) { + const int32_t rotateIdx = repeats - i * nCols; + Plaintext ptDiagonal = cc->MakeCKKSPackedPlaintext( - GenPhiDiag(slots, numCols, numRepeats, i)); - Ciphertext ctRotated = cc->EvalRotate(ctVector, rotateIdx); - Ciphertext ctProd = cc->EvalMult(ctRotated, ptDiagonal); + GenPhiDiag(slots, numCols, repeats, i) + ); + + Ciphertext ctRotated = cc->EvalRotate(ctVector, rotateIdx); + Ciphertext ctProd = cc->EvalMult(ctRotated, ptDiagonal); + if (flag) { ctResult = ctProd; - flag = false; - } - else + flag = false; + } else { cc->EvalAddInPlace(ctResult, ctProd); + } } return ctResult; @@ -423,14 +600,39 @@ Ciphertext EvalLinTransPsi(PrivateKey& secretKey, const std::vector EvalLinTransPsi(const std::vector& input, uint32_t numCols, uint32_t numRepeats) { - return RotateVector>(input, numCols * numRepeats); + return EvalCompactRowShift(input, numCols, numRepeats); } - Ciphertext EvalLinTransPsi(ConstCiphertext& ctVector, uint32_t numCols, uint32_t numRepeats) { CryptoContext cc = ctVector->GetCryptoContext(); - return cc->EvalRotate(ctVector, numCols * numRepeats); -} + const size_t slots = cc->GetEncodingParams()->GetBatchSize(); + ValidateCompactSquareTransform(slots, numCols); + + const int32_t nCols = static_cast(numCols); + const int32_t repeats = static_cast(numRepeats); + const int32_t blockSize = nCols * nCols; + + Plaintext ptNoWrap = cc->MakeCKKSPackedPlaintext( + GenCompactRowShiftDiag(slots, numCols, numRepeats, false) + ); + Ciphertext ctResult = cc->EvalMult( + cc->EvalRotate(ctVector, nCols * repeats), + ptNoWrap + ); + + if (numRepeats != 0) { + Plaintext ptWrap = cc->MakeCKKSPackedPlaintext( + GenCompactRowShiftDiag(slots, numCols, numRepeats, true) + ); + Ciphertext ctWrap = cc->EvalMult( + cc->EvalRotate(ctVector, nCols * repeats - blockSize), + ptWrap + ); + cc->EvalAddInPlace(ctResult, ctWrap); + } + + return ctResult; +} /** * @brief Multiplies two square matrices: CT x CT. * Implementation is based on https://eprint.iacr.org/2018/1041) @@ -549,7 +751,7 @@ Ciphertext EvalTranspose(PrivateKey& secretKey, } Ciphertext EvalTranspose(ConstCiphertext& ciphertext, uint32_t numCols) { try { - if (numCols < 0) { + if (numCols <= 0) { OPENFHE_THROW("numCols must be positive"); } CryptoContext cc = ciphertext->GetCryptoContext(); @@ -559,7 +761,7 @@ Ciphertext EvalTranspose(ConstCiphertext& ciphertext, uint32 const int32_t nCols = static_cast(numCols); for (int32_t index = -nCols + 1; index < nCols; ++index) { - int32_t rotationIndex = (numCols - 1) * index; + int32_t rotationIndex = (nCols - 1) * index; auto diagonalVector = GenTransposeDiag(slots, numCols, index); auto ptDiagonal = cc->MakeCKKSPackedPlaintext(diagonalVector); auto ctRotated = cc->EvalRotate(ciphertext, rotationIndex); diff --git a/openfhe_numpy/cpp/lib/numpy_helper_functions.cpp b/openfhe_numpy/cpp/lib/numpy_helper_functions.cpp index b5b50e8..ed7e86c 100644 --- a/openfhe_numpy/cpp/lib/numpy_helper_functions.cpp +++ b/openfhe_numpy/cpp/lib/numpy_helper_functions.cpp @@ -85,7 +85,8 @@ std::vector MulMatVec(std::vector> mat, std::vector< if (m != k) OPENFHE_THROW("Mismatched vector sizes"); - std::vector result(m, 0); + std::vector result(n, 0.0); + for (uint32_t i = 0; i < n; i++) { for (uint32_t j = 0; j < m; j++) { result[i] += mat[i][j] * vec[j]; diff --git a/openfhe_numpy/cpp/lib/numpy_utils.cpp b/openfhe_numpy/cpp/lib/numpy_utils.cpp index ceb8ebc..be2e817 100644 --- a/openfhe_numpy/cpp/lib/numpy_utils.cpp +++ b/openfhe_numpy/cpp/lib/numpy_utils.cpp @@ -48,27 +48,39 @@ Compute diagonals for the permutation matrix Sigma. B[i,j] = A[i, i +j] */ std::vector GenSigmaDiag(size_t slots, size_t numCols, int32_t k) { - // the cast below is necessary as we don't want to mix signed and unsigned integers in calculations - int32_t nCols = static_cast(numCols); - int32_t n = nCols * nCols; - std::vector diag(slots, 0); - - if (k >= 0) { - for (int32_t i = 0; i < n; i++) { - int32_t tmp = i - nCols * k; - if ((0 <= tmp) && (tmp < nCols - k)) { - diag[i] = 1; - } - } + if (numCols == 0) { + OPENFHE_THROW("numCols must be positive"); } - else { - for (int32_t i = 0; i < n; i++) { - int32_t tmp = i - (nCols + k) * nCols; - if ((-k <= tmp) && (tmp < nCols)) { - diag[i] = 1; + + const int32_t d = static_cast(numCols); + const size_t n = numCols * numCols; + + if (slots < n || slots % n != 0) { + OPENFHE_THROW("slots must be a multiple of numCols * numCols"); + } + + std::vector diag(slots, 0.0); + + for (size_t t = 0; t < slots / n; ++t) { + const size_t base = t * n; + + if (k >= 0) { + for (int32_t i = 0; i < static_cast(n); ++i) { + int32_t tmp = i - d * k; + if ((0 <= tmp) && (tmp < d - k)) { + diag[base + i] = 1.0; + } + } + } else { + for (int32_t i = 0; i < static_cast(n); ++i) { + int32_t tmp = i - (d + k) * d; + if ((-k <= tmp) && (tmp < d)) { + diag[base + i] = 1.0; + } } } } + return diag; } @@ -79,14 +91,29 @@ u_[d.k][k + d*i] = 1 for all 0 <= i < d */ std::vector GenTauDiag(size_t totalSlots, size_t numCols, int32_t k) { - uint32_t n = numCols * numCols; - std::vector diag(totalSlots, 0); + if (numCols == 0) { + OPENFHE_THROW("numCols must be positive"); + } + + const size_t n = numCols * numCols; - for (uint32_t t = 0; t < totalSlots / n; t++) { - for (uint32_t i = 0; i < numCols; i++) { - diag[(t * n) + k + numCols * i] = 1; + if (totalSlots < n || totalSlots % n != 0) { + OPENFHE_THROW("slots must be a multiple of numCols * numCols"); + } + + if (k < 0 || static_cast(k) >= numCols) { + OPENFHE_THROW("Tau diagonal index k is out of range"); + } + + std::vector diag(totalSlots, 0.0); + + for (size_t t = 0; t < totalSlots / n; ++t) { + const size_t base = t * n; + for (size_t i = 0; i < numCols; ++i) { + diag[base + static_cast(k) + numCols * i] = 1.0; } } + return diag; } @@ -98,19 +125,35 @@ std::vector GenTauDiag(size_t totalSlots, size_t numCols, int32_t k) { *diagonal */ std::vector GenPhiDiag(size_t slots, size_t numCols, int32_t k, int type) { - uint32_t n = numCols * numCols; - std::vector diag(slots, 0); + if (numCols == 0) { + OPENFHE_THROW("numCols must be positive"); + } - if (type == 0) { - for (uint32_t i = 0; i < n; i++) - if ((i % numCols >= 0) && ((i % numCols) < numCols - k)) - diag[i] = 1; + const size_t d = numCols; + const size_t n = d * d; + + if (slots < n || slots % n != 0) { + OPENFHE_THROW("slots must be a multiple of numCols * numCols"); } - else { - for (uint32_t i = 0; i < n; i++) - if ((i % numCols >= numCols - k) && (i % numCols < numCols)) { - diag[i] = 1; + + std::vector diag(slots, 0.0); + + for (size_t t = 0; t < slots / n; ++t) { + const size_t base = t * n; + + if (type == 0) { + for (size_t i = 0; i < n; ++i) { + if ((i % d) < d - static_cast(k)) { + diag[base + i] = 1.0; + } + } + } else { + for (size_t i = 0; i < n; ++i) { + if ((i % d) >= d - static_cast(k)) { + diag[base + i] = 1.0; + } } + } } return diag; diff --git a/openfhe_numpy/operations/__init__.py b/openfhe_numpy/operations/__init__.py index 1a56888..ab3e395 100644 --- a/openfhe_numpy/operations/__init__.py +++ b/openfhe_numpy/operations/__init__.py @@ -36,7 +36,6 @@ """ from .broadcast import broadcast_to, generate_broadcast_key -from . import matrix_arithmetic # Import arithmetic operations from .matrix_api import ( @@ -46,14 +45,17 @@ dot, matmul, transpose, - pow, - cumulative_sum, + power, + cumsum, cumulative_reduce, sum, roll, mean, ) +from . import matrix_arithmetic +from . import block_arithmetic + # Import crypto context utilities from .crypto_helper import ( sum_row_keys, @@ -66,8 +68,11 @@ gen_accumulate_rows_key, gen_accumulate_cols_key, generate_slicing_key, + attach_block_matvec_keys, + attach_matvec_keys, ) + # Define public API __all__ = [ # Arithmetic and matrix operations @@ -77,8 +82,8 @@ "dot", "matmul", "transpose", - "pow", - "cumulative_sum", + "power", + "cumsum", "cumulative_reduce", "sum", "roll", @@ -97,4 +102,6 @@ "broadcast_to", "generate_broadcast_key", "generate_slicing_key", + "attach_matvec_keys", + "attach_block_matvec_keys", ] diff --git a/openfhe_numpy/operations/arithmetic_utils.py b/openfhe_numpy/operations/arithmetic_utils.py new file mode 100644 index 0000000..9de7990 --- /dev/null +++ b/openfhe_numpy/operations/arithmetic_utils.py @@ -0,0 +1,144 @@ +# ================================================================================== +# BSD 2-Clause License +# +# Copyright (c) 2014-2025, NJIT, Duality Technologies Inc. and other contributors +# +# All rights reserved. +# +# Author TPOC: contact@openfhe.org +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# ================================================================================== + +""" +Shared arithmetic helpers used by both plain (CTArray/PTArray) and block +tensor operations. +""" + +from __future__ import annotations + +from typing import Any + +from openfhe_numpy.tensor.tensor import BaseTensor +from openfhe_numpy.utils.errors import ONPIncompatibleShapeError, ONPValueError, ONPDimensionError +from openfhe_numpy.utils.packing import _is_row_major, _is_col_major + + +# ------------------------------------------------------------------------------ +# Arithmetic: Convention +# ------------------------------------------------------------------------------ +def _require( + condition: bool, + left: Any, + right: Any, + message: str, + *, + error_cls: type[Exception] = ONPIncompatibleShapeError, +) -> None: + """Raise ``error_cls`` when ``condition`` is false. + + By default this reports an ``ONPIncompatibleShapeError`` built from ``left`` + and ``right``. Pass ``error_cls`` (e.g. ``ONPValueError``) for non-shape + failures such as batch-size or packing-order mismatches; those are raised + with ``message`` only. + """ + if condition: + return + + if error_cls is ONPIncompatibleShapeError: + raise ONPIncompatibleShapeError(left, right, message) + + raise error_cls(message) + + +def _result_cls(a: BaseTensor, b: BaseTensor | None = None) -> type[BaseTensor]: + """Return the encrypted tensor class if either operand is encrypted.""" + if a.is_encrypted: + return type(a) + + if b is not None and b.is_encrypted: + return type(b) + + raise ONPValueError("Expected at least one encrypted tensor operand.") + + +# ------------------------------------------------------------------------------ +# Arithmetic: Validation +# ------------------------------------------------------------------------------ + + +def _require_matvec_order(matrix_order: Any, vector_order: Any) -> None: + """Raise ``ONPValueError`` unless the matrix-vector packing pair is supported. + + Supported pairs are ROW_MAJOR matrix @ COL_MAJOR vector, or COL_MAJOR matrix + @ ROW_MAJOR vector. + """ + supported = (_is_row_major(matrix_order) and _is_col_major(vector_order)) or ( + _is_col_major(matrix_order) and _is_row_major(vector_order) + ) + _require( + supported, + (matrix_order,), + (vector_order,), + "Block matvec requires ROW_MAJOR matrix @ COL_MAJOR vector, " + "or COL_MAJOR matrix @ ROW_MAJOR vector.", + error_cls=ONPValueError, + ) + + +def _get_matvec_key_name(matrix_order: Any) -> str: + """Return the key name required by matrix-vector multiplication.""" + if _is_row_major(matrix_order): + return "colkey" + + if _is_col_major(matrix_order): + return "rowkey" + + raise ONPValueError(f"Unsupported packing order: {matrix_order}") + + +# ------------------------------------------------------------------------------ +# Arithmetic: Numpy axis convention +# ------------------------------------------------------------------------------ +def _normalize_axis(axis, ndim: int) -> int: + """Normalize integer axis. + + Tuple axes and NumPy integer types are not supported. + """ + if type(axis) is not int: + raise ONPDimensionError(f"axis must be an integer, got {type(axis).__name__}.") + + if axis < -ndim or axis >= ndim: + raise ONPDimensionError(f"axis {axis} is out of bounds for tensor with {ndim} dimensions.") + + return axis + ndim if axis < 0 else axis + + +def _normalize_sum_axis(axis, ndim: int) -> int | None: + """Normalize a single reduction axis or None. + + Tuple axes are not supported. + """ + if axis is None: + return None + + return _normalize_axis(axis, ndim) diff --git a/openfhe_numpy/operations/block_arithmetic.py b/openfhe_numpy/operations/block_arithmetic.py new file mode 100644 index 0000000..f2a9834 --- /dev/null +++ b/openfhe_numpy/operations/block_arithmetic.py @@ -0,0 +1,411 @@ +# ================================================================================== +# BSD 2-Clause License +# +# Copyright (c) 2014-2025, NJIT, Duality Technologies Inc. and other contributors +# +# All rights reserved. +# +# Author TPOC: contact@openfhe.org +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# ================================================================================== + +""" +block_arithmetic.py + +This module implements block tensor arithmetic operations for encrypted tensors. +Block operations are registered through the tensor dispatch system and usually +delegate per-block work to CTArray operations. +""" + +# Standard library imports +from typing import Optional + +# Third-party imports +from ..tensor.block_ctarray import BlockCTArray +from .dispatch import register_tensor_function + +from ..utils.errors import ( + ONPNotImplementedError, + ONPValueError, + ONPDimensionError, +) +from .block_arithmetic_utils import ( + _eval_block_binary, + _eval_block_scalar, + _eval_scalar_block, + _eval_block_dot, + _eval_block_matmul, +) +from .matrix_arithmetic import ( + _pow, + _reduce_ct, + sum_ct, +) +from .arithmetic_utils import _normalize_axis, _normalize_sum_axis +from ..utils.matlib import _sum_terms + +# ============================================================================== +# Basic block arithmetic operations +# ============================================================================== + + +# ------------------------------------------------------------------------------ +# Addition Operations +# ------------------------------------------------------------------------------ +@register_tensor_function( + "add", [("BlockCTArray", "BlockCTArray"), ("BlockCTArray", "BlockPTArray")] +) +def add_block_ct(a, b): + """Add two block tensors.""" + return _eval_block_binary(a, b, "add") + + +@register_tensor_function("add", [("BlockCTArray", "scalar")]) +def add_block_ct_scalar(a, scalar): + """Add a scalar to a block ciphertext tensor.""" + return _eval_block_scalar(a, scalar, "add") + + +# ------------------------------------------------------------------------------ +# Subtraction Operations +# ------------------------------------------------------------------------------ +@register_tensor_function( + "subtract", + [ + ("BlockCTArray", "BlockCTArray"), + ("BlockCTArray", "BlockPTArray"), + ("BlockPTArray", "BlockCTArray"), + ], +) +def subtract_block_ct(a, b): + """Subtract two block tensors.""" + return _eval_block_binary(a, b, "subtract") + + +@register_tensor_function("subtract", [("BlockCTArray", "scalar")]) +def subtract_block_ct_scalar(a, scalar): + """Subtract a scalar from a block ciphertext tensor.""" + return _eval_block_scalar(a, scalar, "subtract") + + +@register_tensor_function("subtract", [("scalar", "BlockCTArray")]) +def subtract_scalar_block_ct(scalar, a): + """Subtract a block ciphertext tensor from a scalar.""" + return _eval_scalar_block(scalar, a, "subtract") + + +# ------------------------------------------------------------------------------ +# Multiplication Operations +# ------------------------------------------------------------------------------ +@register_tensor_function( + "multiply", [("BlockCTArray", "BlockCTArray"), ("BlockCTArray", "BlockPTArray")] +) +def multiply_block_ct(a, b): + """Multiply two block tensors element-wise.""" + return _eval_block_binary(a, b, "multiply") + + +@register_tensor_function("multiply", [("BlockCTArray", "scalar")]) +def multiply_block_ct_scalar(a, scalar): + """Multiply a block tensor by a scalar.""" + return _eval_block_scalar(a, scalar, "multiply") + + +# ============================================================================== +# Linear Algebra Operations +# ============================================================================== + + +# ------------------------------------------------------------------------------ +# Matrix Multiplication Operations +# ------------------------------------------------------------------------------ +@register_tensor_function( + "matmul", + [ + ("BlockCTArray", "BlockCTArray"), + ("BlockCTArray", "BlockPTArray"), + ("BlockPTArray", "BlockCTArray"), + ], +) +def matmul_block_ct(a, b): + """Perform block vector dot or block matrix multiplication.""" + return _eval_block_matmul(a, b) + + +# ------------------------------------------------------------------------------ +# Dot Product Operations +# ------------------------------------------------------------------------------ +@register_tensor_function("dot", [("BlockCTArray", "BlockCTArray")]) +def dot_block_ct(a, b): + """Compute dot product of two block vectors.""" + return _eval_block_dot(a, b) + + +# ------------------------------------------------------------------------------ +# Transpose Operations +# ------------------------------------------------------------------------------ +@register_tensor_function("transpose", [("BlockCTArray",)]) +def transpose_block_ct(a): + """Transpose a block ciphertext tensor.""" + if a.ndim == 1: + return a + + if a.ndim != 2: + raise ONPDimensionError(f"transpose requires 1D or 2D tensor, got {a.ndim}D.") + + grid_rows, grid_cols = a.grid_shape + blocks = [] + for j in range(grid_cols): + for i in range(grid_rows): + blocks.append(a.get_block(i, j).transpose()) + + return type(a)( + data=blocks, + grid_shape=(grid_cols, grid_rows), + block_shape=(a.block_shape[1], a.block_shape[0]), + original_shape=(a.original_shape[1], a.original_shape[0]), + batch_size=a.batch_size, + order=a.order, + ) + + +# ============================================================================== +# Advanced block operations +# ============================================================================== + + +# ------------------------------------------------------------------------------ +# Power Operations +# ------------------------------------------------------------------------------ +@register_tensor_function("power", [("BlockCTArray", "int")]) +def power_block_ct(a, exp): + """Raise a block ciphertext tensor to an integer power element-wise. + + Element-wise power (``a_i ** exp``) applies to any block shape; no square + logical/grid constraint is required. + """ + return _pow(a, exp) + + +# ------------------------------------------------------------------------------ +# Cumulative Operations +# ------------------------------------------------------------------------------ +@register_tensor_function( + "cumsum", + [("BlockCTArray",), ("BlockCTArray", "int"), ("BlockCTArray", "int", "bool")], +) +def cumsum_block_ct(obj, axis=None, keepdims=False): + """Compute cumulative sums for encrypted block tensors. + + Currently supports: + - 1D block tensors across multiple blocks. + - 2D block tensors only when the cumulative axis stays inside each block. + """ + if keepdims: + raise ONPNotImplementedError("Block cumsum does not support keepdims=True yet.") + + if obj.ndim == 1: + if axis is None: + axis = 0 + + axis = _normalize_axis(axis, obj.ndim) + + blocks = [] + offset = None + + for block in obj.data: + cumulative = block.cumsum(axis=0) + + if offset is not None: + cumulative = cumulative + offset + + blocks.append(cumulative) + + block_sum = block.sum() + offset = block_sum if offset is None else offset + block_sum + + return obj.clone(data=blocks) + + if obj.ndim == 2: + if axis is None: + raise ONPNotImplementedError("Block cumsum(axis=None) is not implemented yet.") + + axis = _normalize_axis(axis, obj.ndim) + + if axis == 0 and obj.grid_shape[0] != 1: + raise ONPNotImplementedError( + "Block cumsum(axis=0) across multiple block rows is not implemented yet." + ) + + if axis == 1 and obj.grid_shape[1] != 1: + raise ONPNotImplementedError( + "Block cumsum(axis=1) across multiple block columns is not implemented yet." + ) + + return obj.clone(data=[block.cumsum(axis=axis) for block in obj.data]) + + raise ONPDimensionError(f"cumsum requires 1D or 2D tensor, got {obj.ndim}D.") + + +@register_tensor_function( + "cumulative_reduce", + [("BlockCTArray",), ("BlockCTArray", "int"), ("BlockCTArray", "int", "bool")], +) +def cumulative_reduce_block_ct(a, axis=0, keepdims=False): + """Compute cumulative reductions inside each encrypted block.""" + if keepdims: + raise ONPNotImplementedError("Block cumulative_reduce does not support keepdims=True yet.") + + if a.ndim != 2: + raise ONPDimensionError(f"cumulative_reduce requires a 2D block tensor, got {a.ndim}D.") + + axis = _normalize_axis(axis, a.ndim) + + if axis == 0 and a.grid_shape[0] != 1: + raise ONPNotImplementedError( + "Block cumulative_reduce(axis=0) across multiple block rows is not implemented yet." + ) + + if axis == 1 and a.grid_shape[1] != 1: + raise ONPNotImplementedError( + "Block cumulative_reduce(axis=1) across multiple block columns is not implemented yet." + ) + + return a.clone(data=[_reduce_ct(block, axis, keepdims) for block in a.data]) + + +# ------------------------------------------------------------------------------ +# Sum Operations +# ------------------------------------------------------------------------------ +@register_tensor_function( + "sum", + [("BlockCTArray",), ("BlockCTArray", "int"), ("BlockCTArray", "int", "bool")], +) +def sum_block_ct(x: BlockCTArray, axis: Optional[int] = None, keepdims: bool = False): + """Sum encrypted block tensor elements.""" + if keepdims: + raise ONPNotImplementedError("Block sum does not support keepdims=True yet.") + + axis = _normalize_sum_axis(axis, x.ndim) + + if x.ndim == 1: + return _sum_terms(sum_ct(block, None, False) for block in x.data) + + if x.ndim != 2: + raise ONPDimensionError(f"sum requires 1D or 2D tensor, got {x.ndim}D.") + + if axis is None: + return _sum_terms(sum_ct(block, None, False) for block in x.data) + + if axis == 0: + grid_rows, grid_cols = x.grid_shape + blocks = [] + for j in range(grid_cols): + col_sum = sum_ct(x.get_block(0, j), axis=0, keepdims=False) + for i in range(1, grid_rows): + col_sum = col_sum + sum_ct(x.get_block(i, j), axis=0, keepdims=False) + blocks.append(col_sum) + + return type(x)( + data=blocks, + grid_shape=(grid_cols,), + block_shape=(x.block_shape[1],), + original_shape=(x.original_shape[1],), + batch_size=x.batch_size, + order=blocks[0].order, + ) + + if axis == 1: + grid_rows, grid_cols = x.grid_shape + blocks = [] + for i in range(grid_rows): + row_sum = sum_ct(x.get_block(i, 0), axis=1, keepdims=False) + for j in range(1, grid_cols): + row_sum = row_sum + sum_ct(x.get_block(i, j), axis=1, keepdims=False) + blocks.append(row_sum) + + return type(x)( + data=blocks, + grid_shape=(grid_rows,), + block_shape=(x.block_shape[0],), + original_shape=(x.original_shape[0],), + batch_size=x.batch_size, + order=blocks[0].order, + ) + + raise ONPValueError(f"Invalid axis {axis}.") + + +# ------------------------------------------------------------------------------ +# Mean Operations +# ------------------------------------------------------------------------------ +@register_tensor_function( + "mean", + [("BlockCTArray",), ("BlockCTArray", "int"), ("BlockCTArray", "int", "bool")], +) +def mean_block_ct(x: BlockCTArray, axis: Optional[int] = None, keepdims: bool = False): + """Compute the arithmetic mean for encrypted block tensors. + + - axis=None: mean over all entries. + - axis=0: mean down rows and return one value per column. + - axis=1: mean across columns and return one value per row. + + Limitation: + - keepdims=True is not supported for block tensors yet. + """ + if keepdims: + raise ONPNotImplementedError("Block mean does not support keepdims=True yet.") + + axis = _normalize_sum_axis(axis, x.ndim) + mean_x = sum_block_ct(x, axis, keepdims=False) + + if x.ndim == 1: + n = x.original_shape[0] + + elif x.ndim == 2: + nrows, ncols = x.original_shape + if axis is None: + n = nrows * ncols + elif axis == 0: + n = nrows + elif axis == 1: + n = ncols + else: + raise ONPDimensionError(f"Invalid axis {axis} for 2D block tensor.") + + else: + raise ONPDimensionError(f"mean requires 1D or 2D tensor, got {x.ndim}D.") + + return mean_x * (1.0 / n) + + +# ------------------------------------------------------------------------------ +# Rotation Operations +# ------------------------------------------------------------------------------ +@register_tensor_function( + "roll", + [("BlockCTArray", "int"), ("BlockCTArray", "int", "int")], +) +def roll_block_ct(x: BlockCTArray, shift: int, axis: Optional[int] = None) -> BlockCTArray: + """Block tensor roll is not implemented yet.""" + raise ONPNotImplementedError("Block roll is not implemented yet.") diff --git a/openfhe_numpy/operations/block_arithmetic_utils.py b/openfhe_numpy/operations/block_arithmetic_utils.py new file mode 100644 index 0000000..c71669d --- /dev/null +++ b/openfhe_numpy/operations/block_arithmetic_utils.py @@ -0,0 +1,375 @@ +# ================================================================================== +# BSD 2-Clause License +# +# Copyright (c) 2014-2025, NJIT, Duality Technologies Inc. and other contributors +# +# All rights reserved. +# +# Author TPOC: contact@openfhe.org +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# ================================================================================== + +from __future__ import annotations + +import operator +from typing import Any, Callable + +from openfhe_numpy import ArrayEncodingType +from openfhe_numpy.tensor.block_tensor import BlockFHETensor +from openfhe_numpy.tensor.block_ctarray import BlockCTArray +from openfhe_numpy.utils.errors import ONPNotImplementedError, ONPValueError +from openfhe_numpy.utils.typecheck import Number +from openfhe_numpy.utils.matlib import _sum_terms +from openfhe_numpy.operations.arithmetic_utils import ( + _result_cls, + _require, + _require_matvec_order, + _get_matvec_key_name, +) + + +# ------------------------------------------------------------------------------ +# Type aliases +# ------------------------------------------------------------------------------ + +BlockOp = Callable[[Any, Any], Any] + + +# ------------------------------------------------------------------------------ +# Supported element-wise operators +# ------------------------------------------------------------------------------ + +_OPS: dict[str, BlockOp] = { + "add": operator.add, + "subtract": operator.sub, + "multiply": operator.mul, +} + + +def _resolve_op(op_name: str) -> BlockOp: + """Return the blockwise Python operator for ``op_name``.""" + try: + return _OPS[op_name] + except KeyError: + supported = ", ".join(sorted(_OPS)) + raise ONPNotImplementedError( + f"Unsupported block operation {op_name!r}. Supported operations: {supported}." + ) from None + + +# ------------------------------------------------------------------------------ +# Generic block helpers +# ------------------------------------------------------------------------------ + + +def _build_block_result( + reference: BlockFHETensor, + blocks: list[Any], + *, + result_cls: type[BlockFHETensor] | None = None, + grid_shape: tuple[int, ...] | None = None, + block_shape: tuple[int, ...] | None = None, + original_shape: tuple[int, ...] | None = None, + order: Any | None = None, +) -> BlockFHETensor: + """Construct a block tensor result while preserving reference metadata by default.""" + cls = result_cls or type(reference) + + return cls( + data=blocks, + grid_shape=reference.grid_shape if grid_shape is None else grid_shape, + block_shape=reference.block_shape if block_shape is None else block_shape, + original_shape=reference.original_shape if original_shape is None else original_shape, + batch_size=reference.batch_size, + order=reference.order if order is None else order, + ) + + +# ------------------------------------------------------------------------------ +# Element-wise operations +# ------------------------------------------------------------------------------ + + +def _eval_block_binary(a: BlockFHETensor, b: BlockFHETensor, op_name: str) -> BlockFHETensor: + """Evaluate a blockwise binary operation on two block tensors.""" + + _require( + a.same_layout(b), + a.original_shape, + b.original_shape, + f"Block {op_name} requires identical block layout.", + ) + + op = _resolve_op(op_name) + blocks = [op(left, right) for left, right in zip(a.data, b.data)] + + return _build_block_result(a, blocks, result_cls=BlockCTArray) + + +def _eval_block_scalar(a: BlockFHETensor, scalar: Number, op_name: str) -> BlockFHETensor: + """Evaluate ``block_tensor op scalar`` blockwise.""" + op = _resolve_op(op_name) + blocks = [op(block, scalar) for block in a.data] + + return _build_block_result(a, blocks) + + +def _eval_scalar_block(scalar: Number, a: BlockFHETensor, op_name: str) -> BlockFHETensor: + """Evaluate ``scalar op block_tensor`` blockwise.""" + op = _resolve_op(op_name) + blocks = [op(scalar, block) for block in a.data] + + return _build_block_result(a, blocks) + + +# ------------------------------------------------------------------------------ +# Vector @ vector +# ------------------------------------------------------------------------------ + + +def _eval_block_dot(a: BlockFHETensor, b: BlockFHETensor) -> Any: + """Compute a block-vector inner product.""" + + if not a.same_layout(b): + raise ONPValueError("Incompatible layout") + + if not (a.is_encrypted or b.is_encrypted): + raise ONPNotImplementedError("Operation requires at least one encrypted operand.") + + return _sum_terms(left @ right for left, right in zip(a.data, b.data)) + + +# ------------------------------------------------------------------------------ +# Matrix @ vector +# ------------------------------------------------------------------------------ + + +def _assert_block_matvec_compatible(a: BlockFHETensor, b: BlockFHETensor) -> None: + """Validate block matrix-vector multiplication inputs.""" + _require( + a.ndim == 2 and b.ndim == 1, + a.original_shape, + b.original_shape, + "Block matvec requires a 2-D block matrix and a 1-D block vector.", + ) + + if not (a.is_encrypted or b.is_encrypted): + raise ONPNotImplementedError("Operation requires at least one encrypted operand.") + + _require( + a.original_shape[1] == b.original_shape[0], + a.original_shape, + b.original_shape, + "Block matvec dimension mismatch.", + ) + _require( + a.grid_shape[1] == b.grid_shape[0], + a.grid_shape, + b.grid_shape, + "Block matvec requires matching inner block dimension.", + ) + _require( + len(a.block_shape) == 2 and len(b.block_shape) == 1, + a.block_shape, + b.block_shape, + "Matrix blocks must be 2-D, vector blocks 1-D.", + ) + + _, block_cols = a.block_shape + + _require( + b.block_shape == (block_cols,), + a.block_shape, + b.block_shape, + f"Vector block_shape must be ({block_cols},) to match matrix block_shape={a.block_shape}.", + ) + _require( + a.batch_size == b.batch_size, + (a.batch_size,), + (b.batch_size,), + "Block matvec requires equal batch_size.", + error_cls=ONPValueError, + ) + _require_matvec_order(a.order, b.order) + + key_name = _get_matvec_key_name(a.order) + + for index in a.iter_block_indices(): + block = a.get_block(*index) + + # The evaluator always reads the summation key from the matrix block + # (encrypted or not), so validate it on every block regardless of state. + if key_name not in block.extra: + raise ONPValueError( + f"Matrix block {index} is missing extra[{key_name!r}]. " + "Call attach_matvec_keys(...) before block matrix-vector multiplication." + ) + + +def _eval_block_matvec(a: BlockFHETensor, b: BlockFHETensor) -> BlockFHETensor: + """Compute block matrix-vector multiplication. + + For a block matrix A and block vector x, this computes y[i] = sum_k A[i,k] @ x[k]. + """ + _assert_block_matvec_compatible(a, b) + + block_rows, _ = a.block_shape + grid_rows, grid_inner = a.grid_shape + + blocks = [ + _sum_terms(a.get_block(i, k) @ b.get_block(k) for k in range(grid_inner)) + for i in range(grid_rows) + ] + + return _build_block_result( + a, + blocks, + result_cls=BlockCTArray, + grid_shape=(grid_rows,), + block_shape=(block_rows,), + original_shape=(a.original_shape[0],), + order=a.order, + ) + + +# ------------------------------------------------------------------------------ +# Matrix @ matrix +# ------------------------------------------------------------------------------ + + +def _assert_block_matmul_compatible(a: BlockFHETensor, b: BlockFHETensor) -> None: + """Validate block matrix-matrix multiplication inputs.""" + _require( + a.ndim == 2 and b.ndim == 2, + a.original_shape, + b.original_shape, + "Block matmul requires two 2-D block matrices.", + ) + + if not (a.is_encrypted or b.is_encrypted): + raise ONPNotImplementedError("Operation requires at least one encrypted operand.") + + _require( + a.original_shape[1] == b.original_shape[0], + a.original_shape, + b.original_shape, + "Block matmul dimension mismatch.", + ) + _require( + a.grid_shape[1] == b.grid_shape[0], + a.grid_shape, + b.grid_shape, + "Block matmul requires matching inner block dimension.", + ) + _require( + a.block_shape == b.block_shape, + a.block_shape, + b.block_shape, + "Block matmul requires equal block_shape.", + ) + _require( + len(a.block_shape) == 2, + a.block_shape, + b.block_shape, + "Block matmul requires 2-D matrix blocks.", + ) + + block_rows, block_cols = a.block_shape + + _require( + block_rows == block_cols, + a.block_shape, + b.block_shape, + "Block matmul requires square blocks.", + ) + _require( + a.batch_size == b.batch_size, + (a.batch_size,), + (b.batch_size,), + "Block matmul requires equal batch_size.", + error_cls=ONPValueError, + ) + _require( + a.order == b.order == ArrayEncodingType.ROW_MAJOR, + (a.order,), + (b.order,), + "Block matmul currently requires both operands to use ROW_MAJOR packing.", + error_cls=ONPValueError, + ) + + +def _eval_block_matmat(a: BlockFHETensor, b: BlockFHETensor) -> BlockFHETensor: + """Compute block matrix-matrix multiplication. + + For block matrices A and B, this computes C[i,j] = sum_k A[i,k] @ B[k,j]. + + Required logical condition: + - a.shape = (m, k), b.shape = (k, n). + + Current implementation limitations: + - Both tensors must be 2-D block matrices. + - Blocks must have the same square block_shape. + - Both operands must use ROW_MAJOR packing. + - At least one operand must be encrypted. + + """ + _assert_block_matmul_compatible(a, b) + + grid_rows = a.grid_shape[0] + grid_inner = a.grid_shape[1] + grid_cols = b.grid_shape[1] + + blocks = [ + _sum_terms(a.get_block(i, k) @ b.get_block(k, j) for k in range(grid_inner)) + for i in range(grid_rows) + for j in range(grid_cols) + ] + + return _build_block_result( + a, + blocks, + result_cls=_result_cls(a, b), + grid_shape=(grid_rows, grid_cols), + block_shape=a.block_shape, + original_shape=(a.original_shape[0], b.original_shape[1]), + order=a.order, + ) + + +# ------------------------------------------------------------------------------ +# Matmul dispatch +# ------------------------------------------------------------------------------ + + +def _eval_block_matmul(a: BlockFHETensor, b: BlockFHETensor) -> Any: + """Dispatch block matmul by operand ranks.""" + if a.ndim == 1 and b.ndim == 1: + return _eval_block_dot(a, b) + + if a.ndim == 2 and b.ndim == 1: + return _eval_block_matvec(a, b) + + if a.ndim == 2 and b.ndim == 2: + return _eval_block_matmat(a, b) + + raise ONPNotImplementedError(f"Block matmul does not support ndim={a.ndim} @ ndim={b.ndim}.") diff --git a/openfhe_numpy/operations/broadcast.py b/openfhe_numpy/operations/broadcast.py index 13404ac..b852338 100644 --- a/openfhe_numpy/operations/broadcast.py +++ b/openfhe_numpy/operations/broadcast.py @@ -37,7 +37,6 @@ """ # Third-party imports -import math import numpy as np from ..openfhe_numpy import ArrayEncodingType from ..utils.matlib import next_power_of_two diff --git a/openfhe_numpy/operations/crypto_helper.py b/openfhe_numpy/operations/crypto_helper.py index ed5f6e3..714ca27 100644 --- a/openfhe_numpy/operations/crypto_helper.py +++ b/openfhe_numpy/operations/crypto_helper.py @@ -36,9 +36,12 @@ keys needed for various homomorphic operations in OpenFHE-NumPy. """ +from typing import Any import openfhe import openfhe_numpy as backend # Import from cpp source -from openfhe_numpy.utils.matlib import next_power_of_two +from ..utils.packing import _is_row_major, _is_col_major +from ..utils.matlib import next_power_of_two +from ..utils.errors import ONPValueError def accumulation_depth(nrows: int, ncols: int, accumulate_by_rows: bool): @@ -81,7 +84,7 @@ def sum_row_keys(secret_key: openfhe.PrivateKey, ncols: int = 0, slots: int = 0) Generated sum keys """ context = secret_key.GetCryptoContext() - return context.EvalSumRowsKeyGen(secret_key, None, ncols, slots * 4) + return context.EvalSumRowsKeyGen(secret_key, None, ncols, 0) def sum_col_keys(secret_key: openfhe.PrivateKey, ncols: int = 0): @@ -149,7 +152,7 @@ def gen_rotation_keys(secret_key: openfhe.PrivateKey, shifts: list): secret_key : PrivateKey The private key to use for key generation shifts : list - List of rotation indices to generate keys for + List of rotation indices to generate keys for (negated OpenFHE implementation). """ standard_indices = [-x for x in shifts] @@ -241,15 +244,105 @@ def generate_slicing_key(secret_key, original_shape): for i in range(1, max(nrow, ncol)): indices.add(-i) - for res_nrow in range(1, nrow + 1): - for res_ncol in range(1, ncol + 1): - res_pow_2_r = next_power_of_two(res_nrow) - res_pow_2_c = next_power_of_two(res_ncol) - for r in range(res_nrow): - for c in range(res_ncol): - indices.add(-(res_pow_2_c * r + c)) - indices.add(-(res_pow_2_r * c + r)) + # Rotation indices to collapse any sub-matrix result back to slot 0. + # + # The naive enumeration scans every (res_nrow, res_ncol) result shape and + # every (r, c) offset within it -- Theta(nrow^2 * ncol^2). Each added value + # depends on only one of the two result extents (through its power-of-two + # padding), so the identical index set is produced in O(nrow * ncol) by + # grouping result extents that share a padding and taking the widest offset + # range for that group. + def _widest_extent_by_padding(size): + """Map next_power_of_two(k) -> largest k in 1..size sharing that padding.""" + widest = {} + for k in range(1, size + 1): + widest[next_power_of_two(k)] = k # k ascends, so this stays the max + return widest + + # Column-padded offsets: value depends on res_ncol's padding; rows span 0..nrow. + for pad_c, max_c in _widest_extent_by_padding(ncol).items(): + for r in range(nrow): + for c in range(max_c): + indices.add(-(pad_c * r + c)) + + # Row-padded offsets: value depends on res_nrow's padding; cols span 0..ncol. + for pad_r, max_r in _widest_extent_by_padding(nrow).items(): + for c in range(ncol): + for r in range(max_r): + indices.add(-(pad_r * c + r)) + # NOTE: index 0 is intentionally kept. Element extraction rotates the target + # element to slot 0 via EvalRotate(ct, 0) for position-0 elements, which needs + # the identity-rotation key (OpenFHE automorphism index 1). if indices: cc = secret_key.GetCryptoContext() cc.EvalRotateKeyGen(secret_key, sorted(indices)) + + +############################################################################## +# BLOCK ARITHMETIC OPERATIONS +############################################################################## + + +# [CTArray] attach key for mat@vec product +def attach_matvec_keys(matrix, secret_key, emit: bool = False) -> tuple[str, Any] | None: + """Attach the summation key required for matrix-vector multiplication. + + Required keys: + - ROW_MAJOR matrix @ COL_MAJOR vector uses extra["colkey"]. + - COL_MAJOR matrix @ ROW_MAJOR vector uses extra["rowkey"]. + + Limitation: + - The key is generated from one matrix layout. All blocks in a block matrix + are expected to share the same packing order and logical column count. + """ + if matrix.ndim != 2: + raise ONPValueError("attach_matvec_keys expects a 2-D matrix.") + + if _is_row_major(matrix.order): + key_name = "colkey" + key = sum_col_keys(secret_key) + + elif _is_col_major(matrix.order): + key_name = "rowkey" + key = sum_row_keys(secret_key, matrix.ncols, matrix.batch_size) + + else: + raise ONPValueError(f"Unsupported packing order: {matrix.order}") + + matrix.extra[key_name] = key + + if emit: + return key_name, key + + return None + + +# [BlockCTArray] attach key for mat@vec product +def attach_block_matvec_keys( + block_matrix, secret_key, emit: bool = False +) -> tuple[str, Any] | None: + """Attach summation keys to encrypted matrix blocks for block matvec. + + Required keys: + - ROW_MAJOR matrix @ COL_MAJOR vector uses ``extra["colkey"]``. + - COL_MAJOR matrix @ ROW_MAJOR vector uses ``extra["rowkey"]``. + + The key-generation parameters must match the CTArray evaluator: + - EvalSumCols(..., lhs.ncols, ...) + - EvalSumRows(..., lhs.nrows, ..., lhs.batch_size) + """ + + if block_matrix.ndim != 2 or len(block_matrix.data) <= 0: + raise ONPValueError("attach_matvec_keys expects a 2-D block matrix.") + + reference = block_matrix.data[0] + key_name, key = attach_matvec_keys(reference, secret_key, True) + + for block in block_matrix.data: + block.extra[key_name] = key + + if emit: + return key_name, key + + return None diff --git a/openfhe_numpy/operations/matrix_api.py b/openfhe_numpy/operations/matrix_api.py index e428a3b..31396c4 100644 --- a/openfhe_numpy/operations/matrix_api.py +++ b/openfhe_numpy/operations/matrix_api.py @@ -115,32 +115,33 @@ def multiply(a: ArrayLike, b: ArrayLike) -> ArrayLike: pass -@tensor_function_api("pow", binary=True) -def pow(a: ArrayLike, exponent: int) -> ArrayLike: - """ +@tensor_function_api("power", binary=True) +def power(a: ArrayLike, exponent: int) -> ArrayLike: + """Raise each element of a tensor to an integer power (element-wise). - For each element of the tensor, it raises that element to the given power. + Matches ``numpy.power``: every element ``a_i`` is raised to ``exponent`` via + homomorphic element-wise multiplication. This is *not* matrix power. Note ---- - This only supports integer exponents due to homomorphic-encryption constraints. + Only positive integer exponents are supported due to homomorphic-encryption + constraints. ``exponent == 0`` (all-ones) is not supported and raises. Parameters ---------- a : ArrayLike Base tensor. exponent : int - Non-negative integer exponent. + Positive integer exponent (``>= 1``). Returns ------- out : ArrayLike - Element-wise 'a' raised to 'exponent'. + Element-wise ``a`` raised to ``exponent``. See Also -------- - numpy.pow : Corresponding element-wise power function. - numpy.linalg.matrix_power : Repeated matrix multiplication for square matrices. + numpy.power : Corresponding element-wise power function. Examples -------- @@ -256,8 +257,8 @@ def transpose(a: ArrayLike) -> ArrayLike: # =========================== -@tensor_function_api("cumulative_sum", binary=False) -def cumulative_sum(x: ArrayLike, /, *, axis: Optional[int] = None) -> ArrayLike: +@tensor_function_api("cumsum", binary=False) +def cumsum(x: ArrayLike, /, *, axis: Optional[int] = 0) -> ArrayLike: """ Compute the cumulative sum of tensor elements along a specified axis.\ - For 1D inputs, axis must be None. @@ -278,12 +279,12 @@ def cumulative_sum(x: ArrayLike, /, *, axis: Optional[int] = None) -> ArrayLike: See Also -------- - numpy.cumulative_sum : Corresponding NumPy function. + numpy.cumsum : Corresponding NumPy function. Examples -------- >>> import numpy as onp - >>> cumulative_sum(np.array([[1, 2], [3, 4]]), axis=1) + >>> cumsum(np.array([[1, 2], [3, 4]]), axis=1) array([[1, 3], [3, 7]]) """ @@ -312,7 +313,7 @@ def cumulative_reduce(a: ArrayLike, axis: int = 0, keepdims: bool = False) -> Ar See Also -------- - numpy.cumulative_sum : Similar operation for sum. + numpy.cumsum : Similar operation for sum. Examples -------- @@ -334,7 +335,7 @@ def sum(a: ArrayLike, /, *, axis: Optional[int] = None, keepdims: bool = False) a : ArrayLike Input tensor. axis : int, optional - Axis along which to compute the sum. Default is 0. + Axis along which to compute the sum. Default is None. 0: sum over rows 1: sum over cols keepdims : bool, optional @@ -384,7 +385,7 @@ def mean( a : ArrayLike Input tensor. axis : int, optional - Axis along which to compute the mean. Default is 0. + Axis along which to compute the mean. Default is None. keepdims : bool, optional If True, retains reduced dimensions. Default is False. @@ -412,29 +413,30 @@ def mean( @tensor_function_api("roll", binary=False) -def roll(a: ArrayLike, shift, axis: Optional[int] = None) -> ArrayLike: +def roll(a: ArrayLike, shift: int, axis: Optional[int] = None) -> ArrayLike: """ - Roll array elements along a given axis. + Roll packed vector elements. Elements that roll beyond the last position are re-introduced at the first. + Current limitation + ------------------ + This function currently supports only packed vectors with ``axis=None``. + Matrix roll, tuple-axis roll, and block tensor roll are not implemented yet. + Parameters ---------- a : ArrayLike - Input array. - shift : int or tuple of ints - The number of places by which elements are shifted. If a tuple, then - axis must be a tuple of the same size, and each of the given axes is - shifted by the corresponding number. If an int while axis is a tuple - of ints, then the same value is used for all given axes. - axis : int or tuple of ints, optional - Axis or axes along which elements are shifted. By default, the array - is flattened before shifting, after which the original shape is restored. + Input packed vector. + shift : int + Number of positions by which elements are shifted. + axis : None, optional + Must be None in the current implementation. Returns ------- res : ArrayLike - Output array, with the same shape as a. + Output vector with the same shape as ``a``. See Also -------- @@ -446,13 +448,18 @@ def roll(a: ArrayLike, shift, axis: Optional[int] = None) -> ArrayLike: >>> x = np.arange(10) >>> roll(x, 2) array([8, 9, 0, 1, 2, 3, 4, 5, 6, 7]) + >>> roll(x, -2) + array([2, 3, 4, 5, 6, 7, 8, 9, 0, 1]) + >>> roll(x, 0) + array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + + Unsupported examples + -------------------- >>> x2 = np.reshape(x, (2, 5)) >>> roll(x2, 1, axis=0) - array([[5, 6, 7, 8, 9], - [0, 1, 2, 3, 4]]) - >>> roll(x2, (1, 1), axis=(1, 0)) # Multiple axes - array([[9, 5, 6, 7, 8], - [4, 0, 1, 2, 3]]) + Traceback (most recent call last): + ... + ONPNotSupportedError: roll currently supports only packed vectors with axis=None. """ pass diff --git a/openfhe_numpy/operations/matrix_arithmetic.py b/openfhe_numpy/operations/matrix_arithmetic.py index 01faeb0..a0fe7b2 100644 --- a/openfhe_numpy/operations/matrix_arithmetic.py +++ b/openfhe_numpy/operations/matrix_arithmetic.py @@ -110,28 +110,6 @@ def add_cta_scalar(a, scalar): return CTArray(result, a.original_shape, a.batch_size, a.shape, a.order) -@register_tensor_function("add", ("scalar", "CTArray")) -def add_scalar_cta(scalar, a): - """Add a scalar to a x.""" - crypto_context = a.data.GetCryptoContext() - result = crypto_context.EvalAdd(a.data, scalar) - return CTArray(result, a.original_shape, a.batch_size, a.shape, a.order) - - -@register_tensor_function( - "add", [("BlockCTArray", "BlockCTArray"), ("BlockCTArray", "BlockPTArray")] -) -def add_block_ct(a, b): - """Add two block tensors.""" - raise NotImplementedError("BlockPTArray and BlockCTArray addition not implemented yet.") - - -@register_tensor_function("add", [("BlockCTArray", "scalar")]) -def add_block_ct_scalar(a, scalar): - """Add a scalar to a block x.""" - raise NotImplementedError("BlockPTArray and scalar addition not implemented yet.") - - # ------------------------------------------------------------------------------ # Subtraction Operations # ------------------------------------------------------------------------------ @@ -148,7 +126,7 @@ def _eval_sub(lhs, rhs): rhs = rhs.data result = crypto_context.EvalSub(lhs.data, rhs) - return lhs.clone(result) + return CTArray(result, lhs.original_shape, lhs.batch_size, lhs.shape, lhs.order) @register_tensor_function( @@ -237,21 +215,6 @@ def multiply_ct_scalar(a, scalar): return _eval_multiply(a, scalar) -@register_tensor_function( - "multiply", - [("BlockCTArray", "BlockCTArray"), ("BlockCTArray", "BlockPTArray")], -) -def multiply_block_ct(a, b): - """Multiply two block tensors element-wise.""" - raise NotImplementedError("BlockPTArray multiplication not implemented yet.") - - -@register_tensor_function("multiply", [("BlockCTArray", "scalar")]) -def multiply_block_ct_scalar(a, scalar): - """Multiply a block tensor by a scalar.""" - raise NotImplementedError("BlockPTArray and scalar multiplication not implemented yet.") - - ############################################################################## # MATRIX OPERATIONS ############################################################################## @@ -266,8 +229,9 @@ def _eval_matvec_ct(lhs, rhs): if lhs.ndim == 2 and rhs.ndim == 1: if lhs.original_shape[1] != rhs.original_shape[0]: raise ONPIncompatibleShapeError( - lhs.original_shape, rhs.original_shape, - f"Matrix dimension [{lhs.original_shape}] mismatch with vector dimension [{rhs.shape}]" + lhs.original_shape, + rhs.original_shape, + f"Matrix dimension [{lhs.original_shape}] mismatch with vector dimension [{rhs.shape}]", ) if lhs.order == ArrayEncodingType.ROW_MAJOR and rhs.order == ArrayEncodingType.COL_MAJOR: @@ -283,7 +247,7 @@ def _eval_matvec_ct(lhs, rhs): elif lhs.order == ArrayEncodingType.COL_MAJOR and rhs.order == ArrayEncodingType.ROW_MAJOR: ct_mult = cc.EvalMult(lhs.data, rhs.data) - ct_prod = cc.EvalSumRows(ct_mult, lhs.nrows, lhs.extra["rowkey"], lhs.batch_size * 4) + ct_prod = cc.EvalSumRows(ct_mult, lhs.nrows, lhs.extra["rowkey"], 0) return CTArray( ct_prod, (lhs.original_shape[0],), @@ -291,6 +255,7 @@ def _eval_matvec_ct(lhs, rhs): (lhs.shape[1], lhs.shape[0]), ArrayEncodingType.COL_MAJOR, ) + else: raise ONPError( f"Encoding styles of matrix ({lhs.order}) and vector ({rhs.order}) must be complementary (ROW_MAJOR/COL_MAJOR or vice versa)." @@ -371,7 +336,11 @@ def transpose_ct(a): # Power Operations # ------------------------------------------------------------------------------ def _pow(x, exp: int): - """Exponentiate a matrix to power k using homomorphic multiplication.""" + """Raise a tensor to an integer power element-wise (NumPy ``power`` semantics). + + Each element ``x_i`` is raised to ``exp`` via homomorphic element-wise + multiplication (``x * x * ... * x``). + """ if not isinstance(exp, int): raise ONPError(f"Exponent must be integer, got {type(exp).__name__}") @@ -379,57 +348,42 @@ def _pow(x, exp: int): raise ONPError("Negative exponent not supported in homomorphic encryption") if exp == 0: - # return algebra.eye(tensor)) - pass + raise ONPNotSupportedError("Element-wise power with exponent 0 (all-ones) is not supported") if exp == 1: return x.clone() - # Binary exponentiation implementation base = x.clone() result = None while exp: if exp & 1: - result = base if result is None else base @ result - base = base @ base + result = base if result is None else base * result exp >>= 1 + if exp: + base = base * base return result -@register_tensor_function("pow", [("CTArray", "int")]) +@register_tensor_function("power", [("CTArray", "int")]) def pow_ct(a, exp): - """Raise a tensor to an integer power.""" + """Raise a tensor to an integer power element-wise.""" return _pow(a, exp) -@register_tensor_function("pow", [("BlockCTArray", "int")]) -def pow_block_ct(a, exp): - """Raise a block tensor to an integer power.""" - raise NotImplementedError("BlockPTArray power not implemented yet.") - - # ------------------------------------------------------------------------------ # Cumulative Sum Operations # ------------------------------------------------------------------------------ @register_tensor_function( - "cumulative_sum", + "cumsum", [("CTArray",), ("CTArray", "int"), ("CTArray", "int", "bool")], ) -def cumulative_sum_ct(obj, axis=0, keepdims=True): +def cumsum_ct(obj, axis=0, keepdims=True): """Compute cumulative sum of a tensor along specified axis.""" - # return _cumulative_sum_ct(a, axis, keepdims) - return obj.cumulative_sum(axis) - - -@register_tensor_function( - "cumulative_sum", [("BlockCTArray", "int"), ("BlockCTArray", "int", "bool")] -) -def cumulative_sum_block_ct(obj, axis=0, keepdims=True): - """Compute cumulative sum of a block tensor along specified axis.""" - raise NotImplementedError("BlockPTArray cumulative not implemented yet.") + # return _cumsum_ct(a, axis, keepdims) + return obj.cumsum(axis) # ------------------------------------------------------------------------------ @@ -469,12 +423,6 @@ def cumulative_reduce_ct(a, axis=0, keepdims=False): return _reduce_ct(a, axis, keepdims) -@register_tensor_function("cumulative_reduce", [("BlockCTArray", "int")]) -def cumulative_reduce_block_ct(a, axis=0, keepdims=False): - """Compute cumulative reduction of a block tensor along specified axis.""" - raise NotImplementedError("BlockPTArray power not implemented yet.") - - # ------------------------------------------------------------------------------ # Sum Operations # ------------------------------------------------------------------------------ @@ -535,7 +483,7 @@ def _ct_sum_matrix(x: ArrayLike, axis: Optional[int] = None, keepdims: bool = Tr elif axis == 0: # Sum across each row of a packed_encoded matrix ciphertext: fhe_data if order == ArrayEncodingType.ROW_MAJOR: - ct_sum = cc.EvalSumRows(fhe_data, ncols, x.extra["rowkey"], x.batch_size * 4) + ct_sum = cc.EvalSumRows(fhe_data, ncols, x.extra["rowkey"], 0) padded_shape = x.shape order = ArrayEncodingType.COL_MAJOR elif order == ArrayEncodingType.COL_MAJOR: @@ -558,7 +506,7 @@ def _ct_sum_matrix(x: ArrayLike, axis: Optional[int] = None, keepdims: bool = Tr padded_shape = x.shape order = ArrayEncodingType.ROW_MAJOR elif order == ArrayEncodingType.COL_MAJOR: - ct_sum = cc.EvalSumRows(fhe_data, nrows, x.extra["rowkey"], x.batch_size * 4) + ct_sum = cc.EvalSumRows(fhe_data, nrows, x.extra["rowkey"], 0) padded_shape = (ncols, nrows) order = ArrayEncodingType.COL_MAJOR else: diff --git a/openfhe_numpy/tensor/block_ctarray.py b/openfhe_numpy/tensor/block_ctarray.py index 7150270..a99bd32 100644 --- a/openfhe_numpy/tensor/block_ctarray.py +++ b/openfhe_numpy/tensor/block_ctarray.py @@ -32,6 +32,8 @@ from __future__ import annotations +from typing import Any, Callable + import numpy as np from openfhe import Ciphertext diff --git a/openfhe_numpy/tensor/block_tensor.py b/openfhe_numpy/tensor/block_tensor.py index e3e54e1..453de8c 100644 --- a/openfhe_numpy/tensor/block_tensor.py +++ b/openfhe_numpy/tensor/block_tensor.py @@ -548,7 +548,7 @@ def __matmul__(self, other): def __pow__(self, exp): """Return matrix power ``self ** exp`` through tensor dispatch.""" - return self.__tensor_function__("pow", (self, exp)) + return self.__tensor_function__("power", (self, exp)) @property def T(self): @@ -606,7 +606,7 @@ def clone(self, data: list[Any] | None = None) -> BlockFHETensor: ) def same_layout(self, other: Any) -> bool: - """Return ``True`` if layout metadata and encryption status match.""" + """Return ``True`` if block/grid layout metadata match.""" attrs = ( "original_shape", "shape", @@ -614,7 +614,6 @@ def same_layout(self, other: Any) -> bool: "order", "block_shape", "grid_shape", - "is_encrypted", ) if not all(hasattr(other, attr) for attr in attrs): return False @@ -626,12 +625,15 @@ def same_layout(self, other: Any) -> bool: and self._order == other.order and self._block_shape == other.block_shape and self._grid_shape == other.grid_shape - and self.is_encrypted == other.is_encrypted ) def same_metadata(self, other: Any) -> bool: """Return ``True`` if layout metadata, encryption status, and dtype match.""" - return self.same_layout(other) and getattr(other, "dtype", None) == self._dtype + return ( + self.same_layout(other) + and getattr(other, "is_encrypted", None) == self.is_encrypted + and getattr(other, "dtype", None) == self._dtype + ) def __eq__(self, other: Any) -> bool: """Return metadata-only equality. diff --git a/openfhe_numpy/tensor/ctarray.py b/openfhe_numpy/tensor/ctarray.py index 2ab26b3..c7eb9c7 100644 --- a/openfhe_numpy/tensor/ctarray.py +++ b/openfhe_numpy/tensor/ctarray.py @@ -310,20 +310,10 @@ def deserialize(cls, obj: dict) -> "CTArray": obj["order"], ) - def __repr__(self) -> str: - return f"CTArray(metadata={self.metadata})" - def __neg__(self) -> "CTArray": """Return the homomorphic negation of this ciphertext array.""" cc = self.data.GetCryptoContext() return self.clone(cc.EvalNegate(self.data)) - def sum(self, axis=None, keepdims: bool = False) -> "CTArray": - """Sum tensor elements over an axis.""" - return self.__tensor_function__( - "sum", - (self,), - {"axis": axis, "keepdims": keepdims}, - ) def _transpose(self) -> "CTArray": """Internal function to evaluate transpose of an encrypted array.""" @@ -377,7 +367,7 @@ def cumsum(self, axis: int = 0) -> "CTArray": if axis is None: ciphertext = EvalSumCumRows(self.data, self.ncols, self.original_shape[1]) - # cumulative_sum over rows + # cumsum over rows elif axis == 0: if self.order == ArrayEncodingType.ROW_MAJOR: ciphertext = EvalSumCumRows(self.data, self.ncols, self.original_shape[1]) @@ -388,7 +378,7 @@ def cumsum(self, axis: int = 0) -> "CTArray": else: raise ONPError(f"Not support this packing order [{self.order}].") - # cumulative_sum over cols + # cumsum over cols elif axis == 1: if self.order == ArrayEncodingType.ROW_MAJOR: ciphertext = EvalSumCumCols(self.data, self.ncols) @@ -402,20 +392,12 @@ def cumsum(self, axis: int = 0) -> "CTArray": raise ONPError(f"Invalid axis [{axis}].") return CTArray(ciphertext, original_shape, self.batch_size, shape, order) - def gen_sum_row_key(self, secret_key: openfhe.PrivateKey) -> openfhe.EvalKey: - context = secret_key.GetCryptoContext() - if self.order == ArrayEncodingType.ROW_MAJOR: - sum_rows_key = context.EvalSumRowsKeyGen(secret_key, self.ncols, self.batch_size) - elif self.order == ArrayEncodingType.COL_MAJOR: - sum_rows_key = context.EvalSumColsKeyGen(secret_key) - else: - raise ValueError("Invalid order.") - - return sum_rows_key def apply(self, func: Callable, *args: Any, **kwargs: Any) -> "CTArray": """Apply a ciphertext-level function to the underlying ciphertext. + The function must accept ``self.data`` as its first argument and return an OpenFHE ciphertext with the same logical packing layout. + Parameters ---------- func : Callable @@ -424,19 +406,25 @@ def apply(self, func: Callable, *args: Any, **kwargs: Any) -> "CTArray": Additional positional arguments passed to ``func``. **kwargs : Any Additional keyword arguments passed to ``func``. + Returns ------- CTArray A new CTArray with the same shape/metadata but the transformed ciphertext. + Examples -------- Bootstrap to refresh noise level: + ``result = a.apply(cc.EvalBootstrap)`` + Chebyshev-style functions can be wrapped when the ciphertext is not the first argument: + ``result = a.apply(lambda ct: cc.EvalChebyshevSeries(ct, coeffs, -8, 8))`` """ if not callable(func): raise TypeError(f"apply expects a callable, got {type(func).__name__}.") + ct_result = func(self.data, *args, **kwargs) return self.clone(data=ct_result) diff --git a/openfhe_numpy/tensor/ptarray.py b/openfhe_numpy/tensor/ptarray.py index b6dc159..f7e251d 100644 --- a/openfhe_numpy/tensor/ptarray.py +++ b/openfhe_numpy/tensor/ptarray.py @@ -44,6 +44,7 @@ class PTArray(FHETensor[Plaintext]): """Concrete tensor class for OpenFHE plaintexts.""" is_encrypted = False + def clone(self, data=None): return PTArray( data or self.data, @@ -77,6 +78,7 @@ def decode(self, unpack_type: UnpackType = UnpackType.ORIGINAL) -> np.ndarray: if unpack_type == UnpackType.ORIGINAL: return process_packed_data(result, self.info) return np.asarray(result) + def __repr__(self) -> str: return f"PTArray(meta={self.info})" diff --git a/openfhe_numpy/tensor/tensor.py b/openfhe_numpy/tensor/tensor.py index 67c0942..7232434 100644 --- a/openfhe_numpy/tensor/tensor.py +++ b/openfhe_numpy/tensor/tensor.py @@ -49,6 +49,7 @@ # Don't implement anything here class BaseTensor(ABC, Generic[TPL]): is_encrypted: bool = False + @property @abstractmethod def shape(self) -> Tuple[int, ...]: ... @@ -129,7 +130,6 @@ class FHETensor(BaseTensor[TPL], Generic[TPL]): "_dtype", "extra", ) - is_encrypted = False @overload def __init__( @@ -278,10 +278,6 @@ def order(self, order: int): else: raise ONPError(f"Not support order [f{order}]") - @property - def is_encrypted(self) -> int: - return "CT" in self.dtype - @property def info(self) -> Dict[str, Any]: """Metadata dict for serialization or inspection.""" @@ -376,7 +372,6 @@ def __matmul__(self, other): def __pow__(self, exponent): return self.__tensor_function__("power", (self, exponent)) - def sum(self, axis=None, keepdims=False): """ Sum tensor entries over all elements or one axis. @@ -385,10 +380,13 @@ def sum(self, axis=None, keepdims=False): if axis is not None: if not isinstance(axis, int): raise TypeError(f"axis must be None or an integer, got {type(axis).__name__}.") + if axis < 0: axis += self.ndim + if axis < 0 or axis >= self.ndim: raise ValueError(f"Invalid axis {axis} for tensor with {self.ndim} dimensions.") + return self.__tensor_function__( "sum", (self,), @@ -426,47 +424,6 @@ def __getitem__(self, key): """ raise NotImplementedError() - # def ensure_compatible_packing(self, other): - # """ - # Ensure tensors have compatible packing for operations. - - # Returns a version of 'other' with matching packing order. - # """ - # if not isinstance(other, FHETensor): - # return other - - # if self.order == other.order: - # return other - - # return other.convert_packing_order(self.order) - - # def convert_packing_order(self, target_order): - # """ - # Convert tensor to a different packing order. - - # Parameters - # ---------- - # target_order : int - # Desired packing order (ROW_MAJOR or COL_MAJOR) - - # Returns - # ------- - # FHETensor - # New tensor with converted packing order - # """ - # if self.order == target_order: - # return self.clone() - - # # Perform conversion - # if self.dtype == "CTArray": - # # For ciphertexts, use transpose operation - # transposed = self._transpose() - # # Update order flag - # transposed._order = target_order - # return transposed - # else: - # pass - def copy_tensor(tensor: "FHETensor") -> "FHETensor": """ From 4016c7fa8f89a4c8fd6f01e176148986a91e2530 Mon Sep 17 00:00:00 2001 From: tran-go Date: Thu, 16 Jul 2026 20:11:45 +0800 Subject: [PATCH 2/2] Add examples --- .../simple_block_matrix_operation.py | 183 ++++++++++++++++ .../simple_block_matrix_vector_product.py | 206 +++++++++++++++++ .../simple_block_vector_operations.py | 196 +++++++++++++++++ .../simple_square_bmatrix_product.py | 207 ++++++++++++++++++ examples/python/simple_matrix_accumulation.py | 8 +- examples/python/tests.py | 79 +++++++ 6 files changed, 875 insertions(+), 4 deletions(-) create mode 100644 examples/python/block_tensor/simple_block_matrix_operation.py create mode 100644 examples/python/block_tensor/simple_block_matrix_vector_product.py create mode 100644 examples/python/block_tensor/simple_block_vector_operations.py create mode 100644 examples/python/block_tensor/simple_square_bmatrix_product.py create mode 100644 examples/python/tests.py diff --git a/examples/python/block_tensor/simple_block_matrix_operation.py b/examples/python/block_tensor/simple_block_matrix_operation.py new file mode 100644 index 0000000..8ff9ec8 --- /dev/null +++ b/examples/python/block_tensor/simple_block_matrix_operation.py @@ -0,0 +1,183 @@ +import numpy as np +from openfhe import * +import openfhe_numpy as onp + + +def validate_and_print_results(computed, expected, operation_name): + """Helper function to validate and print matrix results.""" + print("\n" + "*" * 60) + print(f"* {operation_name}") + print("*" * 60) + print(f"\nExpected:\n{expected}") + print(f"\nDecrypted Result:\n{computed}") + + is_match, error = onp.check_equality(computed, expected) + print(f"\nMatch: {is_match}, Total Error: {error}") + return is_match, error + + +def print_block_metadata(name, block_array): + """Print block tensor metadata.""" + print(f"\n{name} metadata") + print(f" original_shape : {block_array.original_shape}") + print(f" padded shape : {block_array.shape}") + print(f" block_shape : {block_array.block_shape}") + print(f" grid_shape : {block_array.grid_shape}") + print(f" num_blocks : {block_array.num_blocks}") + print(f" batch_size : {block_array.batch_size}") + + +def main(): + """ + We deliberately pick a tiny CKKS ring dimension (2**6 = 64), which gives only + 32 usable slots per ciphertext. Our 8x8 matrix has 64 entries -- more than 32 -- + so it *cannot* be packed into a single ciphertext. openfhe-numpy tiles it into a + grid of small blocks, each of which fits in one ciphertext. + + Operations shown (all elementwise, so any block shape works): + - block matrix addition (CT + CT) + - block matrix subtraction (CT - CT) + - block matrix elementwise multiplication (CT * CT) + - block matrix scalar multiplication (CT * scalar) + - block matrix mixed addition (CT + PT) + """ + + # --- Cryptographic setup ------------------------------------------------- + ring_dim = 2**6 # 64 + mult_depth = 1 # elementwise mul needs depth 1; increase for chained ops + scale_mod_size = 50 + + # block_shape is OPTIONAL. Leave it None and openfhe-numpy auto-selects the + # largest square block that fits in one ciphertext: + # side = 2 ** floor((batch_size.bit_length() - 1) / 2), batch_size = ring_dim // 2 + # For batch_size = 32 that gives a (4, 4) block. Pass an explicit tuple + # (e.g. (4, 4) or (2, 2)) to choose it yourself. + block_shape = None + + params = CCParamsCKKSRNS() + params.SetRingDim(ring_dim) + params.SetSecurityLevel(HEStd_NotSet) + params.SetMultiplicativeDepth(mult_depth) + params.SetScalingModSize(scale_mod_size) + params.SetFirstModSize(60) + params.SetScalingTechnique(FIXEDAUTO) + + cc = GenCryptoContext(params) + cc.Enable(PKESchemeFeature.PKE) + cc.Enable(PKESchemeFeature.LEVELEDSHE) + cc.Enable(PKESchemeFeature.ADVANCEDSHE) + + keys = cc.KeyGen() + cc.EvalMultKeyGen(keys.secretKey) + + batch_size = cc.GetRingDimension() // 2 + + # --- Why we need blocks -------------------------------------------------- + matrix_a = np.arange(1, 65, dtype=float).reshape(8, 8) + matrix_b = 65.0 - matrix_a + + n_entries = matrix_a.size + print(f"\nCKKS ring dimension : {cc.GetRingDimension()}") + print(f"Slots per ciphertext: {batch_size}") + print(f"Matrix shape : {matrix_a.shape} ({n_entries} entries)") + print( + f"\n{n_entries} entries > {batch_size} slots => the matrix does NOT fit in a\n" + f"single ciphertext, so openfhe-numpy tiles it into blocks that each do." + ) + + print("\nInput matrices") + print("matrix_a:") + print(matrix_a) + print("\nmatrix_b:") + print(matrix_b) + + # --- Build encrypted block matrices -------------------------------------- + # With block_shape=None the 8x8 matrix is auto-tiled into a 2x2 grid of (4,4) + # blocks; each block is one ciphertext. (Non-divisible sizes zero-pad the edge + # blocks.) + ctm_a = onp.block_array( + cc=cc, + data=matrix_a, + batch_size=batch_size, + block_shape=block_shape, + order=onp.ROW_MAJOR, + mode="zero", + fhe_type="C", + public_key=keys.publicKey, + ) + + ctm_b = onp.block_array( + cc=cc, + data=matrix_b, + batch_size=batch_size, + block_shape=block_shape, + order=onp.ROW_MAJOR, + mode="zero", + fhe_type="C", + public_key=keys.publicKey, + ) + + # A plaintext block matrix for the mixed CT + PT test. + ptm_b = onp.block_array( + cc=cc, + data=matrix_b, + batch_size=batch_size, + block_shape=block_shape, + order=onp.ROW_MAJOR, + mode="zero", + fhe_type="P", + ) + + print( + f"\nAuto-selected block_shape = {ctm_a.block_shape}, " + f"grid_shape = {ctm_a.grid_shape} ({ctm_a.num_blocks} ciphertext blocks)" + ) + + print_block_metadata("Encrypted block matrix_a", ctm_a) + print_block_metadata("Encrypted block matrix_b", ctm_b) + print_block_metadata("Plaintext block matrix_b", ptm_b) + + all_ok = True + + # 1) Block matrix addition (CT + CT) + res_add = (ctm_a + ctm_b).decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_add, matrix_a + matrix_b, "Block matrix addition (CT + CT)" + ) + all_ok = all_ok and is_match + + # 2) Block matrix subtraction (CT - CT) + res_sub = (ctm_a - ctm_b).decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_sub, matrix_a - matrix_b, "Block matrix subtraction (CT - CT)" + ) + all_ok = all_ok and is_match + + # 3) Block matrix elementwise multiplication (CT * CT) + res_mul = (ctm_a * ctm_b).decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_mul, matrix_a * matrix_b, "Block matrix elementwise multiplication (CT * CT)" + ) + all_ok = all_ok and is_match + + # 4) Block matrix scalar multiplication (CT * scalar) + res_scalar_mul = (ctm_a * 3.0).decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_scalar_mul, matrix_a * 3.0, "Block matrix scalar multiplication (CT * 3.0)" + ) + all_ok = all_ok and is_match + + # 5) Mixed block matrix addition (CT + PT) + res_add_plain = (ctm_a + ptm_b).decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_add_plain, matrix_a + matrix_b, "Block matrix mixed addition (CT + PT)" + ) + all_ok = all_ok and is_match + + print("\n" + "=" * 60) + print(f"All checks passed: {all_ok}") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/examples/python/block_tensor/simple_block_matrix_vector_product.py b/examples/python/block_tensor/simple_block_matrix_vector_product.py new file mode 100644 index 0000000..7400c9d --- /dev/null +++ b/examples/python/block_tensor/simple_block_matrix_vector_product.py @@ -0,0 +1,206 @@ +import numpy as np +from openfhe import * +import openfhe_numpy as onp + + +def validate_and_print_results(computed, expected, operation_name): + """Helper function to validate and print results.""" + print("\n" + "*" * 60) + print(f"* {operation_name}") + print("*" * 60) + print(f"\nExpected:\n{expected}") + print(f"\nDecrypted Result:\n{computed}") + + is_match, error = onp.check_equality(computed, expected) + print(f"\nMatch: {is_match}, Total Error: {error}") + return is_match, error + + +def print_block_metadata(name, block_array): + """Print block tensor metadata.""" + print(f"\n{name} metadata") + print(f" original_shape : {block_array.original_shape}") + print(f" padded shape : {block_array.shape}") + print(f" block_shape : {block_array.block_shape}") + print(f" grid_shape : {block_array.grid_shape}") + print(f" num_blocks : {block_array.num_blocks}") + print(f" batch_size : {block_array.batch_size}") + + +def main(): + """ + Demonstrate homomorphic block matrix-vector multiplication using OpenFHE-NumPy. + + Two encoding combinations are shown: + + Case 1: ROW_MAJOR matrix @ COL_MAJOR vector -> ROW_MAJOR result + Case 2: COL_MAJOR matrix @ ROW_MAJOR vector -> COL_MAJOR result + + The block_shape=(2, 2) splits the matrix across multiple ciphertext + blocks, each holding a 2x2 sub-matrix in batch_size=4 slots. + compact=True on the vector produces the duplicated, square-compatible + packing required for block matvec; do NOT use compact for plain + vector arithmetic (add/sub/dot). + """ + + # Cryptographic setup + mult_depth = 4 + scale_mod_size = 59 + block_shape = (2, 2) + # Each 2x2 block uses exactly 4 slots. + batch_size = int(np.prod(block_shape)) + + params = CCParamsCKKSRNS() + params.SetMultiplicativeDepth(mult_depth) + params.SetScalingModSize(scale_mod_size) + params.SetFirstModSize(60) + params.SetScalingTechnique(FIXEDAUTO) + + cc = GenCryptoContext(params) + cc.Enable(PKESchemeFeature.PKE) + cc.Enable(PKESchemeFeature.LEVELEDSHE) + cc.Enable(PKESchemeFeature.ADVANCEDSHE) + + keys = cc.KeyGen() + cc.EvalMultKeyGen(keys.secretKey) + cc.EvalSumKeyGen(keys.secretKey) + + ring_dim = cc.GetRingDimension() + print(f"\nCKKS ring dimension: {ring_dim}") + print(f"Block batch size: {batch_size}") + print(f"Block shape: {block_shape}") + + # Sample input + matrix = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + [4.0, 5.0, 6.0], + [4.0, 5.0, 6.0], + [4.0, 5.0, 6.0], + ], + dtype=float, + ) + + vector = np.array([1.0, 1.0, 1.0], dtype=float) + + print("\nInput") + print("\nMatrix:\n", matrix) + print("\nVector:\n", vector) + + expected = matrix @ vector + print(f"\nExpected:\n{expected}") + + all_ok = True + + # ========================================================= + # Case 1: ROW_MAJOR matrix @ COL_MAJOR vector + # ========================================================= + # + # Row-wise packing of a 2-row x 3-col matrix with pad-to-power-of-2: + # + # 1 2 3 0 | 4 5 6 0 (two 2x2 blocks packed as [row0|row1]) + # + # The vector is encoded with compact=True (duplicated packing): + # + # 1 1 1 1 | 0 0 0 0 (compact COL_MAJOR block vector) + # + # The result lands in ROW_MAJOR order, one entry per row of the matrix. + + ctm_rm = onp.block_array( + cc=cc, + data=matrix, + batch_size=batch_size, + block_shape=block_shape, + order=onp.ROW_MAJOR, + mode="zero", + fhe_type="C", + public_key=keys.publicKey, + ) + + # compact=True: produce duplicated block packing required by block matvec. + ctv_cm = onp.block_array( + cc=cc, + data=vector, + batch_size=batch_size, + block_shape=None, + order=onp.COL_MAJOR, + mode="zero", + fhe_type="C", + public_key=keys.publicKey, + compact=True, + ) + + # Attach summation keys to every matrix block before the product. + # ROW_MAJOR matrix requires extra["colkey"] for EvalSumCols. + onp.attach_block_matvec_keys(ctm_rm, keys.secretKey) + + print_block_metadata("Block matrix (ROW_MAJOR)", ctm_rm) + print_block_metadata("Block vector (COL_MAJOR, compact)", ctv_cm) + + ctv_result_rm = ctm_rm @ ctv_cm + res_rm = ctv_result_rm.decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_rm, + expected, + "Block matvec Case 1: ROW_MAJOR matrix @ COL_MAJOR vector", + ) + all_ok = all_ok and is_match + + # ========================================================= + # Case 2: COL_MAJOR matrix @ ROW_MAJOR vector + # ========================================================= + # + # Column-wise packing interleaves the matrix columns across slots. + # The vector uses compact=True (ROW_MAJOR), producing duplicated + # per-element blocks compatible with the COL_MAJOR matrix blocks. + # + # The result lands in COL_MAJOR order, one entry per row of the matrix. + + ctm_cm = onp.block_array( + cc=cc, + data=matrix, + batch_size=batch_size, + block_shape=block_shape, + order=onp.COL_MAJOR, + mode="zero", + fhe_type="C", + public_key=keys.publicKey, + ) + + # compact=True: produce duplicated block packing required by block matvec. + ctv_rm = onp.block_array( + cc=cc, + data=vector, + batch_size=batch_size, + block_shape=None, + order=onp.ROW_MAJOR, + mode="zero", + fhe_type="C", + public_key=keys.publicKey, + compact=True, + ) + + # Attach summation keys to every matrix block before the product. + # COL_MAJOR matrix requires extra["rowkey"] for EvalSumRows. + onp.attach_block_matvec_keys(ctm_cm, keys.secretKey) + + print_block_metadata("Block matrix (COL_MAJOR)", ctm_cm) + print_block_metadata("Block vector (ROW_MAJOR, compact)", ctv_rm) + + ctv_result_cm = ctm_cm @ ctv_rm + res_cm = ctv_result_cm.decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_cm, + expected, + "Block matvec Case 2: COL_MAJOR matrix @ ROW_MAJOR vector", + ) + all_ok = all_ok and is_match + + print("\n" + "=" * 60) + print(f"All checks passed: {all_ok}") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/examples/python/block_tensor/simple_block_vector_operations.py b/examples/python/block_tensor/simple_block_vector_operations.py new file mode 100644 index 0000000..c2ac4b8 --- /dev/null +++ b/examples/python/block_tensor/simple_block_vector_operations.py @@ -0,0 +1,196 @@ +import numpy as np +from openfhe import * +import openfhe_numpy as onp + + +def validate_and_print_results(computed, expected, operation_name): + """Helper function to validate and print vector results.""" + print("\n" + "*" * 60) + print(f"* {operation_name}") + print("*" * 60) + print(f"\nExpected:\n{expected}") + print(f"\nDecrypted Result:\n{computed}") + + is_match, error = onp.check_equality(computed, expected) + print(f"\nMatch: {is_match}, Total Error: {error}") + return is_match, error + + +def print_block_metadata(name, block_array): + """Print block tensor metadata.""" + print(f"\n{name} metadata") + print(f" original_shape : {block_array.original_shape}") + print(f" padded shape : {block_array.shape}") + print(f" block_shape : {block_array.block_shape}") + print(f" grid_shape : {block_array.grid_shape}") + print(f" num_blocks : {block_array.num_blocks}") + print(f" batch_size : {block_array.batch_size}") + + +def main(): + """ + Demonstrate homomorphic block vector operations using OpenFHE-NumPy: + + - block vector addition (CT + CT) + - block vector subtraction (CT - CT) + - block vector elementwise multiplication (CT * CT) + - block vector scalar multiplication (CT * scalar) + - block mixed addition (CT + PT) + - block vector dot product via onp.dot + - block vector dot product via @ + """ + + # Cryptographic setup + # Length-5 vectors with batch_size=4 force two ciphertext blocks each. + mult_depth = 4 + scale_mod_size = 59 + batch_size = 4 + + params = CCParamsCKKSRNS() + params.SetMultiplicativeDepth(mult_depth) + params.SetScalingModSize(scale_mod_size) + params.SetFirstModSize(60) + params.SetScalingTechnique(FIXEDAUTO) + + cc = GenCryptoContext(params) + cc.Enable(PKESchemeFeature.PKE) + cc.Enable(PKESchemeFeature.LEVELEDSHE) + cc.Enable(PKESchemeFeature.ADVANCEDSHE) + + keys = cc.KeyGen() + cc.EvalMultKeyGen(keys.secretKey) + cc.EvalSumKeyGen(keys.secretKey) + + ring_dim = cc.GetRingDimension() + print(f"\nCKKS ring dimension: {ring_dim}") + print(f"Block batch size: {batch_size}") + + # Sample input vectors. + # Length 5 with batch_size=4 forces two ciphertext blocks. + vector_a = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=float) + vector_b = np.array([4.0, 0.0, 1.0, 3.0, 6.0], dtype=float) + + print("\nInput vectors") + print("vector_a:", vector_a) + print("vector_b:", vector_b) + + # Create encrypted block vectors. + # block_shape=None lets the library choose the block size from batch_size. + ctv_a = onp.block_array( + cc=cc, + data=vector_a, + batch_size=batch_size, + block_shape=None, + order=onp.ROW_MAJOR, + mode="zero", + fhe_type="C", + public_key=keys.publicKey, + ) + + ctv_b = onp.block_array( + cc=cc, + data=vector_b, + batch_size=batch_size, + block_shape=None, + order=onp.ROW_MAJOR, + mode="zero", + fhe_type="C", + public_key=keys.publicKey, + ) + + # Also create a plaintext block vector for the mixed CT+PT test. + ptv_b = onp.block_array( + cc=cc, + data=vector_b, + batch_size=batch_size, + block_shape=None, + order=onp.ROW_MAJOR, + mode="zero", + fhe_type="P", + ) + + print_block_metadata("Block encrypted vector_a", ctv_a) + print_block_metadata("Block encrypted vector_b", ctv_b) + print_block_metadata("Block plaintext vector_b", ptv_b) + + all_ok = True + + # 1) Block vector addition (CT + CT) + ctv_add = ctv_a + ctv_b + res_add = ctv_add.decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_add, + vector_a + vector_b, + f"Block vector addition (CT + CT)\n{vector_a}\n+\n{vector_b}", + ) + all_ok = all_ok and is_match + + # 2) Block vector subtraction (CT - CT) + ctv_sub = ctv_a - ctv_b + res_sub = ctv_sub.decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_sub, + vector_a - vector_b, + f"Block vector subtraction (CT - CT)\n{vector_a}\n-\n{vector_b}", + ) + all_ok = all_ok and is_match + + # 3) Block vector elementwise multiplication (CT * CT) + ctv_mul = ctv_a * ctv_b + res_mul = ctv_mul.decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_mul, + vector_a * vector_b, + f"Block vector elementwise multiplication (CT * CT)\n{vector_a}\n*\n{vector_b}", + ) + all_ok = all_ok and is_match + + # 4) Block vector scalar multiplication (CT * scalar) + ctv_scalar_mul = ctv_a * 7.0 + res_scalar_mul = ctv_scalar_mul.decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_scalar_mul, + vector_a * 7.0, + f"Block vector scalar multiplication (CT * 7.0)\n{vector_a} * 7.0", + ) + all_ok = all_ok and is_match + + # 5) Mixed block vector addition (CT + PT) + ctv_add_plain = ctv_a + ptv_b + res_add_plain = ctv_add_plain.decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_add_plain, + vector_a + vector_b, + f"Block vector mixed addition (CT + PT)\n{vector_a}\n+\n{vector_b}", + ) + all_ok = all_ok and is_match + + # 6) Block vector dot product via onp.dot + # + # Expected: 1*4 + 2*0 + 3*1 + 4*3 + 5*6 = 49 + ctv_dot = onp.dot(ctv_a, ctv_b) + res_dot = ctv_dot.decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_dot, + np.dot(vector_a, vector_b), + f"Block vector dot product via onp.dot\n{vector_a}\n·\n{vector_b}", + ) + all_ok = all_ok and is_match + + # 7) Block vector dot product via @ + ctv_dot_op = ctv_a @ ctv_b + res_dot_op = ctv_dot_op.decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results( + res_dot_op, + np.dot(vector_a, vector_b), + f"Block vector dot product via @\n{vector_a}\n@\n{vector_b}", + ) + all_ok = all_ok and is_match + + print("\n" + "=" * 60) + print(f"All checks passed: {all_ok}") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/examples/python/block_tensor/simple_square_bmatrix_product.py b/examples/python/block_tensor/simple_square_bmatrix_product.py new file mode 100644 index 0000000..9523a36 --- /dev/null +++ b/examples/python/block_tensor/simple_square_bmatrix_product.py @@ -0,0 +1,207 @@ +import numpy as np +from openfhe import * +import openfhe_numpy as onp + + +def validate_and_print_results(computed, expected, operation_name): + """Helper function to validate and print results.""" + print("\n" + "*" * 60) + print(f"* {operation_name}") + print("*" * 60) + print(f"\nExpected:\n{expected}") + print(f"\nDecrypted Result:\n{computed}") + + is_match, error = onp.check_equality(computed, expected) + print(f"\nMatch: {is_match}, Total Error: {error}") + return is_match, error + + +def print_block_metadata(name, block_array): + """Print block tensor metadata.""" + print(f"\n{name} metadata") + print(f" original_shape : {block_array.original_shape}") + print(f" padded shape : {block_array.shape}") + print(f" block_shape : {block_array.block_shape}") + print(f" grid_shape : {block_array.grid_shape}") + print(f" num_blocks : {block_array.num_blocks}") + print(f" batch_size : {block_array.batch_size}") + + +def run_block_matmul_example(cc, keys, A, B, block_shape, description): + """Run one homomorphic block matrix multiplication example. + + Both operands must use ROW_MAJOR packing and identical square block_shape. + The rotation keys are generated from the block size (not the full matrix), + since EvalSquareMatMultRotateKeyGen operates on a single block. + """ + print(f"\n--- {description} ---") + print("Input A:\n", A) + print("Input B:\n", B) + + # Each ciphertext block stores one block_shape tile; + # batch_size equals the number of slots consumed by one block. + batch_size = int(np.prod(block_shape)) + + ctm_A = onp.block_array( + cc=cc, + data=A, + batch_size=batch_size, + block_shape=block_shape, + order=onp.ROW_MAJOR, + mode="zero", + fhe_type="C", + public_key=keys.publicKey, + ) + + ctm_B = onp.block_array( + cc=cc, + data=B, + batch_size=batch_size, + block_shape=block_shape, + order=onp.ROW_MAJOR, + mode="zero", + fhe_type="C", + public_key=keys.publicKey, + ) + + print_block_metadata("Block encrypted A", ctm_A) + print_block_metadata("Block encrypted B", ctm_B) + + # Generate rotation keys for one square block of size block_shape[0]. + # The keys cover all rotations needed by EvalMatMulSquare on each block. + block_size = block_shape[0] + onp.gen_square_matmult_key(keys.secretKey, block_size) + + # Compute C[i, j] = sum_k A[i, k] @ B[k, j], + # where each A[i, k] and B[k, j] is an encrypted CTArray block. + ctm_C = ctm_A @ ctm_B + + print_block_metadata("Block encrypted result C", ctm_C) + + res = ctm_C.decrypt(keys.secretKey, unpack_type="original") + is_match, _ = validate_and_print_results(res, A @ B, description) + return is_match + + +def main(): + """ + Demonstrate homomorphic block matrix multiplication using OpenFHE-NumPy: + + 1) power-of-two dimensions: 8x8 @ 8x8 + 2) non-power-of-two dimensions: 3x3 @ 3x3 + + The block_shape is fixed to (2, 2). Each encrypted block stores a 2x2 + sub-matrix in 4 CKKS slots. Non-power-of-two matrices are zero-padded + to a 4x4 block grid and the result is cropped back to the original shape. + Both operands must use ROW_MAJOR packing for block matrix multiplication. + """ + + # Cryptographic setup + # Block matmul calls EvalMatMulSquare on each 2x2 block, which requires + # depth 2 per multiply-accumulate step; mult_depth=8 is sufficient for + # an 8x8 matrix decomposed into 2x2 blocks (4x4 block grid, depth ~4*2). + mult_depth = 8 + scale_mod_size = 59 + + params = CCParamsCKKSRNS() + params.SetMultiplicativeDepth(mult_depth) + params.SetScalingModSize(scale_mod_size) + params.SetFirstModSize(60) + params.SetScalingTechnique(FIXEDAUTO) + + cc = GenCryptoContext(params) + cc.Enable(PKESchemeFeature.PKE) + cc.Enable(PKESchemeFeature.LEVELEDSHE) + cc.Enable(PKESchemeFeature.ADVANCEDSHE) + + keys = cc.KeyGen() + cc.EvalMultKeyGen(keys.secretKey) + cc.EvalSumKeyGen(keys.secretKey) + + ring_dim = cc.GetRingDimension() + block_shape = (2, 2) + print(f"\nCKKS ring dimension: {ring_dim}") + print(f"Block shape: {block_shape}") + + all_ok = True + + # ============================================================ + # Case 1: power-of-two dimensions, 8x8 @ 8x8 + # ============================================================ + A8 = np.array( + [ + [0, 7, 8, 10, 1, 2, 7, 6], + [0, 1, 1, 9, 7, 5, 1, 7], + [8, 8, 4, 5, 8, 2, 6, 1], + [1, 0, 0, 1, 10, 3, 1, 7], + [7, 8, 2, 5, 3, 2, 10, 9], + [0, 3, 4, 10, 10, 5, 2, 5], + [2, 5, 0, 2, 8, 8, 5, 9], + [5, 1, 10, 6, 2, 8, 6, 3], + ], + dtype=float, + ) + + B8 = np.array( + [ + [6, 5, 4, 3, 2, 1, 0, 7], + [7, 1, 1, 2, 7, 5, 9, 3], + [4, 8, 8, 10, 8, 2, 1, 6], + [7, 0, 0, 5, 10, 3, 4, 2], + [9, 3, 2, 8, 3, 2, 1, 0], + [5, 2, 4, 1, 10, 5, 8, 2], + [9, 8, 0, 2, 8, 8, 7, 5], + [3, 6, 10, 1, 2, 8, 4, 0], + ], + dtype=float, + ) + + all_ok = all_ok and run_block_matmul_example( + cc=cc, + keys=keys, + A=A8, + B=B8, + block_shape=block_shape, + description="Block 8x8 Matrix Product (power-of-two)", + ) + + # ============================================================ + # Case 2: non-power-of-two dimensions, 3x3 @ 3x3 + # ============================================================ + # With block_shape=(2, 2), the 3x3 matrix is internally padded to a + # 4x4 block layout (2x2 grid of 2x2 blocks) and the result is + # cropped back to the original 3x3 shape on decryption. + A3 = np.array( + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + [7.0, 8.0, 9.0], + ], + dtype=float, + ) + + B3 = np.array( + [ + [9.0, 8.0, 7.0], + [6.0, 5.0, 4.0], + [3.0, 2.0, 1.0], + ], + dtype=float, + ) + + all_ok = all_ok and run_block_matmul_example( + cc=cc, + keys=keys, + A=A3, + B=B3, + block_shape=block_shape, + description="Block 3x3 Matrix Product (non-power-of-two)", + ) + + print("\n" + "=" * 60) + print(f"All checks passed: {all_ok}") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/examples/python/simple_matrix_accumulation.py b/examples/python/simple_matrix_accumulation.py index c2097fa..01c1969 100644 --- a/examples/python/simple_matrix_accumulation.py +++ b/examples/python/simple_matrix_accumulation.py @@ -27,11 +27,11 @@ def run_row_accumulation_example(cc, keys, ctm_x, matrix): raise ValueError("Invalid order.") # Perform homomorphic row accumulation - ctm_result = onp.cumulative_sum(ctm_x, axis=0) + ctm_result = onp.cumsum(ctm_x, axis=0) result = ctm_result.decrypt(keys.secretKey, unpack_type="original") # Compare with plain result - # NumPy >= 2.0: please use numpy.cumulative_sum instead of numpy.cumsum. + # NumPy >= 2.0: please use numpy.cumsum instead of numpy.cumsum. expected = np.cumsum(matrix, axis=0) # Validate and print results @@ -50,11 +50,11 @@ def run_column_accumulation_example(cc, keys, ctm_x, matrix): raise ValueError("Invalid order.") # Perform homomorphic column accumulation - ctm_result = onp.cumulative_sum(ctm_x, axis=1) + ctm_result = onp.cumsum(ctm_x, axis=1) result = ctm_result.decrypt(keys.secretKey, unpack_type="original") # Compare with plain result - # NumPy >= 2.0: please use numpy.cumulative_sum instead of numpy.cumsum. + # NumPy >= 2.0: please use numpy.cumsum instead of numpy.cumsum. expected = np.cumsum(matrix, axis=1) # Validate and print results diff --git a/examples/python/tests.py b/examples/python/tests.py new file mode 100644 index 0000000..5a2cafe --- /dev/null +++ b/examples/python/tests.py @@ -0,0 +1,79 @@ +import numpy as np +from openfhe import * +import openfhe_numpy as onp + +params = CCParamsCKKSRNS() +params.SetMultiplicativeDepth(4) +params.SetScalingModSize(59) +params.SetFirstModSize(60) +params.SetScalingTechnique(FIXEDAUTO) +cc = GenCryptoContext(params) +cc.Enable(PKESchemeFeature.PKE) +cc.Enable(PKESchemeFeature.LEVELEDSHE) +cc.Enable(PKESchemeFeature.ADVANCEDSHE) +keys = cc.KeyGen() +cc.EvalMultKeyGen(keys.secretKey) +cc.EvalSumKeyGen(keys.secretKey) +batch_size = cc.GetRingDimension() // 2 + +# Non-square: 2 rows, 4 cols +A = np.array( + [ + [1.0, 2.0, 3.0, 4.0], + [10.0, 20.0, 30.0, 40.0], + ] +) +print("A shape:", A.shape) +print("np.cumsum(A, axis=0):\n", np.cumsum(A, axis=0)) + +ctm_a = onp.array( + cc=cc, + data=A, + batch_size=batch_size, + order=onp.ROW_MAJOR, + fhe_type="C", + mode="tile", + public_key=keys.publicKey, +) +print("ctm_a.ncols (padded):", ctm_a.ncols, "original_shape:", ctm_a.original_shape) + +onp.gen_accumulate_rows_key(keys.secretKey, ctm_a.ncols) + +try: + res = ctm_a.cumsum(axis=0) + decrypted = res.decrypt(keys.secretKey, unpack_type="original") + print("\ncumsum(axis=0) result:\n", decrypted) + print("\nMatches numpy?", np.allclose(decrypted, np.cumsum(A, axis=0), atol=1e-1)) +except Exception as e: + print("cumsum(axis=0) FAILED:", type(e).__name__, e) + + +# ------------------------------------------------------------------ +# R1: 1-D vector cumsum -- this is the one that actually crashes. +# ------------------------------------------------------------------ +print("\n--- 1-D vector cumsum ---") +v = np.array([1.0, 2.0, 3.0, 4.0]) +print("v:", v) +print("np.cumsum(v):", np.cumsum(v)) + +ctm_v = onp.array( + cc=cc, + data=v, + batch_size=batch_size, + order=onp.ROW_MAJOR, + fhe_type="C", + mode="tile", + public_key=keys.publicKey, +) +print( + "ctm_v.ndim:", ctm_v.ndim, "ctm_v.ncols:", ctm_v.ncols, "original_shape:", ctm_v.original_shape +) + +onp.gen_accumulate_cols_key(keys.secretKey, ctm_v.ncols) + +try: + res_v = ctm_v.cumsum() # default axis=0 -- same crash as axis=None + decrypted_v = res_v.decrypt(keys.secretKey, unpack_type="original") + print("cumsum() result:", decrypted_v) +except Exception as e: + print("cumsum() FAILED:", type(e).__name__, e)