llvm/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h

//===- FlatLinearValueConstraints.h - Linear Constraints --------*- 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_FLATLINEARVALUECONSTRAINTS_H
#define MLIR_ANALYSIS_FLATLINEARVALUECONSTRAINTS_H

#include "mlir/Analysis/Presburger/IntegerRelation.h"
#include "mlir/Analysis/Presburger/Matrix.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/OpDefinition.h"
#include <optional>

namespace mlir {

class AffineMap;
class IntegerSet;
class MLIRContext;
class Value;
class MemRefType;
struct MutableAffineMap;

namespace presburger {
class MultiAffineFunction;
} // namespace presburger

/// FlatLinearConstraints is an extension of IntegerPolyhedron. It provides an
/// AffineExpr-based API.
class FlatLinearConstraints : public presburger::IntegerPolyhedron {};

/// FlatLinearValueConstraints represents an extension of FlatLinearConstraints
/// where each non-local variable can have an SSA Value attached to it.
class FlatLinearValueConstraints : public FlatLinearConstraints {};

/// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the
/// dimensions, symbols, and additional variables that represent floor divisions
/// of dimensions, symbols, and in turn other floor divisions.  Returns failure
/// if 'expr' could not be flattened (i.e., an unhandled semi-affine was found).
/// 'cst' contains constraints that connect newly introduced local variables
/// to existing dimensional and symbolic variables. See documentation for
/// AffineExprFlattener on how mod's and div's are flattened.
LogicalResult
getFlattenedAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols,
                       SmallVectorImpl<int64_t> *flattenedExpr,
                       FlatLinearConstraints *cst = nullptr,
                       bool addConservativeSemiAffineBounds = false);

/// Flattens the result expressions of the map to their corresponding flattened
/// forms and set in 'flattenedExprs'. Returns failure if any expression in the
/// map could not be flattened (i.e., an unhandled semi-affine was found). 'cst'
/// contains constraints that connect newly introduced local variables to
/// existing dimensional and / symbolic variables. See documentation for
/// AffineExprFlattener on how mod's and div's are flattened. For all affine
/// expressions that share the same operands (like those of an affine map), this
/// method should be used instead of repeatedly calling getFlattenedAffineExpr
/// since local variables added to deal with div's and mod's will be reused
/// across expressions.
LogicalResult
getFlattenedAffineExprs(AffineMap map,
                        std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
                        FlatLinearConstraints *cst = nullptr,
                        bool addConservativeSemiAffineBounds = false);
LogicalResult
getFlattenedAffineExprs(IntegerSet set,
                        std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
                        FlatLinearConstraints *cst = nullptr);

LogicalResult
getMultiAffineFunctionFromMap(AffineMap map,
                              presburger::MultiAffineFunction &multiAff);

/// Re-indexes the dimensions and symbols of an affine map with given `operands`
/// values to align with `dims` and `syms` values.
///
/// Each dimension/symbol of the map, bound to an operand `o`, is replaced with
/// dimension `i`, where `i` is the position of `o` within `dims`. If `o` is not
/// in `dims`, replace it with symbol `i`, where `i` is the position of `o`
/// within `syms`. If `o` is not in `syms` either, replace it with a new symbol.
///
/// Note: If a value appears multiple times as a dimension/symbol (or both), all
/// corresponding dim/sym expressions are replaced with the first dimension
/// bound to that value (or first symbol if no such dimension exists).
///
/// The resulting affine map has `dims.size()` many dimensions and at least
/// `syms.size()` many symbols.
///
/// The SSA values of the symbols of the resulting map are optionally returned
/// via `newSyms`. This is a concatenation of `syms` with the SSA values of the
/// newly added symbols.
///
/// Note: As part of this re-indexing, dimensions may turn into symbols, or vice
/// versa.
AffineMap alignAffineMapWithValues(AffineMap map, ValueRange operands,
                                   ValueRange dims, ValueRange syms,
                                   SmallVector<Value> *newSyms = nullptr);

} // namespace mlir

#endif // MLIR_ANALYSIS_FLATLINEARVALUECONSTRAINTS_H