Skip to content

Commit

Permalink
cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
vchuravy committed Oct 21, 2024
1 parent 5fa3595 commit 102fbbd
Showing 1 changed file with 13 additions and 30 deletions.
43 changes: 13 additions & 30 deletions src/nditeration.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,19 @@ import Base.MultiplicativeInverses: SignedMultiplicativeInverse

# CartesianIndex uses Int instead of Int32

@eval EmptySMI() = $(Expr(:new, SignedMultiplicativeInverse{Int32}, Int32(0), typemax(Int32), 0%Int8, 0%UInt8))
@eval EmptySMI() = $(Expr(:new, SignedMultiplicativeInverse{Int32}, Int32(0), typemax(Int32), 0 % Int8, 0 % UInt8))
SMI(i) = i == 0 ? EmptySMI() : SignedMultiplicativeInverse{Int32}(i)

struct FastCartesianIndices{N} <: AbstractArray{CartesianIndex{N}, N}
inverses::NTuple{N, SignedMultiplicativeInverse{Int32}}
end

function FastCartesianIndices(indices::NTuple{N}) where N
inverses = map(i->SMI(Int32(i)), indices)
function FastCartesianIndices(indices::NTuple{N}) where {N}
inverses = map(i -> SMI(Int32(i)), indices)
FastCartesianIndices(inverses)
end

function Base.size(FCI::FastCartesianIndices{N}) where N
function Base.size(FCI::FastCartesianIndices{N}) where {N}
ntuple(Val(N)) do I
FCI.inverses[I].divisor
end
Expand All @@ -26,36 +26,36 @@ end
return CartesianIndex()
end

@inline function Base.getindex(iter::FastCartesianIndices{N}, I::Vararg{Int, N}) where N
@inline function Base.getindex(iter::FastCartesianIndices{N}, I::Vararg{Int, N}) where {N}
@boundscheck checkbounds(iter, I...)
index = map(iter.inverses, I) do inv, i
@inbounds getindex(Base.OneTo(inv.divisor), i)
end
CartesianIndex(index)
end

_ind2sub_recuse(::Tuple{}, ind) = (ind+1,)
_ind2sub_recuse(::Tuple{}, ind) = (ind + 1,)
function _ind2sub_recurse(indslast::NTuple{1}, ind)
@inline
Base.@_inline_meta
(_lookup(ind, indslast[1]),)
end

function _ind2sub_recurse(inds, ind)
@inline
Base.@_inline_meta
inv = inds[1]
indnext, f, l = _div(ind, inv)
(ind-l*indnext+f, _ind2sub_recurse(Base.tail(inds), indnext)...)
(ind - l * indnext + f, _ind2sub_recurse(Base.tail(inds), indnext)...)
end

_lookup(ind, inv::SignedMultiplicativeInverse) = ind+1
_lookup(ind, inv::SignedMultiplicativeInverse) = ind + 1
function _div(ind, inv::SignedMultiplicativeInverse)
inv.divisor == 0 && throw(DivideError())
div(ind%Int32, inv), 1, inv.divisor
div(ind % Int32, inv), 1, inv.divisor
end

function Base._ind2sub(inv::FastCartesianIndices, ind)
@inline
_ind2sub_recurse(inv.inverses, ind-1)
Base.@_inline_meta
_ind2sub_recurse(inv.inverses, ind - 1)
end

export _Size, StaticSize, DynamicSize, get
Expand Down Expand Up @@ -151,24 +151,7 @@ Base.length(range::NDRange) = length(blocks(range))
end

Base.@propagate_inbounds function expand(ndrange::NDRange{N}, groupidx::Integer, idx::Integer) where {N}
# This causes two sdiv operations, one for each Linear to CartesianIndex
return expand(ndrange, blocks(ndrange)[groupidx], workitems(ndrange)[idx])

# The formulation below saves one sdiv
# but leads to a different index order...
# previous: julia> expand(ndrange, 1, 32*32)
# CartesianIndex(32, 32)
# now: julia> expand(ndrange, 1, 32*32)
# CartesianIndex(1024, 1)
# B = blocks(ndrange)::CartesianIndices
# W = workitems(ndrange)::CartesianIndices
# Ind = ntuple(Val(N)) do I
# Base.@_inline_meta
# b = B.indices[I]
# w = W.indices[I]
# length(b) * length(w)
# end
# CartesianIndices(Ind)[(groupidx-1)* prod(size(W)) + idx]
end

Base.@propagate_inbounds function expand(ndrange::NDRange{N}, groupidx::CartesianIndex{N}, idx::Integer) where {N}
Expand Down

0 comments on commit 102fbbd

Please sign in to comment.