87 lines · cpp
1//===- UniformSupport.cpp - Support utilities for uniform quant -----------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/Quant/Utils/UniformSupport.h"10#include "mlir/IR/BuiltinTypes.h"11#include "llvm/ADT/STLExtras.h"12#include <numeric>13 14using namespace mlir;15using namespace mlir::quant;16 17static bool isQuantizablePrimitiveType(Type inputType) {18 return isa<FloatType>(inputType);19}20 21ExpressedToQuantizedConverter22ExpressedToQuantizedConverter::forInputType(Type inputType) {23 if (isa<TensorType, VectorType>(inputType)) {24 Type elementType = cast<ShapedType>(inputType).getElementType();25 if (!isQuantizablePrimitiveType(elementType))26 return ExpressedToQuantizedConverter{inputType, nullptr};27 return ExpressedToQuantizedConverter{inputType, elementType};28 }29 // Supported primitive type (which just is the expressed type).30 if (isQuantizablePrimitiveType(inputType))31 return ExpressedToQuantizedConverter{inputType, inputType};32 // Unsupported.33 return ExpressedToQuantizedConverter{inputType, nullptr};34}35 36Type ExpressedToQuantizedConverter::convert(QuantizedType elementalType) const {37 assert(expressedType && "convert() on unsupported conversion");38 if (auto tensorType = dyn_cast<RankedTensorType>(inputType))39 return RankedTensorType::get(tensorType.getShape(), elementalType);40 if (isa<UnrankedTensorType>(inputType))41 return UnrankedTensorType::get(elementalType);42 if (auto vectorType = dyn_cast<VectorType>(inputType))43 return VectorType::get(vectorType.getShape(), elementalType);44 45 // If the expressed types match, just use the new elemental type.46 if (elementalType.getExpressedType() == expressedType)47 return elementalType;48 // Unsupported.49 return nullptr;50}51 52ElementsAttr53UniformQuantizedPerAxisValueConverter::convert(Attribute realValue) {54 if (auto attr = dyn_cast<DenseFPElementsAttr>(realValue)) {55 return convert(attr);56 }57 // TODO: handles sparse elements attribute58 return nullptr;59}60 61DenseElementsAttr62UniformQuantizedPerAxisValueConverter::convert(DenseFPElementsAttr attr) {63 // Creates the converter for each chunk. Normally the size of the64 // quantization dim is 3, so we can cache all the converters.65 ShapedType type = attr.getType();66 size_t dimSize = type.getDimSize(quantizationDim);67 if (dimSize != scales.size()) {68 return {};69 }70 SmallVector<UniformQuantizedValueConverter, 4> converters;71 converters.reserve(dimSize);72 for (int i = 0, e = dimSize; i != e; ++i) {73 converters.push_back(getPerChunkConverter(i));74 }75 76 // Scan the elements of the dense elements attributes and quantize them by77 // using the right quantization parameters.78 int64_t flattenIndex = 0;79 auto shape = type.getShape();80 int64_t chunkSize = llvm::product_of(shape.drop_front(quantizationDim + 1));81 Type newElementType = IntegerType::get(attr.getContext(), storageBitWidth);82 return attr.mapValues(newElementType, [&](const APFloat &old) {83 int chunkIndex = (flattenIndex++) / chunkSize;84 return converters[chunkIndex % dimSize].quantizeFloatToInt(old);85 });86}87