CPU backend operator fusion

[zzzzwc]

The CUDA backend already fuses several common op patterns (e.g. RMS_NORM + MUL,RMS_NORM + MUL + ADD),But on the CPU backend these ops are still executed one by one.

I tried to port the same idea to the CPU backend, starting with the simplest and most frequent pattern: RMS_NORM + MUL (the RMSNorm + weight multiplication that appears in basically every LLaMA-style block).

And the benchmark result of the ops-group like:

Note: RMS_NORM_MUL_ADD is a new test case I added to test-backend-ops locally so that the fused pattern can be measured directly

My env:

MacBook Air 2023
CPU: M2
memory: 24GB

Before:

➜  llama.cpp git:(master) ✗ GGML_CPU_DISABLE_FUSION=on ./build/bin/test-backend-ops perf -o RMS_NORM_MUL_ADD -b CPU                
ggml_metal_device_init: tensor API disabled for pre-M5 and pre-A19 devices
ggml_metal_library_init: loading '/Users/zwc/CLionProjects/llama.cpp/build/bin/default.metallib'
ggml_metal_library_init: loaded in 0.002 sec
ggml_metal_rsets_init: creating a residency set collection (keep_alive = 180 s)
ggml_metal_device_init: GPU name:   MTL0
ggml_metal_device_init: GPU family: MTLGPUFamilyApple8  (1008)
ggml_metal_device_init: GPU family: MTLGPUFamilyCommon3 (3003)
ggml_metal_device_init: GPU family: MTLGPUFamilyMetal3  (5001)
ggml_metal_device_init: simdgroup reduction   = true
ggml_metal_device_init: simdgroup matrix mul. = true
ggml_metal_device_init: has unified memory    = true
ggml_metal_device_init: has bfloat            = true
ggml_metal_device_init: has tensor            = false
ggml_metal_device_init: use residency sets    = true
ggml_metal_device_init: use shared buffers    = true
ggml_metal_device_init: recommendedMaxWorkingSetSize  = 17179.89 MB
Testing 3 devices

Backend 1/3: MTL0
  Skipping
Backend 2/3: BLAS
  Skipping
Backend 3/3: CPU
  Device description: Apple M2
  Device memory: 24576 MB (24576 MB free)

  RMS_NORM_MUL_ADD(type=f32,ne=[1024,512,1,1],eps=0.000001,broadcast=0,multi_add=0):                   12365 runs -   808.79 us/run -     6144 kB/run -   14.49 GB/s
  RMS_NORM_MUL_ADD(type=f32,ne=[1024,512,1,1],eps=0.000001,broadcast=0,multi_add=1):                   11626 runs -   860.16 us/run -     6144 kB/run -   13.62 GB/s
  Backend CPU: OK
3/3 backends passed
OK

After:

➜  llama.cpp git:(master) ✗ ./build/bin/test-backend-ops perf -o RMS_NORM_MUL_ADD -b CPU 
ggml_metal_device_init: tensor API disabled for pre-M5 and pre-A19 devices
ggml_metal_library_init: loading '/Users/zwc/CLionProjects/llama.cpp/build/bin/default.metallib'
ggml_metal_library_init: loaded in 0.003 sec
ggml_metal_rsets_init: creating a residency set collection (keep_alive = 180 s)
ggml_metal_device_init: GPU name:   MTL0
ggml_metal_device_init: GPU family: MTLGPUFamilyApple8  (1008)
ggml_metal_device_init: GPU family: MTLGPUFamilyCommon3 (3003)
ggml_metal_device_init: GPU family: MTLGPUFamilyMetal3  (5001)
ggml_metal_device_init: simdgroup reduction   = true
ggml_metal_device_init: simdgroup matrix mul. = true
ggml_metal_device_init: has unified memory    = true
ggml_metal_device_init: has bfloat            = true
ggml_metal_device_init: has tensor            = false
ggml_metal_device_init: use residency sets    = true
ggml_metal_device_init: use shared buffers    = true
ggml_metal_device_init: recommendedMaxWorkingSetSize  = 17179.89 MB
Testing 3 devices

Backend 1/3: MTL0
  Skipping
Backend 2/3: BLAS
  Skipping
Backend 3/3: CPU
  Device description: Apple M2
  Device memory: 24576 MB (24576 MB free)

  RMS_NORM_MUL_ADD(type=f32,ne=[1024,512,1,1],eps=0.000001,broadcast=0,multi_add=0):                   17743 runs -   563.62 us/run -     6144 kB/run -   20.79 GB/s
  RMS_NORM_MUL_ADD(type=f32,ne=[1024,512,1,1],eps=0.000001,broadcast=0,multi_add=1):                   19593 runs -   510.41 us/run -     6144 kB/run -   22.96 GB/s
  Backend CPU: OK
3/3 backends passed
OK

My main question before turning this into a PR: Is there any appetite for operator fusion optimizations on the CPU backend?

Happy to iterate on this based on feedback. Thanks!

[CISC]

Sure, there are plenty of people running CPU only, any significant optimization (but do ensure that you do not negatively impact a CPU target) to that backend is welcome.

[zzzzwc]

@CISC Thanks! PR #22423 is ready for review. It implements the fusion on CPU without perf regressions. Would you mind taking a quick look?

[am17an]

I think CPU supports the graph plan API, fusion is better done at that point rather than checking for nodes at each compute, it might be slower otherwise.

[zzzzwc]

Yes, Graph-plan fusion is likely the better long-term approach, but since benchmarks show the current per-node checks cause no measurable regression (even on simple add), I plan to revisit plan-stage fusion once we have more fusion patterns in place.

[zzzzwc]

I've implemented the changes. Looking forward to your review. #22423

[zzzzwc]

My next focus will be the Gated FFN fusion, which should yield more gains than simple RMS_NORM+MUL fusion

[am17an]

I already have a PR for that #20596, but the speed-up does not seem to universal

[zzzzwc]

Got it, I have an M2 device, I might be able to help by providing some useful perf data for your PR.

[zzzzwc]

Opened a PR for benchmarking operator groups: #22934
Also adds --perf-duration / --threads for more stable bench results.