AffineMap.cpp | Explore in Territory

//===- AffineMap.cpp - MLIR Affine Map Classes ----------------------------===//
//
// 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/IR/AffineMap.h"
#include "AffineMapDetail.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Support/raw_ostream.h"
#include <iterator>
#include <numeric>
#include <optional>
#include <type_traits>

usingnamespacemlir;

divideCeilSigned;
divideFloorSigned;
mod;

namespace {

// AffineExprConstantFolder evaluates an affine expression using constant
// operands passed in 'operandConsts'. Returns an IntegerAttr attribute
// representing the constant value of the affine expression evaluated on
// constant 'operandConsts', or nullptr if it can't be folded.
class AffineExprConstantFolder { … };

} // namespace

/// Returns a single constant result affine map.
AffineMap AffineMap::getConstantMap(int64_t val, MLIRContext *context) { … }

/// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most
/// minor dimensions.
AffineMap AffineMap::getMinorIdentityMap(unsigned dims, unsigned results,
                                         MLIRContext *context) { … }

AffineMap AffineMap::getFilteredIdentityMap(
    MLIRContext *ctx, unsigned numDims,
    llvm::function_ref<bool(AffineDimExpr)> keepDimFilter) { … }

bool AffineMap::isMinorIdentity() const { … }

SmallVector<unsigned> AffineMap::getBroadcastDims() const { … }

/// Returns true if this affine map is a minor identity up to broadcasted
/// dimensions which are indicated by value 0 in the result.
bool AffineMap::isMinorIdentityWithBroadcasting(
    SmallVectorImpl<unsigned> *broadcastedDims) const { … }

/// Return true if this affine map can be converted to a minor identity with
/// broadcast by doing a permute. Return a permutation (there may be
/// several) to apply to get to a minor identity with broadcasts.
/// Ex:
///  * (d0, d1, d2) -> (0, d1) maps to minor identity (d1, 0 = d2) with
///  perm = [1, 0] and broadcast d2
///  * (d0, d1, d2) -> (d0, 0) cannot be mapped to a minor identity by
///  permutation + broadcast
///  * (d0, d1, d2, d3) -> (0, d1, d3) maps to minor identity (d1, 0 = d2, d3)
///  with perm = [1, 0, 2] and broadcast d2
///  * (d0, d1) -> (d1, 0, 0, d0) maps to minor identity (d0, d1) with extra
///  leading broadcat dimensions. The map returned would be (0, 0, d0, d1) with
///  perm = [3, 0, 1, 2]
bool AffineMap::isPermutationOfMinorIdentityWithBroadcasting(
    SmallVectorImpl<unsigned> &permutedDims) const { … }

/// Returns an AffineMap representing a permutation.
AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,
                                       MLIRContext *context) { … }
AffineMap AffineMap::getPermutationMap(ArrayRef<int64_t> permutation,
                                       MLIRContext *context) { … }

AffineMap AffineMap::getMultiDimMapWithTargets(unsigned numDims,
                                               ArrayRef<unsigned> targets,
                                               MLIRContext *context) { … }

/// Creates an affine map each for each list of AffineExpr's in `exprsList`
/// while inferring the right number of dimensional and symbolic inputs needed
/// based on the maximum dimensional and symbolic identifier appearing in the
/// expressions.
template <typename AffineExprContainer>
static SmallVector<AffineMap, 4>
inferFromExprList(ArrayRef<AffineExprContainer> exprsList,
                  MLIRContext *context) { … }

SmallVector<AffineMap, 4>
AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList,
                             MLIRContext *context) { … }

SmallVector<AffineMap, 4>
AffineMap::inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList,
                             MLIRContext *context) { … }

uint64_t AffineMap::getLargestKnownDivisorOfMapExprs() { … }

AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims,
                                            MLIRContext *context) { … }

MLIRContext *AffineMap::getContext() const { … }

bool AffineMap::isIdentity() const { … }

bool AffineMap::isSymbolIdentity() const { … }

bool AffineMap::isEmpty() const { … }

bool AffineMap::isSingleConstant() const { … }

bool AffineMap::isConstant() const { … }

int64_t AffineMap::getSingleConstantResult() const { … }

SmallVector<int64_t> AffineMap::getConstantResults() const { … }

unsigned AffineMap::getNumDims() const { … }
unsigned AffineMap::getNumSymbols() const { … }
unsigned AffineMap::getNumResults() const { … }
unsigned AffineMap::getNumInputs() const { … }
ArrayRef<AffineExpr> AffineMap::getResults() const { … }
AffineExpr AffineMap::getResult(unsigned idx) const { … }

unsigned AffineMap::getDimPosition(unsigned idx) const { … }

std::optional<unsigned> AffineMap::getResultPosition(AffineExpr input) const { … }

/// Folds the results of the application of an affine map on the provided
/// operands to a constant if possible. Returns false if the folding happens,
/// true otherwise.
LogicalResult AffineMap::constantFold(ArrayRef<Attribute> operandConstants,
                                      SmallVectorImpl<Attribute> &results,
                                      bool *hasPoison) const { … }

AffineMap AffineMap::partialConstantFold(ArrayRef<Attribute> operandConstants,
                                         SmallVectorImpl<int64_t> *results,
                                         bool *hasPoison) const { … }

/// Walk all of the AffineExpr's in this mapping. Each node in an expression
/// tree is visited in postorder.
void AffineMap::walkExprs(llvm::function_ref<void(AffineExpr)> callback) const { … }

