Configuration.h

Go to the documentation of this file.

//===- Configuration.h - Configuration object base classes ------*- 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
//
//===----------------------------------------------------------------------===//
//
// The quantizer is relatively agnostic to source and target dialects, with
// the specific represented by configuration policy objects derived from
// classes in this file.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_QUANTIZER_SUPPORT_CONFIGURATION_H
#define MLIR_QUANTIZER_SUPPORT_CONFIGURATION_H

#include <functional>

#include "mlir/Dialect/QuantOps/QuantTypes.h"
#include "mlir/IR/Identifier.h"
#include "mlir/Quantizer/Support/ConstraintAnalysisGraph.h"
#include "mlir/Quantizer/Support/Metadata.h"
#include "mlir/Quantizer/Support/Rules.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/StringSet.h"

namespace mlir {
class Operation;

namespace quantizer {

class CAGSlice;

class TargetConfiguration {
public:
  static constexpr size_t MaxSchemeIndex = 31;
  using OpHandlerFn = std::function<void(Operation *op, CAGSlice &cag)>;

  TargetConfiguration(SolverContext &context);
  virtual ~TargetConfiguration() = default;

  unsigned addCandidateType(quant::AnyQuantizedType quantizedType,
                            CandidateQuantizedType::Scheme scheme) {
    unsigned ordinal = candidateTypes.size();
    assert(allCandidateTypesMask.size() == ordinal);
    CandidateQuantizedType ct{ordinal, quantizedType, scheme};
    candidateTypes.push_back(ct);
    allCandidateTypesMask.push_back(true);
    return ordinal;
  }

  const CandidateQuantizedType &getCandidateType(unsigned index) const {
    assert(index < candidateTypes.size());
    return candidateTypes[index];
  }

  ArrayRef<CandidateQuantizedType> getCandidateTypes() const {
    return candidateTypes;
  }

  llvm::SmallBitVector getAllCandidateTypesMask() const {
    return allCandidateTypesMask;
  }

  llvm::SmallBitVector
  getCandidateTypeDisabledExceptMask(ArrayRef<unsigned> exceptOrdinals) const {
    llvm::SmallBitVector disabled(allCandidateTypesMask);
    for (unsigned ordinal : exceptOrdinals) {
      disabled.reset(ordinal);
    }
    return disabled;
  }

  template <typename OpTy>
  void addOpHandler(OpHandlerFn fn) {
    addOpHandlerByName(OpTy::getOperationName(), fn);
  }

  template <typename OpTy>
  void addRequireStatsOp() {
    addRequireStatsOpByName(OpTy::getOperationName());
  }

  bool isRequireStatsOp(Operation *op) const;

  template <typename OpTy>
  void addValueIdentityOp() {
    addValueIdentityOpByName(OpTy::getOperationName());
  }

  void handleOp(Operation *op, CAGSlice &cag) const;

  virtual void finalizeAnchors(CAGSlice &cag) const {}

  virtual bool isHandledType(Type t) const = 0;

protected:
  virtual void addValueIdentityOpByName(StringRef opName) = 0;
  void addOpHandlerByName(StringRef name, OpHandlerFn fn);

private:
  void addRequireStatsOpByName(StringRef opName);

  std::vector<CandidateQuantizedType> candidateTypes;

  // A SmallBoolVector with bits set for all known candidate types.
  llvm::SmallBitVector allCandidateTypesMask;

  llvm::StringMap<OpHandlerFn> opHandlers;

  llvm::StringSet<> requireStatsOpNames;
};

} // namespace quantizer
} // namespace mlir

#endif // MLIR_QUANTIZER_SUPPORT_CONFIGURATION_H