[mlir][sparse] Implement sparse_tensor.select
authorJim Kitchen <jim22k@gmail.com>
Mon, 3 Oct 2022 19:34:53 +0000 (14:34 -0500)
committerJim Kitchen <jim22k@gmail.com>
Mon, 3 Oct 2022 19:39:26 +0000 (14:39 -0500)
commit791935037b0b3b211bee54fae694aeb5b7b75125
tree9b4f244100afd9c1ef91a232a4605013468b11f3
parent200889fbd91b661a8cf4e5a914d05010300ffee1
[mlir][sparse] Implement sparse_tensor.select

The region within sparse_tensor.select is used as the runtime criteria
for whether to keep the existing value in the sparse tensor.

While the sparse element is provided to the comparison, indices may also
be used to decide on whether to keep the original value. This allows, for
example, to only keep the upper triangle of a matrix.

Reviewed by: aartbik

Differential Revision: https://reviews.llvm.org/D134761
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_select.mlir [new file with mode: 0644]
mlir/unittests/Dialect/SparseTensor/MergerTest.cpp