lib/accy/src/kernel/library/random/block/word.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const base = @import("../root.zig").base;
 2 
 3 const DType = base.DType;
 4 const kernel = base.kernel;
 5 
 6 const uniform_one_bits: u32 = 0x3F800000;
 7 const uniform_mantissa_shift: u5 = 9;
 8 
 9 pub fn uniformFromBits(bits: u32) f32 {
10     const filled = (bits >> uniform_mantissa_shift) | uniform_one_bits;
11     return @as(f32, @bitCast(filled)) - 1.0;
12 }
13 
14 pub fn randomDTypeSupported(dtype: DType) bool {
15     return switch (dtype) {
16         .f32, .i32 => true,
17         else => false,
18     };
19 }
20 
21 pub fn ceilDiv(numerator: u64, denominator: u64) u64 {
22     return numerator / denominator + @intFromBool(numerator % denominator != 0);
23 }
24 
25 pub fn wordConstant(k: anytype, word: u32) !kernel.Value {
26     return k.constantInt(.i32, @as(i32, @bitCast(word)));
27 }
28 
29 pub fn outputWord(k: anytype, comptime dtype: DType, word: kernel.Value) !kernel.Value {
30     return switch (dtype) {
31         .i32 => word,
32         .f32 => blk: {
33             const shifted = try k.ushr(word, try k.constantInt(.i32, uniform_mantissa_shift));
34             const filled = try k.or_(shifted, try wordConstant(k, uniform_one_bits));
35             const unit = try k.bitcast(filled, .f32);
36             break :blk try k.sub(unit, try k.constantFloat(.f32, 1.0));
37         },
38         else => @compileError("random kernels support dtype .i32 or .f32"),
39     };
40 }
41 
42 fn store_guarded_active(inner_builder: anytype, ctx: anytype) !void {
43     try ctx.args.param(.dst).store(inner_builder, ctx.out, ctx.element);
44 }
45 
46 pub fn storeGuarded(
47     k: anytype,
48     args: anytype,
49     out: kernel.Value,
50     element: kernel.Value,
51     count: kernel.Value,
52 ) !void {
53     const active = try k.compare(.lt, element, count);
54     try k.guardDo(active, .{ .args = args, .out = out, .element = element }, store_guarded_active);
55 }