Skip to content

[mlir][vector] Fix invalid IR in ContractionOpLowering #78130

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jan 16, 2024
Merged

[mlir][vector] Fix invalid IR in ContractionOpLowering #78130

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
155 changes: 104 additions & 51 deletions mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,16 +426,8 @@ struct UnrolledOuterProductGenerator
}

FailureOr<Value> outerProd(Value lhs, Value rhs, Value res,
VectorType lhsType, int reductionDim,
VectorType lhsType, int reductionSize,
std::optional<Value> maybeMask = std::nullopt) {
// Unrolling a scalable dimension would be incorrect - bail out.
if (lhsType.getScalableDims()[reductionDim])
return failure();

int reductionSize = lhsType.getDimSize(reductionDim);
assert(reductionSize > 0 &&
"Reduction dim must be a known static size to allow unrolling");

// Incremental support for masking.
if (mask && !maybeMask.has_value())
return failure();
Expand All @@ -458,49 +450,93 @@ struct UnrolledOuterProductGenerator
return res;
}

/// Helper function for `matmat`, `matvec`, `tmatvec`. Returns the size of
/// dimension `reductionDim`. If the dimension is a scalable dimension,
/// returns "nullopt".
std::optional<int64_t> getReductionSize(VectorType vecType,
int64_t reductionDim) {
// Cannot unroll scalable dimension.
if (vecType.getScalableDims()[reductionDim])
return std::nullopt;
int64_t reductionSize = vecType.getDimSize(reductionDim);
assert(reductionSize > 0 &&
"Reduction dim must be a known static size to allow unrolling");
return reductionSize;
}

/// Two outer parallel, one inner reduction (matmat flavor).
FailureOr<Value> matmat() {
if (!iters({Par(), Par(), Red()}))
return failure();
// Set up the parallel/reduction structure in the right form.
AffineExpr m, n, k;
bindDims(rewriter.getContext(), m, n, k);
Value transposedMask = t(mask, {2, 0, 1});

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think keeping the variable is quite useful in terms of readability. Can we keep it?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can keep it, but it must be nested inside the if (auto reductionSize = getReductionSize(...)) check. (Because the implementation creates IR and the if checks will determine if the pattern succeeds or fails.)

// Classical row-major matmul: Just permute the lhs.
if (layout({{m, k}, {k, n}, {m, n}}))
return outerProd(t(lhs), rhs, res, lhsType, /*reductionDim=*/1,
transposedMask);
if (layout({{m, k}, {k, n}, {m, n}})) {
if (auto reductionSize = getReductionSize(lhsType, 1)) {
// Note: `t` creates new IR. It must be nested within this `if` check
// so that no IR is created when then pattern returns "failure".
Value tLhs = t(lhs);
Value tMask = t(mask, {2, 0, 1});
return outerProd(tLhs, rhs, res, lhsType, *reductionSize, tMask);
}
}
// TODO: may be better to fail and use some vector<k> -> scalar reduction.
if (layout({{m, k}, {n, k}, {m, n}})) {
Value tlhs = t(lhs);
return outerProd(tlhs, t(rhs), res, lhsType, /*reductionDim=*/1,
transposedMask);
if (auto reductionSize = getReductionSize(lhsType, 1)) {
Value tLhs = t(lhs);
Value tRhs = t(rhs);
Value tMask = t(mask, {2, 0, 1});
return outerProd(tLhs, tRhs, res, lhsType, *reductionSize, tMask);
}
}
// No need to permute anything.
if (layout({{k, m}, {k, n}, {m, n}}))
return outerProd(lhs, rhs, res, lhsType, /*reductionDim=*/0,
transposedMask);
if (layout({{k, m}, {k, n}, {m, n}})) {
if (auto reductionSize = getReductionSize(lhsType, 0)) {
Value tMask = t(mask, {2, 0, 1});
return outerProd(lhs, rhs, res, lhsType, *reductionSize, tMask);
}
}
// Just permute the rhs.
if (layout({{k, m}, {n, k}, {m, n}}))
return outerProd(lhs, t(rhs), res, lhsType, /*reductionDim=*/0,
transposedMask);
if (layout({{k, m}, {n, k}, {m, n}})) {
if (auto reductionSize = getReductionSize(lhsType, 0)) {
Value tRhs = t(rhs);
Value tMask = t(mask, {2, 0, 1});
return outerProd(lhs, tRhs, res, lhsType, *reductionSize, tMask);
}
}
// Transposed output: swap RHS and LHS.
// Classical row-major matmul: permute the lhs.
if (layout({{m, k}, {k, n}, {n, m}}))
return outerProd(rhs, t(lhs), res, lhsType, /*reductionDim=*/1,
transposedMask);
if (layout({{m, k}, {k, n}, {n, m}})) {
if (auto reductionSize = getReductionSize(lhsType, 1)) {
Value tLhs = t(lhs);
Value tMask = t(mask, {2, 0, 1});
return outerProd(rhs, tLhs, res, lhsType, *reductionSize, tMask);
}
}
// TODO: may be better to fail and use some vector<k> -> scalar reduction.
if (layout({{m, k}, {n, k}, {n, m}})) {
Value trhs = t(rhs);
return outerProd(trhs, t(lhs), res, lhsType, /*reductionDim=*/1,
transposedMask);
if (auto reductionSize = getReductionSize(lhsType, 1)) {
Value tRhs = t(rhs);
Value tLhs = t(lhs);
Value tMask = t(mask, {2, 0, 1});
return outerProd(tRhs, tLhs, res, lhsType, *reductionSize, tMask);
}
}
if (layout({{k, m}, {k, n}, {n, m}})) {
if (auto reductionSize = getReductionSize(lhsType, 0)) {
Value tMask = t(mask, {2, 0, 1});
return outerProd(rhs, lhs, res, lhsType, *reductionSize, tMask);
}
}
if (layout({{k, m}, {n, k}, {n, m}})) {
if (auto reductionSize = getReductionSize(lhsType, 0)) {
Value tRhs = t(rhs);
Value tMask = t(mask, {2, 0, 1});
return outerProd(tRhs, lhs, res, lhsType, *reductionSize, tMask);
}
}
if (layout({{k, m}, {k, n}, {n, m}}))
return outerProd(rhs, lhs, res, lhsType, /*reductionDim=*/0,
transposedMask);
if (layout({{k, m}, {n, k}, {n, m}}))
return outerProd(t(rhs), lhs, res, lhsType, /*reductionDim=*/0,
transposedMask);
return failure();
}

