Fix rocFFT error with multiple tasks

.buildkite/pipeline.yml

@@ -43,7 +43,7 @@ steps:
       rocmgpu: "*"
     if: build.message !~ /\[skip tests\]/
     command: "julia --project -e 'using Pkg; Pkg.update()'"
-    timeout_in_minutes: 60
+    timeout_in_minutes: 90
     env:
       JULIA_NUM_THREADS: 4
       JULIA_AMDGPU_CORE_MUST_LOAD: "1"
@@ -62,7 +62,7 @@ steps:
       rocmgpu: "*"
     if: build.message !~ /\[skip tests\]/
     command: "julia --project -e 'using Pkg; Pkg.update()'"
-    timeout_in_minutes: 60
+    timeout_in_minutes: 90
     env:
       JULIA_NUM_THREADS: 4
       JULIA_AMDGPU_CORE_MUST_LOAD: "1"
@@ -81,7 +81,7 @@ steps:
       rocmgpu: "*"
     if: build.message !~ /\[skip tests\]/
     command: "julia --project -e 'using Pkg; Pkg.update()'"
-    timeout_in_minutes: 60
+    timeout_in_minutes: 90
     env:
       JULIA_NUM_THREADS: 4
       JULIA_AMDGPU_CORE_MUST_LOAD: "1"
@@ -100,7 +100,7 @@ steps:
       rocmgpu: "*"
     if: build.message !~ /\[skip tests\]/
     command: "julia --project -e 'using Pkg; Pkg.update()'"
-    timeout_in_minutes: 60
+    timeout_in_minutes: 90
     soft_fail: true
     env:
       JULIA_NUM_THREADS: 4

src/fft/fft.jl

@@ -15,7 +15,6 @@ function AMDGPU.unsafe_free!(plan::ROCFFTPlan)
         unsafe_free!(plan.buffer)
     end
     unsafe_free!(plan.workarea)
-    rocfft_execution_info_destroy(plan.execution_info)
 end
 
 const ROCFFT_FORWARD = true
@@ -27,9 +26,7 @@ const ROCFFT_INVERSE = false
 # K is flag for forward/inverse
 mutable struct cROCFFTPlan{T,K,inplace,N,R,B} <: ROCFFTPlan{T, K, inplace}
     handle::rocfft_plan
-    stream::HIPStream
     workarea::ROCVector{Int8}
-    execution_info::rocfft_execution_info
     sz::NTuple{N, Int} # Julia size of input array
     osz::NTuple{N, Int} # Julia size of output array
     xtype::rocfft_transform_type
@@ -47,26 +44,14 @@ mutable struct cROCFFTPlan{T,K,inplace,N,R,B} <: ROCFFTPlan{T, K, inplace}
         xtype::rocfft_transform_type, region::NTuple{R,Int},
         buffer::B, input_sz_as_key::Bool, key_T::Type,
     ) where {T,K,inplace,N,R,B}
-        info_ref = Ref{rocfft_execution_info}()
-        rocfft_execution_info_create(info_ref)
-        info = info_ref[]
-
-        # assign to the current stream
-        stream = AMDGPU.stream()
-        rocfft_execution_info_set_stream(info, stream)
-        if length(workarea) > 0
-            rocfft_execution_info_set_work_buffer(info, workarea, length(workarea))
-        end
-        p = new(handle, stream, workarea, info, size(X), sizey, xtype, region, buffer, input_sz_as_key, key_T)
+        p = new(handle, workarea, size(X), sizey, xtype, region, buffer, input_sz_as_key, key_T)
         return finalizer(AMDGPU.unsafe_free!, p)
     end
 end
 
 mutable struct rROCFFTPlan{T,K,inplace,N,R,B} <: ROCFFTPlan{T,K,inplace}
     handle::rocfft_plan
-    stream::HIPStream
     workarea::ROCVector{Int8}
-    execution_info::rocfft_execution_info
     sz::NTuple{N,Int} # Julia size of input array
     osz::NTuple{N,Int} # Julia size of output array
     xtype::rocfft_transform_type
@@ -83,28 +68,27 @@ mutable struct rROCFFTPlan{T,K,inplace,N,R,B} <: ROCFFTPlan{T,K,inplace}
         sizey::Tuple, xtype::rocfft_transform_type, region::NTuple{R,Int},
         buffer::B, input_sz_as_key::Bool, key_T::Type,
     ) where {T,inplace,N,K,R,B}
-        info_ref = Ref{rocfft_execution_info}()
-        rocfft_execution_info_create(info_ref)
-        info = info_ref[]
-
-        stream = AMDGPU.stream()
-        rocfft_execution_info_set_stream(info, stream)
-        if length(workarea) > 0
-            rocfft_execution_info_set_work_buffer(info, workarea, length(workarea))
-        end
-        p = new(handle, stream, workarea, info, size(X), sizey, xtype, region, buffer, input_sz_as_key, key_T)
+        p = new(handle, workarea, size(X), sizey, xtype, region, buffer, input_sz_as_key, key_T)
         return finalizer(unsafe_free!, p)
     end
 end
 
