This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] Introduce new binary and unary op for sparse_tensor
ClosedPublic

Authored by jim22k on Mar 4 2022, 12:06 PM.

Details

Summary

binary performs a sparse binary operation within linalg.generic,
providing the flexibility to do intersection or union or even more
advanced (ex. A-B -> -B when A is missing).

unary performs a sparse unary operation within linalg.generic.
Both the "present" and "missing" values can return a result, allowing
for a simple apply, converting sparse to dense (e.g. A+1), or even
performing a sparse mask inversion.

Diff Detail

Event Timeline

jim22k created this revision.Mar 4 2022, 12:06 PM
Herald added a project: Restricted Project. · View Herald Transcript
jim22k requested review of this revision.Mar 4 2022, 12:06 PM
jim22k updated this revision to Diff 413098.Mar 4 2022, 12:18 PM

Introduce new binary and unary op for sparse_tensor dialect

binary performs a sparse binary operation within linalg.generic,
providing the flexibility to do intersection or union or even more
advanced things (ex. A-B => -B when A is missing)

unary performs a sparse unary operation within linalg.generic.
Both the "present" and "missing" values can return a result, allowing
for a simple apply, converting sparse to dense (e.g. A+1), or even
performing a sparse mask inversion.

jim22k retitled this revision from Minor changes to Introduce new binary and unary op for sparse_tensor.Mar 4 2022, 12:19 PM
jim22k edited the summary of this revision. (Show Details)
jim22k updated this revision to Diff 413123.Mar 4 2022, 1:20 PM

Fix Formatting

What happened to the roundtrip.mlir and invalid.mlir files?

jim22k added a comment.Mar 4 2022, 2:48 PM

They are still there in the history. I did 3 git commits locally: the big one, a minor fix, and then a clang-format one.
While I appreciate that you can view all 3 separate, I don't see a way to view all changes at once, which feels like the most important one to view.
Look at Diff 2 for the real changes.

aartbik added a comment.EditedMar 4 2022, 5:35 PM

They are still there in the history. I did 3 git commits locally: the big one, a minor fix, and then a clang-format one.

I may be wrong, but that is not how phabricator reviews work, or at least I have never seen it being used like this.
One typically uploads a single differential for review (or use a patch series for a progression, but each forms a new differential).
The history in one differential is just to see how comment were addressed. What you see at the last version is what will go in eventually.

(and yes, that is different from typical github PR, where you can just stack commits).

jim22k added a comment.Mar 4 2022, 6:53 PM

Okay, so my understanding was tainted by how a Github PR works. I'll try to update the differential to include all 3 commits, rather than the latest commit (which is what arc diff defaults to).

jim22k updated this revision to Diff 413183.Mar 4 2022, 6:54 PM

Try again with all the changes this time

This looks good, yes, thanks for changing to a single differential.
Please give me a day to or two to review carefully.
But I am very excited about this contribution already!

jim22k added inline comments.Mar 6 2022, 4:15 PM
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
474

Remove SameTypeOperands here. Copy-paste error from binary, but doesn't make sense in the context of unary which only has a single operand.

521

Remove %b. Copy-paste error from binary. Unary only takes a single argument.

Could we split this into multiple independent PRs ?
I find the unary op by itself has enough content and room for design discussion that it should be isolated from the binary op.
Looking at the unary op only for now.

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
479

It would be nice to reshuffle the explanation a bit and add a missing paragraph here that explains the semantics of the op in the absence of linalg.generic.

Based on the explanation, I should be able to provide some advice on the include_index which is currently very implicit.
It would likely be better to connect it to linalg.index operations by passing the SSA values explicitly to sparse_tensor.unary where needed.

Maybe it would help the rephrase the documentation to think how you'd explain the op semantics in the presence of explicit linalg.index operations ?

487

It is unclear what a "set region" is.
Some examples with sparse tensors containing values would be useful.

For instance I do not understand the concept of missing values (and your second example).
This seems to rely on a notional "densified set" that would be [0, max_index(input_tensor)) but I cannot properly infer this from the text.

Taking the example A = [0.0f@1, 2.5f@42] (where @idx represents the index of an element), is "missing" iterating on the values {0, [2..42]} or something else?

494

this seems inaccurate given your 2nd example where the primary region also has indices.

Thanks Jim, here is a first round of feedback.

Nicolas, please note that the original PR was much larger, and I already had ask to break it down into smaller pieces.
Introducing the ops felt like a manageable chunk of code. But please let me know if you feel strongly about splitting this up even more in unary/binary.

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
392

The "Fulfills a need' sounds a bit like a design comment, not documentation of the op.
How about:

Defines a computation within a linalg.generic operation that takes two operands and executes of of the regions depending on whether both operands or either operand is nonzero (i.e., stored explicitly in either sparse storage format).

404–411

