Layouts - combining array structure and indexing #141
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,139 @@ | ||
|
|
||
| _as_index(x) = x | ||
| _as_index(x::OneTo) = static(1):length(x) | ||
| _as_index(x::StepRange) = OptionallyStaticStepRange(x) | ||
| _as_index(x::UnitRange) = OptionallyStaticUnitRange(x) | ||
| _as_index(x::OptionallyStaticRange) = x | ||
|
|
||
| """ | ||
| StrideLayout(A) | ||
|
|
||
| Produces an array whose elements correspond to the linear buffer position of `A`'s elements. | ||
| """ | ||
| struct StrideLayout{N,R,O1,S,A<:Tuple{Vararg{Any,N}}} <: AbstractArray2{Int,N} | ||
| rank::R | ||
| offset1::O1 | ||
| strides::S | ||
| axes::A | ||
| end | ||
|
|
||
| offset1(x::StrideLayout) = getfield(x, :offset1) | ||
| offsets(x::StrideLayout) = map(static_first, axes(x)) | ||
| axes(x::StrideLayout) = getfield(x, :axes) | ||
| @inline function axes(x::StrideLayout{N}, i::Int) where {N} | ||
| if i > N | ||
| return static(1):1 | ||
| else | ||
| return getfield(getfield(x, :axes), i) | ||
| end | ||
| end | ||
| @inline function axes(x::StrideLayout{N}, ::StaticInt{i}) where {N,i} | ||
| if i > N | ||
| return static(1):static(1) | ||
| else | ||
| return getfield(getfield(x, :axes), i) | ||
| end | ||
| end | ||
| strides(x::StrideLayout) = getfield(x, :strides) | ||
| stride_rank(x::StrideLayout) = getfield(x, :rank) | ||
|
|
||
| @inline function StrideLayout(x::DenseArray) | ||
| a = axes(x) | ||
| return StrideLayout( | ||
| stride_rank(x), | ||
| offset1(x), | ||
| size_to_strides(map(static_length, a), static(1)), | ||
| a | ||
| ) | ||
| end | ||
|
|
||
| # TODO optimize this | ||
| @inline function StrideLayout(x) | ||
| return StrideLayout( | ||
| stride_rank(x), | ||
| offset1(x), | ||
| strides(x), | ||
| axes(x) | ||
| ) | ||
| end | ||
|
|
||
| ############## | ||
| ### layout ### | ||
| ############## | ||
| layout(x, i) = layout(x) | ||
| layout(x, i::AbstractVector{<:Integer}) = _maybe_linear_layout(IndexStyle(x), x) | ||
| layout(x, i::Integer) = _maybe_linear_layout(IndexStyle(x), x) | ||
| layout(x, i::AbstractCartesianIndex{1}) = _maybe_linear_layout(IndexStyle(x), x) | ||
| function layout(x, i::AbstractVector{AbstractCartesianIndex{1}}) | ||
| return _maybe_linear_layout(IndexStyle(x), x) | ||
| end | ||
| _maybe_linear_layout(::IndexLinear, x) = _as_index(eachindex(x)) | ||
| _maybe_linear_layout(::IndexStyle, x) = layout(x) | ||
| layout(x::StrideLayout) = x | ||
| layout(x::LinearIndices) = x | ||
| layout(x::CartesianIndices) = x | ||
| function layout(x) | ||
| if defines_strides(x) | ||
| return StrideLayout(x) | ||
| else | ||
| return _layout_indices(IndexStyle(x), axes(x)) | ||
| end | ||
| end | ||
| function layout(x::Transpose) | ||
| if defines_strides(x) | ||
| return StrideLayout(x) | ||
| else | ||
| return Transpose(layout(parent(x))) | ||
| end | ||
| end | ||
| function layout(x::Adjoint{T}) where {T<:Number} | ||
| if defines_strides(x) | ||
| return StrideLayout(x) | ||
| else | ||
| return Transpose(layout(parent(x))) | ||
| end | ||
| end | ||
| function layout(x::PermutedDimsArray{T,N,perm,iperm}) where {T,N,perm,iperm} | ||
| if defines_strides(x) | ||
| return StrideLayout(x) | ||
| else | ||
| p = layout(parent(x)) | ||
| return PermutedDimsArray{eltype(p),ndims(p),perm, iperm,typeof(p)}(p) | ||
| end | ||
| end | ||
| function layout(x::SubArray) | ||
| if defines_strides(x) | ||
| return StrideLayout(x) | ||
| else | ||
| return @inbounds(view(layout(parent(x)), x.indices...)) | ||
| end | ||
| end | ||
| _layout_indices(::IndexStyle, axs) = CartesianIndices(axs) | ||
| _layout_indices(::IndexLinear, axs) = LinearIndices(axs) | ||
|
There was a problem hiding this comment. Whether to use linear or Cartesian indexing is of course context dependent (in the rewrite, I'll add support for switching between representations). Is there a "combine layouts"? I'll get LoopVectorization to be able to transform between these correctly in the rewrite. @avx begin
C .= 0
for n in indices((C,B),2), m in indices((C,A),1), for k in indices((A,B),(2,1))
C[m,n] += A[m,k] * B[k,n]
end
endand to have LoopVectorization be capable of converting the linear indexing of There was a problem hiding this comment. What you're describing is basically what I was getting at with the reference to MLIR. The final "layout" after code gen would be the result of scoping (for summation this would be unordered access to all elements once) which informs how to combine each layer of the arrays. |
||
|
|
||
| """ | ||
| buffer(x) | ||
|
|
||
| Return the raw buffer for `x`, stripping any additional info (structural, indexing, | ||
| metadata, etc.). | ||
| """ | ||
| buffer(x) = x | ||
| @inline buffer(x::PermutedDimsArray) = buffer(parent(x)) | ||
| @inline buffer(x::Transpose) = buffer(parent(x)) | ||
| @inline buffer(x::Adjoint) = buffer(parent(x)) | ||
| @inline buffer(x::SubArray) = buffer(parent(x)) | ||
|
|
||
|
|
||
| """ allocate_memory(::AbstractDevice, ::Type{T}, length::Union{StaticInt,Int}) """ | ||
| allocate_memory(::CPUPointer, ::Type{T}, ::StaticInt{N}) where {T,N} = Ref{NTuple{N,T}} | ||
| allocate_memory(::CPUPointer, ::Type{T}, n::Int) where {T} = Vector{T}(undef, n) | ||
| allocate_memory(::CPUTuple, ::Type{T}, ::StaticInt{N}) where {T,N} = Ref{NTuple{N,T}} | ||
|
|
||
|
|
||
| """ dereference(::AbstractDevice, x) """ | ||
| dereference(::CPUPointer, x) = x | ||
| dereference(::CPUTuple, x::Ref) = x[] | ||
|
|
||
| """ initialize(data, layout) """ | ||
| function initialize end | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
On the VectorizationBase side of things, it'd be nice if
StridedPointercould be aStrideLayout+ aPtr.However,
StrideLayoutis missing some things, most importantly the contiguous axis, and also hasaxeswhich it does not want (although it does want to storemap(static_first, axes(x))).Also,
StridedPointer+axes==PtrArray.StrideLayoutalready subtypesAbstractArray2, so it is almost there itself.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
How about something that is indexable but isn't necessarily an iterator like this
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't have any actual support anywhere for
ArrayInterface.contiguous_batch_sizeyet, so we could add that later, but probably best to start carrying that around now as well.But this looks good.