[draft]: Add initial polynomial dialect

mlir/include/mlir/Dialect/CMakeLists.txt

@@ -28,6 +28,7 @@ add_subdirectory(OpenACCMPCommon)
 add_subdirectory(OpenMP)
 add_subdirectory(PDL)
 add_subdirectory(PDLInterp)
+add_subdirectory(Polynomial)
 add_subdirectory(Quant)
 add_subdirectory(SCF)
 add_subdirectory(Shape)

mlir/include/mlir/Dialect/Polynomial/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

mlir/include/mlir/Dialect/Polynomial/IR/CMakeLists.txt

@@ -0,0 +1,10 @@
+add_mlir_dialect(Polynomial polynomial)
+add_mlir_doc(PolynomialDialect PolynomialDialect Polynomial/ -gen-dialect-doc)
+add_mlir_doc(PolynomialOps PolynomialOps Polynomial/ -gen-op-doc)
+add_mlir_doc(PolynomialAttributes PolynomialAttributes Dialects/ -gen-attrdef-doc)
+add_mlir_doc(PolynomialTypes PolynomialTypes Dialects/ -gen-typedef-doc)
+
+set(LLVM_TARGET_DEFINITIONS Polynomial.td)
+mlir_tablegen(PolynomialAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=polynomial)
+mlir_tablegen(PolynomialAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=polynomial)
+add_public_tablegen_target(MLIRPolynomialAttributesIncGen)

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

