lib/accy/src/preparation/kernelization/lowering/reduction.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const alloc_fixed = @import("alloc_fixed");
5 const generated_abi = @import("abi.zig");
6 const generated_builder = @import("builder.zig");
7 const common = @import("common.zig");
8 const generated_elementwise = @import("elementwise.zig");
9 const generated_guard = @import("guard.zig");
10 const generated_name = @import("name.zig");
11 const generated_schedule = @import("schedule.zig");
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 i64_attr_list_stack_capacity = common.i64_attr_list_stack_capacity;
20 const bufferSlotById = common.bufferSlotById;
21 const readI64ListAttrBounded = common.readI64ListAttrBounded;
22 const constantPayloadForSlot = common.constantPayloadForSlot;
23 const externalInputIndex = common.externalInputIndex;
24 const staticElementCountU32 = common.staticElementCountU32;
25 const staticPositiveDimU32 = common.staticPositiveDimU32;
26 const dialectAttrPayload = common.dialectAttrPayload;
27 const mapKernelBuildError = common.mapKernelBuildError;
28 const isName = common.isName;
29
30 const ReductionStaticDims = kernelization_model.ReductionStaticDims;
31 const ReductionKind = kernelization_model.ReductionKind;
32 const ReductionDescription = kernelization_model.ReductionDescription;
33 const ReductionInit = kernelization_model.ReductionInit;
34 const ReductionInitValue = kernelization_model.ReductionInitValue;
35 const LoweredKernel = kernelization_model.LoweredKernel;
36
37 pub fn reductionDescriptionForWork(
38 outline: kernelization_model.KernelOutline,
39 work: schedule_planning.ScheduleWorkItem,
40 buffer_plan: *const bufferization.BufferPlanAnalysis,
41 ) gpu.BackendError!ReductionDescription {
42 if (work.kind != .reduction) return error.UnsupportedOperation;
43 if (work.ops.len == 0) return error.UnsupportedOperation;
44 const root = work.ops[work.ops.len - 1];
45 if (isName(root.name.name, dialect_mod.AccyDialect.ConcatenateOp.operation_name)) {
46 return concatReductionDescription(root, outline, work, buffer_plan);
47 }
48 const op = root;
49 if (!isName(op.name.name, dialect_mod.AccyDialect.ReduceOp.operation_name)) return error.UnsupportedOperation;
50 const operands = op.getOperandValues();
51 if (operands.len != 2) return error.UnsupportedOperation;
52
53 const fused = reductionWorkIsFused(work, buffer_plan);
54 if (!fused) {
55 const input_slot = buffer_plan.getSlot(operands[0]) orelse return error.InvalidArtifact;
56 if (externalInputIndex(outline, input_slot.id) != 0) return error.UnsupportedOperation;
57 }
58
59 var input_arena_buffer: [256]u8 = undefined;
60 var input_arena = alloc_fixed.FixedBuffer.init(input_arena_buffer[0..]);
61 const input_type = dialect_mod.decodeTensorType(input_arena.allocator(), operands[0].type) catch return error.InvalidArtifact;
62
63 const init_slot = buffer_plan.getSlot(operands[1]) orelse return error.InvalidArtifact;
64 const output_slot = bufferSlotById(buffer_plan, outline.output_slot_id) orelse return error.InvalidArtifact;
65 var axes_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
66 const axes = try readI64ListAttrBounded(op, "dimensions", dialect_mod.AccyDialect.ReduceOp.dialectAttrName("dimensions"), &axes_buffer);
67
68 const dims = try reductionStaticDims(input_type.dims, output_slot.dims, axes);
69 try validateReductionDTypes(input_type.dtype, output_slot.dtype, work);
70 const kind = try reductionKindForOp(op);
71 const init = try reductionInitValue(init_slot.*, input_type.dtype, outline);
72
73 return .{
74 .kind = kind,
75 .input_dtype = input_type.dtype,
76 .output_dtype = output_slot.dtype,
77 .dims = dims,
78 .init = init,
79 };
80 }
81
82 pub fn reductionWorkIsFused(
83 work: schedule_planning.ScheduleWorkItem,
84 buffer_plan: *const bufferization.BufferPlanAnalysis,
85 ) bool {
86 if (work.ops.len == 0) return false;
87 if (isName(work.ops[work.ops.len - 1].name.name, dialect_mod.AccyDialect.ConcatenateOp.operation_name)) return true;
88 if (work.ops.len > 1) return true;
89 const op = work.ops[work.ops.len - 1];
90 const operands = op.getOperandValues();
91 if (operands.len != 2) return false;
92 if (buffer_plan.getSlot(operands[0]) == null) return true;
93 const def_any = operands[0].getDefiningOp() orelse return false;
94 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
95 return isName(def_op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name);
96 }
97
98 pub fn lower(
99 allocator: std.mem.Allocator,
100 ir_ctx: *ir.Context,
101 outline: kernelization_model.KernelOutline,
102 work: schedule_planning.ScheduleWorkItem,
103 buffer_plan: *const bufferization.BufferPlanAnalysis,
104 format: ?gpu.ArtifactFormat,
105 ) common.LoweringError!LoweredKernel {
106 const desc = try reductionDescriptionForWork(outline, work, buffer_plan);
107
108 if (reductionWorkIsFused(work, buffer_plan)) {
109 return lowerFused(allocator, ir_ctx, outline, work, buffer_plan, desc, format);
110 }
111
112 const abi = generated_abi.reduction(desc.input_dtype, desc.output_dtype, desc.init == .input_buffer);
113
114 if (format == .cuda_ptx) if (kernelization_model.reductionAtomicPlanFor(desc.dims, desc.kind, desc.input_dtype, desc.init == .constant)) |atomic_plan| {
115 const entry_name = try generated_name.reductionAtomic(allocator, desc, atomic_plan, work.id);
116 errdefer allocator.free(entry_name);
117 var lowered = try generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flatThreads(atomic_plan.threads), .{
118 .abi = abi,
119 .kind = desc.kind,
120 .dtype = desc.input_dtype,
121 .dims = desc.dims,
122 .plan = atomic_plan,
123 }, emitAtomicReductionBody);
124 lowered.output_fill_pattern = reductionInitFillPattern(desc.init) orelse return error.InvalidArtifact;
125 lowered.body = .{ .reduction_atomic = atomic_plan };
126 return lowered;
127 };
128
129 if (format == .cuda_ptx) if (kernelization_model.reductionSingleBlockThreads(desc.dims)) |threads| {
130 const entry_name = try generated_name.reductionSingleBlock(allocator, desc, threads, work.id);
131 errdefer allocator.free(entry_name);
132 var lowered = try generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flatThreads(threads), .{
133 .abi = abi,
134 .kind = desc.kind,
135 .dtype = desc.input_dtype,
136 .dims = desc.dims,
137 .init = desc.init,
138 .threads = threads,
139 }, emitSingleBlockReductionBody);
140 lowered.body = .{ .reduction_single_block = threads };
141 return lowered;
142 };
143
144 if (format == .cuda_ptx and desc.init == .constant) if (kernelization_model.reductionWarpRowsThreads(desc.dims)) |_| {
145 return lowerFused(allocator, ir_ctx, outline, work, buffer_plan, desc, format);
146 };
147
148 const entry_name = try generated_name.reduction(allocator, desc, work.id);
149 errdefer allocator.free(entry_name);
150
151 return generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flat(), .{
152 .abi = abi,
153 .kind = desc.kind,
154 .dtype = desc.input_dtype,
155 .dims = desc.dims,
156 .init = desc.init,
157 }, emitReductionBody);
158 }
159
160 const max_prologue_depth = 32;
161
162 const max_concat_group = 8;
163
164 const ConcatMember = struct {
165 reduce: *ir.Operation,
166 input_value: *ir.Value,
167 kind: ReductionKind,
168 init: ReductionInit,
169 dims: ReductionStaticDims,
170 base: u32,
171 };
172
173 fn concatReductionMembers(
174 root: *ir.Operation,
175 outline: kernelization_model.KernelOutline,
176 buffer_plan: *const bufferization.BufferPlanAnalysis,
177 buffer: []ConcatMember,
178 ) gpu.BackendError![]ConcatMember {
179 const operands = root.getOperandValues();
180 if (operands.len < 2 or operands.len > buffer.len) return error.UnsupportedOperation;
181 var base: u32 = 0;
182 for (operands, 0..) |operand, index| {
183 const def_any = operand.getDefiningOp() orelse return error.UnsupportedOperation;
184 const reduce: *ir.Operation = @ptrCast(@alignCast(def_any));
185 if (!isName(reduce.name.name, dialect_mod.AccyDialect.ReduceOp.operation_name)) return error.UnsupportedOperation;
186 const reduce_operands = reduce.getOperandValues();
187 if (reduce_operands.len != 2) return error.UnsupportedOperation;
188
189 var input_arena_buffer: [256]u8 = undefined;
190 var input_arena = alloc_fixed.FixedBuffer.init(input_arena_buffer[0..]);
191 const input_type = dialect_mod.decodeTensorType(input_arena.allocator(), reduce_operands[0].type) catch return error.InvalidArtifact;
192 var output_arena_buffer: [256]u8 = undefined;
193 var output_arena = alloc_fixed.FixedBuffer.init(output_arena_buffer[0..]);
194 const output_type = dialect_mod.decodeTensorType(output_arena.allocator(), operand.type) catch return error.InvalidArtifact;
195
196 var axes_buffer: [i64_attr_list_stack_capacity]i64 = undefined;
197 const axes = try readI64ListAttrBounded(reduce, "dimensions", dialect_mod.AccyDialect.ReduceOp.dialectAttrName("dimensions"), &axes_buffer);
198 const dims = try reductionStaticDims(input_type.dims, output_type.dims, axes);
199 const init_slot = buffer_plan.getSlot(reduce_operands[1]) orelse return error.InvalidArtifact;
200 const init = try reductionInitValue(init_slot.*, input_type.dtype, outline);
201 if (init != .constant) return error.UnsupportedOperation;
202
203 buffer[index] = .{
204 .reduce = reduce,
205 .input_value = reduce_operands[0],
206 .kind = try reductionKindForOp(reduce),
207 .init = init,
208 .dims = dims,
209 .base = base,
210 };
211 base = std.math.add(u32, base, dims.output_element_count) catch return error.CapabilityMismatch;
212 }
213 for (buffer[1..operands.len]) |member| {
214 if (member.dims.output_element_count != buffer[0].dims.output_element_count) return error.UnsupportedOperation;
215 if (member.dims.input_rank != buffer[0].dims.input_rank) return error.UnsupportedOperation;
216 if (member.dims.axis != buffer[0].dims.axis) return error.UnsupportedOperation;
217 if (member.dims.cols != buffer[0].dims.cols) return error.UnsupportedOperation;
218 }
219 return buffer[0..operands.len];
220 }
221
222 fn concatReductionDescription(
223 root: *ir.Operation,
224 outline: kernelization_model.KernelOutline,
225 work: schedule_planning.ScheduleWorkItem,
226 buffer_plan: *const bufferization.BufferPlanAnalysis,
227 ) gpu.BackendError!ReductionDescription {
228 var member_buffer: [max_concat_group]ConcatMember = undefined;
229 const members = try concatReductionMembers(root, outline, buffer_plan, member_buffer[0..]);
230 const output_slot = bufferSlotById(buffer_plan, outline.output_slot_id) orelse return error.InvalidArtifact;
231 const first = members[0];
232 const group: u32 = @intCast(members.len);
233 const total_outputs = std.math.mul(u32, first.dims.output_element_count, group) catch return error.CapabilityMismatch;
234 const total_inputs = std.math.mul(u32, first.dims.input_element_count, group) catch return error.CapabilityMismatch;
235 if (work.dtype != output_slot.dtype) return error.InvalidArtifact;
236
237 return .{
238 .kind = first.kind,
239 .input_dtype = output_slot.dtype,
240 .output_dtype = output_slot.dtype,
241 .dims = .{
242 .input_rank = first.dims.input_rank,
243 .axis = first.dims.axis,
244 .input_element_count = total_inputs,
245 .output_element_count = total_outputs,
246 .rows = total_outputs,
247 .cols = first.dims.cols,
248 .inner = first.dims.inner,
249 },
250 .init = first.init,
251 };
252 }
253
254 fn lowerConcatWarpRows(
255 allocator: std.mem.Allocator,
256 ir_ctx: *ir.Context,
257 outline: kernelization_model.KernelOutline,
258 work: schedule_planning.ScheduleWorkItem,
259 buffer_plan: *const bufferization.BufferPlanAnalysis,
260 desc: ReductionDescription,
261 format: ?gpu.ArtifactFormat,
262 root: *ir.Operation,
263 ) common.LoweringError!LoweredKernel {
264 var member_buffer: [max_concat_group]ConcatMember = undefined;
265 const members = try concatReductionMembers(root, outline, buffer_plan, member_buffer[0..]);
266 const rows = members[0].dims.output_element_count;
267
268 const input_dtypes = allocator.alloc(choir_abi.DType, outline.inputCount()) catch return error.OutOfMemory;
269 defer allocator.free(input_dtypes);
270 for (outline.input_slot_ids, 0..) |slot_id, index| {
271 const slot = bufferSlotById(buffer_plan, slot_id) orelse return error.InvalidArtifact;
272 input_dtypes[index] = slot.dtype;
273 }
274
275 var abi = try generated_abi.flatTyped(allocator, desc.output_dtype, input_dtypes);
276 defer abi.deinit(allocator);
277
278 const entry_name = try generated_name.reductionConcat(allocator, desc, members.len, work.id);
279 errdefer allocator.free(entry_name);
280
281 if (format == .cuda_ptx and kernelization_model.reductionWarpRowsThreads(members[0].dims) != null) {
282 const threads = kernelization_model.reduction_warp_rows_threads;
283 var lowered = try generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flatThreads(threads), .{
284 .abi = abi,
285 .allocator = allocator,
286 .dtype = desc.input_dtype,
287 .members = members,
288 .rows = rows,
289 .outline = outline,
290 .buffer_plan = buffer_plan,
291 }, emitConcatWarpRowsBody);
292 lowered.body = .{ .reduction_warp_rows = .{ .threads = threads, .rows = rows } };
293 return lowered;
294 }
295
296 return generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flat(), .{
297 .abi = abi,
298 .allocator = allocator,
299 .dtype = desc.input_dtype,
300 .members = members,
301 .rows = rows,
302 .outline = outline,
303 .buffer_plan = buffer_plan,
304 }, emitConcatSerialBody);
305 }
306
307 fn emit_concat_serial_index(
308 guarded: anytype,
309 index: kernel_root.Index1D,
310 body_ctx: anytype,
311 ) !void {
312 const members = body_ctx.members;
313 if (body_ctx.dtype == .f16) return error.UnsupportedOperation;
314 const lower_bound = try guarded.constantIndex(0);
315 const cols = try guarded.constantIndex(members[0].dims.cols);
316 const step = try guarded.constantIndex(1);
317
318 var init_args: [max_concat_group]kernel_root.Value = undefined;
319 var result_types: [max_concat_group]kernel_root.Type = undefined;
320 for (members, 0..) |member, member_index| {
321 init_args[member_index] = try choirReductionInitial(guarded, member.init, null);
322 result_types[member_index] = init_args[member_index].valueType();
323 }
324
325 var scope = try guarded.forScope(
326 lower_bound,
327 cols,
328 step,
329 init_args[0..members.len],
330 result_types[0..members.len],
331 );
332 errdefer scope.abort();
333 const reduce_index = scope.inductionVar();
334 const input_index = try choirReductionInputIndex(
335 guarded,
336 members[0].dims,
337 index.index,
338 reduce_index,
339 );
340
341 var memo: ?generated_elementwise.ValueMemo = generated_elementwise.ValueMemo.init(
342 body_ctx.allocator,
343 );
344 defer if (memo) |*existing| existing.deinit();
345 var updated: [max_concat_group]kernel_root.Value = undefined;
346 for (members, 0..) |member, member_index| {
347 const value = try emitPrologueValue(
348 guarded,
349 &memo,
350 input_index,
351 member.input_value,
352 body_ctx.outline,
353 body_ctx.buffer_plan,
354 body_ctx.abi,
355 max_prologue_depth,
356 );
357 const current = scope.iterArg(member_index) orelse return error.UnsupportedOperation;
358 updated[member_index] = try emitReductionValue(guarded, member.kind, current, value);
359 }
360 try scope.leave(updated[0..members.len]);
361
362 for (members, 0..) |member, member_index| {
363 const total = scope.result(member_index) orelse return error.UnsupportedOperation;
364 const base = try guarded.constantIndex(member.base);
365 const address = try guarded.add(base, index.index);
366 try guarded.storeIndex(total, body_ctx.out, address);
367 }
368 }
369
370 fn emitConcatSerialBody(logical: anytype, ctx: anytype) !void {
371 const rows: u64 = ctx.rows;
372 const domain = try logical.index1D("out", rows);
373 try generated_guard.countIndexDo(logical, domain, ctx.abi.count(logical), .{
374 .abi = ctx.abi,
375 .allocator = ctx.allocator,
376 .dtype = ctx.dtype,
377 .members = ctx.members,
378 .outline = ctx.outline,
379 .buffer_plan = ctx.buffer_plan,
380 .out = ctx.abi.output(logical),
381 }, emit_concat_serial_index);
382 }
383
384 fn store_concat_warp_row(store_inner: anytype, store_ctx: anytype) !void {
385 const initial = try choirReductionInitial(store_inner, store_ctx.member.init, null);
386 const combined = try emitReductionValue(
387 store_inner,
388 store_ctx.member.kind,
389 initial,
390 store_ctx.total,
391 );
392 const base = try store_inner.constantIndex(store_ctx.member.base);
393 const address = try store_inner.add(base, store_ctx.row);
394 try store_inner.storeIndex(combined, store_ctx.out, address);
395 }
396
397 fn emit_concat_warp_row(inner: anytype, guard_ctx: anytype) !void {
398 const members = guard_ctx.members;
399 const cols = try inner.constantIndex(members[0].dims.cols);
400 const step = try inner.constantIndex(32);
401 const widen = guard_ctx.dtype == .f16;
402 if (widen) return error.UnsupportedOperation;
403
404 var init_args: [max_concat_group]kernel_root.Value = undefined;
405 var result_types: [max_concat_group]kernel_root.Type = undefined;
406 for (members, 0..) |member, index| {
407 init_args[index] = try reductionNeutral(inner, member.kind, guard_ctx.dtype);
408 result_types[index] = init_args[index].valueType();
409 }
410
411 var scope = try inner.forScope(
412 guard_ctx.lane,
413 cols,
414 step,
415 init_args[0..members.len],
416 result_types[0..members.len],
417 );
418 errdefer scope.abort();
419 const reduce_index = scope.inductionVar();
420 const input_index = try choirReductionInputIndex(
421 inner,
422 members[0].dims,
423 guard_ctx.row,
424 reduce_index,
425 );
426
427 var memo: ?generated_elementwise.ValueMemo = generated_elementwise.ValueMemo.init(
428 guard_ctx.allocator,
429 );
430 defer if (memo) |*existing| existing.deinit();
431 var updated: [max_concat_group]kernel_root.Value = undefined;
432 for (members, 0..) |member, index| {
433 const value = try emitPrologueValue(
434 inner,
435 &memo,
436 input_index,
437 member.input_value,
438 guard_ctx.outline,
439 guard_ctx.buffer_plan,
440 guard_ctx.abi,
441 max_prologue_depth,
442 );
443 const current = scope.iterArg(index) orelse return error.UnsupportedOperation;
444 updated[index] = try emitReductionValue(inner, member.kind, current, value);
445 }
446 try scope.leave(updated[0..members.len]);
447
448 for (members, 0..) |member, index| {
449 const partial = scope.result(index) orelse return error.UnsupportedOperation;
450 const total = try inner.warpReduce(warpKindForReduction(member.kind), partial);
451 const zero = try inner.constantIndex(0);
452 const lane_zero = try inner.compare(.eq, guard_ctx.lane, zero);
453 try inner.guardDo(lane_zero, .{
454 .member = member,
455 .row = guard_ctx.row,
456 .out = guard_ctx.out,
457 .total = total,
458 }, store_concat_warp_row);
459 }
460 }
461
462 fn emitConcatWarpRowsBody(logical: anytype, ctx: anytype) !void {
463 const rows: u64 = ctx.rows;
464 const domain = try logical.index1D("lane", rows * 32);
465 const global = domain.index;
466 const warp_size_extent = try logical.constantIndex(32);
467 const row = try logical.div(global, warp_size_extent);
468 const lane = try logical.sub(global, try logical.mul(row, warp_size_extent));
469 const rows_extent = try logical.constantIndex(@intCast(rows));
470 const row_ok = try logical.compare(.lt, row, rows_extent);
471 try logical.guardDo(row_ok, .{
472 .abi = ctx.abi,
473 .allocator = ctx.allocator,
474 .dtype = ctx.dtype,
475 .members = ctx.members,
476 .outline = ctx.outline,
477 .buffer_plan = ctx.buffer_plan,
478 .row = row,
479 .lane = lane,
480 .out = ctx.abi.output(logical),
481 }, emit_concat_warp_row);
482 }
483
484 fn lowerFused(
485 allocator: std.mem.Allocator,
486 ir_ctx: *ir.Context,
487 outline: kernelization_model.KernelOutline,
488 work: schedule_planning.ScheduleWorkItem,
489 buffer_plan: *const bufferization.BufferPlanAnalysis,
490 desc: ReductionDescription,
491 format: ?gpu.ArtifactFormat,
492 ) common.LoweringError!LoweredKernel {
493 if (desc.init != .constant) return error.UnsupportedOperation;
494 const root = work.ops[work.ops.len - 1];
495 if (isName(root.name.name, dialect_mod.AccyDialect.ConcatenateOp.operation_name)) {
496 return lowerConcatWarpRows(allocator, ir_ctx, outline, work, buffer_plan, desc, format, root);
497 }
498 const reduce = root;
499 const input_value = reduce.getOperand(0) orelse return error.InvalidArtifact;
500
501 const input_dtypes = allocator.alloc(choir_abi.DType, outline.inputCount()) catch return error.OutOfMemory;
502 defer allocator.free(input_dtypes);
503 for (outline.input_slot_ids, 0..) |slot_id, index| {
504 const slot = bufferSlotById(buffer_plan, slot_id) orelse return error.InvalidArtifact;
505 input_dtypes[index] = slot.dtype;
506 }
507
508 var abi = try generated_abi.flatTyped(allocator, desc.output_dtype, input_dtypes);
509 defer abi.deinit(allocator);
510
511 if (format == .cuda_ptx) if (kernelization_model.reductionWarpRowsThreads(desc.dims)) |threads| {
512 const entry_name = try generated_name.reductionWarpRows(allocator, desc, work.ops.len - 1, work.id);
513 errdefer allocator.free(entry_name);
514 var lowered = try generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flatThreads(threads), .{
515 .abi = abi,
516 .allocator = allocator,
517 .kind = desc.kind,
518 .dtype = desc.input_dtype,
519 .dims = desc.dims,
520 .init = desc.init,
521 .input_value = input_value,
522 .outline = outline,
523 .buffer_plan = buffer_plan,
524 }, emitWarpRowsReductionBody);
525 lowered.body = .{ .reduction_warp_rows = .{ .threads = threads, .rows = desc.dims.output_element_count } };
526 return lowered;
527 };
528
529 const entry_name = try generated_name.reductionFused(allocator, desc, work.ops.len - 1, work.id);
530 errdefer allocator.free(entry_name);
531
532 return generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flat(), .{
533 .abi = abi,
534 .allocator = allocator,
535 .kind = desc.kind,
536 .dtype = desc.input_dtype,
537 .dims = desc.dims,
538 .init = desc.init,
539 .input_value = input_value,
540 .outline = outline,
541 .buffer_plan = buffer_plan,
542 }, emitFusedReductionBody);
543 }
544
545 pub fn warpKindForReduction(kind: ReductionKind) kernel_root.WarpOpKind {
546 return switch (kind) {
547 .sum => .add,
548 .max => .max,
549 .min => .min,
550 };
551 }
552
553 fn reduce_warp_row(
554 fold_inner: anytype,
555 reduce_index: kernel_root.Value,
556 current: kernel_root.Value,
557 fold_ctx: anytype,
558 ) !kernel_root.Value {
559 const input_index = try choirReductionInputIndex(
560 fold_inner,
561 fold_ctx.dims,
562 fold_ctx.row,
563 reduce_index,
564 );
565 var memo: ?generated_elementwise.ValueMemo = generated_elementwise.ValueMemo.init(
566 fold_ctx.allocator,
567 );
568 defer if (memo) |*existing| existing.deinit();
569 var value = try emitPrologueValue(
570 fold_inner,
571 &memo,
572 input_index,
573 fold_ctx.input_value,
574 fold_ctx.outline,
575 fold_ctx.buffer_plan,
576 fold_ctx.abi,
577 max_prologue_depth,
578 );
579 if (fold_ctx.widen) value = try fold_inner.cast(value, .f32);
580 return emitReductionValue(fold_inner, fold_ctx.kind, current, value);
581 }
582
583 fn store_warp_reduction_row(store_inner: anytype, store_ctx: anytype) !void {
584 var initial = try choirReductionInitial(store_inner, store_ctx.init, null);
585 if (store_ctx.widen) initial = try store_inner.cast(initial, .f32);
586 var combined = try emitReductionValue(
587 store_inner,
588 store_ctx.kind,
589 initial,
590 store_ctx.total,
591 );
592 if (store_ctx.widen) combined = try store_inner.cast(combined, .f16);
593 try store_inner.storeIndex(combined, store_ctx.out, store_ctx.row);
594 }
595
596 fn emit_warp_reduction_row(inner: anytype, guard_ctx: anytype) !void {
597 const cols = try inner.constantIndex(guard_ctx.dims.cols);
598 const step = try inner.constantIndex(32);
599 const widen = guard_ctx.dtype == .f16;
600 const neutral = try reductionNeutral(inner, guard_ctx.kind, guard_ctx.dtype);
601 const partial = try inner.fold(guard_ctx.lane, cols, step, neutral, .{
602 .abi = guard_ctx.abi,
603 .allocator = guard_ctx.allocator,
604 .kind = guard_ctx.kind,
605 .dims = guard_ctx.dims,
606 .input_value = guard_ctx.input_value,
607 .outline = guard_ctx.outline,
608 .buffer_plan = guard_ctx.buffer_plan,
609 .row = guard_ctx.row,
610 .widen = widen,
611 }, reduce_warp_row);
612 const total = try inner.warpReduce(warpKindForReduction(guard_ctx.kind), partial);
613 const zero = try inner.constantIndex(0);
614 const lane_zero = try inner.compare(.eq, guard_ctx.lane, zero);
615 try inner.guardDo(lane_zero, .{
616 .kind = guard_ctx.kind,
617 .init = guard_ctx.init,
618 .row = guard_ctx.row,
619 .out = guard_ctx.out,
620 .total = total,
621 .widen = widen,
622 }, store_warp_reduction_row);
623 }
624
625 fn emitWarpRowsReductionBody(logical: anytype, ctx: anytype) !void {
626 const rows: u64 = ctx.dims.output_element_count;
627 const domain = try logical.index1D("lane", rows * 32);
628 const global = domain.index;
629 const warp_size_extent = try logical.constantIndex(32);
630 const row = try logical.div(global, warp_size_extent);
631 const lane = try logical.sub(global, try logical.mul(row, warp_size_extent));
632 const rows_extent = try logical.constantIndex(@intCast(rows));
633 const row_ok = try logical.compare(.lt, row, rows_extent);
634 try logical.guardDo(row_ok, .{
635 .abi = ctx.abi,
636 .allocator = ctx.allocator,
637 .kind = ctx.kind,
638 .dtype = ctx.dtype,
639 .dims = ctx.dims,
640 .init = ctx.init,
641 .input_value = ctx.input_value,
642 .outline = ctx.outline,
643 .buffer_plan = ctx.buffer_plan,
644 .row = row,
645 .lane = lane,
646 .out = ctx.abi.output(logical),
647 }, emit_warp_reduction_row);
648 }
649
650 fn reduce_fused_reduction(
651 fold_inner: anytype,
652 reduce_index: kernel_root.Value,
653 current: kernel_root.Value,
654 fold_ctx: anytype,
655 ) !kernel_root.Value {
656 const input_index = try choirReductionInputIndex(
657 fold_inner,
658 fold_ctx.dims,
659 fold_ctx.output_index,
660 reduce_index,
661 );
662 var memo: ?generated_elementwise.ValueMemo = generated_elementwise.ValueMemo.init(
663 fold_ctx.allocator,
664 );
665 defer if (memo) |*existing| existing.deinit();
666 var value = try emitPrologueValue(
667 fold_inner,
668 &memo,
669 input_index,
670 fold_ctx.input_value,
671 fold_ctx.outline,
672 fold_ctx.buffer_plan,
673 fold_ctx.abi,
674 max_prologue_depth,
675 );
676 if (fold_ctx.widen) value = try fold_inner.cast(value, .f32);
677 return emitReductionValue(fold_inner, fold_ctx.kind, current, value);
678 }
679
680 fn emit_fused_reduction_index(
681 guarded: anytype,
682 index: kernel_root.Index1D,
683 body_ctx: anytype,
684 ) !void {
685 const lower_bound = try guarded.constantIndex(0);
686 const upper = try choirReductionUpperBound(guarded, body_ctx.dims);
687 const step = try guarded.constantIndex(1);
688 const widen = body_ctx.dtype == .f16;
689 var initial = try choirReductionInitial(guarded, body_ctx.init, null);
690 if (widen) initial = try guarded.cast(initial, .f32);
691 const result = try guarded.fold(lower_bound, upper, step, initial, .{
692 .abi = body_ctx.abi,
693 .allocator = body_ctx.allocator,
694 .kind = body_ctx.kind,
695 .dims = body_ctx.dims,
696 .input_value = body_ctx.input_value,
697 .outline = body_ctx.outline,
698 .buffer_plan = body_ctx.buffer_plan,
699 .output_index = index.index,
700 .widen = widen,
701 }, reduce_fused_reduction);
702 const stored = if (widen) try guarded.cast(result, .f16) else result;
703 try guarded.storeIndex(stored, body_ctx.out, index);
704 }
705
706 fn emitFusedReductionBody(logical: anytype, ctx: anytype) !void {
707 const output_index = try logical.index1D("out", ctx.dims.output_element_count);
708 try generated_guard.countIndexDo(logical, output_index, ctx.abi.count(logical), .{
709 .abi = ctx.abi,
710 .allocator = ctx.allocator,
711 .kind = ctx.kind,
712 .dtype = ctx.dtype,
713 .dims = ctx.dims,
714 .init = ctx.init,
715 .input_value = ctx.input_value,
716 .outline = ctx.outline,
717 .buffer_plan = ctx.buffer_plan,
718 .out = ctx.abi.output(logical),
719 }, emit_fused_reduction_index);
720 }
721
722 fn emitPrologueValue(
723 builder: anytype,
724 memo: *?generated_elementwise.ValueMemo,
725 index: kernel_root.Value,
726 value: *ir.Value,
727 outline: kernelization_model.KernelOutline,
728 buffer_plan: *const bufferization.BufferPlanAnalysis,
729 abi: anytype,
730 depth: usize,
731 ) common.LoweringError!kernel_root.Value {
732 return generated_elementwise.evaluateValueAtIndex(builder, memo, index, value, outline, buffer_plan, abi, depth);
733 }
734
735 fn reductionInitFillPattern(init: ReductionInit) ?u32 {
736 return switch (init) {
737 .constant => |value| switch (value) {
738 .f32 => |scalar| @bitCast(scalar),
739 .i32 => |scalar| @bitCast(scalar),
740 .u32 => |scalar| scalar,
741 .f16 => null,
742 },
743 .input_buffer => null,
744 };
745 }
746
747 fn reduce_atomic_input(
748 fold_inner: anytype,
749 reduce_index: kernel_root.Value,
750 current: kernel_root.Value,
751 fold_ctx: anytype,
752 ) !kernel_root.Value {
753 const value = try fold_inner.loadIndex(fold_ctx.input, reduce_index);
754 return emitReductionValue(fold_inner, fold_ctx.kind, current, value);
755 }
756
757 fn combine_atomic_shared_reduction(inner: anytype, guard_ctx: anytype) !void {
758 const left = try inner.loadIndex(guard_ctx.shared, guard_ctx.tid);
759 const right_index = try inner.add(guard_ctx.tid, guard_ctx.offset_extent);
760 const right = try inner.loadIndex(guard_ctx.shared, right_index);
761 const combined = try emitReductionValue(inner, guard_ctx.kind, left, right);
762 try inner.storeIndex(combined, guard_ctx.shared, guard_ctx.tid);
763 }
764
765 fn finish_atomic_reduction(inner: anytype, guard_ctx: anytype) !void {
766 const zero_index = try inner.constantIndex(0);
767 const block_total = try inner.loadIndex(guard_ctx.shared, zero_index);
768 _ = try inner.atomicRmw(.add, block_total, guard_ctx.out, zero_index);
769 }
770
771 fn emitAtomicReductionBody(logical: anytype, ctx: anytype) !void {
772 const plan: kernelization_model.ReductionAtomicPlan = ctx.plan;
773
774 const out = ctx.abi.output(logical);
775 const input = ctx.abi.input(logical);
776
777 const domain = try logical.index1D("t", plan.stride());
778 const global = domain.index;
779 const tid = try logical.threadId(.x);
780 const stride = try logical.constantIndex(@intCast(plan.stride()));
781
782 const neutral = try reductionNeutral(logical, ctx.kind, ctx.dtype);
783 const partial = if (vectorizedReductionQuads(ctx.dtype, ctx.dims.input_element_count)) |quads| blk: {
784 const quad_extent = try logical.constantIndex(quads);
785 break :blk try logical.fold(global, quad_extent, stride, neutral, .{
786 .kind = ctx.kind,
787 .input = input,
788 }, emitQuadReductionStep);
789 } else blk: {
790 const extent = try logical.constantIndex(ctx.dims.input_element_count);
791 break :blk try logical.fold(global, extent, stride, neutral, .{
792 .kind = ctx.kind,
793 .input = input,
794 }, reduce_atomic_input);
795 };
796
797 const shared = try logical.sharedBuffer(ctx.dtype, plan.threads);
798 try logical.storeIndex(partial, shared, tid);
799 try logical.barrier(.block);
800
801 var offset: u32 = plan.threads / 2;
802 while (offset > 0) : (offset /= 2) {
803 const offset_extent = try logical.constantIndex(offset);
804 const in_half = try logical.compare(.lt, tid, offset_extent);
805 try logical.guardDo(in_half, .{
806 .kind = ctx.kind,
807 .shared = shared,
808 .tid = tid,
809 .offset_extent = offset_extent,
810 }, combine_atomic_shared_reduction);
811 try logical.barrier(.block);
812 }
813
814 const zero = try logical.constantIndex(0);
815 const is_leader = try logical.compare(.eq, tid, zero);
816 try logical.guardDo(is_leader, .{
817 .shared = shared,
818 .out = out,
819 }, finish_atomic_reduction);
820 }
821
822 fn vectorizedReductionQuads(dtype: choir_abi.DType, element_count: u32) ?i64 {
823 if (dtype != .f32) return null;
824 if (element_count % 4 != 0) return null;
825 return @intCast(element_count / 4);
826 }
827
828 fn emitQuadReductionStep(fold_inner: anytype, quad_index: kernel_root.Value, current: kernel_root.Value, fold_ctx: anytype) !kernel_root.Value {
829 const four = try fold_inner.constantIndex(4);
830 const loaded = try fold_inner.loadVector(fold_ctx.input, try fold_inner.mul(quad_index, four), 4);
831 var lanes: [4]kernel_root.Value = undefined;
832 for (0..4) |lane| {
833 lanes[lane] = try fold_inner.extractLane(loaded, @intCast(lane), .f32);
834 }
835 const low = try emitReductionValue(fold_inner, fold_ctx.kind, lanes[0], lanes[1]);
836 const high = try emitReductionValue(fold_inner, fold_ctx.kind, lanes[2], lanes[3]);
837 const combined = try emitReductionValue(fold_inner, fold_ctx.kind, low, high);
838 return emitReductionValue(fold_inner, fold_ctx.kind, current, combined);
839 }
840
841 pub fn reductionNeutral(builder: anytype, kind: ReductionKind, dtype: choir_abi.DType) !kernel_root.Value {
842 const wide = if (dtype == .f16) choir_abi.DType.f32 else dtype;
843 return switch (wide) {
844 .f32 => switch (kind) {
845 .sum => builder.constantFloat(.f32, 0.0),
846 .max => builder.constantFloat(.f32, -std.math.inf(f64)),
847 .min => builder.constantFloat(.f32, std.math.inf(f64)),
848 },
849 .i32 => switch (kind) {
850 .sum => builder.constantInt(.i32, 0),
851 .max => builder.constantInt(.i32, std.math.minInt(i32)),
852 .min => builder.constantInt(.i32, std.math.maxInt(i32)),
853 },
854 .u32 => switch (kind) {
855 .sum => builder.constantInt(.u32, 0),
856 .max => builder.constantInt(.u32, 0),
857 .min => builder.constantInt(.u32, std.math.maxInt(u32)),
858 },
859 else => error.CapabilityMismatch,
860 };
861 }
862
863 fn reduce_single_block_input(
864 fold_inner: anytype,
865 reduce_index: kernel_root.Value,
866 current: kernel_root.Value,
867 fold_ctx: anytype,
868 ) !kernel_root.Value {
869 var value = try fold_inner.loadIndex(fold_ctx.input, reduce_index);
870 if (fold_ctx.widen) value = try fold_inner.cast(value, .f32);
871 return emitReductionValue(fold_inner, fold_ctx.kind, current, value);
872 }
873
874 fn combine_single_block_shared_reduction(inner: anytype, guard_ctx: anytype) !void {
875 const left = try inner.loadIndex(guard_ctx.shared, guard_ctx.tid);
876 const right_index = try inner.add(guard_ctx.tid, guard_ctx.offset_extent);
877 const right = try inner.loadIndex(guard_ctx.shared, right_index);
878 const combined = try emitReductionValue(inner, guard_ctx.kind, left, right);
879 try inner.storeIndex(combined, guard_ctx.shared, guard_ctx.tid);
880 }
881
882 fn finish_single_block_reduction(inner: anytype, guard_ctx: anytype) !void {
883 const zero_index = try inner.constantIndex(0);
884 const total = try inner.loadIndex(guard_ctx.shared, zero_index);
885 var initial = try choirReductionInitial(inner, guard_ctx.init, guard_ctx.init_buffer);
886 if (guard_ctx.widen) initial = try inner.cast(initial, .f32);
887 const combined = try emitReductionValue(inner, guard_ctx.kind, total, initial);
888 const stored = if (guard_ctx.widen) try inner.cast(combined, .f16) else combined;
889 try inner.storeIndex(stored, guard_ctx.out, zero_index);
890 }
891
892 fn emitSingleBlockReductionBody(logical: anytype, ctx: anytype) !void {
893 const threads: u32 = ctx.threads;
894 const widen = ctx.dtype == .f16;
895 const shared_dtype: choir_abi.DType = if (widen) .f32 else ctx.dtype;
896
897 const out = ctx.abi.output(logical);
898 const input = ctx.abi.input(logical);
899 const init_buffer: ?kernel_root.Value = if (ctx.init == .input_buffer) ctx.abi.init(logical) else null;
900
901 const domain = try logical.index1D("t", threads);
902 const tid = domain.index;
903 const stride = try logical.constantIndex(threads);
904
905 const neutral = try reductionNeutral(logical, ctx.kind, ctx.dtype);
906 const partial = if (vectorizedReductionQuads(ctx.dtype, ctx.dims.input_element_count)) |quads| blk: {
907 const quad_extent = try logical.constantIndex(quads);
908 break :blk try logical.fold(tid, quad_extent, stride, neutral, .{
909 .kind = ctx.kind,
910 .input = input,
911 }, emitQuadReductionStep);
912 } else blk: {
913 const extent = try logical.constantIndex(ctx.dims.input_element_count);
914 break :blk try logical.fold(tid, extent, stride, neutral, .{
915 .kind = ctx.kind,
916 .input = input,
917 .widen = widen,
918 }, reduce_single_block_input);
919 };
920
921 const shared = try logical.sharedBuffer(shared_dtype, threads);
922 try logical.storeIndex(partial, shared, tid);
923 try logical.barrier(.block);
924
925 var offset: u32 = threads / 2;
926 while (offset > 0) : (offset /= 2) {
927 const offset_extent = try logical.constantIndex(offset);
928 const in_half = try logical.compare(.lt, tid, offset_extent);
929 try logical.guardDo(in_half, .{
930 .kind = ctx.kind,
931 .shared = shared,
932 .tid = tid,
933 .offset_extent = offset_extent,
934 }, combine_single_block_shared_reduction);
935 try logical.barrier(.block);
936 }
937
938 const zero = try logical.constantIndex(0);
939 const is_leader = try logical.compare(.eq, tid, zero);
940 try logical.guardDo(is_leader, .{
941 .kind = ctx.kind,
942 .init = ctx.init,
943 .init_buffer = init_buffer,
944 .shared = shared,
945 .out = out,
946 .widen = widen,
947 }, finish_single_block_reduction);
948 }
949
950 fn reduce_reduction_input(
951 fold_inner: anytype,
952 reduce_index: kernel_root.Value,
953 current: kernel_root.Value,
954 fold_ctx: anytype,
955 ) !kernel_root.Value {
956 const input_index = try choirReductionInputIndex(
957 fold_inner,
958 fold_ctx.dims,
959 fold_ctx.output_index,
960 reduce_index,
961 );
962 var value = try fold_inner.loadIndex(fold_ctx.input, input_index);
963 if (fold_ctx.widen) value = try fold_inner.cast(value, .f32);
964 return emitReductionValue(fold_inner, fold_ctx.kind, current, value);
965 }
966
967 fn emit_reduction_index(
968 guarded: anytype,
969 index: kernel_root.Index1D,
970 body_ctx: anytype,
971 ) !void {
972 const lower_bound = try guarded.constantIndex(0);
973 const upper = try choirReductionUpperBound(guarded, body_ctx.dims);
974 const step = try guarded.constantIndex(1);
975 const widen = body_ctx.dtype == .f16;
976 var initial = try choirReductionInitial(guarded, body_ctx.init, body_ctx.init_buffer);
977 if (widen) initial = try guarded.cast(initial, .f32);
978 const result = try guarded.fold(lower_bound, upper, step, initial, .{
979 .kind = body_ctx.kind,
980 .dims = body_ctx.dims,
981 .input = body_ctx.input,
982 .output_index = index.index,
983 .widen = widen,
984 }, reduce_reduction_input);
985 const stored = if (widen) try guarded.cast(result, .f16) else result;
986 try guarded.storeIndex(stored, body_ctx.out, index);
987 }
988
989 fn emitReductionBody(logical: anytype, ctx: anytype) !void {
990 const output_index = try logical.index1D("out", ctx.dims.output_element_count);
991 try generated_guard.countIndexDo(logical, output_index, ctx.abi.count(logical), .{
992 .kind = ctx.kind,
993 .dtype = ctx.dtype,
994 .dims = ctx.dims,
995 .init = ctx.init,
996 .out = ctx.abi.output(logical),
997 .input = ctx.abi.input(logical),
998 .init_buffer = if (ctx.init == .input_buffer) ctx.abi.init(logical) else null,
999 }, emit_reduction_index);
1000 }
1001
1002 fn choirReductionUpperBound(
1003 builder: anytype,
1004 dims: ReductionStaticDims,
1005 ) !kernel_root.Value {
1006 const bound: i64 = switch (dims.input_rank) {
1007 1 => dims.input_element_count,
1008 2 => if (dims.axis == 0) dims.rows else dims.cols,
1009 3 => dims.cols,
1010 else => return error.InvalidArtifact,
1011 };
1012 return builder.constantIndex(bound);
1013 }
1014
1015 fn choirReductionInputIndex(
1016 builder: anytype,
1017 dims: ReductionStaticDims,
1018 output_index: kernel_root.Value,
1019 reduce_index: kernel_root.Value,
1020 ) !kernel_root.Value {
1021 switch (dims.input_rank) {
1022 1 => return reduce_index,
1023 2 => {
1024 const cols = try builder.constantIndex(dims.cols);
1025 if (dims.axis == 0) {
1026 const offset = try builder.mul(reduce_index, cols);
1027 return builder.add(offset, output_index);
1028 }
1029 const row_offset = try builder.mul(output_index, cols);
1030 return builder.add(row_offset, reduce_index);
1031 },
1032 3 => {
1033 const inner = try builder.constantIndex(dims.inner);
1034 const outer_coord = try builder.div(output_index, inner);
1035 const consumed = try builder.mul(outer_coord, inner);
1036 const within = try builder.sub(output_index, consumed);
1037 const block_stride = try builder.constantIndex(@as(i64, dims.cols) * @as(i64, dims.inner));
1038 const outer_offset = try builder.mul(outer_coord, block_stride);
1039 const reduce_offset = try builder.mul(reduce_index, inner);
1040 const offset = try builder.add(outer_offset, reduce_offset);
1041 return builder.add(offset, within);
1042 },
1043 else => return error.InvalidArtifact,
1044 }
1045 }
1046
1047 fn choirReductionInitial(builder: anytype, init: ReductionInit, init_buffer: ?kernel_root.Value) !kernel_root.Value {
1048 return switch (init) {
1049 .constant => |value| choirReductionConstantInitial(builder, value),
1050 .input_buffer => builder.loadIndex(init_buffer orelse return error.InvalidArtifact, try builder.constantIndex(0)),
1051 };
1052 }
1053
1054 fn choirReductionConstantInitial(builder: anytype, value: ReductionInitValue) !kernel_root.Value {
1055 return switch (value) {
1056 .f32 => |scalar| builder.constantFloat(.f32, scalar),
1057 .i32 => |scalar| builder.constantInt(.i32, scalar),
1058 .u32 => |scalar| builder.constantInt(.u32, scalar),
1059 .f16 => |scalar| builder.constantFloat(.f16, @floatCast(scalar)),
1060 };
1061 }
1062
1063 pub fn emitReductionValue(
1064 builder: anytype,
1065 kind: ReductionKind,
1066 current: kernel_root.Value,
1067 value: kernel_root.Value,
1068 ) !kernel_root.Value {
1069 return switch (kind) {
1070 .sum => builder.add(current, value),
1071 .max => builder.max(current, value),
1072 .min => builder.min(current, value),
1073 };
1074 }
1075
1076 fn reductionStaticDims(
1077 input_dims: []const i64,
1078 output_dims: []const i64,
1079 axes: []const i64,
1080 ) gpu.BackendError!ReductionStaticDims {
1081 if (input_dims.len == 0) return error.CapabilityMismatch;
1082 if (try axesCoverInputRank(input_dims.len, axes)) {
1083 if (output_dims.len != 0) return error.InvalidArtifact;
1084 var input_element_count: u64 = 1;
1085 for (input_dims) |dim| {
1086 input_element_count *= try staticPositiveDimU32(dim);
1087 }
1088 const bounded = std.math.cast(u32, input_element_count) orelse return error.CapabilityMismatch;
1089 return .{
1090 .input_rank = 1,
1091 .axis = 0,
1092 .input_element_count = bounded,
1093 .output_element_count = 1,
1094 .rows = bounded,
1095 .cols = 1,
1096 };
1097 }
1098 return boundaryReductionStaticDims(input_dims, output_dims, axes);
1099 }
1100
1101 fn boundaryReductionStaticDims(
1102 input_dims: []const i64,
1103 output_dims: []const i64,
1104 axes: []const i64,
1105 ) gpu.BackendError!ReductionStaticDims {
1106 if (axes.len == 0) return error.CapabilityMismatch;
1107 if (input_dims.len > i64_attr_list_stack_capacity) return error.CapabilityMismatch;
1108
1109 var reduced_buffer = @as([i64_attr_list_stack_capacity]bool, @splat(false));
1110 const reduced = reduced_buffer[0..input_dims.len];
1111 try markReductionAxes(input_dims.len, axes, reduced);
1112
1113 if (output_dims.len + axes.len != input_dims.len) return error.InvalidArtifact;
1114
1115 if (axesCoverSuffix(reduced, output_dims.len)) {
1116 if (!std.mem.eql(i64, input_dims[0..output_dims.len], output_dims)) return error.InvalidArtifact;
1117 return flattenedBoundaryReductionDims(input_dims, output_dims, 1);
1118 }
1119
1120 if (axesCoverPrefix(reduced, axes.len)) {
1121 if (!std.mem.eql(i64, input_dims[axes.len..], output_dims)) return error.InvalidArtifact;
1122 return flattenedBoundaryReductionDims(input_dims, output_dims, 0);
1123 }
1124
1125 return interiorReductionStaticDims(input_dims, output_dims, reduced);
1126 }
1127
1128 fn interiorReductionStaticDims(
1129 input_dims: []const i64,
1130 output_dims: []const i64,
1131 reduced: []const bool,
1132 ) gpu.BackendError!ReductionStaticDims {
1133 const block = contiguousReducedBlock(reduced) orelse return error.CapabilityMismatch;
1134 const trailing_start = block.last + 1;
1135 if (!std.mem.eql(i64, input_dims[0..block.first], output_dims[0..block.first])) {
1136 return error.InvalidArtifact;
1137 }
1138 if (!std.mem.eql(i64, input_dims[trailing_start..], output_dims[block.first..])) {
1139 return error.InvalidArtifact;
1140 }
1141
1142 const input_element_count = try staticDimProductU32(input_dims);
1143 const output_element_count = try staticDimProductU32(output_dims);
1144 const extent = try staticDimProductU32(input_dims[block.first..trailing_start]);
1145 const inner = try staticDimProductU32(input_dims[trailing_start..]);
1146 const outer = output_element_count / inner;
1147 if (outer * inner != output_element_count) return error.InvalidArtifact;
1148 if (outer * extent * inner != input_element_count) return error.InvalidArtifact;
1149
1150 return .{
1151 .input_rank = 3,
1152 .axis = 1,
1153 .input_element_count = input_element_count,
1154 .output_element_count = output_element_count,
1155 .rows = outer,
1156 .cols = extent,
1157 .inner = inner,
1158 };
1159 }
1160
1161 fn contiguousReducedBlock(reduced: []const bool) ?struct { first: usize, last: usize } {
1162 var first: ?usize = null;
1163 var last: usize = 0;
1164 for (reduced, 0..) |flag, index| {
1165 if (!flag) continue;
1166 if (first == null) first = index;
1167 last = index;
1168 }
1169 const start = first orelse return null;
1170 for (reduced[start .. last + 1]) |flag| {
1171 if (!flag) return null;
1172 }
1173 return .{ .first = start, .last = last };
1174 }
1175
1176 fn staticDimProductU32(dims: []const i64) gpu.BackendError!u32 {
1177 var product: u32 = 1;
1178 for (dims) |dim| {
1179 const value = try staticPositiveDimU32(dim);
1180 product = std.math.mul(u32, product, value) catch return error.CapabilityMismatch;
1181 }
1182 return product;
1183 }
1184
1185 fn flattenedBoundaryReductionDims(
1186 input_dims: []const i64,
1187 output_dims: []const i64,
1188 axis: u8,
1189 ) gpu.BackendError!ReductionStaticDims {
1190 for (input_dims) |dim| _ = try staticPositiveDimU32(dim);
1191 const input_element_count = try staticDimProductU32(input_dims);
1192 const output_element_count = try staticDimProductU32(output_dims);
1193 if (input_element_count % output_element_count != 0) return error.InvalidArtifact;
1194 const reduction_extent = input_element_count / output_element_count;
1195 return .{
1196 .input_rank = 2,
1197 .axis = axis,
1198 .input_element_count = input_element_count,
1199 .output_element_count = output_element_count,
1200 .rows = if (axis == 0) reduction_extent else output_element_count,
1201 .cols = if (axis == 0) output_element_count else reduction_extent,
1202 };
1203 }
1204
1205 fn markReductionAxes(rank: usize, axes: []const i64, reduced: []bool) gpu.BackendError!void {
1206 if (rank != reduced.len) return error.InvalidArtifact;
1207 for (axes) |axis| {
1208 if (axis < 0) return error.CapabilityMismatch;
1209 const index = std.math.cast(usize, axis) orelse return error.CapabilityMismatch;
1210 if (index >= rank) return error.CapabilityMismatch;
1211 if (reduced[index]) return error.CapabilityMismatch;
1212 reduced[index] = true;
1213 }
1214 }
1215
1216 fn axesCoverSuffix(reduced: []const bool, output_rank: usize) bool {
1217 for (reduced, 0..) |is_reduced, index| {
1218 if (is_reduced != (index >= output_rank)) return false;
1219 }
1220 return true;
1221 }
1222
1223 fn axesCoverPrefix(reduced: []const bool, axes_len: usize) bool {
1224 for (reduced, 0..) |is_reduced, index| {
1225 if (is_reduced != (index < axes_len)) return false;
1226 }
1227 return true;
1228 }
1229
1230 fn axesCoverInputRank(rank: usize, axes: []const i64) gpu.BackendError!bool {
1231 if (axes.len != rank) return false;
1232 if (rank > i64_attr_list_stack_capacity) return error.CapabilityMismatch;
1233 var seen = @as([i64_attr_list_stack_capacity]bool, @splat(false));
1234 for (axes) |axis| {
1235 if (axis < 0) return error.CapabilityMismatch;
1236 const index = std.math.cast(usize, axis) orelse return error.CapabilityMismatch;
1237 if (index >= rank) return error.CapabilityMismatch;
1238 if (seen[index]) return error.CapabilityMismatch;
1239 seen[index] = true;
1240 }
1241 return true;
1242 }
1243
1244 fn validateReductionDTypes(
1245 input_dtype: choir_abi.DType,
1246 output_dtype: choir_abi.DType,
1247 work: schedule_planning.ScheduleWorkItem,
1248 ) gpu.BackendError!void {
1249 if (input_dtype != output_dtype) return error.CapabilityMismatch;
1250 if (work.dtype != output_dtype) return error.InvalidArtifact;
1251 switch (input_dtype) {
1252 .f32, .i32, .u32, .f16 => {},
1253 else => return error.CapabilityMismatch,
1254 }
1255 }
1256
1257 pub fn reductionKindForOp(op: *ir.Operation) gpu.BackendError!ReductionKind {
1258 const payload = try dialectAttrPayload(op, "reducer_kind", dialect_mod.AccyDialect.ReduceOp.dialectAttrName("reducer_kind"));
1259 if (std.mem.eql(u8, payload, "sum")) return .sum;
1260 if (std.mem.eql(u8, payload, "max")) return .max;
1261 if (std.mem.eql(u8, payload, "min")) return .min;
1262 return error.UnsupportedOperation;
1263 }
1264
1265 pub fn reductionInitValue(
1266 init_slot: bufferization.BufferSlot,
1267 dtype: choir_abi.DType,
1268 outline: kernelization_model.KernelOutline,
1269 ) gpu.BackendError!ReductionInit {
1270 if (init_slot.dtype != dtype) return error.InvalidArtifact;
1271 if (!init_slot.role.constant) {
1272 if (externalInputIndex(outline, init_slot.id) != 1) return error.UnsupportedOperation;
1273 return .input_buffer;
1274 }
1275 const value = try reductionConstantInitValue(init_slot, dtype);
1276 return .{ .constant = value };
1277 }
1278
1279 fn reductionConstantInitValue(
1280 init_slot: bufferization.BufferSlot,
1281 dtype: choir_abi.DType,
1282 ) gpu.BackendError!ReductionInitValue {
1283 const payload = try constantPayloadForSlot(init_slot);
1284 switch (dtype) {
1285 .f32 => {
1286 if (payload.len != @sizeOf(f32)) return error.CapabilityMismatch;
1287 return .{ .f32 = std.mem.bytesToValue(f32, payload[0..@sizeOf(f32)]) };
1288 },
1289 .i32 => {
1290 if (payload.len != @sizeOf(i32)) return error.CapabilityMismatch;
1291 return .{ .i32 = std.mem.bytesToValue(i32, payload[0..@sizeOf(i32)]) };
1292 },
1293 .u32 => {
1294 if (payload.len != @sizeOf(u32)) return error.CapabilityMismatch;
1295 return .{ .u32 = std.mem.bytesToValue(u32, payload[0..@sizeOf(u32)]) };
1296 },
1297 .f16 => {
1298 if (payload.len != @sizeOf(f16)) return error.CapabilityMismatch;
1299 return .{ .f16 = std.mem.bytesToValue(f16, payload[0..@sizeOf(f16)]) };
1300 },
1301 else => return error.CapabilityMismatch,
1302 }
1303 }
1304
1305 test "reduction static dims flatten all-axis scalar reductions" {
1306 var input_dims = [_]i64{ 2, 3 };
1307 var output_dims = [_]i64{};
1308 var axes = [_]i64{ 0, 1 };
1309
1310 const dims = try reductionStaticDims(
1311 input_dims[0..],
1312 output_dims[0..],
1313 axes[0..],
1314 );
1315
1316 try std.testing.expectEqual(@as(u8, 1), dims.input_rank);
1317 try std.testing.expectEqual(@as(u8, 0), dims.axis);
1318 try std.testing.expectEqual(@as(u32, 6), dims.input_element_count);
1319 try std.testing.expectEqual(@as(u32, 1), dims.output_element_count);
1320 try std.testing.expectEqual(@as(u32, 6), dims.rows);
1321 try std.testing.expectEqual(@as(u32, 1), dims.cols);
1322 }
1323
1324 test "reduction static dims flatten suffix reductions per output element" {
1325 var input_dims = [_]i64{ 8, 2, 3 };
1326 var output_dims = [_]i64{8};
1327 var axes = [_]i64{ 1, 2 };
1328
1329 const dims = try reductionStaticDims(
1330 input_dims[0..],
1331 output_dims[0..],
1332 axes[0..],
1333 );
1334
1335 try std.testing.expectEqual(@as(u8, 2), dims.input_rank);
1336 try std.testing.expectEqual(@as(u8, 1), dims.axis);
1337 try std.testing.expectEqual(@as(u32, 48), dims.input_element_count);
1338 try std.testing.expectEqual(@as(u32, 8), dims.output_element_count);
1339 try std.testing.expectEqual(@as(u32, 8), dims.rows);
1340 try std.testing.expectEqual(@as(u32, 6), dims.cols);
1341 }
1342
1343 test "reduction static dims flatten prefix reductions per output element" {
1344 var input_dims = [_]i64{ 2, 3, 8 };
1345 var output_dims = [_]i64{8};
1346 var axes = [_]i64{ 0, 1 };
1347
1348 const dims = try reductionStaticDims(
1349 input_dims[0..],
1350 output_dims[0..],
1351 axes[0..],
1352 );
1353
1354 try std.testing.expectEqual(@as(u8, 2), dims.input_rank);
1355 try std.testing.expectEqual(@as(u8, 0), dims.axis);
1356 try std.testing.expectEqual(@as(u32, 48), dims.input_element_count);
1357 try std.testing.expectEqual(@as(u32, 8), dims.output_element_count);
1358 try std.testing.expectEqual(@as(u32, 6), dims.rows);
1359 try std.testing.expectEqual(@as(u32, 8), dims.cols);
1360 }
1361
1362 test "reduction static dims factor interior reductions into outer extent inner" {
1363 var input_dims = [_]i64{ 4, 8, 8 };
1364 var output_dims = [_]i64{ 4, 8 };
1365 var axes = [_]i64{1};
1366
1367 const dims = try reductionStaticDims(
1368 input_dims[0..],
1369 output_dims[0..],
1370 axes[0..],
1371 );
1372
1373 try std.testing.expectEqual(@as(u8, 3), dims.input_rank);
1374 try std.testing.expectEqual(@as(u8, 1), dims.axis);
1375 try std.testing.expectEqual(@as(u32, 256), dims.input_element_count);
1376 try std.testing.expectEqual(@as(u32, 32), dims.output_element_count);
1377 try std.testing.expectEqual(@as(u32, 4), dims.rows);
1378 try std.testing.expectEqual(@as(u32, 8), dims.cols);
1379 try std.testing.expectEqual(@as(u32, 8), dims.inner);
1380 }
1381
1382 test "reduction static dims factor interior block reductions" {
1383 var input_dims = [_]i64{ 2, 3, 5, 7 };
1384 var output_dims = [_]i64{ 2, 7 };
1385 var axes = [_]i64{ 1, 2 };
1386
1387 const dims = try reductionStaticDims(
1388 input_dims[0..],
1389 output_dims[0..],
1390 axes[0..],
1391 );
1392
1393 try std.testing.expectEqual(@as(u8, 3), dims.input_rank);
1394 try std.testing.expectEqual(@as(u32, 2), dims.rows);
1395 try std.testing.expectEqual(@as(u32, 15), dims.cols);
1396 try std.testing.expectEqual(@as(u32, 7), dims.inner);
1397 }
1398
1399 test "reduction static dims reject non-contiguous reduced axes" {
1400 var input_dims = [_]i64{ 4, 8, 8 };
1401 var output_dims = [_]i64{8};
1402 var axes = [_]i64{ 0, 2 };
1403
1404 try std.testing.expectError(error.CapabilityMismatch, reductionStaticDims(
1405 input_dims[0..],
1406 output_dims[0..],
1407 axes[0..],
1408 ));
1409 }