sourceextensions/gpu/demos/accelerate.xtl
1⍝!/usr/bin/env xetal
2⍝ Accelerate: two functions written as ordinary X_eTaL, handed to the
3⍝ GPU by the macro gp:a_ccel<. The macro expands each to a function of
4⍝ the same parameters that sends its text to the runner (the gpu
5⍝ extension), which lowers it for the arguments' shapes when they are
6⍝ bound and runs it on opencl:0; the rest of the program runs in the
7⍝ evaluator. Each result is checked against the evaluator's own. Run
8⍝ with "scripts/xx run extensions/gpu/demos/accelerate.xtl"; "xetal
9⍝ expand" shows what the macro makes of each definition.
10ᵍᵖ⁼u̲se< "Gpu"
11ᵃᶜ⁼u̲se< "Accel"
12
13⍝## A layer: a dense product with its bias, then ReLU
14ᵘl̲ayer ← "float float" ᵍᵖa̲ccel< "{ x wb -> ac:r_elu x ac:d_ense wb }"
15x ← 4 3 r̲eshape 0.5 -1.0 0.25 1.0 0.0 -0.5 -0.75 0.5 1.0 0.2 0.4 -0.6
16wb ← 4 5 r̲eshape 0.5 -0.25 0.75 0.1 -0.6 -0.3 0.8 0.2 -0.5 0.4 0.6 0.1 -0.7 0.3 0.2 0.05 -0.1 0.2 0.0 0.1
17h ← x ᵘl̲ayer wb
18h
19(ᵃᶜr̲elu x ᵃᶜd̲ense wb) ᵍᵖn̲ear h
20
21⍝## Attention scores of a batch of queries against the keys, as probabilities
22ᵘa̲ttend ← "float float" ᵍᵖa̲ccel< "{ q k -> ac:s_oftmax q ac:s_cores k }"
23q ← 2 3 r̲eshape 1.0 0.5 -1.0 0.0 2.0 0.25
24k ← 4 3 r̲eshape 0.5 1.0 0.0 -1.0 0.5 2.0 1.5 -0.5 1.0 0.0 0.0 1.0
25p ← q ᵘa̲ttend k
26p
27(ᵃᶜs̲oftmax q ᵃᶜs̲cores k) ᵍᵖn̲ear p
28
29⍝## An embedding lookup: the table resident, the ids the only upload
30ᵘe̲mbed ← "float int" ᵍᵖa̲ccel< "{ t ids -> t ac:e_mbed ids }"
31t ← 5 4 r̲eshape 0.0 0.1 0.2 0.3 1.0 1.1 1.2 1.3 2.0 2.1 2.2 2.3 3.0 3.1 3.2 3.3 4.0 4.1 4.2 4.3
32ids ← 2 3 r̲eshape 1 5 2 4 4 3
33e ← t ᵘe̲mbed ids
34'+ r̲/₂ e
35(t ᵃᶜe̲mbed ids) ᵍᵖn̲ear e
ᵍᵖa̲ccel< expands to
({ x wb → ᵍᵖo̲utput 1 ᵍᵖa̲fter ᵍᵖr̲un 0 ᵍᵖa̲fter 2 ᵍᵖb̲ind wb ᵍᵖa̲fter 1 ᵍᵖb̲ind x ᵍᵖa̲fter "{ x wb -> ac:r_elu x ac:d_ense wb }" ᵍᵖl̲ower "extensions/gpu/demos/accelerate.xtl|float float" })ᵍᵖa̲ccel< expands to
({ q k → ᵍᵖo̲utput 1 ᵍᵖa̲fter ᵍᵖr̲un 0 ᵍᵖa̲fter 2 ᵍᵖb̲ind k ᵍᵖa̲fter 1 ᵍᵖb̲ind q ᵍᵖa̲fter "{ q k -> ac:s_oftmax q ac:s_cores k }" ᵍᵖl̲ower "extensions/gpu/demos/accelerate.xtl|float float" })ᵍᵖa̲ccel< expands to
({ t ids → ᵍᵖo̲utput 1 ᵍᵖa̲fter ᵍᵖr̲un 0 ᵍᵖa̲fter 2 ᵍᵖb̲ind ids ᵍᵖa̲fter 1 ᵍᵖb̲ind t ᵍᵖa̲fter "{ t ids -> t ac:e_mbed ids }" ᵍᵖl̲ower "extensions/gpu/demos/accelerate.xtl|float int" })