Page MenuHomePhabricator

[mlir][sparse] Breaking up openSparseTensor to better support non-permutations
ClosedPublic

Authored by wrengr on Nov 19 2022, 4:10 PM.

Details

Summary

This commit updates how the SparseTensorConversion pass handles NewOp. It breaks up the underlying openSparseTensor function into two parts (SparseTensorReader::create and SparseTensorReader::readSparseTensor) so that the pass can inject code for constructing lvlSizes between those two parts. Migrating the construction of lvlSizes out of the runtime and into the pass is a necessary first step toward fully supporting non-permutations. (The alternative would be for the pass to generate a FuncOp for performing the construction and then passing that to the runtime; which doesn't seem to have any benefits over the design of this commit.) And since the pass now generates the code to call these two functions, this change also removes the Action::kFromFile value from the enum used by _mlir_ciface_newSparseTensor.

Diff Detail

Event Timeline

wrengr created this revision.Nov 19 2022, 4:10 PM
wrengr requested review of this revision.Nov 19 2022, 4:10 PM
wrengr updated this revision to Diff 476970.Nov 21 2022, 12:13 PM

Fixing a logic bug in SparseTensorReader::readCOO. This bug does not affect permutations (hence why not discovering it earlier), but would affect non-permutations.

aartbik accepted this revision.Nov 30 2022, 7:55 PM
aartbik added inline comments.
mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
656

agreed; being on the boundary of MLIR calling convention and nice C++ templating is really annoying sometimes ;-) it would be even nicer to be able to generate the lib "on demand" so that we only have the cases needed by the compiler (and also generate cases not covered yet)

This revision is now accepted and ready to land.Nov 30 2022, 7:55 PM
wrengr updated this revision to Diff 479495.Dec 1 2022, 6:19 PM

correcting rewriting_for_codegen.mlir post-rebase