simd-squares-of-squares

main.rs

#![feature(stdsimd)]

use std::simd::FromBits;
use std::simd::{i32x8, f32x8};
use std::arch::x86_64::{__m256, _mm256_sqrt_ps, _mm256_floor_ps};
use std::io::{stdout, Write};
use std::thread;

const NUM_THREADS: i32 = 4;

fn main() {
    // to increment a, we add [ 0,  1, -1,  1,  0, -1,  1, -1]
    // to increment b, we add [-1,  1,  0, -1,  1, -1,  0,  1]
    let inc_a = i32x8::new(0, 1, -1, 1, 0, -1, 1, -1);
    let inc_b = i32x8::new(-1, 1, 0, -1, 1, -1, 0, 1);

    let mut children = vec![];

    for thr in 0..NUM_THREADS {
        children.push(thread::spawn(move || {
            for i in 0.. {
                let real_i = NUM_THREADS * i + 3 + thr;
                println!("Starting iteration {}", real_i);
                let c = real_i * real_i;
                let mut base_sq = i32x8::splat(c) + inc_b;

                'b: for b in 1..(c-1) {
                    let mut cur_sq = base_sq + inc_a;
                    for a in 1..(c-b) {
                        if a == b {
                            cur_sq += inc_a;
                            continue;
                        }

                        if check_square(cur_sq) {
                                println!("Found good square: {:?}", cur_sq);
                                println!("  (a, b, c, c-a) = ({}, {}, {})", a, b, c);
                        }
                        // println!("Failed: {:?}", cur_sq);
                        cur_sq += inc_a;
                    }
                    base_sq += inc_b;
                }
            }
        }));
    }

    for child in children {
        child.join().unwrap();
    }
}

fn check_square(sq: i32x8) -> bool {

    if (sq & 2).min_element() == 0 || ((sq & 7) ^ 5).min_element() == 0 || ((sq & 11) ^ 8).min_element() == 0 {
        return false;
    }

    let roots = sqrts(f32x8::from(sq));
    round(roots * roots) == round(floors(roots) * floors(roots))
}

fn round(n: f32x8) -> f32x8 {
    floors(n + 0.5)
}

fn floors(n: f32x8) -> f32x8 {
    unsafe {
        f32x8::from_bits(_mm256_floor_ps(__m256::from_bits(n)))
    }
}

fn sqrts(n: f32x8) -> f32x8 {
    unsafe {
        f32x8::from_bits(_mm256_sqrt_ps(__m256::from_bits(n)))
    }
}