[VectorDistribution] Add layout analysis for distributing multi-dim reduction (2/4) #18800
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
…pport for the case when the reduction is inside scf.for operation Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
bangtianliu
requested review from
antiagainst,
qedawkins and
hanhanW
as code owners
bangtianliu
changed the title
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
bangtianliu
changed the title
bangtianliu
changed the title
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
bangtianliu
changed the title
kuhar
reviewed
Oct 17, 2024
Comment on lines
+110
to
+112
| // Returns a boolean flag indicating whether the input value 'val' is a | ||
| // vector, determined by checking its rank. | ||
| bool isVector(VectorValue val); |
There was a problem hiding this comment.
This strikes me as an odd helper: you give 'vector' a new meaning without introducing a name. Instead, I'd either flip it and add a helper like isRank0(VectorValue val), or just expand the check where you need it.
Comment on lines
+215
to
+219
| bool isVector(VectorValue val) { | ||
| if (val.getType().getRank() != 0) | ||
| return true; | ||
| return false; | ||
| } |
There was a problem hiding this comment.
Suggested change
| bool isVector(VectorValue val) { | |
| if (val.getType().getRank() != 0) | |
| return true; | |
| return false; | |
| } | |
| bool isVector(VectorValue val) { | |
| return val.getType().getRank() != 0; | |
| } |
Comment on lines
+1057
to
+1063
| // Result lattice not has a layout yet. | ||
| if (resultLattices.empty()) | ||
| return; | ||
|
|
||
| // We do not support multiple results yet. | ||
| if (resultLattices.size() != 1) | ||
| return; |
| return; | ||
|
|
||
| for (RegionSuccessor successor : successors) { | ||
| if (auto succ = successor.getSuccessor()) { |
There was a problem hiding this comment.
Can you use the actual type here instead of auto? It's not clear based on the RHS
Comment on lines
712
to
713
| assert(isa<VectorType>(inputType) && | ||
| "Scalar broadcast not supported for now."); |
Comment on lines
+135
to
+136
| if (isa<VectorType>(replacement.getType()) && | ||
| cast<ShapedType>(replacement.getType()).getRank() != 0) { |
There was a problem hiding this comment.
use dyn_cast instead:
if (auto x = dyn_cast<Y>(y)) {
if (x.something() == Z) {| replaceOpWithDistributedValues(rewriter, multiReduceOp, accReduced); | ||
| } else { | ||
| Value accReducedVal = rewriter.create<vector::ExtractOp>( | ||
| loc, accReduction, SmallVector<int64_t>{0}); |
There was a problem hiding this comment.
You don't need a vector here, you can do: ArrayRef{int64_t(0)}
| if (accVector) { | ||
| locallyReduced = dyn_cast<VectorValue>(localReduction.getResult()); | ||
| } else { | ||
| VectorType vecType = VectorType::get(SmallVector<int64_t>{1}, elemTy); |
There was a problem hiding this comment.
Also here, no need to use a vector
Comment on lines
+427
to
+428
| if (accVector) { | ||
| accElemTy = accVector.getType().getElementType(); |
There was a problem hiding this comment.
You can use getElementTypeOrSelf. Also below.
Comment on lines
+311
to
+312
| bool isSrcVector = (srcVector) && (isVector(srcVector)); | ||
| if (!isSrcVector) { |
There was a problem hiding this comment.
This is the only use of isSrcVector, so it makes sense to inline it
Suggested change
| bool isSrcVector = (srcVector) && (isVector(srcVector)); | |
| if (!isSrcVector) { | |
| if (!srcVector || !isVector(srcVector)) { |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Splitting #18519 into four patches.
Depends #18784
This is the second one, adding the corresponding layout analysis and especially supporting the case where reduction is performed inside scf.for operation.
Also, the relevant tests are added.