fold add via constBinaryFold

Author: j2kun

mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h

@@ -175,6 +175,20 @@ class PolynomialBase {
         continue;
       }
 
+      while (it1->getExponent().ult(it2->getExponent())) {
+        newTerms.emplace_back(*it1);
+        it1++;
+        if (it1 == terms.end())
+          break;
+      }
+
+      while (it2->getExponent().ult(it1->getExponent())) {
+        newTerms.emplace_back(*it2);
+        it2++;
+        if (it2 == terms.end())
+          break;
+      }
+
       newTerms.emplace_back(it1->add(*it2));
       it1++;
       it2++;

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

@@ -53,6 +53,7 @@ def Polynomial_Dialect : Dialect {
 
   let useDefaultTypePrinterParser = 1;
   let useDefaultAttributePrinterParser = 1;
+  let hasConstantMaterializer = 1;
 }
 
 class Polynomial_Attr<string name, string attrMnemonic, list<Trait> traits = []>
@@ -83,6 +84,30 @@ def Polynomial_IntPolynomialAttr : Polynomial_Attr<"IntPolynomial", "int_polynom
   let hasCustomAssemblyFormat = 1;
 }
 
+def Polynomial_TypedIntPolynomialAttr : Polynomial_Attr<
+    "TypedIntPolynomial", "typed_int_polynomial", [TypedAttrInterface]> {
+  let summary = "A typed variant of int_polynomial for constant folding.";
+  let parameters = (ins "::mlir::Type":$type, "::mlir::polynomial::IntPolynomialAttr":$value);
+  let assemblyFormat = "`<` struct(params) `>`";
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "Type":$type,
+                                        "const IntPolynomial &":$value), [{
+      return $_get(
+        type.getContext(),
+        type,
+        IntPolynomialAttr::get(type.getContext(), value));
+    }]>,
+    AttrBuilderWithInferredContext<(ins "Type":$type,
+                                        "const Attribute &":$value), [{
+      return $_get(type.getContext(), type, ::llvm::cast<IntPolynomialAttr>(value));
+    }]>
+  ];
+  let extraClassDeclaration = [{
+    // used for constFoldBinaryOp
+    using ValueType = ::mlir::Attribute;
+  }];
+}
+
 def Polynomial_FloatPolynomialAttr : Polynomial_Attr<"FloatPolynomial", "float_polynomial"> {
   let summary = "An attribute containing a single-variable polynomial with double precision floating point coefficients.";
   let description = [{
@@ -105,6 +130,30 @@ def Polynomial_FloatPolynomialAttr : Polynomial_Attr<"FloatPolynomial", "float_p
   let hasCustomAssemblyFormat = 1;
 }
 
+def Polynomial_TypedFloatPolynomialAttr : Polynomial_Attr<
+    "TypedFloatPolynomial", "typed_float_polynomial", [TypedAttrInterface]> {
+  let summary = "A typed variant of float_polynomial for constant folding.";
+  let parameters = (ins "::mlir::Type":$type, "::mlir::polynomial::FloatPolynomialAttr":$value);
+  let assemblyFormat = "`<` struct(params) `>`";
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "Type":$type,
+                                        "const FloatPolynomial &":$value), [{
+      return $_get(
+        type.getContext(),
+        type,
+        FloatPolynomialAttr::get(type.getContext(), value));
+    }]>,
+    AttrBuilderWithInferredContext<(ins "Type":$type,
+                                        "const Attribute &":$value), [{
+      return $_get(type.getContext(), type, ::llvm::cast<FloatPolynomialAttr>(value));
+    }]>
+  ];
+  let extraClassDeclaration = [{
+    // used for constFoldBinaryOp
+    using ValueType = ::mlir::Attribute;
+  }];
+}
+
 def Polynomial_RingAttr : Polynomial_Attr<"Ring", "ring"> {
   let summary = "An attribute specifying a polynomial ring.";
   let description = [{
@@ -221,6 +270,7 @@ def Polynomial_AddOp : Polynomial_BinaryOp<"add", [Commutative]> {
     %2 = polynomial.add %0, %1 : !polynomial.polynomial<#ring>
     ```
   }];
+  let hasFolder = 1;
 }
 
 def Polynomial_SubOp : Polynomial_BinaryOp<"sub"> {
@@ -441,7 +491,7 @@ def Polynomial_AnyPolynomialAttr : AnyAttrOf<[
 ]>;
 
 // Not deriving from Polynomial_Op due to need for custom assembly format
-def Polynomial_ConstantOp : Op<Polynomial_Dialect, "constant", [Pure]> {
+def Polynomial_ConstantOp : Op<Polynomial_Dialect, "constant", [Pure, ConstantLike]> {
   let summary = "Define a constant polynomial via an attribute.";
   let description = [{
     Example:
@@ -458,6 +508,7 @@ def Polynomial_ConstantOp : Op<Polynomial_Dialect, "constant", [Pure]> {
   let arguments = (ins Polynomial_AnyPolynomialAttr:$value);
   let results = (outs Polynomial_PolynomialType:$output);
   let assemblyFormat = "attr-dict `:` type($output)";
+  let hasFolder = 1;
 }
 
 def Polynomial_NTTOp : Polynomial_Op<"ntt", [Pure]> {

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

@@ -48,3 +48,17 @@ void PolynomialDialect::initialize() {
 #include "mlir/Dialect/Polynomial/IR/Polynomial.cpp.inc"
       >();
 }
+
+Operation *PolynomialDialect::materializeConstant(OpBuilder &builder,
+                                                  Attribute value, Type type,
+                                                  Location loc) {
+  auto intPoly = dyn_cast<TypedIntPolynomialAttr>(value);
+  auto floatPoly = dyn_cast<TypedFloatPolynomialAttr>(value);
+  if (!intPoly && !floatPoly)
+    return nullptr;
+
+  Type ty = intPoly ? intPoly.getType() : floatPoly.getType();
+  Attribute valueAttr =
+      intPoly ? (Attribute)intPoly.getValue() : (Attribute)floatPoly.getValue();
+  return builder.create<ConstantOp>(loc, ty, valueAttr);
+}

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

@@ -7,10 +7,13 @@
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "mlir/Dialect/Polynomial/IR/PolynomialTypes.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Support/LogicalResult.h"
@@ -19,6 +22,41 @@
 using namespace mlir;
 using namespace mlir::polynomial;
 
+OpFoldResult ConstantOp::fold(ConstantOp::FoldAdaptor adaptor) {
+  PolynomialType ty = dyn_cast<PolynomialType>(getOutput().getType());
+
+  if (isa<FloatPolynomialAttr>(ty.getRing().getPolynomialModulus()))
+    return TypedFloatPolynomialAttr::get(
+        ty, cast<FloatPolynomialAttr>(getValue()).getPolynomial());
+
+  assert(isa<IntPolynomialAttr>(ty.getRing().getPolynomialModulus()) &&
+         "expected float or integer polynomial");
+  return TypedIntPolynomialAttr::get(
+      ty, cast<IntPolynomialAttr>(getValue()).getPolynomial());
+}
+
+OpFoldResult AddOp::fold(AddOp::FoldAdaptor adaptor) {
+  auto lhsElements = dyn_cast<ShapedType>(getLhs().getType());
+  PolynomialType elementType = cast<PolynomialType>(
+      lhsElements ? lhsElements.getElementType() : getLhs().getType());
+  MLIRContext *context = getContext();
+
+  if (isa<FloatType>(elementType.getRing().getCoefficientType()))
+    return constFoldBinaryOp<TypedFloatPolynomialAttr>(
+      adaptor.getOperands(), elementType, [&](Attribute a, const Attribute &b) {
+        return FloatPolynomialAttr::get(
+            context, cast<FloatPolynomialAttr>(a).getPolynomial().add(
+                         cast<FloatPolynomialAttr>(b).getPolynomial()));
+      });
+
+  return constFoldBinaryOp<TypedIntPolynomialAttr>(
+      adaptor.getOperands(), elementType, [&](Attribute a, const Attribute &b) {
+        return IntPolynomialAttr::get(
+            context, cast<IntPolynomialAttr>(a).getPolynomial().add(
+                         cast<IntPolynomialAttr>(b).getPolynomial()));
+      });
+}
+
 void FromTensorOp::build(OpBuilder &builder, OperationState &result,
                          Value input, RingAttr ring) {
   TensorType tensorType = dyn_cast<TensorType>(input.getType());

mlir/test/Dialect/Polynomial/folding.mlir

@@ -0,0 +1,23 @@
+// RUN: mlir-opt --sccp --canonicalize %s | FileCheck %s
+
+// 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>
+!poly_ty = !polynomial.polynomial<ring=#ring>
+
+// CHECK-LABEL: test_fold_add
+// CHECK-NEXT: polynomial.constant {value = #polynomial.int_polynomial<2 + 7x + x**3>}
+// CHECK-NEXT: return
+func.func @test_fold_add() -> !poly_ty {
+  %0 = polynomial.constant {value=#poly_3t} : !poly_ty
+  %1 = polynomial.constant {value=#poly_t3_plus_4t_plus_2} : !poly_ty
+  %2 = polynomial.add %0, %1 : !poly_ty
+  return %2 : !poly_ty
+}
+
+// Test elementwise folding of add
+// Test float folding of add

mlir/unittests/Dialect/Polynomial/PolynomialMathTest.cpp

@@ -20,9 +20,10 @@ TEST(AddTest, checkSameDegreeAdditionOfIntPolynomial) {
 }
 
 TEST(AddTest, checkDifferentDegreeAdditionOfIntPolynomial) {
-  IntPolynomial x = IntPolynomial::fromCoefficients({1, 2});
+  IntMonomial term2t = IntMonomial(2, 1);
+  IntPolynomial x = IntPolynomial::fromMonomials({term2t}).value();
   IntPolynomial y = IntPolynomial::fromCoefficients({2, 3, 4});
-  IntPolynomial expected = IntPolynomial::fromCoefficients({3, 5, 4});
+  IntPolynomial expected = IntPolynomial::fromCoefficients({2, 5, 4});
   EXPECT_EQ(expected, x.add(y));
   EXPECT_EQ(expected, y.add(x));
 }