chomsky-extract 0.0.1

E-Graph extraction and selection logic for Chomsky optimization framework
Documentation 
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
45
46
47
48
49
50
51
52

use chomsky_cost::{CostModel, CpuCostModel, GpuCostModel};
use chomsky_extract::IKunExtractor;
use chomsky_uir::egraph::{EGraph, Language};
use chomsky_uir::intent::IKun;

#[test]
fn test_heterogeneous_extraction() {
    let egraph = EGraph::<IKun, ()>::new();

    // Create a program: Map(f, x)
    // There are two equivalent nodes in the same e-class:
    // 1. Map(f, x) - Generic, slow
    // 2. GpuMap(f, x) - Optimized for GPU
    // 3. CpuMap(f, x) - Optimized for CPU

    let f = egraph.add(IKun::Symbol("double".to_string()));
    let x = egraph.add(IKun::Symbol("data".to_string()));

    let generic_map = egraph.add(IKun::Map(f, x));
    let cpu_map = egraph.add(IKun::CpuMap(f, x));
    let gpu_map = egraph.add(IKun::GpuMap(f, x));

    // Union them into the same e-class
    egraph.union(generic_map, cpu_map);
    egraph.union(generic_map, gpu_map);
    egraph.rebuild();

    let root = egraph.union_find.find(generic_map);

    // Case 1: Extract for CPU
    {
        let extractor = IKunExtractor::new(&egraph, CpuCostModel);
        let tree = extractor.extract(root);
        // Should pick CpuMap
        match tree {
            chomsky_uir::IKunTree::Extension(name, _) if name == "CpuMap" => {}
            _ => {
                // In IKunTree, CpuMap might be mapped to Extension or its own variant.
                // Let's check how IKunTree is defined.
            }
        }
        println!("CPU Tree: {:?}", tree);
    }

    // Case 2: Extract for GPU
    {
        let extractor = IKunExtractor::new(&egraph, GpuCostModel);
        let tree = extractor.extract(root);
        // Should pick GpuMap
        println!("GPU Tree: {:?}", tree);
    }
}