metagraphoptimization.jl/examples/ABAB_graph.jl

using metagraph_optimization

function main()
    graph = DAG()

    # s to output (exit node)
    d_exit = insert_node(graph, make_node(DataTask(10)))

    # final s compute
    s0 = insert_node(graph, make_node(ComputeTaskS2()))

    # data from v0 and v1 to s0
    d_v0_s0 = insert_node(graph, make_node(DataTask(5)))
    d_v1_s0 = insert_node(graph, make_node(DataTask(5)))

    # v0 and v1 compute
    v0 = insert_node(graph, make_node(ComputeTaskV()))
    v1 = insert_node(graph, make_node(ComputeTaskV()))

    # data from uB, uA, uBp and uAp to v0 and v1
    d_uB_v0 = insert_node(graph, make_node(DataTask(3)))
    d_uA_v0 = insert_node(graph, make_node(DataTask(3)))
    d_uBp_v1 = insert_node(graph, make_node(DataTask(3)))
    d_uAp_v1 = insert_node(graph, make_node(DataTask(3)))

    # uB, uA, uBp and uAp computes
    uB = insert_node(graph, make_node(ComputeTaskU()))
    uA = insert_node(graph, make_node(ComputeTaskU()))
    uBp = insert_node(graph, make_node(ComputeTaskU()))
    uAp = insert_node(graph, make_node(ComputeTaskU()))

    # data from PB, PA, PBp and PAp to uB, uA, uBp and uAp
    d_PB_uB = insert_node(graph, make_node(DataTask(6)))
    d_PA_uA = insert_node(graph, make_node(DataTask(6)))
    d_PBp_uBp = insert_node(graph, make_node(DataTask(6)))
    d_PAp_uAp = insert_node(graph, make_node(DataTask(6)))

    # P computes PB, PA, PBp and PAp
    PB = insert_node(graph, make_node(ComputeTaskP()))
    PA = insert_node(graph, make_node(ComputeTaskP()))
    PBp = insert_node(graph, make_node(ComputeTaskP()))
    PAp = insert_node(graph, make_node(ComputeTaskP()))

    # entry nodes getting data for P computes
    d_PB = insert_node(graph, make_node(DataTask(4)))
    d_PA = insert_node(graph, make_node(DataTask(4)))
    d_PBp = insert_node(graph, make_node(DataTask(4)))
    d_PAp = insert_node(graph, make_node(DataTask(4)))

    # now for all the edgese
    insert_edge(graph, make_edge(d_PB, PB))
    insert_edge(graph, make_edge(d_PA, PA))
    insert_edge(graph, make_edge(d_PBp, PBp))
    insert_edge(graph, make_edge(d_PAp, PAp))

    insert_edge(graph, make_edge(PB, d_PB_uB))
    insert_edge(graph, make_edge(PA, d_PA_uA))
    insert_edge(graph, make_edge(PBp, d_PBp_uBp))
    insert_edge(graph, make_edge(PAp, d_PAp_uAp))
    
    insert_edge(graph, make_edge(d_PB_uB, uB))
    insert_edge(graph, make_edge(d_PA_uA, uA))
    insert_edge(graph, make_edge(d_PBp_uBp, uBp))
    insert_edge(graph, make_edge(d_PAp_uAp, uAp))

    insert_edge(graph, make_edge(uB, d_uB_v0))
    insert_edge(graph, make_edge(uA, d_uA_v0))
    insert_edge(graph, make_edge(uBp, d_uBp_v1))
    insert_edge(graph, make_edge(uAp, d_uAp_v1))

    insert_edge(graph, make_edge(d_uB_v0, v0))
    insert_edge(graph, make_edge(d_uA_v0, v0))
    insert_edge(graph, make_edge(d_uBp_v1, v1))
    insert_edge(graph, make_edge(d_uAp_v1, v1))

    insert_edge(graph, make_edge(v0, d_v0_s0))
    insert_edge(graph, make_edge(v1, d_v1_s0))

    insert_edge(graph, make_edge(d_v0_s0, s0))
    insert_edge(graph, make_edge(d_v1_s0, s0))

    insert_edge(graph, make_edge(s0, d_exit))

    print(graph)

    println("Available optimizations from here:")
    println(generate_options(graph))

    # fuse some things
    fuse1 = node_fusion(graph, PB, d_PB_uB, uB)
    fuse2 = node_fusion(graph, PA, d_PA_uA, uA)
    fuse3 = node_fusion(graph, PBp, d_PBp_uBp, uBp)
    fuse4 = node_fusion(graph, PAp, d_PAp_uAp, uAp)

    # on this graph, nothing can be reduced
    println("Same graph after node fusion")
    print(graph)

    println("Available optimizations from here:")
    println(generate_options(graph))
end

main()