reversible version of weighted norm

[johnnychen94]

I'm trying to do a weighted loss sqeuclidean(X ./ W, ref) in NiLang:

@i function loss(out!::T, W!, X::AbstractArray{T}, ref::AbstractArray{T}) where T
    for i in 1:length(X)
        @routine begin
            tmp ← zero(T)
            tmp += X[i] / W![i]
            tmp -= ref[i]
        end
        out! += abs2(tmp)
        ~@routine
    end
end

sz = (64, 64)
W = Float64.(rand(2:5, sz))
X = rand(sz...)
ref = copy(X)

loss(0.0, W, X, ref)[1] ≈ mapreduce(abs2, +, X./W - ref) # true

but it's not very performant

julia> @btime loss(0.0, W, X, ref);
  104.582 μs (1 allocation: 48 bytes)

julia> @btime mapreduce(abs2, +, X./W - ref);
  5.471 μs (7 allocations: 64.23 KiB)

So I'm trying to do move the entire for loop in the call @routine, but apprarently something is wrong here:

@i function loss(out!::T, W!, X::AbstractArray{T}, ref::AbstractArray{T}) where T
    @routine begin
        tmp ← zero(T)
        for i in 1:length(X)
            tmp += X[i] / W![i]
            tmp -= ref[i]

            out! += abs2(tmp)

            tmp += ref[i]
            tmp -= X[i] / W![i]
        end
    end
    ~@routine
end

sz = (64, 64)
W = Float64.(rand(2:5, sz))
X = rand(sz...)
ref = copy(X)

loss(0.0, W, X, ref)[1] # 3.2666647165058293e-12

# as a reference, it should be non-zero
mapreduce(abs2, +, X./W - ref) # 653.6508272418502

I bet 104.582 μs isn't the best we can achieve but can't figure out a way here.

Any guidance would be very much helpful!

[GiggleLiu]

julia> @i function loss(out!::T, W!, X::AbstractArray{T}, ref::AbstractArray{T}) where T
           @invcheckoff for i in 1:length(X)
               @routine begin
                   tmp ← zero(T)
                   tmp += X[i] / W![i]
                   tmp -= ref[i]
               end
               out! += abs2(tmp)
               ~@routine
           end
       end

julia> @btime loss(0.0, W, X, ref);
  4.385 μs (1 allocation: 48 bytes)

julia> @btime mapreduce(abs2, +, X./W - ref);
  5.053 μs (7 allocations: 64.23 KiB)

This is why we need @invcheckoff