Skip to content

[Unity] Add Relax multi-device e2e cases #15823

Merged
yongwww merged 8 commits intoapache:unityfrom
yongwww:e2e_multi_device
Dec 20, 2023
Merged

[Unity] Add Relax multi-device e2e cases #15823
yongwww merged 8 commits intoapache:unityfrom
yongwww:e2e_multi_device

Conversation

@yongwww
Copy link
Copy Markdown
Member

@yongwww yongwww commented Sep 26, 2023

This Pull Request enables Relax multi-device end-to-end test cases, including multi-CPU, multi-GPU, and heterogeneous CPU-GPU cases.

  • Add the kTarget attribute to prim_func during LegalizeOps if the related vdevice is annotated.
  • Change CallTIRRewrite pass to ModulePass, find and fill the device_index in the emitted alloc_tensor by looking up the global vdevice list of context mod.
  • Update DefaultGPUSchedule to only works for prim_funcs that are supposed to be on GPU, allowing prim_funcs to work for multiple targets.
  • Make the target argument of relax.build optional, and use the annotated target in the vdevice list for compilation if it is not specified.

cc: @tqchen @psrivas2 @junrushao @csullivan @Lunderberg @zxybazh @qelk123

Comment thread src/tir/transforms/default_gpu_schedule.cc
@yongwww yongwww force-pushed the e2e_multi_device branch 2 times, most recently from 01703e1 to 69a0472 Compare October 11, 2023 03:43
@tqchen
Copy link
Copy Markdown
Member

tqchen commented Oct 31, 2023

cc @zxybazh @Hzfengsy

Copy link
Copy Markdown
Member

@zxybazh zxybazh left a comment

Choose a reason for hiding this comment

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

LGTM, thanks @yongwww!

@yongwww yongwww merged commit 1c35c39 into apache:unity Dec 20, 2023
@yongwww yongwww deleted the e2e_multi_device branch January 18, 2024 18:27
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.

3 participants