SDBMExpr.h

Go to the documentation of this file.

//===- SDBMExpr.h - MLIR SDBM Expression ------------------------*- 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
//
//===----------------------------------------------------------------------===//
//
// A striped difference-bound matrix (SDBM) expression is a constant expression,
// an identifier, a binary expression with constant RHS and +, stripe operators
// or a difference expression between two identifiers.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_DIALECT_SDBM_SDBMEXPR_H
#define MLIR_DIALECT_SDBM_SDBMEXPR_H

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

namespace mlir {

class AffineExpr;
class MLIRContext;

enum class SDBMExprKind { Add, Stripe, Diff, Constant, DimId, SymbolId, Neg };

namespace detail {
struct SDBMExprStorage;
struct SDBMBinaryExprStorage;
struct SDBMDiffExprStorage;
struct SDBMTermExprStorage;
struct SDBMConstantExprStorage;
struct SDBMNegExprStorage;
} // namespace detail

class SDBMConstantExpr;
class SDBMDialect;
class SDBMDimExpr;
class SDBMSymbolExpr;
class SDBMTermExpr;

class SDBMExpr {
public:
  using ImplType = detail::SDBMExprStorage;
  SDBMExpr() : impl(nullptr) {}
  /* implicit */ SDBMExpr(ImplType *expr) : impl(expr) {}

  SDBMExpr(const SDBMExpr &) = default;
  SDBMExpr &operator=(const SDBMExpr &) = default;

  bool operator==(const SDBMExpr &other) const { return impl == other.impl; }
  bool operator!=(const SDBMExpr &other) const { return !(*this == other); }

  explicit operator bool() const { return impl != nullptr; }
  bool operator!() const { return !static_cast<bool>(*this); }

  SDBMExpr operator-();

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

  template <typename U> bool isa() const { return U::isClassFor(*this); }
  template <typename U> U dyn_cast() const {
    if (!isa<U>())
      return {};
    return U(const_cast<SDBMExpr *>(this)->impl);
  }
  template <typename U> U cast() const {
    assert(isa<U>() && "cast to incorrect subtype");
    return U(const_cast<SDBMExpr *>(this)->impl);
  }

  ::llvm::hash_code hash_value() const { return ::llvm::hash_value(impl); }

  SDBMExprKind getKind() const;

  MLIRContext *getContext() const;

  SDBMDialect *getDialect() const;

  AffineExpr getAsAffineExpr() const;

  static Optional<SDBMExpr> tryConvertAffineExpr(AffineExpr affine);

protected:
  ImplType *impl;
};

class SDBMConstantExpr : public SDBMExpr {
public:
  using ImplType = detail::SDBMConstantExprStorage;

  using SDBMExpr::SDBMExpr;

  static SDBMConstantExpr get(SDBMDialect *dialect, int64_t value);

  static bool isClassFor(const SDBMExpr &expr) {
    return expr.getKind() == SDBMExprKind::Constant;
  }

  int64_t getValue() const;
};

class SDBMVaryingExpr : public SDBMExpr {
public:
  using ImplType = detail::SDBMExprStorage;
  using SDBMExpr::SDBMExpr;

  static bool isClassFor(const SDBMExpr &expr) {
    return expr.getKind() == SDBMExprKind::DimId ||
           expr.getKind() == SDBMExprKind::SymbolId ||
           expr.getKind() == SDBMExprKind::Neg ||
           expr.getKind() == SDBMExprKind::Stripe ||
           expr.getKind() == SDBMExprKind::Add ||
           expr.getKind() == SDBMExprKind::Diff;
  }
};

class SDBMDirectExpr : public SDBMVaryingExpr {
public:
  using SDBMVaryingExpr::SDBMVaryingExpr;

  SDBMTermExpr getTerm();

  int64_t getConstant();

  static bool isClassFor(const SDBMExpr &expr) {
    return expr.getKind() == SDBMExprKind::DimId ||
           expr.getKind() == SDBMExprKind::SymbolId ||
           expr.getKind() == SDBMExprKind::Stripe ||
           expr.getKind() == SDBMExprKind::Add;
  }
};

class SDBMTermExpr : public SDBMDirectExpr {
public:
  using SDBMDirectExpr::SDBMDirectExpr;

