diff --git a/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h b/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h index 39b05b9d3ad14..3325a6fa3f9fc 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 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 5d8da8399b01b..d3e3ac55677f8 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 + #ring = #polynomial.ring %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 + #ring = #polynomial.ring %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 + #ring = #polynomial.ring %a = polynomial.constant <1 + x**2 - 3x**3> : polynomial.polynomial<#ring> ``` }]; @@ -63,7 +63,21 @@ class Polynomial_Attr traits = []> def Polynomial_PolynomialAttr : Polynomial_Attr<"Polynomial", "polynomial"> { let summary = "An attribute containing a single-variable polynomial."; let description = [{ - #poly = #polynomial.poly + 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 + ``` }]; 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 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; + class Polynomial_Op traits = []> : - Op; + Op { + let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; +} class Polynomial_UnaryOp traits = []> : Polynomial_Op { let arguments = (ins Polynomial_PolynomialType:$operand); let results = (outs Polynomial_PolynomialType:$result); - - let assemblyFormat = "$operand attr-dict `:` qualified(type($result))"; } class Polynomial_BinaryOp traits = []> : - Polynomial_Op { - let arguments = (ins Polynomial_PolynomialType:$lhs, Polynomial_PolynomialType:$rhs); - let results = (outs Polynomial_PolynomialType:$result); + Polynomial_Op { + 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 + #ring = #polynomial.ring + %0 = polynomial.constant #polynomial.polynomial<1 + x**2> : !polynomial.polynomial<#ring> + %1 = polynomial.constant #polynomial.polynomial : !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 + #ring = #polynomial.ring + %0 = polynomial.constant #polynomial.polynomial<1 + x**2> : !polynomial.polynomial<#ring> + %1 = polynomial.constant #polynomial.polynomial : !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 + #ring = #polynomial.ring + %0 = polynomial.constant #polynomial.polynomial<1 + x**2> : !polynomial.polynomial<#ring> + %1 = polynomial.constant #polynomial.polynomial : !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 + #ring = #polynomial.ring + %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 + #ring = #polynomial.ring + %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 + #ring = #polynomial.ring + %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 + #ring = #polynomial.ring + %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 + #ring = #polynomial.ring + %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 + #ring = #polynomial.ring + %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 a672a59b8a465..825b80d70f803 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 96c59a28b8fdc..8e2bb5f27dc6c 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(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 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 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(argType)) { + polyType = cast(shapedPolyType.getElementType()); + } else { + polyType = cast(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 0000000000000..ea1b279fa1ff9 --- /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 +#ring1 = #polynomial.ring +#one_plus_x_squared = #polynomial.polynomial<1 + x**2> + +#ideal = #polynomial.polynomial<-1 + x**1024> +#ring = #polynomial.ring +!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 0000000000000..c34a7de30e5fe --- /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 +!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@below {{is too large to fit in the coefficients}} + // expected-note@below {{rescaled to fit}} + %poly = polynomial.from_tensor %coeffs1 : tensor<2xi32> -> !ty + return +} + +// ----- + +#my_poly = #polynomial.polynomial<1 + x**4> +#ring = #polynomial.ring +!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@below {{input type 'tensor<5xi32>' does not match output type '!polynomial.polynomial<#polynomial.ring>>'}} + // expected-note@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 +!ty = !polynomial.polynomial<#ring> +func.func @test_to_tensor_wrong_output_tensor_type(%arg0 : !ty) { + // expected-error@below {{input type '!polynomial.polynomial<#polynomial.ring>>' does not match output type 'tensor<5xi32>'}} + // expected-note@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 +!ty = !polynomial.polynomial<#ring> + +func.func @test_mul_scalar_wrong_type(%arg0: !ty) -> !ty { + %scalar = arith.constant 2 : i32 // should be i16 + // expected-error@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 64b74d9d36bb1..00296a36e890f 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 +!ty4 = !polynomial.polynomial<#ring4> +func.func @test_negative_leading_1(%0: !ty4) -> !ty4 { + return %0 : !ty4 +} +