Skip to content

Attempt at including offsets in kernel launch

Churavy, Valentin requested to merge github/fork/simone-silvestri/main into main

Created by: simone-silvestri

This PR tries to include offsets in kernel launches so that the Global indices returned by @index(Global, NTuple) and @index(Global, Linear) are offset by an offset argument.

Example:

julia> @kernel function show_index()
             i, j = @index(Global, NTuple)
             @show i, j
       end
show_index (generic function with 6 methods)

julia> show_index(CPU(), (2, 2), (3, 3), (-1, -2))()
(i, j) = (0, -1)
(i, j) = (1, -1)
(i, j) = (0, 0)
(i, j) = (1, 0)
(i, j) = (2, -1)
(i, j) = (2, 0)
(i, j) = (0, 1)
(i, j) = (1, 1)
(i, j) = (2, 1)

where the last argument (-1, -2) is the offsets to the global indices.

This PR constrains the offsetting of global indices on static kernel size at launch. @vchuravy I found it a bit difficult to implement arbitrary indices because of the division in blocks, which would have to be rethought. Aka, this is the easiest (probably not the most general) implementation of offsets. Let me know if you would rather it be implemented in another way.

Merge request reports