fix: don't support pre-moving the data

LuxDL · Oct 3, 2024 · d8c1129 · d8c1129 · avik-pal · Oct 3, 2024
1 parent aafe0d3
commit d8c1129
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/ext/MLDataDevicesMLUtilsExt.jl b/ext/MLDataDevicesMLUtilsExt.jl
@@ -4,7 +4,7 @@ using MLDataDevices: MLDataDevices, AbstractDevice, CPUDevice, CUDADevice, AMDGP
                      MetalDevice, oneAPIDevice, XLADevice, DeviceIterator
 using MLUtils: MLUtils, DataLoader
 
-for dev in (CPUDevice, CUDADevice, AMDGPUDevice, MetalDevice, oneAPIDevice)
+for dev in (CPUDevice, CUDADevice, AMDGPUDevice, MetalDevice, oneAPIDevice, XLADevice)
     @eval function (D::$(dev))(dataloader::DataLoader)
         if dataloader.parallel
             if dataloader.buffer
@@ -28,9 +28,6 @@ for dev in (CPUDevice, CUDADevice, AMDGPUDevice, MetalDevice, oneAPIDevice)
     end
 end
 
-# XXX: Doing it in parallel leads to deadlocks
-(D::XLADevice)(dataloader::DataLoader) = DeviceIterator(D, dataloader)
-
 function eachobsparallel(dev::AbstractDevice, data)
     return MLUtils.Loader(1:MLUtils.numobs(data)) do ch, i
         obs = MLUtils.getobs(data, i)

diff --git a/test/iterator_tests.jl b/test/iterator_tests.jl
@@ -62,8 +62,12 @@ end
     @testset "DataLoader: parallel=$parallel" for parallel in (true, false)
         @info "Testing DataLoader with parallel=$parallel"
         X = rand(Float64, 3, 33)
-        pre = DataLoader(dev(X); batchsize=13, shuffle=false, parallel)
         post = DataLoader(X; batchsize=13, shuffle=false, parallel) |> dev
+        if dev_type === XLADevice
+            pre = post # XXX: deadlocks and other shenanigans
+        else
+            pre = DataLoader(dev(X); batchsize=13, shuffle=false, parallel)
+        end
 
         for epoch in 1:2
             prev_pre, prev_post = nothing, nothing
@@ -86,8 +90,12 @@ end
         end
 
         Y = rand(Float64, 1, 33)
-        pre = DataLoader((; x=dev(X), y=dev(Y)); batchsize=13, shuffle=false, parallel)
         post = DataLoader((; x=X, y=Y); batchsize=13, shuffle=false, parallel) |> dev
+        if dev_type === XLADevice
+            pre = post # XXX: deadlocks and other shenanigans
+        else
+            pre = DataLoader((; x=dev(X), y=dev(Y)); batchsize=13, shuffle=false, parallel)
+        end
 
         for epoch in 1:2
             prev_pre, prev_post = nothing, nothing