tiny.accy.kernel.program.execution

1const std = @import("std");2const choir = @import("choir");34const core = @import("../model/core/root.zig");5const model = @import("../model/program/root.zig");6const plan = @import("../model/plan/root.zig");7const oracle = @import("../oracle/root.zig");89const Cpu = struct {10    machine: ?oracle.Machine = null,11    verified: bool = false,1213    fn deinit(self: *Cpu) void {14        if (self.machine) |*machine| machine.deinit();15        self.* = undefined;16    }1718    fn run(19        self: *Cpu,20        program: *model.Program,21        allocator: std.mem.Allocator,22        args: []const oracle.Argument,23        launch_value: core.schedule.Launch,24    ) !void {25        const machine = try self.ensureMachine(allocator);26        if (!self.verified) {27            try program.verify();28            self.verified = true;29        }30        try machine.run(program.storage.kernel.func().op, args, launch_value);31    }3233    fn runWithDiagnostic(34        self: *Cpu,35        program: *model.Program,36        allocator: std.mem.Allocator,37        args: []const oracle.Argument,38        diagnostic: *oracle.ExecutionDiagnostic,39    ) !void {40        diagnostic.len = 0;41        const machine = try self.ensureMachine(allocator);42        if (!self.verified) {43            try program.verifyWithDiagnostic(diagnostic);44            self.verified = true;45        }46        try machine.runWithDiagnostic(program.storage.kernel.func().op, args, try program.launch(), diagnostic);47    }4849    fn ensureMachine(self: *Cpu, allocator: std.mem.Allocator) !*oracle.Machine {50        if (self.machine) |*machine| {51            const same_ptr = machine.allocator.ptr == allocator.ptr;52            const same_vtable = machine.allocator.vtable == allocator.vtable;53            if (same_ptr and same_vtable) return machine;54            machine.deinit();55            self.machine = null;56            self.verified = false;57        }58        self.machine = oracle.Machine.init(allocator);59        return &self.machine.?;60    }61};6263pub const Executor = struct {64    program: *model.Program,65    cpu: Cpu = .{},6667    pub fn init(program: *model.Program) Executor {68        return .{ .program = program };69    }7071    pub fn deinit(self: *Executor) void {72        self.cpu.deinit();73        self.* = undefined;74    }7576    pub fn runCpu(self: *Executor, allocator: std.mem.Allocator, args: []const oracle.Argument) !void {77        try self.runCpuWithLaunch(allocator, args, try self.program.launch());78    }7980    pub fn runCpuWithLaunch(81        self: *Executor,82        allocator: std.mem.Allocator,83        args: []const oracle.Argument,84        launch_value: core.schedule.Launch,85    ) !void {86        try self.cpu.run(self.program, allocator, args, launch_value);87    }8889    pub fn runCpuWithDiagnostic(90        self: *Executor,91        allocator: std.mem.Allocator,92        args: []const oracle.Argument,93        diagnostic: *oracle.ExecutionDiagnostic,94    ) !void {95        try self.cpu.runWithDiagnostic(self.program, allocator, args, diagnostic);96    }97};9899pub const Program = struct {100    capacity: model.Capacity,101    model_program: model.Program,102    cpu: Cpu = .{},103104    const Self = @This();105106    pub fn init(107        kernel: core.builder.Kernel,108        schedule: core.schedule.Schedule,109        capacity: model.Capacity,110    ) Self {111        return .{112            .capacity = capacity,113            .model_program = model.Program.init(kernel, schedule, capacity),114        };115    }116117    pub fn deinit(self: *Self) void {118        self.cpu.deinit();119        self.model_program.deinit();120        self.* = undefined;121    }122123    pub fn verify(self: *Self) !void {124        try self.model_program.verify();125    }126127    pub fn verifyWithDiagnostic(self: *Self, diagnostic: *core.builder.VerificationDiagnostic) !void {128        try self.model_program.verifyWithDiagnostic(diagnostic);129    }130131    pub fn launch(self: *const Self) core.schedule.ScheduleError!core.schedule.Launch {132        return self.model_program.launch();133    }134135    pub fn scheduleSnapshot(self: *const Self, allocator: std.mem.Allocator) core.schedule.ScheduleError!core.schedule.Snapshot {136        return self.model_program.scheduleSnapshot(allocator);137    }138139    pub fn createPlan(self: *Self, allocator: std.mem.Allocator, options: plan.Options) plan.Error!plan.Plan {140        return self.model_program.createPlan(allocator, options);141    }142143    pub fn kernelModule(self: *const Self) *choir.ir.Operation {144        return self.model_program.kernelModule();145    }146147    pub fn params(self: *const Self) []const core.builder.Param {148        return self.model_program.params();149    }150151    pub fn createCheckedPlan(self: *Self, allocator: std.mem.Allocator, options: plan.Options) !plan.Plan {152        return self.model_program.createCheckedPlan(allocator, options);153    }154155    pub fn checkPlan(self: *Self, allocator: std.mem.Allocator, authored_plan: *const plan.Plan) !void {156        try self.model_program.checkPlan(allocator, authored_plan);157    }158159    pub fn runCpu(self: *Self, allocator: std.mem.Allocator, args: []const oracle.Argument) !void {160        try self.runCpuWithLaunch(allocator, args, try self.launch());161    }162163    pub fn runCpuWithLaunch(164        self: *Self,165        allocator: std.mem.Allocator,166        args: []const oracle.Argument,167        launch_value: core.schedule.Launch,168    ) !void {169        try self.cpu.run(&self.model_program, allocator, args, launch_value);170    }171172    pub fn runCpuWithDiagnostic(173        self: *Self,174        allocator: std.mem.Allocator,175        args: []const oracle.Argument,176        diagnostic: *oracle.ExecutionDiagnostic,177    ) !void {178        try self.cpu.runWithDiagnostic(&self.model_program, allocator, args, diagnostic);179    }180181    pub fn bodyFingerprint(self: *const Self, allocator: std.mem.Allocator) !u64 {182        return self.model_program.bodyFingerprint(allocator);183    }184185    pub fn scheduleFingerprint(self: *const Self) u64 {186        return self.model_program.scheduleFingerprint();187    }188189    pub fn fingerprint(self: *const Self, allocator: std.mem.Allocator) !choir.product.incremental.Fingerprint {190        return self.model_program.fingerprint(allocator);191    }192193    pub fn productStamp(self: *const Self, allocator: std.mem.Allocator) !choir.product.incremental.ProductStamp {194        return self.model_program.productStamp(allocator);195    }196};197198const surface = model.Surface(Program, .{199    .registrations = &choir.backends.gpu.dialect_registrations,200    .preload = &.{"scf"},201});202203pub const Index1D = surface.Index1D;204pub const VectorIndex1D = surface.VectorIndex1D;205pub const DomainAxis = surface.DomainAxis;206pub const Domain2D = surface.Domain2D;207pub const Domain3D = surface.Domain3D;208pub const Index2D = surface.Index2D;209pub const Index3D = surface.Index3D;210pub const Vec2 = surface.Vec2;211pub const Vec3 = surface.Vec3;212pub const TypedValue = surface.TypedValue;213pub const TypedVec2 = surface.TypedVec2;214pub const TypedVec3 = surface.TypedVec3;215pub const domainAxis = surface.domainAxis;216pub const BufferView = surface.BufferView;217pub const Builder = surface.Builder;218pub const Limits = surface.Limits;219pub const Capacity = surface.Capacity;220pub const Guard = surface.Guard;221pub const define = surface.define;222pub const defineWithDiagnostic = surface.defineWithDiagnostic;223224fn buildCopyProgram(225    allocator: std.mem.Allocator,226    name: []const u8,227    extent: u64,228) !Program {229    var builder = try Builder.init(allocator, Limits.testing, name, &.{230        core.builder.dynamicBuffer(.i32),231        core.builder.dynamicBuffer(.i32),232    });233    errdefer builder.deinit();234235    const axis = try builder.axis("i", extent);236    try builder.bind(axis, .thread_x);237    const src = builder.argument(0);238    const dst = builder.argument(1);239    const index = try builder.globalId(.x);240    const value = try builder.load(src, index);241    try builder.store(value, dst, index);242    try builder.return_();243    return builder.finish();244}245246test "Program steady oracle runs make no allocator calls" {247    var program = try buildCopyProgram(std.testing.allocator, "machine_retained_copy", 4);248    defer program.deinit();249250    var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{});251    var input = [_]i32{ 3, -1, 7, 12 };252    var output = [_]i32{ 0, 0, 0, 0 };253    try program.runCpu(failing.allocator(), &.{254        oracle.argumentBuffer(i32, input[0..]),255        oracle.argumentBuffer(i32, output[0..]),256    });257    try std.testing.expectEqualSlices(i32, input[0..], output[0..]);258259    failing.fail_index = failing.alloc_index;260    failing.resize_fail_index = failing.resize_index;261    for (0..8) |_| {262        @memset(output[0..], 0);263        try program.runCpu(failing.allocator(), &.{264            oracle.argumentBuffer(i32, input[0..]),265            oracle.argumentBuffer(i32, output[0..]),266        });267        try std.testing.expectEqualSlices(i32, input[0..], output[0..]);268    }269    try std.testing.expect(!failing.has_induced_failure);270}271272test "Program reacquires the machine when the run allocator changes" {273    var program = try buildCopyProgram(std.testing.allocator, "machine_epoch_copy", 4);274    defer program.deinit();275276    var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{});277    var input = [_]i32{ 5, 6, 7, 8 };278    var output = [_]i32{ 0, 0, 0, 0 };279    try program.runCpu(failing.allocator(), &.{280        oracle.argumentBuffer(i32, input[0..]),281        oracle.argumentBuffer(i32, output[0..]),282    });283    try std.testing.expectEqualSlices(i32, input[0..], output[0..]);284285    @memset(output[0..], 0);286    try program.runCpu(std.testing.allocator, &.{287        oracle.argumentBuffer(i32, input[0..]),288        oracle.argumentBuffer(i32, output[0..]),289    });290    try std.testing.expectEqualSlices(i32, input[0..], output[0..]);291}

1pub const execution = @import("execution.zig");