Vectorize Linalg ops for dynamic cases. #8888
Conversation
hanhanW
added
the
(deprecated) buildkite:benchmark-android
There are lit tests for dynamic and static convolutions. The change and test are included in this PR as well. |
| SmallVectorImpl<int64_t> &reductionSizes) { | ||
| Optional<SmallVector<int64_t, 4>> staticLoopRanges = op.getStaticLoopRanges(); | ||
| for (auto en : | ||
| llvm::enumerate(llvm::zip(*staticLoopRanges, op.iterator_types()))) { |
There was a problem hiding this comment.
Sorry, I'm a bit confused. I thought this would se tile sizes to 1 for dynamic dims but it's doing the opposite? Could you please add some doc to the method? That would be very helpful.
There was a problem hiding this comment.
hmm, the implementation is setting tile sizes to 1 for dynamic dims. If it is not dynamic dim, it will run into continue statement. If it is dynamic dim, the tile size will be set to 1.
There was a problem hiding this comment.
Sorry, I missed the ! :). Thanks for clarifying! In any case, I think it would be good to add a small doc comment to the method. It's not easy to infer that setAlwaysVectorizeSizes is only about dynamic dims.
There was a problem hiding this comment.
oops, it's auto-merged. Good point, I'll send out a PR to add the comment!
| static void setAlwaysVectorizeSizes(linalg::LinalgOp op, | ||
| SmallVectorImpl<int64_t> ¶llelSizes, | ||
| SmallVectorImpl<int64_t> &reductionSizes) { | ||
| Optional<SmallVector<int64_t, 4>> staticLoopRanges = op.getStaticLoopRanges(); |
There was a problem hiding this comment.
Would this crash if the op doesn't have static loop ranges? Maybe we should return gracefully intead?
There was a problem hiding this comment.
The ShapedType::kDynamic (i.e., -1) will be returned if there are dynamic dims. It's very wrong if we can't get ranges. The ranges are inferred based on indexing_maps and operands. It should be an invalid Linalg op if we can't get ranges.
Maybe we should add it to upstream verifier and remove Optional from this method.
There was a problem hiding this comment.
+1 to removing Optional from there. Not sure why it was added to begin with. I think I reviewed that change and suggested dropping the Optional as well.
There was a problem hiding this comment.
Okay, I'll update upstream method and send out fixes to IREE
There was a problem hiding this comment.
You could land this without those changes. If you look at the method today it never returns None
| SmallVectorImpl<int64_t> &reductionSizes) { | ||
| Optional<SmallVector<int64_t, 4>> staticLoopRanges = op.getStaticLoopRanges(); | ||
| for (auto en : | ||
| llvm::enumerate(llvm::zip(*staticLoopRanges, op.iterator_types()))) { |
There was a problem hiding this comment.
Sorry, I missed the ! :). Thanks for clarifying! In any case, I think it would be good to add a small doc comment to the method. It's not easy to infer that setAlwaysVectorizeSizes is only about dynamic dims.
This is a fallback solution which tile dynamic dims with
1. The tiling logic of first level is not changed.This adds more lit tests which include static cases and dynamic cases.