storage_allocation.bend

stelliferous@0.0.2.0 · 57 lines · 5 definitions

1# Keep the native allocation extent constant at each call site. With a dynamic
2# depth, Clang 19 retains a scalar initialization loop with a shift per element;
3# constant depths allow memset or vectorized fills. Dispatch once per buffer,
4# covering every power of two below the native depth limit, independently of
5# model shapes. The fallback preserves Storage.allocate for larger depths.
6import Base
7import stelliferous@0.0.2.0/storage_buffer.bend as Storage
8
9def dispatch(~Element: Data,~allocator: Nat -> Element -> Nat -> Storage.Buffer<Element>,depth: Nat,+value: Element,length: Nat) -> Storage.Buffer<Element>:
10  match depth:
11    case 0n: allocator(0n,value,length)
12    case 1n: allocator(1n,value,length)
13    case 2n: allocator(2n,value,length)
14    case 3n: allocator(3n,value,length)
15    case 4n: allocator(4n,value,length)
16    case 5n: allocator(5n,value,length)
17    case 6n: allocator(6n,value,length)
18    case 7n: allocator(7n,value,length)
19    case 8n: allocator(8n,value,length)
20    case 9n: allocator(9n,value,length)
21    case 10n: allocator(10n,value,length)
22    case 11n: allocator(11n,value,length)
23    case 12n: allocator(12n,value,length)
24    case 13n: allocator(13n,value,length)
25    case 14n: allocator(14n,value,length)
26    case 15n: allocator(15n,value,length)
27    case 16n: allocator(16n,value,length)
28    case 17n: allocator(17n,value,length)
29    case 18n: allocator(18n,value,length)
30    case 19n: allocator(19n,value,length)
31    case 20n: allocator(20n,value,length)
32    case 21n: allocator(21n,value,length)
33    case 22n: allocator(22n,value,length)
34    case 23n: allocator(23n,value,length)
35    case 24n: allocator(24n,value,length)
36    case 25n: allocator(25n,value,length)
37    case 26n: allocator(26n,value,length)
38    case 27n: allocator(27n,value,length)
39    case 28n: allocator(28n,value,length)
40    case 29n: allocator(29n,value,length)
41    case 30n: allocator(30n,value,length)
42    case 31n+rest: allocator(31n+rest,value,length)
43
44def allocated(~Element: Data,depth: Nat,value: Element,length: Nat) -> Storage.Buffer<Element>:
45  Storage.allocate(~Element,depth,value,length)
46
47def allocate(~Element: Data,depth: Nat,value: Element,length: Nat) -> Storage.Buffer<Element>:
48  dispatch(~Element,~allocated(~Element),depth,value,length)
49
50# Internal Array callers can discard the construction certificate after the
51# same constant-depth dispatch; no traversal or clone is introduced.
52def storage(~Element: Data,buffer: Storage.Buffer<Element>) -> Array<Element>:
53  Storage.Buffer{storage,length,certificate} = buffer
54  storage
55
56def array(~Element: Data,depth: Nat,value: Element) -> Array<Element>:
57  storage(~Element,allocate(~Element,depth,value,0n))