nx_linalg_affine_t34.nx source
↩ module page · 197 lines · 12082 B
1// Private additive linear-algebra extension for scale-preserving affine skinning.
2// Uses existing Q30 arithmetic; no new persistent asset format.
3import "nx_linalg.nx"
4const LA_AFFINE_RANGE:i64=0-60
5const LA_AFFINE_SINGULAR:i64=0-61
6const LA_AFFINE_RESIDUAL:i64=0-62
7const LA_AFFINE_MAX:i64=70368744177664
8const LA_AFFINE_I64_MAX:i64=9223372036854775807
9func la_affine_domain(v:i64)->i64{if v<=0-LA_AFFINE_MAX||v>=LA_AFFINE_MAX{return LA_AFFINE_RANGE};return 0}
10// Row-major 3x3 Q30. Caller output stays unchanged on domain/singularity/residual refusal.
11func la_inverse3_q30(a:*Mat,out:*Mat,residual_budget:i64,report:*i64)->i64{
12 if a.rows!=3||a.cols!=3||out.rows!=3||out.cols!=3||residual_budget<0{return LA_AFFINE_RANGE}
13 let cofactor:*i64=sys_mmap(72) as *i64;let inverse:*i64=sys_mmap(72) as *i64
14 var rc:i64=0;var r:i64=0
15 while r<3&&rc==0{var c:i64=0;while c<3&&rc==0{if la_affine_domain(nx_mat_get(a,r,c))!=0{rc=LA_AFFINE_RANGE};c=c+1};r=r+1}
16 r=0;while r<3&&rc==0{var c:i64=0;while c<3&&rc==0{
17 let x:i64=fq_mul(nx_mat_get(a,(r+1)%3,(c+1)%3),nx_mat_get(a,(r+2)%3,(c+2)%3))
18 let y:i64=fq_mul(nx_mat_get(a,(r+1)%3,(c+2)%3),nx_mat_get(a,(r+2)%3,(c+1)%3))
19 // Each product is below 2^62 by the shared primitive's admitted domain.
20 let v:i64=x-y;if la_affine_domain(v)!=0{rc=LA_AFFINE_RANGE}else{cofactor[r*3+c]=v};c=c+1
21 };r=r+1}
22 var determinant:i64=0;var c:i64=0
23 while c<3&&rc==0{let term:i64=fq_mul(nx_mat_get(a,0,c),cofactor[c]);if term>0&&determinant>LA_AFFINE_I64_MAX-term{rc=LA_AFFINE_RANGE}else{if term<0&&determinant<(0-LA_AFFINE_I64_MAX)-term{rc=LA_AFFINE_RANGE}else{determinant=determinant+term}};c=c+1}
24 if determinant==0&&rc==0{rc=LA_AFFINE_SINGULAR}
25 r=0;while r<3&&rc==0{c=0;while c<3&&rc==0{
26 let numerator:i64=cofactor[c*3+r]
27 if fq_abs(numerator)/fq_abs(determinant)>=LA_AFFINE_MAX/FQ_ONE{rc=LA_AFFINE_RANGE}else{let v:i64=fq_div(numerator,determinant);if la_affine_domain(v)!=0{rc=LA_AFFINE_RANGE}else{inverse[r*3+c]=v}};c=c+1
28 };r=r+1}
29 var maximum:i64=0;r=0
30 while r<3&&rc==0{c=0;while c<3&&rc==0{var value:i64=0;var k:i64=0
31 while k<3{let term:i64=fq_mul(nx_mat_get(a,r,k),inverse[k*3+c]);if term>0&&value>LA_AFFINE_I64_MAX-term{rc=LA_AFFINE_RANGE}else{if term<0&&value<(0-LA_AFFINE_I64_MAX)-term{rc=LA_AFFINE_RANGE}else{value=value+term}};k=k+1}
32 var target:i64=0;if r==c{target=FQ_ONE};if value<0-LA_AFFINE_MAX||value>LA_AFFINE_MAX{rc=LA_AFFINE_RANGE}else{let difference:i64=fq_abs(value-target);if difference>maximum{maximum=difference}};c=c+1
33 };r=r+1}
34 report[0]=determinant;report[1]=maximum
35 if rc==0&&maximum>residual_budget{rc=LA_AFFINE_RESIDUAL}
36 if rc==0{r=0;while r<3{c=0;while c<3{nx_mat_set(out,r,c,inverse[r*3+c]);c=c+1};r=r+1}}
37 sys_munmap(cofactor as *u8,72);sys_munmap(inverse as *u8,72);return rc
38}
39
40// Column-major 4x4 Q30; domain follows the existing exact fixed-point multiply owner.
41func la_affine_compose_q30(a:*i64,b:*i64,out:*i64)->i64{
42 var i:i64=0;while i<16{if la_affine_domain(a[i])!=0||la_affine_domain(b[i])!=0{return LA_AFFINE_RANGE};i=i+1}
43 let temporary:*i64=sys_mmap(128) as *i64;var rc:i64=0;var c:i64=0
44 while c<4&&rc==0{var r:i64=0;while r<4&&rc==0{var sum:i64=0;var k:i64=0
45 while k<4&&rc==0{let term:i64=fq_mul(a[k*4+r],b[c*4+k]);if term>0&&sum>LA_AFFINE_I64_MAX-term{rc=LA_AFFINE_RANGE}else{if term<0&&sum<(0-LA_AFFINE_I64_MAX)-term{rc=LA_AFFINE_RANGE}else{sum=sum+term}};k=k+1}
46 if la_affine_domain(sum)!=0{rc=LA_AFFINE_RANGE};temporary[c*4+r]=sum;r=r+1
47 };c=c+1}
48 if rc==0{i=0;while i<16{out[i]=temporary[i];i=i+1}}
49 sys_munmap(temporary as *u8,128);return rc
50}
51
52// Static yaw draw enclosure, binary32 input ABI: bounds loXYZ/hiXYZ, sin, cos,
53// base scale, instance scale, translation XYZ. Deformation is a caller admission
54// prerequisite. The result is a world-axis Q8 AABB, not anatomical geometry.
55// Q30 interval endpoints are outward rounded; source / 2^10 and base * 2^10
56// preserve the exact real expression while staying in fq_mul's proved domain.
57const LA_STATIC_INPUT_WORDS:i64=13
58const LA_STATIC_OUTPUT_WORDS:i64=6
59const LA_STATIC_SCRATCH_WORDS:i64=48
60const LA_BINARY32_FRACTION_BITS:i64=23
61const LA_BINARY32_ULP_DEN:i64=8388608
62const LA_STATIC_SOURCE_SHIFT:i64=10
63const LA_STATIC_PATH_ROUNDS:i64=5
64const LA_STATIC_GAMMA_FACTOR:i64=2*LA_STATIC_PATH_ROUNDS
65const LA_STATIC_Q8_SHIFT:i64=FQ_SHIFT-8
66func la_iv(p:*i64,index:i64)->*i64{return ((p as i64)+index*16) as *i64}
67func la_iv_abs(p:*i64)->i64{let a:i64=fq_abs(p[0]);let b:i64=fq_abs(p[1]);if a>b{return a};return b}
68func la_f32_ordered(a:i64,b:i64)->i64{
69 if (a&2147483647)==0&&(b&2147483647)==0{return 1}
70 var x:i64=2147483648+a;var y:i64=2147483648+b
71 if (a&2147483648)!=0{x=4294967295-a};if (b&2147483648)!=0{y=4294967295-b}
72 return (x<=y) as i64
73}
74func la_iv_f32(bits:i64,power:i64,out:*i64)->i64{
75 if bits<0||bits>4294967295{return LA_AFFINE_RANGE}
76 let exp:i64=(bits>>23)&255;var mant:i64=bits&8388607
77 if exp==255{return LA_AFFINE_RANGE}
78 var shift:i64=exp-120+power
79 if exp==0{shift=0-119+power}else{mant=mant+8388608}
80 var lo:i64=0;var hi:i64=0
81 if mant!=0{
82 if shift>=0{if shift>=63{return LA_AFFINE_RANGE};if mant>((LA_AFFINE_MAX-1)>>shift){return LA_AFFINE_RANGE};lo=mant<<shift;hi=lo}
83 else{let n:i64=0-shift;if n>=63{hi=1}else{lo=mant>>n;hi=lo;if (lo<<n)!=mant{hi=hi+1}}}
84 }
85 if hi>=LA_AFFINE_MAX{return LA_AFFINE_RANGE}
86 if (bits&2147483648)!=0{out[0]=0-hi;out[1]=0-lo}else{out[0]=lo;out[1]=hi}
87 return 0
88}
89func la_iv_mul(a:*i64,b:*i64,out:*i64)->i64{
90 var low:i64=LA_AFFINE_I64_MAX;var high:i64=0-LA_AFFINE_I64_MAX
91 var i:i64=0;while i<2{var j:i64=0;while j<2{
92 if la_affine_domain(a[i])!=0||la_affine_domain(b[j])!=0{return LA_AFFINE_RANGE}
93 let value:i64=fq_mul(a[i],b[j]);if value<low{low=value};if value>high{high=value};j=j+1
94 };i=i+1}
95 // fq_mul truncates toward zero. One Q30 unit on either side encloses its remainder.
96 if low<=0-LA_AFFINE_MAX+1||high>=LA_AFFINE_MAX-1{return LA_AFFINE_RANGE}
97 out[0]=low-1;out[1]=high+1;return 0
98}
99func la_static_storage_disjoint(a:*i64,aw:i64,b:*i64,bw:i64)->i64{
100 let x:i64=a as i64;let y:i64=b as i64
101 if x<=0||y<=0{return 0}
102 if x<=y{if y-x<aw*8{return 0}}else{if x-y<bw*8{return 0}}
103 return 1
104}
105// WGSL permits reassociation and input/result FTZ (15.7.2,15.7.5).
106// Expanded horizontal expression has two four-factor monomials plus translation:
107// 9 leaves and 8 binary operations. Each of these 17 sites may lose less than
108// 2^-126 through FTZ. Its remaining algebraic gain is at most 2*B; remaining
109// rounding amplifies by less than 2. Thus 68*B*2^-126 is an absolute FTZ bound.
110// B bounds every product subset using max(1,abs(each factor group)).
111const LA_BINARY32_MIN_NORMAL_POWER:i64=0-126
112const LA_BINARY32_FINITE_POWER:i64=128
113const LA_STATIC_EXPANDED_MONOMIALS:i64=2
114const LA_STATIC_FACTORS_PER_MONOMIAL:i64=4
115const LA_STATIC_EXPANDED_LEAVES:i64=LA_STATIC_EXPANDED_MONOMIALS*LA_STATIC_FACTORS_PER_MONOMIAL+1
116const LA_STATIC_EXPANDED_OPS:i64=LA_STATIC_EXPANDED_LEAVES-1
117const LA_STATIC_ROUND_GAIN:i64=2
118const LA_STATIC_FTZ_COEFFICIENT:i64=(LA_STATIC_EXPANDED_LEAVES+LA_STATIC_EXPANDED_OPS)*LA_STATIC_EXPANDED_MONOMIALS*LA_STATIC_ROUND_GAIN
119// Three monomials including translation, with rounding gain below2: 6B.
120// 6*2^125 is strictly below maximum finite binary32; derive the power headroom.
121func la_static_sum_headroom()->i64{var n:i64=(LA_STATIC_EXPANDED_MONOMIALS+1)*LA_STATIC_ROUND_GAIN;var p:i64=0;var ceiling:i64=1;while ceiling<n{ceiling=ceiling<<1;p=p+1};return p}
122func la_f32_abs_ge_one_power(bits:i64)->i64{
123 if bits<0||bits>4294967295{return LA_AFFINE_RANGE}
124 let exponent:i64=(bits>>23)&255;if exponent==255{return LA_AFFINE_RANGE}
125 if exponent<127{return 0}
126 var power:i64=exponent-127;if (bits&8388607)!=0{power=power+1};return power
127}
128func la_static_product_subset_power(input:*i64)->i64{
129 var source:i64=0;var yaw:i64=0;var translation:i64=0;var base:i64=0;var instance:i64=0;var i:i64=0
130 while i<LA_STATIC_INPUT_WORDS{
131 let e:i64=la_f32_abs_ge_one_power(input[i]);if e<0{return e}
132 if i<6{if e>source{source=e}}else{if i<8{if e>yaw{yaw=e}}else{if i==8{base=e}else{if i==9{instance=e}else{if e>translation{translation=e}}}}}
133 i=i+1
134 }
135 let total:i64=source+yaw+base+instance+translation
136 if total>LA_BINARY32_FINITE_POWER-la_static_sum_headroom(){return LA_AFFINE_RANGE}
137 return total
138}
139func la_static_ftz_error_q30(power:i64)->i64{
140 let shift:i64=power+FQ_SHIFT+LA_BINARY32_MIN_NORMAL_POWER
141 if shift>=0{return LA_STATIC_FTZ_COEFFICIENT<<shift}
142 let n:i64=0-shift;if n>=63{return 1}
143 return ((LA_STATIC_FTZ_COEFFICIENT-1)>>n)+1
144}
145func la_static_yaw_bounds_f32(input:*i64,input_words:i64,out:*i64,out_words:i64,scratch:*i64,scratch_words:i64)->i64{
146 if input_words!=LA_STATIC_INPUT_WORDS||out_words!=LA_STATIC_OUTPUT_WORDS||scratch_words<LA_STATIC_SCRATCH_WORDS{return LA_AFFINE_RANGE}
147 if la_static_storage_disjoint(input,input_words,out,out_words)==0||la_static_storage_disjoint(input,input_words,scratch,LA_STATIC_SCRATCH_WORDS)==0||la_static_storage_disjoint(out,out_words,scratch,LA_STATIC_SCRATCH_WORDS)==0{return LA_AFFINE_RANGE}
148 let subset_power:i64=la_static_product_subset_power(input);if subset_power<0{return subset_power}
149 let ftz_error:i64=la_static_ftz_error_q30(subset_power)
150 // Normalize the positive binary32 instance scale to [1,2) and move its
151 // exact power of two into the base. Equivalent draw factorizations must not
152 // fail solely because one factor exceeds the Q30 arithmetic operand domain.
153 // Actual products/output still obey the unchanged proved arithmetic bounds.
154 let scale_bits:i64=input[9]
155 if scale_bits<=0||scale_bits>=2147483648{return LA_AFFINE_RANGE}
156 let scale_exp:i64=(scale_bits>>23)&255
157 if scale_exp==255{return LA_AFFINE_RANGE}
158 var scale_power:i64=scale_exp-127
159 if scale_exp==0{
160 var significand:i64=scale_bits&8388607
161 if significand==0{return LA_AFFINE_RANGE}
162 scale_power=0-126
163 while significand<8388608{significand=significand<<1;scale_power=scale_power-1}
164 }
165 var i:i64=0
166 while i<LA_STATIC_INPUT_WORDS{var power:i64=0;if i<6{power=0-LA_STATIC_SOURCE_SHIFT};if i==8{power=LA_STATIC_SOURCE_SHIFT+scale_power};if i==9{power=0-scale_power};let rc:i64=la_iv_f32(input[i],power,la_iv(scratch,i));if rc!=0{return rc};i=i+1}
167 i=0;while i<3{if la_f32_ordered(input[i],input[i+3])==0{return LA_AFFINE_RANGE};let lo:*i64=la_iv(scratch,i);let hi:*i64=la_iv(scratch,i+3);if lo[0]>hi[1]{return LA_AFFINE_RANGE};lo[1]=hi[1];i=i+1}
168 let base:*i64=la_iv(scratch,8);let scale:*i64=la_iv(scratch,9)
169 if base[0]<=0||scale[0]<=0{return LA_AFFINE_RANGE}
170 let combined:*i64=la_iv(scratch,13);if la_iv_mul(base,scale,combined)!=0{return LA_AFFINE_RANGE}
171 var axis:i64=0
172 while axis<3{
173 let first:*i64=la_iv(scratch,14);let second:*i64=la_iv(scratch,15)
174 let term1:*i64=la_iv(scratch,16);let term2:*i64=la_iv(scratch,17)
175 if axis==1{if la_iv_mul(la_iv(scratch,2),combined,term1)!=0{return LA_AFFINE_RANGE};term2[0]=0;term2[1]=0}
176 else{
177 var a:i64=7;var b:i64=6;if axis==2{a=6;b=7}
178 if la_iv_mul(la_iv(scratch,0),la_iv(scratch,a),first)!=0||la_iv_mul(la_iv(scratch,1),la_iv(scratch,b),second)!=0{return LA_AFFINE_RANGE}
179 if la_iv_mul(first,combined,term1)!=0||la_iv_mul(second,combined,term2)!=0{return LA_AFFINE_RANGE}
180 if axis==2{let old:i64=term2[0];term2[0]=0-term2[1];term2[1]=0-old}
181 }
182 let translation:*i64=la_iv(scratch,10+axis)
183 let mass:i64=la_iv_abs(term1)+la_iv_abs(term2)+la_iv_abs(translation)
184 // Each monomial reaches the output through at most 3 multiplies and 2 adds.
185 // With u=2^-23 (either adjacent f32), gamma5=5u/(1-5u)<10u.
186 // Absolute monomial mass also covers cancellation and permitted reassociation.
187 let error:i64=(mass*LA_STATIC_GAMMA_FACTOR+LA_BINARY32_ULP_DEN-1)/LA_BINARY32_ULP_DEN+1+ftz_error
188 let low:i64=term1[0]+term2[0]+translation[0]-error
189 let high:i64=term1[1]+term2[1]+translation[1]+error
190 if la_affine_domain(low)!=0||la_affine_domain(high)!=0{return LA_AFFINE_RANGE}
191 // Arithmetic shifts give floor; negated floor of negative gives ceil.
192 scratch[40+axis]=low>>LA_STATIC_Q8_SHIFT
193 scratch[43+axis]=0-((0-high)>>LA_STATIC_Q8_SHIFT)
194 axis=axis+1
195 }
196 i=0;while i<LA_STATIC_OUTPUT_WORDS{out[i]=scratch[40+i];i=i+1};return 0
197}