[mlir][quantizer] Add gathering of per-axis statistics in quantizer.

[stellaraccident]

mlir/include/mlir/Quantizer/Support/Statistics.h

73

Thanks for picking up this old todo. One bit of debt about this specific class is that it doesn't have any tests in tree (there were some e2e tests in a related pre-llvm repo). I'll have a detailed review of what you wrote later today, but I wanted to mention this in case if you feel like agreeing the testing situation.

We could pretty easily make this unit testable by adding a corresponding (test) op to compute statistics, implementing a constant folder for it, and then lit testing that.

Reviving this code, which has gone somewhat unused for a while will be a project of mine in about a month, so in this specific case, I'm open to accepting this patch without extending the test situation if doing isn't something you can take on, but it would be much better to resolve that.

[merge_guards_bot]

Unit tests: pass. 62268 tests passed, 0 failed and 827 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/lib/Quantizer/Support/Statistics.cpp

165

Remove trivial braces.

[Vooblin]

mlir/include/mlir/Quantizer/Support/Statistics.h

73

Thanks for the quick response and I'm sorry to be late with the reply. I think that testing approach with constant folder is great and I suggest I'll do it in a next patch. It'll be great if it's possible to accept these patches separately.

[Vooblin]

Remove trivial braces

[Vooblin]

mlir/lib/Quantizer/Support/Statistics.cpp

165

Done

[merge_guards_bot]

Unit tests: pass. 62268 tests passed, 0 failed and 827 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[stellaraccident]

Lgtm.

Do you need me to commit this on your behalf?

mlir/include/mlir/Quantizer/Support/Statistics.h

73

Ok, given the pre-existing test infra need and the fact that this is effectively experimental code at this point, we can be pragmatic and upgrade the testing in a follow-up. Thanks.

[Vooblin]

Yes, thanks!

mlir/include/mlir/Quantizer/Support/Statistics.h

 /// Statistics about a tensor axis (or the whole tensor).
 struct TensorAxisStatistics {
   int64_t sampleSize = 0;
   double minValue = 0;
   double maxValue = 0;
   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(); }
 };
 
 /// Base class for querying statistics about a tensor.
 class AbstractTensorStatistics {
 public:
   virtual ~AbstractTensorStatistics() = default;
 

 ///   OpaqueElementsAttr (with Float based type)
 ///   SparseElementAttr  (with Float based type)
 class AttributeTensorStatistics : public AbstractTensorStatistics {
 public:
   AttributeTensorStatistics(Attribute attr) : attr(attr) {}
 
   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;
 };
 
 raw_ostream &operator<<(raw_ostream &os, const TensorAxisStatistics &stats);
 
 } // end namespace quantizer
 } // end namespace mlir
 
 #endif // MLIR_QUANTIZER_SUPPORT_STATISTICS_H

mlir/lib/Quantizer/Support/Statistics.cpp

for (uint64_t i = 0, s = shape[dim]; i < s; ++i) {
     double value = attr.getValue<FloatAttr>(indices).getValueAsDouble();
     statistics.minValue = std::min(statistics.minValue, value);
     statistics.maxValue = std::max(statistics.maxValue, value);
     statistics.mean += 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();
+static void collectElementsStatisticsDimForAxis(
+    unsigned axis, ElementsAttr attr, unsigned numElements,
+    ArrayRef<int64_t> shape, SmallVectorImpl<uint64_t> &indices, uint64_t dim,
+  // 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) {
   ShapedType sType = attr.getType();
   if (!sType.hasStaticShape())
     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();
 
+  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);
   statistics.sampleSize = numElements;
 
   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();
     stats = TensorAxisStatistics(1, value, value, value, 0);
     return true;
   } else if (auto eltAttr = attr.dyn_cast<ElementsAttr>()) {
     return getElementsStatistics(eltAttr, stats);
   }
   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;
 }