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

Jump to bottom

Refine colreduction fusion strategy of kStitch. #1257

Merged
merged 6 commits into from
Oct 9, 2023
Merged

Refine colreduction fusion strategy of kStitch. #1257

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
496147d
Refine colreduction fusion strategy of kStitch.
yunzhongOvO Sep 19, 2023
52cca1e
clang-format
yunzhongOvO Sep 19, 2023
ca6732c
delete value range
yunzhongOvO Sep 20, 2023
d5b67b7
delete outdated tests
yunzhongOvO Sep 21, 2023
986774b
update
yunzhongOvO Sep 25, 2023
8615285
refine fusion rule.
yunzhongOvO Oct 8, 2023
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
27 changes: 0 additions & 27 deletions .github/workflows/pytorch_pre_cpu.yml
View file
Open in desktop

This file was deleted.

28 changes: 0 additions & 28 deletions .github/workflows/pytorch_pre_gpu.yml
View file
Open in desktop

This file was deleted.

28 changes: 20 additions & 8 deletions tao_compiler/mlir/disc/transforms/fusion_utils.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1475,16 +1475,28 @@ bool BaseGpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis,
FusionPattern& lhs, FusionPattern& rhs,
FusionPattern& target) {
// TODO(Yancey): support fusion with different reduction type
bool has_row_reduction = llvm::any_of(target.getOpList(), [](Operation* op) {
return isRank2RowReduction(op);
});
bool has_col_reduciton = llvm::any_of(target.getOpList(), [](Operation* op) {
return isRank2ColReduction(op);
});

if (has_row_reduction && has_col_reduciton) {
bool has_rank2_row_reduction =
llvm::any_of(target.getOpList(),
[](Operation* op) { return isRank2RowReduction(op); });
bool has_rank2_col_reduction =
llvm::any_of(target.getOpList(),
[](Operation* op) { return isRank2ColReduction(op); });

if (has_rank2_row_reduction && has_rank2_col_reduction) {
return false;
}

if (has_rank2_col_reduction) {
const auto& results = target.getResults();
auto ref_shape = getEffectiveShape(target, results[0]);
if (llvm::any_of(results, [&](Value result) {
auto op = target.findLastWriter(result);
return isa<lmhlo::TransposeOp>(op);
})) {
return false;
}
}

return BaseFusionStrategy::tryFuse(shapeAnalysis, lhs, rhs, target);
}

Expand Down
33 changes: 22 additions & 11 deletions tao_compiler/mlir/disc/transforms/fusion_utils_stitch_gpu.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,16 +59,28 @@ bool StitchGpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis,
FusionPattern& lhs, FusionPattern& rhs,
FusionPattern& target) {
// TODO(Yancey): support fusion with different reduction type
bool has_row_reduction = llvm::any_of(target.getOpList(), [](Operation* op) {
return isRank2RowReduction(op);
});
bool has_col_reduciton = llvm::any_of(target.getOpList(), [](Operation* op) {
return isRank2ColReduction(op);
});

if (has_row_reduction && has_col_reduciton) {
bool has_rank2_row_reduction =
llvm::any_of(target.getOpList(),
[](Operation* op) { return isRank2RowReduction(op); });
bool has_rank2_col_reduction =
llvm::any_of(target.getOpList(),
[](Operation* op) { return isRank2ColReduction(op); });

if (has_rank2_row_reduction && has_rank2_col_reduction) {
return false;
}

if (has_rank2_col_reduction) {
const auto& results = target.getResults();
auto ref_shape = getEffectiveShape(target, results[0]);
if (llvm::any_of(results, [&](Value result) {
auto op = target.findLastWriter(result);
return isa<lmhlo::TransposeOp>(op);
})) {
return false;
}
}

return FusionStrategy::tryFuse(shapeAnalysis, lhs, rhs, target);
}

Expand Down Expand Up @@ -428,9 +440,8 @@ bool StitchGpuFusionStrategy::findFusionPatternTypeAndSubroot(
return true;
}
Value shape = getEffectiveShape(fusion_pattern, result);
return isRank2ColReduction(op)
? shapeAnalysis.isShapeEqual(ref_shape, shape)
: shapeAnalysis.isSameNumElements(ref_shape, shape);
return isRank2ColReduction(op) &&
shapeAnalysis.isShapeEqual(ref_shape, shape);
})) {
return false;
}
Expand Down
4 changes: 2 additions & 2 deletions tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -650,7 +650,7 @@ LogicalResult lowerWithScheduleColReduction(

SmallVector<Value, 4> yield_values_for_if;

ValueRange load_index({row_index, col_index});
SmallVector<Value, 2> load_index({row_index, col_index});
b.setInsertionPointToStart(&if_row_valid_op.getThenRegion().front());
int col_red_root_op_idx = 0;
for (auto* root_op : root_ops) {
Expand Down Expand Up @@ -924,7 +924,7 @@ LogicalResult lowerWithScheduleColReductionTileH(

SmallVector<Value, 4> yield_values_for_if;

ValueRange load_index({row_index, col_index});
SmallVector<Value, 2> load_index({row_index, col_index});
b.setInsertionPointToStart(&if_row_valid_op.getThenRegion().front());
int col_red_root_op_idx = 0;
for (auto* root_op : root_ops) {
Expand Down
Loading