Merge pull request #56 from yuehhua/scatter

Improve CPU scatter performance

Author: yuehhua

benchmark/pics/cpu_scatter.png

@@ -0,0 +1,33 @@
+using CSV, DataFrames
+using Gadfly
+import Cairo
+
+julia_bmk_file = joinpath("benchmark", "scatter_julia.tsv")
+python_bmk_file = joinpath("benchmark", "scatter_pytorch.tsv")
+
+bmk_jl = CSV.read(julia_bmk_file; delim='\t')
+bmk_py = CSV.read(python_bmk_file; delim='\t')
+
+bmk_jl[!, :framework] .= "geometricflux"
+bmk_py[!, :framework] .= "pytorch-scatter"
+
+bmk = vcat(bmk_jl, bmk_py)
+
+bmk[!, :min_time] .= bmk[!, :min_time]/1000
+bmk[!, :mean_time] .= bmk[!, :mean_time]/1000
+bmk[!, :max_time] .= bmk[!, :max_time]/1000
+
+function plot_benchmark(device)
+        DEVICE = uppercase(device)
+        p = plot(bmk[bmk[!,:device] .== device, :], x="sample", y="mean_time", color="framework",
+                 Geom.point, Geom.line, Scale.x_log2, Scale.y_log10,
+                 Guide.title("Scatter add performance on $(DEVICE)"),
+                 Guide.xlabel("Matrix Size"), Guide.ylabel("Time (μs)"),
+                 Coord.cartesian(xmin=4, xmax=21, ymin=1, ymax=7))
+
+        draw(SVG(joinpath("benchmark", "pics", "$(device)_scatter.svg"), 9inch, 6inch), p)
+        draw(PNG(joinpath("benchmark", "pics", "$(device)_scatter.png"), 9inch, 6inch), p)
+end
+
+plot_benchmark("cpu")
+plot_benchmark("gpu")

benchmark/pics/cpu_scatter.svg

@@ -0,0 +1,27 @@
+using CUDAdrv
+using GeometricFlux
+using Profile
+using ProfileView
+
+ENV["JULIA_NUM_THREADS"] = 1
+
+d = 50
+nbins = 20
+l = 2^20
+
+hist = zeros(Float32, d, nbins)
+δ = rand(Float32, d, l)
+idx = rand(1:nbins, l)
+scatter_add!(hist, δ, idx)
+
+@profile scatter_add!(hist, δ, idx)
+Profile.print()
+
+@profview scatter_add!(hist, δ, idx)
+
+# sudo nvprof --profile-from-start off julia benchmark/scatter.jl
+# sudo nvprof --profile-from-start off --print-gpu-trace julia --proj benchmark/scatter.jl
+# sudo chown $USER -R $HOME/.julia/
+
+# @profview scatter_add!(hist, δ, idx)
+# CUDAdrv.@profile scatter_add!(hist_gpu, δ_gpu, idx_gpu)

benchmark/plot.jl

@@ -6,7 +6,8 @@ using DataFrames
 using CSV
 using BenchmarkTools
 using BenchmarkTools: Trial, TrialEstimate, median, mean
-# using ProfileView
+
+ENV["JULIA_NUM_THREADS"] = 1
 
 d = 50
 nbins = 20
@@ -53,11 +54,3 @@ CSV.write("benchmark/scatter_julia.tsv", data; delim="\t")
 ## Benchmark
 # @benchmark scatter_add!($hist, $δ, $idx)
 # CuArrays.@time scatter_add!(hist_gpu, δ_gpu, idx_gpu)
-
-## Profiling
-# sudo nvprof --profile-from-start off julia benchmark/scatter.jl
-# sudo nvprof --profile-from-start off --print-gpu-trace julia --proj benchmark/scatter.jl
-# sudo chown yuehhua -R /home/yuehhua/.julia/
-
-# @profview scatter_add!(hist, δ, idx)
-# CUDAdrv.@profile scatter_add!(hist_gpu, δ_gpu, idx_gpu)

benchmark/profile.jl

