Вопрос о последовательном или параллельном выполнении команд передачи и агрегации сообщений DGL

Question

Мой вопрос касается порядка выполнения команд передачи сообщений и агрегации из DGL на графическом процессоре.

Я хотел бы использовать модуль на графике с, возможно, множеством разных типов ребер, и мне интересно, является ли для любого из приведенных ниже вызовов просто планирование задачи на графическом процессоре или вызов, блокирующий оставшиеся итерации цикла до тех пор, пока после того, как его вычисления будут выполнены?

Вот основной цикл forward() метод GatedGraphConv(находится на https://docs.dgl.ai/_modules/dgl/nn/pytorch/conv/gatedgraphconv.html#GatedGraphConv):

      for _ in range(self._n_steps):
    graph.ndata['h'] = feat
    for i in range(self._n_etypes):
        eids = th.nonzero(etypes == i, as_tuple=False).view(-1).type(graph.idtype)
        if len(eids) > 0:
            graph.apply_edges(
                lambda edges: {'W_e*h': self.linears[i](edges.src['h'])},
                eids
            )
    graph.update_all(fn.copy_e('W_e*h', 'm'), fn.sum('m', 'a'))
    a = graph.ndata.pop('a') # (N, D)
    feat = self.gru(a, feat)

Я рассмотрел возможность использования набора потоков cuda на одном графическом процессоре, где каждому из них будет дан свой собственный вызов (с другой строкой "W_e*h" для каждого типа ребра, чтобы они не конфликтовали друг с другом), а затем синхронизируйте их перед вызовом update_all(), но мне пришло в голову, что это могло быть излишним и что, возможно, призывы к apply_edges() просто запланировать задачу на графическом процессоре, а затем предоставить графическому процессору возможность выполнять все запланированные задачи тем способом, который он считает нужным (включая одновременное выполнение нескольких задач, если это возможно).

Спасибо.