Merge pull request #3 from JuliaComputing/dg/ldiv

Feat: Amortize factorization cost for `\`

Author: DhairyaLGandhi

Project.toml

@@ -8,11 +8,18 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 
+[weakdeps]
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+
 [sources]
-SymbolicUtils = {url = "https://github.com/DhairyaLGandhi/SymbolicUtils.jl", rev = "dg/opt_api"}
+SymbolicUtils = {rev = "dg/opt_api", url = "https://github.com/DhairyaLGandhi/SymbolicUtils.jl"}
+
+[extensions]
+SCPLinearSolveExt = ["LinearSolve"]
 
 [compat]
 LinearAlgebra = "1.11.0"
+LinearSolve = "3.53.0"
 PreallocationTools = "0.4.34"
 SymbolicUtils = "4.1.0"
 julia = "1.11"

ext/SCPLinearSolveExt/SCPLinearSolveExt.jl

@@ -0,0 +1,59 @@
+module SCPLinearSolveExt
+
+using SymbolicUtils
+using SymbolicUtils.Code
+using LinearSolve
+using LinearAlgebra
+import SymbolicCompilerPasses: ldiv_transformation, SymbolicCompilerPasses, get_factorization, get_from_cache, FACTORIZATION_CACHE
+
+SymbolicCompilerPasses.LINEARSOLVE_LIB[] = true
+
+function linear_solve(A, B)
+    linsolve = get_factorization(A, B)
+    linsolve.b = B
+    sol = solve!(linsolve)
+    return sol.u
+end
+
+function get_factorization(A, B)
+    get!(FACTORIZATION_CACHE, A) do
+        prob = LinearSolve.LinearProblem(A, B)
+        linsolve = init(prob)
+    end
+end
+
+
+function ldiv_transformation(safe_matches, ::Val{true})
+    @info "Using LinearSolve.jl for in-place backsolve optimizations.
+    In order to opt-out of using LinearSolve, set SymbolicCompilerPasses.LINEARSOLVE_LIB[] = false." maxlog=Inf
+     # Build transformation
+    transformations = Dict{Int, Code.Assignment}()
+
+    rejected_matches = []
+    for match in safe_matches
+        A, B = match.A, match.B
+        result_var = Code.lhs(match.ldiv_candidate)
+        T = Code.vartype(B)
+
+        # Create: result = ldiv!(A, B)
+        if Code.symtype(B) <: AbstractVector
+            ldiv_call = Code.Term{T}(
+                linear_solve,
+                [A, B];
+                type=Code.symtype(B)
+            )
+        else
+            @warn "Skipping LinearSolve optimization for match as B is not a vector." maxlog=Inf
+            push!(rejected_matches, match)
+            continue
+        end
+
+        ldiv_assignment = Code.Assignment(result_var, ldiv_call)
+        transformations[match.ldiv_idx] = ldiv_assignment
+    end
+    fallback_transformations = ldiv_transformation(rejected_matches, Val(false))
+    merge(transformations, fallback_transformations)
+end
+
+
+end

src/SymbolicCompilerPasses.jl

@@ -3,9 +3,11 @@ module SymbolicCompilerPasses
 using LinearAlgebra
 using PreallocationTools
 using SymbolicUtils
-import SymbolicUtils: symtype, vartype, Sym, BasicSymbolic, Term, iscall, operation, arguments, maketerm, Const
+import SymbolicUtils: symtype, vartype, Sym, BasicSymbolic, Term, iscall, operation, arguments, maketerm, Const, shape, isterm, unwrap,
+                    is_function_symbolic, is_called_function_symbolic, getname, Unknown
 import SymbolicUtils.Code: Code, OptimizationRule, substitute_in_ir, apply_optimization_rules, AbstractMatched,
-    Assignment, CSEState, lhs, rhs, apply_substitution_map
+    Assignment, CSEState, lhs, rhs, apply_substitution_map, issym
+import SymbolicUtils: search_variables, search_variables!
 
 function bank(dic, key, value)
     if haskey(dic, key)
@@ -16,5 +18,7 @@ function bank(dic, key, value)
 end
 
 include("matmuladd.jl")
+include("ldiv_opt.jl")
+include("la_opt.jl")
 
 end # module SymbolicCompilerPasses

src/la_opt.jl

@@ -0,0 +1,87 @@
+function detect_f(expr, f, state)
+    args = search_variables(expr)
+    arg_counts = Dict(arg => count_occurrences(arg, expr) for arg in args)
+    triu_candidates_idx = findall(expr.pairs) do x
+        r = rhs(x)
+        iscall(r) || return false
+        op = operation(r)
+        op === f || return false
+
+        args = arguments(r)
+        length(args) == 1 || return false
+
+        arg = args[1]
+        symtype(arg) <: AbstractMatrix || return false
+
+        # The var"##cse#" variables must get counted at least twice (once in LHS and a second time in RHS)
+        # TODO: needs to detect provenance better; see expr7
+        get!(arg_counts, arg) do
+            count_occurrences(arg, expr)
+        end == 1 || return false
+
+        true
+    end
+
+    triu_candidates = expr.pairs[triu_candidates_idx]
+
+    matches = map(triu_candidates_idx, triu_candidates) do idx, candidate
+        A = arguments(rhs(candidate))[1]
+        GenericRuleMatched(A, candidate, idx)
+    end
+
+    f = filter(!isnothing, matches)
+    isempty(f) ? nothing : f
+end
+
+struct GenericRuleMatched{Ta, S} <: Code.AbstractMatched
+    A::Ta
+    candidate::S
+    idx::Int
+end
+
+transform_f(expr, f!, ::Nothing, state) = expr
+function transform_f(expr, f!, matches, state)
+
+    new_pairs = []
+    transformed_idxs = getproperty.(matches, :idx)
+    for (idx, pair) in enumerate(expr.pairs)
+        if idx in transformed_idxs
+            match = matches[findfirst(==(idx), transformed_idxs)]
+            A = match.A
+            candidate = match.candidate
+
+            push!(new_pairs, Assignment(lhs(candidate), term(f!, A,)))
+        else
+            push!(new_pairs, pair)
+        end
+        
+    end
+
+    Code.Let(new_pairs, expr.body, false)
+end
+
+
+function GenericRule(name, f, f!, priority)
+    OptimizationRule(
+        name,
+        (expr, state) -> detect_f(expr, f, state),
+        (expr, matches, state) -> transform_f(expr, f!, matches, state),
+        priority
+    )
+end
+
+const TRIU_RULE = GenericRule("triu", LinearAlgebra.triu, LinearAlgebra.triu!, 8)
+const TRIL_RULE = GenericRule("tril", LinearAlgebra.tril, LinearAlgebra.tril!, 8)
+const NORMALIZE_RULE = GenericRule("normalize", LinearAlgebra.normalize, LinearAlgebra.normalize!, 8)
+# const CONJ_RULE = GenericRule("conj", LinearAlgebra.conj, LinearAlgebra.conj!, 8)
+   
+function triu_opt(expr, state::CSEState)
+    # Try to apply optimization rules
+    optimized = apply_optimization_rules(expr, state, TRIU_RULE)
+    if optimized !== nothing
+        return optimized
+    end
+   
+    # If no optimization applied, return original expression
+    return expr
+end