Skip to content

feat: optimize matmul periodic replication by using maximal valid prefix. #3252

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Aug 12, 2026
Merged

feat: optimize matmul periodic replication by using maximal valid prefix. #3252

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,10 +74,11 @@ std::vector<std::vector<int>> 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<int>(resultLayout, {resultVec}, {m, p});
}

Expand Down
5 changes: 3 additions & 2 deletions lib/Kernel/BicyclicMatmulFuzzTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,10 +55,11 @@ std::vector<std::vector<int>> runBicyclicMatmul(const std::vector<int>& 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<int>(resultLayout, {resultVec}, {m, p});
}

Expand Down
5 changes: 3 additions & 2 deletions lib/Kernel/TricyclicBatchMatmulFuzzTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>(resultLayout, {actualVec}, {h, m, p});

Expand Down
71 changes: 53 additions & 18 deletions lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<RankedTensorType>(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);
Expand Down Expand Up @@ -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<RankedTensorType>(op->getResult(0).getType());
Operation* replicated =
replicateFirstPeriodOfResult(b, addBias->getResult(0), layoutAttr,
dataSemanticResultType.getNumElements());
int64_t reach = period * (steps - 1);
auto ctSemanticResultType =
cast<RankedTensorType>(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);
}
Expand Down Expand Up @@ -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<RankedTensorType>(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<RankedTensorType>(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);
}
Expand Down Expand Up @@ -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<RankedTensorType>(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<RankedTensorType>(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);
}
Expand Down
3 changes: 1 addition & 2 deletions tests/Examples/openfhe/ckks/batch_matmul/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
1 change: 0 additions & 1 deletion tests/Examples/openfhe/ckks/bicyclic_matmul/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
1 change: 0 additions & 1 deletion tests/Examples/openfhe/ckks/bicyclic_matmul_pt/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
2 changes: 1 addition & 1 deletion tests/Transforms/convert_to_ciphertext_semantics/batch_matmul.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion tests/Transforms/convert_to_ciphertext_semantics/matmul.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 4 additions & 3 deletions tests/Transforms/convert_to_ciphertext_semantics/matmul_pt.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<3x5xf32>> {tensor_ext.layout = #layout_ct}, %arg1: tensor<5x7xf32>) -> (!secret.secret<tensor<3x7xf32>> {tensor_ext.layout = #layout_out}) {
%cst = arith.constant dense<0.000000e+00> : tensor<3x7xf32>
%0 = secret.generic(%arg0: !secret.secret<tensor<3x5xf32>> {tensor_ext.layout = #layout_ct}) {
Expand All @@ -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<5x7xf32>> {tensor_ext.layout = #layout_ct2}) -> (!secret.secret<tensor<3x7xf32>> {tensor_ext.layout = #layout_out}) {
%cst = arith.constant dense<0.000000e+00> : tensor<3x7xf32>
%0 = secret.generic(%arg1: !secret.secret<tensor<5x7xf32>> {tensor_ext.layout = #layout_ct2}) {
Expand Down
Loading