librarylibs/Accel/src/Accel.xtl
Accel: the acceleratable subset of X_eTaL: elementwise arithmetic, masks, reductions, products, transposes, windows, lookups, choices, a layer
Import it with an alias of your choice: "ac:" u_se< "Accel". Put libs/Accel/src on XETAL_PATH ("just path"); the reference is libs/Accel/docs.
Every function here is one whole-array expression in the subset a GPU (or an FPGA) can take on directly: static scalar types, pure functions, arithmetic and comparisons item by item with scalar extension, a select by mask, reductions, inner products, tables and transposes. No recursion, no nested arrays, no input or output. The X_eTaL evaluator is the reference: X_eTaL-gpu runs the same functions on a GPU and must give the same answers (the demos are the programs it checks). Names with l: are exported; the others are private.
Elementwise
ˡv̲add : Num a => a -> a -> a
The sum, item by item: the first kernel, one work-item per element.
ᵃᶜ⁼u̲se< "Accel"
1 2 3 4 ᵃᶜv̲add 10 20 30 40 11 22 33 44
ˡv̲add ← { x y → x + y }
ˡr̲elu : Num a => a -> a
The items where positive, 0 elsewhere (a select by comparison with a constant), for Ints or Floats.
ᵃᶜ⁼u̲se< "Accel"
ᵃᶜr̲elu 1 -2 3 -4 1 0 3 0
ˡr̲elu ← { x → x m̲ax 0 }
Masks and select
ˡt̲hreshold : Float -> Float -> Float
t ac:t_hreshold x: x where x > t, 0.0 elsewhere, for Floats: the
mask x > t made Float and multiplied in, the select a GPU does
with a conditional move. A masked-out negative comes back as -0.0.
ᵃᶜ⁼u̲se< "Accel"
0.5 ᵃᶜt̲hreshold 0.25 0.5 0.75 1.0 0.0 0.0 0.75 1.0
ˡt̲hreshold ← { t x → x × f̲loat x > t }
ˡk̲eep : Num a => a -> a -> a
m ac:k_eep x: x where the Int mask m is 1, 0 elsewhere.
ᵃᶜ⁼u̲se< "Accel"
(1 0 1 0) ᵃᶜk̲eep 5 6 7 8 5 0 7 0
ˡk̲eep ← { m x → x × m }
Reductions
ˡs̲um : Num a => a -> a
The sum of the items of a vector: a parallel reduction on a GPU (work-group partials, then a second pass).
ᵃᶜ⁼u̲se< "Accel"
ᵃᶜs̲um 1 2 3 4 10
ˡs̲um ← { x → '+ r̲/ x }
ˡl̲argest : Num a => a -> a
ˡl̲argest ← { x → 'm̲ax r̲/ x }
Along an axis
ˡr̲owSums : Num a => a -> a
The sum of each row of a matrix (along axis 2: one work-item per row on a GPU).
ᵃᶜ⁼u̲se< "Accel"
ᵃᶜr̲owSums 2 3 r̲eshape 1 20 3 4 5 60 24 69
ˡr̲owSums ← { m → '+ r̲/₂ m }
ˡc̲olSums : Num a => a -> a
The sum of each column (along axis 1, the default of r_/).
ᵃᶜ⁼u̲se< "Accel"
ᵃᶜc̲olSums 2 3 r̲eshape 1 20 3 4 5 60 5 25 63
ˡc̲olSums ← { m → '+ r̲/ m }
ˡr̲owMax : Num a => a -> a
ˡr̲owMax ← { m → 'm̲ax r̲/₂ m }
Inner products
ˡm̲atmul : Num a => a -> a -> a
x ac:m_atmul w: the inner product of x's last axis with w's first (APL's +.x): a matrix times a matrix, a matrix times a vector, a vector times a matrix. Each result item is a sum of products, one GPU work-item each (or a tile of them, as the schedule says).
ᵃᶜ⁼u̲se< "Accel"
(2 3 r̲eshape r̲ange 6) ᵃᶜm̲atmul 3 2 r̲eshape r̲ange 6 22 28 49 64
(2 3 r̲eshape r̲ange 6) ᵃᶜm̲atmul 1 0 2 7 16
ˡm̲atmul ← { x w → x '+ '× i̲nner w }
Transposes
ˡs̲cores : Num a => a -> a -> a
q ac:s_cores k: attention scores, every query row against every key row (q k^T): the inner product of q with the transpose of k. On a GPU the transpose is a gather, one work-item per item, and the product follows (research7: attention decomposes into array operations; it is not an intrinsic).
ᵃᶜ⁼u̲se< "Accel"
(2 3 r̲eshape r̲ange 6) ᵃᶜs̲cores 4 3 r̲eshape r̲ange 12 14 32 50 68 32 77 122 167
ˡs̲cores ← { q k → q '+ '× i̲nner o̲\ k }
ˡw̲eightGrad : Num a => a -> a -> a
x ac:w_eightGrad dy: the gradient of a dense layer's weights, for a batch x (a row per example) and the gradient dy at the layer's output: x^T dy, the product every backward pass is made of.
ᵃᶜ⁼u̲se< "Accel"
(4 2 r̲eshape r̲ange 8) ᵃᶜw̲eightGrad 4 3 r̲eshape r̲ange 12 118 134 150 140 160 180
ˡw̲eightGrad ← { x dy → (o̲\ x) '+ '× i̲nner dy }
Lookups
ˡe̲mbed : a -> Int -> a
t ac:e_mbed ids: the rows of the table t at the positions ids (Ints from 1), in the shape of ids: a batch of token ids picks a batch of embeddings (ids s_elect t). On a GPU a gather, one work-item per item of the result, the table resident; the ids are all that a prompt uploads.
ᵃᶜ⁼u̲se< "Accel"
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
t ᵃᶜe̲mbed 4 1 2 3.0 3.1 3.2 0.0 0.1 0.2 1.0 1.1 1.2
ˡe̲mbed ← { t ids → ids s̲elect t }
Windows
ˡw̲indows : a -> a
ac:w_indows x: every 3 by 3 window of every picture of the batch x (n pictures of h by w) as 9 rows, a column per picture and place: the picture's nine shifted copies (a rotation along each of the last two axes, one copy per shift), their wrapped edges dropped, reshaped; the expression is X_eTaL-ML's Conv (cv:w_indows). On a GPU each rotation and each drop is a gather of its own and the reshape is nothing.
ᵃᶜ⁼u̲se< "Accel"
ᵃᶜw̲indows 1 3 4 r̲eshape r̲ange 12 1 2 2 3 3 4 5 6 6 7 7 8 9 10 10 11 11 12
ˡw̲indows ← { x → places ← ((2 s̲elect s̲hape x) − 2) × (3 s̲elect s̲hape x) − 2 (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 }
ˡc̲onv : Num a => a -> a -> a
k ac:c_onv x: the f filters k (f rows of 9, a 3 by 3 filter row by row) over every window of every picture of x, as one matrix product: the maps, n by f by (h - 2) by (w - 2) (a valid convolution without the bias; X_eTaL-ML's Conv adds its biases).
ᵃᶜ⁼u̲se< "Accel"
(1 9 r̲eshape 0 0 0 0 1 0 0 0 0) ᵃᶜc̲onv 1 3 4 r̲eshape r̲ange 12 6 7
ˡc̲onv ← { k x → n ← t̲ally x h ← (2 s̲elect s̲hape x) − 2 w ← (3 s̲elect s̲hape x) − 2 (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 }
Choices
ˡa̲rgmax : Num a => a -> Int
ac:a_rgmax m: for each row of m, the position (from 1) of its largest item, the first on a tie: the decision of a classifier, the next token of a greedy sampler. X_eTaL-ML's nn:a_rgmax, as array operations the IR has: the row's largest spread along the row, a mask where the row equals it, the smallest position under the mask (n + 1 elsewhere). No new GPU operation: a row reduction, a table, a fused mask and a second row reduction.
ᵃᶜ⁼u̲se< "Accel"
ᵃᶜ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 3 1 2
ˡa̲rgmax ← { m → n ← f̲irst -1 t̲ake s̲hape m hit ← f̲loat m = ('m̲ax r̲/₂ m) 'l̲eft t̲able o̲ffsets n 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 }
ˡt̲opk : Int -> Float -> Float
k ac:t_opk m: a mask of each row's k largest items (1.0 there, 0.0 elsewhere): the row's largest taken out k times, as X_eTaL-ML's Sample (sm:t_op) and its mixture-of-experts router choose; items tied with a taken one go with it, so a tie at the k-th keeps all tied. The choice of a router, the candidates of a retrieval. A taken item is pushed below its row by more than the row's range (Sample subtracts 1e300, which is infinity in a GPU's f32 and poisons the row with NaN).
ᵃᶜ⁼u̲se< "Accel"
2 ᵃᶜt̲opk 2 4 r̲eshape 0.1 0.5 0.15 0.25 7.0 8.0 6.0 8.0 0.0 1.0 0.0 1.0 1.0 1.0 0.0 1.0
ˡt̲opk ← { k m → n ← f̲irst -1 t̲ake s̲hape m shift ← (1.0 + ('m̲ax r̲/₂ m) − 'm̲in r̲/₂ m) 'l̲eft t̲able o̲ffsets n left ← k '{ q → q − shift × f̲loat q = ('m̲ax r̲/₂ q) 'l̲eft t̲able o̲ffsets n } p̲ower m f̲loat (m − left) > 0.0 }
A layer of a network
ˡd̲ense : Num a => a -> a -> a
x ac:d_ense wb: the layer x W + b for a batch x (one example per row), the weights W (a row per input, a column per output) with the bias b as one more row at the bottom of wb, as X_eTaL-ML's NN has it. The bias row is spread down the batch by a table (r_ight).
ᵃᶜ⁼u̲se< "Accel"
x ← 2 3 r̲eshape 1.0 2.0 3.0 -1.0 0.5 0.0
wb ← 4 2 r̲eshape 0.5 -1.0 1.0 0.25 -0.5 2.0 0.1 0.2
x ᵃᶜd̲ense wb 1.1 5.7 0.1 1.325
ˡ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 }
ˡs̲oftmax : Num a => a -> Float
Each row of a matrix into probabilities: the row's largest taken off (the same result, no overflow), e to each, each divided by its row's sum; a row's value spread across the row by a table (l_eft).
ᵃᶜ⁼u̲se< "Accel"
ᵃᶜs̲oftmax 2 3 r̲eshape 0.0 0.0 0.0 1.0 2.0 3.0 0.3333333333333333 0.3333333333333333 0.3333333333333333 0.09003057317038046 0.24472847105479764 0.6652409557748218
ˡs̲oftmax ← { m → c ← o̲ffsets f̲irst -1 t̲ake s̲hape m e ← e̲xp m − ('m̲ax r̲/₂ m) 'l̲eft t̲able c e ÷ ('+ r̲/₂ e) 'l̲eft t̲able c }
The acceptance pipeline
ˡp̲ipeline : Num a => a -> a -> a
x ac:p_ipeline w: multiply, add, select, reduce in one expression, the sum of the positive items of x + x * w: the first end-to-end test research7 asks for (arithmetic, map, zip, comparison, selection, reduction), for Ints or Floats.
ᵃᶜ⁼u̲se< "Accel"
1 2 3 4 ᵃᶜp̲ipeline 1 -2 1 -2 8
ˡp̲ipeline ← { x w → ˡs̲um ˡr̲elu x + x × w }