JaxiteWord

‘jaxiteword’ Dialect

The jaxiteword dialect is an exit dialect for generating py code against the jaxiteword library API, using the jaxiteword parameters and encoding scheme.

See https://github.com/google/jaxite/jaxite_word

JaxiteWord types

CryptoContextType

The CryptoContext required to perform homomorphic operations in JaxiteWord.

Syntax: !jaxiteword.crypto_context

EvalKeyType

The evaluation key required to keyswitch/relinearize/rotate/automorphism operation in JaxiteWord.

Syntax: !jaxiteword.eval_key

PrivateKeyType

The private key required to decrypt a ciphertext in JaxiteWord.

Syntax: !jaxiteword.private_key

PublicKeyType

The public key required to encrypt plaintext in JaxiteWord.

Syntax: !jaxiteword.public_key

JaxiteWord ops

jaxiteword.add (heir::jaxiteword::AddOp)

Homomorphic addition of two ciphertexts

Syntax:

operation ::= `jaxiteword.add` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait, Commutative

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
lhsA ciphertext type
rhsA ciphertext type

Results:

ResultDescription
outputA ciphertext type

jaxiteword.add_inplace (heir::jaxiteword::AddInPlaceOp)

In-place homomorphic addition

Syntax:

operation ::= `jaxiteword.add_inplace` operands attr-dict `:` functional-type(operands, results)

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
lhsA ciphertext type
rhsA ciphertext type

jaxiteword.add_plain (heir::jaxiteword::AddPlainOp)

Add plaintext to ciphertext

Syntax:

operation ::= `jaxiteword.add_plain` operands attr-dict `:` functional-type(operands, results)

Traits: AllCiphertextTypesMatch, AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
lhsA ciphertext type or A plaintext type
rhsA ciphertext type or A plaintext type

Results:

ResultDescription
outputA ciphertext type

jaxiteword.decode (heir::jaxiteword::DecodeOp)

Decode a plaintext back to real values

Syntax:

operation ::= `jaxiteword.decode` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
plaintextA plaintext type

Results:

ResultDescription
outputshaped of any type values or floating-point

jaxiteword.decrypt (heir::jaxiteword::DecryptOp)

Decrypt a ciphertext

Syntax:

operation ::= `jaxiteword.decrypt` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
ciphertextA ciphertext type
secretKeyThe private key required to decrypt a ciphertext in JaxiteWord.

Results:

ResultDescription
plaintextA plaintext type

jaxiteword.encode (heir::jaxiteword::EncodeOp)

Encode a real tensor into a plaintext

Syntax:

operation ::= `jaxiteword.encode` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
inputshaped of any type values or floating-point

Results:

ResultDescription
plaintextA plaintext type

jaxiteword.encrypt (heir::jaxiteword::EncryptOp)

Encrypt a plaintext

Syntax:

operation ::= `jaxiteword.encrypt` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
plaintextA plaintext type
publicKeyThe public key required to encrypt plaintext in JaxiteWord.

Results:

ResultDescription
ciphertextA ciphertext type

jaxiteword.gen_keypair (heir::jaxiteword::GenKeyPairOp)

Generate public/private key pair

Syntax:

operation ::= `jaxiteword.gen_keypair` operands attr-dict `:` functional-type(operands, results)

Interfaces: InferTypeOpInterface

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.

Results:

ResultDescription
publicKeyThe public key required to encrypt plaintext in JaxiteWord.
secretKeyThe private key required to decrypt a ciphertext in JaxiteWord.

jaxiteword.gen_mulkey (heir::jaxiteword::GenMulKeyOp)

Generate evaluation key for relinearization

Syntax:

operation ::= `jaxiteword.gen_mulkey` operands attr-dict `:` functional-type(operands, results)

Interfaces: InferTypeOpInterface

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
secretKeyThe private key required to decrypt a ciphertext in JaxiteWord.

Results:

ResultDescription
evalKeyThe evaluation key required to keyswitch/relinearize/rotate/automorphism operation in JaxiteWord.

jaxiteword.gen_params (heir::jaxiteword::GenParamsOp)

Generate crypto context for CROSS

Syntax:

operation ::= `jaxiteword.gen_params` operands attr-dict `:` functional-type(operands, results)

Creates scheme parameters for the CROSS CKKS implementation.

Scheme parameters (from CKKSContext):

  • degree: Ring polynomial degree (N)
  • numSlots: Number of slots (usually N/2)
  • scalingFactor: CKKS scaling factor (Δ)
  • qTowers: Modulus chain for ciphertext (Q)
  • pTowers: Modulus chain for key switching (P)

TPU-specific parameters (from HEMul/HERot):

  • batch: Batch size for TPU
  • r, c: Tiling dimensions (r * c = degree)
  • dnum: Key switching decomposition parameter
  • numEvalMult: Number of evaluation multiplications

Interfaces: InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
degree::mlir::IntegerAttr64-bit signless integer attribute
numSlots::mlir::IntegerAttr64-bit signless integer attribute
scalingFactor::mlir::FloatAttr64-bit float attribute
qTowers::mlir::DenseI64ArrayAttri64 dense array attribute
pTowers::mlir::DenseI64ArrayAttri64 dense array attribute
batch::mlir::IntegerAttr32-bit signless integer attribute
r::mlir::IntegerAttr32-bit signless integer attribute
c::mlir::IntegerAttr32-bit signless integer attribute
dnum::mlir::IntegerAttr32-bit signless integer attribute
numEvalMult::mlir::IntegerAttr32-bit signless integer attribute
compositeDegree::mlir::IntegerAttr32-bit signless integer attribute

