TensorExt

’tensor_ext’ Dialect

The tensor_ext dialect contains operations on plaintext tensors that correspond to the computation model of certain FHE schemes, but are unlikely to be upstreamed to MLIR due to their specificity to FHE.

TensorExt attributes

LayoutAttr

The description of the layout of a data-semantic tensor.

Syntax:

#tensor_ext.layout<
  mlir::StringAttr   # layout
>

This attribute describes how a data-semantic tensor is laid out among a tensor of ciphertexts. The layout is described by an integer relation $(d, s)$, where $d$ is a multiset of data-semantic tensor indices and $s$ is a multiset of slot indices (or coefficient indices). The slot indices are defined by two indices: the ciphertext index and the slot index in that order. The elements of the relation are defined by a set of quasi-affine constraints.

I.e., a point $((2, 3), (7, 0))$ in the relation says that the data entry at index $(2, 3)$ is placed in slot 0 of ciphertext 7. This could be defined as part of the relation by a constraint like row + col + 2 - ct + slot = 0.

The attribute stores a string representation of the integer relation, which follows the ISL syntax for isl_basic_map. For example:

#vec_layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : (i0 - slot) mod 1024 = 7 and i0 >= 0 and 0 >= i0 and slot >= 0 and 1023 >= slot and ct = 0 }">

#mat_layout = #tensor_ext.layout<"{ [row, col] -> [ct, slot] : (slot - row) mod 512 = 0 and (ct + slot - col) mod 512 = 0 and row >= 0 and col >= 0 and ct >= 0 and slot >= 0 and 1023 >= slot and 511 >= ct and 511 >= row and 511 >= col }">

// Example with local (existential) variables.
#layout = #tensor_ext.layout<"{[d0] -> [ct, slot] : exists d3, d4 : -ct + d3 = 0 and d0 - d4 * 1024 = 0 and -d0 + 31 >= 0 and d0 >= 0 and ct >= 0 and slot >= 0 and -slot + 1023 >= 0 and -d0 + d3 * 1024 + 1023 >= 0 and d0 - d3 * 1024 >= 0 and -d0 + d4 * 1024 + 1023 >= 0 and d0 - d4 * 1024 >= 0 }">

Parameters:

OriginalTypeAttr

The original type of a secret tensor whose layout has been converted to ciphertext semantics.

Syntax:

#tensor_ext.original_type<
  ::mlir::Type,   # originalType
  ::mlir::Attribute   # layout
>

This attribute is used to retain the original type of a secret tensor after its conversion to ciphertext semantics, i.e. after applying any padding or alignment to fill ciphertext data types. For example, if a !secret.secret<tensor<32xi8>> is laid out in a ciphertext with 1024 slots, the new type would be !secret.secret<tensor<1024xi8>> with attribute tensor_ext.original_type<!secret.secret<tensor<32xi8>>.

Parameters:

TensorExt ops

tensor_ext.assign_layout (heir::tensor_ext::AssignLayoutOp)

Assign a layout to a plaintext tensor or scalar.

Syntax:

operation ::= `tensor_ext.assign_layout` operands attr-dict `:` type($output)

This op allows the ingestion of a plaintext value into the layout system. For example, ops like linalg.reduce, require a tensor input to represent initial values. These will generally be created by an arith.constant or tensor.empty op, which does not have secret results. Lowerings will convert this to a packed plaintext, so that the subsequent ops can be lowered as ciphertext-plaintext ops.

This op is inserted by layout selection passes.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

tensor_ext.convert_layout (heir::tensor_ext::ConvertLayoutOp)

Convert from one layout to another.

Syntax:

operation ::= `tensor_ext.convert_layout` operands attr-dict `:` type($output)

This op represents the conversion of a value from one packed layout to another. This is implemented via a “shift network” of ciphertext rotations, plaintext masks (ciphertext-plaintext multiplications), and additions.

This op is inserted by layout selection passes.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

tensor_ext.remap (heir::tensor_ext::RemapOp)

