mapping_cpu.art

// Swap functions ------------------------------------------------------------------
fn @cpu_swap_primary_entry(primary: &PrimaryStream, a: i32, b: i32) -> () {
    swap(&mut primary.rays.id(a),    &mut primary.rays.id(b));
    swap(&mut primary.rays.org_x(a), &mut primary.rays.org_x(b));
    swap(&mut primary.rays.org_y(a), &mut primary.rays.org_y(b));
    swap(&mut primary.rays.org_z(a), &mut primary.rays.org_z(b));
    swap(&mut primary.rays.dir_x(a), &mut primary.rays.dir_x(b));
    swap(&mut primary.rays.dir_y(a), &mut primary.rays.dir_y(b));
    swap(&mut primary.rays.dir_z(a), &mut primary.rays.dir_z(b));
    swap(&mut primary.rays.tmin(a),  &mut primary.rays.tmin(b));
    swap(&mut primary.rays.tmax(a),  &mut primary.rays.tmax(b));

    swap(&mut primary.ent_id(a),    &mut primary.ent_id(b));
    swap(&mut primary.prim_id(a),   &mut primary.prim_id(b));
    swap(&mut primary.t(a),         &mut primary.t(b));
    swap(&mut primary.u(a),         &mut primary.u(b));
    swap(&mut primary.v(a),         &mut primary.v(b));
    swap(&mut primary.rnd(a),       &mut primary.rnd(b));

    for c in unroll(0, MaxRayPayloadComponents) {
        swap(&mut primary.user(c)(a), &mut primary.user(c)(b));
    }
}

fn @cpu_swap_secondary_entry(secondary: &SecondaryStream, a: i32, b: i32) -> () {
    swap(&mut secondary.rays.id(a),    &mut secondary.rays.id(b));
    swap(&mut secondary.rays.org_x(a), &mut secondary.rays.org_x(b));
    swap(&mut secondary.rays.org_y(a), &mut secondary.rays.org_y(b));
    swap(&mut secondary.rays.org_z(a), &mut secondary.rays.org_z(b));
    swap(&mut secondary.rays.dir_x(a), &mut secondary.rays.dir_x(b));
    swap(&mut secondary.rays.dir_y(a), &mut secondary.rays.dir_y(b));
    swap(&mut secondary.rays.dir_z(a), &mut secondary.rays.dir_z(b));
    swap(&mut secondary.rays.tmin(a),  &mut secondary.rays.tmin(b));
    swap(&mut secondary.rays.tmax(a),  &mut secondary.rays.tmax(b));
    swap(&mut secondary.color_r(a),    &mut secondary.color_r(b));
    swap(&mut secondary.color_g(a),    &mut secondary.color_g(b));
    swap(&mut secondary.color_b(a),    &mut secondary.color_b(b));
    swap(&mut secondary.mat_id(a),     &mut secondary.mat_id(b));
}

// Sort functions ------------------------------------------------------------------
fn @cpu_sort_primary(primary: &PrimaryStream, ray_begins: &mut[i32], ray_ends: &mut[i32], num_geometries: i32) -> i32 {
    // Count the number of rays per shader
    for i in range(0, num_geometries + 1) {
        ray_ends(i) = 0;
    }
    for i in range(0, primary.size) {
        ray_ends(primary.ent_id(i))++;
    }

    // Compute scan over shader bins
    let mut n = 0;
    for i in range(0, num_geometries + 1) {
        ray_begins(i) = n;
        n += ray_ends(i);
        ray_ends(i) = n;
    }

    // Sort by shader
    for i in range(0, num_geometries) {
        let (begin, end) = (ray_begins(i), ray_ends(i));
        let mut j = begin;
        while j < end {
            let ent_id = primary.ent_id(j);
            if ent_id != i {
                let k = ray_begins(ent_id)++;
                cpu_swap_primary_entry(primary, k, j);
            } else {
                j++;
            }
        }
    }

    // Kill rays that have not intersected anything
    ray_ends(num_geometries - 1)
}

fn @cpu_sort_secondary(secondary: &SecondaryStream) -> i32 {
    fn @map_id(i:i32) = select(secondary.mat_id(i) < 0, 0:i32, 1:i32);

    // Get number of hits
    let mut count = 0;
    for i in range(0, secondary.size) {
        if map_id(i) == 0 { count++; }
    }
    
    let mut counters : [i32*2];
    counters(0) = 0;
    counters(1) = count;

    // Sort
    for i in unroll(0, 2) {
        let (begin, end) = if i == 0 { (0, count) } else { (count, secondary.size) };
        let mut j = begin;
        while j < end {
            let id = map_id(j);
            if id != i {
                let g = counters(id)++;
                cpu_swap_secondary_entry(secondary, g, j);
            } else {
                j++;
            }
        }
    }

    // Number of entries not hitting something
    count
}

