Support non-constant splats in RotationAnalysis

PiperOrigin-RevId: 889451284

Author: j2kun

lib/Analysis/RotationAnalysis/DagBuilder.cpp

@@ -293,12 +293,13 @@ FailureOr<NodePtr> DagBuilder::visit(tensor::SplatOp op) {
   IntegerAttr attr;
   if (matchPattern(op.getInput(), m_Constant(&attr))) {
     LDBG() << "Matched splatted value to constant scalar " << attr;
+    auto dagNode =
+        Node::splat(attr.getInt(), mlirTypeToDagType(op.getResult().getType()));
+    valueToNode[op.getResult()] = dagNode;
+    return dagNode;
   }
 
-  auto dagNode =
-      Node::splat(attr.getInt(), mlirTypeToDagType(op.getResult().getType()));
-  valueToNode[op.getResult()] = dagNode;
-  return dagNode;
+  return findNodeOrMakeNewVariable(op.getResult());
 }
 
 FailureOr<NodePtr> DagBuilder::visit(arith::DivSIOp op) {

lib/Dialect/HEIRInterfaces.cpp

@@ -75,67 +75,68 @@ LogicalResult verifyElementwiseByOperandImpl(
     ElementwiseByOperandOpInterface opInterface) {
   Operation* op = opInterface.getOperation();
 
-  auto typeToShapeStr = [](Type type) {
-    if (auto rankedTensorType = dyn_cast<RankedTensorType>(type)) {
-      std::string shapeStr = "(";
-      for (auto dim : rankedTensorType.getShape()) {
-        shapeStr += std::to_string(dim) + ",";
-      }
-      shapeStr += ")";
-      return shapeStr;
-    }
-    return std::string("(not statically ranked)");
+  auto typeToShapeStr = [](ArrayRef<int64_t> shape) {
+    return "(" +
+           llvm::join(
+               llvm::map_range(shape,
+                               [](int64_t dim) { return std::to_string(dim); }),
+               ", ") +
+           ")";
   };
 
-  std::optional<TensorType> tensorType;
-  int64_t operandIndex = 0;
+  std::optional<SmallVector<int64_t>> mappedShape;
+  int64_t firstMappableOperandIndex = -1;
+
   for (auto [i, operand] : llvm::enumerate(op->getOperands())) {
     auto thisTensorType = dyn_cast<TensorType>(operand.getType());
-    if (!thisTensorType)
-      // Non-tensor types are acceptable, and need not be specified as mappable
-      // or not mappable by the interface.
-      continue;
+    if (!thisTensorType) continue;
 
     if (opInterface.operandIsMappable(i)) {
-      if (!tensorType) {
-        tensorType = thisTensorType;
-        operandIndex = i;
+      SmallVector<int64_t> thisMappedShape;
+      for (int dim : opInterface.mappedDimensionsForOperand(i)) {
+        thisMappedShape.push_back(thisTensorType.getDimSize(dim));
+      }
+
+      if (!mappedShape) {
+        mappedShape = thisMappedShape;
+        firstMappableOperandIndex = i;
         continue;
       }
 
-      if (thisTensorType.getShape() != tensorType->getShape()) {
+      if (thisMappedShape != *mappedShape) {
         return op->emitOpError()
-               << "expected all mappable operands to have the same shape, "
-               << "but found shape " << typeToShapeStr(*tensorType)
-               << " at operand " << operandIndex << " and "
-               << typeToShapeStr(thisTensorType) << " at operand " << i;
+               << "expected all mappable operands to have the same mapped "
+                  "shape, but found mapped shape "
+               << typeToShapeStr(*mappedShape) << " at operand "
+               << firstMappableOperandIndex << " and "
+               << typeToShapeStr(thisMappedShape) << " at operand " << i;
       }
     }
   }
 
   for (auto [i, result] : llvm::enumerate(op->getResults())) {
     auto thisTensorType = dyn_cast<TensorType>(result.getType());
-    if (tensorType && !thisTensorType)
+    if (mappedShape && !mappedShape->empty() && !thisTensorType)
       return op->emitOpError()
-             << "expected all results operands to have the same tensor shape, "
-             << "as the mappable input operands, but found shape "
-             << typeToShapeStr(*tensorType) << " at operand " << operandIndex
-             << " and result " << i << " of non-tensor type "
-             << result.getType();
+             << "expected all results to be tensors with shape "
+             << typeToShapeStr(*mappedShape)
+             << " due to mappable operands, but result " << i
+             << " is of non-tensor type " << result.getType();
 
-    if (!tensorType && thisTensorType)
+    if (!mappedShape && thisTensorType)
       return op->emitOpError()
-             << "No operands were tensor typed, but result at index " << i
-             << " is a tensor of shape " << typeToShapeStr(thisTensorType);
+             << "No operands were mappable, but result at index " << i
+             << " is a tensor of shape "
+             << typeToShapeStr(thisTensorType.getShape());
 
-    if (tensorType && thisTensorType &&
-        thisTensorType.getShape() != tensorType->getShape()) {
+    if (mappedShape && thisTensorType &&
+        thisTensorType.getShape() != ArrayRef<int64_t>(*mappedShape)) {
       return op->emitOpError()
              << "expected all tensor results to have the same shape as "
                 "mappable operands, but found shape "
-             << typeToShapeStr(*tensorType) << " at operand " << operandIndex
-             << " and shape " << typeToShapeStr(thisTensorType) << " at result "
-             << i;
+             << typeToShapeStr(*mappedShape) << " at operand "
+             << firstMappableOperandIndex << " and shape "
+             << typeToShapeStr(thisTensorType.getShape()) << " at result " << i;
     }
   }
 

lib/Dialect/HEIRInterfaces.td

@@ -119,7 +119,8 @@ def ElementwiseByOperandOpInterface : OpInterface<"ElementwiseByOperandOpInterfa
     be replicated and not mapped.
 
     This trait is designed to work with the `convert-elementwise-to-affine`
-    pass.
+    pass. This pass assumes that the mapped dimensions for each operand
+    constitute a contiguous prefix of that operand's dimensions.
   }];
 
   let methods = [
@@ -130,6 +131,24 @@ def ElementwiseByOperandOpInterface : OpInterface<"ElementwiseByOperandOpInterfa
       /*args=*/(ins "unsigned":$operandIndex)
     >,
 
+    InterfaceMethod<
+      /*desc=*/"Return the subset of indices along which this operand is mapped elementise",
+      /*retTy=*/"llvm::SmallVector<int>",
+      /*methodName=*/"mappedDimensionsForOperand",
+      /*args=*/(ins "unsigned":$operandIndex),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        // By default, map over all dimensions.
+        auto operandType = $_op->getOperand(operandIndex).getType();
+        if (auto tensorTy = llvm::dyn_cast<mlir::ShapedType>(operandType)) {
+          auto r = llvm::seq<int>(0, tensorTy.getRank());
+          llvm::SmallVector<int, 4> v(r.begin(), r.end());
+          return v;
+        }
+        return {};
+      }]
+    >,
+
     // It may make sense some day to allow only a subset of op results to be
     // mapped, but for now we assume all results are tensors of the same shape
     // as the mapped operands.

lib/Dialect/LWE/IR/LWEOps.cpp

@@ -10,6 +10,7 @@
 #include "lib/Dialect/Polynomial/IR/PolynomialAttributes.h"
 #include "lib/Dialect/RNS/IR/RNSOps.h"
 #include "lib/Dialect/RNS/IR/RNSTypes.h"
+#include "llvm/include/llvm/ADT/Sequence.h"              // from @llvm-project
 #include "llvm/include/llvm/ADT/TypeSwitch.h"            // from @llvm-project
 #include "llvm/include/llvm/Support/ErrorHandling.h"     // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
@@ -325,6 +326,23 @@ void RMulPlainOp::getCanonicalizationPatterns(RewritePatternSet& results,
   results.add<lwe::PutCiphertextInFirstOperand<RMulPlainOp>>(context);
 }
 
+bool RLWEEncodeOp::operandIsMappable(unsigned operandIndex) {
+  // only `input`
+  return operandIndex == 0;
+}
+
+SmallVector<int> RLWEEncodeOp::mappedDimensionsForOperand(
+    unsigned operandIndex) {
+  auto operandType = getOperation()->getOperand(operandIndex).getType();
+  if (auto tensorTy = dyn_cast<mlir::ShapedType>(operandType)) {
+    // rank - 1 because the trailing dimension is the slot dimension.
+    auto r = llvm::seq<int>(0, tensorTy.getRank() - 1);
+    llvm::SmallVector<int, 4> v(r.begin(), r.end());
+    return v;
+  }
+  return {};
+}
+
 }  // namespace lwe
 }  // namespace heir
 }  // namespace mlir

