tiny.accy.preparation.bufferization

1pub const BufferPlanAnalysis = struct {2    allocator: std.mem.Allocator,3    slots: std.ArrayListUnmanaged(BufferSlot),4    elisions: std.ArrayListUnmanaged(FusionElision),5    value_to_slot: std.AutoHashMap(*ir.Value, usize),6    value_to_elision: std.AutoHashMap(*ir.Value, usize),7    input_slot_count: usize = 0,8    output_slot_count: usize = 0,9    temporary_slot_count: usize = 0,10    constant_slot_count: usize = 0,11    dynamic_slot_count: usize = 0,12    total_static_bytes: u64 = 0,1314    pub fn init(allocator: std.mem.Allocator) BufferPlanAnalysis {15        return .{16            .allocator = allocator,17            .slots = .empty,18            .elisions = .empty,19            .value_to_slot = std.AutoHashMap(*ir.Value, usize).init(allocator),20            .value_to_elision = std.AutoHashMap(*ir.Value, usize).init(allocator),21        };22    }2324    pub fn deinit(self: *BufferPlanAnalysis) void {25        for (self.slots.items) |*slot| {26            slot.deinit(self.allocator);27        }28        self.slots.deinit(self.allocator);29        self.elisions.deinit(self.allocator);30        self.value_to_slot.deinit();31        self.value_to_elision.deinit();32        self.* = undefined;33    }3435    pub fn slotCount(self: BufferPlanAnalysis) usize {36        return self.slots.items.len;37    }3839    pub fn elisionCount(self: BufferPlanAnalysis) usize {40        return self.elisions.items.len;41    }4243    pub fn getSlot(self: *const BufferPlanAnalysis, value: *ir.Value) ?*const BufferSlot {44        const index = self.value_to_slot.get(value) orelse return null;45        return &self.slots.items[index];46    }4748    pub fn getElision(self: *const BufferPlanAnalysis, value: *ir.Value) ?*const FusionElision {49        const index = self.value_to_elision.get(value) orelse return null;50        return &self.elisions.items[index];51    }5253    fn addSlot(54        self: *BufferPlanAnalysis,55        value: *ir.Value,56        producer: ?*ir.Operation,57        function: ?*ir.Operation,58        role: BufferRole,59        info: shape_analysis.TensorInfo,60    ) !usize {61        if (self.value_to_slot.get(value)) |existing| return existing;6263        const total_bytes = try accounting.add(self.total_static_bytes, byteSize(info) orelse 0);64        const id = self.slots.items.len;65        var slot = try BufferSlot.init(66            self.allocator,67            id,68            value,69            producer,70            function,71            role,72            info,73        );74        errdefer slot.deinit(self.allocator);7576        try self.value_to_slot.put(value, id);77        errdefer _ = self.value_to_slot.remove(value);78        try self.slots.append(self.allocator, slot);7980        if (role.input) self.input_slot_count += 1;81        if (role.output) self.output_slot_count += 1;82        if (role.temporary) self.temporary_slot_count += 1;83        if (role.constant) self.constant_slot_count += 1;84        self.total_static_bytes = total_bytes;85        if (slot.byte_size == null) self.dynamic_slot_count += 1;86        return id;87    }8889    fn addAlias(90        self: *BufferPlanAnalysis,91        value: *ir.Value,92        source: *ir.Value,93        role: BufferRole,94    ) !usize {95        if (self.value_to_slot.get(value)) |existing| return existing;9697        const id = self.value_to_slot.get(source) orelse return error.MissingAliasedBufferSlot;98        try self.value_to_slot.put(value, id);99        self.mergeSlotRole(id, role);100        return id;101    }102103    fn mergeSlotRole(self: *BufferPlanAnalysis, id: usize, role: BufferRole) void {104        var slot = &self.slots.items[id];105        if (role.input and !slot.role.input) {106            slot.role.input = true;107            self.input_slot_count += 1;108        }109        if (role.output and !slot.role.output) {110            slot.role.output = true;111            self.output_slot_count += 1;112        }113        if (role.constant and !slot.role.constant) {114            slot.role.constant = true;115            self.constant_slot_count += 1;116        }117        if (role.temporary and !slot.role.temporary and !slot.role.input and !slot.role.output and !slot.role.constant) {118            slot.role.temporary = true;119            self.temporary_slot_count += 1;120        }121        if ((slot.role.input or slot.role.output or slot.role.constant) and slot.role.temporary) {122            slot.role.temporary = false;123            self.temporary_slot_count -= 1;124        }125    }126127    fn addElision(self: *BufferPlanAnalysis, elision: FusionElision) !void {128        if (self.value_to_elision.contains(elision.value)) return;129        const index = self.elisions.items.len;130        try self.value_to_elision.put(elision.value, index);131        errdefer _ = self.value_to_elision.remove(elision.value);132        try self.elisions.append(self.allocator, elision);133    }134};

