FlatLinearValueConstraints.cpp | Explore in Territory

//===- FlatLinearValueConstraints.cpp - Linear Constraint -----------------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "mlir/Analysis//FlatLinearValueConstraints.h"

#include "mlir/Analysis/Presburger/LinearTransform.h"
#include "mlir/Analysis/Presburger/PresburgerSpace.h"
#include "mlir/Analysis/Presburger/Simplex.h"
#include "mlir/Analysis/Presburger/Utils.h"
#include "mlir/IR/AffineExprVisitor.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/IntegerSet.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include <optional>

#define DEBUG_TYPE …

usingnamespacemlir;
usingnamespacepresburger;

//===----------------------------------------------------------------------===//
// AffineExprFlattener
//===----------------------------------------------------------------------===//

namespace {

// See comments for SimpleAffineExprFlattener.
// An AffineExprFlattenerWithLocalVars extends a SimpleAffineExprFlattener by
// recording constraint information associated with mod's, floordiv's, and
// ceildiv's in FlatLinearConstraints 'localVarCst'.
struct AffineExprFlattener : public SimpleAffineExprFlattener { … };

// A SemiAffineExprFlattener is an AffineExprFlattenerWithLocalVars that adds
// conservative bounds for semi-affine expressions (given assumptions hold). If
// the assumptions required to add the semi-affine bounds are found not to hold
// the final constraints set will be empty/inconsistent. If the assumptions are
// never contradicted the final bounds still only will be correct if the
// assumptions hold.
struct SemiAffineExprFlattener : public AffineExprFlattener { … };

} // namespace

// Flattens the expressions in map. Returns failure if 'expr' was unable to be
// flattened. For example two specific cases:
// 1. an unhandled semi-affine expressions is found.
// 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,
                        FlatLinearConstraints *localVarCst,
                        bool addConservativeSemiAffineBounds = false) { … }

// Flattens 'expr' into 'flattenedExpr'. Returns failure if 'expr' was unable to
// be flattened (an unhandled semi-affine was found).
LogicalResult mlir::getFlattenedAffineExpr(
    AffineExpr expr, unsigned numDims, unsigned numSymbols,
    SmallVectorImpl<int64_t> *flattenedExpr, FlatLinearConstraints *localVarCst,
    bool addConservativeSemiAffineBounds) { … }

/// Flattens the expressions in map. Returns failure if 'expr' was unable to be
/// flattened (i.e., an unhandled semi-affine was found).
LogicalResult mlir::getFlattenedAffineExprs(
    AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
    FlatLinearConstraints *localVarCst, bool addConservativeSemiAffineBounds) { … }

LogicalResult mlir::getFlattenedAffineExprs(
    IntegerSet set, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
    FlatLinearConstraints *localVarCst) { … }

//===----------------------------------------------------------------------===//
// FlatLinearConstraints
//===----------------------------------------------------------------------===//

// Similar to `composeMap` except that no Values need be associated with the
// constraint system nor are they looked at -- the dimensions and symbols of
// `other` are expected to correspond 1:1 to `this` system.
LogicalResult FlatLinearConstraints::composeMatchingMap(AffineMap other) { … }

// Determine whether the variable at 'pos' (say var_r) can be expressed as
// modulo of another known variable (say var_n) w.r.t a constant. For example,
// if the following constraints hold true:
// ```
// 0 <= var_r <= divisor - 1
// var_n - (divisor * q_expr) = var_r
// ```
// where `var_n` is a known variable (called dividend), and `q_expr` is an
// `AffineExpr` (called the quotient expression), `var_r` can be written as:
//
// `var_r = var_n mod divisor`.
//
// Additionally, in a special case of the above constaints where `q_expr` is an
// variable itself that is not yet known (say `var_q`), it can be written as a
// floordiv in the following way:
//
// `var_q = var_n floordiv divisor`.
//
// First 'num' dimensional variables starting at 'offset' are
// derived/to-be-derived in terms of the remaining variables. The remaining
// variables are assigned trivial affine expressions in `memo`. For example,
// memo is initilized as follows for a `cst` with 5 dims, when offset=2, num=2:
// memo ==>  d0  d1  .   .   d2 ...
// cst  ==>  c0  c1  c2  c3  c4 ...
//
// Returns true if the above mod or floordiv are detected, updating 'memo' with
// these new expressions. Returns false otherwise.
static bool detectAsMod(const FlatLinearConstraints &cst, unsigned pos,
                        unsigned offset, unsigned num, int64_t lbConst,
                        int64_t ubConst, MLIRContext *context,
                        SmallVectorImpl<AffineExpr> &memo) { … }

