Skip to content

Commit 3e6ae79

Browse files
michalharakalmichalharakal
committed
Add initial export to graphviz..
1 parent be0754f commit 3e6ae79

2 files changed

Lines changed: 286 additions & 0 deletions

File tree

‎skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/graph/utils/graphviz.kt‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
package sk.ainet.lang.graph.utils
2+
3+
import sk.ainet.lang.graph.GraphNode
4+
import sk.ainet.lang.graph.ComputeGraph
5+
6+
/**
7+
* Represents a Graphviz DOT graph output
8+
*/
9+
public data class DotGraph(val content: String)
10+
11+
/**
12+
* Traces the compute graph starting from output nodes to collect all nodes and dependencies
13+
*/
14+
public fun trace(graph: ComputeGraph): Pair<Set<GraphNode>, Set<Pair<GraphNode, GraphNode>>> {
15+
val nodes = graph.nodes.toSet()
16+
val edges = mutableSetOf<Pair<GraphNode, GraphNode>>()
17+
18+
// Build edges based on graph dependencies
19+
for (node in nodes) {
20+
val inputNodes = graph.getInputNodes(node)
21+
for (inputNode in inputNodes) {
22+
edges.add(inputNode to node)
23+
}
24+
}
25+
26+
return nodes to edges
27+
}
28+
29+
/**
30+
* Generates a Graphviz DOT representation of the compute graph
31+
*/
32+
public fun drawDot(graph: ComputeGraph, rankdir: String = "LR"): DotGraph {
33+
require(rankdir in listOf("LR", "TB"))
34+
35+
val (nodes, edges) = trace(graph)
36+
val dotContent = StringBuilder()
37+
38+
dotContent.appendLine("digraph {")
39+
dotContent.appendLine(" rankdir=$rankdir;")
40+
41+
// Add nodes
42+
for (node in nodes) {
43+
val nodeId = node.id.replace("[^a-zA-Z0-9_]".toRegex(), "_")
44+
45+
val labelContent = "${node.operationName} | ${node.id}"
46+
47+
dotContent.appendLine(" $nodeId [label=\"$labelContent\", shape=record];")
48+
49+
// Add operation node if operation has parameters
50+
if (node.operation.parameters.isNotEmpty()) {
51+
val opId = "${nodeId}_op"
52+
val paramStr = node.operation.parameters.entries.joinToString("\\n") { "${it.key}: ${it.value}" }
53+
dotContent.appendLine(" $opId [label=\"$paramStr\", shape=box, style=dashed];")
54+
dotContent.appendLine(" $opId -> $nodeId [style=dotted];")
55+
}
56+
}
57+
58+
// Add edges
59+
for ((from, to) in edges) {
60+
val fromId = from.id.replace("[^a-zA-Z0-9_]".toRegex(), "_")
61+
val toId = to.id.replace("[^a-zA-Z0-9_]".toRegex(), "_")
62+
dotContent.appendLine(" $fromId -> $toId;")
63+
}
64+
65+
dotContent.appendLine("}")
66+
67+
return DotGraph(dotContent.toString())
68+
}
69+
70+
/**
71+
* Generates a Graphviz DOT representation for a subset of nodes in the compute graph
72+
*/
73+
public fun drawDot(graph: ComputeGraph, outputNodes: List<GraphNode>, rankdir: String = "LR"): DotGraph {
74+
require(rankdir in listOf("LR", "TB"))
75+
76+
// Trace from output nodes backward to collect relevant nodes
77+
val relevantNodes = mutableSetOf<GraphNode>()
78+
val edges = mutableSetOf<Pair<GraphNode, GraphNode>>()
79+
80+
fun collectDependencies(node: GraphNode) {
81+
if (node !in relevantNodes) {
82+
relevantNodes.add(node)
83+
val inputNodes = graph.getInputNodes(node)
84+
for (inputNode in inputNodes) {
85+
edges.add(inputNode to node)
86+
collectDependencies(inputNode)
87+
}
88+
}
89+
}
90+
91+
for (outputNode in outputNodes) {
92+
collectDependencies(outputNode)
93+
}
94+
95+
val dotContent = StringBuilder()
96+
97+
dotContent.appendLine("digraph {")
98+
dotContent.appendLine(" rankdir=$rankdir;")
99+
100+
// Add nodes
101+
for (node in relevantNodes) {
102+
val nodeId = node.id.replace("[^a-zA-Z0-9_]".toRegex(), "_")
103+
104+
val labelContent = "${node.operationName} | ${node.id}"
105+
106+
dotContent.appendLine(" $nodeId [label=\"$labelContent\", shape=record];")
107+
108+
// Add operation node if operation has parameters
109+
if (node.operation.parameters.isNotEmpty()) {
110+
val opId = "${nodeId}_op"
111+
val paramStr = node.operation.parameters.entries.joinToString("\\n") { "${it.key}: ${it.value}" }
112+
dotContent.appendLine(" $opId [label=\"$paramStr\", shape=box, style=dashed];")
113+
dotContent.appendLine(" $opId -> $nodeId [style=dotted];")
114+
}
115+
}
116+
117+
// Add edges
118+
for ((from, to) in edges) {
119+
val fromId = from.id.replace("[^a-zA-Z0-9_]".toRegex(), "_")
120+
val toId = to.id.replace("[^a-zA-Z0-9_]".toRegex(), "_")
121+
dotContent.appendLine(" $fromId -> $toId;")
122+
}
123+
124+
dotContent.appendLine("}")
125+
126+
return DotGraph(dotContent.toString())
127+
}

