diff --git a/mlir/include/mlir/Dialect/QuantOps/QuantOps.td b/mlir/include/mlir/Dialect/QuantOps/QuantOps.td --- a/mlir/include/mlir/Dialect/QuantOps/QuantOps.td +++ b/mlir/include/mlir/Dialect/QuantOps/QuantOps.td @@ -236,6 +236,54 @@ }]; } +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:$axis); + let results = (outs quant_RealValueType:$res); + + let hasFolder = 1; + + let verifier = [{ + auto inArg = arg().getType().dyn_cast(); + auto outArg = res().getType().dyn_cast(); + 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."; diff --git a/mlir/lib/Dialect/QuantOps/IR/QuantOps.cpp b/mlir/lib/Dialect/QuantOps/IR/QuantOps.cpp --- a/mlir/lib/Dialect/QuantOps/IR/QuantOps.cpp +++ b/mlir/lib/Dialect/QuantOps/IR/QuantOps.cpp @@ -10,10 +10,12 @@ #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" @@ -22,6 +24,7 @@ using namespace mlir; using namespace mlir::quant; using namespace mlir::quant::detail; +using namespace mlir::quantizer; QuantizationDialect::QuantizationDialect(MLIRContext *context) : Dialect(/*name=*/"quant", context) { @@ -42,5 +45,42 @@ return srcScastOp.arg(); } +OpFoldResult QuantizerStatsOp::fold(ArrayRef operands) { + auto ctx = arg().getContext(); + auto constOp = dyn_cast_or_null(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 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 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" diff --git a/mlir/test/Quantizer/stats.mlir b/mlir/test/Quantizer/stats.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Quantizer/stats.mlir @@ -0,0 +1,37 @@ +// 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> +}