@@ -0,0 +1,130 @@
+//===- Polynomial.h - A data class for polynomials --------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIAL_H_
+#define MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIAL_H_
+
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+
+class MLIRContext;
+
+namespace polynomial {
+
+/// This restricts statically defined polynomials to have at most 64-bit
+/// coefficients. This may be relaxed in the future, but it seems unlikely one
+/// 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 {
+public:
+  Monomial(int64_t coeff, uint64_t expo)
+      : coefficient(apintBitWidth, coeff), exponent(apintBitWidth, expo) {}
+
+  Monomial(const APInt &coeff, const APInt &expo)
+      : coefficient(coeff), exponent(expo) {}
+
+  Monomial() : coefficient(apintBitWidth, 0), exponent(apintBitWidth, 0) {}
+
+  bool operator==(const Monomial &other) const {
+    return other.coefficient == coefficient && other.exponent == exponent;
+  }
+  bool operator!=(const Monomial &other) const {
+    return other.coefficient != coefficient || other.exponent != exponent;
+  }
+
+  /// Monomials are ordered by exponent.
+  bool operator<(const Monomial &other) const {
+    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:
+  APInt coefficient;
+
+  // Always unsigned
+  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 {
+public:
+  Polynomial() = delete;
+
+  explicit Polynomial(ArrayRef<Monomial> terms) : terms(terms){};
+
+  // Returns a Polynomial from a list of monomials.
+  // Fails if two monomials have the same exponent.
+  static FailureOr<Polynomial> fromMonomials(ArrayRef<Monomial> monomials);
+
+  /// 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);
+
+  explicit operator bool() const { return !terms.empty(); }
+  bool operator==(const Polynomial &other) const {
+    return other.terms == terms;
+  }
+  bool operator!=(const Polynomial &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;
+  void dump() const;
+
+  // Prints polynomial so that it can be used as a valid identifier
+  std::string toIdentifier() const;
+
+  unsigned getDegree() const;
+
+  friend ::llvm::hash_code hash_value(const Polynomial &arg);
+
+private:
+  // The monomial terms for this polynomial.
+  SmallVector<Monomial> terms;
+};
+
+// Make Polynomial hashable.
+inline ::llvm::hash_code hash_value(const Polynomial &arg) {
+  return ::llvm::hash_combine_range(arg.terms.begin(), arg.terms.end());
+}
+
+inline ::llvm::hash_code hash_value(const Monomial &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) {
+  polynomial.print(os);
+  return os;
+}
+
+} // namespace polynomial
+} // namespace mlir
+
+#endif // MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIAL_H_

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

@@ -0,0 +1,153 @@
+//===- PolynomialOps.td - Polynomial dialect ---------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef POLYNOMIAL_OPS
+#define POLYNOMIAL_OPS
+
+include "mlir/IR/BuiltinAttributes.td"
+include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+def Polynomial_Dialect : Dialect {
+  let name = "polynomial";
+  let cppNamespace = "::mlir::polynomial";
+  let description = [{
+    The Polynomial dialect defines single-variable polynomial types and
+    operations.
+
+    The simplest use of `polynomial` is to represent mathematical operations in
+    a polynomial ring `R[x]`, where `R` is another MLIR type like `i32`.
+
+    More generally, this dialect supports representing polynomial operations in a
+    quotient ring `R[X]/(f(x))` for some statically fixed polynomial `f(x)`.
+    Two polyomials `p(x), q(x)` are considered equal in this ring if they have the
+    same remainder when dividing by `f(x)`. When a modulus is given, ring operations
+    are performed with reductions modulo `f(x)` and relative to the coefficient ring
+    `R`.
+
+    Examples:
+
+    ```mlir
+    // A constant polynomial in a ring with i32 coefficients and no polynomial modulus
+    #ring = #polynomial.ring<ctype=i32>
+    %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<ctype=i32, ideal=#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<ctype=i32, cmod=17, ideal=#modulus>
+    %a = polynomial.constant <1 + x**2 - 3x**3> : polynomial.polynomial<#ring>
+    ```
+  }];
+
+  let useDefaultTypePrinterParser = 1;
+  let useDefaultAttributePrinterParser = 1;
+}
+
+class Polynomial_Attr<string name, string attrMnemonic, list<Trait> traits = []>
+    : AttrDef<Polynomial_Dialect, name, traits> {
+  let mnemonic = attrMnemonic;
+}
+
+def Polynomial_PolynomialAttr : Polynomial_Attr<"Polynomial", "polynomial"> {
+  let summary = "An attribute containing a single-variable polynomial.";
+  let description = [{
+     #poly = #polynomial.poly<x**1024 + 1>
+  }];
+  let parameters = (ins "Polynomial":$polynomial);
+  let hasCustomAssemblyFormat = 1;
+}
+
+def Polynomial_RingAttr : Polynomial_Attr<"Ring", "ring"> {
+  let summary = "An attribute specifying a polynomial ring.";
+  let description = [{
+    A ring describes the domain in which polynomial arithmetic occurs. The ring
+    attribute in `polynomial` represents the more specific case of polynomials
+    with a single indeterminate; whose coefficients can be represented by
+    another MLIR type (`coefficientType`); and, if the coefficient type is
+    integral, whose coefficients are taken modulo some statically known modulus
+    (`coefficientModulus`).
+
+    Additionally, a polynomial ring can specify an _ideal_, 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
+    via a single polynomial, which we call `polynomialModulus`.
+
+    An expressive example is polynomials with i32 coefficients, whose
+    coefficients are taken modulo `2**32 - 5`, with a polynomial modulus of
+    `x**1024 - 1`.
+
+    ```mlir
+    #poly_mod = #polynomial.polynomial<-1 + x**1024>
+    #ring = #polynomial.ring<coefficientType=i32,
+                             coefficientModulus=4294967291,
+                             polynomialModulus=#poly_mod>
+
+    %0 = ... : polynomial.polynomial<#ring>
+    ```
+
+    In this case, the value of a polynomial is always "converted" to a
+    canonical form by applying repeated reductions by setting `x**1024 = 1`
+    and simplifying.
+
+    The coefficient and polynomial modulus parameters are optional, and the
+    coefficient modulus is only allowed if the coefficient type is integral.
+  }];
+
+  let parameters = (ins
+    "Type": $coefficientType,
+    OptionalParameter<"IntegerAttr">: $coefficientModulus,
+    OptionalParameter<"PolynomialAttr">: $polynomialModulus
+  );
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+class Polynomial_Type<string name, string typeMnemonic>
+    : TypeDef<Polynomial_Dialect, name> {
+  let mnemonic = typeMnemonic;
+}
+
+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 `>`";
+}
+
+class Polynomial_Op<string mnemonic, list<Trait> traits = []> :
+    Op<Polynomial_Dialect, mnemonic, traits # [Pure]>;
+
+class Polynomial_UnaryOp<string mnemonic, list<Trait> traits = []> :
+    Polynomial_Op<mnemonic, traits # [SameOperandsAndResultType]> {
+  let arguments = (ins Polynomial_PolynomialType:$operand);
+  let results = (outs Polynomial_PolynomialType:$result);
+
+  let assemblyFormat = "$operand attr-dict `:` qualified(type($result))";
+}
+
+class Polynomial_BinaryOp<string mnemonic, list<Trait> traits = []> :
+    Polynomial_Op<mnemonic, traits # [SameOperandsAndResultType]> {
+  let arguments = (ins Polynomial_PolynomialType:$lhs, Polynomial_PolynomialType:$rhs);
+  let results = (outs Polynomial_PolynomialType:$result);
+
+  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` qualified(type($result))";
+}
+
+#endif // POLYNOMIAL_OPS

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

@@ -0,0 +1,17 @@
+//===- PolynomialAttributes.h - polynomial dialect attributes ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALATTRIBUTES_H_
+#define MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALATTRIBUTES_H_
+
+#include "Polynomial.h"
+#include "PolynomialDialect.h"
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.h.inc"
+
+#endif // MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALATTRIBUTES_H_

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

@@ -0,0 +1,19 @@
+//===- PolynomialDialect.h - The Polynomial dialect -------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALDIALECT_H_
+#define MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALDIALECT_H_
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/DialectImplementation.h"
+
+// Generated headers (block clang-format from messing up order)
+#include "mlir/Dialect/Polynomial/IR/PolynomialDialect.h.inc"
+
+#endif // MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALDIALECT_H_

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

@@ -0,0 +1,21 @@
+//===- PolynomialOps.h - Ops for the Polynomial dialect ---------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALOPS_H_
+#define MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALOPS_H_
+
+#include "PolynomialDialect.h"
+#include "PolynomialTypes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Polynomial/IR/Polynomial.h.inc"
+
+#endif // MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALOPS_H_

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

@@ -0,0 +1,17 @@
+//===- PolynomialTypes.h - Types for the Polynomial dialect -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALTYPES_H_
+#define MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALTYPES_H_
+
+#include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"
+#include "mlir/Dialect/Polynomial/IR/PolynomialDialect.h"
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/Polynomial/IR/PolynomialTypes.h.inc"
+
+#endif // MLIR_INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALTYPES_H_

mlir/include/mlir/InitAllDialects.h

@@ -61,6 +61,7 @@
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
+#include "mlir/Dialect/Polynomial/IR/PolynomialDialect.h"
 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h"
@@ -131,6 +132,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
                   omp::OpenMPDialect,
                   pdl::PDLDialect,
                   pdl_interp::PDLInterpDialect,
+                  polynomial::PolynomialDialect,
                   quant::QuantizationDialect,
                   ROCDL::ROCDLDialect,
                   scf::SCFDialect,