tiny.accy.tensor.unroll

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

1const std = @import("std");2const interpret = @import("interpret/root.zig");3const program_mod = @import("program.zig");4const trace = @import("trace/root.zig");5const types = @import("type/root.zig");67pub fn apply(allocator: std.mem.Allocator, source: *const program_mod.Program) !program_mod.Program {8    var builder = try trace.Builder.init(allocator, source.name);9    errdefer builder.deinit();1011    var scratch = std.heap.ArenaAllocator.init(allocator);12    defer scratch.deinit();1314    const graph = interpret.Graph{ .builder = &builder };15    return try interpret.run(allocator, source, Semantics(@TypeOf(graph)){16        .next = graph,17        .scratch = scratch.allocator(),18    });19}2021fn Semantics(comptime Next: type) type {22    return struct {23        next: Next,24        scratch: std.mem.Allocator,25        finals: std.AutoHashMapUnmanaged(u32, []trace.Value) = .empty,2627        pub const Value: type = trace.Value;28        pub const Result: type = Next.Result;2930        pub fn operation(self: *@This(), step: *interpret.Step(Value)) !Value {31            switch (step.op.kind) {32                .scan => |scan| {33                    var buffer: [program_mod.max_scan_carries]trace.Value = undefined;34                    for (scan.inits, 0..) |init_id, index| buffer[index] = step.read(init_id);35                    const expanded = try self.expandScan(scan, buffer[0..scan.inits.len]);36                    try self.finals.put(self.scratch, step.op.id.index, expanded);37                    return expanded[0];38                },39                .projection => |projection| {40                    return self.finals.get(projection.source.index).?[projection.index];41                },42                else => {43                    var buffer: [program_mod.max_operation_operands]trace.Value = undefined;44                    return self.next.bind(step.op, interpret.arguments(Value, step.op, step.values, &buffer));45                },46            }47        }4849        fn expandScan(self: *@This(), scan: program_mod.Scan, inits: []const trace.Value) anyerror![]trace.Value {50            const carries = try self.scratch.dupe(trace.Value, inits);51            var iteration: i64 = 0;52            while (iteration < scan.length) : (iteration += 1) {53                try self.expandStep(scan.body, carries);54            }55            return carries;56        }5758        fn expandStep(self: *@This(), body: *const program_mod.Subgraph, carries: []trace.Value) !void {59            const values = try self.scratch.alloc(trace.Value, body.values.len);60            var nested_finals: std.AutoHashMapUnmanaged(u32, []trace.Value) = .empty;61            for (body.operations) |*op| {62                values[op.id.index] = switch (op.kind) {63                    .parameter => |parameter| carries[parameter.index],64                    .scan => |nested| blk: {65                        var buffer: [program_mod.max_scan_carries]trace.Value = undefined;66                        for (nested.inits, 0..) |init_id, index| buffer[index] = values[init_id.index];67                        const expanded = try self.expandScan(nested, buffer[0..nested.inits.len]);68                        try nested_finals.put(self.scratch, op.id.index, expanded);69                        break :blk expanded[0];70                    },71                    .projection => |projection| nested_finals.get(projection.source.index).?[projection.index],72                    else => blk: {73                        var buffer: [program_mod.max_operation_operands]trace.Value = undefined;74                        const args = interpret.arguments(Value, op, values, &buffer);75                        break :blk try self.next.bind(op, args);76                    },77                };78            }79            var buffer: [program_mod.max_scan_carries]trace.Value = undefined;80            for (body.outputs, 0..) |id, index| buffer[index] = values[id.index];81            @memcpy(carries, buffer[0..body.outputs.len]);82        }8384        pub fn finish(self: *@This(), outputs: []const Value) !Result {85            return self.next.finish(outputs);86        }8788        pub fn builderHandle(self: *@This()) *trace.Builder {89            return self.next.builderHandle();90        }91    };92}9394fn incrementScanStep(scan_builder: *trace.Builder, carry: trace.Value) !trace.Value {95    const bump = try scan_builder.full(.f32, .{ .lane = 2 }, 1.0);96    return carry.add(bump);97}9899test "unroll expands scan operations into scan free graphs" {100    var builder = try trace.Builder.init(std.testing.allocator, "unroll_expand");101    defer builder.deinit();102103    const x0 = try builder.input(.f32, .{ .lane = 2 });104    const walked = try builder.scan(.{105        .length = 3,106        .init = x0,107        .body = incrementScanStep,108    });109    var source = try builder.finish(&.{walked});110    defer source.deinit();111112    var expanded = try apply(std.testing.allocator, &source);113    defer expanded.deinit();114115    for (expanded.operations) |op| {116        try std.testing.expect(op.kind != .scan);117        try std.testing.expect(op.kind != .projection);118    }119    try std.testing.expectEqual(@as(usize, 7), expanded.operationCount());120    try types.expectExtents(&.{2}, expanded.typeOf(expanded.outputs[0]));121}122123test "unroll expands across carry arities and lengths including zero" {124    inline for (1..4) |arity| {125        inline for (0..4) |length| {126            var builder = try trace.Builder.init(std.testing.allocator, "unroll_grid");127            defer builder.deinit();128129            var inits: [arity]trace.Value = undefined;130            for (0..arity) |index| {131                inits[index] = try builder.full(.f32, .{ .lane = 2 }, @floatFromInt(index + 1));132            }133            var scope = try builder.scanScope(@intCast(length), inits[0..]);134            var nexts: [arity]trace.Value = undefined;135            for (0..arity) |index| {136                const partner = scope.carry((index + 1) % arity);137                nexts[index] = try scope.carry(index).add(partner);138            }139            const walked = try scope.finish(nexts[0..]);140            var outputs_buffer: [arity]trace.Value = undefined;141            outputs_buffer[0] = walked;142            for (1..arity) |index| {143                outputs_buffer[index] = try builder.projection(walked, index);144            }145            var source = try builder.finish(outputs_buffer[0..]);146            defer source.deinit();147148            var expanded = try apply(std.testing.allocator, &source);149            defer expanded.deinit();150151            for (expanded.operations) |op| {152                try std.testing.expect(op.kind != .scan);153                try std.testing.expect(op.kind != .projection);154            }155            try std.testing.expectEqual(arity + length * arity, expanded.operationCount());156            for (expanded.outputs) |id| {157                try types.expectExtents(&.{2}, expanded.typeOf(id));158            }159        }160    }161}162163fn projectionScanStep(_: *trace.Builder, carry: anytype) !@TypeOf(carry) {164    const doubled = try carry.x.add(carry.x);165    return .{166        .x = doubled,167        .acc = try carry.acc.add(doubled),168    };169}170171test "unroll resolves projections to expanded carries" {172    var builder = try trace.Builder.init(std.testing.allocator, "unroll_projection");173    defer builder.deinit();174175    const x0 = try builder.input(.f32, .{ .lane = 2 });176    const acc0 = try builder.full(.f32, .{ .lane = 2 }, 0.0);177    const walked = try builder.scan(.{178        .length = 2,179        .init = .{ .x = x0, .acc = acc0 },180        .body = projectionScanStep,181    });182    var source = try builder.finish(&.{walked.acc});183    defer source.deinit();184185    var expanded = try apply(std.testing.allocator, &source);186    defer expanded.deinit();187188    for (expanded.operations) |op| {189        try std.testing.expect(op.kind != .scan);190        try std.testing.expect(op.kind != .projection);191    }192    try types.expectExtents(&.{2}, expanded.typeOf(expanded.outputs[0]));193}