Create and lower broadcasted_reduce to get softmax e2e working

lib/Dialect/TensorExt/IR/TensorExtOps.cpp

@@ -405,6 +405,29 @@ LogicalResult RotateAndReduceOp::verify() {
   return success();
 }
 
+LogicalResult BroadcastedReduceOp::verify() {
+  auto tensorType = getTensor().getType();
+  int64_t rank = tensorType.getRank();
+  int64_t dim = getDimension();
+
+  if (dim < 0 || dim >= rank) {
+    return emitOpError() << "dimension " << dim << " is out of bounds for rank "
+                         << rank;
+  }
+
+  if (getReduceOp().has_value()) {
+    StringRef reduceOp = getReduceOp().value();
+    if (reduceOp != "arith.addi" && reduceOp != "arith.addf" &&
+        reduceOp != "arith.muli" && reduceOp != "arith.mulf" &&
+        reduceOp != "addi" && reduceOp != "addf" && reduceOp != "muli" &&
+        reduceOp != "mulf") {
+      return emitOpError() << "unsupported reduceOp: " << reduceOp;
+    }
+  }
+
+  return success();
+}
+
 }  // namespace tensor_ext
 }  // namespace heir
 }  // namespace mlir

lib/Dialect/TensorExt/IR/TensorExtOps.td

@@ -302,5 +302,30 @@ def TensorExt_RotateAndReduceOp : TensorExt_Op<"rotate_and_reduce", [
   // TODO(#2134): Add canonicalization patterns
 }
 
+def TensorExt_BroadcastedReduceOp : TensorExt_Op<"broadcasted_reduce", [
+  Pure,
+  AllTypesMatch<["tensor", "output"]>
+]> {
+  let summary = "Broadcasted reduction of a tensor along a dimension.";
+  let description = [{
+    This op reduces a tensor along a specified dimension and broadcasts the
+    result back to the original shape.
+
+    The reduction operation is specified by the `reduceOp` attribute.
+    The chosen op must be one of `arith.addi`, `arith.addf`, `arith.muli`,
+    or `arith.mulf`.
+
+    This op is layout-preserving.
+  }];
+
+  let arguments = (ins
+    AnyRankedTensor:$tensor,
+    I64Attr:$dimension,
+    OptionalAttr<Builtin_StringAttr>:$reduceOp
+  );
+  let results = (outs AnyRankedTensor:$output);
+  let assemblyFormat = "operands attr-dict `:` type($tensor)";
+  let hasVerifier = 1;
+}
 
 #endif  // LIB_DIALECT_TENSOREXT_IR_TENSOREXTOPS_TD_

lib/Kernel/BUILD

