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: 953525822

Author: wtd2

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",
     ],

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});
 }
 

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});
 }
 

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() {

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(

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();

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,

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;

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>
+}

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>
+}