[Mlir-commits] [mlir] mlir/Presburger: contribute a free-standing parser (PR #94916)

[Kunwar Grover]

================
@@ -0,0 +1,345 @@
+//===- ParseStructs.h - Presburger Parse Structrures ------------*- 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_ANALYSIS_PRESBURGER_PARSER_PARSESTRUCTS_H
+#define MLIR_ANALYSIS_PRESBURGER_PARSER_PARSESTRUCTS_H
+
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
+#include "mlir/Analysis/Presburger/PresburgerSpace.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include <cstdint>
+
+namespace mlir::presburger {
+using llvm::ArrayRef;
+using llvm::SmallVector;
+using llvm::SmallVectorImpl;
+
+/// This structure is central to the parser and flattener, and holds the number
+/// of dimensions, symbols, locals, and the constant term.
+struct ParseInfo {
+  unsigned numDims = 0;
+  unsigned numSymbols = 0;
+  unsigned numExprs = 0;
+  unsigned numDivs = 0;
+
+  constexpr unsigned getDimStartIdx() const { return 0; }
+  constexpr unsigned getSymbolStartIdx() const { return numDims; }
+  constexpr unsigned getLocalVarStartIdx() const {
+    return numDims + numSymbols;
+  }
+  constexpr unsigned getNumCols() const {
+    return numDims + numSymbols + numDivs + 1;
+  }
+  constexpr unsigned getConstantIdx() const { return getNumCols() - 1; }
+
+  constexpr bool isDimIdx(unsigned i) const { return i < getSymbolStartIdx(); }
+  constexpr bool isSymbolIdx(unsigned i) const {
+    return i >= getSymbolStartIdx() && i < getLocalVarStartIdx();
+  }
+  constexpr bool isLocalVarIdx(unsigned i) const {
+    return i >= getLocalVarStartIdx() && i < getConstantIdx();
+  }
+  constexpr bool isConstantIdx(unsigned i) const {
+    return i == getConstantIdx();
+  }
+};
+
+/// Helper for storing coefficients in canonical form: dims followed by symbols,
+/// followed by locals, and finally the constant term.
+///
+/// (x, y)[a, b]: y * 91 + x + 3 * a + 7
+/// coefficients: [1, 91, 3, 0, 7]
+struct CoefficientVector {
+  ParseInfo info;
+  SmallVector<int64_t, 8> coefficients;
+
+  CoefficientVector(const ParseInfo &info, int64_t c = 0) : info(info) {
+    coefficients.resize(info.getNumCols());
+    coefficients[info.getConstantIdx()] = c;
+  }
+
+  // Copyable and movable
+  CoefficientVector(const CoefficientVector &o) = default;
+  CoefficientVector &operator=(const CoefficientVector &o) = default;
+  CoefficientVector(CoefficientVector &&o)
+      : info(o.info), coefficients(std::move(o.coefficients)) {
+    o.coefficients.clear();
+  }
+
+  ArrayRef<int64_t> getCoefficients() const { return coefficients; }
+  int64_t getConstant() const { return coefficients[info.getConstantIdx()]; }
+  size_t size() const { return coefficients.size(); }
+  operator ArrayRef<int64_t>() const { return coefficients; }
+  void resize(size_t size) { coefficients.resize(size); }
+  operator bool() const {
+    return any_of(coefficients, [](int64_t c) { return c; });
+  }
+  int64_t &operator[](unsigned i) {
+    assert(i < coefficients.size());
+    return coefficients[i];
+  }
+  int64_t &back() { return coefficients.back(); }
+  int64_t back() const { return coefficients.back(); }
+  void clear() {
+    for_each(coefficients, [](auto &coeff) { coeff = 0; });
+  }
+
+  CoefficientVector &operator+=(const CoefficientVector &l) {
+    coefficients.resize(l.size());
+    for (auto [idx, c] : enumerate(l.getCoefficients()))
+      coefficients[idx] += c;
+    return *this;
+  }
+  CoefficientVector &operator*=(int64_t c) {
+    for_each(coefficients, [c](auto &coeff) { coeff *= c; });
+    return *this;
+  }
+  CoefficientVector &operator/=(int64_t c) {
+    assert(c && "Division by zero");
+    for_each(coefficients, [c](auto &coeff) { coeff /= c; });
+    return *this;
+  }
+
+  CoefficientVector operator+(const CoefficientVector &l) const {
+    CoefficientVector ret(*this);
+    return ret += l;
+  }
+  CoefficientVector operator*(int64_t c) const {
+    CoefficientVector ret(*this);
+    return ret *= c;
+  }
+  CoefficientVector operator/(int64_t c) const {
+    CoefficientVector ret(*this);
+    return ret /= c;
+  }
+
+  bool isConstant() const {
+    return all_of(drop_end(coefficients), [](int64_t c) { return !c; });
+  }
+  CoefficientVector getPadded(size_t newSize) const {
+    assert(newSize >= size() &&
+           "Padding size should be greater than expr size");
+    CoefficientVector ret(info);
+    ret.resize(newSize);
+
+    // Start constructing the result by taking the dims and symbols of the
+    // coefficients.
+    for (const auto &[col, coeff] : enumerate(drop_end(coefficients)))
+      ret[col] = coeff;
+
+    // Put the constant at the end.
+    ret.back() = back();
+    return ret;
+  }
+
+  uint64_t factorMulFromLinearTerm() const {
+    uint64_t gcd = 1;
+    for (int64_t val : coefficients)
+      gcd = std::gcd(gcd, std::abs(val));
+    return gcd;
+  }
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  bool hasMultipleCoefficients() const {
+    return count_if(coefficients, [](auto &coeff) { return coeff; }) > 1;
+  }
+  LLVM_DUMP_METHOD void dump() const;
+#endif
+};
+
+enum class DimOrSymbolKind {
+  DimId,
+  Symbol,
+};
+
+using DimOrSymbolExpr = std::pair<DimOrSymbolKind, unsigned>;
+enum class DivKind { FloorDiv, Mod };
+
+/// Represents a pure Affine expression. Linear expressions are represented with
+/// divisor = 1, and no nestedDivTerms.
+///
+/// 3 - a * (3 + (3 - x div 3) div 4 + y div 7) div 4
+/// ^ linearDivident = 3, mulFactor = 1, divisor = 1
+///   ^ nest: 1, mulFactor: -a
+///         ^ nest: 1, linearDividend
+///              ^ nest: 2, linearDividend
+///                     ^ nest: 3
+///                          ^ nest: 2
+///                                      ^ nest: 2
+///                               nest: 1, divisor ^
+///
+/// Where div = floordiv|mod; ceildiv is pre-reduced
+struct PureAffineExprImpl {
+  ParseInfo info;
+  DivKind kind = DivKind::FloorDiv;
+  using PureAffineExpr = std::unique_ptr<PureAffineExprImpl>;
+
+  int64_t mulFactor = 1;
+  CoefficientVector linearDividend;
+  int64_t divisor = 1;
+  SmallVector<PureAffineExpr, 4> nestedDivTerms;
----------------
Groverkss wrote:

I don't understand why you need a AST form of an affine expression. Why can't these operations be done directly on Presburger constructs?

https://github.com/llvm/llvm-project/pull/94916