Alias allocations with same size elements despite different dtypes #665
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Assuming it is safe to do this, I am a bit surprised NVRTC doesn't find this opportunity by itself. |
|
!build |
|
Are you saying my claim in #221 about "reusing registers are useless" is wrong? This is 🤯. |
I think so! I was planning to try this for smem tensors then noticed it has an effect for registers. Maybe I am missing some reason this is unsafe so it is not done automatically. |
zasdfgbnm
approved these changes
Aug 1, 2023
jacobhinkle
added a commit
that referenced
this pull request
Oct 2, 2024
See #2934 (comment) PR #665 allowed us to re-use allocations that have different dtypes. We already check that our aliased tensors do not have vectorized accesses larger than those of the original tensors. However, when we have different dtypes we `reinterpret_cast` it to a different `Array` type. Previously we did not specify any alignment in that type's template args, meaning it assumed an alignment of size 1. Since the actual addresses will all still be aligned this does not caused misaligned accesses at runtime. This PR sets the template arg for alignment to be that of the vectorized access width for the alias tensor, so that the compiler could hypothetically do some optimizations knowing the address is aligned.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Currently the
reuseMemoryAllocationslowering pass reusesTensorViews if they have compatible shapes and vectorization, and the same dtype. This PR replaces the dtype condition with a check that the size of each element is the same.It is somewhat rare in our test suite that we find opportunities to re-use one buffer with another of a different type, but it does happen. For example, one case in
FusionWelfordShmoo_CUDAcorresponding totestWelford(DataType::ComplexFloat, 1, 320, 256);currently shows:With the current PR, this becomes
since we re-use two register
TensorViews that change fromint64_ttostd::complex<float>.Re-using shared memory in these cases is likely always beneficial, but the effect of explicitly re-using registers is hard to analyze since the compiler pipeline is so complicated. It is likely rarely harmful, so this behavior is on by default. However, it can be disabled using
DisableOption::ReuseMismatchedTypeRegistersorNVFUSER_DISABLE=reuse_mismatched_type_registers.