tiny.accy.choir.record.memory

1const std = @import("std");2const choir_abi = @import("choir_abi");3const reference = @import("root.zig").reference;45pub const LayoutKind = enum {6    scalar,7    row_major,8    dynamic_row_major,910    pub fn name(self: LayoutKind) []const u8 {11        return switch (self) {12            .scalar => "scalar",13            .row_major => "row_major",14            .dynamic_row_major => "dynamic_row_major",15        };16    }1718    pub fn hasStaticStrides(self: LayoutKind) bool {19        return self == .scalar or self == .row_major;20    }21};2223pub const OutputSource = union(enum) {24    kernel_written: usize,25    aliased: usize,26};2728pub const BoundaryTransfer = enum {29    none,30    host_to_device,31    device_to_host,32    bidirectional,3334    pub fn needsHostInput(self: BoundaryTransfer) bool {35        return self == .host_to_device or self == .bidirectional;36    }3738    pub fn needsHostOutput(self: BoundaryTransfer) bool {39        return self == .device_to_host or self == .bidirectional;40    }41};4243pub const MemoryAccess = enum {44    read_only,45    write_only,46    read_write,47};4849pub const MemorySpace = enum {50    host,51    device_global,52    device_constant,53    device_shared,54    unified,5556    pub fn name(self: MemorySpace) []const u8 {57        return switch (self) {58            .host => "host",59            .device_global => "device_global",60            .device_constant => "device_constant",61            .device_shared => "device_shared",62            .unified => "unified",63        };64    }65};6667pub const BufferRole = struct {68    input: bool = false,69    output: bool = false,70    temporary: bool = false,71    constant: bool = false,7273    pub fn isBoundary(self: BufferRole) bool {74        return self.input or self.output;75    }76};7778pub const Slot = struct {79    id: usize,80    value: reference.Value,81    producer: ?reference.Operation,82    function: ?reference.Operation,83    role: BufferRole,84    dtype: choir_abi.DType,85    dims: []const i64,86    element_count: ?u64,87    row_major_strides: ?[]const u64,88    byte_size: ?u64,89};9091pub const Elision = struct {92    value: reference.Value,93    producer: reference.Operation,94    root: reference.Operation,95    cluster_index: usize,96};9798pub const Buffers = struct {99    slots: []const Slot,100    elisions: []const Elision,101    input_slot_count: usize,102    output_slot_count: usize,103    temporary_slot_count: usize,104    constant_slot_count: usize,105    dynamic_slot_count: usize,106    total_static_bytes: u64,107};108109pub const Assignment = struct {110    slot_id: usize,111    value: reference.Value,112    producer: ?reference.Operation,113    role: BufferRole,114    space: MemorySpace,115    access: MemoryAccess,116    transfer: BoundaryTransfer,117    byte_size: ?u64,118    output_source: ?OutputSource,119};120121pub const Spaces = struct {122    assignments: []const Assignment,123    host_slot_count: usize,124    device_global_slot_count: usize,125    device_constant_slot_count: usize,126    device_shared_slot_count: usize,127    unified_slot_count: usize,128    host_input_transfer_count: usize,129    host_output_transfer_count: usize,130    dynamic_slot_count: usize,131    elided_value_count: usize,132    total_static_bytes: u64,133};134135pub const Layout = struct {136    slot_id: usize,137    value: reference.Value,138    producer: ?reference.Operation,139    role: BufferRole,140    dtype: choir_abi.DType,141    memory_space: MemorySpace,142    kind: LayoutKind,143    rank: usize,144    dims: []const i64,145    element_strides: ?[]const u64,146    minor_to_major: []const usize,147    element_count: ?u64,148    byte_size: ?u64,149    element_size: u64,150    alignment: u64,151    contiguous: bool,152    static_layout: bool,153};154155pub const Layouts = struct {156    assignments: []const Layout,157    scalar_layout_count: usize,158    row_major_layout_count: usize,159    dynamic_row_major_layout_count: usize,160    host_slot_count: usize,161    device_global_slot_count: usize,162    device_constant_slot_count: usize,163    device_shared_slot_count: usize,164    unified_slot_count: usize,165    dynamic_slot_count: usize,166    elided_value_count: usize,167    total_static_bytes: u64,168};169170pub const Binding = struct {171    value: reference.Value,172    slot_id: usize,173174    pub fn lessThan(_: void, left: Binding, right: Binding) bool {175        const lhs = key(left.value);176        const rhs = key(right.value);177        return std.mem.lessThan(u32, &lhs, &rhs);178    }179180    fn key(value: reference.Value) [5]u32 {181        return switch (value) {182            .result => |item| .{ item.operation.ordinal, 0, 0, 0, item.position },183            .argument => |item| .{ item.operation.ordinal, 1, item.region, item.block, item.position },184        };185    }186};187188pub const Record = struct {189    buffers: Buffers,190    bindings: []const Binding,191    spaces: Spaces,192    layouts: Layouts,193};194195pub fn validate(allocator: std.mem.Allocator, value: Record, dispatch: @import("root.zig").dispatch.Record) !void {196    try validateBindings(allocator, value);197    try validateBufferCounts(value.buffers);198    try validateSpaceCounts(value.spaces, value.buffers.elisions.len);199    try validateLayoutCounts(value.layouts, value.buffers.elisions.len);200    const slots = value.buffers.slots.len;201    if (value.spaces.assignments.len != slots or value.layouts.assignments.len != slots) {202        return error.InvalidStageRecord;203    }204    for (value.buffers.slots, 0..) |slot, index| {205        if (slot.id != index) return error.InvalidStageRecord;206        if (slot.row_major_strides) |strides| {207            if (strides.len != slot.dims.len) return error.InvalidStageRecord;208        }209    }210    for (value.buffers.elisions) |elision| {211        if (elision.cluster_index >= dispatch.fusion.clusters.len) return error.InvalidStageRecord;212    }213    for (value.spaces.assignments, 0..) |assignment, index| {214        if (assignment.slot_id != index) return error.InvalidStageRecord;215        if (assignment.output_source) |source| switch (source) {216            .aliased => |slot| if (slot >= slots) return error.InvalidStageRecord,217            .kernel_written => |work| if (work >= dispatch.schedule.work_items.len) {218                return error.InvalidStageRecord;219            },220        };221    }222    for (value.layouts.assignments, 0..) |layout, index| {223        if (layout.slot_id != index or layout.rank != layout.dims.len or224            layout.minor_to_major.len != layout.rank) return error.InvalidStageRecord;225        if (layout.element_strides) |strides| {226            if (strides.len != layout.rank) return error.InvalidStageRecord;227        }228        for (layout.minor_to_major) |axis| {229            if (axis >= layout.rank) return error.InvalidStageRecord;230        }231    }232}233234fn validateBindings(allocator: std.mem.Allocator, value: Record) !void {235    const primary = try allocator.alloc(bool, value.buffers.slots.len);236    defer allocator.free(primary);237    @memset(primary, false);238    for (value.bindings, 0..) |binding, index| {239        if (binding.slot_id >= primary.len) return error.InvalidStageRecord;240        if (index != 0 and !Binding.lessThan({}, value.bindings[index - 1], binding)) {241            return error.InvalidStageRecord;242        }243        if (std.meta.eql(binding.value, value.buffers.slots[binding.slot_id].value)) {244            primary[binding.slot_id] = true;245        }246    }247    for (primary) |present| if (!present) return error.InvalidStageRecord;248}249250const SizeCounts = struct {251    dynamic: usize = 0,252    bytes: u64 = 0,253254    fn add(self: *SizeCounts, bytes: ?u64) !void {255        if (bytes) |size| {256            self.bytes = std.math.add(u64, self.bytes, size) catch257                return error.InvalidStageRecord;258        } else self.dynamic += 1;259    }260261    fn matches(self: SizeCounts, value: anytype) bool {262        return self.dynamic == value.dynamic_slot_count and self.bytes == value.total_static_bytes;263    }264};265266const SpaceCounts = struct {267    values: [5]usize = @splat(0),268269    fn add(self: *SpaceCounts, space: MemorySpace) void {270        const index: usize = switch (space) {271            .host => 0,272            .device_global => 1,273            .device_constant => 2,274            .device_shared => 3,275            .unified => 4,276        };277        self.values[index] += 1;278    }279280    fn matches(self: SpaceCounts, value: anytype) bool {281        return std.mem.eql(usize, &self.values, &.{282            value.host_slot_count,283            value.device_global_slot_count,284            value.device_constant_slot_count,285            value.device_shared_slot_count,286            value.unified_slot_count,287        });288    }289};290291fn validateBufferCounts(value: Buffers) !void {292    var roles: [4]usize = @splat(0);293    var sizes: SizeCounts = .{};294    for (value.slots) |slot| {295        if (slot.role.input) roles[0] += 1;296        if (slot.role.output) roles[1] += 1;297        if (slot.role.temporary) roles[2] += 1;298        if (slot.role.constant) roles[3] += 1;299        try sizes.add(slot.byte_size);300    }301    if (!sizes.matches(value) or !std.mem.eql(usize, &roles, &.{302        value.input_slot_count,303        value.output_slot_count,304        value.temporary_slot_count,305        value.constant_slot_count,306    })) return error.InvalidStageRecord;307}308309fn validateSpaceCounts(value: Spaces, elisions: usize) !void {310    var spaces: SpaceCounts = .{};311    var sizes: SizeCounts = .{};312    var inputs: usize = 0;313    var outputs: usize = 0;314    for (value.assignments) |assignment| {315        spaces.add(assignment.space);316        try sizes.add(assignment.byte_size);317        if (assignment.transfer.needsHostInput()) inputs += 1;318        if (assignment.transfer.needsHostOutput()) outputs += 1;319    }320    if (!spaces.matches(value) or !sizes.matches(value) or321        inputs != value.host_input_transfer_count or outputs != value.host_output_transfer_count or322        elisions != value.elided_value_count) return error.InvalidStageRecord;323}324325fn validateLayoutCounts(value: Layouts, elisions: usize) !void {326    var spaces: SpaceCounts = .{};327    var sizes: SizeCounts = .{};328    var kinds: [3]usize = @splat(0);329    for (value.assignments) |assignment| {330        spaces.add(assignment.memory_space);331        try sizes.add(assignment.byte_size);332        switch (assignment.kind) {333            .scalar => kinds[0] += 1,334            .row_major => kinds[1] += 1,335            .dynamic_row_major => kinds[2] += 1,336        }337    }338    if (!spaces.matches(value) or !sizes.matches(value) or elisions != value.elided_value_count or339        !std.mem.eql(usize, &kinds, &.{340            value.scalar_layout_count,341            value.row_major_layout_count,342            value.dynamic_row_major_layout_count,343        })) return error.InvalidStageRecord;344}

1pub const memory = @import("memory.zig");