[Mlir-commits] [mlir] 327e062 - [mlir][quantizer] Add gathering of per-axis statistics in quantizer.
Stella Laurenzo
llvmlistbot at llvm.org
Sat Feb 8 15:17:56 PST 2020
Author: Dmitry Murygin
Date: 2020-02-08T15:17:37-08:00
New Revision: 327e062a026a521d528244c12b38b8d6a2026b77
URL: https://github.com/llvm/llvm-project/commit/327e062a026a521d528244c12b38b8d6a2026b77
DIFF: https://github.com/llvm/llvm-project/commit/327e062a026a521d528244c12b38b8d6a2026b77.diff
LOG: [mlir][quantizer] Add gathering of per-axis statistics in quantizer.
Reviewers: stellaraccident, nicolasvasilache
Reviewed By: stellaraccident
Subscribers: Joonsoo, merge_guards_bot, denis13
Tags: #llvm
Differential Revision: https://reviews.llvm.org/D73556
Added:
Modified:
mlir/include/mlir/Quantizer/Support/Statistics.h
mlir/lib/Quantizer/Support/Statistics.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Quantizer/Support/Statistics.h b/mlir/include/mlir/Quantizer/Support/Statistics.h
index a0d2417cca65..61f241603ff0 100644
--- a/mlir/include/mlir/Quantizer/Support/Statistics.h
+++ b/mlir/include/mlir/Quantizer/Support/Statistics.h
@@ -27,11 +27,25 @@ struct TensorAxisStatistics {
double mean = 0;
double variance = 0;
+ int64_t sampleSizePerAxis = 0;
+ SmallVector<double, 4> minValuePerAxis;
+ SmallVector<double, 4> maxValuePerAxis;
+ SmallVector<double, 4> meanPerAxis;
+ SmallVector<double, 4> variancePerAxis;
+
TensorAxisStatistics() {}
TensorAxisStatistics(int64_t sampleSize, double minValue, double maxValue,
double mean, double variance)
: sampleSize(sampleSize), minValue(minValue), maxValue(maxValue),
mean(mean), variance(variance) {}
+ TensorAxisStatistics(int64_t sampleSize, ArrayRef<double> minValues,
+ ArrayRef<double> maxValues, ArrayRef<double> means,
+ ArrayRef<double> variances)
+ : sampleSizePerAxis(sampleSize),
+ minValuePerAxis(minValues.begin(), minValues.end()),
+ maxValuePerAxis(maxValues.begin(), maxValues.end()),
+ meanPerAxis(means.begin(), means.end()),
+ variancePerAxis(variances.begin(), variances.end()) {}
void clear() { *this = TensorAxisStatistics(); }
};
@@ -70,7 +84,11 @@ class AttributeTensorStatistics : public AbstractTensorStatistics {
bool get(TensorAxisStatistics &stats) const override;
- // TODO: Implement per-axis.
+ bool supportsPerAxis() const override;
+
+ unsigned getAxisCount() const override;
+
+ bool getForAxis(unsigned axis, TensorAxisStatistics &stats) const override;
private:
Attribute attr;
diff --git a/mlir/lib/Quantizer/Support/Statistics.cpp b/mlir/lib/Quantizer/Support/Statistics.cpp
index dca4c8571098..8a6a00e212d3 100644
--- a/mlir/lib/Quantizer/Support/Statistics.cpp
+++ b/mlir/lib/Quantizer/Support/Statistics.cpp
@@ -50,12 +50,53 @@ static void collectElementsStatisticsDim(ElementsAttr attr,
}
}
+static void collectElementsStatisticsDimForAxis(
+ unsigned axis, ElementsAttr attr, unsigned numElements,
+ ArrayRef<int64_t> shape, SmallVectorImpl<uint64_t> &indices, uint64_t dim,
+ TensorAxisStatistics &statistics) {
+ // Recursive terminating condition.
+ if (dim >= shape.size())
+ return;
+
+ // Axis is passed separately
+ if (dim == axis) {
+ collectElementsStatisticsDimForAxis(axis, attr, numElements, shape, indices,
+ dim + 1, statistics);
+ return;
+ }
+
+ // Go to last not axis dim
+ if (dim < (shape.size() - 2) ||
+ (dim == (shape.size() - 2) && axis != (shape.size() - 1))) {
+ // Recurse past dim.
+ for (uint64_t i = 0, s = shape[dim]; i < s; ++i) {
+ indices[dim] = i;
+ collectElementsStatisticsDimForAxis(axis, attr, numElements, shape,
+ indices, dim + 1, statistics);
+ }
+ return;
+ }
+
+ // Pass axis
+ uint64_t axisSize = shape[axis];
+ for (uint64_t axisIdx = 0; axisIdx < axisSize; ++axisIdx) {
+ indices[axis] = axisIdx;
+ // Collection dim.
+ for (uint64_t i = 0, s = shape[dim]; i < s; ++i) {
+ indices[dim] = i;
+ double value = attr.getValue<FloatAttr>(indices).getValueAsDouble();
+ statistics.minValuePerAxis[axisIdx] =
+ std::min(statistics.minValuePerAxis[axisIdx], value);
+ statistics.maxValuePerAxis[axisIdx] =
+ std::max(statistics.maxValuePerAxis[axisIdx], value);
+ statistics.meanPerAxis[axisIdx] += value / numElements;
+ // TODO: Calculate a running variance.
+ }
+ }
+}
+
static bool getElementsStatistics(ElementsAttr attr,
TensorAxisStatistics &statistics) {
- statistics.clear();
- statistics.minValue = std::numeric_limits<double>::infinity();
- statistics.maxValue = -std::numeric_limits<double>::infinity();
-
ShapedType sType = attr.getType();
if (!sType.hasStaticShape())
return false;
@@ -67,6 +108,11 @@ static bool getElementsStatistics(ElementsAttr attr,
indices.resize(sType.getRank());
ArrayRef<int64_t> shape = sType.getShape();
+ statistics.minValue = std::numeric_limits<double>::infinity();
+ statistics.maxValue = -std::numeric_limits<double>::infinity();
+ statistics.mean = 0;
+ statistics.variance = 0;
+
auto numElements = sType.getNumElements();
collectElementsStatisticsDim(attr, numElements, shape, indices, 0,
statistics);
@@ -75,6 +121,35 @@ static bool getElementsStatistics(ElementsAttr attr,
return true;
}
+static bool getElementsStatisticsForAxis(unsigned axis, ElementsAttr attr,
+ TensorAxisStatistics &statistics) {
+ ShapedType sType = attr.getType();
+ if (!sType.hasStaticShape() || axis >= sType.getRank())
+ return false;
+ Type elementTy = sType.getElementType();
+ if (!elementTy.isa<FloatType>())
+ return false;
+
+ SmallVector<uint64_t, 4> indices;
+ indices.resize(sType.getRank());
+ ArrayRef<int64_t> shape = sType.getShape();
+
+ uint64_t axisSize = shape[axis];
+ statistics.minValuePerAxis.assign(axisSize,
+ std::numeric_limits<double>::infinity());
+ statistics.maxValuePerAxis.assign(axisSize,
+ -std::numeric_limits<double>::infinity());
+ statistics.meanPerAxis.assign(axisSize, 0);
+ statistics.variancePerAxis.assign(axisSize, 0);
+
+ uint64_t numElements = sType.getNumElements() / shape[axis];
+ collectElementsStatisticsDimForAxis(axis, attr, numElements, shape, indices,
+ 0, statistics);
+ statistics.sampleSizePerAxis = numElements;
+
+ return true;
+}
+
bool AttributeTensorStatistics::get(TensorAxisStatistics &stats) const {
if (FloatAttr floatAttr = attr.dyn_cast<FloatAttr>()) {
double value = floatAttr.getValueAsDouble();
@@ -86,10 +161,41 @@ bool AttributeTensorStatistics::get(TensorAxisStatistics &stats) const {
return false;
}
+bool AttributeTensorStatistics::supportsPerAxis() const {
+ if (auto eltAttr = attr.dyn_cast<ElementsAttr>())
+ return eltAttr.getType().getRank() > 1;
+ return false;
+}
+
+unsigned AttributeTensorStatistics::getAxisCount() const {
+ if (!supportsPerAxis())
+ return 0;
+ return attr.cast<ElementsAttr>().getType().getRank();
+}
+
+bool AttributeTensorStatistics::getForAxis(unsigned axis,
+ TensorAxisStatistics &stats) const {
+ if (!supportsPerAxis())
+ return false;
+ auto eltAttr = attr.cast<ElementsAttr>();
+ return getElementsStatisticsForAxis(axis, eltAttr, stats);
+}
+
raw_ostream &mlir::quantizer::operator<<(raw_ostream &os,
const TensorAxisStatistics &stats) {
- os << "STATS[sampleSize=" << stats.sampleSize << ", min=" << stats.minValue
- << ", maxValue=" << stats.maxValue << ", mean=" << stats.mean
- << ", variance=" << stats.variance << "]";
+ os << "STATS[sampleSizeLayer=" << stats.sampleSize
+ << ", minValueLayer=" << stats.minValue
+ << ", maxValueLayer=" << stats.maxValue << ", meanLayer=" << stats.mean
+ << ", varianceLayer=" << stats.variance
+ << ", sampleSizePerAxis=" << stats.sampleSizePerAxis << ", statsPerAxis={";
+ for (unsigned i = 0, n = stats.minValuePerAxis.size(); i < n; ++i) {
+ os << "minValue=" << stats.minValuePerAxis[i]
+ << ", maxValue=" << stats.maxValuePerAxis[i]
+ << ", mean=" << stats.meanPerAxis[i]
+ << ", variance=" << stats.variancePerAxis[i];
+ if (i != n - 1)
+ os << "; ";
+ }
+ os << "}]";
return os;
}
More information about the Mlir-commits
mailing list