Merge pull request #2359 from j2kun:bicyclic-matmul

PiperOrigin-RevId: 826599179

Author: copybara-github

lib/Kernel/BUILD

@@ -73,6 +73,7 @@ cc_library(
         ":AbstractValue",
         ":ArithmeticDag",
         ":Kernel",
+        "@heir//lib/Utils:APIntUtils",
         "@heir//lib/Utils:MathUtils",
         "@llvm-project//mlir:Support",
     ],

lib/Kernel/KernelImplementation.h

@@ -15,6 +15,7 @@
 #include "lib/Kernel/AbstractValue.h"
 #include "lib/Kernel/ArithmeticDag.h"
 #include "lib/Kernel/KernelName.h"
+#include "lib/Utils/APIntUtils.h"
 #include "lib/Utils/MathUtils.h"
 #include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project
 
@@ -178,6 +179,57 @@ implementHaleviShoup(const T& vector, const T& matrix,
   return summedShifts;
 }
 
+// Returns an arithmetic DAG that implements the bicyclic matrix multiplication
+// algorithm.
+//
+// The input matrices packedA and packedB are assumed to be properly packed to
+// meet the conditions for bicyclic multiplication. That is: both matrices are
+// zero-padded so that their dimensions are coprime, they are cyclically
+// repeated to fill all the slots of the ciphertext, and they are packed
+// according to the bicyclic ordering.
+template <typename T>
+std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
+                 std::shared_ptr<ArithmeticDagNode<T>>>
+implementBicyclicMatmul(const T& packedA, const T& packedB, int64_t m,
+                        int64_t n, int64_t p) {
+  using NodeTy = ArithmeticDagNode<T>;
+  auto packedADag = NodeTy::leaf(packedA);
+  auto packedBDag = NodeTy::leaf(packedB);
+
+  // This implements the BMM-I algorithm from https://eprint.iacr.org/2024/1762
+  // with a simplification of the rotation formula in Sec 5.2.1 (equation 9) of
+  // https://eprint.iacr.org/2025/1200
+  //
+  // C = sum_{c=0}^{n-1} rot(A, r1(c)) * rot(B, r2(c))
+  //
+  // where
+  //
+  //  r1(c) = cm(m^{-1} mod n) mod mn
+  //  r2(c) = cp(p^{-1} mod n) mod np
+  //
+  APInt mAPInt = APInt(64, m);
+  APInt nAPInt = APInt(64, n);
+  APInt pAPInt = APInt(64, p);
+
+  APInt mInvModN = multiplicativeInverse(mAPInt, nAPInt);
+  APInt pInvModN = multiplicativeInverse(pAPInt, nAPInt);
+
+  // The part of r1(c), r2(c) that is independent of the loop iterations
+  int64_t r1Const = m * mInvModN.getSExtValue();
+  int64_t r2Const = p * pInvModN.getSExtValue();
+
+  std::shared_ptr<NodeTy> result = nullptr;
+  for (int i = 0; i < n; ++i) {
+    int64_t shiftA = (i * r1Const) % (m * n);
+    int64_t shiftB = (i * r2Const) % (n * p);
+    auto rotA = NodeTy::leftRotate(packedADag, shiftA);
+    auto rotB = NodeTy::leftRotate(packedBDag, shiftB);
+    auto term = NodeTy::mul(rotA, rotB);
+    result = result == nullptr ? term : NodeTy::add(result, term);
+  }
+  return result;
+}
+
 }  // namespace kernel
 }  // namespace heir
 }  // namespace mlir

lib/Kernel/KernelImplementationTest.cpp

@@ -138,6 +138,41 @@ TEST(KernelImplementationTest, Test2DConvWithLayout) {
   EXPECT_EQ(extractedResult, expected);
 }
 
+TEST(KernelImplementationTest, BicyclicMatmul) {
+  MLIRContext context;
+  std::vector<std::vector<int>> matrixA = {
+      {1, 2, 3, 4, 5}, {6, 7, 8, 9, 10}, {11, 12, 13, 14, 15}};
+  std::vector<std::vector<int>> matrixB = {
+      {1, 2}, {3, 4}, {5, 6}, {7, 8}, {9, 10}};
+  int m = 3;
+  int n = 5;
+  int p = 2;
+  int numSlots = m * n * p;
+
+  auto layoutA = getBicyclicLayoutRelation(
+      RankedTensorType::get({m, n}, IndexType::get(&context)), numSlots);
+  auto packedA = evaluateLayoutOnMatrix(layoutA, matrixA);
+
+  auto layoutB = getBicyclicLayoutRelation(
+      RankedTensorType::get({n, p}, IndexType::get(&context)), numSlots);
+  auto packedB = evaluateLayoutOnMatrix(layoutB, matrixB);
+
+  LiteralValue packedAValue = packedA[0];
+  LiteralValue packedBValue = packedB[0];
+
+  auto dag = implementBicyclicMatmul(packedAValue, packedBValue, m, n, p);
+  LiteralValue result = evalKernel(dag);
+  auto resultVec = std::get<std::vector<int>>(result.getTensor());
+
+  auto resultLayout = getBicyclicLayoutRelation(
+      RankedTensorType::get({m, p}, IndexType::get(&context)), numSlots);
+  auto unpackedResult =
+      unpackLayoutToMatrix<int>(resultLayout, {resultVec}, {m, p});
+
+  std::vector<std::vector<int>> expected = {{95, 110}, {220, 260}, {345, 410}};
+  EXPECT_EQ(unpackedResult, expected);
+}
+
 }  // namespace
 }  // namespace kernel
 }  // namespace heir

lib/Utils/Layout/Evaluate.h

@@ -87,6 +87,27 @@ std::vector<std::vector<T>> evaluateLayoutOnMatrix(
   return evaluateLayout(relation, getValueFn);
 }
 
+template <typename T>
+std::vector<std::vector<T>> unpackLayoutToMatrix(
+    const presburger::IntegerRelation& relation,
+    const std::vector<std::vector<T>>& packed,
+    ArrayRef<int64_t> originalShape) {
+  std::vector<std::vector<T>> result(originalShape[0],
+                                     std::vector<T>(originalShape[1], 0));
+
+  // Get all points in the relation.
+  PointPairCollector collector(relation.getNumDomainVars(), /*rangeDims=*/2);
+  enumeratePoints(relation, collector);
+
+  for (const auto& pointPair : collector.points) {
+    std::vector<int64_t> domainPoint = pointPair.first;
+    std::vector<int64_t> rangePoint = pointPair.second;
+    result[domainPoint[0]][domainPoint[1]] =
+        packed[rangePoint[0]][rangePoint[1]];
+  }
+  return result;
+}
+
 }  // namespace heir
 }  // namespace mlir
 

lib/Utils/Layout/Utils.cpp

@@ -208,6 +208,64 @@ presburger::IntegerRelation getDiagonalLayoutRelation(
   return result;
 }
 