/// Check if the pos^th variable can be expressed as a floordiv of an affine
/// function of other variables (where the divisor is a positive constant)
/// given the initial set of expressions in `exprs`. If it can be, the
/// corresponding position in `exprs` is set as the detected affine expr. For
/// eg: 4q <= i + j <= 4q + 3   <=>   q = (i + j) floordiv 4. An equality can
/// also yield a floordiv: eg.  4q = i + j <=> q = (i + j) floordiv 4. 32q + 28
/// <= i <= 32q + 31 => q = i floordiv 32.
static bool detectAsFloorDiv(const FlatLinearConstraints &cst, unsigned pos,
                             MLIRContext *context,
                             SmallVectorImpl<AffineExpr> &exprs) { … }

std::pair<AffineMap, AffineMap> FlatLinearConstraints::getLowerAndUpperBound(
    unsigned pos, unsigned offset, unsigned num, unsigned symStartPos,
    ArrayRef<AffineExpr> localExprs, MLIRContext *context,
    bool closedUB) const { … }

/// Computes the lower and upper bounds of the first 'num' dimensional
/// variables (starting at 'offset') as affine maps of the remaining
/// variables (dimensional and symbolic variables). Local variables are
/// themselves explicitly computed as affine functions of other variables in
/// this process if needed.
void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
                                           MLIRContext *context,
                                           SmallVectorImpl<AffineMap> *lbMaps,
                                           SmallVectorImpl<AffineMap> *ubMaps,
                                           bool closedUB) { … }

LogicalResult FlatLinearConstraints::flattenAlignedMapAndMergeLocals(
    AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
    bool addConservativeSemiAffineBounds) { … }

LogicalResult FlatLinearConstraints::addBound(
    BoundType type, unsigned pos, AffineMap boundMap, bool isClosedBound,
    AddConservativeSemiAffineBounds addSemiAffineBounds) { … }

LogicalResult FlatLinearConstraints::addBound(
    BoundType type, unsigned pos, AffineMap boundMap,
    AddConservativeSemiAffineBounds addSemiAffineBounds) { … }

/// Compute an explicit representation for local vars. For all systems coming
/// from MLIR integer sets, maps, or expressions where local vars were
/// introduced to model floordivs and mods, this always succeeds.
LogicalResult
FlatLinearConstraints::computeLocalVars(SmallVectorImpl<AffineExpr> &memo,
                                        MLIRContext *context) const { … }

IntegerSet FlatLinearConstraints::getAsIntegerSet(MLIRContext *context) const { … }

//===----------------------------------------------------------------------===//
// FlatLinearValueConstraints
//===----------------------------------------------------------------------===//

// Construct from an IntegerSet.
FlatLinearValueConstraints::FlatLinearValueConstraints(IntegerSet set,
                                                       ValueRange operands)
    : … { … }

unsigned FlatLinearValueConstraints::appendDimVar(ValueRange vals) { … }

unsigned FlatLinearValueConstraints::appendSymbolVar(ValueRange vals) { … }

unsigned FlatLinearValueConstraints::insertDimVar(unsigned pos,
                                                  ValueRange vals) { … }

unsigned FlatLinearValueConstraints::insertSymbolVar(unsigned pos,
                                                     ValueRange vals) { … }

unsigned FlatLinearValueConstraints::insertVar(VarKind kind, unsigned pos,
                                               unsigned num) { … }

unsigned FlatLinearValueConstraints::insertVar(VarKind kind, unsigned pos,
                                               ValueRange vals) { … }

