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

Iterative horizontal fusion. #48706

Merged
merged 3 commits into from
Apr 23, 2021
Merged

Iterative horizontal fusion. #48706

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a563c29
Iterative horizontal fusion.
trentlo Sep 16, 2020
bd6d2c8
[XLA/GPU] Cleanup and refactory to address review comments.
trentlo Apr 22, 2021
0e06048
[XLA/GPU] remove const_cast in GetOutputsOfFusible().
trentlo Apr 23, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
2 changes: 1 addition & 1 deletion tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,7 +352,7 @@ Status GpuCompiler::OptimizeHloModule(
fusion.AddPass<HloDCE>();
TF_RETURN_IF_ERROR(fusion.Run(hlo_module).status());

HloPassPipeline horizontal_fusion("horizontal_fusion");
HloPassFix<HloPassPipeline> horizontal_fusion("horizontal_fusion");
horizontal_fusion.AddPass<GpuHorizontalLoopFusion>();
horizontal_fusion.AddPass<GpuHorizontalInputFusion>();
horizontal_fusion.AddPass<HloCSE>(/*is_layout_sensitive=*/true,
Expand Down
30 changes: 30 additions & 0 deletions tensorflow/compiler/xla/service/gpu/gpu_fusible.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -541,5 +541,35 @@ bool IsConsumerTheOnlyNonRootUser(const HloInstruction& instr,
});
}

size_t GetInstrCountOfFusible(const HloInstruction& instr) {
if (instr.opcode() != HloOpcode::kFusion) {
return 1;
} else {
return instr.fused_instruction_count();
}
}

absl::InlinedVector<HloInstruction*, 2> GetOutputsOfFusible(
const HloInstruction& instr) {
if (instr.opcode() != HloOpcode::kFusion) {
return {const_cast<HloInstruction*>(&instr)};
Copy link
Member

Choose a reason for hiding this comment

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

Must we const-cast? Why not return a vector of const pointers?

Copy link
Contributor Author

@trentlo trentlo Apr 23, 2021 •
edited

Choose a reason for hiding this comment

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

The return type has to align with the return type of HloInstruction::operands(), which is of absl::InlinedVector<HloInstruction*, 2>. So, returning absl::InlinedVector<const HloInstruction*, 2> will require new allocation of the vector. const_cast may be a bit less evil.

Let me know if you have any other ideas?

Copy link
Member

Choose a reason for hiding this comment

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

const_cast seems to be technically UB if modifying methods are then called, right? Allocating 2 pointers on a stack (new vector) should be essentially free, right?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Just did some search. It is undefined behavior if the object itself (not the pointer) is indeed const. Then, changing the object is undefined behavior through const_cast.

My previous statement was wrong though--we allocate a new vector (in the stack) anyway. So, let's remove the const_cast. Please help to take a look of the latest commit.

}

HloInstruction* root = instr.fused_expression_root();
if (root->opcode() != HloOpcode::kTuple) {
return {root};
} else {
return root->operands();
}
}

size_t GetOutputSizeOfFusible(const HloInstruction& instr) {
if (!instr.IsMultiOutputFusion()) {
return 1;
}
const HloInstruction* root = instr.fused_expression_root();
return ShapeUtil::TupleElementCount(root->shape());
}

} // namespace gpu
} // namespace xla
11 changes: 11 additions & 0 deletions tensorflow/compiler/xla/service/gpu/gpu_fusible.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,17 @@ HloInstruction::FusionKind ChooseFusionKind(const HloInstruction& producer,
bool IsConsumerTheOnlyNonRootUser(const HloInstruction& instr,
const HloInstruction& consumer);

// Returns number of instructions in the fusible `instr`. If `instr` is not a
// fusion instruction, 1 is returned.
size_t GetInstrCountOfFusible(const HloInstruction& instr);

// Returns the outputs of the fusible `instr`.
absl::InlinedVector<HloInstruction*, 2> GetOutputsOfFusible(
const HloInstruction& instr);

// Returns the output size of the fusible `instr`.
size_t GetOutputSizeOfFusible(const HloInstruction& instr);

} // namespace gpu
} // namespace xla

