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
550

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
444

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.

444

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

446

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

447

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"

466

should be "i"

472

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.

477

ditto

483

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

555

Please clarify what exactly this type is supposed to be.

555

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

560

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)

560

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
764–765

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

792

Level

794

LoopId

816

LoopId

862

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

870

ditto

874

ditto

889

ditto

916

ditto

921

ditto

1864

ditto

1865

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
470

list of loop indices

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

551

set of (tensor, level) pairs.

552

expressions

554

can you break this like

... that

i and j are used ...

556

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?

265

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

824

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

860

\in just in (LaTeX ?)

933

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
447

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
860

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
444

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.

446

"level"

448

"level"

458

This should also be "i"

461

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

462

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

465

"loop"

465

ditto my other comment about using "///"

465

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

470

"which appear"

470

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

470–471

ditto my other comment about using "///"

472

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

476

"lvl"

476

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

477

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

477

"Lvl"

481

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.)

481

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

483

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.)

550

"identifier"

551

"identifier"

552

"subscript"

554

"subscript"

557

"identifier"

559

"its"

562

"identifier"

562

"lvl"

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

fix issue

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

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
470

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

483

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–439

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

452

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
277

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

399

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)

692

typo

1809

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
277

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.