[mlir][quantizer] Support quantizing sparse tensors

[alexeyr]

Summary

Add sparse tensors support for per-axis uniform quantization

[alexeyr]

I've added ShapedType::withElementType here because I noticed there is equivalent code in other places as well, but changes to them would be a separate patch. Alternatively, this method could be just defined locally.

On the other hand, as a comment says, getFlattenedIndex method duplicates a protected method on ElementsAttribute; it could be avoided by moving this one to a static method on ElementsAttribute and making the member method call it.

[alexeyr]

Ping

[stellaraccident]

The core of this change lgtm modulo one bit on where to put the getChunkSize helper. I need river to weigh in on the ShapedType api update.

mlir/include/mlir/IR/StandardTypes.h

239

Nice - thank you for not just continuing to repeat this pattern at all call sites (note: when a lot of the code was written, there was no ShapedType in the hierarchy.

River: any objection to this addition?

mlir/lib/Dialect/QuantOps/Utils/UniformSupport.cpp

86

If just a utility function operating in this translation unit, then this can be moved to a static helper function in this cpp file (versus in the header/on the class), making it take an explicit quantization dim parameter).

98

River: do you have any preference about duplicating this code vs updating the attribute API in some way?

115

I'm trying to determine whether there is a legitimate reason to call this function with incorrect parameters and expect a silent failure. I think the answer is no. In that case, it will yield better error messages to thread a Location into here and produce a warning.

Since this wouldn't be an isolated change, fine to leave as is for now.

[rriddle]

mlir/include/mlir/IR/StandardTypes.h

239

Seems reasonable, thanks.

mlir/lib/Dialect/QuantOps/Utils/UniformSupport.cpp

62

nit: Please remove trivial braces.

98

Yeah, can we move this to Support/MathExtras.h and de-duplicate?

114

nit: please drop trivial braces, here and below.

127

nit: Start comments with a capital letter, and use punctuation. There are a few other places as well.

135

nit: The indices of a SparseElementsAttr are guaranteed to be i64, so you can simply all of these by using getValue<int64_t>.

mlir/lib/IR/StandardTypes.cpp

211

nit: Use llvm_unreachable(...) here instead. This method should cover all of the derived classes.

mlir/include/mlir/Dialect/QuantOps/UniformSupport.h

explicit UniformQuantizedPerAxisValueConverter(
     assert(scales.size() == zeroPoints.size());
   }
 
   /// Quantize an Attribute by the quantization parameters. Return nullptr if
   /// the conversion fails or the input array isn't an ElementsAttr.
   ElementsAttr convert(Attribute realValue);
 
 private:
-  /// Quantize an DenseFPElementsAttr by the quantization parameters.
+  /// Quantize a DenseFPElementsAttr by the quantization parameters.
   DenseElementsAttr convert(DenseFPElementsAttr attr);
 
+  /// Quantize a SparseElementsAttr by the quantization parameters.
+  ElementsAttr convert(SparseElementsAttr attr);
+
   /// Get a uniform converter for the index-th chunk along the quantizationDim.
   /// All the elements in this chunk is quantized by the returned converter.
   UniformQuantizedValueConverter getPerChunkConverter(int index) const {
     UniformQuantizedValueConverter converter(scales[index], zeroPoints[index],
                                              clampMin, clampMax,
                                              storageBitWidth, isSigned);
     return converter;
   }
 
+  int64_t getChunkSize(ArrayRef<int64_t> shape) {
+    return std::accumulate(std::next(shape.begin(), quantizationDim + 1),
+                           shape.end(), 1, std::multiplies<int64_t>());
+  }
+
   const ArrayRef<double> scales;
   const ArrayRef<int64_t> zeroPoints;
   const APFloat clampMin;
   const APFloat clampMax;
   const uint32_t storageBitWidth;
   const bool isSigned;
   int32_t quantizationDim;
 };
 
 } // namespace quant
 } // namespace mlir
 
 #endif // MLIR_DIALECT_QUANTOPS_UNIFORMSUPPORT_H_

mlir/include/mlir/IR/StandardTypes.h

public:
   /// Get the total amount of bits occupied by a value of this type.  This does
   /// not take into account any memory layout or widening constraints, e.g. a
   /// vector<3xi57> is reported to occupy 3x57=171 bit, even though in practice
   /// it will likely be stored as in a 4xi64 vector register.  Fail an assertion
   /// if the size cannot be computed statically, i.e. if the type has a dynamic
   /// shape or if its elemental type does not have a known bit width.
   int64_t getSizeInBits() const;
 
