Implement KaHyPar optimizer

Project.toml

@@ -10,7 +10,9 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 ImmutableArrays = "667c17eb-ab9b-4487-935f-1c621bb82497"
+KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880"
 Memoize = "c03570c3-d221-55d1-a50c-7939bbd78826"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/Optimizers/Hyper.jl

@@ -0,0 +1,109 @@
+using SparseArrays
+using KaHyPar
+
+@kwdef struct HyperGraphPartitioning <: Optimizer
+    metric::Function = removedsize
+    outer::Bool = false
+end
+
+
+function get_hypergraph(path)
+    # assume unique indices across all tensors
+    all_indices = mapreduce(head, ∪, path)
+
+    # Create incidence matrix
+    incidence_matrix = spzeros(Int64, length(path.args), length(all_indices))
+
+    for (i, tensor) in enumerate(path.args)
+        for idx in tensor.head
+            j = findfirst(==(idx), all_indices)
+            incidence_matrix[i, j] = 1
+        end
+    end
+
+    # Vertex weights (assuming equal weight for all tensors)
+    vertex_weights = ones(Int64, length(path.args))
+
+    # Hyperedge weights set to the size of the index
+    edge_weights = [size(path, idx) for idx in all_indices]
+
+    # Create hypergraph
+    hypergraph = KaHyPar.HyperGraph(incidence_matrix, vertex_weights, edge_weights)
+
+    return hypergraph
+end
+
+function partition_hypergraph(hypergraph::KaHyPar.HyperGraph; parts=2, parts_decay=0.5, random_strength=0.01)
+    # Create a context for KaHyPar
+    context = KaHyPar.kahypar_context_new()
+
+    # Load default configuration
+    config_file_path = KaHyPar.default_configuration
+    KaHyPar.kahypar_configure_context_from_file(context, config_file_path)
+
+    # Calculate the relative subgraph size
+    subsize = hypergraph.n_vertices
+    N = hypergraph.n_vertices  # Assuming N is the total number of vertices in the entire graph
+    s = subsize / N
+
+    # Determine the number of partitions based on the relative subgraph size
+    kparts = max(Int(s^parts_decay * parts), 2)
+
+    # Perform the partitioning
+    partitioning_result = KaHyPar.partition(hypergraph, kparts; configuration=config_file_path)
+
+    # Clean up the context
+    KaHyPar.kahypar_context_free(context)
+
+    return partitioning_result
+end
+
+function recursive_partition(expr::EinExpr; parts=2, parts_decay=0.5, random_strength=0.01, cutoff=2, max_iterations=10)
+    # Convert the EinExpr to a hypergraph
+    hypergraph = get_hypergraph(expr)
+    println("Hypergraph size: ", hypergraph.n_vertices)
+
+    # Base case: if the hypergraph is small enough, we stop the recursion
+    if hypergraph.n_vertices <= cutoff
+        println("Base case")
+        return expr
+    end
+
+    if max_iterations == 0
+        println("Max iterations reached")
+        return expr
+    end
+
+    # Partition the hypergraph
+    partitioning_result = partition_hypergraph(hypergraph, parts=parts, parts_decay=parts_decay, random_strength=random_strength)
+
+    # Get unique partitions
+    unique_partitions = unique(partitioning_result)
+
+    # Contract nodes based on the partitioning
+    new_exprs = [partition_to_einexpr(part_id, partitioning_result, expr) for part_id in unique_partitions]
+    combined_expr = sum(new_exprs)
+
+    # Recursively call the function on the combined EinExpr
+    return recursive_partition(combined_expr, parts=parts, parts_decay=parts_decay, random_strength=random_strength, cutoff=cutoff, max_iterations=max_iterations-1)
+end
+
+function partition_to_einexpr(partition_id::Int, partition_result::Vector{Int}, original_path::EinExpr)
+    # Identify the tensor indices that belong to the specific partition_id
+    tensor_indices_in_partition = [idx for (idx, part) in enumerate(partition_result) if part == partition_id]
+
+    # Extract the corresponding tensors from the original EinExpr
+    tensors_in_partition = [original_path.args[idx] for idx in tensor_indices_in_partition]
+
+    # Sum (contract) these tensors to create a new tensor
+    new_tensor = sum(tensors_in_partition)
+
+    # @show tensors_in_partition
+    # @show tensor_indices_in_partition
+    # @show new_tensor
+
+    # Create a new EinExpr with the new tensor's indices in the head, the original tensors in the args, and the size dictionary
+    new_einexpr = EinExpr(new_tensor.head, tensors_in_partition)
+
+    return new_einexpr
+end