polynomial: add dialect shell, attrs, and types

Author: j2kun

mlir/include/mlir/Dialect/CMakeLists.txt

@@ -27,6 +27,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,19 @@
+set(LLVM_TARGET_DEFINITIONS PolynomialOps.td)
+mlir_tablegen(PolynomialDialect.cpp.inc -gen-dialect-defs -dialect=polynomial)
+mlir_tablegen(PolynomialDialect.h.inc -gen-dialect-decls -dialect=polynomial)
+add_public_tablegen_target(MLIRPolynomialDialectIncGen)
+
+mlir_tablegen(PolynomialAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=polynomial)
+mlir_tablegen(PolynomialAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=polynomial)
+mlir_tablegen(PolynomialOps.cpp.inc -gen-op-defs)
+mlir_tablegen(PolynomialOps.h.inc -gen-op-decls)
+mlir_tablegen(PolynomialTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=polynomial)
+mlir_tablegen(PolynomialTypes.h.inc -gen-typedef-decls -typedefs-dialect=polynomial)
+add_public_tablegen_target(MLIRPolynomialAttributesIncGen)
+add_public_tablegen_target(MLIRPolynomialOpsIncGen)
+add_public_tablegen_target(MLIRPolynomialTypesIncGen)
+add_dependencies(mlir-headers MLIRPolynomialOpsIncGen)
+
+add_mlir_doc(PolynoialOps PolynoialOps Dialects/ -gen-dialect-doc -dialect polynomial)
+add_mlir_doc(PolynomialAttributes PolynomialAttributes Dialects/ -gen-attrdef-doc)
+add_mlir_doc(PolynomialTypes PolynomialTypes Dialects/ -gen-typedef-doc)

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

@@ -0,0 +1,157 @@
+//===- Polynomial.h - A storage class for polynomial types --------*- 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 INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIAL_H_
+#define INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIAL_H_
+
+#include <utility>
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/Hashing.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;
+
+namespace detail {
+struct PolynomialStorage;
+} // namespace detail
+
+class Monomial {
+public:
+  Monomial(int64_t coeff, uint64_t expo)
+      : coefficient(apintBitWidth, coeff), exponent(apintBitWidth, expo) {}
+
+  Monomial(APInt coeff, APInt expo)
+      : coefficient(std::move(coeff)), exponent(std::move(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(Monomial arg);
+
+public:
+  APInt coefficient;
+
+  // Always unsigned
+  APInt exponent;
+};
+
+/// A single-variable polynomial with integer coefficients. Polynomials are
+/// immutable and uniqued.
+///
+/// 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:
+  using ImplType = detail::PolynomialStorage;
+
+  constexpr Polynomial() = default;
+  explicit Polynomial(ImplType *terms) : terms(terms) {}
+
+  static Polynomial fromMonomials(ArrayRef<Monomial> monomials,
+                                  MLIRContext *context);
+  /// Returns a polynomial with coefficients given by `coeffs`
+  static Polynomial fromCoefficients(ArrayRef<int64_t> coeffs,
+                                     MLIRContext *context);
+
+  MLIRContext *getContext() const;
+
+  explicit operator bool() const { return terms != nullptr; }
+  bool operator==(Polynomial other) const { return other.terms == terms; }
+  bool operator!=(Polynomial other) const { return !(other.terms == terms); }
+
+  // Prints polynomial to 'os'.
+  void print(raw_ostream &os) const;
+  void print(raw_ostream &os, const std::string &separator,
+             const std::string &exponentiation) const;
+  void dump() const;
+
+  // Prints polynomial so that it can be used as a valid identifier
+  std::string toIdentifier() const;
+
+  // A polynomial's terms are canonically stored in order of increasing degree.
+  ArrayRef<Monomial> getTerms() const;
+
+  unsigned getDegree() const;
+
+  friend ::llvm::hash_code hash_value(Polynomial arg);
+
+private:
+  ImplType *terms{nullptr};
+};
+
+// Make Polynomial hashable.
+inline ::llvm::hash_code hash_value(Polynomial arg) {
+  return ::llvm::hash_value(arg.terms);
+}
+
+inline ::llvm::hash_code hash_value(Monomial arg) {
+  return ::llvm::hash_value(arg.coefficient) ^ ::llvm::hash_value(arg.exponent);
+}
+
+inline raw_ostream &operator<<(raw_ostream &os, Polynomial polynomial) {
+  polynomial.print(os);
+  return os;
+}
+
+} // namespace polynomial
+} // namespace mlir
+
+namespace llvm {
+
+// Polynomials hash just like pointers
+template <>
+struct DenseMapInfo<mlir::polynomial::Polynomial> {
+  static mlir::polynomial::Polynomial getEmptyKey() {
+    auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
+    return mlir::polynomial::Polynomial(
+        static_cast<mlir::polynomial::Polynomial::ImplType *>(pointer));
+  }
+  static mlir::polynomial::Polynomial getTombstoneKey() {
+    auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
+    return mlir::polynomial::Polynomial(
+        static_cast<mlir::polynomial::Polynomial::ImplType *>(pointer));
+  }
+  static unsigned getHashValue(mlir::polynomial::Polynomial val) {
+    return mlir::polynomial::hash_value(val);
+  }
+  static bool isEqual(mlir::polynomial::Polynomial lhs,
+                      mlir::polynomial::Polynomial rhs) {
+    return lhs == rhs;
+  }
+};
+
+} // namespace llvm
+
+#endif // INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIAL_H_

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

@@ -0,0 +1,18 @@
+//===- PolynomialAttributes.h - Attributes 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 INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALATTRIBUTES_H_
+#define 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 // INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALATTRIBUTES_H_

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

@@ -0,0 +1,81 @@
+//===- PolynomialAttributes.td - Attribute definitions for the 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_ATTRIBUTES
+#define POLYNOMIAL_ATTRIBUTES
+
+include "PolynomialDialect.td"
+include "mlir/IR/BuiltinAttributes.td"
+include "mlir/IR/OpBase.td"
+
+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 builders = [
+    AttrBuilderWithInferredContext<(ins "Polynomial":$polynomial), [{
+      return $_get(polynomial.getContext(), polynomial);
+    }]>
+  ];
+
+  let skipDefaultBuilders = 1;
+  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 (`ctype`); and, if the coefficient type is integral,
+    whose coefficients are taken modulo some statically known modulus (`cmod`).
+
+    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`.
+
+    ```
+    #poly_mod = #polynomial.polynomial<-1 + x**1024>
+    #ring = #polynomial.ring<ctype=i32, cmod=4294967291, ideal=#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
+    Builtin_TypeAttr: $coefficientType,
+    OptionalParameter<"std::optional<IntegerAttr>">: $coefficientModulus,
+    OptionalParameter<"std::optional<PolynomialAttr>">: $polynomialModulus
+  );
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+#endif // POLYNOMIAL_ATTRIBUTES

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 INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALDIALECT_H_
+#define 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 // INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALDIALECT_H_

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

@@ -0,0 +1,54 @@
+//===- PolynomialDialect.td - Dialect definition for the 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_DIALECT
+#define POLYNOMIAL_DIALECT
+
+include "mlir/IR/OpBase.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;
+}
+
+#endif // POLYNOMIAL_DIALECT

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 INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALOPS_H_
+#define 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/PolynomialOps.h.inc"
+
+#endif // INCLUDE_MLIR_DIALECT_POLYNOMIAL_IR_POLYNOMIALOPS_H_