AffineExpr.h | Explore in Territory

//===- AffineExpr.h - MLIR Affine Expr 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
//
//===----------------------------------------------------------------------===//
//
// An affine expression is an affine combination of dimension identifiers and
// symbols, including ceildiv/floordiv/mod by a constant integer.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_IR_AFFINEEXPR_H
#define MLIR_IR_AFFINEEXPR_H

#include "mlir/IR/Visitors.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseMapInfo.h"
#include "llvm/ADT/Hashing.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include <type_traits>

namespace mlir {

class MLIRContext;
class AffineMap;
class IntegerSet;

namespace detail {

struct AffineExprStorage;
struct AffineBinaryOpExprStorage;
struct AffineDimExprStorage;
struct AffineConstantExprStorage;

} // namespace detail

enum class AffineExprKind { … };

/// Base type for affine expression.
/// AffineExpr's are immutable value types with intuitive operators to
/// operate on chainable, lightweight compositions.
/// An AffineExpr is an interface to the underlying storage type pointer.
class AffineExpr { … };

/// Affine binary operation expression. An affine binary operation could be an
/// add, mul, floordiv, ceildiv, or a modulo operation. (Subtraction is
/// represented through a multiply by -1 and add.) These expressions are always
/// constructed in a simplified form. For eg., the LHS and RHS operands can't
/// both be constants. There are additional canonicalizing rules depending on
/// the op type: see checks in the constructor.
class AffineBinaryOpExpr : public AffineExpr { … };

/// A dimensional identifier appearing in an affine expression.
class AffineDimExpr : public AffineExpr { … };

/// A symbolic identifier appearing in an affine expression.
class AffineSymbolExpr : public AffineExpr { … };

/// An integer constant appearing in affine expression.
class AffineConstantExpr : public AffineExpr { … };

/// Make AffineExpr hashable.
inline ::llvm::hash_code hash_value(AffineExpr arg) { … }

inline AffineExpr operator+(int64_t val, AffineExpr expr) { … }
inline AffineExpr operator*(int64_t val, AffineExpr expr) { … }
inline AffineExpr operator-(int64_t val, AffineExpr expr) { … }

/// These free functions allow clients of the API to not use classes in detail.
AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context);
AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context);
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context);
SmallVector<AffineExpr> getAffineConstantExprs(ArrayRef<int64_t> constants,
                                               MLIRContext *context);
AffineExpr getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs,
                                 AffineExpr rhs);

/// 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 'eq' is expected to be in the
/// format [dims, symbols, locals, constant term].
AffineExpr getAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
                                     unsigned numDims, unsigned numSymbols,
                                     ArrayRef<AffineExpr> localExprs,
                                     MLIRContext *context);

raw_ostream &operator<<(raw_ostream &os, AffineExpr expr);

template <typename U>
constexpr bool AffineExpr::isa() const { … }
template <typename U>
U AffineExpr::dyn_cast() const { … }
template <typename U>
U AffineExpr::dyn_cast_or_null() const { … }
template <typename U>
U AffineExpr::cast() const { … }

/// Simplify an affine expression by flattening and some amount of simple
/// analysis. This has complexity linear in the number of nodes in 'expr'.
/// Returns the simplified expression, which is the same as the input expression
/// if it can't be simplified. When `expr` is semi-affine, a simplified
/// semi-affine expression is constructed in the sorted order of dimension and
/// symbol positions.
AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims,
                              unsigned numSymbols);

namespace detail {
template <int N>
void bindDims(MLIRContext *ctx) { … }

template <int N, typename AffineExprTy, typename... AffineExprTy2>
void bindDims(MLIRContext *ctx, AffineExprTy &e, AffineExprTy2 &...exprs) { … }

template <int N>
void bindSymbols(MLIRContext *ctx) { … }

template <int N, typename AffineExprTy, typename... AffineExprTy2>
void bindSymbols(MLIRContext *ctx, AffineExprTy &e, AffineExprTy2 &...exprs) { … }

} // namespace detail

/// Bind a list of AffineExpr references to DimExpr at positions:
///   [0 .. sizeof...(exprs)]
template <typename... AffineExprTy>
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs) { … }

template <typename AffineExprTy>
void bindDimsList(MLIRContext *ctx, MutableArrayRef<AffineExprTy> exprs) { … }

/// Bind a list of AffineExpr references to SymbolExpr at positions:
///   [0 .. sizeof...(exprs)]
template <typename... AffineExprTy>
void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs) { … }

template <typename AffineExprTy>
void bindSymbolsList(MLIRContext *ctx, MutableArrayRef<AffineExprTy> exprs) { … }

/// Get a lower or upper (depending on `isUpper`) bound for `expr` while using
/// the constant lower and upper bounds for its inputs provided in
/// `constLowerBounds` and `constUpperBounds`. Return std::nullopt if such a
/// bound can't be computed. This method only handles simple sum of product
/// expressions (w.r.t constant coefficients) so as to not depend on anything
/// heavyweight in `Analysis`. Expressions of the form: c0*d0 + c1*d1 + c2*s0 +
/// ... + c_n are handled. Expressions involving floordiv, ceildiv, mod or
/// semi-affine ones will lead a none being returned.
std::optional<int64_t>
getBoundForAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols,
                      ArrayRef<std::optional<int64_t>> constLowerBounds,
                      ArrayRef<std::optional<int64_t>> constUpperBounds,
                      bool isUpper);

} // namespace mlir

namespace llvm {

// AffineExpr hash just like pointers
template <>
struct DenseMapInfo<mlir::AffineExpr> { … };

/// Add support for llvm style casts. We provide a cast between To and From if
/// From is mlir::AffineExpr or derives from it.
CastInfo<To, From, std::enable_if_t<std::is_same_v<mlir::AffineExpr, std::remove_const_t<From>> || std::is_base_of_v<mlir::AffineExpr, From>>>;

} // namespace llvm

#endif // MLIR_IR_AFFINEEXPR_H
llvm/mlir/include/mlir/IR/AffineExpr.h
