Add test op for quantizer statistics.

[Vooblin]

Summary

Related to https://reviews.llvm.org/D73556.

[stellaraccident]

Nice, thanks. Some small changes requested. Also, I would have expected cmake updates since you added a dep.

mlir/include/mlir/Dialect/QuantOps/QuantOps.td

241

Imo, fine to leave it the "by quantizer" suffix.

Since I don't think this op will be used as-is outside of tests, lets note it as such in the description. Also note what code paths this is exercising.

I think that renaming it "TestComputeStats" is more descriptive (includes "test" and outside parties don't necessarily care that this comes from the quantizer support library.

mlir/lib/Dialect/QuantOps/IR/QuantOps.cpp

18

Note for future refactoring: we currently have some support code in the QuantOps dialect and some in the quantizer. We should pick one and move the others.

mlir/test/Quantizer/stats.mlir

5

We don't usually hardcore individual ssa names in tests (there may still be old ones that do). Prefer something like:
%[[CST:.+]] =

(Here and below)

7

Is there a reason to have this name attribute? (Here and below)

mlir/include/mlir/Dialect/QuantOps/QuantOps.td

if (axisStats()) {
         return emitOpError("axisStats must have shape [N,2] "
                            "where N = the slice size defined by the axis dim");
       }
     }
     return success();
   }];
 }
 
