tiny.accy.kernel.oracle

1pub const KernelArgument = union(enum) {2    memref: []u8,3    scalar: Scalar,4};

1pub const Machine = struct {2    allocator: std.mem.Allocator,3    evaluator: Evaluator,4    bound: std.ArrayListUnmanaged(RuntimeValue) = .empty,5    lanes: std.ArrayListUnmanaged(LaneState) = .empty,6    lane_refs: std.ArrayListUnmanaged(*LaneState) = .empty,78    pub fn init(allocator: std.mem.Allocator) Machine {9        return .{10            .allocator = allocator,11            .evaluator = Evaluator.init(allocator),12        };13    }1415    pub fn deinit(self: *Machine) void {16        for (self.lanes.items) |*lane| lane.deinit();17        self.lanes.deinit(self.allocator);18        self.lane_refs.deinit(self.allocator);19        self.bound.deinit(self.allocator);20        self.evaluator.deinit();21        self.* = undefined;22    }2324    pub fn run(25        self: *Machine,26        func_op: *ir.Operation,27        args: []const KernelArgument,28        launch: anytype,29    ) EvalError!void {30        return self.runWithDiagnostic(func_op, args, launch, null);31    }3233    pub fn runWithDiagnostic(34        self: *Machine,35        func_op: *ir.Operation,36        args: []const KernelArgument,37        launch: anytype,38        diagnostic: ?*Diagnostic,39    ) EvalError!void {40        if (diagnostic) |captured| captured.len = 0;41        return self.runChecking(func_op, args, launch, diagnostic) catch |err| {42            captureFunctionDiagnostic(diagnostic, "accy/kernel/run", func_op, err);43            return err;44        };45    }4647    fn runChecking(48        self: *Machine,49        func_op: *ir.Operation,50        args: []const KernelArgument,51        launch: anytype,52        diagnostic: ?*Diagnostic,53    ) EvalError!void {54        if (!std.mem.eql(u8, func_op.name.name, dialects.FuncDialect.FuncOp.operation_name)) {55            return error.InvalidFunction;56        }57        const region = func_op.getRegion(0) orelse return error.InvalidFunction;58        const entry = region.getEntryBlock() orelse return error.InvalidFunction;59        if (entry.arguments.items.len != args.len) return error.InputCountMismatch;6061        self.evaluator.resetRun();62        self.evaluator.setDiagnostic("accy/kernel/run", diagnostic);6364        try self.bound.resize(self.allocator, args.len);65        for (args, 0..) |arg, i| {66            const block_arg = entry.getArgument(i) orelse return error.InvalidFunction;67            self.bound.items[i] = switch (arg) {68                .memref => |bytes| .{ .memref = try self.evaluator.addBorrowedMemref(block_arg.type, bytes) },69                .scalar => |value| .{ .scalar = try castScalar(try scalarKindOfType(block_arg.type), value) },70            };71        }7273        const lane_count = try launchBlockThreadCount(launch);74        try self.ensureLanes(lane_count);75        const lane_refs = self.lane_refs.items[0..lane_count];7677        for (0..@as(usize, @intCast(launch.grid[2]))) |block_z| {78            for (0..@as(usize, @intCast(launch.grid[1]))) |block_y| {79                for (0..@as(usize, @intCast(launch.grid[0]))) |block_x| {80                    self.evaluator.refillSharedSentinel();81                    var lane_index: usize = 0;82                    for (0..@as(usize, @intCast(launch.block[2]))) |thread_z| {83                        for (0..@as(usize, @intCast(launch.block[1]))) |thread_y| {84                            for (0..@as(usize, @intCast(launch.block[0]))) |thread_x| {85                                const lane = lane_refs[lane_index];86                                lane_index += 1;87                                lane.reset(launchIndex(launch, block_x, block_y, block_z, thread_x, thread_y, thread_z));88                                for (self.bound.items, 0..) |value, i| {89                                    const block_arg = entry.getArgument(i) orelse return error.InvalidFunction;90                                    try lane.setValue(block_arg, value);91                                }92                            }93                        }94                    }95                    std.debug.assert(lane_index == lane_count);96                    try self.evaluator.evalKernelBlock(entry, lane_refs);97                }98            }99        }100    }101102    fn ensureLanes(self: *Machine, lane_count: usize) EvalError!void {103        while (self.lanes.items.len < lane_count) {104            try self.lanes.append(self.allocator, LaneState.init(self.allocator, .{}));105        }106        try self.lane_refs.resize(self.allocator, self.lanes.items.len);107        for (self.lanes.items, 0..) |*lane, index| {108            self.lane_refs.items[index] = lane;109        }110        std.debug.assert(self.lane_refs.items.len >= lane_count);111    }112};

1pub const Scalar = union(ScalarKind) {2    bool: bool,3    index: i64,4    i8: i8,5    i16: i16,6    i32: i32,7    u8: u8,8    u16: u16,9    u32: u32,10    i64: i64,11    u64: u64,12    f16: f16,13    bf16: Bf16,14    f32: f32,15    f64: f64,16};

1pub const Launch = struct {2    grid: [3]u32 = .{ 1, 1, 1 },3    block: [3]u32 = .{ 1, 1, 1 },45    pub fn oneDim(extent: usize) Launch {6        return .{ .grid = .{ @intCast(extent), 1, 1 } };7    }8};

1pub const shared_alloc_sentinel: u8 = 0x7F;

1pub fn tf32RoundF32(value: f32) f32 {2    const bits: u32 = @bitCast(value);3    if (bits & 0x7F80_0000 == 0x7F80_0000) return value;4    const rounded = (bits +% 0x1000) & 0xFFFF_E000;5    return @bitCast(rounded);6}

1pub fn argumentBool(value: bool) Argument {2    return .{ .scalar = .{ .bool = value } };3}

1pub fn argumentBuffer(comptime T: type, values: []T) Argument {2    return .{ .memref = std.mem.sliceAsBytes(values) };3}

1pub fn argumentF32(value: f32) Argument {2    return .{ .scalar = .{ .f32 = value } };3}

1pub fn argumentF64(value: f64) Argument {2    return .{ .scalar = .{ .f64 = value } };3}

1pub fn argumentI32(value: i32) Argument {2    return .{ .scalar = .{ .i32 = value } };3}

1pub fn argumentI64(value: i64) Argument {2    return .{ .scalar = .{ .i64 = value } };3}

1pub fn argumentU32(value: u32) Argument {2    return .{ .scalar = .{ .u32 = value } };3}

1pub fn run(allocator: std.mem.Allocator, kernel: *builder.Kernel, args: []const Argument, launch: Launch) !void {2    try kernel.verify();3    try eval.executeKernelFunction(allocator, kernel.func().op, args, launch);4}

1pub fn run1D(allocator: std.mem.Allocator, kernel: *builder.Kernel, args: []const Argument, extent: usize) !void {2    try run(allocator, kernel, args, Launch.oneDim(extent));3}

1pub fn run1DWithDiagnostic(2    allocator: std.mem.Allocator,3    kernel: *builder.Kernel,4    args: []const Argument,5    extent: usize,6    diagnostic: *ExecutionDiagnostic,7) !void {8    try runWithDiagnostic(allocator, kernel, args, Launch.oneDim(extent), diagnostic);9}

1pub fn runWithDiagnostic(2    allocator: std.mem.Allocator,3    kernel: *builder.Kernel,4    args: []const Argument,5    launch: Launch,6    diagnostic: *ExecutionDiagnostic,7) !void {8    diagnostic.len = 0;9    try kernel.verifyWithDiagnostic(diagnostic);10    try eval.executeKernelFunctionWithDiagnostic(allocator, kernel.func().op, args, launch, diagnostic);11}

1pub const Diagnostic = choir.backends.contract.VerificationDiagnostic;

1pub const enabled = true;

1const machine = @import("machine.zig");23pub const enabled = machine.enabled;4pub const Launch = machine.Launch;5pub const Argument = machine.Argument;6pub const Machine = machine.Machine;7pub const Scalar = machine.Scalar;8pub const ExecutionDiagnostic = machine.ExecutionDiagnostic;9pub const shared_alloc_sentinel = machine.shared_alloc_sentinel;10pub const tf32RoundF32 = machine.tf32RoundF32;11pub const argumentBuffer = machine.argumentBuffer;12pub const argumentBool = machine.argumentBool;13pub const argumentI32 = machine.argumentI32;14pub const argumentU32 = machine.argumentU32;15pub const argumentI64 = machine.argumentI64;16pub const argumentF32 = machine.argumentF32;17pub const argumentF64 = machine.argumentF64;18pub const run = machine.run;19pub const runWithDiagnostic = machine.runWithDiagnostic;20pub const run1D = machine.run1D;21pub const run1DWithDiagnostic = machine.run1DWithDiagnostic;

1pub const oracle = @import("oracle/root.zig");