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.

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

Convert AbstractTensor from struct to class #2967

Merged
merged 3 commits into from
Sep 19, 2024

Conversation

jacobhinkle
Copy link
Collaborator

@jacobhinkle jacobhinkle commented Sep 19, 2024

AbstractTensor now has many methods and as of #2933 it is subclassed into TaggedAbstractTensor which has additional methods. Also, there is no longer just a single vector attribute domain but also info. Instead of requiring the user to manage those vectors and keep them the same length, this PR makes them protected members of a class and uses the existing accessors to adjust them, as suggested in #2963 (comment).

`AbstractTensor` now has many methods and as of #2933 it is subclassed
into `TaggedAbstractTensor` which has additional methods. Also, there is
no longer just a single vector attribute `domain` but also `info`.
Instead of requiring the user to manage those vectors and keep them the
same length, instead this PR makes them `protected` and uses the
existing accessors to adjust them, as suggested in
#2963 (comment).
@jacobhinkle
Copy link
Collaborator Author

!build

@jacobhinkle jacobhinkle marked this pull request as ready for review September 19, 2024 15:07
@zasdfgbnm
Copy link
Collaborator

Is the build failure related?

@jacobhinkle
Copy link
Collaborator Author

Is the build failure related?

I am not sure yet, but I doubt it. I am also seeing this failure on #2913 which doesn't include this change.

Comment on lines +1012 to +1014
std::unordered_set<Tag>& getTags(int64_t i) {
i = wrapDim(i, (int64_t)this->size());
return this->info_[i].tags;
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

I added this non-const accessor so that we can manually insert and erase tags for a single axis.

@zasdfgbnm
Copy link
Collaborator

I am surprised that this PR just work, because I remembered that there are more direct uses of .domain and some requires none-trivial change.

@jacobhinkle
Copy link
Collaborator Author

I am surprised that this PR just work, because I remembered that there are more direct uses of .domain and some requires none-trivial change.

Yeah I thought it was going to be tougher too but the only places AbstractTensor is currently used is in matmul utilities and TMA analysis:

  • MmaSwizzler::scheduleMmaOutputAllocation
  • swizzleSharedMemory
  • collapse_tma_domain::DomainMerger

Only DomainMerger was tinkering with .domain directly. Otherwise, we do make use of AbstractTensor in some tests outside of test_abstract_tensor.cpp but those are not directly manipulating .domain either.

Copy link
Collaborator

@zasdfgbnm zasdfgbnm left a comment

Choose a reason for hiding this comment

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

LGTM, just want to see CI green for double confirming.

@jacobhinkle
Copy link
Collaborator Author

!build

@jacobhinkle
Copy link
Collaborator Author

Rebuilding following merging of #2969

@jacobhinkle jacobhinkle merged commit f2efe2b into main Sep 19, 2024
15 of 16 checks passed
@jacobhinkle jacobhinkle deleted the abstract_tensor_convert_to_class branch September 19, 2024 20:20
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.

2 participants