I think the intersection/union part needs some more explanation.
First, we have four cases now

primary, no left, no right intersection-flavor
primary, left, no right
left union flavor
primary, no left, right right union flavor
primary, left, right
union flavor

Then, independent of that, when either left or right is set, we can "identify" as a shorthand for returning the input parameter "pass through". Using identity in the description of union was a bit confusion, since it is really the presence of a block that determines the action.

I am also wondering whether an extra attribute in the op itself would be useful (an enum that specifies, "union", "left union", "right union", "intersection" and then simply verify whether the expected regions are set). That would perhaps make it more intuitive. WDYT?

470

Have you tried to use a "let assemblyFormat =" description for this
(bit danting given all the possibilities, but I have to ask ;-)

479

Same feedback. The fullfills the need... should be rephrased into a more concrete description of what the op does.

542

This should be more descriptive that it only makes sense without the ops described above. Alternatively, we could simply use linalg.yield for this in the long run (although I am not sure if that would break other stuff, and it would be hard to combine the two for now).

561

no verifier? I would at least expect a test that this indeed appears inside another sparse set op?

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
343

such a helper should be static

386

this block of code can use some more comments

423

Here and below, end comments with a period

jim22k added inline comments.Mar 7 2022, 11:14 AM
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
404–411

I originally had an attribute {how="union"} and made each block optional. I was only thinking of "intersection" and "union", but including "left union" and "right union" fills out the set of possibilities. If we go with an attribute, I would eliminate the special identity token as it would be redundant. I don't know the best name for that attribute -- how works, but feels a little awkward.

Would you still want each region to be required? I would hope to make the left and right regions optional, as the attribute explains what the default behavior of those regions is. That should make the operation more concise.

Another question is whether the attribute would add to confusion when the left or right region is being overridden. For example, A-B. Do I write that as "union" and then override the "right" region? Or do I have to list that as "left union" in order to override the "right" region? In other words, does how only describe default behavior or restrict which regions can be defined?
As another example, if I fully define all 3 regions, what would I list for how?

These questions were why I eliminated an attribute and made each region required, along with an "identity" shortcut. It felt like a simpler approach, but I would be fine with either approach if you feel strongly about it.

470

Yes, I tried hard to use assemblyFormat=, but requiring the name of each region left={} and the special identity token pushed things over the edge to require custom handling.

479

I will let Aart comment on the best approach for linalg.index. I like the idea, as it would simplify sparse_tensor.unary. However, the sparse tensor dialect currently doesn't handle linalg.index, and linalg.index can't be embedded in my the unary blocks because linalg.index expects its parent operation to be linalg.generic.

I see three possible approaches:

  1. Make sparse tensor's lowering of linalg.generic handle linalg.index. Then refer to those SSAs inside the unary block.
  2. Allow linalg.index to live inside the unary block. Handle them during sparse tensor unary lowering.
  3. Leave it as written, still handled during the sparse tensor unary lowering.
487

Your understanding of "missing" is correct. It helps to look at the binary operation which has:

  • primary (present in both sparse tensors)
  • left (present in left sparse tensor, but not right)
  • right (present in right sparse tensor, but not left)

Technically, there is a fourth region:

  • missing (not present in either)

For binary, we usually ignore that "missing" region. For unary, however, the missing region is important.

561

Both the unary and binary verifiers check that each block terminates with sparse_tensor.yield. But I could add a verifier here to check that the parent op is one of the allowed operators.

aartbik added inline comments.Mar 7 2022, 3:39 PM
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
404–411

At first glance, I would say the verification makes sure that the following holds

how=union: primary, left, and right are all non-null
how =left-union:primary and left are not null, right is null
how=right-union: primary and right are not null, left is null
how=intersection: primary is not null, left and right are both null

Then, in all cases where left/right are not-null, you can use identity as shorthand.
Or we can simply say sparse_tensor.yield %arg0 for those cases

470

I was already afraid of that.

479

Of these choices, it seems that 1. is the long term most viable one, since it also adds the ability to use indices to other sparse code. Let me ponder a bit on this....

561

Yes please. Having unary and binary check for the presence of a yield is one side of the verification coin, but *not* having a dangling yield somewhere else is the other side of that verification coin.

aartbik added inline comments.Mar 7 2022, 3:44 PM
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
404–411

Also, let's pick something else for "how".
How about "kind"

CombiningKindAttr:$kind

we can pick "Combining" or "Iterating" or "Running", or something like that.

jim22k added inline comments.Mar 8 2022, 12:51 PM
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
404–411

"kind" will work. Let's go with that.

Then for verification rules, let's try these:

  • For all "kind"s, the primary must be defined, although it could be declared empty (i.e. { })
  • kind=intersection: left and right may not be defined (assumed empty)
  • kind=left-union: right may not be defined (assumed empty); left may be defined (if not defined, assumed to be identity)
  • kind=right-union: left may not be defined (assumed empty); right may be defined (if not defined, assumed to be identity)
  • kind=union: both or either left or right may be defined; if either are not defined, they are assumed to be identity

