Shared Memory and Atomics

A compute shader runs in workgroups of invocations, which run at once. They share the variables after its &shared, wait for each other at a barrier, and change integers that several invocations share with atomics.

Shared variables

After &shared, a compute shader's parameter list names the variables its workgroup shares, (name type), each a value: a scalar, vector, matrix, struct or array of a length. Every invocation of a workgroup reads and writes the same ones, and each workgroup has its own. They start at zero, through Vulkan's shaderZeroInitializeWorkgroupMemory, which every Vulkan 1.4 device has and cl-gfx enables. A write another invocation makes is seen after a barrier.

Barriers

Function	Stage	Does
(barrier)	compute	waits until every invocation of the workgroup reaches it, and makes what each wrote to shared variables before it seen by every other after it
(memory-barrier)	any	orders this invocation's reads and writes of storage buffers, storage images and shared variables: those before it are done, as other invocations see them, before those after it begin
(memory-barrier-buffer), (memory-barrier-image)	any	the same, of storage buffers, or of storage images
(memory-barrier-shared)	compute	the same, of shared variables
(group-memory-barrier)	compute	the same as memory-barrier, as the other invocations of the workgroup see them

Each gives no value. A memory barrier waits for no other invocation. barrier alone orders shared variables: a write to a storage buffer or image is seen across the workgroup after (memory-barrier-buffer), or (memory-barrier-image), and then (barrier).

Every invocation of the workgroup reaches a barrier, so it needs uniform control flow: one in a branch, or a loop, that invocations may take differently is refused, and no declaration relaxes it.

(define-shader early-exit (:compute :workgroup-size (64))
    (&uniform (numbers (:array :uint) :storage :read-write)
     &shared (total :uint))
  (let ((i (x (global-invocation-id))))
    (when (< i (length numbers))
      (atomic-incf total (aref numbers i))
      (barrier)
      (setf (aref numbers i) total))))

In EARLY-EXIT: (BARRIER) waits for every invocation of the workgroup, so it
needs uniform control flow; but it is in a branch of (WHEN (< I (LENGTH
NUMBERS)) ...), whose test depends on (GLOBAL-INVOCATION-ID), which differs
from invocation to invocation.  Move the barrier where every invocation of the
workgroup reaches it, out of the branch.
  in: (BARRIER)

Atomics

An atomic changes an :int or a :uint that invocations share in one step, which no other invocation's change splits, and gives the value it held before:

Function	Stores in place
(atomic-incf place [delta]), (atomic-decf place [delta])	place plus, or minus, delta, 1 by default, as SBCL's sb-ext:atomic-incf and sb-ext:atomic-decf do
(atomic-min place value), (atomic-max place value)	the lesser, or the greater, of place and value
(atomic-logand place value), (atomic-logior place value), (atomic-logxor place value)	(logand place value) and its kin
(atomic-exchange place value)	value
(compare-and-swap place old new)	new, if place holds old, as sb-ext:compare-and-swap does: it gives old if new was stored

place is an :int or a :uint of a storage buffer of :storage :read-write, or of a shared variable, and the values are of its type. Sums wrap, as clsl's integers do. At the default safety, an element out of range is not changed, and gives 0. An atomic of a float, or of a storage buffer of :storage :read, is refused:

(define-shader float-atomic (:compute :workgroup-size (64))
    (&uniform (sums (:array :float) :storage :read-write))
  (atomic-incf (aref sums 0) 1.0))

(define-shader read-only-atomic (:compute :workgroup-size (64))
    (&uniform (sums (:array :uint) :storage :read))
  (atomic-incf (aref sums 0)))

In FLOAT-ATOMIC: ATOMIC-INCF changes an :INT or a :UINT, and (AREF SUMS 0) is
:FLOAT.
  in: (ATOMIC-INCF (AREF SUMS 0) 1.0)

In READ-ONLY-ATOMIC: SUMS is a storage buffer of :STORAGE :READ, which shader
code reads but never writes: give it :STORAGE :READ-WRITE to write it.
  in: SUMS

(define-shader claim (:compute :workgroup-size (64))
    (&uniform (owner (:array :int) :storage :read-write))
  "Let the first invocation to come claim each slot of OWNER, which starts at -1."
  (let ((i (int (x (global-invocation-id)))))
    (compare-and-swap (aref owner (mod i 8)) -1 i)))

A histogram

This counts a thousand numbers into sixteen bins by their low four bits, each workgroup into its shared counts first, and then into the storage buffer bins, on the GPU through cl-gfx:

(asdf:load-system :clsl/vulkan)
(asdf:load-system :cl-gfx/vulkan)

(defpackage #:counting
  (:use #:cl #:clsl)
  (:local-nicknames (#:gfx #:cl-gfx)))
(in-package #:counting)

(define-shader histogram (:compute :workgroup-size (64))
    (&uniform (numbers (:array :uint) :storage :read)
              (bins (:array :uint 16) :storage :read-write)
     &shared (counts (:array :uint 16)))
  "Count NUMBERS into BINS by their low four bits, each workgroup into its
COUNTS first."
  (let ((i (x (global-invocation-id))))
    (when (< i (length numbers))
      (atomic-incf (aref counts (ldb (byte 4 0) (aref numbers i)))))
    (barrier)
    (let ((bin (local-invocation-index)))
      (when (< bin 16)
        (atomic-incf (aref bins bin) (aref counts bin))))))

(program-layout-entries (compile-program '(:compute histogram) :target :vulkan))
;; => (NIL
;;     ((1 :READ-ONLY-STORAGE-BUFFER :VISIBILITY (:COMPUTE))
;;      (2 :STORAGE-BUFFER :VISIBILITY (:COMPUTE))))

(defun words (count &optional (contents (constantly 0)))
  "A vector of COUNT 32-bit words, each CONTENTS of its index."
  (let ((vector (make-array count :element-type '(unsigned-byte 32))))
    (dotimes (i count vector)
      (setf (aref vector i) (funcall contents i)))))

(gfx:with-device (device)
  (let* ((program (compile-program '(:compute histogram) :target :vulkan))
         (module (gfx:create-shader-module device (program-spirv program :compute)))
         (layouts (mapcar (lambda (entries) (gfx:create-bind-group-layout device entries))
                          (program-layout-entries program)))
         (pipeline (gfx:create-compute-pipeline device :module module
                                                       :bind-group-layouts layouts))
         (numbers (gfx:create-buffer device :size (* 4 1000) :usage '(:storage)))
         (bins (gfx:create-buffer device :size (* 4 16) :usage '(:storage :copy-src)))
         (group (gfx:create-bind-group device (nth (uniform-group-index :material) layouts)
                                       (list (list 1 numbers) (list 2 bins)))))
    (gfx:destroy-shader-module module)
    (gfx:write-buffer numbers (words 1000 #'identity))
    (gfx:write-buffer bins (words 16))
    (gfx:with-frame (frame device)
      (gfx:with-compute-pass (pass frame)
        (gfx:set-pipeline pass pipeline)
        (gfx:set-bind-group pass (uniform-group-index :material) group)
        (gfx:dispatch pass (ceiling 1000 64))))
    (gfx:read-buffer bins :element-type '(unsigned-byte 32))))
;; => #(63 63 63 63 63 63 63 63 62 62 62 62 62 62 62 62)

as Lisp counts them:

(let ((bins (make-array 16 :initial-element 0)))
  (dotimes (n 1000 bins)
    (incf (aref bins (ldb (byte 4 0) n)))))
;; => #(63 63 63 63 63 63 63 63 62 62 62 62 62 62 62 62)

The barrier is outside the when, where every invocation of the workgroup reaches it, and counts starts at zero in each workgroup. numbers, which no stage writes, is a read-only storage buffer; bins starts at zero because the program writes zeros into it first.