Skip to content

Commit cc15553

Browse files
hexqnthexqnt
committed
fix bound checks
1 parent ace57d4 commit cc15553

7 files changed

Lines changed: 141 additions & 139 deletions

File tree

‎src/fit/finite_diff.rs‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -96,18 +96,22 @@ where
9696

9797
let denom = 2.0 * step;
9898
let mut column_is_finite = true;
99-
for row in 0..dimension {
100-
let value = (grad_plus[row] - grad_minus[row]) / denom;
99+
for ((&plus, &minus), column_value) in array1_as_slice(&grad_plus)
100+
.iter()
101+
.zip(array1_as_slice(&grad_minus).iter())
102+
.zip(column_values.iter_mut())
103+
{
104+
let value = (plus - minus) / denom;
101105
if !value.is_finite() {
102106
column_is_finite = false;
103107
break;
104108
}
105-
column_values[row] = value;
109+
*column_value = value;
106110
}
107111

108112
if column_is_finite {
109-
for row in 0..dimension {
110-
hessian[[row, column]] = column_values[row];
113+
for (row, &value) in column_values.iter().enumerate() {
114+
hessian[[row, column]] = value;
111115
}
112116
computed = true;
113117
break;

‎src/fit/simd.rs‎

Lines changed: 83 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -62,10 +62,7 @@ pub(super) fn polynomial_cost_scalar(
6262
}
6363

6464
let mut sum = 0.0;
65-
let mut index = 0;
66-
while index < x_values.len() {
67-
let x = x_values[index];
68-
let y = y_values[index];
65+
for (&x, &y) in x_values.iter().zip(y_values.iter()) {
6966
let model = param
7067
.iter()
7168
.copied()
@@ -82,7 +79,6 @@ pub(super) fn polynomial_cost_scalar(
8279
if !sum.is_finite() {
8380
return LARGE_COST;
8481
}
85-
index += 1;
8682
}
8783

8884
sum / x_values.len() as f64
@@ -100,11 +96,13 @@ pub(super) fn inverse_cost_scalar(
10096
}
10197

10298
let mut sum = 0.0;
103-
let mut index = 0;
104-
while index < x_values.len() {
105-
let x = positive_x(x_values[index]);
106-
let y = y_values[index];
107-
let residual = (param[0] + param[1] / x) - y;
99+
let &[a, b, ..] = param else {
100+
unreachable!("inverse model requires two parameters");
101+
};
102+
103+
for (&x, &y) in x_values.iter().zip(y_values.iter()) {
104+
let x = positive_x(x);
105+
let residual = (a + b / x) - y;
108106
if !residual.is_finite() {
109107
return LARGE_COST;
110108
}
@@ -116,7 +114,6 @@ pub(super) fn inverse_cost_scalar(
116114
if !sum.is_finite() {
117115
return LARGE_COST;
118116
}
119-
index += 1;
120117
}
121118

122119
sum / x_values.len() as f64
@@ -131,10 +128,7 @@ pub(super) fn accumulate_polynomial_gradient_scalar(
131128
) {
132129
debug_assert_eq!(x_values.len(), y_values.len());
133130
debug_assert_eq!(gradient.len(), param.len());
134-
let mut index = 0;
135-
while index < x_values.len() {
136-
let x = x_values[index];
137-
let y = y_values[index];
131+
for (&x, &y) in x_values.iter().zip(y_values.iter()) {
138132
let model = param
139133
.iter()
140134
.copied()
@@ -146,7 +140,6 @@ pub(super) fn accumulate_polynomial_gradient_scalar(
146140
*gradient_value += residual * basis;
147141
basis *= x;
148142
}
149-
index += 1;
150143
}
151144
}
152145

@@ -159,14 +152,18 @@ pub(super) fn accumulate_inverse_gradient_scalar(
159152
) {
160153
debug_assert_eq!(x_values.len(), y_values.len());
161154
debug_assert!(gradient.len() >= 2);
162-
let mut index = 0;
163-
while index < x_values.len() {
164-
let x = positive_x(x_values[index]);
165-
let y = y_values[index];
166-
let residual = loss_metric.residual_derivative((param[0] + param[1] / x) - y);
167-
gradient[0] += residual;
168-
gradient[1] += residual / x;
169-
index += 1;
155+
let &[a, b, ..] = param else {
156+
unreachable!("inverse model requires two parameters");
157+
};
158+
let [gradient_0, gradient_1, ..] = gradient else {
159+
unreachable!("inverse gradient requires two parameters");
160+
};
161+
162+
for (&x, &y) in x_values.iter().zip(y_values.iter()) {
163+
let x = positive_x(x);
164+
let residual = loss_metric.residual_derivative((a + b / x) - y);
165+
*gradient_0 += residual;
166+
*gradient_1 += residual / x;
170167
}
171168
}
172169

@@ -226,23 +223,24 @@ pub(super) fn polynomial_cost_simd(
226223

227224
let mut sum = Vf64::splat(0.0);
228225
let mut tail_sum = 0.0;
229-
let mut index = 0;
230-
while index + Vf64::LEN <= x_values.len() {
231-
let x = Vf64::from_slice(&x_values[index..index + Vf64::LEN]);
232-
let y = Vf64::from_slice(&y_values[index..index + Vf64::LEN]);
226+
let (x_chunks, x_tail) = x_values.as_chunks::<{ Vf64::LEN }>();
227+
let (y_chunks, y_tail) = y_values.as_chunks::<{ Vf64::LEN }>();
228+
debug_assert_eq!(x_chunks.len(), y_chunks.len());
229+
debug_assert_eq!(x_tail.len(), y_tail.len());
230+
231+
for (x_chunk, y_chunk) in x_chunks.iter().zip(y_chunks.iter()) {
232+
let x = Vf64::from_array(*x_chunk);
233+
let y = Vf64::from_array(*y_chunk);
233234

234235
let mut model = Vf64::splat(0.0);
235236
for coefficient in param.iter().copied() {
236237
model = model * x + Vf64::splat(coefficient);
237238
}
238239

239240
sum += value_from_residual_simd(loss_metric, model - y);
240-
index += Vf64::LEN;
241241
}
242242

243-
while index < x_values.len() {
244-
let x = x_values[index];
245-
let y = y_values[index];
243+
for (&x, &y) in x_tail.iter().zip(y_tail.iter()) {
246244
let model = param
247245
.iter()
248246
.copied()
@@ -259,7 +257,6 @@ pub(super) fn polynomial_cost_simd(
259257
if !tail_sum.is_finite() {
260258
return LARGE_COST;
261259
}
262-
index += 1;
263260
}
264261

265262
let total = sum.reduce_sum() + tail_sum;
@@ -280,24 +277,29 @@ pub(super) fn inverse_cost_simd(
280277
if x_values.is_empty() {
281278
return 0.0;
282279
}
280+
let &[a_scalar, b_scalar, ..] = param else {
281+
unreachable!("inverse model requires two parameters");
282+
};
283283

284284
let mut sum = Vf64::splat(0.0);
285285
let mut tail_sum = 0.0;
286-
let mut index = 0;
287-
let a = Vf64::splat(param[0]);
288-
let b = Vf64::splat(param[1]);
286+
let (x_chunks, x_tail) = x_values.as_chunks::<{ Vf64::LEN }>();
287+
let (y_chunks, y_tail) = y_values.as_chunks::<{ Vf64::LEN }>();
288+
debug_assert_eq!(x_chunks.len(), y_chunks.len());
289+
debug_assert_eq!(x_tail.len(), y_tail.len());
290+
291+
let a = Vf64::splat(a_scalar);
292+
let b = Vf64::splat(b_scalar);
289293
let eps = Vf64::splat(super::PARAM_EPS);
290-
while index + Vf64::LEN <= x_values.len() {
291-
let x = Vf64::from_slice(&x_values[index..index + Vf64::LEN]).simd_max(eps);
292-
let y = Vf64::from_slice(&y_values[index..index + Vf64::LEN]);
294+
for (x_chunk, y_chunk) in x_chunks.iter().zip(y_chunks.iter()) {
295+
let x = Vf64::from_array(*x_chunk).simd_max(eps);
296+
let y = Vf64::from_array(*y_chunk);
293297
sum += value_from_residual_simd(loss_metric, (a + b / x) - y);
294-
index += Vf64::LEN;
295298
}
296299

297-
while index < x_values.len() {
298-
let x = positive_x(x_values[index]);
299-
let y = y_values[index];
300-
let residual = (param[0] + param[1] / x) - y;
300+
for (&x, &y) in x_tail.iter().zip(y_tail.iter()) {
301+
let x = positive_x(x);
302+
let residual = (a_scalar + b_scalar / x) - y;
301303
if !residual.is_finite() {
302304
return LARGE_COST;
303305
}
@@ -309,7 +311,6 @@ pub(super) fn inverse_cost_simd(
309311
if !tail_sum.is_finite() {
310312
return LARGE_COST;
311313
}
312-
index += 1;
313314
}
314315

315316
let total = sum.reduce_sum() + tail_sum;
@@ -332,10 +333,15 @@ pub(super) fn accumulate_polynomial_gradient_simd(
332333
debug_assert!(gradient.len() <= MAX_POLYNOMIAL_PARAMS);
333334

334335
let mut accum = [Vf64::splat(0.0); MAX_POLYNOMIAL_PARAMS];
335-
let mut index = 0;
336-
while index + Vf64::LEN <= x_values.len() {
337-
let x = Vf64::from_slice(&x_values[index..index + Vf64::LEN]);
338-
let y = Vf64::from_slice(&y_values[index..index + Vf64::LEN]);
336+
let accum = &mut accum[..gradient.len()];
337+
let (x_chunks, x_tail) = x_values.as_chunks::<{ Vf64::LEN }>();
338+
let (y_chunks, y_tail) = y_values.as_chunks::<{ Vf64::LEN }>();
339+
debug_assert_eq!(x_chunks.len(), y_chunks.len());
340+
debug_assert_eq!(x_tail.len(), y_tail.len());
341+
342+
for (x_chunk, y_chunk) in x_chunks.iter().zip(y_chunks.iter()) {
343+
let x = Vf64::from_array(*x_chunk);
344+
let y = Vf64::from_array(*y_chunk);
339345

340346
let mut model = Vf64::splat(0.0);
341347
for coefficient in param.iter().copied() {
@@ -344,20 +350,17 @@ pub(super) fn accumulate_polynomial_gradient_simd(
344350
let residual_derivative = residual_derivative_simd(loss_metric, model - y);
345351

346352
let mut basis = Vf64::splat(1.0);
347-
for gradient_index in (0..gradient.len()).rev() {
348-
accum[gradient_index] += residual_derivative * basis;
353+
for accum_value in accum.iter_mut().rev() {
354+
*accum_value += residual_derivative * basis;
349355
basis *= x;
350356
}
351-
index += Vf64::LEN;
352357
}
353358

354-
for (gradient_index, value) in gradient.iter_mut().enumerate() {
355-
*value += accum[gradient_index].reduce_sum();
359+
for (value, accum_value) in gradient.iter_mut().zip(accum.iter().copied()) {
360+
*value += accum_value.reduce_sum();
356361
}
357362

358-
while index < x_values.len() {
359-
let x = x_values[index];
360-
let y = y_values[index];
363+
for (&x, &y) in x_tail.iter().zip(y_tail.iter()) {
361364
let model = param
362365
.iter()
363366
.copied()
@@ -369,7 +372,6 @@ pub(super) fn accumulate_polynomial_gradient_simd(
369372
*gradient_value += residual * basis;
370373
basis *= x;
371374
}
372-
index += 1;
373375
}
374376
}
375377

@@ -382,32 +384,39 @@ pub(super) fn accumulate_inverse_gradient_simd(
382384
) {
383385
debug_assert_eq!(x_values.len(), y_values.len());
384386
debug_assert!(gradient.len() >= 2);
387+
let &[a_scalar, b_scalar, ..] = param else {
388+
unreachable!("inverse model requires two parameters");
389+
};
390+
let [gradient_scalar_0, gradient_scalar_1, ..] = gradient else {
391+
unreachable!("inverse gradient requires two parameters");
392+
};
385393

386394
let mut gradient_0 = Vf64::splat(0.0);
387395
let mut gradient_1 = Vf64::splat(0.0);
388-
let a = Vf64::splat(param[0]);
389-
let b = Vf64::splat(param[1]);
396+
let a = Vf64::splat(a_scalar);
397+
let b = Vf64::splat(b_scalar);
390398
let eps = Vf64::splat(super::PARAM_EPS);
391399

392-
let mut index = 0;
393-
while index + Vf64::LEN <= x_values.len() {
394-
let x = Vf64::from_slice(&x_values[index..index + Vf64::LEN]).simd_max(eps);
395-
let y = Vf64::from_slice(&y_values[index..index + Vf64::LEN]);
400+
let (x_chunks, x_tail) = x_values.as_chunks::<{ Vf64::LEN }>();
401+
let (y_chunks, y_tail) = y_values.as_chunks::<{ Vf64::LEN }>();
402+
debug_assert_eq!(x_chunks.len(), y_chunks.len());
403+
debug_assert_eq!(x_tail.len(), y_tail.len());
404+
405+
for (x_chunk, y_chunk) in x_chunks.iter().zip(y_chunks.iter()) {
406+
let x = Vf64::from_array(*x_chunk).simd_max(eps);
407+
let y = Vf64::from_array(*y_chunk);
396408
let residual_derivative = residual_derivative_simd(loss_metric, (a + b / x) - y);
397409
gradient_0 += residual_derivative;
398410
gradient_1 += residual_derivative / x;
399-
index += Vf64::LEN;
400411
}
401412

402-
gradient[0] += gradient_0.reduce_sum();
403-
gradient[1] += gradient_1.reduce_sum();
413+
*gradient_scalar_0 += gradient_0.reduce_sum();
414+
*gradient_scalar_1 += gradient_1.reduce_sum();
404415

405-
while index < x_values.len() {
406-
let x = positive_x(x_values[index]);
407-
let y = y_values[index];
408-
let residual = loss_metric.residual_derivative((param[0] + param[1] / x) - y);
409-
gradient[0] += residual;
410-
gradient[1] += residual / x;
411-
index += 1;
416+
for (&x, &y) in x_tail.iter().zip(y_tail.iter()) {
417+
let x = positive_x(x);
418+
let residual = loss_metric.residual_derivative((a_scalar + b_scalar / x) - y);
419+
*gradient_scalar_0 += residual;
420+
*gradient_scalar_1 += residual / x;
412421
}
413422
}

‎src/models/arctangent_step.rs‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -100,9 +100,11 @@ pub(super) fn add_value_grad_raw_hessian(
100100
let mut hessian = Array2::zeros((PARAM_COUNT, PARAM_COUNT));
101101
let params = Params::parse(param);
102102

103-
let mut index = 0;
104-
while index < sample_count {
105-
let x = x_values[index];
103+
for ((&x, &value_first), &value_second) in x_values
104+
.iter()
105+
.zip(value_first.iter())
106+
.zip(value_second.iter())
107+
{
106108
let u = x - params.x0;
107109
let z = params.slope * u;
108110
let atan_z = z.atan();
@@ -113,8 +115,6 @@ pub(super) fn add_value_grad_raw_hessian(
113115
return None;
114116
}
115117

116-
let value_first = value_first[index];
117-
let value_second = value_second[index];
118118
if !value_first.is_finite() || !is_finite_non_negative(value_second) {
119119
return None;
120120
}
@@ -143,7 +143,6 @@ pub(super) fn add_value_grad_raw_hessian(
143143
hessian[[2, 3]] += value_second * jac_c * jac_d;
144144

145145
hessian[[3, 3]] += value_second * jac_d * jac_d;
146-
index += 1;
147146
}
148147

149148
scale_and_mirror_upper_hessian(&mut hessian, sample_scale);

0 commit comments

Comments
 (0)