This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] add merger/topo sort support for slice-based affine sparse index codegen
ClosedPublic

Authored by Peiming on Jan 30 2023, 1:26 PM.

Diff Detail

Event Timeline

Peiming created this revision.Jan 30 2023, 1:26 PM
Herald added a project: Restricted Project. · View Herald Transcript
Peiming requested review of this revision.Jan 30 2023, 1:26 PM
Peiming updated this revision to Diff 505292.Mar 14 2023, 3:12 PM

fix rebase errors.

aartbik added inline comments.Mar 16 2023, 11:37 AM
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
535

I think you will need to expand this comment with more details on what an "dependent" slice is exactly.
Also, I realize this is just bookkeeping for now, so is merger the right place (do you plan to do more with this?)

Peiming updated this revision to Diff 505945.Mar 16 2023, 3:42 PM

rebase + fix variable name

Peiming updated this revision to Diff 505947.Mar 16 2023, 4:06 PM

add more comments.

wrengr added inline comments.Mar 16 2023, 4:46 PM
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
450

Please use the TensorId and LoopId types. Although they're just typedefs for now, by the end of the week I'll have a CL that changes them to be actual types.

450

Elsewhere in this file we use the name "i" for loop identifiers, so please stick with that for consistency.

452

Please use the Level type. Again, although it's just a typedef for now, I'll be changing that soon.

452

should be "i"

453

What does "idx" mean here? I went through a lot of work cleaning up all the ambiguity about several different meanings of "index", so please use a non-ambiguous name that matches how things are named elsewhere

459

should be "i"

465

elsewhere in this file we spell out "lvl" in order to avoid any potential confusion with loop identifiers, so you should stick with that here.

470

ditto

476

Based on the body of this method, that should actually be a TensorLoopId

540

Please clarify what exactly this type is supposed to be.

540

I've avoided using the name "ldx" in this file because it's too confusing. Please change this to "loop" or similar, to match the names used by the other fields

545

