commit
def3a479be
@ -0,0 +1,4 @@ |
||||
node_modules |
||||
build |
||||
data |
||||
models |
||||
@ -0,0 +1,59 @@ |
||||
import { FlatCompat } from '@eslint/eslintrc' |
||||
import js from '@eslint/js' |
||||
import stylistic from '@stylistic/eslint-plugin' |
||||
import stylisticts from '@stylistic/eslint-plugin-ts' |
||||
import ts from '@typescript-eslint/eslint-plugin' |
||||
import parser from '@typescript-eslint/parser' |
||||
import path from 'node:path' |
||||
import { fileURLToPath } from 'node:url' |
||||
|
||||
export default [ |
||||
...new FlatCompat({ |
||||
baseDirectory: path.dirname(fileURLToPath(import.meta.url)), |
||||
recommendedConfig: js.configs.recommended, |
||||
allConfig: js.configs.all, |
||||
}) |
||||
.extends( |
||||
'eslint:recommended', |
||||
'plugin:@typescript-eslint/eslint-recommended', |
||||
'plugin:@typescript-eslint/recommended', |
||||
), |
||||
stylistic.configs['recommended-flat'], |
||||
{ |
||||
files: ['src/**/*.ts', 'test/**/*.ts'], |
||||
languageOptions: { parser }, |
||||
plugins: { |
||||
'@typescript-eslint': ts, |
||||
'@stylistic': stylistic, |
||||
'@stylistic/ts': stylisticts, |
||||
}, |
||||
rules: { |
||||
'arrow-body-style': 'error', |
||||
'no-extra-parens': ['error', 'all'], |
||||
'no-unexpected-multiline': 'off', |
||||
'sort-imports': ['error', { ignoreDeclarationSort: true }], |
||||
'@stylistic/brace-style': ['error', '1tbs', { allowSingleLine: true }], |
||||
'@stylistic/quotes': ['error', 'single'], |
||||
'@stylistic/semi': ['error', 'always'], |
||||
'@stylistic/ts/semi': ['error', 'always'], |
||||
'@stylistic/indent': 'error', |
||||
'@stylistic/arrow-parens': ['error', 'as-needed', { requireForBlockBody: false }], |
||||
'@stylistic/max-statements-per-line': 'off', |
||||
'@stylistic/member-delimiter-style': 'off', |
||||
'@stylistic/ts/member-delimiter-style': 'error', |
||||
'@stylistic/no-mixed-operators': 'off', |
||||
'@typescript-eslint/no-unused-vars': [ |
||||
'error', |
||||
{ |
||||
args: 'all', |
||||
argsIgnorePattern: '^_', |
||||
caughtErrors: 'all', |
||||
caughtErrorsIgnorePattern: '^_', |
||||
destructuredArrayIgnorePattern: '^_', |
||||
varsIgnorePattern: '^_', |
||||
ignoreRestSiblings: true, |
||||
}, |
||||
], |
||||
'@typescript-eslint/no-explicit-any': 'off', |
||||
}, |
||||
}] |
||||
@ -0,0 +1,8 @@ |
||||
import type { Config } from 'jest' |
||||
|
||||
export default { |
||||
testEnvironment: 'node', |
||||
transform: { |
||||
'^.+.tsx?$': ['ts-jest', { tsconfig: 'test/tsconfig.json' }], |
||||
}, |
||||
} as Config |
||||
@ -0,0 +1,40 @@ |
||||
{ |
||||
"name": "nn", |
||||
"version": "1.0.0", |
||||
"type": "module", |
||||
"main": "index.js", |
||||
"author": "Freywar Ulvnaudgari <freywar.ulvnaudgari@gmail.com>", |
||||
"license": "MIT", |
||||
"scripts": { |
||||
"build": "rm -rf build && tsc", |
||||
"start": "yarn build && systemd-inhibit --what=idle:sleep --why='Training a model' node --import=extensionless/register build/index.js", |
||||
"start:debug": "yarn build && node --inspect-brk --import=extensionless/register build/index.js", |
||||
"start:profile": "yarn build && node --prof --import=extensionless/register build/index.js", |
||||
"test": "jest", |
||||
"test:debug": "node --inspect-brk node_modules/.bin/jest --runInBand", |
||||
"lint": "eslint", |
||||
"lint:fix": "eslint --fix" |
||||
}, |
||||
"devDependencies": { |
||||
"@eslint/eslintrc": "^3.1.0", |
||||
"@eslint/js": "^9.8.0", |
||||
"@jest/globals": "^29.7.0", |
||||
"@stylistic/eslint-plugin": "^2.6.1", |
||||
"@stylistic/eslint-plugin-ts": "^2.6.4", |
||||
"@typescript-eslint/eslint-plugin": "^8.0.0", |
||||
"@typescript-eslint/parser": "^8.0.0", |
||||
"eslint": "^9.8.0", |
||||
"extensionless": "^1.9.9", |
||||
"jest": "^29.7.0", |
||||
"ts-jest": "^29.2.3", |
||||
"ts-node": "^10.9.2", |
||||
"typescript": "^5.5.4" |
||||
}, |
||||
"dependencies": { |
||||
"@webgpu/glslang": "^0.0.15", |
||||
"@webgpu/types": "^0.1.60", |
||||
"ts-command-line-args": "^2.5.1", |
||||
"webgpu": "^0.2.11", |
||||
"yargs": "^17.7.2" |
||||
} |
||||
} |
||||
@ -0,0 +1,37 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
struct FVectorBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
struct UVectorBatch { |
||||
size: vec4<u32>, |
||||
data: array<u32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> ps: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read_write> result: UVectorBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size.x = ps.size.y; |
||||
result.size.z = ps.size.z; |
||||
|
||||
var mi = 0u; |
||||
var mp = 0f; |
||||
var pi = global_id.z * ps.size.y * ps.size.x + global_id.x * ps.size.x; |
||||
for (var i = 0u; i < ps.size.x; i++) { |
||||
if (ps.data[pi] > mp) { |
||||
mi = i; |
||||
mp = ps.data[pi]; |
||||
} |
||||
pi++; |
||||
} |
||||
result.data[global_id.z * result.size.x + global_id.x] = mi; |
||||
} |
||||
@ -0,0 +1,33 @@ |
||||
struct PackedMatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<f32>>, |
||||
} |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> e: PackedMatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> tw: PackedMatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> aw: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
aw.size.x = tw.size.y; |
||||
aw.size.y = e.size.y; |
||||
aw.size.z = e.size.z; |
||||
|
||||
var s = 0f; |
||||
let vc = e.size.x / 4; |
||||
let lo = (global_id.z * e.size.y + global_id.y) * vc; |
||||
let ro = (global_id.z * tw.size.y + global_id.x) * vc; |
||||
for (var i = 0u; i < vc; i++) { |
||||
s += dot(e.data[lo + i], tw.data[ro + i]); |
||||
} |
||||
|
||||
aw.data[(global_id.z * aw.size.y + global_id.y) * aw.size.x + global_id.x] = s; |
||||
} |
||||
@ -0,0 +1,32 @@ |
||||
alias number = f32; |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
struct Matrix { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> m: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read_write> result: Matrix; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = m.size; |
||||
result.size.z = 1; |
||||
|
||||
var vc = result.size.x / 4; |
||||
var resi = global_id.y * vc + global_id.x; |
||||
result.data[resi] = vec4(0f); |
||||
|
||||
var mi = resi; |
||||
for (var z = 0u; z < m.size.z; z++) { |
||||
result.data[resi] += m.data[mi] / vec4(f32(m.size.z)); |
||||
mi += m.size.y * vc; |
||||
} |
||||
} |
||||
@ -0,0 +1,18 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
@group(0) @binding(3) |
||||
var<storage, read_write> m: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
for (var z = 0u; z < m.size.z; z++) { |
||||
for (var y = 0u; y < m.size.y; y++) { |
||||
for (var x = 0u; x < m.size.x; x++) { |
||||
m.data[z * m.size.y * m.size.x + y * m.size.x + x] = 0f; |
||||
} |
||||
} |
||||
} |
||||
} |
||||
@ -0,0 +1,26 @@ |
||||
alias number = f32; |
||||
|
||||
struct ScalarBatch { |
||||
size: vec4<u32>, |
||||
data: array<number>, |
||||
} |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> l: ScalarBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> r_or_result: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = r_or_result.size; |
||||
|
||||
var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x; |
||||
result.data[resi] = l.data[global_id.z] * r_or_result.data[resi]; |
||||
} |
||||
@ -0,0 +1,21 @@ |
||||
alias number = f32; |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> l: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> r_or_result: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = l.size; |
||||
|
||||
var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x; |
||||
result.data[resi] = l.data[resi] / r_or_result.data[resi]; |
||||
} |
||||
@ -0,0 +1,21 @@ |
||||
alias number = f32; |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> l: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> r_or_result: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = l.size; |
||||
|
||||
var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x; |
||||
result.data[resi] = l.data[resi] * r_or_result.data[resi]; |
||||
} |
||||
@ -0,0 +1,26 @@ |
||||
alias number = f32; |
||||
|
||||
struct Scalar { |
||||
size: vec4<u32>, |
||||
data: array<number>, |
||||
} |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> l: Scalar; |
||||
@group(0) @binding(1) |
||||
var<storage, read> r_or_result: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = r_or_result.size; |
||||
|
||||
var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x; |
||||
result.data[resi] = l.data[0] * r_or_result.data[resi]; |
||||
} |
||||
@ -0,0 +1,21 @@ |
||||
alias number = f32; |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> l: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> r_or_result: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = l.size; |
||||
|
||||
var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x; |
||||
result.data[resi] = l.data[resi] - r_or_result.data[resi]; |
||||
} |
||||
@ -0,0 +1,21 @@ |
||||
alias number = f32; |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> l: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> r_or_result: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = l.size; |
||||
|
||||
var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x; |
||||
result.data[resi] = l.data[resi] + r_or_result.data[resi]; |
||||
} |
||||
@ -0,0 +1,34 @@ |
||||
alias number = f32; |
||||
|
||||
struct Scalar { |
||||
size: vec4<u32>, |
||||
data: array<number>, |
||||
} |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<number>, |
||||
} |
||||
|
||||
struct UVectorBatch { |
||||
size: vec4<u32>, |
||||
data: array<u32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> m: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> v: UVectorBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read> lr: Scalar; |
||||
@group(0) @binding(3) |
||||
var<storage, read_write> grs: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
grs.size = m.size; |
||||
|
||||
var vi = global_id.z * v.size.x + global_id.y; |
||||
var gi = vi * grs.size.x + global_id.x; |
||||
grs.data[gi] = lr.data[0] * (m.data[gi] - select(0f, 1f, v.data[vi] == global_id.x)); |
||||
} |
||||
@ -0,0 +1,48 @@ |
||||
struct Scalar { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> ds: Scalar; |
||||
@group(0) @binding(1) |
||||
var<storage, read> wa: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read> v: MatrixBatch; |
||||
@group(0) @binding(3) |
||||
var<storage, read> otdldo: MatrixBatch; |
||||
@group(0) @binding(4) |
||||
var<storage, read_write> dlds: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
dlds.size = wa.size; |
||||
|
||||
var s = 0f; |
||||
var li = (global_id.z * v.size.y + global_id.x) * v.size.x; |
||||
var ri = global_id.z * otdldo.size.x * otdldo.size.y + global_id.y; |
||||
|
||||
for (var i = 0u; i < v.size.x; i++) { |
||||
var cs = 0f; |
||||
var lj = global_id.z * wa.size.y * wa.size.x + global_id.y * wa.size.x; |
||||
var mj = global_id.z * v.size.x * v.size.y + i; |
||||
for (var j = 0u; j < wa.size.x; j++) { |
||||
cs += wa.data[lj] * v.data[mj]; |
||||
lj++; |
||||
mj += v.size.x; |
||||
} |
||||
|
||||
s += (v.data[li] - cs) * otdldo.data[ri]; |
||||
|
||||
li++; |
||||
ri += otdldo.size.x; |
||||
} |
||||
|
||||
let resi = (global_id.z * dlds.size.y + global_id.y) * dlds.size.x + global_id.x; |
||||
dlds.data[resi] = ds.data[0] * wa.data[resi] * s; |
||||
} |
||||
@ -0,0 +1,43 @@ |
||||
struct Scalar { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
struct Matrix { |
||||
size: vec4<u32>, |
||||
data: array<vec4<f32>>, |
||||
} |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<f32>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> ebs: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> ows: Matrix; |
||||
@group(0) @binding(2) |
||||
var<storage, read> tmp: Scalar; |
||||
@group(0) @binding(3) |
||||
var<storage, read_write> dts: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
dts.size.x = ows.size.y; |
||||
dts.size.y = ebs.size.y; |
||||
dts.size.z = ebs.size.z; |
||||
|
||||
let vc = ebs.size.x / 4; |
||||
let eo = (global_id.z * ebs.size.y + global_id.y) * vc; |
||||
let oo = 4 * global_id.x * vc; |
||||
var s = vec4(0f); |
||||
for (var i = 0u; i < vc; i++) { |
||||
let eb = ebs.data[eo + i]; |
||||
s.x += dot(eb, ows.data[oo + 0 * vc + i]); |
||||
s.y += dot(eb, ows.data[oo + 1 * vc + i]); |
||||
s.z += dot(eb, ows.data[oo + 2 * vc + i]); |
||||
s.w += dot(eb, ows.data[oo + 3 * vc + i]); |
||||
} |
||||
dts.data[(global_id.z * dts.size.y + global_id.y) * dts.size.x / 4 + global_id.x] = tmp.data[0] * s; |
||||
} |
||||
@ -0,0 +1,30 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> ilogits: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> inp: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read> pos: MatrixBatch; |
||||
@group(0) @binding(3) |
||||
var<storage, read_write> r: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
r.size = pos.size; |
||||
|
||||
var s = 0f; |
||||
var li = (global_id.z * ilogits.size.y + global_id.y) * ilogits.size.x; |
||||
var ri = global_id.z * inp.size.x * inp.size.y + global_id.x; |
||||
for (var i = 0u; i < ilogits.size.x; i++) { |
||||
s += ilogits.data[li] * inp.data[ri]; |
||||
li++; |
||||
ri += inp.size.x; |
||||
} |
||||
|
||||
var resi = (global_id.z * r.size.y + global_id.y) * r.size.x + global_id.x; |
||||
r.data[resi] = s + pos.data[resi]; |
||||
} |
||||
@ -0,0 +1,47 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> ilogits: MatrixBatch; |
||||
@group(0) @binding(0) |
||||
var<storage, read> tilogits: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> dldsq: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read> dldsk: MatrixBatch; |
||||
@group(0) @binding(3) |
||||
var<storage, read> totdldowa: MatrixBatch; |
||||
@group(0) @binding(4) |
||||
var<storage, read> tqry: MatrixBatch; |
||||
@group(0) @binding(5) |
||||
var<storage, read> tkey: MatrixBatch; |
||||
@group(0) @binding(6) |
||||
var<storage, read> tval: MatrixBatch; |
||||
@group(0) @binding(7) |
||||
var<storage, read_write> ginp: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
ginp.size.x = tval.size.x; |
||||
ginp.size.y = ilogits.size.x; |
||||
ginp.size.z = ilogits.size.z; |
||||
|
||||
var ts = 0f; |
||||
var li = global_id.z * ilogits.size.y * ilogits.size.x + global_id.y; |
||||
for (var i = 0u; i < ilogits.size.y; i++) { |
||||
var s = 0f; |
||||
var lj = global_id.z * totdldowa.size.x * totdldowa.size.y + i * totdldowa.size.x; |
||||
var rj = global_id.z * tval.size.x * tval.size.y + global_id.x; |
||||
for (var j = 0u; j < totdldowa.size.x; j++) { |
||||
s += totdldowa.data[lj] * tval.data[rj] + dldsq.data[lj] * tkey.data[rj] + dldsk.data[lj] * tqry.data[rj]; |
||||
lj++; |
||||
rj += tval.size.x; |
||||
} |
||||
ts += ilogits.data[li] * s; |
||||
li += ilogits.size.x; |
||||
} |
||||
|
||||
ginp.data[global_id.z * ginp.size.y * ginp.size.x + global_id.y * ginp.size.x + global_id.x] = ts; |
||||
} |
||||
@ -0,0 +1,28 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> e: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> dlds: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> gtattw: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
gtattw.size.z = e.size.z; |
||||
gtattw.size.y = e.size.x; |
||||
gtattw.size.x = e.size.x; |
||||
|
||||
var s = 0f; |
||||
var li = global_id.z * e.size.x * e.size.y + global_id.x; |
||||
var ri = global_id.z * dlds.size.x * dlds.size.y + global_id.y; |
||||
for (var i = 0u; i < e.size.y; i++) { |
||||
s += e.data[li] * dlds.data[ri]; |
||||
li += e.size.x; |
||||
ri += dlds.size.x; |
||||
} |
||||
gtattw.data[global_id.z * gtattw.size.y * gtattw.size.x + global_id.y * gtattw.size.x + global_id.x] = s; |
||||
} |
||||
@ -0,0 +1,28 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> a: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> dldo: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> gtout: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
gtout.size = a.size; |
||||
gtout.size.y = dldo.size.x; |
||||
|
||||
var s = 0f; |
||||
var ai = global_id.z * a.size.y * a.size.x + global_id.x; |
||||
var di = global_id.z * dldo.size.x * dldo.size.y + global_id.y; |
||||
for (var i = 0u; i < a.size.y; i++) { |
||||
s += a.data[ai] * dldo.data[di]; |
||||
ai += a.size.x; |
||||
di += dldo.size.x; |
||||
} |
||||
|
||||
gtout.data[global_id.z * gtout.size.y * gtout.size.x + global_id.y * gtout.size.x + global_id.x] = s; |
||||
} |
||||
@ -0,0 +1,19 @@ |
||||
struct UVectorBatch { |
||||
size: vec4<u32>, |
||||
data: array<u32>, |
||||
} |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> v: UVectorBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.data[global_id.z * result.size.y * result.size.x + global_id.y * result.size.x + global_id.x] = select(0f, 1f, global_id.x == v.data[global_id.z * v.size.x + global_id.y]); |
||||
} |
||||
@ -0,0 +1,37 @@ |
||||
alias number = f32; |
||||
|
||||
struct VectorBatch { |
||||
size: vec4<u32>, |
||||
data: array<u32>, |
||||
} |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> v: VectorBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> m: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = m.size; |
||||
|
||||
var vc = result.size.x / 4; |
||||
var resi = (global_id.z * result.size.y + global_id.y) * vc + global_id.x; |
||||
result.data[resi] = vec4(0f); |
||||
|
||||
var mi = global_id.z * m.size.y * vc + global_id.x; |
||||
var vi = global_id.z * v.size.x; |
||||
for (var i = 0u; i < v.size.x; i++) { |
||||
if (global_id.y == v.data[vi]) { |
||||
result.data[resi] += m.data[mi]; |
||||
} |
||||
mi += vc; |
||||
vi++; |
||||
} |
||||
} |
||||
@ -0,0 +1,31 @@ |
||||
alias number = f32; |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<number>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> l: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> r: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size.x = r.size.x; |
||||
result.size.y = l.size.y; |
||||
result.size.z = l.size.z; |
||||
|
||||
var s = 0f; |
||||
var li = global_id.z * l.size.x * l.size.y + global_id.y * l.size.x; |
||||
var ri = global_id.z * r.size.x * r.size.y + global_id.x; |
||||
for (var i = 0u; i < l.size.x; i++) { |
||||
s += l.data[li] * r.data[ri]; |
||||
li++; |
||||
ri += r.size.x; |
||||
} |
||||
|
||||
result.data[global_id.z * result.size.y * result.size.x + global_id.y * result.size.x + global_id.x] = s; |
||||
} |
||||
@ -0,0 +1,29 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> dldo: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> tout: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> otdldo: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
otdldo.size.x = dldo.size.y; |
||||
otdldo.size.y = tout.size.x; |
||||
otdldo.size.z = dldo.size.z; |
||||
|
||||
var s = 0f; |
||||
var li = (global_id.z * dldo.size.y + global_id.x) * dldo.size.x; |
||||
var ri = global_id.z * tout.size.x * tout.size.y + global_id.y; |
||||
for (var i = 0u; i < dldo.size.x; i++) { |
||||
s += dldo.data[li] * tout.data[ri]; |
||||
li++; |
||||
ri += tout.size.x; |
||||
} |
||||
|
||||
otdldo.data[(global_id.z * otdldo.size.y + global_id.y) * otdldo.size.x + global_id.x] = s; |
||||
} |
||||
@ -0,0 +1,39 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
struct FVectorBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
struct UVectorBatch { |
||||
size: vec4<u32>, |
||||
data: array<u32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> m: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> r: FVectorBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> result: UVectorBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size.x = m.size.y; |
||||
result.size.z = m.size.z; |
||||
var resi = global_id.z * result.size.x + global_id.x; |
||||
var tp = r.data[global_id.z * r.size.x + global_id.x]; |
||||
var p = 0f; |
||||
var mi = (global_id.z * m.size.y + global_id.x) * m.size.x; |
||||
for (var i = 0u; i < m.size.x; i++) { |
||||
p += m.data[mi]; |
||||
if (p >= tp) { |
||||
result.data[resi] = i; |
||||
break; |
||||
} |
||||
mi++; |
||||
} |
||||
} |
||||
@ -0,0 +1,25 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<f32>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> b: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = b.size; |
||||
|
||||
var vc = result.size.x / 4; |
||||
var resi = (global_id.z * result.size.y + global_id.y) * vc + global_id.x; |
||||
var mv = b.data[resi]; |
||||
var bi = (global_id.z * b.size.y + global_id.y) * vc; |
||||
for (var i = 0u; i < vc; i++) { |
||||
mv = max(mv, b.data[bi]); |
||||
bi++; |
||||
} |
||||
var ms = max(max(mv.x, mv.y), max(mv.z, mv.w)); |
||||
result.data[resi] = exp(b.data[resi] - vec4(ms)); |
||||
} |
||||
@ -0,0 +1,25 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<f32>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> b: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = b.size; |
||||
|
||||
var vc = result.size.x / 4; |
||||
var resi = (global_id.z * result.size.y + global_id.y) * vc + global_id.x; |
||||
var sv = vec4(0f); |
||||
var bi = (global_id.z * b.size.y + global_id.y) * vc; |
||||
for (var i = 0u; i < vc; i++) { |
||||
sv += b.data[bi]; |
||||
bi++; |
||||
} |
||||
var ss = sv.x + sv.y + sv.z + sv.w; |
||||
result.data[resi] = b.data[resi] / vec4(ss); |
||||
} |
||||
@ -0,0 +1,32 @@ |
||||
alias number = f32; |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
struct Matrix { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> m: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read_write> result: Matrix; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
result.size = m.size; |
||||
result.size.z = 1; |
||||
|
||||
var vc = result.size.x / 4; |
||||
var resi = global_id.y * vc + global_id.x; |
||||
result.data[resi] = vec4(0f); |
||||
|
||||
var mi = resi; |
||||
for (var z = 0u; z < m.size.z; z++) { |
||||
result.data[resi] += m.data[mi]; |
||||
mi += m.size.y * vc; |
||||
} |
||||
} |
||||
@ -0,0 +1,29 @@ |
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<f32>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> wa: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read> otdldo: MatrixBatch; |
||||
@group(0) @binding(2) |
||||
var<storage, read_write> totdldowa: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
totdldowa.size.x = otdldo.size.y; |
||||
totdldowa.size.y = wa.size.x; |
||||
totdldowa.size.z = otdldo.size.z; |
||||
|
||||
var s = 0f; |
||||
var li = global_id.z * otdldo.size.x * otdldo.size.y + global_id.x * otdldo.size.x; |
||||
var ri = global_id.z * wa.size.x * wa.size.y + global_id.y; |
||||
for (var i = 0u; i < otdldo.size.x; i++) { |
||||
s += otdldo.data[li] * wa.data[ri]; |
||||
li++; |
||||
ri += wa.size.x; |
||||
} |
||||
|
||||
totdldowa.data[(global_id.z * totdldowa.size.y + global_id.y) * totdldowa.size.x + global_id.x] = s; |
||||
} |
||||
@ -0,0 +1,20 @@ |
||||
alias number = f32; |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<number>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> in: MatrixBatch; |
||||
@group(0) @binding(1) |
||||
var<storage, read_write> out: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) { |
||||
out.size.x = in.size.y; |
||||
out.size.y = in.size.x; |
||||
out.size.z = in.size.z; |
||||
var l = global_id.z * out.size.y * out.size.x; |
||||
out.data[l + global_id.y * out.size.x + global_id.x] = in.data[l + global_id.x * in.size.x + global_id.y]; |
||||
} |
||||
@ -0,0 +1,26 @@ |
||||
alias number = f32; |
||||
|
||||
struct Matrix { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
struct MatrixBatch { |
||||
size: vec4<u32>, |
||||
data: array<vec4<number>>, |
||||
} |
||||
|
||||
@group(0) @binding(0) |
||||
var<storage, read> m: Matrix; |
||||
@group(0) @binding(1) |
||||
var<storage, read_write> result: MatrixBatch; |
||||
|
||||
@compute @workgroup_size(1) |
||||
fn main(@builtin(global_invocation_id) global_id: vec3<u32>, @builtin(num_workgroups) num_workgroups: vec3<u32>) { |
||||
result.size = m.size; |
||||
result.size.z = num_workgroups.z * 2; |
||||
|
||||
var vc = result.size.x / 4; |
||||
var mi = global_id.y * vc + global_id.x; |
||||
result.data[global_id.z * result.size.y * vc + mi] = m.data[mi]; |
||||
} |
||||
@ -0,0 +1,301 @@ |
||||
import fs from 'node:fs'; |
||||
import type { Scalar } from './scalar'; |
||||
|
||||
interface ComputePipeline { |
||||
pipeline: GPUComputePipeline; |
||||
layout: GPUBindGroupLayout; |
||||
} |
||||
|
||||
const PIPELINES: Map<GPUDevice, Map<string, ComputePipeline>> = new Map(); |
||||
|
||||
const COMPWISE_SHADER = { |
||||
sum: fs.readFileSync('shader/csum.wgsl').toString(), |
||||
sub: fs.readFileSync('shader/csub.wgsl').toString(), |
||||
mul: fs.readFileSync('shader/cmul.wgsl').toString(), |
||||
div: fs.readFileSync('shader/cdiv.wgsl').toString(), |
||||
scale: fs.readFileSync('shader/cscale.wgsl').toString(), |
||||
}; |
||||
|
||||
export function createComputePipeline( |
||||
device: GPUDevice, |
||||
name: string, |
||||
shader: () => string, |
||||
entry: string = 'main', |
||||
) { |
||||
if (!PIPELINES.has(device)) { |
||||
PIPELINES.set(device, new Map()); |
||||
} |
||||
|
||||
const d = PIPELINES.get(device)!; |
||||
if (!d.has(name)) { |
||||
const pipeline = device.createComputePipeline({ |
||||
label: name, |
||||
layout: 'auto', |
||||
compute: { |
||||
module: device.createShaderModule({ code: shader() }), |
||||
entryPoint: entry, |
||||
}, |
||||
}); |
||||
const layout = pipeline.getBindGroupLayout(0); |
||||
|
||||
d.set(name, { pipeline, layout }); |
||||
} |
||||
|
||||
return d.get(name)!; |
||||
} |
||||
|
||||
export abstract class GPUStruct<A extends Uint32Array | Int32Array | Float32Array> { |
||||
public static dunit = 8; |
||||
public static lunit = 2; |
||||
|
||||
protected readonly _label: string; |
||||
protected readonly _device: GPUDevice; |
||||
protected readonly _array: new (size: number | ArrayBufferLike | ArrayLike<number>) => A; |
||||
protected readonly _inbuffer: GPUBuffer; |
||||
protected readonly _dummybuffer: GPUBuffer; |
||||
protected readonly _outbuffer: GPUBuffer; |
||||
protected _op: 'reading' | 'writing' | null = null; |
||||
protected _dirty: 'front' | 'back' | null = 'front'; |
||||
|
||||
protected abstract workgroups(pm?: number): [number, number, number]; |
||||
protected abstract workgroup(): [number, number, number]; |
||||
|
||||
public readonly size: Int32Array; |
||||
public readonly data: A; |
||||
|
||||
public get type(): 'u32' | 'i32' | 'f32' { |
||||
switch (this._array as unknown) { |
||||
case Uint32Array: return 'u32'; |
||||
case Int32Array: return 'i32'; |
||||
case Float32Array: return 'f32'; |
||||
default: throw new Error('Unknown array type.'); |
||||
} |
||||
} |
||||
|
||||
public get length(): number { |
||||
return this.data.length; |
||||
} |
||||
|
||||
public get byteLength(): number { |
||||
return this.size.byteLength + this.data.byteLength; |
||||
} |
||||
|
||||
public constructor( |
||||
device: GPUDevice, |
||||
array: new (size: number) => A, |
||||
label: string, |
||||
size: [number, number, number, number], |
||||
) { |
||||
this._device = device; |
||||
this._array = array; |
||||
this._label = label ?? 'struct'; |
||||
this.size = new Int32Array(size); |
||||
this.data = new array(size[0] * size[1] * size[2] * size[3]); |
||||
|
||||
this._inbuffer = device.createBuffer({ |
||||
label: `${this._label}in`, |
||||
size: Math.max(32, this.byteLength), |
||||
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST, |
||||
}); |
||||
|
||||
this._dummybuffer = device.createBuffer({ |
||||
label: `${label}dummy`, |
||||
size: 0, |
||||
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST, |
||||
}); |
||||
|
||||
this._outbuffer = this._device.createBuffer({ |
||||
label: `${label}out`, |
||||
size: Math.max(32, this.byteLength), |
||||
usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ, |
||||
}); |
||||
} |
||||
|
||||
public write(): void { |
||||
if (!this._dirty) { |
||||
return; |
||||
} |
||||
|
||||
if (this._op) { |
||||
throw new Error('Sync in progress.'); |
||||
} |
||||
|
||||
this._op = 'writing'; |
||||
this._device.queue.writeBuffer(this._inbuffer, 0, new Uint8Array(this.size.buffer)); |
||||
this._device.queue.writeBuffer(this._inbuffer, this.size.byteLength, new Uint8Array(this.data.buffer)); |
||||
this._dirty = null; |
||||
this._op = null; |
||||
} |
||||
|
||||
public async read(): Promise<void> { |
||||
if (!this._dirty) { |
||||
return; |
||||
} |
||||
|
||||
if (this._op) { |
||||
throw new Error('Sync in progress.'); |
||||
} |
||||
|
||||
this._op = 'reading'; |
||||
|
||||
const copy = this._device.createCommandEncoder(); |
||||
copy.copyBufferToBuffer(this._inbuffer, 0, this._outbuffer, 0, this.byteLength); |
||||
this._device.queue.submit([copy.finish()]); |
||||
|
||||
await this._device.queue.onSubmittedWorkDone(); |
||||
|
||||
await this._outbuffer.mapAsync(GPUMapMode.READ); |
||||
this.data.set(new this._array(this._outbuffer.getMappedRange()).slice(this.size.length)); |
||||
this._outbuffer.unmap(); |
||||
this._dirty = null; |
||||
this._op = null; |
||||
} |
||||
|
||||
public buffer(): GPUBuffer { |
||||
if (this._dirty === 'front') { |
||||
this.write(); |
||||
} |
||||
return this._inbuffer; |
||||
} |
||||
|
||||
protected workgroupify(shader: string): string { |
||||
const [x, y, z] = this.workgroup(); |
||||
if (!x || !y || !z) { |
||||
throw new Error(`Invalid workgroup size: (${x}, ${y}, ${z}).`); |
||||
} |
||||
return shader.replace(/@workgroup_size\(.*?\)/g, `@workgroup_size(${x}, ${y}, ${z})`); |
||||
} |
||||
|
||||
protected inplacify(shader: string): string { |
||||
let r = ''; |
||||
return shader.replace(/\w+_or_result/g, v => r = r ? 'result' : v); // The first occurrence must be the declaration.
|
||||
} |
||||
|
||||
protected typify(shader: string): string { |
||||
return shader.replace(/alias number = .+?;/g, `alias number = ${this.type};`); |
||||
} |
||||
|
||||
public calculate( |
||||
name: string, |
||||
shader: string, |
||||
args: GPUStruct<Uint32Array | Int32Array | Float32Array>[], |
||||
ip = false, |
||||
pm = 1, |
||||
): typeof this { |
||||
const encoder = this._device.createCommandEncoder(); |
||||
const pass = encoder.beginComputePass(); |
||||
const { pipeline, layout } = createComputePipeline(this._device, `${name}${ip ? '_inplace_' : '_'}${this.type}`, () => this.workgroupify(this.typify(ip ? this.inplacify(shader) : shader))); |
||||
|
||||
pass.setPipeline(pipeline); |
||||
pass.setBindGroup(0, this._device.createBindGroup({ |
||||
layout, entries: [ |
||||
...args.map((s, i) => ({ binding: i, resource: { buffer: s.buffer() } })), |
||||
{ binding: args.length + +ip, resource: { buffer: this.buffer() } }, |
||||
], |
||||
})); |
||||
|
||||
const [x, y, z] = this.workgroups(pm); |
||||
if (!x || !y || !z || Math.floor(x) !== x || Math.floor(y) !== y || Math.floor(z) !== z) { |
||||
throw new Error(`Invalid workgroup count: (${x}, ${y}, ${z}).`); |
||||
} |
||||
pass.dispatchWorkgroups(x, y, z); |
||||
pass.end(); |
||||
|
||||
this._device.queue.submit([encoder.finish()]); |
||||
|
||||
this._dirty = 'back'; |
||||
|
||||
return this; |
||||
} |
||||
|
||||
protected static assertUnique(...ms: GPUStruct<Uint32Array | Int32Array | Float32Array>[]): void { |
||||
const map: Map<GPUStruct<Uint32Array | Int32Array | Float32Array>, boolean> = new Map(); |
||||
for (const m of ms) { |
||||
if (map.has(m)) { |
||||
throw new Error('Can not use the same GPU structure multiple times.'); |
||||
} |
||||
map.set(m, true); |
||||
} |
||||
} |
||||
|
||||
protected static assertColocated(...ms: GPUStruct<Uint32Array | Int32Array | Float32Array>[]): void { |
||||
const map: Map<GPUDevice, boolean> = new Map(); |
||||
for (const m of ms) { |
||||
map.set(m._device, true); |
||||
if (map.size > 1) { |
||||
throw new Error('Can not use GPU structures stored on multiple devices.'); |
||||
} |
||||
} |
||||
} |
||||
|
||||
protected static assertSimilar(...ms: GPUStruct<Uint32Array | Int32Array | Float32Array>[]): void { |
||||
for (const m of ms) { |
||||
if (m.length !== ms[0].length |
||||
|| m.workgroup()[0] !== ms[0].workgroup()[0] |
||||
|| m.workgroup()[1] !== ms[0].workgroup()[1] |
||||
|| m.workgroup()[2] !== ms[0].workgroup()[2]) { |
||||
throw new Error('Can not use GPU structures with different dimensions.'); |
||||
} |
||||
} |
||||
} |
||||
|
||||
protected static assertPackable(pm: number, ...ms: GPUStruct<Uint32Array | Int32Array | Float32Array>[]): void { |
||||
for (const m of ms) { |
||||
if (m.workgroups(pm).some(v => v < 1)) { |
||||
throw new Error(`Struct ${m._label} size must be divisible by ${pm}.`); |
||||
} |
||||
} |
||||
} |
||||
|
||||
public static compwise< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
S extends GPUStruct<A>, |
||||
>(op: keyof typeof COMPWISE_SHADER, l: S, r: S, result: S): S { |
||||
this.assertUnique(l, r); |
||||
this.assertUnique(l, result); |
||||
this.assertColocated(l, r, result); |
||||
this.assertPackable(4, l, r, result); |
||||
|
||||
return result.calculate(op, COMPWISE_SHADER[op], result === r ? [l] : [l, r], result === r, 4); |
||||
} |
||||
|
||||
public static csum< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
S extends GPUStruct<A>, |
||||
>(l: S, r: S, result: S): S { |
||||
return GPUStruct.compwise('sum', l, r, result); |
||||
} |
||||
|
||||
public static csub< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
S extends GPUStruct<A>, |
||||
>(l: S, r: S, result: S): S { |
||||
return GPUStruct.compwise('sub', l, r, result); |
||||
} |
||||
|
||||
public static cmul< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
S extends GPUStruct<A>, |
||||
>(l: S, r: S, result: S): S { |
||||
return GPUStruct.compwise('mul', l, r, result); |
||||
} |
||||
|
||||
public static cdiv< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
S extends GPUStruct<A>, |
||||
>(l: S, r: S, result: S): S { |
||||
return GPUStruct.compwise('div', l, r, result); |
||||
} |
||||
|
||||
public static cscale< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
S extends GPUStruct<A>, |
||||
>(l: Scalar<A>, r: S, result: S): S { |
||||
this.assertUnique(l, r); |
||||
this.assertUnique(l, result); |
||||
this.assertColocated(l, r, result); |
||||
this.assertPackable(4, r, result); |
||||
|
||||
return result.calculate('scale', COMPWISE_SHADER.scale, result === r ? [l] : [l, r], result === r, 4); |
||||
} |
||||
} |
||||
@ -0,0 +1,289 @@ |
||||
import fs from 'node:fs'; |
||||
import { Matrix as CPUMatrix, Vector as CPUVector } from '../math'; |
||||
import { GPUStruct } from './gpu-struct'; |
||||
import { Matrix } from './matrix'; |
||||
import { ScalarBatch } from './scalar-batch'; |
||||
import { VectorBatch } from './vector-batch'; |
||||
|
||||
const SHADER_SCALE = fs.readFileSync('shader/bscale.wgsl').toString(); |
||||
const SHADER_ISUM = fs.readFileSync('shader/isum.wgsl').toString(); |
||||
const SHADER_MUL = fs.readFileSync('shader/mul.wgsl').toString(); |
||||
const SHADER_TRANS = fs.readFileSync('shader/transpose.wgsl').toString(); |
||||
const SHADER_AVG = fs.readFileSync('shader/avg.wgsl').toString(); |
||||
const SHADER_TOT = fs.readFileSync('shader/tot.wgsl').toString(); |
||||
|
||||
export class MatrixBatch< |
||||
L extends number, |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
> extends GPUStruct<A> { |
||||
public get width(): W { |
||||
return this.size[0] as W; |
||||
} |
||||
|
||||
public get height(): H { |
||||
return this.size[1] as H; |
||||
} |
||||
|
||||
public get layers(): L { |
||||
return this.size[2] as L; |
||||
} |
||||
|
||||
protected override workgroups(pm: number = 1): [number, number, number] { |
||||
return [this.width / GPUStruct.dunit / pm, this.height / GPUStruct.dunit, this.layers / MatrixBatch.lunit]; |
||||
} |
||||
|
||||
protected override workgroup(): [number, number, number] { |
||||
return [GPUStruct.dunit, GPUStruct.dunit, MatrixBatch.lunit]; |
||||
} |
||||
|
||||
public constructor(device: GPUDevice, array: new (size: number) => A, data: CPUVector<L, CPUMatrix<H, W, T>>); |
||||
public constructor(device: GPUDevice, array: new (size: number) => A, l: L, h: H, w: W, data?: CPUVector<L, CPUMatrix<H, W, T>>); |
||||
public constructor(device: GPUDevice, array: new (size: number) => A, ...args: unknown[]) { |
||||
let [l, h, w]: [number, number, number] = [1, 1, 1]; |
||||
if (typeof args[0] === 'number') { |
||||
[l, h, w] = args as [L, H, W]; |
||||
} else { |
||||
const [data] = args as [CPUVector<L, CPUMatrix<H, W, T>>]; |
||||
[l, h, w] = [data.length, data[0].length, data[0][0].length, 1]; |
||||
} |
||||
|
||||
super(device, array, `${h}x${w}[${l}]`, [w, h, l, 1]); |
||||
|
||||
MatrixBatch.assertPackable(1, this); |
||||
|
||||
const data = (typeof args[0] === 'number' ? args[3] : args[0]) as CPUVector<L, CPUMatrix<H, W, T>> | undefined; |
||||
if (data) { |
||||
this.set(data); |
||||
} |
||||
} |
||||
|
||||
public index(l: number, y: number, x: number): number { |
||||
return (l * this.height + y) * this.width + x; |
||||
} |
||||
|
||||
public async set(y: number, x: number, v: T): Promise<void>; |
||||
public async set(data: CPUVector<L, CPUMatrix<H, W, T>>): Promise<void>; |
||||
public async set(slice: T[][][], ob?: number, oy?: number, ox?: number): Promise<void>; |
||||
public async set(m: MatrixBatch<L, H, W, A, T>): Promise<void>; |
||||
public async set(...args: unknown[]): Promise<void> { |
||||
if (typeof args[0] === 'number') { |
||||
const [l, y, x, v]: number[] = args as number[]; |
||||
|
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
|
||||
this.data[this.index(l, y, x)] = v; |
||||
this._dirty = 'front'; |
||||
} else if (Array.isArray(args[0])) { |
||||
const [data, ol = 0, oy = 0, ox = 0] = args as [T[][][], number | undefined, number | undefined, number | undefined]; |
||||
|
||||
if (this._dirty === 'back' && (ol || oy || ox || data.length < this.layers || data[0].length < this.height || data[0][0].length < this.width)) { |
||||
await this.read(); |
||||
} |
||||
|
||||
let l = ol; |
||||
let y = oy; |
||||
let x = ox; |
||||
for (const mat of data ?? []) { |
||||
if (l >= this.layers) { |
||||
break; |
||||
} |
||||
for (const row of mat ?? []) { |
||||
if (y >= this.height) { |
||||
break; |
||||
} |
||||
for (const v of row) { |
||||
if (x >= this.width) { |
||||
break; |
||||
} |
||||
this.data[this.index(l, y, x)] = v; |
||||
x++; |
||||
} |
||||
x = ox; |
||||
y++; |
||||
} |
||||
y = oy; |
||||
l++; |
||||
} |
||||
this._dirty = 'front'; |
||||
} else { |
||||
const [batch] = args as [MatrixBatch<L, H, W, A, T>]; |
||||
const copy = this._device.createCommandEncoder(); |
||||
copy.copyBufferToBuffer(batch.buffer(), 0, this.buffer(), 0, this.byteLength); |
||||
this._device.queue.submit([copy.finish()]); |
||||
this._dirty = 'back'; |
||||
} |
||||
} |
||||
|
||||
public async get(l: number, y: number, x: number): Promise<T>; |
||||
public async get(l: number): Promise<CPUMatrix<H, W, T>>; |
||||
public async get(): Promise<CPUVector<L, CPUMatrix<H, W, T>>>; |
||||
public async get(l?: number, y?: number, x?: number): Promise<T | CPUMatrix<H, W, T> | CPUVector<L, CPUMatrix<H, W, T>>> { |
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
if (typeof x === 'number') { |
||||
return this.data[this.index(l!, y!, x!)] as T; |
||||
} else if (typeof l === 'number') { |
||||
const result: CPUMatrix<H, W, T> = [] as unknown as CPUMatrix<H, W, T>; |
||||
for (let y: number = 0; y < this.height; y++) { |
||||
const row: CPUVector<W, T> = [] as unknown as CPUVector<W, T>; |
||||
for (let x: number = 0; x < this.width; x++) { |
||||
row.push(this.data[this.index(l, y, x)] as T); |
||||
} |
||||
result.push(row); |
||||
} |
||||
return result; |
||||
} else { |
||||
const result: CPUVector<L, CPUMatrix<H, W, T>> = [] as unknown as CPUVector<L, CPUMatrix<H, W, T>>; |
||||
for (let l: number = 0; l < this.layers; l++) { |
||||
const layer: CPUMatrix<H, W, T> = [] as unknown as CPUMatrix<H, W, T>; |
||||
for (let y: number = 0; y < this.height; y++) { |
||||
const row: CPUVector<W, T> = [] as unknown as CPUVector<W, T>; |
||||
for (let x: number = 0; x < this.width; x++) { |
||||
row.push(this.data[this.index(l, y, x)] as T); |
||||
} |
||||
layer.push(row); |
||||
} |
||||
result.push(layer); |
||||
} |
||||
return result; |
||||
} |
||||
} |
||||
|
||||
protected static override assertPackable(pm: number, ...ms: MatrixBatch<number, number, number, Uint32Array | Int32Array | Float32Array>[]): void { |
||||
for (const m of ms) { |
||||
if (m.layers % MatrixBatch.lunit !== 0 || m.height % MatrixBatch.dunit !== 0 || m.width % (MatrixBatch.dunit * pm) !== 0) { |
||||
throw new Error(`Batch size ${m.height}x${m.width}[${m.layers}] must be divisible by ${MatrixBatch.dunit}x${MatrixBatch.dunit * pm}[${MatrixBatch.lunit}].`); |
||||
} |
||||
} |
||||
} |
||||
|
||||
public static sum< |
||||
L extends number, |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: MatrixBatch<L, H, W, A, T>, r: MatrixBatch<L, H, W, A, T>, result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, l.height, l.width)): MatrixBatch<L, H, W, A, T> { |
||||
return GPUStruct.csum(l, r, result); |
||||
} |
||||
|
||||
public static sub< |
||||
L extends number, |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: MatrixBatch<L, H, W, A, T>, r: MatrixBatch<L, H, W, A, T>, result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, l.height, l.width)): MatrixBatch<L, H, W, A, T> { |
||||
return GPUStruct.csub(l, r, result); |
||||
} |
||||
|
||||
public static isum< |
||||
L extends number, |
||||
H extends number, |
||||
I extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
lis: VectorBatch<L, I, Int32Array>, |
||||
m: MatrixBatch<L, H, W, A, T>, |
||||
result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(m._device, m._array, m.layers, m.height, m.width), |
||||
): MatrixBatch<L, H, W, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m, result); |
||||
|
||||
return result.calculate('isum', SHADER_ISUM, [lis, m], false, 4); |
||||
} |
||||
|
||||
public static scale< |
||||
L extends number, |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: ScalarBatch<L, A, T>, r: MatrixBatch<L, H, W, A, T>, result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(r._device, r._array, r.layers, r.height, r.width)): MatrixBatch<L, H, W, A, T> { |
||||
this.assertUnique(l, r); |
||||
this.assertUnique(l, result); |
||||
this.assertColocated(l, r, result); |
||||
this.assertPackable(4, r, result); |
||||
|
||||
return result.calculate('scale', SHADER_SCALE, result === r ? [l] : [l, r], result === r, 4); |
||||
} |
||||
|
||||
public static mul< |
||||
L extends number, |
||||
H extends number, |
||||
S extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
l: MatrixBatch<L, H, S, A, T>, |
||||
r: MatrixBatch<L, S, W, A, T>, |
||||
result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, l.height, r.width), |
||||
): MatrixBatch<L, H, W, A, T> { |
||||
this.assertUnique(l, r, result); |
||||
this.assertColocated(l, r, result); |
||||
this.assertPackable(1, l, r, result); |
||||
|
||||
return result.calculate('mul', SHADER_MUL, [l, r]); |
||||
} |
||||
|
||||
public static transpose< |
||||
L extends number, |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
m: MatrixBatch<L, H, W, A, T>, |
||||
result: MatrixBatch<L, W, H, A, T> = new MatrixBatch(m._device, m._array, m.layers, m.width, m.height), |
||||
): MatrixBatch<L, W, H, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(1, m, result); |
||||
|
||||
return result.calculate('transpose', SHADER_TRANS, [m]); |
||||
} |
||||
|
||||
public static avg< |
||||
L extends number, |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
m: MatrixBatch<L, H, W, A, T>, |
||||
result: Matrix<H, W, A, T> = new Matrix(m._device, m._array, m.height, m.width), |
||||
): Matrix<H, W, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m); |
||||
|
||||
return result.calculate('avg', SHADER_AVG, [m], false, 4); |
||||
} |
||||
|
||||
public static tot< |
||||
L extends number, |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
m: MatrixBatch<L, H, W, A, T>, |
||||
result: Matrix<H, W, A, T> = new Matrix(m._device, m._array, m.height, m.width), |
||||
): Matrix<H, W, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m); |
||||
|
||||
return result.calculate('tot', SHADER_TOT, [m], false, 4); |
||||
} |
||||
} |
||||
@ -0,0 +1,233 @@ |
||||
import fs from 'node:fs'; |
||||
import { Matrix as CPUMatrix, Vector as CPUVector } from '../math'; |
||||
import { GPUStruct } from './gpu-struct'; |
||||
import { MatrixBatch } from './matrix-batch'; |
||||
import { Scalar } from './scalar'; |
||||
import { Vector } from './vector'; |
||||
|
||||
const SHADER_ISUM = fs.readFileSync('shader/isum.wgsl').toString(); |
||||
const SHADER_MUL = fs.readFileSync('shader/mul.wgsl').toString(); |
||||
const SHADER_TRANS = fs.readFileSync('shader/transpose.wgsl').toString(); |
||||
const SHADER_UNAVG = fs.readFileSync('shader/unavg.wgsl').toString(); |
||||
|
||||
export class Matrix< |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
> extends GPUStruct<A> { |
||||
public get width(): W { |
||||
return this.size[0] as W; |
||||
} |
||||
|
||||
public get height(): H { |
||||
return this.size[1] as H; |
||||
} |
||||
|
||||
protected override workgroups(pm: number = 1): [number, number, number] { |
||||
return [this.width / GPUStruct.dunit / pm, this.height / GPUStruct.dunit, 1]; |
||||
} |
||||
|
||||
protected override workgroup(): [number, number, number] { |
||||
return [GPUStruct.dunit, GPUStruct.dunit, 1]; |
||||
} |
||||
|
||||
public constructor(device: GPUDevice, array: new (size: number) => A, data: CPUMatrix<H, W, T>); |
||||
public constructor(device: GPUDevice, array: new (size: number) => A, h: H, w: W, data?: CPUMatrix<H, W, T>); |
||||
public constructor(device: GPUDevice, array: new (size: number) => A, ...args: unknown[]) { |
||||
let [h, w]: [number, number] = [1, 1]; |
||||
if (typeof args[0] === 'number') { |
||||
[h, w] = args as [H, W]; |
||||
} else { |
||||
const [data] = args as [CPUMatrix<H, W, T>]; |
||||
[h, w] = [data.length, data[0].length]; |
||||
} |
||||
|
||||
super(device, array, `${h}x${w}`, [w, h, 1, 1]); |
||||
|
||||
Matrix.assertPackable(1, this); |
||||
|
||||
const data = (typeof args[0] === 'number' ? args[2] : args[0]) as CPUMatrix<H, W, T> | undefined; |
||||
if (data) { |
||||
this.set(data); |
||||
} |
||||
} |
||||
|
||||
public index(y: number, x: number): number { |
||||
return y * this.width + x; |
||||
} |
||||
|
||||
public async set(y: number, x: number, v: T): Promise<void>; |
||||
public async set(data: CPUMatrix<H, W, T>): Promise<void>; |
||||
public async set(slice: T[][], oy?: number, ox?: number): Promise<void>; |
||||
public async set(m: Matrix<H, W, A, T>): Promise<void>; |
||||
public async set(...args: unknown[]): Promise<void> { |
||||
if (typeof args[0] === 'number') { |
||||
const [y, x, v]: number[] = args as number[]; |
||||
|
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
|
||||
this.data[this.index(y, x)] = v; |
||||
this._dirty = 'front'; |
||||
} else if (Array.isArray(args[0])) { |
||||
const [data, oy = 0, ox = 0] = args as [T[][], number | undefined, number | undefined]; |
||||
|
||||
if (this._dirty === 'back' && (oy || ox || data.length < this.height || data[0].length < this.width)) { |
||||
await this.read(); |
||||
} |
||||
|
||||
let y = oy; |
||||
let x = ox; |
||||
for (const row of data ?? []) { |
||||
if (y >= this.height) { |
||||
break; |
||||
} |
||||
for (const v of row) { |
||||
if (x >= this.width) { |
||||
break; |
||||
} |
||||
this.data[this.index(y, x)] = v; |
||||
x++; |
||||
} |
||||
x = ox; |
||||
y++; |
||||
} |
||||
this._dirty = 'front'; |
||||
} else { |
||||
const [matrix] = args as [Matrix<H, W, A, T>]; |
||||
const copy = this._device.createCommandEncoder(); |
||||
copy.copyBufferToBuffer(matrix.buffer(), 0, this.buffer(), 0, this.byteLength); |
||||
this._device.queue.submit([copy.finish()]); |
||||
this._dirty = 'back'; |
||||
} |
||||
} |
||||
|
||||
public async get(y: number, x: number): Promise<T>; |
||||
public async get(): Promise<CPUMatrix<H, W, T>>; |
||||
public async get(y?: number, x?: number): Promise<T | CPUMatrix<H, W, T>> { |
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
if (typeof x === 'number') { |
||||
return this.data[this.index(y!, x!)] as T; |
||||
} else { |
||||
const result: CPUMatrix<H, W, T> = [] as unknown as CPUMatrix<H, W, T>; |
||||
for (let y: number = 0; y < this.height; y++) { |
||||
const row: CPUVector<W, T> = [] as unknown as CPUVector<W, T>; |
||||
for (let x: number = 0; x < this.width; x++) { |
||||
row.push(this.data[this.index(y, x)] as T); |
||||
} |
||||
result.push(row); |
||||
} |
||||
return result; |
||||
} |
||||
} |
||||
|
||||
protected static override assertPackable(pm: number, ...ms: Matrix<number, number, Uint32Array | Int32Array | Float32Array>[]): void { |
||||
for (const m of ms) { |
||||
if (m.height % Matrix.dunit !== 0 || m.width % (Matrix.dunit * pm) !== 0) { |
||||
throw new Error(`Matrix size ${m.height}x${m.width} must be divisible by ${Matrix.dunit}x${Matrix.dunit * pm}.`); |
||||
} |
||||
} |
||||
} |
||||
|
||||
public static sum< |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Matrix<H, W, A, T>, r: Matrix<H, W, A, T>, result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.height, l.width)): Matrix<H, W, A, T> { |
||||
return GPUStruct.csum(l, r, result); |
||||
} |
||||
|
||||
public static sub< |
||||
H extends number, |
||||
W extends number, |
||||
|
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Matrix<H, W, A, T>, r: Matrix<H, W, A, T>, result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.height, l.width)): Matrix<H, W, A, T> { |
||||
return GPUStruct.csub(l, r, result); |
||||
} |
||||
|
||||
public static isum< |
||||
H extends number, |
||||
I extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
lis: Vector<I, Int32Array>, |
||||
m: Matrix<H, W, A, T>, |
||||
result: Matrix<H, W, A, T> = new Matrix(m._device, m._array, m.height, m.width), |
||||
): Matrix<H, W, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m, result); |
||||
|
||||
return result.calculate('isum', SHADER_ISUM, [lis, m], false, 4); |
||||
} |
||||
|
||||
public static scale< |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Scalar<A, T>, r: Matrix<H, W, A, T>, result: Matrix<H, W, A, T> = new Matrix(r._device, r._array, r.height, r.width)): Matrix<H, W, A, T> { |
||||
return GPUStruct.cscale(l, r, result); |
||||
} |
||||
|
||||
public static mul< |
||||
H extends number, |
||||
S extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
l: Matrix<H, S, A, T>, |
||||
r: Matrix<S, W, A, T>, |
||||
result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.height, r.width), |
||||
): Matrix<H, W, A, T> { |
||||
this.assertUnique(l, r, result); |
||||
this.assertColocated(l, r, result); |
||||
this.assertPackable(1, l, r, result); |
||||
|
||||
return result.calculate('mul', SHADER_MUL, [l, r]); |
||||
} |
||||
|
||||
public static transpose< |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
m: Matrix<H, W, A, T>, |
||||
result: Matrix<W, H, A, T> = new Matrix(m._device, m._array, m.width, m.height), |
||||
): Matrix<W, H, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(1, m, result); |
||||
|
||||
return result.calculate('transpose', SHADER_TRANS, [m]); |
||||
} |
||||
|
||||
public static unavg< |
||||
L extends number, |
||||
H extends number, |
||||
W extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
l: L, |
||||
m: Matrix<H, W, A, T>, |
||||
result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(m._device, m._array, l, m.height, m.width), |
||||
): MatrixBatch<L, H, W, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m); |
||||
|
||||
return result.calculate('unavg', SHADER_UNAVG, [m], false, 4); |
||||
} |
||||
} |
||||
@ -0,0 +1,176 @@ |
||||
import fs from 'node:fs'; |
||||
import { Vector as CPUVector } from '../math'; |
||||
import { GPUStruct } from './gpu-struct'; |
||||
import { Scalar } from './scalar'; |
||||
|
||||
const SHADER_AVG = fs.readFileSync('shader/avg.wgsl').toString(); |
||||
const SHADER_TOT = fs.readFileSync('shader/tot.wgsl').toString(); |
||||
|
||||
export class ScalarBatch< |
||||
L extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
> extends GPUStruct<A> { |
||||
public get layers(): L { |
||||
return this.size[2] as L; |
||||
} |
||||
|
||||
protected override workgroups(): [number, number, number] { |
||||
return [1, 1, this.layers / GPUStruct.lunit]; |
||||
} |
||||
|
||||
protected override workgroup(): [number, number, number] { |
||||
return [1, 1, GPUStruct.lunit]; |
||||
} |
||||
|
||||
public constructor(device: GPUDevice, array: new (size: number) => A, data: CPUVector<L, T>); |
||||
public constructor(device: GPUDevice, array: new (size: number) => A, l: L, data?: CPUVector<L, T>); |
||||
public constructor(device: GPUDevice, array: new (size: number) => A, ...args: unknown[]) { |
||||
const l: number = typeof args[0] === 'number' |
||||
? args[0] as L |
||||
: (args[0] as CPUVector<L, T>).length; |
||||
|
||||
super(device, array, `[${l}]`, [1, 1, l, 1]); |
||||
|
||||
ScalarBatch.assertPackable(1, this); |
||||
|
||||
const data = (typeof args[0] === 'number' ? args[1] : args[0]) as CPUVector<L, T> | undefined; |
||||
if (data) { |
||||
this.set(data); |
||||
} |
||||
} |
||||
|
||||
public async set(l: number, v: T): Promise<void>; |
||||
public async set(data: CPUVector<L, T>): Promise<void>; |
||||
public async set(slice: T[], ol?: number): Promise<void>; |
||||
public async set(b: ScalarBatch<L, A, T>): Promise<void>; |
||||
public async set(...args: unknown[]): Promise<void> { |
||||
if (typeof args[0] === 'number') { |
||||
const [l, v]: number[] = args as number[]; |
||||
|
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
|
||||
this.data[l] = v; |
||||
this._dirty = 'front'; |
||||
} else if (Array.isArray(args[0])) { |
||||
const [data, ol = 0] = args as [T[], number | undefined]; |
||||
|
||||
if (this._dirty === 'back' && (ol || data.length < this.layers)) { |
||||
await this.read(); |
||||
} |
||||
|
||||
let l = ol; |
||||
for (const v of data ?? []) { |
||||
if (l >= this.layers) { |
||||
break; |
||||
} |
||||
this.data[l] = v; |
||||
l++; |
||||
} |
||||
this._dirty = 'front'; |
||||
} else { |
||||
const [batch] = args as [ScalarBatch<L, A, T>]; |
||||
const copy = this._device.createCommandEncoder(); |
||||
copy.copyBufferToBuffer(batch.buffer(), 0, this.buffer(), 0, this.byteLength); |
||||
this._device.queue.submit([copy.finish()]); |
||||
this._dirty = 'back'; |
||||
} |
||||
} |
||||
|
||||
public async get(l: number): Promise<T>; |
||||
public async get(): Promise<CPUVector<L, T>>; |
||||
public async get(l?: number): Promise<T | CPUVector<L, T>> { |
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
if (typeof l === 'number') { |
||||
return this.data[l] as T; |
||||
} else { |
||||
const result: CPUVector<L, T> = [] as unknown as CPUVector<L, T>; |
||||
for (let l: number = 0; l < this.layers; l++) { |
||||
result.push(this.data[l] as T); |
||||
} |
||||
return result; |
||||
} |
||||
} |
||||
|
||||
protected static override assertPackable(_: number, ...ms: ScalarBatch<number, Uint32Array | Int32Array | Float32Array>[]): void { |
||||
for (const m of ms) { |
||||
if (m.layers % ScalarBatch.dunit !== 0) { |
||||
throw new Error(`Batch size [${m.layers}] must be divisible by [${ScalarBatch.dunit}].`); |
||||
} |
||||
} |
||||
} |
||||
|
||||
public static sum< |
||||
L extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: ScalarBatch<L, A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> { |
||||
return GPUStruct.csum(l, r, result); |
||||
} |
||||
|
||||
public static sub< |
||||
L extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: ScalarBatch<L, A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> { |
||||
return GPUStruct.csub(l, r, result); |
||||
} |
||||
|
||||
public static mul< |
||||
L extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: ScalarBatch<L, A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> { |
||||
return GPUStruct.cmul(l, r, result); |
||||
} |
||||
|
||||
public static div< |
||||
L extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: ScalarBatch<L, A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> { |
||||
return GPUStruct.cdiv(l, r, result); |
||||
} |
||||
|
||||
public static scale< |
||||
L extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Scalar<A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(r._device, r._array, r.layers)): ScalarBatch<L, A, T> { |
||||
return GPUStruct.cscale(l, r, result); |
||||
} |
||||
|
||||
public static avg< |
||||
L extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
m: ScalarBatch<L, A, T>, |
||||
result: Scalar<A, T> = new Scalar(m._device, m._array), |
||||
): Scalar<A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m); |
||||
|
||||
return result.calculate('avg', SHADER_AVG, [m], false, 4); |
||||
} |
||||
|
||||
public static tot< |
||||
L extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
m: ScalarBatch<L, A, T>, |
||||
result: Scalar<A, T> = new Scalar(m._device, m._array), |
||||
): Scalar<A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m); |
||||
|
||||
return result.calculate('tot', SHADER_TOT, [m], false, 4); |
||||
} |
||||
} |
||||
@ -0,0 +1,90 @@ |
||||
import fs from 'node:fs'; |
||||
import { GPUStruct } from './gpu-struct'; |
||||
import { ScalarBatch } from './scalar-batch'; |
||||
|
||||
const SHADER_UNAVG = fs.readFileSync('shader/unavg.wgsl').toString(); |
||||
|
||||
export class Scalar< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
> extends GPUStruct<A> { |
||||
protected override workgroups(): [number, number, number] { |
||||
return [1, 1, 1]; |
||||
} |
||||
|
||||
protected override workgroup(): [number, number, number] { |
||||
return [1, 1, 1]; |
||||
} |
||||
|
||||
public constructor(device: GPUDevice, array: new (size: number) => A, data?: T) { |
||||
super(device, array, 'scalar', [1, 1, 1, 1]); |
||||
this.set(data || 0 as T); |
||||
} |
||||
|
||||
public async set(v: T | Scalar<A, T>): Promise<void> { |
||||
if (typeof v === 'number') { |
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
|
||||
this.data[0] = v; |
||||
this._dirty = 'front'; |
||||
} else { |
||||
const copy = this._device.createCommandEncoder(); |
||||
copy.copyBufferToBuffer(v.buffer(), 0, this.buffer(), 0, this.byteLength); |
||||
this._device.queue.submit([copy.finish()]); |
||||
this._dirty = 'back'; |
||||
} |
||||
} |
||||
|
||||
public async get(): Promise<T> { |
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
return this.data[0] as T; |
||||
} |
||||
|
||||
public static sum< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Scalar<A, T>, r: Scalar<A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> { |
||||
return GPUStruct.csum(l, r, result); |
||||
} |
||||
|
||||
public static sub< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Scalar<A, T>, r: Scalar<A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> { |
||||
return GPUStruct.csub(l, r, result); |
||||
} |
||||
|
||||
public static mul< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Scalar<A, T>, r: Scalar<A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> { |
||||
return GPUStruct.cmul(l, r, result); |
||||
} |
||||
|
||||
public static div< |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Scalar<A, T>, r: Scalar<A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> { |
||||
return GPUStruct.cdiv(l, r, result); |
||||
} |
||||
|
||||
public static unavg< |
||||
L extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
l: L, |
||||
m: Scalar<A, T>, |
||||
result: ScalarBatch<L, A, T> = new ScalarBatch(m._device, m._array, l), |
||||
): ScalarBatch<L, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m); |
||||
|
||||
return result.calculate('unavg', SHADER_UNAVG, [m], false, 4); |
||||
} |
||||
} |
||||
@ -0,0 +1,221 @@ |
||||
import fs from 'node:fs'; |
||||
import { Vector as CPUVector } from '../math'; |
||||
import { GPUStruct } from './gpu-struct'; |
||||
import { ScalarBatch } from './scalar-batch'; |
||||
import { Vector } from './vector'; |
||||
|
||||
const SHADER_AVG = fs.readFileSync('shader/avg.wgsl').toString(); |
||||
const SHADER_TOT = fs.readFileSync('shader/tot.wgsl').toString(); |
||||
const SHADER_SCALE = fs.readFileSync('shader/bscale.wgsl').toString(); |
||||
const SHADER_DOT = fs.readFileSync('shader/mul.wgsl').toString(); |
||||
|
||||
export class VectorBatch< |
||||
L extends number, |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
> extends GPUStruct<A> { |
||||
public get dims(): D { |
||||
return this.size[0] as D; |
||||
} |
||||
|
||||
public get layers(): L { |
||||
return this.size[2] as L; |
||||
} |
||||
|
||||
protected override workgroups(pm: number = 1): [number, number, number] { |
||||
return [this.dims / GPUStruct.dunit / pm, 1, this.layers / GPUStruct.lunit]; |
||||
} |
||||
|
||||
protected override workgroup(): [number, number, number] { |
||||
return [GPUStruct.dunit, 1, GPUStruct.lunit]; |
||||
} |
||||
|
||||
public constructor(device: GPUDevice, array: new (size: number) => A, data: CPUVector<L, CPUVector<D, T>>); |
||||
public constructor(device: GPUDevice, array: new (size: number) => A, l: L, d: D, data?: CPUVector<L, CPUVector<D, T>>); |
||||
public constructor(device: GPUDevice, array: new (size: number) => A, ...args: unknown[]) { |
||||
let [l, d]: [number, number] = [1, 1]; |
||||
if (typeof args[0] === 'number') { |
||||
[l, d] = args as [L, D]; |
||||
} else { |
||||
const [data] = args as [CPUVector<L, CPUVector<D, T>>]; |
||||
[l, d] = [data.length, data[0].length]; |
||||
} |
||||
|
||||
super(device, array, `${d}[${l}]`, [d, 1, l, 1]); |
||||
|
||||
VectorBatch.assertPackable(1, this); |
||||
|
||||
const data = (typeof args[0] === 'number' ? args[2] : args[0]) as CPUVector<L, CPUVector<D, T>> | undefined; |
||||
if (data) { |
||||
this.set(data); |
||||
} |
||||
} |
||||
|
||||
public index(l: number, i: number): number { |
||||
return l * this.dims + i; |
||||
} |
||||
|
||||
public async set(l: number, i: number, v: T): Promise<void>; |
||||
public async set(data: CPUVector<L, CPUVector<D, T>>): Promise<void>; |
||||
public async set(slice: T[][], ol?: number, oi?: number): Promise<void>; |
||||
public async set(b: VectorBatch<L, D, A, T>): Promise<void>; |
||||
public async set(...args: unknown[]): Promise<void> { |
||||
if (typeof args[0] === 'number') { |
||||
const [l, i, v]: number[] = args as number[]; |
||||
|
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
|
||||
this.data[this.index(l, i)] = v; |
||||
this._dirty = 'front'; |
||||
} else if (Array.isArray(args[0])) { |
||||
const [data, ol = 0, oi = 0] = args as [T[][], number | undefined, number | undefined]; |
||||
|
||||
if (this._dirty === 'back' && (ol || oi || data.length < this.layers || data[0].length < this.dims)) { |
||||
await this.read(); |
||||
} |
||||
|
||||
let l = ol; |
||||
let i = oi; |
||||
for (const vec of data ?? []) { |
||||
if (l >= this.layers) { |
||||
break; |
||||
} |
||||
for (const v of vec) { |
||||
if (i >= this.dims) { |
||||
break; |
||||
} |
||||
this.data[this.index(l, i)] = v; |
||||
i++; |
||||
} |
||||
i = oi; |
||||
l++; |
||||
} |
||||
this._dirty = 'front'; |
||||
} else { |
||||
const [batch] = args as [VectorBatch<L, D, A, T>]; |
||||
const copy = this._device.createCommandEncoder(); |
||||
copy.copyBufferToBuffer(batch.buffer(), 0, this.buffer(), 0, this.byteLength); |
||||
this._device.queue.submit([copy.finish()]); |
||||
this._dirty = 'back'; |
||||
} |
||||
} |
||||
|
||||
public async get(l: number, i: number): Promise<T>; |
||||
public async get(l: number): Promise<CPUVector<D, T>>; |
||||
public async get(): Promise<CPUVector<L, CPUVector<D, T>>>; |
||||
public async get(l?: number, i?: number): Promise<T | CPUVector<D, T> | CPUVector<L, CPUVector<D, T>>> { |
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
if (typeof i === 'number') { |
||||
return this.data[this.index(l!, i!)] as T; |
||||
} else if (typeof l === 'number') { |
||||
const result: CPUVector<D, T> = [] as unknown as CPUVector<D, T>; |
||||
for (let i: number = 0; i < this.dims; i++) { |
||||
result.push(this.data[this.index(l, i)] as T); |
||||
} |
||||
return result; |
||||
} else { |
||||
const result: CPUVector<L, CPUVector<D, T>> = [] as unknown as CPUVector<L, CPUVector<D, T>>; |
||||
for (let l: number = 0; l < this.layers; l++) { |
||||
const vec: CPUVector<D, T> = [] as unknown as CPUVector<D, T>; |
||||
for (let i: number = 0; i < this.dims; i++) { |
||||
vec.push(this.data[this.index(l, i)] as T); |
||||
} |
||||
result.push(vec); |
||||
} |
||||
return result; |
||||
} |
||||
} |
||||
|
||||
protected static override assertPackable(pm: number, ...ms: VectorBatch<number, number, Uint32Array | Int32Array | Float32Array>[]): void { |
||||
for (const m of ms) { |
||||
if (m.layers % VectorBatch.lunit !== 0 || m.dims % (VectorBatch.dunit * pm) !== 0) { |
||||
throw new Error(`Batch size ${m.dims}[${m.layers}] must be divisible by ${VectorBatch.dunit * pm}[${VectorBatch.lunit}].`); |
||||
} |
||||
} |
||||
} |
||||
|
||||
public static sum< |
||||
L extends number, |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: VectorBatch<L, D, A, T>, r: VectorBatch<L, D, A, T>, result: VectorBatch<L, D, A, T> = new VectorBatch(l._device, l._array, l.layers, l.dims)): VectorBatch<L, D, A, T> { |
||||
return GPUStruct.csum(l, r, result); |
||||
} |
||||
|
||||
public static sub< |
||||
L extends number, |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: VectorBatch<L, D, A, T>, r: VectorBatch<L, D, A, T>, result: VectorBatch<L, D, A, T> = new VectorBatch(l._device, l._array, l.layers, l.dims)): VectorBatch<L, D, A, T> { |
||||
return GPUStruct.csub(l, r, result); |
||||
} |
||||
|
||||
public static mul< |
||||
L extends number, |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: ScalarBatch<L, A, T>, r: VectorBatch<L, D, A, T>, result: VectorBatch<L, D, A, T> = new VectorBatch(r._device, r._array, r.layers, r.dims)): VectorBatch<L, D, A, T> { |
||||
this.assertUnique(l, r); |
||||
this.assertUnique(l, result); |
||||
this.assertColocated(l, r, result); |
||||
this.assertPackable(4, r, result); |
||||
|
||||
return result.calculate('scale', SHADER_SCALE, result === r ? [l] : [l, r], result === r, 4); |
||||
} |
||||
|
||||
public static dot< |
||||
L extends number, |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: VectorBatch<L, D, A, T>, r: VectorBatch<L, D, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> { |
||||
throw new Error('Not implemented.'); |
||||
|
||||
this.assertUnique(l, r); |
||||
this.assertUnique(l, result); |
||||
this.assertColocated(l, r, result); |
||||
// this.assertPackable(4, r, result);
|
||||
|
||||
return result.calculate('dot', SHADER_DOT, [l, r]); |
||||
} |
||||
|
||||
public static avg< |
||||
L extends number, |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
m: VectorBatch<L, D, A, T>, |
||||
result: Vector<D, A, T> = new Vector(m._device, m._array, m.dims), |
||||
): Vector<D, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m); |
||||
|
||||
return result.calculate('avg', SHADER_AVG, [m], false, 4); |
||||
} |
||||
|
||||
public static tot< |
||||
L extends number, |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
m: VectorBatch<L, D, A, T>, |
||||
result: Vector<D, A, T> = new Vector(m._device, m._array, m.dims), |
||||
): Vector<D, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m); |
||||
|
||||
return result.calculate('tot', SHADER_TOT, [m], false, 4); |
||||
} |
||||
} |
||||
@ -0,0 +1,164 @@ |
||||
import fs from 'node:fs'; |
||||
import { Vector as CPUVector } from '../math'; |
||||
import { GPUStruct } from './gpu-struct'; |
||||
import { Scalar } from './scalar'; |
||||
import { VectorBatch } from './vector-batch'; |
||||
|
||||
const SHADER_UNAVG = fs.readFileSync('shader/unavg.wgsl').toString(); |
||||
|
||||
const SHADER_DOT = fs.readFileSync('shader/mul.wgsl').toString(); |
||||
|
||||
export class Vector< |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
> extends GPUStruct<A> { |
||||
public get dims(): D { |
||||
return this.size[0] as D; |
||||
} |
||||
|
||||
protected override workgroups(pm: number = 1): [number, number, number] { |
||||
return [this.dims / GPUStruct.dunit / pm, 1, 1]; |
||||
} |
||||
|
||||
protected override workgroup(): [number, number, number] { |
||||
return [GPUStruct.dunit, 1, 1]; |
||||
} |
||||
|
||||
public constructor(device: GPUDevice, array: new (size: number) => A, data: CPUVector<D, T>); |
||||
public constructor(device: GPUDevice, array: new (size: number) => A, d: D, data?: CPUVector<D, T>); |
||||
public constructor(device: GPUDevice, array: new (size: number) => A, ...args: unknown[]) { |
||||
const d: number = typeof args[0] === 'number' |
||||
? args[0] as D |
||||
: (args[0] as CPUVector<D, T>).length; |
||||
|
||||
super(device, array, `${d}`, [d, 1, 1, 1]); |
||||
|
||||
Vector.assertPackable(1, this); |
||||
|
||||
const data = (typeof args[0] === 'number' ? args[1] : args[0]) as CPUVector<D, T> | undefined; |
||||
if (data) { |
||||
this.set(data); |
||||
} |
||||
} |
||||
|
||||
public async set(i: number, v: T): Promise<void>; |
||||
public async set(data: CPUVector<D, T>): Promise<void>; |
||||
public async set(slice: T[], oi?: number): Promise<void>; |
||||
public async set(b: Vector<D, A, T>): Promise<void>; |
||||
public async set(...args: unknown[]): Promise<void> { |
||||
if (typeof args[0] === 'number') { |
||||
const [i, v]: number[] = args as number[]; |
||||
|
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
|
||||
this.data[i] = v; |
||||
this._dirty = 'front'; |
||||
} else if (Array.isArray(args[0])) { |
||||
const [data, oi = 0] = args as [T[], number | undefined]; |
||||
|
||||
if (this._dirty === 'back' && (oi || data.length < this.dims)) { |
||||
await this.read(); |
||||
} |
||||
|
||||
let i = oi; |
||||
for (const v of data ?? []) { |
||||
if (i >= this.dims) { |
||||
break; |
||||
} |
||||
this.data[i] = v; |
||||
i++; |
||||
} |
||||
this._dirty = 'front'; |
||||
} else { |
||||
const [vector] = args as [Vector<D, A, T>]; |
||||
const copy = this._device.createCommandEncoder(); |
||||
copy.copyBufferToBuffer(vector.buffer(), 0, this.buffer(), 0, this.byteLength); |
||||
this._device.queue.submit([copy.finish()]); |
||||
this._dirty = 'back'; |
||||
} |
||||
} |
||||
|
||||
public async get(i: number): Promise<T>; |
||||
public async get(): Promise<CPUVector<D, T>>; |
||||
public async get(i?: number): Promise<T | CPUVector<D, T>> { |
||||
if (this._dirty === 'back') { |
||||
await this.read(); |
||||
} |
||||
if (typeof i === 'number') { |
||||
return this.data[i] as T; |
||||
} else { |
||||
const result: CPUVector<D, T> = [] as unknown as CPUVector<D, T>; |
||||
for (let i: number = 0; i < this.dims; i++) { |
||||
result.push(this.data[i] as T); |
||||
} |
||||
return result; |
||||
} |
||||
} |
||||
|
||||
protected static override assertPackable(pm: number, ...ms: Vector<number, Uint32Array | Int32Array | Float32Array>[]): void { |
||||
for (const m of ms) { |
||||
if (m.dims % (Vector.dunit * pm) !== 0) { |
||||
throw new Error(`Vector size ${m.dims} must be divisible by ${Vector.dunit * pm}.`); |
||||
} |
||||
} |
||||
} |
||||
|
||||
public static sum< |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Vector<D, A, T>, r: Vector<D, A, T>, result: Vector<D, A, T> = new Vector(l._device, l._array, l.dims)): Vector<D, A, T> { |
||||
return GPUStruct.csum(l, r, result); |
||||
} |
||||
|
||||
public static sub< |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Vector<D, A, T>, r: Vector<D, A, T>, result: Vector<D, A, T> = new Vector(l._device, l._array, l.dims)): Vector<D, A, T> { |
||||
return GPUStruct.csub(l, r, result); |
||||
} |
||||
|
||||
public static mul< |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Scalar<A, T>, r: Vector<D, A, T>, result: Vector<D, A, T> = new Vector(r._device, r._array, r.dims)): Vector<D, A, T> { |
||||
return GPUStruct.cscale(l, r, result); |
||||
} |
||||
|
||||
public static dot< |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>, |
||||
>(l: Vector<D, A, T>, r: Vector<D, A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> { |
||||
throw new Error('Not implemented.'); |
||||
|
||||
this.assertUnique(l, r); |
||||
this.assertUnique(l, result); |
||||
this.assertColocated(l, r, result); |
||||
// this.assertPackable(4, r, result);
|
||||
|
||||
return result.calculate('dot', SHADER_DOT, [l, r]); |
||||
} |
||||
|
||||
public static unavg< |
||||
L extends number, |
||||
D extends number, |
||||
A extends Uint32Array | Int32Array | Float32Array, |
||||
T extends NonNullable<ReturnType<A['at']>>, |
||||
>( |
||||
l: L, |
||||
m: Vector<D, A, T>, |
||||
result: VectorBatch<L, D, A, T> = new VectorBatch(m._device, m._array, l, m.dims), |
||||
): VectorBatch<L, D, A, T> { |
||||
this.assertUnique(m, result); |
||||
this.assertColocated(m, result); |
||||
this.assertPackable(4, m); |
||||
|
||||
return result.calculate('unavg', SHADER_UNAVG, [m], false, 4); |
||||
} |
||||
} |
||||
@ -0,0 +1,100 @@ |
||||
import { List, ListItem } from './list'; |
||||
|
||||
interface InternalQueueItem<T> { |
||||
value: T; |
||||
self?: ListItem<InternalQueueItem<T>>; |
||||
orphan?: boolean; |
||||
} |
||||
|
||||
export interface QueueItem<T> { |
||||
readonly value: T; |
||||
readonly self: ListItem<QueueItem<T>>; |
||||
readonly orphan?: boolean; |
||||
} |
||||
|
||||
export abstract class GroupQueueBase<K, T> { |
||||
protected _groups: Map<K, List<InternalQueueItem<T>>> = new Map(); |
||||
|
||||
protected abstract hash(v: T): K; |
||||
|
||||
public get length(): number { |
||||
return this._groups.size; |
||||
} |
||||
|
||||
public constructor(values?: Iterable<T>) { |
||||
for (const v of values ?? []) { |
||||
this.insert(v); |
||||
} |
||||
} |
||||
|
||||
public insert(value: T): QueueItem<T> { |
||||
const hash: K = this.hash(value); |
||||
const group: List<InternalQueueItem<T>> = this._groups.get(hash) ?? this._groups.set(hash, new List()).get(hash)!; |
||||
const item: InternalQueueItem<T> = { value }; |
||||
return (item.self = group.append(item)).value as QueueItem<T>; |
||||
} |
||||
|
||||
public remove(item: QueueItem<T>): boolean { |
||||
if (item.orphan) { |
||||
return false; |
||||
} |
||||
const hash: K = this.hash(item.value); |
||||
const group: List<InternalQueueItem<T>> | undefined = this._groups.get(hash); |
||||
if (!group) { |
||||
return false; |
||||
} |
||||
group.delete(item.self); |
||||
if (!group.length) { |
||||
this._groups.delete(hash); |
||||
} |
||||
return true; |
||||
} |
||||
|
||||
public itemOf(value: T): QueueItem<T> | undefined { |
||||
const hash: K = this.hash(value); |
||||
const group: List<InternalQueueItem<T>> | undefined = this._groups.get(hash); |
||||
if (!group) { |
||||
return undefined; |
||||
} |
||||
for (const item of group) { |
||||
if (item.value === value) { |
||||
return item as QueueItem<T>; |
||||
} |
||||
} |
||||
return undefined; |
||||
} |
||||
|
||||
public peek(): List<QueueItem<T>> | undefined { |
||||
if (!this._groups.size) { |
||||
return undefined; |
||||
} |
||||
|
||||
let group: List<InternalQueueItem<T>> | undefined = undefined; |
||||
for (const g of this._groups.values()) { |
||||
if (g.length > (group?.length ?? 0)) { |
||||
group = g; |
||||
} |
||||
} |
||||
return group as List<QueueItem<T>>; |
||||
} |
||||
|
||||
public pop(): List<QueueItem<T>> | undefined { |
||||
const result: List<QueueItem<T>> | undefined = this.peek(); |
||||
for (const item of result ?? []) { |
||||
this.remove(item); |
||||
} |
||||
return result; |
||||
} |
||||
|
||||
public *[Symbol.iterator](): IterableIterator<List<QueueItem<T>>> { |
||||
for (const group of [...this._groups.values()].sort(({ length: ll }, { length: rl }) => rl - ll)) { |
||||
yield group as List<QueueItem<T>>; |
||||
} |
||||
} |
||||
} |
||||
|
||||
export class GroupQueue<T> extends GroupQueueBase<T, T> { |
||||
protected override hash(v: T): T { |
||||
return v; |
||||
} |
||||
} |
||||
@ -0,0 +1,127 @@ |
||||
interface InternalListItem<T> { |
||||
value: T; |
||||
prev: InternalListItem<T>; |
||||
next: InternalListItem<T>; |
||||
orphan?: boolean; |
||||
} |
||||
|
||||
export interface ListItem<T> { |
||||
readonly value: T; |
||||
readonly prev: ListItem<T>; |
||||
readonly next: ListItem<T>; |
||||
readonly orphan?: boolean; |
||||
} |
||||
|
||||
export class List<T> implements Iterable<T> { |
||||
protected _head?: InternalListItem<T>; |
||||
protected _length: number = 0; |
||||
|
||||
public get first(): T | undefined { |
||||
return this._head?.value; |
||||
} |
||||
|
||||
public get last(): T | undefined { |
||||
return this._head?.prev?.value; |
||||
} |
||||
|
||||
public get length(): number { |
||||
return this._length; |
||||
} |
||||
|
||||
public constructor(values?: Iterable<T>) { |
||||
for (const v of values ?? []) { |
||||
this.append(v); |
||||
} |
||||
} |
||||
|
||||
public prepend(value: T): ListItem<T> { |
||||
if (!this._head) { |
||||
this._head = { value } as InternalListItem<T>; |
||||
this._head.prev = this._head.next = this._head; |
||||
} else { |
||||
this._head = { |
||||
value, |
||||
next: this._head, |
||||
prev: this._head.prev, |
||||
}; |
||||
this._head.prev.next = this._head.next.prev = this._head; |
||||
} |
||||
this._length++; |
||||
return this._head; |
||||
} |
||||
|
||||
public append(value: T): ListItem<T> { |
||||
this.prepend(value); |
||||
this._head = this._head?.next; |
||||
return this._head!.prev; |
||||
} |
||||
|
||||
public includes(value: T): boolean { |
||||
return !!this.itemOf(value); |
||||
} |
||||
|
||||
public itemOf(value: T): ListItem<T> | undefined { |
||||
for (const item of this.items()) { |
||||
if (item.value === value) { |
||||
return item; |
||||
} |
||||
} |
||||
return undefined; |
||||
} |
||||
|
||||
public set(item: ListItem<T>, value: T): boolean { |
||||
if (item.orphan) { |
||||
return false; |
||||
} |
||||
|
||||
(item as InternalListItem<T>).value = value; |
||||
|
||||
return true; |
||||
} |
||||
|
||||
public delete(item: ListItem<T>): boolean { |
||||
if (item.orphan) { |
||||
return false; |
||||
} |
||||
|
||||
if (this._head === item) { |
||||
this._head = item.next; |
||||
} |
||||
|
||||
if (this._head?.next === this._head) { |
||||
this._head = undefined; |
||||
} |
||||
|
||||
(item as InternalListItem<T>).prev.next = item.next; |
||||
(item as InternalListItem<T>).next.prev = item.prev; |
||||
(item as InternalListItem<T>).orphan = true; |
||||
|
||||
this._length--; |
||||
|
||||
return true; |
||||
} |
||||
|
||||
public *[Symbol.iterator](): IterableIterator<T> { |
||||
for (const item of this.items()) { |
||||
yield item.value; |
||||
} |
||||
} |
||||
|
||||
public values(): IterableIterator<T> { |
||||
return this[Symbol.iterator](); |
||||
} |
||||
|
||||
public *items(): IterableIterator<ListItem<T>> { |
||||
if (!this._head) { |
||||
return; |
||||
} |
||||
|
||||
let head = this._head; |
||||
let item = this._head; |
||||
do { |
||||
head = this._head; |
||||
yield item; |
||||
item = item.next; |
||||
} while (item !== head); |
||||
} |
||||
} |
||||
@ -0,0 +1,142 @@ |
||||
export type Vector<D extends number, I = number> = I[] & { length: D }; |
||||
export type Matrix<R extends number, C extends number, I = number> = Vector<R, Vector<C, I>>; |
||||
|
||||
export function vec<D extends number, I = number>(s: D, f: (i: number) => I): Vector<D, I> { |
||||
return Array.from({ length: s }, (_, i) => f(i)) as Vector<D, I>; |
||||
} |
||||
|
||||
export function mat<R extends number, C extends number, I = number>(r: R, c: C, f: (x: number, y: number) => number): Matrix<R, C, I> { |
||||
return Array.from({ length: r }, (_, y) => Array.from({ length: c }, (_, x) => f(x, y))) as Matrix<R, C, I>; |
||||
} |
||||
|
||||
export function sum<R extends number, C extends number>(l: Matrix<R, C>, r: Matrix<R, C>): Matrix<R, C> { |
||||
const result: Matrix<R, C> = mat(l.length, l[0].length, () => 0); |
||||
for (let i = 0; i < l.length; i++) { |
||||
for (let j = 0; j < r[0].length; j++) { |
||||
result[i][j] += l[i][j] + r[i][j]; |
||||
} |
||||
} |
||||
return result; |
||||
} |
||||
|
||||
export function mul<R extends number, C extends number>(l: Matrix<R, C>, r: number): Matrix<R, C>; |
||||
export function mul<D0 extends number, D1 extends number>(l: Vector<D0>, r: Matrix<D0, D1>): Vector<D1>; |
||||
export function mul<R extends number, S extends number, C extends number>(l: Matrix<R, S>, r: Matrix<S, C>): Matrix<R, C>; |
||||
export function mul(l: Matrix<number, number> | Vector<number>, r: Matrix<number, number> | number): Matrix<number, number> | Vector<number> { |
||||
if (typeof r === 'number') { |
||||
const lm: Matrix<number, number> = l as Matrix<number, number>; |
||||
const result = Array.from({ length: lm.length }, () => Array(lm[0].length).fill(0)); |
||||
for (let i = 0; i < lm.length; i++) { |
||||
for (let j = 0; j < lm[i].length; j++) { |
||||
result[i][j] += lm[i][j] * r; |
||||
} |
||||
} |
||||
return result; |
||||
} else if (!Array.isArray(l[0])) { |
||||
const lv: Vector<number> = l as Vector<number>; |
||||
const result = Array(r[0].length).fill(0); |
||||
for (let j = 0; j < r[0].length; j++) { |
||||
for (let i = 0; i < lv.length; i++) { |
||||
result[j] += lv[i] * r[i][j]; |
||||
} |
||||
} |
||||
|
||||
return result; |
||||
} else { |
||||
const lm: Matrix<number, number> = l as Matrix<number, number>; |
||||
const result = Array.from({ length: lm.length }, () => Array(r[0].length).fill(0)); |
||||
for (let i = 0; i < lm.length; i++) { |
||||
for (let j = 0; j < r[0].length; j++) { |
||||
for (let k = 0; k < r.length; k++) { |
||||
result[i][j] += lm[i][k] * r[k][j]; |
||||
} |
||||
} |
||||
} |
||||
return result; |
||||
} |
||||
} |
||||
|
||||
export function transpose<R extends number, C extends number>(m: Matrix<R, C>): Matrix<C, R> { |
||||
return mat(m[0].length, m.length, (x, y) => m[x][y]); |
||||
} |
||||
|
||||
export function softmax(row: number[]): number[] { |
||||
const max = Math.max(...row); |
||||
const exps = row.map(v => Math.exp(v - max)); |
||||
const sum = exps.reduce((a, b) => a + b, 0); |
||||
return exps.map(v => v / sum); |
||||
} |
||||
|
||||
export function msoftmax<R extends number, C extends number>(m: Matrix<R, C>): Matrix<R, C> { |
||||
return m.map(softmax) as Matrix<R, C>; |
||||
} |
||||
|
||||
export function norm<R extends number, C extends number>(x: Matrix<R, C>, epsilon = 1e-5): Matrix<R, C> { |
||||
return x.map(vec => { |
||||
const mean = vec.reduce((sum, val) => sum + val, 0) / vec.length; |
||||
const variance = vec.reduce((sum, val) => sum + (val - mean) ** 2, 0) / vec.length; |
||||
return vec.map(val => (val - mean) / Math.sqrt(variance + epsilon)); |
||||
}) as Matrix<R, C>; |
||||
} |
||||
|
||||
export function relu<D extends number>(vec: Vector<D>): Vector<D> { |
||||
return vec.map(v => Math.max(0, v)) as Vector<D>; |
||||
} |
||||
|
||||
export function mean<D extends number>(v: Vector<D>): number { |
||||
return v.reduce((s, c) => s + c, 0) / v.length; |
||||
} |
||||
|
||||
export function drm(data: number[], interval = 0.9): number { |
||||
const sorted = [...data].sort((a, b) => a - b); |
||||
const count = Math.floor(data.length * interval); |
||||
let mi = 0; |
||||
let mv = Infinity; |
||||
|
||||
for (let i = 0; i <= data.length - count; i++) { |
||||
const range = sorted[i + count - 1] - sorted[i]; |
||||
if (range < mv) { |
||||
mv = range; |
||||
mi = i; |
||||
} |
||||
} |
||||
|
||||
return sorted.slice(mi, mi + count).reduce((a, b) => a + b, 0) / count; |
||||
} |
||||
|
||||
export function randseq(n: number): number[] { |
||||
const is = Array.from({ length: n }, (_, i) => i); |
||||
for (let i = n - 1; i > 0; i--) { |
||||
const j = Math.floor(Math.random() * (i + 1)); |
||||
[is[i], is[j]] = [is[j], is[i]]; // swap
|
||||
} |
||||
return is; |
||||
} |
||||
|
||||
export function stddev(xs: number[]): number { |
||||
const m = mean(xs); |
||||
return Math.sqrt(xs.reduce((s, v) => s + (v - m) ** 2, 0) / xs.length); |
||||
} |
||||
|
||||
export function ema(xs: number[], alpha = 0.1): number { |
||||
let result = xs[0]; |
||||
for (let i = 1; i < xs.length; i++) { |
||||
result = alpha * xs[i] + (1 - alpha) * result; |
||||
} |
||||
return result; |
||||
} |
||||
|
||||
export function slope(xs: number[]): number { |
||||
const xMean = (xs.length - 1) / 2; |
||||
const yMean = mean(xs); |
||||
|
||||
let num = 0, den = 0; |
||||
for (let i = 0; i < xs.length; i++) { |
||||
const x = i - xMean; |
||||
const y = xs[i] - yMean; |
||||
num += x * y; |
||||
den += x * x; |
||||
} |
||||
|
||||
return num / den; |
||||
} |
||||
@ -0,0 +1,73 @@ |
||||
export interface InternalTrieNode { |
||||
children: Map<string, InternalTrieNode>; |
||||
token?: string; |
||||
} |
||||
|
||||
export interface TrieNode { |
||||
readonly children: Map<string, TrieNode>; |
||||
readonly token?: string; |
||||
} |
||||
|
||||
export class Trie implements Iterable<string> { |
||||
protected root: InternalTrieNode = { children: new Map() }; |
||||
|
||||
constructor(values?: string[]) { |
||||
for (const value of values ?? []) { |
||||
this.insert(value); |
||||
} |
||||
} |
||||
|
||||
public insert(value: string) { |
||||
let node = this.root; |
||||
for (const char of value) { |
||||
node = node.children.get(char) ?? node.children.set(char, { children: new Map() }).get(char)!; |
||||
} |
||||
node.token = value; |
||||
} |
||||
|
||||
public match(value: string): string | undefined { |
||||
let node: TrieNode = this.root; |
||||
let match: string | undefined = undefined; |
||||
for (let i = 0; i < value.length && node.children.has(value[i]); i++) { |
||||
match = (node = node.children.get(value[i])!).token ?? match; |
||||
} |
||||
return match; |
||||
} |
||||
|
||||
public *[Symbol.iterator](): IterableIterator<string> { |
||||
const queue: InternalTrieNode[] = []; |
||||
let item: InternalTrieNode | undefined = this.root; |
||||
while (item) { |
||||
if (item.token) { |
||||
yield item.token; |
||||
} |
||||
queue.push(...item.children.values()); |
||||
item = queue.pop(); |
||||
} |
||||
} |
||||
|
||||
public toJSON(node: InternalTrieNode = this.root): object { |
||||
return [ |
||||
...node.token ? [node.token] : [], |
||||
...[...node.children.entries()].map(([k, c]) => [k, this.toJSON(c)]), |
||||
]; |
||||
} |
||||
|
||||
protected static nodeFromJSON(element: object): InternalTrieNode { |
||||
if (!Array.isArray(element) || !element.every(e => typeof e === 'string' |
||||
|| Array.isArray(e) && e.length === 2)) { |
||||
throw new Error('Malformed Trie JSON.'); |
||||
} |
||||
|
||||
return { |
||||
children: new Map((typeof element[0] === 'string' ? element.slice(1) : element).map(([k, v]) => [k, this.nodeFromJSON(v)])), |
||||
token: typeof element[0] === 'string' ? element[0] : undefined, |
||||
}; |
||||
} |
||||
|
||||
public static fromJSON(json: object): Trie { |
||||
const result = new Trie(); |
||||
result.root = Trie.nodeFromJSON(json); |
||||
return result; |
||||
} |
||||
} |
||||
@ -0,0 +1,196 @@ |
||||
import fs from 'node:fs'; |
||||
import { create, globals } from 'webgpu'; |
||||
import yargs from 'yargs'; |
||||
import { Matrix } from './data/gpu/matrix'; |
||||
import { Scalar } from './data/gpu/scalar'; |
||||
import { mul, sum, transpose } from './data/math'; |
||||
import { Model } from './model-gpu'; |
||||
import { Tokenizer } from './tokenizer'; |
||||
|
||||
Object.assign(globalThis, globals); |
||||
|
||||
const GPU = create([]); |
||||
|
||||
const device = await GPU.requestAdapter({ powerPreference: 'high-performance' }).then(adapter => adapter?.requestDevice()); |
||||
|
||||
if (!device) { |
||||
process.stderr.write('Can not find device.\n'); |
||||
process.exit(1); |
||||
} |
||||
process.stderr.write(`Running on: ${device.adapterInfo.description}.\n`); |
||||
|
||||
const args = await yargs() |
||||
.option('model', { type: 'string', demandOption: true }) |
||||
.command('test-matrix', 'Test core components') |
||||
.command('train-tokenizer', 'Train new tokenizer', yargs => yargs |
||||
.option('vocab', { type: 'number', demandOption: true })) |
||||
.command('tokenize', 'Tokenize input using existing tokenizer', yargs => yargs |
||||
.option('input', { type: 'string' }) |
||||
.option('output', { type: 'string' })) |
||||
.command('configure', 'Configure new model', yargs => yargs |
||||
.option('window', { type: 'number', demandOption: true }) |
||||
.option('dimensions', { type: 'number', demandOption: true }) |
||||
.option('heads', { type: 'number', demandOption: true }) |
||||
.option('ffs', { type: 'number', demandOption: true })) |
||||
.command('train', 'Train existing model', yargs => yargs |
||||
.option('epochs', { type: 'number' }) |
||||
.option('stats', { type: 'string' }) |
||||
.option('logits', { type: 'string' }) |
||||
.option('samples', { type: 'string' })) |
||||
.command('transform', 'Run single step of existing model', yargs => yargs |
||||
.option('input', { type: 'string' }) |
||||
.option('output', { type: 'string' }) |
||||
.option('temperature', { type: 'number' }) |
||||
.option('logits', { type: 'string' })) |
||||
.command('run', 'Run existing model', yargs => yargs |
||||
.option('input', { type: 'string' }) |
||||
.option('output', { type: 'string' }) |
||||
.option('temperature', { type: 'number' }) |
||||
.option('samples', { type: 'string' }) |
||||
.option('logits', { type: 'string' })) |
||||
.command('benchmark', 'Benchmark existing model') |
||||
.demandCommand(1) |
||||
.parse(process.argv.slice(2)); |
||||
|
||||
if (args._.includes('test-matrix')) { |
||||
(Matrix as any).unit = 1; |
||||
|
||||
const SL = new Matrix(device, Float32Array, 2, 4, [ |
||||
[3, -7, 6, 12], |
||||
[11, -5, 7, 3], |
||||
]); |
||||
const SR = new Matrix(device, Float32Array, 2, 4, [ |
||||
[6, 12, 3, -7], |
||||
[7, 3, 11, -5], |
||||
]); |
||||
|
||||
const SO = new Matrix(device, Float32Array, 2, 4); |
||||
|
||||
const SS = new Scalar(device, Float32Array, 3); |
||||
|
||||
const ML = new Matrix(device, Float32Array, 2, 6, [ |
||||
[3, -7, 6, 12, 4, 2], |
||||
[11, -5, 7, 3, 2, 4], |
||||
]); |
||||
const MR = new Matrix(device, Float32Array, 6, 2, [ |
||||
[6, 12], |
||||
[3, -7], |
||||
[7, 3], |
||||
[11, -5], |
||||
[4, 2], |
||||
[2, 4], |
||||
]); |
||||
|
||||
const MO = new Matrix(device, Float32Array, 2, 2); |
||||
|
||||
console.log('Sum: CPU: ', sum(await SL.get(), await SR.get())); |
||||
console.log('Sum: GPU (new): ', await Matrix.sum(SL, SR, SO).get()); |
||||
console.log('Sum: GPU (in-place): ', await Matrix.sum(SL, SR, SR).get()); |
||||
|
||||
console.log('Scale: CPU: ', mul(await SL.get(), 3)); |
||||
console.log('Scale: GPU (new): ', await Matrix.scale(SS, SL, SR).get()); |
||||
console.log('Scale: GPU (in-place): ', await Matrix.scale(SS, SL, SL).get()); |
||||
|
||||
console.log('Mul: CPU: ', mul(await ML.get(), await MR.get())); |
||||
console.log('Mul: GPU: ', await Matrix.mul(ML, MR, MO).get()); |
||||
|
||||
console.log('Transpose: CPU: ', transpose(await ML.get())); |
||||
console.log('Transpose: GPU: ', await Matrix.transpose(ML, MR).get()); |
||||
} |
||||
|
||||
const dir: string = args.model; |
||||
|
||||
let tokenizer: Tokenizer | undefined; |
||||
let model: Model | undefined; |
||||
if (args._.includes('train-tokenizer')) { |
||||
tokenizer = new Tokenizer(args['vocab'] as number, fs.readFileSync(`${dir}/corpus.txt`).toString()); |
||||
|
||||
fs.writeFileSync(`${dir}/vocab.json`, JSON.stringify(tokenizer.toJSON())); |
||||
fs.writeFileSync(`${dir}/vocab-flat.txt`, tokenizer.tokens.map(s => JSON.stringify(s).slice(1, -1)).join('\n')); |
||||
} |
||||
|
||||
if (args._.includes('tokenize') || args._.includes('configure') || args._.includes('train') || args._.includes('transform') || args._.includes('run') || args._.includes('benchmark')) { |
||||
if (!tokenizer && !fs.existsSync(`${dir}/vocab.json`)) { |
||||
process.stderr.write('Can not find tokenizer data.\n'); |
||||
process.exit(1); |
||||
} |
||||
|
||||
tokenizer ??= new Tokenizer(JSON.parse(fs.readFileSync(`${dir}/vocab.json`).toString())); |
||||
|
||||
if (args._.includes('tokenize')) { |
||||
fs.writeFileSync(args['output'] as string || 1, tokenizer.detokenize(tokenizer.tokenize(fs.readFileSync(args['input'] as string || 0).toString())).map(s => JSON.stringify(s).slice(1, -1)).join('\n')); |
||||
} |
||||
|
||||
if (args._.includes('configure')) { |
||||
model = new Model(device, tokenizer, { |
||||
tokenizer: `${tokenizer.tokens.length}${tokenizer.tokens[0]}`, |
||||
vs: tokenizer.tokens.length, |
||||
ws: args['window'] as number, |
||||
ds: args['dimensions'] as number, |
||||
}); |
||||
fs.writeFileSync(`${dir}/model.json`, JSON.stringify(await model.toJSON())); |
||||
|
||||
if (args['stats']) { fs.writeFileSync(args['stats'] as string, 'index,timestamp,epoch,batch,lr,loss,accuracy,grad_inp,grad_qry,grad_key,grad_val,grad_out,queueing,computing,inferring\n'); } |
||||
if (args['logits']) { fs.writeFileSync(args['logits'] as string, 'sample,token,itid,ttid,t0,o0,p0,t1,o1,p1,t2,o2,p2,tNm2,oNm2,pNm2,tNm1,oNm1,pNm1,tN,oN,pN,sum,stddev,entropy,loss\n'); } |
||||
if (args['samples']) { fs.writeFileSync(args['samples'] as string, ''); } |
||||
} |
||||
|
||||
if (args._.includes('train') || args._.includes('transform') || args._.includes('run') || args._.includes('benchmark')) { |
||||
if (!model && !fs.existsSync(`${dir}/model.json`)) { |
||||
process.stderr.write('Can not find model data.\n'); |
||||
process.exit(1); |
||||
} |
||||
|
||||
model ??= new Model(device, tokenizer, JSON.parse(fs.readFileSync(`${dir}/model.json`).toString())); |
||||
|
||||
if (args._.includes('train')) { |
||||
const snapshot = (json: object) => fs.writeFileSync(`${dir}/model.json`, JSON.stringify(json)); |
||||
|
||||
await model.train( |
||||
args.epochs as number || 'auto', |
||||
fs.readFileSync(`${dir}/corpus.txt`).toString(), |
||||
{ |
||||
logs: process.stderr, |
||||
stats: args['stats'] ? fs.createWriteStream(args['stats'] as string, { flags: 'a' }) : undefined, |
||||
logits: args['logits'] ? fs.createWriteStream(args['logits'] as string, { flags: 'a' }) : undefined, |
||||
samples: args['samples'] ? fs.createWriteStream(args['samples'] as string, { flags: 'a' }) : undefined, |
||||
snapshot, |
||||
}); |
||||
|
||||
snapshot(await model.toJSON()); |
||||
} |
||||
|
||||
if (args._.includes('transform')) { |
||||
const [inp, out] = await model.transform( |
||||
tokenizer.tokenize(fs.readFileSync(args['input'] as string || 0).toString()), |
||||
{ |
||||
temp: args['temperature'] as number || 1, |
||||
logits: args['logits'] ? fs.createWriteStream(args['logits'] as string) : undefined, |
||||
}, |
||||
); |
||||
|
||||
fs.writeFileSync(args['output'] as string || 1, `Input:\n${tokenizer.detokenize(inp).join('')}\nOutput:${tokenizer.detokenize(out).join('\n')}\n`); |
||||
} |
||||
|
||||
if (args._.includes('run')) { |
||||
fs.writeFileSync(args['output'] as string || 1, tokenizer.detokenize(await model.run( |
||||
tokenizer.tokenize(fs.readFileSync(args['input'] as string || 0).toString()), |
||||
{ |
||||
temp: args['temperature'] as number || 1, |
||||
logits: args['logits'] ? fs.createWriteStream(args['logits'] as string) : undefined, |
||||
samples: args['samples'] as number || model.ws, |
||||
}, |
||||
)).join('') + '\n'); |
||||
} |
||||
|
||||
if (args._.includes('benchmark')) { |
||||
throw new Error('Not implemented.'); |
||||
} |
||||
} |
||||
} |
||||
|
||||
await device.queue.onSubmittedWorkDone(); |
||||
|
||||
device.destroy(); |
||||
|
||||
process.stderr.write('All done!\n'); |
||||
@ -0,0 +1,261 @@ |
||||
import { Matrix, Vector, mat, mul, softmax, sum, transpose } from './data/math'; |
||||
import { TID } from './tokenizer'; |
||||
|
||||
export interface ModelInitData< |
||||
V extends number, |
||||
W extends number, |
||||
D extends number, |
||||
H extends number, |
||||
F extends number, |
||||
> { |
||||
readonly vs: V; |
||||
readonly ws: W; |
||||
readonly ds: D; |
||||
readonly hs: H; |
||||
readonly fs: F; |
||||
} |
||||
|
||||
export interface ModelData< |
||||
V extends number, |
||||
W extends number, |
||||
D extends number, |
||||
H extends number, |
||||
HD extends number, |
||||
F extends number, |
||||
FD extends number, |
||||
> extends ModelInitData<V, W, D, H, F> { |
||||
readonly trained: true; |
||||
readonly positions: Matrix<W, D>; |
||||
readonly embeddings: Matrix<V, D>; |
||||
readonly weights: { |
||||
query: Matrix<D, HD>; |
||||
key: Matrix<D, HD>; |
||||
value: Matrix<D, HD>; |
||||
attn: Matrix<D, D>; |
||||
ffi: Matrix<D, FD>; |
||||
ffo: Matrix<FD, D>; |
||||
out: Matrix<D, V>; |
||||
}; |
||||
} |
||||
|
||||
const ES: number = 1; |
||||
|
||||
export class Model< |
||||
V extends number = number, |
||||
W extends number = number, |
||||
D extends number = number, |
||||
H extends number = number, |
||||
HD extends number = number, |
||||
F extends number = number, |
||||
FD extends number = number, |
||||
> { |
||||
protected readonly vs: V; |
||||
protected readonly ws: W; |
||||
protected readonly ds: D; |
||||
protected readonly hs: H; |
||||
protected readonly hds: HD; |
||||
protected readonly fs: F; |
||||
protected readonly fds: FD; |
||||
|
||||
protected positions: Matrix<W, D>; |
||||
protected embeddings: Matrix<V, D>; |
||||
|
||||
protected readonly weights: { |
||||
query: Matrix<D, HD>; |
||||
key: Matrix<D, HD>; |
||||
value: Matrix<D, HD>; |
||||
attn: Matrix<D, D>; |
||||
ffi: Matrix<D, FD>; |
||||
ffo: Matrix<FD, D>; |
||||
out: Matrix<D, V>; |
||||
}; |
||||
|
||||
protected embed(tids: Vector<W, TID>): Matrix<W, D> { |
||||
return sum(tids.map(tid => this.embeddings[tid]) as Matrix<W, D>, this.positions); |
||||
} |
||||
|
||||
// protected attentize<N extends number>(input: Matrix<N, D>): Matrix<N, D> {
|
||||
// const SCALE: number = 1 / Math.sqrt(this.hds);
|
||||
// const Q: Matrix<N, HD> = mul(input, this.weights.query);
|
||||
// const K: Matrix<N, HD> = mul(input, this.weights.key);
|
||||
// const V: Matrix<N, HD> = mul(input, this.weights.value);
|
||||
// const scores: Matrix<N, N> = msoftmax(mul(mul(Q, transpose(K)), SCALE));
|
||||
// const out: Matrix<N, HD> = mul(scores, V);
|
||||
// return mul(out as unknown as Matrix<N, D>, this.weights.attn);
|
||||
// }
|
||||
|
||||
// protected feedforward<N extends number>(input: Matrix<N, D>): Matrix<N, D> {
|
||||
// return input.map((v: Vector<D>) => mul(relu(mul(v, this.weights.ffi)), this.weights.ffo)) as Matrix<N, D>;
|
||||
// }
|
||||
|
||||
// protected transform<N extends number>(input: Matrix<N, D>): Matrix<N, D> {
|
||||
// const attned: Matrix<N, D> = sum(input, this.attentize(input));
|
||||
// const ffned: Matrix<N, D> = sum(input, this.feedforward(attned));
|
||||
// return ffned;
|
||||
// }
|
||||
|
||||
protected logitize(input: Matrix<W, D>): Matrix<W, V> { |
||||
return mul(input, this.weights.out).map(softmax) as Matrix<W, V>; |
||||
} |
||||
|
||||
protected argmax(logit: Vector<V>): TID { |
||||
return logit.reduce((mi, p, i, a) => p > a[mi] ? i : mi, 0); |
||||
} |
||||
|
||||
protected sample(logit: Vector<V>): TID { |
||||
let p: number = 0; |
||||
const t: number = Math.random(); |
||||
for (let i = 0; i < logit.length; i++) { |
||||
p += logit[i]; |
||||
if (t < p) { |
||||
return i; |
||||
} |
||||
} |
||||
return 0; |
||||
} |
||||
|
||||
protected forward(tids: Vector<W, TID>): Matrix<W, V> { |
||||
return this.logitize(this.embed(tids)); |
||||
} |
||||
|
||||
protected loss<N extends number>(probs: Vector<N, Vector<V>>, target: Vector<N, TID>): number { |
||||
return probs.reduce((l: number, ps: Vector<V>, i: number) => l - Math.log(ps[target[i]] + 1e-10), 0) / probs.length; |
||||
} |
||||
|
||||
protected grad(probs: Matrix<W, V>, target: Vector<W, TID>): Matrix<W, V> { |
||||
const grads = mul(probs, 1); |
||||
target.forEach((t, i) => grads[i][t] -= 1); |
||||
return grads; |
||||
} |
||||
|
||||
protected updateWOut<N extends number>(input: Matrix<N, D>, grad: Matrix<N, V>, lr: number): void { |
||||
this.weights.out = sum(mul(mul(transpose(input), grad), -lr), this.weights.out); |
||||
} |
||||
|
||||
public constructor(_device: unknown, data: ModelInitData<V, W, D, H, F> | ModelData<V, W, D, H, HD, F, FD>) { |
||||
this.vs = data.vs; |
||||
this.ws = data.ws; |
||||
this.ds = data.ds; |
||||
this.hs = data.hs; |
||||
this.hds = (this.ds / this.hs) as HD; |
||||
this.fs = data.fs; |
||||
this.fds = (this.ds * this.fs) as FD; |
||||
if (!('trained' in data)) { |
||||
this.positions = mat(this.ws, this.ds, (d, i) => { |
||||
const angle = i / Math.pow(10000, 2 * Math.floor(d / 2) / this.ds); |
||||
return d % 2 === 0 ? Math.sin(angle) * ES : Math.cos(angle) * ES; |
||||
}); |
||||
this.embeddings = mat(this.vs, this.ds, () => Math.random() * 2 * ES - ES); |
||||
this.weights = { |
||||
query: mat(this.ds, this.hds, () => Math.random() * 0.02 - 0.01), |
||||
key: mat(this.ds, this.hds, () => Math.random() * 0.02 - 0.01), |
||||
value: mat(this.ds, this.hds, () => Math.random() * 0.02 - 0.01), |
||||
attn: mat(this.ds, this.ds, () => Math.random() * 0.02 - 0.01), |
||||
ffi: mat(this.ds, this.fds, () => Math.random() * 0.02 - 0.01), |
||||
ffo: mat(this.fds, this.ds, () => Math.random() * 0.02 - 0.01), |
||||
out: mat(this.ds, this.vs, () => Math.random() * 0.2 - 0.1), |
||||
}; |
||||
} else { |
||||
this.positions = data.positions; |
||||
this.embeddings = data.embeddings; |
||||
this.weights = data.weights; |
||||
} |
||||
} |
||||
|
||||
public train(data: TID[], epochs: number): void { |
||||
let log = Date.now(); |
||||
const samples: number[] = []; |
||||
console.log('Training model.'); |
||||
|
||||
const runs = data.length - this.ws - 1; |
||||
for (let epoch = 0; epoch < epochs; epoch++) { |
||||
for (let s = 0; s < runs; s++) { |
||||
const begin: number = Date.now(); |
||||
|
||||
const source = data.slice(s, s + this.ws) as Vector<W>; |
||||
const target = data.slice(s + 1, s + 1 + this.ws) as Vector<W>; |
||||
|
||||
// const encoded = transform(model, embed(model, source));
|
||||
|
||||
const embedded = this.embed(source); |
||||
const logits = this.logitize(embedded); |
||||
|
||||
this.updateWOut(embedded, this.grad(logits, target), 0.001); |
||||
|
||||
const output = logits.map(this.sample); |
||||
|
||||
const a = 100 * target.filter((t, i) => output[i] === t).length / target.length; |
||||
const l = this.loss(logits, target); |
||||
|
||||
const end: number = Date.now(); |
||||
samples.push(end - begin); |
||||
if ((epoch * runs + s) % Math.round(runs * epochs / Math.min(100, epochs)) === 0 || epoch === epochs - 1 && s === runs - 1 || end - log > 5000) { |
||||
console.log(`Done ${samples.length} / ${epochs * runs}, accuracy ${a.toFixed(0)}%, loss ${l}, in ${end - begin}ms, ETA ${Math.round(samples.reduce((a, b) => a + b, 0) / samples.length * (epochs * runs - samples.length) / 1000)}s`); |
||||
log = end; |
||||
} |
||||
} |
||||
} |
||||
} |
||||
|
||||
public run(tids: TID[], samples: number = this.ws): TID[] { |
||||
const inputs: Vector<W> = [ |
||||
...Array(Math.max(0, this.ws - tids.length)).fill(0), |
||||
...tids.slice(-this.ws)] as Vector<W>; |
||||
const result: Vector<W, TID> = [] as unknown as Vector<W, TID>; |
||||
while (result.length < samples) { |
||||
const next: TID = this.sample(this.forward(inputs)[this.ws - 1]); |
||||
result.push(next); |
||||
inputs.shift(); |
||||
inputs.push(next); |
||||
} |
||||
return result; |
||||
} |
||||
|
||||
public benchmark(data: TID[]): void { |
||||
let log = Date.now(); |
||||
const samples: number[] = []; |
||||
console.log('Testing model.'); |
||||
|
||||
let accuracy: number = 0; |
||||
let loss: number = 0; |
||||
|
||||
const runs = data.length - this.ws - 1; |
||||
for (let s = 0; s < runs; s++) { |
||||
const begin: number = Date.now(); |
||||
|
||||
const source = data.slice(s, s + this.ws) as Vector<W>; |
||||
const target = data.slice(s + 1, s + 1 + this.ws) as Vector<W>; |
||||
|
||||
const embedded = this.embed(source); |
||||
const logits = this.logitize(embedded); |
||||
const output = logits.map(this.sample); |
||||
|
||||
const a = 100 * target.filter((t, i) => output[i] === t).length / target.length; |
||||
const l = this.loss(logits, target); |
||||
|
||||
accuracy += a / runs; |
||||
loss += l / runs; |
||||
|
||||
const end: number = Date.now(); |
||||
samples.push(end - begin); |
||||
if (end - log > 5000 || s === runs - 1) { |
||||
console.log(`Done ${samples.length} / ${runs}, accuracy ${a.toFixed(0)}% / ${accuracy.toFixed()}%, loss ${l.toFixed(4)} / ${loss.toFixed(4)}, in ${end - begin}ms, ETA ${Math.round(samples.reduce((a, b) => a + b, 0) / samples.length * (runs - samples.length) / 1000)}s`); |
||||
log = end; |
||||
} |
||||
} |
||||
} |
||||
|
||||
public toJSON(): ModelData<V, W, D, H, HD, F, FD> { |
||||
return { |
||||
vs: this.vs, |
||||
ws: this.ws, |
||||
ds: this.ds, |
||||
hs: this.hs, |
||||
fs: this.fs, |
||||
trained: true, |
||||
positions: this.positions, |
||||
embeddings: this.embeddings, |
||||
weights: this.weights, |
||||
}; |
||||
} |
||||
} |
||||
@ -0,0 +1,662 @@ |
||||
import fs from 'node:fs'; |
||||
import { Writable } from 'node:stream'; |
||||
import { Matrix } from './data/gpu/matrix'; |
||||
import { Scalar } from './data/gpu/scalar'; |
||||
import { Vector } from './data/gpu/vector'; |
||||
import { Matrix as CPUMatrix, Vector as CPUVector, drm, mat, randseq, slope, stddev, vec } from './data/math'; |
||||
import { TID, Tokenizer } from './tokenizer'; |
||||
|
||||
const SHADER_HOTONES = fs.readFileSync('shader/hotones.wgsl').toString(); |
||||
const SHADER_EMBED = fs.readFileSync('shader/embed.wgsl').toString(); |
||||
const SHADER_ATTW = fs.readFileSync('shader/attw.wgsl').toString(); |
||||
const SHADER_SOFTMAX_EXP = fs.readFileSync('shader/softmaxexp.wgsl').toString(); |
||||
const SHADER_SOFTMAX_NORM = fs.readFileSync('shader/softmaxnorm.wgsl').toString(); |
||||
const SHADER_DLDO = fs.readFileSync('shader/dldo.wgsl').toString(); |
||||
const SHADER_OTDLDO = fs.readFileSync('shader/otdldo.wgsl').toString(); |
||||
const SHADER_TOTDLDOWA = fs.readFileSync('shader/totdldowa.wgsl').toString(); |
||||
const SHADER_DLDS = fs.readFileSync('shader/dlds.wgsl').toString(); |
||||
const SHADER_GRAD_TOUT = fs.readFileSync('shader/gradtout.wgsl').toString(); |
||||
const SHADER_GRAD_TATTW = fs.readFileSync('shader/gradtattw.wgsl').toString(); |
||||
const SHADER_GRAD_INP = fs.readFileSync('shader/gradinp.wgsl').toString(); |
||||
const SHADER_ARGMAX = fs.readFileSync('shader/argmax.wgsl').toString(); |
||||
const SHADER_SAMPLES = fs.readFileSync('shader/samples.wgsl').toString(); |
||||
|
||||
export interface ModelInitData< |
||||
V extends number, |
||||
W extends number, |
||||
D extends number, |
||||
> { |
||||
readonly tokenizer: string; |
||||
readonly vs: V; |
||||
readonly ws: W; |
||||
readonly ds: D; |
||||
} |
||||
|
||||
export interface ModelData< |
||||
V extends number, |
||||
W extends number, |
||||
D extends number, |
||||
> extends ModelInitData<V, W, D> { |
||||
stage: 'early' | 'mid' | 'end'; |
||||
tokenizer: string; |
||||
samples: number; |
||||
weights: { |
||||
inp: CPUMatrix<V, D>; |
||||
pos: CPUMatrix<W, D>; |
||||
tqry: CPUMatrix<D, D>; |
||||
tkey: CPUMatrix<D, D>; |
||||
tval: CPUMatrix<D, D>; |
||||
out: CPUMatrix<V, D>; |
||||
}; |
||||
} |
||||
|
||||
interface TrainingPoint { |
||||
index: number; |
||||
timestamp: number; |
||||
epoch: number; |
||||
batch: number; |
||||
lr: number; |
||||
loss: number; |
||||
ginp: number; |
||||
gqry: number; |
||||
gkey: number; |
||||
gval: number; |
||||
gout: number; |
||||
accuracy: number; |
||||
queueing: number; |
||||
computing: number; |
||||
inferring: number; |
||||
} |
||||
|
||||
interface LogitsPoint { |
||||
sample: number; |
||||
token: number; |
||||
itid: string; |
||||
ttid: string; |
||||
t0: string; |
||||
o0: number; |
||||
p0: number; |
||||
t1: string; |
||||
o1: number; |
||||
p1: number; |
||||
t2: string; |
||||
o2: number; |
||||
p2: number; |
||||
tNm2: string; |
||||
oNm2: number; |
||||
pNm2: number; |
||||
tNm1: string; |
||||
oNm1: number; |
||||
pNm1: number; |
||||
tN: string; |
||||
oN: number; |
||||
pN: number; |
||||
sum: number; |
||||
stddev: number; |
||||
entropy: number; |
||||
loss: number; |
||||
} |
||||
|
||||
export class Model< |
||||
V extends number = number, |
||||
W extends number = number, |
||||
D extends number = number, |
||||
> { |
||||
protected readonly device: GPUDevice; |
||||
public readonly tokenizer: Tokenizer; |
||||
public readonly vs: V; |
||||
public readonly ws: W; |
||||
public readonly ds: D; |
||||
|
||||
protected samples: number = 0; |
||||
protected stage: 'early' | 'mid' | 'end' = 'early'; |
||||
|
||||
protected readonly weights: { |
||||
readonly inp: Matrix<V, D, Float32Array>; |
||||
readonly pos: Matrix<W, D, Float32Array>; |
||||
readonly tqry: Matrix<D, D, Float32Array>; |
||||
readonly tkey: Matrix<D, D, Float32Array>; |
||||
readonly tval: Matrix<D, D, Float32Array>; |
||||
readonly tout: Matrix<V, D, Float32Array>; |
||||
}; |
||||
|
||||
protected readonly processing: { |
||||
readonly ds: Scalar<Float32Array>; |
||||
readonly lr: Scalar<Float32Array>; |
||||
readonly itids: Vector<W, Int32Array, TID>; |
||||
readonly ilogits: Matrix<W, V, Float32Array>; |
||||
readonly e: Matrix<W, D, Float32Array>; |
||||
readonly q: Matrix<W, D, Float32Array>; |
||||
readonly k: Matrix<W, D, Float32Array>; |
||||
readonly tk: Matrix<D, W, Float32Array>; |
||||
readonly qtk: Matrix<W, W, Float32Array>; |
||||
readonly v: Matrix<W, D, Float32Array>; |
||||
readonly s: Matrix<W, W, Float32Array>; |
||||
readonly se: Matrix<W, W, Float32Array>; |
||||
readonly wa: Matrix<W, W, Float32Array>; |
||||
readonly a: Matrix<W, D, Float32Array>; |
||||
readonly o: Matrix<W, V, Float32Array>; |
||||
readonly temp: Scalar<Float32Array>; |
||||
readonly tempo: Matrix<W, V, Float32Array>; |
||||
readonly on: Matrix<W, V, Float32Array>; |
||||
readonly ologits: Matrix<W, V, Float32Array>; |
||||
readonly dldo: Matrix<W, V, Float32Array>; |
||||
readonly otdldo: Matrix<D, W, Float32Array>; |
||||
readonly totdldowa: Matrix<W, D, Float32Array>; |
||||
readonly dlds: Matrix<W, W, Float32Array>; |
||||
readonly dldsk: Matrix<W, D, Float32Array>; |
||||
readonly dldsq: Matrix<W, D, Float32Array>; |
||||
readonly ginp: Matrix<V, D, Float32Array>; |
||||
readonly gtqry: Matrix<D, D, Float32Array>; |
||||
readonly gtkey: Matrix<D, D, Float32Array>; |
||||
readonly gtval: Matrix<D, D, Float32Array>; |
||||
readonly gtout: Matrix<V, D, Float32Array>; |
||||
readonly rnd: Vector<W, Float32Array>; |
||||
readonly otids: Vector<W, Int32Array, TID>; |
||||
readonly ttids: Vector<W, Int32Array, TID>; |
||||
}; |
||||
|
||||
protected softmax<H extends number, W extends number>( |
||||
m: Matrix<H, W, Float32Array>, |
||||
e: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width), |
||||
r: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width), |
||||
): Matrix<H, W, Float32Array> { |
||||
e.calculate('softmaxexp', SHADER_SOFTMAX_EXP, [m], false, 4); |
||||
r.calculate('softmaxnorm', SHADER_SOFTMAX_NORM, [e], false, 4); |
||||
return r; |
||||
} |
||||
|
||||
protected async ilogits(itids: Vector<W, Int32Array, TID> = this.processing.itids): Promise<Matrix<W, V, Float32Array>> { |
||||
this.processing.ilogits.calculate('hotones', SHADER_HOTONES, [itids]); |
||||
return this.processing.ilogits; |
||||
} |
||||
|
||||
protected async e(ilogits: Matrix<W, V, Float32Array> = this.processing.ilogits): Promise<Matrix<W, D, Float32Array>> { |
||||
this.processing.e.calculate('embed', SHADER_EMBED, [ilogits, this.weights.inp, this.weights.pos]); |
||||
return this.processing.e; |
||||
} |
||||
|
||||
protected q(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> { |
||||
this.processing.q.calculate('attw', SHADER_ATTW, [e, this.weights.tqry]); |
||||
return this.processing.q; |
||||
} |
||||
|
||||
protected k(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> { |
||||
this.processing.k.calculate('attw', SHADER_ATTW, [e, this.weights.tkey]); |
||||
return this.processing.k; |
||||
} |
||||
|
||||
protected v(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> { |
||||
this.processing.v.calculate('attw', SHADER_ATTW, [e, this.weights.tval]); |
||||
return this.processing.v; |
||||
} |
||||
|
||||
protected s( |
||||
q: Matrix<W, D, Float32Array> = this.processing.q, |
||||
k: Matrix<W, D, Float32Array> = this.processing.k, |
||||
): Matrix<W, W, Float32Array> { |
||||
Matrix.transpose(k, this.processing.tk); |
||||
Matrix.mul(q, this.processing.tk, this.processing.qtk); |
||||
Matrix.scale(this.processing.ds, this.processing.qtk, this.processing.s); |
||||
return this.processing.s; |
||||
} |
||||
|
||||
protected async wa(s: Matrix<W, W, Float32Array> = this.processing.s): Promise<Matrix<W, W, Float32Array>> { |
||||
this.softmax(s, this.processing.se, this.processing.wa); |
||||
return this.processing.wa; |
||||
} |
||||
|
||||
protected a( |
||||
wa: Matrix<W, W, Float32Array> = this.processing.wa, |
||||
v: Matrix<W, D, Float32Array> = this.processing.v, |
||||
): Matrix<W, D, Float32Array> { |
||||
Matrix.mul(wa, v, this.processing.a); |
||||
return this.processing.a; |
||||
} |
||||
|
||||
protected o(a: Matrix<W, D, Float32Array> = this.processing.a): Matrix<W, V, Float32Array> { |
||||
this.processing.o.calculate('attw', SHADER_ATTW, [a, this.weights.tout]); |
||||
return this.processing.o; |
||||
} |
||||
|
||||
protected async ologits(o: Matrix<W, V, Float32Array> = this.processing.o, temp: number = 1): Promise<Matrix<W, V, Float32Array>> { |
||||
await this.processing.temp.set(temp); |
||||
Matrix.scale(this.processing.temp, o, this.processing.tempo); |
||||
this.softmax(this.processing.tempo, this.processing.on, this.processing.ologits); |
||||
return this.processing.ologits; |
||||
} |
||||
|
||||
protected async rnd(): Promise<Vector<W, Float32Array>> { |
||||
await this.processing.rnd.set(vec(this.ws, () => Math.random())); |
||||
return this.processing.rnd; |
||||
} |
||||
|
||||
protected otids(): Vector<W, Int32Array, TID>; |
||||
protected otids(ologits: Matrix<W, V, Float32Array>): Vector<W, Int32Array, TID>; |
||||
protected otids(rnd: Vector<W, Float32Array>, ologits: Matrix<W, V, Float32Array>): Vector<W, Int32Array, TID>; |
||||
protected otids(...args: unknown[]): Vector<W, Int32Array, TID> { |
||||
if (args.length === 1) { |
||||
const [ologits = this.processing.ologits] = args as [Matrix<W, V, Float32Array>]; |
||||
this.processing.otids.calculate('argmax', SHADER_ARGMAX, [ologits]); |
||||
return this.processing.otids; |
||||
} else { |
||||
const [rnd = this.processing.rnd, ologits = this.processing.ologits] = args as [Vector<W, Float32Array> | undefined, Matrix<W, V, Float32Array> | undefined]; |
||||
this.processing.otids.calculate('samples', SHADER_SAMPLES, [ologits, rnd]); |
||||
return this.processing.otids; |
||||
} |
||||
} |
||||
|
||||
protected loss<N extends number>(probs: CPUVector<N, CPUVector<V>>, target: CPUVector<N, TID>): number { |
||||
return probs.reduce((l: number, ps: CPUVector<V>, i: number) => l - Math.log(ps[target[i]] + 1e-10), 0) / probs.length; |
||||
} |
||||
|
||||
protected async upd( |
||||
ilogits: Matrix<W, V, Float32Array> = this.processing.ilogits, |
||||
e: Matrix<W, D, Float32Array> = this.processing.e, |
||||
q: Matrix<W, D, Float32Array> = this.processing.q, |
||||
k: Matrix<W, D, Float32Array> = this.processing.k, |
||||
v: Matrix<W, D, Float32Array> = this.processing.v, |
||||
wa: Matrix<W, W, Float32Array> = this.processing.wa, |
||||
a: Matrix<W, D, Float32Array> = this.processing.a, |
||||
ologits: Matrix<W, V, Float32Array> = this.processing.ologits, |
||||
ttids: Vector<W, Int32Array, TID> = this.processing.ttids, |
||||
lr: number = 0.001, |
||||
): Promise<void> { |
||||
await this.processing.lr.set(-lr); |
||||
|
||||
this.processing.dldo.calculate('dldo', SHADER_DLDO, [ |
||||
ologits, |
||||
ttids, |
||||
this.processing.lr, |
||||
]); |
||||
|
||||
this.processing.otdldo.calculate('otdldo', SHADER_OTDLDO, [this.processing.dldo, this.weights.tout]); |
||||
|
||||
this.processing.dlds.calculate('dlds', SHADER_DLDS, [ |
||||
this.processing.ds, |
||||
wa, |
||||
v, |
||||
this.processing.otdldo, |
||||
]); |
||||
|
||||
Matrix.mul(this.processing.dlds, k, this.processing.dldsk); |
||||
Matrix.mul(Matrix.transpose(this.processing.dlds), q, this.processing.dldsq); |
||||
|
||||
this.processing.ginp.calculate('gradinp', SHADER_GRAD_INP, [ |
||||
ilogits, |
||||
this.processing.dldsq, |
||||
this.processing.dldsk, |
||||
this.processing.totdldowa, |
||||
this.weights.tqry, |
||||
this.weights.tkey, |
||||
this.weights.tval, |
||||
]); |
||||
Matrix.sum(this.processing.ginp, this.weights.inp, this.weights.inp); |
||||
|
||||
this.processing.gtout.calculate('gradtout', SHADER_GRAD_TOUT, [a, this.processing.dldo]); |
||||
Matrix.sum(this.processing.gtout, this.weights.tout, this.weights.tout); |
||||
this.processing.gtqry.calculate('gradtqry', SHADER_GRAD_TATTW, [e, this.processing.dldsk]); |
||||
Matrix.sum(this.processing.gtqry, this.weights.tqry, this.weights.tqry); |
||||
this.processing.gtkey.calculate('gradtkey', SHADER_GRAD_TATTW, [e, this.processing.dldsq]); |
||||
Matrix.sum(this.processing.gtkey, this.weights.tkey, this.weights.tkey); |
||||
|
||||
this.processing.totdldowa.calculate('totdldowa', SHADER_TOTDLDOWA, [wa, this.processing.otdldo]); |
||||
this.processing.gtval.calculate('gradtval', SHADER_GRAD_TATTW, [e, this.processing.totdldowa]); |
||||
Matrix.sum(this.processing.gtval, this.weights.tval, this.weights.tval); |
||||
} |
||||
|
||||
public constructor( |
||||
device: GPUDevice, |
||||
tokenizer: Tokenizer, |
||||
data: ModelInitData<V, W, D> | ModelData<V, W, D>, |
||||
) { |
||||
this.device = device; |
||||
this.tokenizer = tokenizer; |
||||
this.vs = data.vs; |
||||
this.ws = data.ws; |
||||
this.ds = data.ds; |
||||
if (!('stage' in data)) { |
||||
this.weights = { |
||||
inp: new Matrix(this.device, Float32Array, mat(this.vs, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))), |
||||
pos: new Matrix(this.device, Float32Array, mat(this.ws, this.ds, (i, d) => { |
||||
const angle = i / Math.pow(10000, 2 * Math.floor(d / 2) / this.ds); |
||||
return (d % 2 === 0 ? Math.sin(angle) : Math.cos(angle)) * Math.sqrt(1 / this.ds); |
||||
})), |
||||
tqry: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))), |
||||
tkey: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))), |
||||
tval: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))), |
||||
tout: new Matrix(this.device, Float32Array, mat(this.vs, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))), |
||||
}; |
||||
} else { |
||||
this.weights = { |
||||
inp: new Matrix(this.device, Float32Array, data.weights.inp), |
||||
pos: new Matrix(this.device, Float32Array, data.weights.pos), |
||||
tqry: new Matrix(this.device, Float32Array, data.weights.tqry), |
||||
tkey: new Matrix(this.device, Float32Array, data.weights.tkey), |
||||
tval: new Matrix(this.device, Float32Array, data.weights.tval), |
||||
tout: new Matrix(this.device, Float32Array, data.weights.out), |
||||
}; |
||||
this.stage = data.stage; |
||||
this.samples = data.samples; |
||||
if (data.tokenizer !== `${tokenizer.tokens.length}${tokenizer.tokens[0]}`) { |
||||
throw new Error(`The model was trained against a different tokenizer ("${data.tokenizer}").`); |
||||
} |
||||
} |
||||
|
||||
this.processing = { |
||||
ds: new Scalar<Float32Array>(this.device, Float32Array, 1 / Math.sqrt(this.ds)), |
||||
lr: new Scalar<Float32Array>(this.device, Float32Array, 0), |
||||
itids: new Vector<W, Int32Array, TID>(this.device, Int32Array, this.ws), |
||||
ilogits: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs), |
||||
e: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds), |
||||
q: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds), |
||||
k: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds), |
||||
tk: new Matrix<D, W, Float32Array>(this.device, Float32Array, this.ds, this.ws), |
||||
qtk: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws), |
||||
v: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds), |
||||
s: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws), |
||||
se: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws), |
||||
wa: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws), |
||||
a: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds), |
||||
o: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs), |
||||
temp: new Scalar<Float32Array>(this.device, Float32Array, 1), |
||||
tempo: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs), |
||||
on: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs), |
||||
ologits: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs), |
||||
dldo: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs), |
||||
otdldo: new Matrix<D, W, Float32Array>(this.device, Float32Array, this.ds, this.ws), |
||||
totdldowa: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds), |
||||
dlds: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws), |
||||
dldsk: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds), |
||||
dldsq: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds), |
||||
ginp: new Matrix(this.device, Float32Array, this.vs, this.ds), |
||||
gtqry: new Matrix(this.device, Float32Array, this.ds, this.ds), |
||||
gtkey: new Matrix(this.device, Float32Array, this.ds, this.ds), |
||||
gtval: new Matrix(this.device, Float32Array, this.ds, this.ds), |
||||
gtout: new Matrix(this.device, Float32Array, this.vs, this.ds), |
||||
rnd: new Vector<W, Float32Array>(this.device, Float32Array, this.ws), |
||||
otids: new Vector<W, Int32Array, TID>(this.device, Int32Array, this.ws), |
||||
ttids: new Vector<W, Int32Array, TID>(this.device, Int32Array, this.ws), |
||||
}; |
||||
} |
||||
|
||||
public async debug( |
||||
sample: number, |
||||
itids: Vector<W, Int32Array, TID> = this.processing.itids, |
||||
ttids: Vector<W, Int32Array, TID> = this.processing.ttids, |
||||
o: Matrix<W, V, Float32Array> = this.processing.o, |
||||
ologits: Matrix<W, V, Float32Array> = this.processing.ologits, |
||||
file?: Writable, |
||||
llo: boolean = false, |
||||
): Promise<void> { |
||||
if (!file) { |
||||
return; |
||||
} |
||||
|
||||
for (let token = llo ? itids.dims - 1 : 0; token < itids.dims; token++) { |
||||
const os = (await o.get())[token]; |
||||
const ls = (await ologits.get())[token] |
||||
.map((p, tid) => [p, tid]) |
||||
.sort(([l], [r]) => r - l); |
||||
|
||||
const point: LogitsPoint = { |
||||
sample, |
||||
token, |
||||
itid: this.tokenizer.detokenize([await itids.get(token)])[0], |
||||
ttid: this.tokenizer.detokenize([await ttids.get(token)])[0], |
||||
t0: this.tokenizer.detokenize([ls[0][1]])[0], |
||||
o0: os[ls[0][1]], |
||||
p0: ls[0][0], |
||||
t1: this.tokenizer.detokenize([ls[1][1]])[0], |
||||
o1: os[ls[1][1]], |
||||
p1: ls[1][0], |
||||
t2: this.tokenizer.detokenize([ls[2][1]])[0], |
||||
o2: os[ls[2][1]], |
||||
p2: ls[2][0], |
||||
tNm2: this.tokenizer.detokenize([ls[ls.length - 3][1]])[0], |
||||
oNm2: os[ls[ls.length - 3][1]], |
||||
pNm2: ls[ls.length - 3][0], |
||||
tNm1: this.tokenizer.detokenize([ls[ls.length - 2][1]])[0], |
||||
oNm1: os[ls[ls.length - 3][1]], |
||||
pNm1: ls[ls.length - 2][0], |
||||
tN: this.tokenizer.detokenize([ls[ls.length - 1][1]])[0], |
||||
oN: os[ls[ls.length - 1][1]], |
||||
pN: ls[ls.length - 1][0], |
||||
sum: (await ologits.get())[token].reduce((a, b) => a + b, 0), |
||||
stddev: stddev((await o.get())[token]), |
||||
entropy: -(await ologits.get())[token].reduce((sum, p) => sum + (p > 0 ? p * Math.log(p) : 0), 0), |
||||
loss: -Math.log((await ologits.get())[token][await ttids.get(token)] + 1e-10), |
||||
}; |
||||
file.write(`${JSON.stringify(Object.values(point)).slice(1, -1)}\n`); |
||||
} |
||||
} |
||||
|
||||
public async train( |
||||
duration: number | 'auto', |
||||
data: string, |
||||
{ logs = process.stderr, stats, logits, samples, snapshot, statsInterval = 5000, snapshotInterval = 5000 }: { |
||||
logs?: Writable; |
||||
stats?: Writable; |
||||
logits?: Writable; |
||||
samples?: Writable; |
||||
snapshot?: (json: object) => void; |
||||
statsInterval?: number; |
||||
snapshotInterval?: number; |
||||
} = {}, |
||||
): Promise<void> { |
||||
logs.write('Training model...\n'); |
||||
|
||||
let logged = Date.now(); |
||||
let snapshotted = Date.now(); |
||||
|
||||
const aw: number = 10; |
||||
|
||||
const times = []; |
||||
const losses = []; |
||||
|
||||
let epochs = duration === 'auto' ? Infinity : +duration; |
||||
const tids = this.tokenizer.tokenize(data); |
||||
const batches = tids.length - this.ws - 1; |
||||
|
||||
let lr = 0.003; |
||||
|
||||
for (let epoch = 0; epoch < epochs; epoch++) { |
||||
const idxs = randseq(tids.length); |
||||
for (let batch = 0; batch < batches; batch++) { |
||||
const begin: number = Date.now(); |
||||
|
||||
const si = idxs[batch]; |
||||
await this.processing.itids.set(tids.slice(si, si + this.ws) as CPUVector<W>); |
||||
await this.processing.ttids.set(tids.slice(si + 1, si + 1 + this.ws) as CPUVector<W>); |
||||
|
||||
const ilogits = await this.ilogits(this.processing.itids); |
||||
const e = await this.e(ilogits); |
||||
const q = this.q(e); |
||||
const k = this.k(e); |
||||
const v = this.v(e); |
||||
const s = this.s(q, k); |
||||
const wa = await this.wa(s); |
||||
const a = this.a(wa, v); |
||||
const o = this.o(a); |
||||
const ologits = await this.ologits(o); |
||||
await this.upd(ilogits, e, q, k, v, wa, a, ologits, this.processing.ttids, lr); |
||||
|
||||
const queued: number = Date.now(); |
||||
|
||||
await this.device.queue.onSubmittedWorkDone(); |
||||
|
||||
const end: number = Date.now(); |
||||
|
||||
this.samples += 1; |
||||
times.push(end - begin); |
||||
|
||||
if ( |
||||
epoch === 0 && batch === 0 |
||||
|| batch === batches - 1 |
||||
|| end - logged > statsInterval |
||||
) { |
||||
const point: TrainingPoint = { |
||||
index: this.samples, |
||||
timestamp: Date.now(), |
||||
epoch, |
||||
batch, |
||||
lr, |
||||
loss: 0, |
||||
accuracy: 0, |
||||
ginp: (await this.processing.ginp.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr, |
||||
gqry: (await this.processing.gtqry.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr, |
||||
gkey: (await this.processing.gtkey.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr, |
||||
gval: (await this.processing.gtval.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr, |
||||
gout: (await this.processing.gtout.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr, |
||||
queueing: 0, |
||||
computing: 0, |
||||
inferring: 0, |
||||
}; |
||||
point.queueing = Date.now() - begin; |
||||
point.computing = Date.now() - queued; |
||||
|
||||
const out = await this.otids(await this.rnd(), ologits).get(); |
||||
const tgt = await this.processing.ttids.get(); |
||||
|
||||
const inferred: number = Date.now(); |
||||
|
||||
point.inferring = inferred - end; |
||||
|
||||
point.accuracy = tgt.filter((t, i) => out[i] === t).length / tgt.length; |
||||
|
||||
point.loss = this.loss(await ologits.get(), tgt); |
||||
|
||||
losses.push(point.loss); |
||||
|
||||
if (this.stage === 'early') { |
||||
lr = Math.max(0.001, Math.min(0.003, point.loss * 0.001)); |
||||
|
||||
if (losses.length > aw && (slope(losses.slice(-aw)) > -0.005 || stddev(losses.slice(-aw)) < 0.2)) { |
||||
logs?.write('Switching to mid stage.\n'); |
||||
this.stage = 'mid'; |
||||
} |
||||
} else if (this.stage === 'mid') { |
||||
lr = Math.max(0.0003, lr * 0.98); |
||||
|
||||
if (losses.length > aw && Math.abs(slope(losses.slice(-aw))) < 0.002 && stddev(losses.slice(-aw)) < 0.05) { |
||||
logs?.write('Switching to end stage.\n'); |
||||
this.stage = 'end'; |
||||
} |
||||
} else { |
||||
lr = Math.max(1e-4, lr * 0.98); |
||||
} |
||||
|
||||
logs?.write(`Done ${epoch * batches + batch + 1} / ${epochs * batches}, ETA ${Math.round(drm(times) * (epochs * batches - times.length) / 1000)}s\n`); |
||||
stats?.write(`${Object.values(point).join(',')}\n`); |
||||
await this.debug(this.samples, this.processing.itids, this.processing.ttids, o, ologits, logits); |
||||
|
||||
logged = end; |
||||
} |
||||
|
||||
if (batch === 0 && end - snapshotted > snapshotInterval) { |
||||
if (duration === 'auto' && this.stage === 'end' && losses.length > aw && ( |
||||
Math.abs(slope(losses.slice(-aw))) < 0.001 && stddev(losses.slice(-aw)) < 0.001 |
||||
|| slope(losses.slice(-aw)) > 0.001 |
||||
|| stddev(losses.slice(-aw)) > 1)) { |
||||
logs?.write(`Ending training with slope=${slope(losses.slice(-aw))} stddev=${stddev(losses.slice(-aw))}.\n`); |
||||
epochs = 0; |
||||
break; |
||||
} |
||||
|
||||
snapshot?.(await this.toJSON()); |
||||
|
||||
samples?.write(`Epoch ${epoch + 1}:\n`); |
||||
samples?.write('Sample (temp=0.3):\n'); |
||||
samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 0.3 })).join('')); |
||||
samples?.write('\nEnd of sample.\n'); |
||||
samples?.write('Sample (temp=0.7):\n'); |
||||
samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 0.7 })).join('')); |
||||
samples?.write('\nEnd of sample.\n'); |
||||
samples?.write('Sample (temp=1):\n'); |
||||
samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 1 })).join('')); |
||||
samples?.write('\nEnd of sample.\n'); |
||||
samples?.write('Sample (temp=1.3):\n'); |
||||
samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 1.3 })).join('')); |
||||
samples?.write('\nEnd of sample.\n\n'); |
||||
samples?.write('Sample (argmax):\n'); |
||||
samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 'max' })).join('')); |
||||
samples?.write('\nEnd of sample.\n\n'); |
||||
|
||||
snapshotted = end; |
||||
} |
||||
} |
||||
} |
||||
logs?.write('Training completed.\n'); |
||||
} |
||||
|
||||
public async transform( |
||||
tids: TID[], |
||||
{ temp = 1, logits, llo = false }: { |
||||
temp?: number | 'max'; |
||||
logits?: Writable; |
||||
llo?: boolean; |
||||
} = {}, |
||||
): Promise<[TID[], TID[]]> { |
||||
const itids = [ |
||||
...Array(Math.max(0, this.ws - tids.length)).fill(0), |
||||
...tids.slice(-this.ws), |
||||
]; |
||||
|
||||
await this.processing.itids.set(itids); |
||||
await this.processing.ttids.set([...itids.slice(1), 0]); |
||||
const ilogits = await this.ilogits(this.processing.itids); |
||||
const e = await this.e(ilogits); |
||||
const q = this.q(e); |
||||
const k = this.k(e); |
||||
const v = this.v(e); |
||||
const s = this.s(q, k); |
||||
const wa = await this.wa(s); |
||||
const a = this.a(wa, v); |
||||
const o = this.o(a); |
||||
const ologits = await this.ologits(o, temp === 'max' ? 1 : temp); |
||||
const out = temp === 'max' ? this.otids(ologits) : this.otids(await this.rnd(), ologits); |
||||
const otids = await out.get(); |
||||
|
||||
await this.debug(0, this.processing.itids, this.processing.ttids, o, ologits, logits, llo); |
||||
|
||||
return [itids, otids]; |
||||
} |
||||
|
||||
public async run( |
||||
tids: TID[], |
||||
{ temp = 1, samples = this.ws, logits }: { |
||||
temp?: number | 'max'; |
||||
samples?: number; |
||||
logits?: Writable; |
||||
} = {}, |
||||
): Promise<TID[]> { |
||||
const result: TID[] = []; |
||||
while (result.length < samples) { |
||||
const [_, out] = await this.transform(tids, { temp, logits, llo: true }); |
||||
const next: TID = out[this.ws - 1]; |
||||
result.push(next); |
||||
tids.shift(); |
||||
tids.push(next); |
||||
} |
||||
return result; |
||||
} |
||||
|
||||
public async toJSON(): Promise<ModelData<V, W, D>> { |
||||
return { |
||||
tokenizer: `${this.tokenizer.tokens.length}${this.tokenizer.tokens[0]}`, |
||||
vs: this.vs, |
||||
ws: this.ws, |
||||
ds: this.ds, |
||||
stage: this.stage, |
||||
samples: this.samples, |
||||
weights: { |
||||
inp: await this.weights.inp.get(), |
||||
pos: await this.weights.pos.get(), |
||||
tqry: await this.weights.tqry.get(), |
||||
tkey: await this.weights.tkey.get(), |
||||
tval: await this.weights.tval.get(), |
||||
out: await this.weights.tout.get(), |
||||
}, |
||||
}; |
||||
} |
||||
} |
||||
@ -0,0 +1,160 @@ |
||||
import { GroupQueueBase, QueueItem } from './data/group-queue'; |
||||
import { List, ListItem } from './data/list'; |
||||
import { Trie } from './data/trie'; |
||||
|
||||
interface InternalStringPairQueueItem extends QueueItem<ListItem<string>> { |
||||
prev?: StringPairQueueItem; |
||||
next?: StringPairQueueItem; |
||||
} |
||||
|
||||
interface StringPairQueueItem extends QueueItem<ListItem<string>> { |
||||
readonly prev?: StringPairQueueItem; |
||||
readonly next?: StringPairQueueItem; |
||||
} |
||||
|
||||
class StringPairQueue extends GroupQueueBase<string, ListItem<string>> { |
||||
protected items: Map<ListItem<string>, StringPairQueueItem>; |
||||
|
||||
public constructor(values?: Iterable<ListItem<string>>) { |
||||
super(); |
||||
this.items = new Map(); |
||||
for (const v of values ?? []) { |
||||
this.insert(v); |
||||
} |
||||
} |
||||
|
||||
protected override hash(v: ListItem<string>): string { |
||||
return `${v.value}$$$$${v.next.value}`; |
||||
} |
||||
|
||||
override insert(value: ListItem<string>): StringPairQueueItem { |
||||
const item: StringPairQueueItem = super.insert(value); |
||||
|
||||
const prev: StringPairQueueItem | undefined = this.items.get(value.prev); |
||||
const next: StringPairQueueItem | undefined = this.items.get(value.next); |
||||
(item as InternalStringPairQueueItem).prev = prev; |
||||
(item as InternalStringPairQueueItem).next = next; |
||||
|
||||
if (prev) { |
||||
(prev as InternalStringPairQueueItem).next = item; |
||||
} |
||||
if (next) { |
||||
(next as InternalStringPairQueueItem).prev = item; |
||||
} |
||||
|
||||
this.items.set(value, item); |
||||
|
||||
return item; |
||||
} |
||||
|
||||
override remove(item: StringPairQueueItem): boolean { |
||||
if (!super.remove(item)) { |
||||
return false; |
||||
} |
||||
if (item.prev) { |
||||
(item.prev as InternalStringPairQueueItem).next = undefined; |
||||
} |
||||
if (item.next) { |
||||
(item.next as InternalStringPairQueueItem).prev = undefined; |
||||
} |
||||
this.items.delete(item.value); |
||||
return true; |
||||
} |
||||
} |
||||
|
||||
export type TID = number; |
||||
|
||||
export class Tokenizer { |
||||
protected trie: Trie; |
||||
protected list: string[]; |
||||
|
||||
public get tokens(): Readonly<string[]> { |
||||
return this.list; |
||||
} |
||||
|
||||
protected train(size: number, text: string): void { |
||||
if (!text) { |
||||
return; |
||||
} |
||||
|
||||
console.log('Gathering pairs.'); |
||||
|
||||
let log: number = Date.now(); |
||||
const samples: number[] = []; |
||||
|
||||
const tokens = new List(text); |
||||
const pairs: StringPairQueue = new StringPairQueue(tokens.items()); |
||||
const vocab = new Set<string>(tokens); |
||||
|
||||
console.log('Merging pairs.'); |
||||
|
||||
while (vocab.size < size) { |
||||
const begin: number = Date.now(); |
||||
|
||||
const items: List<StringPairQueueItem> | undefined = pairs.peek(); |
||||
if (!items?.length) { |
||||
console.log('No merge found.'); |
||||
break; |
||||
} else { |
||||
const { value: { value: a, next: { value: b } } } = items.first!; |
||||
const merged: string = a + b; |
||||
for (const item of items) { |
||||
if (!item.orphan) { |
||||
if (item.prev) { |
||||
pairs.remove(item.prev); |
||||
} |
||||
if (item.next) { |
||||
pairs.remove(item.next); |
||||
} |
||||
|
||||
pairs.remove(item); |
||||
tokens.set(item.value, merged); |
||||
tokens.delete(item.value.next); |
||||
|
||||
pairs.insert(item.value.prev); |
||||
pairs.insert(item.value); |
||||
} |
||||
} |
||||
vocab.add(merged); |
||||
|
||||
const end: number = Date.now(); |
||||
samples.push(end - begin); |
||||
if (end - log > 5000) { |
||||
console.log(`Merged ${vocab.size} / ${size}, ETA ${Math.round(samples.reduce((a, b) => a + b, 0) / samples.length * (size - vocab.size) / 1000)}s`); |
||||
log = end; |
||||
} |
||||
} |
||||
} |
||||
this.trie = new Trie([...vocab].sort((l, r) => r.length - l.length)); |
||||
} |
||||
|
||||
public constructor(size: number, text: string); |
||||
public constructor(data: object); |
||||
public constructor(sizeOrData: number | object, text?: string) { |
||||
if (text) { |
||||
this.trie = new Trie(); |
||||
this.train(sizeOrData as number, text!); |
||||
} else { |
||||
this.trie = Trie.fromJSON(sizeOrData as object); |
||||
} |
||||
this.list = [...this.trie]; |
||||
} |
||||
|
||||
public tokenize(text: string): TID[] { |
||||
const tokens: string[] = []; |
||||
while (text.length) { |
||||
const match = this.trie.match(text); |
||||
tokens.push(match ?? text[0]); |
||||
text = text.slice(match?.length ?? 1); |
||||
} |
||||
return tokens.map(t => this.list.indexOf(t)); |
||||
} |
||||
|
||||
public detokenize(tids: TID[]): string[] { |
||||
return tids.map(t => this.list[t]); |
||||
} |
||||
|
||||
public toJSON(): object { |
||||
return this.trie.toJSON(); |
||||
} |
||||
} |
||||
@ -0,0 +1,57 @@ |
||||
import { describe, expect, it } from '@jest/globals'; |
||||
import { GroupQueue, QueueItem } from '../src/data/group-queue'; |
||||
import { List } from '../src/data/list'; |
||||
|
||||
describe('GroupQueue', () => { |
||||
function json<T>(list: List<QueueItem<T>>): T[] { |
||||
return [...list].map(({ value }) => value); |
||||
} |
||||
|
||||
describe('from array', () => { |
||||
it('[]', () => expect([...new GroupQueue([])].map(json)).toEqual([])); |
||||
it('[1]', () => expect([...new GroupQueue([1])].map(json)).toEqual([[1]])); |
||||
it('["b", "c", "b", "c"]', () => expect([...new GroupQueue(['b', 'c', 'b', 'c'])].map(json)).toEqual([['b', 'b'], ['c', 'c']])); |
||||
}); |
||||
|
||||
describe('length', () => { |
||||
it('[]', () => expect(new GroupQueue([]).length).toEqual(0)); |
||||
it('[1]', () => expect(new GroupQueue([1]).length).toEqual(1)); |
||||
it('["b", "c", "b", "c"]', () => expect(new GroupQueue(['b', 'c', 'b', 'c']).length).toEqual(2)); |
||||
}); |
||||
|
||||
describe('insert', () => { |
||||
function insert<T>(queue: GroupQueue<T>, value: T): T[][] { |
||||
queue.insert(value); |
||||
return [...queue].map(json); |
||||
} |
||||
it('0 -> []', () => expect([...insert(new GroupQueue<number>([]), 0)]).toEqual([[0]])); |
||||
it('0 -> [1]', () => expect([...insert(new GroupQueue([1]), 0)]).toEqual([[1], [0]])); |
||||
it('1 -> [1]', () => expect([...insert(new GroupQueue([1]), 1)]).toEqual([[1, 1]])); |
||||
it('"d" -> ["b", "c", "b", "c"]', () => expect([...insert(new GroupQueue(['b', 'c', 'b', 'c']), 'd')]).toEqual([['b', 'b'], ['c', 'c'], ['d']])); |
||||
it('"b" -> ["b", "c", "b", "c"]', () => expect([...insert(new GroupQueue(['b', 'c', 'b', 'c']), 'b')]).toEqual([['b', 'b', 'b'], ['c', 'c']])); |
||||
}); |
||||
|
||||
describe('itemOf / remove', () => { |
||||
function remove<T>(queue: GroupQueue<T>, value: T): T[][] { |
||||
const item = queue.itemOf(value); |
||||
if (item) { |
||||
queue.remove(item); |
||||
} |
||||
return [...queue].map(json); |
||||
} |
||||
it('[] - 0', () => expect([...remove(new GroupQueue<number>([]), 0)]).toEqual([])); |
||||
it('[1] - 0', () => expect([...remove(new GroupQueue([1]), 0)]).toEqual([[1]])); |
||||
it('[1] - 1', () => expect([...remove(new GroupQueue([1]), 1)]).toEqual([])); |
||||
it('["b", "c", "b", "c"] - "d"', () => expect([...remove(new GroupQueue(['b', 'c', 'b', 'c']), 'd')]).toEqual([['b', 'b'], ['c', 'c']])); |
||||
it('["b", "c", "b", "c"] - "b"', () => expect([...remove(new GroupQueue(['b', 'c', 'b', 'c']), 'b')]).toEqual([['c', 'c'], ['b']])); |
||||
}); |
||||
|
||||
describe('pop', () => { |
||||
function pop<T>(queue: GroupQueue<T>): [T[], T[][]] { |
||||
return [[...queue.pop() ?? []].map(({ value }) => value), [...queue].map(json)]; |
||||
} |
||||
it('[]', () => expect([...pop(new GroupQueue<number>([]))]).toEqual([[], []])); |
||||
it('[1]', () => expect([...pop(new GroupQueue([1]))]).toEqual([[1], []])); |
||||
it('["b", "c", "b", "c"]', () => expect([...pop(new GroupQueue(['b', 'c', 'b', 'c']))]).toEqual([['b', 'b'], [['c', 'c']]])); |
||||
}); |
||||
}); |
||||
@ -0,0 +1,88 @@ |
||||
import { describe, expect, it } from '@jest/globals'; |
||||
import { List, ListItem } from '../src/data/list'; |
||||
|
||||
describe('List', () => { |
||||
describe('from / to array', () => { |
||||
it('[]', () => expect([...new List([])]).toEqual([])); |
||||
it('[1]', () => expect([...new List([1])]).toEqual([1])); |
||||
it('["b", "c", "b", "c"]', () => expect([...new List(['b', 'c', 'b', 'c'])]).toEqual(['b', 'c', 'b', 'c'])); |
||||
}); |
||||
|
||||
describe('length', () => { |
||||
it('[]', () => expect(new List([]).length).toEqual(0)); |
||||
it('[1]', () => expect(new List([1]).length).toEqual(1)); |
||||
it('["b", "c", "b", "c"]', () => expect(new List(['b', 'c', 'b', 'c']).length).toEqual(4)); |
||||
}); |
||||
|
||||
describe('first', () => { |
||||
it('[]', () => expect(new List([]).first).toBeUndefined()); |
||||
it('[1]', () => expect(new List([1]).first).toEqual(1)); |
||||
it('["b", "c", "b", "c"]', () => expect(new List(['b', 'c', 'b', 'c']).first).toEqual('b')); |
||||
}); |
||||
|
||||
describe('last', () => { |
||||
it('[]', () => expect(new List([]).last).toBeUndefined()); |
||||
it('[1]', () => expect(new List([1]).last).toEqual(1)); |
||||
it('["b", "c", "b", "c"]', () => expect(new List(['b', 'c', 'b', 'c']).last).toEqual('c')); |
||||
}); |
||||
|
||||
describe('prepend', () => { |
||||
function prepend<T>(list: List<T>, value: T): List<T> { |
||||
list.prepend(value); |
||||
return list; |
||||
} |
||||
it('0 -> []', () => expect([...prepend(new List<number>([]), 0)]).toEqual([0])); |
||||
it('0 -> [1]', () => expect([...prepend(new List([1]), 0)]).toEqual([0, 1])); |
||||
it('"d" -> ["b", "c", "b", "c"]', () => expect([...prepend(new List(['b', 'c', 'b', 'c']), 'd')]).toEqual(['d', 'b', 'c', 'b', 'c'])); |
||||
}); |
||||
|
||||
describe('append', () => { |
||||
function append<T>(list: List<T>, value: T): List<T> { |
||||
list.append(value); |
||||
return list; |
||||
} |
||||
it('[] <- 0', () => expect([...append(new List<number>([]), 0)]).toEqual([0])); |
||||
it('[1] <- 0', () => expect([...append(new List([1]), 0)]).toEqual([1, 0])); |
||||
it('["b", "c", "b", "c"] <- "d"', () => expect([...append(new List(['b', 'c', 'b', 'c']), 'd')]).toEqual(['b', 'c', 'b', 'c', 'd'])); |
||||
}); |
||||
|
||||
describe('includes', () => { |
||||
it('[] not includes 0', () => expect(new List<number>([]).includes(0)).toBe(false)); |
||||
it('[1] not includes 0', () => expect(new List([1]).includes(0)).toBe(false)); |
||||
it('[1] includes 1', () => expect(new List([1]).includes(1)).toBe(true)); |
||||
it('["b", "c", "b", "c"] not includes "d"', () => expect(new List(['b', 'c', 'b', 'c']).includes('d')).toBe(false)); |
||||
it('["b", "c", "b", "c"] includes "b"', () => expect(new List(['b', 'c', 'b', 'c']).includes('b')).toBe(true)); |
||||
}); |
||||
|
||||
describe('itemOf / set', () => { |
||||
function update<T>(list: List<T>, from: T, to: T): List<T> { |
||||
const item: ListItem<T> | undefined = list.itemOf(from); |
||||
if (item) { |
||||
list.set(item, to); |
||||
} |
||||
return list; |
||||
} |
||||
|
||||
it('[]: 0 -> 1', () => expect([...update(new List<number>([]), 0, 1)]).toEqual([])); |
||||
it('[1]: 0 -> 1', () => expect([...update(new List([1]), 0, 1)]).toEqual([1])); |
||||
it('[1]: 1 -> 0', () => expect([...update(new List([1]), 1, 0)]).toEqual([0])); |
||||
it('["b", "c", "b", "c"]: "b" -> "d"', () => expect([...update(new List(['b', 'c', 'b', 'c']), 'b', 'd')]).toEqual(['d', 'c', 'b', 'c'])); |
||||
it('["b", "c", "b", "c"]: "d" -> "b"', () => expect([...update(new List(['b', 'c', 'b', 'c']), 'd', 'b')]).toEqual(['b', 'c', 'b', 'c'])); |
||||
}); |
||||
|
||||
describe('itemOf / delete', () => { |
||||
function remove<T>(list: List<T>, value: T): List<T> { |
||||
const item: ListItem<T> | undefined = list.itemOf(value); |
||||
if (item) { |
||||
list.delete(item); |
||||
} |
||||
return list; |
||||
} |
||||
|
||||
it('[] - 0', () => expect([...remove(new List<number>([]), 0)]).toEqual([])); |
||||
it('[1] - 0', () => expect([...remove(new List([1]), 0)]).toEqual([1])); |
||||
it('[1] - 1', () => expect([...remove(new List([1]), 1)]).toEqual([])); |
||||
it('["b", "c", "b", "c"] - "b"', () => expect([...remove(new List(['b', 'c', 'b', 'c']), 'b')]).toEqual(['c', 'b', 'c'])); |
||||
it('["b", "c", "b", "c"] - "d"', () => expect([...remove(new List(['b', 'c', 'b', 'c']), 'd')]).toEqual(['b', 'c', 'b', 'c'])); |
||||
}); |
||||
}); |
||||
@ -0,0 +1,7 @@ |
||||
{ |
||||
"extends": "../tsconfig.json", |
||||
"include": [ |
||||
"../src/**/*.ts", |
||||
"**/*.ts", |
||||
], |
||||
} |
||||
@ -0,0 +1,35 @@ |
||||
#!/bin/bash |
||||
|
||||
model=$1 |
||||
count=${2:-9000} |
||||
|
||||
mkdir -p ${model}/downloads |
||||
rm -rf ${model}/downloads/* |
||||
|
||||
BOOK_IDS=( |
||||
174 1342 11 2701 84 2600 1661 215 1952 98 345 35 1497 2147 4300 |
||||
) |
||||
|
||||
echo "Downloading books..." |
||||
|
||||
i=1 |
||||
for ID in "${BOOK_IDS[@]}"; do |
||||
echo "Fetching $ID..." |
||||
curl -s -L "https://www.gutenberg.org/files/$ID/$ID-0.txt" -o "${model}/downloads/$ID.txt" |
||||
i=$((i+1)) |
||||
if [ ${i} -gt ${count} ]; then |
||||
break; |
||||
fi |
||||
done |
||||
|
||||
echo "Concatenating into corpus.txt..." |
||||
|
||||
# Remove headers/footers and join them |
||||
rm -f ${model}/corpus.txt |
||||
for FILE in ${model}/downloads/*.txt; do |
||||
awk '/\*\*\* START OF/,/\*\*\* END OF/ { if (!/START|END/) print }' "$FILE" >> ${model}/corpus.txt |
||||
echo >> ${model}/corpus.txt |
||||
done |
||||
|
||||
echo "All done. Corpus size:" |
||||
du -h ${model}/corpus.txt |
||||
@ -0,0 +1,178 @@ |
||||
import sys |
||||
import pandas as pd |
||||
import matplotlib.pyplot as plt |
||||
import matplotlib.animation as animation |
||||
import matplotlib.dates as mdates # Add this import at the top |
||||
import select |
||||
from datetime import datetime |
||||
import threading |
||||
import queue |
||||
|
||||
fig = plt.figure(figsize=(10, 6)) |
||||
|
||||
ax_lr = fig.add_axes([0.05, 0.3, 0.9, 0.65]) |
||||
ax_lr.set_facecolor("none") |
||||
ax_lr.set_ylim(bottom=0, top=None) |
||||
ax_lr.yaxis.set_ticks_position('right') |
||||
ax_lr.tick_params(axis='y', labelcolor='tab:brown') |
||||
|
||||
ax_acc = fig.add_axes([0.05, 0.3, 0.9, 0.65], sharex=ax_lr) |
||||
ax_acc.set_facecolor("none") |
||||
ax_acc.set_ylim(0, 100) |
||||
ax_acc.yaxis.set_ticks_position('right') |
||||
ax_acc.tick_params(axis='y', labelcolor='tab:green') |
||||
|
||||
ax_loss = fig.add_axes([0.05, 0.3, 0.9, 0.65], sharex=ax_lr) |
||||
ax_loss.set_facecolor("none") |
||||
ax_loss.set_ylim(bottom=0, top=None) |
||||
ax_loss.tick_params(axis='y', labelcolor='tab:blue') |
||||
|
||||
ax_time = fig.add_axes([0.05, 0.05, 0.9, 0.2]) |
||||
ax_time.set_facecolor("none") |
||||
ax_time.set_ylim(bottom=0, top=None) |
||||
ax_time.yaxis.set_ticks_position('right') |
||||
ax_time.tick_params(axis='y', labelcolor='tab:orange') |
||||
|
||||
columns = ["index", "timestamp", "epoch", "batch", "lr", "loss", "accuracy", "grad_inp", "grad_qry", "grad_key", "grad_val", "grad_out", "queue_ms", "compute_ms", "infer_ms"] |
||||
df = pd.DataFrame(columns=columns) |
||||
ema_alpha = 0.1 |
||||
window_size = 50 |
||||
|
||||
def parse_line(line): |
||||
try: |
||||
parts = line.strip().split(",") |
||||
if parts[0] == "index" or len(parts) < len(columns): |
||||
return None |
||||
row = { |
||||
"index": int(parts[0]), |
||||
"timestamp": datetime.fromtimestamp(int(parts[1]) / 1000), |
||||
"epoch": int(parts[2]), |
||||
"batch": int(parts[3]), |
||||
"lr": float(parts[4]), |
||||
"loss": float(parts[5]), |
||||
"accuracy": float(parts[6]) * 100, |
||||
"grad_inp": float(parts[7]), |
||||
"grad_qry": float(parts[8]), |
||||
"grad_key": float(parts[9]), |
||||
"grad_val": float(parts[10]), |
||||
"grad_out": float(parts[11]), |
||||
"queue_ms": float(parts[12]), |
||||
"compute_ms": float(parts[13]), |
||||
"infer_ms": float(parts[14]), |
||||
} |
||||
return row |
||||
except Exception as e: |
||||
print(f"Bad line: {line.strip()} ({e})", file=sys.stderr) |
||||
return None |
||||
|
||||
line_queue = queue.Queue() |
||||
|
||||
def stdin_reader(): |
||||
for line in sys.stdin: |
||||
line_queue.put(line) |
||||
|
||||
threading.Thread(target=stdin_reader, daemon=True).start() |
||||
|
||||
def get_new_lines(): |
||||
lines = [] |
||||
while not line_queue.empty(): |
||||
lines.append(line_queue.get()) |
||||
return lines |
||||
|
||||
last_seen_index = None |
||||
|
||||
def animate(_): |
||||
global df |
||||
|
||||
lines = get_new_lines() |
||||
if not lines: |
||||
return |
||||
|
||||
new_rows = [parse_line(line) for line in lines] |
||||
new_rows = [r for r in new_rows if r is not None] |
||||
|
||||
if not new_rows: |
||||
return |
||||
|
||||
global last_seen_index |
||||
|
||||
new_index = new_rows[-1]["index"] |
||||
if last_seen_index and (new_index < last_seen_index): |
||||
print("Index reset, restarting plot") |
||||
df = pd.DataFrame(columns=columns) |
||||
|
||||
last_seen_index = new_index |
||||
|
||||
df_new = pd.DataFrame(new_rows) |
||||
df_new["total_ms"] = df_new["queue_ms"] + df_new["compute_ms"] + df_new["infer_ms"] |
||||
df = pd.concat([df, df_new], ignore_index=True) |
||||
|
||||
df["ema_loss"] = df["loss"].ewm(alpha=ema_alpha).mean().combine_first(df["loss"]) |
||||
df["avg_loss"] = df["loss"].rolling(window=window_size).mean().combine_first(df["loss"]) |
||||
df["avg_accuracy"] = df["accuracy"].rolling(window=window_size).mean().combine_first(df["accuracy"]) |
||||
df["avg_total_ms"] = df["total_ms"].rolling(window=window_size).mean().combine_first(df["total_ms"]) |
||||
|
||||
ax_lr.clear() |
||||
ax_lr.set_ylabel("LR") |
||||
ax_lr.yaxis.set_label_position("right") |
||||
|
||||
ax_acc.clear() |
||||
ax_acc.set_ylabel("Accuracy (%)") |
||||
|
||||
ax_loss.clear() |
||||
ax_loss.set_title("Live Training Metrics") |
||||
ax_loss.set_xlabel("Epoch / Batch") |
||||
ax_loss.set_ylabel("Loss") |
||||
ax_acc.yaxis.set_label_position("right") |
||||
|
||||
ax_time.clear() |
||||
ax_time.set_xlabel("Time") |
||||
ax_time.set_ylabel("Processing Time (ms)") |
||||
ax_time.yaxis.set_label_position("right") |
||||
|
||||
ax_lr.plot(df["index"], df["lr"], label="LR", linestyle='--', alpha=0.3, color="tab:brown") |
||||
|
||||
ax_acc.plot(df["index"], df["accuracy"], label="Accuracy", linestyle='--', alpha=0.3, color="tab:green") |
||||
ax_acc.plot(df["index"], df["avg_accuracy"], label="Avg Accuracy", color="tab:green") |
||||
|
||||
ax_loss.plot(df["index"], df["loss"], label="Loss", linestyle=':', alpha=0.3, color="tab:blue") |
||||
ax_loss.plot(df["index"], df["ema_loss"], label="EMA Loss", linestyle='--', alpha=0.7, color="tab:blue") |
||||
ax_loss.plot(df["index"], df["avg_loss"], label=f"Avg Loss", color="tab:blue") |
||||
|
||||
ax_time.plot(df["index"], df["queue_ms"], label="Queueing", linestyle=':', alpha=0.3, color="#ffa64d") |
||||
ax_time.plot(df["index"], df["compute_ms"], label="Computing", linestyle=':', alpha=0.3, color="#cc6600") |
||||
ax_time.plot(df["index"], df["infer_ms"], label="Inferring", linestyle=':', alpha=0.3, color="#ffb84d") |
||||
ax_time.plot(df["index"], df["total_ms"], label="Total Time",linestyle='--', alpha=0.7, color="tab:orange") |
||||
ax_time.plot(df["index"], df["avg_total_ms"], label="Avg Total Time", color="tab:orange") |
||||
|
||||
# Add epoch/batch label ticks to secondary X axis |
||||
tick_locs = ax_time.get_xticks() |
||||
tick_labels = [] |
||||
|
||||
if not df.empty: |
||||
for loc in tick_locs: |
||||
idx = (abs(df["index"] - loc)).argmin() |
||||
e, b = int(df.iloc[idx]["epoch"]), int(df.iloc[idx]["batch"]) |
||||
tick_labels.append(f"{e + 1}/{b + 1}") |
||||
else: |
||||
tick_labels = [""] * len(tick_locs) |
||||
|
||||
ax_loss.set_xlim(ax_time.get_xlim()) |
||||
ax_loss.set_xticks(tick_locs) |
||||
ax_loss.set_xticklabels(tick_labels) |
||||
|
||||
# Gather all legend entries |
||||
lines = [] |
||||
labels = [] |
||||
|
||||
for ax in [ax_lr, ax_acc, ax_loss]: |
||||
l, lab = ax.get_legend_handles_labels() |
||||
lines += l |
||||
labels += lab |
||||
|
||||
# Add combined legend to ax_loss |
||||
ax_loss.legend(lines, labels, loc='upper left') # or 'upper right', etc. |
||||
ax_time.legend(loc='upper left') |
||||
|
||||
ani = animation.FuncAnimation(fig, animate, interval=1000) |
||||
plt.show() |
||||
@ -0,0 +1,24 @@ |
||||
{ |
||||
"compilerOptions": { |
||||
"baseUrl": ".", |
||||
"outDir": "build", |
||||
"target": "ES2022", |
||||
"module": "ES2022", |
||||
"moduleResolution": "Node", |
||||
"esModuleInterop": true, |
||||
"lib": [ |
||||
"DOM", "ES2022" |
||||
], |
||||
"noImplicitAny": true, |
||||
"noImplicitOverride": true, |
||||
"noImplicitReturns": true, |
||||
"noImplicitThis": true, |
||||
"strictNullChecks": true, |
||||
"strictPropertyInitialization": true, |
||||
"declaration": false, |
||||
"sourceMap": true, |
||||
}, |
||||
"include": [ |
||||
"src/**/*.ts", |
||||
], |
||||
} |
||||
Loading…
Reference in new issue