tiny.accy.kernel.library.random.block

1pub fn philoxBlock(rounds: u32, counter: [4]u32, key: [2]u32) [4]u32 {2    var state = counter;3    var bumped = key;4    var round: u32 = 0;5    while (round < rounds) : (round += 1) {6        if (round > 0) {7            bumped[0] +%= philox_w0;8            bumped[1] +%= philox_w1;9        }10        const product0 = @as(u64, philox_m0) * @as(u64, state[0]);11        const product1 = @as(u64, philox_m1) * @as(u64, state[2]);12        const hi0: u32 = @truncate(product0 >> 32);13        const lo0: u32 = @truncate(product0);14        const hi1: u32 = @truncate(product1 >> 32);15        const lo1: u32 = @truncate(product1);16        state = .{ hi1 ^ state[1] ^ bumped[0], lo1, hi0 ^ state[3] ^ bumped[1], lo0 };17    }18    return state;19}

1pub fn philoxWords(2    k: anytype,3    rounds: u32,4    counter: kernel.Value,5    counter_fold: kernel.Value,6    seed_lo: kernel.Value,7    seed_hi: kernel.Value,8) ![philox_lanes]kernel.Value {9    const zero = try k.constantInt(.i32, 0);10    const m0 = try wordConstant(k, philox_m0);11    const m1 = try wordConstant(k, philox_m1);12    const w0 = try wordConstant(k, philox_w0);13    const w1 = try wordConstant(k, philox_w1);14    var state = [philox_lanes]kernel.Value{ counter, counter_fold, zero, zero };15    var key = [2]kernel.Value{ seed_lo, seed_hi };16    var round: u32 = 0;17    while (round < rounds) : (round += 1) {18        if (round > 0) {19            key[0] = try k.add(key[0], w0);20            key[1] = try k.add(key[1], w1);21        }22        const hi0 = try k.umulhi(m0, state[0]);23        const lo0 = try k.mul(m0, state[0]);24        const hi1 = try k.umulhi(m1, state[2]);25        const lo1 = try k.mul(m1, state[2]);26        state = .{27            try k.xor(try k.xor(hi1, state[1]), key[0]),28            lo1,29            try k.xor(try k.xor(hi0, state[3]), key[1]),30            lo0,31        };32    }33    return state;34}

1pub const philox_default_rounds: u32 = 10;

1pub const philox_lanes: u64 = 4;

1pub const philox_max_rounds: u32 = 16;

1pub fn rotateWide32(k: anytype, value: kernel.Value) !kernel.Value {2    const half = try k.constantInt(.i64, 32);3    const low = try k.ushr(value, half);4    const high = try k.shl(value, half);5    return k.or_(low, high);6}

1pub fn squaresBlock(counter: u64, key: u64) u32 {2    var x = counter *% key;3    const y = x;4    const z = y +% key;5    x = x *% x +% y;6    x = std.math.rotr(u64, x, 32);7    x = x *% x +% z;8    x = std.math.rotr(u64, x, 32);9    x = x *% x +% y;10    x = std.math.rotr(u64, x, 32);11    return @truncate((x *% x +% z) >> 32);12}

1pub fn squaresKeyValue(k: anytype, key_lo: kernel.Value, key_hi: kernel.Value) !kernel.Value {2    const lo = try zeroExtendWord(k, key_lo);3    const hi = try zeroExtendWord(k, key_hi);4    const half = try k.constantInt(.i64, 32);5    return k.or_(try k.shl(hi, half), lo);6}

1pub fn squaresWord(2    k: anytype,3    counter: kernel.Value,4    key: kernel.Value,5) !kernel.Value {6    var x = try k.mul(counter, key);7    const y = x;8    const z = try k.add(y, key);9    x = try k.add(try k.mul(x, x), y);10    x = try rotateWide32(k, x);11    x = try k.add(try k.mul(x, x), z);12    x = try rotateWide32(k, x);13    x = try k.add(try k.mul(x, x), y);14    x = try rotateWide32(k, x);15    const final = try k.add(try k.mul(x, x), z);16    const half = try k.constantInt(.i64, 32);17    const high = try k.ushr(final, half);18    return k.cast(high, .i32);19}

1pub const squares_default_key: u64 = 0x123456789abcdef0;

1pub const squares_lanes: u64 = 1;

1pub fn zeroExtendWord(k: anytype, word: kernel.Value) !kernel.Value {2    const wide = try k.cast(word, .i64);3    const mask = try k.constantInt(.i64, 0xFFFFFFFF);4    return k.and_(wide, mask);5}

1pub fn rotateLeftWord(k: anytype, value: kernel.Value, amount: u5) !kernel.Value {2    const left = try k.shl(value, try k.constantInt(.i32, amount));3    const right = try k.ushr(value, try k.constantInt(.i32, 32 - @as(i64, amount)));4    return k.or_(left, right);5}