fn @cpu_sort_secondary_with_materials(secondary: &SecondaryStream, ray_begins: &mut[i32], ray_ends: &mut[i32], num_materials: i32) -> i32 {
    fn @map_id(i:i32) -> i32 {
        let id = secondary.mat_id(i); // Is +1
        select(id < 0, -id, num_materials + id) - 1
    }

    // Reset
    for i in range(0, 2 * num_materials) {
        ray_ends(i) = 0;
    }

    // Count the number of rays per shader
    for i in range(0, secondary.size) {
        ray_ends(map_id(i))++;
    }

    // Compute scan over shader bins
    let mut n = 0;
    for i in range(0, 2 * num_materials) {
        ray_begins(i) = n;
        n += ray_ends(i);
        ray_ends(i) = n;
    }

    // Sort by shader
    for i in range(0, 2 * num_materials) {
        let (begin, end) = (ray_begins(i), ray_ends(i));
        let mut j = begin;
        while j < end {
            let id = map_id(j);
            if id != i {
                let g = ray_begins(id)++;
                cpu_swap_secondary_entry(secondary, g, j);
            } else {
                j++;
            }
        }
    }

    // Number of entries not hitting something
    ray_ends(num_materials-1)
}

// Compact functions ------------------------------------------------------------------
fn @cpu_compact_ray_stream(rays: RayStream, i: i32, j: i32, mask: bool) -> () {
    rays.org_x(i) = rv_compact(rays.org_x(j), mask);
    rays.org_y(i) = rv_compact(rays.org_y(j), mask);
    rays.org_z(i) = rv_compact(rays.org_z(j), mask);
    rays.dir_x(i) = rv_compact(rays.dir_x(j), mask);
    rays.dir_y(i) = rv_compact(rays.dir_y(j), mask);
    rays.dir_z(i) = rv_compact(rays.dir_z(j), mask);
    rays.tmin(i)  = rv_compact(rays.tmin(j),  mask);
    rays.tmax(i)  = rv_compact(rays.tmax(j),  mask);
}

fn @cpu_move_ray_stream(rays: RayStream, i: i32, j: i32) -> () {
    rays.org_x(i) = rays.org_x(j);
    rays.org_y(i) = rays.org_y(j);
    rays.org_z(i) = rays.org_z(j);
    rays.dir_x(i) = rays.dir_x(j);
    rays.dir_y(i) = rays.dir_y(j);
    rays.dir_z(i) = rays.dir_z(j);
    rays.tmin(i)  = rays.tmin(j);
    rays.tmax(i)  = rays.tmax(j);
}

fn @cpu_compact_primary(primary: &PrimaryStream, vector_width: i32, vector_compact: bool) -> i32 {
    fn cpu_compact_primary_specialized(primary2: &PrimaryStream) -> i32 {
        let mut k = 0;
        if vector_compact {
            for i in range_step(0, primary2.size, vector_width) {
                vectorize(vector_width, |j| { 
                    let id = primary.rays.id(i + j);
                    let mask = (id >= 0) & (i + j < primary2.size);

                    primary2.rays.id(k + j) = bitcast[i32](rv_compact(bitcast[f32](id), mask));

                    cpu_compact_ray_stream(primary2.rays, k + j, i + j, mask);

                    primary2.rnd(k + j) = bitcast[u32](rv_compact(bitcast[f32](primary2.rnd(i + j)), mask));
                    for c in unroll(0, MaxRayPayloadComponents) {
                        primary2.user(c)(k + j) = rv_compact(primary2.user(c)(i + j), mask);
                    }

                    k += cpu_popcount32(rv_ballot(mask));
                });
            }
        } else {
            for i in range(0, primary2.size) {
                let id = primary2.rays.id(i);
                if id >= 0 {
                    primary2.rays.id(k) = id;
                    cpu_move_ray_stream(primary2.rays, k, i);
                    primary2.rnd(k) = primary2.rnd(i);

                    for c in unroll(0, MaxRayPayloadComponents) {
                        primary2.user(c)(k) = primary2.user(c)(i);
                    }
                    k++;
                }
            } 
        }
        k
    }
    $cpu_compact_primary_specialized(primary)
}

fn @cpu_compact_secondary(secondary: &SecondaryStream, vector_width: i32, vector_compact: bool) -> i32 {
    fn cpu_compact_secondary_specialized(secondary2: &SecondaryStream) -> i32 {
        let mut k = 0;
        if vector_compact {
            for i in range_step(0, secondary2.size, vector_width) {
                vectorize(vector_width, |j| {
                    let id = secondary2.rays.id(i + j);
                    let mask = (id >= 0) & (i + j < secondary2.size);

                    secondary2.rays.id(k + j) = bitcast[i32](rv_compact(bitcast[f32](id), mask));

                    cpu_compact_ray_stream(secondary2.rays, k + j, i + j, mask);

                    secondary2.mat_id(k + j)  = bitcast[i32](rv_compact(bitcast[f32](secondary2.mat_id(i + j)), mask));
                    secondary2.color_r(k + j) = rv_compact(secondary2.color_r(i + j), mask);
                    secondary2.color_g(k + j) = rv_compact(secondary2.color_g(i + j), mask);
                    secondary2.color_b(k + j) = rv_compact(secondary2.color_b(i + j), mask);

                    k += cpu_popcount32(rv_ballot(mask));
                });
            }
        } else {
            for i in range(0, secondary2.size) {
                let id = secondary2.rays.id(i);
                if id >= 0 {
                    secondary2.rays.id(k) = id;
                    cpu_move_ray_stream(secondary2.rays, k, i);
                    secondary2.mat_id(k)  = secondary2.mat_id(i);
                    secondary2.color_r(k) = secondary2.color_r(i);
                    secondary2.color_g(k) = secondary2.color_g(i);
                    secondary2.color_b(k) = secondary2.color_b(i);
                    k++;
                }
            }
        }
        k
    }
    $cpu_compact_secondary_specialized(secondary)
}

// Ray generator ------------------------------------------------------------------
fn @cpu_generate_rays( primary: PrimaryStream
                     , capacity: i32
                     , emitter: RayEmitter
                     , id: &mut i32
                     , xmin: i32
                     , ymin: i32
                     , xmax: i32
                     , ymax: i32
                     , film_width: i32
                     , film_height: i32
                     , spi: i32
                     , vector_width: i32
                     ) -> i32 {
    let write_ray = make_ray_stream_writer(primary.rays, 1);
    let write_rnd = make_primary_stream_rnd_state_writer(primary, 1);
    let write_payload = make_primary_stream_payload_writer(primary, 1);
    let first_id = *id;
    let (tile_width, tile_height) = (xmax - xmin, ymax - ymin);
    let num_rays = cpu_intrinsics.min(spi * tile_width * tile_height - first_id, capacity - primary.size);
    let tile_div = make_fast_div(tile_width as u32);
    for i, _ in vectorized_range(vector_width, 0, num_rays) {
        let in_tile_id = first_id + i;

        // Compute x, y of ray within tile
        let sample = in_tile_id % spi;
        let in_tile_pixel = in_tile_id / spi;
        let in_tile_y = fast_div(tile_div, in_tile_pixel as u32) as i32;
        let in_tile_x = in_tile_pixel - in_tile_y * tile_width;
        let x = xmin + in_tile_x;
        let y = ymin + in_tile_y;
        let cur_ray = primary.size + i;
        let (ray, rnd, payload) = @emitter(sample, x, y, film_width, film_height);
        write_ray(cur_ray, 0, ray);
        write_rnd(cur_ray, 0, rnd);
        write_payload(cur_ray, 0, payload);
        primary.rays.id(cur_ray) = y * film_width + x;
    }
    
    *id = first_id + num_rays;
    primary.size + num_rays
}

fn @cpu_generate_rays_handler(size: i32
                            , capacity: i32
                            , emitter: RayEmitter
                            , id: &mut i32
                            , xmin: i32
                            , ymin: i32
                            , xmax: i32
                            , ymax: i32
                            , spi: i32
                            , vector_width: i32) -> i32 {
    let work_info = get_work_info();

    let mut primary : PrimaryStream;
    ignis_cpu_get_primary_stream(&mut primary, capacity);
    primary.size = size;
    cpu_generate_rays(primary, capacity, emitter, id, xmin, ymin, xmax, ymax, work_info.width, work_info.height, spi, vector_width)
}

// Traverse functions ------------------------------------------------------------------
fn @cpu_traverse_primary(scene: SceneGeometry, min_max: MinMax, primary: &PrimaryStream, single: bool, vector_width: i32) -> () {
    fn cpu_traverse_primary_specialized(scene2: SceneGeometry, primary2: &PrimaryStream) -> () {
        cpu_traverse(
            min_max,
            scene2,
            make_ray_stream_reader(primary2.rays, vector_width),
            make_primary_stream_hit_writer(*primary2, vector_width, scene.info.num_entities),
            vector_width /*packet_size*/,
            primary2.size / vector_width + select(primary2.size % vector_width != 0, 1, 0),
            single,
            false /*any_hit*/
        );
    }
    $cpu_traverse_primary_specialized(scene, primary);
}

fn @cpu_traverse_secondary(scene: SceneGeometry, min_max: MinMax, secondary: &SecondaryStream, single: bool, vector_width: i32) -> () {
    fn cpu_traverse_secondary_specialized(scene2: SceneGeometry, secondary2: &SecondaryStream) -> () {
        cpu_traverse(
            min_max,
            scene2,
            make_ray_stream_reader(secondary2.rays, vector_width),
            make_secondary_stream_hit_writer(*secondary2, vector_width),
            vector_width /*packet_size*/,
            secondary2.size / vector_width + select(secondary2.size % vector_width != 0, 1, 0),
            single,
            true /*any_hit*/
        );
    }
    $cpu_traverse_secondary_specialized(scene, secondary);
}

