lib/accy/src/kernel/model/logical/schedule/root.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const kernel_program = @import("../../program/root.zig");
2 const domain = @import("../domain/root.zig");
3
4 pub const Threads = struct {
5 x: u32 = 256,
6 y: u32 = 4,
7 z: u32 = 1,
8 };
9
10 pub const ThreadBlocks = struct {
11 threads: Threads,
12
13 pub fn index1D(self: ThreadBlocks, owner: anytype, axis_value: domain.Axis) !kernel_program.Index1D {
14 return owner.index1D(axis_value.name, axis_value.extent, self.threads.x);
15 }
16
17 pub fn index2D(self: ThreadBlocks, owner: anytype, shape: domain.Domain2D) !kernel_program.Index2D {
18 return owner.index2D(.{
19 .x = kernel_program.domainAxis(shape.x.name, shape.x.extent, self.threads.x),
20 .y = kernel_program.domainAxis(shape.y.name, shape.y.extent, self.threads.y),
21 });
22 }
23
24 pub fn index3D(self: ThreadBlocks, owner: anytype, shape: domain.Domain3D) !kernel_program.Index3D {
25 return owner.index3D(.{
26 .x = kernel_program.domainAxis(shape.x.name, shape.x.extent, self.threads.x),
27 .y = kernel_program.domainAxis(shape.y.name, shape.y.extent, self.threads.y),
28 .z = kernel_program.domainAxis(shape.z.name, shape.z.extent, self.threads.z),
29 });
30 }
31 };
32
33 pub const UseThreads = struct {
34 threads: Threads,
35
36 pub fn index1D(self: *@This(), ctx: anytype) !kernel_program.Index1D {
37 return ctx.useThreads(self.threads.x);
38 }
39
40 pub fn index2D(self: *@This(), ctx: anytype) !kernel_program.Index2D {
41 return ctx.useThreads(self.threads);
42 }
43
44 pub fn index3D(self: *@This(), ctx: anytype) !kernel_program.Index3D {
45 return ctx.useThreads(self.threads);
46 }
47 };
48
49 pub fn with(impl: anytype) With(@TypeOf(impl)) {
50 return .{ .impl = impl };
51 }
52
53 pub fn useThreads(threads: Threads) With(UseThreads) {
54 return with(UseThreads{ .threads = threads });
55 }
56
57 pub fn With(comptime Impl: type) type {
58 return struct {
59 impl: Impl,
60
61 pub fn attach(self: @This(), next: anytype) Layer(@TypeOf(next), Impl) {
62 return .{
63 .next = next,
64 .impl = self.impl,
65 };
66 }
67 };
68 }
69
70 pub fn stack(specs: anytype) StackAttach(0, @TypeOf(specs), ThreadBlocks) {
71 return attachFrom(0, specs, default());
72 }
73
74 pub fn Layer(comptime Next: type, comptime Impl: type) type {
75 return struct {
76 next: Next,
77 impl: Impl,
78
79 const Self = @This();
80
81 pub fn index1D(self: *Self, owner: anytype, axis_value: domain.Axis) !kernel_program.Index1D {
82 if (comptime @hasDecl(Impl, "index1D")) {
83 var ctx = Index1DContext(PointerChild(@TypeOf(owner)), Next){
84 .owner = owner,
85 .next = &self.next,
86 .axis = axis_value,
87 };
88 return self.impl.index1D(&ctx);
89 }
90 return self.next.index1D(owner, axis_value);
91 }
92
93 pub fn index2D(self: *Self, owner: anytype, shape: domain.Domain2D) !kernel_program.Index2D {
94 if (comptime @hasDecl(Impl, "index2D")) {
95 var ctx = Index2DContext(PointerChild(@TypeOf(owner)), Next){
96 .owner = owner,
97 .next = &self.next,
98 .shape = shape,
99 };
100 return self.impl.index2D(&ctx);
101 }
102 return self.next.index2D(owner, shape);
103 }
104
105 pub fn index3D(self: *Self, owner: anytype, shape: domain.Domain3D) !kernel_program.Index3D {
106 if (comptime @hasDecl(Impl, "index3D")) {
107 var ctx = Index3DContext(PointerChild(@TypeOf(owner)), Next){
108 .owner = owner,
109 .next = &self.next,
110 .shape = shape,
111 };
112 return self.impl.index3D(&ctx);
113 }
114 return self.next.index3D(owner, shape);
115 }
116 };
117 }
118
119 pub fn Index1DContext(comptime Owner: type, comptime Next: type) type {
120 return struct {
121 owner: *Owner,
122 next: *Next,
123 axis: domain.Axis,
124
125 pub fn default(self: *@This()) !kernel_program.Index1D {
126 return self.next.index1D(self.owner, self.axis);
127 }
128
129 pub fn useThreads(self: *@This(), threads_per_block: u32) !kernel_program.Index1D {
130 return self.owner.index1D(self.axis.name, self.axis.extent, threads_per_block);
131 }
132 };
133 }
134
135 pub fn Index2DContext(comptime Owner: type, comptime Next: type) type {
136 return struct {
137 owner: *Owner,
138 next: *Next,
139 shape: domain.Domain2D,
140
141 pub fn default(self: *@This()) !kernel_program.Index2D {
142 return self.next.index2D(self.owner, self.shape);
143 }
144
145 pub fn useThreads(self: *@This(), threads: Threads) !kernel_program.Index2D {
146 return self.owner.index2D(.{
147 .x = kernel_program.domainAxis(self.shape.x.name, self.shape.x.extent, threads.x),
148 .y = kernel_program.domainAxis(self.shape.y.name, self.shape.y.extent, threads.y),
149 });
150 }
151 };
152 }
153
154 pub fn Index3DContext(comptime Owner: type, comptime Next: type) type {
155 return struct {
156 owner: *Owner,
157 next: *Next,
158 shape: domain.Domain3D,
159
160 pub fn default(self: *@This()) !kernel_program.Index3D {
161 return self.next.index3D(self.owner, self.shape);
162 }
163
164 pub fn useThreads(self: *@This(), threads: Threads) !kernel_program.Index3D {
165 return self.owner.index3D(.{
166 .x = kernel_program.domainAxis(self.shape.x.name, self.shape.x.extent, threads.x),
167 .y = kernel_program.domainAxis(self.shape.y.name, self.shape.y.extent, threads.y),
168 .z = kernel_program.domainAxis(self.shape.z.name, self.shape.z.extent, threads.z),
169 });
170 }
171 };
172 }
173
174 pub fn threadBlocks(threads: Threads) ThreadBlocks {
175 return .{ .threads = threads };
176 }
177
178 pub fn default() ThreadBlocks {
179 return threadBlocks(.{});
180 }
181
182 fn stackLength(comptime Specs: type) comptime_int {
183 const info = @typeInfo(Specs);
184 if (info != .@"struct" or !info.@"struct".is_tuple) {
185 @compileError("kernel.logical.schedule.stack expects a tuple of schedule layers");
186 }
187 if (info.@"struct".field_names.len == 0) {
188 @compileError("kernel.logical.schedule.stack expects at least one schedule layer");
189 }
190 return info.@"struct".field_names.len;
191 }
192
193 fn StackAttach(comptime index: usize, comptime Specs: type, comptime Next: type) type {
194 const len = stackLength(Specs);
195 if (index == len) return Next;
196 const Spec = @typeInfo(Specs).@"struct".field_types[index];
197 return @TypeOf(@as(Spec, undefined).attach(@as(StackAttach(index + 1, Specs, Next), undefined)));
198 }
199
200 fn attachFrom(comptime index: usize, specs: anytype, next: anytype) StackAttach(index, @TypeOf(specs), @TypeOf(next)) {
201 if (comptime index == stackLength(@TypeOf(specs))) return next;
202 const inner = attachFrom(index + 1, specs, next);
203 return @field(specs, @typeInfo(@TypeOf(specs)).@"struct".field_names[index]).attach(inner);
204 }
205
206 fn PointerChild(comptime Pointer: type) type {
207 return switch (@typeInfo(Pointer)) {
208 .pointer => |info| info.child,
209 else => @compileError("logical schedules attach to builder pointers"),
210 };
211 }