/// This method substitutes any uses of dimensions and symbols (e.g.
/// dim#0 with dimReplacements[0]) in subexpressions and returns the modified
/// expression mapping.  Because this can be used to eliminate dims and
/// symbols, the client needs to specify the number of dims and symbols in
/// the result.  The returned map always has the same number of results.
AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
                                           ArrayRef<AffineExpr> symReplacements,
                                           unsigned numResultDims,
                                           unsigned numResultSyms) const { … }

/// Sparse replace method. Apply AffineExpr::replace(`expr`, `replacement`) to
/// each of the results and return a new AffineMap with the new results and
/// with the specified number of dims and symbols.
AffineMap AffineMap::replace(AffineExpr expr, AffineExpr replacement,
                             unsigned numResultDims,
                             unsigned numResultSyms) const { … }

/// Sparse replace method. Apply AffineExpr::replace(`map`) to each of the
/// results and return a new AffineMap with the new results and with the
/// specified number of dims and symbols.
AffineMap AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map,
                             unsigned numResultDims,
                             unsigned numResultSyms) const { … }

AffineMap
AffineMap::replace(const DenseMap<AffineExpr, AffineExpr> &map) const { … }

AffineMap AffineMap::dropResults(const llvm::SmallBitVector &positions) const { … }

AffineMap AffineMap::compose(AffineMap map) const { … }

SmallVector<int64_t, 4> AffineMap::compose(ArrayRef<int64_t> values) const { … }

bool AffineMap::isProjectedPermutation(bool allowZeroInResults) const { … }

bool AffineMap::isPermutation() const { … }

AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) const { … }

AffineMap AffineMap::getSliceMap(unsigned start, unsigned length) const { … }

AffineMap AffineMap::getMajorSubMap(unsigned numResults) const { … }

AffineMap AffineMap::getMinorSubMap(unsigned numResults) const { … }

/// Implementation detail to compress multiple affine maps with a compressionFun
/// that is expected to be either compressUnusedDims or compressUnusedSymbols.
/// The implementation keeps track of num dims and symbols across the different
/// affine maps.
static SmallVector<AffineMap> compressUnusedListImpl(
    ArrayRef<AffineMap> maps,
    llvm::function_ref<AffineMap(AffineMap)> compressionFun) { … }

AffineMap mlir::compressDims(AffineMap map,
                             const llvm::SmallBitVector &unusedDims) { … }

AffineMap mlir::compressUnusedDims(AffineMap map) { … }

SmallVector<AffineMap> mlir::compressUnusedDims(ArrayRef<AffineMap> maps) { … }

AffineMap mlir::compressSymbols(AffineMap map,
                                const llvm::SmallBitVector &unusedSymbols) { … }

AffineMap mlir::compressUnusedSymbols(AffineMap map) { … }

SmallVector<AffineMap> mlir::compressUnusedSymbols(ArrayRef<AffineMap> maps) { … }

AffineMap mlir::foldAttributesIntoMap(Builder &b, AffineMap map,
                                      ArrayRef<OpFoldResult> operands,
                                      SmallVector<Value> &remainingValues) { … }

AffineMap mlir::simplifyAffineMap(AffineMap map) { … }

AffineMap mlir::removeDuplicateExprs(AffineMap map) { … }

AffineMap mlir::inversePermutation(AffineMap map) { … }

AffineMap mlir::inverseAndBroadcastProjectedPermutation(AffineMap map) { … }

AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) { … }

/// Common implementation to project out dimensions or symbols from an affine
/// map based on the template type.
/// Additionally, if 'compress' is true, the projected out dimensions or symbols
/// are also dropped from the resulting map.
template <typename AffineDimOrSymExpr>
static AffineMap projectCommonImpl(AffineMap map,
                                   const llvm::SmallBitVector &toProject,
                                   bool compress) { … }

AffineMap mlir::projectDims(AffineMap map,
                            const llvm::SmallBitVector &projectedDimensions,
                            bool compressDimsFlag) { … }

AffineMap mlir::projectSymbols(AffineMap map,
                               const llvm::SmallBitVector &projectedSymbols,
                               bool compressSymbolsFlag) { … }

AffineMap mlir::getProjectedMap(AffineMap map,
                                const llvm::SmallBitVector &projectedDimensions,
                                bool compressDimsFlag,
                                bool compressSymbolsFlag) { … }

llvm::SmallBitVector mlir::getUnusedDimsBitVector(ArrayRef<AffineMap> maps) { … }

llvm::SmallBitVector mlir::getUnusedSymbolsBitVector(ArrayRef<AffineMap> maps) { … }

AffineMap
mlir::expandDimsToRank(AffineMap map, int64_t rank,
                       const llvm::SmallBitVector &projectedDimensions) { … }

//===----------------------------------------------------------------------===//
// MutableAffineMap.
//===----------------------------------------------------------------------===//

MutableAffineMap::MutableAffineMap(AffineMap map)
    : … { … }

void MutableAffineMap::reset(AffineMap map) { … }

bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const { … }

// Simplifies the result affine expressions of this map. The expressions
// have to be pure for the simplification implemented.
void MutableAffineMap::simplify() { … }

AffineMap MutableAffineMap::getAffineMap() const { … }
llvm/mlir/lib/IR/AffineMap.cpp
