Skip to content

Fix memory fence for ARM - #753

Closed
aurianer wants to merge 4 commits into
NVIDIA:mainfrom
aurianer:fix_barrier_for_aarch64
Closed

Fix memory fence for ARM#753
aurianer wants to merge 4 commits into
NVIDIA:mainfrom
aurianer:fix_barrier_for_aarch64

Conversation

@aurianer

@aurianeraurianer commented Apr 5, 2024

Copy link
Copy Markdown
Contributor

TransformerEngine doesn't compile for aarch64, this PR:

  • Conditionally use data memory barrier for ARM architecture.
  • Fix a redefinition error coming from gdrapi.h defining the GPU_PAGE_* macros and already included in userbuffers.h. This adds an explicit include (which is guarded and prevent redefinition) since the GPU_PAGE_* variables are used in this .cpp file.

@aurianer
aurianerforce-pushed the fix_barrier_for_aarch64 branch from 5882c94 to 842dc17CompareApril 5, 2024 14:44
timmoon10
timmoon10 approved these changes Apr 5, 2024
View reviewed changes

@timmoon10timmoon10 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Overall LGTM

Comment threadtransformer_engine/pytorch/csrc/userbuffers/userbuffers-host.cpp Outdated
Comment threadtransformer_engine/pytorch/csrc/userbuffers/userbuffers-host.cpp Outdated
@timmoon10

timmoon10 commented Apr 5, 2024

Copy link
Copy Markdown
Member

/te-ci pytorch

@aurianer
aurianerforce-pushed the fix_barrier_for_aarch64 branch from 8a06c92 to df72d04CompareApril 8, 2024 13:37
@aurianer

aurianer commented Apr 8, 2024

Copy link
Copy Markdown
ContributorAuthor

Out of curiosity, why is the sleep(1) needed? Would it make sense to use a memory barrier in assembly using the nop instruction (longer than _mm_mfence I believe)

henrique added a commit to henrique/TransformerEngine that referenced this pull request Apr 10, 2024
@henrique
aurianer: Fix memory fence for ARMNVIDIA#753
8cc3859
henrique
henrique reviewed Apr 10, 2024
View reviewed changes
Comment threadtransformer_engine/pytorch/csrc/userbuffers/userbuffers-host.cpp Outdated
aurianerand others added 3 commits April 10, 2024 13:08
@aurianer
Use another barrier for ARM
4214dd0
Signed-off-by: aurianer <aurianer@cscs.ch>
@aurianer
Fix redefinition of GPU variables already defined in gdrapi
68a43d5
Signed-off-by: aurianer <aurianer@cscs.ch>
@timmoon10@aurianer
Update transformer_engine/pytorch/csrc/userbuffers/userbuffers-host.cpp
3a382af
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
@aurianer
aurianerforce-pushed the fix_barrier_for_aarch64 branch from df72d04 to cf93f1bCompareApril 10, 2024 11:09
@aurianer
Add msvc variant for the barrier
a7d5357
Signed-off-by: aurianer <aurianer@cscs.ch>
@aurianer
aurianerforce-pushed the fix_barrier_for_aarch64 branch from cf93f1b to a7d5357CompareApril 10, 2024 16:52
@timmoon10

timmoon10 commented Apr 10, 2024

Copy link
Copy Markdown
Member

/te-ci pytorch

shamisp
shamisp reviewed Apr 10, 2024
View reviewed changes
Comment threadtransformer_engine/pytorch/csrc/userbuffers/userbuffers-host.cpp
@aurianer

aurianer commented Apr 12, 2024

Copy link
Copy Markdown
ContributorAuthor

closing in favour of #757

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
+2 more reviewers
@henriquehenriquehenrique left review comments
@shamispshamispshamisp left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@aurianer@timmoon10@henrique@shamisp