@@ -7,7 +7,6 @@ using BenchmarkTools: Trial, TrialEstimate, median, mean
 py"""
 import torch
 import torch_scatter as sc
-torch.set_num_threads(12)
 cuda = torch.device("cuda:0")
 d = 50
 nbins = 20

benchmark/scatter.jl

@@ -8,7 +8,9 @@ for op = [:add, :sub, :mul, :div]
     @eval function $fn(ys::Array{T}, us::Array{T}, xs::Array{<:IntOrTuple}) where {T<:Real}
         @simd for k = 1:length(xs)
             k = CartesianIndices(xs)[k]
-            @inbounds ys[:, xs[k]...] .= $(name2op[op]).(view(ys, :, xs[k]...), view(us, :, k))
+            ys_v = view(ys, :, xs[k]...)
+            us_v = view(us, :, k)
+            @inbounds ys_v .= $(name2op[op]).(ys_v, us_v)
         end
         ys
     end
@@ -17,15 +19,19 @@ end
 function scatter_max!(ys::Array{T}, us::Array{T}, xs::Array{<:IntOrTuple}) where {T<:Real}
     @simd for k = 1:length(xs)
         k = CartesianIndices(xs)[k]
-        @inbounds ys[:, xs[k]...] .= max.(view(ys, :, xs[k]...), view(us, :, k))
+        ys_v = view(ys, :, xs[k]...)
+        us_v = view(us, :, k)
+        @inbounds view(ys, :, xs[k]...) .= max.(ys_v, us_v)
     end
     ys
 end
 
 function scatter_min!(ys::Array{T}, us::Array{T}, xs::Array{<:IntOrTuple}) where {T<:Real}
     @simd for k = 1:length(xs)
         k = CartesianIndices(xs)[k]
-        @inbounds ys[:, xs[k]...] .= min.(view(ys, :, xs[k]...), view(us, :, k))
+        ys_v = view(ys, :, xs[k]...)
+        us_v = view(us, :, k)
+        @inbounds ys_v .= min.(ys_v, us_v)
     end
     ys
 end
@@ -48,7 +54,9 @@ for op = [:add, :sub, :mul, :div]
                        xs::StaticArray{<:Tuple,<:IntOrTuple}) where {T<:Real}
         @simd for k = 1:length(xs)
             k = CartesianIndices(xs)[k]
-            @inbounds ys[:, xs[k]...] .= $(name2op[op]).(view(ys, :, xs[k]...), view(us, :, k))
+            ys_v = view(ys, :, xs[k]...)
+            us_v = view(us, :, k)
+            @inbounds ys_v .= $(name2op[op]).(ys_v, us_v)
         end
         ys
     end
@@ -58,7 +66,9 @@ function scatter_max!(ys::StaticArray{<:Tuple,T}, us::StaticArray{<:Tuple,T},
                       xs::StaticArray{<:Tuple,<:IntOrTuple}) where {T<:Real}
     @simd for k = 1:length(xs)
         k = CartesianIndices(xs)[k]
-        @inbounds ys[:, xs[k]...] .= max.(view(ys, :, xs[k]...), view(us, :, k))
+        ys_v = view(ys, :, xs[k]...)
+        us_v = view(us, :, k)
+        @inbounds ys_v .= max.(ys_v, us_v)
     end
     ys
 end
@@ -67,7 +77,9 @@ function scatter_min!(ys::StaticArray{<:Tuple,T}, us::StaticArray{<:Tuple,T},
                       xs::StaticArray{<:Tuple,<:IntOrTuple}) where {T<:Real}
     @simd for k = 1:length(xs)
         k = CartesianIndices(xs)[k]
-        @inbounds ys[:, xs[k]...] .= min.(view(ys, :, xs[k]...), view(us, :, k))
+        ys_v = view(ys, :, xs[k]...)
+        us_v = view(us, :, k)
+        @inbounds ys_v .= min.(ys_v, us_v)
     end
     ys
 end
@@ -163,7 +175,7 @@ end
             counts += sum(xs.==i) * (xs.==i)
         end
         @inbounds for ind = CartesianIndices(counts)
-            Δu[:, ind] ./= counts[ind]
+            view(Δu, :, ind) ./= counts[ind]
         end
         (Δ, Δu, nothing)
     end