[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

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.

[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

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.

[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

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]

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]

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]

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]

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.

[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

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]

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

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]

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]

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]

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]

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]

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]

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]

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

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]

Remove SameTypeOperands trait for binary

[aartbik]

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]

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/CMakeLists.txt

 add_mlir_dialect(SparseTensorOps sparse_tensor)
 add_mlir_doc(SparseTensorOps SparseTensorOps Dialects/ -gen-dialect-doc)
 
+set(LLVM_TARGET_DEFINITIONS SparseTensorOps.td)
+mlir_tablegen(SparseTensorOpsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(SparseTensorOpsEnums.cpp.inc -gen-enum-defs)
+
 set(LLVM_TARGET_DEFINITIONS SparseTensorAttrDefs.td)
 mlir_tablegen(SparseTensorAttrDefs.h.inc -gen-attrdef-decls)
 mlir_tablegen(SparseTensorAttrDefs.cpp.inc -gen-attrdef-defs)
 add_public_tablegen_target(MLIRSparseTensorAttrDefsIncGen)

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TensorEncoding.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsEnums.h.inc"
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.h.inc"
 
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsDialect.h.inc"
 

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

 //===- SparseTensorOps.td - Sparse tensor dialect ops ------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SPARSETENSOR_OPS
 #define SPARSETENSOR_OPS
 
 include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td"
 include "mlir/Dialect/SparseTensor/IR/SparseTensorBase.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/EnumAttr.td"
 
 //===----------------------------------------------------------------------===//
 // Base class.
 //===----------------------------------------------------------------------===//
 
 class SparseTensor_Op<string mnemonic, list<Trait> traits = []>
   : Op<SparseTensor_Dialect, mnemonic, traits>;
 
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 OverlapKindIntersection : I32EnumAttrCase<"Intersection", 0, "intersection">;
+def OverlapKindLeftUnion    : I32EnumAttrCase<"LeftUnion", 1, "left_union">;
+def OverlapKindRightUnion   : I32EnumAttrCase<"RightUnion", 2, "right_union">;
+def OverlapKindUnion        : I32EnumAttrCase<"Union", 3, "union">;
+
+/// Enum attribute of the different kinds of overlap for binary regions.
+def OverlapKindAttr : I32EnumAttr<"OverlapKind", "sparse_tensor binary overlap kind",
+  [OverlapKindIntersection, OverlapKindLeftUnion, OverlapKindRightUnion, OverlapKindUnion]> {
+  let cppNamespace = "::mlir::sparse_tensor";
+}
+
+def SparseTensor_BinaryOp : SparseTensor_Op<"binary", [NoSideEffect, SameTypeOperands]>,
+    Arguments<(ins AnyType:$x, AnyType:$y, OverlapKindAttr:$kind)>,
+    Results<(outs AnyType:$output)> {
+  let summary = "Binary set operation utilized within linalg.generic";
+  let description = [{
+      Defines a computation within `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 sparse storage format).
+
+      Three regions are defined for the operation and must appear in this order (if present):
+      - primary (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.
+      The block must end with 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={ }`, implying that the output is a missing value.
+
+      The `kind` attribute provides restrictions and default behaviors for the `left` and
+      `right` regions, which are both optional.
+      - intersection (left and right must be empty)
+      - left_union (left region returns identity if not defined; right must be empty)
+      - right_union (right region returns identity if not defined; left must be empty)
+      - union (left and right return identity if not defined)
+
+      Example of isEqual applied for 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 intersection %a, %b : f64 to i8 {
+              ^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
+            }
+          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 union %a, %b : f64 to f64 {
+              ^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
+            } 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>
+      ```
+  }];
+
+  let regions = (region AnyRegion:$primaryRegion, AnyRegion:$leftRegion, AnyRegion:$rightRegion);
+  let assemblyFormat = [{
+        $kind $x `,` $y `:` attr-dict type($x) `to` type($output) $primaryRegion (`left` `=` $leftRegion^)? (`right` `=` $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 (if present):
+      - primary (elements present in the sparse tensor)
+      - missing (elements not present in the sparse tensor)
+
+      Each region contains a single block describing the computation and result.
+      The block must end with 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 empty, implying that the output is a missing value.
+
+      The primary region is required.
+      The missing region is optional and is assumed to be empty if not defined.
+
+      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 {
+              ^bb0(%arg0: f64):
+                %cf1 = arith.constant 1.0 : f64
+                %ret = arith.addf %arg0, %cf1 : f64
+                sparse_tensor.yield %ret : f64
+            }
+          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 i64 {
+          ^bb0(%x: f64):
+            %ret = arith.constant 1 : i64
+            sparse_tensor.yield %ret : i64
+        } missing={
+          %ret = arith.constant -1 : i64
+          sparse_tensor.yield %ret : i64
+        }
+      ```
+
+      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 {
+        } missing={
+          %ret = arith.constant 1 : i64
+          sparse_tensor.yield %ret : i64
+        }
+      ```
+  }];
+
+  let regions = (region AnyRegion:$primaryRegion, AnyRegion:$missingRegion);
+  let assemblyFormat = [{
+        $x attr-dict `:` type($x) `to` type($output) $primaryRegion (`missing` `=` $missingRegion^)?
+  }];
+  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 = [{
+      Yield 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.
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsEnums.cpp.inc"
+
 #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:

 
 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,
+                   unsigned expectedNum, Type inputType, Type outputType) {
+  unsigned numArgs = region.getNumArguments();
+  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 != inputType)
+      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 inputType = x().getType();
+  Type outputType = output().getType();
+  OverlapKind kind = attrs.get("kind").cast<OverlapKindAttr>().getValue();
+  Region &primary = primaryRegion();
+  Region &left = leftRegion();
+  Region &right = rightRegion();
+
+  // Verify that expected empty region (based on kind) are actually empty
+  if (kind == OverlapKind::Intersection) {
+    if (!left.empty() || !right.empty())
+      return emitError("left and right region must be empty for intersection");
+  } else if (kind == OverlapKind::LeftUnion) {
+    if (!right.empty())
+      return emitError("right region must be empty for left_union");
+  } else if (kind == OverlapKind::RightUnion) {
+    if (!left.empty())
+      return emitError("left region must be empty for right_union");
+  }
+
+  // Check correct number of arguments and return type for each non-empty region
+  LogicalResult regionResult = success();
+  if (!primary.empty()) {
+    regionResult =
+        verifyNumBlockArgs(this, primary, "primary", 2, inputType, outputType);
+    if (failed(regionResult))
+      return regionResult;
+  }
+  if (!left.empty()) {
+    regionResult =
+        verifyNumBlockArgs(this, left, "left", 1, inputType, outputType);
+    if (failed(regionResult))
+      return regionResult;
+  }
+  if (!right.empty()) {
+    regionResult =
+        verifyNumBlockArgs(this, right, "right", 1, inputType, outputType);
+    if (failed(regionResult))
+      return regionResult;
+  }
+
+  return success();
+}
+
+LogicalResult UnaryOp::verify() {
+  Type inputType = x().getType();
+  Type outputType = output().getType();
+  LogicalResult regionResult = success();
+
+  // Check the number of block arguments and return type for all non-empty
+  // regions
+  Region &primary = primaryRegion();
+  if (!primary.empty()) {
+    regionResult =
+        verifyNumBlockArgs(this, primary, "primary", 1, inputType, outputType);
+    if (failed(regionResult))
+      return regionResult;
+  }
+  Region &missing = missingRegion();
+  if (!missing.empty()) {
+    regionResult =
+        verifyNumBlockArgs(this, missing, "missing", 0, inputType, 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();
+
+  // NOTE: Return type check is performed in each parent op's verify method
+
+  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_kind(%arg0: f64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{custom op 'sparse_tensor.binary' expected string or keyword containing one of the following enum values for attribute 'kind' [intersection, left_union, right_union, union]}}
+  %r = sparse_tensor.binary %arg0, %arg1 : f64 to f64 {
+      ^bb0(%x: f64):
+        sparse_tensor.yield %x : f64
+    }
+  return %r : f64
+}
+
+// -----
+
+func @invalid_binary_num_args_mismatch(%arg0: f64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{primary region must have exactly 2 arguments}}
+  %r = sparse_tensor.binary intersection %arg0, %arg1 : f64 to f64 {
+      ^bb0(%x: f64):
+        sparse_tensor.yield %x : f64
+    }
+  return %r : f64
+}
+
+// -----
+
+func @invalid_binary_argtype_mismatch(%arg0: f64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{primary region argument 2 type mismatch}}
+  %r = sparse_tensor.binary left_union %arg0, %arg1 : f64 to f64 {
+      ^bb0(%x: f64, %y: index):
+        sparse_tensor.yield %x : f64
+    }
+  return %r : f64
+}
+
+// -----
+
+func @invalid_binary_region_override(%arg0: f64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{right region must be empty for left_union}}
+  %r = sparse_tensor.binary left_union %arg0, %arg1 : f64 to f64 {
+    } left={
+    } right={
+      ^bb0(%y: f64):
+        sparse_tensor.yield %y : f64
+    }
+  return %r : f64
+}
+
+// -----
+
+func @invalid_binary_wrong_return_type(%arg0: f64, %arg1: f64) -> f64 {
+  // expected-error@+1 {{right region yield type mismatch}}
+  %0 = sparse_tensor.binary right_union %arg0, %arg1 : f64 to f64 {
+    } right={
+      ^bb0(%x: f64):
+        %1 = arith.constant 0.0 : f32
+        sparse_tensor.yield %1 : f32
+    }
+  return %0 : f64
+}
+
+// -----
+
+func @invalid_unary_argtype_mismatch(%arg0: f64) -> f64 {
+  // expected-error@+1 {{primary region argument 1 type mismatch}}
+  %r = sparse_tensor.unary %arg0 : f64 to f64 {
+      ^bb0(%x: index):
+        sparse_tensor.yield %x : index
+    }
+  return %r : f64
+}
+
+// -----
+
+func @invalid_unary_num_args_mismatch(%arg0: f64) -> f64 {
+  // expected-error@+1 {{missing region must have exactly 0 arguments}}
+  %r = sparse_tensor.unary %arg0 {include_index=true} : f64 to f64 {
+    } missing={
+      ^bb0(%x: f64):
+        sparse_tensor.yield %x : f64
+    }
+  return %r : f64
+}
+
+// -----
+
+func @invalid_unary_wrong_return_type(%arg0: f64) -> f64 {
+  // expected-error@+1 {{primary region yield type mismatch}}
+  %0 = sparse_tensor.unary %arg0 : f64 to f64 {
+      ^bb0(%x: f64):
+        %1 = arith.constant 0.0 : f32
+        sparse_tensor.yield %1 : f32
+    }
+  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:.*]]: f64) -> f64 {
+//       CHECK:   %[[C1:.*]] = sparse_tensor.binary right_union %[[A]], %[[B]] : f64 to f64 {
+//       CHECK:     ^bb0(%[[A1:.*]]: f64, %[[B1:.*]]: f64):
+//       CHECK:       sparse_tensor.yield %[[A1]] : f64
+//       CHECK:   } right = {
+//       CHECK:     ^bb0(%[[A2:.*]]: f64):
+//       CHECK:       sparse_tensor.yield %[[A2]] : f64
+//       CHECK:   }
+//       CHECK:   return %[[C1]] : f64
+//       CHECK: }
+func @sparse_binary(%arg0: f64, %arg1: f64) -> f64 {
+  %r = sparse_tensor.binary right_union %arg0, %arg1 : f64 to f64 {
+      ^bb0(%x: f64, %y: f64):
+        sparse_tensor.yield %x : f64
+    } right={
+      ^bb0(%y: f64):
+        sparse_tensor.yield %y : 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:     ^bb0(%[[A1:.*]]: f64):
+//       CHECK:       sparse_tensor.yield %[[A1]] : f64
+//       CHECK:   } missing = {
+//       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 {
+      ^bb0(%x: f64):
+        sparse_tensor.yield %x : f64
+    } missing={
+      ^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:     ^bb0(%[[A1:.*]]: f64):
+//       CHECK:       %[[R:.*]] = arith.fptosi %[[A1]] : f64 to i64
+//       CHECK:       sparse_tensor.yield %[[R]] : i64
+//       CHECK:   }
+//       CHECK:   return %[[C1]] : i64
+//       CHECK: }
+func @sparse_unary(%arg0: f64) -> i64 {
+  %r = sparse_tensor.unary %arg0 : f64 to i64 {
+      ^bb0(%x: f64):
+        %ret = arith.fptosi %x : f64 to i64
+        sparse_tensor.yield %ret : i64
+    }
+  return %r : i64
+}