GPU Design

[rcoreilly]

Goals of the current GPU implementation

• Wide availability: works on every platform and GPU hardware (especially Mac!).
  • Vulkan is the most widely supported GPU framework, powers axon GUI already. CUDA only runs on NVIDIA.
• High performance.
  • Vulkan is very low-level, "close to the metal", high-performance.
• Handle complexity of axon: ~5K LOC for core algorithm code.
  • Auto-convert from existing Go code: gosl.
• GPU notoriously difficult to debug & error-prone.
  • Built-in validation checks in Go conversion; minimal hand-written code; side-by-side CPU / GPU implementation: gosl.
• Easy to modify for algorithm changes.
  • Write and validate in Go first, typically auto-generates correct GPU version: gosl.
• Support multiple different GPU "back ends" (Vulkan and CUDA).
  • gosl can generate either -- both are essentially C-like and should be largely the same.

Drawbacks of the Vulkan and gosl approach:

• Vulkan not as widely used for compute as CUDA.
  • But graphics algorithms and games are increasingly leveraging GPU compute.
• Vulkan does not have optimized libraries like in CUDA.
  • But a very small portion of the axon code is potentially appropriate for standard matrix operations (SendSpike), and per the more detailed performance considerations, it is beneficial to use existing data structures, rather than copying back and forth between std matrix structures (especially at the Synapse level). Also it is unclear if the extreme sparseness in connectivity and activity in SendSpike would really benefit from these existing algorithms.
• Vulkan does not support advanced GPU hardware like NVLink.
  • The CUDA port via gosl could support that, and it is also unclear how soon we would exhaust the ~80Gb RAM avail on current-day server GPUs.

[rcoreilly]

Overall Compute Logic and Performance Considerations

GPU Constraints

• GPU requires all data in global arrays (Storage Buffers), no pointers, no malloc..

  • CPU uses local pointers in Layer, Prjn to simply and efficiently access data -- not possible on GPU.
  • NetworkBase has global arrays of all data structures (Neurons, Synapses, LayerVals, Pools, and various indexes) that are sync'd up to the GPU.
  • All data is accessed via indexes, not pointers, in GPU.

• GPU (HLSL) only allows pointers as args to methods / functions (and they look like values!).

  • Use in, out, inout to determine whether changes affect source input (pass by reference: out or inout) or not (pass by value: in).
  • All shared CPU / GPU code (gosl) is thus methods with pointer args for all data of relevance.
  • CPU vs. GPU feed these methods via separate code paths: GPU in gpu_hlsl/gpu_*.hlsl shaders, CPU in layer_compute.go and pjrn_compute.go.

• Each GPU computation is automatically parallel over some data structure (Neurons, Synapses, Pools), is a separate "kernel" (shader) call.

  • Can't arbitrarily mix-and-match different levels of parallelism.
  • Any memory write operation must be at or below given parallel level, or use atomics.
  • It was actually a big win to use atomic add for SendSpike ("scatter" -- sender based only when a neuron spikes) vs. a fully parallel RecvSpikes receiver-Neuron-based method ("gather") without need for atomic.

• There is significant overhead for vk.QueueSubmit call that launches a CommandBuffer of GPU commands (this is a GPU hardware thing and true for any such system).

  • Batch as much computation per submit as possible: now running 10 Cycles per submit (sped up small ra25 model by factor of 6).
  • Requires that the entire computation at Cycle level is on the GPU -- warts and all -- encourages use of "native" neural computation; explicit TD algorithm math requires annoying special code. Gated flag in Matrix is special too.
  • Really pushes use of gosl to get everything running on the GPU with minimum maintenance effort.
  • GPU speedup vs. CPU increases as network size increases: amortizing the overhead.
  • Data parallelism is a great way to amortize the overhead in smaller models: compute ~16 networks at the same time.

• StorageBuffer structs must have 16 byte alignment & sizes (including start of struct fields inside of other structs).

  • Misalignment between CPU and GPU is a major source of hard-to-find bugs in general.
  • gosl automatically checks the alignment and flags issues -- no more alignment bugs!
  • pad* int32 non-visible members manually added to all GPU-accessed structs -- run go generate and monitor output when modifying algorithm structs.

[rcoreilly]

Performance Results

examples/bench/run_gpu.sh

Results for my MacBook Pro M1 Max with 64 Gb RAM

Note: CPU results are significantly faster than in current master branch because SynCa has been sped up significantly (closed-form expression and updated per-neuron).

==============================================================
HUGE Network (5 x 1024 units)
Total Secs:   5.51