// Hit shader ------------------------------------------------------------------
fn @cpu_hit_shade(entity_id: i32, primary: &PrimaryStream, secondary: &SecondaryStream, shader: Shader, scene: Scene, path_tracer: Technique, accumulate: FilmAccumulator, begin: i32, end: i32, vector_width: i32) -> () {
    fn cpu_shade_specialized(entity_id2: i32, primary2: &PrimaryStream, secondary2: &SecondaryStream, begin2: i32, end2: i32) -> () {
        if begin == end { return() }

        let read_primary_ray        = make_ray_stream_reader(primary2.rays, 1);
        let read_primary_hit        = make_primary_stream_hit_reader(*primary2, 1);
        let read_primary_rnd_state  = make_primary_stream_rnd_state_reader(*primary2, 1);
        let read_primary_payload    = make_primary_stream_payload_reader(*primary2, 1);
        let write_primary_ray       = make_ray_stream_writer(primary2.rays, 1);
        let write_secondary_ray     = make_ray_stream_writer(secondary2.rays, 1);
        let write_primary_rnd_state = make_primary_stream_rnd_state_writer(*primary2, 1);
        let write_primary_payload   = make_primary_stream_payload_writer(*primary2, 1);
        
        let entities  = scene.database.entities;
        let shapes    = scene.database.shapes;
        let on_hit    = path_tracer.on_hit;
        let on_shadow = path_tracer.on_shadow;
        let on_bounce = path_tracer.on_bounce;
        let entity    = entities(entity_id2);
        let shape     = shapes(entity.shape_id);
        for i, r_vector_width in vectorized_range(vector_width, begin2, end2) {
            let ray     = read_primary_ray(i, 0);
            let hit     = read_primary_hit(i, 0);
            let mut rnd = read_primary_rnd_state(i, 0);
            let payload = read_primary_payload(i, 0);
            let ray_id  = primary2.rays.id(i);
            let pixel   = ray_id;

            let local_ray = transform_ray(ray, entity.local_mat);
            let lcl_surf  = shape.surface_element(local_ray, hit);
            let glb_surf  = map_surface_element(lcl_surf, entity.global_mat, entity.normal_mat);
            
            // Execute hit point shading, and add the contribution of each lane to the frame buffer
            let mat       = @shader(ray, hit, glb_surf);
            let hit_color = if let Option[Color]::Some(color) = @on_hit(ray, pixel, hit, payload, glb_surf, mat) { color } else { color_builtins::black };
            
            for lane in unroll(0, r_vector_width) {
                let j = bitcast[i32](rv_extract(bitcast[f32](pixel), lane));
                accumulate(j,
                    make_color(
                        rv_extract(hit_color.r, lane),
                        rv_extract(hit_color.g, lane),
                        rv_extract(hit_color.b, lane),
                        1
                    )
                );
            }

            // Compute shadow rays
            match @on_shadow(ray, pixel, hit, &mut rnd, payload, glb_surf, mat) {
                ShadowRay::Simple(new_ray, color) => {
                    write_secondary_ray(i, 0, new_ray);
                    secondary2.mat_id(i)  = mat.id + 1;
                    secondary2.color_r(i) = color.r;
                    secondary2.color_g(i) = color.g;
                    secondary2.color_b(i) = color.b;
                    secondary2.rays.id(i) = ray_id;
                },
                ShadowRay::Advanced(new_ray, color, mat_id) => {
                    write_secondary_ray(i, 0, new_ray);
                    secondary2.mat_id(i)  = mat_id + 1;
                    secondary2.color_r(i) = color.r;
                    secondary2.color_g(i) = color.g;
                    secondary2.color_b(i) = color.b;
                    secondary2.rays.id(i) = ray_id;
                },
                _ => { /* None */
                    secondary2.rays.id(i) = -1;
                }
            }

            // Sample new rays
            if let Option[(Ray, RayPayload)]::Some(new_ray, new_payload) = @on_bounce(ray, pixel, hit, &mut rnd, payload, glb_surf, mat) {
                write_primary_ray(i, 0, new_ray);
                write_primary_rnd_state(i, 0, rnd);
                write_primary_payload(i, 0, new_payload);
            } else {
                primary2.rays.id(i) = -1;
            }
        }
    }
    $cpu_shade_specialized(entity_id, primary, secondary, begin, end);
}