-function update_stream!(plan::ROCFFTPlan)
-    new_stream = AMDGPU.stream()
-    if plan.stream != new_stream
-        plan.stream = new_stream
-        info = plan.execution_info
-        rocfft_execution_info_set_stream(info, new_stream)
+# Create a per-call execution_info with the current task's stream and workbuffer.
+# rocfft_execution_info_set_stream must be called on the same OS thread as rocfft_execute;
+# creating it fresh per-call avoids any OS-thread binding issues from task migration.
+function with_execution_info(f::F, plan::ROCFFTPlan) where {F}
+    info_ref = Ref{rocfft_execution_info}()
+    rocfft_execution_info_create(info_ref)
+    info = info_ref[]
+    try
+        rocfft_execution_info_set_stream(info, AMDGPU.stream())
+        if length(plan.workarea) > 0
+            rocfft_execution_info_set_work_buffer(info, plan.workarea, length(plan.workarea))
+        end
+        return f(info)
+    finally
+        rocfft_execution_info_destroy(info)
     end
-    return
 end
 
 const xtypenames = ("complex forward", "complex inverse", "real forward", "real inverse")
@@ -258,33 +242,44 @@ function assert_applicable(p::ROCFFTPlan{T,K}, X::ROCArray{T}, Y::ROCArray{Ty})
 end
 
 function unsafe_execute!(plan::cROCFFTPlan{T,K,true,N}, X::ROCArray{T,N}) where {T,K,N}
-    update_stream!(plan)
-    rocfft_execute(plan, [pointer(X),], C_NULL, plan.execution_info)
+    in_buffer = [pointer(X)]
+    with_execution_info(plan) do info
+        rocfft_execute(plan, in_buffer, C_NULL, info)
+    end
 end
 
 function unsafe_execute!(
     plan::cROCFFTPlan{T,K,false,N}, X::ROCArray{T,N}, Y::ROCArray{T},
 ) where {T,N,K}
-    update_stream!(plan)
-    rocfft_execute(plan, [pointer(X),], [pointer(Y),], plan.execution_info)
+    in_buffer = [pointer(X)]
+    out_buffer = [pointer(Y)]
+    with_execution_info(plan) do info
+        rocfft_execute(plan, in_buffer, out_buffer, info)
+    end
 end
 
 function unsafe_execute!(
     plan::rROCFFTPlan{T,ROCFFT_FORWARD,false,N},
     X::ROCArray{T,N}, Y::ROCArray{<:rocfftComplexes,N},
 ) where {T<:rocfftReals,N}
     @assert plan.xtype == rocfft_transform_type_real_forward
-    update_stream!(plan)
-    rocfft_execute(plan, [pointer(X),], [pointer(Y),], plan.execution_info)
+    in_buffer = [pointer(X)]
+    out_buffer = [pointer(Y)]
+    with_execution_info(plan) do info
+        rocfft_execute(plan, in_buffer, out_buffer, info)
+    end
 end
 
 function unsafe_execute!(
     plan::rROCFFTPlan{T,ROCFFT_INVERSE,false,N},
     X::ROCArray{T,N}, Y::ROCArray{<:rocfftReals,N},
 ) where {T<:rocfftComplexes,N}
     @assert plan.xtype == rocfft_transform_type_real_inverse
-    update_stream!(plan)
-    rocfft_execute(plan, [pointer(X),], [pointer(Y),], plan.execution_info)
+    in_buffer = [pointer(X)]
+    out_buffer = [pointer(Y)]
+    with_execution_info(plan) do info
+        rocfft_execute(plan, in_buffer, out_buffer, info)
+    end
 end
 
 function LinearAlgebra.mul!(y::ROCArray{Ty}, p::ROCFFTPlan{T,K,false}, x::ROCArray{T}) where {T,Ty,K}

test/hip_rocarray/fft.jl

@@ -357,10 +357,15 @@ end
 
     Y = p * X
 
-    task = Threads.@spawn d_p * d_X  # executes FFT on separate AMDGPU stream
-    d_Y = fetch(task)
+    for _ in 1:10
+        task = Threads.@spawn begin # executes FFT on separate AMDGPU stream
+            yield()
+            d_p * d_X
+        end
+        d_Y = fetch(task)
 
-    @test isapprox(collect(d_Y), Y; rtol=MYRTOL, atol=MYATOL)
+        @test isapprox(collect(d_Y), Y; rtol=MYRTOL, atol=MYATOL)
+    end
 end
 
 end # testset FFT