From 2c08fd3a9e55937fd79fd80c4d66a14fe74423f1 Mon Sep 17 00:00:00 2001 From: Manuel Berkemeier Date: Thu, 5 Jan 2023 18:17:28 +0100 Subject: [PATCH 1/2] expression based specialization function `specialize_methods`, apparently working for matrix-vector and matrix-matrix products. --- Project.toml | 3 +- src/Symbolics.jl | 4 ++ src/init.jl | 2 + src/specialize_methods.jl | 36 +++++++++++++++ src/wrapper-types.jl | 92 ++++++++++++++++++++++++--------------- 5 files changed, 102 insertions(+), 35 deletions(-) create mode 100644 src/specialize_methods.jl diff --git a/Project.toml b/Project.toml index a13f0d49d..b3f40f3c4 100644 --- a/Project.toml +++ b/Project.toml @@ -11,6 +11,7 @@ DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf" +ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04" Groebner = "0b43b601-686d-58a3-8a1c-6623616c7cd4" IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173" LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f" @@ -46,6 +47,7 @@ DomainSets = "0.5" Groebner = "0.1, 0.2" IfElse = "0.1" LaTeXStrings = "1.3" +LambertW = "0.4.5" Latexify = "0.11, 0.12, 0.13, 0.14, 0.15" MacroTools = "0.5" Metatheory = "1.2.0" @@ -62,7 +64,6 @@ StaticArrays = "1.1" SymbolicUtils = "0.18, 0.19" TermInterface = "0.2, 0.3" TreeViews = "0.3" -LambertW = "0.4.5" julia = "1.6" [extras] diff --git a/src/Symbolics.jl b/src/Symbolics.jl index cd8efd2d2..209c9007e 100644 --- a/src/Symbolics.jl +++ b/src/Symbolics.jl @@ -47,6 +47,10 @@ using MacroTools import MacroTools: splitdef, combinedef, postwalk, striplines include("wrapper-types.jl") +import ExprTools +export specialize_methods +include("specialize_methods.jl") + include("num.jl") include("complex.jl") diff --git a/src/init.jl b/src/init.jl index 73dab9a0b..0db87ebe4 100644 --- a/src/init.jl +++ b/src/init.jl @@ -23,4 +23,6 @@ function __init__() end end # SymPy + + specialize_methods() end diff --git a/src/specialize_methods.jl b/src/specialize_methods.jl new file mode 100644 index 000000000..4870ad2ad --- /dev/null +++ b/src/specialize_methods.jl @@ -0,0 +1,36 @@ +""" + specialize_methods(func, abstract_arg_types, inner_func, mods=nothing) + +For any method that implements `func` with signature +fitting `abstract_arg_types`, define methods for corresponding +symbolic types that pass all arguments to `inner_func`. +`mods` is an optional list of modules to look for methods in. +""" +function specialize_methods(func, abstract_arg_types, inner_func, mods=nothing) + ms = isnothing(mods) ? methods(func, abstract_arg_types) : methods(func, abstract_arg_types, mods) + for m in ms + mod = m.module + if mod != @__MODULE__ # do not overwrite method definitions from within this module itself, else: precompilation warnings + sig = ExprTools.signature(m; extra_hygiene=true) + fname = sig[:name] + args = sig[:args] + kwargs = get(sig, :kwargs, Symbol[]) + whereparams = get(sig, :whereparams, Symbol[]) + args_names = expr_argname.(args) + kwargs_names = expr_kwargname.(kwargs) + body = :($(inner_func)($(args_names...); $(kwargs_names...))) + Base.eval( + @__MODULE__, + wrap_func_expr( + mod, fname, args, kwargs, args_names, kwargs_names, whereparams, body; + abstract_arg_types + ) + ) + end#of `mod != @__MODULE__` + end#of `for m in ms` +end + +function specialize_methods(mods=nothing) + specialize_methods(Base.:(*), (AbstractMatrix, AbstractVector), _matvec, mods) + specialize_methods(Base.:(*), (AbstractMatrix, AbstractMatrix), _matmul, mods) +end \ No newline at end of file diff --git a/src/wrapper-types.jl b/src/wrapper-types.jl index c36fc9eca..55767773b 100644 --- a/src/wrapper-types.jl +++ b/src/wrapper-types.jl @@ -57,58 +57,82 @@ function wraps_type end has_symwrapper(::Type) = false is_wrapper_type(::Type) = false +# helper function to extract keyword argument names from expressions +function expr_kwargname(kwarg) + if kwarg isa Expr && kwarg.head == :kw + kwarg.args[1] + elseif kwarg isa Expr && kwarg.head == :(...) + kwarg.args[1] + else + kwarg + end +end + +# helper function to extract argument names from expressions +function expr_argname(arg) + if arg isa Expr && (arg.head == :(::) || arg.head == :(...)) + arg.args[1] + elseif arg isa Expr + error("$arg not supported as an argument") + else + arg + end +end + function wrap_func_expr(mod, expr) @assert expr.head == :function || (expr.head == :(=) && expr.args[1] isa Expr && expr.args[1].head == :call) def = splitdef(expr) - - sig = expr.args[1] body = def[:body] - fname = def[:name] args = get(def, :args, []) kwargs = get(def, :kwargs, []) + args_names = expr_argname.(args) + kwargs_names = expr_kwargname.(kwargs) + + wrap_func_expr(mod, fname, args, kwargs, args_names, kwargs_names, Symbol[], body) +end - impl_name = Symbol(fname,"_", hash(string(args)*string(kwargs))) - - function kwargname(kwarg) - if kwarg isa Expr && kwarg.head == :kw - kwarg.args[1] - elseif kwarg isa Expr && kwarg.head == :(...) - kwarg.args[1] - else - kwarg - end - end - - function argname(arg) - if arg isa Expr && (arg.head == :(::) || arg.head == :(...)) - arg.args[1] - elseif arg isa Expr - error("$arg not supported as an argument") - else - arg - end - end - - names = vcat(argname.(args), kwargname.(kwargs)) - - function type_options(arg) +function wrap_func_expr( + mod, fname, args, kwargs, args_names, kwargs_names, whereparams, body; + abstract_arg_types=nothing +) + names = vcat(args_names, kwargs_names) + + function type_options(wparams, arg, arg_ind) + pmod = parentmodule(mod) + atype = isnothing(abstract_arg_types) ? Any : abstract_arg_types[arg_ind] if arg isa Expr && arg.head == :(::) - T = Base.eval(mod, arg.args[2]) + T = Base.eval(mod, quote + let $(Symbol(pmod)) = $(pmod); # make name of parent module available in eval scope + #= + NOTE + `typeintersect` is important here for consecutive calls to `specialize_methods` + with conceptually different super types. + E.g.: Consider we first specialize `*(::AbstractMatrix, ::AbstractVector)` to + redirect to `_matvec`, and then `*(::AbstractMatrix, ::AbstractMatrix)` to + redirect to `_matmul`. If we encounter some existing method for `*` which accepts + an `AbstractMatrix` and an `VecOrMat` (type union), then we accidentally redirect + a matrix-vector-product to `_matmul` without `typeintersect`. + =# + typeintersect($(atype), $(arg.args[2]) where {$(wparams...)}) + end + end) has_symwrapper(T) ? (T, :(SymbolicUtils.Symbolic{<:$T}), wrapper_type(T)) : - (T,:(SymbolicUtils.Symbolic{<:$T})) + (T, :(SymbolicUtils.Symbolic{<:$T})) elseif arg isa Expr && arg.head == :(...) - Ts = type_options(arg.args[1]) + Ts = type_options(wparams, arg.args[1], arg_ind) map(x->Vararg{x},Ts) else (Any,) end end - types = map(type_options, args) + types = [type_options(whereparams, arg, arg_ind) for (arg_ind, arg)=enumerate(args)] + + impl_name = Symbol(fname,"_", hash(string(args)*string(kwargs)*string(types))) impl = :(function $impl_name($(names...)) $body @@ -139,9 +163,9 @@ function wrap_func_expr(mod, expr) quote $impl $(methods...) - end |> esc + end end macro wrapped(expr) - wrap_func_expr(__module__, expr) + esc(wrap_func_expr(__module__, expr)) end From 8107d12efa532293f7a6e1984db3b2f13243a650 Mon Sep 17 00:00:00 2001 From: Manuel Berkemeier Date: Fri, 17 Feb 2023 10:27:38 +0100 Subject: [PATCH 2/2] Restrict specialization to `*` in `LinearAlgebra`. --- src/init.jl | 2 +- src/specialize_methods.jl | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/init.jl b/src/init.jl index 0db87ebe4..ead8373ff 100644 --- a/src/init.jl +++ b/src/init.jl @@ -24,5 +24,5 @@ function __init__() end # SymPy - specialize_methods() + specialize_methods((LinearAlgebra,)) end diff --git a/src/specialize_methods.jl b/src/specialize_methods.jl index 4870ad2ad..29d520eb0 100644 --- a/src/specialize_methods.jl +++ b/src/specialize_methods.jl @@ -30,6 +30,15 @@ function specialize_methods(func, abstract_arg_types, inner_func, mods=nothing) end#of `for m in ms` end +""" + specialize_methods(mods=nothing) + +Define specialized methods accepting symbolic types for the following functions and +signatures found in modules `mods` via `methods(...)`: + +* `Base.:(*)` for arguments of `(AbstractMatrix, AbstractVector)` to redirect to `_matvec`. +* `Base.:(*)` for arguments of `(AbstractMatrix, AbstractMetrax)` to redirect to `_matmul`. +""" function specialize_methods(mods=nothing) specialize_methods(Base.:(*), (AbstractMatrix, AbstractVector), _matvec, mods) specialize_methods(Base.:(*), (AbstractMatrix, AbstractMatrix), _matmul, mods)