feat: concurrent sum (for huge numbers) (#1495)

* concurrent sum module

* compute shader blueprint

* endeavour to run compute shader

* introduced std log and std log2

* log now supports vectors

* removed redundant files

* log in vectorops

* added tests

* complete tests

* updated to main

* Fixes - compute shader runs

* alg fixes

* working presum

* introduced concurrent-sum example

* concurrent sum in math examples

* finally somewhat parallel

* Fixed exceeding workgroup size

* One function to wrap the compute shader

* packed whole concurrent-sum

* Added quicker concurrentSum algorithm - shared memory

* each thread copies two values into shared memory, instead of one

* Found workGroupSize buffer

* Introduced sums scan

* Compute shader used for sums

* export fixes

* incrementSHader

* Found copy bug

* alg fix

* Dynamic buffer length support

* Now supporting optional buffer as output

* no more consts

* Divided into upPass and downPass

* Refmt

* work in progress

* Fixed

* reintroduced down pass

* warning fixes

* Fully working impl

* example for testing

* clean layouts

* fully working recursive step

* Dispatching multiple workgroups

* supporting optional buffer

* stash the cache

* buffer cache

* nits

* performance inprovements

* alternative cache

* gpu time

* refactor of types

* 🦕

* basic support for any operation

* concurrent sum supports all operators now

* example tidy up

* removed unnecessary flush

* example improvements

* types

* Support for single element sum

* example sanity changes'

* filenames

* 🦕

* fixes

* 🦕

* 🦕

* flex example

* lines

* speed up isntead of time

* some improvements

* tailwind baby

* 3 values

* working example

* pr fixes

* html simplifiaction

* other operators fix

* typo

* length is now chooseable

* fixed on mobile

* tgpu callable

* profile pic mod

* copilot

* buttons cleanup

* introduced calculate index

* got rid of ugly casts

* Major rename

* Removed old concurrent sum example

* 🦕

* typo

* better cache

* oppsies

* html changes

* height control

* looking goood on mobile

* 🦕

* dynamic y axis

* rename

* 🦕

* PR fixes 1

* fix

* single element scan fix

* 🦕

* proper private values

* docs and tests

* farewell buffer cache

* 🦕

* Merge fixes

* docs: Tests for prefix scan (#1804)

* jsdocs example

* implicit single scan conv

* fix timestamp queries

* 🦕🦕🦕

* tgpu bump 0.6 -> 0.8

* API option 4

* 🦕🦕🦕

* 🦕🦕🦕🦕

* tgpu version bump

* PR fixes 1

* Code duplication with scan (only difference is true vs false), this body can be extracted to a helper function.

* outputBuffer does not work in scan (there is always a new one-element buffer returned)

* fix faulty test

* ordering of lines

* removed number cast

* simplified array init

* bars opacity in loop

* readme change

* .with update

* better pipeline if

* removed tgpu callable

* uniformAdd -> uniformOp

* fix, simplify, modernize (#2164)

* simplify (#2169)

* COMBINE (#2182)

* add -> op

* export computer for *advanced* usage I guess :P

* two level cache

* fixes

* better wording imp

---------

Co-authored-by: Iwo Plaza <iwoplaza@gmail.com>
Co-authored-by: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com>
Co-authored-by: Aleksander Katan <aleksander.katan@swmansion.com>
Co-authored-by: Konrad Reczko <reczkok@gmail.com>
Co-authored-by: Konrad Reczko <66403540+reczkok@users.noreply.github.com>

Author: 5 people

apps/typegpu-docs/package.json

@@ -29,6 +29,7 @@
     "@typegpu/color": "workspace:*",
     "@typegpu/geometry": "workspace:*",
     "@typegpu/noise": "workspace:*",
+    "@typegpu/concurrent-scan": "workspace:*",
     "@typegpu/three": "workspace:*",
     "@typegpu/sdf": "workspace:*",
     "three": "catalog:example",

apps/typegpu-docs/src/examples/algorithms/concurrent-chart/calculator.ts

@@ -0,0 +1,78 @@
+import { initCache, prefixScan } from '@typegpu/concurrent-scan';
+import * as d from 'typegpu/data';
+import type { TgpuRoot } from 'typegpu';
+import * as std from 'typegpu/std';
+
+type SumResult = {
+  success: boolean;
+  jsTime: number;
+  gpuTime: number;
+  gpuShaderTime: number;
+};
+
+function prefixSumOnJS(arr: number[]) {
+  for (let i = 1; i < arr.length; i++) {
+    arr[i] += arr[i - 1];
+  }
+  // In Blelloch scan, the result starts with identity element
+  arr.unshift(0);
+  arr.pop();
+  return arr;
+}
+
+function arraysEqual(a: number[], b: number[]): boolean {
+  if (a.length !== b.length) {
+    return false;
+  }
+  for (let i = 0; i < a.length; i++) {
+    if (a[i] !== b[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+export async function performCalculationsWithTime(
+  root: TgpuRoot,
+  inputArray: number[],
+): Promise<SumResult> {
+  const arraySize = inputArray.length;
+  const inputBuffer = root
+    .createBuffer(d.arrayOf(d.f32, arraySize))
+    .$usage('storage');
+  inputBuffer.write(inputArray);
+
+  // JS version
+  const jsStartTime = performance.now();
+  const jsResult = prefixSumOnJS(inputArray);
+  const jsTime = performance.now() - jsStartTime;
+
+  // GPU version
+  initCache(root, { operation: std.add, identityElement: 0 });
+  const querySet = root.createQuerySet('timestamp', 2);
+  const gpuStartTime = performance.now();
+  const calcResult = prefixScan(
+    root,
+    {
+      inputBuffer: inputBuffer,
+      outputBuffer: inputBuffer,
+      operation: std.add,
+      identityElement: 0,
+    },
+    querySet,
+  );
+  querySet.resolve();
+  await root.device.queue.onSubmittedWorkDone();
+  const gpuTime = performance.now() - gpuStartTime;
+
+  const gpuResult = await calcResult.read();
+  const timestamps = await querySet.read();
+  const gpuShaderTime = Number(timestamps[1] - timestamps[0]) / 1_000_000;
+
+  return {
+    success: arraysEqual(jsResult, gpuResult),
+    jsTime,
+    gpuTime,
+    gpuShaderTime,
+  };
+}