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. 4 import Base 5 import ./machine_limits.bend as Machine 6 7 def bounded_result(valid: Bool,value: Nat) -> Maybe<&2,Nat>: 8 match valid: 9 case False{}: None{} 10 case True{}: Some{value} 11 12 def bounded(+value: Nat,limit: Nat) -> Maybe<&2,Nat>: 13 bounded_result(Nat.is_le(value,limit),value) 14 15 def 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 20 def 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 25 def 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 30 def 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. 36 def 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 41 def 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 46 def multiply(+left: U32,+right: U32) -> Maybe<&2,U32>: 47 multiply_left(U32.is_eq(left,1),left,right)