lib/accy/src/kernel/library/reduction.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3
4 const entry = @import("entry.zig");
5 const kernel = @import("../root.zig");
6
7 pub const OneDimensional = struct {
8 extent: u64,
9 threads: u32,
10 input_axis: []const u8 = "i",
11 output_axis: []const u8 = "out",
12 };
13
14 fn foldUpper(comptime extent: u64) i64 {
15 if (extent > @as(u64, @intCast(std.math.maxInt(i64)))) {
16 @compileError("kernel library reduction extent overflows index fold range");
17 }
18 return @intCast(extent);
19 }
20
21 fn sumSpecialization(comptime spec: OneDimensional) entry.Specialization {
22 return .{
23 .dtype = .f32,
24 .operation = .{ .reduction = .sum },
25 .equation = "i->",
26 .inputs = &.{entry.shape1D(spec.input_axis, spec.extent)},
27 .outputs = &.{entry.shapeScalar()},
28 .reductions = &.{entry.reduction("sum", .sum, entry.shape1D(spec.input_axis, spec.extent))},
29 .launch = entry.launch1D(1, spec.threads),
30 .schedule = entry.threadBlocks1D(spec.output_axis, 1, spec.threads),
31 };
32 }
33
34 fn dotSpecialization(comptime spec: OneDimensional) entry.Specialization {
35 return .{
36 .dtype = .f32,
37 .operation = .{ .reduction = .dot_product },
38 .equation = "i,i->",
39 .inputs = &.{
40 entry.shape1D(spec.input_axis, spec.extent),
41 entry.shape1D(spec.input_axis, spec.extent),
42 },
43 .outputs = &.{entry.shapeScalar()},
44 .reductions = &.{entry.reduction("dot", .dot_product, entry.shape1D(spec.input_axis, spec.extent))},
45 .launch = entry.launch1D(1, spec.threads),
46 .schedule = entry.threadBlocks1D(spec.output_axis, 1, spec.threads),
47 };
48 }
49
50 fn sum_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {
51 const zero = try inner.constantFloat(.f32, 0.0);
52 const sum = try inner.foldRange(0, foldUpper(ctx.spec.extent), 1, zero, .{
53 .src = ctx.args.param(.src),
54 }, sum_step);
55 try ctx.args.param(.dst).store(inner, sum, index);
56 }
57
58 fn sum_step(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
59 const value = try ctx.src.load(fold_inner, offset);
60 return fold_inner.add(acc, value.raw());
61 }
62
63 fn sumProgram(comptime spec: OneDimensional) type {
64 const Body = struct {
65 fn run(k: anytype, args: anytype) !void {
66 _ = try k.forEach1D(spec.output_axis, 1, .{ .spec = spec, .args = args }, sum_each);
67 }
68 };
69
70 return kernel.logical.Program(.{
71 .name = std.fmt.comptimePrint("accy_kernel_reduction_sum{}x{}_f32", .{ spec.extent, spec.threads }),
72 .parameters = .{
73 .dst = kernel.dynamicBuffer(.f32),
74 .src = kernel.dynamicBuffer(.f32),
75 },
76 .body = Body.run,
77 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
78 }
79
80 fn dot_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {
81 const zero = try inner.constantFloat(.f32, 0.0);
82 const sum = try inner.foldRange(0, foldUpper(ctx.spec.extent), 1, zero, .{
83 .lhs = ctx.args.param(.lhs),
84 .rhs = ctx.args.param(.rhs),
85 }, dot_step);
86 try ctx.args.param(.dst).store(inner, sum, index);
87 }
88
89 fn dot_step(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
90 const lhs = try ctx.lhs.load(fold_inner, offset);
91 const rhs = try ctx.rhs.load(fold_inner, offset);
92 const product = try lhs.mul(fold_inner, rhs);
93 return fold_inner.add(acc, product.raw());
94 }
95
96 fn dotProgram(comptime spec: OneDimensional) type {
97 const Body = struct {
98 fn run(k: anytype, args: anytype) !void {
99 _ = try k.forEach1D(spec.output_axis, 1, .{ .spec = spec, .args = args }, dot_each);
100 }
101 };
102
103 return kernel.logical.Program(.{
104 .name = std.fmt.comptimePrint("accy_kernel_reduction_dot{}x{}_f32", .{ spec.extent, spec.threads }),
105 .parameters = .{
106 .dst = kernel.dynamicBuffer(.f32),
107 .lhs = kernel.dynamicBuffer(.f32),
108 .rhs = kernel.dynamicBuffer(.f32),
109 },
110 .body = Body.run,
111 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
112 }
113
114 pub fn sumF32(comptime spec: OneDimensional) type {
115 return entry.Entry(sumProgram(spec), .{
116 .target = std.fmt.comptimePrint("accy.kernel.reduction.sum{}x{}_f32", .{ spec.extent, spec.threads }),
117 .layer = .logical,
118 .category = .reduction,
119 .specialization = sumSpecialization(spec),
120 });
121 }
122
123 pub fn dotF32(comptime spec: OneDimensional) type {
124 return entry.Entry(dotProgram(spec), .{
125 .target = std.fmt.comptimePrint("accy.kernel.reduction.dot{}x{}_f32", .{ spec.extent, spec.threads }),
126 .layer = .logical,
127 .category = .reduction,
128 .specialization = dotSpecialization(spec),
129 });
130 }
131
132 pub const Sum8F32 = sumF32(.{ .extent = 8, .threads = 1 });
133 pub const Dot8F32 = dotF32(.{ .extent = 8, .threads = 1 });
134
135 test "reduction sum entry runs on CPU" {
136 var src = [_]f32{ 1.0, -2.0, 3.0, 4.0, 5.5, 0.5, -1.0, 8.0 };
137 var dst = [_]f32{0.0};
138
139 try Sum8F32.runCpu(std.testing.allocator, Sum8F32.Limits.testing, &.{
140 kernel.argumentBuffer(f32, dst[0..]),
141 kernel.argumentBuffer(f32, src[0..]),
142 });
143 try std.testing.expectEqual(@as(f32, 19.0), dst[0]);
144
145 const launch_value = try Sum8F32.launch(std.testing.allocator, Sum8F32.Limits.testing);
146 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
147 try std.testing.expectEqual(@as(u32, 1), launch_value.block[0]);
148 }
149
150 test "reduction dot entry runs on CPU" {
151 var lhs = [_]f32{ 1.0, -2.0, 3.0, 4.0, 0.5, -1.5, 2.0, -3.0 };
152 var rhs = [_]f32{ 2.0, 3.0, -1.0, 0.25, 8.0, -2.0, 0.5, 4.0 };
153 var dst = [_]f32{0.0};
154
155 try Dot8F32.runCpu(std.testing.allocator, Dot8F32.Limits.testing, &.{
156 kernel.argumentBuffer(f32, dst[0..]),
157 kernel.argumentBuffer(f32, lhs[0..]),
158 kernel.argumentBuffer(f32, rhs[0..]),
159 });
160 try std.testing.expectEqual(@as(f32, -10.0), dst[0]);
161
162 const launch_value = try Dot8F32.launch(std.testing.allocator, Dot8F32.Limits.testing);
163 try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);
164 try std.testing.expectEqual(@as(u32, 1), launch_value.block[0]);
165 }
166
167 test "reduction constructor creates independent shape-specialized entries" {
168 const Sum16F32 = sumF32(.{ .extent = 16, .threads = 1 });
169 const Dot16F32 = dotF32(.{ .extent = 16, .threads = 1 });
170
171 try std.testing.expectEqualStrings("accy.kernel.reduction.sum8x1_f32", Sum8F32.target);
172 try std.testing.expectEqualStrings("accy.kernel.reduction.sum16x1_f32", Sum16F32.target);
173 try std.testing.expectEqualStrings("accy.kernel.reduction.dot8x1_f32", Dot8F32.target);
174 try std.testing.expectEqualStrings("accy.kernel.reduction.dot16x1_f32", Dot16F32.target);
175 try std.testing.expectEqual(@as(u64, 8), Sum8F32.specialization.inputs[0].elementCount().?);
176 try std.testing.expectEqualStrings("sum", Sum16F32.specialization.reductions[0].name);
177 try std.testing.expectEqual(entry.ReductionOperator.sum, Sum16F32.specialization.reductions[0].operator);
178 try std.testing.expect(Sum16F32.specialization.operationIs(.{ .reduction = .sum }));
179 try std.testing.expectEqual(@as(u64, 16), Sum16F32.specialization.reductions[0].shape.elementCount().?);
180 try std.testing.expectEqual(@as(u64, 1), Sum16F32.specialization.outputs[0].elementCount().?);
181 try std.testing.expectEqualStrings("dot", Dot16F32.specialization.reductions[0].name);
182 try std.testing.expectEqual(entry.ReductionOperator.dot_product, Dot16F32.specialization.reductions[0].operator);
183 try std.testing.expect(Dot16F32.specialization.operationIs(.{ .reduction = .dot_product }));
184 try std.testing.expectEqual(@as(u64, 16), Dot16F32.specialization.reductions[0].shape.elementCount().?);
185 try std.testing.expect(Dot16F32.specialization.outputHasExtents(0, &.{}));
186 try std.testing.expectEqualDeep(Sum16F32.specialization.launch.?, Sum16F32.specialization.schedule.?.launch());
187 try std.testing.expectEqualDeep(Dot16F32.specialization.launch.?, Dot16F32.specialization.schedule.?.launch());
188
189 var snapshot = try Sum16F32.scheduleSnapshot(std.testing.allocator, Sum16F32.Limits.testing);
190 defer snapshot.deinit(std.testing.allocator);
191 try std.testing.expect(Sum16F32.specialization.schedule.?.matchesSnapshot(&snapshot));
192
193 var dot_snapshot = try Dot16F32.scheduleSnapshot(std.testing.allocator, Dot16F32.Limits.testing);
194 defer dot_snapshot.deinit(std.testing.allocator);
195 try std.testing.expect(Dot16F32.specialization.schedule.?.matchesSnapshot(&dot_snapshot));
196
197 var src = @as([16]f32, @splat(1.0));
198 var dst = [_]f32{0.0};
199 try Sum16F32.runCpu(std.testing.allocator, Sum16F32.Limits.testing, &.{
200 kernel.argumentBuffer(f32, dst[0..]),
201 kernel.argumentBuffer(f32, src[0..]),
202 });
203 try std.testing.expectEqual(@as(f32, 16.0), dst[0]);
204
205 var dot_dst = [_]f32{0.0};
206 try Dot16F32.runCpu(std.testing.allocator, Dot16F32.Limits.testing, &.{
207 kernel.argumentBuffer(f32, dot_dst[0..]),
208 kernel.argumentBuffer(f32, src[0..]),
209 kernel.argumentBuffer(f32, src[0..]),
210 });
211 try std.testing.expectEqual(@as(f32, 16.0), dot_dst[0]);
212 }
213
214 test "reduction sum entry creates registry-ready artifact" {
215 const allocator = std.testing.allocator;
216 var state = gpu.recording.BackendState{
217 .allocator = allocator,
218 .kind = .cuda,
219 .format = .cuda_ptx,
220 };
221
222 var call_artifact = try Sum8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = Sum8F32.Limits.testing });
223 defer call_artifact.deinit();
224
225 const artifact = call_artifact.registry().find(Sum8F32.target, Sum8F32.version, .cuda_ptx) orelse {
226 return error.TestExpectedKernelCallArtifact;
227 };
228 try std.testing.expectEqualStrings(Sum8F32.name, artifact.entry_name);
229 try std.testing.expectEqual(@as(u32, 2), artifact.argument_count);
230 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
231 switch (artifact.launch) {
232 .fixed => |geometry| {
233 try std.testing.expectEqual(Sum8F32.specialization.launch.?.grid[0], geometry.grid[0]);
234 try std.testing.expectEqual(Sum8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
235 },
236 else => return error.TestExpectedFixedLaunch,
237 }
238 }
239
240 test "reduction dot entry creates registry-ready artifact" {
241 const allocator = std.testing.allocator;
242 var state = gpu.recording.BackendState{
243 .allocator = allocator,
244 .kind = .cuda,
245 .format = .cuda_ptx,
246 };
247
248 var call_artifact = try Dot8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = Dot8F32.Limits.testing });
249 defer call_artifact.deinit();
250
251 const artifact = call_artifact.registry().find(Dot8F32.target, Dot8F32.version, .cuda_ptx) orelse {
252 return error.TestExpectedKernelCallArtifact;
253 };
254 try std.testing.expectEqualStrings(Dot8F32.name, artifact.entry_name);
255 try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);
256 try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);
257 switch (artifact.launch) {
258 .fixed => |geometry| {
259 try std.testing.expectEqual(Dot8F32.specialization.launch.?.grid[0], geometry.grid[0]);
260 try std.testing.expectEqual(Dot8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);
261 },
262 else => return error.TestExpectedFixedLaunch,
263 }
264 }