[mlir][sparse] Add lowering for unary and binary ops
authorJim Kitchen <jim22k@gmail.com>
Tue, 3 May 2022 20:50:26 +0000 (15:50 -0500)
committerJim Kitchen <jim22k@gmail.com>
Tue, 3 May 2022 20:50:26 +0000 (15:50 -0500)
commit2c33266084602a046736bfa0caf441de313a6682
tree47abe64b50909e62696b0ad17f7ab18f41fb83ac
parentc4546091ed29763df3e4649327679fec5dcb5ff6
[mlir][sparse] Add lowering for unary and binary ops

Adding lowering for Unary and Binary required several changes due to
their unique nature of containing custom code for different "regions"
of the sparse structure being operated on. Along with a Kind, a pointer
to the Operation is passed along to be merged once the lattice
structure is figured out.

The original operation is maintained, as it is required for subsequent
lattice decisions. However, sparse_tensor.binary has some branches
are considered as fully handled and therefore are marked with as
kBinaryBranch to distinguish them.

A unique aspect of the custom code is that sometimes the desired result
is no result at all -- i.e. a user wants overlapping sparse entries to
become empty in the output. The solution to this is to return an
uninitialized Value(), which is checked and handled elsewhere in the
code and results in nothing being written to the output tensor for that
case.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D123057
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir [new file with mode: 0644]
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir [new file with mode: 0644]
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir