`mapreducedim!` is super slow

[simeonschaub]

Reductions on CLArrays seem to be almost 100x slower than Base (This is with the pocl CPU backend):

julia> using OpenCL, pocl_jll

julia> X = rand(Float32, 1000, 1000);

julia> X′ = CLArray(X);

julia> @benchmark sum(X; dims = 1)
BenchmarkTools.Trial: 10000 samples with 1 evaluation per sample.
 Range (min … max):   58.380 μs … 922.617 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):      90.710 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   103.658 μs ±  35.143 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

        ▄▃▁ ▃▇█▃                                                 
  ▁▁▁▂▄█████████▆▆▆▅▅▄▄▃▂▃▃▃▃▄▃▃▂▂▂▂▃▃▃▃▃▂▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁ ▃
  58.4 μs          Histogram: frequency by time          212 μs <

 Memory estimate: 4.02 KiB, allocs estimate: 3.

julia> @benchmark OpenCL.synchronize(sum(X′; dims = 1))
BenchmarkTools.Trial: 653 samples with 1 evaluation per sample.
 Range (min … max):  5.585 ms … 12.908 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     7.424 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):   7.604 ms ±  1.126 ms  ┊ GC (mean ± σ):  0.00% ± 0.00%

       ▁▁▁▂▃▁▄█ ▁▇▇▃▄▄▄▃▄▂▃▂▃▁▁▃▁                             
  ▃▂▃▅▆██████████████████████████▇▆█▇▄▇▇█▆▅▇▄▁▄▅▄▄▃▃▅▄▃▂▂▂▂▃ ▅
  5.59 ms        Histogram: frequency by time        10.6 ms <

 Memory estimate: 22.52 KiB, allocs estimate: 247.

Is there any low-hanging fruit in terms of optimizations here?

[maleadt]

That's curious, I've seen it outperform Base when doing some initial testing last year.

The mapreduce implementation is taken from CUDA.jl, so given PoCL's 'portable performance' promise we shouldn't be seeing those kinds of slowdowns. If anybody wants to look into this, I'd try and see which parallelization method PoCL is using here (try forcing one using POCL_WORK_GROUP_METHOD) and see if the performance changes. There's definitely improvements to be made to our mapreduce implementation, but nothing that should yield a 100x speed-up.

[simeonschaub]

Here are some results with different POCL_WORK_GROUP_METHODs:

Default method (loopvec)

julia> @benchmark OpenCL.synchronize(sum(X′; dims = 1))
BenchmarkTools.Trial: 1298 samples with 1 evaluation per sample.
 Range (min … max):  2.214 ms …   5.297 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     3.794 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   3.841 ms ± 295.490 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

                                 ▁▄▆▇██▆▄▂                     
  ▂▁▁▁▁▂▂▂▂▂▂▁▁▁▁▁▁▂▁▁▁▁▁▁▂▂▂▃▄▅▇██████████▆▆▅▅▅▄▄▄▃▄▃▃▃▃▃▂▂▂ ▃
  2.21 ms         Histogram: frequency by time        4.77 ms <

 Memory estimate: 22.19 KiB, allocs estimate: 248.

cbs

BenchmarkTools.Trial: 661 samples with 1 evaluation per sample.
 Range (min … max):  4.527 ms …  10.373 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     7.426 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   7.548 ms ± 610.995 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

                              ▁▁▅▄▇▇█▇▄▅▁▂                     
  ▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▂▁▁▁▁▁▁▁▂▃▁▅▅████████████▇▆█▆▇▅▆▅▄▄▄▃▄▃▄▃▂▂▃ ▄
  4.53 ms         Histogram: frequency by time        9.33 ms <

 Memory estimate: 22.19 KiB, allocs estimate: 248.

loops

BenchmarkTools.Trial: 1846 samples with 1 evaluation per sample.
 Range (min … max):  1.490 ms …   5.119 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     2.652 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   2.698 ms ± 320.992 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

                         ▂▂▆▇█▇▇▅▅▅▃▃▂                         
  ▂▂▂▂▂▃▂▂▁▁▁▂▂▂▂▂▁▃▂▄▄▆▇██████████████▇▅▅▅▅▆▄▄▄▄▃▃▂▃▂▂▂▂▁▂▁▂ ▄
  1.49 ms         Histogram: frequency by time        3.76 ms <

 Memory estimate: 22.19 KiB, allocs estimate: 248.

I also tried it with the cpu-minimal device

cpu-minimal

BenchmarkTools.Trial: 386 samples with 1 evaluation per sample.
 Range (min … max):  11.950 ms … 26.817 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     12.595 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):   12.971 ms ±  1.400 ms  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▄▁ █▆▂▂▃▄▂▂▁                                                 
  ██▇██████████▇▅▆▇▁▄▅▄▁▁▅▁▁▁▁▁▁▁▄▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▄▁▁▁▁▁▁▄ ▆
  12 ms        Histogram: log(frequency) by time      20.9 ms <

 Memory estimate: 22.19 KiB, allocs estimate: 248.

And just for comparison:

Base.mapreducedim!

julia> @benchmark sum(X; dims = 1)
BenchmarkTools.Trial: 10000 samples with 1 evaluation per sample.
 Range (min … max):  33.613 μs … 334.269 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     42.350 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   42.562 μs ±   6.059 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

                         ▅█▃▃▅▁                                 
  ▂▂▂▁▂▂▃▆▆▄▅▇▆▅▆▅▄▄▄▆█▆▇████████▇▆▅▅▄▄▄▄▄▃▃▃▃▃▃▂▃▂▂▂▂▂▂▂▂▂▂▂▂ ▄
  33.6 μs         Histogram: frequency by time         54.7 μs <

 Memory estimate: 4.02 KiB, allocs estimate: 3.

So overall, the loops method seems to be a little bit faster than the default, but it's still orders of magnitude slower than the method from Base. Do note that this is on my laptop (Ryzen AI 7 PRO 360), so timings may vary. Could this be something device specific, or can others reproduce this as well? Is there anything else I could look at in pocl?

[maleadt]

Interesting. Maybe @vchuravy can shed some light on the benchmarks he ran a while ago.

Sadly, I don't have the time to look into this right now. Maybe you can vibe-port a known-good OpenCL sum implementation to see how that fares? Our much more generic mapreduce (N-dimensional, using Cartesian indices necessitating expensive integer divisions at run time now that #348 is reverted) is known to perform badly on certain GPUs, maybe the same applies to CPUs.

Also consider that there may be an API-level issue, so it may be worth to get PoCL to dump kernel times as well, just to be sure (e.g. by using POCL_TRACING).

[simeonschaub]

It doesn't seem like it's an API-level issue, in this instance the kernel times were over 3 ms:

Let me try with a different sum implementation

[simeonschaub]

The following implementation is really fast:

using OpenCL, pocl_jll, BenchmarkTools

const source = """
__kernel void sum_columns_subgroup(
    __global const float* X,     // Input matrix (column-major)
    __global float* result,      // Output (length N)
    const int M,                 // Rows
    const int N                  // Columns
) {
    int col = get_global_id(0);
    int row_thread = get_global_id(1);
    int row_stride = get_global_size(1);

    if (col >= N)
        return;

    float partial = 0.0f;

    // Column-major access: X[i + j*M]
    for (int row = row_thread; row < M; row += row_stride) {
        int idx = col * M + row;
        partial += X[idx];
    }

    // Subgroup-level reduction
    partial = sub_group_reduce_add(partial);

    // Only one thread per subgroup writes to result
    if (get_sub_group_local_id() == 0) {
        __atomic_add_fetch(&result[col], partial, __ATOMIC_RELAXED);
    }
}
"""
p = cl.Program(; source) |> cl.build!
k = cl.Kernel(p, "sum_columns_subgroup")

X = OpenCL.rand(Float32, 1000, 1000)
out = OpenCL.zeros(Float32, 1000)
@benchmark begin
    clcall(k, Tuple{CLPtr{Float32}, CLPtr{Float32}, Cint, Cint}, X, out, 1000, 1000; global_size=(1000,64), local_size=(1,64))
    OpenCL.synchronize(out)
end

With results:

BenchmarkTools.Trial: 10000 samples with 9 evaluations per sample.
 Range (min … max):  2.392 μs … 34.693 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     4.447 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):   4.452 μs ±  1.086 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

                     ▁▄▅███▆▅▄▂                               
  ▁▁▂▄▅▆▆▅▅▄▃▃▃▄▃▃▄▅▆███████████▆▅▅▄▃▃▃▃▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁ ▃
  2.39 μs        Histogram: frequency by time        7.43 μs <

 Memory estimate: 1.16 KiB, allocs estimate: 33.

so it's probably not a general pocl issue. I wasn't able to translate this into a native Julia kernel though since we seem to be missing a sub_group_reduce_add intrinsic.

The actual example I care about is also a more complex reduction (logsumexp! from LogExpFunctions), so it would be nice if we could speed up the generic mapreduce as well

[maleadt]

We do support subgroup intrinsics:

OpenCL.jl/lib/intrinsics/src/work_item.jl

Lines 10 to 20 in a387890

	# indices
	:get_global_linear_id => (:BuiltInGlobalLinearId, Csize_t, 1),
	:get_local_linear_id => (:BuiltInLocalInvocationIndex, Csize_t, 1),
	:get_sub_group_id => (:BuiltInSubgroupId, UInt32, 1),
	:get_sub_group_local_id => (:BuiltInSubgroupLocalInvocationId, UInt32, 1),
	# sizes
	:get_work_dim => (:BuiltInWorkDim, UInt32, 0),
	:get_sub_group_size => (:BuiltInSubgroupSize, UInt32, 0),
	:get_max_sub_group_size => (:BuiltInSubgroupMaxSize, UInt32, 0),
	:get_num_sub_groups => (:BuiltInNumSubgroups, UInt32, 0),
	:get_enqueued_num_sub_groups => (:BuiltInNumEnqueuedSubgroups, UInt32, 0)]

But sub_group_reduce_add is a high level one that doesn't work generally for Julia anyway (since operator and type semantics do not match), so it'd be better to compare against a kernel that doesn't simply call a subgroup reductor.

[simeonschaub]

so it'd be better to compare against a kernel that doesn't simply call a subgroup reductor.

What kind of strategy are you proposing? Do you mean something similar to the builtin mapreduce in https://github.com/JuliaGPU/OpenCL.jl/blob/master/src/mapreduce.jl and then try to optimize that?

I am also not really sure what you mean by non-matching operator and type semantics, is it just the lack of genericness?

[maleadt]

I am also not really sure what you mean by non-matching operator and type semantics, is it just the lack of genericness?

Yes, the semantics of add often don't match Julia's (notably wrt. NaNs), and there's of course the limited type support.

What kind of strategy are you proposing? Do you mean something similar to the builtin mapreduce in https://github.com/JuliaGPU/OpenCL.jl/blob/master/src/mapreduce.jl and then try to optimize that?

It's fine keeping it in OpenCL C for now, which would be useful to determine if it's our SPIR-V codegen that introduces a performance penalty, but instead of calling sub_group_reduce_add you could use shuffle intrinsics (sub_group_shuffle I presume, but I haven't looked closely) or a tree-style reduction (such as used in our current implementation).

[simeonschaub]

Here's a sub_group_shuffle implementation:

using OpenCL, pocl_jll, BenchmarkTools

const source = """
#pragma OPENCL EXTENSION cl_khr_subgroups : enable

float warp_reduce(float val) {
    uint lane = get_sub_group_local_id();
    uint width = get_sub_group_size();

    for (uint offset = 1; offset < width; offset <<= 1) {
        if (lane >= offset) {
            float other = sub_group_shuffle(val, lane - offset);
            val += other;
        }
    }
    return val;
}

__kernel void sum_columns_subgroup(
    __global const float* X,     // Input matrix (column-major)
    __global float* result,      // Output (length N)
    const int M,                 // Rows
    const int N                  // Columns
) {
    int col = get_global_id(0);
    int row_thread = get_global_id(1);
    int row_stride = get_global_size(1);

    if (col >= N)
        return;

    float partial = 0.0f;

    // Column-major indexing: X[i + j*M]
    for (int row = row_thread; row < M; row += row_stride) {
        int idx = col * M + row;
        partial += X[idx];
    }

    // Subgroup reduction using shuffle
    partial = warp_reduce(partial);

    // Only the first lane writes result
    if (get_sub_group_local_id() == 0) {
        __atomic_add_fetch(&result[col], partial, __ATOMIC_RELAXED);
    }
}
"""
p = cl.Program(; source) |> cl.build!
k = cl.Kernel(p, "sum_columns_subgroup")

X = OpenCL.rand(Float32, 1000, 1000)
out = OpenCL.zeros(Float32, 1000)
@benchmark begin
    clcall(k, Tuple{CLPtr{Float32}, CLPtr{Float32}, Cint, Cint}, X, out, 1000, 1000; global_size=(1000,64), local_size=(1,64))
    OpenCL.synchronize(out)
end

This is now 4x slower than Base, but still a lot faster than the current mapreduce implementation:

BenchmarkTools.Trial: 10000 samples with 1 evaluation per sample.
 Range (min … max):  207.792 μs …  2.056 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     263.337 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):   272.909 μs ± 57.939 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

          ▁▃▄▆██▆▅▅▅▃▃▂                                         
  ▁▁▁▁▂▃▅▇██████████████▇▇▆▅▄▄▄▄▃▃▃▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▃
  208 μs          Histogram: frequency by time          414 μs <

 Memory estimate: 1.16 KiB, allocs estimate: 33.

A tree-style reduction is quite a bit slower yet:

const source2 = """
__kernel void sum_columns_mapreduce(
    __global const float* X,     // Input matrix in column-major layout
    __global float* result,      // Output sums (length N)
    const int M,                 // Number of rows
    const int N                  // Number of columns
) {
    // Which column is this workgroup working on?
    int col = get_group_id(0);
    if (col >= N) return;

    int local_id = get_local_id(0);
    int local_size = get_local_size(0);
    int global_id = get_global_id(0);

    // Allocate local memory for partial sums per work-item
    __local float shared[64];

    // Each work-item sums multiple rows (group-stride loop)
    float partial = 0.0f;
    for (int row = local_id; row < M; row += local_size) {
        int idx = col * M + row;
        partial += X[idx];
    }

    // Store partial sum in local memory
    shared[local_id] = partial;
    barrier(CLK_LOCAL_MEM_FENCE);

    // Perform reduction in local memory (tree reduction)
    for (int offset = local_size >> 1; offset > 0; offset >>= 1) {
        if (local_id < offset) {
            shared[local_id] += shared[local_id + offset];
        }
        barrier(CLK_LOCAL_MEM_FENCE);
    }

    // Work-item 0 writes final sum for this column
    if (local_id == 0) {
        result[col] = shared[0];
    }
}
"""
p2 = cl.Program(; source = source2) |> cl.build!
k2 = cl.Kernel(p, "sum_columns_subgroup")

X = OpenCL.rand(Float32, 1000, 1000)
out = OpenCL.zeros(Float32, 1000)
local_size = 64
global_size = (1000 * local_size,)  # Launch one workgroup per column
@benchmark begin
    clcall(k2, Tuple{CLPtr{Float32}, CLPtr{Float32}, Cint, Cint}, X, out, 1000, 1000; global_size, local_size=(local_size,))
    OpenCL.synchronize(out)
end

BenchmarkTools.Trial: 4320 samples with 1 evaluation per sample.
 Range (min … max):  768.730 μs …   3.685 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):       1.197 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):     1.153 ms ± 298.562 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

   ▂▅▇▇▇▆▆▅▄▄▂▂▁                                 ▂▅▇█▇▆▅▄▂▁     ▂
  ▇██████████████▇▇▇▇▇▄▇▅▇▆▇▆▆▅▅▇▅▆▆▇▆▆▄▄▄▄▇▅▆▆▇▆███████████▇▇▅ █
  769 μs        Histogram: log(frequency) by time       1.57 ms <

 Memory estimate: 1.17 KiB, allocs estimate: 36.

But still faster than the current Julia implementation. So it might be a combination of factors?

[maleadt]

I wouldn't expect such drastic slowdowns, recuce_add -> shuffle -> tree going from 2us -> 200us -> 750us is much worse than I've seen on any GPU platform. The additional 2x slowdown to the current mapreducedim!'s 1.5ms doesn't seem that outrageous anymore then, so yeah this is probably a combination of factors.

I wonder if the reduce_add results in SIMD CPU code where the shuffle / tree operations don't?
We could probably consider having optimized kernels for known types / reduce operators.