Reorganize the entries of a tensor by a static mapping.

Syntax:

operation ::= `tensor_ext.remap` operands attr-dict `:` type($input)

This op represents a remapping of entries of a tensor.

This op is primarily inserted as a lowered form of convert_layout op, to represent a ciphertext slot-repacking operation. However, it can more generally express any partial mapping from source tensor indices to target indices within a tensor. Unmapped entries from the input tensor are unmodified.

This operation differs from a tensor.gather/scatter operation in that it can replicate values from the source tensor or omit values in the destination tensor. It differs from a “lane shuffle” in that the mapping need not be a permutation. It differs from a “swizzle” in that it cannot change the dimension of the result tensor.

In the slot form of a ciphertext, this op is lowered to a “shift network” of ciphertext-plaintext mask, rotate, and sum operations by the implement-shift-network pass.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

tensor_ext.rotate_and_reduce (heir::tensor_ext::RotateAndReduceOp)

Performs a reduction of a periodically rotated tensor.

Syntax:

operation ::= `tensor_ext.rotate_and_reduce` operands attr-dict `:` functional-type(operands, results)

This op reduces products of a plaintext with a periodically rotated tensor.

In almost full generality, the reduction performed is

[ \sum_{i \in [0, n]} p(P, Ti) \cdot rotate(v, Ti) ]

where $f$ is a function, $p(P, Ti)$ is a function of a plaintext $P$ and $rotate(v, Ti)$ is a rotation of the ciphertext $v$ with period $T$. The operation takes as input the ciphertext vector $v$, the period $T$, the number of reductions $n$, and a tensor of plaintext values [p(P, 0), p(P, T), ..., p(P, T*(n-1))].

This can be used to implement a matrix vector product that uses a Halevi-Shoup diagonalization of the plaintext matrix. In this case, the reduction is

[ \sum_{i \in [0, n]} P(i) \cdot rotate(v, i) ]

where $P(i)$ is the $i$th diagonal of the plaintext matrix and the period $T$ is $1$.

An accumulation of the ciphertext slots is also handled via this operation by omitting the plaintext $p(P, T*i)$ argument and using a period of 1 with n = |v| so that the reduction is simply a sum of all rotation of the ciphertext.

If reduceOp is set to an MLIR operation name (e.g., arith.mulf), then the reduction operation is modified to use that operation instead of a sum. The chosen op must be one of arith.muli, arith.mulf, arith.addi, or arith.addf.

Efficient lowerings of this operation can use the Baby-Step / Giant-Step approach from Faster Homomorphic Linear Transformations in HElib to reduce the number of ciphertext rotations.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

tensor_ext.rotate (heir::tensor_ext::RotateOp)

Rotate a tensor some number of indices left.

Syntax:

operation ::= `tensor_ext.rotate` operands attr-dict `:` qualified(type($tensor)) `,` type($shift)

This op represents a left-rotation of a tensor by given number of indices. Negative shift values are interpreted as right-rotations.

This corresponds to the rotate operation in arithmetic FHE schemes like BGV. This op currently only supports 1D rotations of the last axis of a tensor. A tensor<4x64xi32> is interpreted as 4 ciphertexts each with 64 slots, and a rotation on a value of this type rotates each ciphertext by the given amount.

// In the future, the op will be adjusted to support rotations of general // multi-dimensional tensors with a vector of rotation indices for each // dimension. The lowering will implement the correct operations to rotate // the tensor along the indices given its packing.

Examples:

%0 = ... : tensor<16xi32>
%c7 = arith.constant 7 : i32
%1 = tensor_ext.rotate %0, %c7 : tensor<16xi32>, i32

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

tensor_ext.unpack (heir::tensor_ext::UnpackOp)

Unpack data from a ciphertext-semantic tensor.

Syntax:

operation ::= `tensor_ext.unpack` operands attr-dict `:` functional-type(operands, results)

This op extracts the underlying cleartext data from a ciphertext-semantic tensor.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results: