search

NVIDIA / Fuser

A Fusion Code Generator for NVIDIA GPUs (commonly known as "nvFuser")
Other
229 stars 42 forks source link

DeviceMesh helper functions #2228

Closed cowanmeg closed 1 month ago

cowanmeg commented 1 month ago

This PR introduces some common convenience helper functions for DeviceMesh include size, at, etc. so that we rely less on accessing the underlying vector of DeviceIdxType. There are still a few places in csrc/multidevice/lower_communications.cpp that access the vector directly.

Eventually, these helper functions will be parameterized by the axis of the DeviceMesh.

cowanmeg commented 1 month ago

Note: I will update this after https://github.com/NVIDIA/Fuser/pull/2185 is merged.

cowanmeg commented 1 month ago

!build --dist

cowanmeg commented 1 month ago

Failures are unrelated from the missing imports.