SelectLeafContext

fn SelectLeafContext(comptime Variant: type) type

File

Code

fn SelectLeafContext(comptime Variant: type) type {
    const cv_size = Variant.cv_size;
    const Result = BatchResult(Variant);

    return struct {
        view: *const MultiSliceView,
        batch_idx: usize,
        start_offset: usize,
        num_leaves: usize,

        fn process(ctx: @This()) Result {
            var result: Result = .{
                .batch_idx = ctx.batch_idx,
                .cv_len = ctx.num_leaves * cv_size,
                .cvs = undefined,
            };

            var leaf_buffer: [bytes_per_batch]u8 align(cache_line_size) = undefined;
            var leaves_processed: usize = 0;
            var byte_offset = ctx.start_offset;
            var cv_offset: usize = 0;
            const simd_batch_bytes = optimal_vector_len * chunk_size;
            while (leaves_processed + optimal_vector_len <= ctx.num_leaves) {
                if (ctx.view.tryGetSlice(byte_offset, byte_offset + simd_batch_bytes)) |leaf_data| {
                    var leaf_cvs: [optimal_vector_len * Variant.cv_size]u8 = undefined;
                    processLeaves(Variant, optimal_vector_len, leaf_data, &leaf_cvs);
                    @memcpy(result.cvs[cv_offset..][0..leaf_cvs.len], &leaf_cvs);
                } else {
                    ctx.view.copyRange(byte_offset, byte_offset + simd_batch_bytes, leaf_buffer[0..simd_batch_bytes]);
                    var leaf_cvs: [optimal_vector_len * Variant.cv_size]u8 = undefined;
                    processLeaves(Variant, optimal_vector_len, leaf_buffer[0..simd_batch_bytes], &leaf_cvs);
                    @memcpy(result.cvs[cv_offset..][0..leaf_cvs.len], &leaf_cvs);
                }
                leaves_processed += optimal_vector_len;
                byte_offset += optimal_vector_len * chunk_size;
                cv_offset += optimal_vector_len * cv_size;
            }

            while (leaves_processed < ctx.num_leaves) {
                const leaf_end = byte_offset + chunk_size;
                var cv_buffer: [64]u8 = undefined;

                if (ctx.view.tryGetSlice(byte_offset, leaf_end)) |leaf_data| {
                    const cv_slice = MultiSliceView.init(leaf_data, &[_]u8{}, &[_]u8{});
                    Variant.turboShakeToBuffer(&cv_slice, 0x0B, cv_buffer[0..cv_size]);
                } else {
                    ctx.view.copyRange(byte_offset, leaf_end, leaf_buffer[0..chunk_size]);
                    const cv_slice = MultiSliceView.init(leaf_buffer[0..chunk_size], &[_]u8{}, &[_]u8{});
                    Variant.turboShakeToBuffer(&cv_slice, 0x0B, cv_buffer[0..cv_size]);
                }
                @memcpy(result.cvs[cv_offset..][0..cv_size], cv_buffer[0..cv_size]);

                leaves_processed += 1;
                byte_offset += chunk_size;
                cv_offset += cv_size;
            }

            return result;
        }
    };
}