try constraints

Author: j2kun

mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td

@@ -62,7 +62,7 @@ class Polynomial_Attr<string name, string attrMnemonic, list<Trait> traits = []>
 }
 
 def Polynomial_IntPolynomialAttr : Polynomial_Attr<"IntPolynomial", "int_polynomial"> {
-  let summary = "An attribute containing a single-variable polynomial with integer coefficients.";
+  let summary = "an attribute containing a single-variable polynomial with integer coefficients";
   let description = [{
     A polynomial attribute represents a single-variable polynomial with integer
     coefficients, which is used to define the modulus of a `RingAttr`, as well
@@ -109,7 +109,7 @@ def Polynomial_TypedIntPolynomialAttr : Polynomial_Attr<
 }
 
 def Polynomial_FloatPolynomialAttr : Polynomial_Attr<"FloatPolynomial", "float_polynomial"> {
-  let summary = "An attribute containing a single-variable polynomial with double precision floating point coefficients.";
+  let summary = "an attribute containing a single-variable polynomial with double precision floating point coefficients";
   let description = [{
     A polynomial attribute represents a single-variable polynomial with double
     precision floating point coefficients.
@@ -489,6 +489,25 @@ def Polynomial_AnyPolynomialAttr : AnyAttrOf<[
   Polynomial_FloatPolynomialAttr,
   Polynomial_IntPolynomialAttr
 ]>;
+def Polynomial_PolynomialElementsAttr :
+    ElementsAttrBase<And<[//CPred<"::llvm::isa<::mlir::ElementsAttr>($_self)">,
+                          CPred<[{
+                              isa<::mlir::polynomial::PolynomialType>(
+                                ::llvm::cast<::mlir::ElementsAttr>($_self)
+                                                                    .getShapedType()
+                                                                    .getElementType())
+                              }]>]>,
+                     "an elements attribute containing polynomial attributes"> {
+  let storageType = [{ ::mlir::ElementsAttr }];
+  let returnType = [{ ::mlir::ElementsAttr }];
+  let convertFromStorage = "$_self";
+}
+
+def Polynomial_PolynomialOrElementsAttr : AnyAttrOf<[
+  Polynomial_FloatPolynomialAttr,
+  Polynomial_IntPolynomialAttr,
+  Polynomial_PolynomialElementsAttr,
+]>;
 
 // Not deriving from Polynomial_Op due to need for custom assembly format
 def Polynomial_ConstantOp : Op<Polynomial_Dialect, "constant", [Pure, ConstantLike]> {
@@ -505,8 +524,8 @@ def Polynomial_ConstantOp : Op<Polynomial_Dialect, "constant", [Pure, ConstantLi
     %0 = polynomial.constant #polynomial.float_polynomial<0.5 + 1.3e06 x**2> : !polynomial.polynomial<#float_ring>
     ```
   }];
-  let arguments = (ins Polynomial_AnyPolynomialAttr:$value);
-  let results = (outs Polynomial_PolynomialType:$output);
+  let arguments = (ins Polynomial_PolynomialOrElementsAttr:$value);
+  let results = (outs PolynomialLike:$output);
   let assemblyFormat = "attr-dict `:` type($output)";
   let hasFolder = 1;
 }

mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp

@@ -7,7 +7,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Polynomial/IR/PolynomialOps.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/Dialect/Polynomial/IR/Polynomial.h"
 #include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"

mlir/test/Dialect/Polynomial/folding.mlir

@@ -2,11 +2,9 @@
 
 // Tests for folding
 
-#my_poly = #polynomial.int_polynomial<1 + x**1024>
 #poly_3t = #polynomial.int_polynomial<3t>
 #poly_t3_plus_4t_plus_2 = #polynomial.int_polynomial<t**3 + 4t + 2>
-#modulus = #polynomial.int_polynomial<-1 + x**1024>
-#ring = #polynomial.ring<coefficientType=i32, coefficientModulus=256, polynomialModulus=#modulus, primitiveRoot=193>
+#ring = #polynomial.ring<coefficientType=i32>
 !poly_ty = !polynomial.polynomial<ring=#ring>
 
 // CHECK-LABEL: test_fold_add
@@ -19,5 +17,38 @@ func.func @test_fold_add() -> !poly_ty {
   return %2 : !poly_ty
 }
 
+// CHECK-LABEL: test_fold_add_elementwise
+// CHECK-NEXT: polynomial.constant {value = dense<
+// CHECK-SAME:  #polynomial.typed_int_polynomial<type=
+// CHECK-SAME:     value = <2 + 7x + x**3>>,
+// CHECK-SAME:  #polynomial.typed_int_polynomial<type=
+// CHECK-SAME:     value = <2 + 7x + x**3>>,
+// CHECK-SAME: ]>}
+// CHECK-NEXT: return
+#typed_poly1 = #polynomial.typed_int_polynomial<type=!poly_ty, value=<3t>>
+#typed_poly2 = #polynomial.typed_int_polynomial<type=!poly_ty, value=<t**3 + 4t + 2>>
+!tensor_ty = tensor<2x!poly_ty>
+func.func @test_fold_add_elementwise() -> !tensor_ty {
+  %0 = polynomial.constant {value=[#typed_poly1, #typed_poly2]} : !tensor_ty
+  %1 = polynomial.constant {value=[#typed_poly2, #typed_poly1]} : !tensor_ty
+  %2 = polynomial.add %0, %1 : !tensor_ty
+  return %2 : !tensor_ty
+}
+
+
+#fpoly_1 = #polynomial.float_polynomial<3.5t>
+#fpoly_2 = #polynomial.float_polynomial<1.0t**3 + 1.25t + 2.0>
+#fring = #polynomial.ring<coefficientType=f32>
+!fpoly_ty = !polynomial.polynomial<ring=#fring>
+
+// CHECK-LABEL: test_fold_add_float
+// CHECK-NEXT: polynomial.constant {value = #polynomial.float_polynomial<2 + 4.75x + x**3>}
+// CHECK-NEXT: return
+func.func @test_fold_add_float() -> !fpoly_ty {
+  %0 = polynomial.constant {value=#fpoly_1} : !fpoly_ty
+  %1 = polynomial.constant {value=#fpoly_2} : !fpoly_ty
+  %2 = polynomial.add %0, %1 : !fpoly_ty
+  return %2 : !fpoly_ty
+}
+
 // Test elementwise folding of add
-// Test float folding of add