-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Feature] add metis partitioning to DGL (dmlc#1308)
* add metis. * add test. * construct partition id. * link to METIS github repo. * update metis. * add a tool for partitioning a graph. * update metis. * update. * update. * fix metis. * fix lint * fix indent. * another way of building metis. * disable metis in windows. * test windows * fix. * disable metis for windows properly. * fix for tensorflow. * skip test for gpu. * make graph symmetric * address comments. * more comments. * fix compile * fix a bug. * add test. * change the default #hops of HALO nodes. Co-authored-by: Ubuntu <[email protected]>
- Loading branch information
Showing
9 changed files
with
247 additions
and
6 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,86 @@ | ||
| /*! | ||
| * Copyright (c) 2020 by Contributors | ||
| * \file graph/metis_partition.cc | ||
| * \brief Call Metis partitioning | ||
| */ | ||
|
|
||
| #include <metis.h> | ||
| #include <dgl/graph_op.h> | ||
| #include <dgl/packed_func_ext.h> | ||
| #include "../c_api_common.h" | ||
|
|
||
| using namespace dgl::runtime; | ||
|
|
||
| namespace dgl { | ||
|
|
||
| #if !defined(_WIN32) | ||
|
|
||
| IdArray GraphOp::MetisPartition(GraphPtr g, int k) { | ||
| // The index type of Metis needs to be compatible with DGL index type. | ||
| CHECK_EQ(sizeof(idx_t), sizeof(dgl_id_t)); | ||
| ImmutableGraphPtr ig = std::dynamic_pointer_cast<ImmutableGraph>(g); | ||
| CHECK(ig) << "The input graph must be an immutable graph."; | ||
| // This is a symmetric graph, so in-csr and out-csr are the same. | ||
| const auto mat = ig->GetInCSR()->ToCSRMatrix(); | ||
|
|
||
| idx_t nvtxs = g->NumVertices(); | ||
| idx_t ncon = 1; // # balacing constraints. | ||
| idx_t *xadj = static_cast<idx_t*>(mat.indptr->data); | ||
| idx_t *adjncy = static_cast<idx_t*>(mat.indices->data); | ||
| idx_t nparts = k; | ||
| IdArray part_arr = aten::NewIdArray(nvtxs); | ||
| idx_t objval = 0; | ||
| idx_t *part = static_cast<idx_t*>(part_arr->data); | ||
| int ret = METIS_PartGraphKway(&nvtxs, // The number of vertices | ||
| &ncon, // The number of balancing constraints. | ||
| xadj, // indptr | ||
| adjncy, // indices | ||
| NULL, // the weights of the vertices | ||
| NULL, // The size of the vertices for computing | ||
| // the total communication volume | ||
| NULL, // The weights of the edges | ||
| &nparts, // The number of partitions. | ||
| NULL, // the desired weight for each partition and constraint | ||
| NULL, // the allowed load imbalance tolerance | ||
| NULL, // the array of options | ||
| &objval, // the edge-cut or the total communication volume of | ||
| // the partitioning solution | ||
| part); | ||
| LOG(INFO) << "Partition a graph with " << g->NumVertices() | ||
| << " nodes and " << g->NumEdges() | ||
| << " edges into " << k | ||
| << " parts and get " << objval << " edge cuts"; | ||
| switch (ret) { | ||
| case METIS_OK: | ||
| return part_arr; | ||
| case METIS_ERROR_INPUT: | ||
| LOG(FATAL) << "Error in Metis partitioning: input error"; | ||
| case METIS_ERROR_MEMORY: | ||
| LOG(FATAL) << "Error in Metis partitioning: cannot allocate memory"; | ||
| default: | ||
| LOG(FATAL) << "Error in Metis partitioning: other errors"; | ||
| } | ||
| // return an array of 0 elements to indicate the error. | ||
| return aten::NullArray(); | ||
| } | ||
|
|
||
| DGL_REGISTER_GLOBAL("transform._CAPI_DGLMetisPartition") | ||
| .set_body([] (DGLArgs args, DGLRetValue* rv) { | ||
| GraphRef g = args[0]; | ||
| int k = args[1]; | ||
| *rv = GraphOp::MetisPartition(g.sptr(), k); | ||
| }); | ||
|
|
||
| #else | ||
|
|
||
| DGL_REGISTER_GLOBAL("transform._CAPI_DGLMetisPartition") | ||
| .set_body([] (DGLArgs args, DGLRetValue* rv) { | ||
| GraphRef g = args[0]; | ||
| int k = args[1]; | ||
| LOG(WARNING) << "DGL doesn't support METIS partitioning in Windows"; | ||
| *rv = aten::NullArray(); | ||
| }); | ||
|
|
||
| #endif // !defined(_WIN32) | ||
|
|
||
| } // namespace dgl |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,56 @@ | ||
| import numpy as np | ||
| import argparse | ||
| import signal | ||
| import dgl | ||
| from dgl import backend as F | ||
| from dgl.data.utils import load_graphs, save_graphs | ||
|
|
||
| def main(): | ||
| parser = argparse.ArgumentParser(description='Partition a graph') | ||
| parser.add_argument('--data', required=True, type=str, | ||
| help='The file path of the input graph in the DGL format.') | ||
| parser.add_argument('-k', '--num-parts', required=True, type=int, | ||
| help='The number of partitions') | ||
| parser.add_argument('--num-hops', type=int, default=1, | ||
| help='The number of hops of HALO nodes we include in a partition') | ||
| parser.add_argument('-m', '--method', required=True, type=str, | ||
| help='The partitioning method: random, metis') | ||
| parser.add_argument('-o', '--output', required=True, type=str, | ||
| help='The output directory of the partitioned results') | ||
| args = parser.parse_args() | ||
| data_path = args.data | ||
| num_parts = args.num_parts | ||
| num_hops = args.num_hops | ||
| method = args.method | ||
| output = args.output | ||
|
|
||
| glist, _ = load_graphs(data_path) | ||
| g = glist[0] | ||
|
|
||
| if args.method == 'metis': | ||
| part_dict = dgl.transform.metis_partition(g, num_parts, num_hops) | ||
| elif args.method == 'random': | ||
| node_parts = np.random.choice(num_parts, g.number_of_nodes()) | ||
| part_dict = dgl.transform.partition_graph_with_halo(g, node_parts, num_hops) | ||
| else: | ||
| raise Exception('unknown partitioning method: ' + args.method) | ||
|
|
||
| tot_num_inner_edges = 0 | ||
| for part_id in part_dict: | ||
| part = part_dict[part_id] | ||
|
|
||
| num_inner_nodes = len(np.nonzero(F.asnumpy(part.ndata['inner_node']))[0]) | ||
| num_inner_edges = len(np.nonzero(F.asnumpy(part.edata['inner_edge']))[0]) | ||
| print('part {} has {} nodes and {} edges. {} nodes and {} edges are inside the partition'.format( | ||
| part_id, part.number_of_nodes(), part.number_of_edges(), | ||
| num_inner_nodes, num_inner_edges)) | ||
| tot_num_inner_edges += num_inner_edges | ||
|
|
||
| # TODO I duplicate some node features. | ||
| part.copy_from_parent() | ||
| save_graphs(output + '/' + str(part_id) + '.dgl', [part]) | ||
| print('there are {} edges in the graph and {} edge cuts for {} partitions.'.format( | ||
| g.number_of_edges(), g.number_of_edges() - tot_num_inner_edges, len(part_dict))) | ||
|
|
||
| if __name__ == '__main__': | ||
| main() |