fn @cpu_hit_shade_handler(entity_id: i32, shader: Shader, scene: Scene, path_tracer: Technique, begin: i32, end: i32, spi: i32, use_framebuffer: bool, vector_width: i32) -> () {
    let (film_pixels, _, _) = cpu_get_film_data();

    let mut primary : PrimaryStream;
    ignis_cpu_get_primary_stream_const(&mut primary);
    let mut secondary : SecondaryStream;
    ignis_cpu_get_secondary_stream_const(&mut secondary);

    let accumulator = if !use_framebuffer { make_null_accumulator() } else { make_standard_accumulator(film_pixels, spi) };
    cpu_hit_shade(entity_id, primary, secondary, shader, scene, path_tracer, accumulator, begin, end, vector_width);
}

// Miss shader ------------------------------------------------------------------
fn @cpu_miss_shade(primary: &PrimaryStream, path_tracer: Technique, accumulate: FilmAccumulator, begin: i32, end: i32, vector_width: i32) -> () {
    fn cpu_miss_shade_specialized(primary2: &PrimaryStream, begin2: i32, end2: i32) -> () {
        if begin == end { return() }

        let read_primary_ray     = make_ray_stream_reader(primary2.rays, 1);
        let read_primary_payload = make_primary_stream_payload_reader(*primary2, 1);
        
        let on_miss = path_tracer.on_miss;

        for i, r_vector_width in vectorized_range(vector_width, begin2, end2) {
            let ray     = read_primary_ray(i, 0);
            let payload = read_primary_payload(i, 0);
            let ray_id  = primary2.rays.id(i);
            let pixel   = ray_id;

            // Execute hit point shading, and add the contribution of each lane to the frame buffer
            let hit_color = if let Option[Color]::Some(color) = @on_miss(ray, pixel, payload) { color } else { color_builtins::black };

            for lane in unroll(0, r_vector_width) {
                let j = bitcast[i32](rv_extract(bitcast[f32](pixel), lane));
                accumulate(j,
                    make_color(
                        rv_extract(hit_color.r, lane),
                        rv_extract(hit_color.g, lane),
                        rv_extract(hit_color.b, lane),
                        1
                    )
                );
            }

            primary2.rays.id(i) = -1;
        }
    }
    $cpu_miss_shade_specialized(primary, begin, end);
}

fn @cpu_miss_shade_handler(path_tracer: Technique, begin: i32, end: i32, spi: i32, use_framebuffer: bool, vector_width: i32) -> () {
    let (film_pixels, _, _) = cpu_get_film_data();

    let mut primary : PrimaryStream;
    ignis_cpu_get_primary_stream_const(&mut primary);

    let accumulator = if !use_framebuffer { make_null_accumulator() } else { make_standard_accumulator(film_pixels, spi) };
    cpu_miss_shade(primary, path_tracer, accumulator, begin, end, vector_width);
}

// Advanced shadow shader ------------------------------------------------------------------
fn @cpu_advanced_shadow(is_hit: bool, shader: Shader, secondary: &SecondaryStream, path_tracer: Technique, accumulate: FilmAccumulator, begin: i32, end: i32, vector_width: i32) -> () {
    fn cpu_advanced_shadow_specialized(secondary2: &SecondaryStream, begin2: i32, end2: i32) -> () {
        if begin == end { return() }

        let read_secondary_ray   = make_ray_stream_reader(secondary2.rays, 1);
        let read_secondary_color = make_secondary_stream_color_reader(*secondary2, 1);
        
        let on_hit  = path_tracer.on_shadow_hit;
        let on_miss = path_tracer.on_shadow_miss;

        let callback = if is_hit { on_hit } else { on_miss };

        for i, r_vector_width in vectorized_range(vector_width, begin2, end2) {
            let ray   = read_secondary_ray(i, 0);
            let color = read_secondary_color(i, 0);
            let pixel = secondary2.rays.id(i);

            // Execute hit point shading, and add the contribution of each lane to the frame buffer
            if let Option[Color]::Some(new_color) = @callback(ray, pixel, shader, color) { 
                for lane in unroll(0, r_vector_width) {
                    let j = bitcast[i32](rv_extract(bitcast[f32](pixel), lane));
                    accumulate(j,
                        make_color(
                            rv_extract(new_color.r, lane),
                            rv_extract(new_color.g, lane),
                            rv_extract(new_color.b, lane),
                            1
                        )
                    );
                }
            }
        }
    }
    $cpu_advanced_shadow_specialized(secondary, begin, end);
}

fn @cpu_advanced_shadow_handler(shader: Shader, path_tracer: Technique, begin: i32, end: i32, spi: i32, use_framebuffer: bool, is_hit: bool, vector_width: i32) -> () {
    let (film_pixels, _, _) = cpu_get_film_data();

    let mut secondary : SecondaryStream;
    ignis_cpu_get_secondary_stream_const(&mut secondary);

    let accumulator = if !use_framebuffer { make_null_accumulator() } else { make_standard_accumulator(film_pixels, spi) };
    cpu_advanced_shadow(is_hit, shader, secondary, path_tracer, accumulator, begin, end, vector_width);
}

