-
Notifications
You must be signed in to change notification settings - Fork 13
/
cugraph.go
44 lines (39 loc) · 1.49 KB
/
cugraph.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
// All CuGraph operations.
package main
import "os"
// When KITE_ENABLE_CUDA=yes is set, these NetworKit operations
// will instead be passed to CuGraph for execution.
var preferCUDAOverNetworKit = map[string]bool{
"EstimateBetweenness": true,
"KatzCentrality": true,
"PLM": true,
"CoreDecomposition": true,
"ForceAtlas2": true,
}
func cudaEnabled() bool {
return os.Getenv("KITE_ENABLE_CUDA") == "yes"
}
// Set ifCudaDoesNotTakePrecedence as canCompute for NetworKit operations.
func ifCudaDoesNotTakePrecedence(operationDescription OperationDescription) bool {
// When CUDA is enabled, we allow the CUDA implementation to take precedence.
op := operationDescription.Data["op"].(string)
if cudaEnabled() && preferCUDAOverNetworKit[op] {
return false
}
return true
}
func init() {
if !cudaEnabled() {
return
}
diskOperationRepository["ConnectedComponents"] = pythonOperation("cugraph_connected_components")
diskOperationRepository["PageRank"] = pythonOperation("cugraph_pagerank")
insteadOfNK := pythonOperation("cugraph_instead_of_networkit")
insteadOfNK.canCompute = func(od OperationDescription) bool {
return preferCUDAOverNetworKit[od.Data["op"].(string)]
}
diskOperationRepository["NetworKitComputeDoubleAttribute"] = insteadOfNK
diskOperationRepository["NetworKitComputeDoubleEdgeAttribute"] = insteadOfNK
diskOperationRepository["NetworKitCommunityDetection"] = insteadOfNK
diskOperationRepository["NetworKitComputeVectorAttribute"] = insteadOfNK
}