This will allow for compact code like:

%result = sparse_tensor.binary %x, %y {kind="left-union"} : f64 to f64 {
  ^bb0(%arg0: f64, %arg1: f64):
    sparse_tensor.yield %arg1: f64
}

There is no need to say left=identity because it is already implied by "kind". This will make writing the parser simpler.

Sounds good.

Note that in https://reviews.llvm.org/D121251 I added support for linalg.index in the sparse code generator.
You should be able to connect that easily later (even though you will need to inspect the contents of your opaque blocks for that).
Connecting with linalg.index has my preference too over any special treatment with parameters and such,
so let's drop that part from the new ops altogether!

jim22k updated this revision to Diff 414233.EditedMar 9 2022, 3:59 PM

Incorporate comments into design of binary and unary

  • Remove include_index (will use linalg.index instead in the future)
  • Remove custom parser and printing
  • Simplify assembly format (to allow for non-custom parsing)
  • Add OverlapKindAttr for binary kind attribute

I simply *love* what we are converging on!

On last nit, I feel the unary op should not follow the same approach with a "kind" for absolute symmetry?
WDYT?

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
403

within a linalg.generic operation (add "a")

405

in the sparse storage format (add "the")

417

is this past 80-cols?

450

this is neat!

we need to make sure we don't associate a lattice with just pre-blocks of course (since that would iterate ovre all indices)

;-)

472

Yeah \O/

So much better. Clean DSL based parsing/printing with mimimal logic in the verifier.
I like it!

488

I feel we should use the same approach for the unary case and let a "kind" attribute define the behavior, and verify that blocks are as expected

kindPresent : only primary
kindAbsent: only missing
kindBoth: both primary and missing

WDYT?

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
24

this should not be in the section define at L21.

Please add a new section

===----------------------------------------------------------------------===
TensorDialect Enum Methods.
===----------------------------------------------------------------------===//

aartbik retitled this revision from Introduce new binary and unary op for sparse_tensor to [mlir][sparse] Introduce new binary and unary op for sparse_tensor.Mar 9 2022, 4:47 PM

On last nit, I feel the unary op should not follow the same approach with a "kind" for absolute symmetry?
WDYT?

This should read: I feel the unary op should *now* follow ....
I make this mistaking typo quite often, and it confuses everyone a lot ;-)

jim22k marked 5 inline comments as done.Mar 10 2022, 7:36 AM

@aartbik Let me know your thoughts about my response to adding kind to unary. Once we have agreement, I will update the diff.

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
417

Yes, by quite a bit. I will shorten the lines. My editor has a gutter line at 120, so I often forget how long the lines are getting.

488

I'm not in favor of adding kind to unary. For binary, it serves to indicate which regions are "active" even if they are not defined. For example, a union with only the primary region defined also has the left and right regions set to identity implicitly. The kind also restricts which regions the user is allowed to override, but I consider that a less important feature than the implied default behavior of unspecified regions.

For unary, there is no concept of implicit behavior, so the kind would only serve the role of restricting which regions the user is allowed to override. And with only 2 allowable regions, adding the kind feels unnecessary. If the user defines a region, they clearly meant to override it.

And while I see your point of bringing a unified approach to these two operations, remember that I am planning to introduce reduce and mask in a future PR. Neither of those will have a kind.

aartbik added inline comments.Mar 10 2022, 11:31 AM
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
417

I believe you make a distinction between

(1) declared empty, as in `left={ }
(2) not defined, i.e. not present at all

but for intersection, you say both must be "empty", when I think you mean must not be defined at all, unless you assume that is the default value when not defined. We could even have the semantics that empty means "drop the value, no output" and not defined (when it should), means "identity output". That would mean that using a union, but with left/right present, but empty, would really be an intersection again.

Given all these choices, I think you need to refine this documentation to really zoom in on the intended semantics (i.e. what empty / not-defined really means). Maybe a table of possibilities for the three ops would be useful. Also, for theoretical completeness, we could even have a version that runs when *neither* are present, so perhaps we should even say something about that case.

417

Please be very precise on pass-through "id" behavior (ie. value appears as output) and "output is missing value", i.e. no output at all for sparsity purposes.

488

I still don't like that for the binary case we define-and-verify the meaning and now for the unary case we infer the meaning. Also, how do you specify running only something for missing (since "primary" is documented as mandatory). I would see the case of "inverting" a sparse tensor by only running for the zeros, and returning a nonzero for example (so a 40% sparse would become 60% sparse and vv). Perhaps we should call the branches "present" and "absent" just to allow that?

jim22k marked an inline comment as done.Mar 10 2022, 2:15 PM
jim22k added inline comments.
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
417

The distinction between "declared empty" and "not defined" is a little tricky. "Declared empty" (i.e. left={}) and "not defined" are identical to the standard parser assembly format.

(`left` `=` $leftRegion^)?`

