Accel.xtl

1⍝# Accel: the acceleratable subset of X_eTaL: elementwise arithmetic, masks, reductions, products, transposes, windows, lookups, choices, a layer
2⍝#
3⍝# Import it with an alias of your choice: "ac:" u_se< "Accel".
4⍝# Put libs/Accel/src on XETAL_PATH ("just path"); the reference is libs/Accel/docs.
5⍝#
6⍝# Every function here is one whole-array expression in the subset a
7⍝# GPU (or an FPGA) can take on directly: static scalar types, pure
8⍝# functions, arithmetic and comparisons item by item with scalar
9⍝# extension, a select by mask, reductions, inner products, tables and
10⍝# transposes.
11⍝# No recursion, no nested arrays, no input or output. The X_eTaL
12⍝# evaluator is the reference: X_eTaL-gpu runs the same functions on a
13⍝# GPU and must give the same answers (the demos are the programs it
14⍝# checks). Names with l: are exported; the others are private.
15
16⍝## Elementwise
17
18⍝# The sum, item by item: the first kernel, one work-item per element.
19⍝# >> "ac:" u_se< "Accel"
20⍝# >> 1 2 3 4 ac:v_add 10 20 30 40
21⍝# 11 22 33 44
22ˡv̲add ← { x y → x + y }
23
24⍝# Every item of the right times the number on the left.
25⍝# >> "ac:" u_se< "Accel"
26⍝# >> 3 ac:s_cale 1 2 3
27⍝# 3 6 9
28ˡs̲cale ← { a x → a × x }
29
30⍝# The items where positive, 0 elsewhere (a select by comparison with
31⍝# a constant), for Ints or Floats.
32⍝# >> "ac:" u_se< "Accel"
33⍝# >> ac:r_elu 1 -2 3 -4
34⍝# 1 0 3 0
35ˡr̲elu ← { x → x m̲ax 0 }
36
37⍝## Masks and select
38
39⍝# t ac:t_hreshold x: x where x > t, 0.0 elsewhere, for Floats: the
40⍝# mask x > t made Float and multiplied in, the select a GPU does
41⍝# with a conditional move. A masked-out negative comes back as -0.0.
42⍝# >> "ac:" u_se< "Accel"
43⍝# >> 0.5 ac:t_hreshold 0.25 0.5 0.75 1.0
44⍝# 0.0 0.0 0.75 1.0
45ˡt̲hreshold ← { t x → x × f̲loat x > t }
46
47⍝# m ac:k_eep x: x where the Int mask m is 1, 0 elsewhere.
48⍝# >> "ac:" u_se< "Accel"
49⍝# >> (1 0 1 0) ac:k_eep 5 6 7 8
50⍝# 5 0 7 0
51ˡk̲eep ← { m x → x × m }
52
53⍝## Reductions
54
55⍝# The sum of the items of a vector: a parallel reduction on a GPU
56⍝# (work-group partials, then a second pass).
57⍝# >> "ac:" u_se< "Accel"
58⍝# >> ac:s_um 1 2 3 4
59⍝# 10
60ˡs̲um ← { x → '+ r̲/ x }
61
62⍝# The largest item of a vector.
63⍝# >> "ac:" u_se< "Accel"
64⍝# >> ac:l_argest 3 9 2
65⍝# 9
66ˡl̲argest ← { x → 'm̲ax r̲/ x }
67
68⍝# The inner product of two vectors: a multiply, then a reduce.
69⍝# >> "ac:" u_se< "Accel"
70⍝# >> 1 2 3 ac:d_ot 4 5 6
71⍝# 32
72ˡd̲ot ← { x y → '+ r̲/ x × y }
73
74⍝## Along an axis
75
76⍝# The sum of each row of a matrix (along axis 2: one work-item per
77⍝# row on a GPU).
78⍝# >> "ac:" u_se< "Accel"
79⍝# >> ac:r_owSums 2 3 r_eshape 1 20 3 4 5 60
80⍝# 24 69
81ˡr̲owSums ← { m → '+ r̲/₂ m }
82
83⍝# The sum of each column (along axis 1, the default of r_/).
84⍝# >> "ac:" u_se< "Accel"
85⍝# >> ac:c_olSums 2 3 r_eshape 1 20 3 4 5 60
86⍝# 5 25 63
87ˡc̲olSums ← { m → '+ r̲/ m }
88
89⍝# The largest item of each row.
90⍝# >> "ac:" u_se< "Accel"
91⍝# >> ac:r_owMax 2 3 r_eshape 1 20 3 4 5 60
92⍝# 20 60
93ˡr̲owMax ← { m → 'm̲ax r̲/₂ m }
94
95⍝## Inner products
96
97⍝# x ac:m_atmul w: the inner product of x's last axis with w's first
98⍝# (APL's +.x): a matrix times a matrix, a matrix times a vector, a
99⍝# vector times a matrix. Each result item is a sum of products, one
100⍝# GPU work-item each (or a tile of them, as the schedule says).
101⍝# >> "ac:" u_se< "Accel"
102⍝# >> (2 3 r_eshape r_ange 6) ac:m_atmul 3 2 r_eshape r_ange 6
103⍝# 22 28
104⍝# 49 64
105⍝# >> (2 3 r_eshape r_ange 6) ac:m_atmul 1 0 2
106⍝# 7 16
107ˡm̲atmul ← { x w → x '+ '× i̲nner w }
108
109⍝## Transposes
110
111⍝# q ac:s_cores k: attention scores, every query row against every
112⍝# key row (q k^T): the inner product of q with the transpose of k. On
113⍝# a GPU the transpose is a gather, one work-item per item, and the
114⍝# product follows (research7: attention decomposes into array
115⍝# operations; it is not an intrinsic).
116⍝# >> "ac:" u_se< "Accel"
117⍝# >> (2 3 r_eshape r_ange 6) ac:s_cores 4 3 r_eshape r_ange 12
118⍝# 14 32  50  68
119⍝# 32 77 122 167
120ˡs̲cores ← { q k → q '+ '× i̲nner o̲\ k }
121
122⍝# x ac:w_eightGrad dy: the gradient of a dense layer's weights, for
123⍝# a batch x (a row per example) and the gradient dy at the layer's
124⍝# output: x^T dy, the product every backward pass is made of.
125⍝# >> "ac:" u_se< "Accel"
126⍝# >> (4 2 r_eshape r_ange 8) ac:w_eightGrad 4 3 r_eshape r_ange 12
127⍝# 118 134 150
128⍝# 140 160 180
129ˡw̲eightGrad ← { x dy → (o̲\ x) '+ '× i̲nner dy }
130
131⍝## Lookups
132
133⍝# t ac:e_mbed ids: the rows of the table t at the positions ids
134⍝# (Ints from 1), in the shape of ids: a batch of token ids picks a
135⍝# batch of embeddings (ids s_elect t). On a GPU a gather, one
136⍝# work-item per item of the result, the table resident; the ids are
137⍝# all that a prompt uploads.
138⍝# >> "ac:" u_se< "Accel"
139⍝# >> t := 4 3 r_eshape 0.0 0.1 0.2 1.0 1.1 1.2 2.0 2.1 2.2 3.0 3.1 3.2
140⍝# >> t ac:e_mbed 4 1 2
141⍝# 3.0 3.1 3.2
142⍝# 0.0 0.1 0.2
143⍝# 1.0 1.1 1.2
144ˡe̲mbed ← { t ids → ids s̲elect t }
145
146⍝## Windows
147
148⍝# ac:w_indows x: every 3 by 3 window of every picture of the batch x
149⍝# (n pictures of h by w) as 9 rows, a column per picture and place:
150⍝# the picture's nine shifted copies (a rotation along each of the
151⍝# last two axes, one copy per shift), their wrapped edges dropped,
152⍝# reshaped; the expression is X_eTaL-ML's Conv (cv:w_indows). On a
153⍝# GPU each rotation and each drop is a gather of its own and the
154⍝# reshape is nothing.
155⍝# >> "ac:" u_se< "Accel"
156⍝# >> ac:w_indows 1 3 4 r_eshape r_ange 12
157⍝#  1  2
158⍝#  2  3
159⍝#  3  4
160⍝#  5  6
161⍝#  6  7
162⍝#  7  8
163⍝#  9 10
164⍝# 10 11
165⍝# 11 12
166ˡw̲indows ← { x →
167  places ← ((2 s̲elect s̲hape x) − 2) × (3 s̲elect s̲hape x) − 2
168  (9 c̲at places × t̲ally x) r̲eshape 1 d̲rop₅ -1 d̲rop₅ 1 d̲rop₄ -1 d̲rop₄ -1 0 1 o̲-₃ -1 0 1 o̲-₃ x
169}
170
171⍝# k ac:c_onv x: the f filters k (f rows of 9, a 3 by 3 filter row by
172⍝# row) over every window of every picture of x, as one matrix
173⍝# product: the maps, n by f by (h - 2) by (w - 2) (a valid
174⍝# convolution without the bias; X_eTaL-ML's Conv adds its biases).
175⍝# >> "ac:" u_se< "Accel"
176⍝# >> (1 9 r_eshape 0 0 0 0 1 0 0 0 0) ac:c_onv 1 3 4 r_eshape r_ange 12
177⍝# 6 7
178ˡc̲onv ← { k x →
179  n ← t̲ally x
180  h ← (2 s̲elect s̲hape x) − 2
181  w ← (3 s̲elect s̲hape x) − 2
182  (n c̲at (t̲ally k) c̲at h c̲at w) r̲eshape 2 1 3 t̲ranspose ((t̲ally k) c̲at n c̲at h × w) r̲eshape k '+ '× i̲nner ˡw̲indows x
183}
184
185⍝## Choices
186
187⍝# ac:a_rgmax m: for each row of m, the position (from 1) of its
188⍝# largest item, the first on a tie: the decision of a classifier,
189⍝# the next token of a greedy sampler. X_eTaL-ML's nn:a_rgmax, as
190⍝# array operations the IR has: the row's largest spread along the
191⍝# row, a mask where the row equals it, the smallest position under
192⍝# the mask (n + 1 elsewhere). No new GPU operation: a row reduction,
193⍝# a table, a fused mask and a second row reduction.
194⍝# >> "ac:" u_se< "Accel"
195⍝# >> ac:a_rgmax 3 4 r_eshape 5.0 1.0 9.0 9.0 2.0 2.0 2.0 2.0 7.0 8.0 6.0 8.0
196⍝# 3 1 2
197ˡa̲rgmax ← { m →
198  n ← f̲irst -1 t̲ake s̲hape m
199  hit ← f̲loat m = ('m̲ax r̲/₂ m) 'l̲eft t̲able o̲ffsets n
200  f̲loor 'm̲in r̲/₂ (hit × (o̲ffsets t̲ally m) 'r̲ight t̲able f̲loat r̲ange n) + (1.0 − hit) × f̲loat n + 1
201}
202
203⍝# k ac:t_opk m: a mask of each row's k largest items (1.0 there,
204⍝# 0.0 elsewhere): the row's largest taken out k times, as X_eTaL-ML's
205⍝# Sample (sm:t_op) and its mixture-of-experts router choose; items
206⍝# tied with a taken one go with it, so a tie at the k-th keeps all
207⍝# tied. The choice of a router, the candidates of a retrieval. A
208⍝# taken item is pushed below its row by more than the row's range
209⍝# (Sample subtracts 1e300, which is infinity in a GPU's f32 and
210⍝# poisons the row with NaN).
211⍝# >> "ac:" u_se< "Accel"
212⍝# >> 2 ac:t_opk 2 4 r_eshape 0.1 0.5 0.15 0.25 7.0 8.0 6.0 8.0
213⍝# 0.0 1.0 0.0 1.0
214⍝# 1.0 1.0 0.0 1.0
215ˡt̲opk ← { k m →
216  n ← f̲irst -1 t̲ake s̲hape m
217  shift ← (1.0 + ('m̲ax r̲/₂ m) − 'm̲in r̲/₂ m) 'l̲eft t̲able o̲ffsets n
218  left ← k '{ q → q − shift × f̲loat q = ('m̲ax r̲/₂ q) 'l̲eft t̲able o̲ffsets n } p̲ower m
219  f̲loat (m − left) > 0.0
220}
221
222⍝## A layer of a network
223
224⍝# x ac:d_ense wb: the layer x W + b for a batch x (one example per
225⍝# row), the weights W (a row per input, a column per output) with the
226⍝# bias b as one more row at the bottom of wb, as X_eTaL-ML's NN has
227⍝# it. The bias row is spread down the batch by a table (r_ight).
228⍝# >> "ac:" u_se< "Accel"
229⍝# >> x := 2 3 r_eshape 1.0 2.0 3.0 -1.0 0.5 0.0
230⍝# >> wb := 4 2 r_eshape 0.5 -1.0 1.0 0.25 -0.5 2.0 0.1 0.2
231⍝# >> x ac:d_ense wb
232⍝# 1.1   5.7
233⍝# 0.1 1.325
234ˡd̲ense ← { x wb → (x '+ '× i̲nner -1 d̲rop wb) + (o̲ffsets t̲ally x) 'r̲ight t̲able r̲avel -1 t̲ake wb }
235
236⍝# Each row of a matrix into probabilities: the row's largest taken
237⍝# off (the same result, no overflow), e to each, each divided by its
238⍝# row's sum; a row's value spread across the row by a table (l_eft).
239⍝# >> "ac:" u_se< "Accel"
240⍝# >> ac:s_oftmax 2 3 r_eshape 0.0 0.0 0.0 1.0 2.0 3.0
241⍝#  0.3333333333333333  0.3333333333333333 0.3333333333333333
242⍝# 0.09003057317038046 0.24472847105479764 0.6652409557748218
243ˡs̲oftmax ← { m →
244  c ← o̲ffsets f̲irst -1 t̲ake s̲hape m
245  e ← e̲xp m − ('m̲ax r̲/₂ m) 'l̲eft t̲able c
246  e ÷ ('+ r̲/₂ e) 'l̲eft t̲able c
247}
248
249⍝## The acceptance pipeline
250
251⍝# x ac:p_ipeline w: multiply, add, select, reduce in one expression,
252⍝# the sum of the positive items of x + x * w: the first end-to-end
253⍝# test research7 asks for (arithmetic, map, zip, comparison,
254⍝# selection, reduction), for Ints or Floats.
255⍝# >> "ac:" u_se< "Accel"
256⍝# >> 1 2 3 4 ac:p_ipeline 1 -2 1 -2
257⍝# 8
258ˡp̲ipeline ← { x w → ˡs̲um ˡr̲elu x + x × w }