Expand Down
25 changes: 18 additions & 7 deletions tensorflow/compiler/xla/service/gpu/horizontal_input_fusion.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,12 +78,12 @@ std::vector<HloInstruction*> FindAndSortFusionCandidates(
HloInstruction* consumer) {
absl::flat_hash_set<HloInstruction*> fusion_instr_set;
std::vector<HloInstruction*> fusion_instrs;
for (auto opnd : consumer->operands()) {
for (HloInstruction* opnd : consumer->operands()) {
HloInstruction* predecessor = opnd->LatestNonGteAncestor();
// Find out the input fusion instructions whose only consumer is `consumer`.
// This guarantees that fusing these candidates will never create cycles, as
// there is no back edge.
if (IsReduceInputFusion(*predecessor) &&
if (IsInputFusibleReduction(*predecessor) &&
IsConsumerTheOnlyNonRootUser(*predecessor, *consumer)) {
if (fusion_instr_set.insert(predecessor).second) {
fusion_instrs.push_back(predecessor);
Expand All @@ -102,8 +102,7 @@ std::vector<HloInstruction*> FindAndSortFusionCandidates(
}
// Sort `fusion_instrs` according to instruction counts, because
// we'd like to fuse together computations of similar sizes.
return a->fused_instruction_count() <
b->fused_instruction_count();
return GetInstrCountOfFusible(*a) < GetInstrCountOfFusible(*b);
});

return fusion_instrs;
Expand All @@ -116,12 +115,24 @@ StatusOr<bool> HorizontalInputFusionImpl::Run() {
// Using def-to-use order is sound since we do not modify users.
std::vector<HloInstruction*> def_to_use_order =
computation_->MakeInstructionPostOrder();
for (auto consumer : def_to_use_order) {
for (HloInstruction* consumer : def_to_use_order) {
auto candidates = FindAndSortFusionCandidates(consumer);
if (candidates.empty()) {
if (candidates.size() <= 1) {
continue;
}

// Convert candidates into fusions if needed.
for (size_t j = 0; j < candidates.size(); ++j) {
if (candidates[j]->opcode() != HloOpcode::kFusion) {
TF_ASSIGN_OR_RETURN(
HloInstruction * fusion_instr,
MakeFusionInstruction(candidates[j],
HloInstruction::FusionKind::kInput));
candidates[j] = fusion_instr;
changed = true;
}
}

size_t fusion_anchor_id = 0;
for (size_t j = 1; j < candidates.size(); ++j) {
HloInstruction* fusion_anchor = candidates[fusion_anchor_id];
Expand Down Expand Up @@ -155,7 +166,7 @@ StatusOr<bool> GpuHorizontalInputFusion::RunOnComputation(
StatusOr<bool> GpuHorizontalInputFusion::Run(HloModule* module) {
bool changed = false;
VLOG(2) << "Run horizontal input fusion.";
for (auto* comp : module->MakeNonfusionComputations()) {
for (HloComputation* comp : module->MakeNonfusionComputations()) {
TF_ASSIGN_OR_RETURN(changed, RunOnComputation(comp));
}

Expand Down
33 changes: 33 additions & 0 deletions tensorflow/compiler/xla/service/gpu/horizontal_input_fusion_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,39 @@ TEST_F(HorizontalInputFusionTest, MultiOutputFusionTest) {
EXPECT_TRUE(GpuHorizontalInputFusion().Run(module.get()).ValueOrDie());
}

TEST_F(HorizontalInputFusionTest, NonfusionInstrs) {
auto module = ParseAndReturnVerifiedModule(R"(
HloModule NonfusionInstrs

%add_f16 {
%x = f16[] parameter(0)
%y = f16[] parameter(1)
ROOT %add = f16[] add(%x, %y)
}

ENTRY entry_computation {
arg.0 = f16[1024]{0} parameter(0)
arg.1 = f16[1024]{0} parameter(1)
constant0 = f16[] constant(0)
reduce.0 = f16[] reduce(arg.0, constant0), dimensions={0}, to_apply=%add_f16
reduce.1 = f16[] reduce(arg.1, constant0), dimensions={0}, to_apply=%add_f16
ROOT tuple.0 = (f16[], f16[]) tuple(reduce.0, reduce.1)
}
)").ValueOrDie();

EXPECT_TRUE(GpuHorizontalInputFusion().Run(module.get()).ValueOrDie());

const HloInstruction* entry_root =
module->entry_computation()->root_instruction();
EXPECT_THAT(entry_root, op::Tuple((op::GetTupleElement(op::Fusion())),
(op::GetTupleElement(op::Fusion()))));

const HloInstruction* fusion = entry_root->operand(0)->operand(0);
ASSERT_TRUE(fusion->IsMultiOutputFusion());
EXPECT_THAT(fusion->fused_expression_root(),
op::Tuple(op::Reduce(), op::Reduce()));
}

} // namespace
} // namespace gpu
} // namespace xla