This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] support dynamic sparse tensor slices.
ClosedPublic

Authored by Peiming on Jan 11 2023, 11:25 AM.

Diff Detail

Event Timeline

Peiming created this revision.Jan 11 2023, 11:25 AM
Herald added a reviewer: aartbik. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
Peiming requested review of this revision.Jan 11 2023, 11:25 AM
Peiming updated this revision to Diff 488378.Jan 11 2023, 2:03 PM

add check test

Peiming updated this revision to Diff 488384.Jan 11 2023, 2:24 PM

add matmul dynamic slice integrate test

Peiming updated this revision to Diff 497162.Feb 13 2023, 5:32 PM

fix crash.

Peiming updated this revision to Diff 503949.Mar 9 2023, 3:22 PM

fix typos

aartbik added inline comments.Mar 10 2023, 9:56 AM
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
569

while you are here:

that do not alter (not does)

580–581

can we make this / an and
(it reads like a division right now)

582

I would continue this comment on the same line L585, since it just continues the reasoning

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h
364

Generates code

369

Generates code

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.cpp
98

put this before the if

(so that First, Second, Third are at same leve)

111

put this Third before the if

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h
193

Generates

272

related

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
100

an attribute

107

directly return

in fact, write this as

if (Value v = ...)

return v

...

718

xxx -> op?

1181–1185

we cannot generate

Peiming updated this revision to Diff 504239.Mar 10 2023, 11:49 AM
Peiming marked 13 inline comments as done.

address comments.

Peiming updated this revision to Diff 504240.Mar 10 2023, 11:52 AM

revert unintended changes.

A bunch of repetitive comments about using the proper types. If you rebase this after D145756, then that handles a bunch of them.

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
676

ditto

686

ditto

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h
367

If this is actually a dimension, then use the Dimension type. (Otherwise, should rename the variable)

372

If this is actually a dimension, then use the Dimension type. (Otherwise, should rename the variable)

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.cpp
47

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

53

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

62

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

71–75

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

84

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

91

I've noticed you use this variable name a lot in this file. What exactly do you mean by it? Seems like there should be a more informative name

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h
197

Is this actually (1) a memref of all the coordinates associated with a single element; in which case it should be "coords"; or (2) a single coordinate value; in which case it should be "crd"?

198

If this is actually a level, then use the Level type. (Otherwise, should rename the variable)

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
102

I'm pretty sure there's already something that does this in StaticValueUtils.h

wrengr added inline comments.Mar 10 2023, 12:31 PM
mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
680

This should just check the offset local variable, rather than calling the function again. And for brevity it ought to use the syntax if (auto offset = ...) return ...

690

ditto

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.cpp
49

You should add a FIXME comment about how toOrigDim is deprecated.

55

ditto

120

Should just use a foreach loop instead; i.e., for (auto cond : ArrayRef(conds).drop_front())

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h
193

"are"

194

"in"

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
1163

Instead of doing this manually, it would be a lot easier to follow if you use for (... : llvm::enumerate(llvm::zip(...))) instead

Also, should use the Dimension type

1176–1177

yes please :)

Peiming updated this revision to Diff 504279.Mar 10 2023, 2:16 PM
Peiming marked 22 inline comments as done.

address comments.

Peiming added inline comments.Mar 10 2023, 2:18 PM
mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.cpp
91

it means transformed coordinates.

120

conds is an SmallVector, and it does not provide a drop_front, I prefer to start with 1 here instead of casting it to ArrayRef before iteration.

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h
198

The header for Level is not included yet, will do it later to avoid more conflicts.

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
102

Thanks for letting me know.

aartbik accepted this revision.Mar 10 2023, 2:19 PM
This revision is now accepted and ready to land.Mar 10 2023, 2:19 PM
Peiming updated this revision to Diff 504281.Mar 10 2023, 2:25 PM

add bazel dependence.

wrengr added inline comments.Mar 10 2023, 2:44 PM
mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.cpp
120

I still think the foreach loop is cleaner, since it captures the intention directly and avoids introducing extraneous details like the index variable and getting the size.

243–249

Elsewhere the code always uses the name "pHi", so that name should be used here too

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
1163

Elsewhere throughout the MLIR code folks name this "it". Also, you should make it auto& to avoid an unnecessary copy

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h
414

Did you mean to add this back in?

Back when I renamed stuff in this file I cleaned all those up (since at that point, at least, the StorageSpecifierKind argument always took a Level), so the "[CLARIFY_DIM_LVL]" note no longer exists. If this CL reintroduces ambiguity/confusion issues, then you should give a longer comment explaining what the problem is

Peiming updated this revision to Diff 504286.Mar 10 2023, 2:54 PM
Peiming marked 3 inline comments as done.

address comments.

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.cpp
120

Okay. Changed

243–249

This is slightly different here actually, maybe I will just call it lvlSz to be more precise.

wrengr added inline comments.Mar 10 2023, 2:59 PM
mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.cpp
243–249

hmm, but it comes from createOrFoldDimOp(d) so shouldn't it be dimSz? Then again, since it's being stored into lvlSizes, that does suggest it ought to be lvlSz... So can you add a comment explaining why createOrFoldDimOp(d) is safe/correct to use here? (And presumably it'll be a TODO/FIXME, since I'm guessing this is just a stand-in for once I finish up the proper dim<->lvl conversion class)

Peiming updated this revision to Diff 504289.Mar 10 2023, 3:08 PM

address comment.

This revision was landed with ongoing or failed builds.Mar 10 2023, 3:12 PM
This revision was automatically updated to reflect the committed changes.