lib/Dialect/LWE/IR/LWEOps.td

@@ -73,7 +73,7 @@ class LWE_BinOp<string mnemonic, list<Trait> traits = []> :
 }
 
 class LWE_CiphertextPlaintextOp<string mnemonic, list<Trait> traits = []> :
-        LWE_Op<mnemonic, traits #  [IsCiphertextPlaintextOp, InferTypeOpAdaptor]> {
+        LWE_Op<mnemonic, traits #  [IsCiphertextPlaintextOp, InferTypeOpAdaptor, ElementwiseMappable]> {
   let arguments = (ins LWEPlaintextOrCiphertextLike:$lhs, LWEPlaintextOrCiphertextLike:$rhs);
   let results = (outs LWECiphertextLike:$output);
 }
@@ -183,12 +183,21 @@ def LWE_MulScalarOp : LWE_Op<"mul_scalar", [ElementwiseMappable,
   let hasFolder = 1;
 }
 
-def LWE_RLWEEncodeOp : LWE_Op<"rlwe_encode", [HasEncoding<"output", "encoding", "LWEPlaintextType">]> {
+def LWE_RLWEEncodeOp : LWE_Op<"rlwe_encode", [
+    HasEncoding<"output", "encoding", "LWEPlaintextType">,
+    DeclareOpInterfaceMethods<ElementwiseByOperandOpInterface, [
+      "operandIsMappable",
+      // When encoding a tensor<3x1024xi16>, for example, the leading dimension
+      // is mapped over and the trailing dimension is encoded into slots, so the
+      // result type is tensor<3x!pt>.
+      "mappedDimensionsForOperand",
+    ]>,
+]> {
   let summary = "Encode an integer to yield an RLWE plaintext";
   let description = [{
-    Encode an integer to yield an RLWE plaintext.
+    Encode a packed cleartext tensor to yield an RLWE plaintext.
 
-    This op uses a an encoding attribute to encode the bits of the integer into
+    This op uses a an encoding attribute to encode the bits of the input into
     an RLWE plaintext value that can then be encrypted. CKKS cleartext inputs may
     be floating points, and a scaling factor described by the encoding will be
     applied.

lib/Dialect/LWE/Transforms/ImplementTrivialEncryptionAsAddition.cpp

@@ -146,20 +146,19 @@ func::FuncOp getOrCreateEncryptionOfZerosFunc(func::FuncOp parentFunc,
 
   Type coeffType =
       plaintextType.getPlaintextSpace().getRing().getCoefficientType();
-  Type cleartextElementType = builder.getIntegerType(
-      llvm::TypeSwitch<Type, int>(coeffType)
-          .Case<IntegerType, FloatType>(
-              [&](auto ty) { return ty.getIntOrFloatBitWidth(); })
-          .Case<mod_arith::ModArithType>([&](auto ty) {
-            return ty.getModulus().getType().getIntOrFloatBitWidth();
-          })
+  Type cleartextElementType =
+      llvm::TypeSwitch<Type, Type>(coeffType)
+          .Case<IntegerType, FloatType>([&](auto ty) { return ty; })
+          .Case<mod_arith::ModArithType>(
+              [&](auto ty) { return ty.getModulus().getType(); })
           .Default([&](auto ty) {
             originalOp->emitOpError()
                 << "has unsupported plaintext coefficient type; can't "
                    "determine what constant type to use for encryption of "
                    "zero.";
-            return 0;
-          }));
+            return Type();
+          });
+  if (!cleartextElementType) return nullptr;
 
   RankedTensorType zeroType =
       RankedTensorType::get({numSlots}, cleartextElementType);