Skip to content

Add constant propagation for polynomial ops #91655

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 3 commits into
base: main
Choose a base branch
from

Conversation

j2kun
Copy link
Contributor

@j2kun j2kun commented May 9, 2024

This PR adds an initial boilerplate setup for constant folding of polynomial ops via compile-time polynomial math. It includes:

  • Adding a CRTP derived type for Monomial and Polynomial
  • Implementation of polynomial addition
  • Constant folding for polynomial.add
  • Constant materialization for the dialect

@j2kun j2kun requested a review from rupprecht as a code owner May 9, 2024 21:07
@j2kun j2kun marked this pull request as draft May 9, 2024 21:07
@llvmbot llvmbot added mlir bazel "Peripheral" support tier build system: utils/bazel labels May 9, 2024
@llvmbot
Copy link
Member

llvmbot commented May 9, 2024

@llvm/pr-subscribers-mlir

Author: Jeremy Kun (j2kun)

Changes

Rebased over


Patch is 73.03 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/91655.diff

17 Files Affected:

  • (modified) mlir/include/mlir/Dialect/CommonFolders.h (+2-1)
  • (modified) mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h (+183-43)
  • (modified) mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td (+126-52)
  • (modified) mlir/lib/Dialect/Polynomial/IR/CMakeLists.txt (+4)
  • (modified) mlir/lib/Dialect/Polynomial/IR/Polynomial.cpp (+27-51)
  • (modified) mlir/lib/Dialect/Polynomial/IR/PolynomialAttributes.cpp (+76-96)
  • (added) mlir/lib/Dialect/Polynomial/IR/PolynomialCanonicalization.td (+37)
  • (modified) mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp (+59)
  • (modified) mlir/test/Dialect/Polynomial/attributes.mlir (+8-14)
  • (added) mlir/test/Dialect/Polynomial/canonicalization.mlir (+50)
  • (modified) mlir/test/Dialect/Polynomial/ops.mlir (+33-31)
  • (modified) mlir/test/Dialect/Polynomial/ops_errors.mlir (+33-33)
  • (modified) mlir/test/Dialect/Polynomial/types.mlir (+37-28)
  • (modified) mlir/unittests/Dialect/CMakeLists.txt (+1)
  • (added) mlir/unittests/Dialect/Polynomial/CMakeLists.txt (+8)
  • (added) mlir/unittests/Dialect/Polynomial/PolynomialMathTest.cpp (+43)
  • (modified) utils/bazel/llvm-project-overlay/mlir/BUILD.bazel (+18)