+presburger::IntegerRelation getBicyclicLayoutRelation(
+    RankedTensorType matrixType, int64_t numSlots) {
+  unsigned int rows = matrixType.getDimSize(0);
+  unsigned int cols = matrixType.getDimSize(1);
+
+  assert(std::gcd(rows, cols) == 1 &&
+         "bicyclic layout requires coprime dimensions");
+
+  IntegerRelation result(PresburgerSpace::getRelationSpace(
+      matrixType.getRank(), /*numRange=*/2, /*numSymbol=*/0,
+      /*numLocals=*/0));
+
+  // Add bounds for the data matrix dimensions.
+  int domainOffset = result.getVarKindOffset(VarKind::Domain);
+  int rangeOffset = result.getVarKindOffset(VarKind::Range);
+  int rowVarIndex = domainOffset;
+  int colVarIndex = domainOffset + 1;
+  int ctVarIndex = rangeOffset;
+  int slotVarIndex = rangeOffset + 1;
+
+  addBounds(result, rowVarIndex, 0, rows - 1);
+  addBounds(result, colVarIndex, 0, cols - 1);
+  addBounds(result, ctVarIndex, 0,
+            std::ceil((float)matrixType.getNumElements() / numSlots) - 1);
+  addBounds(result, slotVarIndex, 0, numSlots - 1);
+
+  // Let k = ct * numSlots + slot.
+  // We need to add constraints for:
+  // row = k % rows
+  // col = k % cols
+
+  // k_mod_rows = (ct * numSlots + slot) % rows
+  SmallVector<int64_t> kCoeffs(result.getNumCols(), 0);
+  kCoeffs[ctVarIndex] = numSlots;
+  kCoeffs[slotVarIndex] = 1;
+  auto kModRows = addModConstraint(result, kCoeffs, rows);
+
+  // row = k_mod_rows
+  SmallVector<int64_t> rowEquality(result.getNumCols(), 0);
+  rowEquality[rowVarIndex] = 1;
+  rowEquality[kModRows] = -1;
+  result.addEquality(rowEquality);
+
+  // k_mod_cols = (ct * numSlots + slot) % cols
+  kCoeffs.resize(result.getNumCols(), 0);
+  kCoeffs[ctVarIndex] = numSlots;
+  kCoeffs[slotVarIndex] = 1;
+  auto kModCols = addModConstraint(result, kCoeffs, cols);
+
+  // col = k_mod_cols
+  SmallVector<int64_t> colEquality(result.getNumCols(), 0);
+  colEquality[colVarIndex] = 1;
+  colEquality[kModCols] = -1;
+  result.addEquality(colEquality);
+
+  return result;
+}
+
 presburger::IntegerRelation getPerRowLayoutRelation(RankedTensorType matrixType,
                                                     int64_t ciphertextSize) {
   auto domainSize = matrixType.getRank();

lib/Utils/Layout/Utils.h

@@ -48,6 +48,11 @@ presburger::IntegerRelation getRowMajorLayoutRelation(
 presburger::IntegerRelation getDiagonalLayoutRelation(
     RankedTensorType matrixType, int64_t ciphertextSize);
 
+// Returns an IntegerRelation that represents a bicyclic layout for a matrix.
+// See https://eprint.iacr.org/2024/1762 for details.
+presburger::IntegerRelation getBicyclicLayoutRelation(
+    RankedTensorType matrixType, int64_t numSlots);
+
 // 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

@@ -160,6 +160,45 @@ TEST(UtilsTest, SquatDiagonalLayout) {
   }
 }
 
+TEST(UtilsTest, BicyclicLayout3x5) {
+  MLIRContext context;
+  int64_t numSlots = 15;
+  RankedTensorType matrixType =
+      RankedTensorType::get({3, 5}, IndexType::get(&context));
+  IntegerRelation bicyclicRelation =
+      getBicyclicLayoutRelation(matrixType, numSlots);
+
+  std::vector<std::vector<int>> matrix = {
+      {1, 2, 3, 4, 5}, {6, 7, 8, 9, 10}, {11, 12, 13, 14, 15}};
+  std::vector<std::vector<int>> packedMatrix =
+      evaluateLayoutOnMatrix(bicyclicRelation, matrix);
+
+  std::vector<std::vector<int>> expected = {
+      {1, 7, 13, 4, 10, 11, 2, 8, 14, 5, 6, 12, 3, 9, 15}};
+  EXPECT_EQ(packedMatrix, expected);
+}
+
+TEST(UtilsTest, BicyclicLayout3x5Repeated) {
+  MLIRContext context;
+
+  int64_t numSlots = 32;
+  RankedTensorType matrixType =
+      RankedTensorType::get({3, 5}, IndexType::get(&context));
+  IntegerRelation bicyclicRelation =
+      getBicyclicLayoutRelation(matrixType, numSlots);
+
+  std::vector<std::vector<int>> matrix = {
+      {1, 2, 3, 4, 5}, {6, 7, 8, 9, 10}, {11, 12, 13, 14, 15}};
+  std::vector<std::vector<int>> packedMatrix =
+      evaluateLayoutOnMatrix(bicyclicRelation, matrix);
+
+  std::vector<std::vector<int>> expected = {
+      {1, 7, 13, 4, 10, 11, 2, 8, 14, 5, 6, 12, 3, 9, 15,
+       // Cyclically repeated to fill 32 slots
+       1, 7, 13, 4, 10, 11, 2, 8, 14, 5, 6, 12, 3, 9, 15, 1, 7}};
+  EXPECT_EQ(packedMatrix, expected);
+}
+
 TEST(UtilsTest, TestGetRangePoints) {
   MLIRContext context;
   auto rel = getIntegerRelationFromIslStr(