RNS

‘rns’ Dialect

The rns dialect represents types and ops related to residue number system (RNS) representations of ring-like types, such as integers or polynomials decomposed from high-bit width to lower-bit-width prime moduli. Sometimes RNS is referred to as CRT, for “Chinese Remainder Theorem.”

This dialect is intended to be as generic as possible in terms of its interaction with standard MLIR. However, because of upstream MLIR constraints, we do not have the ability to override, say, arith.addi to operate on an rns type. So such situations require dedicated ops, canonicalization patterns, etc.

RNS attributes

RNSAttr

A typed RNS value

A typed RNS value with one integer per basis limb.

The type parameter is expected to be an RNS coefficient type, and values stores one residue for each limb.

Example:

#v1 = #mod_arith.value<17 : !mod_arith.int<256 : i32>>
#v2 = #mod_arith.value<19 : !mod_arith.int<256 : i32>>
#v = #rns.value<[#v1, #v2]>

Parameters:

RNS types

RNSType

A residue number system representation

Syntax:

!rns.rns<
  ::llvm::ArrayRef<mlir::Type>   # basisTypes
>

Parameters:

RNS ops

rns.convert_basis (heir::rns::ConvertBasisOp)

Change an RNS value’s basis

Syntax:

operation ::= `rns.convert_basis` $value attr-dict `:` type($value) `->` type($output)

Given x \in Z_Q represented as {x_i mod q_i} where Q=\prod q_i, convertBasis returns the target-basis representation obtained by:

1. Compute x’s centered integer representative in [-Q/2, Q/2]
2. Reduce the integer representation mod p_i for each p_i in the target basis.

This performs the conversion directly from the RNS representation of x, without first reconstructing x over the integers. Note that both bases must be composed entirely of odd moduli; this algorithm does not work for bases with even moduli.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

rns.extract_residue (heir::rns::ExtractResidueOp)

Extract one RNS residue as the type of the limb

Syntax:

operation ::= `rns.extract_residue` $input attr-dict `:` type($input) `->` type($output)

Given an RNS-typed value, extract the basis limb at index and return it as the corresponding type.

For shaped inputs, this operation is applied elementwise, preserving the input container shape while replacing each RNS element with the selected limb.

Traits: AlwaysSpeculatableImplTrait, Elementwise, Scalarizable, Tensorizable, Vectorizable

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

rns.extract_slice (heir::rns::ExtractSliceOp)

Extracts a slice of RNS limbs

Syntax:

operation ::= `rns.extract_slice` $input attr-dict `:` type($input) `->` type($output)

Given an RNS-typed value with $k$ basis types (limbs), extract the slice of RNS components starting at start and having size size.

The result type is an RNS type containing the subset of basis types corresponding to the extracted slice. This is useful for operations like truncating or partitioning a modulus chain.

Traits: AlwaysSpeculatableImplTrait, Elementwise, Scalarizable, Tensorizable, Vectorizable

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

rns.pack (heir::rns::PackOp)

Pack multiple residues into one RNS value

Syntax:

operation ::= `rns.pack` operands attr-dict `:` type($input)

rns.pack converts a list of mod_arith values into a single value with RNS type.

The output RNS basis is formed by the input operand types in operand order, so operand i becomes limb i of the packed result.

Examples:

!zp17 = !mod_arith.int<17 : i32>
!zp257 = !mod_arith.int<257 : i32>
!rns = !rns.rns<!zp17, !zp257>

%packed = rns.pack %x, %y : !zp17, !zp257
return %packed : !rns

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

RNS additional definitions

TypeInterface definitions

RNSBasisTypeInterface (RNSBasisTypeInterface)

This interface is required for a type to be used as a parameter to an rns type.

Methods:

isCompatibleWith

bool isCompatibleWith(::mlir::Type otherRnsBasisType);

Returns true if this type is compatible with another type in the same RNS basis. In particular, the set of types used for a single RNS basis are never equal as types, but instead have some common attribute that must be checked here. For example, an RNS type where the basis types are polynomials would return true if the two types are both polynomial types, even if they have different coefficient moduli.

Another example is using mod arith types as the basis types, where by the nature of chinese reminder theorem, it is required that the modulus of them must be mutually coprime.

isCompatibleWith must be commutative, in the sense that type1.isCompatibleWith(type2) if and only if type2.isCompatibleWith(type1).

NOTE: This method must be implemented by the user.