robustness

Author: luraess

src/fft/fft.jl

@@ -76,15 +76,19 @@ end
 # 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 make_execution_info(plan::ROCFFTPlan)
+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[]
-    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))
+    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 info
 end
 
 const xtypenames = ("complex forward", "complex inverse", "real forward", "real inverse")
@@ -238,37 +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}
-    info = make_execution_info(plan)
-    rocfft_execute(plan, [pointer(X),], C_NULL, info)
-    rocfft_execution_info_destroy(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}
-    info = make_execution_info(plan)
-    rocfft_execute(plan, [pointer(X),], [pointer(Y),], info)
-    rocfft_execution_info_destroy(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
-    info = make_execution_info(plan)
-    rocfft_execute(plan, [pointer(X),], [pointer(Y),], info)
-    rocfft_execution_info_destroy(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
-    info = make_execution_info(plan)
-    rocfft_execute(plan, [pointer(X),], [pointer(Y),], info)
-    rocfft_execution_info_destroy(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