ConstraintAnalysisGraph.h

Go to the documentation of this file.

//===- ConstraintAnalysisGraph.h - Graphs type for constraints --*- 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
//
//===----------------------------------------------------------------------===//
//
// This file provides graph-based data structures for representing anchors
// and constraints between them.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_QUANTIZER_SUPPORT_CONSTRAINTANALYSISGRAPH_H
#define MLIR_QUANTIZER_SUPPORT_CONSTRAINTANALYSISGRAPH_H

#include <utility>
#include <vector>

#include "mlir/IR/Function.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Module.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/Types.h"
#include "mlir/Quantizer/Support/Metadata.h"
#include "llvm/ADT/DenseMap.h"

namespace mlir {
namespace quantizer {

class CAGNode;
class CAGSlice;
class TargetConfiguration;

class CAGNode {
public:
  enum class Kind {
    Anchor,
    OperandAnchor,
    ResultAnchor,
    LastAnchor = ResultAnchor,

    Constraint,
    SolveUniformConstraint,
    UniformPropagateExplicitScale,
    LastConstraint = UniformPropagateExplicitScale,
  };

  // Vector and iterator over nodes.
  using node_vector = SmallVector<CAGNode *, 1>;
  using iterator = node_vector::iterator;
  using const_iterator = node_vector::const_iterator;

  virtual ~CAGNode() = default;

  Kind getKind() const { return kind; }

  int getNodeId() const { return nodeId; }

  bool isDirty() const { return dirty; }
  void markDirty() { dirty = true; }
  void clearDirty() { dirty = false; }

  const_iterator begin() const { return outgoing.begin(); }
  const_iterator end() const { return outgoing.end(); }
  iterator begin() { return outgoing.begin(); }
  iterator end() { return outgoing.end(); }

  const_iterator incoming_begin() const { return incoming.begin(); }
  const_iterator incoming_end() const { return incoming.end(); }
  iterator incoming_begin() { return incoming.begin(); }
  iterator incoming_end() { return incoming.end(); }

  virtual void propagate(SolverContext &solverContext,
                         const TargetConfiguration &config) {}

  virtual void printLabel(raw_ostream &os) const;

  template <typename T> void findChildrenOfKind(SmallVectorImpl<T *> &found) {
    for (CAGNode *child : *this) {
      T *ofKind = dyn_cast<T>(child);
      if (ofKind) {
        found.push_back(ofKind);
      }
    }
  }

  void replaceIncoming(CAGNode *otherNode);

  void addOutgoing(CAGNode *toNode);

  bool isOrphan() const { return incoming.empty() && outgoing.empty(); }

protected:
  CAGNode(Kind kind) : kind(kind) {}

private:
  Kind kind;
  int nodeId = -1;
  node_vector outgoing;
  node_vector incoming;
  bool dirty = false;

  friend class CAGSlice;
};

class CAGAnchorNode : public CAGNode {
public:
  enum class TypeTransformRule {
    Direct,

    DirectStorage,

    ExpressedOnly,
  };

  CAGUniformMetadata &getUniformMetadata() { return uniformMetadata; }
  const CAGUniformMetadata &getUniformMetadata() const {
    return uniformMetadata;
  }

  virtual Operation *getOp() const = 0;
  virtual Value getValue() const = 0;

  static bool classof(const CAGNode *n) {
    return n->getKind() >= Kind::Anchor && n->getKind() <= Kind::LastAnchor;
  }

  void propagate(SolverContext &solverContext,
                 const TargetConfiguration &config) override;

  void printLabel(raw_ostream &os) const override;

  Type getTransformedType();

  TypeTransformRule getTypeTransformRule() const { return typeTransformRule; }

  void setTypeTransformRule(TypeTransformRule r) { typeTransformRule = r; }

  Type getOriginalType() const { return originalType; }

protected:
  CAGAnchorNode(Kind kind, Type originalType)
      : CAGNode(kind), originalType(originalType) {}

private:
  CAGUniformMetadata uniformMetadata;
  Type originalType;
  TypeTransformRule typeTransformRule = TypeTransformRule::Direct;
};

class CAGOperandAnchor : public CAGAnchorNode {
public:
  CAGOperandAnchor(Operation *op, unsigned operandIdx);

  Operation *getOp() const final { return op; }
  unsigned getOperandIdx() const { return operandIdx; }

  static bool classof(const CAGNode *n) {
    return n->getKind() == Kind::Anchor || n->getKind() == Kind::OperandAnchor;
  }

