This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] Add runtime support for reading a COO tensor and writing the data to the given indices and values buffers.
ClosedPublic

Authored by bixia on Feb 12 2023, 2:23 PM.

Diff Detail

Event Timeline

bixia created this revision.Feb 12 2023, 2:23 PM
Herald added a project: Restricted Project. · View Herald Transcript
bixia requested review of this revision.Feb 12 2023, 2:23 PM
wrengr added inline comments.Feb 13 2023, 12:25 PM
mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
114–132

It makes me squeamish to have so much code duplication between these macros and the basic ones.

There isn't a great way to handle this due to (intentional) limitations in CPP. But the best option is to rephrase the original x-macros to be variadic: i.e., something like MLIR_SPARSETENSOR_FOREVERY_V(DO,...) DO(F64, double, __VA_ARGS__) etc. Which will at least save you from needing to duplicate the innermost x-macro. However, IIRC Peiming had some issues getting this to work on Windows last time I raised this concern, so you should check with him about what problems he ran into.

If you cannot get the variadic macros to work, then at the very least you should file a bug (assigned to me) and add a FIXME comment with the bug number.

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
394

Please rename to dimRank to match the terminology used elsewhere (and as per D143800 and https://github.com/openxla/stablehlo/pull/1143).

399

Just use !IsSymmetric

403–404

Why use new temp variables rather than just reusing the function parameter variables?

406–410

I like how this helps save code duplication, but I feel it's at the wrong level of abstraction. E.g., if you use the indices and values variables directly rather than using the pIndices and pValues variables, then you could have this function advance the two pointers after writing to them —thus maintaining the invariant that the pointers are always pointing to the next free entry (or one-past-the-end). Doing so keeps both halves of the pointer logic (writing to them and advancing them) together; whereas the current code mixes some of the pointer logic with the isSorted logic.

Having this function advance the pointers after writing to them would also help in clarifying things in the if constexpr (IsSymmetric) conditional, since it would be doing indices[0,1] = indices[-2,-1]; values[0] = values[-1] which (to my mind) helps clarify that it's writing the *next* entry based on the *previous* entry.

435–436

Since IsSymmetric can only be true when dimRank == 2, the pIndices[2] and pIndices[3] memory addresses look like they should be invalid. Yes, I know why the code actually works, but it would be prudent to add a comment explaining why it's okay so that other readers won't get confused.

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
646

please rename to dimRank

bixia updated this revision to Diff 497143.Feb 13 2023, 4:11 PM
bixia marked 7 inline comments as done.

Address review comments.

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
114–132

Thanks. Done.

aartbik added inline comments.Feb 14 2023, 11:03 AM
mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
38

does it really say, "function-liked"?

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
252

Can we make the NNZ -> NSE or NOE change

(number of stored entries, rather than nonzeros, since we may store a few zeros ;-)

434

I feels like the way we handle symmetry (with the expand option) is starting to convolute our implantation quite a bit. Here we have the whole complication of knowing the number of entries in advance, but still have to compute how many are on the diagonal and thus not counted twice.

perhaps we should abandon the "expand symmetry" option altogether, and simply read in symmetric matrices as lower/upper triangular and start working on adding the proper "property" to our dim/level types

bixia updated this revision to Diff 498141.Feb 16 2023, 2:29 PM
bixia marked 2 inline comments as done.

Update after removing expand_symmetry property. Rename nnz to nse.

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
38

see here

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
434
wrengr added inline comments.Feb 17 2023, 1:01 PM
mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
114–132

Thanks for doing that :)

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
252

Fwiw I've been using "NSE" in my renaming patch. (I like that that specifically has "stored" in it.)

255

I find it confusing that this method has the same name as readCOO<I> but very different arguments. You should rename this method to readIntoBuffers (or readInto if you're feeling succinct).

Also, you should have this method take the lvlRank argument in addition to dim2lvl so that it matches the API of readCOO<I>. Also, I follow the style of (a) providing buffer lengths before the buffers themselves, and (b) placing out-parameters last. Thus, the arguments should be in the order (lvlRank, dim2lvl, maxNse, indices, values)

300

Ditto my comments at readCOO<I,V>, this method should:

  • be named readIntoBuffersLoop.
  • take the lvlRank for API symmetry (and for the future-looking case, since we may/will eventually want to read files in via a non-permutation).
  • use the argument order (lvlRank, dim2lvl, maxNse, indices, values).

In addition, you should change the type of dim2lvl to be detail::PermutationRef, so that it matches the API of readCOOLoop<I>. When you move the construction of the PermutationRef to readIntoBuffers<I,V> be sure to also add the assert(lvlRank == dimRank && "Rank mismatch"); as well.

Ideally both readCOO<I> and readIntoBuffers<I,V> should take PermutationRef rather than a pointer: because we want to push construction of PermutationRef as far up the callstack as possible, since doing so helps make more of the code typesafe and improves the C++ API. (Honestly I'm not entirely sure why I didn't do that in the first place; but I can make that change in a separate CL.)

387

You should add an assertion here that nnz <= maxNse, to ensure we catch errors as early as possible.

You should also rename the local variable to nse so that it matches the function parameter. Don't worry about the mismatch between the local variable and the method name; I'll be fixing that in my big renaming CL

395

Why split this out as a separate function? It would be cleaner to combine it with getOneElement so that it maintains the invariant that indices/values are always pointing to free memory.

I'm guessing you're worried about the sortedness check. If so, then the cleanest thing would be to move that sortedness check into getOneElement so that the function maintains that invariant as well! To make that work all you have to do is set isSorted=false before getting the first element (so that the conditional fails and all we do is read and advance pointers), and then set isSorted=true before entering the loop.

Rewriting getOneElement to include both advanceOneElement and the sortedness check will greatly help readability and reasoning. It'll help not just because it's maintaining invariants (which always helps), but because it's using the proper inductive hypothesis. It's precisely because that's the correct IH that (a) the function maintains a bunch of invariants, and (b) the rest of the code is just a simple loop calling the function plus a basis case outside of the loop.

408

Once you add the assert(nse <= maxNse) above, then you don't need this assertion anymore

bixia updated this revision to Diff 498516.Feb 17 2023, 3:56 PM
bixia marked 4 inline comments as done.

Reorder parameters. Rename local nnz to nse. Combines two lambdas.

bixia added inline comments.Feb 17 2023, 3:56 PM
mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
255

Rename to readToBuffers.
Remove maxNse, and move the assert to the client. Add doc to specify precondition.
Didn't add lvlRank and dim2lvl to this routine through. Unlike readCOO, which allocate the COO and needs those two values, the routine here doesn't need to allocate the COO and doesn't need those parameters. Let me know if you don't agree on this.

300

Rename to readToBuffersLoop.
Use PermutationRef instead.

387

Rename nnz to nse.
Assertion is added to caller.

bixia updated this revision to Diff 498848.Feb 20 2023, 7:12 AM

Trying to fix the build.

aartbik added inline comments.Feb 21 2023, 12:52 PM
mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
297

we don't template over isSymmetric?

402

this is only safe when we already have one element.

You do have the unrolled for loop below, but I don't see how that helps?

405

it feels like

 if (prevIndices[d] != indices[d]) { 
    if (prevIndices[d] >= indices[d])
       isSorted = false;
    break;
}

avoids writing true to the isSorted field for every element.
In fact, the == case denotes a duplicate

414

is this one reading outside the buffer?

bixia updated this revision to Diff 499523.Feb 22 2023, 8:28 AM
bixia marked an inline comment as done.

Avoid writing isSorted with the same value. Add comments to the code.

bixia updated this revision to Diff 499899.Feb 23 2023, 9:31 AM

Split a macro into variadic version and non-variadic version to avoid warnings.

I think you missed some of my comments?

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
256

Technically this should be a *Level, since it maps dims (as index now) to levels in the array.

@wrengr do we want a typedef for dim2lvl and lvl2dim too just to be complete?

297

did you see my comment. The comment on L297 still mentions IsSymmetric, which is gone now

bixia updated this revision to Diff 500002.Feb 23 2023, 3:42 PM
bixia marked an inline comment as done.

Remove IsSymmetric in comment.

wrengr added inline comments.Feb 23 2023, 3:43 PM
mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
252

please use "coordinates" to match the new naming convention

255

Since readToBuffers does still construct a PermutationRef, it should have the same assert(lvlRank == dimRank && "Rank mismatch") assertion that readCOO does. While the PermutationRef ctor checks that dim2lvl is in fact a permutation, having this assertion before calling the ctor does two things: First, it ensures consistent error messages between the runtime/codegen passes, which is important for UX. Second, it helps catch errors —dim2lvl can be a permutation of dimRank even if dimRank != lvlRank, so we need to be sure to catch the rank mismatch in order to ensure correctness.

Also, as I mention in other comments below, the readToBuffersLoop does in fact want the lvlRank argument

255

Please rename the I template parameter to C to match the new naming convention.

256

We can, though that'd best be done as a separate CL, just to keep things clean.

Fwiw, I'm not too worried about it at the moment. Once I wrap up my current CL, I'll be getting back to updating the compiler to support non-permutations; and once the compiler is in a position to generate the dim2lvl/lvl2dim functions, that's when I was going to go through and update everything in the RT to take function-pointers in lieu of the current vectors/memrefs/pointers. Thus, so long as the new methods follow the same coding conventions as the pre-existing methods, then it should be easy enough to change at that point.

256

Please rename to lvlCoordinates to match the new naming convention.

(N.B., we do not use the shorter lvlCoords here because this is a shared pool of coords, rather than just the coords for a single element)

299

Please rename the I template parameter to C to match the new naming convention.

300

Please rename to lvlCoordinates to match the new naming convention.

(N.B., we do not use the shorter lvlCoords here because this is a shared pool of coords, rather than just the coords for a single element)

366

Please rename the I template parameter to C to match the new naming convention.

367

Please rename to lvlCoordinates to match the new naming convention.

(N.B., we do not use the shorter lvlCoords here because this is a shared pool of coords, rather than just the coords for a single element)

385

Please rename the I template parameter to C to match the new naming convention.

387

Please rename to lvlCoordinates to match the new naming convention.

(N.B., we do not use the shorter lvlCoords here because this is a shared pool of coords, rather than just the coords for a single element)

391

Please rename to dimCoords to match the new naming convention.

402–403

This comment should be moved up to where you first declare/define the isSorted variable, rather than here

404

This should be lvlRank.

Once we change dim2lvl to be an arbitrary mapping rather than a permutation, then each element will have lvlRank-many coordinates, therefore we need to make sure to advance the lvlCoordinates pointer by lvlRank for each element.

404

Please rename to prevLvlCoords to match the new naming convention

405

This should be lvlRank.

405

This should be l

413

This should be lvlRank.

420

These braces aren't needed for so simple a loop body

mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h
286

please say "coordinates" to match the new naming convention

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
634

use CNAME and C to match the new naming convention

aartbik added inline comments.Feb 23 2023, 3:49 PM
mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
256

We can, though that'd best be done as a separate CL, just to keep things clean.

Yeah, I was not suggesting to add that here, but this was really more a note to you on perhaps adding two more types for the two mappings.
Using something like

void foo(Dim2Lvl dim2lvl)

would be nice ;-=)

