diff --git a/src/macro.jl b/src/macro.jl index 98005cb..e149965 100644 --- a/src/macro.jl +++ b/src/macro.jl @@ -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 diff --git a/test/parsing.jl b/test/parsing.jl index 7b8dce6..83c18d7 100644 --- a/test/parsing.jl +++ b/test/parsing.jl @@ -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) @@ -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]