Skip to content
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

[VectorDistribution] Fix 0-rank vector.broadcast distribution #19007

Merged
merged 3 commits into from
Nov 11, 2024
Merged

[VectorDistribution] Fix 0-rank vector.broadcast distribution #19007

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7030b8a
0-d vector fix
Groverkss Nov 6, 2024
8848fd8
Remove revert
Groverkss Nov 8, 2024
6bef0ea
Add test for zero-rank broadcast
Groverkss Nov 11, 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
6 changes: 3 additions & 3 deletions compiler/src/iree/compiler/Codegen/Common/GPU/GPUNestedLayoutDistributionPatterns.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -305,9 +305,9 @@ struct DistributeBroadcast final : OpDistributionPattern<vector::BroadcastOp> {
auto vectorType = VectorType::get(distShape, elementType);

VectorValue srcVector = dyn_cast<VectorValue>(broadcastOp.getSource());
// If the srcVector is a scalar (like f32) or a rank-0 vector (like
// vector<f32>), we proceed with the scalar distribution branch.
if (!srcVector || !isNonZeroRank(srcVector)) {
// If the srcVector is a scalar (like f32) we proceed with the scalar
// distribution branch.
if (!srcVector) {
// The way distribution currently works, there is no partial thread
// distribution, so a scalar is available to all threads. Scalar
// distribution is simply a broadcast from scalar to the distributed
Expand Down
18 changes: 8 additions & 10 deletions compiler/src/iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,16 +132,14 @@ void DistributionPattern::replaceOpWithDistributedValues(
for (auto [opResult, replacement] :
llvm::zip_equal(op->getOpResults(), values)) {
// If this value is a vector type, it must be converted back to simd.
if (auto replacementType = dyn_cast<VectorType>(replacement.getType())) {
if (replacementType.getRank() != 0) {
auto oldResult = cast<VectorValue>(opResult);
// Create a toSIMD op to convert the value back to the simd.
rewriter.setInsertionPointAfterValue(oldResult);
Value toSIMD = rewriter.create<IREE::VectorExt::ToSIMDOp>(
oldResult.getLoc(), oldResult.getType(), replacement);
// Add to replacements.
replacement = toSIMD;
}
if (isa<VectorType>(replacement.getType())) {
auto oldResult = cast<VectorValue>(opResult);
// Create a toSIMD op to convert the value back to the simd.
rewriter.setInsertionPointAfterValue(oldResult);
Value toSIMD = rewriter.create<IREE::VectorExt::ToSIMDOp>(
oldResult.getLoc(), oldResult.getType(), replacement);
// Add to replacements.
replacement = toSIMD;
}
replacements.push_back(replacement);
}
Expand Down
41 changes: 41 additions & 0 deletions ...iler/src/iree/compiler/Codegen/Common/GPU/test/gpu_nested_layout_vector_distribution.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -783,6 +783,47 @@ builtin.module attributes { transform.with_named_sequence } {

// -----

#layout = #iree_vector_ext.nested_layout<
subgroup_tile = [2, 2, 2],
batch_tile = [2, 2, 1],
outer_tile = [2, 1, 1],
thread_tile = [4, 16, 8],
element_tile = [1, 4, 4],
subgroup_strides = [4, 2, 1],
thread_strides = [128, 8, 1]
>

func.func @zero_rank_broadcast(%src: vector<f16>) -> (vector<32x256x64xf16>) {
%bcast = vector.broadcast %src : vector<f16> to vector<32x256x64xf16>
%bcastl = iree_vector_ext.to_layout %bcast to layout(#layout) : vector<32x256x64xf16>
return %bcastl : vector<32x256x64xf16>
}

builtin.module attributes { transform.with_named_sequence } {
transform.named_sequence @__transform_main(%variant_op: !transform.any_op {transform.readonly}) {
%top_level_func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op
transform.iree.test_gpu_vector_distribution %top_level_func : !transform.any_op
transform.yield
}
}

// CHECK-LABEL: func @zero_rank_broadcast
// CHECK-SAME: (%[[SRC:.*]]: vector<f16>)
// CHECK: %[[SRC_SIMT:.*]] = iree_vector_ext.to_simt %[[SRC]] : vector<f16>
// CHECK: %[[EXTRACT:.*]] = vector.extract %[[SRC_SIMT]]
// CHECK: %[[BCAST:.*]] = vector.broadcast %[[EXTRACT]] : f16 to vector<1x4x4xf16>
// CHECK: vector.insert %[[BCAST]], %{{.*}}
// CHECK: vector.insert %[[BCAST]], %{{.*}}
// CHECK: vector.insert %[[BCAST]], %{{.*}}
// CHECK: vector.insert %[[BCAST]], %{{.*}}
// CHECK: vector.insert %[[BCAST]], %{{.*}}
// CHECK: vector.insert %[[BCAST]], %{{.*}}
// CHECK: vector.insert %[[BCAST]], %{{.*}}
// CHECK: %[[OUT:.*]] = vector.insert %[[BCAST]], %{{.*}}
// CHECK: iree_vector_ext.to_simd %[[OUT]] : vector<2x2x1x2x1x1x1x4x4xf16> -> vector<32x256x64xf16>

// -----

#layout = #iree_vector_ext.nested_layout<
subgroup_tile = [2, 2, 2],
batch_tile = [2, 2, 1],
Expand Down
2 changes: 1 addition & 1 deletion third_party/llvm-project
Open in desktop
Submodule llvm-project updated 4 files
+0 −21 mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
+2 −2 mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
+10 −10 mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir
+2 −2 mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
Loading