Results:

ResultDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.

jaxiteword.gen_rot_key (heir::jaxiteword::GenRotKeyOp)

Generate rotation keys

Syntax:

operation ::= `jaxiteword.gen_rot_key` operands attr-dict `:` functional-type(operands, results)

Interfaces: InferTypeOpInterface

Attributes:

AttributeMLIR TypeDescription
indices::mlir::DenseI64ArrayAttri64 dense array attribute

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
secretKeyThe private key required to decrypt a ciphertext in JaxiteWord.

Results:

ResultDescription
rotKeysThe evaluation key required to keyswitch/relinearize/rotate/automorphism operation in JaxiteWord.

jaxiteword.mod_reduce (heir::jaxiteword::ModReduceOp)

Rescale/modulus reduce a ciphertext (CKKS)

Syntax:

operation ::= `jaxiteword.mod_reduce` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
ciphertextA ciphertext type

Results:

ResultDescription
outputA ciphertext type

jaxiteword.mul (heir::jaxiteword::MulOp)

Homomorphic multiplication of two ciphertexts with relinearization

Syntax:

operation ::= `jaxiteword.mul` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
lhsA ciphertext type
rhsA ciphertext type
evalKeyThe evaluation key required to keyswitch/relinearize/rotate/automorphism operation in JaxiteWord.

Results:

ResultDescription
outputA ciphertext type

jaxiteword.mul_no_relin (heir::jaxiteword::MulNoRelinOp)

Homomorphic multiplication without relinearization

Syntax:

operation ::= `jaxiteword.mul_no_relin` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
lhsA ciphertext type
rhsA ciphertext type

Results:

ResultDescription
outputA ciphertext type

jaxiteword.mul_plain (heir::jaxiteword::MulPlainOp)

Multiply ciphertext with plaintext

Syntax:

operation ::= `jaxiteword.mul_plain` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
ciphertextA ciphertext type
plaintextA plaintext type

Results:

ResultDescription
outputA ciphertext type

jaxiteword.negate (heir::jaxiteword::NegateOp)

Negate a ciphertext

Syntax:

operation ::= `jaxiteword.negate` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
ciphertextA ciphertext type

Results:

ResultDescription
outputA ciphertext type

jaxiteword.program_initialization (heir::jaxiteword::ProgramInitializationOp)

Initialize CROSS program context with HEMul/HERot accessors

Syntax:

operation ::= `jaxiteword.program_initialization` operands attr-dict `:` functional-type(operands, results)

Calls ctx.program_initialization() which pre-builds level-indexed HEMul and HERot accessors and loads all keys onto the context. Replaces separate GenMulKeyOp + GenRotKeyOp for the new CROSS API.

Attributes:

AttributeMLIR TypeDescription
totalHemulLevels::mlir::IntegerAttr64-bit signless integer attribute
totalRotationIndices::mlir::DenseI64ArrayAttri64 dense array attribute
dnum::mlir::IntegerAttr32-bit signless integer attribute
r::mlir::IntegerAttr32-bit signless integer attribute
c::mlir::IntegerAttr32-bit signless integer attribute
batch::mlir::IntegerAttr32-bit signless integer attribute

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
secretKeyThe private key required to decrypt a ciphertext in JaxiteWord.

jaxiteword.relin (heir::jaxiteword::RelinOp)

Relinearize a ciphertext

Syntax:

operation ::= `jaxiteword.relin` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
ciphertextA ciphertext type
evalKeyThe evaluation key required to keyswitch/relinearize/rotate/automorphism operation in JaxiteWord.

Results:

ResultDescription
outputA ciphertext type

jaxiteword.rot (heir::jaxiteword::RotOp)

Rotate slots in a ciphertext

Syntax:

operation ::= `jaxiteword.rot` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

AttributeMLIR TypeDescription
index::mlir::IntegerAttr64-bit signless integer attribute

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
ciphertextA ciphertext type
evalKeyThe evaluation key required to keyswitch/relinearize/rotate/automorphism operation in JaxiteWord.

Results:

ResultDescription
outputA ciphertext type

jaxiteword.square (heir::jaxiteword::SquareOp)

Square a ciphertext

Syntax:

operation ::= `jaxiteword.square` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
ciphertextA ciphertext type

Results:

ResultDescription
outputA ciphertext type

jaxiteword.sub (heir::jaxiteword::SubOp)

Homomorphic subtraction of two ciphertexts

Syntax:

operation ::= `jaxiteword.sub` operands attr-dict `:` functional-type(operands, results)

Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultRings

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
lhsA ciphertext type
rhsA ciphertext type

Results:

ResultDescription
outputA ciphertext type

jaxiteword.sub_inplace (heir::jaxiteword::SubInPlaceOp)

In-place homomorphic subtraction

Syntax:

operation ::= `jaxiteword.sub_inplace` operands attr-dict `:` functional-type(operands, results)

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
lhsA ciphertext type
rhsA ciphertext type

jaxiteword.sub_plain (heir::jaxiteword::SubPlainOp)

Subtract plaintext from ciphertext

Syntax:

operation ::= `jaxiteword.sub_plain` operands attr-dict `:` functional-type(operands, results)

Traits: AllCiphertextTypesMatch, AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

OperandDescription
cryptoContextThe CryptoContext required to perform homomorphic operations in JaxiteWord.
lhsA ciphertext type or A plaintext type
rhsA ciphertext type or A plaintext type

Results:

ResultDescription
outputA ciphertext type