+  /// Returns the same kind of type with the same shape and new element type.
+  ShapedType withElementType(Type newElementType) const;
+
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type) {
     return type.getKind() == StandardTypes::Vector ||
            type.getKind() == StandardTypes::RankedTensor ||
            type.getKind() == StandardTypes::UnrankedTensor ||
            type.getKind() == StandardTypes::UnrankedMemRef ||
            type.getKind() == StandardTypes::MemRef;
   }

mlir/lib/Dialect/QuantOps/Utils/UniformSupport.cpp

ExpressedToQuantizedConverter::forInputType(Type inputType) {
   case StandardTypes::RankedTensor:
   case StandardTypes::UnrankedTensor:
   case StandardTypes::Vector: {
     Type elementType = inputType.cast<ShapedType>().getElementType();
     if (!isQuantizablePrimitiveType(elementType)) {
       // Unsupported.
       return ExpressedToQuantizedConverter{inputType, nullptr};
     }
-    return ExpressedToQuantizedConverter{
-        inputType, inputType.cast<ShapedType>().getElementType()};
+    return ExpressedToQuantizedConverter{inputType, elementType};
   }
   }
 }
 
 Type ExpressedToQuantizedConverter::convert(QuantizedType elementalType) const {
   assert(expressedType && "convert() on unsupported conversion");
 
-  switch (inputType.getKind()) {
-  default:
-    if (isQuantizablePrimitiveType(elementalType)) {
-      // For primitives, just use the new elemental type.
-      return elementalType;
-    }
-    // Unsupported.
-    return nullptr;
-  case StandardTypes::RankedTensor:
-    return RankedTensorType::get(inputType.cast<RankedTensorType>().getShape(),
-                                 elementalType);
-  case StandardTypes::UnrankedTensor:
-    return UnrankedTensorType::get(elementalType);
-  case StandardTypes::Vector:
-    return VectorType::get(inputType.cast<VectorType>().getShape(),
-                           elementalType);
-  }
+  if (auto shapedInputType = inputType.dyn_cast<ShapedType>())
+    return shapedInputType.withElementType(elementalType);
+
 }
 
 ElementsAttr
 UniformQuantizedPerAxisValueConverter::convert(Attribute realValue) {
   if (auto attr = realValue.dyn_cast<DenseFPElementsAttr>()) {
     return convert(attr);
   }
-  // TODO(fengliuai): handles sparse elements attribute
+  if (auto attr = realValue.dyn_cast<SparseElementsAttr>()) {
+    return convert(attr);
+  }
   return nullptr;
 }
 
 DenseElementsAttr
 UniformQuantizedPerAxisValueConverter::convert(DenseFPElementsAttr attr) {
   // Creates the converter for each chunk. Normally the size of the
   // quantization dim is 3, so we can cache all the converters.
   ShapedType type = attr.getType();
   size_t dimSize = type.getDimSize(quantizationDim);
   if (dimSize != scales.size()) {
     return {};
   }
   SmallVector<UniformQuantizedValueConverter, 4> converters;
   converters.reserve(dimSize);
   for (int i = 0, e = dimSize; i != e; ++i) {
     converters.push_back(getPerChunkConverter(i));
   }
 
   // Scan the elements of the dense elements attributes and quantize them by
   // using the right quantization parameters.
   int64_t flattenIndex = 0;
-  auto shape = type.getShape();
-  int64_t chunkSize =
-      std::accumulate(std::next(shape.begin(), quantizationDim + 1),
-                      shape.end(), 1, std::multiplies<int64_t>());
+  int64_t chunkSize = getChunkSize(type.getShape());
   Type newElementType = IntegerType::get(storageBitWidth, attr.getContext());
   return attr.mapValues(newElementType, [&](const APFloat &old) {
     int chunkIndex = (flattenIndex++) / chunkSize;
     return converters[chunkIndex % dimSize].quantizeFloatToInt(old);
   });
 }
