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
573

while you are here:

that do not alter (not does)

585

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

586

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

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

Generates code

372

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
205

Generates

288

related

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

an attribute

107

directly return

in fact, write this as

if (Value v = ...)

return v

...

719

xxx -> op?

1179–1183

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
699

ditto

709

ditto

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

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

375

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)

74–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
209

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"?

210

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
703

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

713

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
205

"are"

206

"in"

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

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

1174–1175

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
210

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.

320–326

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

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

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
408

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

320–326

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
320–326

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.