From f23b772333e8a97d1b0a415c2ff4bf44a09bbcc7 Mon Sep 17 00:00:00 2001 From: Wenhao Zhang Date: Tue, 11 Aug 2026 13:56:03 -0700 Subject: [PATCH] fix: implement physical slot replication for bicyclic and tricyclic layout after matrix multiplication. In the bicyclic layout, we require for each 0 <= slot < numSlot, (ct, slot) is mapped to (slot % n, slot % m). However, after computing through BSGS we can only guarantee for all 0 <= slot < n*m this property preserves (indeed we can guarantee more, but for the tail part it is not correct if n*m does not divide numSlot). In this commit, we add a new relation called `periodic replication relation` that replicate the first copy to all slots periodically. After each matrix multiplication with CRT layouts (bicyclic, tricyclic), we compose this relation to derive the result. The cost of the layout switching will be up to logN rotations. PiperOrigin-RevId: 962982741 --- lib/Kernel/BUILD | 6 + lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp | 10 +- lib/Kernel/BicyclicMatmulFuzzTest.cpp | 9 +- lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp | 30 +++-- .../ConvertToCiphertextSemantics.cpp | 47 ++++++- lib/Utils/Layout/Utils.cpp | 36 +++++ lib/Utils/Layout/Utils.h | 6 + lib/Utils/Layout/UtilsTest.cpp | 18 +++ .../Transforms/implement_shift_network.mlir | 12 ++ .../common/bicyclic_matmul_chain.mlir | 7 + .../Examples/openfhe/ckks/batch_matmul/BUILD | 3 +- .../openfhe/ckks/bicyclic_matmul/BUILD | 1 + .../openfhe/ckks/bicyclic_matmul_pt/BUILD | 1 + .../Examples/plaintext/bicyclic_matmul/BUILD | 9 ++ .../bicyclic_matmul_chain_test.cpp | 125 ++++++++++++++++++ .../batch_matmul.mlir | 3 + .../matmul.mlir | 3 + .../matmul_pt.mlir | 6 + 18 files changed, 316 insertions(+), 16 deletions(-) create mode 100644 tests/Examples/common/bicyclic_matmul_chain.mlir create mode 100644 tests/Examples/plaintext/bicyclic_matmul/bicyclic_matmul_chain_test.cpp diff --git a/lib/Kernel/BUILD b/lib/Kernel/BUILD index 0b587f5249..4c1af391bb 100644 --- a/lib/Kernel/BUILD +++ b/lib/Kernel/BUILD @@ -211,8 +211,10 @@ cc_test( ":KernelImplementation", "@fuzztest//fuzztest", "@googletest//:gtest_main", + "@heir//lib/Utils:MathUtils", "@heir//lib/Utils/Layout:Evaluate", "@heir//lib/Utils/Layout:Utils", + "@llvm-project//mlir:Analysis", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], @@ -228,8 +230,10 @@ cc_test( ":KernelImplementation", "@fuzztest//fuzztest", "@googletest//:gtest_main", + "@heir//lib/Utils:MathUtils", "@heir//lib/Utils/Layout:Evaluate", "@heir//lib/Utils/Layout:Utils", + "@llvm-project//mlir:Analysis", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], @@ -245,8 +249,10 @@ cc_test( ":KernelImplementation", "@fuzztest//fuzztest", "@googletest//:gtest_main", + "@heir//lib/Utils:MathUtils", "@heir//lib/Utils/Layout:Evaluate", "@heir//lib/Utils/Layout:Utils", + "@llvm-project//mlir:Analysis", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], diff --git a/lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp b/lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp index 1f864cc196..03f379ed2b 100644 --- a/lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp +++ b/lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp @@ -12,6 +12,8 @@ #include "lib/Kernel/KernelImplementation.h" #include "lib/Utils/Layout/Evaluate.h" #include "lib/Utils/Layout/Utils.h" +#include "lib/Utils/MathUtils.h" +#include "mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h" // from @llvm-project #include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project @@ -29,7 +31,8 @@ std::vector> runDiagonalMatmul(bool isCtPt, int64_t m, int64_t n, int64_t p, bool unroll = true) { MLIRContext context; - int64_t numSlots = 2 * m * n * p; + int64_t minSlots = (isCtPt ? m * n : n * p) + m * p; + int64_t numSlots = nextPowerOfTwo(minSlots); int64_t rowsCt = isCtPt ? m : n; int64_t colsCt = isCtPt ? n : p; @@ -71,7 +74,10 @@ std::vector> runDiagonalMatmul(bool isCtPt, auto resultLayout = getBicyclicLayoutRelation( RankedTensorType::get({m, p}, mlir::IndexType::get(&context)), numSlots); - + // Restrict the unpacking to the first output period. + addBounds(resultLayout, + resultLayout.getVarKindOffset(presburger::VarKind::Range) + 1, 0, + m * p - 1); return unpackLayoutToMatrix(resultLayout, {resultVec}, {m, p}); } diff --git a/lib/Kernel/BicyclicMatmulFuzzTest.cpp b/lib/Kernel/BicyclicMatmulFuzzTest.cpp index 4c1c5feff0..cf372f9d7f 100644 --- a/lib/Kernel/BicyclicMatmulFuzzTest.cpp +++ b/lib/Kernel/BicyclicMatmulFuzzTest.cpp @@ -11,6 +11,8 @@ #include "lib/Kernel/KernelImplementation.h" #include "lib/Utils/Layout/Evaluate.h" #include "lib/Utils/Layout/Utils.h" +#include "lib/Utils/MathUtils.h" +#include "mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h" // from @llvm-project #include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project @@ -28,7 +30,8 @@ std::vector> runBicyclicMatmul(const std::vector& vecA, int64_t m, int64_t n, int64_t p) { MLIRContext context; - int64_t numSlots = m * n * p; + int64_t minSlots = m * n + n * p + m * p; + int64_t numSlots = nextPowerOfTwo(minSlots); auto layoutA = getBicyclicLayoutRelation( RankedTensorType::get({m, n}, mlir::IndexType::get(&context)), numSlots); @@ -52,6 +55,10 @@ std::vector> runBicyclicMatmul(const std::vector& vecA, auto resultLayout = getBicyclicLayoutRelation( RankedTensorType::get({m, p}, mlir::IndexType::get(&context)), numSlots); + // Restrict the unpacking to the first output period. + addBounds(resultLayout, + resultLayout.getVarKindOffset(presburger::VarKind::Range) + 1, 0, + m * p - 1); return unpackLayoutToMatrix(resultLayout, {resultVec}, {m, p}); } diff --git a/lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp b/lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp index 334f1a8520..3c254294ec 100644 --- a/lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp +++ b/lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp @@ -11,6 +11,8 @@ #include "lib/Kernel/KernelImplementation.h" #include "lib/Utils/Layout/Evaluate.h" #include "lib/Utils/Layout/Utils.h" +#include "lib/Utils/MathUtils.h" +#include "mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h" // from @llvm-project #include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project @@ -34,7 +36,8 @@ void tricyclicBatchMatmulMatchesNaive( } MLIRContext context; - int64_t numSlots = h * m * n * p; + int64_t minSlots = h * (m * n + n * p + m * p); + int64_t numSlots = nextPowerOfTwo(minSlots); RankedTensorType typeA = RankedTensorType::get({h, m, n}, mlir::IndexType::get(&context)); @@ -65,18 +68,27 @@ void tricyclicBatchMatmulMatchesNaive( RankedTensorType resultType = RankedTensorType::get({h, m, p}, mlir::IndexType::get(&context)); auto resultLayout = getTricyclicLayoutRelation(resultType, numSlots); - auto expectedPacked = - evaluateLayout(resultLayout, [&](const std::vector& pt) { - int64_t ih = pt[0], im = pt[1], ip = pt[2]; - int sum = 0; + // Restrict the unpacking to the first output period. + addBounds(resultLayout, + resultLayout.getVarKindOffset(presburger::VarKind::Range) + 1, 0, + h * m * p - 1); + auto actual = + unpackLayoutTo3DTensor(resultLayout, {actualVec}, {h, m, p}); + + std::vector>> expected( + h, std::vector>(m, std::vector(p, 0))); + for (int64_t ih = 0; ih < h; ++ih) { + for (int64_t im = 0; im < m; ++im) { + for (int64_t ip = 0; ip < p; ++ip) { for (int64_t in = 0; in < n; ++in) { - sum += + expected[ih][im][ip] += vecA[ih * m * n + im * n + in] * vecB[ih * n * p + in * p + ip]; } - return sum; - }); + } + } + } - EXPECT_EQ(expectedPacked[0], actualVec); + EXPECT_EQ(expected, actual); } auto tricyclicShapeAndTensors() { diff --git a/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp b/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp index aba5f7265f..b4cdd5fcb7 100644 --- a/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp +++ b/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp @@ -155,6 +155,23 @@ Operation* remapAndExtractResult(ImplicitLocOpBuilder& builder, Value input, return extractRemap; } +// Rebuilds the full periodic layout of a kernel's output from its first valid +// period. Assumes the first period is uncorrupted by wrap-around bounds. +Operation* replicateFirstPeriodOfResult(ImplicitLocOpBuilder& b, Value input, + LayoutAttr resultLayout, + int64_t period) { + auto ctSemanticType = cast(input.getType()); + int64_t numCiphertexts = ctSemanticType.getDimSize(0); + int64_t numSlots = ctSemanticType.getDimSize(1); + IntegerRelation replication = + getPeriodicReplicationRelation(numCiphertexts, numSlots, period); + LayoutAttr replicationMapping = + LayoutAttr::getFromIntegerRelation(b.getContext(), replication); + auto remapOp = tensor_ext::RemapOp::create(b, input, replicationMapping); + remapOp->setAttr(kLayoutAttrName, resultLayout); + return remapOp; +} + } // namespace // An unset value of a permutation as it's being built up. @@ -2807,7 +2824,15 @@ struct ConvertLinalgMatmul makeAppropriatelyTypedAddOp(b, op->getLoc(), finalOutput, result); addBias->setAttr(kLayoutAttrName, layoutAttr); setMaterializedAttr(addBias); - rewriter.replaceOp(op, addBias); + + // Rebuild the full periodic output layout from the first period. + auto dataSemanticResultType = + cast(op->getResult(0).getType()); + Operation* replicated = + replicateFirstPeriodOfResult(b, addBias->getResult(0), layoutAttr, + dataSemanticResultType.getNumElements()); + setMaterializedAttr(replicated); + rewriter.replaceOp(op, replicated); } bool supportsBicyclic(linalg::MatmulOp op, OpAdaptor adaptor) const { @@ -2856,7 +2881,15 @@ struct ConvertLinalgMatmul makeAppropriatelyTypedAddOp(b, op->getLoc(), finalOutput, result); addBias->setAttr(kLayoutAttrName, layoutAttr); setMaterializedAttr(addBias); - rewriter.replaceOp(op, addBias); + + // Rebuild the full periodic output layout from the first period. + auto dataSemanticResultType = + cast(op->getResult(0).getType()); + Operation* replicated = + replicateFirstPeriodOfResult(b, addBias->getResult(0), layoutAttr, + dataSemanticResultType.getNumElements()); + setMaterializedAttr(replicated); + rewriter.replaceOp(op, replicated); } LogicalResult matchAndRewrite( @@ -2934,7 +2967,15 @@ struct ConvertLinalgBatchMatmul makeAppropriatelyTypedAddOp(b, op->getLoc(), finalOutput, result); addBias->setAttr(kLayoutAttrName, layoutAttr); setMaterializedAttr(addBias); - rewriter.replaceOp(op, addBias); + + // Rebuild the full periodic output layout from the first period. + auto dataSemanticResultType = + cast(op->getResult(0).getType()); + Operation* replicated = + replicateFirstPeriodOfResult(b, addBias->getResult(0), layoutAttr, + dataSemanticResultType.getNumElements()); + setMaterializedAttr(replicated); + rewriter.replaceOp(op, replicated); } LogicalResult matchAndRewrite( diff --git a/lib/Utils/Layout/Utils.cpp b/lib/Utils/Layout/Utils.cpp index 94e0e6f1e0..b05d686640 100644 --- a/lib/Utils/Layout/Utils.cpp +++ b/lib/Utils/Layout/Utils.cpp @@ -469,6 +469,42 @@ presburger::IntegerRelation getTricyclicLayoutRelation( return result; } +presburger::IntegerRelation getPeriodicReplicationRelation( + int64_t numCiphertexts, int64_t numSlots, int64_t period) { + assert(numCiphertexts == 1 && "only support single ciphertext layout"); + assert(period > 0 && period <= numSlots && + "period must be positive and at most numSlots"); + + IntegerRelation result(PresburgerSpace::getRelationSpace( + /*numDomain=*/2, /*numRange=*/2, /*numSymbol=*/0, + /*numLocals=*/0)); + + int domainOffset = result.getVarKindOffset(VarKind::Domain); + int rangeOffset = result.getVarKindOffset(VarKind::Range); + int sourceCtIndex = domainOffset; + int sourceSlotIndex = domainOffset + 1; + int targetCtIndex = rangeOffset; + int targetSlotIndex = rangeOffset + 1; + + addBounds(result, sourceCtIndex, 0, numCiphertexts - 1); + addBounds(result, sourceSlotIndex, 0, period - 1); + addBounds(result, targetSlotIndex, 0, numSlots - 1); + + addConstraint(result, {{sourceCtIndex, 1}, {targetCtIndex, -1}}, + /*equality=*/true); + + // source_slot = target_slot % period + SmallVector targetSlotCoeffs(result.getNumCols(), 0); + targetSlotCoeffs[targetSlotIndex] = 1; + auto targetSlotMod = addModConstraint(result, targetSlotCoeffs, period); + SmallVector sourceEquality(result.getNumCols(), 0); + sourceEquality[sourceSlotIndex] = 1; + sourceEquality[targetSlotMod] = -1; + result.addEquality(sourceEquality); + + return result; +} + presburger::IntegerRelation getPerRowLayoutRelation(RankedTensorType matrixType, int64_t minSlotCount) { auto domainSize = matrixType.getRank(); diff --git a/lib/Utils/Layout/Utils.h b/lib/Utils/Layout/Utils.h index bba0e06ae0..0fdd1ea102 100644 --- a/lib/Utils/Layout/Utils.h +++ b/lib/Utils/Layout/Utils.h @@ -96,6 +96,12 @@ presburger::IntegerRelation getBicyclicDiagonalRelation( RankedTensorType matrixType, int64_t contractionDim, int64_t stride, int64_t numSlots); +// Returns an IntegerRelation with domain and range space both (ct, slot) that +// maps each slot s in [0, period) of a ciphertext to every slot s' in [0, +// numSlots) with s' equiv s (mod period). Excepts numCiphertexts == 1. +presburger::IntegerRelation getPeriodicReplicationRelation( + int64_t numCiphertexts, int64_t numSlots, int64_t period); + // Returns an IntegerRelation that represents a per-row layout for a matrix // such that each row of the matrix is in a separate ciphertext. presburger::IntegerRelation getPerRowLayoutRelation(RankedTensorType matrixType, diff --git a/lib/Utils/Layout/UtilsTest.cpp b/lib/Utils/Layout/UtilsTest.cpp index b29904265d..445d17c699 100644 --- a/lib/Utils/Layout/UtilsTest.cpp +++ b/lib/Utils/Layout/UtilsTest.cpp @@ -240,6 +240,24 @@ TEST(UtilsTest, BicyclicLayout3x5Repeated) { EXPECT_EQ(packedMatrix, expected); } +TEST(UtilsTest, PeriodicReplicationRelation) { + int64_t numSlots = 10; + int64_t period = 3; + IntegerRelation replication = + getPeriodicReplicationRelation(/*numCiphertexts=*/1, numSlots, period); + + // Every target slot t is reached exactly from source slot t % period. + for (int64_t t = 0; t < numSlots; ++t) { + for (int64_t s = 0; s < period; ++s) { + EXPECT_EQ(replication.containsPointNoLocal({0, s, 0, t}).has_value(), + s == t % period); + } + } + + // Source slots outside the first period are not in the domain. + EXPECT_FALSE(replication.containsPointNoLocal({0, period, 0, period})); +} + TEST(UtilsTest, BicyclicCtPtDiagonal3x5x7) { MLIRContext context; int64_t numSlots = 105; diff --git a/tests/Dialect/TensorExt/Transforms/implement_shift_network.mlir b/tests/Dialect/TensorExt/Transforms/implement_shift_network.mlir index 161c4314ea..73d4baa3e8 100644 --- a/tests/Dialect/TensorExt/Transforms/implement_shift_network.mlir +++ b/tests/Dialect/TensorExt/Transforms/implement_shift_network.mlir @@ -86,3 +86,15 @@ func.func @multi_ciphertext_complex(%0: tensor<4x64xi32>) -> tensor<4x64xi32> { %1 = tensor_ext.remap %0 {permutation = #map5} : tensor<4x64xi32> return %1 : tensor<4x64xi32> } + +// CHECK: func.func @periodic_replication +// CHECK-NOT: tensor_ext.remap +// CHECK: tensor.extract_slice +// CHECK: tensor_ext.rotate +// CHECK: tensor.insert_slice +#layout_bicyclic = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (4i0 + 5i1 + slot) mod 30 = 0 and 0 <= i0 <= 2 and 0 <= i1 <= 1 and 0 <= slot <= 1023 }"> +#replication = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : i0 = 0 and ct = 0 and (-i1 + slot) mod 6 = 0 and 0 <= i1 <= 5 and 0 <= slot <= 1023 }"> +func.func @periodic_replication(%arg0: tensor<1x1024xi16> {tensor_ext.layout = #layout_bicyclic}) -> (tensor<1x1024xi16> {tensor_ext.layout = #layout_bicyclic}) { + %0 = tensor_ext.remap %arg0 {permutation = #replication} : tensor<1x1024xi16> + return %0 : tensor<1x1024xi16> +} diff --git a/tests/Examples/common/bicyclic_matmul_chain.mlir b/tests/Examples/common/bicyclic_matmul_chain.mlir new file mode 100644 index 0000000000..f5db4f4459 --- /dev/null +++ b/tests/Examples/common/bicyclic_matmul_chain.mlir @@ -0,0 +1,7 @@ +func.func @bicyclic_matmul_chain(%arg0: tensor<13x18xf32> {secret.secret}, %arg1: tensor<18x16xf32>, %arg2: tensor<16x9xf32> {secret.secret}) -> tensor<13x9xf32> { + %cst0 = arith.constant dense<0.000000e+00> : tensor<13x16xf32> + %cst1 = arith.constant dense<0.000000e+00> : tensor<13x9xf32> + %0 = linalg.matmul ins(%arg0, %arg1 : tensor<13x18xf32>, tensor<18x16xf32>) outs(%cst0 : tensor<13x16xf32>) -> tensor<13x16xf32> + %1 = linalg.matmul ins(%0, %arg2 : tensor<13x16xf32>, tensor<16x9xf32>) outs(%cst1 : tensor<13x9xf32>) -> tensor<13x9xf32> + return %1 : tensor<13x9xf32> +} diff --git a/tests/Examples/openfhe/ckks/batch_matmul/BUILD b/tests/Examples/openfhe/ckks/batch_matmul/BUILD index f2d16e68d3..87ef2a076c 100644 --- a/tests/Examples/openfhe/ckks/batch_matmul/BUILD +++ b/tests/Examples/openfhe/ckks/batch_matmul/BUILD @@ -4,10 +4,11 @@ package(default_applicable_licenses = ["@heir//:license"]) openfhe_end_to_end_test( name = "batch_matmul_test", + size = "medium", generated_lib_header = "batch_matmul_lib.h", heir_opt_flags = [ "--annotate-module=backend=openfhe scheme=ckks", - "--mlir-to-ckks=min-slot-count=4096 experimental-disable-loop-unroll=true greedy-level-budget=40 first-mod-bits=60 scaling-mod-bits=50", + "--mlir-to-ckks=min-slot-count=4096 experimental-disable-loop-unroll=true greedy-level-budget=40 first-mod-bits=60 scaling-mod-bits=50 greedy-bootstrap-waterline=20", "--scheme-to-openfhe", ], mlir_src = "@heir//tests/Examples/common:batch_matmul.mlir", diff --git a/tests/Examples/openfhe/ckks/bicyclic_matmul/BUILD b/tests/Examples/openfhe/ckks/bicyclic_matmul/BUILD index 2129fa6cab..de53464a30 100644 --- a/tests/Examples/openfhe/ckks/bicyclic_matmul/BUILD +++ b/tests/Examples/openfhe/ckks/bicyclic_matmul/BUILD @@ -4,6 +4,7 @@ package(default_applicable_licenses = ["@heir//:license"]) openfhe_end_to_end_test( name = "bicyclic_matmul_test", + size = "medium", generated_lib_header = "bicyclic_matmul_lib.h", heir_opt_flags = [ "--annotate-module=backend=openfhe scheme=ckks", diff --git a/tests/Examples/openfhe/ckks/bicyclic_matmul_pt/BUILD b/tests/Examples/openfhe/ckks/bicyclic_matmul_pt/BUILD index 5017e87c79..744b0c27bd 100644 --- a/tests/Examples/openfhe/ckks/bicyclic_matmul_pt/BUILD +++ b/tests/Examples/openfhe/ckks/bicyclic_matmul_pt/BUILD @@ -4,6 +4,7 @@ package(default_applicable_licenses = ["@heir//:license"]) openfhe_end_to_end_test( name = "bicyclic_matmul_pt_test", + size = "medium", generated_lib_header = "bicyclic_matmul_pt_lib.h", heir_opt_flags = [ "--annotate-module=backend=openfhe scheme=ckks", diff --git a/tests/Examples/plaintext/bicyclic_matmul/BUILD b/tests/Examples/plaintext/bicyclic_matmul/BUILD index 7840d95b3f..c5b117bb15 100644 --- a/tests/Examples/plaintext/bicyclic_matmul/BUILD +++ b/tests/Examples/plaintext/bicyclic_matmul/BUILD @@ -10,3 +10,12 @@ llvm_runner_test( main_c_src = "bicyclic_matmul_test.cpp", mlir_src = "@heir//tests/Examples/common:bicyclic_matmul.mlir", ) + +llvm_runner_test( + name = "bicyclic_matmul_chain_test", + heir_opt_flags = [ + "--mlir-to-plaintext-backend=plaintext-size=2048", + ], + main_c_src = "bicyclic_matmul_chain_test.cpp", + mlir_src = "@heir//tests/Examples/common:bicyclic_matmul_chain.mlir", +) diff --git a/tests/Examples/plaintext/bicyclic_matmul/bicyclic_matmul_chain_test.cpp b/tests/Examples/plaintext/bicyclic_matmul/bicyclic_matmul_chain_test.cpp new file mode 100644 index 0000000000..f0b80a28bf --- /dev/null +++ b/tests/Examples/plaintext/bicyclic_matmul/bicyclic_matmul_chain_test.cpp @@ -0,0 +1,125 @@ +#include +#include +#include + +#include "gtest/gtest.h" // from @googletest +#include "tests/llvm_runner/memref_types.h" + +#if defined(__has_feature) +#if __has_feature(memory_sanitizer) +#include +#define HEIR_MSAN_UNPOISON(p, s) __msan_unpoison((p), (s)) +#else +#define HEIR_MSAN_UNPOISON(p, s) +#endif +#else +#define HEIR_MSAN_UNPOISON(p, s) +#endif + +extern "C" { +void _mlir_ciface_bicyclic_matmul_chain(StridedMemRefType* res, + StridedMemRefType* arg0, + StridedMemRefType* arg1, + StridedMemRefType* arg2); + +void _mlir_ciface_bicyclic_matmul_chain__encrypt__arg0( + StridedMemRefType* res, StridedMemRefType* arg); +void _mlir_ciface_bicyclic_matmul_chain__encrypt__arg2( + StridedMemRefType* res, StridedMemRefType* arg); + +void _mlir_ciface_bicyclic_matmul_chain__decrypt__result0( + StridedMemRefType* res, StridedMemRefType* arg); +} + +TEST(ChainedMatmulBicyclicPlaintextTest, Test1) { + std::vector arg0(13 * 18, 0.0); + std::vector arg1(18 * 16, 0.0); + std::vector arg2(16 * 9, 0.0); + + // A[i][j] = i + j + for (int i = 0; i < 13; ++i) { + for (int j = 0; j < 18; ++j) { + arg0[i * 18 + j] = (i + j) / 100.0; + } + } + + // B[j][k] = j - k + for (int j = 0; j < 18; ++j) { + for (int k = 0; k < 16; ++k) { + arg1[j * 16 + k] = (j - k) / 100.0; + } + } + + // V[k][l] = k - 2*l + for (int k = 0; k < 16; ++k) { + for (int l = 0; l < 9; ++l) { + arg2[k * 9 + l] = (k - 2 * l) / 100.0; + } + } + + // Y[i][k] = sum_j A[i][j] * B[j][k] + std::vector intermediate(13 * 16, 0.0); + for (int i = 0; i < 13; ++i) { + for (int k = 0; k < 16; ++k) { + for (int j = 0; j < 18; ++j) { + intermediate[i * 16 + k] += arg0[i * 18 + j] * arg1[j * 16 + k]; + } + } + } + + // Z[i][l] = sum_k Y[i][k] * V[k][l] + std::vector expected(13 * 9, 0.0); + for (int i = 0; i < 13; ++i) { + for (int l = 0; l < 9; ++l) { + for (int k = 0; k < 16; ++k) { + expected[i * 9 + l] += intermediate[i * 16 + k] * arg2[k * 9 + l]; + } + } + } + + int64_t sizes0[2] = {13, 18}; + int64_t strides0[2] = {18, 1}; + StridedMemRefType inputs0(arg0.data(), arg0.data(), 0, sizes0, + strides0); + + int64_t sizes1[2] = {18, 16}; + int64_t strides1[2] = {16, 1}; + StridedMemRefType inputs1(arg1.data(), arg1.data(), 0, sizes1, + strides1); + + int64_t sizes2[2] = {16, 9}; + int64_t strides2[2] = {9, 1}; + StridedMemRefType inputs2(arg2.data(), arg2.data(), 0, sizes2, + strides2); + + StridedMemRefType encArg0; + _mlir_ciface_bicyclic_matmul_chain__encrypt__arg0(&encArg0, &inputs0); + HEIR_MSAN_UNPOISON(&encArg0, sizeof(StridedMemRefType)); + + StridedMemRefType encArg2; + _mlir_ciface_bicyclic_matmul_chain__encrypt__arg2(&encArg2, &inputs2); + HEIR_MSAN_UNPOISON(&encArg2, sizeof(StridedMemRefType)); + + StridedMemRefType packedRes; + _mlir_ciface_bicyclic_matmul_chain(&packedRes, &encArg0, &inputs1, &encArg2); + HEIR_MSAN_UNPOISON(&packedRes, sizeof(StridedMemRefType)); + + StridedMemRefType outRef; + _mlir_ciface_bicyclic_matmul_chain__decrypt__result0(&outRef, &packedRes); + HEIR_MSAN_UNPOISON(&outRef, sizeof(StridedMemRefType)); + HEIR_MSAN_UNPOISON(outRef.basePtr, 13 * 9 * sizeof(float)); + + float errorThreshold = 1e-3; + for (int i = 0; i < 13; ++i) { + for (int l = 0; l < 9; ++l) { + float actual = outRef.data[i * outRef.strides[0] + l * outRef.strides[1]]; + EXPECT_NEAR(expected[i * 9 + l], actual, errorThreshold) + << "mismatch at (" << i << ", " << l << ")"; + } + } + + free(encArg0.basePtr); + free(encArg2.basePtr); + free(packedRes.basePtr); + free(outRef.basePtr); +} diff --git a/tests/Transforms/convert_to_ciphertext_semantics/batch_matmul.mlir b/tests/Transforms/convert_to_ciphertext_semantics/batch_matmul.mlir index c872f5c34f..2541f93dae 100644 --- a/tests/Transforms/convert_to_ciphertext_semantics/batch_matmul.mlir +++ b/tests/Transforms/convert_to_ciphertext_semantics/batch_matmul.mlir @@ -5,8 +5,11 @@ #layout2 = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : ct = 0 and (399i0 - 210i1 - 190i2 + slot) mod 798 = 0 and 0 <= i0 <= 1 and 0 <= i1 <= 18 and 0 <= i2 <= 20 and 0 <= slot <= 8191 }"> #layout3 = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : ct = 0 and (-323i0 - 19i1 - i2 + slot) mod 1024 = 0 and 0 <= i0 <= 1 and 0 <= i1 <= 16 and 0 <= i2 <= 8191 - 323i0 - 19i1 and i2 <= 18 and 0 <= slot <= 8191 and 8192*floor((-1024 + 323i0 + 19i1 + i2)/8192) <= -8192 + 323i0 + 19i1 + i2 }"> module { + // CHECK: #[[replication:layout[0-9]*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : i0 = 0 and ct = 0 and (-i1 + slot) mod 714 = 0 and 0 <= i1 <= 713 and 0 <= slot <= 8191 }"> // CHECK: @batch_matmul_secret_secret // CHECK-NOT: linalg.batch_matmul + // CHECK: tensor_ext.remap + // CHECK-SAME: permutation = #[[replication]] func.func @batch_matmul_secret_secret(%arg0: !secret.secret> {tensor_ext.layout = #layout1}, %arg1: !secret.secret> {tensor_ext.layout = #layout2}) -> (!secret.secret> {tensor_ext.layout = #layout}) { %cst = arith.constant dense<0.000000e+00> : tensor<2x17x21xf32> %0 = secret.generic(%arg0: !secret.secret> {tensor_ext.layout = #layout1}, %arg1: !secret.secret> {tensor_ext.layout = #layout2}) { diff --git a/tests/Transforms/convert_to_ciphertext_semantics/matmul.mlir b/tests/Transforms/convert_to_ciphertext_semantics/matmul.mlir index 14876af5e5..ea01b4becd 100644 --- a/tests/Transforms/convert_to_ciphertext_semantics/matmul.mlir +++ b/tests/Transforms/convert_to_ciphertext_semantics/matmul.mlir @@ -5,8 +5,11 @@ #layout2 = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (4i0 + 5i1 + slot) mod 10 = 0 and 0 <= i0 <= 4 and 0 <= i1 <= 1 and 0 <= slot <= 1023 }"> #layout3 = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (-5i0 - i1 + slot) mod 16 = 0 and 0 <= i0 <= 2 and 0 <= i1 <= 1023 - 5i0 and i1 <= 4 and 0 <= slot <= 1023 and 1024*floor((-16 + 5i0 + i1)/1024) <= -1024 + 5i0 + i1 }"> module { + // CHECK: #[[replication:layout[0-9]*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : i0 = 0 and ct = 0 and (-i1 + slot) mod 6 = 0 and 0 <= i1 <= 5 and 0 <= slot <= 1023 }"> // CHECK: @matmul_secret_secret // CHECK-NOT: linalg.matmul + // CHECK: tensor_ext.remap + // CHECK-SAME: permutation = #[[replication]] func.func @matmul_secret_secret(%arg0: !secret.secret> {tensor_ext.layout = #layout1}, %arg1: !secret.secret> {tensor_ext.layout = #layout2}) -> (!secret.secret> {tensor_ext.layout = #layout}) { %cst = arith.constant dense<0.000000e+00> : tensor<3x2xf32> %0 = secret.generic(%arg0: !secret.secret> {tensor_ext.layout = #layout1}, %arg1: !secret.secret> {tensor_ext.layout = #layout2}) { diff --git a/tests/Transforms/convert_to_ciphertext_semantics/matmul_pt.mlir b/tests/Transforms/convert_to_ciphertext_semantics/matmul_pt.mlir index 69604b9836..81e099d544 100644 --- a/tests/Transforms/convert_to_ciphertext_semantics/matmul_pt.mlir +++ b/tests/Transforms/convert_to_ciphertext_semantics/matmul_pt.mlir @@ -10,10 +10,14 @@ #layout_ct2 = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (14i0 - 15i1 + slot) mod 35 = 0 and 0 <= i0 <= 4 and 0 <= i1 <= 6 and 0 <= slot <= 63 }"> module { + // CHECK: #[[replication:layout[0-9]*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : i0 = 0 and ct = 0 and (-i1 + slot) mod 21 = 0 and 0 <= i1 <= 20 and 0 <= slot <= 1023 }"> + // CHECK: @matmul_ctpt // CHECK-NOT: linalg.matmul // CHECK: tensor_ext.rotate // CHECK: arith.mulf + // CHECK: tensor_ext.remap + // CHECK-SAME: permutation = #[[replication]] func.func @matmul_ctpt(%arg0: !secret.secret> {tensor_ext.layout = #layout_ct}, %arg1: tensor<5x7xf32>) -> (!secret.secret> {tensor_ext.layout = #layout_out}) { %cst = arith.constant dense<0.000000e+00> : tensor<3x7xf32> %0 = secret.generic(%arg0: !secret.secret> {tensor_ext.layout = #layout_ct}) { @@ -30,6 +34,8 @@ module { // CHECK-NOT: linalg.matmul // CHECK: tensor_ext.rotate // CHECK: arith.mulf + // CHECK: tensor_ext.remap + // CHECK-SAME: permutation = #[[replication]] func.func @matmul_ptct(%arg0: tensor<3x5xf32>, %arg1: !secret.secret> {tensor_ext.layout = #layout_ct2}) -> (!secret.secret> {tensor_ext.layout = #layout_out}) { %cst = arith.constant dense<0.000000e+00> : tensor<3x7xf32> %0 = secret.generic(%arg1: !secret.secret> {tensor_ext.layout = #layout_ct2}) {