From e580c0d455692dce233a6b6dafa9ac7f6b57d14f Mon Sep 17 00:00:00 2001 From: Dhairya Gandhi Date: Wed, 4 Feb 2026 23:24:49 +0000 Subject: [PATCH 1/8] chore: use jacobian to extract stencil --- src/jacobian_vectorize.jl | 238 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 238 insertions(+) create mode 100644 src/jacobian_vectorize.jl diff --git a/src/jacobian_vectorize.jl b/src/jacobian_vectorize.jl new file mode 100644 index 0000000..14bbacf --- /dev/null +++ b/src/jacobian_vectorize.jl @@ -0,0 +1,238 @@ +function find_definition(target, expr::Code.Let) + idx = findfirst(expr.pairs) do p + Code.lhs(p) === target + end + if !isnothing(idx) + return (; idx, def = expr.pairs[idx]) + else + return (; idx = nothing, def = nothing) + end +end + +construct_template(ass::Code.Assignment) = construct_template(ass.rhs) +function construct_template(expr) + seq = [] + construct_template!(seq, expr) +end + +function construct_template!(seq, expr) + if iscall(expr) + push!(seq, operation(expr)) + args = arguments(expr) + for a in args + construct_template!(seq, a) + end + seq + else + seq + end + seq +end + +function basic_groupby(temps, expr) + groups = Dict() + for (i, t) in enumerate(temps) + SC.bank(groups, t, (i, expr.pairs[i])) + end + groups +end + +function code_flow(expr) + g = basic_groupby(construct_template.(expr.pairs), expr) + return g +end + +function hold_eq(eq) + if SU.isconst(eq) # || SU.issym(eq) + return true + else + if SU.iscall(eq) && operation(eq) == getindex + return true + else + return false + end + end + return false + +end + +function expand_eq(du_expr, expr) + args = arguments(du_expr) + defs = [find_definition(arg, expr).def.rhs for arg in args] + sub_map = Dict(arg => def for (arg, def) in zip(args, defs)) + + new_eq = Code.substitute_in_ir(du_expr, sub_map) + + new_args = reduce(vcat, arguments.(defs)) + expand_eq(new_eq, new_args, expr) +end +function expand_eq(du_expr, args, expr) + if all(hold_eq(a) for a in args) + return du_expr + end + + args_to_sub = filter(a -> !hold_eq(a), args) + defs = [find_definition(arg, expr).def.rhs for arg in args_to_sub] + sub_map = Dict(arg => def for (arg, def) in zip(args_to_sub, defs)) + new_eq = Code.substitute_in_ir(du_expr, sub_map) + new_args = reduce(vcat, hold_eq(def) ? def : arguments(def) for def in defs) + expand_eq(new_eq, new_args, expr) +end + +function get_expanded_args(expr) + args = [] + get_expanded_args!(args, expr) + return args +end +function get_expanded_args!(args, expr) + if SU.iscall(expr) + if operation(expr) == getindex + push!(args, expr) + return + end + arg_list = arguments(expr) + for a in arg_list + get_expanded_args!(args, a) + end + end +end + +function replace_idx_vars_with_vals(eq, expr) + if SU.iscall(eq) + if operation(eq) == getindex + idx_vars = arguments(eq)[2:end] + if all(SU.isconst, idx_vars) + return eq + end + idx_vals = map(idx_vars) do iv + def = find_definition(iv, expr).def + def_rhs = Code.rhs(def) + end + new_eq = Code.substitute_in_ir(eq, Dict(zip(idx_vars, idx_vals))) + return new_eq + else + arg_list = arguments(eq) + new_args = map(arg_list) do arg + replace_idx_vars_with_vals(arg, expr) + end + return operation(eq)(new_args...) + end + end + eq +end + +function contains_variable(expr, var) + if SU.issym(expr) + return expr === var + elseif SU.isconst(expr) + return false, nothing + elseif SU.iscall(expr) + if operation(expr) === var + return true, var + elseif operation(expr) === getindex + arr, _ = arguments(expr) + c = arr === var + if c + return true, arr + else + return false, nothing + end + else + arg_list = arguments(expr) + for a in arg_list + if contains_variable(a, var)[1] + arg_list_filtered = filter(x -> !(x === a), arg_list) + if all(SU.isconst, arg_list_filtered) + return true, expr + else + return true, a + end + end + end + end + end + false, nothing +end + +function separate_linear_nonlinear(eq, U) + has_U = true + nls = [] + while has_U + has_U, ex = contains_variable(eq, U) + if isnothing(ex) + break + end + eq = eq - ex + push!(nls, ex) + # break + end + + return eq, isempty(nls) ? Num(0) : +(nls...) +end + +function separate_linear_nonlinear(J::AbstractArray, U) + L = zeros(Symbolics.Num, size(J)) + NL = zeros(Symbolics.Num, size(J)) + + for i = 1:size(J, 1), j = 1:size(J, 2) + SU.isconst(J[i, j]) && continue + lin, nlin = separate_linear_nonlinear(J[i, j], U) + L[i, j] = lin + NL[i, j] = nlin + end + + return L, NL +end + +function detect_jacobian_trace(f, DU, U, P, t) + traced = f(DU, U, P, t) + # traced_let = Code.cse(traced) + # eqs = candidate_templates(traced_let) + + # # J[i,j] = ∂(du[i])/∂(u[j]) + # expanded_eqs = expand_eq.(eqs, Ref(traced_let)) + + # full_expanded_eqs = replace_idx_vars_with_vals.(expanded_eqs, Ref(traced_let)) + # J = Symbolics.jacobian(vec(full_expanded_eqs), vec(collect(U))) + J = Symbolics.jacobian(vec(traced), vec(collect(U))) + + L, NL = separate_linear_nonlinear(J, U) + VectorizedJacobianMatched(J, L, NL, U, size(U)) +end + +struct VectorizedJacobianMatched{TJ, TL, TNL, TU, S} <: AbsctractMatched + J::TJ + L::TL + NL::TNL + U::TU + shape::S +end + +function detect_jacobian(expr::Code.Let, state) + eqs = candidate_templates(expr) + # J[i,j] = ∂(du[i])/∂(u[j]) + expanded_eqs = expand_eq.(eqs, Ref(expr)) + full_expanded_eqs = replace_idx_vars_with_vals.(expanded_eqs, Ref(expr)) + J = Symbolics.jacobian(vec(full_expanded_eqs), vec(collect(U))) + L, NL = separate_linear_nonlinear(J, U) + VectorizedJacobianMatched(J, L, NL, U) +end + +function transform_jacobian(expr::Code.Let, match_data::VectorizedJacobianMatched, state) + # L * U + NL(U) + + linear_terms = match_data.L * vec(match_data.U) + nonlinear_terms = diag(match_data.NL) + val = linear_terms + nonlinear_terms + T = vartype(match_data.L) + res = Term{T}(reshape, [val, match_data.shape]; type = symtype(match_data.U)) + + Code.Let([], res, expr.let_block) +end + +const JACOBIAN_VECTORIZE_RULE = OptimizationRule( + "JacobianVectorize", + detect_jacobian, + transform_jacobian, + 10 # Medium priority +) From 187abac694def04bee01a539636dc22a51faa06a Mon Sep 17 00:00:00 2001 From: Dhairya Gandhi Date: Thu, 5 Feb 2026 12:40:56 +0000 Subject: [PATCH 2/8] chore: handle nonlinear cases better --- Project.toml | 4 ++++ src/SymbolicCompilerPasses.jl | 4 +++- src/jacobian_vectorize.jl | 38 ++++++++++++++++++++++++++++------- 3 files changed, 38 insertions(+), 8 deletions(-) diff --git a/Project.toml b/Project.toml index 04a7197..7f13e8c 100644 --- a/Project.toml +++ b/Project.toml @@ -7,8 +7,10 @@ version = "0.1.0" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b" +Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" [weakdeps] LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" @@ -27,8 +29,10 @@ LinearAlgebra = "1.11.0, 1.10" LinearSolve = "3.53.0" PreallocationTools = "0.4.34" Rotations = "1.7.1" +SparseArrays = "1.11.0" StaticArrays = "1.9.15" SymbolicUtils = "4.1.0" +Symbolics = "7.2.0" julia = "1.11" [extras] diff --git a/src/SymbolicCompilerPasses.jl b/src/SymbolicCompilerPasses.jl index 3b9c3bc..42dc48c 100644 --- a/src/SymbolicCompilerPasses.jl +++ b/src/SymbolicCompilerPasses.jl @@ -9,7 +9,9 @@ import SymbolicUtils.Code: Code, OptimizationRule, substitute_in_ir, apply_optim Assignment, CSEState, lhs, rhs, apply_substitution_map, IfElse, issym, isterm, toexpr, _is_array_of_symbolics, MakeArray, shape import SymbolicUtils: search_variables, search_variables! +import SymbolicUtils as SU using StaticArrays +using Symbolics using DataStructures @@ -28,7 +30,7 @@ include("ldiv_opt.jl") include("la_opt.jl") include("mb_opt.jl") -include("scalar_to_vec_opt.jl") +include("jacobian_vectorize.jl") # function apply_optimizations(ir, state, rules) # for rule in sort(rules, by = x -> x.priority) diff --git a/src/jacobian_vectorize.jl b/src/jacobian_vectorize.jl index 14bbacf..5b95065 100644 --- a/src/jacobian_vectorize.jl +++ b/src/jacobian_vectorize.jl @@ -42,9 +42,24 @@ function code_flow(expr) return g end +function candidate_templates(expr) + dus = find_definition.(expr.body, Ref(expr)) + map(dus) do du + du.def.rhs + end +end + function hold_eq(eq) if SU.isconst(eq) # || SU.issym(eq) return true + elseif SU.issym(eq) + # TODO: some variables seem to get used but don't have a definition + # How do you handle those? + if eq === Main.t || eq === Main.P + return true + else + return false + end else if SU.iscall(eq) && operation(eq) == getindex return true @@ -72,6 +87,14 @@ function expand_eq(du_expr, args, expr) end args_to_sub = filter(a -> !hold_eq(a), args) + # for a in args_to_sub + # fd = @show find_definition(a, expr) + # if isnothing(fd.idx) + # @warn a + # @show SU.issym(a) + # @warn du_expr + # end + # end defs = [find_definition(arg, expr).def.rhs for arg in args_to_sub] sub_map = Dict(arg => def for (arg, def) in zip(args_to_sub, defs)) new_eq = Code.substitute_in_ir(du_expr, sub_map) @@ -101,6 +124,7 @@ function replace_idx_vars_with_vals(eq, expr) if SU.iscall(eq) if operation(eq) == getindex idx_vars = arguments(eq)[2:end] + idx_vars = filter(!SU.isconst, idx_vars) if all(SU.isconst, idx_vars) return eq end @@ -200,7 +224,7 @@ function detect_jacobian_trace(f, DU, U, P, t) VectorizedJacobianMatched(J, L, NL, U, size(U)) end -struct VectorizedJacobianMatched{TJ, TL, TNL, TU, S} <: AbsctractMatched +struct VectorizedJacobianMatched{TJ, TL, TNL, TU, S} <: AbstractMatched J::TJ L::TL NL::TNL @@ -213,18 +237,18 @@ function detect_jacobian(expr::Code.Let, state) # J[i,j] = ∂(du[i])/∂(u[j]) expanded_eqs = expand_eq.(eqs, Ref(expr)) full_expanded_eqs = replace_idx_vars_with_vals.(expanded_eqs, Ref(expr)) - J = Symbolics.jacobian(vec(full_expanded_eqs), vec(collect(U))) - L, NL = separate_linear_nonlinear(J, U) - VectorizedJacobianMatched(J, L, NL, U) + J = Symbolics.jacobian(vec(full_expanded_eqs), vec(collect(Main.U))) + L, NL = separate_linear_nonlinear(J, Main.U) + VectorizedJacobianMatched(J, L, NL, Main.U, size(Main.U)) end function transform_jacobian(expr::Code.Let, match_data::VectorizedJacobianMatched, state) - # L * U + NL(U) + # DU = L * U + NL(U) linear_terms = match_data.L * vec(match_data.U) - nonlinear_terms = diag(match_data.NL) + nonlinear_terms = dropdims(sum(match_data.NL, dims = 2), dims = 2) val = linear_terms + nonlinear_terms - T = vartype(match_data.L) + T = SymReal res = Term{T}(reshape, [val, match_data.shape]; type = symtype(match_data.U)) Code.Let([], res, expr.let_block) From 83e26e0945e9deebf45cb07f98283a42993ee4e8 Mon Sep 17 00:00:00 2001 From: Dhairya Gandhi Date: Thu, 5 Feb 2026 14:17:58 +0000 Subject: [PATCH 3/8] chore: handle inputs better in let block --- src/jacobian_vectorize.jl | 34 ++++++++++++++++++++++++---------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/src/jacobian_vectorize.jl b/src/jacobian_vectorize.jl index 5b95065..b59f735 100644 --- a/src/jacobian_vectorize.jl +++ b/src/jacobian_vectorize.jl @@ -49,13 +49,13 @@ function candidate_templates(expr) end end -function hold_eq(eq) +function hold_eq(eq, inputs) if SU.isconst(eq) # || SU.issym(eq) return true elseif SU.issym(eq) # TODO: some variables seem to get used but don't have a definition # How do you handle those? - if eq === Main.t || eq === Main.P + if eq in inputs return true else return false @@ -71,7 +71,7 @@ function hold_eq(eq) end -function expand_eq(du_expr, expr) +function expand_eq(du_expr, inputs, expr) args = arguments(du_expr) defs = [find_definition(arg, expr).def.rhs for arg in args] sub_map = Dict(arg => def for (arg, def) in zip(args, defs)) @@ -79,14 +79,14 @@ function expand_eq(du_expr, expr) new_eq = Code.substitute_in_ir(du_expr, sub_map) new_args = reduce(vcat, arguments.(defs)) - expand_eq(new_eq, new_args, expr) + expand_eq(new_eq, new_args, inputs, expr) end -function expand_eq(du_expr, args, expr) - if all(hold_eq(a) for a in args) +function expand_eq(du_expr, args, inputs, expr) + if all(hold_eq(a, inputs) for a in args) return du_expr end - args_to_sub = filter(a -> !hold_eq(a), args) + args_to_sub = filter(a -> !hold_eq(a, inputs), args) # for a in args_to_sub # fd = @show find_definition(a, expr) # if isnothing(fd.idx) @@ -98,8 +98,8 @@ function expand_eq(du_expr, args, expr) defs = [find_definition(arg, expr).def.rhs for arg in args_to_sub] sub_map = Dict(arg => def for (arg, def) in zip(args_to_sub, defs)) new_eq = Code.substitute_in_ir(du_expr, sub_map) - new_args = reduce(vcat, hold_eq(def) ? def : arguments(def) for def in defs) - expand_eq(new_eq, new_args, expr) + new_args = reduce(vcat, hold_eq(def, inputs) ? def : arguments(def) for def in defs) + expand_eq(new_eq, new_args, inputs, expr) end function get_expanded_args(expr) @@ -232,10 +232,24 @@ struct VectorizedJacobianMatched{TJ, TL, TNL, TU, S} <: AbstractMatched shape::S end +function improve_inputs(inputs) + improved = Set() + for inp in inputs + if SU.iscall(inp) && operation(inp) == getindex + arr, _ = arguments(inp) + push!(improved, arr) + else + push!(improved, inp) + end + end + improved +end + function detect_jacobian(expr::Code.Let, state) + inputs = improve_inputs(search_variables(expr)) eqs = candidate_templates(expr) # J[i,j] = ∂(du[i])/∂(u[j]) - expanded_eqs = expand_eq.(eqs, Ref(expr)) + expanded_eqs = expand_eq.(eqs, Ref(inputs), Ref(expr)) full_expanded_eqs = replace_idx_vars_with_vals.(expanded_eqs, Ref(expr)) J = Symbolics.jacobian(vec(full_expanded_eqs), vec(collect(Main.U))) L, NL = separate_linear_nonlinear(J, Main.U) From 630ce61c5a799cf03afd6b577aa9e4fc19e772f4 Mon Sep 17 00:00:00 2001 From: Dhairya Gandhi Date: Fri, 6 Feb 2026 01:08:53 +0000 Subject: [PATCH 4/8] chore: working bruss + nonlinearities --- src/jacobian_vectorize.jl | 198 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 188 insertions(+), 10 deletions(-) diff --git a/src/jacobian_vectorize.jl b/src/jacobian_vectorize.jl index b59f735..aaabb69 100644 --- a/src/jacobian_vectorize.jl +++ b/src/jacobian_vectorize.jl @@ -154,11 +154,19 @@ function contains_variable(expr, var) if operation(expr) === var return true, var elseif operation(expr) === getindex - arr, _ = arguments(expr) + arr, expr_idxs = arguments(expr) c = arr === var if c return true, arr else + if issym(var) + return false, nothing + elseif operation(var) === getindex + arr_var, var_idxs = arguments(var) + if arr === arr_var && isempty(setdiff(expr_idxs, var_idxs)) + return true, expr + end + end return false, nothing end else @@ -188,15 +196,98 @@ function separate_linear_nonlinear(eq, U) end eq = eq - ex push!(nls, ex) - # break end return eq, isempty(nls) ? Num(0) : +(nls...) end +""" +constant here means constant in Us; could depend on other variables +should not be recursive +""" +function find_constant_terms!(const_terms, expr, U) + # @show expr + # if U isa AbstractArray + # U = first(U) + # end + # has_var, _ = contains_variable(expr, U) + # # @warn has_var + # if !has_var + # push!(const_terms, expr) + # return const_terms + # end + # args = arguments(simplify(expr)) + # @show args + + # for a in args + # has_variable, t = contains_variable(a, U) + # if !has_variable + # push!(const_terms, a) + # end + # end + zero_dict = Dict(collect(U) .=> 0) + const_expr = simplify(Symbolics.substitute(expr, zero_dict)) + push!(const_terms, const_expr) + const_terms +end + +function find_linear_terms!(linear_terms, expr, U) + args = collect(Set(get_expanded_args(expr))) + args_filtered = filter(args) do a + has_var, t = contains_variable(a, U) + if has_var + return true + else + return false + end + end + # TODO: why is args_filtered empty + # TODO: handle vector of U + args_filtered = U + zero_dict = Dict(collect(U) .=> 0) + gs = Symbolics.gradient(expr, args_filtered) + for (a, g) in zip(args_filtered, gs) + # Trying to catch linear terms here is very error prone + # substitute all variables with 0 should yield all constants + linear_terms[a] = simplify(Symbolics.substitute(g, zero_dict)) + # linear_terms[cartesian_to_linear(a)] = simplify(Symbolics.substitute(g, zero_dict)) + end + linear_terms +end + +function find_nonlinear_terms!(nonlinear_terms, expr, U) + # Needs to take care of trig fns + # Needs to take care of coupled terms u[i] * u[j] + # Needs to take care of squuared/ cubed/ polynomial terms + + # @show expr + + push!(nonlinear_terms, simplify(expr)) + nonlinear_terms +end + +function separate_linear_nonlinear_constant(eq, U) + const_terms = [] + linear_terms = Dict() + nonlinear_terms = [] + find_constant_terms!(const_terms, eq, U) + const_sum = isempty(const_terms) ? Num(0) : +(const_terms...) + eq = simplify(eq - const_sum) + + find_linear_terms!(linear_terms, eq, U) + linear_sum = isempty(linear_terms) ? Num(0) : +([key * val for (key, val) in pairs(linear_terms)]...) + eq = simplify(eq - linear_sum) + + find_nonlinear_terms!(nonlinear_terms, eq, U) + nonlinear_sum = isempty(nonlinear_terms) ? Num(0) : +(nonlinear_terms...) + + return const_terms, linear_terms, nonlinear_sum +end + function separate_linear_nonlinear(J::AbstractArray, U) L = zeros(Symbolics.Num, size(J)) NL = zeros(Symbolics.Num, size(J)) + C = zeros(Symbolics.Num, size(J)) for i = 1:size(J, 1), j = 1:size(J, 2) SU.isconst(J[i, j]) && continue @@ -219,15 +310,16 @@ function detect_jacobian_trace(f, DU, U, P, t) # full_expanded_eqs = replace_idx_vars_with_vals.(expanded_eqs, Ref(traced_let)) # J = Symbolics.jacobian(vec(full_expanded_eqs), vec(collect(U))) J = Symbolics.jacobian(vec(traced), vec(collect(U))) + C = J L, NL = separate_linear_nonlinear(J, U) - VectorizedJacobianMatched(J, L, NL, U, size(U)) + VectorizedJacobianMatched(L, NL, C, U, size(U)) end -struct VectorizedJacobianMatched{TJ, TL, TNL, TU, S} <: AbstractMatched - J::TJ +struct VectorizedJacobianMatched{TL, TNL, TC, TU, S} <: AbstractMatched L::TL NL::TNL + C::C U::TU shape::S end @@ -245,23 +337,109 @@ function improve_inputs(inputs) improved end +function construct_jacobian(eqs, vars) + J = Symbolics.jacobian(eqs, vars) + return J +end + +function cartesian_to_linear(x) + if !(iscall(x) && operation(x) == getindex) + error("Input must be a getindex call") + end + + arr, idx... = arguments(x) + cartesian_to_linear((idx...,), size(arr)) +end + +function cartesian_to_linear(indices::Tuple, dims::Tuple) + @assert length(indices) == length(dims) "Indices and dimensions must have same length" + + linear_idx = indices[1] + stride = 1 + + for d in 2:length(indices) + stride *= dims[d-1] + linear_idx += (indices[d] - 1) * stride + end + + return linear_idx +end + +function linear_to_cartesian(idx, x) + x[linear_to_cartesian(idx, size(x))...] +end +function linear_to_cartesian(idx, dims::Tuple) + idx = Symbolics.value(idx) + + n_dims = length(dims) + indices = Vector{Int}(undef, n_dims) + + remaining = idx - 1 + + for d in 1:n_dims + indices[d] = mod(remaining, dims[d]) + 1 + remaining = div(remaining, dims[d]) + end + + return Tuple(indices) +end + +function extract_coeffs(eq, vars, expr) + eq = simplify(eq) + vars = Set(collect(vars)) + eq_args = get_expanded_args(simplify(eq)) + eq_args_filtered = filter(eq_args) do a + if issym(a) || SU.isconst(a) + false + elseif SU.iscall(a) && operation(a) == getindex + return a in vars + end + false + end + eq_args_filtered = collect(Set(eq_args_filtered)) + js = cartesian_to_linear.(eq_args_filtered) + + const_terms, linear_terms, nonlinear_sum = separate_linear_nonlinear_constant(eq, eq_args_filtered) + js, const_terms, linear_terms, nonlinear_sum +end + +""" +Construct sparse jacobian using expression expansion and not full dense jacobian +""" +function construct_jacobian(eqs, vars, expr) + # J[i,j] = ∂(du[i])/∂(u[j]) + A = zeros(Symbolics.Num, length(eqs), length(vars)) + NL = zeros(Symbolics.Num, length(eqs)) + C = zeros(Symbolics.Num, length(eqs)) + + for i in 1:length(eqs) + js, const_terms, linear_terms, nonlinear_sum = extract_coeffs(eqs[i], vars, expr) + for (k, j) in enumerate(js) + A[i, Symbolics.value(j)] = linear_terms[linear_to_cartesian(j, vars)] + end + NL[i] = nonlinear_sum + C[i] = isempty(const_terms) ? Num(0) : +(const_terms...) + end + return A, NL, C +end + function detect_jacobian(expr::Code.Let, state) inputs = improve_inputs(search_variables(expr)) eqs = candidate_templates(expr) # J[i,j] = ∂(du[i])/∂(u[j]) expanded_eqs = expand_eq.(eqs, Ref(inputs), Ref(expr)) full_expanded_eqs = replace_idx_vars_with_vals.(expanded_eqs, Ref(expr)) - J = Symbolics.jacobian(vec(full_expanded_eqs), vec(collect(Main.U))) - L, NL = separate_linear_nonlinear(J, Main.U) - VectorizedJacobianMatched(J, L, NL, Main.U, size(Main.U)) + L, NL, C = construct_jacobian(vec(full_expanded_eqs), Main.U, expr) + VectorizedJacobianMatched(L, NL, C, Main.U, size(Main.U)) end function transform_jacobian(expr::Code.Let, match_data::VectorizedJacobianMatched, state) # DU = L * U + NL(U) linear_terms = match_data.L * vec(match_data.U) - nonlinear_terms = dropdims(sum(match_data.NL, dims = 2), dims = 2) - val = linear_terms + nonlinear_terms + nonlinear_terms = match_data.NL + const_terms = match_data.C + val = linear_terms + nonlinear_terms + const_terms T = SymReal res = Term{T}(reshape, [val, match_data.shape]; type = symtype(match_data.U)) From 1235f1a1a060ddf56274d73f80f7fbc569a653b6 Mon Sep 17 00:00:00 2001 From: Dhairya Gandhi Date: Fri, 6 Feb 2026 01:12:07 +0000 Subject: [PATCH 5/8] docs: docstring --- src/jacobian_vectorize.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/jacobian_vectorize.jl b/src/jacobian_vectorize.jl index aaabb69..9985e2d 100644 --- a/src/jacobian_vectorize.jl +++ b/src/jacobian_vectorize.jl @@ -324,6 +324,10 @@ struct VectorizedJacobianMatched{TL, TNL, TC, TU, S} <: AbstractMatched shape::S end +""" +search_variables adds things like U[i, j, k] to the variables but not U +this only keeps the array variables +""" function improve_inputs(inputs) improved = Set() for inp in inputs @@ -434,7 +438,7 @@ function detect_jacobian(expr::Code.Let, state) end function transform_jacobian(expr::Code.Let, match_data::VectorizedJacobianMatched, state) - # DU = L * U + NL(U) + # DU = L * U + NL(U) + C linear_terms = match_data.L * vec(match_data.U) nonlinear_terms = match_data.NL From 26af0768c101595daefffd98e0897beeb474ab9b Mon Sep 17 00:00:00 2001 From: Dhairya Gandhi Date: Fri, 6 Feb 2026 15:13:28 +0000 Subject: [PATCH 6/8] test: add brusselator vectorized version --- test/runtests.jl | 2 ++ test/scalar_to_vec.jl | 58 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+) create mode 100644 test/scalar_to_vec.jl diff --git a/test/runtests.jl b/test/runtests.jl index c13257a..2ee8e40 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,4 +6,6 @@ using Pkg, Test, SafeTestsets @safetestset "Literal Small Array Allocation" begin include("array_literal.jl") end @safetestset "Ldiv Factorization Optimization" begin include("ldiv_opt.jl") end @safetestset "Apply Multiple Rules" begin include("multiple.jl") end + + @safetestset "Vectorize Scalar Equations" begin include("scalar_to_vec.jl") end end \ No newline at end of file diff --git a/test/scalar_to_vec.jl b/test/scalar_to_vec.jl new file mode 100644 index 0000000..9bf3a95 --- /dev/null +++ b/test/scalar_to_vec.jl @@ -0,0 +1,58 @@ +using SymbolicUtils +using SymbolicUtils.Code +import SymbolicUtils as SU +import SymbolicCompilerPasses as SC +using LinearAlgebra +using Test + +brusselator_f(x, y, t) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.0 +limit(a, N) = a == N + 1 ? 1 : a == 0 ? N : a +function brusselator_2d_loop(du, u, p, t) + A, B, alpha, dx = p + alpha = alpha / dx^2 + @inbounds for I in CartesianIndices((N, N)) + i, j = Tuple(I) + x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]] + ip1, im1, jp1, + jm1 = limit(i + 1, N), limit(i - 1, N), limit(j + 1, N), + limit(j - 1, N) + du[i, j, 1] = alpha * (u[im1, j, 1] + u[ip1, j, 1] + u[i, jp1, 1] + u[i, jm1, 1] - + 4u[i, j, 1]) + + B + u[i, j, 1]^2 * u[i, j, 2] - (A + 1) * u[i, j, 1] + + brusselator_f(x, y, t) + du[i, j, 2] = alpha * (u[im1, j, 2] + u[ip1, j, 2] + u[i, jp1, 2] + u[i, jm1, 2] - + 4u[i, j, 2]) + + A * u[i, j, 1] - u[i, j, 1]^2 * u[i, j, 2] + end + du +end + +@testset "Brusselator Vectorized" begin + N = 32 + xyd_brusselator = range(0, stop = 1, length = N) + @syms DU[1:N, 1:N, 1:2]::Real U[1:N, 1:N, 1:2]::Real P[1:4]::Real t::Real + traced_bruss = brusselator_2d_loop(collect(DU), collect(U), P, t) + current = Code.cse(traced_bruss) + + u_test = rand(size(U)...) + du_test = zeros(size(U)...) + p_test = [3.4, 1.0, 1.0, 0.1] + t_test = 0.5 + brusselator_2d_loop(du_test, u_test, p_test, t_test) + + d = Dict(collect(U) .=> u_test) + d[P[1]] = p_test[1] + d[P[2]] = p_test[2] + d[P[3]] = p_test[3] + d[P[4]] = p_test[4] + d[t] = t_test + + let_res = SC.detect_jacobian(current, Code.CSEState()) + optimized = SU.Code.apply_optimization_rule(current, Code.CSEState(), SC.JACOBIAN_VECTORIZE_RULE) + + optimized_expr = Func([U, P, t], [], optimized) + + optimized_f = eval(toexpr(optimized_expr)); + du_optimized = @invokelatest optimized_f(u_test, p_test, t_test) + @test isapprox(du_test, du_optimized, rtol = 1e-10, atol = 1e-10) +end \ No newline at end of file From d713829d4906e0305f40706ad1289c049b720abc Mon Sep 17 00:00:00 2001 From: Dhairya Gandhi Date: Fri, 6 Feb 2026 15:24:22 +0000 Subject: [PATCH 7/8] chore: improve generated code --- src/jacobian_vectorize.jl | 37 +++++++++++++++++++++++++++++++++---- 1 file changed, 33 insertions(+), 4 deletions(-) diff --git a/src/jacobian_vectorize.jl b/src/jacobian_vectorize.jl index 9985e2d..5a8f9d1 100644 --- a/src/jacobian_vectorize.jl +++ b/src/jacobian_vectorize.jl @@ -319,7 +319,7 @@ end struct VectorizedJacobianMatched{TL, TNL, TC, TU, S} <: AbstractMatched L::TL NL::TNL - C::C + C::TC U::TU shape::S end @@ -439,15 +439,44 @@ end function transform_jacobian(expr::Code.Let, match_data::VectorizedJacobianMatched, state) # DU = L * U + NL(U) + C + T = SymReal + counter = 1 + linear_sym = Symbol("##jacobian_vectorize##", counter) linear_terms = match_data.L * vec(match_data.U) + # linear_terms = Term{T}(*, [match_data.L, vec(match_data.U)]; + # type = symtype(match_data.U), + # shape = match_data.shape) + linear_var = Sym{T}(linear_sym; type = symtype(match_data.U)) + linear_ass = Assignment(linear_var, linear_terms) + + counter += 1 + nonlinear_sym = Symbol("##jacobian_vectorize##", counter) nonlinear_terms = match_data.NL + nonlinear_var = Sym{T}(nonlinear_sym; type = symtype(match_data.U)) + nonlinear_ass = Assignment(nonlinear_var, nonlinear_terms) + + counter += 1 + const_sym = Symbol("##jacobian_vectorize##", counter) const_terms = match_data.C - val = linear_terms + nonlinear_terms + const_terms - T = SymReal + const_var = Sym{T}(const_sym; type = symtype(match_data.U)) + const_ass = Assignment(const_var, const_terms) + + counter += 1 + val_sym = Symbol("##jacobian_vectorize##", counter) + # val = linear_terms + nonlinear_terms + const_terms + val = term(+, linear_terms, nonlinear_terms, const_terms, type = symtype(match_data.U)) + val_var = Sym{T}(val_sym; type = symtype(match_data.U)) + val_ass = Assignment(val_var, val) + + counter += 1 + res_sym = Symbol("##jacobian_vectorize##", counter) res = Term{T}(reshape, [val, match_data.shape]; type = symtype(match_data.U)) + res_var = Sym{T}(res_sym; type = symtype(match_data.U)) + res_ass = Assignment(res_var, res) - Code.Let([], res, expr.let_block) + new_pairs = [linear_ass, nonlinear_ass, const_ass, val_ass, res_ass] + Code.Let(new_pairs, res_var, expr.let_block) end const JACOBIAN_VECTORIZE_RULE = OptimizationRule( From 296f2acfc92c1f9428329119961b8250a9c37b09 Mon Sep 17 00:00:00 2001 From: Dhairya Gandhi Date: Mon, 9 Feb 2026 10:37:21 +0000 Subject: [PATCH 8/8] test: update tagbot --- .github/workflows/TagBot.yml | 6 ++++-- src/jacobian_vectorize.jl | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/.github/workflows/TagBot.yml b/.github/workflows/TagBot.yml index 5df3b8f..b6fee3d 100644 --- a/.github/workflows/TagBot.yml +++ b/.github/workflows/TagBot.yml @@ -1,9 +1,11 @@ name: TagBot on: - schedule: - - cron: 0 * * * * + issue_comment: + types: + - created jobs: TagBot: + if: github.event_name == 'workflow_dispatch' || github.actor == 'JuliaTagBot' runs-on: ubuntu-latest steps: - uses: JuliaRegistries/TagBot@v1 diff --git a/src/jacobian_vectorize.jl b/src/jacobian_vectorize.jl index 5a8f9d1..63ef2f6 100644 --- a/src/jacobian_vectorize.jl +++ b/src/jacobian_vectorize.jl @@ -418,7 +418,7 @@ function construct_jacobian(eqs, vars, expr) for i in 1:length(eqs) js, const_terms, linear_terms, nonlinear_sum = extract_coeffs(eqs[i], vars, expr) - for (k, j) in enumerate(js) + for j in js A[i, Symbolics.value(j)] = linear_terms[linear_to_cartesian(j, vars)] end NL[i] = nonlinear_sum