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.

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

Teach @polly to simplify range-based loops on AST-level #17965

Open
Changes from all commits
Commits
File filter...
Filter file types
Jump to…
Jump to file or symbol
Failed to load files and symbols.

Always

Just for now

@@ -145,15 +145,6 @@ macro propagate_inbounds(ex)
end
end

"""
@polly
Tells the compiler to apply the polyhedral optimizer Polly to a function.
"""
macro polly(ex)
esc(isa(ex, Expr) ? pushmeta!(ex, :polly) : ex)
end

## some macro utilities ##

find_vars(e) = find_vars(e, [])
@@ -0,0 +1,152 @@
# This file is a part of Julia. License is MIT: http://julialang.org/license

# Support for @polly

module Polly

export @polly

import Base: start, next, done

"""
Tells the compiler to apply the polyhedral optimizer Polly to a function.
"""
macro polly(func)
(isa(func, Expr) && func.head == :function) || throw(ArgumentError("@polly can only be applied to functions"))
canonicalize!(func)
return esc(Base.pushmeta!(func, :polly))
end

# This range type only differs from `Base.UnitRange` in its representation of
# emptiness. An empty `Base.UnitRange` will always have `stop == start - 1`.
# For example, constructing a range `5:2` will actually result in `5:4`.
# `Polly.UnitRange` drops this requirement, i.e. `5:2` would be used as is,
# which allows for a simpler constructor. When iterating over a
# `Polly.UnitRange` loop bounds will therefore be more obvious to Polly than
# with a `Base.UnitRange`.
immutable UnitRange{T<:Real} <: AbstractUnitRange{T}
start::T
stop::T
UnitRange(start, stop) = new(start, stop)
end
UnitRange{T<:Real}(start::T, stop::T) = Polly.UnitRange{T}(start, stop)

# This method was directly adopted from `Base.UnitRange`.
start{T}(r::Polly.UnitRange{T}) = oftype(r.start + one(T), r.start)

# This has to be different than for `Base.UnitRange` to reflect the different
# behavior of the `Polly.UnitRange` constructor.
done{T}(r::Polly.UnitRange{T}, i) = (i < oftype(i, r.start)) | (i > oftype(i, r.stop))

# `Base.StepRange` uses the same representation of emptiness as described above
# for `Base.UnitRange` with `stop == start - 1` but additionally, in the case of
# non-emptiness, its constructor will precompute the last value that is actually
# part of the range. For example, `5:2:8` would actually result in `5:2:7`. In
# `Polly.StepRange` we simplify construction by dropping these requirements,
# i.e. `5:2:8` would also be used as is. When iterating over a `Polly.StepRange`
# loop bounds will therefore be more obvious to Polly than with a
# `Base.StepRange`.
immutable StepRange{T,S} <: OrdinalRange{T,S}
start::T
step::S
stop::T
StepRange(start::T, step::S, stop::T) = new(start, step, stop)
end
StepRange{T,S}(start::T, step::S, stop::T) = Polly.StepRange{T,S}(start, step, stop)

# This method was directly adopted from `Base.StepRange`.
start(r::Polly.StepRange) = oftype(r.start + r.step, r.start)

# This method was directly adopted from `Base.StepRange`.
next{T}(r::Polly.StepRange{T}, i) = (convert(T,i), i + r.step)

# We have to use a simpler condition as for `Base.StepRange` in order to
# be able to derive the loop bounds in Polly. For now it also ignores
# wrap-arounds which could for example happen for `1:1:typemax(Int64)` which we
# consider a rare application.
done{T,S}(r::Polly.StepRange{T,S}, i) = (r.step > zero(r.step)) ? (i > oftype(i, r.stop)) :
(i < oftype(i, r.stop))

# This was directly adopted from `Base.UnitRange` and `Base.StepRange` to avoid
# overflows for smaller `Integer` types.
let smallint = (Int === Int64 ?
Union{Int8,UInt8,Int16,UInt16,Int32,UInt32} :
Union{Int8,UInt8,Int16,UInt16})
global start
global next
start{T<:smallint}(r::Polly.StepRange{T}) = convert(Int, r.start)
next{T<:smallint}(r::Polly.StepRange{T}, i) = (i % T, i + r.step)
start{T<:smallint}(r::Polly.UnitRange{T}) = convert(Int, r.start)
end

# Find assigments of the form `i = start:stop` and `i = start:step:stop` that
# occur in `for`-loop headers in `func` and replace them by
# `i = Polly.UnitRange(start,stop)` and `i = Polly.StepRange(start,step,stop)`.
function canonicalize!(func)
worklist = [func]
while !isempty(worklist)
expr = pop!(worklist)
if expr.head == :for
loop_header = expr.args[1]

This comment has been minimized.

Copy link
@singam-sanjay

singam-sanjay May 26, 2017

Contributor

Wow ! What's this feature of Julia that allows functions to be accessed like a parse tree ?

This comment has been minimized.

Copy link
@singam-sanjay

singam-sanjay May 26, 2017

Contributor

Thank You !

canonicalize_ranges_in_loop_header!(loop_header)
# The loop body might contain further loops that should be
# canonicalized, so push it to the worklist for later examination.
loop_body = expr.args[2]
push!(worklist, loop_body)
else
# If `Expr` isn't a `for`-loop, it might contain nested expressions
# which themselves contain `for`-loops.
for arg in expr.args
if isa(arg, Expr)
push!(worklist, arg)
end
end
end
end
end

