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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Sun Jun 9 14:20:37 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-presburger

Author: Ramkumar Ramachandra (artagnon)

<details>
<summary>Changes</summary>

The Presburger library is already quite independent of MLIR, with the exception of the MLIR Support library. There is, however, one major exception: the test suite for the library depends on the core MLIR parser. To free it of this dependency, extract the parts of the core MLIR parser that are applicable to the Presburger test suite, author custom parsing data structures, and adapt the new parser to parse into these structures.

This patch is part of a project to move the Presburger library into LLVM.

---

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


29 Files Affected:

- (renamed) mlir/include/mlir/Analysis/Presburger/Parser.h (+10-33) 
- (modified) mlir/lib/Analysis/Presburger/CMakeLists.txt (+2) 
- (added) mlir/lib/Analysis/Presburger/Parser/CMakeLists.txt (+6) 
- (added) mlir/lib/Analysis/Presburger/Parser/Flattener.cpp (+410) 
- (added) mlir/lib/Analysis/Presburger/Parser/Flattener.h (+248) 
- (added) mlir/lib/Analysis/Presburger/Parser/Lexer.cpp (+165) 
- (added) mlir/lib/Analysis/Presburger/Parser/Lexer.h (+62) 
- (added) mlir/lib/Analysis/Presburger/Parser/ParseStructs.cpp (+271) 
- (added) mlir/lib/Analysis/Presburger/Parser/ParseStructs.h (+300) 
- (added) mlir/lib/Analysis/Presburger/Parser/ParserImpl.cpp (+825) 
- (added) mlir/lib/Analysis/Presburger/Parser/ParserImpl.h (+242) 
- (added) mlir/lib/Analysis/Presburger/Parser/ParserState.h (+48) 
- (added) mlir/lib/Analysis/Presburger/Parser/Token.cpp (+68) 
- (added) mlir/lib/Analysis/Presburger/Parser/Token.h (+95) 
- (added) mlir/lib/Analysis/Presburger/Parser/TokenKinds.def (+73) 
- (modified) mlir/unittests/Analysis/Presburger/BarvinokTest.cpp (+2-2) 
- (modified) mlir/unittests/Analysis/Presburger/CMakeLists.txt (+1-3) 
- (modified) mlir/unittests/Analysis/Presburger/FractionTest.cpp (-1) 
- (modified) mlir/unittests/Analysis/Presburger/GeneratingFunctionTest.cpp (+1-1) 
- (modified) mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp (+1-1) 
- (modified) mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp (+1-1) 
- (modified) mlir/unittests/Analysis/Presburger/MatrixTest.cpp (+2-1) 
- (modified) mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp (+1-4) 
- (modified) mlir/unittests/Analysis/Presburger/ParserTest.cpp (+1-1) 
- (modified) mlir/unittests/Analysis/Presburger/PresburgerRelationTest.cpp (+1-1) 
- (modified) mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp (+1-2) 
- (modified) mlir/unittests/Analysis/Presburger/QuasiPolynomialTest.cpp (+2-2) 
- (modified) mlir/unittests/Analysis/Presburger/SimplexTest.cpp (+1-4) 
- (modified) mlir/unittests/Analysis/Presburger/Utils.h (-5) 