+
+/// Returns the 1 dimensional flattened row-major index from the given
+/// multi-dimensional index.
+static uint64_t getFlattenedIndex(ArrayRef<int64_t> shape,
+                                  ArrayRef<uint64_t> index) {
+  // Duplicates ElementsAttr::getFlattenedIndex logic
+  size_t rank = shape.size();
+  uint64_t valueIndex = 0;
+  uint64_t dimMultiplier = 1;
+  for (int i = rank - 1; i >= 0; --i) {
+    valueIndex += index[i] * dimMultiplier;
+    dimMultiplier *= shape[i];
+  }
+  return valueIndex;
+}
+
+ElementsAttr
+UniformQuantizedPerAxisValueConverter::convert(SparseElementsAttr attr) {
+  if (auto valuesAttr = attr.getValues().dyn_cast<DenseFPElementsAttr>()) {
+    ShapedType type = attr.getType();
+    size_t dimSize = type.getDimSize(quantizationDim);
+    if (dimSize != scales.size()) {
+      return {};
+    }
+    SmallVector<UniformQuantizedValueConverter, 4> converters;
+    converters.reserve(dimSize);
+    for (int i = 0, e = dimSize; i != e; ++i) {
+      converters.push_back(getPerChunkConverter(i));
+    }
+
+    Type newElementType = IntegerType::get(storageBitWidth, attr.getContext());
+    auto numValues = static_cast<size_t>(valuesAttr.getType().getNumElements());
+    DenseIntElementsAttr indicesAttr = attr.getIndices();
+    SmallVector<APInt, 32> newValues;
+    // if all zero-points are 0, emit a sparse attribute with the same indices
+    if (llvm::all_of(zeroPoints,
+                     [](int64_t zeroPoint) { return zeroPoint == 0; })) {
+      newValues.reserve(numValues);
+      for (size_t i = 0; i < numValues; ++i) {
+        auto value = valuesAttr.getValue<APFloat>(i);
+        uint64_t indexInQuantDim =
+            indicesAttr
+                .getValue<APInt>({i, static_cast<uint64_t>(quantizationDim)})
+                .getZExtValue();
+        APInt quantized = converters[indexInQuantDim].quantizeFloatToInt(value);
+        newValues.push_back(quantized);
+      }
+
+      auto newValuesAttr = DenseIntElementsAttr::get(
+          valuesAttr.getType().withElementType(newElementType), newValues);
+      return SparseElementsAttr::get(type.withElementType(newElementType),
+                                     indicesAttr, newValuesAttr);
+    }
+
+    // there are non-0 zero-points
+    // TODO optimize size in this case (e.g. when _most_ zero-points are 0)
+    auto numElements = static_cast<size_t>(type.getNumElements());
+    newValues.reserve(numElements);
+    ArrayRef<int64_t> shape = type.getShape();
+    size_t chunkSize = getChunkSize(shape);
+    size_t numChunks = numElements / chunkSize;
+    // initialize zero-points vector
+    SmallVector<APInt, 4> apZeroPoints;
+    apZeroPoints.reserve(dimSize);
+    APFloat zero =
+        APFloat::getZero(valuesAttr.getValue<APFloat>(0).getSemantics());
+    for (size_t i = 0; i < dimSize; ++i) {
+      // can't just use APInt(zeroPoints[i]) because it can have wrong bitwidth
+      apZeroPoints.push_back(converters[i].quantizeFloatToInt(zero));
+    }
+    // initialize newValues vector from zero-points
+    for (size_t chunkIndex = 0; chunkIndex < numChunks; ++chunkIndex) {
+      APInt zeroPoint = apZeroPoints[chunkIndex % dimSize];
+      for (size_t i = 0; i < chunkSize; ++i)
+        newValues.push_back(zeroPoint);
+    }
+    // finally update from valuesAttr
+    size_t rank = shape.size();
+    SmallVector<uint64_t, 4> index(rank);
+    auto indicesIter = indicesAttr.getValues<APInt>().begin();
+    for (size_t i = 0; i < numValues; ++i) {
+      auto value = valuesAttr.getValue<APFloat>(i);
+      for (size_t j = 0; j < rank; ++j) {
+        index[j] = (*indicesIter++).getZExtValue();
+      }
+      APInt quantized =
+          converters[index[quantizationDim]].quantizeFloatToInt(value);
+      newValues[getFlattenedIndex(shape, index)] = quantized;
+    }
+
+    return DenseIntElementsAttr::get(type.withElementType(newElementType),
+                                     newValues);
+  }
+
+  return {};
+}

mlir/lib/IR/StandardTypes.cpp

 bool ShapedType::hasStaticShape() const {
   return hasRank() && llvm::none_of(getShape(), isDynamic);
 }
 
 bool ShapedType::hasStaticShape(ArrayRef<int64_t> shape) const {
   return hasStaticShape() && getShape() == shape;
 }
 
