4242#include " llvm/ADT/SmallVector.h"
4343#include " llvm/ADT/StringSet.h"
4444#include " llvm/ADT/TypeSwitch.h"
45+ #include " llvm/Support/FormatVariadic.h"
4546
4647#include < cassert>
4748#include < cstdint>
@@ -6172,49 +6173,57 @@ class FoldTransposeCreateMask final : public OpRewritePattern<TransposeOp> {
61726173// / The algorithm works by partitioning dimensions into groups that can be
61736174// / locally permuted while preserving order, and checks that the transpose
61746175// / only permutes within these groups.
6176+ // /
6177+ // / Groups are either contiguous sequences of 1s, or non-1s (1-element groups).
6178+ // / Consider broadcasting 4x1x1x7 to 2x3x4x5x6x7. This is equivalent to
6179+ // / broadcasting from 1x1x4x1x1x7.
6180+ // / ^^^ ^ ^^^ ^
6181+ // / groups: 0 1 2 3
6182+ // / Order preserving permutations for this example are ones that only permute
6183+ // / within the groups [0,1] and [3,4], like (1 0 2 4 3 5 6).
61756184class FoldTransposeBroadcast : public OpRewritePattern <vector::TransposeOp> {
61766185public:
61776186 using OpRewritePattern::OpRewritePattern;
61786187 FoldTransposeBroadcast (MLIRContext *context, PatternBenefit benefit = 1 )
61796188 : OpRewritePattern<vector::TransposeOp>(context, benefit) {}
61806189
6181- static bool canFoldIntoPrecedingBroadcast (vector::TransposeOp transpose) {
6190+ LogicalResult matchAndRewrite (vector::TransposeOp transpose,
6191+ PatternRewriter &rewriter) const override {
61826192
61836193 vector::BroadcastOp broadcast =
61846194 transpose.getVector ().getDefiningOp <vector::BroadcastOp>();
6185- if (!broadcast)
6186- return false ;
6195+ if (!broadcast) {
6196+ return rewriter.notifyMatchFailure (transpose,
6197+ " not preceded by a broadcast"
6198+ }
61876199
61886200 auto inputType = dyn_cast<VectorType>(broadcast.getSourceType ());
6201+
6202+ // transpose(broadcast(scalar)) -> broadcast(scalar) is always valid
61896203 bool inputIsScalar = !inputType;
6204+ if (inputIsScalar) {
6205+ rewriter.replaceOpWithNewOp <vector::BroadcastOp>(
6206+ transpose, transpose.getResultVectorType (), transpose.getVector ());
6207+ return success ();
6208+ }
6209+
6210+ ArrayRef<int64_t > permutation = transpose.getPermutation ();
61906211 ArrayRef<int64_t > inputShape = inputType.getShape ();
61916212 int64_t inputRank = inputType.getRank ();
61926213 int64_t outputRank = transpose.getType ().getRank ();
61936214 int64_t deltaRank = outputRank - inputRank;
61946215
6195- // transpose(broadcast(scalar)) -> broadcast(scalar) is always valid
6196- if (inputIsScalar)
6197- return true ;
6198-
61996216 // Return true if all permutation destinations for indices in [low, high)
62006217 // are in [low, high), so the permutation is local to the group.
6201- auto isGroupBound = [&](int low, int high) {
6202- ArrayRef<int64_t > permutation = transpose.getPermutation ();
6203- for (int j = low; j < high; ++j) {
6204- if (permutation[j] < low || permutation[j] >= high) {
6218+ auto isGroupBound = [permutation](int low, int high) {
6219+ for (int i = low; i < high; ++i) {
6220+ if (permutation[i] < low || permutation[i] >= high) {
62056221 return false ;
62066222 }
62076223 }
62086224 return true ;
62096225 };
62106226
6211- // Groups are either contiguous sequences of 1s and non-1s (1-element
6212- // groups). Consider broadcasting 4x1x1x7 to 2x3x4x5x6x7. This is equivalent
6213- // to broadcasting from 1x1x4x1x1x7.
6214- // ^^^ ^ ^^^ ^
6215- // groups: 0 1 2 3
6216- // Order preserving permutations for this example are ones that only permute
6217- // within the groups [0,1] and [3,4], like (1 0 2 4 3 5 6).
62186227 int low = 0 ;
62196228 for (int inputIndex = 0 ; inputIndex < inputRank; ++inputIndex) {
62206229 bool notOne = inputShape[inputIndex] != 1 ;
@@ -6223,32 +6232,29 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
62236232 if (groupEndFound) {
62246233 int high = inputIndex + deltaRank;
62256234 if (!isGroupBound (low, high)) {
6226- return false ;
6235+ return rewriter.notifyMatchFailure (
6236+ transpose, llvm::formatv (" output dimensions in interval [{0}, "
6237+ " {1}) aren't locally permuted."
6238+ low, high));
62276239 }
62286240 low = high;
62296241 }
62306242 }
62316243 if (!isGroupBound (low, outputRank)) {
6232- return false ;
6244+ return rewriter.notifyMatchFailure (
6245+ transpose,
6246+ llvm::formatv (" output dimensions in final interval [{0}, {1}) "
6247+ " aren't locally permuted."
6248+ low, outputRank));
62336249 }
62346250
6235- // The preceding logic ensures that by this point, the ouutput of the
6236- // transpose is definitely broadcastable from the input shape. So we don't
6237- // need to call 'vector::isBroadcastableTo', but asserting here just as a
6238- // sanity check:
6251+ // The preceding logic ensures that at this point, the output of the
6252+ // transpose is definitely broadcastable from the input shape. We confirm
6253+ // this as a sanity check:
62396254 bool isBroadcastable =
62406255 vector::isBroadcastableTo (inputType, transpose.getResultVectorType ()) ==
62416256 vector::BroadcastableToResult::Success;
6242- assert (isBroadcastable &&
6243- " (I think) it must be broadcastable at this point."
6244-
6245- return true ;
6246- }
6247-
6248- LogicalResult matchAndRewrite (vector::TransposeOp transpose,
6249- PatternRewriter &rewriter) const override {
6250- if (!canFoldIntoPrecedingBroadcast (transpose))
6251- return failure ();
6257+ assert (isBroadcastable && " It must be broadcastable at this point."
62526258
62536259 rewriter.replaceOpWithNewOp <vector::BroadcastOp>(
62546260 transpose, transpose.getResultVectorType (), transpose.getVector ());
0 commit comments