AffineMap.h

Go to the documentation of this file.

//===- AffineMap.h - MLIR Affine Map 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
//
//===----------------------------------------------------------------------===//
//
// Affine maps are mathematical functions which map a list of dimension
// identifiers and symbols, to multidimensional affine expressions.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_IR_AFFINE_MAP_H
#define MLIR_IR_AFFINE_MAP_H

#include "mlir/Support/LLVM.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMapInfo.h"

namespace mlir {

namespace detail {
struct AffineMapStorage;
} // end namespace detail

class AffineExpr;
class Attribute;
struct LogicalResult;
class MLIRContext;

class AffineMap {
public:
  using ImplType = detail::AffineMapStorage;

  AffineMap() : map(nullptr) {}
  explicit AffineMap(ImplType *map) : map(map) {}
  AffineMap(const AffineMap &other) : map(other.map) {}
  AffineMap &operator=(const AffineMap &other) = default;

  static AffineMap get(MLIRContext *context);

  static AffineMap get(unsigned dimCount, unsigned symbolCount,
                       ArrayRef<AffineExpr> results);

  static AffineMap getConstantMap(int64_t val, MLIRContext *context);

  static AffineMap getMultiDimIdentityMap(unsigned numDims,
                                          MLIRContext *context);

  static AffineMap getPermutationMap(ArrayRef<unsigned> permutation,
                                     MLIRContext *context);

  MLIRContext *getContext() const;

  explicit operator bool() { return map != nullptr; }
  bool operator==(AffineMap other) const { return other.map == map; }
  bool operator!=(AffineMap other) const { return !(other.map == map); }

  bool isIdentity() const;

  bool isEmpty() const;

  bool isSingleConstant() const;

  int64_t getSingleConstantResult() const;

  // Prints affine map to 'os'.
  void print(raw_ostream &os) const;
  void dump() const;

  unsigned getNumDims() const;
  unsigned getNumSymbols() const;
  unsigned getNumResults() const;
  unsigned getNumInputs() const;

  ArrayRef<AffineExpr> getResults() const;
  AffineExpr getResult(unsigned idx) const;

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

  AffineMap replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
                                  ArrayRef<AffineExpr> symReplacements,
                                  unsigned numResultDims,
                                  unsigned numResultSyms);

  LogicalResult constantFold(ArrayRef<Attribute> operandConstants,
                             SmallVectorImpl<Attribute> &results) const;

  AffineMap compose(AffineMap map);

  bool isProjectedPermutation();

  bool isPermutation();

  AffineMap getSubMap(ArrayRef<unsigned> resultPos);

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

private:
  ImplType *map;

  static AffineMap getImpl(unsigned dimCount, unsigned symbolCount,
                           ArrayRef<AffineExpr> results, MLIRContext *context);
};

// Make AffineExpr hashable.
inline ::llvm::hash_code hash_value(AffineMap arg) {
  return ::llvm::hash_value(arg.map);
}

AffineMap simplifyAffineMap(AffineMap map);

AffineMap inversePermutation(AffineMap map);

AffineMap concatAffineMaps(ArrayRef<AffineMap> maps);

inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) {
  map.print(os);
  return os;
}
} // end namespace mlir

namespace llvm {

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

} // namespace llvm

#endif // MLIR_IR_AFFINE_MAP_H