1pub const BufferSlot = struct {2    id: usize,3    value: *ir.Value,4    producer: ?*ir.Operation,5    function: ?*ir.Operation,6    role: BufferRole,7    dtype: choir_abi.DType,8    dims: []i64,9    element_count: ?u64,10    row_major_strides: ?[]u64,11    byte_size: ?u64,1213    fn init(14        allocator: std.mem.Allocator,15        id: usize,16        value: *ir.Value,17        producer: ?*ir.Operation,18        function: ?*ir.Operation,19        role: BufferRole,20        info: shape_analysis.TensorInfo,21    ) !BufferSlot {22        const dims = try allocator.dupe(i64, info.dims);23        errdefer allocator.free(dims);2425        var strides: ?[]u64 = null;26        if (info.row_major_strides) |existing| {27            strides = try allocator.dupe(u64, existing);28            errdefer if (strides) |owned| allocator.free(owned);29        }3031        return .{32            .id = id,33            .value = value,34            .producer = producer,35            .function = function,36            .role = role,37            .dtype = info.dtype,38            .dims = dims,39            .element_count = info.element_count,40            .row_major_strides = strides,41            .byte_size = byteSize(info),42        };43    }4445    pub fn hasStaticSize(self: BufferSlot) bool {46        return self.byte_size != null;47    }4849    fn deinit(self: *BufferSlot, allocator: std.mem.Allocator) void {50        allocator.free(self.dims);51        if (self.row_major_strides) |strides| allocator.free(strides);52        self.* = undefined;53    }54};

1pub const FusionElision = struct {2    value: *ir.Value,3    producer: *ir.Operation,4    root: *ir.Operation,5    cluster_index: usize,6};

1pub const buffer_plan_analysis_descriptor = passes.AnalysisDescriptor{2    .id = passes.analysisId(buffer_plan_analysis_name),3    .name = buffer_plan_analysis_name,4    .work_contract = .{5        .identity = .{ .name = buffer_plan_analysis_name, .version = 1 },6        .estimate = bufferAnalysisWork,7    },8};

1pub const buffer_plan_analysis_name = "accy-choir-buffer-plan";

1pub fn bufferizationPlanningPass() passes.Pass {2    return .{3        .name = bufferization_planning_pass_name,4        .description = bufferization_planning_pass_description,5        .run_fn = runBufferizationPlanningPass,6        .work_contract = .{7            .identity = .{ .name = bufferization_planning_pass_name, .version = 1 },8            .estimate = bufferPassWork,9        },10    };11}

1pub const bufferization_planning_pass_description =2    "Plan Accy Choir tensor buffer slots before lowering";

1pub const bufferization_planning_pass_name = "accy-choir-plan-buffers";

1pub fn getBufferPlanAnalysis(2    pass_ctx: *passes.PassContext,3    op: *ir.Operation,4) !*BufferPlanAnalysis {5    const ptr = try pass_ctx.getAnalysis(6        op,7        &buffer_plan_analysis_descriptor,8        computeBufferPlanAnalysis,9        cleanupBufferPlanAnalysis,10    );11    return @ptrCast(@alignCast(ptr));12}

1pub const BufferRole = accy_choir.record.memory.BufferRole;

1const pass = @import("pass.zig");23pub const buffer_plan_analysis_name = pass.buffer_plan_analysis_name;4pub const bufferization_planning_pass_name = pass.bufferization_planning_pass_name;5pub const bufferization_planning_pass_description = pass.bufferization_planning_pass_description;6pub const BufferRole = pass.BufferRole;7pub const BufferSlot = pass.BufferSlot;8pub const FusionElision = pass.FusionElision;9pub const BufferPlanAnalysis = pass.BufferPlanAnalysis;10pub const buffer_plan_analysis_descriptor = pass.buffer_plan_analysis_descriptor;11pub const getBufferPlanAnalysis = pass.getBufferPlanAnalysis;12pub const bufferizationPlanningPass = pass.bufferizationPlanningPass;

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