Keep Galley Plans Per Approximate Sparsity Pattern

benchmark/benchmarks.jl

@@ -94,13 +94,12 @@ for (scheduler_name, scheduler) in [
         end
 
         let
-            A = Tensor(Dense(SparseList(Element(0.0))), fsprand(1, 1, 1))
-            x = rand(1)
             SUITE["high-level"]["einsum_spmv_call_overhead"][scheduler_name] = @benchmarkable(
                 begin
-                    A, x = ($A, $x)
                     @einsum y[i] += A[i, j] * x[j]
                 end,
+                setup = (A = Tensor(Dense(SparseList(Element(0.0))), fsprand(1, 1, 1)); x = rand(1)),
+                evals = 1
             )
         end
 

src/Finch.jl

@@ -280,7 +280,7 @@ include("interface/einsum.jl")
 include("Galley/Galley.jl")
 using .Galley
 
-export galley_scheduler
+export galley_scheduler, GalleyOptimizer, GalleyExecutorCode, GalleyExecutor
 
 @deprecate default fill_value
 @deprecate redefault! set_fill_value!

src/FinchLogic/nodes.jl

@@ -352,7 +352,7 @@ function Base.:(==)(a::LogicNode, b::LogicNode)
     elseif a.kind === immediate
         return b.kind === immediate && a.val === b.val
     elseif a.kind === deferred
-        return b.kind === deferred && a.val === b.val && a.type === b.type
+        return b.kind === deferred && a.ex === b.ex && a.type === b.type
     elseif a.kind === field
         return b.kind === field && a.name == b.name
     elseif a.kind === alias
@@ -370,7 +370,7 @@ function Base.hash(a::LogicNode, h::UInt)
     elseif istree(a)
         return hash(a.kind, hash(a.children, h))
     elseif a.kind === deferred
-        return hash(a.kind, hash(a.val, hash(a.type, h)))
+        return hash(a.kind, hash(a.ex, hash(a.type, h)))
     else
         error("unimplemented")
     end

src/Galley/FinchCompat/executor.jl

@@ -41,16 +41,100 @@
     julia_prgm
 end
 
-function Finch.set_options(ctx::GalleyOptimizer; estimator=DCStats)
+function Finch.set_options(ctx::GalleyOptimizer; estimator=DCStats, verbose=false)
     ctx.estimator=estimator
+    ctx.verbose=verbose
     return ctx
 end
 
+"""
+    get_stats_dict(ctx::GalleyOptimizer, prgm)
+
+Returns a dictionary mapping the location of input tensors in the program to their statistics objects.
+"""
+function get_stats_dict(ctx::GalleyOptimizer, prgm)
+    deferred_prgm = Finch.defer_tables(:prgm, prgm)
+    expr_stats_dict = Dict()
+    for node in PostOrderDFS(deferred_prgm)
+        if node.kind == table
+            expr_stats_dict[node.tns.ex] = ctx.estimator(node.tns.imm, [i.name for i in node.idxs])
+        end
+    end
+    return expr_stats_dict
+end
+
+"""
+    GalleyExecutor(ctx::GalleyOptimizer, tag=:global, verbose=false)
+
+Executes a logic program by compiling it with the given compiler `ctx`. Compiled
+codes are cached for each program structure. If the 'tag' argument is ':global', it maintains a set of plans 
+for inputs with different sparsity structures. In this case, it first checks the cache for a plan that
+was compiled for similar inputs and only compiles if it doesn't find one. If the `tag` argument is anything else,
+it will only compile once for that tag and will skip this search process.
+"""
+@kwdef struct GalleyExecutor
+    ctx::GalleyOptimizer
+    tag
+    verbose
+end
+
+Base.:(==)(a::GalleyExecutor, b::GalleyExecutor) = a.ctx == b.ctx && a.verbose == b.verbose
+Base.hash(a::GalleyExecutor, h::UInt) = hash(GalleyExecutor, hash(a.ctx, hash(a.verbose, h)))
+
+GalleyExecutor(ctx::GalleyOptimizer; tag = :global, verbose = false) = GalleyExecutor(ctx, tag, verbose)
+function Finch.set_options(ctx::GalleyExecutor; tag = ctx.tag, verbose = ctx.verbose, kwargs...)
+    GalleyExecutor(Finch.set_options(ctx.ctx; verbose=verbose, kwargs...), tag, verbose)
+end
+
+galley_codes = Dict()
+function (ctx::GalleyExecutor)(prgm)
+    (f, code) = if ctx.tag == :global
+        cur_stats_dict = get_stats_dict(ctx.ctx, prgm)
+        stats_list = get!(galley_codes, (ctx.ctx, ctx.tag, Finch.get_structure(prgm)), [])
+        valid_match = nothing
+        for (stats_dict, f_code) in stats_list
+            if all(issimilar(cur_stats, stats_dict[cur_expr], 4) for (cur_expr, cur_stats) in cur_stats_dict)
+                valid_match = f_code
+            end
+        end
+        if isnothing(valid_match)
+            thunk = Finch.logic_executor_code(ctx.ctx, prgm)
+            valid_match = (eval(thunk), thunk)
+            push!(stats_list, (cur_stats_dict, valid_match))
+        end
+        valid_match
+    else
+        get!(galley_codes, (ctx.ctx, ctx.tag, Finch.get_structure(prgm))) do
+            thunk = Finch.logic_executor_code(ctx.ctx, prgm)
+            (eval(thunk), thunk)
+        end
+    end
+    if ctx.verbose
+        println("Executing:")
+        display(code)
+    end
+    return Base.invokelatest(f, prgm)
+end
+
+"""
+    GalleyExecutorCode(ctx)
+
+Return the code that would normally be used by the GalleyExecutor to run a program.
+"""
+struct GalleyExecutorCode
+    ctx
+end
+
+function (ctx::GalleyExecutorCode)(prgm)
+    return Finch.logic_executor_code(ctx.ctx, prgm)
+end
+
 """
     galley_scheduler(verbose = false, estimator=DCStats)
 
 The galley scheduler uses the sparsity patterns of the inputs to optimize the computation.
 The first set of inputs given to galley is used to optimize, and the `estimator` is used to
 estimate the sparsity of intermediate computations during optimization.
 """
-galley_scheduler(; verbose = false, estimator=DCStats) = Finch.LogicExecutor(GalleyOptimizer(verbose=verbose, estimator=estimator); verbose=verbose)
+galley_scheduler(;verbose=false) = GalleyExecutor(GalleyOptimizer(;verbose=false); verbose=false)
+

src/Galley/Galley.jl

@@ -25,7 +25,7 @@ export PlanNode, Value, Index, Alias, Input, MapJoin, Aggregate, Materialize, Qu
 export Scalar, Σ, Mat, Agg
 export DCStats, NaiveStats, TensorDef, DC, insert_statistics
 export naive, greedy, pruned, exact
-export GalleyOptimizer, galley_scheduler
+export GalleyOptimizer, GalleyExecutor, GalleyExecutorCode, galley_scheduler
 
 IndexExpr = Symbol
 TensorId = Symbol