lib/accy/src/kernel/logical/family.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const kernel = @import("../root.zig");
3 const logical_builder = @import("../model/logical/root.zig");
4 const parameter = @import("../program/root.zig").parameter;
5 const schedule_mod = logical_builder.schedule;
6
7 pub fn Family(comptime definition: anytype) type {
8 return struct {
9 const FamilySource: type = kernel.Family(FamilyDefinition(definition));
10
11 pub const name = definition.name;
12 pub const parameters = definition.parameters;
13 pub const Layout: type = FamilySource.Layout;
14 pub const Instance: type = definition.Instance;
15 pub const Limits: type = FamilySource.Limits;
16
17 pub const arg = FamilySource.arg;
18 pub const schema = FamilySource.schema;
19 pub const build = FamilySource.build;
20 pub const buildNamed = FamilySource.buildNamed;
21 pub const interpret = FamilySource.interpret;
22 pub const launch = FamilySource.launch;
23 pub const scheduleSnapshot = FamilySource.scheduleSnapshot;
24 pub const createPlan = FamilySource.createPlan;
25 pub const createCheckedPlan = FamilySource.createCheckedPlan;
26 pub const compileFragment = FamilySource.compileFragment;
27 pub const createKernelArtifact = FamilySource.createKernelArtifact;
28 pub const createKernelCallArtifact = FamilySource.createKernelCallArtifact;
29 pub const runCpu = FamilySource.runCpu;
30 pub const runCpuWithDiagnostic = FamilySource.runCpuWithDiagnostic;
31 pub const verify = FamilySource.verify;
32 };
33 }
34
35 fn FamilyDefinition(comptime definition: anytype) type {
36 return struct {
37 pub const name = definition.name;
38 pub const parameters = definition.parameters;
39 pub const Instance: type = definition.Instance;
40
41 pub fn body(raw_builder: anytype, instance: definition.Instance, _: anytype) !void {
42 var logical = logical_builder.wrap(raw_builder, scheduleValue(definition, instance));
43 try call(definition, instance, &logical);
44 }
45 };
46 }
47
48 fn ScheduleValue(comptime definition: anytype) type {
49 if (@hasField(@TypeOf(definition), "schedule")) {
50 return @TypeOf(definition.schedule(@as(definition.Instance, undefined)));
51 }
52 return schedule_mod.ThreadBlocks;
53 }
54
55 fn scheduleValue(comptime definition: anytype, instance: definition.Instance) ScheduleValue(definition) {
56 if (comptime @hasField(@TypeOf(definition), "schedule")) {
57 return definition.schedule(instance);
58 }
59 return schedule_mod.default();
60 }
61
62 fn call(comptime definition: anytype, instance: definition.Instance, k: anytype) !void {
63 if (comptime parameter.named(definition.parameters)) {
64 try definition.body(k, instance, parameter.Args(definition.parameters, PointerChild(@TypeOf(k))){ .builder = k });
65 } else {
66 const count = comptime parameter.arity(definition.parameters);
67 var args: [count]kernel.Value = undefined;
68 inline for (0..count) |index_value| {
69 args[index_value] = k.argument(index_value);
70 }
71 try definition.body(k, instance, args[0..]);
72 }
73 }
74
75 fn PointerChild(comptime Pointer: type) type {
76 return switch (@typeInfo(Pointer)) {
77 .pointer => |info| info.child,
78 else => @compileError("logical kernel Family bodies receive builder pointers"),
79 };
80 }
81
82 const ScaleInstance = struct {
83 extent: u64,
84 threads: u32,
85 };
86
87 fn scaleFamilySchedule(instance: ScaleInstance) schedule_mod.ThreadBlocks {
88 return schedule_mod.threadBlocks(.{ .x = instance.threads });
89 }
90
91 fn scaleFamilyBody(k: anytype, instance: ScaleInstance, args: anytype) !void {
92 _ = try k.forEach1D("i", instance.extent, args, struct {
93 fn each(inner: anytype, index: kernel.Index1D, each_args: anytype) !void {
94 const value = try each_args.param(.src).load(inner, index);
95 const scaled = try value.mul(inner, each_args.param(.scale));
96 try each_args.param(.dst).store(inner, scaled, index);
97 }
98 }.each);
99 }
100
101 const ScaleFamily = Family(.{
102 .name = "kernel_logical_scale_family_f32",
103 .parameters = .{
104 .src = kernel.dynamicBuffer(.f32),
105 .dst = kernel.dynamicBuffer(.f32),
106 .scale = kernel.scalar(.f32),
107 },
108 .Instance = ScaleInstance,
109 .schedule = scaleFamilySchedule,
110 .body = scaleFamilyBody,
111 });
112
113 test "logical Family schedules runtime instances through thread block policies" {
114 const compact_launch = try ScaleFamily.launch(std.testing.allocator, ScaleFamily.Limits.testing, .{ .extent = 5, .threads = 5 });
115 try std.testing.expectEqual(@as(u32, 1), compact_launch.grid[0]);
116 try std.testing.expectEqual(@as(u32, 5), compact_launch.block[0]);
117
118 const tiled_launch = try ScaleFamily.launch(std.testing.allocator, ScaleFamily.Limits.testing, .{ .extent = 6, .threads = 2 });
119 try std.testing.expectEqual(@as(u32, 3), tiled_launch.grid[0]);
120 try std.testing.expectEqual(@as(u32, 2), tiled_launch.block[0]);
121
122 var tiled_snapshot = try ScaleFamily.scheduleSnapshot(std.testing.allocator, ScaleFamily.Limits.testing, .{ .extent = 6, .threads = 2 });
123 defer tiled_snapshot.deinit(std.testing.allocator);
124 try std.testing.expectEqual(@as(usize, 2), tiled_snapshot.allAxes().len);
125 try std.testing.expectEqualStrings("i_tile", tiled_snapshot.allAxes()[0].name);
126 try std.testing.expectEqualStrings("i_lane", tiled_snapshot.allAxes()[1].name);
127 try std.testing.expectEqual(kernel.BindTarget.block_x, tiled_snapshot.allAxes()[0].bind.?);
128 try std.testing.expectEqual(kernel.BindTarget.thread_x, tiled_snapshot.allAxes()[1].bind.?);
129 }
130
131 test "logical Family executes runtime instances on the CPU oracle" {
132 var input = [_]f32{ 1.0, -2.0, 3.5, 4.0, -0.25, 6.0 };
133 var output = [_]f32{ 0, 0, 0, 0, 0, 0 };
134 try ScaleFamily.runCpu(std.testing.allocator, ScaleFamily.Limits.testing, .{ .extent = 6, .threads = 2 }, &.{
135 kernel.argumentBuffer(f32, input[0..]),
136 kernel.argumentBuffer(f32, output[0..]),
137 kernel.argumentF32(2.0),
138 });
139 try std.testing.expectEqualSlices(f32, &.{ 2.0, -4.0, 7.0, 8.0, -0.5, 12.0 }, output[0..]);
140
141 var head_output = [_]f32{ 0, 0, 0, 0, 0, 0 };
142 try ScaleFamily.runCpu(std.testing.allocator, ScaleFamily.Limits.testing, .{ .extent = 3, .threads = 4 }, &.{
143 kernel.argumentBuffer(f32, input[0..]),
144 kernel.argumentBuffer(f32, head_output[0..]),
145 kernel.argumentF32(3.0),
146 });
147 try std.testing.expectEqualSlices(f32, &.{ 3.0, -6.0, 10.5, 0.0, 0.0, 0.0 }, head_output[0..]);
148 }
149
150 const LoopCarryInstance = struct {
151 steps: u32,
152 };
153
154 fn loopCarryFamilyBody(k: anytype, instance: LoopCarryInstance, args: anytype) !void {
155 const zero_index = try k.constantIndex(0);
156 const one_index = try k.constantIndex(1);
157 const upper = try k.constantIndex(instance.steps);
158 const zero_i32 = try k.constantInt(.i32, 0);
159 const one_i32 = try k.constantInt(.i32, 1);
160
161 var init_values = [_]kernel.Value{ zero_i32, zero_i32 };
162 var result_types = [_]kernel.Type{ zero_i32.valueType(), zero_i32.valueType() };
163 var scope = try k.forScope(zero_index, upper, one_index, init_values[0..], result_types[0..]);
164 errdefer scope.abort();
165
166 const sum = scope.iterArg(0).?;
167 const count = scope.iterArg(1).?;
168 const step_i32 = try k.cast(scope.inductionVar(), .i32);
169 const next_sum = try k.add(sum, step_i32);
170 const next_count = try k.add(count, one_i32);
171 try scope.leave(&.{ next_sum, next_count });
172
173 const total = try k.add(scope.result(0).?, scope.result(1).?);
174 try args.param(.dst).store(k, total, zero_index);
175 }
176
177 const LoopCarryFamily = Family(.{
178 .name = "kernel_logical_loop_carry_i32",
179 .parameters = .{
180 .dst = kernel.dynamicBuffer(.i32),
181 },
182 .Instance = LoopCarryInstance,
183 .body = loopCarryFamilyBody,
184 });
185
186 const ForScopeCounter = struct {
187 for_scope_count: usize = 0,
188 iter_arg_count: usize = 0,
189 leave_count: usize = 0,
190 result_count: usize = 0,
191
192 pub const Result = struct {
193 graph: kernel.Graph,
194 for_scope_count: usize,
195 iter_arg_count: usize,
196 leave_count: usize,
197 result_count: usize,
198
199 pub fn deinit(self: *@This()) void {
200 self.graph.deinit();
201 self.* = undefined;
202 }
203 };
204
205 pub fn forScope(self: *@This(), ctx: anytype) !@TypeOf(ctx.*).ScopeType {
206 self.for_scope_count += 1;
207 return ctx.default();
208 }
209
210 pub fn forScopeIterArg(self: *@This(), ctx: anytype) ?kernel.Value {
211 self.iter_arg_count += 1;
212 return ctx.default();
213 }
214
215 pub fn forScopeLeave(self: *@This(), ctx: anytype) !void {
216 self.leave_count += 1;
217 try ctx.default();
218 }
219
220 pub fn forScopeResult(self: *@This(), ctx: anytype) ?kernel.Value {
221 self.result_count += 1;
222 return ctx.default();
223 }
224
225 pub fn finish(self: *@This(), ctx: anytype) !Result {
226 return .{
227 .graph = try ctx.default(),
228 .for_scope_count = self.for_scope_count,
229 .iter_arg_count = self.iter_arg_count,
230 .leave_count = self.leave_count,
231 .result_count = self.result_count,
232 };
233 }
234 };
235
236 test "logical Family forwards forScope through interpret layers" {
237 var result = try LoopCarryFamily.interpret(std.testing.allocator, LoopCarryFamily.Limits.testing, .{ .steps = 4 }, ForScopeCounter{});
238 defer result.deinit();
239
240 try std.testing.expectEqual(@as(usize, 1), result.for_scope_count);
241 try std.testing.expectEqual(@as(usize, 2), result.iter_arg_count);
242 try std.testing.expectEqual(@as(usize, 1), result.leave_count);
243 try std.testing.expectEqual(@as(usize, 2), result.result_count);
244 try result.graph.verify();
245
246 var output = [_]i32{0};
247 try result.graph.runCpu(std.testing.allocator, &.{
248 kernel.argumentBuffer(i32, output[0..]),
249 });
250 try std.testing.expectEqual(@as(i32, 10), output[0]);
251 }