Add elementwise mappable variation for rlwe_encode

This change adds a new ElementwiseByOperandOpInterface method that allows an op to specify which dimensions of a tensor-typed operand should be mapped over. This is necessary to support having an elementwise op like lwe.rlwe_encode, whose input is a tensor<k x slots> and whose output is a tensor<k x lwe.lwe_plaintext<...>>. This is needed in turn because server-side secret.conceal lowers to an elementwise trivial_encrypt, which lowers to an elementwise encode + encrypt.

The verifier is updated to ensure that the "mapped shape" (subset of dimensions specified by the interface as mappable) aligns across operands and the output.

Then the ElementwiseToAffine pass is updated to use this interface to support lwe.rlwe_encode.

PiperOrigin-RevId: 888110533

Author: j2kun

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
@@ -72,8 +73,12 @@ LogicalResult RMulRingEltOp::verify() {
 }
 
 LogicalResult TrivialEncryptOp::verify() {
-  auto plaintextSpace = this->getInput().getType().getPlaintextSpace();
-  auto outPlaintextSpace = this->getOutput().getType().getPlaintextSpace();
+  auto plaintextSpace =
+      cast<LWEPlaintextType>(getElementTypeOrSelf(this->getInput().getType()))
+          .getPlaintextSpace();
+  auto outPlaintextSpace =
+      cast<LWECiphertextType>(getElementTypeOrSelf(this->getOutput().getType()))
+          .getPlaintextSpace();
 
   if (plaintextSpace != outPlaintextSpace) {
     return this->emitOpError()
@@ -174,12 +179,8 @@ LogicalResult EncodeOp::verify() {
 }
 
 LogicalResult RLWEEncodeOp::verify() {
-  if (auto tensorTy = dyn_cast<ShapedType>(getInput().getType())) {
-    if (tensorTy.getRank() > 1) {
-      return emitOpError() << "RLWEEncodeOp only supports 1D tensors";
-    }
-  }
-  return verifyEncodingAndTypeMatch(getInput().getType(), getEncoding());
+  return verifyEncodingAndTypeMatch(getElementTypeOrSelf(getInput().getType()),
+                                    getEncoding());
 }
 
 LogicalResult RLWEDecodeOp::verify() {
@@ -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

@@ -23,7 +23,7 @@ class HasEncoding<
     "the first arg's type's encoding matches the given encoding",
     CPred<
       comparator # "(" #
-      "::llvm::cast<lwe::" # ty # ">($" # encodingHolder # ".getType()).getPlaintextSpace().getEncoding(), " #
+      "::llvm::cast<lwe::" # ty # ">(getElementTypeOrSelf($" # encodingHolder # ".getType())).getPlaintextSpace().getEncoding(), " #
       "$" # encoding # ")"
     >
   >;
@@ -38,8 +38,8 @@ class EncodingsMatch<
     "the first arg's type's encoding matches the given encoding",
     CPred<
       comparator # "(" #
-      "::llvm::cast<lwe::" # ty1 # ">($" # encodingHolder1 # ".getType()).getPlaintextSpace().getEncoding(), " #
-      "::llvm::cast<lwe::" # ty2 # ">($" # encodingHolder2 # ".getType()).getPlaintextSpace().getEncoding())"
+      "::llvm::cast<lwe::" # ty1 # ">(getElementTypeOrSelf($" # encodingHolder1 # ".getType())).getPlaintextSpace().getEncoding(), " #
+      "::llvm::cast<lwe::" # ty2 # ">(getElementTypeOrSelf($" # encodingHolder2 # ".getType())).getPlaintextSpace().getEncoding())"
     >
   >;
 
@@ -102,16 +102,14 @@ def LWE_EncodeOp : LWE_Op<"encode"> {
 
   let results = (outs LWEPlaintext:$output);
   let assemblyFormat = "$input attr-dict `:` qualified(type($input)) `to` qualified(type($output))";
-
-  // Verify that the input type and the application data are compatible.
   let hasVerifier = 1;
 }
 
 def LWE_TrivialEncryptOp: LWE_Op<"trivial_encrypt", [
     EncodingsMatch<"input", "LWEPlaintextType", "output", "LWECiphertextType">]> {
   let summary = "Create a trivial encryption of a plaintext.";
-  let arguments = (ins LWEPlaintext:$input);
-  let results = (outs LWECiphertext:$output);
+  let arguments = (ins LWEPlaintextLike:$input);
+  let results = (outs LWECiphertextLike:$output);
   let assemblyFormat = [{
     $input attr-dict `:`  type(operands) `->` type(results)
   }];
@@ -183,20 +181,29 @@ 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.
 
     Examples:
 
     ```
-    %Y = lwe.rlwe_encode %value {encoding = #enc, ring = #ring}: i1 to !lwe.rlwe_plaintext<encoding = #enc, ring = #ring>
+    %Y = lwe.rlwe_encode %value {encoding = #enc, ring = #ring}: i16 to !lwe.rlwe_plaintext<encoding = #enc, ring = #ring>
     ```
   }];
 
@@ -206,10 +213,8 @@ def LWE_RLWEEncodeOp : LWE_Op<"rlwe_encode", [HasEncoding<"output", "encoding",
     Polynomial_RingAttr:$ring
   );
 
-  let results = (outs LWEPlaintext:$output);
+  let results = (outs LWEPlaintextLike:$output);
   let assemblyFormat = "$input attr-dict `:` qualified(type($input)) `->` qualified(type($output))";
-
-  // Verify that the input type and the encoding are compatible.
   let hasVerifier = true;
 }
 

lib/Dialect/Secret/Conversions/Patterns.cpp

@@ -17,6 +17,7 @@
 #include "lib/Utils/ContextAwareTypeConversion.h"
 #include "llvm/include/llvm/ADT/STLExtras.h"           // from @llvm-project
 #include "llvm/include/llvm/Support/Debug.h"           // from @llvm-project
+#include "llvm/include/llvm/Support/DebugLog.h"        // from @llvm-project
 #include "llvm/include/llvm/Support/FormatVariadic.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
@@ -62,13 +63,47 @@ Value insertKeyArgument(func::FuncOp parentFunc, Type encryptionKeyType,
   return keyBlockArg;
 }
 
+LogicalResult ConvertClientConceal::lowerToTrivialEncryption(
+    secret::ConcealOp op, OpAdaptor adaptor,
+    ContextAwareConversionPatternRewriter& rewriter) const {
+  LDBG() << "Lowering conceal op as a trivial encryption";
+  auto mgmtAttrResult = getTypeConverter()->getContextualAttr(op.getResult());
+  if (failed(mgmtAttrResult)) {
+    return rewriter.notifyMatchFailure(op, "found no mgmt attr");
+  }
+  Type resultTy = getTypeConverter()->convertType(op.getResult().getType(),
+                                                  mgmtAttrResult.value());
+  auto ctTy = cast<lwe::LWECiphertextType>(getElementTypeOrSelf(resultTy));
+  auto plaintextTy =
+      lwe::LWEPlaintextType::get(op.getContext(), ctTy.getPlaintextSpace());
+
+  Type encodeOpResultTy = plaintextTy;
+  if (auto resultTensorTy = cast<RankedTensorType>(resultTy)) {
+    encodeOpResultTy =
+        RankedTensorType::get(resultTensorTy.getShape(), plaintextTy);
+  }
+
+  // Intentionally use op.getCleartext() because we don't want to type-convert
+  // the input to a ciphertext.
+  auto encoded = lwe::RLWEEncodeOp::create(
+      rewriter, op.getLoc(), encodeOpResultTy, op.getCleartext(),
+      ctTy.getPlaintextSpace().getEncoding(),
+      ctTy.getPlaintextSpace().getRing());
+  auto newOp =
+      lwe::TrivialEncryptOp::create(rewriter, op.getLoc(), resultTy, encoded);
+  newOp->setAttrs(op->getAttrs());
+  rewriter.replaceOp(op, newOp);
+  return success();
+}
+
 LogicalResult ConvertClientConceal::matchAndRewrite(
     secret::ConcealOp op, OpAdaptor adaptor,
     ContextAwareConversionPatternRewriter& rewriter) const {
   func::FuncOp parentFunc = op->getParentOfType<func::FuncOp>();
-  if (!parentFunc || !parentFunc->hasAttr(kClientEncFuncAttrName)) {
-    return op->emitError() << "expected to be inside a function with attribute "
-                           << kClientEncFuncAttrName;
+  if (!parentFunc) return failure();
+
+  if (!isClientHelper(parentFunc)) {
+    return lowerToTrivialEncryption(op, adaptor, rewriter);
   }
 
   // The encryption func encrypts a single value, so it must have a single

lib/Dialect/Secret/Conversions/Patterns.h

@@ -39,6 +39,10 @@ struct ConvertClientConceal
       ContextAwareConversionPatternRewriter& rewriter) const override;
 
  private:
+  LogicalResult lowerToTrivialEncryption(
+      secret::ConcealOp op, OpAdaptor adaptor,
+      ContextAwareConversionPatternRewriter& rewriter) const;
+
   bool usePublicKey;
   polynomial::RingAttr ring;
 };