feat: in-layout strided reduce for bicyclic/tricyclic layout

lib/Dialect/Secret/IR/SecretOps.cpp

@@ -589,6 +589,8 @@ GenericOp GenericOp::extractOpBeforeGeneric(Operation* opToExtract,
         })) {
       this->setOperandAttrsAttr(
           ArrayAttr::get(this->getContext(), newGenericArgAttrs));
+    } else {
+      this->removeAllOperandAttrsAttr();
     }
   });
   rewriter.replaceOp(opToExtract, oldGenericNewBlockArgs);

lib/Dialect/Secret/IR/SecretPatterns.cpp

@@ -176,19 +176,28 @@ LogicalResult RemoveUnusedGenericArgs::matchAndRewrite(
     if (arg.use_empty()) {
       LLVM_DEBUG(llvm::dbgs() << arg << " has no uses; removing\n");
       hasUnusedOps = true;
+      // Read the operand attrs BEFORE erasing the operand: the accessor
+      // treats a size-mismatched (stale) array as absent, so reading after
+      // the erase would both skip the rebuild and leave the stale array on
+      // the op.
+      auto attrs = op.getAllOperandAttrsAttr();
       rewriter.modifyOpInPlace(op, [&]() {
         body->eraseArgument(i);
         op.getOperation()->eraseOperand(i);
       });
-      auto attrs = op.getAllOperandAttrsAttr();
       if (attrs) {
         SmallVector<Attribute> attrList;
         for (auto [j, attr] : llvm::enumerate(attrs)) {
           if (j != i) {
             attrList.push_back(attr);
           }
         }
-        op.setOperandAttrsAttr(ArrayAttr::get(op.getContext(), attrList));
+        // An empty array is not "no attrs"; remove it outright so later
+        // operand appends don't see a stale array.
+        if (attrList.empty())
+          op.removeAllOperandAttrsAttr();
+        else
+          op.setOperandAttrsAttr(ArrayAttr::get(op.getContext(), attrList));
       }
 
       // Ensure the next iteration uses the right arg number
@@ -263,19 +272,24 @@ LogicalResult RemoveNonSecretGenericArgs::matchAndRewrite(
       BlockArgument correspondingArg = body->getArgument(i);
 
       rewriter.replaceAllUsesWith(correspondingArg, op->getOperand(i));
+      // Read the operand attrs BEFORE erasing the operand; see
+      // RemoveUnusedGenericArgs.
+      auto attrs = op.getAllOperandAttrsAttr();
       rewriter.modifyOpInPlace(op, [&]() {
         body->eraseArgument(i);
         op.getOperation()->eraseOperand(i);
       });
-      auto attrs = op.getAllOperandAttrsAttr();
       if (attrs) {
         SmallVector<Attribute> attrList;
         for (auto [j, attr] : llvm::enumerate(attrs)) {
           if (j != i) {
             attrList.push_back(attr);
           }
         }
-        op.setOperandAttrsAttr(ArrayAttr::get(op.getContext(), attrList));
+        if (attrList.empty())
+          op.removeAllOperandAttrsAttr();
+        else
+          op.setOperandAttrsAttr(ArrayAttr::get(op.getContext(), attrList));
       }
       i--;
     }

lib/Dialect/TensorExt/Transforms/BUILD

@@ -109,7 +109,9 @@ cc_library(
         ":pass_inc_gen",
         "@heir//lib/Dialect/TensorExt/IR:Dialect",
         "@heir//lib/Kernel:AbstractValue",
+        "@heir//lib/Kernel:ArithmeticDag",
         "@heir//lib/Kernel:IRMaterializingVisitor",
+        "@heir//lib/Utils",
         "@heir//lib/Utils:MathUtils",
         "@heir//lib/Utils/ADT:FrozenVector",
         "@heir//lib/Utils/Graph",

lib/Dialect/TensorExt/Transforms/ImplementShiftNetwork.cpp

@@ -2,11 +2,15 @@
 
 #include <algorithm>
 #include <cassert>
+#include <cmath>
 #include <cstddef>
 #include <cstdint>
 #include <iterator>
 #include <limits>
+#include <map>
+#include <memory>
 #include <random>
+#include <tuple>
 #include <unordered_map>
 #include <utility>
 #include <vector>
@@ -16,15 +20,16 @@
 #include "lib/Dialect/TensorExt/Transforms/RotationGroupKernel.h"
 #include "lib/Dialect/TensorExt/Transforms/ShiftScheme.h"
 #include "lib/Kernel/AbstractValue.h"
+#include "lib/Kernel/ArithmeticDag.h"
 #include "lib/Kernel/IRMaterializingVisitor.h"
 #include "lib/Utils/ADT/FrozenVector.h"
 #include "lib/Utils/Graph/Graph.h"
 #include "lib/Utils/Layout/Utils.h"
+#include "lib/Utils/Utils.h"
+#include "llvm/include/llvm/ADT/DenseSet.h"              // from @llvm-project
 #include "llvm/include/llvm/ADT/STLExtras.h"             // from @llvm-project
 #include "llvm/include/llvm/Support/Debug.h"             // from @llvm-project
-#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
-#include "mlir/include/mlir/IR/Attributes.h"             // from @llvm-project
 #include "mlir/include/mlir/IR/Builders.h"               // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
@@ -141,7 +146,7 @@ ShiftScheme VosVosErkinShiftNetworks::findBestShiftScheme(
     // iteration.
     SmallVector<int64_t> shiftOrder = initShiftOrder;
 
-    std::ranges::shuffle(shiftOrder.begin(), shiftOrder.end(), g);
+    std::shuffle(shiftOrder.begin(), shiftOrder.end(), g);
 
     ShiftStrategy strategy = evaluateShiftStrategy(mapping, shiftOrder);
 
@@ -242,12 +247,6 @@ LogicalResult convertRemapOp(RemapOp op,
               "DenseIntElementsAttr";
   }
 
-  ShiftScheme scheme = shiftNetworks.findShiftScheme(mapping);
-  auto rotationGroups = scheme.rotationGroups;
-
-  assert(!rotationGroups.empty() &&
-         "Shift network must have at least one group");
-
   b.setInsertionPointAfter(op);
 
   // Could add a special case here if the numCiphertexts == 1, using
@@ -272,8 +271,80 @@ LogicalResult convertRemapOp(RemapOp op,
     ciphertexts.push_back(kernel::SSAValue(slice.getResult()));
   }
 
-  auto resultNodes = implementShiftNetwork(ciphertexts, mapping, scheme,
-                                           minSlotCount, dagElemType);
+  using NodeTy = kernel::ArithmeticDagNode<kernel::SSAValue>;
+  SmallVector<std::shared_ptr<NodeTy>> resultNodes;
+
+  const auto targetToSource = mapping.getTargetToSource();
+
+  // Count the rotations the direct path would emit.
+  llvm::DenseSet<std::pair<int64_t, int64_t>> uniqueShifts;
+  for (const auto& [target, source] : targetToSource) {
+    int64_t r = (source.slot - target.slot) % minSlotCount;
+    if (r < 0) r += minSlotCount;
+    if (r != 0) uniqueShifts.insert({source.ct, r});
+  }
+
+  // If the number of distinct rotation offsets is <= log2(N_slots), direct
+  // depth-1 rotation + plaintext masking uses at most `uniqueShifts` rotations
+  // at multiplicative depth 1, vs VVE's multi-stage log2(N) rotations and
+  // ~2·log2(N) masks.
+  int64_t maxDirectRotations =
+      static_cast<int64_t>(std::ceil(std::log2(minSlotCount)));
+
+  if (static_cast<int64_t>(uniqueShifts.size()) <= maxDirectRotations) {
+    // Build rotation-mask groups.
+    std::map<std::tuple<int64_t, int64_t, int64_t>, std::vector<double>> groups;
+    for (const auto& [target, source] : targetToSource) {
+      int64_t r = (source.slot - target.slot) % minSlotCount;
+      if (r < 0) r += minSlotCount;
+      auto& mask = groups
+                       .try_emplace(std::make_tuple(target.ct, source.ct, r),
+                                    std::vector<double>(minSlotCount, 0.0))
+                       .first->second;
+      mask[target.slot] = 1.0;
+    }
+
+    SmallVector<std::shared_ptr<NodeTy>> cts;
+    for (const auto& ct : ciphertexts) {
+      cts.push_back(NodeTy::leaf(ct));
+    }
+
+    std::map<std::pair<int64_t, int64_t>, std::shared_ptr<NodeTy>>
+        rotationCache;
+    auto maskType = makeTensorType(dagElemType, {1, minSlotCount});
+    resultNodes.assign(numCiphertexts, nullptr);
+    for (auto& [key, mask] : groups) {
+      auto [targetCt, sourceCt, r] = key;
+      std::shared_ptr<NodeTy> rotated;
+      if (r == 0) {
+        rotated = cts[sourceCt];
+      } else {
+        auto [it, inserted] = rotationCache.try_emplace({sourceCt, r}, nullptr);
+        if (inserted) {
+          it->second = NodeTy::leftRotate(cts[sourceCt], r);
+        }
+        rotated = it->second;
+      }
+      auto [allZero, allOne] = allZeroAllOne(mask);
+      (void)allZero;
+      std::shared_ptr<NodeTy> term =
+          allOne ? rotated
+                 : NodeTy::mul(rotated, NodeTy::constantTensor(mask, maskType));
+      resultNodes[targetCt] = resultNodes[targetCt]
+                                  ? NodeTy::add(resultNodes[targetCt], term)
+                                  : term;
+    }
+    std::vector<double> zeros(minSlotCount, 0.0);
+    for (auto& node : resultNodes) {
+      if (!node) node = NodeTy::constantTensor(zeros, maskType);
+    }
+  } else {
+    ShiftScheme scheme = shiftNetworks.findShiftScheme(mapping);
+    assert(!scheme.rotationGroups.empty() &&
+           "Shift network must have at least one group");
+    resultNodes = implementShiftNetwork(ciphertexts, mapping, scheme,
+                                        minSlotCount, dagElemType);
+  }
 
   kernel::IRMaterializingVisitor visitor(singleCiphertextType);
   auto resultVectors = visitor.process(resultNodes, b);

lib/Kernel/BUILD

@@ -138,11 +138,13 @@ cc_test(
         ":KernelImplementation",
         ":RotationCountVisitor",
         "@googletest//:gtest_main",
+        "@heir//lib/Utils:RotationUtils",
         "@heir//lib/Utils/Layout:Codegen",
         "@heir//lib/Utils/Layout:Convolution",
         "@heir//lib/Utils/Layout:ConvolutionTestUtil",
         "@heir//lib/Utils/Layout:Evaluate",
         "@heir//lib/Utils/Layout:Utils",
+        "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Support",
     ],

lib/Kernel/Kernel.cpp

@@ -36,6 +36,8 @@ std::string kernelNameAsStr(const KernelName& kernelName) {
       return "MatmulBicyclic";
     case KernelName::MatmulBicyclicDiagonal:
       return "MatmulBicyclicDiagonal";
+    case KernelName::BatchMatmulTricyclicDiagonal:
+      return "BatchMatmulTricyclicDiagonal";
     case KernelName::BatchMatmulTricyclic:
       return "BatchMatmulTricyclic";
     case KernelName::Dot:

lib/Kernel/Kernel.h

@@ -31,6 +31,8 @@ struct FieldParser<heir::KernelName> {
     if (kernelName == "MatmulBicyclic") return heir::KernelName::MatmulBicyclic;
     if (kernelName == "MatmulBicyclicDiagonal")
       return heir::KernelName::MatmulBicyclicDiagonal;
+    if (kernelName == "BatchMatmulTricyclicDiagonal")
+      return heir::KernelName::BatchMatmulTricyclicDiagonal;
     if (kernelName == "BatchMatmulTricyclic")
       return heir::KernelName::BatchMatmulTricyclic;
     if (kernelName == "Dot") return heir::KernelName::Dot;

lib/Kernel/KernelImplementation.h

@@ -77,13 +77,35 @@ std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
 implementRotateAndReduceAccumulation(
     std::shared_ptr<ArithmeticDagNode<T>> vectorDag, int64_t period,
     int64_t steps, DagReducer<T> reduceFunc) {
+  assert(steps >= 1 && "rotate-and-reduce needs at least one step");
   using NodeTy = ArithmeticDagNode<T>;
-  for (int64_t shiftSize = steps / 2; shiftSize > 0; shiftSize /= 2) {
-    auto rotated = NodeTy::leftRotate(vectorDag, shiftSize * period);
-    auto reduced = reduceFunc(vectorDag, rotated);
-    vectorDag = reduced;
+  if ((steps & (steps - 1)) == 0) {
+    for (int64_t shiftSize = steps / 2; shiftSize > 0; shiftSize /= 2) {
+      auto rotated = NodeTy::leftRotate(vectorDag, shiftSize * period);
+      auto reduced = reduceFunc(vectorDag, rotated);
+      vectorDag = reduced;
+    }
+    return vectorDag;
   }
-  return vectorDag;
+  std::shared_ptr<ArithmeticDagNode<T>> acc = nullptr;
+  std::shared_ptr<ArithmeticDagNode<T>> span = vectorDag;
+  int64_t spanLen = 1;
+  int64_t offset = 0;
+  int64_t remaining = steps;
+  while (remaining > 0) {
+    if (remaining & 1) {
+      auto term =
+          offset == 0 ? span : NodeTy::leftRotate(span, offset * period);
+      acc = acc ? reduceFunc(acc, term) : term;
+      offset += spanLen;
+    }
+    remaining >>= 1;
+    if (remaining > 0) {
+      span = reduceFunc(span, NodeTy::leftRotate(span, spanLen * period));
+      spanLen <<= 1;
+    }
+  }
+  return acc;
 }
 
 template <typename T>
@@ -106,6 +128,13 @@ implementRotateAndReduceAccumulationRolled(
   using NodeTy = ArithmeticDagNode<T>;
   using NodePtr = std::shared_ptr<NodeTy>;
 
+  // Non-power-of-two steps fall back to the unrolled binary span-doubling
+  // form because the variable shift and rotation sequence cannot be expressed
+  // as a uniform loop with a single halved shift per iteration.
+  if ((steps & (steps - 1)) != 0) {
+    return implementRotateAndReduceAccumulation<T>(vectorDag, period, steps,
+                                                   reduceFunc);
+  }
   int64_t numIterations = static_cast<int64_t>(std::log2(steps));
   if (numIterations <= 0) return vectorDag;
 

lib/Kernel/KernelName.h

@@ -39,6 +39,10 @@ enum KernelName : int {
   // plaintext-lhs-secret-rhs cases.
   MatmulBicyclicDiagonal,
 
+  // Ciphertext-plaintext batch matmul by mapping the plaintext into n
+  // generalized diagonals.
+  BatchMatmulTricyclicDiagonal,
+
   // Product and sum of two vectors, using a log2 rotate-and-reduce approach.
   Dot,
 };

lib/Kernel/RotateAndReduceFuzzTest.cpp

@@ -177,7 +177,8 @@ FUZZ_TEST(RotateAndReduceFuzzTest, rotateAndReduceWithoutPlaintexts)
             .WithMinSize(1)
             .WithMaxSize(32),
         /*period=*/fuzztest::InRange(1L, 4L),
-        /*steps=*/fuzztest::ElementOf({1L, 2L, 4L, 8L, 16L}),
+        /*steps=*/
+        fuzztest::ElementOf({1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L, 16L}),
         fuzztest::Arbitrary<bool>());
 
 // Fuzz test for rotate and reduce with plaintexts and zero diagonals

lib/Kernel/RotateAndReduceImplTest.cpp

@@ -1,13 +1,17 @@
 #include <cstdint>
 #include <memory>
 #include <optional>
+#include <unordered_set>
+#include <variant>
 #include <vector>
 
 #include "gtest/gtest.h"  // from @googletest
 #include "lib/Kernel/AbstractValue.h"
 #include "lib/Kernel/ArithmeticDag.h"
 #include "lib/Kernel/EvalVisitor.h"
 #include "lib/Kernel/KernelImplementation.h"
+#include "lib/Utils/RotationUtils.h"
+#include "llvm/include/llvm/ADT/DenseSet.h"  // from @llvm-project
 
 namespace mlir {
 namespace heir {
@@ -313,6 +317,53 @@ TEST(RotateAndReduceImplTest, BroadcastedReduce_Masked_Stride) {
   }
 }
 
+void collectRotations(
+    const std::shared_ptr<ArithmeticDagNode<LiteralValue>>& node,
+    llvm::DenseSet<int64_t>& rotations,
+    std::unordered_set<const ArithmeticDagNode<LiteralValue>*>& visited) {
+  if (!node || !visited.insert(node.get()).second) return;
+  if (auto* rotate =
+          std::get_if<LeftRotateNode<LiteralValue>>(&node->node_variant)) {
+    if (auto* scalar =
+            std::get_if<ConstantScalarNode>(&rotate->shift->node_variant)) {
+      rotations.insert(static_cast<int64_t>(scalar->value));
+    }
+    collectRotations(rotate->operand, rotations, visited);
+    collectRotations(rotate->shift, rotations, visited);
+    return;
+  }
+  if (auto* add = std::get_if<AddNode<LiteralValue>>(&node->node_variant)) {
+    collectRotations(add->left, rotations, visited);
+    collectRotations(add->right, rotations, visited);
+    return;
+  }
+}
+
+TEST(RotateAndReduceImplTest, TestPredictorSync) {
+  std::vector<int> dummyVec = {1};
+  LiteralValue val(dummyVec);
+  auto leaf = ArithmeticDagNode<LiteralValue>::leaf(val);
+  auto addReducer = [](std::shared_ptr<ArithmeticDagNode<LiteralValue>> a,
+                       std::shared_ptr<ArithmeticDagNode<LiteralValue>> b) {
+    return ArithmeticDagNode<LiteralValue>::add(a, b);
+  };
+
+  for (int64_t period : {1, 3, 7}) {
+    for (int64_t steps = 1; steps <= 100; ++steps) {
+      auto dag = implementRotateAndReduceAccumulation<LiteralValue>(
+          leaf, period, steps, addReducer);
+      llvm::DenseSet<int64_t> actualRotations;
+      std::unordered_set<const ArithmeticDagNode<LiteralValue>*> visited;
+      collectRotations(dag, actualRotations, visited);
+
+      llvm::DenseSet<int64_t> predicted = rotateAndReduceRotationIndices(
+          period, steps, /*hasPlaintexts=*/false);
+      EXPECT_EQ(actualRotations, predicted)
+          << "Mismatch for period=" << period << ", steps=" << steps;
+    }
+  }
+}
+
 }  // namespace
 }  // namespace kernel
 }  // namespace heir