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

AlignmentAttr

An attribute describing padding and alignment of a tensor.

Syntax:

#tensor_ext.alignment<
  ::mlir::DenseI64ArrayAttr,   # in
  ::mlir::DenseI64ArrayAttr,   # out
  ::mlir::DenseI64ArrayAttr,   # insertedDims
  ::mlir::DenseI64ArrayAttr,   # padding
  TypedAttr   # paddingValue
>

This attribute is used to describe how a data-semantic tensor is padded and replicated to align its size before applying a ciphertext layout (see tensor_ext.layout).

It describes transformations to be applied to an input tensor.

The in attribute describes the shape of the original tensor. The following transformations are applied to the input tensor.

1. New unit dimensions are inserted to match the number of dimensions of the output tensor. The insertedDims attribute is an array specifying which axes in the output tensor are new unit dimensions. By default no new dimensions are inserted.
2. Padding is applied. The padding attribute is an array with the same size as the output tensor shape, which each entry denoting the number of values to pad at the end of that axis. The paddingValue attribute describes the value used for padding, and must be set if padding is set.
3. The padded result is replicated to fill the output tensor shape.

Example:

#repl = #tensor_ext.alignment<
  in = [7],
  padding = [1],
  paddingValue = 0:i32,
  out = [16],
>

This indicates that the original input tensor is a tensor<7xi32>, and it is being padded and replicated into a tensor<16xi32>. If the input tensor has values [1, 2, 3, 4, 5, 6, 7] then the packed tensor contains the data [1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0].

Example:

#repl = #tensor_ext.alignment<
  in = [3, 10],
  padding = [1, 6],
  paddingValue = 0:i32,
  out = [32, 32],
>

This indicates that the original input tensor is a tensor<3x10xi32>, and it is being padded and replicated into a tensor<32x32xi32>. One row of zeros and six columns of zeros are added to the input tensor, and then it is repeated once along the column dimension and four times along the row dimension to fill the 32x32 output tensor.

Example:

#repl = #tensor_ext.alignment<
  in = [10],
  insertedDims = [0],
  padding = [0, 6],
  paddingValue = 0:i32,
  out = [32, 32],
>

This indicates that the original input tensor is a tensor<10xi32>, and it is being padded and replicated into a tensor<32x32xi32>. First a unit dimension is inserted as the first axis, then zeros are padded along the second axis to make a tensor<1x16xi32>. Finally, it is replicated twice along the columns and 32 times along the rows to fill the 32x32 output tensor.

Parameters:

LayoutAttr

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

Syntax:

#tensor_ext.layout<
  ::mlir::AffineMap,   # map
  ::mlir::heir::tensor_ext::AlignmentAttr   # alignment
>

This attribute describes how a data-semantic tensor is laid out among a tensor of ciphertexts. The layout is described by an affine map mapping the data-semantic tensor indices to the ciphertext tensor indices, where the trailing affine map’s result expressions mark the slot index (or indices) of the ciphertext.

The layout may include optional padding and alignment of the data-semantic tensor before the layout is applied. This would be required if, for example, a tensor<4xi32> is to be laid out in a tensor<16xi32> ciphertext tensor. The data-semantic tensor must be extended somehow to fill the 16 slots. If an alignment attribute is not provided, then lowerings may raise errors if there is no unambiguous way to align the tensor.

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::heir::tensor_ext::LayoutAttr   # 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:

SIMDPackingAttr

An attribute describing the SIMD packing of a tensor.

Syntax:

#tensor_ext.simd_packing<
  ::mlir::DenseI64ArrayAttr,   # in
  ::mlir::DenseI64ArrayAttr,   # padding
  ::mlir::DenseI64ArrayAttr,   # out
  int64_t   # padding_value
>

This attribute is used as the encoding attribute on a tensor. It describes the transformations that were applied to an input tensor to pack it into the given tensor.

The in attribute describes the shape of the original tensor. The following transformations are applied to the input tensor.

1. Padding is applied first. The padding attribute is an array with the same size as the input tensor shape. Padding is applied at the end of the array using the padding_value attribute (default zero). The result after zero padding should be a power of two.

2. The padded result is replicated or split to fill the output tensor shape.

For example,

#packing = #tensor_ext.simd_packing<
  in = [7],
  padding = [1],
  padding_value = 0,
  out = [16],
>

may be used on a tensor type like

tensor<16xi32>

This indicates that the original input tensor was a tensor<7xi32>, and if it had values [1, 2, 3, 4, 5, 6, 7] then the packed tensor contains the data [1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0].

Parameters:

TensorExt ops

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

Assign a layout to a plaintext tensor.

Syntax:

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

This op allows the ingestion of a plaintext tensor 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 tensor 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.permute (heir::tensor_ext::PermuteOp)

Permute a tensor by a static permutation.

Syntax:

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

This op represents a permutation of a tensor.

This is lowered from a convert_layout op, and is implemented in terms of rotate operations.

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 operation’s current behavior allows rotating multi-dimensional tensors by rotating along the tensor’s only non-unit dimension. This assumes the tensor is packed along the non-unit dimension.

// 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: