-
Notifications
You must be signed in to change notification settings - Fork 46
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.
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
Make Communications IRs inheriting from Expr. #2185
Conversation
3d79609
to
3e6fddf
Compare
3e6fddf
to
4560e21
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM in general! I made a few comments to address before I merge this.
} else { | ||
assertBufferCount(params_.dst_bufs, 0); | ||
// TODO add checking symbolic representation of src and dst buffers | ||
bool Communication::sameAs(const Statement* other) const { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I failed to see why we couldn't reuse the "default" implementation:
Line 342 in 4e83569
bool Expr::sameAs(const Statement* other) const { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Two Communications could be "Expr::sameAs" without being the same, for example if one is Allgather and the other Allreduce. Right?
My goal here is to mimick what's done for other IRs, such as
Line 2674 in 4e83569
bool IterDomain::sameAs(const Statement* other) const { |
But I don't actually use sameAs -- I simply thought implementing it was necessary for matching the IR specs.
So I'm open to any suggestion about this implementation, including removing it
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Expr::sameAs
calls Expr::sameOp
, which checks equality of all attributes. I believe the communication type is one of the attributes.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Correct me if Im wrong, but in the current implementation it is not an attribute. The reason is that, IIUC, all attributes need to be Statement*
s, and CommunicationParams
is not.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There is addDataAttribute as used in
Lines 2423 to 2424 in 730abb5
addDataAttribute(op_type); | |
addDataAttribute(cache_op); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM assuming that the post and validate logic is just moved around! If there is changes in those let me know so I can look more closely.
Also it is great to see Send/Recv finally removed from communicator!
!build --dist |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM! I'll try to merge this. I may need to split this up into multiple PRs, but will let you know.
I'm in the process of resolving conflicts. I should be able to merge this tomorrow or Monday. |
Thank you! I'm waiting for this one to be merge before going on with host ir dev. Let me know if I can help! |
3ebd7d4
to
0f7fa53
Compare
!build --dist |
This PR makes the "Communication" class a proper IR inheriting from Expr. This patch is needed for implementing Host Irs. It is also one step towards making the Communications (and more generally the multidevice module) fully symbolic.
By the way, we proceed to a couple of refactoring, and remove
Communicator::sendRecv
method.Remarks:
Communication
and one derived class per collective type (Allgather, Allreduce, Broadcast, etc.). Now, there is only the classCommunication
, and the collective type is encoded through an enum classCollectiveType
added to the parameter memberCommParams
Communication::post
method was replaced by a standalone functionpostCollective
. The motivation is to scoop out the runtime execution from the symbolic representation of the collective.Collective
is instantiated with concrete device Idx and concrete at::Tensor, while it should be instantiated with symbolic representations and binded to actual device indices and Aten buffers at runtime. This will be added in a future PR.CI: https://gitlab-master.nvidia.com/dl/pytorch/fuser-gh-mirror/-/pipelines/14923305