please rename to "lvlToDependantLoop" or similar. (We never use "tlvl" anywhere. "depended" is misspelled. And I've intentionally avoided using "ldx" elsewhere in this file)

545

Please clarify exactly what this type is supposed to be

wrengr added inline comments.Mar 16 2023, 4:46 PM
mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
584–759

please clarify exactly what all these unsigned are actually supposed to be.

612

Level

614

LoopId

636

LoopId

682

Please clarify exactly what this type is supposed to be. Is it Dimension or something else?

690

ditto

694

ditto

709

ditto

736

ditto

741

ditto

1858

ditto

1859

ditto

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
224

ditto

226

ditto

227

ditto

aartbik added inline comments.Mar 16 2023, 4:54 PM
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
463

list of loop indices

(you use plural when saying something like set of XXXs)

536

set of (tensor, level) pairs.

537

expressions

539

can you break this like

... that

i and j are used ...

541

dependent

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
113

this comment does not match the implementation
this is just a collector, right, adding it to the dims vector

what you describe seems to match intended use elsewhere?

267

Can you give a small example of that goes in and out here; just giving a basic example will go a long way in helping the reader

644

"relax" is a bit easier to read (I think), same for function name
tryRelaxAffineConstraints

680

\in just in (LaTeX ?)

753

For this revision, can you not change the comment of this method

/ Computes a topologically sorted iteration graph for the linalg operation.
/ Ensures all tensors are visited in natural coordinate order. This is
/// essential for sparse storage formats since these only support access

so that the DIFF will be easier to line up the method with the original code?

Peiming marked 39 inline comments as done.Mar 16 2023, 5:25 PM
Peiming added inline comments.
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
453

it means tensor expression like T_{i+j}, that uses a non-trivial expression to index the tensor.

Changed it to isLvlWithNonTrivialIdxExp to be more complete.

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
680

Yeah

Peiming updated this revision to Diff 505958.Mar 16 2023, 5:25 PM
Peiming marked 2 inline comments as done.

address comments.

Peiming updated this revision to Diff 505960.Mar 16 2023, 5:34 PM

fix error.

wrengr added inline comments.Mar 16 2023, 6:38 PM
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
450

If you're factoring this out, then you might as well factor out the getLvl(b) and getDimLevelType(b) expressions too...

Originally I was going to call that a minor nit, but on reflection I've noticed that the current implementation allows the first callback to mutate those fields (via setLevelAndType) before the second callback. That's an extremely bug-prone design, so I can't imagine that it's intentional. But if for some reason it is intentional, then it needs to be explicitly documented that this is expected behavior and why it's intended.

451

This should also be "i"

452

"level"

454

"level"

454

Please rename this to loopToDependencies. Again, I do not use "ldx" anywhere in this file.

455

Please rename this to levelToDependentLoops. Again, I've done a lot of work to avoid the ambiguous use of "idx".

458

"loop"

458

ditto my other comment about using "///"

458

Should add a verb here (e.g., "Checks whether").

463

"which appear"

463

Although this matches the name of the getMatchingIndexingMap method, I think it would be a lot clearer to call these "subscript expressions".

463–464

ditto my other comment about using "///"

465

Should be "dependent", here and everywhere else in this file.

469

"lvl"

469

When providing documentation this must be "///" so that it's picked up by the program that generates the documentation.

470

This doesn't match the documentation. If the documentation is correct then this should be LoopId i. Otherwise, the documentation needs updating.

470

"Lvl"

474

This is inconsistent with the rest of the file. A LatPoint contains a whole set of TensorLoopIds.

(If what is currently called TensorLoopId is actually supposed to be a "lattice point", then we need to rename the LatPoint and LatPointId types to something else to avoid confusion. In addition to updating all the documentation, method names, etc.)

474

Ditto, although this matches the name of the getMatchingIndexingMap method, it would be clearer to call these "subscript expressions" throughout this file

476

Should rename this to hasDependencies.

(Because: this has nothing to do with "levels" so far as I can see. Calling it hasDependencies matches the name of the field it's querying. And it's always the case that "dependencies" correspond with "non-trivial subscripting expressions", so there's no need to spell out the latter in the method name.)

535

"identifier"

536

"identifier"

537

"subscript"

539

"subscript"

542

"identifier"

544

"its"

547

"identifier"

547

"lvl"

Peiming updated this revision to Diff 505972.Mar 16 2023, 7:36 PM

fix issue

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

OMG, you are SO right, i forget the else!

Peiming updated this revision to Diff 506109.Mar 17 2023, 9:23 AM
Peiming marked 28 inline comments as done.

address comments.

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

I will keep it as TACO paper uses the term index variable.

476

I will keep it for now, as it matches the technique report that I am writing, we can revisit the term later ;-)

Peiming updated this revision to Diff 506119.Mar 17 2023, 9:48 AM

fix names.

aartbik added inline comments.Mar 20 2023, 12:42 PM
mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
438–444

please document the parameters of the expected callback (since the bool is no longer "trivial")

458

in the description of the callback above, use names, and then say

/*name=*/false here and above for true, so that it becomes more readable

mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h
102

hmm, I actually liked that I was able to hide the topsort completely in this new datastructure, but now it is leaking again
(actually, the slice was already leaking)

but we have plans for an iteration graph class, then it would make sense perhaps to hide it there..

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
279

also, not now but perhaps later, if feels like all the affine analysis could also be in its own class, or perhaps with a new iteration graph class

398

this method also does a lot more analysis than before, perhaps update the method doc for now a bit to reflect this
(having a new affine analyzer could be next)

687

typo

1803

yeah, avoiding too many special cases, but making the general case and special case more uniform would be nice in the longer run

Peiming updated this revision to Diff 506691.Mar 20 2023, 1:19 PM
Peiming marked 7 inline comments as done.

address comments.

Peiming added inline comments.Mar 20 2023, 1:21 PM
mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h
102

I think I can hide it again in the next commit, because I move some code there into CodegenEnv.

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
279

Yeah, that would be better.

aartbik accepted this revision.Mar 20 2023, 2:05 PM

Wait to see if Wren is also okay with how her comments were addressed.

This revision is now accepted and ready to land.Mar 20 2023, 2:05 PM
This revision was landed with ongoing or failed builds.Mar 20 2023, 2:24 PM
This revision was automatically updated to reflect the committed changes.