If I write left={}, my verify method sees leftRegion as empty. If I don't declare "left" at all in my MLIR statement, the fact that $leftRegion is optional in the assembly format means that leftRegion is also empty in the verify method. I don't have a way to distinguish between those two cases without custom parsing.

This leads to the very unfortunate situation where an empty region means different things depending on the kind.
For example:

sparse_tensor.binary left_union %a, %b : f64 to f64 {
  ^bb0(...)
} left={
} right={
}

In this left_union case, the left region's emptiness will be handled as an identity function, while the right region's emptiness will be handled as not contributing to the output.
We do restrict the user from declaring right to be anything other than empty. But we have no way to restrict the user from defining left as empty because the parser sees it the same as if it weren't defined at all.

jim22k marked an inline comment as not done.Mar 10 2022, 2:38 PM
jim22k added inline comments.
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
417

Okay, I found a potential solution to the dilemma. I can use UnitAttr to have the standard parser differentiate between left=identity and left={}.

So here is my new proposal (which is actually my old proposal :D )

  1. Remove the kind on binary. Instead, make the primary, left, and right regions required to be defined, but rename primary to overlap
  2. Allow the special token "identity" for left and right regions in binary
  3. In the documentation, explain that an empty region {} denotes no output
  4. For unary, require the primary and missing regions to both be defined, but rename them as present and absent

Here is what a binary intersection looks like:

sparse_tensor.binary %a, %b : f64 to f64
  overlap={
    ^bb0(%x: f64, %y: f64):
      sparse_tensor.yield %x : f64
  }
  left={}
  right={}

Here is a binary right_union:

sparse_tensor.binary %a, %b : f64 to f64
  overlap={
    ^bb0(%x: f64, %y: f64):
      sparse_tensor.yield %x : f64
  }
  left={}
  right=identity

Here is a different binary right-union where the right region has custom code

sparse_tensor.binary %a, %b : f64 to f64
  overlap={
    ^bb0(%x: f64, %y: f64):
      sparse_tensor.yield %x : f64
  }
  left={}
  right={
    ^bb0(%y: f64):
      sparse_tensor.yield ...  // do something custom here
  }

And here is a unary:

sparse_tensor.unary %a : f64 to f64
  present={
    ^bb0(%x: f64):
      sparse_tensor.yield %x : f64
  }
  absent={}

What do you think?

aartbik added inline comments.Mar 10 2022, 2:57 PM
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
417

I love it! Very clean semantics now, and easy to describe.

This also enables us to do weird stuff, such as an unary that defines both branches ;-)
For binary, we just miss the theoretical "neither side" case, but that is okay.

Last concern, can you make this work with the DSL parser/printer?

jim22k marked 22 inline comments as done.Mar 10 2022, 3:06 PM
jim22k added inline comments.
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
417

Yes, I can make this work with the DSL parser/printer. I will get a new version out tonight with the changes.

aartbik added inline comments.Mar 10 2022, 4:54 PM
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
417

Very nice! Yeah, this will be very intuitive and concise. Thanks for working with us converging on this solution!

jim22k updated this revision to Diff 414558.EditedMar 10 2022, 6:59 PM
jim22k marked an inline comment as done.

Change binary and unary signature

  • Remove kind for binary
  • Make all regions required
  • Allow identity token for binary left and right regions
  • Rename regions (binary=overlap, left, right) (unary=present, absent)
jim22k added inline comments.Mar 14 2022, 3:49 PM
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
387

I just realized I don't want the SameTypeOperands trait here. It should be allowed to have any type as long as the rank and shape match. The output type of each region must match the output type, but we should allow any code within that region to manipulate the inputs without placing too many assumptions on those inputs.

jim22k updated this revision to Diff 415535.Mar 15 2022, 12:11 PM
Remove SameTypeOperands trait for binary

last few nits, but looks good!

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
403

More precise: Every non-empty block must end with a ....

419

applied *to* intersecting...

469

Returns a copy of ...

511

A non-empty block ....

576

Yields a ....

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
443–444

remove this, that is pretty standard MLIR stuff

jim22k updated this revision to Diff 415884.Mar 16 2022, 9:57 AM
Minor text updates

@aartbik Unless you find some more updates to the descriptions, I think this is ready. What is the next step? I don't think I have commit rights, so you will need to commit on my behalf.

aartbik accepted this revision.Mar 16 2022, 11:01 AM

Ship it! Let me know if you encounter any issues submitting the code.

This revision is now accepted and ready to land.Mar 16 2022, 11:01 AM
This revision was automatically updated to reflect the committed changes.