``````````diff
diff --git a/mlir/unittests/Analysis/Presburger/Parser.h b/mlir/include/mlir/Analysis/Presburger/Parser.h
similarity index 67%
rename from mlir/unittests/Analysis/Presburger/Parser.h
rename to mlir/include/mlir/Analysis/Presburger/Parser.h
index 75842fb054e2b..67f6766fe5bd2 100644
--- a/mlir/unittests/Analysis/Presburger/Parser.h
+++ b/mlir/include/mlir/Analysis/Presburger/Parser.h
@@ -11,32 +11,25 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_UNITTESTS_ANALYSIS_PRESBURGER_PARSER_H
-#define MLIR_UNITTESTS_ANALYSIS_PRESBURGER_PARSER_H
+#ifndef MLIR_ANALYSIS_PRESBURGER_PARSER_H
+#define MLIR_ANALYSIS_PRESBURGER_PARSER_H
 
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/PWMAFunction.h"
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
-#include "mlir/AsmParser/AsmParser.h"
-#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
-#include "mlir/IR/AffineExpr.h"
-#include "mlir/IR/AffineMap.h"
-#include "mlir/IR/IntegerSet.h"
 
 namespace mlir {
 namespace presburger {
+/// Parses an IntegerPolyhedron from a StringRef.
+IntegerPolyhedron parseIntegerPolyhedron(StringRef str);
 
-/// Parses an IntegerPolyhedron from a StringRef. It is expected that the string
-/// represents a valid IntegerSet.
-inline IntegerPolyhedron parseIntegerPolyhedron(StringRef str) {
-  MLIRContext context(MLIRContext::Threading::DISABLED);
-  return affine::FlatAffineValueConstraints(parseIntegerSet(str, &context));
-}
+/// Parses a MultiAffineFunction from a StringRef.
+MultiAffineFunction parseMultiAffineFunction(StringRef str);
 
 /// Parse a list of StringRefs to IntegerRelation and combine them into a
-/// PresburgerSet by using the union operation. It is expected that the strings
-/// are all valid IntegerSet representation and that all of them have compatible
-/// spaces.
+/// PresburgerSet by using the union operation. It is expected that the
+/// strings are all valid IntegerSet representation and that all of them have
+/// compatible spaces.
 inline PresburgerSet parsePresburgerSet(ArrayRef<StringRef> strs) {
   assert(!strs.empty() && "strs should not be empty");
 
@@ -47,25 +40,10 @@ inline PresburgerSet parsePresburgerSet(ArrayRef<StringRef> strs) {
   return result;
 }
 
-inline MultiAffineFunction parseMultiAffineFunction(StringRef str) {
-  MLIRContext context(MLIRContext::Threading::DISABLED);
-
-  // TODO: Add default constructor for MultiAffineFunction.
-  MultiAffineFunction multiAff(PresburgerSpace::getRelationSpace(),
-                               IntMatrix(0, 1));
-  if (getMultiAffineFunctionFromMap(parseAffineMap(str, &context), multiAff)
-          .failed())
-    llvm_unreachable(
-        "Failed to parse MultiAffineFunction because of semi-affinity");
-  return multiAff;
-}
-
 inline PWMAFunction
 parsePWMAF(ArrayRef<std::pair<StringRef, StringRef>> pieces) {
   assert(!pieces.empty() && "At least one piece should be present.");
 
-  MLIRContext context(MLIRContext::Threading::DISABLED);
-
   IntegerPolyhedron initDomain = parseIntegerPolyhedron(pieces[0].first);
   MultiAffineFunction initMultiAff = parseMultiAffineFunction(pieces[0].second);
 
@@ -100,8 +78,7 @@ parsePresburgerRelationFromPresburgerSet(ArrayRef<StringRef> strs,
   result.convertVarKind(VarKind::SetDim, 0, numDomain, VarKind::Domain, 0);
   return result;
 }
-
 } // namespace presburger
 } // namespace mlir
 
-#endif // MLIR_UNITTESTS_ANALYSIS_PRESBURGER_PARSER_H
+#endif // MLIR_ANALYSIS_PRESBURGER_PARSER_H
diff --git a/mlir/lib/Analysis/Presburger/CMakeLists.txt b/mlir/lib/Analysis/Presburger/CMakeLists.txt
index 83d0514c9e7d1..4a89e3929a709 100644
--- a/mlir/lib/Analysis/Presburger/CMakeLists.txt
+++ b/mlir/lib/Analysis/Presburger/CMakeLists.txt
@@ -1,3 +1,5 @@
+add_subdirectory(Parser)
+
 add_mlir_library(MLIRPresburger
   Barvinok.cpp
   IntegerRelation.cpp
diff --git a/mlir/lib/Analysis/Presburger/Parser/CMakeLists.txt b/mlir/lib/Analysis/Presburger/Parser/CMakeLists.txt
new file mode 100644
index 0000000000000..f708a5c8db949
--- /dev/null
+++ b/mlir/lib/Analysis/Presburger/Parser/CMakeLists.txt
@@ -0,0 +1,6 @@
+add_mlir_library(MLIRPresburgerParser
+  Flattener.cpp
+  Lexer.cpp
+  ParserImpl.cpp
+  ParseStructs.cpp
+  Token.cpp)
diff --git a/mlir/lib/Analysis/Presburger/Parser/Flattener.cpp b/mlir/lib/Analysis/Presburger/Parser/Flattener.cpp
new file mode 100644
index 0000000000000..97507b64b28e7
--- /dev/null
+++ b/mlir/lib/Analysis/Presburger/Parser/Flattener.cpp
@@ -0,0 +1,410 @@
+//===- Flattener.cpp - Presburger ParseStruct Flattener ---------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the Flattener class for flattening the parse tree
+// produced by the parser for the Presburger library.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Flattener.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+namespace presburger {
+namespace detail {
+using llvm::SmallVector;
+
+AffineExpr AffineExprFlattener::getAffineExprFromFlatForm(
+    ArrayRef<int64_t> flatExprs, unsigned numDims, unsigned numSymbols) {
+  assert(flatExprs.size() - numDims - numSymbols - 1 == localExprs.size() &&
+         "unexpected number of local expressions");
+
+  // Dimensions and symbols.
+  AffineExpr expr = std::make_unique<AffineConstantExpr>(0);
+  for (unsigned j = 0; j < getLocalVarStartIndex(); ++j) {
+    if (flatExprs[j] == 0)
+      continue;
+    if (j < numDims)
+      expr =
+          std::move(expr) + std::make_unique<AffineDimExpr>(j) * flatExprs[j];
+    else
+      expr = std::move(expr) +
+             std::make_unique<AffineSymbolExpr>(j - numDims) * flatExprs[j];
+  }
+
+  // Local identifiers.
+  for (unsigned j = getLocalVarStartIndex(); j < flatExprs.size() - 1; ++j) {
+    if (flatExprs[j] == 0)
+      continue;
+    // It is safe to move out of the localExprs vector, since no expr is used
+    // more than once.
+    AffineExpr term =
+        std::move(localExprs[j - getLocalVarStartIndex()]) * flatExprs[j];
+    expr = std::move(expr) + std::move(term);
+  }
+
+  // Constant term.
+  int64_t constTerm = flatExprs[flatExprs.size() - 1];
+  if (constTerm != 0)
+    return std::move(expr) + constTerm;
+  return expr;
+}
+
+// In pure affine t = expr * c, we multiply each coefficient of lhs with c.
+//
+// In case of semi affine multiplication expressions, t = expr * symbolic_expr,
+// introduce a local variable p (= expr * symbolic_expr), and the affine
+// expression expr * symbolic_expr is added to `localExprs`.
+LogicalResult AffineExprFlattener::visitMulExpr(const AffineBinOpExpr &expr) {
+  assert(operandExprStack.size() >= 2);
+  SmallVector<int64_t, 8> rhs = operandExprStack.back();
+  operandExprStack.pop_back();
+  SmallVector<int64_t, 8> &lhs = operandExprStack.back();
+
+  // Flatten semi-affine multiplication expressions by introducing a local
+  // variable in place of the product; the affine expression
+  // corresponding to the quantifier is added to `localExprs`.
+  if (!isa<AffineConstantExpr>(expr.getRHS())) {
+    AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols);
+    AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols);
+    addLocalVariableSemiAffine(std::move(a) * std::move(b), lhs, lhs.size());
+    return success();
+  }
+
+  // Get the RHS constant.
+  int64_t rhsConst = rhs[getConstantIndex()];
+  for (int64_t &lhsElt : lhs)
+    lhsElt *= rhsConst;
+
+  return success();
+}
+
+LogicalResult AffineExprFlattener::visitAddExpr(const AffineBinOpExpr &expr) {
+  assert(operandExprStack.size() >= 2);
+  const auto &rhs = operandExprStack.back();
+  auto &lhs = operandExprStack[operandExprStack.size() - 2];
+  assert(lhs.size() == rhs.size());
+  // Update the LHS in place.
+  for (unsigned i = 0; i < rhs.size(); ++i)
+    lhs[i] += rhs[i];
+  // Pop off the RHS.
+  operandExprStack.pop_back();
+  return success();
+}
+
+//
+// t = expr mod c   <=>  t = expr - c*q and c*q <= expr <= c*q + c - 1
+//
+// A mod expression "expr mod c" is thus flattened by introducing a new local
+// variable q (= expr floordiv c), such that expr mod c is replaced with
+// 'expr - c * q' and c * q <= expr <= c * q + c - 1 are added to localVarCst.
+//
+// In case of semi-affine modulo expressions, t = expr mod symbolic_expr,
+// introduce a local variable m (= expr mod symbolic_expr), and the affine
+// expression expr mod symbolic_expr is added to `localExprs`.
+LogicalResult AffineExprFlattener::visitModExpr(const AffineBinOpExpr &expr) {
+  assert(operandExprStack.size() >= 2);
+
+  SmallVector<int64_t, 8> rhs = operandExprStack.back();
+  operandExprStack.pop_back();
+  SmallVector<int64_t, 8> &lhs = operandExprStack.back();
+
+  // Flatten semi affine modulo expressions by introducing a local
+  // variable in place of the modulo value, and the affine expression
+  // corresponding to the quantifier is added to `localExprs`.
+  if (!isa<AffineConstantExpr>(expr.getRHS())) {
+    AffineExpr dividendExpr =
+        getAffineExprFromFlatForm(lhs, numDims, numSymbols);
+    AffineExpr divisorExpr =
+        getAffineExprFromFlatForm(rhs, numDims, numSymbols);
+    AffineExpr modExpr = std::move(dividendExpr) % std::move(divisorExpr);
+    addLocalVariableSemiAffine(std::move(modExpr), lhs, lhs.size());
+    return success();
+  }
+
+  int64_t rhsConst = rhs[getConstantIndex()];
+  if (rhsConst <= 0)
+    return failure();
+
+  // Check if the LHS expression is a multiple of modulo factor.
+  unsigned i;
+  for (i = 0; i < lhs.size(); ++i)
+    if (lhs[i] % rhsConst != 0)
+      break;
+  // If yes, modulo expression here simplifies to zero.
+  if (i == lhs.size()) {
+    std::fill(lhs.begin(), lhs.end(), 0);
+    return success();
+  }
+
+  // Add a local variable for the quotient, i.e., expr % c is replaced by
+  // (expr - q * c) where q = expr floordiv c. Do this while canceling out
+  // the GCD of expr and c.
+  SmallVector<int64_t, 8> floorDividend(lhs);
+  uint64_t gcd = rhsConst;
+  for (int64_t lhsElt : lhs)
+    gcd = std::gcd(gcd, (uint64_t)std::abs(lhsElt));
+  // Simplify the numerator and the denominator.
+  if (gcd != 1) {
+    for (int64_t &floorDividendElt : floorDividend)
+      floorDividendElt = floorDividendElt / static_cast<int64_t>(gcd);
+  }
+  int64_t floorDivisor = rhsConst / static_cast<int64_t>(gcd);
+
+  // Construct the AffineExpr form of the floordiv to store in localExprs.
+
+  AffineExpr dividendExpr =
+      getAffineExprFromFlatForm(floorDividend, numDims, numSymbols);
+  AffineExpr divisorExpr = std::make_unique<AffineConstantExpr>(floorDivisor);
+  AffineExpr floorDivExpr =
+      floorDiv(std::move(dividendExpr), std::move(divisorExpr));
+  int loc;
+  if ((loc = findLocalId(floorDivExpr)) == -1) {
+    addLocalFloorDivId(floorDividend, floorDivisor, std::move(floorDivExpr));
+    // Set result at top of stack to "lhs - rhsConst * q".
+    lhs[getLocalVarStartIndex() + numLocals - 1] = -rhsConst;
+  } else {
+    // Reuse the existing local id.
+    lhs[getLocalVarStartIndex() + loc] = -rhsConst;
+  }
+  return success();
+}
+
+LogicalResult
+AffineExprFlattener::visitCeilDivExpr(const AffineBinOpExpr &expr) {
+  return visitDivExpr(expr, /*isCeil=*/true);
+}
+LogicalResult
+AffineExprFlattener::visitFloorDivExpr(const AffineBinOpExpr &expr) {
+  return visitDivExpr(expr, /*isCeil=*/false);
+}
+
+LogicalResult AffineExprFlattener::visitDimExpr(const AffineDimExpr &expr) {
+  operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
+  auto &eq = operandExprStack.back();
+  assert(expr.getPosition() < numDims && "Inconsistent number of dims");
+  eq[getDimStartIndex() + expr.getPosition()] = 1;
+  return success();
+}
+
+LogicalResult
+AffineExprFlattener::visitSymbolExpr(const AffineSymbolExpr &expr) {
+  operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
+  auto &eq = operandExprStack.back();
+  assert(expr.getPosition() < numSymbols && "inconsistent number of symbols");
+  eq[getSymbolStartIndex() + expr.getPosition()] = 1;
+  return success();
+}
+
+LogicalResult
+AffineExprFlattener::visitConstantExpr(const AffineConstantExpr &expr) {
+  operandExprStack.emplace_back(SmallVector<int64_t, 32>(getNumCols(), 0));
+  auto &eq = operandExprStack.back();
+  eq[getConstantIndex()] = expr.getValue();
+  return success();
+}
+
+void AffineExprFlattener::addLocalVariableSemiAffine(
+    AffineExpr &&expr, SmallVectorImpl<int64_t> &result,
+    unsigned long resultSize) {
+  assert(result.size() == resultSize && "result vector size mismatch");
+  int loc;
+  if ((loc = findLocalId(expr)) == -1)
+    addLocalIdSemiAffine(std::move(expr));
+  std::fill(result.begin(), result.end(), 0);
+  if (loc == -1)
+    result[getLocalVarStartIndex() + numLocals - 1] = 1;
+  else
+    result[getLocalVarStartIndex() + loc] = 1;
+}
+
+// t = expr floordiv c   <=> t = q, c * q <= expr <= c * q + c - 1
+// A floordiv is thus flattened by introducing a new local variable q, and
+// replacing that expression with 'q' while adding the constraints
+// c * q <= expr <= c * q + c - 1 to localVarCst (done by
+// IntegerRelation::addLocalFloorDiv).
+//
+// A ceildiv is similarly flattened:
+// t = expr ceildiv c   <=> t =  (expr + c - 1) floordiv c
+//
+// In case of semi affine division expressions, t = expr floordiv symbolic_expr
+// or t = expr ceildiv symbolic_expr, introduce a local variable q (= expr
+// floordiv/ceildiv symbolic_expr), and the affine floordiv/ceildiv is added to
+// `localExprs`.
+LogicalResult AffineExprFlattener::visitDivExpr(const AffineBinOpExpr &expr,
+                                                bool isCeil) {
+  assert(operandExprStack.size() >= 2);
+
+  SmallVector<int64_t, 8> rhs = operandExprStack.back();
+  operandExprStack.pop_back();
+  SmallVector<int64_t, 8> &lhs = operandExprStack.back();
+
+  // Flatten semi affine division expressions by introducing a local
+  // variable in place of the quotient, and the affine expression corresponding
+  // to the quantifier is added to `localExprs`.
+  if (!isa<AffineConstantExpr>(expr.getRHS())) {
+    AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols);
+    AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols);
+    AffineExpr divExpr = isCeil ? ceilDiv(std::move(a), std::move(b))
+                                : floorDiv(std::move(a), std::move(b));
+    addLocalVariableSemiAffine(std::move(divExpr), lhs, lhs.size());
+    return success();
+  }
+
+  // This is a pure affine expr; the RHS is a positive constant.
+  int64_t rhsConst = rhs[getConstantIndex()];
+  if (rhsConst <= 0)
+    return failure();
+
+  // Simplify the floordiv, ceildiv if possible by canceling out the greatest
+  // common divisors of the numerator and denominator.
+  uint64_t gcd = std::abs(rhsConst);
+  for (int64_t lhsElt : lhs)
+    gcd = std::gcd(gcd, (uint64_t)std::abs(lhsElt));
+  // Simplify the numerator and the denominator.
+  if (gcd != 1) {
+    for (int64_t &lhsElt : lhs)
+      lhsElt = lhsElt / static_cast<int64_t>(gcd);
+  }
+  int64_t divisor = rhsConst / static_cast<int64_t>(gcd);
+  // If the divisor becomes 1, the updated LHS is the result. (The
+  // divisor can't be negative since rhsConst is positive).
+  if (divisor == 1)
+    return success();
+
+  // If the divisor cannot be simplified to one, we will have to retain
+  // the ceil/floor expr (simplified up until here). Add an existential
+  // quantifier to express its result, i.e., expr1 div expr2 is replaced
+  // by a new identifier, q.
+  AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols);
+  AffineExpr b = std::make_unique<AffineConstantExpr>(divisor);
+
+  int loc;
+  AffineExpr divExpr = isCeil ? ceilDiv(std::move(a), std::move(b))
+                              : floorDiv(std::move(a), std::move(b));
+  if ((loc = findLocalId(divExpr)) == -1) {
+    if (!isCeil) {
+      SmallVector<int64_t, 8> dividend(lhs);
+      addLocalFloorDivId(dividend, divisor, std::move(divExpr));
+    } else {
+      // lhs ceildiv c <=>  (lhs + c - 1) floordiv c
+      SmallVector<int64_t, 8> dividend(lhs);
+      dividend.back() += divisor - 1;
+      addLocalFloorDivId(dividend, divisor, std::move(divExpr));
+    }
+  }
+  // Set the expression on stack to the local var introduced to capture the
+  // result of the division (floor or ceil).
+  std::fill(lhs.begin(), lhs.end(), 0);
+  if (loc == -1)
+    lhs[getLocalVarStartIndex() + numLocals - 1] = 1;
+  else
+    lhs[getLocalVarStartIndex() + loc] = 1;
+  return success();
+}
+
+void AffineExprFlattener::addLocalFloorDivId(ArrayRef<int64_t> dividend,
+                                             int64_t divisor,
+                                             AffineExpr &&localExpr) {
+  assert(divisor > 0 && "positive constant divisor expected");
+  for (SmallVector<int64_t, 8> &subExpr : operandExprStack)
+    subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
+  localExprs.emplace_back(std::move(localExpr));
+  ++numLocals;
+  // Update localVarCst.
+  localVarCst.addLocalFloorDiv(dividend, divisor);
+}
+
+void AffineExprFlattener::addLocalIdSemiAffine(AffineExpr &&localExpr) {
+  for (SmallVector<int64_t, 8> &subExpr : operandExprStack)
+    subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
+  localExprs.emplace_back(std::move(localExpr));
+  ++numLocals;
+}
+
+int AffineExprFlattener::findLocalId(const AffineExpr &localExpr) {
+  auto *it = llvm::find(localExprs, localExpr);
+  if (it == localExprs.end())
+    return -1;
+  return it - localExprs.begin();
+}
+
+AffineExprFlattener::AffineExprFlattener(unsigned numDims, unsigned numSymbols)
+    : numDims(numDims), numSymbols(numSymbols), numLocals(0),
+      localVarCst(PresburgerSpace::getSetSpace(numDims, numSymbols)) {
+  operandExprStack.reserve(8);
+}
+
+// Flattens the expressions in map. Returns failure if 'expr' was unable to be
+// flattened. For example two specific cases:
+// 1. semi-affine expressions not handled yet.
+// 2. has poison expression (i.e., division by zero).
+static LogicalResult
+getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
+                        unsigned numSymbols,
+                        std::vector<SmallVector<int64_t, 8>> &flattenedExprs,
+                        IntegerPolyhedron &localVarCst) {
+  if (exprs.empty()) {
+    localVarCst = IntegerPolyhedron(
+        0, 0, numDims + numSymbols + 1,
+        presburger::PresburgerSpace::getSetSpace(numDims, numSymbols, 0));
+    return success();
+  }
+
+  AffineExprFlattener flattener(numDims, numSymbols);
+  // Use the same flattener to simplify each expression successively. This way
+  // local variables / expressions are shared.
+  for (const AffineExpr &expr : exprs) {
+    if (!expr->isPureAffine())
+      return failure();
+    // has poison expression
+    LogicalResult flattenResult = flattener.walkPostOrder(*expr);
+    if (failed(flattenResult))
+      return failure();
+  }
+
+  assert(flattener.operandExprStack.size() == exprs.size());
+  flattenedExprs.clear();
+  flattenedExprs.assign(flattener.operandExprStack.begin(),
+                        flattener.operandExprStack.end());
+
+  localVarCst.clearAndCopyFrom(flattener.localVarCst);
+
+  return success();
+}
+
+LogicalResult
+getFlattenedAffineExprs(const AffineMap &map,
+                        std::vector<SmallVector<int64_t, 8>> &flattenedExprs,
+                        IntegerPolyhedron &cst) {
+  if (map.getNumExprs() == 0) {
+    cst = IntegerPolyhedron(0, 0, map.getNumDims() + map.getNumSymbols() + 1,
+                            presburger::PresburgerSpace::getSetSpace(
+                                map.getNumDims(), map.getNumSymbols(), 0));
+    return success();
+  }
+  return getFlattenedAffineExprs(map.getExprs(), map.getNumDims(),
+                                 map.getNumSymbols(), flattenedExprs, cst);
+}
+
+LogicalResult
+getFlattenedAffineExprs(const IntegerSet &set,
+                        std::vector<SmallVector<int64_t, 8>> &flattenedExprs,
+                   ...
[truncated]

``````````

</details>


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


More information about the Mlir-commits mailing list