lib/accy/src/preparation/schedule/pass.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

   1 const std = @import("std");
   2 const choir_abi = @import("choir_abi");
   3 const choir = @import("choir");
   4 const accy_root = @import("../../root.zig");
   5 const accy_choir = @import("../../choir/root.zig");
   6 const kernel_library = @import("../../kernel/library/root.zig");
   7 const dialect_mod = accy_choir.dialect;
   8 const fusion = @import("../fusion/root.zig");
   9 const shape_analysis = @import("../shape/root.zig");
  10 
  11 const semantics = accy_choir.semantics;
  12 const ir = choir.ir;
  13 const passes = choir.passes;
  14 const accounting = passes.pass.work;
  15 
  16 pub const schedule_plan_analysis_name = "accy-choir-schedule-plan";
  17 pub const schedule_planning_pass_name = "accy-choir-plan-schedule";
  18 pub const schedule_planning_pass_description =
  19     "Plan static Accy Choir tensor work items before kernel outlining";
  20 
  21 pub const ScheduleWorkKind = accy_choir.record.dispatch.ScheduleWorkKind;
  22 
  23 pub const ScheduleResourceEstimate = accy_choir.record.dispatch.ScheduleResourceEstimate;
  24 
  25 pub const ScheduleWorkItem = struct {
  26     id: usize,
  27     kind: ScheduleWorkKind,
  28     root: *ir.Operation,
  29     ops: []*ir.Operation,
  30     output_value: *ir.Value,
  31     dtype: choir_abi.DType,
  32     rank: usize,
  33     element_count: u64,
  34     resources: ScheduleResourceEstimate,
  35 
  36     pub fn opCount(self: ScheduleWorkItem) usize {
  37         return self.ops.len;
  38     }
  39 
  40     fn deinit(self: *ScheduleWorkItem, allocator: std.mem.Allocator) void {
  41         allocator.free(self.ops);
  42         self.* = undefined;
  43     }
  44 };
  45 
  46 pub const SchedulePlanAnalysis = struct {
  47     allocator: std.mem.Allocator,
  48     work_items: std.ArrayListUnmanaged(ScheduleWorkItem),
  49     root_to_item: std.AutoHashMap(*ir.Operation, usize),
  50     single_work_count: usize = 0,
  51     fusion_work_count: usize = 0,
  52     kernel_call_work_count: usize = 0,
  53     scheduled_op_count: usize = 0,
  54     total_static_elements: u64 = 0,
  55 
  56     pub fn init(allocator: std.mem.Allocator) SchedulePlanAnalysis {
  57         return .{
  58             .allocator = allocator,
  59             .work_items = .empty,
  60             .root_to_item = std.AutoHashMap(*ir.Operation, usize).init(allocator),
  61         };
  62     }
  63 
  64     pub fn deinit(self: *SchedulePlanAnalysis) void {
  65         for (self.work_items.items) |*item| {
  66             item.deinit(self.allocator);
  67         }
  68         self.work_items.deinit(self.allocator);
  69         self.root_to_item.deinit();
  70         self.* = undefined;
  71     }
  72 
  73     pub fn workItemCount(self: SchedulePlanAnalysis) usize {
  74         return self.work_items.items.len;
  75     }
  76 
  77     pub fn getWorkForRoot(
  78         self: *const SchedulePlanAnalysis,
  79         root: *ir.Operation,
  80     ) ?*const ScheduleWorkItem {
  81         const index = self.root_to_item.get(root) orelse return null;
  82         return &self.work_items.items[index];
  83     }
  84 
  85     fn addWorkItem(
  86         self: *SchedulePlanAnalysis,
  87         kind: ScheduleWorkKind,
  88         root: *ir.Operation,
  89         ops: []const *ir.Operation,
  90         output_value: *ir.Value,
  91         info: shape_analysis.TensorInfo,
  92         shapes: *const shape_analysis.ShapeLayoutAnalysis,
  93     ) !void {
  94         if (ops.len == 0) return;
  95         if (self.root_to_item.contains(root)) return;
  96         const element_count = info.element_count orelse return;
  97         const total_elements = try accounting.add(self.total_static_elements, element_count);
  98 
  99         const owned_ops = try self.allocator.alloc(*ir.Operation, ops.len);
 100         errdefer self.allocator.free(owned_ops);
 101         @memcpy(owned_ops, ops);
 102 
 103         const id = self.work_items.items.len;
 104         const resources = try resourceEstimateForWork(
 105             self.allocator,
 106             ops,
 107             info,
 108             shapes,
 109         );
 110         try self.root_to_item.put(root, id);
 111         errdefer _ = self.root_to_item.remove(root);
 112         try self.work_items.append(self.allocator, .{
 113             .id = id,
 114             .kind = kind,
 115             .root = root,
 116             .ops = owned_ops,
 117             .output_value = output_value,
 118             .dtype = info.dtype,
 119             .rank = info.rank(),
 120             .element_count = element_count,
 121             .resources = resources,
 122         });
 123 
 124         switch (kind) {
 125             .elementwise_single => self.single_work_count += 1,
 126             .elementwise_fusion => self.fusion_work_count += 1,
 127             .shape => self.single_work_count += 1,
 128             .dot_general => self.single_work_count += 1,
 129             .reduction => self.single_work_count += 1,
 130             .kernel_call => self.kernel_call_work_count += 1,
 131             .row_pipeline => self.fusion_work_count += 1,
 132             .iterate => self.single_work_count += 1,
 133             .flash_attention => self.fusion_work_count += 1,
 134             .scan => self.single_work_count += 1,
 135         }
 136         self.scheduled_op_count += ops.len;
 137         self.total_static_elements = total_elements;
 138     }
 139 };
 140 
 141 const ScheduleWork = struct {
 142     input: accounting.Census,
 143     catalog: u64,
 144 
 145     fn inspect(op: *ir.Operation) !ScheduleWork {
 146         var calls: u64 = 0;
 147         _ = try op.walk(.{ .order = .pre_order }, &calls, countCatalogCalls);
 148         return .{
 149             .input = try accounting.Census.inspect(op),
 150             .catalog = if (calls == 0) 0 else try accounting.multiply(calls, try catalogUnits()),
 151         };
 152     }
 153 
 154     fn countCatalogCalls(count: *u64, op: *ir.Operation) !ir.Operation.WalkResult {
 155         if (isKernelCallOp(op)) count.* = try accounting.add(count.*, 1);
 156         return .advance;
 157     }
 158 
 159     fn catalogUnits() !u64 {
 160         var units: u64 = 1;
 161         for (kernel_library.catalog.descriptors) |descriptor| {
 162             units = try accounting.add(
 163                 units,
 164                 try accounting.add(1, descriptor.metadata.target.len),
 165             );
 166             const specialization = descriptor.metadata.specialization;
 167             for (specialization.outputs) |shape| {
 168                 units = try accounting.add(units, try accounting.add(1, shape.axes.len));
 169             }
 170             var reuses: u64 = 1;
 171             for (specialization.reduction_reuse) |reuse| {
 172                 reuses = try accounting.add(
 173                     reuses,
 174                     try accounting.add(reuse.reduction.len, reuse.shape.axes.len),
 175                 );
 176             }
 177             for (specialization.reductions) |reduction| {
 178                 const own = try accounting.add(reduction.name.len, reduction.shape.axes.len);
 179                 units = try accounting.add(
 180                     units,
 181                     try accounting.add(try accounting.add(1, own), reuses),
 182                 );
 183             }
 184         }
 185         return units;
 186     }
 187 
 188     fn storage(self: ScheduleWork) !u64 {
 189         const count = self.input.operations;
 190         const pairs = try accounting.multiply(count, count);
 191         const external = try accounting.add(self.input.values, self.input.operands);
 192         var bytes: u64 = @sizeOf(SchedulePlanAnalysis) + @alignOf(SchedulePlanAnalysis);
 193         bytes = try accounting.add(bytes, try accounting.arrayListGrowth(ScheduleWorkItem, count));
 194         const indexes = try accounting.hashMapGrowth(*ir.Operation, usize, count);
 195         bytes = try accounting.add(bytes, try accounting.multiply(2, indexes));
 196         bytes = try accounting.add(bytes, try accounting.hashMapGrowth(*ir.Operation, void, count));
 197         bytes = try accounting.add(bytes, try accounting.hashMapGrowth(*ir.Value, usize, count));
 198         const unique = try accounting.hashMapGrowth(*ir.Value, void, external);
 199         bytes = try accounting.add(bytes, try accounting.multiply(count, unique));
 200         bytes = try accounting.add(bytes, try accounting.multiply(pairs, @sizeOf(*ir.Operation)));
 201         bytes = try accounting.add(bytes, try accounting.multiply(count, @alignOf(*ir.Operation)));
 202         const sort_element = @sizeOf(usize) + 2 * @sizeOf(ScheduleWorkItem) + @sizeOf(bool);
 203         bytes = try accounting.add(bytes, try accounting.multiply(count, sort_element));
 204         bytes = try accounting.add(bytes, 4 * @alignOf(ScheduleWorkItem));
 205         bytes = try accounting.add(
 206             bytes,
 207             try accounting.arrayListGrowth(ScheduleDependency, pairs),
 208         );
 209         if (bytes > std.math.maxInt(usize)) return error.WorkOverflow;
 210         return bytes;
 211     }
 212 
 213     fn bounds(self: ScheduleWork) !accounting.Bounds {
 214         const count = self.input.operations;
 215         const pairs = try accounting.multiply(count, count);
 216         const storage_bytes = try self.storage();
 217         const bytes = try accounting.add(self.input.input_bytes, self.catalog);
 218         const units = try accounting.add(try accounting.add(self.input.atoms, bytes), 1);
 219         const entries = try accounting.add(
 220             count,
 221             try accounting.add(self.input.values, self.input.operands),
 222         );
 223         const probes = try accounting.hashMapCapacity(entries);
 224         const inner = try accounting.add(
 225             try accounting.add(pairs, count),
 226             try accounting.add(probes, units),
 227         );
 228         const visits = try accounting.multiply(
 229             try accounting.multiply(units, try accounting.add(count, 1)),
 230             inner,
 231         );
 232         return .{
 233             .work = .{
 234                 .input_bytes = bytes,
 235                 .structural_visits = try accounting.multiply(64, visits),
 236                 .analysis_computations = 1,
 237                 .allocation_capacity = storage_bytes,
 238             },
 239             .workspace = storage_bytes,
 240             .retained_storage = storage_bytes,
 241         };
 242     }
 243 };
 244 
 245 fn scheduleAnalysisWork(input: accounting.Input) !accounting.Bounds {
 246     return (try ScheduleWork.inspect(input.operation)).bounds();
 247 }
 248 
 249 fn schedulePassWork(_: accounting.Input) !accounting.Bounds {
 250     return .{ .work = .{ .structural_visits = 1 } };
 251 }
 252 
 253 pub const schedule_plan_analysis_descriptor = passes.AnalysisDescriptor{
 254     .id = passes.analysisId(schedule_plan_analysis_name),
 255     .name = schedule_plan_analysis_name,
 256     .work_contract = .{
 257         .identity = .{ .name = schedule_plan_analysis_name, .version = 1 },
 258         .estimate = scheduleAnalysisWork,
 259     },
 260 };
 261 
 262 pub fn getSchedulePlanAnalysis(
 263     pass_ctx: *passes.PassContext,
 264     op: *ir.Operation,
 265 ) !*SchedulePlanAnalysis {
 266     const ptr = try pass_ctx.getAnalysis(
 267         op,
 268         &schedule_plan_analysis_descriptor,
 269         computeSchedulePlanAnalysis,
 270         cleanupSchedulePlanAnalysis,
 271     );
 272     return @ptrCast(@alignCast(ptr));
 273 }
 274 
 275 pub fn schedulePlanningPass() passes.Pass {
 276     return .{
 277         .name = schedule_planning_pass_name,
 278         .description = schedule_planning_pass_description,
 279         .run_fn = runSchedulePlanningPass,
 280         .work_contract = .{
 281             .identity = .{ .name = schedule_planning_pass_name, .version = 1 },
 282             .estimate = schedulePassWork,
 283         },
 284     };
 285 }
 286 
 287 fn runSchedulePlanningPass(pass_ctx: *passes.PassContext) passes.PassResult {
 288     _ = getSchedulePlanAnalysis(pass_ctx, pass_ctx.op) catch return .failure;
 289     pass_ctx.preserveAllAnalyses();
 290     return .success;
 291 }
 292 
 293 fn computeSchedulePlanAnalysis(
 294     pass_ctx: *passes.PassContext,
 295     op: *ir.Operation,
 296 ) anyerror!*anyopaque {
 297     const shapes = try shape_analysis.getShapeLayoutAnalysis(pass_ctx, op);
 298     const fusion_plan = try fusion.getFusionPlanAnalysis(pass_ctx, op);
 299 
 300     const analysis = try pass_ctx.allocator.create(SchedulePlanAnalysis);
 301     analysis.* = SchedulePlanAnalysis.init(pass_ctx.allocator);
 302     errdefer {
 303         analysis.deinit();
 304         pass_ctx.allocator.destroy(analysis);
 305     }
 306 
 307     var claimed = std.AutoHashMap(*ir.Operation, void).init(pass_ctx.allocator);
 308     defer claimed.deinit();
 309 
 310     try addFusionWorkItems(fusion_plan, shapes, analysis, &claimed);
 311     try addStandaloneWorkItems(op, shapes, analysis, &claimed);
 312     try addBroadcastWorkItems(op, shapes, analysis, &claimed);
 313     try orderWorkItemsByDependencies(pass_ctx.allocator, analysis);
 314 
 315     return @ptrCast(analysis);
 316 }
 317 
 318 fn cleanupSchedulePlanAnalysis(ptr: *anyopaque, allocator: std.mem.Allocator) void {
 319     const analysis: *SchedulePlanAnalysis = @ptrCast(@alignCast(ptr));
 320     analysis.deinit();
 321     allocator.destroy(analysis);
 322 }
 323 
 324 fn addFusionWorkItems(
 325     fusion_plan: *const fusion.FusionPlanAnalysis,
 326     shapes: *const shape_analysis.ShapeLayoutAnalysis,
 327     analysis: *SchedulePlanAnalysis,
 328     claimed: *std.AutoHashMap(*ir.Operation, void),
 329 ) !void {
 330     for (fusion_plan.clusters.items) |cluster| {
 331         const root = cluster.root() orelse continue;
 332         const info = staticResultInfo(root, shapes) orelse continue;
 333         const result = root.getResult(0) orelse continue;
 334 
 335         try analysis.addWorkItem(
 336             switch (cluster.kind) {
 337                 .elementwise => .elementwise_fusion,
 338                 .dot_epilogue => .dot_general,
 339                 .reduction_input => .reduction,
 340                 .row_pipeline => .row_pipeline,
 341                 .flash_attention => .flash_attention,
 342             },
 343             root,
 344             cluster.ops,
 345             result,
 346             info,
 347             shapes,
 348         );
 349 
 350         for (cluster.ops) |op| {
 351             try claimed.put(op, {});
 352         }
 353     }
 354 
 355     for (fusion_plan.elided.items) |op| {
 356         try claimed.put(op, {});
 357     }
 358 }
 359 
 360 fn addStandaloneWorkItems(
 361     op: *ir.Operation,
 362     shapes: *const shape_analysis.ShapeLayoutAnalysis,
 363     analysis: *SchedulePlanAnalysis,
 364     claimed: *std.AutoHashMap(*ir.Operation, void),
 365 ) !void {
 366     if (fusion.isFusableElementwiseOp(op) and !claimed.contains(op)) {
 367         if (staticResultInfo(op, shapes)) |info| {
 368             if (op.getResult(0)) |result| {
 369                 try analysis.addWorkItem(
 370                     .elementwise_single,
 371                     op,
 372                     &.{op},
 373                     result,
 374                     info,
 375                     shapes,
 376                 );
 377                 try claimed.put(op, {});
 378             }
 379         }
 380     }
 381 
 382     if (isDotGeneralOp(op) and !claimed.contains(op)) {
 383         if (staticResultInfo(op, shapes)) |info| {
 384             if (op.getResult(0)) |result| {
 385                 try analysis.addWorkItem(
 386                     .dot_general,
 387                     op,
 388                     &.{op},
 389                     result,
 390                     info,
 391                     shapes,
 392                 );
 393                 try claimed.put(op, {});
 394             }
 395         }
 396     }
 397 
 398     if (isShapeKernelOp(op) and !claimed.contains(op)) {
 399         if (staticResultInfo(op, shapes)) |info| {
 400             if (op.getResult(0)) |result| {
 401                 try analysis.addWorkItem(
 402                     .shape,
 403                     op,
 404                     &.{op},
 405                     result,
 406                     info,
 407                     shapes,
 408                 );
 409                 try claimed.put(op, {});
 410             }
 411         }
 412     }
 413 
 414     if (isReduceOp(op) and !claimed.contains(op)) {
 415         if (staticResultInfo(op, shapes)) |info| {
 416             if (op.getResult(0)) |result| {
 417                 try analysis.addWorkItem(
 418                     .reduction,
 419                     op,
 420                     &.{op},
 421                     result,
 422                     info,
 423                     shapes,
 424                 );
 425                 try claimed.put(op, {});
 426             }
 427         }
 428     }
 429 
 430     if (isCumsumOp(op) and !claimed.contains(op)) {
 431         if (staticResultInfo(op, shapes)) |info| {
 432             if (op.getResult(0)) |result| {
 433                 try analysis.addWorkItem(
 434                     .scan,
 435                     op,
 436                     &.{op},
 437                     result,
 438                     info,
 439                     shapes,
 440                 );
 441                 try claimed.put(op, {});
 442             }
 443         }
 444     }
 445 
 446     if (isIterateOp(op) and !claimed.contains(op)) {
 447         if (staticResultInfo(op, shapes)) |info| {
 448             if (iterateOutputResult(op)) |result| {
 449                 try analysis.addWorkItem(
 450                     .iterate,
 451                     op,
 452                     &.{op},
 453                     result,
 454                     info,
 455                     shapes,
 456                 );
 457                 try claimed.put(op, {});
 458             }
 459         }
 460     }
 461 
 462     if (isKernelCallOp(op) and !claimed.contains(op)) {
 463         if (op.getNumResults() == 1) {
 464             if (staticResultInfo(op, shapes)) |info| {
 465                 if (op.getResult(0)) |result| {
 466                     try analysis.addWorkItem(
 467                         .kernel_call,
 468                         op,
 469                         &.{op},
 470                         result,
 471                         info,
 472                         shapes,
 473                     );
 474                     try claimed.put(op, {});
 475                 }
 476             }
 477         }
 478     }
 479 
 480     if (fusion.regionsAreOpaque(op)) return;
 481 
 482     for (op.regions.items) |*region| {
 483         var block_iter = region.getBlocks();
 484         while (block_iter.next()) |block| {
 485             var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
 486             while (current) |current_op| {
 487                 try addStandaloneWorkItems(
 488                     current_op,
 489                     shapes,
 490                     analysis,
 491                     claimed,
 492                 );
 493                 current = current_op.next_op;
 494             }
 495         }
 496     }
 497 }
 498 
 499 fn isDotGeneralOp(op: *ir.Operation) bool {
 500     return std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.DotGeneralOp.operation_name);
 501 }
 502 
 503 fn isReduceOp(op: *ir.Operation) bool {
 504     return std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ReduceOp.operation_name);
 505 }
 506 
 507 fn iterateOutputResult(op: *ir.Operation) ?*ir.Value {
 508     var result_index: usize = 0;
 509     while (op.getResult(result_index)) |result| : (result_index += 1) {
 510         if (result.first_use != null) return result;
 511     }
 512     return op.getResult(0);
 513 }
 514 
 515 fn isCumsumOp(op: *ir.Operation) bool {
 516     return std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.CumsumOp.operation_name);
 517 }
 518 
 519 fn isIterateOp(op: *ir.Operation) bool {
 520     return std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.IterateOp.operation_name);
 521 }
 522 
 523 fn isKernelCallOp(op: *ir.Operation) bool {
 524     return std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.KernelCallOp.operation_name);
 525 }
 526 
 527 fn isShapeKernelOp(op: *ir.Operation) bool {
 528     return std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.IotaOp.operation_name) or
 529         std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ReshapeOp.operation_name) or
 530         std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.TransposeOp.operation_name) or
 531         std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.SliceOp.operation_name) or
 532         std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.PadOp.operation_name) or
 533         std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ConcatenateOp.operation_name) or
 534         std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.GatherOp.operation_name) or
 535         std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ScatterOp.operation_name);
 536 }
 537 
 538 fn addBroadcastWorkItems(
 539     op: *ir.Operation,
 540     shapes: *const shape_analysis.ShapeLayoutAnalysis,
 541     analysis: *SchedulePlanAnalysis,
 542     claimed: *std.AutoHashMap(*ir.Operation, void),
 543 ) !void {
 544     const initial_work_count = analysis.workItemCount();
 545     for (0..initial_work_count) |index| {
 546         const work = analysis.work_items.items[index];
 547         for (work.ops) |member| {
 548             for (member.getOperandValues()) |operand| {
 549                 const input = kernelInputValueForOperand(work, operand);
 550                 try materializeBroadcast(input, shapes, analysis, claimed);
 551             }
 552         }
 553     }
 554     var state = ReturnedBroadcastState{
 555         .shapes = shapes,
 556         .analysis = analysis,
 557         .claimed = claimed,
 558     };
 559     _ = try op.walk(.{ .order = .pre_order }, &state, ReturnedBroadcastState.visit);
 560 }
 561 
 562 const ReturnedBroadcastState = struct {
 563     shapes: *const shape_analysis.ShapeLayoutAnalysis,
 564     analysis: *SchedulePlanAnalysis,
 565     claimed: *std.AutoHashMap(*ir.Operation, void),
 566 
 567     fn visit(self: *ReturnedBroadcastState, op: *ir.Operation) !ir.WalkResult {
 568         if (fusion.regionsAreOpaque(op)) return .skip;
 569         if (std.mem.eql(u8, op.name.name, "func.return") or
 570             std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ReturnOp.operation_name))
 571         {
 572             for (op.getOperandValues()) |result| {
 573                 try materializeBroadcast(result, self.shapes, self.analysis, self.claimed);
 574             }
 575         }
 576         return .advance;
 577     }
 578 };
 579 
 580 fn materializeBroadcast(
 581     value: *ir.Value,
 582     shapes: *const shape_analysis.ShapeLayoutAnalysis,
 583     analysis: *SchedulePlanAnalysis,
 584     claimed: *std.AutoHashMap(*ir.Operation, void),
 585 ) !void {
 586     var current = value;
 587     while (current.getDefiningOp()) |defining| {
 588         const op: *ir.Operation = @ptrCast(@alignCast(defining));
 589         const broadcast_name = dialect_mod.AccyDialect.BroadcastInDimOp.operation_name;
 590         if (!std.mem.eql(u8, op.name.name, broadcast_name)) {
 591             return;
 592         }
 593         if (claimed.contains(op)) return;
 594         const info = staticResultInfo(op, shapes) orelse return;
 595         try analysis.addWorkItem(.shape, op, &.{op}, current, info, shapes);
 596         try claimed.put(op, {});
 597         current = op.getOperand(0) orelse return error.InvalidBroadcast;
 598     }
 599 }
 600 
 601 fn staticResultInfo(
 602     op: *ir.Operation,
 603     shapes: *const shape_analysis.ShapeLayoutAnalysis,
 604 ) ?shape_analysis.TensorInfo {
 605     const result = op.getResult(0) orelse return null;
 606     const info = shapes.get(result) orelse return null;
 607     if (!info.hasStaticLayout()) return null;
 608     return info;
 609 }
 610 
 611 fn resourceEstimateForWork(
 612     allocator: std.mem.Allocator,
 613     ops: []const *ir.Operation,
 614     info: shape_analysis.TensorInfo,
 615     shapes: *const shape_analysis.ShapeLayoutAnalysis,
 616 ) !ScheduleResourceEstimate {
 617     const element_count = info.element_count orelse return error.UnsupportedOperation;
 618     const resource_element_count = reductionInputElementCount(ops, shapes) orelse element_count;
 619     const element_size: u64 = @intCast(info.dtype.sizeOf());
 620     var resources = ScheduleResourceEstimate{
 621         .element_count = resource_element_count,
 622         .element_size = element_size,
 623         .op_count = ops.len,
 624         .estimated_element_ops = estimatedElementOpsForWork(ops, element_count, shapes),
 625     };
 626 
 627     var unique_external_values = std.AutoHashMap(*ir.Value, void).init(allocator);
 628     defer unique_external_values.deinit();
 629 
 630     if (tensorByteSize(info)) |bytes| {
 631         resources.static_write_bytes = saturatedAdd(resources.static_write_bytes, bytes);
 632     } else {
 633         resources.static_bytes_complete = false;
 634     }
 635 
 636     for (ops) |op| {
 637         for (op.getOperandValues(), 0..) |operand, operand_index| {
 638             if (isConstantReduceInitOperand(op, operand, operand_index)) continue;
 639             if (valueProducedByWork(operand, ops)) {
 640                 resources.chain_operand_count += 1;
 641                 continue;
 642             }
 643             resources.external_operand_count += 1;
 644             try putUniqueExternalValue(&unique_external_values, operand, &resources);
 645             if (!operandReads(op, operand_index)) continue;
 646             const operand_info = shapes.get(operand) orelse {
 647                 resources.static_bytes_complete = false;
 648                 continue;
 649             };
 650             if (tensorByteSize(operand_info)) |bytes| {
 651                 resources.static_read_bytes = saturatedAdd(resources.static_read_bytes, bytes);
 652             } else {
 653                 resources.static_bytes_complete = false;
 654             }
 655         }
 656     }
 657 
 658     resources.static_total_bytes = saturatedAdd(resources.static_read_bytes, resources.static_write_bytes);
 659     return resources;
 660 }
 661 
 662 fn operandReads(op: *ir.Operation, operand_index: usize) bool {
 663     if (!isKernelCallOp(op)) return true;
 664     const attr = op.getAttr("operand_effects") orelse return true;
 665     if (!std.mem.eql(u8, attr.abstract.name, dialect_mod.AccyDialect.KernelCallOp.operand_effects_attr_name)) return true;
 666     const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return true;
 667     if (operand_index >= dialect_attr.payload.len) return true;
 668     const effect = semantics.KernelOperandEffect.fromByte(dialect_attr.payload[operand_index]) orelse return true;
 669     return switch (effect) {
 670         .read, .read_write, .unknown => true,
 671         .none, .write => false,
 672     };
 673 }
 674 
 675 fn isReduceInitOperand(op: *ir.Operation, operand_index: usize) bool {
 676     return isReduceOp(op) and operand_index == 1;
 677 }
 678 
 679 fn isConstantReduceInitOperand(
 680     op: *ir.Operation,
 681     operand: *ir.Value,
 682     operand_index: usize,
 683 ) bool {
 684     if (!isReduceInitOperand(op, operand_index)) return false;
 685     const def_any = operand.getDefiningOp() orelse return false;
 686     const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
 687     return std.mem.eql(u8, def_op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name);
 688 }
 689 
 690 fn estimatedElementOpsForWork(
 691     ops: []const *ir.Operation,
 692     fallback_element_count: u64,
 693     shapes: *const shape_analysis.ShapeLayoutAnalysis,
 694 ) u64 {
 695     if (ops.len == 1 and isDotGeneralOp(ops[0])) {
 696         return dotGeneralElementOps(ops[0], shapes) orelse fallback_element_count;
 697     }
 698     if (ops.len == 1 and isReduceOp(ops[0])) {
 699         return reductionInputElementCount(ops, shapes) orelse fallback_element_count;
 700     }
 701     if (ops.len == 1 and isKernelCallOp(ops[0])) {
 702         return kernelCallElementOps(ops[0]) orelse fallback_element_count;
 703     }
 704     return saturatedMul(fallback_element_count, @intCast(ops.len));
 705 }
 706 
 707 fn kernelCallElementOps(op: *ir.Operation) ?u64 {
 708     const descriptor = kernelCallLibraryDescriptor(op) orelse return null;
 709     return descriptor.metadata.specialization.estimatedElementOps();
 710 }
 711 
 712 fn kernelCallLibraryDescriptor(op: *ir.Operation) ?kernel_library.CatalogDescriptor {
 713     const target = kernelCallTarget(op) orelse return null;
 714     const version = kernelCallVersion(op) orelse return null;
 715     return kernel_library.findEntry(target, version);
 716 }
 717 
 718 fn kernelCallTarget(op: *ir.Operation) ?[]const u8 {
 719     const attr = op.getAttr("target") orelse return null;
 720     if (!std.mem.eql(u8, attr.abstract.name, dialect_mod.AccyDialect.KernelCallOp.target_attr_name)) return null;
 721     const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return null;
 722     return dialect_attr.payload;
 723 }
 724 
 725 fn kernelCallVersion(op: *ir.Operation) ?u32 {
 726     const attr = op.getAttrAs(ir.Attribute.IntegerAttr, "version") orelse return null;
 727     const value = attr.getValue();
 728     if (value < 0 or value > std.math.maxInt(u32)) return null;
 729     return @intCast(value);
 730 }
 731 
 732 fn reductionInputElementCount(
 733     ops: []const *ir.Operation,
 734     shapes: *const shape_analysis.ShapeLayoutAnalysis,
 735 ) ?u64 {
 736     if (ops.len != 1 or !isReduceOp(ops[0])) return null;
 737     const operands = ops[0].getOperandValues();
 738     if (operands.len < 1) return null;
 739     const input_info = shapes.get(operands[0]) orelse return null;
 740     return input_info.element_count;
 741 }
 742 
 743 fn dotGeneralElementOps(
 744     op: *ir.Operation,
 745     shapes: *const shape_analysis.ShapeLayoutAnalysis,
 746 ) ?u64 {
 747     if (op.getNumOperands() != 2) return null;
 748     const operands = op.getOperandValues();
 749     if (operands.len != 2) return null;
 750     const lhs = shapes.get(operands[0]) orelse return null;
 751     const rhs = shapes.get(operands[1]) orelse return null;
 752     const output = op.getResult(0) orelse return null;
 753     const output_info = shapes.get(output) orelse return null;
 754     if (lhs.dims.len != 2 or rhs.dims.len != 2 or output_info.dims.len != 2) return null;
 755     if (lhs.dims[0] < 0 or
 756         lhs.dims[1] < 0 or
 757         rhs.dims[0] < 0 or
 758         rhs.dims[1] < 0 or
 759         output_info.dims[0] < 0 or
 760         output_info.dims[1] < 0)
 761     {
 762         return null;
 763     }
 764     const m: u64 = @intCast(lhs.dims[0]);
 765     const k: u64 = @intCast(lhs.dims[1]);
 766     const n: u64 = @intCast(rhs.dims[1]);
 767     return saturatedMul(saturatedMul(saturatedMul(m, n), k), 2);
 768 }
 769 
 770 fn valueProducedByWork(
 771     value: *ir.Value,
 772     ops: []const *ir.Operation,
 773 ) bool {
 774     const def_any = value.getDefiningOp() orelse return false;
 775     const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
 776     for (ops) |op| {
 777         if (op == def_op) return true;
 778     }
 779     return false;
 780 }
 781 
 782 fn tensorByteSize(info: shape_analysis.TensorInfo) ?u64 {
 783     const element_count = info.element_count orelse return null;
 784     const element_size: u64 = @intCast(info.dtype.sizeOf());
 785     return std.math.mul(u64, element_count, element_size) catch null;
 786 }
 787 
 788 fn putUniqueExternalValue(
 789     unique_external_values: *std.AutoHashMap(*ir.Value, void),
 790     value: *ir.Value,
 791     resources: *ScheduleResourceEstimate,
 792 ) !void {
 793     if (unique_external_values.contains(value)) return;
 794     try unique_external_values.put(value, {});
 795     resources.external_input_value_count += 1;
 796 }
 797 
 798 fn saturatedAdd(lhs: u64, rhs: u64) u64 {
 799     return std.math.add(u64, lhs, rhs) catch std.math.maxInt(u64);
 800 }
 801 
 802 fn saturatedMul(lhs: u64, rhs: u64) u64 {
 803     return std.math.mul(u64, lhs, rhs) catch std.math.maxInt(u64);
 804 }
 805 
 806 const ScheduleDependency = struct {
 807     from: usize,
 808     to: usize,
 809 };
 810 
 811 fn orderWorkItemsByDependencies(
 812     allocator: std.mem.Allocator,
 813     analysis: *SchedulePlanAnalysis,
 814 ) !void {
 815     const work_items = analysis.work_items.items;
 816     if (work_items.len < 2) return;
 817 
 818     var output_to_item = std.AutoHashMap(*ir.Value, usize).init(allocator);
 819     defer output_to_item.deinit();
 820     for (work_items, 0..) |work, index| {
 821         try output_to_item.put(work.output_value, index);
 822     }
 823 
 824     const indegree = try allocator.alloc(usize, work_items.len);
 825     defer allocator.free(indegree);
 826     @memset(indegree, 0);
 827 
 828     var edges: std.ArrayListUnmanaged(ScheduleDependency) = .empty;
 829     defer edges.deinit(allocator);
 830 
 831     for (work_items, 0..) |work, consumer_index| {
 832         for (work.ops) |op| {
 833             for (op.getOperandValues()) |operand| {
 834                 const producer_value = kernelInputValueForOperand(work, operand);
 835                 const producer_index = output_to_item.get(producer_value) orelse continue;
 836                 if (producer_index == consumer_index) continue;
 837                 try appendScheduleDependency(
 838                     allocator,
 839                     &edges,
 840                     indegree,
 841                     producer_index,
 842                     consumer_index,
 843                 );
 844             }
 845         }
 846     }
 847 
 848     if (edges.items.len == 0) return;
 849 
 850     const original = try allocator.alloc(ScheduleWorkItem, work_items.len);
 851     defer allocator.free(original);
 852     @memcpy(original, work_items);
 853 
 854     const ordered = try allocator.alloc(ScheduleWorkItem, work_items.len);
 855     defer allocator.free(ordered);
 856 
 857     const emitted = try allocator.alloc(bool, work_items.len);
 858     defer allocator.free(emitted);
 859     @memset(emitted, false);
 860 
 861     var ordered_count: usize = 0;
 862     while (ordered_count < work_items.len) {
 863         const next_index = nextReadyWorkItem(indegree, emitted) orelse {
 864             return error.CyclicScheduleDependency;
 865         };
 866         ordered[ordered_count] = original[next_index];
 867         ordered_count += 1;
 868         emitted[next_index] = true;
 869 
 870         for (edges.items) |edge| {
 871             if (edge.from != next_index) continue;
 872             std.debug.assert(indegree[edge.to] > 0);
 873             indegree[edge.to] -= 1;
 874         }
 875     }
 876 
 877     @memcpy(work_items, ordered);
 878     try rebuildWorkIndexMap(analysis);
 879 }
 880 
 881 pub fn kernelInputValueForOperand(
 882     work: ScheduleWorkItem,
 883     operand: *ir.Value,
 884 ) *ir.Value {
 885     switch (work.kind) {
 886         .elementwise_single,
 887         .elementwise_fusion,
 888         .dot_general,
 889         .reduction,
 890         .row_pipeline,
 891         .iterate,
 892         .flash_attention,
 893         => {
 894             const def_any = operand.getDefiningOp() orelse return operand;
 895             const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
 896             const broadcast_name = dialect_mod.AccyDialect.BroadcastInDimOp.operation_name;
 897             if (!std.mem.eql(u8, def_op.name.name, broadcast_name)) {
 898                 return operand;
 899             }
 900             const operands = def_op.getOperandValues();
 901             if (operands.len != 1) return operand;
 902             return operands[0];
 903         },
 904         else => return operand,
 905     }
 906 }
 907 
 908 fn appendScheduleDependency(
 909     allocator: std.mem.Allocator,
 910     edges: *std.ArrayListUnmanaged(ScheduleDependency),
 911     indegree: []usize,
 912     from: usize,
 913     to: usize,
 914 ) !void {
 915     for (edges.items) |edge| {
 916         if (edge.from == from and edge.to == to) return;
 917     }
 918     try edges.append(allocator, .{
 919         .from = from,
 920         .to = to,
 921     });
 922     indegree[to] += 1;
 923 }
 924 
 925 fn nextReadyWorkItem(indegree: []const usize, emitted: []const bool) ?usize {
 926     for (indegree, 0..) |incoming, index| {
 927         if (!emitted[index] and incoming == 0) return index;
 928     }
 929     return null;
 930 }
 931 
 932 fn rebuildWorkIndexMap(analysis: *SchedulePlanAnalysis) !void {
 933     var rebuilt = std.AutoHashMap(*ir.Operation, usize).init(analysis.allocator);
 934     errdefer rebuilt.deinit();
 935     for (analysis.work_items.items, 0..) |work, index| {
 936         try rebuilt.put(work.root, index);
 937     }
 938     analysis.root_to_item.deinit();
 939     analysis.root_to_item = rebuilt;
 940 }
 941 
 942 const testing = std.testing;
 943 const semantic = accy_choir.semantic;
 944 
 945 fn findOpNamedInBlock(block: *ir.Block, name: []const u8) ?*ir.Operation {
 946     return findNthOpNamedInBlock(block, name, 0);
 947 }
 948 
 949 fn findNthOpNamedInBlock(block: *ir.Block, name: []const u8, needle_index: usize) ?*ir.Operation {
 950     var seen: usize = 0;
 951     var iter = block.operations.head;
 952     while (iter) |op_ptr| {
 953         const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
 954         if (std.mem.eql(u8, op.name.name, name)) {
 955             if (seen == needle_index) return op;
 956             seen += 1;
 957         }
 958         iter = op.next_op;
 959     }
 960     return null;
 961 }
 962 
 963 test "schedule planning emits one work item for an elementwise fusion cluster" {
 964     const allocator = testing.allocator;
 965 
 966     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
 967     defer builder.deinit();
 968     const f32_4 = try builder.tensor(.f32, &.{4});
 969     var fb = try builder.beginFunction("schedule_fused_add_mul", &.{ f32_4, f32_4, f32_4 }, &.{f32_4});
 970     const sum = try fb.add(fb.parameter(0), fb.parameter(1));
 971     const product = try fb.mul(sum, fb.parameter(2));
 972     try fb.return_(&.{product});
 973     try fb.finish();
 974     const module = try builder.finish();
 975     defer module.deinit();
 976 
 977     const choir_mod = module.choir_module;
 978     const ctx = module.context();
 979     const ledger = try scheduleTestLedger();
 980     defer ledger.destroy();
 981     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
 982     defer cache.deinit();
 983     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
 984     defer pass_ctx.deinit();
 985 
 986     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
 987     const func = ir.inspection.functionByNameInBlock(body, "schedule_fused_add_mul") orelse return error.TestExpectedFunc;
 988     const entry = func.getRegion(0).?.getEntryBlock().?;
 989     const mul = findOpNamedInBlock(entry, dialect_mod.AccyDialect.MulOp.operation_name) orelse return error.TestExpectedMul;
 990 
 991     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
 992     try ledger.producersComplete();
 993     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
 994     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
 995     try testing.expectEqual(@as(usize, 0), analysis.single_work_count);
 996     try testing.expectEqual(@as(usize, 1), analysis.fusion_work_count);
 997     try testing.expectEqual(@as(usize, 2), analysis.scheduled_op_count);
 998     try testing.expectEqual(@as(u64, 4), analysis.total_static_elements);
 999 
1000     const item = analysis.getWorkForRoot(mul) orelse return error.TestExpectedWorkItem;
1001     try testing.expectEqual(ScheduleWorkKind.elementwise_fusion, item.kind);
1002     try testing.expectEqual(@as(usize, 2), item.opCount());
1003     try testing.expectEqual(choir_abi.DType.f32, item.dtype);
1004     try testing.expectEqual(@as(usize, 1), item.rank);
1005     try testing.expectEqual(@as(u64, 4), item.element_count);
1006     try testing.expectEqual(@as(u64, 4), item.resources.element_count);
1007     try testing.expectEqual(@as(u64, 4), item.resources.element_size);
1008     try testing.expectEqual(@as(usize, 2), item.resources.op_count);
1009     try testing.expectEqual(@as(usize, 3), item.resources.external_input_value_count);
1010     try testing.expectEqual(@as(usize, 3), item.resources.external_operand_count);
1011     try testing.expectEqual(@as(usize, 1), item.resources.chain_operand_count);
1012     try testing.expectEqual(@as(u64, 48), item.resources.static_read_bytes);
1013     try testing.expectEqual(@as(u64, 16), item.resources.static_write_bytes);
1014     try testing.expectEqual(@as(u64, 64), item.resources.static_total_bytes);
1015     try testing.expectEqual(@as(u64, 8), item.resources.estimated_element_ops);
1016     try testing.expectEqual(@as(u64, 128), item.resources.elementOpsPerKiB());
1017     try testing.expect(item.resources.static_bytes_complete);
1018 
1019     try testing.expectEqual(mul.getResult(0).?, item.output_value);
1020 }
1021 
1022 test "schedule planning emits standalone work when fusion is rejected" {
1023     const allocator = testing.allocator;
1024 
1025     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1026     defer builder.deinit();
1027     const f32_4 = try builder.tensor(.f32, &.{4});
1028     var fb = try builder.beginFunction("schedule_escape", &.{ f32_4, f32_4, f32_4 }, &.{ f32_4, f32_4 });
1029     const sum = try fb.add(fb.parameter(0), fb.parameter(1));
1030     const product = try fb.mul(sum, fb.parameter(2));
1031     try fb.return_(&.{ sum, product });
1032     try fb.finish();
1033     const module = try builder.finish();
1034     defer module.deinit();
1035 
1036     const choir_mod = module.choir_module;
1037     const ctx = module.context();
1038     const ledger = try scheduleTestLedger();
1039     defer ledger.destroy();
1040     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1041     defer cache.deinit();
1042     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1043     defer pass_ctx.deinit();
1044 
1045     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1046     const func = ir.inspection.functionByNameInBlock(body, "schedule_escape") orelse return error.TestExpectedFunc;
1047     const entry = func.getRegion(0).?.getEntryBlock().?;
1048     const add = findOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name) orelse return error.TestExpectedAdd;
1049     const mul = findOpNamedInBlock(entry, dialect_mod.AccyDialect.MulOp.operation_name) orelse return error.TestExpectedMul;
1050 
1051     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1052     try ledger.producersComplete();
1053     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1054     try testing.expectEqual(@as(usize, 2), analysis.workItemCount());
1055     try testing.expectEqual(@as(usize, 2), analysis.single_work_count);
1056     try testing.expectEqual(@as(usize, 0), analysis.fusion_work_count);
1057     try testing.expectEqual(@as(usize, 2), analysis.scheduled_op_count);
1058     try testing.expectEqual(@as(u64, 8), analysis.total_static_elements);
1059     const add_work = analysis.getWorkForRoot(add) orelse return error.TestExpectedWorkItem;
1060     const mul_work = analysis.getWorkForRoot(mul) orelse return error.TestExpectedWorkItem;
1061     try testing.expectEqual(ScheduleWorkKind.elementwise_single, add_work.kind);
1062     try testing.expectEqual(ScheduleWorkKind.elementwise_single, mul_work.kind);
1063     try testing.expectEqual(@as(usize, 2), add_work.resources.external_input_value_count);
1064     try testing.expectEqual(@as(usize, 2), add_work.resources.external_operand_count);
1065     try testing.expectEqual(@as(u64, 32), add_work.resources.static_read_bytes);
1066     try testing.expectEqual(@as(u64, 16), add_work.resources.static_write_bytes);
1067     try testing.expectEqual(@as(u64, 48), add_work.resources.static_total_bytes);
1068     try testing.expectEqual(@as(u64, 4), add_work.resources.estimated_element_ops);
1069     try testing.expectEqual(@as(usize, 2), mul_work.resources.external_input_value_count);
1070     try testing.expectEqual(@as(usize, 2), mul_work.resources.external_operand_count);
1071     try testing.expectEqual(@as(usize, 0), mul_work.resources.chain_operand_count);
1072     try testing.expectEqual(@as(u64, 48), mul_work.resources.static_total_bytes);
1073 }
1074 
1075 test "schedule planning emits dot_general work for static rank-2 matmul" {
1076     const allocator = testing.allocator;
1077 
1078     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1079     defer builder.deinit();
1080     const f32_16x16 = try builder.tensor(.f32, &.{ 16, 16 });
1081     var fb = try builder.beginFunction("schedule_dot_general_matmul", &.{ f32_16x16, f32_16x16 }, &.{f32_16x16});
1082     const dot = try fb.dotGeneral(
1083         fb.parameter(0),
1084         fb.parameter(1),
1085         f32_16x16,
1086         &.{1},
1087         &.{0},
1088         &.{},
1089         &.{},
1090     );
1091     try fb.return_(&.{dot});
1092     try fb.finish();
1093     const module = try builder.finish();
1094     defer module.deinit();
1095 
1096     const choir_mod = module.choir_module;
1097     const ctx = module.context();
1098     const ledger = try scheduleTestLedger();
1099     defer ledger.destroy();
1100     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1101     defer cache.deinit();
1102     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1103     defer pass_ctx.deinit();
1104 
1105     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1106     const func = ir.inspection.functionByNameInBlock(body, "schedule_dot_general_matmul") orelse return error.TestExpectedFunc;
1107     const entry = func.getRegion(0).?.getEntryBlock().?;
1108     const dot_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.DotGeneralOp.operation_name) orelse return error.TestExpectedDotGeneral;
1109 
1110     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1111     try ledger.producersComplete();
1112     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1113     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
1114     try testing.expectEqual(@as(usize, 1), analysis.single_work_count);
1115     try testing.expectEqual(@as(usize, 0), analysis.fusion_work_count);
1116     try testing.expectEqual(@as(usize, 1), analysis.scheduled_op_count);
1117     try testing.expectEqual(@as(u64, 256), analysis.total_static_elements);
1118 
1119     const item = analysis.getWorkForRoot(dot_op) orelse return error.TestExpectedWorkItem;
1120     try testing.expectEqual(ScheduleWorkKind.dot_general, item.kind);
1121     try testing.expectEqual(@as(usize, 1), item.opCount());
1122     try testing.expectEqual(choir_abi.DType.f32, item.dtype);
1123     try testing.expectEqual(@as(usize, 2), item.rank);
1124     try testing.expectEqual(@as(u64, 256), item.element_count);
1125     try testing.expectEqual(@as(u64, 256), item.resources.element_count);
1126     try testing.expectEqual(@as(u64, 4), item.resources.element_size);
1127     try testing.expectEqual(@as(usize, 1), item.resources.op_count);
1128     try testing.expectEqual(@as(usize, 2), item.resources.external_input_value_count);
1129     try testing.expectEqual(@as(usize, 2), item.resources.external_operand_count);
1130     try testing.expectEqual(@as(usize, 0), item.resources.chain_operand_count);
1131     try testing.expectEqual(@as(u64, 2048), item.resources.static_read_bytes);
1132     try testing.expectEqual(@as(u64, 1024), item.resources.static_write_bytes);
1133     try testing.expectEqual(@as(u64, 3072), item.resources.static_total_bytes);
1134     try testing.expectEqual(@as(u64, 8192), item.resources.estimated_element_ops);
1135     try testing.expectEqual(@as(u64, 2730), item.resources.elementOpsPerKiB());
1136     try testing.expect(item.resources.static_bytes_complete);
1137 }
1138 
1139 test "schedule planning emits reduction work for static rank-1 reduce" {
1140     const allocator = testing.allocator;
1141 
1142     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1143     defer builder.deinit();
1144     const f32_256 = try builder.tensor(.f32, &.{256});
1145     const f32_scalar = try builder.tensor(.f32, &.{});
1146     var fb = try builder.beginFunction("schedule_reduce_sum_rank1", &.{f32_256}, &.{f32_scalar});
1147     const zero_value: f32 = 0.0;
1148     const zero = try fb.constant(f32_scalar, std.mem.asBytes(&zero_value));
1149     const reduced = try fb.reduce(fb.parameter(0), zero, f32_scalar, "sum", &.{0});
1150     try fb.return_(&.{reduced});
1151     try fb.finish();
1152     const module = try builder.finish();
1153     defer module.deinit();
1154 
1155     const choir_mod = module.choir_module;
1156     const ctx = module.context();
1157     const ledger = try scheduleTestLedger();
1158     defer ledger.destroy();
1159     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1160     defer cache.deinit();
1161     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1162     defer pass_ctx.deinit();
1163 
1164     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1165     const func = ir.inspection.functionByNameInBlock(body, "schedule_reduce_sum_rank1") orelse return error.TestExpectedFunc;
1166     const entry = func.getRegion(0).?.getEntryBlock().?;
1167     const reduce_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.ReduceOp.operation_name) orelse return error.TestExpectedReduce;
1168 
1169     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1170     try ledger.producersComplete();
1171     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1172     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
1173     try testing.expectEqual(@as(usize, 1), analysis.single_work_count);
1174     try testing.expectEqual(@as(usize, 0), analysis.fusion_work_count);
1175     try testing.expectEqual(@as(usize, 1), analysis.scheduled_op_count);
1176     try testing.expectEqual(@as(u64, 1), analysis.total_static_elements);
1177 
1178     const item = analysis.getWorkForRoot(reduce_op) orelse return error.TestExpectedWorkItem;
1179     try testing.expectEqual(ScheduleWorkKind.reduction, item.kind);
1180     try testing.expectEqual(@as(usize, 1), item.opCount());
1181     try testing.expectEqual(choir_abi.DType.f32, item.dtype);
1182     try testing.expectEqual(@as(usize, 0), item.rank);
1183     try testing.expectEqual(@as(u64, 1), item.element_count);
1184     try testing.expectEqual(@as(u64, 256), item.resources.element_count);
1185     try testing.expectEqual(@as(u64, 4), item.resources.element_size);
1186     try testing.expectEqual(@as(usize, 1), item.resources.op_count);
1187     try testing.expectEqual(@as(usize, 1), item.resources.external_input_value_count);
1188     try testing.expectEqual(@as(usize, 1), item.resources.external_operand_count);
1189     try testing.expectEqual(@as(usize, 0), item.resources.chain_operand_count);
1190     try testing.expectEqual(@as(u64, 1024), item.resources.static_read_bytes);
1191     try testing.expectEqual(@as(u64, 4), item.resources.static_write_bytes);
1192     try testing.expectEqual(@as(u64, 1028), item.resources.static_total_bytes);
1193     try testing.expectEqual(@as(u64, 256), item.resources.estimated_element_ops);
1194     try testing.expectEqual(@as(u64, 255), item.resources.elementOpsPerKiB());
1195     try testing.expect(item.resources.static_bytes_complete);
1196 }
1197 
1198 test "schedule planning emits kernel_call work with operand effects" {
1199     const allocator = testing.allocator;
1200 
1201     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1202     defer builder.deinit();
1203     const f32_4 = try builder.tensor(.f32, &.{4});
1204     var fb = try builder.beginFunction("schedule_kernel_call", &.{ f32_4, f32_4 }, &.{f32_4});
1205     const call = try fb.kernelCall(
1206         &.{ fb.parameter(0), fb.parameter(1) },
1207         &.{f32_4},
1208         .{
1209             .target = "accy.custom.scale",
1210             .operand_effects = &.{ .read, .write },
1211             .result_aliases = &.{null},
1212         },
1213     );
1214     try fb.return_(&.{call.getFirstResult()});
1215     try fb.finish();
1216     const module = try builder.finish();
1217     defer module.deinit();
1218 
1219     const choir_mod = module.choir_module;
1220     const ctx = module.context();
1221     const ledger = try scheduleTestLedger();
1222     defer ledger.destroy();
1223     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1224     defer cache.deinit();
1225     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1226     defer pass_ctx.deinit();
1227 
1228     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1229     const func = ir.inspection.functionByNameInBlock(body, "schedule_kernel_call") orelse return error.TestExpectedFunc;
1230     const entry = func.getRegion(0).?.getEntryBlock().?;
1231     const call_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse return error.TestExpectedKernelCall;
1232 
1233     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1234     try ledger.producersComplete();
1235     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1236     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
1237     try testing.expectEqual(@as(usize, 0), analysis.single_work_count);
1238     try testing.expectEqual(@as(usize, 0), analysis.fusion_work_count);
1239     try testing.expectEqual(@as(usize, 1), analysis.kernel_call_work_count);
1240     try testing.expectEqual(@as(usize, 1), analysis.scheduled_op_count);
1241     try testing.expectEqual(@as(u64, 4), analysis.total_static_elements);
1242 
1243     const item = analysis.getWorkForRoot(call_op) orelse return error.TestExpectedWorkItem;
1244     try testing.expectEqual(ScheduleWorkKind.kernel_call, item.kind);
1245     try testing.expectEqual(@as(usize, 1), item.opCount());
1246     try testing.expectEqual(choir_abi.DType.f32, item.dtype);
1247     try testing.expectEqual(@as(usize, 1), item.rank);
1248     try testing.expectEqual(@as(u64, 4), item.element_count);
1249     try testing.expectEqual(@as(u64, 4), item.resources.element_count);
1250     try testing.expectEqual(@as(u64, 4), item.resources.element_size);
1251     try testing.expectEqual(@as(usize, 1), item.resources.op_count);
1252     try testing.expectEqual(@as(usize, 2), item.resources.external_input_value_count);
1253     try testing.expectEqual(@as(usize, 2), item.resources.external_operand_count);
1254     try testing.expectEqual(@as(usize, 0), item.resources.chain_operand_count);
1255     try testing.expectEqual(@as(u64, 16), item.resources.static_read_bytes);
1256     try testing.expectEqual(@as(u64, 16), item.resources.static_write_bytes);
1257     try testing.expectEqual(@as(u64, 32), item.resources.static_total_bytes);
1258     try testing.expectEqual(@as(u64, 4), item.resources.estimated_element_ops);
1259     try testing.expectEqual(call.getFirstResult(), item.output_value);
1260 }
1261 
1262 test "schedule planning estimates catalog kernel_call reductions" {
1263     const allocator = testing.allocator;
1264 
1265     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1266     defer builder.deinit();
1267     const lhs_ty = try builder.tensor(.f32, &.{ 4, 8 });
1268     const rhs_ty = try builder.tensor(.f32, &.{ 8, 16 });
1269     const out_ty = try builder.tensor(.f32, &.{ 4, 16 });
1270     var fb = try builder.beginFunction("schedule_catalog_kernel_call_matmul", &.{ lhs_ty, rhs_ty }, &.{out_ty});
1271     const call = try fb.kernelCall(
1272         &.{ fb.parameter(0), fb.parameter(1) },
1273         &.{out_ty},
1274         .{
1275             .target = kernel_library.linalg.MatrixProduct4x16x8F32.target,
1276             .operand_effects = &.{ .read, .read },
1277             .result_aliases = &.{null},
1278         },
1279     );
1280     try fb.return_(&.{call.getFirstResult()});
1281     try fb.finish();
1282     const module = try builder.finish();
1283     defer module.deinit();
1284 
1285     const choir_mod = module.choir_module;
1286     const ctx = module.context();
1287     const ledger = try scheduleTestLedger();
1288     defer ledger.destroy();
1289     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1290     defer cache.deinit();
1291     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1292     defer pass_ctx.deinit();
1293 
1294     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1295     const func = ir.inspection.functionByNameInBlock(body, "schedule_catalog_kernel_call_matmul") orelse return error.TestExpectedFunc;
1296     const entry = func.getRegion(0).?.getEntryBlock().?;
1297     const call_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse return error.TestExpectedKernelCall;
1298 
1299     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1300     try ledger.producersComplete();
1301     const declared = try ScheduleWork.inspect(choir_mod);
1302     try testing.expect(declared.catalog >= kernel_library.catalog.descriptors.len);
1303     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1304     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
1305     try testing.expectEqual(@as(usize, 1), analysis.kernel_call_work_count);
1306 
1307     const item = analysis.getWorkForRoot(call_op) orelse return error.TestExpectedWorkItem;
1308     try testing.expectEqual(ScheduleWorkKind.kernel_call, item.kind);
1309     try testing.expectEqual(@as(u64, 64), item.element_count);
1310     try testing.expectEqual(@as(u64, 64), item.resources.element_count);
1311     try testing.expectEqual(@as(u64, 640), item.resources.static_read_bytes);
1312     try testing.expectEqual(@as(u64, 256), item.resources.static_write_bytes);
1313     try testing.expectEqual(@as(u64, 896), item.resources.static_total_bytes);
1314     try testing.expectEqual(@as(u64, 1024), item.resources.estimated_element_ops);
1315     try testing.expectEqual(@as(u64, 1170), item.resources.elementOpsPerKiB());
1316 }
1317 
1318 test "schedule planning emits shape work for static transpose" {
1319     const allocator = testing.allocator;
1320 
1321     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1322     defer builder.deinit();
1323     const f32_16x32 = try builder.tensor(.f32, &.{ 16, 32 });
1324     const f32_32x16 = try builder.tensor(.f32, &.{ 32, 16 });
1325     var fb = try builder.beginFunction("schedule_transpose", &.{f32_16x32}, &.{f32_32x16});
1326     const transposed = try fb.transpose(fb.parameter(0), f32_32x16, &.{ 1, 0 });
1327     try fb.return_(&.{transposed});
1328     try fb.finish();
1329     const module = try builder.finish();
1330     defer module.deinit();
1331 
1332     const choir_mod = module.choir_module;
1333     const ctx = module.context();
1334     const ledger = try scheduleTestLedger();
1335     defer ledger.destroy();
1336     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1337     defer cache.deinit();
1338     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1339     defer pass_ctx.deinit();
1340 
1341     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1342     const func = ir.inspection.functionByNameInBlock(body, "schedule_transpose") orelse return error.TestExpectedFunc;
1343     const entry = func.getRegion(0).?.getEntryBlock().?;
1344     const transpose_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.TransposeOp.operation_name) orelse return error.TestExpectedTranspose;
1345 
1346     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1347     try ledger.producersComplete();
1348     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1349     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
1350     try testing.expectEqual(@as(usize, 1), analysis.single_work_count);
1351     try testing.expectEqual(@as(usize, 0), analysis.fusion_work_count);
1352 
1353     const item = analysis.getWorkForRoot(transpose_op) orelse return error.TestExpectedWorkItem;
1354     try testing.expectEqual(ScheduleWorkKind.shape, item.kind);
1355     try testing.expectEqual(@as(usize, 1), item.opCount());
1356     try testing.expectEqual(choir_abi.DType.f32, item.dtype);
1357     try testing.expectEqual(@as(usize, 2), item.rank);
1358     try testing.expectEqual(@as(u64, 512), item.element_count);
1359     try testing.expectEqual(@as(u64, 2048), item.resources.static_read_bytes);
1360     try testing.expectEqual(@as(u64, 2048), item.resources.static_write_bytes);
1361     try testing.expectEqual(@as(u64, 4096), item.resources.static_total_bytes);
1362 }
1363 
1364 test "schedule planning emits shape work for static slice" {
1365     const allocator = testing.allocator;
1366 
1367     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1368     defer builder.deinit();
1369     const f32_18x18 = try builder.tensor(.f32, &.{ 18, 18 });
1370     const f32_16x16 = try builder.tensor(.f32, &.{ 16, 16 });
1371     var fb = try builder.beginFunction("schedule_slice", &.{f32_18x18}, &.{f32_16x16});
1372     const sliced = try fb.slice(fb.parameter(0), f32_16x16, &.{ 1, 1 }, &.{ 17, 17 }, &.{ 1, 1 });
1373     try fb.return_(&.{sliced});
1374     try fb.finish();
1375     const module = try builder.finish();
1376     defer module.deinit();
1377 
1378     const choir_mod = module.choir_module;
1379     const ctx = module.context();
1380     const ledger = try scheduleTestLedger();
1381     defer ledger.destroy();
1382     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1383     defer cache.deinit();
1384     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1385     defer pass_ctx.deinit();
1386 
1387     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1388     const func = ir.inspection.functionByNameInBlock(body, "schedule_slice") orelse return error.TestExpectedFunc;
1389     const entry = func.getRegion(0).?.getEntryBlock().?;
1390     const slice_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.SliceOp.operation_name) orelse return error.TestExpectedSlice;
1391 
1392     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1393     try ledger.producersComplete();
1394     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1395     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
1396     try testing.expectEqual(@as(usize, 1), analysis.single_work_count);
1397     try testing.expectEqual(@as(usize, 0), analysis.fusion_work_count);
1398 
1399     const item = analysis.getWorkForRoot(slice_op) orelse return error.TestExpectedWorkItem;
1400     try testing.expectEqual(ScheduleWorkKind.shape, item.kind);
1401     try testing.expectEqual(@as(usize, 1), item.opCount());
1402     try testing.expectEqual(choir_abi.DType.f32, item.dtype);
1403     try testing.expectEqual(@as(usize, 2), item.rank);
1404     try testing.expectEqual(@as(u64, 256), item.element_count);
1405     try testing.expectEqual(@as(u64, 1296), item.resources.static_read_bytes);
1406     try testing.expectEqual(@as(u64, 1024), item.resources.static_write_bytes);
1407     try testing.expectEqual(@as(u64, 2320), item.resources.static_total_bytes);
1408 }
1409 
1410 test "schedule planning emits shape work for static pad" {
1411     const allocator = testing.allocator;
1412 
1413     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1414     defer builder.deinit();
1415     const f32_scalar = try builder.tensor(.f32, &.{});
1416     const f32_4x4 = try builder.tensor(.f32, &.{ 4, 4 });
1417     const f32_6x6 = try builder.tensor(.f32, &.{ 6, 6 });
1418     var fb = try builder.beginFunction("schedule_pad", &.{ f32_4x4, f32_scalar }, &.{f32_6x6});
1419     const padded = try fb.pad(fb.parameter(0), fb.parameter(1), f32_6x6, &.{ 1, 1 }, &.{ 1, 1 }, &.{ 0, 0 });
1420     try fb.return_(&.{padded});
1421     try fb.finish();
1422     const module = try builder.finish();
1423     defer module.deinit();
1424 
1425     const choir_mod = module.choir_module;
1426     const ctx = module.context();
1427     const ledger = try scheduleTestLedger();
1428     defer ledger.destroy();
1429     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1430     defer cache.deinit();
1431     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1432     defer pass_ctx.deinit();
1433 
1434     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1435     const func = ir.inspection.functionByNameInBlock(body, "schedule_pad") orelse return error.TestExpectedFunc;
1436     const entry = func.getRegion(0).?.getEntryBlock().?;
1437     const pad_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.PadOp.operation_name) orelse return error.TestExpectedPad;
1438 
1439     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1440     try ledger.producersComplete();
1441     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1442     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
1443     try testing.expectEqual(@as(usize, 1), analysis.single_work_count);
1444     try testing.expectEqual(@as(usize, 0), analysis.fusion_work_count);
1445 
1446     const item = analysis.getWorkForRoot(pad_op) orelse return error.TestExpectedWorkItem;
1447     try testing.expectEqual(ScheduleWorkKind.shape, item.kind);
1448     try testing.expectEqual(@as(usize, 1), item.opCount());
1449     try testing.expectEqual(choir_abi.DType.f32, item.dtype);
1450     try testing.expectEqual(@as(usize, 2), item.rank);
1451     try testing.expectEqual(@as(u64, 36), item.element_count);
1452     try testing.expectEqual(@as(usize, 2), item.resources.external_input_value_count);
1453     try testing.expectEqual(@as(usize, 2), item.resources.external_operand_count);
1454     try testing.expectEqual(@as(usize, 0), item.resources.chain_operand_count);
1455     try testing.expectEqual(@as(u64, 68), item.resources.static_read_bytes);
1456     try testing.expectEqual(@as(u64, 144), item.resources.static_write_bytes);
1457     try testing.expectEqual(@as(u64, 212), item.resources.static_total_bytes);
1458 }
1459 
1460 test "schedule planning emits shape work for static iota" {
1461     const allocator = testing.allocator;
1462 
1463     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1464     defer builder.deinit();
1465     const f32_4x8 = try builder.tensor(.f32, &.{ 4, 8 });
1466     var fb = try builder.beginFunction("schedule_iota", &.{}, &.{f32_4x8});
1467     const iota = try fb.iota(f32_4x8, 1);
1468     try fb.return_(&.{iota});
1469     try fb.finish();
1470     const module = try builder.finish();
1471     defer module.deinit();
1472 
1473     const choir_mod = module.choir_module;
1474     const ctx = module.context();
1475     const ledger = try scheduleTestLedger();
1476     defer ledger.destroy();
1477     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1478     defer cache.deinit();
1479     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1480     defer pass_ctx.deinit();
1481 
1482     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1483     const func = ir.inspection.functionByNameInBlock(body, "schedule_iota") orelse return error.TestExpectedFunc;
1484     const entry = func.getRegion(0).?.getEntryBlock().?;
1485     const iota_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.IotaOp.operation_name) orelse return error.TestExpectedIota;
1486 
1487     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1488     try ledger.producersComplete();
1489     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1490     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
1491     try testing.expectEqual(@as(usize, 1), analysis.single_work_count);
1492     try testing.expectEqual(@as(usize, 0), analysis.fusion_work_count);
1493 
1494     const item = analysis.getWorkForRoot(iota_op) orelse return error.TestExpectedWorkItem;
1495     try testing.expectEqual(ScheduleWorkKind.shape, item.kind);
1496     try testing.expectEqual(@as(usize, 1), item.opCount());
1497     try testing.expectEqual(choir_abi.DType.f32, item.dtype);
1498     try testing.expectEqual(@as(usize, 2), item.rank);
1499     try testing.expectEqual(@as(u64, 32), item.element_count);
1500     try testing.expectEqual(@as(usize, 0), item.resources.external_input_value_count);
1501     try testing.expectEqual(@as(u64, 0), item.resources.static_read_bytes);
1502     try testing.expectEqual(@as(u64, 128), item.resources.static_write_bytes);
1503     try testing.expectEqual(@as(u64, 128), item.resources.static_total_bytes);
1504 }
1505 
1506 test "schedule planning emits shape work for static concatenate" {
1507     const allocator = testing.allocator;
1508 
1509     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1510     defer builder.deinit();
1511     const f32_4x4 = try builder.tensor(.f32, &.{ 4, 4 });
1512     const f32_4x8 = try builder.tensor(.f32, &.{ 4, 8 });
1513     var fb = try builder.beginFunction("schedule_concatenate", &.{ f32_4x4, f32_4x4 }, &.{f32_4x8});
1514     const concatenated = try fb.concatenate(&.{ fb.parameter(0), fb.parameter(1) }, f32_4x8, 1);
1515     try fb.return_(&.{concatenated});
1516     try fb.finish();
1517     const module = try builder.finish();
1518     defer module.deinit();
1519 
1520     const choir_mod = module.choir_module;
1521     const ctx = module.context();
1522     const ledger = try scheduleTestLedger();
1523     defer ledger.destroy();
1524     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1525     defer cache.deinit();
1526     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1527     defer pass_ctx.deinit();
1528 
1529     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1530     const func = ir.inspection.functionByNameInBlock(body, "schedule_concatenate") orelse return error.TestExpectedFunc;
1531     const entry = func.getRegion(0).?.getEntryBlock().?;
1532     const concat_op = findOpNamedInBlock(entry, dialect_mod.AccyDialect.ConcatenateOp.operation_name) orelse return error.TestExpectedConcatenate;
1533 
1534     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1535     try ledger.producersComplete();
1536     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1537     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
1538     try testing.expectEqual(@as(usize, 1), analysis.single_work_count);
1539     try testing.expectEqual(@as(usize, 0), analysis.fusion_work_count);
1540 
1541     const item = analysis.getWorkForRoot(concat_op) orelse return error.TestExpectedWorkItem;
1542     try testing.expectEqual(ScheduleWorkKind.shape, item.kind);
1543     try testing.expectEqual(@as(usize, 1), item.opCount());
1544     try testing.expectEqual(choir_abi.DType.f32, item.dtype);
1545     try testing.expectEqual(@as(usize, 2), item.rank);
1546     try testing.expectEqual(@as(u64, 32), item.element_count);
1547     try testing.expectEqual(@as(usize, 2), item.resources.external_input_value_count);
1548     try testing.expectEqual(@as(usize, 2), item.resources.external_operand_count);
1549     try testing.expectEqual(@as(u64, 128), item.resources.static_read_bytes);
1550     try testing.expectEqual(@as(u64, 128), item.resources.static_write_bytes);
1551     try testing.expectEqual(@as(u64, 256), item.resources.static_total_bytes);
1552 }
1553 
1554 test "schedule planning orders fusion after external producers" {
1555     const allocator = testing.allocator;
1556 
1557     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1558     defer builder.deinit();
1559     const f32_4 = try builder.tensor(.f32, &.{4});
1560     var fb = try builder.beginFunction("schedule_branch_fusion_order", &.{ f32_4, f32_4, f32_4, f32_4 }, &.{f32_4});
1561     const shared = try fb.add(fb.parameter(0), fb.parameter(1));
1562     const left = try fb.mul(shared, fb.parameter(2));
1563     const right = try fb.add(shared, fb.parameter(3));
1564     const product = try fb.mul(left, right);
1565     try fb.return_(&.{product});
1566     try fb.finish();
1567     const module = try builder.finish();
1568     defer module.deinit();
1569 
1570     const choir_mod = module.choir_module;
1571     const ctx = module.context();
1572     const ledger = try scheduleTestLedger();
1573     defer ledger.destroy();
1574     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1575     defer cache.deinit();
1576     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1577     defer pass_ctx.deinit();
1578 
1579     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1580     const func = ir.inspection.functionByNameInBlock(body, "schedule_branch_fusion_order") orelse return error.TestExpectedFunc;
1581     const entry = func.getRegion(0).?.getEntryBlock().?;
1582     const shared_op = findNthOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name, 0) orelse return error.TestExpectedAdd;
1583     const left_op = findNthOpNamedInBlock(entry, dialect_mod.AccyDialect.MulOp.operation_name, 0) orelse return error.TestExpectedMul;
1584     const right_op = findNthOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name, 1) orelse return error.TestExpectedAdd;
1585     const product_op = findNthOpNamedInBlock(entry, dialect_mod.AccyDialect.MulOp.operation_name, 1) orelse return error.TestExpectedMul;
1586 
1587     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1588     try ledger.producersComplete();
1589     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1590     try testing.expectEqual(@as(usize, 1), analysis.workItemCount());
1591     try testing.expectEqual(@as(usize, 0), analysis.single_work_count);
1592     try testing.expectEqual(@as(usize, 1), analysis.fusion_work_count);
1593     try testing.expectEqual(product_op, analysis.work_items.items[0].root);
1594 
1595     const fused = analysis.getWorkForRoot(product_op) orelse return error.TestExpectedWorkItem;
1596     try testing.expectEqual(ScheduleWorkKind.elementwise_fusion, fused.kind);
1597     try testing.expectEqual(@as(usize, 4), fused.opCount());
1598     try testing.expectEqual(shared_op, fused.ops[0]);
1599     try testing.expectEqual(left_op, fused.ops[1]);
1600     try testing.expectEqual(right_op, fused.ops[2]);
1601     try testing.expectEqual(product_op, fused.ops[3]);
1602     try testing.expectEqual(@as(usize, 4), fused.resources.external_input_value_count);
1603     try testing.expectEqual(@as(usize, 4), fused.resources.external_operand_count);
1604     try testing.expectEqual(@as(usize, 4), fused.resources.chain_operand_count);
1605     try testing.expectEqual(@as(u64, 80), fused.resources.static_total_bytes);
1606 }
1607 
1608 test "schedule planning orders fusion after broadcast source producers" {
1609     const allocator = testing.allocator;
1610 
1611     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1612     defer builder.deinit();
1613     const f32_4 = try builder.tensor(.f32, &.{4});
1614     const f32_4x4 = try builder.tensor(.f32, &.{ 4, 4 });
1615     var fb = try builder.beginFunction("schedule_broadcast_source_order", &.{ f32_4x4, f32_4 }, &.{f32_4x4});
1616     const row_product = try fb.mul(fb.parameter(1), fb.parameter(1));
1617     const row_broadcast = try fb.broadcastInDim(row_product, f32_4x4, &.{ 4, 4 }, &.{0});
1618     const sum = try fb.add(fb.parameter(0), row_broadcast);
1619     const product = try fb.mul(sum, fb.parameter(0));
1620     try fb.return_(&.{product});
1621     try fb.finish();
1622     const module = try builder.finish();
1623     defer module.deinit();
1624 
1625     const choir_mod = module.choir_module;
1626     const ctx = module.context();
1627     const ledger = try scheduleTestLedger();
1628     defer ledger.destroy();
1629     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1630     defer cache.deinit();
1631     var pass_ctx = passes.PassContext.init(choir_mod, ctx, allocator, &cache);
1632     defer pass_ctx.deinit();
1633 
1634     const body = choir_mod.getRegion(0).?.getEntryBlock().?;
1635     const func = ir.inspection.functionByNameInBlock(body, "schedule_broadcast_source_order") orelse return error.TestExpectedFunc;
1636     const entry = func.getRegion(0).?.getEntryBlock().?;
1637     const row_product_op = findNthOpNamedInBlock(entry, dialect_mod.AccyDialect.MulOp.operation_name, 0) orelse return error.TestExpectedMul;
1638     const sum_op = findNthOpNamedInBlock(entry, dialect_mod.AccyDialect.AddOp.operation_name, 0) orelse return error.TestExpectedAdd;
1639     const product_op = findNthOpNamedInBlock(entry, dialect_mod.AccyDialect.MulOp.operation_name, 1) orelse return error.TestExpectedMul;
1640 
1641     const analysis = try getSchedulePlanAnalysis(&pass_ctx, choir_mod);
1642     try ledger.producersComplete();
1643     try checkScheduleStorage(choir_mod, ctx, &cache, analysis);
1644     try testing.expectEqual(@as(usize, 2), analysis.workItemCount());
1645     try testing.expectEqual(@as(usize, 1), analysis.single_work_count);
1646     try testing.expectEqual(@as(usize, 1), analysis.fusion_work_count);
1647     try testing.expectEqual(row_product_op, analysis.work_items.items[0].root);
1648     try testing.expectEqual(product_op, analysis.work_items.items[1].root);
1649 
1650     const fused = analysis.getWorkForRoot(product_op) orelse return error.TestExpectedWorkItem;
1651     try testing.expectEqual(ScheduleWorkKind.elementwise_fusion, fused.kind);
1652     try testing.expectEqual(@as(usize, 2), fused.opCount());
1653     try testing.expectEqual(sum_op, fused.ops[0]);
1654     try testing.expectEqual(product_op, fused.ops[1]);
1655 }
1656 
1657 test "schedule planning pass preserves IR" {
1658     const allocator = testing.allocator;
1659 
1660     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1661     defer builder.deinit();
1662     const f32_4 = try builder.tensor(.f32, &.{4});
1663     var fb = try builder.beginFunction("schedule_pass_add4", &.{ f32_4, f32_4 }, &.{f32_4});
1664     const sum = try fb.add(fb.parameter(0), fb.parameter(1));
1665     try fb.return_(&.{sum});
1666     try fb.finish();
1667     const module = try builder.finish();
1668     defer module.deinit();
1669 
1670     const choir_mod = module.choir_module;
1671     const ctx = module.context();
1672     var pm = passes.PassManager.init(allocator);
1673     defer pm.deinit();
1674     try pm.addPass(schedulePlanningPass());
1675 
1676     const ledger = try choir.product.revision.AccountingV1.create(allocator, .{
1677         .allowance = choir.product.revision.WorkVector.uniform(std.math.maxInt(u64)),
1678         .workspace = std.math.maxInt(u64),
1679         .events = 8,
1680     }, &.{.{ .name = schedule_planning_pass_name, .version = 1 }});
1681     defer ledger.destroy();
1682     var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 3);
1683     defer cache.deinit();
1684     try testing.expectEqual(
1685         passes.PassResult.success,
1686         pm.runWithAnalysisCache(choir_mod, ctx, &cache, .{}),
1687     );
1688     try ledger.producersComplete();
1689     try testing.expectEqual(@as(u64, 1), pm.stats.pass_runs);
1690     try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
1691 }
1692 
1693 fn scheduleTestLedger() !*choir.product.revision.AccountingV1 {
1694     const revision = choir.product.revision;
1695     return revision.AccountingV1.create(testing.allocator, .{
1696         .allowance = revision.WorkVector.uniform(std.math.maxInt(u64)),
1697         .workspace = std.math.maxInt(u64),
1698         .events = 8,
1699     }, &.{});
1700 }
1701 
1702 fn checkScheduleStorage(
1703     op: *ir.Operation,
1704     ctx: *ir.Context,
1705     cache: *passes.AnalysisCache,
1706     expected: *const SchedulePlanAnalysis,
1707 ) !void {
1708     const bounds = try scheduleAnalysisWork(.{ .operation = op });
1709     const bytes = try testing.allocator.alloc(u8, @intCast(bounds.workspace));
1710     defer testing.allocator.free(bytes);
1711     var storage = @import("alloc_fixed").Tracked.init(bytes);
1712     var pass_ctx = passes.PassContext.init(op, ctx, storage.allocator(), cache);
1713     defer pass_ctx.deinit();
1714     const ptr = try computeSchedulePlanAnalysis(&pass_ctx, op);
1715     defer cleanupSchedulePlanAnalysis(ptr, storage.allocator());
1716     const actual: *SchedulePlanAnalysis = @ptrCast(@alignCast(ptr));
1717     try testing.expectEqual(expected.workItemCount(), actual.workItemCount());
1718     try testing.expectEqual(expected.scheduled_op_count, actual.scheduled_op_count);
1719     try testing.expectEqual(expected.total_static_elements, actual.total_static_elements);
1720     for (expected.work_items.items, actual.work_items.items) |left, right| {
1721         try testing.expectEqual(left.id, right.id);
1722         try testing.expectEqual(left.kind, right.kind);
1723         try testing.expectEqual(left.root, right.root);
1724         try testing.expectEqual(left.output_value, right.output_value);
1725         try testing.expectEqualSlices(*ir.Operation, left.ops, right.ops);
1726         try testing.expectEqualDeep(left.resources, right.resources);
1727         try testing.expectEqual(right.id, actual.getWorkForRoot(right.root).?.id);
1728     }
1729     try testing.expect(!storage.exhausted);
1730     try testing.expect(storage.status().high_water_bytes <= bounds.workspace);
1731     try testing.expect(storage.status().high_water_bytes >= @sizeOf(SchedulePlanAnalysis));
1732 }
1733 
1734 test "schedule planning work contract covers growing dependency chains" {
1735     for ([_]usize{ 1, 2, 6, 7, 16, 64 }) |count| try checkScheduleChain(count);
1736     try testing.expectError(error.WorkOverflow, (ScheduleWork{
1737         .input = .{ .operations = std.math.maxInt(u64) },
1738         .catalog = 0,
1739     }).bounds());
1740     try testing.expectError(error.WorkOverflow, (ScheduleWork{
1741         .input = .{},
1742         .catalog = std.math.maxInt(u64),
1743     }).bounds());
1744 }
1745 
1746 fn checkScheduleChain(count: usize) !void {
1747     var builder = try semantic.Builder.init(
1748         testing.allocator,
1749         semantic.Builder.ContextLimits.standard,
1750     );
1751     defer builder.deinit();
1752     const typ = try builder.tensor(.f32, &.{ 2, 3 });
1753     var function = try builder.beginFunction("schedule_chain", &.{typ}, &.{typ});
1754     var result = function.parameter(0);
1755     for (0..count) |_| result = try function.reshape(result, typ, &.{ 2, 3 });
1756     try function.return_(&.{result});
1757     try function.finish();
1758     const module = try builder.finish();
1759     defer module.deinit();
1760     const ledger = try scheduleTestLedger();
1761     defer ledger.destroy();
1762     var cache = try passes.AnalysisCache.initAccounted(testing.allocator, null, ledger, .{}, 3);
1763     defer cache.deinit();
1764     var pass_ctx = passes.PassContext.init(
1765         module.choir_module,
1766         module.context(),
1767         testing.allocator,
1768         &cache,
1769     );
1770     defer pass_ctx.deinit();
1771     const analysis = try getSchedulePlanAnalysis(&pass_ctx, module.choir_module);
1772     try ledger.producersComplete();
1773     try testing.expectEqual(count, analysis.workItemCount());
1774     for (analysis.work_items.items, 0..) |item, index| {
1775         try testing.expectEqual(index, item.id);
1776         try testing.expectEqual(ScheduleWorkKind.shape, item.kind);
1777         if (index != 0) {
1778             const previous = analysis.work_items.items[index - 1];
1779             try testing.expectEqual(previous.output_value, item.root.getOperand(0).?);
1780         }
1781     }
1782     try checkScheduleStorage(module.choir_module, module.context(), &cache, analysis);
1783     try checkScheduleAdmission(module.choir_module, module.context());
1784 }
1785 
1786 fn checkScheduleAdmission(op: *ir.Operation, ctx: *ir.Context) !void {
1787     const revision = choir.product.revision;
1788     var charge: u64 = 1;
1789     for ([_]passes.AnalysisDescriptor{
1790         shape_analysis.shape_layout_analysis_descriptor,
1791         fusion.fusion_plan_analysis_descriptor,
1792         schedule_plan_analysis_descriptor,
1793     }) |descriptor| {
1794         const bounds = try descriptor.work_contract.?.estimate(.{ .operation = op });
1795         charge = try accounting.add(charge, bounds.work.structural_visits);
1796     }
1797     for ([_]i8{ -1, 0, 1 }) |offset| {
1798         var allowance = revision.WorkVector.uniform(std.math.maxInt(u64));
1799         allowance.structural_visits = @intCast(@as(i128, charge) + offset);
1800         const ledger = try revision.AccountingV1.create(testing.allocator, .{
1801             .allowance = allowance,
1802             .workspace = std.math.maxInt(u64),
1803             .events = 8,
1804         }, &.{.{ .name = schedule_planning_pass_name, .version = 1 }});
1805         defer ledger.destroy();
1806         var cache = try passes.AnalysisCache.initAccounted(
1807             testing.allocator,
1808             null,
1809             ledger,
1810             .{},
1811             3,
1812         );
1813         defer cache.deinit();
1814         var manager = passes.PassManager.init(testing.allocator);
1815         defer manager.deinit();
1816         try manager.addPass(schedulePlanningPass());
1817         const result = manager.runWithAnalysisCache(op, ctx, &cache, .{});
1818         if (offset < 0) {
1819             try testing.expectEqual(passes.PassResult.failure, result);
1820             try testing.expectEqual(revision.receipt.Outcome.exhausted, ledger.view().outcome);
1821             try testing.expectEqual(@as(usize, 1), cache.entries.count());
1822         } else {
1823             try testing.expectEqual(passes.PassResult.success, result);
1824             try ledger.producersComplete();
1825             try testing.expectEqual(@as(usize, 3), cache.entries.count());
1826         }
1827     }
1828 }
1829 
1830 test "schedule planning reports aggregate element overflow without caching a plan" {
1831     var builder = try semantic.Builder.init(
1832         testing.allocator,
1833         semantic.Builder.ContextLimits.standard,
1834     );
1835     defer builder.deinit();
1836     const extent = std.math.maxInt(i64);
1837     const typ = try builder.tensor(.f32, &.{extent});
1838     var function = try builder.beginFunction("schedule_overflow", &.{typ}, &.{typ});
1839     const first = try function.reshape(function.parameter(0), typ, &.{extent});
1840     const second = try function.reshape(first, typ, &.{extent});
1841     const third = try function.reshape(second, typ, &.{extent});
1842     try function.return_(&.{third});
1843     try function.finish();
1844     const module = try builder.finish();
1845     defer module.deinit();
1846     const ledger = try scheduleTestLedger();
1847     defer ledger.destroy();
1848     var cache = try passes.AnalysisCache.initAccounted(
1849         testing.allocator,
1850         null,
1851         ledger,
1852         .{},
1853         3,
1854     );
1855     defer cache.deinit();
1856     var pass_ctx = passes.PassContext.init(
1857         module.choir_module,
1858         module.context(),
1859         testing.allocator,
1860         &cache,
1861     );
1862     defer pass_ctx.deinit();
1863     try testing.expectError(
1864         error.WorkOverflow,
1865         getSchedulePlanAnalysis(&pass_ctx, module.choir_module),
1866     );
1867     try testing.expectEqual(
1868         choir.product.revision.receipt.Outcome.exhausted,
1869         ledger.view().outcome,
1870     );
1871     try testing.expectEqual(@as(usize, 2), cache.entries.count());
1872     try testing.expectError(
1873         error.TerminalWorkOutcome,
1874         getSchedulePlanAnalysis(&pass_ctx, module.choir_module),
1875     );
1876 }