// Framebuffer & AOVs ------------------------------------------------------------------
fn @cpu_get_film_data() -> (&mut [f32], i32, i32) {
    let mut film_pixels : &mut [f32];
    let mut film_width  : i32;
    let mut film_height : i32;
    ignis_get_film_data(0, &mut film_pixels, &mut film_width, &mut film_height);
    (film_pixels, film_width, film_height)
}

fn @cpu_get_aov_image(id: i32, spi: i32) -> AOVImage {
    // Width & height always the same as film_width, film_height
    let mut ptr : &mut [f32];
    ignis_get_aov_image(0, id, &mut ptr);
    let accumulate = make_standard_accumulator(ptr, spi);

    AOVImage {
        splat = @|pixel, color| -> () { 
            for lane in unroll(0, rv_num_lanes()) {
                let j = bitcast[i32](rv_extract(bitcast[f32](pixel), lane));
                accumulate(j,
                    make_color(
                        rv_extract(color.r, lane),
                        rv_extract(color.g, lane),
                        rv_extract(color.b, lane),
                        1
                    )
                );
            }
        },
        get = @|pixel| -> Color {
            // TODO: Make sure this is correct
            let mut color = color_builtins::black;
            for lane in unroll(0, rv_num_lanes()) {
                let j = bitcast[i32](rv_extract(bitcast[f32](pixel), lane));
                color.r = rv_insert(color.r, lane, ptr(j * 3 + 0));
                color.g = rv_insert(color.g, lane, ptr(j * 3 + 1));
                color.b = rv_insert(color.b, lane, ptr(j * 3 + 2));
            }
            color
        }
    }
}

// Main shader ------------------------------------------------------------------
fn @cpu_get_stream_capacity(spi: i32, tile_size: i32) = spi * tile_size * tile_size;

fn @cpu_trace( scene: SceneGeometry
             , pipeline: Pipeline
             , min_max: MinMax
             , single: bool
             , tile_size: i32
             , spi: i32
             , num_cores: i32
             , vector_width: i32
             , vector_compact: bool
             ) -> () {
    let (film_pixels, _film_width, _film_height) = cpu_get_film_data();

    let work_info  = get_work_info();
    let accumulate = make_standard_accumulator(film_pixels, spi);

    for xmin, ymin, xmax, ymax in cpu_parallel_tiles(work_info.width, work_info.height, tile_size, tile_size, num_cores) {
        ignis_register_thread();
        
        // Get ray streams/states from the CPU driver
        let mut primary   : PrimaryStream;
        let mut secondary : SecondaryStream;
        let capacity = cpu_get_stream_capacity(spi, tile_size);
        ignis_cpu_get_primary_stream(&mut primary,     capacity);
        ignis_cpu_get_secondary_stream(&mut secondary, capacity);

        let mut temp : TemporaryStorageHost;
        ignis_get_temporary_storage(0, &mut temp);

        let mut id = 0;
        let num_rays = spi * (ymax - ymin) * (xmax - xmin);
        while id < num_rays || primary.size > 0 {
            // (Re-)generate primary rays
            if primary.size < capacity && id < num_rays {
                let before_s = primary.size;
                primary.size = pipeline.on_generate(&mut id, primary.size, xmin, ymin, xmax, ymax);
                stats::add_quantity(stats::Quantity::CameraRayCount, primary.size - before_s);
            }

            if scene.info.num_entities == 0 {
                pipeline.on_miss_shade(0, primary.size);
                primary.size = 0;
            } else {
                // Trace primary rays
                cpu_traverse_primary(scene, min_max, primary, single, vector_width);

                // Sort hits by shader id, and filter invalid hits
                primary.size = cpu_sort_primary(primary, temp.ray_begins, temp.ray_ends, scene.info.num_entities);

                // Perform (vectorized) shading
                let mut begin = 0;
                for ent_id in range(0, scene.info.num_entities) {
                    let end = temp.ray_ends(ent_id);
                    if begin < end {
                        pipeline.on_hit_shade(ent_id, begin, end);
                    }
                    begin = end;
                }

                // Shade misses as well
                let last = temp.ray_ends(scene.info.num_entities);
                if begin < last {
                    pipeline.on_miss_shade(begin, last);
                }

                // Filter terminated rays
                secondary.size = primary.size;
                primary.size   = cpu_compact_primary(primary, vector_width, vector_compact);
                stats::add_quantity(stats::Quantity::BounceRayCount, primary.size);

                // Compact and trace secondary rays
                secondary.size = cpu_compact_secondary(secondary, vector_width, vector_compact);
                if likely(secondary.size > 0) {
                    cpu_traverse_secondary(scene, min_max, secondary, single, vector_width);
                    stats::add_quantity(stats::Quantity::ShadowRayCount, secondary.size);

                    // Add the contribution for secondary rays to the frame buffer
                    if work_info.advanced_shadows {
                        let hit_start = cpu_sort_secondary(secondary);
                        if hit_start != 0 {
                            // Call valids (miss)
                            pipeline.on_advanced_shadow(0, 0, hit_start, false);
                        }

                        if hit_start < secondary.size {
                            // Call invalids (hits)
                            pipeline.on_advanced_shadow(0, hit_start, secondary.size, true);
                        }
                    } else if work_info.advanced_shadows_with_materials {
                        let hit_start = cpu_sort_secondary_with_materials(secondary, temp.ray_begins, temp.ray_ends, scene.info.num_materials);

                        let mut sbegin = 0;
                        if hit_start != 0 {
                            // Call valids (miss)
                            for mat_id in range(0, scene.info.num_materials) {
                                let end = temp.ray_ends(mat_id);
                                if sbegin < end {
                                    pipeline.on_advanced_shadow(mat_id, sbegin, end, false);
                                }
                                sbegin = end;
                            }
                        }

                        if hit_start < secondary.size {
                            // Call invalids (hits)
                            for mat_id in range(0, scene.info.num_materials) {
                                let end = temp.ray_ends(mat_id + scene.info.num_materials);
                                if sbegin < end {
                                    pipeline.on_advanced_shadow(mat_id, sbegin, end, true);
                                }
                                sbegin = end;
                            }
                        }
                    } else if !work_info.framebuffer_locked {    
                        for i in range(0, secondary.size) {
                            if secondary.mat_id(i) < 0 {
                                let j = secondary.rays.id(i);
                                accumulate(j,
                                    make_color(
                                        secondary.color_r(i),
                                        secondary.color_g(i),
                                        secondary.color_b(i),
                                        1
                                    )
                                );
                            }
                        }
                    }
                }
            }
        }

        ignis_unregister_thread();
    }
}

