Skip to content

Change iteration order #102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dannys4 merged 4 commits into from
Feb 4, 2026
Merged

Change iteration order #102

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
03005ba
change iteration order
wheeheee Feb 3, 2026
fa3aa90
another `_mul_loop!` for function barrier
wheeheee Feb 3, 2026
e9410af
add some common files to .gitignore
wheeheee Feb 3, 2026
de3eb54
Revert "add some common files to .gitignore"
wheeheee Feb 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 19 additions & 16 deletions src/plan.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,10 +145,17 @@ function LinearAlgebra.mul!(y::AbstractArray{U,N}, p::FFTAPlan_cx{T,1}, x::Abstr
end
Rpre = CartesianIndices(size(x)[1:p.region[]-1])
Rpost = CartesianIndices(size(x)[p.region[]+1:end])
for Ipre in Rpre
for Ipost in Rpost
@views fft!(y[Ipre,:,Ipost], x[Ipre,:,Ipost], 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1)
end
_mul_loop!(y, x, Rpre, Rpost, p)
end

function _mul_loop!(
y::AbstractArray{U,N},
x::AbstractArray{T,N},
Rpre::CartesianIndices,
Rpost::CartesianIndices,
p::FFTAPlan_cx{T,1}) where {T,U,N}
for Ipost in Rpost, Ipre in Rpre
@views fft!(y[Ipre,:,Ipost], x[Ipre,:,Ipost], 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1)
end
end

Expand All @@ -172,10 +179,10 @@ function LinearAlgebra.mul!(y::AbstractArray{U,N}, p::FFTAPlan_cx{T,2}, x::Abstr
# Introduce function barrier here since the variables used in the loop ranges aren't inferred. This
# is partly because the region field of the plan is abstractly typed but even if that wasn't the case,
# it might be a bit tricky to construct the Rxs in an inferred way.
_mul_loop(y_tmp, y, x, p, R1, R2, R3, rows, cols)
_mul_loop!(y_tmp, y, x, p, R1, R2, R3, rows, cols)
end

function _mul_loop(
function _mul_loop!(
y_tmp::AbstractArray,
y::AbstractArray,
x::AbstractArray,
Expand All @@ -186,17 +193,13 @@ function _mul_loop(
rows::Int,
cols::Int
)
for I1 in R1
for I2 in R2
for I3 in R3
for k in 1:cols
@views fft!(y_tmp[:,k], x[I1,:,I2,k,I3], 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1)
end
for I3 in R3, I2 in R2, I1 in R1
for k in 1:cols
@views fft!(y_tmp[:,k], x[I1,:,I2,k,I3], 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1)
end

for k in 1:rows
@views fft!(y[I1,k,I2,:,I3], y_tmp[k,:], 1, 1, p.dir, p.callgraph[2][1].type, p.callgraph[2], 1)
end
end
for k in 1:rows
@views fft!(y[I1,k,I2,:,I3], y_tmp[k,:], 1, 1, p.dir, p.callgraph[2][1].type, p.callgraph[2], 1)
end
end
end
Expand Down
Loading