tiny.accy.preparation.schedule

1pub const SchedulePlanAnalysis = struct {2    allocator: std.mem.Allocator,3    work_items: std.ArrayListUnmanaged(ScheduleWorkItem),4    root_to_item: std.AutoHashMap(*ir.Operation, usize),5    single_work_count: usize = 0,6    fusion_work_count: usize = 0,7    kernel_call_work_count: usize = 0,8    scheduled_op_count: usize = 0,9    total_static_elements: u64 = 0,1011    pub fn init(allocator: std.mem.Allocator) SchedulePlanAnalysis {12        return .{13            .allocator = allocator,14            .work_items = .empty,15            .root_to_item = std.AutoHashMap(*ir.Operation, usize).init(allocator),16        };17    }1819    pub fn deinit(self: *SchedulePlanAnalysis) void {20        for (self.work_items.items) |*item| {21            item.deinit(self.allocator);22        }23        self.work_items.deinit(self.allocator);24        self.root_to_item.deinit();25        self.* = undefined;26    }2728    pub fn workItemCount(self: SchedulePlanAnalysis) usize {29        return self.work_items.items.len;30    }3132    pub fn getWorkForRoot(33        self: *const SchedulePlanAnalysis,34        root: *ir.Operation,35    ) ?*const ScheduleWorkItem {36        const index = self.root_to_item.get(root) orelse return null;37        return &self.work_items.items[index];38    }3940    fn addWorkItem(41        self: *SchedulePlanAnalysis,42        kind: ScheduleWorkKind,43        root: *ir.Operation,44        ops: []const *ir.Operation,45        output_value: *ir.Value,46        info: shape_analysis.TensorInfo,47        shapes: *const shape_analysis.ShapeLayoutAnalysis,48    ) !void {49        if (ops.len == 0) return;50        if (self.root_to_item.contains(root)) return;51        const element_count = info.element_count orelse return;52        const total_elements = try accounting.add(self.total_static_elements, element_count);5354        const owned_ops = try self.allocator.alloc(*ir.Operation, ops.len);55        errdefer self.allocator.free(owned_ops);56        @memcpy(owned_ops, ops);5758        const id = self.work_items.items.len;59        const resources = try resourceEstimateForWork(60            self.allocator,61            ops,62            info,63            shapes,64        );65        try self.root_to_item.put(root, id);66        errdefer _ = self.root_to_item.remove(root);67        try self.work_items.append(self.allocator, .{68            .id = id,69            .kind = kind,70            .root = root,71            .ops = owned_ops,72            .output_value = output_value,73            .dtype = info.dtype,74            .rank = info.rank(),75            .element_count = element_count,76            .resources = resources,77        });7879        switch (kind) {80            .elementwise_single => self.single_work_count += 1,81            .elementwise_fusion => self.fusion_work_count += 1,82            .shape => self.single_work_count += 1,83            .dot_general => self.single_work_count += 1,84            .reduction => self.single_work_count += 1,85            .kernel_call => self.kernel_call_work_count += 1,86            .row_pipeline => self.fusion_work_count += 1,87            .iterate => self.single_work_count += 1,88            .flash_attention => self.fusion_work_count += 1,89            .scan => self.single_work_count += 1,90        }91        self.scheduled_op_count += ops.len;92        self.total_static_elements = total_elements;93    }94};

1pub const ScheduleWorkItem = struct {2    id: usize,3    kind: ScheduleWorkKind,4    root: *ir.Operation,5    ops: []*ir.Operation,6    output_value: *ir.Value,7    dtype: choir_abi.DType,8    rank: usize,9    element_count: u64,10    resources: ScheduleResourceEstimate,1112    pub fn opCount(self: ScheduleWorkItem) usize {13        return self.ops.len;14    }1516    fn deinit(self: *ScheduleWorkItem, allocator: std.mem.Allocator) void {17        allocator.free(self.ops);18        self.* = undefined;19    }20};

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

1pub fn kernelInputValueForOperand(2    work: ScheduleWorkItem,3    operand: *ir.Value,4) *ir.Value {5    switch (work.kind) {6        .elementwise_single,7        .elementwise_fusion,8        .dot_general,9        .reduction,10        .row_pipeline,11        .iterate,12        .flash_attention,13        => {14            const def_any = operand.getDefiningOp() orelse return operand;15            const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));16            const broadcast_name = dialect_mod.AccyDialect.BroadcastInDimOp.operation_name;17            if (!std.mem.eql(u8, def_op.name.name, broadcast_name)) {18                return operand;19            }20            const operands = def_op.getOperandValues();21            if (operands.len != 1) return operand;22            return operands[0];23        },24        else => return operand,25    }26}

1pub fn schedulePlanningPass() passes.Pass {2    return .{3        .name = schedule_planning_pass_name,4        .description = schedule_planning_pass_description,5        .run_fn = runSchedulePlanningPass,6        .work_contract = .{7            .identity = .{ .name = schedule_planning_pass_name, .version = 1 },8            .estimate = schedulePassWork,9        },10    };11}

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

1pub const schedule_plan_analysis_name = "accy-choir-schedule-plan";

1pub const schedule_planning_pass_description =2    "Plan static Accy Choir tensor work items before kernel outlining";

1pub const ScheduleResourceEstimate = accy_choir.record.dispatch.ScheduleResourceEstimate;

1pub const ScheduleWorkKind = accy_choir.record.dispatch.ScheduleWorkKind;

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

1const pass = @import("pass.zig");23pub const schedule_plan_analysis_name = pass.schedule_plan_analysis_name;4pub const schedule_planning_pass_name = pass.schedule_planning_pass_name;5pub const schedule_planning_pass_description = pass.schedule_planning_pass_description;6pub const ScheduleWorkKind = pass.ScheduleWorkKind;7pub const ScheduleResourceEstimate = pass.ScheduleResourceEstimate;8pub const ScheduleWorkItem = pass.ScheduleWorkItem;9pub const SchedulePlanAnalysis = pass.SchedulePlanAnalysis;10pub const schedule_plan_analysis_descriptor = pass.schedule_plan_analysis_descriptor;11pub const getSchedulePlanAnalysis = pass.getSchedulePlanAnalysis;12pub const schedulePlanningPass = pass.schedulePlanningPass;13pub const kernelInputValueForOperand = pass.kernelInputValueForOperand;