+ShapedType ShapedType::withElementType(Type newElementType) const {
+  switch (getKind()) {
+  case StandardTypes::Kind::Vector:
+    return VectorType::get(cast<VectorType>().getShape(), newElementType);
+  case StandardTypes::Kind::RankedTensor:
+    return RankedTensorType::get(cast<RankedTensorType>().getShape(),
+                                 newElementType);
+  case StandardTypes::Kind::UnrankedTensor:
+    return UnrankedTensorType::get(newElementType);
+  case StandardTypes::Kind::MemRef:
+    return MemRefType::Builder(cast<MemRefType>())
+        .setElementType(newElementType);
+  case StandardTypes::Kind::UnrankedMemRef:
+    return UnrankedMemRefType::get(newElementType,
+                                   cast<UnrankedMemRefType>().getMemorySpace());
+  default:
+    return {};
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // VectorType
 //===----------------------------------------------------------------------===//
 
 VectorType VectorType::get(ArrayRef<int64_t> shape, Type elementType) {
   return Base::get(elementType.getContext(), StandardTypes::Vector, shape,
                    elementType);
 }

mlir/test/Dialect/QuantOps/convert-const.mlir

// CHECK: "quant.scast"(%cst_0) : (tensor<2x3xi8>) -> tensor<2x3x!quant.uniform<i8:f32:1, {7.812500e-03:128,1.000000e+00,1.000000e+00:1}>>
   %2 = "quant.dcast"(%1) : (tensor<2x3x!quant.uniform<i8:f32:0, {7.812500e-03:128, 1.0}>>) -> (tensor<2x3xf32>)
 
   %cst0 = constant dense<[[-2.0, -0.5, 0.0], [0.0, 1.0, 2.0]]> : tensor<2x3xf32>
   %3 = "quant.qcast"(%cst0) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32:1, {7.812500e-03:128, 1.0, 1.0:1}>>
   %4 = "quant.dcast"(%3) : (tensor<2x3x!quant.uniform<i8:f32:1, {7.812500e-03:128, 1.0, 1.0:1}>>) -> (tensor<2x3xf32>)
 
   return %2, %4 : tensor<2x3xf32>, tensor<2x3xf32>
 }
+
+// -----
+// Verifies per-axis quantization results for sparse.
+// CHECK-LABEL: per_axis_sparse_quantization
+func @per_axis_sparse_quantization() -> (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) {
+
+// CHECK-NEXT: %[[cst:.*]] = constant sparse<{{\[}}[0, 0], [0, 1], [1, 1], [1, 2]], [-128, -64, 1, 2]> : tensor<2x3xi8>
+// CHECK-NEXT: %[[cst0:.*]] = constant dense<{{\[}}[-128, -1, 1], [127, 1, 3]]> : tensor<2x3xi8>
+// CHECK-NEXT: %[[cst1:.*]] = constant dense<{{\[}}[-128, 64, 127], [0, 1, 2]]> : tensor<2x3xi8>
+// CHECK: "quant.scast"(%[[cst]]) : (tensor<2x3xi8>) -> tensor<2x3x!quant.uniform<i8:f32:0, {7.812500e-03,1.000000e+00}>>
+// CHECK: "quant.scast"(%cst_0) : (tensor<2x3xi8>) -> tensor<2x3x!quant.uniform<i8:f32:1, {7.812500e-03:128,1.000000e+00,1.000000e+00:1}>>
+// CHECK: "quant.scast"(%[[cst1]]) : (tensor<2x3xi8>) -> tensor<2x3x!quant.uniform<i8:f32:0, {7.812500e-03:128,1.000000e+00}>>
+
+  %cst = constant sparse<[[0, 0], [0, 1], [1, 1], [1, 2]], [-2.0, -0.5, 1.0, 2.0]> : tensor<2x3xf32>
+  %1 = "quant.qcast"(%cst) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32:0, {7.812500e-03:0, 1.0}>>
+  %2 = "quant.dcast"(%1) : (tensor<2x3x!quant.uniform<i8:f32:0, {7.812500e-03, 1.0}>>) -> (tensor<2x3xf32>)
+
+  %cst0 = constant sparse<[[0, 0], [0, 1], [1, 1], [1, 2]], [-2.0, -0.5, 1.0, 2.0]> : tensor<2x3xf32>
+  %3 = "quant.qcast"(%cst0) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32:1, {7.812500e-03:128, 1.0, 1.0:1}>>
+  %4 = "quant.dcast"(%3) : (tensor<2x3x!quant.uniform<i8:f32:1, {7.812500e-03:128, 1.0, 1.0:1}>>) -> (tensor<2x3xf32>)
+
+  %cst1 = constant sparse<[[0, 0], [0, 1], [1, 1], [1, 2]], [-2.0, -0.5, 1.0, 2.0]> : tensor<2x3xf32>
+  %5 = "quant.qcast"(%cst1) : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8:f32:0, {7.812500e-03:128, 1.0}>>
+  %6 = "quant.dcast"(%5) : (tensor<2x3x!quant.uniform<i8:f32:0, {7.812500e-03:128, 1.0}>>) -> (tensor<2x3xf32>)
+
+  return %2, %4, %6 : tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>
+}