/
PolynomialOps.td
183 lines (149 loc) · 7.29 KB
/
PolynomialOps.td
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
#ifndef INCLUDE_DIALECT_POLYNOMIAL_IR_POLYNOMIALOPS_TD_
#define INCLUDE_DIALECT_POLYNOMIAL_IR_POLYNOMIALOPS_TD_
include "include/Dialect/Polynomial/IR/PolynomialAttributes.td"
include "include/Dialect/Polynomial/IR/PolynomialDialect.td"
include "include/Dialect/Polynomial/IR/PolynomialTypes.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
class Polynomial_Op<string mnemonic, list<Trait> traits = []> :
Op<Polynomial_Dialect, mnemonic, traits> {
let assemblyFormat = [{
operands attr-dict `:` `(` type(operands) `)` `->` type(results)
}];
let cppNamespace = "::mlir::heir::polynomial";
}
class Polynomial_BinOp<string mnemonic, list<Trait> traits = []> :
Polynomial_Op<mnemonic, !listconcat(traits, [Pure, SameOperandsAndResultType, ElementwiseMappable])> {
let arguments = (ins PolynomialLike:$lhs, PolynomialLike:$rhs);
let results = (outs PolynomialLike:$output);
let assemblyFormat = "`(` operands `)` attr-dict `:` qualified(type($output))" ;
}
def Polynomial_AddOp : Polynomial_BinOp<"add", [Commutative]> {
let summary = "Addition operation between polynomials.";
}
def Polynomial_SubOp : Polynomial_BinOp<"sub"> {
let summary = "Subtraction operation between polynomials.";
let hasCanonicalizer = 1;
}
def Polynomial_MulOp : Polynomial_BinOp<"mul", [Commutative]> {
let summary = "Multiplication operation between polynomials.";
}
def Polynomial_MulScalarOp : Polynomial_Op<"mul_scalar", [
ElementwiseMappable, AllTypesMatch<["polynomial", "output"]>]> {
let summary = "Multiplication by a scalar of the field.";
let arguments = (ins
PolynomialLike:$polynomial,
AnyInteger:$scalar
);
let results = (outs
PolynomialLike:$output
);
let assemblyFormat = "operands attr-dict `:` qualified(type($polynomial)) `,` type($scalar)";
}
def Polynomial_LeadingTermOp: Polynomial_Op<"leading_term"> {
let summary = "Compute the leading term of the polynomial.";
let description = [{
The degree of a polynomial is the largest $k$ for which the coefficient
$a_k$ of $x^k$ is nonzero. The leading term is the term $a_k x^k$, which
this op represents as a pair of results.
}];
let arguments = (ins Polynomial:$input);
let results = (outs Index:$degree, AnyInteger:$coefficient);
let assemblyFormat = "operands attr-dict `:` qualified(type($input)) `->` `(` type($degree) `,` type($coefficient) `)`";
}
def Polynomial_MonomialOp: Polynomial_Op<"monomial"> {
let summary = "Create a polynomial that consists of a single monomial.";
let arguments = (ins AnyInteger:$coefficient, Index:$degree);
let results = (outs Polynomial:$output);
}
def Polynomial_MonomialMulOp: Polynomial_Op<"monomial_mul", [AllTypesMatch<["input", "output"]>]> {
let summary = "Multiply a polynomial by a monic monomial.";
let description = [{
In the ring of polynomials mod $x^n - 1$, `monomial_mul` can be interpreted
as a cyclic shift of the coefficients of the polynomial. For some rings,
this results in optimized lowerings that involve rotations and rescaling
of the coefficients of the input.
}];
let arguments = (ins Polynomial:$input, Index:$monomialDegree);
let results = (outs Polynomial:$output);
let hasVerifier = 1;
}
def Polynomial_FromTensorOp : Polynomial_Op<"from_tensor", [Pure]> {
let summary = "Creates a polynomial from integer coefficients stored in a tensor.";
let description = [{
`polynomial.from_tensor` creates a polynomial value from a tensor of coefficients.
The input tensor must list the coefficients in degree-increasing order.
The input one-dimensional tensor may have size at most the degree of the
ring's ideal generator polynomial, with smaller dimension implying that
all higher-degree terms have coefficient zero.
}];
let arguments = (ins RankedTensorOf<[AnyInteger]>:$input);
let results = (outs Polynomial:$output);
let assemblyFormat = "$input attr-dict `:` type($input) `->` qualified(type($output))";
let builders = [
// Builder that infers coefficient modulus from tensor bit width,
// and uses whatever input ring is provided by the caller.
OpBuilder<(ins "::mlir::Value":$input, "RingAttr":$ring)>
];
let hasVerifier = 1;
}
def Polynomial_ToTensorOp : Polynomial_Op<"to_tensor", [Pure]> {
let summary = "Creates a tensor containing the coefficients of a polynomial.";
let description = [{
`polynomial.to_tensor` creates a tensor value containing the coefficients of the
input polynomial. The output tensor contains the coefficients in
degree-increasing order.
Operations that act on the coefficients of a polynomial, such as extracting
a specific coefficient or extracting a range of coefficients, should be
implemented by composing `to_tensor` with the relevant `tensor` dialect
ops.
The output tensor has shape equal to the degree of the ring's ideal
generator polynomial, including zeroes.
}];
let arguments = (ins Polynomial:$input);
let results = (outs RankedTensorOf<[AnyInteger]>:$output);
let assemblyFormat = "$input attr-dict `:` qualified(type($input)) `->` type($output)";
let hasVerifier = 1;
}
def Polynomial_ConstantOp : Polynomial_Op<"constant", [Pure]> {
let summary = "Define a constant polynomial via an attribute.";
let arguments = (ins Polynomial_Attr:$input);
let results = (outs Polynomial:$output);
let assemblyFormat = "$input attr-dict `:` qualified(type($output))";
}
def Polynomial_NTTOp : Polynomial_Op<"ntt", [Pure]> {
let summary = "Computes point-value tensor representation of a polynomial.";
let description = [{
`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);
let results = (outs RankedTensorOf<[AnyInteger]>:$output);
let assemblyFormat = "$input attr-dict `:` qualified(type($input)) `->` type($output)";
let hasVerifier = 1;
}
def Polynomial_INTTOp : Polynomial_Op<"intt", [Pure]> {
let summary = "Computes the reverse integer Number Theoretic Transform (NTT).";
let description = [{
`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);
let results = (outs Polynomial:$output);
let assemblyFormat = "$input attr-dict `:` qualified(type($input)) `->` type($output)";
let hasVerifier = 1;
}
#endif // INCLUDE_DIALECT_POLYNOMIAL_IR_POLYNOMIALOPS_TD_