@@ -174,6 +174,20 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "BroadcastedReduceFuzzTest",
+    srcs = ["BroadcastedReduceFuzzTest.cpp"],
+    deps = [
+        ":AbstractValue",
+        ":ArithmeticDag",
+        ":EvalVisitor",
+        ":KernelImplementation",
+        ":TestingUtils",
+        "@fuzztest//fuzztest",
+        "@fuzztest//fuzztest:fuzztest_gtest_main",
+    ],
+)
+
 cc_test(
     name = "MatvecZeroDiagonalsFuzzTest",
     srcs = ["MatvecZeroDiagonalsFuzzTest.cpp"],

lib/Kernel/BroadcastedReduceFuzzTest.cpp

@@ -0,0 +1,127 @@
+#include <cstdint>
+#include <memory>
+#include <optional>
+#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"
+
+// copybara hack: avoid reordering include
+#include "fuzztest/fuzztest.h"  // from @fuzztest
+
+namespace mlir {
+namespace heir {
+namespace kernel {
+namespace {
+
+std::vector<int> runNaiveBroadcastedReduce(const std::vector<int>& vec,
+                                           int64_t period, int64_t steps) {
+  int64_t n = vec.size();
+  int64_t B = steps;
+  int64_t blockSize = B * period;
+  std::vector<int> result(n, 0);
+
+  int64_t numBlocks = n / blockSize;
+
+  for (int64_t k = 0; k < numBlocks; ++k) {
+    for (int64_t offset = 0; offset < period; ++offset) {
+      int sum = 0;
+      for (int64_t i = 0; i < B; ++i) {
+        sum += vec[k * blockSize + i * period + offset];
+      }
+      for (int64_t i = 0; i < B; ++i) {
+        result[k * blockSize + i * period + offset] = sum;
+      }
+    }
+  }
+  return result;
+}
+
+std::vector<int> generateCleanupMask(int64_t numSlots, int64_t period,
+                                     int64_t steps) {
+  std::vector<int> mask(numSlots, 0);
+  int64_t B = steps;
+  int64_t blockSize = B * period;
+  int64_t numBlocks = numSlots / blockSize;
+  for (int64_t k = 0; k < numBlocks; ++k) {
+    for (int64_t offset = 0; offset < period; ++offset) {
+      mask[k * blockSize + (B - 1) * period + offset] = 1;
+    }
+  }
+  return mask;
+}
+
+void broadcastedReduceMatchesNaive(int logN, int logB, int logPeriod,
+                                   const std::vector<int>& inputTemplate,
+                                   bool unroll) {
+  int64_t numSlots = 1 << logN;
+  int64_t steps = 1 << logB;
+  int64_t period = 1 << logPeriod;
+
+  if (steps * period > numSlots) return;
+
+  // Resize inputTemplate to numSlots
+  std::vector<int> vec(numSlots);
+  for (int64_t i = 0; i < numSlots; ++i) {
+    vec[i] = inputTemplate[i % inputTemplate.size()];
+  }
+
+  std::vector<int> expected = runNaiveBroadcastedReduce(vec, period, steps);
+
+  using NodeTy = ArithmeticDagNode<LiteralValue>;
+  using NodePtr = std::shared_ptr<NodeTy>;
+
+  LiteralValue vectorInput(vec);
+  auto vectorDag = NodeTy::leaf(vectorInput);
+
+  std::optional<NodePtr> cleanupMaskDag = std::nullopt;
+  if (steps * period < numSlots) {
+    auto mask = generateCleanupMask(numSlots, period, steps);
+    cleanupMaskDag = NodeTy::leaf(LiteralValue(mask));
+  }
+
+  auto result = implementBroadcastedReduce<LiteralValue>(
+      vectorDag, cleanupMaskDag, period, steps, numSlots,
+      DagType::intTensor(32, {numSlots}), "arith.addi", unroll);
+
+  std::vector<int> actual =
+      std::get<std::vector<int>>(evalKernel(result)[0].get());
+
+  EXPECT_EQ(expected, actual);
+}
+
+auto ValidParameters() {
+  return fuzztest::FlatMap(
+      [](int logN) {
+        return fuzztest::FlatMap(
+            [logN](int logB) {
+              return fuzztest::TupleOf(fuzztest::Just(logN),
+                                       fuzztest::Just(logB),
+                                       fuzztest::InRange(0, logN - logB));
+            },
+            fuzztest::InRange(1, logN));
+      },
+      fuzztest::InRange(3, 7)  // N from 8 to 128
+  );
+}
+
+void BroadcastedReduceFuzz(const std::tuple<int, int, int>& params,
+                           const std::vector<int>& inputTemplate, bool unroll) {
+  auto [logN, logB, logPeriod] = params;
+  broadcastedReduceMatchesNaive(logN, logB, logPeriod, inputTemplate, unroll);
+}
+
+FUZZ_TEST(BroadcastedReduceFuzzTest, BroadcastedReduceFuzz)
+    .WithDomains(ValidParameters(),
+                 fuzztest::VectorOf(fuzztest::InRange(-100, 100))
+                     .WithMinSize(1)
+                     .WithMaxSize(128),
+                 fuzztest::Arbitrary<bool>());
+
+}  // namespace
+}  // namespace kernel
+}  // namespace heir
+}  // namespace mlir

lib/Kernel/KernelImplementation.h

@@ -100,14 +100,12 @@ implementRotateAndReduceAccumulation(const T& vector, int64_t period,
 template <typename T>
 std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                  std::shared_ptr<ArithmeticDagNode<T>>>
-implementRotateAndReduceAccumulationRolled(const T& vector, int64_t period,
-                                           int64_t steps,
-                                           DagReducer<T> reduceFunc,
-                                           const DagType& baseType) {
+implementRotateAndReduceAccumulationRolled(
+    std::shared_ptr<ArithmeticDagNode<T>> vectorDag, int64_t period,
+    int64_t steps, DagReducer<T> reduceFunc, const DagType& baseType) {
   using NodeTy = ArithmeticDagNode<T>;
   using NodePtr = std::shared_ptr<NodeTy>;
 
-  auto vectorDag = NodeTy::leaf(vector);
   int64_t numIterations = static_cast<int64_t>(std::log2(steps));
   if (numIterations <= 0) return vectorDag;
 
@@ -132,6 +130,19 @@ implementRotateAndReduceAccumulationRolled(const T& vector, int64_t period,
   return NodeTy::resultAt(loopNode, 0);
 }
 
+// Rolled version of implementRotateAndReduceAccumulation.
+template <typename T>
+std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
+                 std::shared_ptr<ArithmeticDagNode<T>>>
+implementRotateAndReduceAccumulationRolled(const T& vector, int64_t period,
+                                           int64_t steps,
+                                           DagReducer<T> reduceFunc,
+                                           const DagType& baseType) {
+  using NodeTy = ArithmeticDagNode<T>;
+  return implementRotateAndReduceAccumulationRolled<T>(
+      NodeTy::leaf(vector), period, steps, reduceFunc, baseType);
+}
+
 // A function that generalizes the choice of rotation for the "baby stepped
 // operand" of a baby-step giant-step algorithm. This is required because
 // the rotation used in Halevi-Shoup matvec differs from that of bicyclic
@@ -499,6 +510,62 @@ implementDot(const T& lhs, const T& rhs, int64_t steps,
                                                  NodeTy::add);
 }
 
+// Returns an arithmetic DAG that implements a broadcasted reduce kernel.
+template <typename T>
+std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
+                 std::shared_ptr<ArithmeticDagNode<T>>>
+implementBroadcastedReduce(
+    std::shared_ptr<ArithmeticDagNode<T>> vectorDag,
+    std::optional<std::shared_ptr<ArithmeticDagNode<T>>> cleanupMaskDag,
+    int64_t period, int64_t steps, int64_t numSlots, const DagType& dagType,
+    const std::string& reduceOp = "arith.addi", bool unroll = true) {
+  using NodeTy = ArithmeticDagNode<T>;
+  using NodePtr = std::shared_ptr<NodeTy>;
+
+  DagReducer<T> reduceFunc = [&](NodePtr lhs, NodePtr rhs) {
+    if (reduceOp == "arith.addi" || reduceOp == "arith.addf") {
+      return NodeTy::add(lhs, rhs);
+    }
+    if (reduceOp == "arith.muli" || reduceOp == "arith.mulf") {
+      return NodeTy::mul(lhs, rhs);
+    }
+    return NodeTy::add(lhs, rhs);
+  };
+
+  NodePtr reduced;
+  if (unroll) {
+    reduced = implementRotateAndReduceAccumulation<T>(vectorDag, period, steps,
+                                                      reduceFunc);
+  } else {
+    reduced = implementRotateAndReduceAccumulationRolled<T>(
+        vectorDag, period, steps, reduceFunc, dagType);
+  }
+
+  // Check Natural Replication
+  if (steps * period == numSlots) {
+    return reduced;
+  }
+
+  // Shift to last slots
+  int64_t shiftToLast = numSlots - (steps - 1) * period;
+  auto shifted = NodeTy::leftRotate(reduced, shiftToLast);
+
+  NodePtr current = shifted;
+  // Cleanup Mask
+  if (cleanupMaskDag.has_value()) {
+    current = NodeTy::mul(current, cleanupMaskDag.value());
+  }
+
+  // Replication Tree (Left rotations only)
+  for (int64_t rep_shift = 1; rep_shift < steps; rep_shift *= 2) {
+    int64_t rotateAmount = rep_shift * period;
+    auto rotated = NodeTy::leftRotate(current, rotateAmount);
+    current = reduceFunc(current, rotated);
+  }
+
+  return current;
+}
+
 // Returns an arithmetic DAG that implements a baby-step-giant-step between
 // ciphertexts.
 //

lib/Kernel/RotateAndReduceImplTest.cpp

@@ -241,6 +241,78 @@ TEST(RotateAndReduceImplTest, RegressionTest) {
   EXPECT_EQ(expected, actual);
 }
 
+std::vector<int> runBroadcastedReduceImpl(
+    const std::vector<int>& vec, std::optional<std::vector<int>> cleanupMask,
+    int64_t period, int64_t steps, bool unroll = true) {
+  using NodeTy = ArithmeticDagNode<LiteralValue>;
+  using NodePtr = std::shared_ptr<NodeTy>;
+
+  LiteralValue vectorInput(vec);
+  auto vectorDag = NodeTy::leaf(vectorInput);
+
+  std::optional<NodePtr> cleanupMaskDag = std::nullopt;
+  if (cleanupMask.has_value()) {
+    cleanupMaskDag = NodeTy::leaf(LiteralValue(cleanupMask.value()));
+  }
+
+  auto result = implementBroadcastedReduce<LiteralValue>(
+      vectorDag, cleanupMaskDag, period, steps, vec.size(),
+      DagType::intTensor(32, {static_cast<int64_t>(vec.size())}), "arith.addi",
+      unroll);
+
+  return std::get<std::vector<int>>(evalKernel(result)[0].get());
+}
+
+TEST(RotateAndReduceImplTest, BroadcastedReduce_Natural_PowerOfTwo) {
+  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
+  std::vector<int> expected(8, 28);
+
+  for (bool unroll : {true, false}) {
+    std::vector<int> actual =
+        runBroadcastedReduceImpl(vector, std::nullopt, 1, 8, unroll);
+    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
+  }
+}
+
+TEST(RotateAndReduceImplTest, BroadcastedReduce_Natural_Stride) {
+  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
+  std::vector<int> expected = {12, 16, 12, 16, 12, 16, 12, 16};
+
+  for (bool unroll : {true, false}) {
+    std::vector<int> actual =
+        runBroadcastedReduceImpl(vector, std::nullopt, 2, 4, unroll);
+    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
+  }
+}
+
+TEST(RotateAndReduceImplTest, BroadcastedReduce_Masked_Contiguous) {
+  std::vector<int> vector = {0,  1,  2,  3,  4,  5,  6,  7,
+                             10, 11, 12, 13, 14, 15, 16, 17};
+  std::vector<int> mask = {0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1};
+  std::vector<int> expected = {28,  28,  28,  28,  28,  28,  28,  28,
+                               108, 108, 108, 108, 108, 108, 108, 108};
+
+  for (bool unroll : {true, false}) {
+    std::vector<int> actual =
+        runBroadcastedReduceImpl(vector, mask, 1, 8, unroll);
+    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
+  }
+}
+
+TEST(RotateAndReduceImplTest, BroadcastedReduce_Masked_Stride) {
+  std::vector<int> vector = {0,  1,  2,  3,  4,  5,  6,  7,
+                             10, 11, 12, 13, 14, 15, 16, 17};
+  std::vector<int> mask = {0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1};
+  std::vector<int> expected = {12, 16, 12, 16, 12, 16, 12, 16,
+                               52, 56, 52, 56, 52, 56, 52, 56};
+
+  for (bool unroll : {true, false}) {
+    std::vector<int> actual =
+        runBroadcastedReduceImpl(vector, mask, 2, 4, unroll);
+    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
+  }
+}
+
 }  // namespace
 }  // namespace kernel
 }  // namespace heir