+def quant_QuantizerStatsOp : quant_Op<"quantizer_stats", []> {
+  let summary =
+      "Extracts per-layer or per-axis tensor statistics by quantizer.";
+
+  let description = [{
+    Extracts per-layer or per-axis statistics over tensor by quantizer.
+    Statistics include min value, max value, mean and variance of the whole
+    tensor or every slice along the particular axis (attribute 'axis').
+
+    Returns dense tensor with statistics.
+    For per-layer it is 4-elements tensor: [min, max, mean, variance].
+    For per-axis case it is tensor with the shape (N, 4), where N is the
+    number of 'axis' slices.
+  }];
+
+  let arguments = (ins
+    quant_RealValueType:$arg,
+    OptionalAttr<NonNegativeI64Attr>:$axis);
+  let results = (outs quant_RealValueType:$res);
+
+  let hasFolder = 1;
+
+  let verifier = [{
+    auto inArg = arg().getType().dyn_cast<TensorType>();
+    auto outArg = res().getType().dyn_cast<TensorType>();
+    if (!inArg) return emitOpError("arg needs to be tensor type.");
+    if (!outArg) return emitOpError("res needs to be tensor type.");
+
+    // Verify input and output shapes
+    auto inShape = inArg.getShape();
+    auto outShape = outArg.getShape();
+    if (axis()) {
+      uint64_t statAxis = axis().getValue().getLimitedValue();
+      if (inShape.size() < 2)
+        return emitOpError("per-axis: input rank must be more than 1");
+      if (statAxis >= inShape.size())
+        return emitOpError("per-axis: axis must be less then input rank");
+      if ((outShape.size() != 2) || (outShape[0] != inShape[statAxis]) ||
+          (outShape[1] != 4))
+        return emitOpError("per-axis: result shape must be [axisSize, 4]");
+    } else {
+      if ((outShape.size() != 1) || (outShape[0] != 4))
+        return emitOpError("per-layer: result shape must be [4]");
+    }
+    return success();
+  }];
+}
+
 def quant_CoupledRefOp : quant_Op<"coupled_ref", [SameOperandsAndResultType]> {
   let summary =
       "Indicates that one point of the computation is coupled to another.";
 
   let description = [{
     Ordinarily, relationships between ops for the purposes of determining
     compatible quantized types is explicit based on the use-def chain. However,
     in some situations, a use may be separated from its def by arbitrary

mlir/lib/Dialect/QuantOps/IR/QuantOps.cpp

 //===- QuantOps.cpp - Quantization Type and Ops Implementation --*- C++ -*-===//
 //
 // Part of the LLVM 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/QuantOps/QuantOps.h"
 #include "TypeDetail.h"
 
 #include "mlir/Dialect/QuantOps/QuantTypes.h"
+#include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
+#include "mlir/Quantizer/Support/Statistics.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/MathExtras.h"
 #include <numeric>
 
 using namespace mlir;
 using namespace mlir::quant;
 using namespace mlir::quant::detail;
+using namespace mlir::quantizer;
 
 QuantizationDialect::QuantizationDialect(MLIRContext *context)
     : Dialect(/*name=*/"quant", context) {
   addTypes<AnyQuantizedType, UniformQuantizedType,
            UniformQuantizedPerAxisType>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/QuantOps/QuantOps.cpp.inc"
       >();
 }
 
 OpFoldResult StorageCastOp::fold(ArrayRef<Attribute> operands) {
   /// Matches x -> [scast -> scast] -> y, replacing the second scast with the
   /// value of x if the casts invert each other.
   auto srcScastOp = dyn_cast_or_null<StorageCastOp>(arg().getDefiningOp());
   if (!srcScastOp || srcScastOp.arg().getType() != getType())
     return OpFoldResult();
   return srcScastOp.arg();
 }
 
+OpFoldResult QuantizerStatsOp::fold(ArrayRef<Attribute> operands) {
+  auto ctx = arg().getContext();
+  auto constOp = dyn_cast_or_null<ConstantOp>(arg().getDefiningOp());
+  if (!constOp)
+    return OpFoldResult();
+  auto constValue = constOp.getValue();
+  AttributeTensorStatistics stats(constValue);
+  TensorAxisStatistics layerStats;
+  if (axis()) {
+    if (!stats.getForAxis(axis().getValue().getLimitedValue(), layerStats))
+      return OpFoldResult();
+    int64_t axisSize = layerStats.minValuePerAxis.size();
+    auto shape = RankedTensorType::get({axisSize, 4}, FloatType::getF64(ctx));
+    SmallVector<APFloat, 12> resultValues(
+        axisSize * 4, APFloat::getZero(APFloat::IEEEdouble()));
+    for (int64_t i = 0; i < axisSize; ++i) {
+      resultValues[i * 4] = APFloat(layerStats.minValuePerAxis[i]);
+      resultValues[i * 4 + 1] = APFloat(layerStats.maxValuePerAxis[i]);
+      resultValues[i * 4 + 2] = APFloat(layerStats.meanPerAxis[i]);
+      resultValues[i * 4 + 3] = APFloat(layerStats.variancePerAxis[i]);
+    }
+    auto result = DenseElementsAttr::get(shape, resultValues);
+    return result;
+  }
+  if (!stats.get(layerStats))
+    return OpFoldResult();
+  auto shape = RankedTensorType::get(4, FloatType::getF64(ctx));
+  SmallVector<APFloat, 4> resultValues(4,
+                                       APFloat::getZero(APFloat::IEEEdouble()));
+  resultValues[0] = APFloat(layerStats.minValue);
+  resultValues[1] = APFloat(layerStats.maxValue);
+  resultValues[2] = APFloat(layerStats.mean);
+  resultValues[3] = APFloat(layerStats.variance);
+  auto result = DenseElementsAttr::get(shape, resultValues);
+  return result;
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/QuantOps/QuantOps.cpp.inc"

mlir/test/Quantizer/stats.mlir

+// RUN: mlir-opt %s -test-constant-fold -split-input-file | FileCheck %s
+
+// -----
+// CHECK-LABEL: @per_layer
+// CHECK: %cst = constant dense<[1.000000e+00, 1.200000e+01, 6.500000e+00, 0.000000e+00]> : tensor<4xf64>
+func @per_layer() -> tensor<4xf64> {
+  %cst = constant {name = "layer"} dense<[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]]]> : tensor<2x3x2xf32>
+  %result = "quant.quantizer_stats"(%cst) : (tensor<2x3x2xf32>) -> tensor<4xf64>
+  return %result : tensor<4xf64>
+}
+
+// -----
+// CHECK-LABEL: @per_axis0
+// CHECK: %cst = constant dense<{{\[\[}}1.000000e+00, 6.000000e+00, 3.500000e+00, 0.000000e+00], [7.000000e+00, 1.200000e+01, 9.500000e+00, 0.000000e+00]]> : tensor<2x4xf64>
+func @per_axis0() -> tensor<2x4xf64> {
+  %cst = constant {name = "axis0"} dense<[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]]]> : tensor<2x3x2xf32>
+  %result = "quant.quantizer_stats"(%cst){axis = 0} : (tensor<2x3x2xf32>) -> tensor<2x4xf64>
+  return %result : tensor<2x4xf64>
+}
+
+// -----
+// CHECK-LABEL: @per_axis1
+// CHECK: %cst = constant dense<{{\[\[}}1.000000e+00, 8.000000e+00, 4.500000e+00, 0.000000e+00], [3.000000e+00, 1.000000e+01, 6.500000e+00, 0.000000e+00], [5.000000e+00, 1.200000e+01, 8.500000e+00, 0.000000e+00]]> : tensor<3x4xf64>
+func @per_axis1() -> tensor<3x4xf64> {
+  %cst = constant {name = "axis1"} dense<[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]]]> : tensor<2x3x2xf32>
+  %result = "quant.quantizer_stats"(%cst){axis = 1} : (tensor<2x3x2xf32>) -> tensor<3x4xf64>
+  return %result : tensor<3x4xf64>
+}
+
+// -----
+// CHECK-LABEL: @per_axis2
+// CHECK: %cst = constant dense<{{\[\[}}1.000000e+00, 1.100000e+01, 6.000000e+00, 0.000000e+00], [2.000000e+00, 1.200000e+01, 7.000000e+00, 0.000000e+00]]> : tensor<2x4xf64>
+func @per_axis2() -> tensor<2x4xf64> {
+  %cst = constant {name = "axis2"} dense<[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]]]> : tensor<2x3x2xf32>
+  %result = "quant.quantizer_stats"(%cst){axis = 2} : (tensor<2x3x2xf32>) -> tensor<2x4xf64>
+  return %result : tensor<2x4xf64>
+}