  Value getValue() const final { return op->getOperand(operandIdx); }

  void printLabel(raw_ostream &os) const override;

private:
  Operation *op;
  unsigned operandIdx;
};

class CAGResultAnchor : public CAGAnchorNode {
public:
  CAGResultAnchor(Operation *op, unsigned resultIdx);

  static bool classof(const CAGNode *n) {
    return n->getKind() == Kind::Anchor || n->getKind() == Kind::ResultAnchor;
  }

  Operation *getOp() const final { return resultValue->getDefiningOp(); }
  Value getValue() const final { return resultValue; }

  void printLabel(raw_ostream &os) const override;

private:
  Value resultValue;
};

class CAGConstraintNode : public CAGNode {
public:
  CAGConstraintNode(Kind kind) : CAGNode(kind) {}

  static bool classof(const CAGNode *n) {
    return n->getKind() >= Kind::Constraint &&
           n->getKind() <= Kind::LastConstraint;
  }
};

class CAGSlice {
public:
  CAGSlice(SolverContext &context);
  ~CAGSlice();

  using node_vector = std::vector<CAGNode *>;
  using iterator = node_vector::iterator;
  using const_iterator = node_vector::const_iterator;

  iterator begin() { return allNodes.begin(); }
  iterator end() { return allNodes.end(); }
  const_iterator begin() const { return allNodes.begin(); }
  const_iterator end() const { return allNodes.end(); }

  CAGOperandAnchor *getOperandAnchor(Operation *op, unsigned operandIdx);

  CAGResultAnchor *getResultAnchor(Operation *op, unsigned resultIdx);

  template <typename T, typename... Args>
  T *addUniqueConstraint(ArrayRef<CAGAnchorNode *> anchors, Args... args) {
    static_assert(std::is_convertible<T *, CAGConstraintNode *>(),
                  "T must be a CAGConstraingNode");
    T *constraintNode = addNode(std::make_unique<T>(args...));
    for (auto *anchor : anchors)
      anchor->addOutgoing(constraintNode);
    return constraintNode;
  }

  template <typename T, typename... Args>
  T *addUnidirectionalConstraint(CAGAnchorNode *fromAnchor,
                                 ArrayRef<CAGAnchorNode *> toAnchors,
                                 Args... args) {
    static_assert(std::is_convertible<T *, CAGConstraintNode *>(),
                  "T must be a CAGConstraingNode");
    T *constraintNode = addNode(std::make_unique<T>(args...));
    fromAnchor->addOutgoing(constraintNode);
    for (auto *toAnchor : toAnchors) {
      constraintNode->addOutgoing(toAnchor);
    }
    return constraintNode;
  }

  template <typename T>
  T *addClusteredConstraint(ArrayRef<CAGAnchorNode *> anchors) {
    static_assert(std::is_convertible<T *, CAGConstraintNode *>(),
                  "T must be a CAGConstraingNode");
    SmallVector<T *, 8> cluster;
    for (auto *anchor : anchors) {
      anchor->findChildrenOfKind<T>(cluster);
    }

    T *constraintNode;
    if (cluster.empty()) {
      // Create new.
      constraintNode = addNode(std::make_unique<T>());
    } else {
      // Merge existing.
      constraintNode = cluster[0];
      for (size_t i = 1, e = cluster.size(); i < e; ++i) {
        cluster[i]->replaceIncoming(constraintNode);
      }
    }
    for (auto *anchor : anchors) {
      anchor->addOutgoing(constraintNode);
    }
    return constraintNode;
  }

  void enumerateImpliedConnections(
      std::function<void(CAGAnchorNode *from, CAGAnchorNode *to)> callback);

  unsigned propagate(const TargetConfiguration &config);

private:
  template <typename T>
  T *addNode(std::unique_ptr<T> node) {
    node->nodeId = allNodes.size();
    T *unownedNode = node.release();
    allNodes.push_back(unownedNode);
    return unownedNode;
  }

  SolverContext &context;
  std::vector<CAGNode *> allNodes;
  DenseMap<std::pair<Operation *, unsigned>, CAGOperandAnchor *> operandAnchors;
  DenseMap<std::pair<Operation *, unsigned>, CAGResultAnchor *> resultAnchors;
};

inline raw_ostream &operator<<(raw_ostream &os, const CAGNode &node) {
  node.printLabel(os);
  return os;
}

} // namespace quantizer
} // namespace mlir

#endif // MLIR_QUANTIZER_SUPPORT_CONSTRAINTANALYSISGRAPH_H