How to implement n-dimensional functions?

[roflmaostc]

Hey,

I'm trying to figure out how we could use Tullio to generalize some n-dimensional functions.
Let's say I've got a n-dimensional array and I want to use Tullio for that. At the moment I need to treat dimensions differently.

using Tullio
using Random

Random.seed!(42)

function spatial_grad_square_1D(arr)
    resi = let 
         @tullio resi = abs2(arr[i + 1] - arr[i - 1])
    end
    return resi
end

function spatial_grad_square_2D(arr)
    resi = let 
         @tullio resi = abs2(arr[i + 1, j] - arr[i - 1, j])
    end
    resj = let 
         @tullio resj = abs2(arr[i, j + 1] - arr[i, j - 1])
    end
    return resi + resj
end

arr_1D = randn((10))
arr_2D = randn((10, 10))
spatial_grad_square_1D(arr_1D) 
spatial_grad_square_2D(arr_2D)

Do you have any hints on how to achieve this in a more elegant way? Using the Julia metaprogramming documentation I couldn't figure it out 😞

Thanks,

Felix

[mcabbott]

There is nothing built-in to deal with n-dimensional arrays. But you can do something like this: (No, in fact you can't)

@generated function spatial_grad_square_ND(arr)
    out, add = [], []
    for d in 1:ndims(arr)
        ind = :i # @gensym ind
        inds1 = map(1:ndims(arr)) do di
            i = Symbol(ind, di)
            di == d ? :($i+1) : i
        end
        inds2 = map(1:ndims(arr)) do di
            i = Symbol(ind, di)
            di == d ? :($i-1) : i
        end
        res = Symbol(:res_,d) # @gensym res
        push!(out, :($res = let
            @tullio $res = abs2(arr[$(inds1...)] - arr[$(inds2...)])
        end))
        push!(add, res)
    end
    push!(out, :(+($(add...))))
    quote
        $(out...)
    end
end

[roflmaostc]

Thanks for your code, I think I can follow it conceptually.
Without the @generated I also can inspect the generated code, it looks exactly like what I want to have.

However, with @generated I can't call it. The following output:

julia> spatial_grad_square_ND(arr_2D)
ERROR: The function body AST defined by this @generated function is not pure. This likely means it contains a closure or comprehension.
Stacktrace:
 [1] top-level scope at REPL[5]:1

What's the problem here?

[mcabbott]

Oh I'm sorry, I thought I'd tried this but I messed up. There are indeed closures.

It might be possible to make this work with https://github.com/thautwarm/GeneralizedGenerated.jl, however right now it objects to type parameters -- @tullio has functions like act!(res::AbstractArray{T}, arg) where T, but it could perhaps say T = eltype(res) instead.

Or it might be possible to avoid closures, I'm not sure. I wanted something like this for #4 too but had not looked recently.

[mcabbott]

What you can do is to define a method for each ndims that you may need. Not as elegant but should work:

for N in 1:10
    inds1 = # work out pieces as before
    @eval function spatial_grad_square_ND(arr::AbstractArray{T,$N}) where T
        $(out...)
    end
end

[roflmaostc]

Thanks for your help, maybe it's not the best solution, but at least my idea is working and I learnt a lot about metaprogramming from your snippet.

For the time being I will use that, I could exchange it once there is a smarter solution...

Should I close the issue?

[mcabbott]

As you wish... I did a bit more fiddling but don't have a solution yet.

@tullio needs to define some functions, and this appears to be disallowed within the quote returned by @generated. But perhaps it can be done up front somehow.

[AriMKatz]

Can IRTools.jl help here?

[mcabbott]

That I don't know, have not really built any feeling of what it can & can't do. Can you say more?

(This comment #4 (comment) has a sketch of what @tullio produces, & why @generated doesn't like it.)

[AriMKatz]

It's basically magic. @MasonProtter would know better, but I have the feeling alot of this package would benefit from being pushed down to IR transforms.

[MasonProtter]

While possible, doing this at the IR level would be incredibly painful I think.

[roflmaostc]

Hey,

not really new information for that issue, but rather problems which appeared using the information from here.
See this discussion.

Might be helpful if someone runs into similar problems.

Cheers,

Felix