EnzymeAD / EnzymeAD/Reactant.jl

HaloArrays

Open
#1,112 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
370
Forks
74
Avg merge
18h 47m
Merged PRs (30d)
30

Description

Crazy implementation of halo arrays ***absolutely not tested and very likely inefficient***

```julia
abstract type AbstractHaloArray{T,N} <: AbstractArray{T,N} end

######################## 1D
#
# | |xxxxxx| |

struct HaloArray1D{T,A1<:AbstractArray{T,1},A2<:AbstractArray{T,1},A3<:AbstractArray{T,1}} <: AbstractHaloArray{T,1}
left::A1
inner::A2
right::A3
halo::Int
function HaloArray1D(left::A1, inner::A2, right::A3, halo=size(left, 1)) where {T,A1<:AbstractArray{T,1},A2<:AbstractArray{T,1},A3<:AbstractArray{T,1}}
@assert size(left, 1) == size(right, 1) == halo
new{T,A1,A2,A3}(left, inner, right, halo)
end
end

function Base.getindex(v::HaloArray1D, idx::Int)
if idx <= v.halo
return v.left[idx]
elseif idx - v.halo <= size(v.inner, 1)
return v.inner[idx - v.halo]
else
return v.right[idx - v.halo - size(v.inner, 1)]
end
end

function Base.setindex!(v::HaloArray1D{T}, value::T, idx::Int) where {T}
if 0 < idx - v.halo <= size(v.inner, 1)
v.inner[idx - v.halo] = value
end
end

######################## 2D
# ______________
# | | | |
# | |______| |
# | |xxxxxx| |
# | |xxxxxx| |
# | |______| |
# | | | |
# |__|______|__|

struct HaloArray2D{T,A1<:AbstractArray{T,2},A2<:AbstractArray{T,2},A3<:AbstractArray{T,2},A4<:AbstractArray{T,2},A5<:AbstractArray{T,2}} <: AbstractHaloArray{T,2}
top::A1
left::A2
inner::A3
right::A4
bottom::A5
halo::Int
function HaloArray2D(top::A1, left::A2, inner::A3, right::A4, bottom::A5, halo=size(top, 1)) where {T,A1<:AbstractArray{T,2},A2<:AbstractArray{T,2},A3<:AbstractArray{T,2},A4<:AbstractArray{T,2},A5<:AbstractArray{T,2}}
@assert size(top) == size(bottom) == (halo, size(inner, 2))
@assert size(left) == size(right) == (size(inner, 1) + 2 * halo, halo)
new{T,A1,A2,A3,A4,A5}(top, left, inner, right, bottom, halo)
end
end

function Base.getindex(m::HaloArray2D, idx1::Int, idx2::Int)
if idx2 <= m.halo
return m.left[idx1, idx2]
elseif idx2 - m.halo <= size(m.inner, 2)
if idx1 <= m.halo
return m.top[idx1, idx2 - m.halo]
elseif idx1 - m.halo <= size(m.inner, 1)
return m.inner[idx1 - m.halo, idx2 - m.halo]
else
return m.bottom[idx1 - m.halo - size(m.inner, 1), idx2 - m.halo]
end
else
return m.right[idx1, idx - m.halo - size(m.inner, 2)]
end
end

function Base.setindex!(v::HaloArray2D{T}, value::T, idx1::Int, idx2::Int) where {T}
if 0 < idx1 - m.halo <= size(m.inner, 1) && 0 < idx2 - m.halo <= size(m.inner, 2)
m.inner[idx1 - m.halo, idx2 - m.halo] = value
end
end

######################## 3D
# _____________
# / / / /|
# / /______/ / |
# / /xxxxxx/ / |
# / /xxxxxx/ / |
# / /______/ / /| |
# / / / / /x| |
# /__/______/__/ /xx/ |
# | | | | |x/ /
# | |______| | |/ /
# | |xxxxxx| | | /
# | |xxxxxx| | /
# | |______| | /
# | | | |/
# |__|______|__|

struct HaloArray3D{T,A1<:AbstractArray{T,3},A2<:AbstractArray{T,3},A3<:AbstractArray{T,3},A4<:AbstractArray{T,3},A5<:AbstractArray{T,3},A6<:AbstractArray{T,3},A7<:AbstractArray{T,3}} <: AbstractHaloArray{T,3}
top::A1
left::A2
front::A4
inner::A4
back::A5
right::A6
bottom::A7
halo::Int
function HaloArray3D(top::A1, left::A2, front::A3, inner::A4, back::A5, right::A6, bottom::A7, halo=size(top, 1)) where {T,A1<:AbstractArray{T,3},A2<:AbstractArray{T,3},A3<:AbstractArray{T,3},A4<:AbstractArray{T,3},A5<:AbstractArray{T,3},A6<:AbstractArray{T,3},A7<:AbstractArray{T,3}}
@assert size(top) == size(bottom) == (halo, size(inner, 2), size(inner, 3) + 2 * halo)
@assert size(front) == size(back) == (size(inner, 1) + 2 * halo, size(inner, 2) + 2 * halo, halo)
@assert size(left) == size(right) == (size(inner, 1) + 2 * halo, halo, size(inner, 3) + 2 * halo)
new{T,A1,A2,A3,A4,A5}(top, left, inner, right, bottom, halo)
end
end

function Base.getindex(t::HaloArray3D, idx1::Int, idx2::Int, idx3::Int)
if idx2 <= t.halo
return t.left[idx1, idx2, idx3]
elseif idx2 - t.halo <= size(t.inner, 2)
if idx1 <= t.halo
return t.top[idx1, idx2 - t.halo, idx3]
elseif idx1 - t.halo <= size(t.inner, 1)
if idx3 < t.halo
return t.front[idx1 - t.halo, idx2 - t.halo, idx]
elseif idx3 - t.halo <= size(t.inner, 3)
return t.inner[idx1 - t.halo, idx2 - t.halo, idx3 - t.halo]
else
return t.back[idx1 - t.halo, idx2 - t.halo, idx - t.halo - size(t.inner, 3)]
end
else
return t.bottom[idx1 - t.halo - size(t.inner, 1), idx2 - t.halo, idx3]
end

else
return t.right[idx1, idx - t.halo - size(t.inner, 2), idx3]
end
end

function Base.setindex!(t::HaloArray3D{T}, value::T, idx1::Int, idx2::Int, idx3::Int) where {T}
if 0 < idx1 - t.halo <= size(t.inner, 1) && 0 < idx2 - t.halo <= size(t.inner, 2) && 0 < idx3 - t.halo <= size(t.inner, 3)
t.inner[idx1 - t.halo, idx2 - t.halo, idx3 - t.halo] = value
end
end

```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.