‎skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/graph/utils/GraphVizExportTest.kt‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
package sk.ainet.lang.graph.utils
2+
3+
import sk.ainet.lang.graph.*
4+
import kotlin.test.Test
5+
import kotlin.test.assertNotNull
6+
import kotlin.test.assertTrue
7+
import kotlin.test.assertEquals
8+
9+
/**
10+
* Test for GraphViz export functionality
11+
*/
12+
class GraphVizExportTest {
13+
14+
private class TestOperation(
15+
override val name: String,
16+
override val type: String,
17+
override val parameters: Map<String, Any> = emptyMap()
18+
) : BaseOperation(name, type, parameters) {
19+
20+
override fun <T : sk.ainet.lang.types.DType, V> execute(inputs: List<sk.ainet.lang.tensor.Tensor<T, V>>): List<sk.ainet.lang.tensor.Tensor<T, V>> {
21+
return inputs // Pass through for testing
22+
}
23+
24+
override fun validateInputs(inputs: List<TensorSpec>): ValidationResult {
25+
return ValidationResult.Valid
26+
}
27+
28+
override fun inferOutputs(inputs: List<TensorSpec>): List<TensorSpec> {
29+
return inputs // Pass through for testing
30+
}
31+
32+
override fun clone(newParameters: Map<String, Any>): Operation {
33+
return TestOperation(name, type, newParameters)
34+
}
35+
}
36+
37+
@Test
38+
fun testBasicGraphVizExport() {
39+
println("[DEBUG_LOG] Testing basic GraphViz export functionality")
40+
41+
// Create a simple compute graph
42+
val graph = DefaultComputeGraph()
43+
44+
// Create test operations
45+
val inputOp = TestOperation("input", "input")
46+
val processOp = TestOperation("process", "compute", mapOf("kernel_size" to 3, "stride" to 1))
47+
val outputOp = TestOperation("output", "output")
48+
49+
// Create test nodes
50+
val inputNode = GraphNode(
51+
id = "input_node",
52+
operation = inputOp,
53+
inputs = emptyList(),
54+
outputs = listOf(TensorSpec("input_out", listOf(1, 10), "FP32"))
55+
)
56+
57+
val processNode = GraphNode(
58+
id = "process_node",
59+
operation = processOp,
60+
inputs = listOf(TensorSpec("process_in", listOf(1, 10), "FP32")),
61+
outputs = listOf(TensorSpec("process_out", listOf(1, 5), "FP32"))
62+
)
63+
64+
val outputNode = GraphNode(
65+
id = "output_node",
66+
operation = outputOp,
67+
inputs = listOf(TensorSpec("output_in", listOf(1, 5), "FP32")),
68+
outputs = listOf(TensorSpec("output_out", listOf(1, 5), "FP32"))
69+
)
70+
71+
// Add nodes to graph
72+
graph.addNode(inputNode)
73+
graph.addNode(processNode)
74+
graph.addNode(outputNode)
75+
76+
// Add edges to connect them
77+
graph.addEdge(GraphEdge("edge1", inputNode, processNode, 0, 0, inputNode.outputs.first()))
78+
graph.addEdge(GraphEdge("edge2", processNode, outputNode, 0, 0, processNode.outputs.first()))
79+
80+
println("[DEBUG_LOG] Created graph with ${graph.nodes.size} nodes and ${graph.edges.size} edges")
81+
82+
// Test full graph export
83+
val dotGraph = drawDot(graph)
84+
assertNotNull(dotGraph, "DOT graph should not be null")
85+
assertNotNull(dotGraph.content, "DOT content should not be null")
86+
assertTrue(dotGraph.content.isNotEmpty(), "DOT content should not be empty")
87+
88+
println("[DEBUG_LOG] Full graph DOT export:")
89+
println(dotGraph.content)
90+
91+
// Verify DOT content structure
92+
assertTrue(dotGraph.content.contains("digraph {"), "Should contain digraph declaration")
93+
assertTrue(dotGraph.content.contains("rankdir=LR"), "Should contain rankdir setting")
94+
assertTrue(dotGraph.content.contains("input_node"), "Should contain input node")
95+
assertTrue(dotGraph.content.contains("process_node"), "Should contain process node")
96+
assertTrue(dotGraph.content.contains("output_node"), "Should contain output node")
97+
assertTrue(dotGraph.content.contains("->"), "Should contain edges")
98+
assertTrue(dotGraph.content.contains("}"), "Should contain closing brace")
99+
100+
// Test parameters are included
101+
assertTrue(dotGraph.content.contains("kernel_size"), "Should include operation parameters")
102+
103+
println("[DEBUG_LOG] Basic GraphViz export test passed")
104+
}
105+
106+
@Test
107+
fun testSubsetGraphVizExport() {
108+
println("[DEBUG_LOG] Testing subset GraphViz export functionality")
109+
110+
// Create a compute graph
111+
val graph = DefaultComputeGraph()
112+
113+
// Create test nodes
114+
val node1 = GraphNode("node1", TestOperation("op1", "type1"), emptyList(), listOf(TensorSpec("out1", listOf(1), "FP32")))
115+
val node2 = GraphNode("node2", TestOperation("op2", "type2"), listOf(TensorSpec("in2", listOf(1), "FP32")), listOf(TensorSpec("out2", listOf(1), "FP32")))
116+
val node3 = GraphNode("node3", TestOperation("op3", "type3"), listOf(TensorSpec("in3", listOf(1), "FP32")), listOf(TensorSpec("out3", listOf(1), "FP32")))
117+
118+
graph.addNode(node1)
119+
graph.addNode(node2)
120+
graph.addNode(node3)
121+
122+
graph.addEdge(GraphEdge("edge1", node1, node2, 0, 0, node1.outputs.first()))
123+
graph.addEdge(GraphEdge("edge2", node2, node3, 0, 0, node2.outputs.first()))
124+
125+
// Test subset export (only from node3 backward)
126+
val dotGraphSubset = drawDot(graph, listOf(node3))
127+
assertNotNull(dotGraphSubset, "Subset DOT graph should not be null")
128+
assertTrue(dotGraphSubset.content.isNotEmpty(), "Subset DOT content should not be empty")
129+
130+
println("[DEBUG_LOG] Subset DOT export:")
131+
println(dotGraphSubset.content)
132+
133+
// All nodes should be included since node3 depends on all previous nodes
134+
assertTrue(dotGraphSubset.content.contains("node1"), "Should contain node1 in subset")
135+
assertTrue(dotGraphSubset.content.contains("node2"), "Should contain node2 in subset")
136+
assertTrue(dotGraphSubset.content.contains("node3"), "Should contain node3 in subset")
137+
138+
println("[DEBUG_LOG] Subset GraphViz export test passed")
139+
}
140+
141+
@Test
142+
fun testDifferentRankDirections() {
143+
println("[DEBUG_LOG] Testing different rank directions")
144+
145+
val graph = DefaultComputeGraph()
146+
val node = GraphNode("test", TestOperation("test", "test"), emptyList(), listOf(TensorSpec("out", listOf(1), "FP32")))
147+
graph.addNode(node)
148+
149+
// Test LR direction (default)
150+
val dotLR = drawDot(graph, "LR")
151+
assertTrue(dotLR.content.contains("rankdir=LR"), "Should set LR rank direction")
152+
153+
// Test TB direction
154+
val dotTB = drawDot(graph, "TB")
155+
assertTrue(dotTB.content.contains("rankdir=TB"), "Should set TB rank direction")
156+
157+
println("[DEBUG_LOG] Rank direction test passed")
158+
}
159+
}

0 commit comments

Comments
 (0)