google · copybara-service · Aug 11, 2026 · Aug 11, 2026
diff --git 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
@@ -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<std::vector<int>> 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<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.
+  addBounds(resultLayout,
+            resultLayout.getVarKindOffset(presburger::VarKind::Range) + 1, 0,
+            m * p - 1);
   return unpackLayoutToMatrix<int>(resultLayout, {resultVec}, {m, p});
 }
 

diff --git 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<std::vector<int>> runBicyclicMatmul(const std::vector<int>& 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<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.
+  addBounds(resultLayout,
+            resultLayout.getVarKindOffset(presburger::VarKind::Range) + 1, 0,
+            m * p - 1);
   return unpackLayoutToMatrix<int>(resultLayout, {resultVec}, {m, p});
 }
 

diff --git 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<int>(resultLayout, [&](const std::vector<int64_t>& 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<int>(resultLayout, {actualVec}, {h, m, p});
+
+  std::vector<std::vector<std::vector<int>>> expected(
+      h, std::vector<std::vector<int>>(m, std::vector<int>(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
@@ -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<RankedTensorType>(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<RankedTensorType>(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<RankedTensorType>(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<RankedTensorType>(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
@@ -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<int64_t> targetSlotCoeffs(result.getNumCols(), 0);
+  targetSlotCoeffs[targetSlotIndex] = 1;
+  auto targetSlotMod = addModConstraint(result, targetSlotCoeffs, period);
+  SmallVector<int64_t> 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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",
+)