// CPU device ----------------------------------------------------------------------
fn @make_cpu_device(vector_compact: bool, single: bool, min_max: MinMax, vector_width: i32, num_cores: i32, tile_size: i32) = Device {
    id    = 0,
    trace = @ |scene, pipeline, spi| {
        cpu_trace(
            scene,
            pipeline,
            min_max,
            single,
            tile_size,
            spi,
            num_cores,
            vector_width,
            vector_compact
        )
    },
    generate_rays = @ | emitter, id, size, xmin, ymin, xmax, ymax, spi | -> i32 {
        cpu_generate_rays_handler(size, @cpu_get_stream_capacity(spi, tile_size), emitter, id, xmin, ymin, xmax, ymax, spi, vector_width)
    },
    handle_miss_shader = @ | path_tracer, first, last, spi, use_framebuffer | {
        cpu_miss_shade_handler(path_tracer, first, last, spi, use_framebuffer, vector_width);
    },
    handle_hit_shader = @ | entity_id, shader, scene, path_tracer, first, last, spi, use_framebuffer | {
        cpu_hit_shade_handler(entity_id, shader, scene, path_tracer, first, last, spi, use_framebuffer, vector_width);
    },
    handle_advanced_shadow_shader = @ | shader, path_tracer, first, last, spi, use_framebuffer, is_hit | {
        cpu_advanced_shadow_handler(shader, path_tracer, first, last, spi, use_framebuffer, is_hit, vector_width);
    },
    present = @ || ignis_present(0),
    sync    = @ || {},
    parallel_range = @ |body| {
        @|start, end| {
            for i in parallel(num_cores, start, end) {
                @ body(i)
            }
        }
    },
    parallel_range_2d = @ |body| {
        @|start_x, end_x, start_y, end_y| {
            for xmin, ymin, xmax, ymax in cpu_parallel_tiles(end_x - start_x, end_y - start_y, tile_size, tile_size, num_cores) {
                for x in range(xmin + start_x, xmax + start_x) {
                    for y in range(ymin + start_y, ymax + start_y) {
                        @ body(x,y)
                    }   
                }
            }
        }
    },
    parallel_reduce_i32 = @|n, elem, op| cpu_reduce[i32](n, elem, op),
    parallel_reduce_f32 = @|n, elem, op| cpu_reduce[f32](n, elem, op),
    get_device_buffer_accessor = @|| { @|ptr| make_cpu_buffer(ptr, 0) },
    load_scene_bvh = @ || {
        if vector_width >= 8 {
            let mut nodes: &[Node8];
            let mut objs: &[EntityLeaf1];
            ignis_load_bvh8_ent(0, &mut nodes, &mut objs);
            make_cpu_bvh8_ent(nodes, objs)
        } else {
            let mut nodes: &[Node4];
            let mut objs: &[EntityLeaf1];
            ignis_load_bvh4_ent(0, &mut nodes, &mut objs);
            make_cpu_bvh4_ent(nodes, objs)
        }
    },
    load_scene_database = @ || {
        let mut database: SceneDatabase;
        ignis_load_scene(0, &mut database);
        database
    },
    load_scene_info       = @ || {
        let mut info: SceneInfo;
        ignis_load_scene_info(0, &mut info);
        info
    },
    load_entity_table     = @ |dtb| make_entity_table(dtb, @|ptr| make_cpu_buffer(ptr, 0)),
    load_shape_table      = @ |dtb| make_shape_table(dtb, @|ptr| make_cpu_buffer(ptr, 0)),
    load_specific_shape   = @ |num_face, num_vert, num_norm, num_tex, off, dtb| load_specific_shape_from_table(num_face, num_vert, num_norm, num_tex, off, dtb, @|ptr| make_cpu_buffer(ptr, 0)),
    load_custom_dyntable  = @ |name| -> DynTable {
        let mut table: DynTable;
        ignis_load_custom_dyntable(0, name, &mut table);
        table
    },
    load_bvh_table        = @ |dtb, _| -> BVHTable {
        @ |id| {
            let entry  = get_lookup_entry(id as u64,   dtb, @|ptr| make_cpu_buffer(ptr, 0)); 
            let header = get_table_entry(entry.offset, dtb, @|ptr| make_cpu_buffer(ptr, 0));
            let leaf_offset = header.load_i32(0) as u64;
    
            if vector_width >= 8 {
                let nodes = get_table_ptr(entry.offset + 16 , dtb) as &[Node8];
                let tris  = get_table_ptr(entry.offset + 16 + leaf_offset * (sizeof[Node8]() as u64), dtb) as &[Tri4];
                make_cpu_bvh8_tri4(nodes, tris)
            } else {
                let nodes = get_table_ptr(entry.offset + 16, dtb) as &[Node4];
                let tris  = get_table_ptr(entry.offset + 16 + leaf_offset * (sizeof[Node4]() as u64), dtb) as &[Tri4];
                make_cpu_bvh4_tri4(nodes, tris)
            }
        } 
    },
    load_image = @ |filename| {
        let mut pixel_data : &[f32];
        let mut width      : i32;
        let mut height     : i32;
        ignis_load_image(0, filename, &mut pixel_data, &mut width, &mut height);
        make_image_rgba32(@ |x, y| cpu_load_vec4(pixel_data, y * width + x),
                          width, height)
    },
    load_packed_image = @ |filename, hint_opaque| {
        let mut pixel_data : &[u32];
        let mut width      : i32;
        let mut height     : i32;
        ignis_load_packed_image(0, filename, &mut pixel_data, &mut width, &mut height);
        make_image_rgba32(@ |x, y| image_rgba_unpack(pixel_data(y * width + x), hint_opaque),
                          width, height)
    },
    load_aov_image = @|id, spi| { @cpu_get_aov_image(id, spi) },
    load_rays = @ || {
        let mut rays: &[StreamRay];
        ignis_load_rays(0, &mut rays);
        rays
    },
    load_host_buffer = load_cpu_buffer,
    load_buffer = load_cpu_buffer,
    request_buffer = @ |name, size, flags| {
        let mut ptr : &[u8];
        ignis_request_buffer(0, name, &mut ptr, size * sizeof[i32]() as i32, flags);
        make_cpu_buffer(ptr, size)
    },
    make_buffer = @ |ptr, size| make_cpu_buffer(ptr, size),
    dump_buffer = @ |id, filename| ignis_dbg_dump_buffer(0, id, filename),
    request_debug_output = @|| {
        if vector_width != 1 {
            // FIXME: Fix this for vectorized systems!
            make_fallback_debug_output()
        } else {
            let mut ptr : &[u8];
            let size = 4096;
            ignis_request_buffer(0, "__dbg_output", &mut ptr, size * sizeof[i32]() as i32, 0);
            let buffer = make_cpu_buffer(ptr, size);
            make_debug_output(buffer)
        }
    }
};

fn @make_avx512_device()             = make_cpu_device( true,  true, make_cpu_int_min_max(),16, 0, 32); // Not tested
fn @make_avx2_device()               = make_cpu_device( true,  true, make_cpu_int_min_max(), 8, 0, 16);
fn @make_avx_device()                = make_cpu_device( true,  true, make_cpu_int_min_max(), 8, 0, 16);
fn @make_sse42_device()              = make_cpu_device(false,  true, make_cpu_int_min_max(), 4, 0, 16);
fn @make_asimd_device()              = make_cpu_device(false, false, make_cpu_int_min_max(), 4, 0, 16);
fn @make_cpu_default_device()        = make_cpu_device(false, false, make_default_min_max(), 1, 0, 16);
fn @make_cpu_singlethreaded_device() = make_cpu_device(false, false, make_default_min_max(), 1, 1, 16);