Statistics.h

Go to the documentation of this file.

//===- Statistics.h - Collects statistics over tensors ----------*- 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 defines adapters for extracting various (per layer and per axis)
// statistics over tensors.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_QUANTIZER_SUPPORT_STATISTICS_H
#define MLIR_QUANTIZER_SUPPORT_STATISTICS_H

#include "mlir/IR/Attributes.h"

namespace mlir {
namespace quantizer {

struct TensorAxisStatistics {
  int64_t sampleSize = 0;
  double minValue = 0;
  double maxValue = 0;
  double mean = 0;
  double variance = 0;

  TensorAxisStatistics() {}
  TensorAxisStatistics(int64_t sampleSize, double minValue, double maxValue,
                       double mean, double variance)
      : sampleSize(sampleSize), minValue(minValue), maxValue(maxValue),
        mean(mean), variance(variance) {}
  void clear() { *this = TensorAxisStatistics(); }
};

class AbstractTensorStatistics {
public:
  virtual ~AbstractTensorStatistics() = default;

  virtual bool get(TensorAxisStatistics &stats) const { return false; }

  virtual bool supportsPerAxis() const { return false; }

  virtual unsigned getAxisCount() const { return 0; }

  virtual bool getForAxis(unsigned axis, TensorAxisStatistics &stats) const {
    return false;
  }
};

class AttributeTensorStatistics : public AbstractTensorStatistics {
public:
  AttributeTensorStatistics(Attribute attr) : attr(attr) {}

  bool get(TensorAxisStatistics &stats) const override;

  // TODO: Implement per-axis.

private:
  Attribute attr;
};

raw_ostream &operator<<(raw_ostream &os, const TensorAxisStatistics &stats);

} // end namespace quantizer
} // end namespace mlir

#endif // MLIR_QUANTIZER_SUPPORT_STATISTICS_H