Skip to content

[mlir][vector] Update castAwayContractionLeadingOneDim to omit transposes solely on leading unit dims. #85694

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 2 commits into from
Apr 3, 2024
Merged

[mlir][vector] Update castAwayContractionLeadingOneDim to omit transposes solely on leading unit dims. #85694

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7975d0b
only tranpose non leading unit dims
KoolJBlack Mar 18, 2024
d36279f
review comments
KoolJBlack Mar 21, 2024
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
21 changes: 19 additions & 2 deletions mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -398,13 +398,30 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
transposeResults.push_back(targetExpr);
}
}

// Checks if only the outer, unit dimensions (of size 1) are permuted.
// Such transposes do not materially effect the underlying vector and can
// be omitted. EG: perm [1, 0, 2] applied to vector<1x1x8xi32>
bool transposeNonOuterUnitDims = false;
auto operandShape = operands[it.index()].getType().cast<ShapedType>();
for (auto [index, dim] :
llvm::enumerate(ArrayRef<int64_t>(perm).drop_back(1))) {
if (dim != static_cast<int64_t>(index) &&
operandShape.getDimSize(index) != 1) {
transposeNonOuterUnitDims = true;
break;
}
}

// Do the tranpose now if needed so that we can drop the
// correct dim using extract later.
if (tranposeNeeded) {
map = AffineMap::get(map.getNumDims(), 0, transposeResults,
contractOp.getContext());
operands[it.index()] = rewriter.create<vector::TransposeOp>(
loc, operands[it.index()], perm);
if (transposeNonOuterUnitDims) {
operands[it.index()] = rewriter.createOrFold<vector::TransposeOp>(
loc, operands[it.index()], perm);
}
}
}
// We have taken care to have the dim to be dropped be
Expand Down
12 changes: 11 additions & 1 deletion mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,17 @@ func.func @cast_away_contraction_leading_one_dims_nonleadingunitdim_rank4_acctra
return %0: vector<1x1x2x16xf32>
}

// -----

// CHECK-LABEL: func.func @cast_away_contraction_does_not_transpose_leading_unit_dims
// CHECK-NOT vector.transpose
// CHECK: vector.contract
func.func @cast_away_contraction_does_not_transpose_leading_unit_dims(%lhs: vector<1x1x8xi32>,
%rhs: vector<1x8x8xi32>,
%acc: vector<1x8xi32>) -> vector<1x8xi32> {
%result = vector.contract {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d1, d2)>], iterator_types = ["parallel", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %rhs, %acc : vector<1x1x8xi32>, vector<1x8x8xi32> into vector<1x8xi32>
return %result : vector<1x8xi32>
}

// -----
// CHECK-LABEL: func @cast_away_extract_strided_slice_leading_one_dims
Expand Down Expand Up @@ -663,4 +674,3 @@ func.func @drop_unit_dims_scalar_cond_select(%cond: i1, %arg0: vector<1x16xi1>,
%sel = arith.select %cond, %arg0, %arg1 : vector<1x16xi1>
return %sel : vector<1x16xi1>
}