This is an archive of the discontinued LLVM Phabricator instance.

Support sparse tensor output.
ClosedPublic

Authored by bixia on Dec 10 2021, 3:26 PM.

Details

Summary

Add convertFromMLIRSparseTensor to the supporting C shared library to convert
SparseTensorStorage to COO-flavor format.

Add Python routine sparse_tensor_to_coo_tensor to convert sparse tensor storage
pointer to numpy values for COO-flavor format tensor.

Add a Python test for sparse tensor output.

Diff Detail

Event Timeline

bixia created this revision.Dec 10 2021, 3:26 PM
bixia requested review of this revision.Dec 10 2021, 3:26 PM
bixia updated this revision to Diff 393611.Dec 10 2021, 3:38 PM

Fix build.

bixia updated this revision to Diff 393618.Dec 10 2021, 3:53 PM

Exclude the tools directory from tests.

bixia updated this revision to Diff 393619.Dec 10 2021, 4:00 PM

Misc fixes.

aartbik added inline comments.Dec 10 2021, 4:03 PM
mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
1075

Make it a little more clear that these are the *output* parameters (so we expect pointers that will be populated by this method)

1088

also add TODO on generalizing type, just like L1051

1096

does allocation leak?

why not just use a std::vector<uint64_t> for this

std::vector<uint64_t> perm(rank);
std::iota(perm.begin(), perm.end(), 0);

mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py
23

Make this a TODO since we will have to work on allowing this

117

actually, in Python land we don't need to CHECK output, we can actually just verify that the values are as expected, see the other py test in this dir

mlir/test/Integration/Dialect/SparseTensor/python/tools/np_to_sparse_tensor.py
29

is there a way to do this only once?

aartbik added inline comments.Dec 10 2021, 4:13 PM
mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
1113

how about using a "base" for this to avoid the * as was done above
(not trusting compiler optimization ;-)

bixia updated this revision to Diff 393631.Dec 10 2021, 4:28 PM
bixia marked 4 inline comments as done.

Address review comment.

bixia added inline comments.Dec 10 2021, 4:29 PM
mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
1096

Good catch, add delete.
The function toCOO only accept c style array.

bixia updated this revision to Diff 393635.Dec 10 2021, 4:35 PM

Fixed test.

aartbik added inline comments.Dec 10 2021, 5:57 PM
mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
1096

Yes, but you can pass perm.data() for that. I have a slight preference for that since that is the idiom already used at line 1057 and further. This way we have a single way of doing the permutation setup.

1113

did you see this one?

bixia updated this revision to Diff 393676.Dec 11 2021, 7:18 AM
bixia marked 4 inline comments as done.

Replace the use of an allocated array with std::vector.

bixia added inline comments.Dec 11 2021, 7:29 AM
mlir/test/Integration/Dialect/SparseTensor/python/tools/np_to_sparse_tensor.py
29

We probably need to make c_lib a global, and then user set it to None to trigger this initialization.
Shall we do this?

bixia updated this revision to Diff 393690.Dec 11 2021, 12:47 PM

Put the code to look up the c shared library in a function and decorate the function with functools.lru_cache.

bixia added inline comments.Dec 11 2021, 1:21 PM
mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
1113

Use "base".

mlir/test/Integration/Dialect/SparseTensor/python/tools/np_to_sparse_tensor.py
29

Style guide recommend to avoid global variable. I outlined the look up of c_lib to a routine and added decorator lru_cache.

bixia updated this revision to Diff 393697.Dec 11 2021, 2:22 PM

Fixed python function docstrings.

bixia updated this revision to Diff 393699.Dec 11 2021, 2:32 PM

Fixed docstrings.

bixia updated this revision to Diff 393701.Dec 11 2021, 2:57 PM

Fixed a typo.

bixia marked an inline comment as done.Dec 13 2021, 10:04 AM
aartbik added inline comments.Dec 13 2021, 11:24 AM
mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py
116

We will have to check the return value, i.e. something like

if np.allclose(....):

  pass
else:
  quit(f'FAILURE')
bixia updated this revision to Diff 393980.Dec 13 2021, 11:44 AM

Use the result of np.allclose.

bixia marked an inline comment as done.Dec 13 2021, 11:45 AM
aartbik accepted this revision.Dec 13 2021, 12:00 PM

One last suggestion on testing the other output variable as well. but other than that, solid work. Ship it!

mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py
117

perhaps also test the other output parameters, either with a similar test or by printing and CHECKing the value. E.g. the rank is of course also reflected in the indices, but still nice to add an explicit check in this early phase of developing the tests

This revision is now accepted and ready to land.Dec 13 2021, 12:00 PM
This revision was automatically updated to reflect the committed changes.