lib/accy/src/preparation/kernelization/lowering/row.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_name = @import("name.zig");
10 const generated_schedule = @import("schedule.zig");
11 const preparation = @import("../../root.zig");
12 const reduction = @import("reduction.zig");
13
14 const ir = common.ir;
15 const dialect_mod = common.dialect_mod;
16 const kernel_root = common.kernel_root;
17 const bufferization = common.bufferization;
18 const kernelization_model = @import("../model/root.zig");
19 const schedule_planning = common.schedule_planning;
20 const target_facts = preparation.target;
21 const bufferSlotById = common.bufferSlotById;
22 const externalInputIndex = common.externalInputIndex;
23 const mapKernelBuildError = common.mapKernelBuildError;
24 const isName = common.isName;
25
26 const LoweredKernel = kernelization_model.LoweredKernel;
27 const ReductionKind = kernelization_model.ReductionKind;
28 const Schedule = target_facts.GeneratedRowPipelineSchedule;
29
30 const warp_size: u32 = 32;
31 const max_reduces = 4;
32 const max_leaves = 6;
33 const max_quads = 8;
34
35 /// A pass declares its own costs before it runs as its work bound so compilation can refuse a pass
36 /// that would exceed the caller's limits, and the kernel stage calls this once per try at emitting
37 /// a kernel under one candidate schedule to charge row-kernel scratch. Block-wide generation keeps
38 /// one memo table for each reduction and output lane and one for shared statistics, and one-thread
39 /// generation adds one more table: each is a map from a source value to the kernel value already
40 /// emitted for it. Each table is sized for `value_count` source values. The parameter list and
41 /// input arrays belong to the caller and are charged separately.
42 pub fn scratchStorageBound(value_count: u64) !u64 {
43 const accounting = @import("choir").passes.pass.work;
44 const memos = (max_reduces + 1) * max_quads * 4 + 2;
45 const table = try accounting.hashMapGrowth(*ir.Value, kernel_root.Value, value_count);
46 const bytes = try accounting.multiply(memos, table);
47 if (bytes > std.math.maxInt(usize)) return error.WorkOverflow;
48 return bytes;
49 }
50
51 const block_schedules = [_]u32{ 256, 512, 128 };
52
53 pub const schedule_version: u32 = 1;
54 pub const max_schedule_candidates = block_schedules.len;
55
56 pub fn scheduleCandidates(
57 cols: u64,
58 format: ?gpu.ArtifactFormat,
59 buffer: *[max_schedule_candidates]Schedule,
60 ) []const Schedule {
61 if (format != .cuda_ptx) return buffer[0..0];
62 var count: usize = 0;
63 for (block_schedules) |candidate_threads| {
64 if (!scheduleViable(candidate_threads, cols)) continue;
65 buffer[count] = .{ .threads = candidate_threads };
66 count += 1;
67 }
68 return buffer[0..count];
69 }
70
71 fn scheduleViable(candidate_threads: u32, cols: u64) bool {
72 const quad_width = @as(u64, candidate_threads) * 4;
73 if (cols % quad_width != 0) return false;
74 return cols / quad_width <= max_quads;
75 }
76
77 const PipelineReduce = struct {
78 op: *ir.Operation,
79 kind: ReductionKind,
80 init: f32,
81 input: *ir.Value,
82 };
83
84 const PipelineLeafKind = enum {
85 full,
86 column,
87 };
88
89 const PipelineLeaf = struct {
90 value: *ir.Value,
91 input_index: usize,
92 kind: PipelineLeafKind,
93 };
94
95 const PipelineDescription = struct {
96 rows: u32,
97 cols: u32,
98 root: *ir.Operation,
99 reduces: [max_reduces]PipelineReduce,
100 reduce_count: usize,
101 leaves: [max_leaves]PipelineLeaf,
102 leaf_count: usize,
103 };
104
105 fn BlockRowPipelineBodyType(comptime pipeline_threads: u32) type {
106 return struct {
107 fn emit(logical: anytype, ctx: anytype) !void {
108 try emitBlockRowPipelineBodyScheduled(pipeline_threads, logical, ctx);
109 }
110 };
111 }
112
113 pub fn lower(
114 allocator: std.mem.Allocator,
115 ir_ctx: *ir.Context,
116 outline: kernelization_model.KernelOutline,
117 work: schedule_planning.ScheduleWorkItem,
118 buffer_plan: *const bufferization.BufferPlanAnalysis,
119 format: ?gpu.ArtifactFormat,
120 row_pipeline_schedules: ?[]const u8,
121 ) common.LoweringError!LoweredKernel {
122 const desc = try pipelineDescriptionForWork(outline, work, buffer_plan);
123
124 const input_dtypes = allocator.alloc(choir_abi.DType, outline.inputCount()) catch return error.OutOfMemory;
125 defer allocator.free(input_dtypes);
126 for (outline.input_slot_ids, 0..) |slot_id, index| {
127 const slot = bufferSlotById(buffer_plan, slot_id) orelse return error.InvalidArtifact;
128 input_dtypes[index] = slot.dtype;
129 }
130
131 var abi = try generated_abi.flatTyped(allocator, .f32, input_dtypes);
132 defer abi.deinit(allocator);
133
134 if (format == .cuda_ptx) {
135 if (row_pipeline_schedules) |encoded| {
136 if (try target_facts.resolveGeneratedRowPipelineSchedule(encoded, desc.rows, desc.cols)) |schedule| {
137 if (scheduleViable(schedule.threads, desc.cols)) {
138 if (try lowerBlockRowPipeline(schedule, allocator, ir_ctx, outline, work, buffer_plan, desc, abi)) |lowered| {
139 return lowered;
140 }
141 }
142 }
143 }
144 var candidate_buffer: [max_schedule_candidates]Schedule = undefined;
145 for (scheduleCandidates(desc.cols, format, &candidate_buffer)) |schedule| {
146 if (try lowerBlockRowPipeline(schedule, allocator, ir_ctx, outline, work, buffer_plan, desc, abi)) |lowered| {
147 return lowered;
148 }
149 }
150 }
151
152 const entry_name = try generated_name.rowPipelineSerial(allocator, desc.reduce_count, work.ops.len, work.id);
153 errdefer allocator.free(entry_name);
154 return generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flat(), .{
155 .abi = abi,
156 .allocator = allocator,
157 .desc = desc,
158 .work = work,
159 .outline = outline,
160 .buffer_plan = buffer_plan,
161 }, emitSerialRowPipelineBody);
162 }
163
164 fn pipelineDescriptionForWork(
165 outline: kernelization_model.KernelOutline,
166 work: schedule_planning.ScheduleWorkItem,
167 buffer_plan: *const bufferization.BufferPlanAnalysis,
168 ) common.LoweringError!PipelineDescription {
169 if (work.kind != .row_pipeline) return error.UnsupportedOperation;
170 if (work.ops.len == 0) return error.UnsupportedOperation;
171 const root = work.ops[work.ops.len - 1];
172 const root_result = root.getResult(0) orelse return error.InvalidArtifact;
173
174 var dims_arena_buffer: [256]u8 = undefined;
175 var dims_arena = alloc_fixed.FixedBuffer.init(dims_arena_buffer[0..]);
176 const root_type = dialect_mod.decodeTensorType(dims_arena.allocator(), root_result.type) catch return error.InvalidArtifact;
177 if (root_type.dtype != .f32 or root_type.dims.len != 2) return error.UnsupportedOperation;
178 const rows = std.math.cast(u32, root_type.dims[0]) orelse return error.UnsupportedOperation;
179 const cols = std.math.cast(u32, root_type.dims[1]) orelse return error.UnsupportedOperation;
180 if (rows == 0 or cols == 0) return error.UnsupportedOperation;
181
182 var desc = PipelineDescription{
183 .rows = rows,
184 .cols = cols,
185 .root = root,
186 .reduces = undefined,
187 .reduce_count = 0,
188 .leaves = undefined,
189 .leaf_count = 0,
190 };
191
192 for (work.ops) |op| {
193 if (!isName(op.name.name, dialect_mod.AccyDialect.ReduceOp.operation_name)) continue;
194 if (desc.reduce_count >= max_reduces) return error.UnsupportedOperation;
195 const kind = try reduction.reductionKindForOp(op);
196 const init_operand = op.getOperand(1) orelse return error.InvalidArtifact;
197 const init_slot = buffer_plan.getSlot(init_operand) orelse return error.InvalidArtifact;
198 const init = switch (try reduction.reductionInitValue(init_slot.*, .f32, outline)) {
199 .constant => |value| switch (value) {
200 .f32 => |scalar| scalar,
201 else => return error.UnsupportedOperation,
202 },
203 .input_buffer => return error.UnsupportedOperation,
204 };
205 const input = op.getOperand(0) orelse return error.InvalidArtifact;
206 desc.reduces[desc.reduce_count] = .{ .op = op, .kind = kind, .init = init, .input = input };
207 desc.reduce_count += 1;
208 }
209 if (desc.reduce_count == 0) return error.UnsupportedOperation;
210
211 var member_set = MemberSet{};
212 for (work.ops) |op| member_set.put(op);
213
214 for (work.ops) |op| {
215 if (isName(op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) {
216 var dims_buffer: [8]i64 = undefined;
217 const dims = readBroadcastDimsBounded(op, dims_buffer[0..]) orelse return error.InvalidArtifact;
218 if (dims.len != 1) return error.UnsupportedOperation;
219 if (dims[0] != 1) continue;
220 const source = op.getOperand(0) orelse return error.InvalidArtifact;
221 if (leafIndexFor(&desc, source) != null) continue;
222 if (desc.leaf_count >= max_leaves) return error.UnsupportedOperation;
223 const slot = buffer_plan.getSlot(source) orelse return error.UnsupportedOperation;
224 const input_index = externalInputIndex(outline, slot.id) orelse return error.UnsupportedOperation;
225 desc.leaves[desc.leaf_count] = .{ .value = source, .input_index = input_index, .kind = .column };
226 desc.leaf_count += 1;
227 continue;
228 }
229 const operand_limit: usize = if (isName(op.name.name, dialect_mod.AccyDialect.ReduceOp.operation_name)) 1 else op.getOperandValues().len;
230 for (op.getOperandValues()[0..operand_limit]) |operand| {
231 if (operand.getDefiningOp()) |def_any| {
232 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
233 if (member_set.contains(def_op)) continue;
234 if (isName(def_op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) continue;
235 if (broadcastOfConstant(def_op)) continue;
236 }
237 if (leafIndexFor(&desc, operand) != null) continue;
238 var operand_arena_buffer: [256]u8 = undefined;
239 var operand_arena = alloc_fixed.FixedBuffer.init(operand_arena_buffer[0..]);
240 const operand_type = dialect_mod.decodeTensorType(operand_arena.allocator(), operand.type) catch return error.InvalidArtifact;
241 if (operand_type.dims.len != 2) return error.UnsupportedOperation;
242 if (desc.leaf_count >= max_leaves) return error.UnsupportedOperation;
243 const slot = buffer_plan.getSlot(operand) orelse return error.UnsupportedOperation;
244 const input_index = externalInputIndex(outline, slot.id) orelse return error.UnsupportedOperation;
245 desc.leaves[desc.leaf_count] = .{ .value = operand, .input_index = input_index, .kind = .full };
246 desc.leaf_count += 1;
247 }
248 }
249
250 return desc;
251 }
252
253 fn broadcastOfConstant(op: *ir.Operation) bool {
254 if (!isName(op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) return false;
255 const operands = op.getOperandValues();
256 if (operands.len != 1) return false;
257 const source_any = operands[0].getDefiningOp() orelse return false;
258 const source: *ir.Operation = @ptrCast(@alignCast(source_any));
259 return isName(source.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name);
260 }
261
262 fn readBroadcastDimsBounded(op: *ir.Operation, buffer: []i64) ?[]const i64 {
263 const payload = payloadFor(op, "broadcast_dims") orelse return null;
264 if (payload.len % @sizeOf(i64) != 0) return null;
265 const count = payload.len / @sizeOf(i64);
266 if (count > buffer.len) return null;
267 for (buffer[0..count], 0..) |*slot, index| {
268 @memcpy(std.mem.asBytes(slot), payload[index * @sizeOf(i64) ..][0..@sizeOf(i64)]);
269 }
270 return buffer[0..count];
271 }
272
273 fn payloadFor(op: *ir.Operation, name: []const u8) ?[]const u8 {
274 if (op.getAttrAs(ir.Attribute.DialectAttr, name)) |attr| return attr.payload;
275 if (op.getAttrAs(ir.Attribute.DialectAttr, dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("broadcast_dims"))) |attr| return attr.payload;
276 return null;
277 }
278
279 const MemberSet = struct {
280 ops: [fusionMaxOps()]*ir.Operation = undefined,
281 count: usize = 0,
282
283 fn put(self: *MemberSet, op: *ir.Operation) void {
284 if (self.count < self.ops.len) {
285 self.ops[self.count] = op;
286 self.count += 1;
287 }
288 }
289
290 fn contains(self: *const MemberSet, op: *ir.Operation) bool {
291 for (self.ops[0..self.count]) |member| {
292 if (member == op) return true;
293 }
294 return false;
295 }
296 };
297
298 fn fusionMaxOps() usize {
299 return 32;
300 }
301
302 fn leafIndexFor(desc: *const PipelineDescription, value: *ir.Value) ?usize {
303 for (desc.leaves[0..desc.leaf_count], 0..) |leaf, index| {
304 if (leaf.value == value) return index;
305 }
306 return null;
307 }
308
309 fn reduceIndexFor(desc: *const PipelineDescription, op: *ir.Operation) ?usize {
310 for (desc.reduces[0..desc.reduce_count], 0..) |entry, index| {
311 if (entry.op == op) return index;
312 }
313 return null;
314 }
315
316 const ResolveContext = struct {
317 desc: *const PipelineDescription,
318 scalars: [max_reduces]?kernel_root.Value,
319 leaf_values: [max_leaves]kernel_root.Value,
320 memo: std.AutoHashMap(*ir.Value, kernel_root.Value),
321 stats: ?*std.AutoHashMap(*ir.Value, kernel_root.Value) = null,
322
323 fn init(allocator: std.mem.Allocator, desc: *const PipelineDescription) ResolveContext {
324 return .{
325 .desc = desc,
326 .scalars = @splat(null),
327 .leaf_values = undefined,
328 .memo = std.AutoHashMap(*ir.Value, kernel_root.Value).init(allocator),
329 };
330 }
331
332 fn deinit(self: *ResolveContext) void {
333 self.memo.deinit();
334 }
335 };
336
337 fn resolveValue(
338 builder: anytype,
339 ctx: *ResolveContext,
340 value: *ir.Value,
341 ) common.LoweringError!kernel_root.Value {
342 if (leafIndexFor(ctx.desc, value)) |leaf_index| return ctx.leaf_values[leaf_index];
343 if (ctx.memo.get(value)) |cached| return cached;
344
345 const def_any = value.getDefiningOp() orelse return error.UnsupportedOperation;
346 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
347
348 if (reduceIndexFor(ctx.desc, def_op)) |reduce_index| {
349 return ctx.scalars[reduce_index] orelse error.UnsupportedOperation;
350 }
351
352 if (isName(def_op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) {
353 return common.splatConstantValue(builder, def_op) catch |err|
354 return common.mapGeneratedKernelError(err);
355 }
356
357 if (isName(def_op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) {
358 var dims_buffer: [8]i64 = undefined;
359 const dims = readBroadcastDimsBounded(def_op, dims_buffer[0..]) orelse return error.InvalidArtifact;
360 const source = def_op.getOperand(0) orelse return error.InvalidArtifact;
361 if (dims.len == 0) {
362 return resolveValue(builder, ctx, source);
363 }
364 if (dims.len != 1) return error.UnsupportedOperation;
365 if (dims[0] == 1) {
366 return resolveValue(builder, ctx, source);
367 }
368 if (ctx.stats) |stats| {
369 if (stats.get(value)) |cached| return cached;
370 }
371 const resolved = try resolveStatValue(builder, ctx, source);
372 if (ctx.stats) |stats| {
373 stats.put(value, resolved) catch return error.OutOfMemory;
374 }
375 return resolved;
376 }
377
378 const kind = generated_elementwise.kernelForOperation(def_op) orelse return error.UnsupportedOperation;
379 const operands = def_op.getOperandValues();
380 if (operands.len > 3) return error.UnsupportedOperation;
381 var inputs: [3]kernel_root.Value = undefined;
382 for (operands, 0..) |operand, operand_index| {
383 inputs[operand_index] = try resolveValue(builder, ctx, operand);
384 }
385 const resolved = generated_elementwise.emitElementwiseValue(builder, kind, inputs[0..operands.len], .f32, def_op) catch |err| return common.mapGeneratedKernelError(err);
386 ctx.memo.put(value, resolved) catch return error.OutOfMemory;
387 return resolved;
388 }
389
390 fn resolveStatValue(
391 builder: anytype,
392 ctx: *ResolveContext,
393 value: *ir.Value,
394 ) common.LoweringError!kernel_root.Value {
395 if (ctx.stats) |stats| {
396 if (stats.get(value)) |cached| return cached;
397 }
398 const def_any = value.getDefiningOp() orelse return error.UnsupportedOperation;
399 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
400
401 if (reduceIndexFor(ctx.desc, def_op)) |reduce_index| {
402 return ctx.scalars[reduce_index] orelse error.UnsupportedOperation;
403 }
404 if (isName(def_op.name.name, dialect_mod.AccyDialect.ConstantOp.operation_name)) {
405 return common.splatConstantValue(builder, def_op) catch |err|
406 return common.mapGeneratedKernelError(err);
407 }
408 const kind = generated_elementwise.kernelForOperation(def_op) orelse return error.UnsupportedOperation;
409 const operands = def_op.getOperandValues();
410 if (operands.len > 3) return error.UnsupportedOperation;
411 var inputs: [3]kernel_root.Value = undefined;
412 for (operands, 0..) |operand, operand_index| {
413 inputs[operand_index] = try resolveStatValue(builder, ctx, operand);
414 }
415 const resolved = generated_elementwise.emitElementwiseValue(builder, kind, inputs[0..operands.len], .f32, def_op) catch |err| return common.mapGeneratedKernelError(err);
416 if (ctx.stats) |stats| {
417 stats.put(value, resolved) catch return error.OutOfMemory;
418 }
419 return resolved;
420 }
421
422 fn lowerBlockRowPipeline(
423 schedule: Schedule,
424 allocator: std.mem.Allocator,
425 ir_ctx: *ir.Context,
426 outline: kernelization_model.KernelOutline,
427 work: schedule_planning.ScheduleWorkItem,
428 buffer_plan: *const bufferization.BufferPlanAnalysis,
429 desc: PipelineDescription,
430 abi: generated_abi.Flat,
431 ) common.LoweringError!?LoweredKernel {
432 inline for (block_schedules) |candidate_threads| {
433 if (candidate_threads == schedule.threads) {
434 const entry_name = try generated_name.rowPipeline(allocator, desc.reduce_count, candidate_threads, work.ops.len, work.id);
435 errdefer allocator.free(entry_name);
436 var lowered = try generated_builder.withoutLaunch(allocator, ir_ctx, work.id, entry_name, abi.params(), generated_schedule.flatThreads(candidate_threads), .{
437 .abi = abi,
438 .allocator = allocator,
439 .desc = desc,
440 .work = work,
441 .outline = outline,
442 .buffer_plan = buffer_plan,
443 }, BlockRowPipelineBodyType(candidate_threads).emit);
444 lowered.body = .{ .row_pipeline = .{
445 .threads = candidate_threads,
446 .rows = desc.rows,
447 .cols = desc.cols,
448 .warps = candidate_threads / warp_size,
449 } };
450 return lowered;
451 }
452 }
453 return null;
454 }
455
456 fn store_row_pipeline_partial(inner: anytype, store_ctx: anytype) !void {
457 try inner.storeIndex(store_ctx.acc, store_ctx.partials, store_ctx.warp);
458 }
459
460 fn emitBlockRowPipelineBodyScheduled(comptime pipeline_threads: u32, logical: anytype, ctx: anytype) !void {
461 const desc: PipelineDescription = ctx.desc;
462 const quads: u32 = desc.cols / (pipeline_threads * 4);
463 const warps: u32 = pipeline_threads / warp_size;
464
465 const domain = try logical.index1D("lane", @as(u64, desc.rows) * pipeline_threads);
466 const global = domain.index;
467 const threads_extent = try logical.constantIndex(pipeline_threads);
468 const row = try logical.div(global, threads_extent);
469 const tid = try logical.sub(global, try logical.mul(row, threads_extent));
470 const cols_extent = try logical.constantIndex(desc.cols);
471 const row_base = try logical.mul(row, cols_extent);
472
473 const warp_extent = try logical.constantIndex(warp_size);
474 const warp = try logical.div(tid, warp_extent);
475 const lane = try logical.sub(tid, try logical.mul(warp, warp_extent));
476
477 const partials = try logical.sharedBuffer(.f32, warps);
478
479 var leaf_quads: [max_leaves][max_quads]kernel_root.Value = undefined;
480 for (desc.leaves[0..desc.leaf_count], 0..) |leaf, leaf_index| {
481 const input = ctx.abi.input(logical, leaf.input_index);
482 var quad: u32 = 0;
483 while (quad < quads) : (quad += 1) {
484 const offset = try logical.constantIndex(@as(i64, quad) * pipeline_threads);
485 const element = try logical.mul(try logical.add(offset, tid), try logical.constantIndex(4));
486 const index = switch (leaf.kind) {
487 .full => try logical.add(row_base, element),
488 .column => element,
489 };
490 leaf_quads[leaf_index][quad] = try logical.loadVector(input, index, 4);
491 }
492 }
493
494 var stat_values = std.AutoHashMap(*ir.Value, kernel_root.Value).init(ctx.allocator);
495 defer stat_values.deinit();
496
497 var scalars: [max_reduces]?kernel_root.Value = @splat(null);
498 for (desc.reduces[0..desc.reduce_count], 0..) |entry, reduce_index| {
499 var acc = try reduction.reductionNeutral(logical, entry.kind, .f32);
500 var quad: u32 = 0;
501 while (quad < quads) : (quad += 1) {
502 var lane_index: u32 = 0;
503 while (lane_index < 4) : (lane_index += 1) {
504 var resolve_ctx = ResolveContext.init(ctx.allocator, &desc);
505 defer resolve_ctx.deinit();
506 resolve_ctx.scalars = scalars;
507 resolve_ctx.stats = &stat_values;
508 for (0..desc.leaf_count) |leaf_index| {
509 resolve_ctx.leaf_values[leaf_index] = try logical.extractLane(leaf_quads[leaf_index][quad], lane_index, .f32);
510 }
511 const value = try resolveValue(logical, &resolve_ctx, entry.input);
512 acc = try reduction.emitReductionValue(logical, entry.kind, acc, value);
513 }
514 }
515 acc = try logical.warpReduce(reduction.warpKindForReduction(entry.kind), acc);
516
517 if (reduce_index > 0) try logical.barrier(.block);
518 const zero = try logical.constantIndex(0);
519 const lane_zero = try logical.compare(.eq, lane, zero);
520 try logical.guardDo(
521 lane_zero,
522 .{ .partials = partials, .warp = warp, .acc = acc },
523 store_row_pipeline_partial,
524 );
525 try logical.barrier(.block);
526
527 var total = try logical.constantFloat(.f32, entry.init);
528 var warp_index: u32 = 0;
529 while (warp_index < warps) : (warp_index += 1) {
530 const loaded = try logical.loadIndex(partials, try logical.constantIndex(warp_index));
531 total = try reduction.emitReductionValue(logical, entry.kind, total, loaded);
532 }
533 scalars[reduce_index] = total;
534 }
535
536 const out = ctx.abi.output(logical);
537 const root_result = desc.root.getResult(0) orelse return error.InvalidArtifact;
538 var quad: u32 = 0;
539 while (quad < quads) : (quad += 1) {
540 const offset = try logical.constantIndex(@as(i64, quad) * pipeline_threads);
541 const element = try logical.mul(try logical.add(offset, tid), try logical.constantIndex(4));
542 const base = try logical.add(row_base, element);
543 var lane_index: u32 = 0;
544 while (lane_index < 4) : (lane_index += 1) {
545 var resolve_ctx = ResolveContext.init(ctx.allocator, &desc);
546 defer resolve_ctx.deinit();
547 resolve_ctx.scalars = scalars;
548 resolve_ctx.stats = &stat_values;
549 for (0..desc.leaf_count) |leaf_index| {
550 resolve_ctx.leaf_values[leaf_index] = try logical.extractLane(leaf_quads[leaf_index][quad], lane_index, .f32);
551 }
552 const value = try resolveValue(logical, &resolve_ctx, root_result);
553 const index = try logical.add(base, try logical.constantIndex(lane_index));
554 try logical.storeIndex(value, out, index);
555 }
556 }
557 }
558
559 fn store_serial_row_pipeline_value(
560 store_inner: anytype,
561 col: kernel_root.Value,
562 store_ctx: anytype,
563 ) !void {
564 const value = try serialResolve(
565 store_inner,
566 store_ctx,
567 store_ctx.scalars,
568 store_ctx.row_base,
569 col,
570 store_ctx.root_result,
571 );
572 const index = try store_inner.add(store_ctx.row_base, col);
573 try store_inner.storeIndex(value, store_ctx.out, index);
574 }
575
576 fn emit_serial_row_pipeline_row(inner: anytype, guard_ctx: anytype) !void {
577 const desc_inner: PipelineDescription = guard_ctx.desc;
578 const cols_extent = try inner.constantIndex(desc_inner.cols);
579 const row_base = try inner.mul(guard_ctx.row, cols_extent);
580 const one = try inner.constantIndex(1);
581 const zero = try inner.constantIndex(0);
582
583 var scalars: [max_reduces]?kernel_root.Value = @splat(null);
584 for (desc_inner.reduces[0..desc_inner.reduce_count], 0..) |entry, reduce_index| {
585 const init_value = try inner.constantFloat(.f32, entry.init);
586 var scope = try inner.forScope(
587 zero,
588 cols_extent,
589 one,
590 &.{init_value},
591 &.{init_value.valueType()},
592 );
593 errdefer scope.abort();
594 const col = scope.inductionVar();
595 const current = scope.iterArg(0) orelse return error.UnsupportedOperation;
596 const value = try serialResolve(inner, guard_ctx, scalars, row_base, col, entry.input);
597 const updated = try reduction.emitReductionValue(inner, entry.kind, current, value);
598 try scope.leave(&.{updated});
599 scalars[reduce_index] = scope.result(0) orelse return error.UnsupportedOperation;
600 }
601
602 const out = guard_ctx.abi.output(inner);
603 const root_result = desc_inner.root.getResult(0) orelse return error.InvalidArtifact;
604 _ = try inner.forDo(zero, cols_extent, one, .{
605 .abi = guard_ctx.abi,
606 .allocator = guard_ctx.allocator,
607 .desc = desc_inner,
608 .row_base = row_base,
609 .scalars = scalars,
610 .root_result = root_result,
611 .out = out,
612 }, store_serial_row_pipeline_value);
613 }
614
615 fn emitSerialRowPipelineBody(logical: anytype, ctx: anytype) !void {
616 const desc: PipelineDescription = ctx.desc;
617
618 const domain = try logical.index1D("row", desc.rows);
619 const row = domain.index;
620 const rows_extent = try logical.constantIndex(desc.rows);
621 const row_ok = try logical.compare(.lt, row, rows_extent);
622 try logical.guardDo(row_ok, .{
623 .abi = ctx.abi,
624 .allocator = ctx.allocator,
625 .desc = desc,
626 .row = row,
627 }, emit_serial_row_pipeline_row);
628 }
629
630 fn serialResolve(
631 inner: anytype,
632 guard_ctx: anytype,
633 scalars: [max_reduces]?kernel_root.Value,
634 row_base: kernel_root.Value,
635 col: kernel_root.Value,
636 value: *ir.Value,
637 ) common.LoweringError!kernel_root.Value {
638 const desc: PipelineDescription = guard_ctx.desc;
639 var resolve_ctx = ResolveContext.init(guard_ctx.allocator, &desc);
640 defer resolve_ctx.deinit();
641 resolve_ctx.scalars = scalars;
642 for (desc.leaves[0..desc.leaf_count], 0..) |leaf, leaf_index| {
643 const input = guard_ctx.abi.input(inner, leaf.input_index);
644 const index = switch (leaf.kind) {
645 .full => inner.add(row_base, col) catch |err| return mapKernelBuildError(err),
646 .column => col,
647 };
648 resolve_ctx.leaf_values[leaf_index] = inner.loadIndex(input, index) catch |err| return mapKernelBuildError(err);
649 }
650 return resolveValue(inner, &resolve_ctx, value);
651 }