AffineExpr.h

Go to the documentation of this file.

//===- AffineExpr.h - MLIR Affine Expr Class --------------------*- C++ -*-===//
//
// Part of the MLIR 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_AFFINE_EXPR_H
#define MLIR_IR_AFFINE_EXPR_H

#include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseMapInfo.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 AffineSymbolExprStorage;
struct AffineConstantExprStorage;

} // namespace detail

enum class AffineExprKind {
  Add,
  Mul,
  Mod,
  FloorDiv,
  CeilDiv,

  LAST_AFFINE_BINARY_OP = CeilDiv,

  Constant,
  DimId,
  SymbolId,
};

class AffineExpr {
public:
  using ImplType = detail::AffineExprStorage;

  AffineExpr() : expr(nullptr) {}
  /* implicit */ AffineExpr(const ImplType *expr)
      : expr(const_cast<ImplType *>(expr)) {}

  AffineExpr(const AffineExpr &other) : expr(other.expr) {}
  AffineExpr &operator=(AffineExpr other) {
    expr = other.expr;
    return *this;
  }

  bool operator==(AffineExpr other) const { return expr == other.expr; }
  bool operator!=(AffineExpr other) const { return !(*this == other); }
  bool operator==(int64_t v) const;
  bool operator!=(int64_t v) const { return !(*this == v); }
  explicit operator bool() const { return expr; }

  bool operator!() const { return expr == nullptr; }

  template <typename U> bool isa() const;
  template <typename U> U dyn_cast() const;
  template <typename U> U dyn_cast_or_null() const;
  template <typename U> U cast() const;

  MLIRContext *getContext() const;

  AffineExprKind getKind() const;

  void print(raw_ostream &os) const;
  void dump() const;

  bool isSymbolicOrConstant() const;

  bool isPureAffine() const;

  int64_t getLargestKnownDivisor() const;

  bool isMultipleOf(int64_t factor) const;

  bool isFunctionOfDim(unsigned position) const;

  void walk(std::function<void(AffineExpr)> callback) const;

  AffineExpr replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
                                   ArrayRef<AffineExpr> symReplacements) const;

  AffineExpr operator+(int64_t v) const;
  AffineExpr operator+(AffineExpr other) const;
  AffineExpr operator-() const;
  AffineExpr operator-(int64_t v) const;
  AffineExpr operator-(AffineExpr other) const;
  AffineExpr operator*(int64_t v) const;
  AffineExpr operator*(AffineExpr other) const;
  AffineExpr floorDiv(uint64_t v) const;
  AffineExpr floorDiv(AffineExpr other) const;
  AffineExpr ceilDiv(uint64_t v) const;
  AffineExpr ceilDiv(AffineExpr other) const;
  AffineExpr operator%(uint64_t v) const;
  AffineExpr operator%(AffineExpr other) const;

  AffineExpr compose(AffineMap map) const;

  friend ::llvm::hash_code hash_value(AffineExpr arg);

protected:
  ImplType *expr;
};

class AffineBinaryOpExpr : public AffineExpr {
public:
  using ImplType = detail::AffineBinaryOpExprStorage;
  /* implicit */ AffineBinaryOpExpr(AffineExpr::ImplType *ptr);
  AffineExpr getLHS() const;
  AffineExpr getRHS() const;
};

class AffineDimExpr : public AffineExpr {
public:
  using ImplType = detail::AffineDimExprStorage;
  /* implicit */ AffineDimExpr(AffineExpr::ImplType *ptr);
  unsigned getPosition() const;
};

class AffineSymbolExpr : public AffineExpr {
public:
  using ImplType = detail::AffineDimExprStorage;
  /* implicit */ AffineSymbolExpr(AffineExpr::ImplType *ptr);
  unsigned getPosition() const;
};

class AffineConstantExpr : public AffineExpr {
public:
  using ImplType = detail::AffineConstantExprStorage;
  /* implicit */ AffineConstantExpr(AffineExpr::ImplType *ptr);
  int64_t getValue() const;
};

inline ::llvm::hash_code hash_value(AffineExpr arg) {
  return ::llvm::hash_value(arg.expr);
}

inline AffineExpr operator+(int64_t val, AffineExpr expr) { return expr + val; }
inline AffineExpr operator*(int64_t val, AffineExpr expr) { return expr * val; }
inline AffineExpr operator-(int64_t val, AffineExpr expr) {
  return expr * (-1) + val;
}

AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context);
AffineExpr getAffineSymbolExpr(unsigned position, MLIRContext *context);
AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context);
AffineExpr getAffineBinaryOpExpr(AffineExprKind kind, AffineExpr lhs,
                                 AffineExpr rhs);

AffineExpr toAffineExpr(ArrayRef<int64_t> eq, unsigned numDims,
                        unsigned numSymbols, ArrayRef<AffineExpr> localExprs,
                        MLIRContext *context);

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

template <typename U> bool AffineExpr::isa() const {
  if (std::is_same<U, AffineBinaryOpExpr>::value)
    return getKind() <= AffineExprKind::LAST_AFFINE_BINARY_OP;
  if (std::is_same<U, AffineDimExpr>::value)
    return getKind() == AffineExprKind::DimId;
  if (std::is_same<U, AffineSymbolExpr>::value)
    return getKind() == AffineExprKind::SymbolId;
  if (std::is_same<U, AffineConstantExpr>::value)
    return getKind() == AffineExprKind::Constant;
}
template <typename U> U AffineExpr::dyn_cast() const {
  if (isa<U>())
    return U(expr);
  return U(nullptr);
}
template <typename U> U AffineExpr::dyn_cast_or_null() const {
  return (!*this || !isa<U>()) ? U(nullptr) : U(expr);
}
template <typename U> U AffineExpr::cast() const {
  assert(isa<U>());
  return U(expr);
}

AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims,
                              unsigned numSymbols);

bool getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,
                            unsigned numSymbols,
                            SmallVectorImpl<int64_t> *flattenedExpr);

bool getFlattenedAffineExprs(
    AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs);
bool getFlattenedAffineExprs(
    IntegerSet set, std::vector<SmallVector<int64_t, 8>> *flattenedExprs);

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

template <int N, typename AffineExprTy, typename... AffineExprTy2>
void bindDims(MLIRContext *ctx, AffineExprTy &e, AffineExprTy2 &... exprs) {
  e = getAffineDimExpr(N, ctx);
  bindDims<N + 1, AffineExprTy2 &...>(ctx, exprs...);
}
} // namespace detail

template <typename... AffineExprTy>
void bindDims(MLIRContext *ctx, AffineExprTy &... exprs) {
  detail::bindDims<0>(ctx, exprs...);
}

} // namespace mlir

namespace llvm {

// AffineExpr hash just like pointers
template <> struct DenseMapInfo<mlir::AffineExpr> {
  static mlir::AffineExpr getEmptyKey() {
    auto pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
    return mlir::AffineExpr(static_cast<mlir::AffineExpr::ImplType *>(pointer));
  }
  static mlir::AffineExpr getTombstoneKey() {
    auto pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
    return mlir::AffineExpr(static_cast<mlir::AffineExpr::ImplType *>(pointer));
  }
  static unsigned getHashValue(mlir::AffineExpr val) {
    return mlir::hash_value(val);
  }
  static bool isEqual(mlir::AffineExpr LHS, mlir::AffineExpr RHS) {
    return LHS == RHS;
  }
};

} // namespace llvm

#endif // MLIR_IR_AFFINE_EXPR_H