bixia updated this revision to Diff 500008.Feb 23 2023, 4:15 PM
bixia marked 4 inline comments as done.

Rename indices to lvlCoordinates, add assert(lvlRank == dimRank && "Rank mismatch").

bixia updated this revision to Diff 500012.Feb 23 2023, 4:27 PM
bixia marked 12 inline comments as done.

Rename I to C, indices to coordinates.

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
256

Will use the new typedef here when we have it.

256

Ok. Leaving this as it is.

297

We remove the expansion for isSymmetric in https://reviews.llvm.org/D144059 and now only record the property.

297

Sorry for missing this.

402

I add a comment to the code to explain this.

405

Agree. changed to
if (prevIndices[d] > indices[d])

isSorted = false;
408

we don't have maxNse anymore.

414

This reads the first element with isSorted = false, to prepare for the reading of the rest, with isSorted initialized to true. Add a comment to the code.

Just a few more nits re naming. Once these and the couple remaining ones from before (dimInd->dimCoords; d->l) are finished, then this should be good to go! :)

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
252

Shouldn't that say "reads from the file" or something similar? (Since this isn't using SparseTensorCOO)

394

You already asserted this in readToBuffers, so don't need to do it again here (since readToBuffersLoop is private).

411

This should be just prevLvlCoords since it's intended to refer to the coordinates of a single element (despite the fact that those "coords" live inside the larger "coordinates" buffer). Yeah, I know it looks idiosyncratic here, since the function parameter is being used both for the entire collection of "coordinates" (when viewed from outside the function) as well as being reused to mean just the "coords" of the most-recent element (when viewed from inside the loop).

