//===- AffineMap.h - MLIR Affine Map Class ----------------------*- 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 // //===----------------------------------------------------------------------===// // // Affine maps are mathematical functions which map a list of dimension // identifiers and symbols, to multidimensional affine expressions. // //===----------------------------------------------------------------------===// #ifndef MLIR_IR_AFFINEMAP_H #define MLIR_IR_AFFINEMAP_H #include "mlir/IR/AffineExpr.h" #include "mlir/IR/Value.h" #include "mlir/Support/LLVM.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseMapInfo.h" #include "llvm/ADT/SmallBitVector.h" #include "llvm/ADT/SmallVectorExtras.h" #include <optional> namespace llvm { class SmallBitVector; } // namespace llvm namespace mlir { namespace detail { struct AffineMapStorage; } // namespace detail class Attribute; class Builder; class OpFoldResult; class MLIRContext; /// A multi-dimensional affine map /// Affine map's are immutable like Type's, and they are uniqued. /// Eg: (d0, d1) -> (d0/128, d0 mod 128, d1) /// The names used (d0, d1) don't matter - it's the mathematical function that /// is unique to this affine map. class AffineMap { … }; // Make AffineExpr hashable. inline ::llvm::hash_code hash_value(AffineMap arg) { … } /// A mutable affine map. Its affine expressions are however unique. struct MutableAffineMap { … }; /// Simplifies an affine map by simplifying its underlying AffineExpr results. AffineMap simplifyAffineMap(AffineMap map); /// Drop the dims that are listed in `unusedDims`. AffineMap compressDims(AffineMap map, const llvm::SmallBitVector &unusedDims); /// Drop the dims that are not used. AffineMap compressUnusedDims(AffineMap map); /// Drop the dims that are not used by any of the individual maps in `maps`. /// Asserts that all maps in `maps` are normalized to the same number of /// dims and symbols. SmallVector<AffineMap> compressUnusedDims(ArrayRef<AffineMap> maps); /// Drop the symbols that are listed in `unusedSymbols`. AffineMap compressSymbols(AffineMap map, const llvm::SmallBitVector &unusedSymbols); /// Drop the symbols that are not used. AffineMap compressUnusedSymbols(AffineMap map); /// Drop the symbols that are not used by any of the individual maps in `maps`. /// Asserts that all maps in `maps` are normalized to the same number of /// dims and symbols. SmallVector<AffineMap> compressUnusedSymbols(ArrayRef<AffineMap> maps); /// Fold all attributes among the given operands into the affine map. Return the /// folded affine map. Return all remaining values via `remainingValues`. AffineMap foldAttributesIntoMap(Builder &b, AffineMap map, ArrayRef<OpFoldResult> operands, SmallVector<Value> &remainingValues); /// Returns a map with the same dimension and symbol count as `map`, but whose /// results are the unique affine expressions of `map`. AffineMap removeDuplicateExprs(AffineMap map); /// Returns a map of codomain to domain dimensions such that the first codomain /// dimension for a particular domain dimension is selected. /// Returns an empty map if the input map is empty. /// Returns null map (not empty map) if `map` is not invertible (i.e. `map` does /// not contain a subset that is a permutation of full domain rank). /// /// Prerequisites: /// 1. `map` has no symbols. /// /// Example 1: /// /// ```mlir /// (d0, d1, d2) -> (d1, d1, d0, d2, d1, d2, d1, d0) /// 0 2 3 /// ``` /// /// returns: /// /// ```mlir /// (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3) /// ``` /// /// Example 2: /// /// ```mlir /// (d0, d1, d2) -> (d1, d0 + d1, d0, d2, d1, d2, d1, d0) /// 0 2 3 /// ``` /// /// returns: /// /// ```mlir /// (d0, d1, d2, d3, d4, d5, d6, d7) -> (d2, d0, d3) /// ``` AffineMap inversePermutation(AffineMap map); /// Return the reverse map of a projected permutation where the projected /// dimensions are transformed into 0s. /// /// Prerequisites: `map` must be a projected permutation. /// /// Example 1: /// /// ```mlir /// affine_map<(d0, d1, d2, d3) -> (d2, d0)> /// ``` /// /// returns: /// /// ```mlir /// affine_map<(d0, d1) -> (d1, 0, d0, 0)> /// ``` /// /// Example 2: /// /// ```mlir /// affine_map<(d0, d1, d2, d3) -> (d0, d3)> /// ``` /// /// returns: /// /// ```mlir /// affine_map<(d0, d1) -> (d0, 0, 0, d1)> /// ``` /// /// Example 3: /// /// ```mlir /// affine_map<(d0, d1, d2, d3) -> (d2)> /// ``` /// /// returns: /// /// ```mlir /// affine_map<(d0) -> (0, 0, d0, 0)> /// ``` /// Example 4: /// /// ```mlir /// affine_map<(d0, d1, d2) -> (d0, 0)> /// ``` /// /// returns: /// /// ```mlir /// affine_map<(d0, d1) -> (d0, 0, 0)> /// ``` AffineMap inverseAndBroadcastProjectedPermutation(AffineMap map); /// Concatenates a list of `maps` into a single AffineMap, stepping over /// potentially empty maps. Assumes each of the underlying map has 0 symbols. /// The resulting map has a number of dims equal to the max of `maps`' dims and /// the concatenated results as its results. /// Returns an empty map if all input `maps` are empty. /// /// Example: /// When applied to the following list of 3 affine maps, /// /// ```mlir /// { /// (i, j, k) -> (i, k), /// (i, j, k) -> (k, j), /// (i, j, k) -> (i, j) /// } /// ``` /// /// Returns the map: /// /// ```mlir /// (i, j, k) -> (i, k, k, j, i, j) /// ``` AffineMap concatAffineMaps(ArrayRef<AffineMap> maps); /// Returns the map that results from projecting out the dimensions specified in /// `projectedDimensions`. The projected dimensions are set to 0. /// /// Example: /// 1) map : affine_map<(d0, d1, d2) -> (d0, d1)> /// projected_dimensions : {2} /// result : affine_map<(d0, d1) -> (d0, d1)> /// /// 2) map : affine_map<(d0, d1) -> (d0 + d1)> /// projected_dimensions : {1} /// result : affine_map<(d0) -> (d0)> /// /// 3) map : affine_map<(d0, d1, d2) -> (d0, d1)> /// projected_dimensions : {1} /// result : affine_map<(d0, d1) -> (d0, 0)> /// /// This function also compresses the dims when the boolean flag is true. AffineMap projectDims(AffineMap map, const llvm::SmallBitVector &projectedDimensions, bool compressDimsFlag = false); /// Symbol counterpart of `projectDims`. /// This function also compresses the symbols when the boolean flag is true. AffineMap projectSymbols(AffineMap map, const llvm::SmallBitVector &projectedSymbols, bool compressSymbolsFlag = false); /// Calls `projectDims(map, projectedDimensions, compressDimsFlag)`. /// If `compressSymbolsFlag` is true, additionally call `compressUnusedSymbols`. AffineMap getProjectedMap(AffineMap map, const llvm::SmallBitVector &projectedDimensions, bool compressDimsFlag = true, bool compressSymbolsFlag = true); // Return a bitvector where each bit set indicates a dimension that is not used // by any of the maps in the input array `maps`. llvm::SmallBitVector getUnusedDimsBitVector(ArrayRef<AffineMap> maps); // Return a bitvector where each bit set indicates a symbol that is not used // by any of the maps in the input array `maps`. llvm::SmallBitVector getUnusedSymbolsBitVector(ArrayRef<AffineMap> maps); /// Expand `map` to operate on `rank` dims while projecting out the dims in /// `projectedDimensions`. This amounts to composing `map` with /// `id(rank).dropResults(projectedDimensions)`. AffineMap expandDimsToRank(AffineMap map, int64_t rank, const llvm::SmallBitVector &projectedDimensions); inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) { … } //===----------------------------------------------------------------------===// // Templated helper functions. //===----------------------------------------------------------------------===// /// Apply a permutation from `map` to `source` and return the result. template <typename T> SmallVector<T> applyPermutationMap(AffineMap map, llvm::ArrayRef<T> source) { … } /// Calculates maximum dimension and symbol positions from the expressions /// in `exprsLists` and stores them in `maxDim` and `maxSym` respectively. template <typename AffineExprContainer> static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList, int64_t &maxDim, int64_t &maxSym) { … } } // namespace mlir namespace llvm { // AffineExpr hash just like pointers template <> struct DenseMapInfo<mlir::AffineMap> { … }; } // namespace llvm #endif // MLIR_IR_AFFINEMAP_H
Codebase Browser
llvm