lib/accy/src/preparation/kernelization/lowering/dot.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

   1 const std = @import("std");
   2 const gpu = @import("gpu");
   3 const choir_abi = @import("choir_abi");
   4 const generated_abi = @import("abi.zig");
   5 const generated_builder = @import("builder.zig");
   6 const common = @import("common.zig");
   7 const generated_elementwise = @import("elementwise.zig");
   8 const generated_name = @import("name.zig");
   9 const generated_schedule = @import("schedule.zig");
  10 
  11 const fusion = common.fusion;
  12 
  13 const ir = common.ir;
  14 const dialect_mod = common.dialect_mod;
  15 const kernel_root = common.kernel_root;
  16 const bufferization = common.bufferization;
  17 const kernelization_model = @import("../model/root.zig");
  18 const schedule_planning = common.schedule_planning;
  19 const bufferSlotById = common.bufferSlotById;
  20 const staticPositiveDimU32 = common.staticPositiveDimU32;
  21 const isName = common.isName;
  22 
  23 const DotGeneralStaticDims = kernelization_model.DotGeneralStaticDims;
  24 const DotGeneralDescription = kernelization_model.DotGeneralDescription;
  25 const LoweredKernel = kernelization_model.LoweredKernel;
  26 
  27 fn dotInitialValue(builder: anytype, dtype: choir_abi.DType) !kernel_root.Value {
  28     return switch (dtype) {
  29         .f32 => builder.constantFloat(.f32, 0.0),
  30         .f16 => builder.constantFloat(.f32, 0.0),
  31         .i32 => builder.constantInt(.i32, 0),
  32         .u32 => builder.constantInt(.u32, 0),
  33         else => error.CapabilityMismatch,
  34     };
  35 }
  36 
  37 fn hostLoopFormat(format: ?gpu.ArtifactFormat) bool {
  38     const artifact_format = format orelse return false;
  39     return gpu.artifactFormatUsesHostLoopLaunch(artifact_format);
  40 }
  41 
  42 pub fn dotGeneralDescriptionForWork(
  43     outline: kernelization_model.KernelOutline,
  44     work: schedule_planning.ScheduleWorkItem,
  45     buffer_plan: *const bufferization.BufferPlanAnalysis,
  46 ) gpu.BackendError!DotGeneralDescription {
  47     if (work.kind != .dot_general) return error.UnsupportedOperation;
  48     if (work.ops.len == 0) return error.UnsupportedOperation;
  49     if (!isName(work.ops[0].name.name, dialect_mod.AccyDialect.DotGeneralOp.operation_name)) return error.UnsupportedOperation;
  50     if (work.ops.len == 1 and outline.inputCount() != 2) return error.UnsupportedOperation;
  51     if (outline.inputCount() < 2) return error.UnsupportedOperation;
  52 
  53     const lhs_slot = bufferSlotById(buffer_plan, outline.input_slot_ids[0]) orelse return error.InvalidArtifact;
  54     const rhs_slot = bufferSlotById(buffer_plan, outline.input_slot_ids[1]) orelse return error.InvalidArtifact;
  55     const output_slot = bufferSlotById(buffer_plan, outline.output_slot_id) orelse return error.InvalidArtifact;
  56     try validateCanonicalDotLayout(work.ops[0], lhs_slot.dims.len);
  57     const dims = try dotGeneralStaticDims(lhs_slot.*, rhs_slot.*, output_slot.*);
  58     try validateDotGeneralDTypes(lhs_slot.*, rhs_slot.*, output_slot.*, work);
  59 
  60     return .{
  61         .input_dtype = lhs_slot.dtype,
  62         .output_dtype = output_slot.dtype,
  63         .dims = dims,
  64     };
  65 }
  66 
  67 pub const max_epilogue_inputs = 8;
  68 
  69 const EpilogueBroadcastOperand = struct {
  70     param: usize,
  71     op: *ir.Operation,
  72     slot: *const bufferization.BufferSlot,
  73 };
  74 
  75 const EpilogueOperand = union(enum) {
  76     chain,
  77     full: usize,
  78     broadcast: EpilogueBroadcastOperand,
  79 };
  80 
  81 const EpilogueOp = struct {
  82     kind: kernelization_model.ElementwiseKernel,
  83     op: *ir.Operation,
  84     operands: [3]EpilogueOperand,
  85     operand_count: u8,
  86 };
  87 
  88 const EpiloguePlan = struct {
  89     ops: [fusion.max_dot_epilogue_ops]EpilogueOp = undefined,
  90     count: usize = 0,
  91     param_dtypes: [max_epilogue_inputs]choir_abi.DType = undefined,
  92     param_count: usize = 0,
  93 
  94     fn chain(self: *const EpiloguePlan) []const EpilogueOp {
  95         return self.ops[0..self.count];
  96     }
  97 };
  98 
  99 fn buildEpiloguePlan(
 100     work: schedule_planning.ScheduleWorkItem,
 101     buffer_plan: *const bufferization.BufferPlanAnalysis,
 102 ) common.LoweringError!EpiloguePlan {
 103     var plan = EpiloguePlan{};
 104     if (work.ops.len <= 1) return plan;
 105     if (work.ops.len - 1 > plan.ops.len) return error.UnsupportedOperation;
 106 
 107     var param_values: [max_epilogue_inputs]*ir.Value = undefined;
 108 
 109     const dot = work.ops[0];
 110     for (dot.getOperandValues()) |operand| {
 111         _ = try epilogueParamIndex(&plan, &param_values, buffer_plan, operand);
 112     }
 113 
 114     var previous = dot.getResult(0) orelse return error.InvalidArtifact;
 115     for (work.ops[1..]) |op| {
 116         const kind = generated_elementwise.kernelForOperation(op) orelse return error.UnsupportedOperation;
 117         const operands = op.getOperandValues();
 118         if (operands.len > 3) return error.UnsupportedOperation;
 119         var entry = EpilogueOp{
 120             .kind = kind,
 121             .op = op,
 122             .operands = undefined,
 123             .operand_count = @intCast(operands.len),
 124         };
 125         for (operands, 0..) |operand, operand_index| {
 126             if (operand == previous) {
 127                 entry.operands[operand_index] = .chain;
 128                 continue;
 129             }
 130             if (fusion.broadcastLeafSource(operand)) |source| {
 131                 const def_any = operand.getDefiningOp() orelse return error.InvalidArtifact;
 132                 const broadcast_op: *ir.Operation = @ptrCast(@alignCast(def_any));
 133                 const param = try epilogueParamIndex(&plan, &param_values, buffer_plan, source);
 134                 const slot = buffer_plan.getSlot(source) orelse return error.InvalidArtifact;
 135                 entry.operands[operand_index] = .{ .broadcast = .{ .param = param, .op = broadcast_op, .slot = slot } };
 136                 continue;
 137             }
 138             const param = try epilogueParamIndex(&plan, &param_values, buffer_plan, operand);
 139             entry.operands[operand_index] = .{ .full = param };
 140         }
 141         plan.ops[plan.count] = entry;
 142         plan.count += 1;
 143         previous = op.getResult(0) orelse return error.InvalidArtifact;
 144     }
 145     return plan;
 146 }
 147 
 148 fn epilogueParamIndex(
 149     plan: *EpiloguePlan,
 150     param_values: *[max_epilogue_inputs]*ir.Value,
 151     buffer_plan: *const bufferization.BufferPlanAnalysis,
 152     value: *ir.Value,
 153 ) common.LoweringError!usize {
 154     for (param_values[0..plan.param_count], 0..) |existing, index| {
 155         if (existing == value) return index;
 156     }
 157     if (plan.param_count >= max_epilogue_inputs) return error.UnsupportedOperation;
 158     const slot = buffer_plan.getSlot(value) orelse return error.UnsupportedOperation;
 159     param_values[plan.param_count] = value;
 160     plan.param_dtypes[plan.param_count] = slot.dtype;
 161     plan.param_count += 1;
 162     return plan.param_count - 1;
 163 }
 164 
 165 fn applyEpilogueChain(
 166     logical: anytype,
 167     epilogue: []const EpilogueOp,
 168     abi: anytype,
 169     initial: kernel_root.Value,
 170     flat_index: kernel_root.Value,
 171 ) !kernel_root.Value {
 172     var current = initial;
 173     for (epilogue) |entry| {
 174         var inputs: [3]kernel_root.Value = undefined;
 175         for (entry.operands[0..entry.operand_count], 0..) |operand, index| {
 176             inputs[index] = switch (operand) {
 177                 .chain => current,
 178                 .full => |param| try logical.loadIndex(abi.input(logical, param), flat_index),
 179                 .broadcast => |broadcast| blk: {
 180                     const source_index = try generated_elementwise.broadcastInDimSourceIndex(logical, flat_index, broadcast.op, broadcast.slot);
 181                     break :blk try logical.loadIndex(abi.input(logical, broadcast.param), source_index);
 182                 },
 183             };
 184         }
 185         current = try generated_elementwise.emitElementwiseValue(logical, entry.kind, inputs[0..entry.operand_count], .f32, entry.op);
 186     }
 187     return current;
 188 }
 189 
 190 pub fn lower(
 191     allocator: std.mem.Allocator,
 192     ir_ctx: *ir.Context,
 193     outline: kernelization_model.KernelOutline,
 194     work: schedule_planning.ScheduleWorkItem,
 195     buffer_plan: *const bufferization.BufferPlanAnalysis,
 196     format: ?gpu.ArtifactFormat,
 197     math_tier: gpu.BackendMathTier,
 198 ) common.LoweringError!LoweredKernel {
 199     const desc = try dotGeneralDescriptionForWork(outline, work, buffer_plan);
 200     if (!dotLoweringSupported(desc.input_dtype, desc.output_dtype)) return error.CapabilityMismatch;
 201     var mma_tile: ?kernelization_model.DotGeneralMmaTile = null;
 202     if (math_tier != .exact) {
 203         if (math_tier != .tf32_tensor or format != .cuda_ptx) return error.UnsupportedOperation;
 204         mma_tile = kernelization_model.dotGeneralMmaTileFor(desc) orelse return error.UnsupportedOperation;
 205     }
 206 
 207     const epilogue_plan = try buildEpiloguePlan(work, buffer_plan);
 208     if (epilogue_plan.count != 0) {
 209         if (desc.output_dtype != .f32) return error.CapabilityMismatch;
 210         if (epilogue_plan.param_count != outline.inputCount()) return error.InvalidArtifact;
 211         var abi = try generated_abi.dotFused(allocator, desc.output_dtype, epilogue_plan.param_dtypes[0..epilogue_plan.param_count]);
 212         defer abi.deinit(allocator);
 213 
 214         if (mma_tile) |picked_mma| {
 215             var tile_value = picked_mma;
 216             tile_value.splits = 1;
 217             const entry_name = try generated_name.dotGeneralFusedMma(allocator, desc, work, tile_value, work.id);
 218             errdefer allocator.free(entry_name);
 219             var lowered = try generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.matrixThreads(32, tile_value.warps()), .{
 220                 .abi = abi,
 221                 .dims = desc.dims,
 222                 .tile = tile_value,
 223                 .epilogue = epilogue_plan.chain(),
 224             }, emitMmaDotBody);
 225             lowered.body = .{ .dot_mma_tile = tile_value };
 226             return lowered;
 227         }
 228 
 229         if (format == .cuda_ptx) if (kernelization_model.dotGeneralBlockTileFor(desc)) |picked_tile| {
 230             var tile = picked_tile;
 231             tile.splits = 1;
 232             const entry_name = try generated_name.dotGeneralFused(allocator, desc, work, tile, work.id);
 233             errdefer allocator.free(entry_name);
 234             var lowered = try generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.matrix(), .{
 235                 .abi = abi,
 236                 .dims = desc.dims,
 237                 .tile = tile,
 238                 .epilogue = epilogue_plan.chain(),
 239             }, emitTiledDotBody);
 240             lowered.body = .{ .dot_block_tile = tile };
 241             return lowered;
 242         };
 243 
 244         const entry_name = try generated_name.dotGeneralFused(allocator, desc, work, null, work.id);
 245         errdefer allocator.free(entry_name);
 246         if (hostLoopFormat(format)) {
 247             return generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flat(), .{
 248                 .abi = abi,
 249                 .dims = desc.dims,
 250                 .dtype = desc.output_dtype,
 251                 .epilogue = epilogue_plan.chain(),
 252             }, emitFlatDotBody);
 253         }
 254         return generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.matrix(), .{
 255             .abi = abi,
 256             .dims = desc.dims,
 257             .dtype = desc.output_dtype,
 258             .epilogue = epilogue_plan.chain(),
 259         }, emitDotBody);
 260     }
 261 
 262     const abi = generated_abi.dot(desc.input_dtype, desc.output_dtype);
 263     const empty_epilogue: []const EpilogueOp = &.{};
 264 
 265     if (mma_tile) |tile_value| {
 266         const entry_name = try generated_name.dotGeneralMma(allocator, desc, tile_value, work.id);
 267         errdefer allocator.free(entry_name);
 268         var lowered = try generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.matrixThreads(32, tile_value.warps()), .{
 269             .abi = abi,
 270             .dims = desc.dims,
 271             .tile = tile_value,
 272             .epilogue = empty_epilogue,
 273         }, emitMmaDotBody);
 274         lowered.body = .{ .dot_mma_tile = tile_value };
 275         if (tile_value.splits > 1) lowered.output_fill_pattern = 0;
 276         return lowered;
 277     }
 278 
 279     if (format == .cuda_ptx) if (kernelization_model.dotGeneralBlockTileFor(desc)) |tile| {
 280         const entry_name = try generated_name.dotGeneralTiled(allocator, desc, tile, work.id);
 281         errdefer allocator.free(entry_name);
 282         var lowered = try generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.matrix(), .{
 283             .abi = abi,
 284             .dims = desc.dims,
 285             .tile = tile,
 286             .epilogue = empty_epilogue,
 287         }, emitTiledDotBody);
 288         lowered.body = .{ .dot_block_tile = tile };
 289         if (tile.splits > 1) lowered.output_fill_pattern = 0;
 290         return lowered;
 291     };
 292 
 293     if (hostLoopFormat(format)) {
 294         const entry_name = try generated_name.dotGeneral(allocator, desc, work.id);
 295         errdefer allocator.free(entry_name);
 296 
 297         return generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flat(), .{
 298             .abi = abi,
 299             .dims = desc.dims,
 300             .dtype = desc.output_dtype,
 301             .epilogue = empty_epilogue,
 302         }, emitFlatDotBody);
 303     }
 304 
 305     const entry_name = try generated_name.dotGeneral(allocator, desc, work.id);
 306     errdefer allocator.free(entry_name);
 307 
 308     return generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.matrix(), .{
 309         .abi = abi,
 310         .dims = desc.dims,
 311         .dtype = desc.output_dtype,
 312         .epilogue = empty_epilogue,
 313     }, emitDotBody);
 314 }
 315 
 316 const max_block_tile_accumulators = 64;
 317 const max_block_tile_fragments = 8;
 318 
 319 const mma_shape = kernel_root.MmaShape{ .m = 16, .n = 8, .k = 8 };
 320 const mma_lane_acc_count = 4;
 321 const max_mma_fragments = 8;
 322 
 323 fn emitMmaDotBody(logical: anytype, ctx: anytype) !void {
 324     const tile: kernelization_model.DotGeneralMmaTile = ctx.tile;
 325     const dims: DotGeneralStaticDims = ctx.dims;
 326     const m_frags = tile.warpTileM() / mma_shape.m;
 327     const n_frags = tile.warpTileN() / mma_shape.n;
 328     const k_steps = tile.bk / mma_shape.k;
 329     const threads = tile.threads();
 330     const accumulator_count: usize = @intCast(m_frags * n_frags * mma_lane_acc_count);
 331     if (accumulator_count > max_block_tile_accumulators) return error.UnsupportedOperation;
 332     if (m_frags > max_mma_fragments or n_frags > max_mma_fragments) return error.UnsupportedOperation;
 333     if ((tile.bm * tile.bk) % (threads * 4) != 0) return error.UnsupportedOperation;
 334     if ((tile.bk * tile.bn) % (threads * 4) != 0) return error.UnsupportedOperation;
 335 
 336     if (tile.splits > 1 and dims.batch != 1) return error.UnsupportedOperation;
 337     if (tile.splits > 1 and dims.k % tile.splits != 0) return error.UnsupportedOperation;
 338     const split_chunk: u64 = dims.k / tile.splits;
 339 
 340     const tiles_m = dims.m / tile.bm;
 341     const tiles_n = dims.n / tile.bn;
 342     _ = try logical.index3D(.{
 343         .x = kernel_root.logical.axis("lane", @as(u64, tiles_n) * 32),
 344         .y = kernel_root.logical.axis("row", @as(u64, tiles_m) * tile.warps()),
 345         .z = kernel_root.logical.axis("batch", dims.batch * tile.splits),
 346     });
 347 
 348     const out = ctx.abi.output(logical);
 349     const lhs = ctx.abi.lhs(logical);
 350     const rhs = ctx.abi.rhs(logical);
 351 
 352     const a_pitch: i64 = @as(i64, tile.bk) + 4;
 353     const b_pitch: i64 = @as(i64, tile.bn) + 4;
 354     const tile_a = try logical.sharedBuffer(.f32, @as(u64, tile.bm) * @as(u64, @intCast(a_pitch)));
 355     const tile_b = try logical.sharedBuffer(.f32, @as(u64, tile.bk) * @as(u64, @intCast(b_pitch)));
 356 
 357     const lane = try logical.threadId(.x);
 358     const warp = try logical.threadId(.y);
 359     const bx = try logical.blockId(.x);
 360     const by = try logical.blockId(.y);
 361     const bz = try logical.blockId(.z);
 362 
 363     const k_extent = try logical.constantIndex(dims.k);
 364     const n_extent = try logical.constantIndex(dims.n);
 365     const a_pitch_extent = try logical.constantIndex(a_pitch);
 366     const b_pitch_extent = try logical.constantIndex(b_pitch);
 367     const four = try logical.constantIndex(4);
 368     const eight = try logical.constantIndex(8);
 369 
 370     const zero_base = try logical.constantIndex(0);
 371     const lhs_base = if (tile.splits > 1) zero_base else try logical.mul(bz, try logical.constantIndex(@as(i64, dims.m) * dims.k));
 372     const rhs_base = if (tile.splits > 1) zero_base else try logical.mul(bz, try logical.constantIndex(@as(i64, dims.k) * dims.n));
 373     const out_base = if (tile.splits > 1) zero_base else try logical.mul(bz, try logical.constantIndex(@as(i64, dims.m) * dims.n));
 374     const split_lower = if (tile.splits > 1) try logical.mul(bz, try logical.constantIndex(@as(i64, @intCast(split_chunk)))) else zero_base;
 375     const split_upper = if (tile.splits > 1) try logical.add(split_lower, try logical.constantIndex(@as(i64, @intCast(split_chunk)))) else k_extent;
 376 
 377     const row_block = try logical.mul(by, try logical.constantIndex(tile.bm));
 378     const col_block = try logical.mul(bx, try logical.constantIndex(tile.bn));
 379 
 380     const group = try logical.div(lane, four);
 381     const tid = try logical.sub(lane, try logical.mul(group, four));
 382     const group_hi = try logical.add(group, eight);
 383     const tid_hi = try logical.add(tid, four);
 384     const warps_n_extent = try logical.constantIndex(tile.warps_n);
 385     const warp_m = try logical.div(warp, warps_n_extent);
 386     const warp_n = try logical.sub(warp, try logical.mul(warp_m, warps_n_extent));
 387     const thread_id = try logical.add(try logical.mul(warp, try logical.constantIndex(32)), lane);
 388 
 389     const warp_row = try logical.mul(warp_m, try logical.constantIndex(tile.warpTileM()));
 390     const warp_col = try logical.mul(warp_n, try logical.constantIndex(tile.warpTileN()));
 391 
 392     const zero_scalar = try logical.constantFloat(.f32, 0.0);
 393     var init_args: [max_block_tile_accumulators]kernel_root.Value = undefined;
 394     var result_types: [max_block_tile_accumulators]kernel_root.Type = undefined;
 395     for (init_args[0..accumulator_count], result_types[0..accumulator_count]) |*arg, *ty_slot| {
 396         arg.* = zero_scalar;
 397         ty_slot.* = zero_scalar.valueType();
 398     }
 399 
 400     const stage_ctx = MmaStageContext{
 401         .tile = tile,
 402         .threads = threads,
 403         .thread_id = thread_id,
 404         .row_block = row_block,
 405         .col_block = col_block,
 406         .lhs = lhs,
 407         .rhs = rhs,
 408         .a_pitch_extent = a_pitch_extent,
 409         .b_pitch_extent = b_pitch_extent,
 410         .n_extent = n_extent,
 411         .k_extent = k_extent,
 412         .lhs_base = lhs_base,
 413         .rhs_base = rhs_base,
 414         .four = four,
 415     };
 416     var slice_ctx = MmaSliceContext{
 417         .m_frags = m_frags,
 418         .n_frags = n_frags,
 419         .k_steps = k_steps,
 420         .a_pitch_extent = a_pitch_extent,
 421         .b_pitch_extent = b_pitch_extent,
 422         .warp_row = warp_row,
 423         .warp_col = warp_col,
 424         .group = group,
 425         .group_hi = group_hi,
 426         .tid = tid,
 427         .tid_hi = tid_hi,
 428         .round_frags = false,
 429     };
 430 
 431     const multistage = tile.splits == 1 and tile.stages > 2 and tile.stages <= max_pipeline_stages and
 432         split_chunk % (@as(u64, tile.stages) * tile.bk) == 0 and
 433         split_chunk >= 3 * @as(u64, tile.stages) * tile.bk;
 434     const double_buffered = !multistage and tile.stages >= 2 and split_chunk % (2 * @as(u64, tile.bk)) == 0 and split_chunk >= 4 * @as(u64, tile.bk);
 435     const loop_lower = split_lower;
 436 
 437     if (multistage) {
 438         slice_ctx.round_frags = true;
 439         const stage_count = tile.stages;
 440         var slots_a: [max_pipeline_stages]kernel_root.Value = undefined;
 441         var slots_b: [max_pipeline_stages]kernel_root.Value = undefined;
 442         slots_a[0] = tile_a;
 443         slots_b[0] = tile_b;
 444         var slot: u32 = 1;
 445         while (slot < stage_count) : (slot += 1) {
 446             slots_a[slot] = try logical.sharedBuffer(.f32, @as(u64, tile.bm) * @as(u64, @intCast(a_pitch)));
 447             slots_b[slot] = try logical.sharedBuffer(.f32, @as(u64, tile.bk) * @as(u64, @intCast(b_pitch)));
 448         }
 449         slot = 0;
 450         while (slot + 1 < stage_count) : (slot += 1) {
 451             const k_stage = try logical.constantIndex(@as(i64, slot) * tile.bk);
 452             try stageMmaTilesAsync(logical, stage_ctx, k_stage, slots_a[slot], slots_b[slot], false);
 453             try logical.asyncCopyCommit();
 454         }
 455 
 456         const loop_step = try logical.constantIndex(@as(i64, stage_count) * tile.bk);
 457         var scope = try logical.forScope(loop_lower, split_upper, loop_step, init_args[0..accumulator_count], result_types[0..accumulator_count]);
 458         errdefer scope.abort();
 459         const k0 = scope.inductionVar();
 460 
 461         var accumulators: [max_block_tile_accumulators]kernel_root.Value = undefined;
 462         for (accumulators[0..accumulator_count], 0..) |*accumulator, index| {
 463             accumulator.* = scope.iterArg(index) orelse return error.UnsupportedOperation;
 464         }
 465 
 466         var sub: u32 = 0;
 467         while (sub < stage_count) : (sub += 1) {
 468             try logical.asyncCopyWait(stage_count - 2);
 469             try logical.barrier(.block);
 470             const k_ahead = try logical.add(k0, try logical.constantIndex((@as(i64, sub) + @as(i64, stage_count) - 1) * tile.bk));
 471             const write_slot = (sub + stage_count - 1) % stage_count;
 472             try stageMmaTilesAsync(logical, stage_ctx, k_ahead, slots_a[write_slot], slots_b[write_slot], true);
 473             try logical.asyncCopyCommit();
 474             try emitMmaComputeSlice(logical, slice_ctx, slots_a[sub], slots_b[sub], accumulators[0..accumulator_count]);
 475         }
 476         try scope.leave(accumulators[0..accumulator_count]);
 477         try emitMmaEpilogue(logical, ctx, &scope, accumulator_count, .{
 478             .m_frags = m_frags,
 479             .n_frags = n_frags,
 480             .splits = tile.splits,
 481             .out = out,
 482             .out_base = out_base,
 483             .row_block = row_block,
 484             .col_block = col_block,
 485             .warp_row = warp_row,
 486             .warp_col = warp_col,
 487             .group = group,
 488             .group_hi = group_hi,
 489             .tid = tid,
 490             .n_extent = n_extent,
 491         });
 492         return;
 493     }
 494 
 495     if (double_buffered) {
 496         slice_ctx.round_frags = true;
 497         const tile_a_next = try logical.sharedBuffer(.f32, @as(u64, tile.bm) * @as(u64, @intCast(a_pitch)));
 498         const tile_b_next = try logical.sharedBuffer(.f32, @as(u64, tile.bk) * @as(u64, @intCast(b_pitch)));
 499         const bk_extent = try logical.constantIndex(tile.bk);
 500 
 501         try stageMmaTilesAsync(logical, stage_ctx, loop_lower, tile_a, tile_b, false);
 502         try logical.asyncCopyCommit();
 503 
 504         const loop_step = try logical.constantIndex(2 * @as(i64, tile.bk));
 505         var scope = try logical.forScope(loop_lower, split_upper, loop_step, init_args[0..accumulator_count], result_types[0..accumulator_count]);
 506         errdefer scope.abort();
 507         const k0 = scope.inductionVar();
 508 
 509         var accumulators: [max_block_tile_accumulators]kernel_root.Value = undefined;
 510         for (accumulators[0..accumulator_count], 0..) |*accumulator, index| {
 511             accumulator.* = scope.iterArg(index) orelse return error.UnsupportedOperation;
 512         }
 513 
 514         try logical.barrier(.block);
 515         const k_ahead = try logical.add(k0, bk_extent);
 516         try stageMmaTilesAsync(logical, stage_ctx, k_ahead, tile_a_next, tile_b_next, false);
 517         try logical.asyncCopyCommit();
 518         try logical.asyncCopyWait(1);
 519         try logical.barrier(.block);
 520         try emitMmaComputeSlice(logical, slice_ctx, tile_a, tile_b, accumulators[0..accumulator_count]);
 521         try logical.barrier(.block);
 522         const k_wrap = try logical.add(k0, try logical.constantIndex(2 * @as(i64, tile.bk)));
 523         try stageMmaTilesAsync(logical, stage_ctx, k_wrap, tile_a, tile_b, true);
 524         try logical.asyncCopyCommit();
 525         try logical.asyncCopyWait(1);
 526         try logical.barrier(.block);
 527         try emitMmaComputeSlice(logical, slice_ctx, tile_a_next, tile_b_next, accumulators[0..accumulator_count]);
 528         try scope.leave(accumulators[0..accumulator_count]);
 529         try emitMmaEpilogue(logical, ctx, &scope, accumulator_count, .{
 530             .m_frags = m_frags,
 531             .n_frags = n_frags,
 532             .splits = tile.splits,
 533             .out = out,
 534             .out_base = out_base,
 535             .row_block = row_block,
 536             .col_block = col_block,
 537             .warp_row = warp_row,
 538             .warp_col = warp_col,
 539             .group = group,
 540             .group_hi = group_hi,
 541             .tid = tid,
 542             .n_extent = n_extent,
 543         });
 544         return;
 545     }
 546 
 547     const loop_step = try logical.constantIndex(tile.bk);
 548     var scope = try logical.forScope(loop_lower, split_upper, loop_step, init_args[0..accumulator_count], result_types[0..accumulator_count]);
 549     errdefer scope.abort();
 550     const k0 = scope.inductionVar();
 551 
 552     try logical.barrier(.block);
 553     try stageMmaTiles(logical, stage_ctx, k0, tile_a, tile_b, false);
 554     try logical.barrier(.block);
 555 
 556     var accumulators: [max_block_tile_accumulators]kernel_root.Value = undefined;
 557     for (accumulators[0..accumulator_count], 0..) |*accumulator, index| {
 558         accumulator.* = scope.iterArg(index) orelse return error.UnsupportedOperation;
 559     }
 560 
 561     try emitMmaComputeSlice(logical, slice_ctx, tile_a, tile_b, accumulators[0..accumulator_count]);
 562 
 563     try scope.leave(accumulators[0..accumulator_count]);
 564 
 565     try emitMmaEpilogue(logical, ctx, &scope, accumulator_count, .{
 566         .m_frags = m_frags,
 567         .n_frags = n_frags,
 568         .splits = tile.splits,
 569         .out = out,
 570         .out_base = out_base,
 571         .row_block = row_block,
 572         .col_block = col_block,
 573         .warp_row = warp_row,
 574         .warp_col = warp_col,
 575         .group = group,
 576         .group_hi = group_hi,
 577         .tid = tid,
 578         .n_extent = n_extent,
 579     });
 580 }
 581 
 582 const MmaStageContext = struct {
 583     tile: kernelization_model.DotGeneralMmaTile,
 584     threads: u32,
 585     thread_id: kernel_root.Value,
 586     row_block: kernel_root.Value,
 587     col_block: kernel_root.Value,
 588     lhs: kernel_root.Value,
 589     rhs: kernel_root.Value,
 590     a_pitch_extent: kernel_root.Value,
 591     b_pitch_extent: kernel_root.Value,
 592     n_extent: kernel_root.Value,
 593     k_extent: kernel_root.Value,
 594     lhs_base: kernel_root.Value,
 595     rhs_base: kernel_root.Value,
 596     four: kernel_root.Value,
 597 };
 598 
 599 fn stageMmaTiles(
 600     logical: anytype,
 601     ctx: MmaStageContext,
 602     k_slice: kernel_root.Value,
 603     tile_a: kernel_root.Value,
 604     tile_b: kernel_root.Value,
 605     clamp: bool,
 606 ) !void {
 607     const tile = ctx.tile;
 608     const threads = ctx.threads;
 609     const k_quad_limit = try logical.sub(ctx.k_extent, ctx.four);
 610     const k_row_limit = try logical.sub(ctx.k_extent, try logical.constantIndex(1));
 611 
 612     const a_quads_per_thread = (tile.bm * tile.bk) / (threads * 4);
 613     const a_quads_per_row = tile.bk / 4;
 614     var quad: u32 = 0;
 615     while (quad < a_quads_per_thread) : (quad += 1) {
 616         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, quad) * threads));
 617         const quads_per_row_extent = try logical.constantIndex(a_quads_per_row);
 618         const local_row = try logical.div(flat, quads_per_row_extent);
 619         const quad_in_row = try logical.sub(flat, try logical.mul(local_row, quads_per_row_extent));
 620         const local_k = try logical.mul(quad_in_row, ctx.four);
 621         const global_row = try logical.add(ctx.row_block, local_row);
 622         var global_k = try logical.add(k_slice, local_k);
 623         if (clamp) global_k = try logical.min(global_k, k_quad_limit);
 624         const global_index = try logical.add(try logical.add(ctx.lhs_base, try logical.mul(global_row, ctx.k_extent)), global_k);
 625         const quad_value = try logical.loadVector(ctx.lhs, global_index, 4);
 626         const shared_base = try logical.add(try logical.mul(local_row, ctx.a_pitch_extent), local_k);
 627         var lane_index: u32 = 0;
 628         while (lane_index < 4) : (lane_index += 1) {
 629             const scalar = try logical.extractLane(quad_value, lane_index, .f32);
 630             const rounded = try logical.tf32Round(scalar);
 631             const shared_index = try logical.add(shared_base, try logical.constantIndex(lane_index));
 632             try logical.storeIndex(rounded, tile_a, shared_index);
 633         }
 634     }
 635 
 636     const b_quads_per_thread = (tile.bk * tile.bn) / (threads * 4);
 637     const b_quads_per_row = tile.bn / 4;
 638     quad = 0;
 639     while (quad < b_quads_per_thread) : (quad += 1) {
 640         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, quad) * threads));
 641         const quads_per_row_extent = try logical.constantIndex(b_quads_per_row);
 642         const local_k = try logical.div(flat, quads_per_row_extent);
 643         const quad_in_row = try logical.sub(flat, try logical.mul(local_k, quads_per_row_extent));
 644         const local_col = try logical.mul(quad_in_row, ctx.four);
 645         var global_k = try logical.add(k_slice, local_k);
 646         if (clamp) global_k = try logical.min(global_k, k_row_limit);
 647         const global_col = try logical.add(ctx.col_block, local_col);
 648         const global_index = try logical.add(try logical.add(ctx.rhs_base, try logical.mul(global_k, ctx.n_extent)), global_col);
 649         const quad_value = try logical.loadVector(ctx.rhs, global_index, 4);
 650         const shared_base = try logical.add(try logical.mul(local_k, ctx.b_pitch_extent), local_col);
 651         var lane_index: u32 = 0;
 652         while (lane_index < 4) : (lane_index += 1) {
 653             const scalar = try logical.extractLane(quad_value, lane_index, .f32);
 654             const rounded = try logical.tf32Round(scalar);
 655             const shared_index = try logical.add(shared_base, try logical.constantIndex(lane_index));
 656             try logical.storeIndex(rounded, tile_b, shared_index);
 657         }
 658     }
 659 }
 660 
 661 fn stageMmaTilesAsync(
 662     logical: anytype,
 663     ctx: MmaStageContext,
 664     k_slice: kernel_root.Value,
 665     tile_a: kernel_root.Value,
 666     tile_b: kernel_root.Value,
 667     clamp: bool,
 668 ) !void {
 669     const tile = ctx.tile;
 670     const threads = ctx.threads;
 671     const k_quad_limit = try logical.sub(ctx.k_extent, ctx.four);
 672     const k_row_limit = try logical.sub(ctx.k_extent, try logical.constantIndex(1));
 673 
 674     const a_quads_per_thread = (tile.bm * tile.bk) / (threads * 4);
 675     const a_quads_per_row = tile.bk / 4;
 676     var quad: u32 = 0;
 677     while (quad < a_quads_per_thread) : (quad += 1) {
 678         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, quad) * threads));
 679         const quads_per_row_extent = try logical.constantIndex(a_quads_per_row);
 680         const local_row = try logical.div(flat, quads_per_row_extent);
 681         const quad_in_row = try logical.sub(flat, try logical.mul(local_row, quads_per_row_extent));
 682         const local_k = try logical.mul(quad_in_row, ctx.four);
 683         const global_row = try logical.add(ctx.row_block, local_row);
 684         var global_k = try logical.add(k_slice, local_k);
 685         if (clamp) global_k = try logical.min(global_k, k_quad_limit);
 686         const global_index = try logical.add(try logical.add(ctx.lhs_base, try logical.mul(global_row, ctx.k_extent)), global_k);
 687         const shared_base = try logical.add(try logical.mul(local_row, ctx.a_pitch_extent), local_k);
 688         try logical.asyncCopyShared(tile_a, shared_base, ctx.lhs, global_index, 16);
 689     }
 690 
 691     const b_quads_per_thread = (tile.bk * tile.bn) / (threads * 4);
 692     const b_quads_per_row = tile.bn / 4;
 693     quad = 0;
 694     while (quad < b_quads_per_thread) : (quad += 1) {
 695         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, quad) * threads));
 696         const quads_per_row_extent = try logical.constantIndex(b_quads_per_row);
 697         const local_k = try logical.div(flat, quads_per_row_extent);
 698         const quad_in_row = try logical.sub(flat, try logical.mul(local_k, quads_per_row_extent));
 699         const local_col = try logical.mul(quad_in_row, ctx.four);
 700         var global_k = try logical.add(k_slice, local_k);
 701         if (clamp) global_k = try logical.min(global_k, k_row_limit);
 702         const global_col = try logical.add(ctx.col_block, local_col);
 703         const global_index = try logical.add(try logical.add(ctx.rhs_base, try logical.mul(global_k, ctx.n_extent)), global_col);
 704         const shared_base = try logical.add(try logical.mul(local_k, ctx.b_pitch_extent), local_col);
 705         try logical.asyncCopyShared(tile_b, shared_base, ctx.rhs, global_index, 16);
 706     }
 707 }
 708 
 709 const MmaSliceContext = struct {
 710     m_frags: u32,
 711     n_frags: u32,
 712     k_steps: u32,
 713     a_pitch_extent: kernel_root.Value,
 714     b_pitch_extent: kernel_root.Value,
 715     warp_row: kernel_root.Value,
 716     warp_col: kernel_root.Value,
 717     group: kernel_root.Value,
 718     group_hi: kernel_root.Value,
 719     tid: kernel_root.Value,
 720     tid_hi: kernel_root.Value,
 721     round_frags: bool,
 722 };
 723 
 724 fn mmaFragValue(logical: anytype, ctx: MmaSliceContext, buffer: kernel_root.Value, index: kernel_root.Value) !kernel_root.Value {
 725     const loaded = try logical.loadIndex(buffer, index);
 726     if (!ctx.round_frags) return loaded;
 727     return logical.tf32Round(loaded);
 728 }
 729 
 730 fn emitMmaComputeSlice(
 731     logical: anytype,
 732     ctx: MmaSliceContext,
 733     tile_a: kernel_root.Value,
 734     tile_b: kernel_root.Value,
 735     accumulators: []kernel_root.Value,
 736 ) !void {
 737     const m_frags = ctx.m_frags;
 738     const n_frags = ctx.n_frags;
 739     var step: u32 = 0;
 740     while (step < ctx.k_steps) : (step += 1) {
 741         const kk = try logical.constantIndex(@as(i64, step) * mma_shape.k);
 742         const col_lo = try logical.add(kk, ctx.tid);
 743         const col_hi = try logical.add(kk, ctx.tid_hi);
 744 
 745         var a_frags: [max_mma_fragments][4]kernel_root.Value = undefined;
 746         var mf: u32 = 0;
 747         while (mf < m_frags) : (mf += 1) {
 748             const frag_row = try logical.add(ctx.warp_row, try logical.constantIndex(@as(i64, mf) * mma_shape.m));
 749             const row_lo = try logical.mul(try logical.add(frag_row, ctx.group), ctx.a_pitch_extent);
 750             const row_hi = try logical.mul(try logical.add(frag_row, ctx.group_hi), ctx.a_pitch_extent);
 751             a_frags[mf] = .{
 752                 try mmaFragValue(logical, ctx, tile_a, try logical.add(row_lo, col_lo)),
 753                 try mmaFragValue(logical, ctx, tile_a, try logical.add(row_hi, col_lo)),
 754                 try mmaFragValue(logical, ctx, tile_a, try logical.add(row_lo, col_hi)),
 755                 try mmaFragValue(logical, ctx, tile_a, try logical.add(row_hi, col_hi)),
 756             };
 757         }
 758 
 759         const k_row_lo = try logical.mul(col_lo, ctx.b_pitch_extent);
 760         const k_row_hi = try logical.mul(col_hi, ctx.b_pitch_extent);
 761         var b_frags: [max_mma_fragments][2]kernel_root.Value = undefined;
 762         var nf: u32 = 0;
 763         while (nf < n_frags) : (nf += 1) {
 764             const frag_col = try logical.add(ctx.warp_col, try logical.constantIndex(@as(i64, nf) * mma_shape.n));
 765             const col_offset = try logical.add(frag_col, ctx.group);
 766             b_frags[nf] = .{
 767                 try mmaFragValue(logical, ctx, tile_b, try logical.add(k_row_lo, col_offset)),
 768                 try mmaFragValue(logical, ctx, tile_b, try logical.add(k_row_hi, col_offset)),
 769             };
 770         }
 771 
 772         mf = 0;
 773         while (mf < m_frags) : (mf += 1) {
 774             nf = 0;
 775             while (nf < n_frags) : (nf += 1) {
 776                 const slot = (@as(usize, mf) * n_frags + nf) * mma_lane_acc_count;
 777                 const updated = try logical.mmaSync(mma_shape, a_frags[mf], b_frags[nf], .{
 778                     accumulators[slot],
 779                     accumulators[slot + 1],
 780                     accumulators[slot + 2],
 781                     accumulators[slot + 3],
 782                 });
 783                 accumulators[slot] = updated[0];
 784                 accumulators[slot + 1] = updated[1];
 785                 accumulators[slot + 2] = updated[2];
 786                 accumulators[slot + 3] = updated[3];
 787             }
 788         }
 789     }
 790 }
 791 
 792 const MmaEpilogueContext = struct {
 793     m_frags: u32,
 794     n_frags: u32,
 795     splits: u32 = 1,
 796     out: kernel_root.Value,
 797     out_base: kernel_root.Value,
 798     row_block: kernel_root.Value,
 799     col_block: kernel_root.Value,
 800     warp_row: kernel_root.Value,
 801     warp_col: kernel_root.Value,
 802     group: kernel_root.Value,
 803     group_hi: kernel_root.Value,
 804     tid: kernel_root.Value,
 805     n_extent: kernel_root.Value,
 806 };
 807 
 808 fn emitMmaEpilogue(
 809     logical: anytype,
 810     ctx: anytype,
 811     scope: anytype,
 812     accumulator_count: usize,
 813     epi: MmaEpilogueContext,
 814 ) !void {
 815     var results: [max_block_tile_accumulators]kernel_root.Value = undefined;
 816     for (results[0..accumulator_count], 0..) |*result, index| {
 817         result.* = scope.result(index) orelse return error.UnsupportedOperation;
 818     }
 819 
 820     const two = try logical.constantIndex(2);
 821     const one = try logical.constantIndex(1);
 822     const tid2 = try logical.mul(epi.tid, two);
 823     var mf: u32 = 0;
 824     while (mf < epi.m_frags) : (mf += 1) {
 825         const frag_row = try logical.add(epi.row_block, try logical.add(epi.warp_row, try logical.constantIndex(@as(i64, mf) * mma_shape.m)));
 826         const row_lo = try logical.add(frag_row, epi.group);
 827         const row_hi = try logical.add(frag_row, epi.group_hi);
 828         const row_lo_offset = try logical.add(epi.out_base, try logical.mul(row_lo, epi.n_extent));
 829         const row_hi_offset = try logical.add(epi.out_base, try logical.mul(row_hi, epi.n_extent));
 830         var nf: u32 = 0;
 831         while (nf < epi.n_frags) : (nf += 1) {
 832             const frag_col = try logical.add(epi.col_block, try logical.add(epi.warp_col, try logical.constantIndex(@as(i64, nf) * mma_shape.n)));
 833             const col_lo = try logical.add(frag_col, tid2);
 834             const col_hi = try logical.add(col_lo, one);
 835             const slot = (@as(usize, mf) * epi.n_frags + nf) * mma_lane_acc_count;
 836             const addresses = [mma_lane_acc_count]kernel_root.Value{
 837                 try logical.add(row_lo_offset, col_lo),
 838                 try logical.add(row_lo_offset, col_hi),
 839                 try logical.add(row_hi_offset, col_lo),
 840                 try logical.add(row_hi_offset, col_hi),
 841             };
 842             for (addresses, 0..) |address, acc_index| {
 843                 const value = try applyEpilogueChain(logical, ctx.epilogue, ctx.abi, results[slot + acc_index], address);
 844                 if (epi.splits > 1) {
 845                     _ = try logical.atomicRmw(.add, value, epi.out, address);
 846                 } else {
 847                     try logical.storeIndex(value, epi.out, address);
 848                 }
 849             }
 850         }
 851     }
 852 }
 853 
 854 fn tiledDotVectorizable(tile: kernelization_model.DotGeneralBlockTile, dims: DotGeneralStaticDims) bool {
 855     if (!tile.exact(dims)) return false;
 856     if (dims.k % 4 != 0 or dims.n % 4 != 0) return false;
 857     if (tile.bk % 4 != 0 or tile.bn % 4 != 0) return false;
 858     if (tile.tm % 4 != 0 or tile.tn % 4 != 0) return false;
 859     const threads = tile.threadsX() * tile.threadsY();
 860     if ((tile.bm * tile.bk) % (threads * 4) != 0) return false;
 861     if ((tile.bk * tile.bn) % 4 != 0) return false;
 862     return true;
 863 }
 864 
 865 fn emitTiledDotBody(logical: anytype, ctx: anytype) !void {
 866     const tile: kernelization_model.DotGeneralBlockTile = ctx.tile;
 867     const dims: DotGeneralStaticDims = ctx.dims;
 868     const exact = tile.exact(dims);
 869     if (tile.tm * tile.tn > max_block_tile_accumulators) return error.UnsupportedOperation;
 870     if (tile.tm > max_block_tile_fragments or tile.tn > max_block_tile_fragments) return error.UnsupportedOperation;
 871     const staging_threads = tile.threadsX() * tile.threadsY();
 872     if ((tile.bm * tile.bk) % staging_threads != 0) return error.UnsupportedOperation;
 873     if ((tile.bk * tile.bn) % staging_threads != 0) return error.UnsupportedOperation;
 874     if (tile.splits > 1 and dims.batch != 1) return error.UnsupportedOperation;
 875     if (tile.splits > 1 and dims.k % tile.splits != 0) return error.UnsupportedOperation;
 876     const split_chunk: u64 = dims.k / tile.splits;
 877     const vectorized = tiledDotVectorizable(tile, dims);
 878     const multistage = vectorized and tile.stages > 2 and tile.stages <= max_pipeline_stages and
 879         split_chunk % (@as(u64, tile.stages) * tile.bk) == 0 and
 880         split_chunk >= 3 * @as(u64, tile.stages) * tile.bk;
 881     const double_buffered = !multistage and vectorized and split_chunk % (2 * tile.bk) == 0 and split_chunk >= 4 * tile.bk;
 882 
 883     const tiles_m = (dims.m + tile.bm - 1) / tile.bm;
 884     const tiles_n = (dims.n + tile.bn - 1) / tile.bn;
 885     _ = try logical.index3D(.{
 886         .x = kernel_root.logical.axis("lane", @as(u64, tiles_n) * tile.threadsX()),
 887         .y = kernel_root.logical.axis("row", @as(u64, tiles_m) * tile.threadsY()),
 888         .z = kernel_root.logical.axis("batch", dims.batch * tile.splits),
 889     });
 890 
 891     const out = ctx.abi.output(logical);
 892     const lhs = ctx.abi.lhs(logical);
 893     const rhs = ctx.abi.rhs(logical);
 894 
 895     const a_pitch: i64 = @as(i64, tile.bm) + if (vectorized) @as(i64, 4) else @as(i64, 1);
 896     const tile_a = try logical.sharedBuffer(.f32, @as(u64, tile.bk) * @as(u64, @intCast(a_pitch)));
 897     const tile_b = try logical.sharedBuffer(.f32, @as(u64, tile.bk) * tile.bn);
 898     const tile_a_next = if (double_buffered) try logical.sharedBuffer(.f32, @as(u64, tile.bk) * @as(u64, @intCast(a_pitch))) else tile_a;
 899     const tile_b_next = if (double_buffered) try logical.sharedBuffer(.f32, @as(u64, tile.bk) * tile.bn) else tile_b;
 900 
 901     const tx = try logical.threadId(.x);
 902     const ty = try logical.threadId(.y);
 903     const bx = try logical.blockId(.x);
 904     const by = try logical.blockId(.y);
 905     const bz = try logical.blockId(.z);
 906 
 907     const m_edge = try logical.constantIndex(@as(i64, dims.m) - 1);
 908     const n_edge = try logical.constantIndex(@as(i64, dims.n) - 1);
 909     const k_edge = try logical.constantIndex(@as(i64, dims.k) - 1);
 910     const m_extent = try logical.constantIndex(dims.m);
 911     const n_extent = try logical.constantIndex(dims.n);
 912     const k_extent = try logical.constantIndex(dims.k);
 913 
 914     const zero_index = try logical.constantIndex(0);
 915     const lhs_base = if (tile.splits > 1) zero_index else try logical.mul(bz, try logical.constantIndex(@as(i64, dims.m) * dims.k));
 916     const rhs_base = if (tile.splits > 1) zero_index else try logical.mul(bz, try logical.constantIndex(@as(i64, dims.k) * dims.n));
 917     const out_base = if (tile.splits > 1) zero_index else try logical.mul(bz, try logical.constantIndex(@as(i64, dims.m) * dims.n));
 918     const split_lower = if (tile.splits > 1) try logical.mul(bz, try logical.constantIndex(@as(i64, @intCast(split_chunk)))) else zero_index;
 919     const split_upper = if (tile.splits > 1) try logical.add(split_lower, try logical.constantIndex(@as(i64, @intCast(split_chunk)))) else k_extent;
 920 
 921     const row_block = try logical.mul(by, try logical.constantIndex(tile.bm));
 922     const col_block = try logical.mul(bx, try logical.constantIndex(tile.bn));
 923     const row_local = try logical.mul(ty, try logical.constantIndex(if (vectorized) 4 else @as(i64, tile.tm)));
 924     const col_local = try logical.mul(tx, try logical.constantIndex(if (vectorized) 4 else @as(i64, tile.tn)));
 925     const row_base = try logical.add(row_block, row_local);
 926     const col_base = try logical.add(col_block, col_local);
 927 
 928     const threads = staging_threads;
 929     const thread_id = try logical.add(try logical.mul(ty, try logical.constantIndex(tile.threadsX())), tx);
 930     const a_pitch_extent = try logical.constantIndex(a_pitch);
 931     const bk_extent = try logical.constantIndex(tile.bk);
 932     const bn_extent = try logical.constantIndex(tile.bn);
 933 
 934     const stage_ctx = StageContext{
 935         .tile = tile,
 936         .threads = threads,
 937         .thread_id = thread_id,
 938         .row_block = row_block,
 939         .col_block = col_block,
 940         .lhs = lhs,
 941         .rhs = rhs,
 942         .a_pitch_extent = a_pitch_extent,
 943         .bn_extent = bn_extent,
 944         .n_extent = n_extent,
 945         .k_extent = k_extent,
 946         .lhs_base = lhs_base,
 947         .rhs_base = rhs_base,
 948     };
 949     const slice_ctx = SliceContext{
 950         .tile = tile,
 951         .a_pitch = a_pitch,
 952         .row_local = row_local,
 953         .col_local = col_local,
 954     };
 955 
 956     const zero_scalar = try logical.constantFloat(.f32, 0.0);
 957     var init_args: [max_block_tile_accumulators]kernel_root.Value = undefined;
 958     var result_types: [max_block_tile_accumulators]kernel_root.Type = undefined;
 959     const accumulator_count: usize = @intCast(tile.tm * tile.tn);
 960     for (init_args[0..accumulator_count], result_types[0..accumulator_count]) |*arg, *ty_slot| {
 961         arg.* = zero_scalar;
 962         ty_slot.* = zero_scalar.valueType();
 963     }
 964 
 965     if (multistage) {
 966         const stage_count = tile.stages;
 967         var slots_a: [max_pipeline_stages]kernel_root.Value = undefined;
 968         var slots_b: [max_pipeline_stages]kernel_root.Value = undefined;
 969         slots_a[0] = tile_a;
 970         slots_b[0] = tile_b;
 971         var slot: u32 = 1;
 972         while (slot < stage_count) : (slot += 1) {
 973             slots_a[slot] = try logical.sharedBuffer(.f32, @as(u64, tile.bk) * @as(u64, @intCast(a_pitch)));
 974             slots_b[slot] = try logical.sharedBuffer(.f32, @as(u64, tile.bk) * tile.bn);
 975         }
 976         slot = 0;
 977         while (slot + 1 < stage_count) : (slot += 1) {
 978             const k_stage = try logical.add(split_lower, try logical.constantIndex(@as(i64, slot) * tile.bk));
 979             try stageBTileAsync(logical, stage_ctx, k_stage, slots_b[slot], false);
 980             try logical.asyncCopyCommit();
 981             const staged = try loadATileQuads(logical, stage_ctx, k_stage, false);
 982             try storeATileQuads(logical, stage_ctx, staged, slots_a[slot]);
 983         }
 984 
 985         const loop_step = try logical.constantIndex(@as(i64, stage_count) * tile.bk);
 986         var scope = try logical.forScope(split_lower, split_upper, loop_step, init_args[0..accumulator_count], result_types[0..accumulator_count]);
 987         errdefer scope.abort();
 988         const k0 = scope.inductionVar();
 989 
 990         var accumulators: [max_block_tile_accumulators]kernel_root.Value = undefined;
 991         for (accumulators[0..accumulator_count], 0..) |*accumulator, index| {
 992             accumulator.* = scope.iterArg(index) orelse return error.UnsupportedOperation;
 993         }
 994 
 995         var sub: u32 = 0;
 996         while (sub < stage_count) : (sub += 1) {
 997             try logical.asyncCopyWait(stage_count - 2);
 998             try logical.barrier(.block);
 999             const k_ahead = try logical.add(k0, try logical.constantIndex((@as(i64, sub) + @as(i64, stage_count) - 1) * tile.bk));
1000             const write_slot = (sub + stage_count - 1) % stage_count;
1001             try stageBTileAsync(logical, stage_ctx, k_ahead, slots_b[write_slot], true);
1002             try logical.asyncCopyCommit();
1003             const staged = try loadATileQuads(logical, stage_ctx, k_ahead, true);
1004             try emitComputeSlice(logical, slice_ctx, slots_a[sub], slots_b[sub], accumulators[0..accumulator_count], true);
1005             try storeATileQuads(logical, stage_ctx, staged, slots_a[write_slot]);
1006         }
1007         try scope.leave(accumulators[0..accumulator_count]);
1008         try emitTiledEpilogue(logical, ctx, &scope, accumulator_count, out, out_base, row_base, col_base, m_extent, n_extent, exact, true);
1009         return;
1010     }
1011 
1012     if (double_buffered) {
1013         try stageBTileAsync(logical, stage_ctx, split_lower, tile_b, false);
1014         try logical.asyncCopyCommit();
1015         const prologue_quads = try loadATileQuads(logical, stage_ctx, split_lower, false);
1016         try storeATileQuads(logical, stage_ctx, prologue_quads, tile_a);
1017 
1018         const loop_step = try logical.constantIndex(2 * @as(i64, tile.bk));
1019         var scope = try logical.forScope(split_lower, split_upper, loop_step, init_args[0..accumulator_count], result_types[0..accumulator_count]);
1020         errdefer scope.abort();
1021         const k0 = scope.inductionVar();
1022 
1023         var accumulators: [max_block_tile_accumulators]kernel_root.Value = undefined;
1024         for (accumulators[0..accumulator_count], 0..) |*accumulator, index| {
1025             accumulator.* = scope.iterArg(index) orelse return error.UnsupportedOperation;
1026         }
1027 
1028         try logical.barrier(.block);
1029         const k_ahead = try logical.add(k0, bk_extent);
1030         try stageBTileAsync(logical, stage_ctx, k_ahead, tile_b_next, false);
1031         try logical.asyncCopyCommit();
1032         const quads_ahead = try loadATileQuads(logical, stage_ctx, k_ahead, false);
1033         try logical.asyncCopyWait(1);
1034         try logical.barrier(.block);
1035         try emitComputeSlice(logical, slice_ctx, tile_a, tile_b, accumulators[0..accumulator_count], true);
1036         try logical.barrier(.block);
1037         const k_wrap = try logical.add(k0, try logical.constantIndex(2 * @as(i64, tile.bk)));
1038         try stageBTileAsync(logical, stage_ctx, k_wrap, tile_b, true);
1039         try logical.asyncCopyCommit();
1040         const quads_wrap = try loadATileQuads(logical, stage_ctx, k_wrap, true);
1041         try storeATileQuads(logical, stage_ctx, quads_ahead, tile_a_next);
1042         try logical.asyncCopyWait(1);
1043         try logical.barrier(.block);
1044         try emitComputeSlice(logical, slice_ctx, tile_a_next, tile_b_next, accumulators[0..accumulator_count], true);
1045         try storeATileQuads(logical, stage_ctx, quads_wrap, tile_a);
1046         try scope.leave(accumulators[0..accumulator_count]);
1047         try emitTiledEpilogue(logical, ctx, &scope, accumulator_count, out, out_base, row_base, col_base, m_extent, n_extent, exact, true);
1048         return;
1049     }
1050 
1051     const loop_step = try logical.constantIndex(tile.bk);
1052     var scope = try logical.forScope(split_lower, split_upper, loop_step, init_args[0..accumulator_count], result_types[0..accumulator_count]);
1053     errdefer scope.abort();
1054     const k0 = scope.inductionVar();
1055 
1056     if (vectorized) {
1057         try stageVectorized(logical, stage_ctx, k0, tile_a, tile_b, false);
1058     } else {
1059         try emitScalarStaging(logical, .{
1060             .tile = tile,
1061             .threads = threads,
1062             .thread_id = thread_id,
1063             .k0 = k0,
1064             .row_block = row_block,
1065             .col_block = col_block,
1066             .lhs = lhs,
1067             .rhs = rhs,
1068             .tile_a = tile_a,
1069             .tile_b = tile_b,
1070             .a_pitch_extent = a_pitch_extent,
1071             .bk_extent = bk_extent,
1072             .bn_extent = bn_extent,
1073             .m_edge = m_edge,
1074             .n_edge = n_edge,
1075             .k_edge = k_edge,
1076             .m_extent = m_extent,
1077             .n_extent = n_extent,
1078             .k_extent = k_extent,
1079             .lhs_base = lhs_base,
1080             .rhs_base = rhs_base,
1081             .zero_scalar = zero_scalar,
1082             .exact = exact,
1083         });
1084     }
1085     try logical.barrier(.block);
1086 
1087     var accumulators: [max_block_tile_accumulators]kernel_root.Value = undefined;
1088     for (accumulators[0..accumulator_count], 0..) |*accumulator, index| {
1089         accumulator.* = scope.iterArg(index) orelse return error.UnsupportedOperation;
1090     }
1091 
1092     try emitComputeSlice(logical, slice_ctx, tile_a, tile_b, accumulators[0..accumulator_count], vectorized);
1093     try logical.barrier(.block);
1094     try scope.leave(accumulators[0..accumulator_count]);
1095     try emitTiledEpilogue(logical, ctx, &scope, accumulator_count, out, out_base, row_base, col_base, m_extent, n_extent, exact, vectorized);
1096 }
1097 
1098 const StageContext = struct {
1099     tile: kernelization_model.DotGeneralBlockTile,
1100     threads: u32,
1101     thread_id: kernel_root.Value,
1102     row_block: kernel_root.Value,
1103     col_block: kernel_root.Value,
1104     lhs: kernel_root.Value,
1105     rhs: kernel_root.Value,
1106     a_pitch_extent: kernel_root.Value,
1107     bn_extent: kernel_root.Value,
1108     n_extent: kernel_root.Value,
1109     k_extent: kernel_root.Value,
1110     lhs_base: kernel_root.Value,
1111     rhs_base: kernel_root.Value,
1112 };
1113 
1114 const SliceContext = struct {
1115     tile: kernelization_model.DotGeneralBlockTile,
1116     a_pitch: i64,
1117     row_local: kernel_root.Value,
1118     col_local: kernel_root.Value,
1119 };
1120 
1121 fn load_partial_b_tile_quad(inner: anytype, guard_ctx: anytype) !void {
1122     const quad_value = try inner.loadVector(guard_ctx.rhs, guard_ctx.global_index, 4);
1123     try inner.storeIndex(quad_value, guard_ctx.tile_b, guard_ctx.shared_index);
1124 }
1125 
1126 fn stageVectorized(
1127     logical: anytype,
1128     ctx: StageContext,
1129     k_slice: kernel_root.Value,
1130     tile_a: kernel_root.Value,
1131     tile_b: kernel_root.Value,
1132     clamp: bool,
1133 ) !void {
1134     const tile = ctx.tile;
1135     const k_quad_limit = try logical.sub(ctx.k_extent, try logical.constantIndex(4));
1136 
1137     const a_quads_per_thread = (tile.bm * tile.bk) / (ctx.threads * 4);
1138     const a_quads_per_row = tile.bk / 4;
1139     var quad: u32 = 0;
1140     while (quad < a_quads_per_thread) : (quad += 1) {
1141         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, quad) * ctx.threads));
1142         const quads_per_row_extent = try logical.constantIndex(a_quads_per_row);
1143         const local_row = try logical.div(flat, quads_per_row_extent);
1144         const quad_in_row = try logical.sub(flat, try logical.mul(local_row, quads_per_row_extent));
1145         const local_k = try logical.mul(quad_in_row, try logical.constantIndex(4));
1146         const global_row = try logical.add(ctx.row_block, local_row);
1147         var global_k = try logical.add(k_slice, local_k);
1148         if (clamp) global_k = try logical.min(global_k, k_quad_limit);
1149         const global_index = try logical.add(try logical.add(ctx.lhs_base, try logical.mul(global_row, ctx.k_extent)), global_k);
1150         const quad_value = try logical.loadVector(ctx.lhs, global_index, 4);
1151         var lane: u32 = 0;
1152         while (lane < 4) : (lane += 1) {
1153             const scalar = try logical.extractLane(quad_value, lane, .f32);
1154             const shared_row = try logical.add(local_k, try logical.constantIndex(lane));
1155             const shared_index = try logical.add(try logical.mul(shared_row, ctx.a_pitch_extent), local_row);
1156             try logical.storeIndex(scalar, tile_a, shared_index);
1157         }
1158     }
1159 
1160     const total_b_quads = (tile.bk * tile.bn) / 4;
1161     const b_quads_per_row = tile.bn / 4;
1162     var issued: u32 = 0;
1163     quad = 0;
1164     while (issued < total_b_quads) : (quad += 1) {
1165         const pass_quads = @min(ctx.threads, total_b_quads - issued);
1166         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, quad) * ctx.threads));
1167         const quads_per_row_extent = try logical.constantIndex(b_quads_per_row);
1168         const local_k = try logical.div(flat, quads_per_row_extent);
1169         const quad_in_row = try logical.sub(flat, try logical.mul(local_k, quads_per_row_extent));
1170         const local_col = try logical.mul(quad_in_row, try logical.constantIndex(4));
1171         var global_k = try logical.add(k_slice, local_k);
1172         if (clamp) global_k = try logical.min(global_k, k_quad_limit);
1173         const global_col = try logical.add(ctx.col_block, local_col);
1174         const global_index = try logical.add(try logical.add(ctx.rhs_base, try logical.mul(global_k, ctx.n_extent)), global_col);
1175         if (pass_quads == ctx.threads) {
1176             const quad_value = try logical.loadVector(ctx.rhs, global_index, 4);
1177             const shared_index = try logical.add(try logical.mul(local_k, ctx.bn_extent), local_col);
1178             try logical.storeIndex(quad_value, tile_b, shared_index);
1179         } else {
1180             const shared_index = try logical.add(try logical.mul(local_k, ctx.bn_extent), local_col);
1181             const in_pass = try logical.compare(.lt, flat, try logical.constantIndex(total_b_quads));
1182             try logical.guardDo(in_pass, .{
1183                 .tile_b = tile_b,
1184                 .shared_index = shared_index,
1185                 .rhs = ctx.rhs,
1186                 .global_index = global_index,
1187             }, load_partial_b_tile_quad);
1188         }
1189         issued += pass_quads;
1190     }
1191 }
1192 
1193 const max_stage_quads = 4;
1194 const max_pipeline_stages = 5;
1195 
1196 const StagedQuads = struct {
1197     values: [max_stage_quads]kernel_root.Value,
1198     count: u32,
1199 };
1200 
1201 fn loadATileQuads(logical: anytype, ctx: StageContext, k_slice: kernel_root.Value, clamp: bool) !StagedQuads {
1202     const tile = ctx.tile;
1203     const k_quad_limit = try logical.sub(ctx.k_extent, try logical.constantIndex(4));
1204     const a_quads_per_thread = (tile.bm * tile.bk) / (ctx.threads * 4);
1205     const a_quads_per_row = tile.bk / 4;
1206     if (a_quads_per_thread > max_stage_quads) return error.UnsupportedOperation;
1207     var staged = StagedQuads{ .values = undefined, .count = a_quads_per_thread };
1208     var quad: u32 = 0;
1209     while (quad < a_quads_per_thread) : (quad += 1) {
1210         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, quad) * ctx.threads));
1211         const quads_per_row_extent = try logical.constantIndex(a_quads_per_row);
1212         const local_row = try logical.div(flat, quads_per_row_extent);
1213         const quad_in_row = try logical.sub(flat, try logical.mul(local_row, quads_per_row_extent));
1214         const local_k = try logical.mul(quad_in_row, try logical.constantIndex(4));
1215         const global_row = try logical.add(ctx.row_block, local_row);
1216         var global_k = try logical.add(k_slice, local_k);
1217         if (clamp) global_k = try logical.min(global_k, k_quad_limit);
1218         const global_index = try logical.add(try logical.add(ctx.lhs_base, try logical.mul(global_row, ctx.k_extent)), global_k);
1219         staged.values[quad] = try logical.loadVector(ctx.lhs, global_index, 4);
1220     }
1221     return staged;
1222 }
1223 
1224 fn storeATileQuads(logical: anytype, ctx: StageContext, staged: StagedQuads, tile_a: kernel_root.Value) !void {
1225     const tile = ctx.tile;
1226     const a_quads_per_row = tile.bk / 4;
1227     var quad: u32 = 0;
1228     while (quad < staged.count) : (quad += 1) {
1229         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, quad) * ctx.threads));
1230         const quads_per_row_extent = try logical.constantIndex(a_quads_per_row);
1231         const local_row = try logical.div(flat, quads_per_row_extent);
1232         const quad_in_row = try logical.sub(flat, try logical.mul(local_row, quads_per_row_extent));
1233         const local_k = try logical.mul(quad_in_row, try logical.constantIndex(4));
1234         var lane: u32 = 0;
1235         while (lane < 4) : (lane += 1) {
1236             const scalar = try logical.extractLane(staged.values[quad], lane, .f32);
1237             const shared_row = try logical.add(local_k, try logical.constantIndex(lane));
1238             const shared_index = try logical.add(try logical.mul(shared_row, ctx.a_pitch_extent), local_row);
1239             try logical.storeIndex(scalar, tile_a, shared_index);
1240         }
1241     }
1242 }
1243 
1244 fn copy_partial_b_tile_quad(inner: anytype, guard_ctx: anytype) !void {
1245     try inner.asyncCopyShared(
1246         guard_ctx.tile_b,
1247         guard_ctx.shared_index,
1248         guard_ctx.rhs,
1249         guard_ctx.global_index,
1250         16,
1251     );
1252 }
1253 
1254 fn stageBTileAsync(logical: anytype, ctx: StageContext, k_slice: kernel_root.Value, tile_b: kernel_root.Value, clamp: bool) !void {
1255     const tile = ctx.tile;
1256     const k_row_limit = try logical.sub(ctx.k_extent, try logical.constantIndex(1));
1257     const total_quads = (tile.bk * tile.bn) / 4;
1258     const b_quads_per_row = tile.bn / 4;
1259     var issued: u32 = 0;
1260     var quad: u32 = 0;
1261     while (issued < total_quads) : (quad += 1) {
1262         const pass_quads = @min(ctx.threads, total_quads - issued);
1263         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, quad) * ctx.threads));
1264         const quads_per_row_extent = try logical.constantIndex(b_quads_per_row);
1265         const local_k = try logical.div(flat, quads_per_row_extent);
1266         const quad_in_row = try logical.sub(flat, try logical.mul(local_k, quads_per_row_extent));
1267         const local_col = try logical.mul(quad_in_row, try logical.constantIndex(4));
1268         var global_k = try logical.add(k_slice, local_k);
1269         if (clamp) global_k = try logical.min(global_k, k_row_limit);
1270         const global_col = try logical.add(ctx.col_block, local_col);
1271         const global_index = try logical.add(try logical.add(ctx.rhs_base, try logical.mul(global_k, ctx.n_extent)), global_col);
1272         const shared_index = try logical.add(try logical.mul(local_k, ctx.bn_extent), local_col);
1273         if (pass_quads == ctx.threads) {
1274             try logical.asyncCopyShared(tile_b, shared_index, ctx.rhs, global_index, 16);
1275         } else {
1276             const in_pass = try logical.compare(.lt, flat, try logical.constantIndex(total_quads));
1277             try logical.guardDo(in_pass, .{
1278                 .tile_b = tile_b,
1279                 .shared_index = shared_index,
1280                 .rhs = ctx.rhs,
1281                 .global_index = global_index,
1282             }, copy_partial_b_tile_quad);
1283         }
1284         issued += pass_quads;
1285     }
1286 }
1287 
1288 fn emitComputeSlice(
1289     logical: anytype,
1290     ctx: SliceContext,
1291     tile_a: kernel_root.Value,
1292     tile_b: kernel_root.Value,
1293     accumulators: []kernel_root.Value,
1294     vectorized: bool,
1295 ) !void {
1296     const tile = ctx.tile;
1297     var kk: u32 = 0;
1298     while (kk < tile.bk) : (kk += 1) {
1299         var fragment_a: [max_block_tile_fragments]kernel_root.Value = undefined;
1300         var fragment_b: [max_block_tile_fragments]kernel_root.Value = undefined;
1301         const a_row_offset = try logical.constantIndex(@as(i64, kk) * ctx.a_pitch);
1302         const b_row_offset = try logical.constantIndex(@as(i64, kk) * tile.bn);
1303         if (vectorized) {
1304             const m_quad_stride = @divExact(@as(i64, tile.bm) * 4, tile.tm);
1305             const n_quad_stride = @divExact(@as(i64, tile.bn) * 4, tile.tn);
1306             var quad: u32 = 0;
1307             while (quad < tile.tm / 4) : (quad += 1) {
1308                 const local = try logical.add(ctx.row_local, try logical.constantIndex(@as(i64, quad) * m_quad_stride));
1309                 const quad_value = try logical.loadVector(tile_a, try logical.add(a_row_offset, local), 4);
1310                 var lane: u32 = 0;
1311                 while (lane < 4) : (lane += 1) {
1312                     fragment_a[quad * 4 + lane] = try logical.extractLane(quad_value, lane, .f32);
1313                 }
1314             }
1315             quad = 0;
1316             while (quad < tile.tn / 4) : (quad += 1) {
1317                 const local = try logical.add(ctx.col_local, try logical.constantIndex(@as(i64, quad) * n_quad_stride));
1318                 const quad_value = try logical.loadVector(tile_b, try logical.add(b_row_offset, local), 4);
1319                 var lane: u32 = 0;
1320                 while (lane < 4) : (lane += 1) {
1321                     fragment_b[quad * 4 + lane] = try logical.extractLane(quad_value, lane, .f32);
1322                 }
1323             }
1324         } else {
1325             var i_frag: u32 = 0;
1326             while (i_frag < tile.tm) : (i_frag += 1) {
1327                 const local = try logical.add(ctx.row_local, try logical.constantIndex(i_frag));
1328                 fragment_a[i_frag] = try logical.loadIndex(tile_a, try logical.add(a_row_offset, local));
1329             }
1330             var j_frag: u32 = 0;
1331             while (j_frag < tile.tn) : (j_frag += 1) {
1332                 const local = try logical.add(ctx.col_local, try logical.constantIndex(j_frag));
1333                 fragment_b[j_frag] = try logical.loadIndex(tile_b, try logical.add(b_row_offset, local));
1334             }
1335         }
1336         var i: u32 = 0;
1337         while (i < tile.tm) : (i += 1) {
1338             var jj: u32 = 0;
1339             while (jj < tile.tn) : (jj += 1) {
1340                 const slot = i * tile.tn + jj;
1341                 const product = try logical.mul(fragment_a[i], fragment_b[jj]);
1342                 accumulators[slot] = try logical.add(accumulators[slot], product);
1343             }
1344         }
1345     }
1346 }
1347 
1348 fn store_block_tile_result(inner: anytype, nested_ctx: anytype) !void {
1349     try inner.storeIndex(nested_ctx.value, nested_ctx.out, nested_ctx.address);
1350 }
1351 
1352 fn guard_block_tile_store_column(inner: anytype, outer_ctx: anytype) !void {
1353     try inner.guardDo(outer_ctx.col_ok, outer_ctx, store_block_tile_result);
1354 }
1355 
1356 fn emitTiledEpilogue(
1357     logical: anytype,
1358     ctx: anytype,
1359     scope: anytype,
1360     accumulator_count: usize,
1361     out: kernel_root.Value,
1362     out_base: kernel_root.Value,
1363     row_base: kernel_root.Value,
1364     col_base: kernel_root.Value,
1365     m_extent: kernel_root.Value,
1366     n_extent: kernel_root.Value,
1367     exact: bool,
1368     vectorized: bool,
1369 ) !void {
1370     const tile: kernelization_model.DotGeneralBlockTile = ctx.tile;
1371     const m_quad_stride = @divExact(@as(i64, tile.bm) * 4, tile.tm);
1372     const n_quad_stride = @divExact(@as(i64, tile.bn) * 4, tile.tn);
1373     var results: [max_block_tile_accumulators]kernel_root.Value = undefined;
1374     for (results[0..accumulator_count], 0..) |*result, index| {
1375         result.* = scope.result(index) orelse return error.UnsupportedOperation;
1376     }
1377 
1378     var i: u32 = 0;
1379     while (i < tile.tm) : (i += 1) {
1380         const row_step: i64 = if (vectorized) @as(i64, i / 4) * m_quad_stride + i % 4 else @as(i64, i);
1381         const row = try logical.add(row_base, try logical.constantIndex(row_step));
1382         const row_offset = try logical.add(out_base, try logical.mul(row, n_extent));
1383         var j: u32 = 0;
1384         while (j < tile.tn) : (j += 1) {
1385             const col_step: i64 = if (vectorized) @as(i64, j / 4) * n_quad_stride + j % 4 else @as(i64, j);
1386             const col = try logical.add(col_base, try logical.constantIndex(col_step));
1387             const address = try logical.add(row_offset, col);
1388             const value = try applyEpilogueChain(logical, ctx.epilogue, ctx.abi, results[i * tile.tn + j], address);
1389             if (exact) {
1390                 if (tile.splits > 1) {
1391                     _ = try logical.atomicRmw(.add, value, out, address);
1392                 } else {
1393                     try logical.storeIndex(value, out, address);
1394                 }
1395             } else {
1396                 const row_ok = try logical.compare(.lt, row, m_extent);
1397                 const col_ok = try logical.compare(.lt, col, n_extent);
1398                 try logical.guardDo(row_ok, .{
1399                     .col_ok = col_ok,
1400                     .value = value,
1401                     .out = out,
1402                     .address = address,
1403                 }, guard_block_tile_store_column);
1404             }
1405         }
1406     }
1407 }
1408 
1409 fn emitScalarStaging(logical: anytype, ctx: anytype) !void {
1410     const tile: kernelization_model.DotGeneralBlockTile = ctx.tile;
1411     const a_loads_per_thread = (tile.bm * tile.bk) / ctx.threads;
1412     var stage: u32 = 0;
1413     while (stage < a_loads_per_thread) : (stage += 1) {
1414         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, stage) * ctx.threads));
1415         const local_row = try logical.div(flat, ctx.bk_extent);
1416         const local_k = try logical.sub(flat, try logical.mul(local_row, ctx.bk_extent));
1417         const global_row = try logical.add(ctx.row_block, local_row);
1418         const global_k = try logical.add(ctx.k0, local_k);
1419         const value = try stagedLoad(logical, ctx.lhs, .{
1420             .base = ctx.lhs_base,
1421             .major = global_row,
1422             .major_edge = ctx.m_edge,
1423             .major_extent = ctx.m_extent,
1424             .pitch = ctx.k_extent,
1425             .minor = global_k,
1426             .minor_edge = ctx.k_edge,
1427             .minor_extent = ctx.k_extent,
1428             .zero = ctx.zero_scalar,
1429             .exact = ctx.exact,
1430         });
1431         const shared_index = try logical.add(try logical.mul(local_k, ctx.a_pitch_extent), local_row);
1432         try logical.storeIndex(value, ctx.tile_a, shared_index);
1433     }
1434 
1435     const b_loads_per_thread = (tile.bk * tile.bn) / ctx.threads;
1436     stage = 0;
1437     while (stage < b_loads_per_thread) : (stage += 1) {
1438         const flat = try logical.add(ctx.thread_id, try logical.constantIndex(@as(i64, stage) * ctx.threads));
1439         const local_k = try logical.div(flat, ctx.bn_extent);
1440         const local_col = try logical.sub(flat, try logical.mul(local_k, ctx.bn_extent));
1441         const global_k = try logical.add(ctx.k0, local_k);
1442         const global_col = try logical.add(ctx.col_block, local_col);
1443         const value = try stagedLoad(logical, ctx.rhs, .{
1444             .base = ctx.rhs_base,
1445             .major = global_k,
1446             .major_edge = ctx.k_edge,
1447             .major_extent = ctx.k_extent,
1448             .pitch = ctx.n_extent,
1449             .minor = global_col,
1450             .minor_edge = ctx.n_edge,
1451             .minor_extent = ctx.n_extent,
1452             .zero = ctx.zero_scalar,
1453             .exact = ctx.exact,
1454         });
1455         const shared_index = try logical.add(try logical.mul(local_k, ctx.bn_extent), local_col);
1456         try logical.storeIndex(value, ctx.tile_b, shared_index);
1457     }
1458 }
1459 
1460 fn stagedLoad(logical: anytype, buffer: kernel_root.Value, spec: anytype) !kernel_root.Value {
1461     if (spec.exact) {
1462         const offset = try logical.mul(spec.major, spec.pitch);
1463         const address = try logical.add(try logical.add(spec.base, offset), spec.minor);
1464         return logical.loadIndex(buffer, address);
1465     }
1466     const major_ok = try logical.compare(.lt, spec.major, spec.major_extent);
1467     const minor_ok = try logical.compare(.lt, spec.minor, spec.minor_extent);
1468     const major_clamped = try logical.min(spec.major, spec.major_edge);
1469     const minor_clamped = try logical.min(spec.minor, spec.minor_edge);
1470     const offset = try logical.mul(major_clamped, spec.pitch);
1471     const address = try logical.add(try logical.add(spec.base, offset), minor_clamped);
1472     const loaded = try logical.loadIndex(buffer, address);
1473     const minor_masked = try logical.select(minor_ok, loaded, spec.zero);
1474     return logical.select(major_ok, minor_masked, spec.zero);
1475 }
1476 
1477 fn accumulate_dot_product(
1478     fold_inner: anytype,
1479     k_index: kernel_root.Value,
1480     current: kernel_root.Value,
1481     fold_ctx: anytype,
1482 ) !kernel_root.Value {
1483     const m_stride = try fold_inner.constantIndex(fold_ctx.dims.m);
1484     const k_stride = try fold_inner.constantIndex(fold_ctx.dims.k);
1485     const n_stride = try fold_inner.constantIndex(fold_ctx.dims.n);
1486     const lhs_batch_rows = try fold_inner.mul(fold_ctx.batch, m_stride);
1487     const lhs_row = try fold_inner.add(lhs_batch_rows, fold_ctx.row);
1488     const lhs_row_offset = try fold_inner.mul(lhs_row, k_stride);
1489     const lhs_index = try fold_inner.add(lhs_row_offset, k_index);
1490     const rhs_batch_rows = try fold_inner.mul(fold_ctx.batch, k_stride);
1491     const rhs_row = try fold_inner.add(rhs_batch_rows, k_index);
1492     const rhs_row_offset = try fold_inner.mul(rhs_row, n_stride);
1493     const rhs_index = try fold_inner.add(rhs_row_offset, fold_ctx.col);
1494     var lhs_value = try fold_inner.loadIndex(fold_ctx.lhs, lhs_index);
1495     var rhs_value = try fold_inner.loadIndex(fold_ctx.rhs, rhs_index);
1496     if (fold_ctx.widen) {
1497         lhs_value = try fold_inner.cast(lhs_value, .f32);
1498         rhs_value = try fold_inner.cast(rhs_value, .f32);
1499     }
1500     const product = try fold_inner.mul(lhs_value, rhs_value);
1501     return fold_inner.add(current, product);
1502 }
1503 
1504 fn emit_dot_index(inner: anytype, index: kernel_root.Index3D, each_ctx: anytype) !void {
1505     const lower_bound = try inner.constantIndex(0);
1506     const upper = try inner.constantIndex(each_ctx.dims.k);
1507     const step = try inner.constantIndex(1);
1508     const widen = each_ctx.dtype == .f16;
1509     const initial = try dotInitialValue(inner, each_ctx.dtype);
1510     const result = try inner.fold(lower_bound, upper, step, initial, .{
1511         .dims = each_ctx.dims,
1512         .lhs = each_ctx.lhs,
1513         .rhs = each_ctx.rhs,
1514         .row = index.y.index,
1515         .col = index.x.index,
1516         .batch = index.z.index,
1517         .widen = widen,
1518     }, accumulate_dot_product);
1519     const stored = if (widen) try inner.cast(result, .f16) else result;
1520     const m_stride = try inner.constantIndex(each_ctx.dims.m);
1521     const n_stride = try inner.constantIndex(each_ctx.dims.n);
1522     const out_batch_rows = try inner.mul(index.z.index, m_stride);
1523     const out_row = try inner.add(out_batch_rows, index.y.index);
1524     const out_row_offset = try inner.mul(out_row, n_stride);
1525     const out_index = try inner.add(out_row_offset, index.x.index);
1526     const finished = try applyEpilogueChain(
1527         inner,
1528         each_ctx.epilogue,
1529         each_ctx.abi,
1530         stored,
1531         out_index,
1532     );
1533     try inner.storeIndex(finished, each_ctx.out, out_index);
1534 }
1535 
1536 fn emitDotBody(logical: anytype, ctx: anytype) !void {
1537     const output_index = try logical.index3D(.{
1538         .x = kernel_root.logical.axis("col", ctx.dims.n),
1539         .y = kernel_root.logical.axis("row", ctx.dims.m),
1540         .z = kernel_root.logical.axis("batch", ctx.dims.batch),
1541     });
1542     try logical.guardIndex3DDo(output_index, .{
1543         .dims = ctx.dims,
1544         .dtype = ctx.dtype,
1545         .out = ctx.abi.output(logical),
1546         .lhs = ctx.abi.lhs(logical),
1547         .rhs = ctx.abi.rhs(logical),
1548         .abi = ctx.abi,
1549         .epilogue = ctx.epilogue,
1550     }, emit_dot_index);
1551 }
1552 
1553 fn emitFlatDotBody(logical: anytype, ctx: anytype) !void {
1554     const output_elements = @as(u64, ctx.dims.m) * @as(u64, ctx.dims.n) * @as(u64, ctx.dims.batch);
1555     const output_index = try logical.index1D("out", output_elements);
1556     const flat = output_index.index;
1557     const n_extent = try logical.constantIndex(ctx.dims.n);
1558     const mn_extent = try logical.constantIndex(@as(i64, ctx.dims.m) * @as(i64, ctx.dims.n));
1559     const batch = try logical.div(flat, mn_extent);
1560     const batch_base = try logical.mul(batch, mn_extent);
1561     const in_batch = try logical.sub(flat, batch_base);
1562     const row = try logical.div(in_batch, n_extent);
1563     const row_base = try logical.mul(row, n_extent);
1564     const col = try logical.sub(in_batch, row_base);
1565 
1566     const widen = ctx.dtype == .f16;
1567     var result = try dotInitialValue(logical, ctx.dtype);
1568     var k: u32 = 0;
1569     while (k < ctx.dims.k) : (k += 1) {
1570         result = try accumulateDotProduct(
1571             logical,
1572             ctx.dims,
1573             ctx.abi.lhs(logical),
1574             ctx.abi.rhs(logical),
1575             row,
1576             col,
1577             batch,
1578             result,
1579             try logical.constantIndex(k),
1580             widen,
1581         );
1582     }
1583     const stored = if (widen) try logical.cast(result, .f16) else result;
1584     const finished = try applyEpilogueChain(logical, ctx.epilogue, ctx.abi, stored, flat);
1585     try logical.storeIndex(finished, ctx.abi.output(logical), flat);
1586 }
1587 
1588 fn accumulateDotProduct(
1589     builder: anytype,
1590     dims: DotGeneralStaticDims,
1591     lhs: kernel_root.Value,
1592     rhs: kernel_root.Value,
1593     row: kernel_root.Value,
1594     col: kernel_root.Value,
1595     batch: kernel_root.Value,
1596     current: kernel_root.Value,
1597     k_index: kernel_root.Value,
1598     widen: bool,
1599 ) !kernel_root.Value {
1600     const m_stride = try builder.constantIndex(dims.m);
1601     const k_stride = try builder.constantIndex(dims.k);
1602     const n_stride = try builder.constantIndex(dims.n);
1603     const lhs_batch_rows = try builder.mul(batch, m_stride);
1604     const lhs_row = try builder.add(lhs_batch_rows, row);
1605     const lhs_row_offset = try builder.mul(lhs_row, k_stride);
1606     const lhs_index = try builder.add(lhs_row_offset, k_index);
1607     const rhs_batch_rows = try builder.mul(batch, k_stride);
1608     const rhs_row = try builder.add(rhs_batch_rows, k_index);
1609     const rhs_row_offset = try builder.mul(rhs_row, n_stride);
1610     const rhs_index = try builder.add(rhs_row_offset, col);
1611     var lhs_value = try builder.loadIndex(lhs, lhs_index);
1612     var rhs_value = try builder.loadIndex(rhs, rhs_index);
1613     if (widen) {
1614         lhs_value = try builder.cast(lhs_value, .f32);
1615         rhs_value = try builder.cast(rhs_value, .f32);
1616     }
1617     const product = try builder.mul(lhs_value, rhs_value);
1618     return builder.add(current, product);
1619 }
1620 
1621 fn validateCanonicalDotLayout(op: *ir.Operation, lhs_rank: usize) gpu.BackendError!void {
1622     var lhs_contract_buffer: [8]i64 = undefined;
1623     const lhs_contract = try common.readI64ListAttrBounded(op, "lhs_contract", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("lhs_contract"), &lhs_contract_buffer);
1624     var rhs_contract_buffer: [8]i64 = undefined;
1625     const rhs_contract = try common.readI64ListAttrBounded(op, "rhs_contract", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("rhs_contract"), &rhs_contract_buffer);
1626     var lhs_batch_buffer: [8]i64 = undefined;
1627     const lhs_batch = try common.readI64ListAttrBounded(op, "lhs_batch", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("lhs_batch"), &lhs_batch_buffer);
1628     var rhs_batch_buffer: [8]i64 = undefined;
1629     const rhs_batch = try common.readI64ListAttrBounded(op, "rhs_batch", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("rhs_batch"), &rhs_batch_buffer);
1630 
1631     switch (lhs_rank) {
1632         2 => {
1633             if (lhs_batch.len != 0 or rhs_batch.len != 0) return error.UnsupportedOperation;
1634             if (lhs_contract.len != 1 or rhs_contract.len != 1) return error.UnsupportedOperation;
1635             if (lhs_contract[0] != 1 or rhs_contract[0] != 0) return error.UnsupportedOperation;
1636         },
1637         3 => {
1638             if (lhs_batch.len != 1 or rhs_batch.len != 1) return error.UnsupportedOperation;
1639             if (lhs_batch[0] != 0 or rhs_batch[0] != 0) return error.UnsupportedOperation;
1640             if (lhs_contract.len != 1 or rhs_contract.len != 1) return error.UnsupportedOperation;
1641             if (lhs_contract[0] != 2 or rhs_contract[0] != 1) return error.UnsupportedOperation;
1642         },
1643         else => return error.UnsupportedOperation,
1644     }
1645 }
1646 
1647 fn dotGeneralStaticDims(
1648     lhs_slot: bufferization.BufferSlot,
1649     rhs_slot: bufferization.BufferSlot,
1650     output_slot: bufferization.BufferSlot,
1651 ) gpu.BackendError!DotGeneralStaticDims {
1652     if (lhs_slot.dims.len == 2 and rhs_slot.dims.len == 2 and output_slot.dims.len == 2) {
1653         return dotGeneralRank2StaticDims(lhs_slot, rhs_slot, output_slot);
1654     }
1655     if (lhs_slot.dims.len == 3 and rhs_slot.dims.len == 3 and output_slot.dims.len == 3) {
1656         return dotGeneralRank3StaticDims(lhs_slot, rhs_slot, output_slot);
1657     }
1658     return error.UnsupportedOperation;
1659 }
1660 
1661 fn dotGeneralRank2StaticDims(
1662     lhs_slot: bufferization.BufferSlot,
1663     rhs_slot: bufferization.BufferSlot,
1664     output_slot: bufferization.BufferSlot,
1665 ) gpu.BackendError!DotGeneralStaticDims {
1666     const m = try staticPositiveDimU32(lhs_slot.dims[0]);
1667     const k = try staticPositiveDimU32(lhs_slot.dims[1]);
1668     const rhs_k = try staticPositiveDimU32(rhs_slot.dims[0]);
1669     const n = try staticPositiveDimU32(rhs_slot.dims[1]);
1670     const output_m = try staticPositiveDimU32(output_slot.dims[0]);
1671     const output_n = try staticPositiveDimU32(output_slot.dims[1]);
1672     if (rhs_k != k or output_m != m or output_n != n) return error.InvalidArtifact;
1673     return .{ .m = m, .n = n, .k = k };
1674 }
1675 
1676 fn dotGeneralRank3StaticDims(
1677     lhs_slot: bufferization.BufferSlot,
1678     rhs_slot: bufferization.BufferSlot,
1679     output_slot: bufferization.BufferSlot,
1680 ) gpu.BackendError!DotGeneralStaticDims {
1681     const batch = try staticPositiveDimU32(lhs_slot.dims[0]);
1682     const m = try staticPositiveDimU32(lhs_slot.dims[1]);
1683     const k = try staticPositiveDimU32(lhs_slot.dims[2]);
1684     const rhs_batch = try staticPositiveDimU32(rhs_slot.dims[0]);
1685     const rhs_k = try staticPositiveDimU32(rhs_slot.dims[1]);
1686     const n = try staticPositiveDimU32(rhs_slot.dims[2]);
1687     const output_batch = try staticPositiveDimU32(output_slot.dims[0]);
1688     const output_m = try staticPositiveDimU32(output_slot.dims[1]);
1689     const output_n = try staticPositiveDimU32(output_slot.dims[2]);
1690     if (rhs_batch != batch or output_batch != batch or rhs_k != k or output_m != m or output_n != n) return error.InvalidArtifact;
1691     return .{ .m = m, .n = n, .k = k, .batch = batch };
1692 }
1693 
1694 fn validateDotGeneralDTypes(
1695     lhs_slot: bufferization.BufferSlot,
1696     rhs_slot: bufferization.BufferSlot,
1697     output_slot: bufferization.BufferSlot,
1698     work: schedule_planning.ScheduleWorkItem,
1699 ) gpu.BackendError!void {
1700     if (lhs_slot.dtype != rhs_slot.dtype) return error.CapabilityMismatch;
1701     switch (lhs_slot.dtype) {
1702         .f32 => if (output_slot.dtype != .f32) return error.CapabilityMismatch,
1703         .f16 => if (output_slot.dtype != .f16) return error.CapabilityMismatch,
1704         .i32 => if (output_slot.dtype != .i32) return error.CapabilityMismatch,
1705         .u32 => if (output_slot.dtype != .u32) return error.CapabilityMismatch,
1706         else => return error.CapabilityMismatch,
1707     }
1708     if (work.dtype != output_slot.dtype) return error.InvalidArtifact;
1709 }
1710 
1711 fn dotLoweringSupported(input_dtype: choir_abi.DType, output_dtype: choir_abi.DType) bool {
1712     if (input_dtype == .f32 and output_dtype == .f32) return true;
1713     if (input_dtype == .f16 and output_dtype == .f16) return true;
1714     if (input_dtype == .u32 and output_dtype == .u32) return true;
1715     return input_dtype == .i32 and output_dtype == .i32;
1716 }
1717 
1718 fn testBufferSlot(dtype: choir_abi.DType, dims: []i64) bufferization.BufferSlot {
1719     return .{
1720         .id = 0,
1721         .value = undefined,
1722         .producer = null,
1723         .function = null,
1724         .role = .{},
1725         .dtype = dtype,
1726         .dims = dims,
1727         .element_count = null,
1728         .row_major_strides = null,
1729         .byte_size = null,
1730     };
1731 }
1732 
1733 test "dot general static dims accept edge tile shapes" {
1734     var lhs_dims = [_]i64{ 1, 32 };
1735     var rhs_dims = [_]i64{ 32, 24 };
1736     var out_dims = [_]i64{ 1, 24 };
1737 
1738     const dims = try dotGeneralStaticDims(
1739         testBufferSlot(.f32, lhs_dims[0..]),
1740         testBufferSlot(.f32, rhs_dims[0..]),
1741         testBufferSlot(.f32, out_dims[0..]),
1742     );
1743 
1744     try std.testing.expectEqual(@as(u32, 1), dims.m);
1745     try std.testing.expectEqual(@as(u32, 24), dims.n);
1746     try std.testing.expectEqual(@as(u32, 32), dims.k);
1747 }
1748 
1749 test "dot general static dims accept batched matmul shapes" {
1750     var lhs_dims = [_]i64{ 8, 2, 4 };
1751     var rhs_dims = [_]i64{ 8, 4, 3 };
1752     var out_dims = [_]i64{ 8, 2, 3 };
1753 
1754     const dims = try dotGeneralStaticDims(
1755         testBufferSlot(.f32, lhs_dims[0..]),
1756         testBufferSlot(.f32, rhs_dims[0..]),
1757         testBufferSlot(.f32, out_dims[0..]),
1758     );
1759 
1760     try std.testing.expectEqual(@as(u32, 8), dims.batch);
1761     try std.testing.expectEqual(@as(u32, 2), dims.m);
1762     try std.testing.expectEqual(@as(u32, 3), dims.n);
1763     try std.testing.expectEqual(@as(u32, 4), dims.k);
1764 }