  static bool isClassFor(const SDBMExpr &expr) {
    return expr.getKind() == SDBMExprKind::DimId ||
           expr.getKind() == SDBMExprKind::SymbolId ||
           expr.getKind() == SDBMExprKind::Stripe;
  }
};

class SDBMSumExpr : public SDBMDirectExpr {
public:
  using ImplType = detail::SDBMBinaryExprStorage;
  using SDBMDirectExpr::SDBMDirectExpr;

  static SDBMSumExpr get(SDBMTermExpr lhs, SDBMConstantExpr rhs);

  static bool isClassFor(const SDBMExpr &expr) {
    SDBMExprKind kind = expr.getKind();
    return kind == SDBMExprKind::Add;
  }

  SDBMTermExpr getLHS() const;
  SDBMConstantExpr getRHS() const;
};

class SDBMDiffExpr : public SDBMVaryingExpr {
public:
  using ImplType = detail::SDBMDiffExprStorage;
  using SDBMVaryingExpr::SDBMVaryingExpr;

  static SDBMDiffExpr get(SDBMDirectExpr lhs, SDBMTermExpr rhs);

  static bool isClassFor(const SDBMExpr &expr) {
    return expr.getKind() == SDBMExprKind::Diff;
  }

  SDBMDirectExpr getLHS() const;
  SDBMTermExpr getRHS() const;
};

class SDBMStripeExpr : public SDBMTermExpr {
public:
  using ImplType = detail::SDBMBinaryExprStorage;
  using SDBMTermExpr::SDBMTermExpr;

  static bool isClassFor(const SDBMExpr &expr) {
    return expr.getKind() == SDBMExprKind::Stripe;
  }

  static SDBMStripeExpr get(SDBMDirectExpr var, SDBMConstantExpr stripeFactor);

  SDBMDirectExpr getLHS() const;
  SDBMConstantExpr getStripeFactor() const;
};

class SDBMInputExpr : public SDBMTermExpr {
public:
  using ImplType = detail::SDBMTermExprStorage;
  using SDBMTermExpr::SDBMTermExpr;

  static bool isClassFor(const SDBMExpr &expr) {
    return expr.getKind() == SDBMExprKind::DimId ||
           expr.getKind() == SDBMExprKind::SymbolId;
  }

  unsigned getPosition() const;
};

class SDBMDimExpr : public SDBMInputExpr {
public:
  using ImplType = detail::SDBMTermExprStorage;
  using SDBMInputExpr::SDBMInputExpr;

  static SDBMDimExpr get(SDBMDialect *dialect, unsigned position);

  static bool isClassFor(const SDBMExpr &expr) {
    return expr.getKind() == SDBMExprKind::DimId;
  }
};

class SDBMSymbolExpr : public SDBMInputExpr {
public:
  using ImplType = detail::SDBMTermExprStorage;
  using SDBMInputExpr::SDBMInputExpr;

  static SDBMSymbolExpr get(SDBMDialect *dialect, unsigned position);

  static bool isClassFor(const SDBMExpr &expr) {
    return expr.getKind() == SDBMExprKind::SymbolId;
  }
};

class SDBMNegExpr : public SDBMVaryingExpr {
public:
  using ImplType = detail::SDBMNegExprStorage;
  using SDBMVaryingExpr::SDBMVaryingExpr;

  static SDBMNegExpr get(SDBMDirectExpr var);

  static bool isClassFor(const SDBMExpr &expr) {
    return expr.getKind() == SDBMExprKind::Neg;
  }

