Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Inplace binary operation class for the OpenFHE IR + sum and sub inplace operations #1339

Merged
merged 2 commits into from
Feb 7, 2025
+40 −0
Merged

Inplace binary operation class for the OpenFHE IR + sum and sub inplace operations #1339

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
654d5a7
Inplace binary operation class for the OpenFHE IR + sum and sub inpla…
ShirChoiEX Jan 29, 2025
e44d36b
conflict fix
ShirChoiEX Jan 31, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions lib/Dialect/Openfhe/IR/OpenfheOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,21 @@ class Openfhe_BinaryOp<string mnemonic, list<Trait> traits = []>
let results = (outs NewLWECiphertext:$output);
}

class Openfhe_BinaryInPlaceOp<string mnemonic, list<Trait> traits = []>
: Openfhe_Op<mnemonic, traits # [
AllTypesMatch<["lhs", "rhs"]>,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is a recent PR (#1334) that relaxes the AllTypesMatch requirement to [SameOperandsAndResultRings, InferTypeOpAdaptor]

Discussion for maintainers: We might also need to register a MemoryEffects<[MemWrite]> to make liveness analysis happy, though I found --cse --canonicalize does not eliminate these inplace ops.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is there a need to add [SameOperandsAndResultRings] trait for the inplace operations? It seems to be analyzing both inputs and the output, but there is no output for an inplace function. There seem to be no alternative trait like [SameOperandsRings], or does the [SameOperandsAndResultRings] ignore the result check if it is not created? Also, regarding the #1334 PR, he uses [SameOperandsAndResultRings] in AddOp and SubOp but not in the MulOp why is it so?

Copy link
Collaborator

@AlexanderViand-Intel AlexanderViand-Intel Jan 30, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, regarding the #1334 PR, he uses [SameOperandsAndResultRings] in AddOp and SubOp but not in the MulOp why is it so?

If the question is why no SameOperandsAndResultRings is needed for openfhe.mul it's because
AllTypesMatch<["lhs", "rhs", "output"]> implies SameOperandsAndResultRings, but if the question is instead why I didn't "relax" the constraints on openfhe.mul to "only" SameOperandsAndResultRings, that's just because the op isn't used in our default lowering pipeline so there was no need to.

EDIT: Note that the AllTypesMatch<["lhs", "rhs", "output"]> was already applied to openfhe.mul before my PR, it was just previously inherited from Openfhe_BinaryOp

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

or does the [SameOperandsAndResultRings] ignore the result check if it is not created?

From the following line we can know that if there is no result, the trait only checks the equivalence among operands.

heir/lib/Dialect/LWE/IR/LWETraits.h

Lines 19 to 50 in ad59e8f

static LogicalResult verifyTrait(Operation *op) {
::mlir::heir::polynomial::RingAttr rings = nullptr;
for (auto rTy : op->getResultTypes()) {
auto ct = dyn_cast<lwe::NewLWECiphertextType>(rTy);
if (!ct) continue;
if (rings == nullptr) {
rings = ct.getCiphertextSpace().getRing();
continue;
}
if (rings != ct.getCiphertextSpace().getRing()) {
return op->emitOpError()
<< "requires all operands and results to have the same rings";
}
}
for (auto oTy : op->getOperandTypes()) {
auto ct = dyn_cast<lwe::NewLWECiphertextType>(oTy);
if (!ct) continue; // Check only ciphertexts
if (rings == nullptr) {
rings = ct.getCiphertextSpace().getRing();
continue;
}
if (rings != ct.getCiphertextSpace().getRing()) {
return op->emitOpError()
<< "requires all operands and results to have the same rings";
}
}
return success();
}
};

]> {

let summary = "In-place binary operation for OpenFHE";

let arguments = (ins
Openfhe_CryptoContext:$cryptoContext,
NewLWECiphertext:$lhs,
NewLWECiphertext:$rhs
);
}


def GenParamsOp : Openfhe_Op<"gen_params"> {
let description = [{
Generates the parameters for the OpenFHE scheme.
Expand Down Expand Up @@ -165,6 +180,17 @@ def SubOp : Openfhe_BinaryOp<"sub",
let summary = "OpenFHE sub operation of two ciphertexts.";
}

// In-Place Addition
def AddInPlaceOp : Openfhe_BinaryInPlaceOp<"add_inplace"> {
let summary = "Performs in-place homomorphic addition, modifying lhs.";
}

// In-Place Subtraction
def SubInPlaceOp : Openfhe_BinaryInPlaceOp<"sub_inplace"> {
let summary = "Performs in-place homomorphic subtraction, modifying lhs.";
}


def AddPlainOp : Openfhe_Op<"add_plain",[
Pure,
AllTypesMatch<["ciphertext", "output"]>
Expand Down
14 changes: 14 additions & 0 deletions tests/Dialect/Openfhe/IR/ops.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,20 @@ module {
%out = openfhe.add %cc, %c1, %c2: (!cc, !ct, !ct) -> !ct
return
}
// CHECK-LABEL: func @test_inplace_add
func.func @test_inplace_add(%cc: !cc, %pt : !pt, %pk : !pk) {
%c1 = openfhe.encrypt %cc, %pt, %pk : (!cc, !pt, !pk) -> !ct
%c2 = openfhe.encrypt %cc, %pt, %pk : (!cc, !pt, !pk) -> !ct
openfhe.add_inplace %cc, %c1, %c2: (!cc, !ct, !ct) -> ()
return
}
// CHECK-LABEL: func @test_inplace_sub
func.func @test_inplace_sub(%cc: !cc, %pt : !pt, %pk : !pk) {
%c1 = openfhe.encrypt %cc, %pt, %pk : (!cc, !pt, !pk) -> !ct
%c2 = openfhe.encrypt %cc, %pt, %pk : (!cc, !pt, !pk) -> !ct
openfhe.sub_inplace %cc, %c1, %c2: (!cc, !ct, !ct) -> ()
return
}

// CHECK-LABEL: func @test_sub
func.func @test_sub(%cc : !cc, %pt : !pt, %pk: !pk) {
Expand Down
Loading