diff --git a/Project.toml b/Project.toml index cb63978..433a51d 100644 --- a/Project.toml +++ b/Project.toml @@ -1,9 +1,10 @@ name = "SymbolicCompilerPasses" uuid = "3384d301-0fbe-4b40-9ae0-b0e68bedb069" -version = "0.1.2" authors = ["Dhairya Gandhi "] +version = "0.1.2" [deps] +DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" @@ -11,22 +12,29 @@ SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b" [weakdeps] LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" +Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc" [extensions] SCPLinearSolveExt = ["LinearSolve"] +SymbolicCompilerPassesRotationsExt = ["Rotations"] [compat] -LinearAlgebra = "1" +DataStructures = "0.19.3" +LinearAlgebra = "1.11.0, 1.10" LinearSolve = "3.53.0" PreallocationTools = "0.4.34, 1" +Rotations = "1.7.1" StaticArrays = "1.9.15, 1" SymbolicUtils = "4.1.0" +Symbolics = "7.2" julia = "1" [extras] Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" +Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc" [targets] -test = ["Pkg", "SafeTestsets", "Test"] +test = ["Pkg", "SafeTestsets", "Test", "Symbolics", "Rotations"] diff --git a/ext/SCPLinearSolveExt/SCPLinearSolveExt.jl b/ext/SCPLinearSolveExt/SCPLinearSolveExt.jl index 1bf1e94..d5b3f85 100644 --- a/ext/SCPLinearSolveExt/SCPLinearSolveExt.jl +++ b/ext/SCPLinearSolveExt/SCPLinearSolveExt.jl @@ -4,17 +4,37 @@ using SymbolicUtils using SymbolicUtils.Code using LinearSolve using LinearAlgebra -import SymbolicCompilerPasses: ldiv_transformation, SymbolicCompilerPasses, get_factorization, get_from_cache, FACTORIZATION_CACHE, LINEARSOLVE_LIB +import SymbolicCompilerPasses: ldiv_transformation, SymbolicCompilerPasses, get_factorization, get_from_cache, FACTORIZATION_CACHE +using StaticArrays __init__() = SymbolicCompilerPasses.LINEARSOLVE_LIB[] = true +const LINSOLVEPROB_CACHE = Dict() + +function get_linear_prob(A::StaticArray, B::StaticArray) + prob = LinearSolve.LinearProblem(A, B) +end + +function get_linear_prob(A::TA, B::TB) where {TA, TB} + get!(LINSOLVEPROB_CACHE, A) do + prob = LinearSolve.LinearProblem(A, B) + init(prob) + end# ::Base.promote_op(init, Tuple{Base.promote_op(LinearSolve.LinearProblem, Tuple{TA, TB})}) +end + function linear_solve(A, B) - linsolve = get_factorization(A, B) + linsolve = get_linear_prob(A, B) linsolve.b = B sol = solve!(linsolve) return sol.u end +function linear_solve(A::StaticArray, B::StaticArray) + linsolve = get_linear_prob(A, B) + sol = solve(linsolve) + return sol.u +end + function get_factorization(A, B) get!(FACTORIZATION_CACHE, A) do prob = LinearSolve.LinearProblem(A, B) @@ -25,7 +45,7 @@ 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 + In order to opt-out of using LinearSolve, set SymbolicCompilerPasses.LINEARSOLVE_LIB[] = false." maxlog=1 # Build transformation transformations = Dict{Int, Code.Assignment}() diff --git a/ext/SymbolicCompilerPassesRotationsExt/SymbolicCompilerPassesRotationsExt.jl b/ext/SymbolicCompilerPassesRotationsExt/SymbolicCompilerPassesRotationsExt.jl new file mode 100644 index 0000000..93a258b --- /dev/null +++ b/ext/SymbolicCompilerPassesRotationsExt/SymbolicCompilerPassesRotationsExt.jl @@ -0,0 +1,9 @@ +module SymbolicCompilerPassesRotationsExt + +using SymbolicCompilerPasses +import SymbolicCompilerPasses: is_orthogonal_type + +using Rotations +is_orthogonal_type(::Rotations.Rotation) = true + +end \ No newline at end of file diff --git a/src/SymbolicCompilerPasses.jl b/src/SymbolicCompilerPasses.jl index 3848403..2d56c5e 100644 --- a/src/SymbolicCompilerPasses.jl +++ b/src/SymbolicCompilerPasses.jl @@ -4,13 +4,15 @@ using LinearAlgebra using PreallocationTools using SymbolicUtils import SymbolicUtils: symtype, vartype, Sym, BasicSymbolic, Term, iscall, operation, arguments, maketerm, Const, shape, isterm, unwrap, - is_function_symbolic, is_called_function_symbolic, getname, Unknown + is_function_symbolic, is_called_function_symbolic, getname, Unknown, search_variables!, search_variables import SymbolicUtils.Code: Code, OptimizationRule, substitute_in_ir, apply_optimization_rules, AbstractMatched, Assignment, CSEState, lhs, rhs, apply_substitution_map, issym, isterm, toexpr, _is_array_of_symbolics, MakeArray, shape import SymbolicUtils: search_variables, search_variables! using StaticArrays +using DataStructures + function bank(dic, key, value) if haskey(dic, key) dic[key] = vcat(dic[key], value) @@ -20,8 +22,11 @@ function bank(dic, key, value) end include("matmuladd.jl") +include("ortho_inv_opt.jl") include("hvncat_static_opt.jl") include("ldiv_opt.jl") include("la_opt.jl") +include("mb_opt.jl") + end # module SymbolicCompilerPasses diff --git a/src/hvncat_static_opt.jl b/src/hvncat_static_opt.jl index 659bde3..8783580 100644 --- a/src/hvncat_static_opt.jl +++ b/src/hvncat_static_opt.jl @@ -204,21 +204,16 @@ function transform_hvncat_to_static(expr::Code.Let, match_data::Vector{HvncatMat # Column vector: SVector{n}(elements...) n = dims[1] t = term(Core.apply_type, StaticArrays.SVector, n; type = Any) - static_ctor = Term{T}( - t, - elements; - type=symtype(lhs_var) - ) else # Matrix: SMatrix{m,n}(elements...) m, n = dims t = term(Core.apply_type, StaticArrays.SMatrix, m, n; type = Any) - static_ctor = Term{T}( - t, - elements; - type=symtype(lhs_var) - ) end + static_ctor = Term{T}( + t, + elements; + type=symtype(lhs_var) + ) new_assignment = Assignment(lhs_var, static_ctor) transformations[match.assignment_idx] = new_assignment diff --git a/src/ldiv_opt.jl b/src/ldiv_opt.jl index e51cab3..fc98683 100644 --- a/src/ldiv_opt.jl +++ b/src/ldiv_opt.jl @@ -1,4 +1,4 @@ -const FACTORIZATION_CACHE = WeakKeyDict() +const FACTORIZATION_CACHE = Dict() struct LdivMatch{Ta, Tb, S <: Assignment, P <: AbstractString} <: AbstractMatched A::Ta @@ -149,6 +149,7 @@ function get_factorization(A) qr_A = get!(FACTORIZATION_CACHE, A) do qr(A) end + # qr_A = qr(A) qr_A end @@ -160,7 +161,7 @@ ldiv_transformation(x, ::Nothing) = ldiv_transformation(x, Val(false)) function ldiv_transformation(safe_matches, ::Val{false}) @warn "Backsolve may be sped up by adding LinearSolve.jl. In order to enable this optimization, add LinearSolve.jl to your environment. - To opt-out of using LinearSolve, set SymbolicCompilerPasses.LINEARSOLVE_LIB[] = false." maxlog=Inf + To opt-out of using LinearSolve, set SymbolicCompilerPasses.LINEARSOLVE_LIB[] = false." maxlog=1 # Build transformation transformations = Dict{Int, Code.Assignment}() diff --git a/src/mb_opt.jl b/src/mb_opt.jl new file mode 100644 index 0000000..f9d0efe --- /dev/null +++ b/src/mb_opt.jl @@ -0,0 +1,63 @@ +function detect_small_views(expr::Code.Let, state) + matches = [] + for (i, p) in enumerate(expr.pairs) + r = rhs(p) + iscall(r) || continue + if operation(r) === view + arr, inds... = arguments(r) + myt = find_term(inds[1], expr) + is_small_hvncat(size(Code.rhs(myt))...) || continue + push!(matches, (idx = i, expr = r)) + end + end + matches +end + +function construct_type(dims) + # if length(dims) == 1 + # return Core.apply_type(SVector, dims[1]) + # else + # return Core.apply_type(SVector, Tuple(dims)) + # end + Core.apply_type(SVector, length(dims)) +end + +function find_term(target, expr::Code.Let) + filter(expr.pairs) do p + Code.lhs(p) === target + end |> only +end + +function transform_view(expr, match_data, state) + new_pairs = [] + idxs = Set(getproperty.(match_data, :idx)) + transformations = Dict() + for match in match_data + idx = match.idx + r = match.expr + T = symtype(r) + V = vartype(r) + arr, inds... = arguments(r) + t = term(construct_type, inds[1]) + transformations[idx] = Term{V}(t, [r], type = T) + end + + for (i, p) in enumerate(expr.pairs) + if i in idxs + new_rhs = transformations[i] + push!(new_pairs, Code.Assignment(lhs(p), new_rhs)) + else + push!(new_pairs, p) + end + end + + Code.Let(new_pairs, expr.body, expr.let_block) +end + + +const MB_VIEW_RULE = OptimizationRule( + "MB_VIEW_RULE", + detect_small_views, + transform_view, + 10, +) \ No newline at end of file diff --git a/src/ortho_inv_opt.jl b/src/ortho_inv_opt.jl new file mode 100644 index 0000000..a9ba1db --- /dev/null +++ b/src/ortho_inv_opt.jl @@ -0,0 +1,101 @@ +is_orthogonal_matrix(A) = A * A' ≈ I(size(A, 1)) +is_orthogonal_type(x) = begin + if issym(x) + return getmetadata(x, IsOrthogonal, false) + else + return false + end + false +end + +struct IsOrthogonal end + +struct InvMatched{T, C} <: AbstractMatched + A::T + candidate::C + idx::Int +end + +function detect_orthogonal_matrix(expr, state::Code.CSEState) + # var_to_ortho = Dict() + # ortho_to_var = Dict() + # for v in search_variables(expr) + # @show symtype(v) + # bank(var_to_ortho, v, getmetadata(v, IsOrthogonal, false)) + # bank(ortho_to_var, getmetadata(v, IsOrthogonal, false), v) + # end + + # if !haskey(ortho_to_var, true) + # @warn "not found any metadata arrays" + # return nothing + # end + + idxs = findall(expr.pairs) do p + r = rhs(p) + iscall(r) || return false + op = operation(r) + if op === inv + args = arguments(r) + length(args) == 1 || return false + getmetadata(args[1], IsOrthogonal, false) == true || return false + return true + end + false + end + + candidates = expr.pairs[idxs] + + matches = map(idxs, candidates) do idx, candidate + A = arguments(rhs(candidate))[1] + InvMatched(A, candidate, idx) + end + + f = filter(!isnothing, matches) + isempty(f) ? nothing : f +end + +function transform_inv_optimization(expr, matches, state::Code.CSEState) + + expr_copy = deepcopy(expr) + map(matches) do match + A = match.A + if getmetadata(A, IsOrthogonal, false) == true + expr_copy.pairs[match.idx] = Code.Assignment( + lhs(match.candidate), + transpose(A) + ) + else + t = term(is_orthogonal_type, A) + # code = IfElse( + # t, + # transpose(A), + # inv(A) + # ) + code = ifelse(t, transpose(A), inv(A)) + expr_copy.pairs[match.idx] = Code.Assignment( + lhs(match.candidate), + code + ) + end + end + expr_copy +end + +const ORTHO_INV_RULE = OptimizationRule( + "Ortho_Inv", + detect_orthogonal_matrix, + transform_inv_optimization, + 10 +) + +function ortho_inv_opt(expr, state::Code.CSEState) + + # Try to apply optimization rules + optimized = apply_optimization_rules(expr, state, [ORTHO_INV_RULE]) + if optimized !== nothing + return optimized + end + + # If no optimization applied, return original expression + return expr +end \ No newline at end of file diff --git a/test/multiple.jl b/test/multiple.jl new file mode 100644 index 0000000..088f6f3 --- /dev/null +++ b/test/multiple.jl @@ -0,0 +1,55 @@ +using SymbolicUtils +using SymbolicUtils.Code +import SymbolicUtils as SU +import SymbolicCompilerPasses as SC +using LinearAlgebra +using Symbolics +using Test + +using Rotations + +function test_codegen(expr, rules, args...) + current = SU.Code.cse(expr) + optimized = SU.Code.apply_optimization_rules(current, SU.Code.CSEState(), rules) + + current_expr = Func([args...], [], current) + optimized_expr = Func([args...], [], optimized) + + current_f = eval(toexpr(current_expr)) + optimized_f = eval(toexpr(optimized_expr)) + + N = 3 + test_args = map(args) do arg + if getmetadata(arg, SC.IsOrthogonal, false) == true + RotXYZ((rand([0,1]) for x in 1:size(arg, 1))...) + else + rand(size(arg)...) + end + end + + current_res = @invokelatest current_f(test_args...) + optimized_res = @invokelatest optimized_f(test_args...) + + @test isapprox(current_res, optimized_res, rtol=1e-10, atol=1e-10) +end + + +@testset "Combined Optimizations" begin + @syms A[1:3, 1:3] B[1:3, 1:3] C[1:3, 1:3] D[1:3, 1:3] E[1:3, 1:3] + + P = A \ B + expr = P + C * D + test_codegen(expr, [SC.LDIV_RULE, SC.MATMUL_ADD_RULE], A, B, C, D) + + # note that LDIV_RULE should not apply because A is reused + P = A \ B + expr2 = P * E + A + test_codegen(expr2, [SC.LDIV_RULE, SC.MATMUL_ADD_RULE], A, B, E) + + @syms R[1:3, 1:3] + Ro = setmetadata(R, SC.IsOrthogonal, true) + P = A \ B + # expr3 = P + C * inv(Ro) + expr3 = tril(P) + C * inv(Ro) + test_codegen(expr3, [SC.LDIV_RULE, SC.MATMUL_ADD_RULE, SC.ORTHO_INV_RULE], A, B, C, Ro) +end diff --git a/test/ortho_inv_opt.jl b/test/ortho_inv_opt.jl new file mode 100644 index 0000000..9c32485 --- /dev/null +++ b/test/ortho_inv_opt.jl @@ -0,0 +1,74 @@ +using SymbolicUtils +using SymbolicUtils.Code +import SymbolicUtils as SU +import SymbolicCompilerPasses as SC +using Symbolics + +using LinearAlgebra +using Rotations +using Test + +function has_ortho_opt(expr::Code.Let) + return any(expr.pairs) do assignment + rhs_expr = Code.rhs(assignment) + has_ortho_opt(rhs_expr) + end +end +# has_ortho_opt(expr::Code.IfElse) = has_ortho_opt(expr.ifbody) || has_ortho_opt(expr.elsebody) + +function has_ortho_opt(expr) + if SU.iscall(expr) + op = SU.operation(expr) + return op === transpose || op === SC.is_orthogonal_type + end + false +end + +function check_ortho_opt(expr, A, B; expected_ortho = true) + current = SU.Code.cse(expr) + toexpr(current) + + # optimized = SC.ortho_inv_opt(current, SU.Code.CSEState()) + optimized = Code.apply_optimization_rule(current, SU.Code.CSEState(), SC.ORTHO_INV_RULE) + # return optimized + # return toexpr(optimized) + + @test has_ortho_opt(optimized) == expected_ortho + + current_fun = Func([A, B], [], current) + optimized_fun = Func([A, B], [], optimized) + + current_f = eval(toexpr(current_fun)) + optimized_f = eval(toexpr(optimized_fun)) + + a = rand(3,3) + b = rand(3,3) + R_euler = RotMatrix{3, Float64}(rand(3, 3)) + + # res1 = invokelatest(current_f, a, b) + # res2 = invokelatest(optimized_f, a, b) + # @test isapprox(res1, res2, rtol=1e-10, atol=1e-10) + + res1 = invokelatest(current_f, R_euler, b) + res2 = invokelatest(optimized_f, R_euler, b) + @test isapprox(res1, res2, rtol=1e-10, atol=1e-10) +end + + +@testset "Orthogonal Matrices: inv -> transpose" begin + @syms A[1:3, 1:3] B[1:3, 1:3] C[1:3, 1:3] D[1:3, 1:3] E[1:3, 1:3] + Ao = SU.setmetadata(A, SC.IsOrthogonal, true) + + + expr = inv(Ao) * B + check_ortho_opt(expr, Ao, B) + + expr2 = inv(A) * B + check_ortho_opt(expr2, A, B; expected_ortho = false) + + expr3 = inv(A) * B * inv(Ao) + check_ortho_opt(expr3, A, B) + + expr4 = inv(Ao * B) + B + check_ortho_opt(expr4, A, B; expected_ortho = false) # Can't optimize inv(Ao * B) without more info +end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index b82d2a6..22632ca 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,5 +5,6 @@ using Pkg, Test, SafeTestsets @safetestset "MatmulAdd Optimization" begin include("mul5_opt.jl") end @safetestset "Literal Small Array Allocation" begin include("array_literal.jl") end @safetestset "Ldiv Factorization Optimization" begin include("ldiv_opt.jl") end -end - + @safetestset "Orthogonal Inverse Optimization" begin include("ortho_inv_opt.jl") end + @safetestset "Apply Multiple Rules" begin include("multiple.jl") end +end \ No newline at end of file