Merge pull request #2163 from j2kun:rotate-and-reduce-indivisible

PiperOrigin-RevId: 800611385

Author: copybara-github

lib/Dialect/TensorExt/Transforms/ImplementRotateAndReduce.cpp

@@ -49,21 +49,6 @@ LogicalResult convertRotateAndReduceOp(RotateAndReduceOp op) {
   }
 
   TypedValue<RankedTensorType> plaintexts = op.getPlaintexts();
-
-  // Validate divisibility of step size
-  auto babySteps = static_cast<int64_t>(std::floor(std::sqrt(steps)));
-  unsigned giantSteps = steps / babySteps;
-  if (giantSteps * babySteps != steps) {
-    return op.emitOpError()
-           << "requires steps to be a multiple of sqrt(steps), but found "
-              "steps="
-           << steps << " and babySteps=" << babySteps;
-  }
-  LLVM_DEBUG(llvm::dbgs()
-             << "Using baby-step / giant-step decomposition of sum of size "
-             << steps << " with babySteps= " << babySteps
-             << " and giantSteps= " << giantSteps << "\n");
-
   auto plaintextsLeaf = std::optional<SSAValue>(plaintexts);
   implementedKernel =
       implementRotateAndReduce(vectorLeaf, plaintextsLeaf, period, steps);

lib/Kernel/KernelImplementation.h

@@ -149,6 +149,26 @@ implementRotateAndReduce(const T& vector, std::optional<T> plaintexts,
 
   // Use a value of sqrt(n) as the baby step / giant step size.
   int64_t numBabySteps = static_cast<int64_t>(std::floor(std::sqrt(steps)));
+  if (steps % numBabySteps != 0) {
+    // Find the nearest divisible number to use for baby step
+    // TODO(#2162): determine the right tradeoff here
+    int lower = numBabySteps;
+    int upper = numBabySteps;
+
+    while (steps % lower != 0 && steps % upper != steps) {
+      lower--;
+      upper++;
+    }
+
+    if (steps % lower == 0 && lower > 1) {
+      numBabySteps = lower;
+    } else if (steps % upper == 0) {
+      numBabySteps = upper;
+    } else {
+      numBabySteps = steps;
+    }
+  }
+
   int64_t giantStepSize = numBabySteps;
   int64_t numGiantSteps = steps / numBabySteps;
 

lib/Kernel/RotateAndReduceImplTest.cpp

@@ -153,6 +153,44 @@ TEST(RotateAndReduceImplTest, TestPeriod2WithNoPlaintext) {
   EXPECT_EQ(expected, actual);
 }
 
+TEST(RotateAndReduceImplTest, TestUnitPeriodWithIndivisibleN10) {
+  int64_t n = 10;
+  int64_t period = 1;
+
+  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
+  std::vector<std::vector<int>> plaintexts = {
+      {1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
+      {2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
+      {3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
+      {4, 5, 6, 7, 8, 9, 10, 11, 12, 13},
+      {5, 6, 7, 8, 9, 10, 11, 12, 13, 14},
+      {6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
+      {7, 8, 9, 10, 11, 12, 13, 14, 15, 16},
+      {8, 9, 10, 11, 12, 13, 14, 15, 16, 17},
+      {9, 10, 11, 12, 13, 14, 15, 16, 17, 18},
+      {10, 11, 12, 13, 14, 15, 16, 17, 18, 19},
+  };
+
+  std::vector<int> expected = runNaive(vector, plaintexts, period, n);
+  std::vector<int> actual = runImpl(vector, plaintexts, period, n);
+  EXPECT_EQ(expected, actual);
+}
+
+TEST(RotateAndReduceImplTest, TestUnitPeriodWithIndivisibleN6) {
+  int64_t n = 6;
+  int64_t period = 1;
+
+  std::vector<int> vector = {0, 1, 2, 3, 4, 5};
+  std::vector<std::vector<int>> plaintexts = {
+      {1, 2, 3, 4, 5, 6}, {2, 3, 4, 5, 6, 7},  {3, 4, 5, 6, 7, 8},
+      {4, 5, 6, 7, 8, 9}, {5, 6, 7, 8, 9, 10}, {6, 7, 8, 9, 10, 11},
+  };
+
+  std::vector<int> expected = runNaive(vector, plaintexts, period, n);
+  std::vector<int> actual = runImpl(vector, plaintexts, period, n);
+  EXPECT_EQ(expected, actual);
+}
+
 }  // namespace
 }  // namespace kernel
 }  // namespace heir

lib/Kernel/TestingUtils.cpp

@@ -86,6 +86,9 @@ LiteralValue EvalVisitor::operator()(const LeftRotateNode<LiteralValue>& node) {
   auto operand = this->process(node.operand);
   auto dim = operand.getShape()[0];
   int amount = node.shift;
+  // Normalize amount to be in [0, dim)
+  amount = ((amount % dim) + dim) % dim;
+
   const auto& oVal = operand.getTensor();
   const auto* oVec = std::get_if<std::vector<int>>(&oVal);
   assert(oVec && "unsupported rotate operand");