Skip to content

Tree decomposition implements a wrong heuristics, and unclear what a correct and good heuristics is. #960

Description

@yihozhang

In the last step of tree decomposition, we merge bags in a tree decomposition to force a path decomposition. Here's the current heuristics implemented:

    let mut best_pathwidth = vec![usize::MAX; bags.len()];
    let mut full = vec![HashSet::default(); bags.len()];
    let mut choice = vec![usize::MAX; bags.len()];
    for i in 0..bags.len() {
        let mut full_i: HashSet<AtomId> =
            bags[i].atoms.iter().map(|(atom_id, _)| atom_id).collect();
        for child in all_children_list[i].iter() {
            full_i.extend(full[*child].iter().copied());
        }
        full[i] = full_i;
        best_pathwidth[i] = full[i].len();
        for chain_child in all_children_list[i].iter() {
            let mut chain_score: HashSet<_> =
                bags[i].atoms.iter().map(|(atom_id, _)| atom_id).collect();
            chain_score.extend(
                all_children_list[*chain_child]
                    .iter()
                    .filter(|child| *child != chain_child)
                    .flat_map(|child| full[*child].iter().copied()),
            );
            let s = chain_score.len().max(best_pathwidth[*chain_child]);
            if s <= best_pathwidth[i] {
                best_pathwidth[i] = s;
                choice[i] = *chain_child;
            }
        }

        // Find the parent of this bag, which must be the lowerest-numbered bag
        // that shares the most variables with it.
        let parent = bags
            .iter()
            .enumerate()
            .skip(i + 1)
            .map(|(j, b)| (j, b.common_vars_with(&bags[i]).count()))
            .filter(|(_, count)| *count > 0)
            .max_by_key(|(j, count)| (*count, -(*j as isize)));
        if let Some((j, _count)) = parent {
            all_children_list[j].push(i);
        }
    }

   // ...

    for i in (0..bags_opt.len()).rev() {
        if visited[i] {
            continue;
        }
        stack.push((i, None));
        visited[i] = true;

        while let Some((bag_id, parent)) = stack.pop() {
            let bag = mem::take(&mut bags_opt[bag_id]).unwrap();

            let this;
            if let Some(parent) = parent {
                bags_topo[parent].merge_bag(&bag);
                this = parent;
            } else {
                this = bags_topo.len();
            }

            let all_children = &mut all_children_list[bag_id];

            if parent.is_some() {
                // This bag is being absorbed into `bags_topo[this]`. To keep the
                // result a chain, every descendant of this bag is also absorbed —
                // none of them get to spawn a new chain node.
                for &i in all_children.iter() {
                    visited[i] = true;
                    stack.push((i, Some(this)));
                }
            } else {
                // This bag is a chain node. The child that minimizes pathwidth continues the
                // chain; the rest (and all their descendants, via the branch above)
                // are absorbed into this chain node.
                if !all_children.is_empty() {
                    for &i in all_children[1..].iter() {
                        if i == choice[bag_id] {
                            continue;
                        }
                        visited[i] = true;
                        stack.push((i, Some(this)));
                    }
                    visited[choice[bag_id]] = true;
                    stack.push((choice[bag_id], None));
                }
            }

            if parent.is_none() {
                bags_topo.push(bag);
            }
        }
    }

It tries to minimize the pathwidth (the size of the biggest bag). But it has two bugs:

First, the chain score for picking a child chain_child of node i should be the union of scores of all children of i that does not include this chain_child, not the children of chain_child, which would be the grandchildren of i.

            chain_score.extend(
-                all_children_list[*chain_child]
+                all_children_list[*i] 
                    .iter()
                    .filter(|child| *child != chain_child)
                    .flat_map(|child| full[*child].iter().copied()),
            );

Second, when actually merging the children bags to parents, the code wrongly excludes the first child from being considered (leftover from previous heuristics). It should consider all children and only exclude the child that is excluded.

                if !all_children.is_empty() {
-                    for &i in all_children[1..].iter() {
+                    for &i in all_children.iter() {
                        if i == choice[bag_id] {
                            continue;
                        }
                        visited[i] = true;
                        stack.push((i, Some(this)));
                    }
                    visited[choice[bag_id]] = true;
                    stack.push((choice[bag_id], None));
                }

However, after fixing them, minimizing the pathwidth leads to a significant drop in performance on some benchmarks, and it times out on paged_llama. My hypothesis is that it does not constrain the message variables between bags (the separator size), which leads to big intermediate results. Another hypothesis is that maybe we only minimize the size of the bag, disregarding the actual cardinality.

Another heuristic I tried is to minimize the separator size, i.e., the largest size of message variables passed between bags in a chain. The result is mixed:


  Benchmark                                  Before (s)   After (s)     Δ (s)      Δ %
  ────────────────────────────────────────────────────────────────────────────────────────
  hardboiled_conv1d_32.egg                        0.105       0.103   -0.003    -2.4%  ▼ faster
  python_array_optimize.egg                       0.195       0.195  +  0.000  +   0.1%  ·
  paged_llama.egg                                 1.301       0.427   -0.874   -67.2%  ▼ faster
  qwen3_moe.egg                                   0.487       0.678  +  0.190  +  39.1%  ▲ slower
  whisper.egg                                     1.013       2.614  +  1.601  + 158.0%  ▲ slower

  Summary: 2 faster  ·  2 slower  ·  1 unchanged  ·  0 missing
  Overall average Δ: +25.51%

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions