[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:

In my mind, I wanted an isl-like syntax for presburger constructs, which allows building divs using a flat form. But I understand, that it is not an easy to do change. I'm okay with accepting an AST for at start, with a TODO to move to an ISL-like syntax which allows directly parsing Presburger structures.

Re: inefficiency, I would prefer a simpler but inefficient parser. The parser is only used for tests and playing around with Presburger constructs, not for critical code. Please try to reuse things and keep the Parser simple.

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