[PATCH] D74705: [mlir][quantizer] Support quantizing sparse tensors

[Stella Laurenzo via Phabricator via llvm-commits]

stellaraccident requested changes to this revision.
stellaraccident added a comment.
This revision now requires changes to proceed.

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.



================
Comment at: mlir/include/mlir/IR/StandardTypes.h:239
+  /// Returns the same kind of type with the same shape and new element type.
+  ShapedType withElementType(Type newElementType) const;
+
----------------
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?


================
Comment at: mlir/lib/Dialect/QuantOps/Utils/UniformSupport.cpp:86
   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());
----------------
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).


================
Comment at: mlir/lib/Dialect/QuantOps/Utils/UniformSupport.cpp:98
+                                  ArrayRef<uint64_t> index) {
+  // Duplicates ElementsAttr::getFlattenedIndex logic
+  size_t rank = shape.size();
----------------
River: do you have any preference about duplicating this code vs updating the attribute API in some way?


================
Comment at: mlir/lib/Dialect/QuantOps/Utils/UniformSupport.cpp:115
+    if (dimSize != scales.size()) {
+      return {};
+    }
----------------
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.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D74705/new/

https://reviews.llvm.org/D74705