JuliaGPU / JuliaGPU/AcceleratedKernels.jl

Use matrix input in `searchsortedfirst`

Open
#61 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
204
Forks
15
Avg merge
10d 5h
Merged PRs (30d)
6

Description

It may be useful to make `AcceleratedKernels.searchsortedfirst` working also with `x` as a matrix input (while `v` still a vector input).

The minimum code change required would be:
1. change the input type signature in `searchsortedfirst!` and `searchsortedfirst`;
2. change the `@argcheck` in `searchsortedfirst!`.

The complete example is:

```
using AcceleratedKernels: foreachindex
using KernelAbstractions: get_backend, Backend
using ArgCheck: @argcheck

import AcceleratedKernels as AK

function _searchsortedfirst(v, x, lo::T, hi::T, comp) where T<:Integer
hi = hi + T(1)
len = hi - lo
@inbounds while len != 0x0
half_len = len >>> 0x1
m = lo + half_len
if comp(v[m], x)
lo = m + 0x1
len -= half_len + 0x1
else
hi = m
len = half_len
end
end
return lo
end

function _searchsortedfirst(v, x, lo::T, hi::T, ord::Base.Order.Ordering) where T<:Integer
hi = hi + T(1)
len = hi - lo
@inbounds while len != 0x0
half_len = len >>> 0x1
m = lo + half_len
if Base.Order.lt(ord, v[m], x)
lo = m + 0x1
len -= half_len + 0x1
else
hi = m
len = half_len
end
end
return lo
end

function searchsortedfirst!(
ix::AbstractVecOrMat, # INSTEAD OF AbstractVector
v::AbstractVector,
x::AbstractVecOrMat, # INSTEAD OF AbstractVector
backend::Backend=get_backend(x);

by=identity, lt=isless, rev::Bool=false,

# CPU settings with different default from `foreachindex`
min_elems::Int=1000,
kwargs...
)
# Simple sanity checks
@argcheck size(ix) == size(x) # INSTEAD OF lenght(ix) == lenght(x)

# Construct comparator
ord = Base.Order.ord(lt, by, rev)
comp = (x, y) -> Base.Order.lt(ord, x, y)

foreachindex(
x, backend;
min_elems, kwargs...
) do i
@inbounds ix[i] = _searchsortedfirst(v, x[i], firstindex(v), lastindex(v), comp)
end
end

function searchsortedfirst(
v::AbstractVector,
x::AbstractVecOrMat, # INSTEAD OF AbstractVector
backend::Backend=get_backend(x);
kwargs...
)
ix = similar(x, Int)
searchsortedfirst!(
ix, v, x, backend;
kwargs...
)
ix
end

# `x` vector input (case already implemented)

x_vec = range(0.25, 0.75, length=300) |> collect
v = range(0., 1., length=100) |> collect

ix1 = AK.searchsortedfirst(v, x_vec)
ix2 = Base.searchsortedfirst.(Ref(v), x_vec)
ix3 = searchsortedfirst(v, x_vec)

@show ix1 == ix2 == ix3

# `x` matrix input

x_mat = Matrix(reduce(hcat, [x_vec for _ in 1:5])')

# ix1 = AK.searchsortedfirst(v,x) not working
ix2_mat = Base.searchsortedfirst.(Ref(v), x_mat)
ix3_mat = searchsortedfirst(v, x_mat)

@show ix2_mat == ix3_mat
```

Output:

```
ix1 == ix2 == ix3 = true
ix2_mat == ix3_mat = true
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Locate the definitions of searchsortedfirst! and searchsortedfirst, then compare their current signatures and argument checks with the provided example. Run the vector and matrix examples against Base.searchsortedfirst; done means matrix input works and the displayed comparisons remain true.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend
Issue type
Feature
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.