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

Author: wtd2

lib/Kernel/BicyclicDiagonalMatmulFuzzTest.cpp

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

lib/Kernel/BicyclicMatmulFuzzTest.cpp

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

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<int>(resultLayout, {actualVec}, {h, m, p});
 

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

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

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

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

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

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

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<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}) {
@@ -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}) {