math_compare.bend
stelliferous@0.0.3.0 · 110 lines · 0 definitions
1# IEEE comparisons expressed as 32-bit masks, including unordered inputs.
2import Base
3import stelliferous@0.0.3.0/math_fp32.bend as FP32
4
5# One for a nonzero word, zero otherwise, without a Bool match.
6def nonzero(+word: U32) -> U32:
7 U32.shrn(U32.or(word,(0 - word : U32)),31n)
8
9def number(word: U32) -> U32:
10 U32.xor(1,U32.shrn((2139095040 - U32.and(word,2147483647) : U32),31n))
11
12def ordered(+left: U32,+right: U32) -> U32:
13 U32.and(number(left),number(right))
14
15# Magnitudes fit in 31 bits, so the top bit of their difference is the
16# unsigned comparison. Reverse it for negatives and handle opposite signs.
17def less_bit(+left: U32,+right: U32) -> U32:
18 +a = U32.and(left,2147483647)
19 +b = U32.and(right,2147483647)
20 +negative = U32.shrn(left,31n)
21 +different = U32.shrn(U32.xor(left,right),31n)
22 +below = U32.shrn((a - b : U32),31n)
23 +above = U32.shrn((b - a : U32),31n)
24 +same = U32.xor(below,U32.and(negative,U32.xor(below,above)))
25 U32.and(U32.xor(same,U32.and(different,U32.xor(same,negative))),nonzero(U32.or(a,b)))
26
27def equal_bit(+left: U32,+right: U32) -> U32:
28 U32.or(U32.xor(1,nonzero(U32.xor(left,right))),U32.xor(1,nonzero(U32.and(U32.or(left,right),2147483647))))
29
30def mask_value(bit: U32) -> F32:
31 FP32.from_bits(U32.and((0 - bit : U32),1065353216))
32
33def less_words(+left: U32,+right: U32) -> F32:
34 mask_value(U32.and(ordered(left,right),less_bit(left,right)))
35
36def less_equal_words(+left: U32,+right: U32) -> F32:
37 mask_value(U32.and(ordered(left,right),U32.xor(1,less_bit(right,left))))
38
39def equal_words(+left: U32,+right: U32) -> F32:
40 mask_value(U32.and(ordered(left,right),equal_bit(left,right)))
41
42def not_equal_words(+left: U32,+right: U32) -> F32:
43 mask_value(U32.xor(1,U32.and(ordered(left,right),equal_bit(left,right))))
44
45# Context comes first in Operations.lanes; the tensor element comes second.
46def less(value: F32,element: F32) -> F32:
47 less_words(F32.bits(element),F32.bits(value))
48
49def less_equal(value: F32,element: F32) -> F32:
50 less_equal_words(F32.bits(element),F32.bits(value))
51
52def greater(value: F32,element: F32) -> F32:
53 less_words(F32.bits(value),F32.bits(element))
54
55def greater_equal(value: F32,element: F32) -> F32:
56 less_equal_words(F32.bits(value),F32.bits(element))
57
58def equal(value: F32,element: F32) -> F32:
59 equal_words(F32.bits(element),F32.bits(value))
60
61def not_equal(value: F32,element: F32) -> F32:
62 not_equal_words(F32.bits(element),F32.bits(value))
63
64type Bounds is Data:
65 Bounds{low: F32,high: F32}
66
67def clamp(bounds: Bounds,element: F32) -> F32:
68 Bounds{low,high} = bounds
69 F32.min(F32.max(element,low),high)
70
71# Binary element order is left then right.
72def less_elements(left: F32,right: F32) -> F32:
73 less(right,left)
74
75def less_equal_elements(left: F32,right: F32) -> F32:
76 less_equal(right,left)
77
78def greater_elements(left: F32,right: F32) -> F32:
79 greater(right,left)
80
81def greater_equal_elements(left: F32,right: F32) -> F32:
82 greater_equal(right,left)
83
84def equal_elements(left: F32,right: F32) -> F32:
85 equal(right,left)
86
87def not_equal_elements(left: F32,right: F32) -> F32:
88 not_equal(right,left)
89
90# Selection only changes words; the chosen payload is never FP arithmetic.
91# selected is 0 or all ones, so selected + 1 is 1 or 0 modulo 2^32.
92# This masked integer product keeps the native selection in value lanes.
93def select(mask: F32,yes: F32,no: F32) -> F32:
94 F32{+mask_word} = mask
95 F32{+yes_word} = yes
96 F32{+no_word} = no
97 +magnitude = U32.and(U32{mask_word},2147483647)
98 +selected = (0 - U32.shrn((magnitude + 2147483647 : U32),31n) : U32)
99 FP32.from_bits((U32.and(U32{yes_word},selected) + U32{no_word} * (selected + 1 : U32) : U32))
100
101def selected(value: F32,mask: F32) -> F32:
102 select(mask,value,0.0)
103
104def rejected(value: F32,mask: F32) -> F32:
105 select(mask,0.0,value)
106
107def union(left: F32,right: F32) -> F32:
108 F32{left_word} = left
109 F32{right_word} = right
110 FP32.from_bits(U32.or(U32{left_word},U32{right_word}))