Fix square_view perf regression

JuliaSIMD · Apr 23, 2024 · 520cf84 · 520cf84 · chriselrod · Apr 23, 2024
1 parent f42feb9
commit 520cf84
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/src/ptr_array.jl b/src/ptr_array.jl
@@ -527,7 +527,8 @@ ArrayInterface.device(
   ::Type{<:AbstractStrideArray{<:Any,<:Any,R,<:Any,X}}
 ) where {R,X}
   i = findfirst(isone, R)
-  C = i === nothing ? -1 : (X.parameters[i] === Nothing ? i : -1)
+  C = i === nothing ? -1 : (((X.parameters[i] === Nothing) || (X.parameters[i] === One)) ? i : -1)
+  # C = i === nothing ? -1 : ((X.parameters[i] === Nothing) ? i : -1)
   StaticInt{C}()
 end
 ArrayInterface.contiguous_batch_size(::Type{<:AbstractStrideArray}) =

diff --git a/src/views.jl b/src/views.jl
@@ -96,11 +96,13 @@ function rank_to_sortperm(R::NTuple{N,Int}) where {N}
   sp
 end
 rank_to_sortperm(R) = sortperm(R)
-
-Base.@propagate_inbounds function square_view(A::PtrMatrix, i)
+_one_to_nothing(x::I) where {I} = StrideReset(x)
+_one_to_nothing(::One) = nothing
+Base.@propagate_inbounds function square_view(A::PtrMatrix{T,R,S,X,O}, i::I) where {T,R,S,X,O,I}
   # sizes = size(A)
   # @boundscheck i <= min(sizes[1], sizes[2]) || throw(BoundsError(A, (i, i)))
-  SquarePtrMatrix(pointer(A), i, static_strides(A), offsets(A))
+  x = map(_one_to_nothing, static_strides(A))
+  SquarePtrMatrix{T,R,I,typeof(x),O}(pointer(A), i, x, offsets(A))
 end
 # Base.@propagate_inbounds function square_view(A::AbstractMatrix, i)
 #   StrideArray(square_view(PtrArray(A), i), preserve_buffer(A))