Repository navigation
Introduce a isa_wrapped_array function #460
Description
Most details are in EnzymeAD/Reactant.jl#369 (comment). I will copy over the important parts.
We introduce a function isa_wrapped_array that downstream packages can use to mark that their array type wraps another array. Using a union type from Adapt doesn't solve this problem, because that fundamentally doesn't extend to new array types.
With this function, we can override functions inside our custom interpreter. Consider this simple example of extending LinearAlgebra.diag
Base.Experimental.@overlay REACTANT_METHOD_TABLE function LinearAlgebra.diag(
x::AbstractArray{T,2}, k::Integer=0
) where {T}
if isa_wrapped_array(x) && ancestor(x) isa TracedRArray
y = materialize_traced_array(x) # convert it to a known type
return diag(y, k)
else
# invoke diag(x) on NativeInterpreter
end
end
Activity
rafaqz commentedon Dec 15, 2024
How doesmaterialize_traced_arraywork? is the idea to get the wrappedTracedArrayout ofx?Can that work through multiple layers of wrappers?Ok think I get it.
ancestorwill recurse down through wrapped arrays to find the first non-wrapper array?(and
is_wrapped_arraymust mean wraps without changing values or size?)avik-pal commentedon Dec 15, 2024
Correct. I am still debating if it should be
ancestororancestors. For example,Tridiagonalstores 3 arrays underneath to it returns a 3-tuple. For aDiagonalorSubArraywe return a 1-tuple.Broader than that. Any array type that wraps over "primitive" array types (examples being
Array,JLArray,CuArray, etc.) would count. For example:Tridiagonalwould count as a wrapper even though it's size is not same as the size of the underlying buffersrafaqz commentedon Dec 15, 2024
Doesn't that risk losing information from the wrapper? E.g. would
SubArraybe a wrapper?avik-pal commentedon Dec 15, 2024
Do you mean in the pseudo-code implementation of
diagI shared?rafaqz commentedon Dec 15, 2024
I meant this kind of method generally, but yes, I assume getting the ancestor through a SubArray would be broken in your example.
ChrisRackauckas commentedon Dec 15, 2024
Makes sense to me.
avik-pal commentedon Dec 16, 2024
Yes it is lossy and that is intentional. Let me summarize the functions first and then I can clarify why it being lossy is fine.
Functions Introduced
isa_wrapper_typeis being used to extend what a wrapper type is by external packages. For most cases this is essentiallyparent(x) !== x, but this doesn't hold for Tridiagonal, Diagonal, etc. (as demonstrated above)ancestorsis a function giving us the leaves of a wrapper type. An example:Calling
ancestors(z)should return(z.dl.v, z.d.v, z.du.v). Why is this useful?TracedRArrayand then use that information for a specialized dispatch.AnyCuArraybut if the user has aHermitian(CuArray(rand(4, 4))then this returns an incorrect result (CPUDevice()) since AnyCuArray is a Union type.Why is lossy information still useful?
Consider any array type that doesn't support
fast_scalar_indexing. If we don't detect the ancestors correctly it will fallback to Base implementations which will often simply use a loop over the values.If we detect that this is a special ancestor type (say TracedRArray), we do a fallback implementation where we materialize it into a dense array. While this is not going to be optimal for performance (in general), it is almost certainly not going to error. Note that in cases where a loop doesn't error, materializing a dense array will still be more performant than falling back to a loop.
I will open a draft PR on this. It might help us give correct results for https://github.com/JuliaArrays/ArrayInterface.jl/blob/master/ext/ArrayInterfaceGPUArraysCoreExt.jl#L8 for cases not covered in the union type
rafaqz commentedon Dec 16, 2024
It could also do
ancestors(TracedArray, A)to be more general. That would allow getting closer ancestors like Diagonal withancestors(Diagonal, A)rather than the vector it wraps.