//===- AffineExpr.cpp - MLIR Affine Expr 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 <cmath> #include <cstdint> #include <limits> #include <utility> #include "AffineExprDetail.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/IntegerSet.h" #include "mlir/Support/TypeID.h" #include "llvm/ADT/STLExtras.h" #include "llvm/Support/MathExtras.h" #include <numeric> #include <optional> usingnamespacemlir; usingnamespacemlir::detail; divideCeilSigned; divideFloorSigned; divideSignedWouldOverflow; mod; MLIRContext *AffineExpr::getContext() const { … } AffineExprKind AffineExpr::getKind() const { … } /// Walk all of the AffineExprs in `e` in postorder. This is a private factory /// method to help handle lambda walk functions. Users should use the regular /// (non-static) `walk` method. template <typename WalkRetTy> WalkRetTy mlir::AffineExpr::walk(AffineExpr e, function_ref<WalkRetTy(AffineExpr)> callback) { … } // Explicitly instantiate for the two supported return types. template void mlir::AffineExpr::walk(AffineExpr e, function_ref<void(AffineExpr)> callback); template WalkResult mlir::AffineExpr::walk(AffineExpr e, function_ref<WalkResult(AffineExpr)> callback); // Dispatch affine expression construction based on kind. AffineExpr mlir::getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs, AffineExpr rhs) { … } /// This method substitutes any uses of dimensions and symbols (e.g. /// dim#0 with dimReplacements[0]) and returns the modified expression tree. AffineExpr AffineExpr::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements, ArrayRef<AffineExpr> symReplacements) const { … } AffineExpr AffineExpr::replaceDims(ArrayRef<AffineExpr> dimReplacements) const { … } AffineExpr AffineExpr::replaceSymbols(ArrayRef<AffineExpr> symReplacements) const { … } /// Replace dims[offset ... numDims) /// by dims[offset + shift ... shift + numDims). AffineExpr AffineExpr::shiftDims(unsigned numDims, unsigned shift, unsigned offset) const { … } /// Replace symbols[offset ... numSymbols) /// by symbols[offset + shift ... shift + numSymbols). AffineExpr AffineExpr::shiftSymbols(unsigned numSymbols, unsigned shift, unsigned offset) const { … } /// Sparse replace method. Return the modified expression tree. AffineExpr AffineExpr::replace(const DenseMap<AffineExpr, AffineExpr> &map) const { … } /// Sparse replace method. Return the modified expression tree. AffineExpr AffineExpr::replace(AffineExpr expr, AffineExpr replacement) const { … } /// Returns true if this expression is made out of only symbols and /// constants (no dimensional identifiers). bool AffineExpr::isSymbolicOrConstant() const { … } /// Returns true if this is a pure affine expression, i.e., multiplication, /// floordiv, ceildiv, and mod is only allowed w.r.t constants. bool AffineExpr::isPureAffine() const { … } // Returns the greatest known integral divisor of this affine expression. int64_t AffineExpr::getLargestKnownDivisor() const { … } bool AffineExpr::isMultipleOf(int64_t factor) const { … } bool AffineExpr::isFunctionOfDim(unsigned position) const { … } bool AffineExpr::isFunctionOfSymbol(unsigned position) const { … } AffineBinaryOpExpr::AffineBinaryOpExpr(AffineExpr::ImplType *ptr) : … { … } AffineExpr AffineBinaryOpExpr::getLHS() const { … } AffineExpr AffineBinaryOpExpr::getRHS() const { … } AffineDimExpr::AffineDimExpr(AffineExpr::ImplType *ptr) : … { … } unsigned AffineDimExpr::getPosition() const { … } /// Returns true if the expression is divisible by the given symbol with /// position `symbolPos`. The argument `opKind` specifies here what kind of /// division or mod operation called this division. It helps in implementing the /// commutative property of the floordiv and ceildiv operations. If the argument ///`exprKind` is floordiv and `expr` is also a binary expression of a floordiv /// operation, then the commutative property can be used otherwise, the floordiv /// operation is not divisible. The same argument holds for ceildiv operation. static bool isDivisibleBySymbol(AffineExpr expr, unsigned symbolPos, AffineExprKind opKind) { … } /// Divides the given expression by the given symbol at position `symbolPos`. It /// considers the divisibility condition is checked before calling itself. A /// null expression is returned whenever the divisibility condition fails. static AffineExpr symbolicDivide(AffineExpr expr, unsigned symbolPos, AffineExprKind opKind) { … } /// Populate `result` with all summand operands of given (potentially nested) /// addition. If the given expression is not an addition, just populate the /// expression itself. /// Example: Add(Add(7, 8), Mul(9, 10)) will return [7, 8, Mul(9, 10)]. static void getSummandExprs(AffineExpr expr, SmallVector<AffineExpr> &result) { … } /// Return "true" if `candidate` is a negated expression, i.e., Mul(-1, expr). /// If so, also return the non-negated expression via `expr`. static bool isNegatedAffineExpr(AffineExpr candidate, AffineExpr &expr) { … } /// Return "true" if `lhs` % `rhs` is guaranteed to evaluate to zero based on /// the fact that `lhs` contains another modulo expression that ensures that /// `lhs` is divisible by `rhs`. This is a common pattern in the resulting IR /// after loop peeling. /// /// Example: lhs = ub - ub % step /// rhs = step /// => (ub - ub % step) % step is guaranteed to evaluate to 0. static bool isModOfModSubtraction(AffineExpr lhs, AffineExpr rhs, unsigned numDims, unsigned numSymbols) { … } /// Simplify a semi-affine expression by handling modulo, floordiv, or ceildiv /// operations when the second operand simplifies to a symbol and the first /// operand is divisible by that symbol. It can be applied to any semi-affine /// expression. Returned expression can either be a semi-affine or pure affine /// expression. static AffineExpr simplifySemiAffine(AffineExpr expr, unsigned numDims, unsigned numSymbols) { … } static AffineExpr getAffineDimOrSymbol(AffineExprKind kind, unsigned position, MLIRContext *context) { … } AffineExpr mlir::getAffineDimExpr(unsigned position, MLIRContext *context) { … } AffineSymbolExpr::AffineSymbolExpr(AffineExpr::ImplType *ptr) : … { … } unsigned AffineSymbolExpr::getPosition() const { … } AffineExpr mlir::getAffineSymbolExpr(unsigned position, MLIRContext *context) { … } AffineConstantExpr::AffineConstantExpr(AffineExpr::ImplType *ptr) : … { … } int64_t AffineConstantExpr::getValue() const { … } bool AffineExpr::operator==(int64_t v) const { … } AffineExpr mlir::getAffineConstantExpr(int64_t constant, MLIRContext *context) { … } SmallVector<AffineExpr> mlir::getAffineConstantExprs(ArrayRef<int64_t> constants, MLIRContext *context) { … } /// Simplify add expression. Return nullptr if it can't be simplified. static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) { … } AffineExpr AffineExpr::operator+(int64_t v) const { … } AffineExpr AffineExpr::operator+(AffineExpr other) const { … } /// Simplify a multiply expression. Return nullptr if it can't be simplified. static AffineExpr simplifyMul(AffineExpr lhs, AffineExpr rhs) { … } AffineExpr AffineExpr::operator*(int64_t v) const { … } AffineExpr AffineExpr::operator*(AffineExpr other) const { … } // Unary minus, delegate to operator*. AffineExpr AffineExpr::operator-() const { … } // Delegate to operator+. AffineExpr AffineExpr::operator-(int64_t v) const { … } AffineExpr AffineExpr::operator-(AffineExpr other) const { … } static AffineExpr simplifyFloorDiv(AffineExpr lhs, AffineExpr rhs) { … } AffineExpr AffineExpr::floorDiv(uint64_t v) const { … } AffineExpr AffineExpr::floorDiv(AffineExpr other) const { … } static AffineExpr simplifyCeilDiv(AffineExpr lhs, AffineExpr rhs) { … } AffineExpr AffineExpr::ceilDiv(uint64_t v) const { … } AffineExpr AffineExpr::ceilDiv(AffineExpr other) const { … } static AffineExpr simplifyMod(AffineExpr lhs, AffineExpr rhs) { … } AffineExpr AffineExpr::operator%(uint64_t v) const { … } AffineExpr AffineExpr::operator%(AffineExpr other) const { … } AffineExpr AffineExpr::compose(AffineMap map) const { … } raw_ostream &mlir::operator<<(raw_ostream &os, AffineExpr expr) { … } /// Constructs an affine expression from a flat ArrayRef. If there are local /// identifiers (neither dimensional nor symbolic) that appear in the sum of /// products expression, `localExprs` is expected to have the AffineExpr /// for it, and is substituted into. The ArrayRef `flatExprs` is expected to be /// in the format [dims, symbols, locals, constant term]. AffineExpr mlir::getAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs, unsigned numDims, unsigned numSymbols, ArrayRef<AffineExpr> localExprs, MLIRContext *context) { … } /// Constructs a semi-affine expression from a flat ArrayRef. If there are /// local identifiers (neither dimensional nor symbolic) that appear in the sum /// of products expression, `localExprs` is expected to have the AffineExprs for /// it, and is substituted into. The ArrayRef `flatExprs` is expected to be in /// the format [dims, symbols, locals, constant term]. The semi-affine /// expression is constructed in the sorted order of dimension and symbol /// position numbers. Note: local expressions/ids are used for mod, div as well /// as symbolic RHS terms for terms that are not pure affine. static AffineExpr getSemiAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs, unsigned numDims, unsigned numSymbols, ArrayRef<AffineExpr> localExprs, MLIRContext *context) { … } SimpleAffineExprFlattener::SimpleAffineExprFlattener(unsigned numDims, unsigned numSymbols) : … { … } // 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 SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) { … } LogicalResult SimpleAffineExprFlattener::visitAddExpr(AffineBinaryOpExpr expr) { … } // // 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 SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) { … } LogicalResult SimpleAffineExprFlattener::visitCeilDivExpr(AffineBinaryOpExpr expr) { … } LogicalResult SimpleAffineExprFlattener::visitFloorDivExpr(AffineBinaryOpExpr expr) { … } LogicalResult SimpleAffineExprFlattener::visitDimExpr(AffineDimExpr expr) { … } LogicalResult SimpleAffineExprFlattener::visitSymbolExpr(AffineSymbolExpr expr) { … } LogicalResult SimpleAffineExprFlattener::visitConstantExpr(AffineConstantExpr expr) { … } LogicalResult SimpleAffineExprFlattener::addLocalVariableSemiAffine( ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs, AffineExpr localExpr, SmallVectorImpl<int64_t> &result, unsigned long resultSize) { … } // 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 SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr, bool isCeil) { … } // Add a local identifier (needed to flatten a mod, floordiv, ceildiv expr). // The local identifier added is always a floordiv of a pure add/mul affine // function of other identifiers, coefficients of which are specified in // dividend and with respect to a positive constant divisor. localExpr is the // simplified tree expression (AffineExpr) corresponding to the quantifier. void SimpleAffineExprFlattener::addLocalFloorDivId(ArrayRef<int64_t> dividend, int64_t divisor, AffineExpr localExpr) { … } LogicalResult SimpleAffineExprFlattener::addLocalIdSemiAffine( ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs, AffineExpr localExpr) { … } int SimpleAffineExprFlattener::findLocalId(AffineExpr localExpr) { … } /// Simplify the affine expression by flattening it and reconstructing it. AffineExpr mlir::simplifyAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols) { … } std::optional<int64_t> mlir::getBoundForAffineExpr( AffineExpr expr, unsigned numDims, unsigned numSymbols, ArrayRef<std::optional<int64_t>> constLowerBounds, ArrayRef<std::optional<int64_t>> constUpperBounds, bool isUpper) { … }
Codebase Browser
llvm