math_checked_word.bend

stelliferous@0.0.2.0 · 47 lines · 9 definitions

1# Exact overflow-checked multiplication for metadata. Each Horner step starts
2# below word_max and forms at most 3*word_max, safely below the native Nat cap.
3# The loop has exactly 32 bit steps and never touches tensor payloads.
4import Base
5import stelliferous@0.0.2.0/machine_limits.bend as Machine
6
7def bounded_result(valid: Bool,value: Nat) -> Maybe<&2,Nat>:
8  match valid:
9    case False{}: None{}
10    case True{}: Some{value}
11
12def bounded(+value: Nat,limit: Nat) -> Maybe<&2,Nat>:
13  bounded_result(Nat.is_le(value,limit),value)
14
15def multiply_step(bit: Bool,right: Nat,previous: Maybe<&2,Nat>,limit: Nat) -> Maybe<&2,Nat>:
16  match previous:
17    case None{}: None{}
18    case Some{value}: bounded(Nat.add(Nat.double(value),Bool.pick(Nat,bit,right,0n)),limit)
19
20def multiply_bits(width: Nat,bits: Word(width),+right: Nat,+limit: Nat) -> Maybe<&2,Nat>:
21  match width bits:
22    case 0n WNil{}: Some{0n}
23    case 1n+rest WCon{bit,tail}: multiply_step(bit,right,multiply_bits(rest,tail,right,limit),limit)
24
25def encoded(result: Maybe<&2,Nat>) -> Maybe<&2,U32>:
26  match result:
27    case None{}: None{}
28    case Some{value}: Some{U32.from_nat(value)}
29
30def multiply_bounded(left: U32,right: U32,limit: U32) -> Maybe<&2,U32>:
31  U32{bits} = left
32  encoded(multiply_bits(32n,bits,U32.to_nat(right),U32.to_nat(limit)))
33
34# Unit products already fit a word. Shapes and contiguous strides frequently
35# use them; returning the operand avoids the 32-step natural Horner traversal.
36def multiply_right(unit: Bool,left: U32,right: U32) -> Maybe<&2,U32>:
37  match unit:
38    case True{}: Some{left}
39    case False{}: multiply_bounded(left,right,Machine.word_max())
40
41def multiply_left(unit: Bool,+left: U32,+right: U32) -> Maybe<&2,U32>:
42  match unit:
43    case True{}: Some{right}
44    case False{}: multiply_right(U32.is_eq(right,1),left,right)
45
46def multiply(+left: U32,+right: U32) -> Maybe<&2,U32>:
47  multiply_left(U32.is_eq(left,1),left,right)