This first number is total run time for GPU version in optimized mode (10 cycles per submit) which does not allow for any detailed report on the time for each step of computation.

      BenchNet:	 Neurons: 5120	 NeurMem: 1.8 MB 	 Syns: 7340032 	 SynMem: 532.0 MB

	Function Name 	   Secs	    Pct
	GPU:ApplyExts 	  0.003	    0.0
	GPU:BetweenGi 	  0.455	    4.2
	    GPU:Cycle 	  0.577	    5.4
	GPU:CyclePost 	  0.499	    4.7
	      GPU:DWt 	  0.032	    0.3
	GPU:GatherSpikes 	  0.663	    6.2
	    GPU:LayGi 	  0.487	    4.5
	GPU:MinusPhase 	  0.005	    0.0
	 GPU:NewState 	  0.022	    0.2
	GPU:PlusPhase 	  0.004	    0.0
	   GPU:PoolGi 	  0.463	    4.3
	GPU:SendSpike 	  1.908	   17.8
	GPU:SynCaRecv 	  2.841	   26.6
	GPU:SynCaSend 	  2.713	   25.4
	  GPU:WtFmDWt 	  0.030	    0.3
	        Total 	 10.702

Note: 10.7 seconds vs. 5.5 seconds is difference between batching all Submit into one call vs. separate calls required to get timing for each one (thus each call above includes overhead).

CPU version:

Took  8.973 secs for 2 epochs, avg per epc:  4.487
OS Threads (=GOMAXPROCS): 10. Gorountines: 10 (Neurons) 10 (SendSpike) 4 (SynCa)
	Function Name 	   Secs	    Pct
	  CycleNeuron 	  0.420	    4.7
	    CyclePost 	  0.000	    0.0
	          DWt 	  1.744	   19.5
	   DWtSubMean 	  0.000	    0.0
	 GatherSpikes 	  0.160	    1.8
	   GiFmSpikes 	  0.114	    1.3
	PoolGiFmSpikes 	  0.001	    0.0
	    SendSpike 	  1.605	   17.9
	    SynCaRecv 	  0.957	   10.7
	    SynCaSend 	  3.530	   39.4
	      WtFmDWt 	  0.431	    4.8
	        Total 	  8.963

GPU speedup = 1.6x

Larger nets scale well:

==============================================================
GINORMOUS Network (5 x 2048 units)
GPU: Total Secs:   7.52
      BenchNet:	 Neurons: 10125	 NeurMem: 3.6 MB 	 Syns: 28704375 	 SynMem: 2.0 GB

CPU: Took  14.74 secs for 1 epochs, avg per epc:  14.74

GPU speedup = 1.96x

NVIDIA A100 Results

• GPU: Total Secs: 18.9 -- 3.4x WORSE than M1 -- clearly there is something wrong here!
  • is the NVIDA much more sensitive to sequential data access than the M1?
  • note: intrinsic exp call in SynCa and DWt is slightly faster than our FastExp, on NVIDIA and M1.
  • Interestingly, the separate submit version with timing is nearly as fast -- not as much impact of submit overhead.

	Function Name 	   Secs	    Pct
	GPU:ApplyExts 	  0.002	    0.0
	GPU:BetweenGi 	  0.285	    1.4
	    GPU:Cycle 	  0.264	    1.3
	GPU:CyclePost 	  0.241	    1.2
	      GPU:DWt 	  3.715	   18.4
	GPU:GatherSpikes 	  0.300	    1.5
	    GPU:LayGi 	  0.298	    1.5
	GPU:MinusPhase 	  0.005	    0.0
	 GPU:NewState 	  0.040	    0.2
	GPU:PlusPhase 	  0.127	    0.6
	   GPU:PoolGi 	  0.339	    1.7
	GPU:SendSpike 	  1.110	    5.5
	GPU:SynCaRecv 	  6.861	   33.9
	GPU:SynCaSend 	  6.599	   32.7
	  GPU:WtFmDWt 	  0.023	    0.1
	        Total 	 20.210

• CPU -- AMD EPYC 7502 32-Core

OS Threads (=GOMAXPROCS): 32. Gorountines: 21 (Neurons) 21 (SendSpike) 4 (SynCa)
Took 21.5 secs for 2 epochs, avg per epc:  10.75

Would be worth experimenting with these thread counts here -- possibly could be faster with lower usage. Also the 4 SynCa is actually vestigial and it is using the 21 Neurons count actually now.

==============================================================
GINORMOUS Network (5 x 2048 units)
Total Secs:   21.3

Same good scaling for GPU, CPU is now 28 for this, so GPU starting to show advantages.

ccnlab/lvis/sims/lvis_cu3d100_te16deg_axon

This is the "full size" LVis model, which is most representative of how things should scale in protobrain etc.

Neurons: 47,872	 NeurMem: 16.8 MB 	 Syns: 31,316,128 	 SynMem: 2.2 GB

PerTrlMSec = wall clock time in msec for each 200 msec of simulated time:

• GPU: 1280
• CPU: 2200

= 1.7x GPU speedup

With data parallel, should easily be able to exceed real time, for 31M synapses -- pretty good overall.