  SDBMDirectExpr getVar() const;
};

template <typename Derived, typename Result = void> class SDBMVisitor {
public:
  Result visit(SDBMExpr expr) {
    auto *derived = static_cast<Derived *>(this);
    switch (expr.getKind()) {
    case SDBMExprKind::Add:
    case SDBMExprKind::Diff:
    case SDBMExprKind::DimId:
    case SDBMExprKind::SymbolId:
    case SDBMExprKind::Neg:
    case SDBMExprKind::Stripe:
      return derived->visitVarying(expr.cast<SDBMVaryingExpr>());
    case SDBMExprKind::Constant:
      return derived->visitConstant(expr.cast<SDBMConstantExpr>());
    }

    llvm_unreachable("unsupported SDBM expression kind");
  }

  void walkPreorder(SDBMExpr expr) { return walk</*isPreorder=*/true>(expr); }

  void walkPostorder(SDBMExpr expr) { return walk</*isPreorder=*/false>(expr); }

protected:
  void visitSum(SDBMSumExpr) {}
  void visitDiff(SDBMDiffExpr) {}
  void visitStripe(SDBMStripeExpr) {}
  void visitDim(SDBMDimExpr) {}
  void visitSymbol(SDBMSymbolExpr) {}
  void visitNeg(SDBMNegExpr) {}
  void visitConstant(SDBMConstantExpr) {}

  Result visitDirect(SDBMDirectExpr expr) {
    auto *derived = static_cast<Derived *>(this);
    if (auto sum = expr.dyn_cast<SDBMSumExpr>())
      return derived->visitSum(sum);
    else
      return derived->visitTerm(expr.cast<SDBMTermExpr>());
  }

  Result visitTerm(SDBMTermExpr expr) {
    auto *derived = static_cast<Derived *>(this);
    if (expr.getKind() == SDBMExprKind::Stripe)
      return derived->visitStripe(expr.cast<SDBMStripeExpr>());
    else
      return derived->visitInput(expr.cast<SDBMInputExpr>());
  }

  Result visitInput(SDBMInputExpr expr) {
    auto *derived = static_cast<Derived *>(this);
    if (expr.getKind() == SDBMExprKind::DimId)
      return derived->visitDim(expr.cast<SDBMDimExpr>());
    else
      return derived->visitSymbol(expr.cast<SDBMSymbolExpr>());
  }

  Result visitVarying(SDBMVaryingExpr expr) {
    auto *derived = static_cast<Derived *>(this);
    if (auto var = expr.dyn_cast<SDBMDirectExpr>())
      return derived->visitDirect(var);
    else if (auto neg = expr.dyn_cast<SDBMNegExpr>())
      return derived->visitNeg(neg);
    else if (auto diff = expr.dyn_cast<SDBMDiffExpr>())
      return derived->visitDiff(diff);

    llvm_unreachable("unhandled subtype of varying SDBM expression");
  }

  template <bool isPreorder> void walk(SDBMExpr expr) {
    if (isPreorder)
      visit(expr);
    if (auto sumExpr = expr.dyn_cast<SDBMSumExpr>()) {
      walk<isPreorder>(sumExpr.getLHS());
      walk<isPreorder>(sumExpr.getRHS());
    } else if (auto diffExpr = expr.dyn_cast<SDBMDiffExpr>()) {
      walk<isPreorder>(diffExpr.getLHS());
      walk<isPreorder>(diffExpr.getRHS());
    } else if (auto stripeExpr = expr.dyn_cast<SDBMStripeExpr>()) {
      walk<isPreorder>(stripeExpr.getLHS());
      walk<isPreorder>(stripeExpr.getStripeFactor());
    } else if (auto negExpr = expr.dyn_cast<SDBMNegExpr>()) {
      walk<isPreorder>(negExpr.getVar());
    }
    if (!isPreorder)
      visit(expr);
  }
};

namespace ops_assertions {

SDBMExpr operator+(SDBMExpr lhs, SDBMExpr rhs);
inline SDBMExpr operator+(SDBMExpr lhs, int64_t rhs) {
  return lhs + SDBMConstantExpr::get(lhs.getDialect(), rhs);
}
inline SDBMExpr operator+(int64_t lhs, SDBMExpr rhs) {
  return SDBMConstantExpr::get(rhs.getDialect(), lhs) + rhs;
}

SDBMExpr operator-(SDBMExpr lhs, SDBMExpr rhs);
inline SDBMExpr operator-(SDBMExpr lhs, int64_t rhs) {
  return lhs - SDBMConstantExpr::get(lhs.getDialect(), rhs);
}
inline SDBMExpr operator-(int64_t lhs, SDBMExpr rhs) {
  return SDBMConstantExpr::get(rhs.getDialect(), lhs) - rhs;
}

SDBMExpr stripe(SDBMExpr expr, SDBMExpr factor);
inline SDBMExpr stripe(SDBMExpr expr, int64_t factor) {
  return stripe(expr, SDBMConstantExpr::get(expr.getDialect(), factor));
}
} // namespace ops_assertions

} // end namespace mlir

