[Mlir-commits] [mlir] 145176d - polynomial: Add basic ops (#89525)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Apr 28 14:58:48 PDT 2024
Author: Jeremy Kun
Date: 2024-04-28T14:58:44-07:00
New Revision: 145176dc0c93566ce4aef721044d49ab8ba50f87
URL: https://github.com/llvm/llvm-project/commit/145176dc0c93566ce4aef721044d49ab8ba50f87
DIFF: https://github.com/llvm/llvm-project/commit/145176dc0c93566ce4aef721044d49ab8ba50f87.diff
LOG: polynomial: Add basic ops (#89525)
Adds a few basic polynomial ops.
- add, sub, mul
- mul_scalar
- leading_term, monomial_mul, monomial (useful for lowering `mul` to
standard MLIR)
- from_tensor, to_tensor, constant
---------
Co-authored-by: Jeremy Kun <j2kun at users.noreply.github.com>
Co-authored-by: Oleksandr "Alex" Zinenko <ftynse at gmail.com>
Added:
mlir/test/Dialect/Polynomial/ops.mlir
mlir/test/Dialect/Polynomial/ops_errors.mlir
Modified:
mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h
mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td
mlir/lib/Dialect/Polynomial/IR/PolynomialDialect.cpp
mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp
mlir/test/Dialect/Polynomial/types.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h b/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h
index 39b05b9d3ad14b8..3325a6fa3f9fcf0 100644
--- a/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h
+++ b/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h
@@ -51,9 +51,6 @@ class Monomial {
return (exponent.ult(other.exponent));
}
- // Prints polynomial to 'os'.
- void print(raw_ostream &os) const;
-
friend ::llvm::hash_code hash_value(const Monomial &arg);
public:
@@ -102,6 +99,8 @@ class Polynomial {
unsigned getDegree() const;
+ ArrayRef<Monomial> getTerms() const { return terms; }
+
friend ::llvm::hash_code hash_value(const Polynomial &arg);
private:
diff --git a/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td b/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td
index 5d8da8399b01b54..d3e3ac55677f861 100644
--- a/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td
+++ b/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td
@@ -35,18 +35,18 @@ def Polynomial_Dialect : Dialect {
```mlir
// A constant polynomial in a ring with i32 coefficients and no polynomial modulus
- #ring = #polynomial.ring<ctype=i32>
+ #ring = #polynomial.ring<coefficientType=i32>
%a = polynomial.constant <1 + x**2 - 3x**3> : polynomial.polynomial<#ring>
// A constant polynomial in a ring with i32 coefficients, modulo (x^1024 + 1)
#modulus = #polynomial.polynomial<1 + x**1024>
- #ring = #polynomial.ring<ctype=i32, ideal=#modulus>
+ #ring = #polynomial.ring<coefficientType=i32, polynomialModulus=#modulus>
%a = polynomial.constant <1 + x**2 - 3x**3> : polynomial.polynomial<#ring>
// A constant polynomial in a ring with i32 coefficients, with a polynomial
// modulus of (x^1024 + 1) and a coefficient modulus of 17.
#modulus = #polynomial.polynomial<1 + x**1024>
- #ring = #polynomial.ring<ctype=i32, cmod=17, ideal=#modulus>
+ #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=17, polynomialModulus=#modulus>
%a = polynomial.constant <1 + x**2 - 3x**3> : polynomial.polynomial<#ring>
```
}];
@@ -63,7 +63,21 @@ class Polynomial_Attr<string name, string attrMnemonic, list<Trait> traits = []>
def Polynomial_PolynomialAttr : Polynomial_Attr<"Polynomial", "polynomial"> {
let summary = "An attribute containing a single-variable polynomial.";
let description = [{
- #poly = #polynomial.poly<x**1024 + 1>
+ A polynomial attribute represents a single-variable polynomial, which
+ is used to define the modulus of a `RingAttr`, as well as to define constants
+ and perform constant folding for `polynomial` ops.
+
+ The polynomial must be expressed as a list of monomial terms, with addition
+ or subtraction between them. The choice of variable name is arbitrary, but
+ must be consistent across all the monomials used to define a single
+ attribute. The order of monomial terms is arbitrary, each monomial degree
+ must occur at most once.
+
+ Example:
+
+ ```mlir
+ #poly = #polynomial.polynomial<x**1024 + 1>
+ ```
}];
let parameters = (ins "Polynomial":$polynomial);
let hasCustomAssemblyFormat = 1;
@@ -79,10 +93,10 @@ def Polynomial_RingAttr : Polynomial_Attr<"Ring", "ring"> {
integral, whose coefficients are taken modulo some statically known modulus
(`coefficientModulus`).
- Additionally, a polynomial ring can specify an _ideal_, which converts
+ Additionally, a polynomial ring can specify a _polynomialModulus_, which converts
polynomial arithmetic to the analogue of modular integer arithmetic, where
each polynomial is represented as its remainder when dividing by the
- modulus. For single-variable polynomials, an "ideal" is always specificed
+ modulus. For single-variable polynomials, an "polynomialModulus" is always specificed
via a single polynomial, which we call `polynomialModulus`.
An expressive example is polynomials with i32 coefficients, whose
@@ -122,32 +136,284 @@ class Polynomial_Type<string name, string typeMnemonic>
def Polynomial_PolynomialType : Polynomial_Type<"Polynomial", "polynomial"> {
let summary = "An element of a polynomial ring.";
-
let description = [{
A type for polynomials in a polynomial quotient ring.
}];
-
let parameters = (ins Polynomial_RingAttr:$ring);
let assemblyFormat = "`<` $ring `>`";
}
+def PolynomialLike: TypeOrContainer<Polynomial_PolynomialType, "polynomial-like">;
+
class Polynomial_Op<string mnemonic, list<Trait> traits = []> :
- Op<Polynomial_Dialect, mnemonic, traits # [Pure]>;
+ Op<Polynomial_Dialect, mnemonic, traits # [Pure]> {
+ let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)";
+}
class Polynomial_UnaryOp<string mnemonic, list<Trait> traits = []> :
Polynomial_Op<mnemonic, traits # [SameOperandsAndResultType]> {
let arguments = (ins Polynomial_PolynomialType:$operand);
let results = (outs Polynomial_PolynomialType:$result);
-
- let assemblyFormat = "$operand attr-dict `:` qualified(type($result))";
}
class Polynomial_BinaryOp<string mnemonic, list<Trait> traits = []> :
- Polynomial_Op<mnemonic, traits # [SameOperandsAndResultType]> {
- let arguments = (ins Polynomial_PolynomialType:$lhs, Polynomial_PolynomialType:$rhs);
- let results = (outs Polynomial_PolynomialType:$result);
+ Polynomial_Op<mnemonic, !listconcat(traits, [Pure, SameOperandsAndResultType, ElementwiseMappable])> {
+ let arguments = (ins PolynomialLike:$lhs, PolynomialLike:$rhs);
+ let results = (outs PolynomialLike:$result);
+ let assemblyFormat = "operands attr-dict `:` type($result)";
+}
+
+def Polynomial_AddOp : Polynomial_BinaryOp<"add", [Commutative]> {
+ let summary = "Addition operation between polynomials.";
+ let description = [{
+ Performs polynomial addition on the operands. The operands may be single
+ polynomials or containers of identically-typed polynomials, i.e., polynomials
+ from the same underlying ring with the same coefficient types.
+
+ Addition is defined to occur in the ring defined by the ring attribute of
+ the two operands, meaning the addition is taken modulo the coefficientModulus
+ and the polynomialModulus of the ring.
+
+ Example:
+
+ ```mlir
+ // add two polynomials modulo x^1024 - 1
+ #poly = #polynomial.polynomial<x**1024 - 1>
+ #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536, polynomialModulus=#poly>
+ %0 = polynomial.constant #polynomial.polynomial<1 + x**2> : !polynomial.polynomial<#ring>
+ %1 = polynomial.constant #polynomial.polynomial<x**5 - x + 1> : !polynomial.polynomial<#ring>
+ %2 = polynomial.add %0, %1 : !polynomial.polynomial<#ring>
+ ```
+ }];
+}
+
+def Polynomial_SubOp : Polynomial_BinaryOp<"sub"> {
+ let summary = "Subtraction operation between polynomials.";
+ let description = [{
+ Performs polynomial subtraction on the operands. The operands may be single
+ polynomials or containers of identically-typed polynomials, i.e., polynomials
+ from the same underlying ring with the same coefficient types.
+
+ Subtraction is defined to occur in the ring defined by the ring attribute of
+ the two operands, meaning the subtraction is taken modulo the coefficientModulus
+ and the polynomialModulus of the ring.
+
+ Example:
- let assemblyFormat = "$lhs `,` $rhs attr-dict `:` qualified(type($result))";
+ ```mlir
+ // subtract two polynomials modulo x^1024 - 1
+ #poly = #polynomial.polynomial<x**1024 - 1>
+ #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536, polynomialModulus=#poly>
+ %0 = polynomial.constant #polynomial.polynomial<1 + x**2> : !polynomial.polynomial<#ring>
+ %1 = polynomial.constant #polynomial.polynomial<x**5 - x + 1> : !polynomial.polynomial<#ring>
+ %2 = polynomial.sub %0, %1 : !polynomial.polynomial<#ring>
+ ```
+ }];
+}
+
+def Polynomial_MulOp : Polynomial_BinaryOp<"mul", [Commutative]> {
+ let summary = "Multiplication operation between polynomials.";
+ let description = [{
+ Performs polynomial multiplication on the operands. The operands may be single
+ polynomials or containers of identically-typed polynomials, i.e., polynomials
+ from the same underlying ring with the same coefficient types.
+
+ Multiplication is defined to occur in the ring defined by the ring attribute of
+ the two operands, meaning the multiplication is taken modulo the coefficientModulus
+ and the polynomialModulus of the ring.
+
+ Example:
+
+ ```mlir
+ // multiply two polynomials modulo x^1024 - 1
+ #poly = #polynomial.polynomial<x**1024 - 1>
+ #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536, polynomialModulus=#poly>
+ %0 = polynomial.constant #polynomial.polynomial<1 + x**2> : !polynomial.polynomial<#ring>
+ %1 = polynomial.constant #polynomial.polynomial<x**5 - x + 1> : !polynomial.polynomial<#ring>
+ %2 = polynomial.mul %0, %1 : !polynomial.polynomial<#ring>
+ ```
+ }];
+}
+
+def Polynomial_MulScalarOp : Polynomial_Op<"mul_scalar", [
+ ElementwiseMappable, AllTypesMatch<["polynomial", "output"]>]> {
+ let summary = "Multiplication by a scalar of the field.";
+ let description = [{
+ Multiplies the polynomial operand's coefficients by a given scalar value.
+ The operation is defined to occur in the ring defined by the ring attribute
+ of the two operands, meaning the multiplication is taken modulo the
+ coefficientModulus of the ring.
+
+ The `scalar` input must have the same type as the polynomial ring's
+ coefficientType.
+
+ Example:
+
+ ```mlir
+ // multiply two polynomials modulo x^1024 - 1
+ #poly = #polynomial.polynomial<x**1024 - 1>
+ #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536, polynomialModulus=#poly>
+ %0 = polynomial.constant #polynomial.polynomial<1 + x**2> : !polynomial.polynomial<#ring>
+ %1 = arith.constant 3 : i32
+ %2 = polynomial.mul_scalar %0, %1 : !polynomial.polynomial<#ring>, i32
+ ```
+ }];
+
+ let arguments = (ins
+ PolynomialLike:$polynomial,
+ AnyInteger:$scalar
+ );
+ let results = (outs
+ PolynomialLike:$output
+ );
+ let assemblyFormat = "operands attr-dict `:` type($polynomial) `,` type($scalar)";
+ let hasVerifier = 1;
+}
+
+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. The first is the degree `k` as an
+ index, and the second is the coefficient, whose type matches the
+ coefficient type of the polynomial's ring attribute.
+
+ Example:
+
+ ```mlir
+ #poly = #polynomial.polynomial<x**1024 - 1>
+ #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536, polynomialModulus=#poly>
+ %0 = polynomial.constant #polynomial.polynomial<1 + x**2> : !polynomial.polynomial<#ring>
+ %1, %2 = polynomial.leading_term %0 : !polynomial.polynomial<#ring> -> (index, i32)
+ ```
+ }];
+ let arguments = (ins Polynomial_PolynomialType:$input);
+ let results = (outs Index:$degree, AnyInteger:$coefficient);
+ let assemblyFormat = "operands attr-dict `:` type($input) `->` `(` type($degree) `,` type($coefficient) `)`";
+}
+
+def Polynomial_MonomialOp: Polynomial_Op<"monomial"> {
+ let summary = "Create a polynomial that consists of a single monomial.";
+ let description = [{
+ Construct a polynomial that consists of a single monomial term, from its
+ degree and coefficient as dynamic inputs.
+
+ The coefficient type of the output polynomial's ring attribute must match
+ the `coefficient` input type.
+
+ Example:
+
+ ```mlir
+ #poly = #polynomial.polynomial<x**1024 - 1>
+ #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536, polynomialModulus=#poly>
+ %deg = arith.constant 1023 : index
+ %five = arith.constant 5 : i32
+ %0 = polynomial.monomial %five, %deg : (i32, index) -> !polynomial.polynomial<#ring>
+ ```
+ }];
+ let arguments = (ins AnyInteger:$coefficient, Index:$degree);
+ let results = (outs Polynomial_PolynomialType:$output);
+}
+
+def Polynomial_MonicMonomialMulOp: Polynomial_Op<"monic_monomial_mul", [AllTypesMatch<["input", "output"]>]> {
+ let summary = "Multiply a polynomial by a monic monomial.";
+ let description = [{
+ Multiply a polynomial by a monic monomial, meaning a polynomial of the form
+ `1 * x^k` for an index operand `k`.
+
+ In some special rings of polynomials, such as a ring of polynomials
+ modulo `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 PolynomialLike:$input, Index:$monomialDegree);
+ let results = (outs PolynomialLike:$output);
+}
+
+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 polynomialModulus generator polynomial, with smaller dimension implying that
+ all higher-degree terms have coefficient zero.
+
+ Example:
+
+ ```mlir
+ #poly = #polynomial.polynomial<x**1024 - 1>
+ #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536, polynomialModulus=#poly>
+ %two = arith.constant 2 : i32
+ %five = arith.constant 5 : i32
+ %coeffs = tensor.from_elements %two, %two, %five : tensor<3xi32>
+ %poly = polynomial.from_tensor %coeffs : tensor<3xi32> -> !polynomial.polynomial<#ring>
+ ```
+ }];
+ let arguments = (ins RankedTensorOf<[AnyInteger]>:$input);
+ let results = (outs Polynomial_PolynomialType:$output);
+
+ let assemblyFormat = "$input attr-dict `:` type($input) `->` 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, "::mlir::polynomial::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 dense 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 polynomial ring
+ attribute's polynomialModulus, including zeroes.
+
+ Example:
+
+ ```mlir
+ #poly = #polynomial.polynomial<x**1024 - 1>
+ #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536, polynomialModulus=#poly>
+ %two = arith.constant 2 : i32
+ %five = arith.constant 5 : i32
+ %coeffs = tensor.from_elements %two, %two, %five : tensor<3xi32>
+ %poly = polynomial.from_tensor %coeffs : tensor<3xi32> -> !polynomial.polynomial<#ring>
+ %tensor = polynomial.to_tensor %poly : !polynomial.polynomial<#ring> -> tensor<1024xi32>
+ ```
+ }];
+ let arguments = (ins Polynomial_PolynomialType:$input);
+ let results = (outs RankedTensorOf<[AnyInteger]>:$output);
+ let assemblyFormat = "$input attr-dict `:` type($input) `->` type($output)";
+
+ let hasVerifier = 1;
+}
+
+def Polynomial_ConstantOp : Polynomial_Op<"constant", [Pure]> {
+ let summary = "Define a constant polynomial via an attribute.";
+ let description = [{
+ Example:
+
+ ```mlir
+ #poly = #polynomial.polynomial<x**1024 - 1>
+ #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536, polynomialModulus=#poly>
+ %0 = polynomial.constant #polynomial.polynomial<1 + x**2> : !polynomial.polynomial<#ring>
+ ```
+ }];
+ let arguments = (ins Polynomial_PolynomialAttr:$input);
+ let results = (outs Polynomial_PolynomialType:$output);
+ let assemblyFormat = "$input attr-dict `:` type($output)";
}
#endif // POLYNOMIAL_OPS
diff --git a/mlir/lib/Dialect/Polynomial/IR/PolynomialDialect.cpp b/mlir/lib/Dialect/Polynomial/IR/PolynomialDialect.cpp
index a672a59b8a465d9..825b80d70f80326 100644
--- a/mlir/lib/Dialect/Polynomial/IR/PolynomialDialect.cpp
+++ b/mlir/lib/Dialect/Polynomial/IR/PolynomialDialect.cpp
@@ -8,9 +8,18 @@
#include "mlir/Dialect/Polynomial/IR/Polynomial.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"
#include "mlir/Dialect/Polynomial/IR/PolynomialOps.h"
#include "mlir/Dialect/Polynomial/IR/PolynomialTypes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/APInt.h"
#include "llvm/ADT/TypeSwitch.h"
using namespace mlir;
diff --git a/mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp b/mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp
index 96c59a28b8fdcea..8e2bb5f27dc6cc5 100644
--- a/mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp
+++ b/mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp
@@ -6,10 +6,101 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/Polynomial/IR/PolynomialOps.h"
#include "mlir/Dialect/Polynomial/IR/Polynomial.h"
+#include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"
+#include "mlir/Dialect/Polynomial/IR/PolynomialTypes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/APInt.h"
using namespace mlir;
using namespace mlir::polynomial;
-#define GET_OP_CLASSES
-#include "mlir/Dialect/Polynomial/IR/Polynomial.cpp.inc"
+void FromTensorOp::build(OpBuilder &builder, OperationState &result,
+ Value input, RingAttr ring) {
+ TensorType tensorType = dyn_cast<TensorType>(input.getType());
+ auto bitWidth = tensorType.getElementTypeBitWidth();
+ APInt cmod(1 + bitWidth, 1);
+ cmod = cmod << bitWidth;
+ Type resultType = PolynomialType::get(builder.getContext(), ring);
+ build(builder, result, resultType, input);
+}
+
+LogicalResult FromTensorOp::verify() {
+ ArrayRef<int64_t> tensorShape = getInput().getType().getShape();
+ RingAttr ring = getOutput().getType().getRing();
+ unsigned polyDegree = ring.getPolynomialModulus().getPolynomial().getDegree();
+ bool compatible = tensorShape.size() == 1 && tensorShape[0] <= polyDegree;
+ if (!compatible) {
+ InFlightDiagnostic diag = emitOpError()
+ << "input type " << getInput().getType()
+ << " does not match output type "
+ << getOutput().getType();
+ diag.attachNote() << "the input type must be a tensor of shape [d] where d "
+ "is at most the degree of the polynomialModulus of "
+ "the output type's ring attribute";
+ return diag;
+ }
+
+ APInt coefficientModulus = ring.getCoefficientModulus().getValue();
+ unsigned cmodBitWidth = coefficientModulus.ceilLogBase2();
+ unsigned inputBitWidth = getInput().getType().getElementTypeBitWidth();
+
+ if (inputBitWidth > cmodBitWidth) {
+ InFlightDiagnostic diag = emitOpError()
+ << "input tensor element type "
+ << getInput().getType().getElementType()
+ << " is too large to fit in the coefficients of "
+ << getOutput().getType();
+ diag.attachNote() << "the input tensor's elements must be rescaled"
+ " to fit before using from_tensor";
+ return diag;
+ }
+
+ return success();
+}
+
+LogicalResult ToTensorOp::verify() {
+ ArrayRef<int64_t> tensorShape = getOutput().getType().getShape();
+ unsigned polyDegree = getInput()
+ .getType()
+ .getRing()
+ .getPolynomialModulus()
+ .getPolynomial()
+ .getDegree();
+ bool compatible = tensorShape.size() == 1 && tensorShape[0] == polyDegree;
+
+ if (compatible)
+ return success();
+
+ InFlightDiagnostic diag =
+ emitOpError() << "input type " << getInput().getType()
+ << " does not match output type " << getOutput().getType();
+ diag.attachNote() << "the output type must be a tensor of shape [d] where d "
+ "is at most the degree of the polynomialModulus of "
+ "the input type's ring attribute";
+ return diag;
+}
+
+LogicalResult MulScalarOp::verify() {
+ Type argType = getPolynomial().getType();
+ PolynomialType polyType;
+
+ if (auto shapedPolyType = dyn_cast<ShapedType>(argType)) {
+ polyType = cast<PolynomialType>(shapedPolyType.getElementType());
+ } else {
+ polyType = cast<PolynomialType>(argType);
+ }
+
+ Type coefficientType = polyType.getRing().getCoefficientType();
+
+ if (coefficientType != getScalar().getType())
+ return emitOpError() << "polynomial coefficient type " << coefficientType
+ << " does not match scalar type "
+ << getScalar().getType();
+
+ return success();
+}
diff --git a/mlir/test/Dialect/Polynomial/ops.mlir b/mlir/test/Dialect/Polynomial/ops.mlir
new file mode 100644
index 000000000000000..ea1b279fa1ff96a
--- /dev/null
+++ b/mlir/test/Dialect/Polynomial/ops.mlir
@@ -0,0 +1,82 @@
+// RUN: mlir-opt %s | FileCheck %s
+
+// This simply tests for syntax.
+
+#my_poly = #polynomial.polynomial<1 + x**1024>
+#my_poly_2 = #polynomial.polynomial<2>
+#my_poly_3 = #polynomial.polynomial<3x>
+#my_poly_4 = #polynomial.polynomial<t**3 + 4t + 2>
+#ring1 = #polynomial.ring<coefficientType=i32, coefficientModulus=2837465, polynomialModulus=#my_poly>
+#one_plus_x_squared = #polynomial.polynomial<1 + x**2>
+
+#ideal = #polynomial.polynomial<-1 + x**1024>
+#ring = #polynomial.ring<coefficientType=i32, coefficientModulus=18, polynomialModulus=#ideal>
+!poly_ty = !polynomial.polynomial<#ring>
+
+module {
+ func.func @test_multiply() -> !polynomial.polynomial<#ring1> {
+ %c0 = arith.constant 0 : index
+ %two = arith.constant 2 : i16
+ %five = arith.constant 5 : i16
+ %coeffs1 = tensor.from_elements %two, %two, %five : tensor<3xi16>
+ %coeffs2 = tensor.from_elements %five, %five, %two : tensor<3xi16>
+
+ %poly1 = polynomial.from_tensor %coeffs1 : tensor<3xi16> -> !polynomial.polynomial<#ring1>
+ %poly2 = polynomial.from_tensor %coeffs2 : tensor<3xi16> -> !polynomial.polynomial<#ring1>
+
+ %3 = polynomial.mul %poly1, %poly2 : !polynomial.polynomial<#ring1>
+
+ return %3 : !polynomial.polynomial<#ring1>
+ }
+
+ func.func @test_elementwise(%p0 : !polynomial.polynomial<#ring1>, %p1: !polynomial.polynomial<#ring1>) {
+ %tp0 = tensor.from_elements %p0, %p1 : tensor<2x!polynomial.polynomial<#ring1>>
+ %tp1 = tensor.from_elements %p1, %p0 : tensor<2x!polynomial.polynomial<#ring1>>
+
+ %c = arith.constant 2 : i32
+ %mul_const_sclr = polynomial.mul_scalar %tp0, %c : tensor<2x!polynomial.polynomial<#ring1>>, i32
+
+ %add = polynomial.add %tp0, %tp1 : tensor<2x!polynomial.polynomial<#ring1>>
+ %sub = polynomial.sub %tp0, %tp1 : tensor<2x!polynomial.polynomial<#ring1>>
+ %mul = polynomial.mul %tp0, %tp1 : tensor<2x!polynomial.polynomial<#ring1>>
+
+ return
+ }
+
+ func.func @test_to_from_tensor(%p0 : !polynomial.polynomial<#ring1>) {
+ %c0 = arith.constant 0 : index
+ %two = arith.constant 2 : i16
+ %coeffs1 = tensor.from_elements %two, %two : tensor<2xi16>
+ // CHECK: from_tensor
+ %poly = polynomial.from_tensor %coeffs1 : tensor<2xi16> -> !polynomial.polynomial<#ring1>
+ // CHECK: to_tensor
+ %tensor = polynomial.to_tensor %poly : !polynomial.polynomial<#ring1> -> tensor<1024xi16>
+
+ return
+ }
+
+ func.func @test_degree(%p0 : !polynomial.polynomial<#ring1>) {
+ %0, %1 = polynomial.leading_term %p0 : !polynomial.polynomial<#ring1> -> (index, i32)
+ return
+ }
+
+ func.func @test_monomial() {
+ %deg = arith.constant 1023 : index
+ %five = arith.constant 5 : i16
+ %0 = polynomial.monomial %five, %deg : (i16, index) -> !polynomial.polynomial<#ring1>
+ return
+ }
+
+ func.func @test_monic_monomial_mul() {
+ %five = arith.constant 5 : index
+ %0 = polynomial.constant #one_plus_x_squared : !polynomial.polynomial<#ring1>
+ %1 = polynomial.monic_monomial_mul %0, %five : (!polynomial.polynomial<#ring1>, index) -> !polynomial.polynomial<#ring1>
+ return
+ }
+
+ func.func @test_constant() {
+ %0 = polynomial.constant #one_plus_x_squared : !polynomial.polynomial<#ring1>
+ %1 = polynomial.constant <1 + x**2> : !polynomial.polynomial<#ring1>
+ return
+ }
+}
diff --git a/mlir/test/Dialect/Polynomial/ops_errors.mlir b/mlir/test/Dialect/Polynomial/ops_errors.mlir
new file mode 100644
index 000000000000000..c34a7de30e5fe5b
--- /dev/null
+++ b/mlir/test/Dialect/Polynomial/ops_errors.mlir
@@ -0,0 +1,53 @@
+// RUN: mlir-opt --split-input-file --verify-diagnostics %s
+
+#my_poly = #polynomial.polynomial<1 + x**1024>
+#ring = #polynomial.ring<coefficientType=i16, coefficientModulus=256, polynomialModulus=#my_poly>
+!ty = !polynomial.polynomial<#ring>
+
+func.func @test_from_tensor_too_large_coeffs() {
+ %two = arith.constant 2 : i32
+ %coeffs1 = tensor.from_elements %two, %two : tensor<2xi32>
+ // expected-error at below {{is too large to fit in the coefficients}}
+ // expected-note at below {{rescaled to fit}}
+ %poly = polynomial.from_tensor %coeffs1 : tensor<2xi32> -> !ty
+ return
+}
+
+// -----
+
+#my_poly = #polynomial.polynomial<1 + x**4>
+#ring = #polynomial.ring<coefficientType=i32, coefficientModulus=256, polynomialModulus=#my_poly>
+!ty = !polynomial.polynomial<#ring>
+func.func @test_from_tensor_wrong_tensor_type() {
+ %two = arith.constant 2 : i32
+ %coeffs1 = tensor.from_elements %two, %two, %two, %two, %two : tensor<5xi32>
+ // expected-error at below {{input type 'tensor<5xi32>' does not match output type '!polynomial.polynomial<#polynomial.ring<coefficientType=i32, coefficientModulus=256 : i32, polynomialModulus=#polynomial.polynomial<1 + x**4>>>'}}
+ // expected-note at below {{at most the degree of the polynomialModulus of the output type's ring attribute}}
+ %poly = polynomial.from_tensor %coeffs1 : tensor<5xi32> -> !ty
+ return
+}
+
+// -----
+
+#my_poly = #polynomial.polynomial<1 + x**4>
+#ring = #polynomial.ring<coefficientType=i32, coefficientModulus=256, polynomialModulus=#my_poly>
+!ty = !polynomial.polynomial<#ring>
+func.func @test_to_tensor_wrong_output_tensor_type(%arg0 : !ty) {
+ // expected-error at below {{input type '!polynomial.polynomial<#polynomial.ring<coefficientType=i32, coefficientModulus=256 : i32, polynomialModulus=#polynomial.polynomial<1 + x**4>>>' does not match output type 'tensor<5xi32>'}}
+ // expected-note at below {{at most the degree of the polynomialModulus of the input type's ring attribute}}
+ %tensor = polynomial.to_tensor %arg0 : !ty -> tensor<5xi32>
+ return
+}
+
+// -----
+
+#my_poly = #polynomial.polynomial<1 + x**1024>
+#ring = #polynomial.ring<coefficientType=i16, coefficientModulus=256, polynomialModulus=#my_poly>
+!ty = !polynomial.polynomial<#ring>
+
+func.func @test_mul_scalar_wrong_type(%arg0: !ty) -> !ty {
+ %scalar = arith.constant 2 : i32 // should be i16
+ // expected-error at below {{polynomial coefficient type 'i16' does not match scalar type 'i32'}}
+ %poly = polynomial.mul_scalar %arg0, %scalar : !ty, i32
+ return %poly : !ty
+}
diff --git a/mlir/test/Dialect/Polynomial/types.mlir b/mlir/test/Dialect/Polynomial/types.mlir
index 64b74d9d36bb1c0..00296a36e890f99 100644
--- a/mlir/test/Dialect/Polynomial/types.mlir
+++ b/mlir/test/Dialect/Polynomial/types.mlir
@@ -40,3 +40,17 @@ func.func @test_non_x_variable_64_bit(%0: !ty2) -> !ty2 {
func.func @test_linear_poly(%0: !ty3) -> !ty3 {
return %0 : !ty3
}
+
+// CHECK-LABEL: func @test_negative_leading_1
+// CHECK-SAME: !polynomial.polynomial<
+// CHECK-SAME: #polynomial.ring<
+// CHECK-SAME: coefficientType=i32,
+// CHECK-SAME: coefficientModulus=2837465 : i32,
+// CHECK-SAME: polynomialModulus=#polynomial.polynomial<-1 + x**1024>>>
+#my_poly_4 = #polynomial.polynomial<-1 + x**1024>
+#ring4 = #polynomial.ring<coefficientType=i32, coefficientModulus=2837465, polynomialModulus=#my_poly_4>
+!ty4 = !polynomial.polynomial<#ring4>
+func.func @test_negative_leading_1(%0: !ty4) -> !ty4 {
+ return %0 : !ty4
+}
+
More information about the Mlir-commits
mailing list