[XLA] Pass the module to HloDataflowAnalysis by const reference.

PiperOrigin-RevId: 186072673
This commit is contained in:
Michael Kuperstein 2018-02-16 18:13:53 -08:00 committed by TensorFlower Gardener
parent a189502cc3
commit 090bb9168c
9 changed files with 22 additions and 23 deletions

View File

@ -1156,7 +1156,7 @@ bool IsWhileBody(const HloComputation* computation,
HloModule* module) { HloModule* module) {
std::unique_ptr<CallGraph> call_graph = CallGraph::Build(module); std::unique_ptr<CallGraph> call_graph = CallGraph::Build(module);
TF_ASSIGN_OR_RETURN(std::unique_ptr<HloDataflowAnalysis> dataflow, TF_ASSIGN_OR_RETURN(std::unique_ptr<HloDataflowAnalysis> dataflow,
HloDataflowAnalysis::Run(module)); HloDataflowAnalysis::Run(*module));
bool changed = false; bool changed = false;

View File

@ -49,7 +49,7 @@ StatusOr<bool> GpuCopyInsertion::Run(HloModule* module) {
TF_ASSIGN_OR_RETURN(bool changed, generic_copy_insertion.Run(module)); TF_ASSIGN_OR_RETURN(bool changed, generic_copy_insertion.Run(module));
TF_ASSIGN_OR_RETURN(std::unique_ptr<HloDataflowAnalysis> dataflow, TF_ASSIGN_OR_RETURN(std::unique_ptr<HloDataflowAnalysis> dataflow,
HloDataflowAnalysis::Run(module)); HloDataflowAnalysis::Run(*module));
// Make sure all operands of a library call are in memory instead of constants // Make sure all operands of a library call are in memory instead of constants
// in IR. // in IR.

View File

@ -419,7 +419,7 @@ StatusOr<std::unique_ptr<HloAliasAnalysis>> HloAliasAnalysis::Run(
auto alias_analysis = WrapUnique(new HloAliasAnalysis(module)); auto alias_analysis = WrapUnique(new HloAliasAnalysis(module));
TF_ASSIGN_OR_RETURN( TF_ASSIGN_OR_RETURN(
alias_analysis->dataflow_analysis_, alias_analysis->dataflow_analysis_,
HloDataflowAnalysis::Run(module, /*ssa_form=*/true, HloDataflowAnalysis::Run(*module, /*ssa_form=*/true,
/*bitcast_defines_value=*/false)); /*bitcast_defines_value=*/false));
BufferValueMap buffer_map(alias_analysis->dataflow_analysis()); BufferValueMap buffer_map(alias_analysis->dataflow_analysis());

View File

@ -38,12 +38,12 @@ namespace xla {
using ::tensorflow::strings::StrAppend; using ::tensorflow::strings::StrAppend;
using ::tensorflow::strings::StrCat; using ::tensorflow::strings::StrCat;
HloDataflowAnalysis::HloDataflowAnalysis(HloModule* module, bool ssa_form, HloDataflowAnalysis::HloDataflowAnalysis(const HloModule& module, bool ssa_form,
bool bitcast_defines_value) bool bitcast_defines_value)
: module_(module), : module_(module),
ssa_form_(ssa_form), ssa_form_(ssa_form),
bitcast_defines_value_(bitcast_defines_value), bitcast_defines_value_(bitcast_defines_value),
call_graph_(CallGraph::Build(module)) {} call_graph_(CallGraph::Build(&module)) {}
bool HloDataflowAnalysis::ValueIsDefinedAt(const HloInstruction* instruction, bool HloDataflowAnalysis::ValueIsDefinedAt(const HloInstruction* instruction,
const ShapeIndex& index) const { const ShapeIndex& index) const {
@ -115,9 +115,9 @@ void HloDataflowAnalysis::DeleteMarkedValues() {
} }
string HloDataflowAnalysis::ToString() const { string HloDataflowAnalysis::ToString() const {
string out = StrCat("HloDataflowAnalysis, module ", module_->name(), "\n"); string out = StrCat("HloDataflowAnalysis, module ", module_.name(), "\n");
StrAppend(&out, " Instruction value sets:\n"); StrAppend(&out, " Instruction value sets:\n");
for (const HloComputation* computation : module_->computations()) { for (const HloComputation* computation : module_.computations()) {
for (const HloInstruction* instruction : computation->instructions()) { for (const HloInstruction* instruction : computation->instructions()) {
StrAppend(&out, " ", instruction->name(), ":\n"); StrAppend(&out, " ", instruction->name(), ":\n");
if (ShapeUtil::IsTuple(instruction->shape())) { if (ShapeUtil::IsTuple(instruction->shape())) {
@ -592,7 +592,7 @@ void HloDataflowAnalysis::Propagate() {
} }
}; };
for (HloComputation* computation : module_->computations()) { for (HloComputation* computation : module_.computations()) {
for (HloInstruction* instruction : computation->instructions()) { for (HloInstruction* instruction : computation->instructions()) {
add_to_worklist(instruction); add_to_worklist(instruction);
} }
@ -686,7 +686,7 @@ InstructionValueSet& HloDataflowAnalysis::GetInstructionValueSet(
} }
Status HloDataflowAnalysis::InitializeInstructionValueSets() { Status HloDataflowAnalysis::InitializeInstructionValueSets() {
for (const HloComputation* computation : module_->computations()) { for (const HloComputation* computation : module_.computations()) {
const CallGraphNode& call_graph_node = call_graph_->GetNode(computation); const CallGraphNode& call_graph_node = call_graph_->GetNode(computation);
for (HloInstruction* instruction : computation->instructions()) { for (HloInstruction* instruction : computation->instructions()) {
// Create an empty shape tree. // Create an empty shape tree.
@ -787,9 +787,9 @@ Status HloDataflowAnalysis::InitializeInstructionValueSets() {
/* static */ /* static */
StatusOr<std::unique_ptr<HloDataflowAnalysis>> HloDataflowAnalysis::Run( StatusOr<std::unique_ptr<HloDataflowAnalysis>> HloDataflowAnalysis::Run(
HloModule* module, bool ssa_form, bool bitcast_defines_value) { const HloModule& module, bool ssa_form, bool bitcast_defines_value) {
VLOG(1) << "HloDataflowAnalysis::Run on module " << module->name(); VLOG(1) << "HloDataflowAnalysis::Run on module " << module.name();
XLA_VLOG_LINES(2, module->ToString()); XLA_VLOG_LINES(2, module.ToString());
auto dataflow_analysis = WrapUnique( auto dataflow_analysis = WrapUnique(
new HloDataflowAnalysis(module, ssa_form, bitcast_defines_value)); new HloDataflowAnalysis(module, ssa_form, bitcast_defines_value));
@ -806,7 +806,7 @@ StatusOr<std::unique_ptr<HloDataflowAnalysis>> HloDataflowAnalysis::Run(
// lookup is faster. // lookup is faster.
std::vector<std::vector<HloPosition>> value_positions( std::vector<std::vector<HloPosition>> value_positions(
dataflow_analysis->next_value_id_); dataflow_analysis->next_value_id_);
for (const HloComputation* computation : module->computations()) { for (const HloComputation* computation : module.computations()) {
for (HloInstruction* instruction : computation->instructions()) { for (HloInstruction* instruction : computation->instructions()) {
for (const auto& pair : for (const auto& pair :
dataflow_analysis->GetInstructionValueSet(instruction)) { dataflow_analysis->GetInstructionValueSet(instruction)) {
@ -858,7 +858,7 @@ Status HloDataflowAnalysis::Verify() const {
// For each value in each value set, verify that the value set's position // For each value in each value set, verify that the value set's position
// appears in the value's positions(). // appears in the value's positions().
for (const auto& computation : module_->computations()) { for (const auto& computation : module_.computations()) {
for (const auto& instruction : computation->instructions()) { for (const auto& instruction : computation->instructions()) {
for (const auto& pair : GetInstructionValueSet(instruction)) { for (const auto& pair : GetInstructionValueSet(instruction)) {
const ShapeIndex& index = pair.first; const ShapeIndex& index = pair.first;

View File

@ -60,7 +60,7 @@ class HloDataflowAnalysis {
// a new HLO value in the analysis. If false then Bitcast forwards the // a new HLO value in the analysis. If false then Bitcast forwards the
// value of its operand. // value of its operand.
static StatusOr<std::unique_ptr<HloDataflowAnalysis>> Run( static StatusOr<std::unique_ptr<HloDataflowAnalysis>> Run(
HloModule* module, bool ssa_form = false, const HloModule& module, bool ssa_form = false,
bool bitcast_defines_value = false); bool bitcast_defines_value = false);
// Returns true if 'instruction' defines an HLO value at the given shape index // Returns true if 'instruction' defines an HLO value at the given shape index
@ -119,7 +119,7 @@ class HloDataflowAnalysis {
string ToString() const; string ToString() const;
protected: protected:
HloDataflowAnalysis(HloModule* module, bool ssa_form, HloDataflowAnalysis(const HloModule& module, bool ssa_form,
bool bitcast_defines_value = false); bool bitcast_defines_value = false);
// Returns a new HloValue defined at the given instruction and shape index. // Returns a new HloValue defined at the given instruction and shape index.
@ -180,7 +180,7 @@ class HloDataflowAnalysis {
// Verify various invariants of the dataflow analysis. // Verify various invariants of the dataflow analysis.
Status Verify() const; Status Verify() const;
HloModule* const module_; const HloModule& module_;
const bool ssa_form_; const bool ssa_form_;
const bool bitcast_defines_value_; const bool bitcast_defines_value_;

View File

@ -50,7 +50,7 @@ class HloDataflowAnalysisTest : public HloTestBase,
bool bitcast_defines_value = false) { bool bitcast_defines_value = false) {
hlo_graph_dumper::MaybeDumpHloModule(*module_, "Before dataflow analysis"); hlo_graph_dumper::MaybeDumpHloModule(*module_, "Before dataflow analysis");
analysis_ = analysis_ =
HloDataflowAnalysis::Run(module_.get(), ssa_form, bitcast_defines_value) HloDataflowAnalysis::Run(*module_, ssa_form, bitcast_defines_value)
.ConsumeValueOrDie(); .ConsumeValueOrDie();
return *analysis_; return *analysis_;
} }

View File

@ -262,8 +262,8 @@ TEST_F(HloOrderingTest, ValuesInWhileComputations) {
scalar_shape, HloOpcode::kAdd, constant, xla_while)); scalar_shape, HloOpcode::kAdd, constant, xla_while));
module->AddEntryComputation(builder.Build()); module->AddEntryComputation(builder.Build());
TF_ASSERT_OK_AND_ASSIGN( TF_ASSERT_OK_AND_ASSIGN(auto dataflow,
auto dataflow, HloDataflowAnalysis::Run(module.get(), /*ssa_form=*/true)); HloDataflowAnalysis::Run(*module, /*ssa_form=*/true));
DependencyHloOrdering ordering(module.get()); DependencyHloOrdering ordering(module.get());
// Init value is defined before the while, but live range is not before the // Init value is defined before the while, but live range is not before the

View File

@ -35,8 +35,7 @@ class PointsToAnalysisTestBase : public HloTestBase {
CHECK_NOTNULL(module_.get()); CHECK_NOTNULL(module_.get());
points_to_analysis_ = points_to_analysis_ =
TuplePointsToAnalysis::Run(module_.get()).ConsumeValueOrDie(); TuplePointsToAnalysis::Run(module_.get()).ConsumeValueOrDie();
dataflow_analysis_ = dataflow_analysis_ = HloDataflowAnalysis::Run(*module_).ConsumeValueOrDie();
HloDataflowAnalysis::Run(module_.get()).ConsumeValueOrDie();
} }
void BuildModuleAndRunAnalysis(std::unique_ptr<HloComputation> computation) { void BuildModuleAndRunAnalysis(std::unique_ptr<HloComputation> computation) {

View File

@ -287,7 +287,7 @@ StatusOr<std::unique_ptr<Literal>> MakeFakeLiteral(const Shape& shape) {
StatusOr<std::vector<std::unique_ptr<Literal>>> MakeFakeArguments( StatusOr<std::vector<std::unique_ptr<Literal>>> MakeFakeArguments(
HloModule* const module) { HloModule* const module) {
TF_ASSIGN_OR_RETURN(auto dataflow, HloDataflowAnalysis::Run(module)); TF_ASSIGN_OR_RETURN(auto dataflow, HloDataflowAnalysis::Run(*module));
const auto params = module->entry_computation()->parameter_instructions(); const auto params = module->entry_computation()->parameter_instructions();
std::minstd_rand0 engine; std::minstd_rand0 engine;
std::vector<std::unique_ptr<Literal>> arguments(params.size()); std::vector<std::unique_ptr<Literal>> arguments(params.size());