Skip to content

Commit

Permalink
review comments
Browse files Browse the repository at this point in the history
  • Loading branch information
inbelic committed Mar 29, 2024
1 parent 2e996c8 commit 5bd3bdc
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 16 deletions.
31 changes: 21 additions & 10 deletions include/Dialect/Polynomial/IR/PolynomialOps.td
Expand Up @@ -137,12 +137,18 @@ def Polynomial_ConstantOp : Polynomial_Op<"constant", [Pure]> {
def Polynomial_NTTOp : Polynomial_Op<"ntt", [Pure]> {
let summary = "Computes point-value tensor representation of a polynomial.";
let description = [{
`polynomial.ntt` creates a tensor value containing the point-value
representation of the input polynomial. The polynomial's RingAttr is
embedded as the encoding attribute.

The output tensor has shape equal to the degree of the ring's ideal
generator polynomial, including zeroes.
`polynomial.ntt` computes the forward integer Number Theoretic Transform
(NTT) on the input polynomial. It returns a tensor containing a point-value
representation of the input polynomial. The output tensor has shape equal to
the degree of the ring's ideal generation polynomial. The polynomial's
RingAttr is embedded as the encoding attribute of the output tensor.

Given an input polynomial $F(x)$ (over a ring with degree $n$) and a
primitive $n$-th root of unity $\omega_n$, the output is the list of $n$
evaluations

$f_k = F(\omega_n^k) ; k \in [0, n)$
The choice of primitive root is determined by subsequent lowerings.
}];

let arguments = (ins Polynomial:$input);
Expand All @@ -154,11 +160,16 @@ def Polynomial_NTTOp : Polynomial_Op<"ntt", [Pure]> {
}

def Polynomial_INTTOp : Polynomial_Op<"intt", [Pure]> {
let summary = "Computes the polynomial of a point-value representation tensor";
let summary = "Computes the reverse integer Number Theoretic Transform (NTT).";
let description = [{
`polynomial.intt` creates a polynomial value from the input tensor
containing the point-value representation. The ring of the polynomial is
taken from the required encoding attribute of the tensor.
`polynomial.intt` computes the reverse integer Number Theoretic Transform
(INTT) on the input tensor. This is the inverse operation of the
`polynomial.ntt` operation.

The input tensor is interpreted as a point-value representation of the
output polynomial at powers of a primitive $n$-th root of unity (see
`polynomial.ntt`). The ring of the polynomial is taken from the required
encoding attribute of the tensor.
}];

let arguments = (ins RankedTensorOf<[AnyInteger]>:$input);
Expand Down
11 changes: 5 additions & 6 deletions lib/Dialect/Polynomial/IR/PolynomialOps.cpp
Expand Up @@ -83,18 +83,17 @@ LogicalResult MonomialMulOp::verify() {
"must be of the form (x^n - 1) for some n";
}

template <typename Op>
static LogicalResult verifyNTTOp(Op *op, RingAttr ring,
static LogicalResult verifyNTTOp(Operation *op, RingAttr ring,
RankedTensorType tensorType) {
auto encoding = tensorType.getEncoding();
if (!encoding) {
return op->emitOpError()
<< "a ring encoding was not provided to the tensor output";
<< "a ring encoding was not provided to the tensor.";
}
auto encodedRing = dyn_cast<RingAttr>(encoding);
if (!encodedRing) {
return op->emitOpError()
<< "the provided tensor output encoding is not a ring attribute";
<< "the provided tensor encoding is not a ring attribute.";
}

if (encodedRing != ring) {
Expand All @@ -119,13 +118,13 @@ static LogicalResult verifyNTTOp(Op *op, RingAttr ring,
LogicalResult NTTOp::verify() {
auto ring = getInput().getType().getRing();
auto tensorType = getOutput().getType();
return verifyNTTOp<NTTOp>(this, ring, tensorType);
return verifyNTTOp(this->getOperation(), ring, tensorType);
}

LogicalResult INTTOp::verify() {
auto tensorType = getInput().getType();
auto ring = getOutput().getType().getRing();
return verifyNTTOp<INTTOp>(this, ring, tensorType);
return verifyNTTOp(this->getOperation(), ring, tensorType);
}

void SubOp::getCanonicalizationPatterns(RewritePatternSet &results,
Expand Down

0 comments on commit 5bd3bdc

Please sign in to comment.