ModArith

ModArith attributes

ModArithAttr

a typed mod_arith attribute

Syntax:

#mod_arith.int<
  ::mlir::heir::mod_arith::ModArithType,   # type
  mlir::IntegerAttr   # value
>

Example:

#attr = 123:i32
#attr_verbose = #mod_arith.int<123:i32>

Parameters:

ModArith types

ModArithType

Integer type with modular arithmetic

Syntax:

!mod_arith.int<
  ::mlir::IntegerAttr   # modulus
>

mod_arith.int<p> represents an element of the ring of integers modulo $p$. The modulus attribute is the ring modulus, and mod_arith operations lower to arith operations that produce results in the range [0, modulus), often called the canonical representative.

modulus is specified with an integer type suffix, for example, mod_arith.int<65537 : i32>. This corresponds to the storage type for the modulus, and is i64 by default.

It is required that the underlying integer type should be larger than twice the modulus (have one extra bit of storage space) to avoid signedness issues. For example, when modulus == 2 ** 32 - 1, the underlying type for the modulus should be at least i33, though i64 is a natural choice.

Passes may allow intermediate values that do not always produce a canonical representative in [0, modulus). For example, if the machine storage type is i64, but the modulus fits within an i32, a lowering could allow intermediate arithmetic values to grow to as large as an i64 before reducing them. However, all passes must ensure that values used outside the local scope (e.g., function return values or arguments to calls to linked functions) are appropriately reduced to the canonical representative. modulus is the modulus the arithmetic working with.

Examples:

!Zp1 = !mod_arith.int<7> // implicitly being i64
!Zp2 = !mod_arith.int<65537 : i32>
!Zp3 = !mod_arith.int<536903681 : i64>

Parameters:

ModArith ops

mod_arith.add (heir::mod_arith::AddOp)

Modular addition operation

Syntax:

operation ::= `mod_arith.add` operands attr-dict `:` type($output)

Computes modular addition.

Unless otherwise specified, the operation assumes both inputs are canonical representatives and guarantees the output being canonical representative.

Traits: AlwaysSpeculatableImplTrait, Commutative, Elementwise, SameOperandsAndResultType, Scalarizable, Tensorizable, Vectorizable

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

mod_arith.barrett_reduce (heir::mod_arith::BarrettReduceOp)

Compute the first step of the Barrett reduction.

Syntax:

operation ::= `mod_arith.barrett_reduce` operands attr-dict `:` qualified(type($input))

Let $q$ denote a statically known modulus and $b = 4^{w}$, where $w$ is the smallest bit-width that contains the range $[0, q)$. The Barrett reduce operation computes barret_reduce x = x - floor(x * floor(b / q) / b) * q.

Given $0 <= x < q^2$, then this will compute $(x \mod q)$ or $(x \mod q) + q$.

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface

Attributes:

Operands:

Results:

mod_arith.constant (heir::mod_arith::ConstantOp)

Define a constant value via an attribute.

Example:

%0 = mod_arith.constant 123 : !mod_arith.int<65537:i32>

Traits: AlwaysSpeculatableImplTrait, InferTypeOpAdaptor

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Results:

mod_arith.encapsulate (heir::mod_arith::EncapsulateOp)

Encapsulate an integer into a mod_arith type

Syntax:

operation ::= `mod_arith.encapsulate` operands attr-dict `:` type($input) `->` type($output)

mod_arith.encapsulate converts the integer to be of mod_arith type.

Examples:

mod_arith.encapsulate %c0 : i32 -> mod_arith.int<65537 : i32>
mod_arith.encapsulate %c1 : i64 -> mod_arith.int<65537>

Traits: AlwaysSpeculatableImplTrait, Elementwise, Scalarizable, Tensorizable, Vectorizable

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

mod_arith.extract (heir::mod_arith::ExtractOp)

Extract the integer stored inside mod_arith type

Syntax:

operation ::= `mod_arith.extract` operands attr-dict `:` type($input) `->` type($output)

mod_arith.extract extracts the integer inside the mod_arith type.

It is required that the bitwidth of the output integer type is the same as that of the storage type of the input mod_arith type.

Examples:

%m0 = mod_arith.encapsulate %c0 : i32 -> mod_arith.int<65537 : i32>
%m1 = mod_arith.encapsulate %c1 : i64 -> mod_arith.int<65537>
%c2 = mod_arith.extract %m0 : mod_arith.int<65537 : i32> -> i32
%c3 = mod_arith.extract %m1 : mod_arith.int<65537> -> i64

Traits: AlwaysSpeculatableImplTrait, Elementwise, Scalarizable, Tensorizable, Vectorizable

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

mod_arith.mac (heir::mod_arith::MacOp)

Modular multiplication-and-accumulation operation

Syntax:

operation ::= `mod_arith.mac` operands attr-dict `:` type($output)

mod_arith.mac x, y, z computes $(x * y) + z$

Unless otherwise specified, the operation assumes all inputs are canonical representatives and guarantees the output being canonical representative.

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, Scalarizable, Tensorizable, Vectorizable

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

mod_arith.mod_switch (heir::mod_arith::ModSwitchOp)

Change the modulus of a mod_arith

Syntax:

operation ::= `mod_arith.mod_switch` $input attr-dict `:` type($input) `to` type($output)

“mod_switch” operation to change the modulus of a mod_arith type to a bigger space.

Examples:

`mod_arith.mod_switch %c0 : mod_arith.int<65537 : i32> to mod_arith.int<65539 : i32>`

Traits: AlwaysSpeculatableImplTrait, SameOperandsAndResultShape

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

mod_arith.mul (heir::mod_arith::MulOp)

Modular multiplication operation

Syntax:

operation ::= `mod_arith.mul` operands attr-dict `:` type($output)

Computes modular multiplication.

Unless otherwise specified, the operation assumes both inputs are canonical representatives and guarantees the output being canonical representative.

Traits: AlwaysSpeculatableImplTrait, Commutative, Elementwise, SameOperandsAndResultType, Scalarizable, Tensorizable, Vectorizable

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

mod_arith.reduce (heir::mod_arith::ReduceOp)

Reduce the mod arith type to its canonical representative

Syntax:

operation ::= `mod_arith.reduce` operands attr-dict `:` type($output)

mod_arith.reduce x produces $y$, the canonical representative in $[0, q)$ such that $x \equiv y \mod q$.

Examples:

%c0 = arith.constant 65538 : i32
%m0 = mod_arith.encapsulate %c0 : i32 -> mod_arith.int<65537 : i32>
// mod_arith.extract %m0 produces 65538
%m1 = mod_arith.reduce %m0 : mod_arith.int<65537: i32>
// mod_arith.extract %m1 produces 1

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, Scalarizable, Tensorizable, Vectorizable

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

mod_arith.subifge (heir::mod_arith::SubIfGEOp)

Compute (x >= y) ? x - y : x.

Syntax:

operation ::= `mod_arith.subifge` operands attr-dict `:` qualified(type($output))

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface

Operands:

Results:

mod_arith.sub (heir::mod_arith::SubOp)

Modular subtraction operation

Syntax:

operation ::= `mod_arith.sub` operands attr-dict `:` type($output)

Computes modular subtraction.

Unless otherwise specified, the operation assumes both inputs are canonical representatives and guarantees the output being canonical representative.

Traits: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType, Scalarizable, Tensorizable, Vectorizable

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results: