| import Base |
|
|
| type Candidates is Data: |
| Leaf{index: U32, score: F32} |
| Fork{left: Candidates, right: Candidates} |
|
|
| def pick(flag: Bool, a: Candidates, b: Candidates) -> Candidates: |
| match flag: |
| case True{}: |
| a |
| case False{}: |
| b |
|
|
| def better(a: Candidates, b: Candidates) -> Candidates: |
| match a b: |
| case Leaf{+i, +x} Leaf{+j, +y}: |
| pick(F32.is_gt(x, y) || (F32.is_eq(x, y) && U32.is_le(i, j)), Leaf{i, x}, Leaf{j, y}) |
| case _ _: |
| a |
|
|
| def winner(tree: Candidates) -> Candidates: |
| match tree: |
| case Leaf{i, x}: |
| Leaf{i, x} |
| case Fork{left, right}: |
| a b = winner(left) winner(right) |
| better(a, b) |
|
|
| def input() -> IO(Candidates): |
| import "./bridge.c" |
|
|
| def output(tree: Candidates) -> IO(Unit): |
| import "./bridge.c" |
|
|
| # Transformer math uses the same affine trees. Forks expose independent work. |
| def exponentials(tree: Candidates, +maximum: F32) -> Candidates: |
| match tree: |
| case Leaf{i, x}: |
| Leaf{i, F32.exp(F32.sub(x, maximum))} |
| case Fork{left, right}: |
| a b = exponentials(left, maximum) exponentials(right, maximum) |
| Fork{a, b} |
|
|
| def total(tree: Candidates) -> F32: |
| match tree: |
| case Leaf{i, x}: |
| x |
| case Fork{left, right}: |
| a b = total(left) total(right) |
| F32.add(a, b) |
|
|
| def divide(tree: Candidates, +sum: F32) -> Candidates: |
| match tree: |
| case Leaf{i, x}: |
| Leaf{i, F32.div(x, sum)} |
| case Fork{left, right}: |
| a b = divide(left, sum) divide(right, sum) |
| Fork{a, b} |
|
|
| def softmax_best(tree: Candidates, best: Candidates) -> Candidates: |
| match best: |
| case Leaf{i, maximum}: |
| +values = exponentials(tree, maximum) |
| divide(values, total(values)) |
| case Fork{left, right}: |
| tree |
|
|
| def softmax(+tree: Candidates) -> Candidates: |
| softmax_best(tree, winner(tree)) |
|
|
| type Features is Data: |
| Feature{x: F32, gamma: F32, beta: F32} |
| Branch{left: Features, right: Features} |
|
|
| def feature_sum(tree: Features) -> F32: |
| match tree: |
| case Feature{x, gamma, beta}: |
| x |
| case Branch{left, right}: |
| a b = feature_sum(left) feature_sum(right) |
| F32.add(a, b) |
|
|
| def variance(tree: Features, +mean: F32) -> F32: |
| match tree: |
| case Feature{x, gamma, beta}: |
| F32.square(F32.sub(x, mean)) |
| case Branch{left, right}: |
| a b = variance(left, mean) variance(right, mean) |
| F32.add(a, b) |
|
|
| def affine(tree: Features, +mean: F32, +scale: F32) -> Features: |
| match tree: |
| case Feature{x, +gamma, +beta}: |
| Feature{F32.add(F32.mul(F32.mul(F32.sub(x, mean), scale), gamma), beta), gamma, beta} |
| case Branch{left, right}: |
| a b = affine(left, mean, scale) affine(right, mean, scale) |
| Branch{a, b} |
|
|
| def layernorm(+tree: Features, +count: F32, epsilon: F32) -> Features: |
| +mean = F32.div(feature_sum(tree), count) |
| scale = F32.div(1.0, F32.sqrt(F32.add(F32.div(variance(tree, mean), count), epsilon))) |
| affine(tree, mean, scale) |
|
|
| def features() -> IO(Features): |
| import "./bridge.c" |
|
|
| def write_features(tree: Features) -> IO(Unit): |
| import "./bridge.c" |
|
|
| # Coarse parallelism: fork over independent rows, flat tail loops within a row. |
| # Scalar-level tasks cost more to schedule than the arithmetic they perform. |
| type Vector is Data: |
| VNil{} |
| VCons{x: F32, gamma: F32, beta: F32, next: Vector} |
|
|
| type NormRows is Data: |
| NLeaf{values: Vector} |
| NFork{left: NormRows, right: NormRows} |
|
|
| def vector_sum(xs: Vector, acc: F32) -> F32: |
| match xs: |
| case VNil{}: |
| acc |
| case VCons{x, g, b, next}: |
| vector_sum(next, F32.add(acc, x)) |
|
|
| def vector_variance(xs: Vector, +mean: F32, acc: F32) -> F32: |
| match xs: |
| case VNil{}: |
| acc |
| case VCons{x, g, b, next}: |
| vector_variance(next, mean, F32.add(acc, F32.square(F32.sub(x, mean)))) |
|
|
| def vector_affine(xs: Vector, +mean: F32, +scale: F32, out: Vector) -> Vector: |
| match xs: |
| case VNil{}: |
| out |
| case VCons{x, +g, +b, next}: |
| value = F32.add(F32.mul(F32.mul(F32.sub(x, mean), scale), g), b) |
| vector_affine(next, mean, scale, VCons{value, g, b, out}) |
|
|
| def norm_vector(+xs: Vector, +count: F32, epsilon: F32) -> Vector: |
| +mean = F32.div(vector_sum(xs, 0.0), count) |
| scale = F32.div(1.0, F32.sqrt(F32.add(F32.div(vector_variance(xs, mean, 0.0), count), epsilon))) |
| vector_affine(xs, mean, scale, VNil{}) |
|
|
| def norm_rows(xs: NormRows, +count: F32, +epsilon: F32) -> NormRows: |
| match xs: |
| case NLeaf{values}: |
| NLeaf{norm_vector(values, count, epsilon)} |
| case NFork{left, right}: |
| a b = norm_rows(left, count, epsilon) norm_rows(right, count, epsilon) |
| NFork{a, b} |
|
|
| def batch_features() -> IO(NormRows): |
| import "./bridge.c" |
|
|
| def write_batch(tree: NormRows) -> IO(Unit): |
| import "./bridge.c" |
|
|
| # Resident dense projections. Each leaf computes four output channels in one |
| # flat loop; only independent output tiles fork. Inputs and weights are borrowed. |
| type DenseWeights is Data: |
| JWEnd{} |
| JWVal{w0: F32, w1: F32, w2: F32, w3: F32, w4: F32, w5: F32, w6: F32, next: DenseWeights} |
|
|
| type JuliaMatrix is Data: |
| JMLeaf{weights: DenseWeights} |
| JMFork{left: JuliaMatrix, right: JuliaMatrix} |
|
|
| def dense_weights() -> IO(JuliaMatrix): |
| import "./bridge.c" |
|
|
| # A 4x7 tile fits seven weights plus a link in one eight-word heap allocation. |
| # Each weight is reused across four tokens; no half-empty 16-word allocations. |
| type TileInput is Data: |
| JTEnd{} |
| JTVal{a: F32, b: F32, c: F32, d: F32, next: TileInput} |
|
|
| type TileResult is Data: |
| JTLeaf{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32, v8: F32, v9: F32, v10: F32, v11: F32, v12: F32, v13: F32, v14: F32, v15: F32, v16: F32, v17: F32, v18: F32, v19: F32, v20: F32, v21: F32, v22: F32, v23: F32, v24: F32, v25: F32, v26: F32, v27: F32} |
| JTFork{left: TileResult, right: TileResult} |
|
|
| def dot_tile(xs: TileInput, ws: DenseWeights, s0: F32, s1: F32, s2: F32, s3: F32, s4: F32, s5: F32, s6: F32, s7: F32, s8: F32, s9: F32, s10: F32, s11: F32, s12: F32, s13: F32, s14: F32, s15: F32, s16: F32, s17: F32, s18: F32, s19: F32, s20: F32, s21: F32, s22: F32, s23: F32, s24: F32, s25: F32, s26: F32, s27: F32) -> TileResult: |
| match xs ws: |
| case JTVal{+x0, +x1, +x2, +x3, xt} JWVal{+w0, +w1, +w2, +w3, +w4, +w5, +w6, wt}: |
| dot_tile(xt, wt, F32.add(s0, F32.mul(x0, w0)), |
| F32.add(s1, F32.mul(x0, w1)), |
| F32.add(s2, F32.mul(x0, w2)), |
| F32.add(s3, F32.mul(x0, w3)), |
| F32.add(s4, F32.mul(x0, w4)), |
| F32.add(s5, F32.mul(x0, w5)), |
| F32.add(s6, F32.mul(x0, w6)), |
| F32.add(s7, F32.mul(x1, w0)), |
| F32.add(s8, F32.mul(x1, w1)), |
| F32.add(s9, F32.mul(x1, w2)), |
| F32.add(s10, F32.mul(x1, w3)), |
| F32.add(s11, F32.mul(x1, w4)), |
| F32.add(s12, F32.mul(x1, w5)), |
| F32.add(s13, F32.mul(x1, w6)), |
| F32.add(s14, F32.mul(x2, w0)), |
| F32.add(s15, F32.mul(x2, w1)), |
| F32.add(s16, F32.mul(x2, w2)), |
| F32.add(s17, F32.mul(x2, w3)), |
| F32.add(s18, F32.mul(x2, w4)), |
| F32.add(s19, F32.mul(x2, w5)), |
| F32.add(s20, F32.mul(x2, w6)), |
| F32.add(s21, F32.mul(x3, w0)), |
| F32.add(s22, F32.mul(x3, w1)), |
| F32.add(s23, F32.mul(x3, w2)), |
| F32.add(s24, F32.mul(x3, w3)), |
| F32.add(s25, F32.mul(x3, w4)), |
| F32.add(s26, F32.mul(x3, w5)), |
| F32.add(s27, F32.mul(x3, w6))) |
| case _ _: |
| JTLeaf{s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, s16, s17, s18, s19, s20, s21, s22, s23, s24, s25, s26, s27} |
|
|
| def dense_tile(+xs: TileInput, weights: JuliaMatrix) -> TileResult: |
| match weights: |
| case JMLeaf{ws}: |
| dot_tile(xs, ws, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) |
| case JMFork{left, right}: |
| a b = dense_tile(xs, left) dense_tile(xs, right) |
| JTFork{a, b} |
|
|
| def tile_input() -> IO(TileInput): |
| import "./bridge.c" |
|
|
| def tile_output(result: TileResult) -> IO(Unit): |
| import "./bridge.c" |
|
|
| # One fork tree spans token tiles and output tiles; one runtime dispatch per matrix. |
| type InputTiles is Data: |
| JILeaf{tile: TileInput} |
| JIFork{left: InputTiles, right: InputTiles} |
|
|
| type OutputTiles is Data: |
| JOLeaf{tile: TileResult} |
| JOFork{left: OutputTiles, right: OutputTiles} |
|
|
| def dense_batch(tiles: InputTiles, +weights: JuliaMatrix) -> OutputTiles: |
| match tiles: |
| case JILeaf{tile}: |
| JOLeaf{dense_tile(tile, weights)} |
| case JIFork{left, right}: |
| a b = dense_batch(left, weights) dense_batch(right, weights) |
| JOFork{a, b} |
|
|
| def input_tiles() -> IO(InputTiles): |
| import "./bridge.c" |
|
|
| def output_tiles(result: OutputTiles) -> IO(Unit): |
| import "./bridge.c" |
|
|
| # Packed tiles expose contiguous lanes to the native compiler. All arithmetic |
| # stays in Bend. Shared inputs are read-only; output tiles are disjoint. |
| # Unsafe loops terminate at the positive, validated dimensions in the bridge. |
| type Input8 is Data: |
| JInput8{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32} |
|
|
| type Weight8 is Data: |
| JWeight8{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32} |
|
|
| type Output64 is Data: |
| JOutput64{v0: F32, v1: F32, v2: F32, v3: F32, v4: F32, v5: F32, v6: F32, v7: F32, v8: F32, v9: F32, v10: F32, v11: F32, v12: F32, v13: F32, v14: F32, v15: F32, v16: F32, v17: F32, v18: F32, v19: F32, v20: F32, v21: F32, v22: F32, v23: F32, v24: F32, v25: F32, v26: F32, v27: F32, v28: F32, v29: F32, v30: F32, v31: F32, v32: F32, v33: F32, v34: F32, v35: F32, v36: F32, v37: F32, v38: F32, v39: F32, v40: F32, v41: F32, v42: F32, v43: F32, v44: F32, v45: F32, v46: F32, v47: F32, v48: F32, v49: F32, v50: F32, v51: F32, v52: F32, v53: F32, v54: F32, v55: F32, v56: F32, v57: F32, v58: F32, v59: F32, v60: F32, v61: F32, v62: F32, v63: F32} |
|
|
| type PackedBuffers is Type: |
| JBuffers{x: Array<Input8>, w: Array<Weight8>, o: Array<Output64>} |
|
|
| @unsafe |
| def packed_dot(x: Array<Input8>, w: Array<Weight8>, o: Array<Output64>, +row: U32, +col: U32, +cols: U32, +inner: U32, +end: U32, +j: U32, +s0: F32, +s1: F32, +s2: F32, +s3: F32, +s4: F32, +s5: F32, +s6: F32, +s7: F32, +s8: F32, +s9: F32, +s10: F32, +s11: F32, +s12: F32, +s13: F32, +s14: F32, +s15: F32, +s16: F32, +s17: F32, +s18: F32, +s19: F32, +s20: F32, +s21: F32, +s22: F32, +s23: F32, +s24: F32, +s25: F32, +s26: F32, +s27: F32, +s28: F32, +s29: F32, +s30: F32, +s31: F32, +s32: F32, +s33: F32, +s34: F32, +s35: F32, +s36: F32, +s37: F32, +s38: F32, +s39: F32, +s40: F32, +s41: F32, +s42: F32, +s43: F32, +s44: F32, +s45: F32, +s46: F32, +s47: F32, +s48: F32, +s49: F32, +s50: F32, +s51: F32, +s52: F32, +s53: F32, +s54: F32, +s55: F32, +s56: F32, +s57: F32, +s58: F32, +s59: F32, +s60: F32, +s61: F32, +s62: F32, +s63: F32, active: Bool) -> PackedBuffers: |
| match active: |
| case True{}: |
| packed_read_x(Array.get(Input8, x, U32.add(U32.mul(row, inner), j)), w, o, row, col, cols, inner, end, j, s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, s16, s17, s18, s19, s20, s21, s22, s23, s24, s25, s26, s27, s28, s29, s30, s31, s32, s33, s34, s35, s36, s37, s38, s39, s40, s41, s42, s43, s44, s45, s46, s47, s48, s49, s50, s51, s52, s53, s54, s55, s56, s57, s58, s59, s60, s61, s62, s63) |
| case False{}: |
| o = Array.set(Output64, o, U32.add(U32.mul(row, cols), col), JOutput64{s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, s16, s17, s18, s19, s20, s21, s22, s23, s24, s25, s26, s27, s28, s29, s30, s31, s32, s33, s34, s35, s36, s37, s38, s39, s40, s41, s42, s43, s44, s45, s46, s47, s48, s49, s50, s51, s52, s53, s54, s55, s56, s57, s58, s59, s60, s61, s62, s63}) |
| packed_step(x, w, o, U32.add(U32.add(U32.mul(row, cols), col), 1), end, cols, inner, U32.is_lt(U32.add(U32.add(U32.mul(row, cols), col), 1), end)) |
|
|
| @unsafe |
| def packed_read_x(pair: Array<Input8> & Input8, w: Array<Weight8>, o: Array<Output64>, +row: U32, +col: U32, +cols: U32, +inner: U32, +end: U32, +j: U32, +s0: F32, +s1: F32, +s2: F32, +s3: F32, +s4: F32, +s5: F32, +s6: F32, +s7: F32, +s8: F32, +s9: F32, +s10: F32, +s11: F32, +s12: F32, +s13: F32, +s14: F32, +s15: F32, +s16: F32, +s17: F32, +s18: F32, +s19: F32, +s20: F32, +s21: F32, +s22: F32, +s23: F32, +s24: F32, +s25: F32, +s26: F32, +s27: F32, +s28: F32, +s29: F32, +s30: F32, +s31: F32, +s32: F32, +s33: F32, +s34: F32, +s35: F32, +s36: F32, +s37: F32, +s38: F32, +s39: F32, +s40: F32, +s41: F32, +s42: F32, +s43: F32, +s44: F32, +s45: F32, +s46: F32, +s47: F32, +s48: F32, +s49: F32, +s50: F32, +s51: F32, +s52: F32, +s53: F32, +s54: F32, +s55: F32, +s56: F32, +s57: F32, +s58: F32, +s59: F32, +s60: F32, +s61: F32, +s62: F32, +s63: F32) -> PackedBuffers: |
| (x, input) = pair |
| JInput8{+x0, +x1, +x2, +x3, +x4, +x5, +x6, +x7} = input |
| packed_read_w(Array.get(Weight8, w, U32.add(U32.mul(col, inner), j)), x, o, row, col, cols, inner, end, j, s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, s16, s17, s18, s19, s20, s21, s22, s23, s24, s25, s26, s27, s28, s29, s30, s31, s32, s33, s34, s35, s36, s37, s38, s39, s40, s41, s42, s43, s44, s45, s46, s47, s48, s49, s50, s51, s52, s53, s54, s55, s56, s57, s58, s59, s60, s61, s62, s63, x0, x1, x2, x3, x4, x5, x6, x7) |
|
|
| @unsafe |
| def packed_read_w(pair: Array<Weight8> & Weight8, x: Array<Input8>, o: Array<Output64>, +row: U32, +col: U32, +cols: U32, +inner: U32, +end: U32, +j: U32, +s0: F32, +s1: F32, +s2: F32, +s3: F32, +s4: F32, +s5: F32, +s6: F32, +s7: F32, +s8: F32, +s9: F32, +s10: F32, +s11: F32, +s12: F32, +s13: F32, +s14: F32, +s15: F32, +s16: F32, +s17: F32, +s18: F32, +s19: F32, +s20: F32, +s21: F32, +s22: F32, +s23: F32, +s24: F32, +s25: F32, +s26: F32, +s27: F32, +s28: F32, +s29: F32, +s30: F32, +s31: F32, +s32: F32, +s33: F32, +s34: F32, +s35: F32, +s36: F32, +s37: F32, +s38: F32, +s39: F32, +s40: F32, +s41: F32, +s42: F32, +s43: F32, +s44: F32, +s45: F32, +s46: F32, +s47: F32, +s48: F32, +s49: F32, +s50: F32, +s51: F32, +s52: F32, +s53: F32, +s54: F32, +s55: F32, +s56: F32, +s57: F32, +s58: F32, +s59: F32, +s60: F32, +s61: F32, +s62: F32, +s63: F32, +x0: F32, +x1: F32, +x2: F32, +x3: F32, +x4: F32, +x5: F32, +x6: F32, +x7: F32) -> PackedBuffers: |
| (w, weight) = pair |
| JWeight8{+w0, +w1, +w2, +w3, +w4, +w5, +w6, +w7} = weight |
| packed_dot(x, w, o, row, col, cols, inner, end, U32.add(j, 1), F32.add(s0, F32.mul(x0, w0)), F32.add(s1, F32.mul(x0, w1)), F32.add(s2, F32.mul(x0, w2)), F32.add(s3, F32.mul(x0, w3)), F32.add(s4, F32.mul(x0, w4)), F32.add(s5, F32.mul(x0, w5)), F32.add(s6, F32.mul(x0, w6)), F32.add(s7, F32.mul(x0, w7)), F32.add(s8, F32.mul(x1, w0)), F32.add(s9, F32.mul(x1, w1)), F32.add(s10, F32.mul(x1, w2)), F32.add(s11, F32.mul(x1, w3)), F32.add(s12, F32.mul(x1, w4)), F32.add(s13, F32.mul(x1, w5)), F32.add(s14, F32.mul(x1, w6)), F32.add(s15, F32.mul(x1, w7)), F32.add(s16, F32.mul(x2, w0)), F32.add(s17, F32.mul(x2, w1)), F32.add(s18, F32.mul(x2, w2)), F32.add(s19, F32.mul(x2, w3)), F32.add(s20, F32.mul(x2, w4)), F32.add(s21, F32.mul(x2, w5)), F32.add(s22, F32.mul(x2, w6)), F32.add(s23, F32.mul(x2, w7)), F32.add(s24, F32.mul(x3, w0)), F32.add(s25, F32.mul(x3, w1)), F32.add(s26, F32.mul(x3, w2)), F32.add(s27, F32.mul(x3, w3)), F32.add(s28, F32.mul(x3, w4)), F32.add(s29, F32.mul(x3, w5)), F32.add(s30, F32.mul(x3, w6)), F32.add(s31, F32.mul(x3, w7)), F32.add(s32, F32.mul(x4, w0)), F32.add(s33, F32.mul(x4, w1)), F32.add(s34, F32.mul(x4, w2)), F32.add(s35, F32.mul(x4, w3)), F32.add(s36, F32.mul(x4, w4)), F32.add(s37, F32.mul(x4, w5)), F32.add(s38, F32.mul(x4, w6)), F32.add(s39, F32.mul(x4, w7)), F32.add(s40, F32.mul(x5, w0)), F32.add(s41, F32.mul(x5, w1)), F32.add(s42, F32.mul(x5, w2)), F32.add(s43, F32.mul(x5, w3)), F32.add(s44, F32.mul(x5, w4)), F32.add(s45, F32.mul(x5, w5)), F32.add(s46, F32.mul(x5, w6)), F32.add(s47, F32.mul(x5, w7)), F32.add(s48, F32.mul(x6, w0)), F32.add(s49, F32.mul(x6, w1)), F32.add(s50, F32.mul(x6, w2)), F32.add(s51, F32.mul(x6, w3)), F32.add(s52, F32.mul(x6, w4)), F32.add(s53, F32.mul(x6, w5)), F32.add(s54, F32.mul(x6, w6)), F32.add(s55, F32.mul(x6, w7)), F32.add(s56, F32.mul(x7, w0)), F32.add(s57, F32.mul(x7, w1)), F32.add(s58, F32.mul(x7, w2)), F32.add(s59, F32.mul(x7, w3)), F32.add(s60, F32.mul(x7, w4)), F32.add(s61, F32.mul(x7, w5)), F32.add(s62, F32.mul(x7, w6)), F32.add(s63, F32.mul(x7, w7)), U32.is_lt(U32.add(j, 1), inner)) |
|
|
| def packed_join(a: PackedBuffers, b: PackedBuffers) -> PackedBuffers: |
| JBuffers{xa, wa, oa} = a |
| JBuffers{xb, wb, ob} = b |
| JBuffers{Array.join(Input8, xa, xb), Array.join(Weight8, wa, wb), Array.join(Output64, oa, ob)} |
|
|
| @unsafe |
| def packed_step(x: Array<Input8>, w: Array<Weight8>, o: Array<Output64>, |
| +lo: U32, +hi: U32, +tiles_col: U32, +inner: U32, active: Bool) -> PackedBuffers: |
| match active: |
| case False{}: |
| JBuffers{x, w, o} |
| case True{}: |
| row = U32.div(lo, tiles_col) |
| col = U32.mod(lo, tiles_col) |
| packed_dot(x, w, o, row, col, tiles_col, inner, hi, 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, True{}) |
|
|
| @unsafe |
| def packed_split(x: Array<Input8> & Array<Input8>, w: Array<Weight8> & Array<Weight8>, o: Array<Output64> & Array<Output64>, |
| +lo: U32, +hi: U32, +tiles_col: U32, +grain: U32, +inner: U32) -> PackedBuffers: |
| (xa, xb) = x |
| (wa, wb) = w |
| (oa, ob) = o |
| +mid = U32.add(lo, U32.div(U32.sub(hi, lo), 2)) |
| left right = packed_gemm(xa, wa, oa, lo, mid, tiles_col, grain, inner, U32.is_le(U32.sub(mid, lo), grain)) packed_gemm(xb, wb, ob, mid, hi, tiles_col, grain, inner, U32.is_le(U32.sub(hi, mid), grain)) |
| packed_join(left, right) |
|
|
| @unsafe |
| def packed_gemm(x: Array<Input8>, w: Array<Weight8>, o: Array<Output64>, |
| +lo: U32, +hi: U32, +tiles_col: U32, +grain: U32, +inner: U32, leaf: Bool) -> PackedBuffers: |
| match leaf: |
| case True{}: |
| row = U32.div(lo, tiles_col) |
| col = U32.mod(lo, tiles_col) |
| packed_dot(x, w, o, row, col, tiles_col, inner, hi, 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, True{}) |
| case False{}: |
| packed_split(Array.fork(Input8, x), Array.fork(Weight8, w), Array.fork(Output64, o), lo, hi, tiles_col, grain, inner) |
|
|
| def packed_input() -> IO(Array<Input8>): |
| import "./bridge.c" |
|
|
| def packed_weights() -> IO(Array<Weight8>): |
| import "./bridge.c" |
|
|
| def packed_out() -> IO(Array<Output64>): |
| import "./bridge.c" |
|
|
| def packed_output(result: PackedBuffers) -> IO(Unit): |
| import "./bridge.c" |
|
|
| def main() -> IO(Unit): |
| do IO<Unit>: |
| +tree : Candidates <- input() |
| output(winner(tree)) |
| output(softmax(tree)) |
| vector : Features <- features() |
| write_features(layernorm(vector, 384.0, 0.00001)) |
| rows : NormRows <- batch_features() |
| write_batch(norm_rows(rows, 384.0, 0.00001)) |
| +tile : TileInput <- tile_input() |
| +w : JuliaMatrix <- dense_weights() |
| tile_output(dense_tile(tile, w)) |
| tile_output(dense_tile(tile, w)) |
| +tiles : InputTiles <- input_tiles() |
| +matrix : JuliaMatrix <- dense_weights() |
| output_tiles(dense_batch(tiles, matrix)) |
| output_tiles(dense_batch(tiles, matrix)) |
| px : Array<Input8> <- packed_input() |
| pw : Array<Weight8> <- packed_weights() |
| po : Array<Output64> <- packed_out() |
| packed_output(packed_gemm(px, pw, po, 0, 1, 1, 1, 1, True{})) |
|
|
