Support tensor reshape op in secret-to-ckks

PiperOrigin-RevId: 800485015

Author: asraa

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

@@ -38,6 +38,7 @@
 #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/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
 #include "mlir/include/mlir/IR/Diagnostics.h"            // from @llvm-project
 #include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
@@ -153,7 +154,7 @@ class SecretToCKKSTypeConverter
 
     auto ciphertext = lwe::LWECiphertextType::get(
         ctx,
-        lwe::ApplicationDataAttr::get(ctx, type.getValueType(),
+        lwe::ApplicationDataAttr::get(ctx, valueTy,
                                       lwe::NoOverflowAttr::get(ctx)),
         lwe::PlaintextSpaceAttr::get(
             ctx, plaintextRing,
@@ -164,25 +165,60 @@ class SecretToCKKSTypeConverter
         lwe::KeyAttr::get(ctx, 0),
         lwe::ModulusChainAttr::get(ctx, moduliChain, level));
 
-    // Return a single ciphertext if inputs are packed into a single
-    // ciphertext SIMD slot or the secret value type is a scalar.
-    if (this->packTensorInSlots_ || !isa<TensorType>(valueTy)) {
-      return ciphertext;
+    // Return a single ciphertext if the input is a scalar.
+    if (!isa<TensorType>(valueTy)) return ciphertext;
+
+    // The input is a tensor type.
+    assert(dyn_cast<RankedTensorType>(valueTy) &&
+           "expected ranked tensor type");
+    auto tensorTy = cast<RankedTensorType>(valueTy);
+    // If the input is packed into a ciphertext SIMD slots (i.e. it is a tensor
+    // of shape NxciphertextSize) then return Nxciphertext.
+    if (this->packTensorInSlots_) {
+      Type underlyingTy;
+      if (tensorTy.getRank() == 1) {
+        underlyingTy = valueTy;
+        auto ciphertext = lwe::LWECiphertextType::get(
+            ctx,
+            lwe::ApplicationDataAttr::get(ctx, underlyingTy,
+                                          lwe::NoOverflowAttr::get(ctx)),
+            lwe::PlaintextSpaceAttr::get(
+                ctx, plaintextRing,
+                lwe::InverseCanonicalEncodingAttr::get(ctx, scale)),
+            lwe::CiphertextSpaceAttr::get(
+                ctx, getRlweRNSRingWithLevel(ring_, level),
+                lwe::LweEncryptionType::mix, dimension),
+            lwe::KeyAttr::get(ctx, 0),
+            lwe::ModulusChainAttr::get(ctx, moduliChain, level));
+        return ciphertext;
+      }
+      assert(tensorTy.getRank() == 2 && "expected rank 1 or 2 tensor");
+      underlyingTy = RankedTensorType::get(tensorTy.getShape().drop_front(),
+                                           tensorTy.getElementType());
+      auto ciphertext = lwe::LWECiphertextType::get(
+          ctx,
+          lwe::ApplicationDataAttr::get(ctx, underlyingTy,
+                                        lwe::NoOverflowAttr::get(ctx)),
+          lwe::PlaintextSpaceAttr::get(
+              ctx, plaintextRing,
+              lwe::InverseCanonicalEncodingAttr::get(ctx, scale)),
+          lwe::CiphertextSpaceAttr::get(ctx,
+                                        getRlweRNSRingWithLevel(ring_, level),
+                                        lwe::LweEncryptionType::mix, dimension),
+          lwe::KeyAttr::get(ctx, 0),
+          lwe::ModulusChainAttr::get(ctx, moduliChain, level));
+      return RankedTensorType::get(tensorTy.getShape().drop_back(), ciphertext);
     }
     // If the input IR does not contain aligned ciphertexts, we will not
     // pack tensors into ciphertext SIMD slots, so tensors are converted
     // into tensors of RLWE ciphertexts.
-    assert(dyn_cast<RankedTensorType>(valueTy) &&
-           "expected ranked tensor type");
-    auto scalarType = cast<RankedTensorType>(valueTy).getElementType();
     ciphertext = lwe::LWECiphertextType::get(
         ctx,
-        lwe::ApplicationDataAttr::get(ctx, scalarType,
+        lwe::ApplicationDataAttr::get(ctx, getElementTypeOrSelf(valueTy),
                                       lwe::NoOverflowAttr::get(ctx)),
         ciphertext.getPlaintextSpace(), ciphertext.getCiphertextSpace(),
         ciphertext.getKey(), ciphertext.getModulusChain());
-    return RankedTensorType::get(cast<RankedTensorType>(valueTy).getShape(),
-                                 ciphertext);
+    return RankedTensorType::get(tensorTy.getShape(), ciphertext);
   }
 
  private:
@@ -265,6 +301,83 @@ class SecretGenericTensorInsertConversion
   }
 };
 
+class SecretGenericTensorExpandConversion
+    : public SecretGenericOpConversion<tensor::ExpandShapeOp,
+                                       tensor::ExpandShapeOp> {
+ public:
+  using SecretGenericOpConversion<
+      tensor::ExpandShapeOp, tensor::ExpandShapeOp>::SecretGenericOpConversion;
+
+  FailureOr<Operation*> matchAndRewriteInner(
+      secret::GenericOp op, TypeRange outputTypes, ValueRange inputs,
+      ArrayRef<NamedAttribute> attributes,
+      ContextAwareConversionPatternRewriter& rewriter) const override {
+    // We expect this operation to occur when dropping unit dimensions in order
+    // to allow rotation ops to operate on 1-D tensors.
+    SliceVerificationResult res = isRankReducedType(
+        cast<ShapedType>(
+            cast<secret::SecretType>(op.getResultTypes()[0]).getValueType()),
+        cast<ShapedType>(
+            cast<secret::SecretType>(op.getOperandTypes()[0]).getValueType()));
+    if (res != SliceVerificationResult::Success) {
+      return rewriter.notifyMatchFailure(
+          op, "expected input type to be a rank reduced type of the output");
+    }
+    if (!isa<lwe::LWECiphertextType>(inputs[0].getType())) {
+      return rewriter.notifyMatchFailure(
+          op, "expected input that was expanded to be of type RLWE ciphertext");
+    }
+
+    if (!isa<RankedTensorType>(outputTypes[0])) {
+      return rewriter.notifyMatchFailure(
+          op, "expected expanded output to be a ranked tensor");
+    }
+    return rewriter
+        .replaceOpWithNewOp<tensor::FromElementsOp>(op, outputTypes, inputs)
+        .getOperation();
+  }
+};
+
+class SecretGenericTensorCollapseConversion
+    : public SecretGenericOpConversion<tensor::CollapseShapeOp,
+                                       tensor::CollapseShapeOp> {
+ public:
+  using SecretGenericOpConversion<
+      tensor::CollapseShapeOp,
+      tensor::CollapseShapeOp>::SecretGenericOpConversion;
+
+  FailureOr<Operation*> matchAndRewriteInner(
+      secret::GenericOp op, TypeRange outputTypes, ValueRange inputs,
+      ArrayRef<NamedAttribute> attributes,
+      ContextAwareConversionPatternRewriter& rewriter) const override {
+    // We expect this operation to occur when dropping unit dimensions in order
+    // to allow rotation ops to operate on 1-D tensors.
+    SliceVerificationResult res = isRankReducedType(
+        cast<ShapedType>(
+            cast<secret::SecretType>(op.getOperandTypes()[0]).getValueType()),
+        cast<ShapedType>(
+            cast<secret::SecretType>(op.getResultTypes()[0]).getValueType()));
+    if (res != SliceVerificationResult::Success) {
+      return rewriter.notifyMatchFailure(
+          op, "expected input type to be a rank reduced type of the output");
+    }
+    if (!isa<RankedTensorType>(inputs[0].getType())) {
+      return rewriter.notifyMatchFailure(
+          op, "expected input that was collapsed to be a ranked tensor");
+    }
+    if (!isa<lwe::LWECiphertextType>(outputTypes[0])) {
+      return rewriter.notifyMatchFailure(
+          op, "expected collapsed output to be of type RLWE ciphertext");
+    }
+
+    Value idx = rewriter.create<arith::ConstantIndexOp>(op.getLoc(), 0);
+    return rewriter
+        .replaceOpWithNewOp<tensor::ExtractOp>(op, outputTypes[0], inputs[0],
+                                               idx)
+        .getOperation();
+  }
+};
+
 bool hasSecretOperandsOrResults(Operation* op) {
   return llvm::any_of(op->getOperands(),
                       [](Value operand) {
@@ -354,7 +467,8 @@ struct SecretToCKKS : public impl::SecretToCKKSBase<SecretToCKKS> {
     target.addDynamicallyLegalOp<func::CallOp>(
         [&](Operation* op) { return typeConverter.isLegal(op); });
     target.addDynamicallyLegalOp<tensor::ExtractOp, tensor::ExtractSliceOp,
-                                 tensor::InsertOp>(
+                                 tensor::InsertOp, tensor::ExpandShapeOp,
+                                 tensor::CollapseShapeOp>(
         [&](Operation* op) { return typeConverter.isLegal(op); });
 
     target.markUnknownOpDynamicallyLegal(
@@ -394,6 +508,8 @@ struct SecretToCKKS : public impl::SecretToCKKSBase<SecretToCKKS> {
         SecretGenericOpLevelReduceConversion<ckks::LevelReduceOp>,
         SecretGenericTensorExtractConversion,
         SecretGenericTensorInsertConversion,
+        SecretGenericTensorCollapseConversion,
+        SecretGenericTensorExpandConversion,
         ConvertAnyContextAware<affine::AffineForOp>,
         ConvertAnyContextAware<affine::AffineYieldOp>,
         ConvertAnyContextAware<tensor::ExtractSliceOp>,

lib/Pipelines/ArithmeticPipelineRegistration.cpp

@@ -22,6 +22,7 @@
 #include "lib/Dialect/TensorExt/Conversions/TensorExtToTensor/TensorExtToTensor.h"
 #include "lib/Dialect/TensorExt/Transforms/CollapseInsertionChains.h"
 #include "lib/Dialect/TensorExt/Transforms/FoldConvertLayoutIntoAssignLayout.h"
+#include "lib/Dialect/TensorExt/Transforms/ImplementRotateAndReduce.h"
 #include "lib/Dialect/TensorExt/Transforms/InsertRotate.h"
 #include "lib/Dialect/TensorExt/Transforms/RotateAndReduce.h"
 #include "lib/Pipelines/PipelineRegistration.h"
@@ -152,6 +153,7 @@ void mlirToSecretArithmeticPipelineBuilder(
   convertToCiphertextSemanticsOptions.ciphertextSize = options.ciphertextDegree;
   pm.addPass(
       createConvertToCiphertextSemantics(convertToCiphertextSemanticsOptions));
+  pm.addPass(tensor_ext::createImplementRotateAndReduce());
 
   mathToPolynomialApproximationBuilder(pm);
 

lib/Pipelines/BUILD

@@ -119,6 +119,7 @@ cc_library(
         "@heir//lib/Dialect/TensorExt/Conversions/TensorExtToTensor",
         "@heir//lib/Dialect/TensorExt/Transforms:CollapseInsertionChains",
         "@heir//lib/Dialect/TensorExt/Transforms:FoldConvertLayoutIntoAssignLayout",
+        "@heir//lib/Dialect/TensorExt/Transforms:ImplementRotateAndReduce",
         "@heir//lib/Dialect/TensorExt/Transforms:InsertRotate",
         "@heir//lib/Dialect/TensorExt/Transforms:RotateAndReduce",
         "@heir//lib/Transforms/AddClientInterface",

lib/Transforms/ConvertToCiphertextSemantics/BUILD

@@ -16,6 +16,7 @@ cc_library(
         ":pass_inc_gen",
         "@heir//lib/Dialect/Secret/IR:SecretPatterns",
         "@heir//lib/Dialect/TensorExt/IR:Dialect",
+        "@heir//lib/Transforms/DropUnitDims",
         "@heir//lib/Utils",
         "@heir//lib/Utils:AffineMapUtils",
         "@heir//lib/Utils:AttributeUtils",

lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp

@@ -16,6 +16,7 @@
 #include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
 #include "lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.h"
 #include "lib/Transforms/ConvertToCiphertextSemantics/TypeConversion.h"
+#include "lib/Transforms/DropUnitDims/DropUnitDims.h"
 #include "lib/Utils/AffineMapUtils.h"
 #include "lib/Utils/AttributeUtils.h"
 #include "lib/Utils/ContextAwareConversionUtils.h"
@@ -27,6 +28,7 @@
 #include "lib/Utils/Utils.h"
 #include "llvm/include/llvm/ADT/ArrayRef.h"         // from @llvm-project
 #include "llvm/include/llvm/ADT/STLExtras.h"        // from @llvm-project
+#include "llvm/include/llvm/ADT/SmallVector.h"      // from @llvm-project
 #include "llvm/include/llvm/ADT/StringExtras.h"     // from @llvm-project
 #include "llvm/include/llvm/Support/Debug.h"        // from @llvm-project
 #include "llvm/include/llvm/Support/raw_ostream.h"  // from @llvm-project
@@ -47,13 +49,13 @@
 #include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
 #include "mlir/include/mlir/IR/ImplicitLocOpBuilder.h"   // from @llvm-project
+#include "mlir/include/mlir/IR/OpDefinition.h"           // from @llvm-project
 #include "mlir/include/mlir/IR/OperationSupport.h"       // from @llvm-project
 #include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
 #include "mlir/include/mlir/Support/LLVM.h"              // from @llvm-project
 #include "mlir/include/mlir/Support/LogicalResult.h"     // from @llvm-project
 #include "mlir/include/mlir/Transforms/DialectConversion.h"  // from @llvm-project
 #include "mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
-#include "mlir/include/mlir/Transforms/WalkPatternRewriteDriver.h"  // from @llvm-project
 
 #define DEBUG_TYPE "convert-to-ciphertext-semantics"
 
@@ -1357,7 +1359,6 @@ class ConvertExpandShape
     if (!sourceLayout) {
       return op.emitError() << "failed to fetch new layout attribute for input";
     }
-    op.dump();
 
     if (resultType != srcType) {
       return rewriter.notifyMatchFailure(
@@ -1385,6 +1386,99 @@ class ConvertExpandShape
   }
 };
 
+struct DropRotateUnitDims : OpRewritePattern<tensor_ext::RotateOp> {
+  using OpRewritePattern<tensor_ext::RotateOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor_ext::RotateOp rotateOp,
+                                PatternRewriter& rewriter) const override {
+    SmallVector<int64_t> operandUnitDims =
+        getUnitDims(rotateOp.getTensor().getType());
+    if (operandUnitDims.empty()) {
+      LLVM_DEBUG(llvm::dbgs() << "no unit dims to drop");
+      return failure();
+    }
+
+    SmallVector<Value> collapsedOperands =
+        collapseOperands(rewriter, {rotateOp.getTensor()}, operandUnitDims);
+
+    tensor_ext::RotateOp collapsedOp = tensor_ext::RotateOp::create(
+        rewriter, rotateOp.getLoc(), collapsedOperands[0], rotateOp.getShift());
+    rewriter.replaceOp(rotateOp, expandResult(rewriter, collapsedOp.getResult(),
+                                              rotateOp.getOutput().getType(),
+                                              operandUnitDims));
+    return success();
+  }
+};
+
+struct DropRotateAndReduceUnitDims
+    : OpRewritePattern<tensor_ext::RotateAndReduceOp> {
+  using OpRewritePattern<tensor_ext::RotateAndReduceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor_ext::RotateAndReduceOp rotateOp,
+                                PatternRewriter& rewriter) const override {
+    SmallVector<int64_t> operandUnitDims =
+        getUnitDims(rotateOp.getTensor().getType());
+    if (operandUnitDims.empty()) {
+      LLVM_DEBUG(llvm::dbgs() << "no unit dims to drop");
+      return failure();
+    }
+
+    SmallVector<Value> collapsedOperands =
+        collapseOperands(rewriter, {rotateOp.getTensor()}, operandUnitDims);
+
+    auto collapsedOp = tensor_ext::RotateAndReduceOp::create(
+        rewriter, rotateOp.getLoc(), collapsedOperands[0],
+        rotateOp.getPlaintexts(), rotateOp.getPeriod(), rotateOp.getSteps());
+    rewriter.replaceOp(rotateOp, expandResult(rewriter, collapsedOp.getResult(),
+                                              rotateOp.getOutput().getType(),
+                                              operandUnitDims));
+    return success();
+  }
+};
+
+struct DropElementwiseUnitDims : OpTraitRewritePattern<OpTrait::Elementwise> {
+  explicit DropElementwiseUnitDims(MLIRContext* context)
+      : OpTraitRewritePattern(context) {}
+
+  LogicalResult matchAndRewrite(mlir::Operation* op,
+                                PatternRewriter& rewriter) const override {
+    // Ensure that all operands and results have the same type.
+    SmallVector<Type> operandAndResultTypes =
+        llvm::to_vector(op->getOperandTypes());
+    operandAndResultTypes.append(op->getResultTypes().begin(),
+                                 op->getResultTypes().end());
+    if (!llvm::all_equal(operandAndResultTypes) || op->getNumOperands() == 0 ||
+        op->getNumResults() != 1) {
+      return failure();
+    }
+
+    auto tensorType = dyn_cast<RankedTensorType>(op->getOperand(0).getType());
+    if (!tensorType) {
+      return failure();
+    }
+
+    SmallVector<int64_t> operandUnitDims = getUnitDims(tensorType);
+    if (operandUnitDims.empty()) {
+      LLVM_DEBUG(llvm::dbgs() << "no unit dims to drop");
+      return failure();
+    }
+
+    SmallVector<Value> collapsedOperands = collapseOperands(
+        rewriter, llvm::to_vector(op->getOperands()), operandUnitDims);
+
+    Type resultType = collapsedOperands[0].getType();
+    Operation* collapsedOp = rewriter.create(OperationState(
+        op->getLoc(), op->getName().getStringRef(), collapsedOperands,
+        resultType, op->getAttrs(), op->getSuccessors()));
+
+    rewriter.replaceOp(
+        op, expandResult(rewriter, collapsedOp->getResults()[0],
+                         cast<RankedTensorType>(op->getResult(0).getType()),
+                         operandUnitDims));
+    return success();
+  }
+};
+
 struct ConvertToCiphertextSemantics
     : impl::ConvertToCiphertextSemanticsBase<ConvertToCiphertextSemantics> {
   using ConvertToCiphertextSemanticsBase::ConvertToCiphertextSemanticsBase;
@@ -1425,6 +1519,10 @@ struct ConvertToCiphertextSemantics
     // Note ConvertAssignLayout generates tensor.concat
     RewritePatternSet cleanupPatterns2(context);
     tensor::populateDecomposeTensorConcatPatterns(cleanupPatterns2);
+    // Drop unit dimensions for tensor_ext ops that require 1-D tensors (i.e.
+    // rotation ops) and elementwise ops.
+    cleanupPatterns2.add<DropRotateUnitDims, DropRotateAndReduceUnitDims,
+                         DropElementwiseUnitDims>(context);
     // Folding here will remove any unrealized conversion cast ops that were
     // inserted to persist new layouts.
     if (failed(applyPatternsGreedily(module, std::move(cleanupPatterns2)))) {