/// Checks if two constraint systems are in the same space, i.e., if they are
/// associated with the same set of variables, appearing in the same order.
static bool areVarsAligned(const FlatLinearValueConstraints &a,
                           const FlatLinearValueConstraints &b) { … }

/// Calls areVarsAligned to check if two constraint systems have the same set
/// of variables in the same order.
bool FlatLinearValueConstraints::areVarsAlignedWithOther(
    const FlatLinearConstraints &other) { … }

/// Checks if the SSA values associated with `cst`'s variables in range
/// [start, end) are unique.
static bool LLVM_ATTRIBUTE_UNUSED areVarsUnique(
    const FlatLinearValueConstraints &cst, unsigned start, unsigned end) { … }

/// Checks if the SSA values associated with `cst`'s variables are unique.
static bool LLVM_ATTRIBUTE_UNUSED
areVarsUnique(const FlatLinearValueConstraints &cst) { … }

/// Checks if the SSA values associated with `cst`'s variables of kind `kind`
/// are unique.
static bool LLVM_ATTRIBUTE_UNUSED
areVarsUnique(const FlatLinearValueConstraints &cst, VarKind kind) { … }

/// Merge and align the variables of A and B starting at 'offset', so that
/// both constraint systems get the union of the contained variables that is
/// dimension-wise and symbol-wise unique; both constraint systems are updated
/// so that they have the union of all variables, with A's original
/// variables appearing first followed by any of B's variables that didn't
/// appear in A. Local variables in B that have the same division
/// representation as local variables in A are merged into one. We allow A
/// and B to have non-unique values for their variables; in such cases, they are
/// still aligned with the variables appearing first aligned with those
/// appearing first in the other system from left to right.
//  E.g.: Input: A has ((%i, %j) [%M, %N]) and B has (%k, %j) [%P, %N, %M])
//        Output: both A, B have (%i, %j, %k) [%M, %N, %P]
static void mergeAndAlignVars(unsigned offset, FlatLinearValueConstraints *a,
                              FlatLinearValueConstraints *b) { … }

// Call 'mergeAndAlignVars' to align constraint systems of 'this' and 'other'.
void FlatLinearValueConstraints::mergeAndAlignVarsWithOther(
    unsigned offset, FlatLinearValueConstraints *other) { … }

/// Merge and align symbols of `this` and `other` such that both get union of
/// of symbols. Existing symbols need not be unique; they will be aligned from
/// left to right with duplicates aligned in the same order. Symbols with Value
/// as `None` are considered to be inequal to all other symbols.
void FlatLinearValueConstraints::mergeSymbolVars(
    FlatLinearValueConstraints &other) { … }

void FlatLinearValueConstraints::removeVarRange(VarKind kind, unsigned varStart,
                                                unsigned varLimit) { … }

AffineMap
FlatLinearValueConstraints::computeAlignedMap(AffineMap map,
                                              ValueRange operands) const { … }

bool FlatLinearValueConstraints::findVar(Value val, unsigned *pos,
                                         unsigned offset) const { … }

bool FlatLinearValueConstraints::containsVar(Value val) const { … }

void FlatLinearValueConstraints::addBound(BoundType type, Value val,
                                          int64_t value) { … }

void FlatLinearConstraints::printSpace(raw_ostream &os) const { … }

void FlatLinearValueConstraints::printSpace(raw_ostream &os) const { … }

void FlatLinearValueConstraints::projectOut(Value val) { … }

LogicalResult FlatLinearValueConstraints::unionBoundingBox(
    const FlatLinearValueConstraints &otherCst) { … }

//===----------------------------------------------------------------------===//
// Helper functions
//===----------------------------------------------------------------------===//

AffineMap mlir::alignAffineMapWithValues(AffineMap map, ValueRange operands,
                                         ValueRange dims, ValueRange syms,
                                         SmallVector<Value> *newSyms) { … }

LogicalResult
mlir::getMultiAffineFunctionFromMap(AffineMap map,
                                    MultiAffineFunction &multiAff) { … }
llvm/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