If you want to avoid the idiosyncrasy, then feel free to (re)introduce a local variable lvlCoords for the pointer that we advance by lvlRank each time this lambda is called. (Yeah, you used to have such a local variable in earlier versions of this CL; but when I suggested removing it, that was before I started on the Big Renaming CL to try and clean all this stuff up ;)

bixia updated this revision to Diff 500220.Feb 24 2023, 8:32 AM
bixia marked 9 inline comments as done.

Rename d->l, dimInd->dimCoords, prevlvlCoordinates->prevLvlCoords.

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
252

add "from the file".

wrengr added inline comments.Feb 24 2023, 2:01 PM
mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
376–378

personally I'd rather this be phrased as const bool isSorted = IsPattern ? ... : ...; (or const bool isSorted = isPattern() ? ... : ...;) since that's a lot more succinct and keeps things const

402

should be dimCoords here too

407–418

To avoid the code repetition, it'd be nice to replace this whole conditional and loopnest with isSorted &&= elemLT(Element<_>(prevLvlCoords,_), Element<_>(lvlCoordinates,_)) where const ElementLT<_> elemLT(lvlRank); is defined above the lambda. Though I'm hesitant to suggest that wholeheartedly, because there's the issue of filling in those underscores. The natural thing to do would be to use V and *values/*(values-1), but that incurs the cost of copying V and hence would slow things down. (The unnatural thing to do would be to define struct Empty {}; or similar and then use that instead of V; but that cure is worse than the disease imo.)

So, rather than making that change per se, could you add a TODO comment saying that we should define a new CoordsLT which is like ElementLT but doesn't have the V parameter? That way I can remember to fix that in the future since it'd be trivial enough to define the new comparator. Alternatively, could you file a ticket with that same comment and assign it to me?

409

should be prevLvlCoords here too

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
637

rename to cref (or more ideally lvlCoordinatesRef) to match the new naming scheme

638

You don't need to assert these, since that's already handled by the ASSERT_NO_STRIDE macro.

643–644

These variables should use the full "size" rather than just "s"

646

Does this really need to be ==? The more liberal thing would be to assert(lvlCoordinatesSize >= reader.getNNZ() * lvlRank); which also matches the liberalness of assert(valuesSize >= reader.getNNZ()). If you really do want equality, then you should use the ASSERT_USIZE_EQ macro instead. (And if you really want >= then you should define a new ASSERT_USIZE_GE/ASSERT_USIZE_LE macro that mirrors ASSERT_USIZE_EQ except for the comparison operator used.)

646

This check should be using lvlRank instead, since that's what the buffer actually needs to store. Of course, there's no way to get the lvlRank from the reader itself, so you'll need to pass it in as an additional argument to this function.

647

"enough"

648–649

Since you don't need the ps variable anywhere other than this assertion, you should use the ASSERT_USIZE_EQ macro instead.

654

should be lvlCoordinates to match the name in the readToBuffers method

bixia updated this revision to Diff 500651.Feb 26 2023, 7:54 PM
bixia marked 12 inline comments as done.

Address review comments: use ASSERT_USIZE_EQ, rename dimRank to lvlRank, fix comments.

bixia added inline comments.Feb 27 2023, 12:57 PM
mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
407–418

Add a TODO.

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
637

use cref.

646

Relax the check to use <=, and use replace dimRank with lvlRank.

aartbik accepted this revision.Feb 27 2023, 8:07 PM
aartbik added inline comments.
mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
405

The comment "We set isSorted to false to read the first element" feels like you are actively setting it here instead of above (where you also comment on that). I think here it should read

"// Note that isSorted was set to false while reading the first element to...."

This revision is now accepted and ready to land.Feb 27 2023, 8:07 PM
bixia updated this revision to Diff 501140.Feb 28 2023, 7:46 AM
bixia marked an inline comment as done.

Fix comment.