# How to write a kernel to evaluate a function on an arbitrary dimensional grid?

**URL:** <https://discourse.julialang.org/t/how-to-write-a-kernel-to-evaluate-a-function-on-an-arbitrary-dimensional-grid/123359>\
**Category:** General Usage\
**Created:** [December 2, 2024, 2:33pm UTC](https://discourse.julialang.org/t/how-to-write-a-kernel-to-evaluate-a-function-on-an-arbitrary-dimensional-grid/123359 "2024-12-02T14:33:38Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![marcsgil](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/marcsgil/32/33908_2.png) [@marcsgil](https://discourse.julialang.org/u/marcsgil)\
**Post date:** [December 2, 2024, 2:33pm UTC](https://discourse.julialang.org/t/how-to-write-a-kernel-to-evaluate-a-function-on-an-arbitrary-dimensional-grid/123359/1 "2024-12-02T14:33:38Z")

</div>

Hi, I’m trying to write a kernel (using KernelAbstractions.jl) to evaluate a function `f` over an arbitrary dimensional Cartesian grid, which is represented by vectors `rs...`. To give a concrete example, for a 2D grid I can write the following:

```julia
using KernelAbstractions

@kernel function map_grid_kernel!(dest, f, x, y)
    J = @index(Global, NTuple)
    dest[J...] = f(x[J[1]], y[J[2]])
end

function map_grid!(dest, f, x, y)
    backend = get_backend(dest)
    f! = map_grid_kernel!(backend)
    f!(dest, f, x, y; ndrange=size(dest))
end

```

Then the following works just fine:

```julia
x = LinRange(-3, 3, 128)
y = copy(x)
dest = Array{Float64}(undef, length(x), length(y))

f(x...) = exp(-sum(abs2, x))

map_grid!(dest, f, x, y)

```

I’m having problems generalizing this to accept grids of arbitrary dimension. Although I’m able to write

```julia
@kernel function map_grid_arbitrary!(dest, f, rs...)
    J = @index(Global, NTuple)
    dest[J...] = f(ntuple(n -> rs[n][J[n]], ndims(dest))...)
end

function map_grid_arbitrary!(dest, f, rs...)
    backend = get_backend(dest)
    f! = map_grid_arbitrary!(backend)
    f!(dest, f, rs...; ndrange=size(dest))
end

```

this version allocates more then the previous one when called with the same arguments, and is also much slower. I couldn’t find a solution that matches the speed of the explicit 2D case. Does anyone has an idea of what can be done here? Thanks in advance!
