Fix allocations with pow (#2308)

Author: charleskawczynski

src/Fields/broadcast.jl

@@ -313,7 +313,7 @@ end
 Base.Broadcast.broadcasted(
     ::typeof(Base.literal_pow),
     ::typeof(^),
-    f::Field,
+    f::Union{Field, Base.Broadcast.Broadcasted{<:AbstractFieldStyle}},
     ::Val{n},
 ) where {n} = Base.Broadcast.broadcasted(x -> Base.literal_pow(^, x, Val(n)), f)
 

test/Fields/unit_field.jl

@@ -147,6 +147,10 @@ function pow_n(f)
     @. f.x = f.x^2
     return nothing
 end
+function pow_n_bc(f)
+    @. f.x = (f.x * 2)^2
+    return nothing
+end
 @testset "Broadcasting with ^n" begin
     FT = Float32
     device = ClimaComms.CPUSingleThreaded() # fill is broken on gpu
@@ -160,6 +164,13 @@ end
         else
             @test p_allocated == 0 broken = (device isa ClimaComms.CUDADevice)
         end
+        pow_n_bc(f) # Compile first
+        p_allocated = @allocated pow_n_bc(f)
+        if space isa Spaces.SpectralElementSpace1D
+            @test p_allocated == 0
+        else
+            @test p_allocated == 0 broken = (device isa ClimaComms.CUDADevice)
+        end
     end
 end