lib/accy/src/kernel/model/core/view.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const choir_abi = @import("choir_abi");
  2 
  3 const builder = @import("builder.zig");
  4 const typed = @import("typed.zig");
  5 
  6 const DType = choir_abi.DType;
  7 
  8 pub const TypedValue = typed.Value;
  9 pub const KeyValue = typed.KeyValue;
 10 
 11 pub fn BufferView(comptime dtype: DType) type {
 12     if (dtype == .key) return KeyBufferView;
 13     return ElementBufferView(dtype);
 14 }
 15 
 16 pub const KeyBufferView = struct {
 17     value: builder.Value,
 18 
 19     const Self = @This();
 20 
 21     pub const element_dtype: DType = .key;
 22 
 23     pub fn init(value: builder.Value) Self {
 24         return .{ .value = value };
 25     }
 26 
 27     pub fn raw(self: Self) builder.Value {
 28         return self.value;
 29     }
 30 
 31     pub fn load(self: Self, k: anytype, index: anytype) !KeyValue {
 32         const base = try wordBase(k, index);
 33         const lo = try k.loadIndex(self.value, base);
 34         const hi = try k.loadIndex(self.value, try k.add(base, try k.constantIndex(1)));
 35         return .{ .lo = lo, .hi = hi };
 36     }
 37 
 38     pub fn store(self: Self, k: anytype, value: KeyValue, index: anytype) !void {
 39         const base = try wordBase(k, index);
 40         try k.storeIndex(value.lo, self.value, base);
 41         try k.storeIndex(value.hi, self.value, try k.add(base, try k.constantIndex(1)));
 42     }
 43 
 44     fn wordBase(k: anytype, index: anytype) !builder.Value {
 45         return k.mul(try k.linearIndex(index), try k.constantIndex(2));
 46     }
 47 };
 48 
 49 fn ElementBufferView(comptime dtype: DType) type {
 50     return struct {
 51         value: builder.Value,
 52 
 53         const Self = @This();
 54 
 55         pub const element_dtype = dtype;
 56 
 57         pub fn init(value: builder.Value) Self {
 58             return .{ .value = value };
 59         }
 60 
 61         pub fn raw(self: Self) builder.Value {
 62             return self.value;
 63         }
 64 
 65         pub fn load(self: Self, k: anytype, index: anytype) !TypedValue(dtype) {
 66             return k.typedValue(dtype, try k.loadIndex(self.value, index));
 67         }
 68 
 69         pub fn loadVector(self: Self, k: anytype, index: anytype, width: u32) !builder.Value {
 70             return k.loadVectorIndex(self.value, index, width);
 71         }
 72 
 73         pub fn store(self: Self, k: anytype, value: anytype, index: anytype) !void {
 74             try k.storeIndex((try TypedValue(dtype).from(k, value)).raw(), self.value, index);
 75         }
 76 
 77         pub fn storeVector(self: Self, k: anytype, value: builder.Value, index: anytype) !void {
 78             try k.storeIndex(value, self.value, index);
 79         }
 80 
 81         pub fn atomicRmw(
 82             self: Self,
 83             k: anytype,
 84             kind: builder.AtomicRmwKind,
 85             value: anytype,
 86             index: anytype,
 87         ) !TypedValue(dtype) {
 88             const operand = (try TypedValue(dtype).from(k, value)).raw();
 89             return k.typedValue(dtype, try k.atomicRmwIndex(kind, operand, self.value, index));
 90         }
 91 
 92         pub fn atomicCas(
 93             self: Self,
 94             k: anytype,
 95             expected: anytype,
 96             desired: anytype,
 97             index: anytype,
 98         ) !TypedValue(dtype) {
 99             const expected_operand = (try TypedValue(dtype).from(k, expected)).raw();
100             const desired_operand = (try TypedValue(dtype).from(k, desired)).raw();
101             return k.typedValue(dtype, try k.atomicCasIndex(expected_operand, desired_operand, self.value, index));
102         }
103     };
104 }