b1eeb934494ef1eee20cf2d35b718790cb9cb550,python/dgl/runtime/scheduler.py,,schedule_update_all,#Any#Any#Any#Any#,258

Before Change

    reduced_feat = _gen_send_reduce(call_type, graph,
            message_func, reduce_func, (var_src, var_dst, var_eid), recv_nodes)
    // generate optional apply
    if apply_func:
        // To avoid writing reduced features back to node frame and reading
        // it again for apply phase. Instead, we first read the the node
        // features and "merge" it with the reduced features.
        v_nf = ir.READ_ROW(var_nf, var_recv_nodes)
        v_nf = ir.UPDATE_DICT(v_nf, reduced_feat)
        def _afunc_wrapper(node_data):
            nb = NodeBatch(graph, recv_nodes, node_data)
            return apply_func(nb)
        afunc = var.FUNC(_afunc_wrapper)
        applied_feat = ir.NODE_UDF(afunc, v_nf)
        final_feat = ir.UPDATE_DICT(reduced_feat, applied_feat)
    else:
        final_feat = reduced_feat
    ir.WRITE_DICT_(var_nf, final_feat)

def schedule_apply_nodes(graph, v, apply_func):
    get apply nodes schedule

After Change

        var_dst = var.IDX(dst)
        var_eid = var.IDX(eid)
        // generate send + reduce
        reduced_feat = _gen_send_reduce(call_type, graph,
                message_func, reduce_func, (var_src, var_dst, var_eid), recv_nodes)
        // generate optional apply
        final_feat = _apply_with_accum(graph, var_recv_nodes, var_nf, reduced_feat, apply_func)
        ir.WRITE_DICT_(var_nf, final_feat)

def schedule_apply_nodes(graph, v, apply_func):
    get apply nodes schedule

Instances