Skip to content

Conversation

@yongwww
Copy link
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

@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
Member

tqchen commented Oct 31, 2023

cc @zxybazh @Hzfengsy

Copy link
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