1pub fn threefryBlock(rounds: u32, counter: [2]u32, key: [2]u32) [2]u32 {2    const subkeys = [3]u32{ key[0], key[1], threefry_parity ^ key[0] ^ key[1] };3    var x0 = counter[0] +% subkeys[0];4    var x1 = counter[1] +% subkeys[1];5    var round: u32 = 0;6    while (round < rounds) : (round += 1) {7        x0 +%= x1;8        x1 = std.math.rotl(u32, x1, threefry_rotations[round % 8]);9        x1 ^= x0;10        if (round % 4 == 3) {11            const injection = round / 4 + 1;12            x0 +%= subkeys[injection % 3];13            x1 +%= subkeys[(injection + 1) % 3];14            x1 +%= injection;15        }16    }17    return .{ x0, x1 };18}

1pub fn threefryWords(2    k: anytype,3    rounds: u32,4    counter: kernel.Value,5    counter_fold: kernel.Value,6    seed_lo: kernel.Value,7    seed_hi: kernel.Value,8) ![threefry_lanes]kernel.Value {9    const parity = try wordConstant(k, threefry_parity);10    const subkeys = [3]kernel.Value{11        seed_lo,12        seed_hi,13        try k.xor(try k.xor(parity, seed_lo), seed_hi),14    };15    var x0 = try k.add(counter, subkeys[0]);16    var x1 = try k.add(counter_fold, subkeys[1]);17    var round: u32 = 0;18    while (round < rounds) : (round += 1) {19        x0 = try k.add(x0, x1);20        x1 = try rotateLeftWord(k, x1, threefry_rotations[round % 8]);21        x1 = try k.xor(x1, x0);22        if (round % 4 == 3) {23            const injection = round / 4 + 1;24            x0 = try k.add(x0, subkeys[injection % 3]);25            x1 = try k.add(x1, subkeys[(injection + 1) % 3]);26            x1 = try k.add(x1, try k.constantInt(.i32, injection));27        }28    }29    return .{ x0, x1 };30}

1pub const threefry_default_rounds: u32 = 20;

1pub const threefry_lanes: u64 = 2;

1pub const threefry_max_rounds: u32 = 32;

1pub fn ceilDiv(numerator: u64, denominator: u64) u64 {2    return numerator / denominator + @intFromBool(numerator % denominator != 0);3}

1pub fn outputWord(k: anytype, comptime dtype: DType, word: kernel.Value) !kernel.Value {2    return switch (dtype) {3        .i32 => word,4        .f32 => blk: {5            const shifted = try k.ushr(word, try k.constantInt(.i32, uniform_mantissa_shift));6            const filled = try k.or_(shifted, try wordConstant(k, uniform_one_bits));7            const unit = try k.bitcast(filled, .f32);8            break :blk try k.sub(unit, try k.constantFloat(.f32, 1.0));9        },10        else => @compileError("random kernels support dtype .i32 or .f32"),11    };12}

1pub fn randomDTypeSupported(dtype: DType) bool {2    return switch (dtype) {3        .f32, .i32 => true,4        else => false,5    };6}

1pub fn storeGuarded(2    k: anytype,3    args: anytype,4    out: kernel.Value,5    element: kernel.Value,6    count: kernel.Value,7) !void {8    const active = try k.compare(.lt, element, count);9    try k.guardDo(active, .{ .args = args, .out = out, .element = element }, store_guarded_active);10}

1pub fn uniformFromBits(bits: u32) f32 {2    const filled = (bits >> uniform_mantissa_shift) | uniform_one_bits;3    return @as(f32, @bitCast(filled)) - 1.0;4}

1pub fn wordConstant(k: anytype, word: u32) !kernel.Value {2    return k.constantInt(.i32, @as(i32, @bitCast(word)));3}

1const common = @import("word.zig");2const philox = @import("philox.zig");3const squares = @import("squares.zig");4const threefry = @import("threefry.zig");56pub const philox_lanes = philox.philox_lanes;7pub const philox_default_rounds = philox.philox_default_rounds;8pub const philox_max_rounds = philox.philox_max_rounds;9pub const philoxBlock = philox.philoxBlock;10pub const philoxWords = philox.philoxWords;1112pub const threefry_lanes = threefry.threefry_lanes;13pub const threefry_default_rounds = threefry.threefry_default_rounds;14pub const threefry_max_rounds = threefry.threefry_max_rounds;15pub const threefryBlock = threefry.threefryBlock;16pub const rotateLeftWord = threefry.rotateLeftWord;17pub const threefryWords = threefry.threefryWords;1819pub const squares_lanes = squares.squares_lanes;20pub const squares_default_key = squares.squares_default_key;21pub const squaresBlock = squares.squaresBlock;22pub const zeroExtendWord = squares.zeroExtendWord;23pub const rotateWide32 = squares.rotateWide32;24pub const squaresWord = squares.squaresWord;25pub const squaresKeyValue = squares.squaresKeyValue;2627pub const uniformFromBits = common.uniformFromBits;28pub const randomDTypeSupported = common.randomDTypeSupported;29pub const ceilDiv = common.ceilDiv;30pub const wordConstant = common.wordConstant;31pub const outputWord = common.outputWord;32pub const storeGuarded = common.storeGuarded;

1pub const block = @import("block/root.zig");