Expand All @@ -514,24 +550,37 @@ struct UnrolledOuterProductGenerator
return failure();
AffineExpr m, k;
bindDims(rewriter.getContext(), m, k);
Value transposedMask = t(mask);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See above, must be nested.


// Case mat-vec: transpose.
if (layout({{m, k}, {k}, {m}}))
return outerProd(t(lhs), rhs, res, lhsType, /*reductionDim=*/1,
transposedMask);
if (layout({{m, k}, {k}, {m}})) {
if (auto reductionSize = getReductionSize(lhsType, 1)) {
Value tLhs = t(lhs);
Value tMask = t(mask);
return outerProd(tLhs, rhs, res, lhsType, *reductionSize, tMask);
}
}
// Case mat-trans-vec: ready to go.
if (layout({{k, m}, {k}, {m}}))
return outerProd(lhs, rhs, res, lhsType, /*reductionDim=*/0,
transposedMask);
if (layout({{k, m}, {k}, {m}})) {
if (auto reductionSize = getReductionSize(lhsType, 0)) {
Value tMask = t(mask);
return outerProd(lhs, rhs, res, lhsType, *reductionSize, tMask);
}
}
// Case vec-mat: swap and transpose.
if (layout({{k}, {m, k}, {m}}))
return outerProd(t(rhs), lhs, res, lhsType, /*reductionDim=*/0,
transposedMask);
if (layout({{k}, {m, k}, {m}})) {
if (auto reductionSize = getReductionSize(lhsType, 0)) {
Value tRhs = t(rhs);
Value tMask = t(mask);
return outerProd(tRhs, lhs, res, lhsType, *reductionSize, tMask);
}
}
// Case vec-mat-trans: swap and ready to go.
if (layout({{k}, {k, m}, {m}}))
return outerProd(rhs, lhs, res, lhsType, /*reductionDim=*/0,
transposedMask);
if (layout({{k}, {k, m}, {m}})) {
if (auto reductionSize = getReductionSize(lhsType, 0)) {
Value tMask = t(mask);
return outerProd(rhs, lhs, res, lhsType, *reductionSize, tMask);
}
}
return failure();
}

Expand All @@ -547,16 +596,20 @@ struct UnrolledOuterProductGenerator

