From d581622e76b592ddbba14e228dbc4739055130bb Mon Sep 17 00:00:00 2001 From: Wenhao Zhang Date: Wed, 12 Aug 2026 09:54:02 -0700 Subject: [PATCH] feat: optimize matmul periodic replication by using maximal valid prefix. Replaces the single-period replication with a dynamically calculated maximal valid prefix (the actual uncorrupted output slots before replication). The replication period is therefore based on the valid prefix. PiperOrigin-RevId: 963505841 --- lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp | 5 +- lib/Kernel/BicyclicMatmulFuzzTest.cpp | 5 +- lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp | 5 +- .../ConvertToCiphertextSemantics.cpp | 71 ++++++++++++++----- .../Examples/openfhe/ckks/batch_matmul/BUILD | 3 +- .../openfhe/ckks/bicyclic_matmul/BUILD | 1 - .../openfhe/ckks/bicyclic_matmul_pt/BUILD | 1 - .../batch_matmul.mlir | 2 +- .../matmul.mlir | 2 +- .../matmul_pt.mlir | 7 +- 10 files changed, 69 insertions(+), 33 deletions(-) diff --git a/lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp b/lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp index 03f379ed2b..bb1261c600 100644 --- a/lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp +++ b/lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp @@ -74,10 +74,11 @@ 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. + // Restrict the unpacking to the reach-derived valid prefix. + int64_t validPrefix = numSlots - period * (steps - 1); addBounds(resultLayout, resultLayout.getVarKindOffset(presburger::VarKind::Range) + 1, 0, - m * p - 1); + validPrefix - 1); return unpackLayoutToMatrix(resultLayout, {resultVec}, {m, p}); } diff --git a/lib/Kernel/BicyclicMatmulFuzzTest.cpp b/lib/Kernel/BicyclicMatmulFuzzTest.cpp index cf372f9d7f..f9f7eb24b7 100644 --- a/lib/Kernel/BicyclicMatmulFuzzTest.cpp +++ b/lib/Kernel/BicyclicMatmulFuzzTest.cpp @@ -55,10 +55,11 @@ 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. + // Restrict the unpacking to the reach-derived valid prefix. + int64_t validPrefix = numSlots - (n * p - 1 + m * (n - 1)); addBounds(resultLayout, resultLayout.getVarKindOffset(presburger::VarKind::Range) + 1, 0, - m * p - 1); + validPrefix - 1); return unpackLayoutToMatrix(resultLayout, {resultVec}, {m, p}); } diff --git a/lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp b/lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp index 3c254294ec..d3d642ce4a 100644 --- a/lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp +++ b/lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp @@ -68,10 +68,11 @@ void tricyclicBatchMatmulMatchesNaive( RankedTensorType resultType = RankedTensorType::get({h, m, p}, mlir::IndexType::get(&context)); auto resultLayout = getTricyclicLayoutRelation(resultType, numSlots); - // Restrict the unpacking to the first output period. + // Restrict the unpacking to the reach-derived valid prefix. + int64_t validPrefix = numSlots - (h * n * p - 1 + h * m * (n - 1)); addBounds(resultLayout, resultLayout.getVarKindOffset(presburger::VarKind::Range) + 1, 0, - h * m * p - 1); + validPrefix - 1); auto actual = unpackLayoutTo3DTensor(resultLayout, {actualVec}, {h, m, p}); diff --git a/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp b/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp index b4cdd5fcb7..48231b1f7e 100644 --- a/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp +++ b/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp @@ -155,16 +155,20 @@ 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, +// Rebuilds the full periodic layout of a kernel's output from the valid prefix +// of the periodic pattern. Greedily replicates the valid periodic prefix of the +// layout until the all of the slot count is covered. +Operation* replicateValidPrefixOfResult(ImplicitLocOpBuilder& b, Value input, LayoutAttr resultLayout, - int64_t period) { + int64_t inputPeriod, + int64_t validPrefix) { auto ctSemanticType = cast(input.getType()); int64_t numCiphertexts = ctSemanticType.getDimSize(0); int64_t numSlots = ctSemanticType.getDimSize(1); - IntegerRelation replication = - getPeriodicReplicationRelation(numCiphertexts, numSlots, period); + int64_t actualValidPrefix = + validPrefix > 0 ? (validPrefix / inputPeriod) * inputPeriod : 0; + IntegerRelation replication = getPeriodicReplicationRelation( + numCiphertexts, numSlots, actualValidPrefix); LayoutAttr replicationMapping = LayoutAttr::getFromIntegerRelation(b.getContext(), replication); auto remapOp = tensor_ext::RemapOp::create(b, input, replicationMapping); @@ -2825,12 +2829,20 @@ struct ConvertLinalgMatmul addBias->setAttr(kLayoutAttrName, layoutAttr); setMaterializedAttr(addBias); - // Rebuild the full periodic output layout from the first period. + // Rebuild the full periodic output layout from the widest valid + // period-aligned window. The rotation reach of rotate-and-reduce with + // `steps` iterations of stride `period` is period * (steps - 1). auto dataSemanticResultType = cast(op->getResult(0).getType()); - Operation* replicated = - replicateFirstPeriodOfResult(b, addBias->getResult(0), layoutAttr, - dataSemanticResultType.getNumElements()); + int64_t reach = period * (steps - 1); + auto ctSemanticResultType = + cast(addBias->getResult(0).getType()); + int64_t validPrefix = ctSemanticResultType.getDimSize(1) - reach; + LLVM_DEBUG(llvm::dbgs() << "Bicyclic diagonal matmul valid prefix: " + << validPrefix << "\n"); + Operation* replicated = replicateValidPrefixOfResult( + b, addBias->getResult(0), layoutAttr, + dataSemanticResultType.getNumElements(), validPrefix); setMaterializedAttr(replicated); rewriter.replaceOp(op, replicated); } @@ -2882,12 +2894,23 @@ struct ConvertLinalgMatmul addBias->setAttr(kLayoutAttrName, layoutAttr); setMaterializedAttr(addBias); - // Rebuild the full periodic output layout from the first period. + // Rebuild the full periodic output layout from the widest valid + // period-aligned window. For (m x n) * (n x p), the BSGS rotation reach + // is n * p - 1 + m * (n - 1). auto dataSemanticResultType = cast(op->getResult(0).getType()); - Operation* replicated = - replicateFirstPeriodOfResult(b, addBias->getResult(0), layoutAttr, - dataSemanticResultType.getNumElements()); + int64_t m = lhsType.getDimSize(0); + int64_t n = lhsType.getDimSize(1); + int64_t p = rhsType.getDimSize(1); + int64_t reach = n * p - 1 + m * (n - 1); + auto ctSemanticResultType = + cast(addBias->getResult(0).getType()); + int64_t validPrefix = ctSemanticResultType.getDimSize(1) - reach; + LLVM_DEBUG(llvm::dbgs() + << "Bicyclic matmul valid prefix: " << validPrefix << "\n"); + Operation* replicated = replicateValidPrefixOfResult( + b, addBias->getResult(0), layoutAttr, + dataSemanticResultType.getNumElements(), validPrefix); setMaterializedAttr(replicated); rewriter.replaceOp(op, replicated); } @@ -2968,12 +2991,24 @@ struct ConvertLinalgBatchMatmul addBias->setAttr(kLayoutAttrName, layoutAttr); setMaterializedAttr(addBias); - // Rebuild the full periodic output layout from the first period. + // Rebuild the full periodic output layout from the widest valid + // period-aligned window. For (h x m x n) * (h x n x p), the BSGS rotation + // reach is h * n * p - 1 + h * m * (n - 1). auto dataSemanticResultType = cast(op->getResult(0).getType()); - Operation* replicated = - replicateFirstPeriodOfResult(b, addBias->getResult(0), layoutAttr, - dataSemanticResultType.getNumElements()); + int64_t h = lhsType.getShape()[0]; + int64_t m = lhsType.getShape()[1]; + int64_t n = lhsType.getShape()[2]; + int64_t p = rhsType.getShape()[2]; + int64_t reach = h * n * p - 1 + h * m * (n - 1); + auto ctSemanticResultType = + cast(addBias->getResult(0).getType()); + int64_t validPrefix = ctSemanticResultType.getDimSize(1) - reach; + LLVM_DEBUG(llvm::dbgs() << "Tricyclic batch matmul valid prefix: " + << validPrefix << "\n"); + Operation* replicated = replicateValidPrefixOfResult( + b, addBias->getResult(0), layoutAttr, + dataSemanticResultType.getNumElements(), validPrefix); setMaterializedAttr(replicated); rewriter.replaceOp(op, replicated); } diff --git a/tests/Examples/openfhe/ckks/batch_matmul/BUILD b/tests/Examples/openfhe/ckks/batch_matmul/BUILD index 87ef2a076c..f2d16e68d3 100644 --- a/tests/Examples/openfhe/ckks/batch_matmul/BUILD +++ b/tests/Examples/openfhe/ckks/batch_matmul/BUILD @@ -4,11 +4,10 @@ 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 greedy-bootstrap-waterline=20", + "--mlir-to-ckks=min-slot-count=4096 experimental-disable-loop-unroll=true greedy-level-budget=40 first-mod-bits=60 scaling-mod-bits=50", "--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 de53464a30..2129fa6cab 100644 --- a/tests/Examples/openfhe/ckks/bicyclic_matmul/BUILD +++ b/tests/Examples/openfhe/ckks/bicyclic_matmul/BUILD @@ -4,7 +4,6 @@ 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 744b0c27bd..5017e87c79 100644 --- a/tests/Examples/openfhe/ckks/bicyclic_matmul_pt/BUILD +++ b/tests/Examples/openfhe/ckks/bicyclic_matmul_pt/BUILD @@ -4,7 +4,6 @@ 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/Transforms/convert_to_ciphertext_semantics/batch_matmul.mlir b/tests/Transforms/convert_to_ciphertext_semantics/batch_matmul.mlir index 2541f93dae..d4c2df0906 100644 --- a/tests/Transforms/convert_to_ciphertext_semantics/batch_matmul.mlir +++ b/tests/Transforms/convert_to_ciphertext_semantics/batch_matmul.mlir @@ -5,7 +5,7 @@ #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: #[[replication:layout[0-9]*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : i0 = 0 and ct = 0 and (-i1 + slot) mod 6426 = 0 and 0 <= i1 <= 6425 and 0 <= slot <= 8191 }"> // CHECK: @batch_matmul_secret_secret // CHECK-NOT: linalg.batch_matmul // CHECK: tensor_ext.remap diff --git a/tests/Transforms/convert_to_ciphertext_semantics/matmul.mlir b/tests/Transforms/convert_to_ciphertext_semantics/matmul.mlir index ea01b4becd..3c85b291bc 100644 --- a/tests/Transforms/convert_to_ciphertext_semantics/matmul.mlir +++ b/tests/Transforms/convert_to_ciphertext_semantics/matmul.mlir @@ -5,7 +5,7 @@ #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: #[[replication:layout[0-9]*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : i0 = 0 and ct = 0 and (-i1 + slot) mod 1002 = 0 and 0 <= i1 <= 1001 and 0 <= slot <= 1023 }"> // CHECK: @matmul_secret_secret // CHECK-NOT: linalg.matmul // CHECK: tensor_ext.remap diff --git a/tests/Transforms/convert_to_ciphertext_semantics/matmul_pt.mlir b/tests/Transforms/convert_to_ciphertext_semantics/matmul_pt.mlir index 81e099d544..b5a61a8ebe 100644 --- a/tests/Transforms/convert_to_ciphertext_semantics/matmul_pt.mlir +++ b/tests/Transforms/convert_to_ciphertext_semantics/matmul_pt.mlir @@ -10,14 +10,15 @@ #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: #[[replication1:layout[0-9]*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : i0 = 0 and ct = 0 and (-i1 + slot) mod 1008 = 0 and 0 <= i1 <= 1007 and 0 <= slot <= 1023 }"> + // CHECK: #[[replication2:layout[0-9]*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : i0 = 0 and ct = 0 and (-i1 + slot) mod 987 = 0 and 0 <= i1 <= 986 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]] + // CHECK-SAME: permutation = #[[replication1]] 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}) { @@ -35,7 +36,7 @@ module { // CHECK: tensor_ext.rotate // CHECK: arith.mulf // CHECK: tensor_ext.remap - // CHECK-SAME: permutation = #[[replication]] + // CHECK-SAME: permutation = #[[replication2]] 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}) {