Preserve linear_transform in Lattigo backend

lib/Dialect/LWE/Conversions/LWEToLattigo/LWEToLattigo.cpp

@@ -697,6 +697,97 @@ struct ConvertOrionChebyshevOp
   }
 };
 
+struct ConvertKernelLinearTransformOp
+    : public OpConversionPattern<kernel::LinearTransformOp> {
+  using OpConversionPattern<kernel::LinearTransformOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      kernel::LinearTransformOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter& rewriter) const override {
+    LLVM_DEBUG(llvm::dbgs() << "Lowering Kernel LinearTransformOp\n");
+
+    FailureOr<Value> evaluatorResult =
+        getContextualEvaluator<lattigo::CKKSEvaluatorType>(op.getOperation());
+    if (failed(evaluatorResult)) {
+      return rewriter.notifyMatchFailure(
+          op, "CKKS evaluator not found in function context");
+    }
+    Value evaluator = evaluatorResult.value();
+
+    FailureOr<Value> encoderResult =
+        getContextualEvaluator<lattigo::CKKSEncoderType>(op.getOperation());
+    if (failed(encoderResult)) {
+      return rewriter.notifyMatchFailure(
+          op, "CKKS encoder not found in function context");
+    }
+    Value encoder = encoderResult.value();
+
+    // Extract level from input LWE ciphertext type
+    auto lweType = dyn_cast<lwe::LWECiphertextType>(op.getInput().getType());
+    if (!lweType) {
+      return rewriter.notifyMatchFailure(op, "input is not LWE ciphertext");
+    }
+    auto modulusChain = lweType.getModulusChain();
+    if (!modulusChain) {
+      return rewriter.notifyMatchFailure(op,
+                                         "input LWE type has no modulus chain");
+    }
+    int64_t levelQ =
+        modulusChain.getElements().size() - 1 - modulusChain.getCurrent();
+
+    // Convert diagonal_indices from I64 to I32 (Lattigo CKKSLinearTransformOp
+    // expects I32)
+    auto diagonalIndicesAttr = op.getDiagonalIndices();
+    std::vector<int32_t> diagonalIndicesI32;
+    for (auto val : diagonalIndicesAttr) {
+      diagonalIndicesI32.push_back(static_cast<int32_t>(val));
+    }
+    auto diagonalIndicesI32Attr =
+        rewriter.getDenseI32ArrayAttr(diagonalIndicesI32);
+
+    // logBabyStepGiantStepRatio
+    // For now default to 0.
+    int64_t logBSGSRatio = 0;
+
+    auto levelQAttr = rewriter.getI64IntegerAttr(levelQ);
+    auto logBSGSRatioAttr = rewriter.getI64IntegerAttr(logBSGSRatio);
+
+    auto diagonalsAttr = op.getDiagonals();
+    Value diagonalsValue =
+        rewriter.create<arith::ConstantOp>(op.getLoc(), diagonalsAttr);
+
+    auto linearTransformOp = rewriter.create<lattigo::CKKSLinearTransformOp>(
+        op.getLoc(), adaptor.getInput().getType(), evaluator, encoder,
+        adaptor.getInput(), diagonalsValue, diagonalIndicesI32Attr, levelQAttr,
+        logBSGSRatioAttr);
+
+    auto outputLweType =
+        dyn_cast<lwe::LWECiphertextType>(op.getResult().getType());
+    if (!outputLweType) {
+      return rewriter.notifyMatchFailure(op, "output is not LWE ciphertext");
+    }
+    auto outputModulusChain = outputLweType.getModulusChain();
+    if (!outputModulusChain) {
+      return rewriter.notifyMatchFailure(
+          op, "output LWE type has no modulus chain");
+    }
+
+    Value result = linearTransformOp.getResult();
+    if (outputModulusChain.getCurrent() < modulusChain.getCurrent()) {
+      int64_t diff =
+          modulusChain.getCurrent() - outputModulusChain.getCurrent();
+      for (int64_t i = 0; i < diff; ++i) {
+        auto rescaleOp = rewriter.create<lattigo::CKKSRescaleNewOp>(
+            op.getLoc(), result.getType(), evaluator, result);
+        result = rescaleOp.getResult();
+      }
+    }
+
+    rewriter.replaceOp(op, result);
+    return success();
+  }
+};
+
 struct ConvertKernelEvalChebyshevOp
     : public OpConversionPattern<kernel::EvalChebyshevOp> {
   using OpConversionPattern<kernel::EvalChebyshevOp>::OpConversionPattern;
@@ -918,7 +1009,7 @@ struct LWEToLattigo : public impl::LWEToLattigoBase<LWEToLattigo> {
         .addIllegalOp<lwe::RLWEEncryptOp, lwe::RLWEDecryptOp, lwe::RLWEEncodeOp,
                       lwe::RLWEDecodeOp, lwe::RAddOp, lwe::RSubOp, lwe::RMulOp,
                       lwe::RMulPlainOp, lwe::RSubPlainOp, lwe::RAddPlainOp,
-                      kernel::EvalChebyshevOp>();
+                      kernel::EvalChebyshevOp, kernel::LinearTransformOp>();
 
     RewritePatternSet patterns(context);
     addStructuralConversionPatterns(typeConverter, patterns, target);
@@ -1098,8 +1189,8 @@ struct LWEToLattigo : public impl::LWEToLattigoBase<LWEToLattigo> {
           ConvertCKKSEncryptOp, ConvertCKKSDecryptOp, ConvertCKKSEncodeOp,
           ConvertCKKSDecodeOp, ConvertCKKSLevelReduceOp,
           ConvertCKKSBootstrappingOp, ConvertOrionLinearTransformOp,
-          ConvertOrionChebyshevOp, ConvertKernelEvalChebyshevOp>(typeConverter,
-                                                                 context);
+          ConvertOrionChebyshevOp, ConvertKernelEvalChebyshevOp,
+          ConvertKernelLinearTransformOp>(typeConverter, context);
     }
     // Misc
 

lib/Dialect/Secret/Conversions/SecretToBGV/SecretToBGV.cpp

@@ -100,7 +100,10 @@ class SecretToBGVTypeConverter
         ring(rlweRing),
         plaintextModulus(ptm),
         isBFV(isBFV) {
-    addConversion([](Type type, Attribute attr) { return type; });
+    addConversion([](Type type, Attribute attr) -> std::optional<Type> {
+      if (isa<secret::SecretType>(type)) return std::nullopt;
+      return type;
+    });
     addConversion(
         [this](RankedTensorType type, mgmt::MgmtAttr mgmtAttr) -> Type {
           // For cases like tensor.empty + mgmt.init, we need to convert this
@@ -217,22 +220,16 @@ struct SecretToBGV : public impl::SecretToBGVBase<SecretToBGV> {
     bool usePublicKey =
         schemeParamAttr.getEncryptionType() == bgv::BGVEncryptionType::pk;
 
-    // NOTE: 2 ** logN != minSlotCount
-    // they have different semantic
-    // auto logN = schemeParamAttr.getLogN();
     auto plaintextModulus = schemeParamAttr.getPlaintextModulus();
-
-    // pass option minSlotCount is actually the number of slots
-    // TODO(#1402): use a proper name for BGV
     auto rlweRing = getRlweRNSRing(context, schemeParamAttr.getQ().asArrayRef(),
-                                   minSlotCount);
+                                   1 << schemeParamAttr.getLogN());
     if (failed(rlweRing)) {
       return signalPassFailure();
     }
     // Ensure that all secret types are uniform and have last dimension
-    // matching the ring parameter size. In other words, this asserts that any
-    // data-semantic tensors have been converted to ciphertext-semantic tensors
-    // with the correct shape.
+    // less than or equal to the ring parameter size. In other words, this
+    // asserts that any data-semantic tensors have been converted to
+    // ciphertext-semantic tensors with the correct shape.
     Operation* foundOp = walkAndDetect(module, [&](Operation* op) {
       ValueRange valuesToCheck = op->getOperands();
       if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
@@ -241,7 +238,7 @@ struct SecretToBGV : public impl::SecretToBGVBase<SecretToBGV> {
       for (auto value : valuesToCheck) {
         if (auto secretTy = dyn_cast<secret::SecretType>(value.getType())) {
           auto tensorTy = dyn_cast<RankedTensorType>(secretTy.getValueType());
-          if (tensorTy && tensorTy.getDimSize(tensorTy.getRank() - 1) !=
+          if (tensorTy && tensorTy.getDimSize(tensorTy.getRank() - 1) >
                               rlweRing.value()
                                   .getPolynomialModulus()
                                   .getPolynomial()
@@ -255,7 +252,7 @@ struct SecretToBGV : public impl::SecretToBGVBase<SecretToBGV> {
     if (foundOp != nullptr) {
       foundOp->emitError(
           "expected secret types to be tensors with last dimension "
-          "matching ring parameter");
+          "less than or equal to ring parameter");
       signalPassFailure();
       return;
     }

lib/Dialect/Secret/Conversions/SecretToCKKS/SecretToCKKS.cpp

@@ -1,7 +1,9 @@
 #include "lib/Dialect/Secret/Conversions/SecretToCKKS/SecretToCKKS.h"
 
+#include <algorithm>
 #include <cassert>
 #include <cstdint>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -100,7 +102,10 @@ class SecretToCKKSTypeConverter
   SecretToCKKSTypeConverter(MLIRContext* ctx, polynomial::RingAttr rlweRing)
       : UniquelyNamedAttributeAwareTypeConverter(
             mgmt::MgmtDialect::kArgMgmtAttrName) {
-    addConversion([](Type type, Attribute attr) { return type; });
+    addConversion([](Type type, Attribute attr) -> std::optional<Type> {
+      if (isa<secret::SecretType>(type)) return std::nullopt;
+      return type;
+    });
     addConversion(
         [this](RankedTensorType type, mgmt::MgmtAttr mgmtAttr) -> Type {
           // For cases like tensor.empty + mgmt.init, we need to convert this
@@ -244,6 +249,118 @@ class SecretGenericPlaintextDivision
   }
 };
 
+struct LinearTransformOpConversion
+    : public ContextAwareOpConversionPattern<secret::GenericOp> {
+  LinearTransformOpConversion(const ContextAwareTypeConverter& typeConverter,
+                              MLIRContext* context, int64_t ringDim,
+                              PatternBenefit benefit = 1)
+      : ContextAwareOpConversionPattern<secret::GenericOp>(typeConverter,
+                                                           context, benefit),
+        ringDim(ringDim) {}
+
+  LogicalResult matchAndRewrite(
+      secret::GenericOp op, OpAdaptor adaptor,
+      ContextAwareConversionPatternRewriter& rewriter) const override {
+    if (op.getBody()->getOperations().size() > 2) {
+      return failure();
+    }
+
+    auto& innerOp = op.getBody()->getOperations().front();
+    auto ltOp = dyn_cast<kernel::LinearTransformOp>(innerOp);
+    if (!ltOp) {
+      return failure();
+    }
+
+    // Convert inputs
+    SmallVector<Value> inputs;
+    for (Value operand : ltOp->getOperands()) {
+      if (auto* secretArg = op.getOpOperandForBlockArgument(operand)) {
+        inputs.push_back(adaptor.getInputs()[secretArg->getOperandNumber()]);
+      } else {
+        inputs.push_back(operand);
+      }
+    }
+
+    // Convert result types
+    SmallVector<Type> resultTypes;
+    if (failed(getTypeConverter()->convertTypes(op.getResultTypes(),
+                                                op.getResults(), resultTypes)))
+      return failure();
+
+    // Preserve attributes (similar to SecretGenericOpConversion)
+    SmallVector<NamedAttribute> attrsToPreserve;
+    for (auto& namedAttr : ltOp->getDialectAttrs()) {
+      attrsToPreserve.push_back(namedAttr);
+    }
+    for (auto attrName : ltOp.getAttributeNames()) {
+      if (attrName == "diagonals")
+        continue;  // We will handle diagonals separately
+      if (auto attr = ltOp->getAttr(attrName)) {
+        attrsToPreserve.push_back(rewriter.getNamedAttr(attrName, attr));
+      }
+    }
+
+    // Pad diagonals
+    auto diagonalsAttr = cast<DenseElementsAttr>(ltOp.getDiagonals());
+    auto diagonalsType = cast<RankedTensorType>(diagonalsAttr.getType());
+    auto shape = diagonalsType.getShape();
+    int64_t numDiagonals = shape[0];
+    int64_t numCols = shape[1];
+
+    int64_t actualSlots = ringDim / 2;  // CKKS assumption
+
+    DenseElementsAttr newDiagonalsAttr;
+    if (numCols == actualSlots) {
+      newDiagonalsAttr = diagonalsAttr;
+    } else {
+      if (numCols > actualSlots) {
+        return ltOp.emitOpError("diagonals slot size (")
+               << numCols << ") is larger than actual slots (" << actualSlots
+               << ")";
+      }
+      SmallVector<Attribute> paddedValues;
+      auto elementValues = diagonalsAttr.getValues<Attribute>();
+      auto elemType = diagonalsType.getElementType();
+      Attribute zeroAttr = rewriter.getZeroAttr(elemType);
+
+      for (int64_t i = 0; i < numDiagonals; ++i) {
+        for (int64_t j = 0; j < numCols; ++j) {
+          paddedValues.push_back(elementValues[i * numCols + j]);
+        }
+        for (int64_t j = numCols; j < actualSlots; ++j) {
+          paddedValues.push_back(zeroAttr);
+        }
+      }
+
+      auto newDiagonalsType =
+          RankedTensorType::get({numDiagonals, actualSlots}, elemType);
+      newDiagonalsAttr = DenseElementsAttr::get(newDiagonalsType, paddedValues);
+    }
+    attrsToPreserve.push_back(
+        rewriter.getNamedAttr("diagonals", newDiagonalsAttr));
+
+    // Handle mgmt attrs
+    convertArrayOfDicts(op.getAllResultAttrsAttr(), attrsToPreserve);
+    convertArrayOfDicts(op.getAllOperandAttrsAttr(), attrsToPreserve);
+    DenseSet<StringRef> seenNames;
+    SmallVector<NamedAttribute> dedupedAttrsToPreserve;
+    for (auto attr : llvm::reverse(attrsToPreserve)) {
+      if (seenNames.insert(attr.getName().getValue()).second) {
+        dedupedAttrsToPreserve.push_back(attr);
+      }
+    }
+    std::reverse(dedupedAttrsToPreserve.begin(), dedupedAttrsToPreserve.end());
+    auto newLtOp = kernel::LinearTransformOp::create(
+        rewriter, ltOp.getLoc(), resultTypes, inputs, dedupedAttrsToPreserve);
+
+    rewriter.replaceOp(op, newLtOp->getResults());
+    return success();
+  }
+
+ private:
+  int64_t ringDim;
+};
+
 struct SecretToCKKS : public impl::SecretToCKKSBase<SecretToCKKS> {
   using SecretToCKKSBase::SecretToCKKSBase;
 
@@ -266,7 +383,7 @@ struct SecretToCKKS : public impl::SecretToCKKSBase<SecretToCKKS> {
     // pass option minSlotCount is actually the number of slots
     // TODO(#1402): use a proper name for CKKS
     auto rlweRing = getRlweRNSRing(context, schemeParamAttr.getQ().asArrayRef(),
-                                   minSlotCount);
+                                   1 << schemeParamAttr.getLogN());
     if (failed(rlweRing)) {
       return signalPassFailure();
     }
@@ -310,6 +427,9 @@ struct SecretToCKKS : public impl::SecretToCKKSBase<SecretToCKKS> {
         SecretGenericOpLevelReduceConversion<ckks::LevelReduceOp>>(
         typeConverter, context);
 
+    int64_t ringDim = 1 << schemeParamAttr.getLogN();
+    patterns.add<LinearTransformOpConversion>(typeConverter, context, ringDim);
+
     patterns.add<ConvertClientConceal>(typeConverter, context, usePublicKey,
                                        rlweRing.value());
     patterns.add<ConvertClientReveal>(typeConverter, context, rlweRing.value());