namespace llvm {
// SDBMExpr hash just like pointers.
template <> struct DenseMapInfo<mlir::SDBMExpr> {
  static mlir::SDBMExpr getEmptyKey() {
    auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
    return mlir::SDBMExpr(static_cast<mlir::SDBMExpr::ImplType *>(pointer));
  }
  static mlir::SDBMExpr getTombstoneKey() {
    auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
    return mlir::SDBMExpr(static_cast<mlir::SDBMExpr::ImplType *>(pointer));
  }
  static unsigned getHashValue(mlir::SDBMExpr expr) {
    return expr.hash_value();
  }
  static bool isEqual(mlir::SDBMExpr lhs, mlir::SDBMExpr rhs) {
    return lhs == rhs;
  }
};

// SDBMDirectExpr hash just like pointers.
template <> struct DenseMapInfo<mlir::SDBMDirectExpr> {
  static mlir::SDBMDirectExpr getEmptyKey() {
    auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
    return mlir::SDBMDirectExpr(
        static_cast<mlir::SDBMExpr::ImplType *>(pointer));
  }
  static mlir::SDBMDirectExpr getTombstoneKey() {
    auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
    return mlir::SDBMDirectExpr(
        static_cast<mlir::SDBMExpr::ImplType *>(pointer));
  }
  static unsigned getHashValue(mlir::SDBMDirectExpr expr) {
    return expr.hash_value();
  }
  static bool isEqual(mlir::SDBMDirectExpr lhs, mlir::SDBMDirectExpr rhs) {
    return lhs == rhs;
  }
};

// SDBMTermExpr hash just like pointers.
template <> struct DenseMapInfo<mlir::SDBMTermExpr> {
  static mlir::SDBMTermExpr getEmptyKey() {
    auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
    return mlir::SDBMTermExpr(static_cast<mlir::SDBMExpr::ImplType *>(pointer));
  }
  static mlir::SDBMTermExpr getTombstoneKey() {
    auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
    return mlir::SDBMTermExpr(static_cast<mlir::SDBMExpr::ImplType *>(pointer));
  }
  static unsigned getHashValue(mlir::SDBMTermExpr expr) {
    return expr.hash_value();
  }
  static bool isEqual(mlir::SDBMTermExpr lhs, mlir::SDBMTermExpr rhs) {
    return lhs == rhs;
  }
};

// SDBMConstantExpr hash just like pointers.
template <> struct DenseMapInfo<mlir::SDBMConstantExpr> {
  static mlir::SDBMConstantExpr getEmptyKey() {
    auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
    return mlir::SDBMConstantExpr(
        static_cast<mlir::SDBMExpr::ImplType *>(pointer));
  }
  static mlir::SDBMConstantExpr getTombstoneKey() {
    auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
    return mlir::SDBMConstantExpr(
        static_cast<mlir::SDBMExpr::ImplType *>(pointer));
  }
  static unsigned getHashValue(mlir::SDBMConstantExpr expr) {
    return expr.hash_value();
  }
  static bool isEqual(mlir::SDBMConstantExpr lhs, mlir::SDBMConstantExpr rhs) {
    return lhs == rhs;
  }
};
} // namespace llvm

#endif // MLIR_DIALECT_SDBM_SDBMEXPR_H