Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use unflatten_long_ops in build_function again #801

Merged
merged 2 commits into from
Feb 19, 2021
Merged

use unflatten_long_ops in build_function again #801

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ddcba6c
avoid fixpoint in unflatten_long_ops
shashi Feb 19, 2021
7d5839d
Fix unflatten_long_ops
YingboMa Feb 19, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 14 additions & 9 deletions src/build_function.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,14 +52,19 @@ function build_function(args...;target = JuliaTarget(),kwargs...)
_build_function(target,args...;kwargs...)
end

function unflatten_long_ops(op, N=4)
rule1 = @rule((+)((~~x)) => length(~~x) > N ?
+(+((~~x)[1:N]...) + (+)((~~x)[N+1:end]...)) : nothing)
rule2 = @rule((*)((~~x)) => length(~~x) > N ?
*(*((~~x)[1:N]...) * (*)((~~x)[N+1:end]...)) : nothing)
function unflatten_args(f, args, N=4)
length(args) < N && return Term{Real}(f, args)
unflatten_args(f, [Term{Real}(f, group)
for group in Iterators.partition(args, N)], N)
end

function unflatten_long_ops(op, N=4)
YingboMa marked this conversation as resolved.
Show resolved Hide resolved
op = value(op)
Rewriters.Fixpoint(Rewriters.Postwalk(Rewriters.Chain([rule1, rule2])))(op)
!istree(op) && return Num(op)
rule1 = @rule((+)(~~x) => length(~~x) > N ? unflatten_args(+, ~~x, 4) : nothing)
rule2 = @rule((*)(~~x) => length(~~x) > N ? unflatten_args(*, ~~x, 4) : nothing)

Num(Rewriters.Postwalk(Rewriters.Chain([rule1, rule2]))(op))
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand why this fixed it...

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lol. It fixed the sparse array case. We need Num to make sure that zero is defined.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ahh ok

end


Expand All @@ -76,7 +81,7 @@ function _build_function(target::JuliaTarget, op, args...;
linenumbers = true)

dargs = map(destructure_arg, [args...])
expr = toexpr(Func(dargs, [], op))
expr = toexpr(Func(dargs, [], unflatten_long_ops(op)))

if expression == Val{true}
expr
Expand Down Expand Up @@ -247,7 +252,7 @@ function _make_array(rhss::AbstractArray, similarto)
end
end

_make_array(x, similarto) = x
_make_array(x, similarto) = unflatten_long_ops(x)

## In-place version

Expand Down Expand Up @@ -298,7 +303,7 @@ function _set_array(out, outputidxs, rhss::AbstractArray, checkbounds, skipzeros
end)
end

_set_array(out, outputidxs, rhs, checkbounds, skipzeros) = rhs
_set_array(out, outputidxs, rhs, checkbounds, skipzeros) = unflatten_long_ops(rhs)


function vars_to_pairs(name,vs::Union{Tuple, AbstractArray}, symsdict=Dict())
Expand Down