def get_bandwidth_matrix()

in optimum/amd/topology_utils.py [0:0]

• 34 lines of code
• 20 McCabe index (conditional complexity)

def get_bandwidth_matrix():
    """
    Returns a matrix of bandwidths between all GPU devices in the system.
    """
    amdsmi.amdsmi_init()
    devices = amdsmi.amdsmi_get_device_handles()

    num_devices = len(devices)
    bandwidth_matrix = [[None for _ in range(num_devices)] for _ in range(num_devices)]

    # direct bandwidth
    for i, src_device in enumerate(devices):
        for j, dst_device in enumerate(devices):
            if i == j:
                bandwidth_matrix[i][j] = float("inf")
            else:
                try:
                    curr_bandwidth = amdsmi.amdsmi_get_minmax_bandwidth(src_device, dst_device)["max_bandwidth"]
                    if curr_bandwidth != 0:
                        bandwidth_matrix[i][j] = curr_bandwidth
                except Exception:
                    pass

    # indirect bandwidth
    for i in range(num_devices):
        for j in range(num_devices):
            if bandwidth_matrix[i][j] is None:
                maxmin_bandwidth = 0
                for k in range(num_devices):
                    if k == i or k == j:
                        continue
                    elif bandwidth_matrix[i][k] is not None and bandwidth_matrix[k][j] is not None:
                        min_bandwidth = min(bandwidth_matrix[i][k], bandwidth_matrix[k][j])
                        if min_bandwidth > maxmin_bandwidth:
                            bandwidth_matrix[i][j] = min_bandwidth
                            maxmin_bandwidth = min_bandwidth

    # fill missing values
    for i in range(num_devices):
        for j in range(num_devices):
            if bandwidth_matrix[i][j] is None:
                bandwidth_matrix[i][j] = 0

    amdsmi.amdsmi_shut_down()

    return bandwidth_matrix