|
| 1 | +//===- AffineExprBounds.h - Compute bounds of affine expressions *- C++ -*-===// |
| 2 | +// |
| 3 | +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| 4 | +// See https://llvm.org/LICENSE.txt for license information. |
| 5 | +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| 6 | +// |
| 7 | +//===----------------------------------------------------------------------===// |
| 8 | +// |
| 9 | +// This header file defines an analysis of affine expressions to compute their |
| 10 | +// ranges (lower/upper bounds) in a given context. |
| 11 | +// |
| 12 | +//===----------------------------------------------------------------------===// |
| 13 | +#ifndef MLIR_ANALYSIS_AFFINEEXPRBOUNDS_H |
| 14 | +#define MLIR_ANALYSIS_AFFINEEXPRBOUNDS_H |
| 15 | + |
| 16 | +#include "mlir/IR/AffineExprVisitor.h" |
| 17 | +#include "mlir/IR/Attributes.h" |
| 18 | +#include "mlir/IR/BuiltinAttributes.h" |
| 19 | +#include "mlir/Interfaces/InferIntRangeInterface.h" |
| 20 | + |
| 21 | +#include "mlir/IR/AffineExpr.h" |
| 22 | +#include "mlir/IR/AffineMap.h" |
| 23 | +#include "mlir/Support/LogicalResult.h" |
| 24 | + |
| 25 | +using namespace mlir; |
| 26 | + |
| 27 | +/// This visitor computes the bounds of affine expressions, using as context the |
| 28 | +/// bounds of the dimensions of the expression. |
| 29 | +/// |
| 30 | +/// Example: |
| 31 | +/// Given bounds 0 <= d0 <= 99 and 0 <= d1 <= 199, we can compute the bounds |
| 32 | +/// of the following expression: |
| 33 | +/// lb(2 * d0 + 3 * d1) = 0 |
| 34 | +/// ub(2 * d0 + 3 * d1) = 795 |
| 35 | +/// |
| 36 | +/// * The bounds given in the context are inclusive, and the bounds returned |
| 37 | +/// are also inclusive. |
| 38 | +/// * If bounds are not available for a dimension, std::nullopt can be used |
| 39 | +/// instead. The bounds of an expression that involves it will be std::nullopt. |
| 40 | +/// * Limitations: |
| 41 | +/// - Parametric expressions (using symbols) are not supported. |
| 42 | +/// - Unsigned FloorDiv is currently not supported. |
| 43 | +class AffineExprBoundsVisitor |
| 44 | + : public AffineExprVisitor<AffineExprBoundsVisitor, LogicalResult> { |
| 45 | +public: |
| 46 | + /// Initialize the context (bounds) with APInt. All bounds must have the same |
| 47 | + /// signedness and bit width. |
| 48 | + AffineExprBoundsVisitor(ArrayRef<std::optional<APInt>> constLowerBounds, |
| 49 | + ArrayRef<std::optional<APInt>> constUpperBounds, |
| 50 | + bool boundsSigned, uint64_t bitWidth, |
| 51 | + MLIRContext *context); |
| 52 | + |
| 53 | + /// Initialize the context (bounds) with 64-bit signed integers. This allows |
| 54 | + /// to directly map index-type values such as Linalg op bounds, which are |
| 55 | + /// represented as int64_t. |
| 56 | + AffineExprBoundsVisitor(ArrayRef<std::optional<int64_t>> constLowerBounds, |
| 57 | + ArrayRef<std::optional<int64_t>> constUpperBounds, |
| 58 | + MLIRContext *context); |
| 59 | + |
| 60 | + /// Get the upper bound of \p expr using the context bounds. |
| 61 | + std::optional<APInt> getUpperBound(AffineExpr expr); |
| 62 | + std::optional<int64_t> getIndexUpperBound(AffineExpr expr); |
| 63 | + |
| 64 | + /// Get the lower bound of \p expr using the context bounds. |
| 65 | + std::optional<APInt> getLowerBound(AffineExpr expr); |
| 66 | + std::optional<int64_t> getIndexLowerBound(AffineExpr expr); |
| 67 | + |
| 68 | + // These methods are directly called by the AffineExprVisitor base class. |
| 69 | + LogicalResult visitMulExpr(AffineBinaryOpExpr expr); |
| 70 | + LogicalResult visitAddExpr(AffineBinaryOpExpr expr); |
| 71 | + LogicalResult visitDimExpr(AffineDimExpr expr); |
| 72 | + LogicalResult visitSymbolExpr(AffineSymbolExpr expr); |
| 73 | + LogicalResult visitConstantExpr(AffineConstantExpr expr); |
| 74 | + LogicalResult visitCeilDivExpr(AffineBinaryOpExpr expr); |
| 75 | + LogicalResult visitFloorDivExpr(AffineBinaryOpExpr expr); |
| 76 | + LogicalResult visitModExpr(AffineBinaryOpExpr expr); |
| 77 | + |
| 78 | +private: |
| 79 | + bool boundsSigned; |
| 80 | + uint64_t bitWidth; |
| 81 | + void inferBinOpRange( |
| 82 | + AffineBinaryOpExpr expr, |
| 83 | + const std::function<ConstantIntRanges(ArrayRef<ConstantIntRanges>)> |
| 84 | + &opInference); |
| 85 | + |
| 86 | + /// Bounds that have been computed for subexpressions are memoized and reused. |
| 87 | + llvm::DenseMap<AffineExpr, APInt> lb; |
| 88 | + llvm::DenseMap<AffineExpr, APInt> ub; |
| 89 | +}; |
| 90 | + |
| 91 | +#endif // MLIR_ANALYSIS_AFFINEEXPRBOUNDS_H |
0 commit comments