# Find assigments of the form `i = start:stop` and `i = start:step:stop` in
# the given `loop_header` and replace them by `i = Polly.UnitRange(start,stop)`
# and `i = Polly.StepRange(start,step,stop)`.
function canonicalize_ranges_in_loop_header!(loop_header)
if loop_header.head == :block
# If `loop_header` is a `:block` expression we are dealing with a loop
# of the form `for i1 = ..., i2 = ..., ...` which uses multiple
# iteration variables.
for assignment in loop_header.args
canonicalize_range_in_assignment!(assignment)
end
else
# If `loop_header` is not a `:block` expression we have just a simple
# `for i = ...` with a single iteration variable.
canonicalize_range_in_assignment!(loop_header)
end
end

# If the given assignment has the form `i = start:stop` or `i = start:step:stop`
# then rewrite it to `i = Polly.UnitRange(start,stop)` or
# `i = Polly.StepRange(start,step,stop)`.
function canonicalize_range_in_assignment!(assignment)
@assert(assignment.head == :(=))
rhs = assignment.args[2]
new_rhs = nothing

if rhs.head == :(:)
if length(rhs.args) == 2
start = rhs.args[1]
stop = rhs.args[2]
new_rhs = :(Base.Polly.UnitRange($start,$stop))
elseif length(rhs.args) == 3
start = rhs.args[1]
step = rhs.args[2]
stop = rhs.args[3]
new_rhs = :(Base.Polly.StepRange($start,$step,$stop))
end
end

if new_rhs != nothing
assignment.args[2] = new_rhs
end
end

end # module Polly
@@ -119,6 +119,10 @@ include("arraymath.jl")
include("simdloop.jl")
importall .SimdLoop

# The polyhedral optimizer Polly
include("polly.jl")
importall .Polly

# map-reduce operators
include("reduce.jl")

@@ -34,7 +34,7 @@ function choosetests(choices = [])
"enums", "cmdlineargs", "i18n", "workspace", "libdl", "int",
"checked", "intset", "floatfuncs", "compile", "parallel", "inline",
"boundscheck", "error", "ambiguous", "cartesian", "asmvariant",
"channels"
"channels", "polly"
]
profile_skipped = false
if startswith(string(Sys.ARCH), "arm")
@@ -0,0 +1,156 @@
# This file is a part of Julia. License is MIT: http://julialang.org/license

This comment has been minimized.

Copy link
@tkelman

tkelman Aug 11, 2016

Contributor

you'll need to add this file to the list in test/choosetests.jl if you want it to run done


function is_syntax_equal(x::Expr, y::Expr)
if x.head === y.head
if x.head === :line
# `:line` expressions are treated as syntactically equivalent
# regardless of their actual arguments
return true
else
return is_syntax_of_args_equal(x, y)
end
end
return false
end

is_syntax_equal(x, y) = x == y

function is_syntax_of_args_equal(x::Expr, y::Expr)
if length(x.args) != length(y.args)
return false
end

for i in eachindex(x.args)
if !is_syntax_equal(x.args[i], y.args[i])
return false
end
end

return true
end

# Test whether `Base.Polly.canonicalize!()` works for a `UnitRange`-based loop.
let single_unit_range_loop = quote
for i = 1:10
end
end

expected = quote
for i = Base.Polly.UnitRange(1,10)
end
end

Base.Polly.canonicalize!(single_unit_range_loop)
@test is_syntax_equal(single_unit_range_loop, expected)
end

# Test whether `Base.Polly.canonicalize!()` works for a `StepRange`-based loop.
let single_step_range_loop = quote
for i = 1:2:10
end
end

expected = quote
for i = Base.Polly.StepRange(1,2,10)
end
end

Base.Polly.canonicalize!(single_step_range_loop)
@test is_syntax_equal(single_step_range_loop, expected)
end

# Test whether `Base.Polly.canonicalize!()` works for nested range-based loops.
let nested_loops = quote
for i = 1:10, j = i:3:20
for k = i:j
end
end
end

expected = quote
for i = Base.Polly.UnitRange(1,10), j = Base.Polly.StepRange(i,3,20)
for k = Base.Polly.UnitRange(i,j)
end
end
end

Base.Polly.canonicalize!(nested_loops)
@test is_syntax_equal(nested_loops, expected)
end

# Test whether `Base.Polly.canonicalize!()` works for successive range-based
# loops.
let successive_loops = quote
for i = 1:10
end

for j = 1:2:10
end
end

expected = quote
for i = Base.Polly.UnitRange(1,10)
end

for j = Base.Polly.StepRange(1,2,10)
end
end

Base.Polly.canonicalize!(successive_loops)
@test is_syntax_equal(successive_loops, expected)
end

# Test whether `Base.Polly.canonicalize!()` works for loops nested inside
# `if`-statements
let loops_inside_if = quote
if some_condition
for i = 1:10
end
else
for j = 1:2:10
end
end
end

expected = quote
if some_condition
for i = Base.Polly.UnitRange(1,10)
end
else
for j = Base.Polly.StepRange(1,2,10)
end
end
end

Base.Polly.canonicalize!(loops_inside_if)
@test is_syntax_equal(loops_inside_if, expected)
end

# Test whether `Base.Polly.canonicalize!()` works for a more complex AST.
let trmm = quote
function trmm(alpha, A, B)
m,n = size(B)
for i = 1:m, j = 1:n
for k = (i+1):m
B[i,j] += A[k,i] * B[k,j]
end
B[i,j] = alpha * B[i,j]
end
end
end

expected = quote
function trmm(alpha, A, B)
m,n = size(B)
for i = Base.Polly.UnitRange(1,m), j = Base.Polly.UnitRange(1,n)
for k = Base.Polly.UnitRange((i+1),m)
B[i,j] += A[k,i] * B[k,j]
end
B[i,j] = alpha * B[i,j]
end
end
end

Base.Polly.canonicalize!(trmm)
@test is_syntax_equal(trmm, expected)
end
ProTip! Use n and p to navigate between commits in a pull request.
You can’t perform that action at this time.