Page MenuHomePhabricator

[mlir][quantizer] Support quantizing sparse tensors
Needs RevisionPublic

Authored by alexeyr on Feb 17 2020, 5:25 AM.

Details

Summary

Add sparse tensors support for per-axis uniform quantization

Diff Detail

Event Timeline

alexeyr created this revision.Feb 17 2020, 5:25 AM

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.

stellaraccident requested changes to this revision.Mar 10 2020, 5:58 AM

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.

This revision now requires changes to proceed.Mar 10 2020, 5:58 AM
rriddle added inline comments.Mar 10 2020, 3:27 PM
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.