[mlir][sparse] Updating the `Merger::{exp,lat,set}` methods to return const

[wrengr]

Summary

This helps the Merger maintain invariants, as well as clarifying the immutability of the underlying objects (with the one exception of TensorExp::val).

Depends On: D146559

[wrengr]

Split D146559 off into a separate commit

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp

133

If any reviewers know the answer to this, please let me know.

[wrengr]

rebase

[wrengr]

rebase

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

516

I would remove L468-471 as a note to self (unless you have an immediate solution)
mainly because I want to make it very clear that L452-464 apply to all three methods here

550

would it be an option to enforce that update is only used when (1) old value IS set, and (2) new value IS set?
That way, we have a clear set->update->update->clear behavioral chain?

(not saying we should do this, but the question comes to mind)

mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp

133

Ah, that answers my question above ;-)

Even for the reductions, we should have a clear set->update->update->clear chain
if that is not the case, we must have a second look at this

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

516

I think your line numbers are off vs the current version (e.g., L468 is the definition of setLoopDependentTensorLevel...). If you mean the FIXME comment attached to the set method, then I can remove it (and stash it among my local notes-to-self). Just let me know if that's the lines you mean

Afaict, the problem occurs because the entire inner vector is unpacked and stored directly within the outer vector, hence whenever the outer vector is resized then it copies the contents of the inner vectors over to the new buffer. So the obvious solution is to instead have the outer vector only store a pointer to the data of the inner vector (perhaps along with other header info like the vector length); so that resizing the outer vector only copies the pointers themselves around, but leaves the data in place. As for how to implement that, I'm not sure if switching the inner vector to std::vector would be sufficient to do that, or not. If not, then we'd have to do some manual code to explicitly allocate things on the heap and free them at the end. Definitely doable, but not the highest priority at the moment

550

Originally I'd been hoping to have this method do (the equivalent of) { clearExprValue(e); setExprValue(e, v); }, but that didn't work out for CodegenEnv::genLoopBoundary... So then I was hoping to have the method do (the equivalent of) { clearExprValue(e); if (v) setExprValue(e, v); }, but that also didn't work for CodegenEnv::genLoopBoundary...

mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp

133

Yeah. Personally I think the implementation of this method is buggy, precisely because this updateExprValue has to do a unilateral change rather than either of the well-behaved protocols I discussed above (clear+set, or clear+if(v)set).

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

516

Ah, yes, I meant remove the FIXME, so that the NOTE: above clearly applies to all three

mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp

133

Yeah, just leave this FIXME and I will have a look later

[wrengr]

Addressing nits, and rebasing

[wrengr]

rebase

[wrengr]

rebase