In the paper they also describe a potential optimization of modulo arithmetic by removing the use of a modulo operator (Rem[S|U]Op) for multiplication/addition of polynomials in the NTT domain. They introduce an operation to represent the first Barrett reduction step, say arith_ext.barrett_step, and, arith_ext.subifge x y that denotes z = (x >= y) ? x - y : y.
For instance:

mul = arith.muli x y
res = arith.remui mul cmod

becomes

mul = arith.muli x y
barret = arith_ext.barrett_reduce mul cmod
res = arith_ext.subifge barret cmod

We are able to avoid the use of potential division used in the remainder operation and only have runtime multiplication and bitshift as the Barrett ratio is able to be statically computed. We would be able to directly use this optimization in the current NTT lowering.

Further, they describe the use of a data-flow analysis to be able to reduce the number of arith_ext.subifge when there are subsequent uses from muls/adds/subs. These optimizations require the assumption that the input polynomial coefficients are in the range [0, cmod), however this restriction is not currently required in the polynomial (can be negative). This could be addressed with an operation poly.normalise that will ensure all the polynomial coefficients are in the range [0, cmod).

So I would propose the following steps to implement the papers optimizations:

Create a new dialect arith_ext which has the operations barret_reduce and subifge.
Create a pass that converts arith.muli + arith.remui into arith.muli + artih_ext.barret_reduce + arith_ext.subifge when the operands are in the range [0, cmod).
Create a pass that converts arith.addi + arith.remui into arith.addi + arith_ext.subifge when the operands are in the range [0, cmod).
Create a pass that converts arith.subi + arith.remui into arith.subi + arith.addi cmod + arith.subifge when the operands are in the range [0, cmod).
Add a new poly.normalise operation to poly to provide a fixed range [0, cod) for use in the above passes and modelling the ranges in the data-flow analysis.
Implement the papers data-flow analysis of reducing the number of arith.subifge operations.

Looking for feedback in all aspects of the proposed solution, especially operation names.

Port optimizations from HEaaN.mlir paper #635

Description

Metadata

Metadata

Assignees

Labels

Type

Fields

Projects

Milestone

Relationships

Development

Issue actions