try constBinaryFold

Author: j2kun

mlir/include/mlir/Dialect/CommonFolders.h

@@ -82,7 +82,8 @@ Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
     if (!elementResult)
       return {};
 
-    return DenseElementsAttr::get(cast<ShapedType>(resultType), *elementResult);
+    return DenseElementsAttr::get(cast<ShapedType>(resultType),
+                                  llvm::ArrayRef(*elementResult));
   }
 
   if (isa<ElementsAttr>(operands[0]) && isa<ElementsAttr>(operands[1])) {

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

@@ -83,6 +83,23 @@ 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::IntPolynomial":$value);
+  let assemblyFormat = "`<` struct(params) `>`";
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "Type":$type,
+                                        "const IntPolynomial &":$value), [{
+      return $_get(type.getContext(), type, value);
+    }]>
+  ];
+  let extraClassDeclaration = [{
+    // used for constFoldBinaryOp
+    using ValueType = ::mlir::polynomial::IntPolynomial;
+  }];
+}
+
 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 +122,23 @@ 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::FloatPolynomial":$value);
+  let assemblyFormat = "`<` struct(params) `>`";
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "Type":$type,
+                                        "const FloatPolynomial &":$value), [{
+      return $_get(type.getContext(), type, value);
+    }]>
+  ];
+  let extraClassDeclaration = [{
+    // used for constFoldBinaryOp
+    using ValueType = ::mlir::polynomial::FloatPolynomial;
+  }];
+}
+
 def Polynomial_RingAttr : Polynomial_Attr<"Ring", "ring"> {
   let summary = "An attribute specifying a polynomial ring.";
   let description = [{
@@ -221,6 +255,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"> {
@@ -442,7 +477,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:
@@ -459,6 +494,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/PolynomialOps.cpp

@@ -8,10 +8,12 @@
 
 #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/IR/PatternMatch.h"
@@ -21,6 +23,38 @@
 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) {
+  // Folded input attributes can either be typed_int_polynomial or
+  // typed_float_polynomial, and those require different invocations of
+  // constFoldBinaryOp.
+  PolynomialType ty = dyn_cast<PolynomialType>(getLhs().getType());
+  if (!ty) {
+    ShapedType shapedTy = dyn_cast<ShapedType>(getLhs().getType());
+    assert(shapedTy && "lhs must be a polynomial or a shaped type");
+    ty = cast<PolynomialType>(shapedTy.getElementType());
+  }
+
+  if (isa<FloatPolynomialAttr>(ty.getRing().getPolynomialModulus()))
+    return constFoldBinaryOp<TypedFloatPolynomialAttr>(
+        adaptor.getOperands(), getLhs().getType(),
+        [](FloatPolynomial a, const FloatPolynomial &b) { return a.add(b); });
+
+  return constFoldBinaryOp<TypedIntPolynomialAttr>(
+      adaptor.getOperands(), getLhs().getType(),
+      [](IntPolynomial a, const IntPolynomial &b) { return a.add(b); });
+}
+
 void FromTensorOp::build(OpBuilder &builder, OperationState &result,
                          Value input, RingAttr ring) {
   TensorType tensorType = dyn_cast<TensorType>(input.getType());