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

TFactor: bulkerify GEMV (both MC and GPU) #1214

Merged
merged 60 commits into from
Feb 12, 2025
Merged

Conversation

albestro
Copy link
Collaborator

@albestro albestro commented Nov 15, 2024

In this PR we try to increase parallelisation of GEMV step of TFactor, which is by construction serial (all results goes into the same block tile), by using workspaces for storing partial results and then reducing them before the final TRMV step.

Algorithmically, the main change is that stepGEMV loop has been replaced with a single stepGEMVAll, and the concept has been applied in a similar way for both backends, MC and GPU:

  • MC: pika::bulk splits input tiles over multiple task, each one stores the partial results in their workspace and just one at the end does the reduction.
  • GPU: similarly to the CPU, the work is forked over different pika tasks, each one computing a partial result on a different GPU stream. These tasks then join into a single task which performs the reduction.

In order to implement this solution, workspaces for intermediate results have been added (there is another option which does not require any additional workspace nor reduction, and that we might explore for MC in another PR).

TODO:

Close #798.

EDIT: since this is conceptually similar to #798 and it is going to be closed as soon as this gets merged, I migrated here the doc fixes happened there.

@albestro albestro added this to the Optimizations milestone Nov 15, 2024
@albestro albestro self-assigned this Nov 15, 2024
@albestro albestro force-pushed the alby/tfactor-optim/bulk branch from 6558780 to bf17fcc Compare November 18, 2024 13:16
Base automatically changed from alby/tfactor-optim/no-gemv-divergence to master November 22, 2024 11:28
@albestro albestro force-pushed the alby/tfactor-optim/bulk branch 4 times, most recently from 37ea75b to 2845339 Compare December 2, 2024 16:48
@albestro
Copy link
Collaborator Author

albestro commented Dec 2, 2024

cscs-ci run

just to check if there is any major problems

@albestro albestro marked this pull request as ready for review December 3, 2024 16:08
Copy link
Collaborator

@rasolca rasolca left a comment

Choose a reason for hiding this comment

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

LGTM.
Please name new functions in snake case, and rename existing internal functions if possible.

Comment on lines 268 to 269
matrix::Matrix<T, Device::GPU> ws_T({nworkspaces * nrefls_step, nrefls_step},
{nrefls_step, nrefls_step});
Copy link
Collaborator

Choose a reason for hiding this comment

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

All Ts are allocated at scheduling. Better reuse it.

@albestro albestro requested a review from rasolca December 5, 2024 11:22
Copy link
Collaborator

@msimberg msimberg left a comment

Choose a reason for hiding this comment

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

Only a few questions, nothing blocking.

@albestro albestro force-pushed the alby/tfactor-optim/bulk branch from 05dc48d to f83f317 Compare December 5, 2024 15:00
@albestro albestro force-pushed the alby/tfactor-optim/bulk branch from 0397c9e to d370894 Compare December 9, 2024 09:47
@albestro albestro requested a review from rasolca December 9, 2024 09:47
@albestro
Copy link
Collaborator Author

cscs-ci run

@albestro albestro marked this pull request as draft December 11, 2024 16:15
@albestro
Copy link
Collaborator Author

Converted to draft in order to prevent merging, since I'm still doing some checks.

@albestro albestro force-pushed the alby/tfactor-optim/bulk branch 2 times, most recently from 3749abc to ffba9a4 Compare February 5, 2025 09:28
@@ -1063,7 +1075,14 @@ Matrix<T, Device::CPU> ReductionToBand<B, D, T>::call(Matrix<T, D>& mat_a, const
// TODO probably the first one in any panel is ok?
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Random thought: just saw this comment and wondering if we should have a look and see if we can "merge" this workspace into the other one? @rasolca

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 will leave it for another PR 😉

@albestro albestro force-pushed the alby/tfactor-optim/bulk branch 4 times, most recently from aaba87f to 0d2def8 Compare February 10, 2025 17:13
@albestro albestro force-pushed the alby/tfactor-optim/bulk branch from 0d2def8 to c306078 Compare February 10, 2025 17:19
@albestro albestro requested review from msimberg and rasolca February 11, 2025 10:02
@albestro
Copy link
Collaborator Author

cscs-ci run

Copy link
Collaborator

@msimberg msimberg left a comment

Choose a reason for hiding this comment

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

Only a minor include-nit, otherwise this looks good to me.

@albestro make sure to merge latest master since it's got the new CI changes now.

@albestro
Copy link
Collaborator Author

cscs-ci run

@albestro
Copy link
Collaborator Author

cscs-ci run

@rasolca rasolca merged commit c5c5280 into master Feb 12, 2025
5 checks passed
@rasolca rasolca deleted the alby/tfactor-optim/bulk branch February 12, 2025 14:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
Archived in project
Development

Successfully merging this pull request may close these issues.

4 participants