Skip to content

Better handling of slices f(A[i,:]) etc. #20

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

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open

Better handling of slices f(A[i,:]) etc. #20

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a795509
colon indices?
Jul 30, 2020
909e7ce
KA tweaks
Jul 30, 2020
6decc26
further KA tweaks
Jul 30, 2020
cbd9476
Merge branch 'master' into slices
Aug 15, 2020
c34d213
Merge branch 'master' into slices
Aug 24, 2020
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
15 changes: 11 additions & 4 deletions src/macro.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -958,14 +958,21 @@ function make_many_actors(act!, args, ex1, outer::Vector, ex3, inner::Vector, ex
outer = [Symbol(EPS, 1)] # fake index name, only appears in @index
sizes = [:(one(Int))] # iterate over 1:1
end
kex1 = macroexpand(store.mod, quote
ka_args = map(args) do x # mark everything const, except for ZED::AbstractArray{TYP} and gradients
x isa Symbol || return x
startswith(string(x), string(DEL)) && x != Symbol(DEL, ZED) && return x
:(@Const($x))
end
kex0 = quote

KernelAbstractions.@kernel function $kernel($(args...), $KEEP, $FINAL) where {$TYP}
KernelAbstractions.@kernel function $kernel($(ka_args...), @Const($KEEP), @Const($FINAL)) where {$TYP}
($(outer...),) = @index(Global, NTuple)
($ex1; $ex3; $ex4; $ex6)
@views ($ex1; $ex3; $ex4; $ex6)
end

end)
end
store.verbose==2 && @show verbosetidy(kex0)
kex1 = macroexpand(store.mod, kex0)
push!(store.outpre, kex1)
if isdefined(store.mod, :CUDA) && isdefined(store.mod, :CuArray) # new-style, CUDA.jl, with CUDADevice()
info2 = store.verbose>0 ? :(@info "running KernelAbstractions + CUDA actor $($note)") : nothing
Expand Down
15 changes: 15 additions & 0 deletions test/parsing.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,12 @@ using Tullio, Test, LinearAlgebra
@test A == @tullio P[i'] := A[i′]
@test [1,4,9] == @tullio Q[i'] := (i′)^2 (i' in 1:3)

# colons
@tullio R[i] := prod(M[:,i]) avx=false
@test R ≈ vec(prod(M, dims=1))

# @tullio S[:,i] := cumsum(M[:,i]) avx=false

# non-numeric array
@tullio Y[i] := (ind=i, val=A[i])
@test Y[2] === (ind = 2, val = 4)
Expand Down Expand Up @@ -205,6 +211,15 @@ end

@test_throws LoadError @eval @tullio [i,j] = A[i] + 100

# colons
R = similar(A)
@tullio R[i] = sum(D[:,i]) avx=false
@test R == vec(sum(D, dims=1))

S = similar(D);
@tullio S[i,:] = cumsum(D[:,i]) avx=false
@test_broken S == cumsum(D, dims=1)

# zero off-diagonal? not now, but maybe it should?
@tullio D[i,i] = A[i]

Expand Down