// Case mat-vec: transpose.
if (layout({{m, k}, {k}, {m}}))
return outerProd(t(lhs), rhs, res, lhsType, /*reductionDim=*/1, mask);
if (auto reductionSize = getReductionSize(lhsType, 1))
return outerProd(t(lhs), rhs, res, lhsType, *reductionSize, mask);
// Case mat-trans-vec: ready to go.
if (layout({{k, m}, {k}, {m}}))
return outerProd(lhs, rhs, res, lhsType, /*reductionDim=*/0, mask);
if (auto reductionSize = getReductionSize(lhsType, 0))
return outerProd(lhs, rhs, res, lhsType, *reductionSize, mask);
// Case vec-mat: swap and transpose.
if (layout({{k}, {m, k}, {m}}))
return outerProd(t(rhs), lhs, res, lhsType, /*reductionDim=*/0, mask);
if (auto reductionSize = getReductionSize(lhsType, 0))
return outerProd(t(rhs), lhs, res, lhsType, *reductionSize, mask);
// Case vec-mat-trans: swap and ready to go.
if (layout({{k}, {k, m}, {m}}))
return outerProd(rhs, lhs, res, lhsType, /*reductionDim=*/0, mask);
if (auto reductionSize = getReductionSize(lhsType, 0))
return outerProd(rhs, lhs, res, lhsType, *reductionSize, mask);
return failure();
}

Expand Down
16 changes: 14 additions & 2 deletions mlir/test/Dialect/Vector/vector-contract-to-outerproduct-matvec-transforms.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -320,8 +320,8 @@ func.func @masked_matvec_k_mk_m(%A: vector<4x2xf32>,
%x: vector<2xf32>,
%b: vector<4xf32>,
%mask: vector<4x2xi1>) -> vector<4xf32> {
// CHECK: vector.transpose %[[MASK]]
// CHECK: vector.transpose %[[A]]
// CHECK: vector.transpose %[[MASK]]
// CHECK-COUNT-2: vector.mask %{{.*}} { vector.outerproduct %{{.*}}, %{{.*}}, %{{.*}} {kind = #vector.kind<add>} : vector<4xf32>, f32 }
%res = vector.mask %mask {
vector.contract #matvec_trait_3 %x, %A, %b
Expand All @@ -339,8 +339,8 @@ func.func @masked_matvec_k_mk_m_scalable_parallel_dim(%A: vector<[4]x2xf32>,
%x: vector<2xf32>,
%b: vector<[4]xf32>,
%mask: vector<[4]x2xi1>) -> vector<[4]xf32> {
// CHECK: vector.transpose %[[MASK]]
// CHECK: vector.transpose %[[A]]
// CHECK: vector.transpose %[[MASK]]
// CHECK-COUNT-2: vector.mask %{{.*}} { vector.outerproduct %{{.*}}, %{{.*}}, %{{.*}} {kind = #vector.kind<add>} : vector<[4]xf32>, f32 }
%res = vector.mask %mask {
vector.contract #matvec_trait_3 %x, %A, %b
Expand Down Expand Up @@ -641,6 +641,18 @@ func.func @masked_tmatvec_k_km_m_scalable_parallel_dim(%A: vector<2x[4]xf32>,
return %res : vector<[4]xf32>
}

// Unrolling scalable reduction dim is not supported - bail out
// CHECK-LABEL: @masked_extract_contract2_scalable_reduction_dim(
// CHECK: vector.contract {{.*}} : vector<[2]x[3]xf32>, vector<[3]xf32> into vector<[2]xf32>
func.func @masked_extract_contract2_scalable_reduction_dim(%arg0: vector<[2]x[3]xf32>,
%arg1: vector<[3]xf32>,
%arg2: vector<[2]xf32>,
%m: vector<[2]x[3]xi1>) -> vector<[2]xf32> {
%0 = vector.mask %m { vector.contract #matvec_trait_1 %arg0, %arg1, %arg2
: vector<[2]x[3]xf32>, vector<[3]xf32> into vector<[2]xf32> } : vector<[2]x[3]xi1> -> vector<[2]xf32>
return %0 : vector<[2]xf32>
}

// ============================================================================
// TD sequence
// ============================================================================
Expand Down
35 changes: 0 additions & 35 deletions mlir/test/Dialect/Vector/vector-contract-to-outerproduct-transforms-unsupported.mlir
View file
Open in desktop

This file was deleted.