This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] Factoring out readAllElements

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



This change enables the SparseTensorCodegen pass to reuse the majority of the readCOO/readSparseTensor code without relying on the SparseTensorCOO/SparseTensorStorage classes. It does however introduce a dynamic function call in the for-loop, so we should perform benchmarking to make sure this refactoring doesn't slow readCOO down too much.

Depends On D138365

Diff Detail

Event Timeline

wrengr created this revision.Nov 19 2022, 4:16 PM
wrengr requested review of this revision.Nov 19 2022, 4:16 PM
wrengr updated this revision to Diff 476980.Nov 21 2022, 1:00 PM

Rebasing to incorporate the bugfix of Also doing a bit of redesign of readAllElements since that bugfix means we can perform the dim2lvl mapping in the callback rather than in readAllElements.

(As disscussed offline, the design of readAllElements may not actually work out for the codegen pass, since the type of the callback means we need to be able to form closures in order to close over the COO or other object being inserted into. Nevertheless, the general approach of D138365 can still be used by the codegen pass in order to do the LICM to avoid unnecessary conditionals in the for-loop.)

wrengr updated this revision to Diff 477366.Nov 22 2022, 6:36 PM


wrengr abandoned this revision.Dec 8 2022, 1:14 PM

Abandoning because this is only meaningfully usable if the caller can produce closures to pass to readAllElements (or if they use global variables or similar); hence it'll be difficult for the codegen pass to actually make use of this abstraction. So to solve the code-reuse problem, this should instead be redesigned in a more iterator-style approach.