Page MenuHomePhabricator

[mlir][sparse] support for negation and subtractions
ClosedPublic

Authored by aartbik on Jul 1 2021, 12:10 PM.

Details

Summary

This revision extends the sparse compiler support from fp/int addition and multiplication to fp/int negation and subtraction, thereby increasing the scope of sparse kernels that can be compiled.

Diff Detail

Event Timeline

aartbik created this revision.Jul 1 2021, 12:10 PM
aartbik requested review of this revision.Jul 1 2021, 12:10 PM
aartbik retitled this revision from [mlir][sparse] add sparse compiler support for unary and binary minus to [mlir][sparse] support for negation and subtractions.Jul 2 2021, 10:10 AM
aartbik edited the summary of this revision. (Show Details)

Initial set of comments. Not yet done with review, but I figured I should give you some comments to read through!

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
127
127–128

I'm not sure I understand the name and doc comment for this function, though that's likely just because I've never heard the term "identity merge" before. It seems like this function takes a set of points and creates a new set of points, with the expression at each new point having type kind and having the corresponding point from s as its only child. Is that correct?

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
639–647

I understand what's happening here, but why isn't there an equivalent NegIOp?

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
26–31

This may matter less after the union-refactoring, but why did you choose to store the negation's child expression in e1 and not e0?

102–104

should there be some assertion here that kind == kNegI || kin == kNegF? It seems like that is an implicit assumption of this function.

199

It seems like this function is only meant to be used for binary ops. If that's the case, could you just assert that kind is a binary operator kind? Additionally, maybe it should be called something like kindToOpSymbol?

304

I'm unsure what this comment means! (Perhaps related to the fact that I don't understand the use of "Id" in the function name, related to the comment above)

aartbik marked 6 inline comments as done.Jul 2 2021, 11:27 AM
aartbik added inline comments.
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
127–128

Exactly. The identity applies to the lattices themselves, but it adds an operator to the expression.
Let me try to come up with a better name (suggestions also welcome!)

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
639–647

The short answer: just to keep code concise and consistent for binary ops only.

Slightly longer:
I had initially special cases for unary NegFOp and NegIOp, but that adds a bit of special cases, especially at the end since std does not support negi (it does negf). So this way, once we "parsed" the input, all logic is consistently dealing with binary ops. We could, when generating code again, map the 0-x back to -x for negf but that is also done by backend optimizations anyway.

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
26–31

Yeah, this code is gone, but I did it so that 0-y and -y are "the same" ;-)

102–104

Agreed. Added!

gussmith23 added inline comments.Jul 2 2021, 11:43 AM
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
127–128

To me it feels almost like a map operator in functional programming. You're mapping an expression constructor over a lattice set, to produce a new lattice set with the given expression type. So perhaps something involving map? Alternatively, if this is only meant to work with unary ops (or even just negations), you could make it even more specific (e.g. constructNegationsFromSet or mapUnaryOpOverSet)

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
639–647

Interesting that std does not support NegI. Makes sense!

gussmith23 added inline comments.Jul 2 2021, 12:48 PM
mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
351

It seems like this function just fails silently/returns None at the moment, e.g. if we encounter an op with one operand that doesn't produce a value on this line. Is that expected behavior?

Finished reading through everything. Just have some questions, but for the most part everything looks good!

aartbik marked 6 inline comments as done.Jul 2 2021, 12:59 PM
aartbik added inline comments.
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
127–128

Since we will (may) use this in the future for other ops too (division for example), and its main goal is to introduce a zero at places where we cannot omit the op (e.g. 0+a = a but 0-a must keep it), how about simply mapZero()?

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
351

Yes, anything that does not build (also unknown binary ops by the way) falls into None

Then in Sparsification we have

// Builds the tensor expression for the Linalg operation in SSA form.
Optional<unsigned> exp = merger.buildTensorExpFromLinalg(op);
if (!exp.hasValue())
  return failure();
aartbik updated this revision to Diff 356254.Jul 2 2021, 1:00 PM
aartbik marked 2 inline comments as done.

first set of comments addressed

gussmith23 requested changes to this revision.Jul 2 2021, 1:06 PM

Just small changes left.

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
127

"a lattice set" maybe? or "lattice sets"?

127–128

I like mapZero!

This revision now requires changes to proceed.Jul 2 2021, 1:06 PM
aartbik updated this revision to Diff 356261.Jul 2 2021, 1:23 PM

renamed into mapZero

gussmith23 accepted this revision.Jul 2 2021, 2:28 PM
This revision is now accepted and ready to land.Jul 2 2021, 2:28 PM
This revision was automatically updated to reflect the committed changes.