diff --git a/mlir/include/mlir/Dialect/CommonFolders.h b/mlir/include/mlir/Dialect/CommonFolders.h
index 7dabc781cd595..29e6fccdf2553 100644
--- a/mlir/include/mlir/Dialect/CommonFolders.h
+++ b/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])) {
diff --git a/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h b/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h
index 3325a6fa3f9fc..f0e5bdf16036c 100644
--- a/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h
+++ b/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h
@@ -11,10 +11,13 @@
 
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Hashing.h"
-#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 
@@ -27,98 +30,235 @@ namespace polynomial {
 /// would want to specify 128-bit polynomials statically in the source code.
 constexpr unsigned apintBitWidth = 64;
 
-/// A class representing a monomial of a single-variable polynomial with integer
-/// coefficients.
-class Monomial {
+template <class Derived, typename CoefficientType>
+class MonomialBase {
 public:
-  Monomial(int64_t coeff, uint64_t expo)
-      : coefficient(apintBitWidth, coeff), exponent(apintBitWidth, expo) {}
-
-  Monomial(const APInt &coeff, const APInt &expo)
+  MonomialBase(const CoefficientType &coeff, const APInt &expo)
       : coefficient(coeff), exponent(expo) {}
+  virtual ~MonomialBase() = 0;
 
-  Monomial() : coefficient(apintBitWidth, 0), exponent(apintBitWidth, 0) {}
+  const CoefficientType &getCoefficient() const { return coefficient; }
+  CoefficientType &getMutableCoefficient() { return coefficient; }
+  const APInt &getExponent() const { return exponent; }
+  void setCoefficient(const CoefficientType &coeff) { coefficient = coeff; }
+  void setExponent(const APInt &exp) { exponent = exp; }
 
-  bool operator==(const Monomial &other) const {
+  bool operator==(const MonomialBase &other) const {
     return other.coefficient == coefficient && other.exponent == exponent;
   }
-  bool operator!=(const Monomial &other) const {
+  bool operator!=(const MonomialBase &other) const {
     return other.coefficient != coefficient || other.exponent != exponent;
   }
 
   /// Monomials are ordered by exponent.
-  bool operator<(const Monomial &other) const {
+  bool operator<(const MonomialBase &other) const {
     return (exponent.ult(other.exponent));
   }
 
-  friend ::llvm::hash_code hash_value(const Monomial &arg);
+  Derived add(const Derived &other) {
+    assert(exponent == other.exponent);
+    CoefficientType newCoeff = coefficient + other.coefficient;
+    Derived result;
+    result.setCoefficient(newCoeff);
+    result.setExponent(exponent);
+    return result;
+  }
 
-public:
-  APInt coefficient;
+  virtual bool isMonic() const = 0;
+  virtual void
+  coefficientToString(llvm::SmallString<16> &coeffString) const = 0;
 
-  // Always unsigned
+  template <class D, typename T>
+  friend ::llvm::hash_code hash_value(const MonomialBase<D, T> &arg);
+
+protected:
+  CoefficientType coefficient;
   APInt exponent;
 };
 
-/// A single-variable polynomial with integer coefficients.
-///
-/// Eg: x^1024 + x + 1
-///
-/// The symbols used as the polynomial's indeterminate don't matter, so long as
-/// it is used consistently throughout the polynomial.
-class Polynomial {
+/// A class representing a monomial of a single-variable polynomial with integer
+/// coefficients.
+class IntMonomial : public MonomialBase<IntMonomial, APInt> {
 public:
-  Polynomial() = delete;
+  IntMonomial(int64_t coeff, uint64_t expo)
+      : MonomialBase(APInt(apintBitWidth, coeff), APInt(apintBitWidth, expo)) {}
 
-  explicit Polynomial(ArrayRef<Monomial> terms) : terms(terms){};
+  IntMonomial()
+      : MonomialBase(APInt(apintBitWidth, 0), APInt(apintBitWidth, 0)) {}
 
-  // Returns a Polynomial from a list of monomials.
-  // Fails if two monomials have the same exponent.
-  static FailureOr<Polynomial> fromMonomials(ArrayRef<Monomial> monomials);
+  ~IntMonomial() override = default;
 
-  /// Returns a polynomial with coefficients given by `coeffs`. The value
-  /// coeffs[i] is converted to a monomial with exponent i.
-  static Polynomial fromCoefficients(ArrayRef<int64_t> coeffs);
+  bool isMonic() const override { return coefficient == 1; }
+
+  void coefficientToString(llvm::SmallString<16> &coeffString) const override {
+    coefficient.toStringSigned(coeffString);
+  }
+};
+
+/// A class representing a monomial of a single-variable polynomial with integer
+/// coefficients.
+class FloatMonomial : public MonomialBase<FloatMonomial, APFloat> {
+public:
+  FloatMonomial(double coeff, uint64_t expo)
+      : MonomialBase(APFloat(coeff), APInt(apintBitWidth, expo)) {}
+
+  FloatMonomial() : MonomialBase(APFloat((double)0), APInt(apintBitWidth, 0)) {}
+
+  ~FloatMonomial() override = default;
+
+  bool isMonic() const override { return coefficient == APFloat(1.0); }
+
+  void coefficientToString(llvm::SmallString<16> &coeffString) const override {
+    coefficient.toString(coeffString);
+  }
+};
+
+template <class Derived, typename Monomial>
+class PolynomialBase {
+public:
+  PolynomialBase() = delete;
+
+  explicit PolynomialBase(ArrayRef<Monomial> terms) : terms(terms){};
 
   explicit operator bool() const { return !terms.empty(); }
-  bool operator==(const Polynomial &other) const {
+  bool operator==(const PolynomialBase &other) const {
     return other.terms == terms;
   }
-  bool operator!=(const Polynomial &other) const {
+  bool operator!=(const PolynomialBase &other) const {
     return !(other.terms == terms);
   }
 
-  // Prints polynomial to 'os'.
-  void print(raw_ostream &os) const;
   void print(raw_ostream &os, ::llvm::StringRef separator,
-             ::llvm::StringRef exponentiation) const;
+             ::llvm::StringRef exponentiation) const {
+    bool first = true;
+    for (const Monomial &term : getTerms()) {
+      if (first) {
+        first = false;
+      } else {
+        os << separator;
+      }
+      std::string coeffToPrint;
+      if (term.isMonic() && term.getExponent().uge(1)) {
+        coeffToPrint = "";
+      } else {
+        llvm::SmallString<16> coeffString;
+        term.coefficientToString(coeffString);
+        coeffToPrint = coeffString.str();
+      }
+
+      if (term.getExponent() == 0) {
+        os << coeffToPrint;
+      } else if (term.getExponent() == 1) {
+        os << coeffToPrint << "x";
+      } else {
+        llvm::SmallString<16> expString;
+        term.getExponent().toStringSigned(expString);
+        os << coeffToPrint << "x" << exponentiation << expString;
+      }
+    }
+  }
+
+  Derived add(const Derived &other) {
+    SmallVector<Monomial> newTerms;
+    auto it1 = terms.begin();
+    auto it2 = other.terms.begin();
+    while (it1 != terms.end() || it2 != other.terms.end()) {
+      if (it1 == terms.end()) {
+        newTerms.emplace_back(*it2);
+        it2++;
+        continue;
+      }
+
+      if (it2 == other.terms.end()) {
+        newTerms.emplace_back(*it1);
+        it1++;
+        continue;
+      }
+
+      newTerms.emplace_back(it1->add(*it2));
+      it1++;
+      it2++;
+    }
+    return Derived(newTerms);
+  }
+
+  // Prints polynomial to 'os'.
+  void print(raw_ostream &os) const { print(os, " + ", "**"); }
+
   void dump() const;
 
   // Prints polynomial so that it can be used as a valid identifier
-  std::string toIdentifier() const;
+  std::string toIdentifier() const {
+    std::string result;
+    llvm::raw_string_ostream os(result);
+    print(os, "_", "");
+    return os.str();
+  }
 
-  unsigned getDegree() const;
+  unsigned getDegree() const {
+    return terms.back().getExponent().getZExtValue();
+  }
 
   ArrayRef<Monomial> getTerms() const { return terms; }
 
-  friend ::llvm::hash_code hash_value(const Polynomial &arg);
+  template <class D, typename T>
+  friend ::llvm::hash_code hash_value(const PolynomialBase<D, T> &arg);
 
 private:
   // The monomial terms for this polynomial.
   SmallVector<Monomial> terms;
 };
 
-// Make Polynomial hashable.
-inline ::llvm::hash_code hash_value(const Polynomial &arg) {
+/// A single-variable polynomial with integer coefficients.
+///
+/// Eg: x^1024 + x + 1
+class IntPolynomial : public PolynomialBase<IntPolynomial, IntMonomial> {
+public:
+  explicit IntPolynomial(ArrayRef<IntMonomial> terms) : PolynomialBase(terms) {}
+
+  // Returns a Polynomial from a list of monomials.
+  // Fails if two monomials have the same exponent.
+  static FailureOr<IntPolynomial>
+  fromMonomials(ArrayRef<IntMonomial> monomials);
+
+  /// Returns a polynomial with coefficients given by `coeffs`. The value
+  /// coeffs[i] is converted to a monomial with exponent i.
+  static IntPolynomial fromCoefficients(ArrayRef<int64_t> coeffs);
+};
+
+/// A single-variable polynomial with double coefficients.
+///
+/// Eg: 1.0 x^1024 + 3.5 x + 1e-05
+class FloatPolynomial : public PolynomialBase<FloatPolynomial, FloatMonomial> {
+public:
+  explicit FloatPolynomial(ArrayRef<FloatMonomial> terms)
+      : PolynomialBase(terms) {}
+
+  // Returns a Polynomial from a list of monomials.
+  // Fails if two monomials have the same exponent.
+  static FailureOr<FloatPolynomial>
+  fromMonomials(ArrayRef<FloatMonomial> monomials);
+
+  /// Returns a polynomial with coefficients given by `coeffs`. The value
+  /// coeffs[i] is converted to a monomial with exponent i.
+  static FloatPolynomial fromCoefficients(ArrayRef<double> coeffs);
+};
+
+// Make Polynomials hashable.
+template <class D, typename T>
+inline ::llvm::hash_code hash_value(const PolynomialBase<D, T> &arg) {
   return ::llvm::hash_combine_range(arg.terms.begin(), arg.terms.end());
 }
 
-inline ::llvm::hash_code hash_value(const Monomial &arg) {
+template <class D, typename T>
+inline ::llvm::hash_code hash_value(const MonomialBase<D, T> &arg) {
   return llvm::hash_combine(::llvm::hash_value(arg.coefficient),
                             ::llvm::hash_value(arg.exponent));
 }
 
-inline raw_ostream &operator<<(raw_ostream &os, const Polynomial &polynomial) {
+template <class D, typename T>
+inline raw_ostream &operator<<(raw_ostream &os,
+                               const PolynomialBase<D, T> &polynomial) {
   polynomial.print(os);
   return os;
 }
diff --git a/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td b/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td
index ed1f4ce8b7e59..82b7d0c82952e 100644
--- a/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td
+++ b/mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.td
@@ -39,14 +39,14 @@ def Polynomial_Dialect : Dialect {
     %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>
+    #modulus = #polynomial.int_polynomial<1 + x**1024>
     #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<coefficientType=i32, coefficientModulus=17, polynomialModulus=#modulus>
+    #modulus = #polynomial.int_polynomial<1 + x**1024>
+    #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=17:i32, polynomialModulus=#modulus>
     %a = polynomial.constant <1 + x**2 - 3x**3> : polynomial.polynomial<#ring>
     ```
   }];
@@ -60,12 +60,12 @@ class Polynomial_Attr<string name, string attrMnemonic, list<Trait> traits = []>
   let mnemonic = attrMnemonic;
 }
 
-def Polynomial_PolynomialAttr : Polynomial_Attr<"Polynomial", "polynomial"> {
-  let summary = "An attribute containing a single-variable polynomial.";
+def Polynomial_IntPolynomialAttr : Polynomial_Attr<"IntPolynomial", "int_polynomial"> {
+  let summary = "An attribute containing a single-variable polynomial with integer coefficients.";
   let description = [{
-    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.
+    A polynomial attribute represents a single-variable polynomial with integer
+    coefficients, 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
@@ -76,13 +76,69 @@ def Polynomial_PolynomialAttr : Polynomial_Attr<"Polynomial", "polynomial"> {
     Example:
 
     ```mlir
-    #poly = #polynomial.polynomial<x**1024 + 1>
+    #poly = #polynomial.int_polynomial<x**1024 + 1>
     ```
   }];
-  let parameters = (ins "::mlir::polynomial::Polynomial":$polynomial);
+  let parameters = (ins "::mlir::polynomial::IntPolynomial":$polynomial);
   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 = [{
+    A polynomial attribute represents a single-variable polynomial with double
+    precision floating point coefficients.
+
+    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.float_polynomial<0.5 x**7 + 1.5>
+    ```
+  }];
+  let parameters = (ins "FloatPolynomial":$polynomial);
+  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 = [{
@@ -104,9 +160,9 @@ def Polynomial_RingAttr : Polynomial_Attr<"Ring", "ring"> {
     `x**1024 - 1`.
 
     ```mlir
-    #poly_mod = #polynomial.polynomial<-1 + x**1024>
+    #poly_mod = #polynomial.int_polynomial<-1 + x**1024>
     #ring = #polynomial.ring<coefficientType=i32,
-                             coefficientModulus=4294967291,
+                             coefficientModulus=4294967291:i32,
                              polynomialModulus=#poly_mod>
 
     %0 = ... : polynomial.polynomial<#ring>
@@ -123,19 +179,24 @@ def Polynomial_RingAttr : Polynomial_Attr<"Ring", "ring"> {
   let parameters = (ins
     "Type": $coefficientType,
     OptionalParameter<"::mlir::IntegerAttr">: $coefficientModulus,
-    OptionalParameter<"::mlir::polynomial::PolynomialAttr">: $polynomialModulus,
+    OptionalParameter<"::mlir::polynomial::IntPolynomialAttr">: $polynomialModulus,
     OptionalParameter<"::mlir::IntegerAttr">: $primitiveRoot
   );
-
+  let assemblyFormat = "`<` struct(params) `>`";
   let builders = [
-    AttrBuilder<
+    AttrBuilderWithInferredContext<
         (ins "::mlir::Type":$coefficientTy,
-             "::mlir::IntegerAttr":$coefficientModulusAttr,
-             "::mlir::polynomial::PolynomialAttr":$polynomialModulusAttr), [{
-      return $_get($_ctxt, coefficientTy, coefficientModulusAttr, polynomialModulusAttr, nullptr);
-    }]>
+              CArg<"::mlir::IntegerAttr", "nullptr"> :$coefficientModulusAttr,
+              CArg<"::mlir::polynomial::IntPolynomialAttr", "nullptr"> :$polynomialModulusAttr,
+              CArg<"::mlir::IntegerAttr", "nullptr"> :$primitiveRootAttr), [{
+      return $_get(
+        coefficientTy.getContext(),
+        coefficientTy,
+        coefficientModulusAttr,
+        polynomialModulusAttr,
+        primitiveRootAttr);
+    }]>,
   ];
-  let hasCustomAssemblyFormat = 1;
 }
 
 class Polynomial_Type<string name, string typeMnemonic>
@@ -149,7 +210,7 @@ def Polynomial_PolynomialType : Polynomial_Type<"Polynomial", "polynomial"> {
     A type for polynomials in a polynomial quotient ring.
   }];
   let parameters = (ins Polynomial_RingAttr:$ring);
-  let assemblyFormat = "`<` $ring `>`";
+  let assemblyFormat = "`<` struct(params) `>`";
 }
 
 def PolynomialLike: TypeOrContainer<Polynomial_PolynomialType, "polynomial-like">;
@@ -187,13 +248,14 @@ def Polynomial_AddOp : Polynomial_BinaryOp<"add", [Commutative]> {
 
     ```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>
+    #poly = #polynomial.int_polynomial<x**1024 - 1>
+    #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536:i32, polynomialModulus=#poly>
+    %0 = polynomial.constant #polynomial.int_polynomial<1 + x**2> : !polynomial.polynomial<#ring>
+    %1 = polynomial.constant #polynomial.int_polynomial<x**5 - x + 1> : !polynomial.polynomial<#ring>
     %2 = polynomial.add %0, %1 : !polynomial.polynomial<#ring>
     ```
   }];
+  let hasFolder = 1;
 }
 
 def Polynomial_SubOp : Polynomial_BinaryOp<"sub"> {
@@ -211,13 +273,14 @@ def Polynomial_SubOp : Polynomial_BinaryOp<"sub"> {
 
     ```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>
+    #poly = #polynomial.int_polynomial<x**1024 - 1>
+    #ring = #polynomial.ring<coefficientType=i32, coefficientModulus=65536:i32, polynomialModulus=#poly>
+    %0 = polynomial.constant #polynomial.int_polynomial<1 + x**2> : !polynomial.polynomial<#ring>
+    %1 = polynomial.constant #polynomial.int_polynomial<x**5 - x + 1> : !polynomial.polynomial<#ring>
     %2 = polynomial.sub %0, ...
[truncated]

Copy link

github-actions bot commented May 9, 2024

⚠️ We detected that you are using a GitHub private e-mail address to contribute to the repo.
Please turn off Keep my email addresses private setting in your account.
See LLVM Discourse for more information.

Copy link

github-actions bot commented May 9, 2024

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:
git-clang-format --diff 19008d32182ebbe421aaa222ee8af5c3e134e550 4c6c10dea5dfb4cdf0d6ca875b11b23bfadbfc36 -- mlir/unittests/Dialect/Polynomial/PolynomialMathTest.cpp mlir/include/mlir/Dialect/Polynomial/IR/Polynomial.h mlir/lib/Dialect/Polynomial/IR/PolynomialDialect.cpp mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp
View the diff from clang-format here.
diff --git a/mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp b/mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp
index 8cbc3b4615..36ae3c151c 100644
--- a/mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp
+++ b/mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp
@@ -42,11 +42,12 @@ OpFoldResult AddOp::fold(AddOp::FoldAdaptor adaptor) {
 
   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()));
-      });
+        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) {

@j2kun j2kun force-pushed the poly-constant-propagation branch from 04614be to 3cc066c Compare May 13, 2024 23:12
@j2kun j2kun force-pushed the poly-constant-propagation branch from 8f9045a to 4c6c10d Compare May 14, 2024 19:43
j2kun added a commit that referenced this pull request May 17, 2024
A change extracted from #91655,
where I'm still trying to get the attributes working for elementwise
constant folding of polynomial ops. This piece is self-contained.

- use CRTP for base classes
- Add unit test

---------

Co-authored-by: Jeremy Kun <j2kun@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bazel "Peripheral" support tier build system: utils/bazel mlir
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants