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

[jim22k]

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.

[jim22k]

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]

Fix Formatting

[aartbik]

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

[jim22k]

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]

In D121018#3361164, @jim22k wrote:

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]

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]

Try again with all the changes this time

[aartbik]

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]

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

473

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

520

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

[nicolasvasilache]

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

478

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 ?

486

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?

493

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

[aartbik]

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

391

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).

403–410

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?

469

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

478

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

541

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).

560

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

341

such a helper should be static

384

this block of code can use some more comments

421

Here and below, end comments with a period

[jim22k]

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

403–410

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.

469

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.

478

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.

486

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.

560

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]

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

403–410

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

469

I was already afraid of that.

478

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....

560

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]

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

403–410

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]

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

403–410

"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.

[aartbik]

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]

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

[aartbik]

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

402

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

404

in the sparse storage format (add "the")

416

is this past 80-cols?

449

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)

;-)

471

Yeah \O/

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

487

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]

In D121018#3371487, @aartbik wrote:

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]

@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

416

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.

487

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]

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

416

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.

416

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.

487

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]

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

416

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]

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

416

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]

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

416

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]

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

416

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

[aartbik]

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

416

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

[jim22k]

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]

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

386

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]

Remove SameTypeOperands trait for binary

[aartbik]

last few nits, but looks good!

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

402

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

418

applied *to* intersecting...

468

Returns a copy of ...

510

A non-empty block ....

575

Yields a ....

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

441–442

remove this, that is pretty standard MLIR stuff

[jim22k]

Minor text updates

[jim22k]

@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]

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

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

string description = [{
     ```mlir
     sparse_tensor.out %t, %dest : tensor<1024x1024xf64, #CSR>, !Dest
     ```
   }];
   let assemblyFormat = "$tensor `,` $dest attr-dict `:` type($tensor) `,` type($dest)";
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Sparse Tensor Custom Linalg.Generic Operations.
+//===----------------------------------------------------------------------===//
+
+def SparseTensor_BinaryOp : SparseTensor_Op<"binary", [NoSideEffect]>,
+    Arguments<(ins AnyType:$x, AnyType:$y, UnitAttr:$left_identity, UnitAttr:$right_identity)>,
+    Results<(outs AnyType:$output)> {
+  let summary = "Binary set operation utilized within linalg.generic";
+  let description = [{
+      Defines a computation within a `linalg.generic` operation that takes two
+      operands and executes one of the regions depending on whether both operands
+      or either operand is nonzero (i.e. stored explicitly in the sparse storage
+      format).
+
+      Three regions are defined for the operation and must appear in this order:
+      - overlap (elements present in both sparse tensors)
+      - left (elements only present in the left sparse tensor)
+      - right (element only present in the right sparse tensor)
+
+      Each region contains a single block describing the computation and result.
+      Every non-empty block must end with a sparse_tensor.yield and the return
+      type must match the type of `output`. The primary region's block has two
+      arguments, while the left and right region's block has only one argument.
+
+      A region may also be declared empty (i.e. `left={}`), indicating that the
+      region does not contribute to the output. For example, setting both
+      `left={}` and `right={}` is equivalent to the intersection of the two
+      inputs as only the overlap region will contribute values to the output.
+
+      As a convenience, there is also a special token `identity` which can be
+      used in place of the left or right region. This token indicates that
+      the return value is the input value (i.e. func(%x) => return %x).
+      As a practical example, setting `left=identity` and `right=identity`
+      would be equivalent to a union operation where non-overlapping values
+      in the inputs are copied to the output unchanged.
+
+      Example of isEqual applied to intersecting elements only:
+      ```mlir
+      %C = sparse_tensor.init...
+      %0 = linalg.generic #trait
+        ins(%A: tensor<?xf64, #SparseVec>, %B: tensor<?xf64, #SparseVec>)
+        outs(%C: tensor<?xi8, #SparseVec>) {
+        ^bb0(%a: f64, %b: f64, %c: i8) :
+          %result = sparse_tensor.binary %a, %b : f64, f64 to i8
+            overlap={
+              ^bb0(%arg0: f64, %arg1: f64):
+                %cmp = arith.cmpf "oeq", %arg0, %arg1 : f64
+                %ret_i8 = arith.extui %cmp : i1 to i8
+                sparse_tensor.yield %ret_i8 : i8
+            }
+            left={}
+            right={}
+          linalg.yield %result : i8
+      } -> tensor<?xi8, #SparseVec>
+      ```
+
+      Example of A+B in upper triangle, A-B in lower triangle:
+      ```mlir
+      %C = sparse_tensor.init...
+      %1 = linalg.generic #trait
+        ins(%A: tensor<?x?xf64, #CSR>, %B: tensor<?x?xf64, #CSR>
+        outs(%C: tensor<?x?xf64, #CSR> {
+        ^bb0(%a: f64, %b: f64, %c: f64) :
+          %row = linalg.index 0 : index
+          %col = linalg.index 1 : index
+          %result = sparse_tensor.binary %a, %b : f64, f64 to f64
+            overlap={
+              ^bb0(%x: f64, %y: f64):
+                %cmp = arith.cmpi "uge", %column, %row : index
+                %upperTriangleResult = arith.addf %x, %y : f64
+                %lowerTriangleResult = arith.subf %x, %y : f64
+                %ret = arith.select %cmp, %upperTriangleResult, %lowerTriangleResult : f64
+                sparse_tensor.yield %ret : f64
+            }
+            left=identity
+            right={
+              ^bb0(%y: f64):
+                %cmp = arith.cmpi "uge", %column, %row : index
+                %lowerTriangleResult = arith.negf %y : f64
+                %ret = arith.select %cmp, %y, %lowerTriangleResult
+                sparse_tensor.yield %ret : f64
+            }
+          linalg.yield %result : f64
+      } -> tensor<?x?xf64, #CSR>
+      ```
+
+      Example of set difference. Returns a copy of A where its sparse structure
+      is *not* overlapped by B. The element type of B can be different than A
+      because we never use its values, only its sparse structure.
+      ```mlir
+      %C = sparse_tensor.init...
+      %2 = linalg.generic #trait
+        ins(%A: tensor<?x?xf64, #CSR>, %B: tensor<?x?xi32, #CSR>
+        outs(%C: tensor<?x?xf64, #CSR> {
+        ^bb0(%a: f64, %b: i32, %c: f64) :
+          %result = sparse_tensor.binary %a, %b : f64, i32 to f64
+            overlap={}
+            left=identity
+            right={}
+          linalg.yield %result : f64
+      } -> tensor<?x?xf64, #CSR>
+      ```
+  }];
+
+  let regions = (region AnyRegion:$overlapRegion, AnyRegion:$leftRegion, AnyRegion:$rightRegion);
+  let assemblyFormat = [{
+        $x `,` $y `:` attr-dict type($x) `,` type($y) `to` type($output) `\n`
+        `overlap` `=` $overlapRegion `\n`
+        `left` `=` (`identity` $left_identity^):($leftRegion)? `\n`
+        `right` `=` (`identity` $right_identity^):($rightRegion)?
+  }];
+  let hasVerifier = 1;
+}
+
+def SparseTensor_UnaryOp : SparseTensor_Op<"unary", [NoSideEffect]>,
+    Arguments<(ins AnyType:$x)>,
+    Results<(outs AnyType:$output)> {
+  let summary = "Unary set operation utilized within linalg.generic";
+  let description = [{
+      Defines a computation with a `linalg.generic` operation that takes a single
+      operand and executes one of two regions depending on whether the operand is
+      nonzero (i.e. stored explicitly in the sparse storage format).
+
+      Two regions are defined for the operation must appear in this order:
+      - present (elements present in the sparse tensor)
+      - absent (elements not present in the sparse tensor)
+
+      Each region contains a single block describing the computation and result.
+      A non-empty block must end with a sparse_tensor.yield and the return type
+      must match the type of `output`. The primary region's block has one
+      argument, while the missing region's block has zero arguments.
+
+      A region may also be declared empty (i.e. `absent={}`), indicating that the
+      region does not contribute to the output.
+
+      Example of A+1, restricted to existing elements:
+      ```mlir
+      %C = sparse_tensor.init...
+      %0 = linalg.generic #trait
+        ins(%A: tensor<?xf64, #SparseVec>)
+        outs(%C: tensor<?xf64, #SparseVec>) {
+        ^bb0(%a: f64, %c: f64) :
+          %result = sparse_tensor.unary %a : f64 to f64
+            present={
+              ^bb0(%arg0: f64):
+                %cf1 = arith.constant 1.0 : f64
+                %ret = arith.addf %arg0, %cf1 : f64
+                sparse_tensor.yield %ret : f64
+            }
+            absent={}
+          linalg.yield %result : f64
+      } -> tensor<?xf64, #SparseVec>
+      ```
+
+      Example returning +1 for existing values and -1 for missing values:
+      ```mlir
+      %result = sparse_tensor.unary %a : f64 to i32
+        present={
+          ^bb0(%x: f64):
+            %ret = arith.constant 1 : i32
+            sparse_tensor.yield %ret : i32
+        }
+        absent={
+          %ret = arith.constant -1 : i32
+          sparse_tensor.yield %ret : i32
+        }
+      ```
+
+      Example showing a structural inversion (existing values become missing in
+      the output, while missing values are filled with 1):
+      ```mlir
+      %result = sparse_tensor.unary %a : f64 to i64
+        present={}
+        absent={
+          %ret = arith.constant 1 : i64
+          sparse_tensor.yield %ret : i64
+        }
+      ```
+  }];
+
+  let regions = (region AnyRegion:$presentRegion, AnyRegion:$absentRegion);
+  let assemblyFormat = [{
+        $x attr-dict `:` type($x) `to` type($output) `\n`
+        `present` `=` $presentRegion `\n`
+        `absent` `=` $absentRegion
+  }];
+  let hasVerifier = 1;
+}
+
+def SparseTensor_YieldOp : SparseTensor_Op<"yield", [NoSideEffect, Terminator]>,
+    Arguments<(ins AnyType:$result)> {
+  let summary = "Yield from sparse_tensor set-like operations";
+  let description = [{
+      Yields a value from within a `binary` or `unary` block.
+
+      Example:
+      ```
+      %0 = sparse_tensor.unary %a : i64 to i64 {
+        ^bb0(%arg0: i64):
+          %cst = arith.constant 1 : i64
+          %ret = arith.addi %arg0, %cst : i64
+          sparse_tensor.yield %ret : i64
+      }
+      ```
+  }];
+
+  let assemblyFormat = [{
+        $result attr-dict `:` type($result)
+  }];
+  let hasVerifier = 1;
+}
+
 #endif // SPARSETENSOR_OPS

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

 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
 // TensorDialect Attribute Methods.
 //===----------------------------------------------------------------------===//
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc"
 
 static bool acceptBitWidth(unsigned bitWidth) {
   switch (bitWidth) {
   case 0:
   case 8:
   case 16:
   case 32:

 
 LogicalResult OutOp::verify() {
   if (!getSparseTensorEncoding(tensor().getType()))
     return emitError("expected a sparse tensor for output");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// TensorDialect Linalg.Generic Operations.
+//===----------------------------------------------------------------------===//
+
+template <class T>
+static LogicalResult verifyNumBlockArgs(T *op, Region &region,
+                                        const char *regionName,
+                                        TypeRange inputTypes, Type outputType) {
+  unsigned numArgs = region.getNumArguments();
+  unsigned expectedNum = inputTypes.size();
+  if (numArgs != expectedNum)
+    return op->emitError() << regionName << " region must have exactly "
+                           << expectedNum << " arguments";
+
+  for (unsigned i = 0; i < numArgs; i++) {
+    Type typ = region.getArgument(i).getType();
+    if (typ != inputTypes[i])
+      return op->emitError() << regionName << " region argument " << (i + 1)
+                             << " type mismatch";
+  }
+  Operation *term = region.front().getTerminator();
+  YieldOp yield = dyn_cast<YieldOp>(term);
+  if (!yield)
+    return op->emitError() << regionName
+                           << " region must end with sparse_tensor.yield";
+  if (yield.getOperand().getType() != outputType)
+    return op->emitError() << regionName << " region yield type mismatch";
+
+  return success();
+}
+
+LogicalResult BinaryOp::verify() {
+  NamedAttrList attrs = (*this)->getAttrs();
+  Type leftType = x().getType();
+  Type rightType = y().getType();
+  Type outputType = output().getType();
+  Region &overlap = overlapRegion();
+  Region &left = leftRegion();
+  Region &right = rightRegion();
+
+  // Check correct number of block arguments and return type for each
+  // non-empty region.
+  LogicalResult regionResult = success();
+  if (!overlap.empty()) {
+    regionResult = verifyNumBlockArgs(
+        this, overlap, "overlap", TypeRange{leftType, rightType}, outputType);
+    if (failed(regionResult))
+      return regionResult;
+  }
+  if (!left.empty()) {
+    regionResult =
+        verifyNumBlockArgs(this, left, "left", TypeRange{leftType}, outputType);
+    if (failed(regionResult))
+      return regionResult;
+  } else if (left_identity()) {
+    if (leftType != outputType)
+      return emitError("left=identity requires first argument to have the same "
+                       "type as the output");
+  }
+  if (!right.empty()) {
+    regionResult = verifyNumBlockArgs(this, right, "right",
+                                      TypeRange{rightType}, outputType);
+    if (failed(regionResult))
+      return regionResult;
+  } else if (right_identity()) {
+    if (rightType != outputType)
+      return emitError("right=identity requires second argument to have the "
+                       "same type as the output");
+  }
+
+  return success();
+}
+
+LogicalResult UnaryOp::verify() {
+  Type inputType = x().getType();
+  Type outputType = output().getType();
+  LogicalResult regionResult = success();
+
+  // Check correct number of block arguments and return type for each
+  // non-empty region.
+  Region &present = presentRegion();
+  if (!present.empty()) {
+    regionResult = verifyNumBlockArgs(this, present, "present",
+                                      TypeRange{inputType}, outputType);
+    if (failed(regionResult))
+      return regionResult;
+  }
+  Region &absent = absentRegion();
+  if (!absent.empty()) {
+    regionResult =
+        verifyNumBlockArgs(this, absent, "absent", TypeRange{}, outputType);
+    if (failed(regionResult))
+      return regionResult;
+  }
+
+  return success();
+}
+
+LogicalResult YieldOp::verify() {
+  // Check for compatible parent.
+  auto *parentOp = (*this)->getParentOp();
+  if (auto binaryOp = dyn_cast<BinaryOp>(parentOp))
+    return success();
+  if (auto unaryOp = dyn_cast<UnaryOp>(parentOp))
+    return success();
+
+  return emitOpError("expected parent op to be sparse_tensor binary or unary");
+}
+
+//===----------------------------------------------------------------------===//
 // TensorDialect Methods.
 //===----------------------------------------------------------------------===//
 
 void SparseTensorDialect::initialize() {
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc"
       >();

mlir/test/Dialect/SparseTensor/invalid.mlir

 
 // -----
 
 func @invalid_out_dense(%arg0: tensor<10xf64>, %arg1: !llvm.ptr<i8>) {
   // expected-error@+1 {{expected a sparse tensor for output}}
   sparse_tensor.out %arg0, %arg1 : tensor<10xf64>, !llvm.ptr<i8>
   return
 }
+
+// -----
+
+func @invalid_binary_num_args_mismatch_overlap(%arg0: f64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{overlap region must have exactly 2 arguments}}
+  %r = sparse_tensor.binary %arg0, %arg1 : f64, f64 to f64
+    overlap={
+      ^bb0(%x: f64):
+        sparse_tensor.yield %x : f64
+    }
+    left={}
+    right={}
+  return %r : f64
+}
+
+// -----
+
+func @invalid_binary_num_args_mismatch_right(%arg0: f64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{right region must have exactly 1 arguments}}
+  %r = sparse_tensor.binary %arg0, %arg1 : f64, f64 to f64
+    overlap={}
+    left={}
+    right={
+      ^bb0(%x: f64, %y: f64):
+        sparse_tensor.yield %y : f64
+    }
+  return %r : f64
+}
+
+// -----
+
+func @invalid_binary_argtype_mismatch(%arg0: f64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{overlap region argument 2 type mismatch}}
+  %r = sparse_tensor.binary %arg0, %arg1 : f64, f64 to f64
+    overlap={
+      ^bb0(%x: f64, %y: f32):
+        sparse_tensor.yield %x : f64
+    }
+    left=identity
+    right=identity
+  return %r : f64
+}
+
+// -----
+
+func @invalid_binary_wrong_return_type(%arg0: f64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{left region yield type mismatch}}
+  %0 = sparse_tensor.binary %arg0, %arg1 : f64, f64 to f64
+    overlap={}
+    left={
+      ^bb0(%x: f64):
+        %1 = arith.constant 0.0 : f32
+        sparse_tensor.yield %1 : f32
+    }
+    right=identity
+  return %0 : f64
+}
+
+// -----
+
+func @invalid_binary_wrong_identity_type(%arg0: i64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{left=identity requires first argument to have the same type as the output}}
+  %0 = sparse_tensor.binary %arg0, %arg1 : i64, f64 to f64
+    overlap={}
+    left=identity
+    right=identity
+  return %0 : f64
+}
+
+// -----
+
+func @invalid_unary_argtype_mismatch(%arg0: f64) -> f64 {
+  // expected-error@+1 {{present region argument 1 type mismatch}}
+  %r = sparse_tensor.unary %arg0 : f64 to f64
+    present={
+      ^bb0(%x: index):
+        sparse_tensor.yield %x : index
+    }
+    absent={}
+  return %r : f64
+}
+
+// -----
+
+func @invalid_unary_num_args_mismatch(%arg0: f64) -> f64 {
+  // expected-error@+1 {{absent region must have exactly 0 arguments}}
+  %r = sparse_tensor.unary %arg0 : f64 to f64
+    present={}
+    absent={
+      ^bb0(%x: f64):
+        sparse_tensor.yield %x : f64
+    }
+  return %r : f64
+}
+
+// -----
+
+func @invalid_unary_wrong_return_type(%arg0: f64) -> f64 {
+  // expected-error@+1 {{present region yield type mismatch}}
+  %0 = sparse_tensor.unary %arg0 : f64 to f64
+    present={
+      ^bb0(%x: f64):
+        %1 = arith.constant 0.0 : f32
+        sparse_tensor.yield %1 : f32
+    }
+    absent={}
+  return %0 : f64
+}

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //  CHECK-SAME: %[[A:.*]]: tensor<?x?xf64, #sparse_tensor.encoding<{{.*}}>>,
 //  CHECK-SAME: %[[B:.*]]: !llvm.ptr<i8>)
 //       CHECK: sparse_tensor.out %[[A]], %[[B]] : tensor<?x?xf64, #sparse_tensor.encoding<{{.*}}>>, !llvm.ptr<i8>
 //       CHECK: return
 func @sparse_out(%arg0: tensor<?x?xf64, #SparseMatrix>, %arg1: !llvm.ptr<i8>) {
   sparse_tensor.out %arg0, %arg1 : tensor<?x?xf64, #SparseMatrix>, !llvm.ptr<i8>
   return
 }
+
+// -----
+
+#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @sparse_binary(
+//  CHECK-SAME:   %[[A:.*]]: f64, %[[B:.*]]: i64) -> f64 {
+//       CHECK:   %[[Z:.*]] = arith.constant 0.000000e+00 : f64
+//       CHECK:   %[[C1:.*]] = sparse_tensor.binary %[[A]], %[[B]] : f64, i64 to f64
+//       CHECK:     overlap = {
+//       CHECK:       ^bb0(%[[A1:.*]]: f64, %[[B1:.*]]: i64):
+//       CHECK:         sparse_tensor.yield %[[A1]] : f64
+//       CHECK:     }
+//       CHECK:     left = identity
+//       CHECK:     right = {
+//       CHECK:       ^bb0(%[[A2:.*]]: i64):
+//       CHECK:         sparse_tensor.yield %[[Z]] : f64
+//       CHECK:     }
+//       CHECK:   return %[[C1]] : f64
+//       CHECK: }
+func @sparse_binary(%arg0: f64, %arg1: i64) -> f64 {
+  %cf0 = arith.constant 0.0 : f64
+  %r = sparse_tensor.binary %arg0, %arg1 : f64, i64 to f64
+    overlap={
+      ^bb0(%x: f64, %y: i64):
+        sparse_tensor.yield %x : f64
+    }
+    left=identity
+    right={
+      ^bb0(%y: i64):
+        sparse_tensor.yield %cf0 : f64
+    }
+  return %r : f64
+}
+
+// -----
+
+#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @sparse_unary(
+//  CHECK-SAME:   %[[A:.*]]: f64) -> f64 {
+//       CHECK:   %[[C1:.*]] = sparse_tensor.unary %[[A]] : f64 to f64
+//       CHECK:     present = {
+//       CHECK:       ^bb0(%[[A1:.*]]: f64):
+//       CHECK:         sparse_tensor.yield %[[A1]] : f64
+//       CHECK:     }
+//       CHECK:     absent = {
+//       CHECK:       %[[R:.*]] = arith.constant -1.000000e+00 : f64
+//       CHECK:       sparse_tensor.yield %[[R]] : f64
+//       CHECK:     }
+//       CHECK:   return %[[C1]] : f64
+//       CHECK: }
+func @sparse_unary(%arg0: f64) -> f64 {
+  %r = sparse_tensor.unary %arg0 : f64 to f64
+    present={
+      ^bb0(%x: f64):
+        sparse_tensor.yield %x : f64
+    } absent={
+      ^bb0:
+        %cf1 = arith.constant -1.0 : f64
+        sparse_tensor.yield %cf1 : f64
+    }
+  return %r : f64
+}
+
+// -----
+
+#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @sparse_unary(
+//  CHECK-SAME:   %[[A:.*]]: f64) -> i64 {
+//       CHECK:   %[[C1:.*]] = sparse_tensor.unary %[[A]] : f64 to i64
+//       CHECK:     present = {
+//       CHECK:       ^bb0(%[[A1:.*]]: f64):
+//       CHECK:         %[[R:.*]] = arith.fptosi %[[A1]] : f64 to i64
+//       CHECK:         sparse_tensor.yield %[[R]] : i64
+//       CHECK:     }
+//       CHECK:     absent = {
+//       CHECK:     }
+//       CHECK:   return %[[C1]] : i64
+//       CHECK: }
+func @sparse_unary(%arg0: f64) -> i64 {
+  %r = sparse_tensor.unary %arg0 : f64 to i64
+    present={
+      ^bb0(%x: f64):
+        %ret = arith.fptosi %x : f64 to i64
+        sparse_tensor.yield %ret : i64
+    }
+    absent={}
+  return %r : i64
+}