tiny.accy.kernel.logical.schedule

1pub const schedule = model.schedule;

1const kernel_program = @import("../../program/root.zig");2const domain = @import("../domain/root.zig");34pub const Threads = struct {5    x: u32 = 256,6    y: u32 = 4,7    z: u32 = 1,8};910pub const ThreadBlocks = struct {11    threads: Threads,1213    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    }1617    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    }2324    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};3233pub const UseThreads = struct {34    threads: Threads,3536    pub fn index1D(self: *@This(), ctx: anytype) !kernel_program.Index1D {37        return ctx.useThreads(self.threads.x);38    }3940    pub fn index2D(self: *@This(), ctx: anytype) !kernel_program.Index2D {41        return ctx.useThreads(self.threads);42    }4344    pub fn index3D(self: *@This(), ctx: anytype) !kernel_program.Index3D {45        return ctx.useThreads(self.threads);46    }47};4849pub fn with(impl: anytype) With(@TypeOf(impl)) {50    return .{ .impl = impl };51}5253pub fn useThreads(threads: Threads) With(UseThreads) {54    return with(UseThreads{ .threads = threads });55}5657pub fn With(comptime Impl: type) type {58    return struct {59        impl: Impl,6061        pub fn attach(self: @This(), next: anytype) Layer(@TypeOf(next), Impl) {62            return .{63                .next = next,64                .impl = self.impl,65            };66        }67    };68}6970pub fn stack(specs: anytype) StackAttach(0, @TypeOf(specs), ThreadBlocks) {71    return attachFrom(0, specs, default());72}7374pub fn Layer(comptime Next: type, comptime Impl: type) type {75    return struct {76        next: Next,77        impl: Impl,7879        const Self = @This();8081        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        }9293        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        }104105        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}118119pub fn Index1DContext(comptime Owner: type, comptime Next: type) type {120    return struct {121        owner: *Owner,122        next: *Next,123        axis: domain.Axis,124125        pub fn default(self: *@This()) !kernel_program.Index1D {126            return self.next.index1D(self.owner, self.axis);127        }128129        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}134135pub fn Index2DContext(comptime Owner: type, comptime Next: type) type {136    return struct {137        owner: *Owner,138        next: *Next,139        shape: domain.Domain2D,140141        pub fn default(self: *@This()) !kernel_program.Index2D {142            return self.next.index2D(self.owner, self.shape);143        }144145        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}153154pub fn Index3DContext(comptime Owner: type, comptime Next: type) type {155    return struct {156        owner: *Owner,157        next: *Next,158        shape: domain.Domain3D,159160        pub fn default(self: *@This()) !kernel_program.Index3D {161            return self.next.index3D(self.owner, self.shape);162        }163164        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}173174pub fn threadBlocks(threads: Threads) ThreadBlocks {175    return .{ .threads = threads };176}177178pub fn default() ThreadBlocks {179    return threadBlocks(.{});180}181182fn 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}192193fn 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}199200fn 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}205206fn 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}