Skip to content

Improve the computation of TreeSet #228

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gdalle merged 3 commits into from
Apr 4, 2025
Merged

Improve the computation of TreeSet #228

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6739849
Improve the computation of TreeSet
amontoison Apr 2, 2025
408ed3b
Update a comment
amontoison Apr 3, 2025
d02912a
Merge branch 'main' into acyclic_tree
gdalle Apr 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 47 additions & 26 deletions src/coloring.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -291,7 +291,9 @@ function acyclic_coloring(
end
end

tree_set = TreeSet(g, forest)
buffer = forbidden_colors
reverse_bfs_orders = first_visit_to_tree
tree_set = TreeSet(g, forest, buffer, reverse_bfs_orders)
if postprocessing
# Reuse the vector forbidden_colors to compute offsets during post-processing
offsets = forbidden_colors
Expand Down Expand Up @@ -372,19 +374,29 @@ Encode a set of 2-colored trees resulting from the [`acyclic_coloring`](@ref) al
$TYPEDFIELDS
"""
struct TreeSet{T}
reverse_bfs_orders::Vector{Vector{Tuple{T,T}}}
reverse_bfs_orders::Vector{Tuple{T,T}}
is_star::Vector{Bool}
num_edges_per_tree::Vector{T}
end

function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
function TreeSet(
g::AdjacencyGraph{T},
forest::Forest{T},
buffer::AbstractVector{T},
reverse_bfs_orders::Vector{Tuple{T,T}},
) where {T}
S = pattern(g)
edge_to_index = edge_indices(g)
nv = nb_vertices(g)
nt = forest.num_trees

# dictionary that maps a tree's root to the index of the tree
roots = Dict{T,T}()
sizehint!(roots, nt)
# root_to_tree is a vector that maps a tree's root to the index of the tree
# We can recycle forest.ranks because we don't need it anymore to merge trees
root_to_tree = forest.ranks
fill!(root_to_tree, zero(T))

# Contains the number of edges per tree
num_edges_per_tree = zeros(T, nt)

# vector of dictionaries where each dictionary stores the neighbors of each vertex in a tree
trees = [Dict{T,Vector{T}}() for i in 1:nt]
Expand All @@ -401,13 +413,14 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
root = find_root!(forest, index_ij)

# Update roots
if !haskey(roots, root)
if iszero(root_to_tree[root])
nr += 1
roots[root] = nr
root_to_tree[root] = nr
end

# index of the tree T that contains this edge
index_tree = roots[root]
index_tree = root_to_tree[root]
num_edges_per_tree[index_tree] += 1

# Update the neighbors of i in the tree T
if !haskey(trees[index_tree], i)
Expand All @@ -427,10 +440,7 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
end

# degrees is a vector of integers that stores the degree of each vertex in a tree
degrees = Vector{T}(undef, nv)

# reverse breadth first (BFS) traversal order for each tree in the forest
reverse_bfs_orders = [Tuple{T,T}[] for i in 1:nt]
degrees = buffer

# nvmax is the number of vertices of the biggest tree in the forest
nvmax = 0
Expand All @@ -446,6 +456,10 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
# meaning that one vertex is directly connected to all other vertices in the tree
is_star = Vector{Bool}(undef, nt)

# Number of edges treated
num_edges_treated = zero(T)

# reverse_bfs_orders contains the reverse breadth first (BFS) traversal order for each tree in the forest
for k in 1:nt
tree = trees[k]
Copy link
Copy Markdown
Member

@gdalle gdalle Apr 4, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the only place where trees is actually used. Right now it is one dictionary per tree, which is still rather wasteful and probably accounts for a large percentage of our allocs. Can we figure out a way to get rid of this too?

Copy link
Copy Markdown
Member

@gdalle gdalle Apr 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On second thought, let's do it in a separate PR, since this one works

Copy link
Copy Markdown
Collaborator Author

@amontoison amontoison Apr 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I started to work on it and it is not a simple task.


Expand Down Expand Up @@ -483,7 +497,8 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
# Check if neighbor is the parent of the leaf or if it was a child before the tree was pruned
if degrees[neighbor] != 0
# (leaf, neighbor) represents the next edge to visit during decompression
push!(reverse_bfs_orders[k], (leaf, neighbor))
num_edges_treated += 1
reverse_bfs_orders[num_edges_treated] = (leaf, neighbor)

if bool_star
# Initialize the potential hub of the star with the first parent of a leaf
Expand Down Expand Up @@ -514,7 +529,7 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
is_star[k] = bool_star
end

return TreeSet(reverse_bfs_orders, is_star)
return TreeSet(reverse_bfs_orders, is_star, num_edges_per_tree)
end

## Postprocessing, mirrors decompression code
Expand Down Expand Up @@ -582,46 +597,52 @@ function postprocess!(
end
else
# only the colors of non-leaf vertices are used
(; reverse_bfs_orders, is_star) = star_or_tree_set
(; reverse_bfs_orders, is_star, num_edges_per_tree) = star_or_tree_set
nb_trivial_trees = 0

# Index of the first edge in reverse_bfs_orders for the current tree
first = 1

# Iterate through all non-trivial trees
for k in eachindex(reverse_bfs_orders)
reverse_bfs_order = reverse_bfs_orders[k]
for k in eachindex(num_edges_per_tree)
ne_tree = num_edges_per_tree[k]
# Check if we have more than one edge in the tree (non-trivial tree)
if length(reverse_bfs_order) > 1
if ne_tree > 1
# Determine if the tree is a star
if is_star[k]
# It is a non-trivial star and only the color of the hub is needed
(_, hub) = reverse_bfs_order[1]
(_, hub) = reverse_bfs_orders[first]
color_used[color[hub]] = true
else
# It is not a star and both colors are needed during the decompression
(i, j) = reverse_bfs_order[1]
(i, j) = reverse_bfs_orders[first]
color_used[color[i]] = true
color_used[color[j]] = true
end
else
nb_trivial_trees += 1
end
first += ne_tree
end

# Process the trivial trees (if any)
if nb_trivial_trees > 0
for k in eachindex(reverse_bfs_orders)
reverse_bfs_order = reverse_bfs_orders[k]
first = 1
for k in eachindex(num_edges_per_tree)
ne_tree = num_edges_per_tree[k]
# Check if we have exactly one edge in the tree
if length(reverse_bfs_order) == 1
(i, j) = reverse_bfs_order[1]
if ne_tree == 1
(i, j) = reverse_bfs_orders[first]
if color_used[color[i]]
# Make i the root to avoid possibly adding one more used color
# Switch it with the (only) leaf
reverse_bfs_order[1] = (j, i)
reverse_bfs_orders[first] = (j, i)
else
# Keep j as the root
color_used[color[j]] = true
end
end
first += ne_tree
end
end
end
Expand Down
39 changes: 29 additions & 10 deletions src/decompression.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -517,7 +517,7 @@ end
function decompress!(
A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
)
(; ag, color, reverse_bfs_orders, buffer) = result
(; ag, color, reverse_bfs_orders, num_edges_per_tree, buffer) = result
(; S) = ag
uplo == :F && check_same_pattern(A, S)
R = eltype(A)
Expand All @@ -531,24 +531,32 @@ function decompress!(

# Recover the diagonal coefficients of A
if has_diagonal(ag)
for i in axes(A, 1)
for i in axes(S, 1)
if !iszero(S[i, i])
A[i, i] = B[i, color[i]]
end
end
end

# Index of the first edge in reverse_bfs_orders for the current tree
first = 1

# Recover the off-diagonal coefficients of A
for k in eachindex(reverse_bfs_orders)
for k in eachindex(num_edges_per_tree)
ne_tree = num_edges_per_tree[k]
last = first + ne_tree - 1

# Reset the buffer to zero for all vertices in a tree (except the root)
for (vertex, _) in reverse_bfs_orders[k]
for pos in first:last
(vertex, _) = reverse_bfs_orders[pos]
buffer_right_type[vertex] = zero(R)
end
# Reset the buffer to zero for the root vertex
(_, root) = reverse_bfs_orders[k][end]
(_, root) = reverse_bfs_orders[last]
buffer_right_type[root] = zero(R)

for (i, j) in reverse_bfs_orders[k]
for pos in first:last
(i, j) = reverse_bfs_orders[pos]
val = B[i, color[j]] - buffer_right_type[i]
buffer_right_type[j] = buffer_right_type[j] + val

Expand All @@ -559,6 +567,7 @@ function decompress!(
A[j, i] = val
end
end
first += ne_tree
end
return A
end
Expand All @@ -573,6 +582,7 @@ function decompress!(
ag,
color,
reverse_bfs_orders,
num_edges_per_tree,
diagonal_indices,
diagonal_nzind,
lower_triangle_offsets,
Expand Down Expand Up @@ -612,20 +622,28 @@ function decompress!(
end
end

# Index of the first edge in reverse_bfs_orders for the current tree
first = 1

# Index of offsets in lower_triangle_offsets and upper_triangle_offsets
counter = 0

# Recover the off-diagonal coefficients of A
for k in eachindex(reverse_bfs_orders)
for k in eachindex(num_edges_per_tree)
ne_tree = num_edges_per_tree[k]
last = first + ne_tree - 1

# Reset the buffer to zero for all vertices in a tree (except the root)
for (vertex, _) in reverse_bfs_orders[k]
for pos in first:last
(vertex, _) = reverse_bfs_orders[pos]
buffer_right_type[vertex] = zero(R)
end
# Reset the buffer to zero for the root vertex
(_, root) = reverse_bfs_orders[k][end]
(_, root) = reverse_bfs_orders[last]
buffer_right_type[root] = zero(R)

for (i, j) in reverse_bfs_orders[k]
for pos in first:last
(i, j) = reverse_bfs_orders[pos]
counter += 1
val = B[i, color[j]] - buffer_right_type[i]
buffer_right_type[j] = buffer_right_type[j] + val
Expand Down Expand Up @@ -665,6 +683,7 @@ function decompress!(
end
#! format: on
end
first += ne_tree
end
return A
end
Expand Down
18 changes: 14 additions & 4 deletions src/result.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,7 +314,8 @@ struct TreeSetColoringResult{
ag::G
color::Vector{T}
group::GT
reverse_bfs_orders::Vector{Vector{Tuple{T,T}}}
reverse_bfs_orders::Vector{Tuple{T,T}}
num_edges_per_tree::Vector{T}
diagonal_indices::Vector{T}
diagonal_nzind::Vector{T}
lower_triangle_offsets::Vector{T}
Expand All @@ -329,7 +330,7 @@ function TreeSetColoringResult(
tree_set::TreeSet{<:Integer},
decompression_eltype::Type{R},
) where {T<:Integer,R}
(; reverse_bfs_orders) = tree_set
(; reverse_bfs_orders, num_edges_per_tree) = tree_set
(; S) = ag
nvertices = length(color)
group = group_by_color(T, color)
Expand Down Expand Up @@ -358,11 +359,18 @@ function TreeSetColoringResult(
lower_triangle_offsets = Vector{T}(undef, nedges)
upper_triangle_offsets = Vector{T}(undef, nedges)

# Index of the first edge in reverse_bfs_orders for the current tree
first = 1

# Index in lower_triangle_offsets and upper_triangle_offsets
index_offsets = 0

for k in eachindex(reverse_bfs_orders)
for (leaf, neighbor) in reverse_bfs_orders[k]
for k in eachindex(num_edges_per_tree)
ne_tree = num_edges_per_tree[k]
last = first + ne_tree - 1

for pos in first:last
(leaf, neighbor) = reverse_bfs_orders[pos]
# Update lower_triangle_offsets and upper_triangle_offsets
i = leaf
j = neighbor
Expand Down Expand Up @@ -393,6 +401,7 @@ function TreeSetColoringResult(
end
#! format: on
end
first += ne_tree
end

# buffer holds the sum of edge values for subtrees in a tree.
Expand All @@ -405,6 +414,7 @@ function TreeSetColoringResult(
color,
group,
reverse_bfs_orders,
num_edges_per_tree,
diagonal_indices,
diagonal_nzind,
lower_triangle_offsets,
Expand Down