Skip to content

Flatten Computation classes#5230

Merged
will-cromar merged 12 commits into
masterfrom
wcromar/flatten-computation
Sep 20, 2023
Merged

Flatten Computation classes#5230
will-cromar merged 12 commits into
masterfrom
wcromar/flatten-computation

Conversation

@will-cromar
Copy link
Copy Markdown
Collaborator

@will-cromar will-cromar commented Jun 22, 2023

  • Remove the wrapper torch_xla::Computation class now that we can mix PyTorch and OpenXLA in the same build.
  • Remove unused fields and make them explicitly unimplemented.
  • Always set name and hash, even if name is unused.
  • Update comments from computation.h (please make sure my understanding is correct).

To be honest, I still don't think the difference between the three uses of Computation is clear, so I'll likely follow up with more refactoring to make the differences explicit.

@will-cromar will-cromar force-pushed the wcromar/flatten-computation branch from 2f1c52f to 20d4061 Compare September 11, 2023 18:20
@will-cromar will-cromar changed the title [WIP] Flatten Computation classes Flatten Computation classes Sep 13, 2023
@will-cromar will-cromar marked this pull request as ready for review September 13, 2023 21:29
@JackCaoG
Copy link
Copy Markdown
Collaborator

Let me take a look now..

Comment on lines -83 to -85
: computation_(std::move(computation)), devices_(std::move(devices)) {
program_shape_ = ConsumeValue(computation_.GetProgramShape());
}
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

any reason program_shape_ is only being inited on one of the constructors?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

All of the other constructors delegate to the one that sets program_shape_ and hash_.

https://secure.wikimedia.org/wikipedia/en/wiki/C++11#Object_construction_improvement

}

const std::vector<std::string>& parameter_names() const override {
return program_shape().parameter_names();
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

seems like we can't alwasy reply on that program_shape exists

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

All of the constructors initialize program shape, as long as its available from the XlaComputation. See my other comment.

Comment thread torch_xla/csrc/runtime/computation_client.h
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG left a comment

Choose a reason for hiding this comment

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

mostly lgtm

Copy link
Copy Markdown
Collaborator

@alanwaketan alanwaketan left a comment

Choose a reason for hiding this comment

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

This makes our code a lot cleaner!

@will-cromar will-cromar force-pushed the wcromar/flatten-computation branch from 2ab0297 to 661896e Compare September 15, 2023 23:18
@will-cromar will-cromar merged commit caa96de into master Sep 20, 2023
jonb377 added a commit that referenced this pull request Sep 22, 2023
jonb377 added a commit that referenced this pull request Sep 22, 2023
zpcore pushed a commit that referenced this pull request Sep 28, 2023
* Flatten `Computation` classes

* Remove compat_logging hack

* Remove wrap/unwrap

* Remove alias

* Remove unused fields

* Don't set hash when device is passed?

* Fix test build issue

* Move hash initialization

* formatting

* Consolidate constructors and add comments back

* formatting

* Update comments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants