Add multithread optimization to remap_ts

GridTools · Aug 23, 2024 · 9a8b90b · 9a8b90b
1 parent b1f539e
commit 9a8b90b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/GridTools.jl b/src/GridTools.jl
@@ -10,6 +10,7 @@ using Base: @propagate_inbounds
 using MacroTools
 using OffsetArrays: IdOffsetRange
 using Debugger
+using Polyester: @batch
 
 import Base.Broadcast: Extruded, Style, BroadcastStyle, ArrayStyle, Broadcasted
 
@@ -441,7 +442,7 @@ function remap_ts(
     #    end
     #end, CartesianIndices(map(len -> Base.OneTo(len), out_field_size)))
     out_field = zeros(eltype(field.data), out_field_size)
-    for position in eachindex(IndexCartesian(), out_field)
+    @batch for position in eachindex(IndexCartesian(), out_field)
         neighbor_exists, new_position =
             remap_position(Tuple(position), out_field_dims, offset, nb_ind, conn)
         if neighbor_exists