tiny.accy.tensor.transform

1pub const transform = @import("transform.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 const Operand = struct {8    id: program_mod.Id,9    value: trace.Value,10    ty: trace.Type,11};1213pub const Context = struct {14    source: *const program_mod.Program,15    builder: *trace.Builder,16    op: *const program_mod.Operation,17    args: []const trace.Value,1819    pub fn arg(self: Context, index: usize) trace.Value {20        return self.args[index];21    }2223    pub fn operand(self: Context, index: usize) Operand {24        const id = self.operandId(index).?;25        const value = self.arg(index);26        return .{27            .id = id,28            .value = value,29            .ty = value.ty,30        };31    }3233    pub fn isZero(self: Context, index: usize) bool {34        const id = self.operandId(index) orelse return false;35        if (id.index >= self.source.valueCount()) return false;36        return self.source.isZeroConstant(id);37    }3839    pub fn constantPayload(self: Context, index: usize) ?[]const u8 {40        const id = self.operandId(index) orelse return null;41        if (id.index >= self.source.valueCount()) return null;42        return self.source.constantPayload(id);43    }4445    fn operandId(self: Context, index: usize) ?program_mod.Id {46        return switch (self.op.kind) {47            .parameter, .constant, .iota => null,48            .unary => |unary| if (index == 0) unary.input else null,49            .binary => |binary| switch (index) {50                0 => binary.lhs,51                1 => binary.rhs,52                else => null,53            },54            .broadcast => |broadcast| if (index == 0) broadcast.input else null,55            .broadcast_in_dim => |broadcast| if (index == 0) broadcast.input else null,56            .reshape => |reshape| if (index == 0) reshape.input else null,57            .transpose => |transpose| if (index == 0) transpose.input else null,58            .reduce => |reduce| switch (index) {59                0 => reduce.input,60                1 => reduce.init,61                else => null,62            },63            .gather => |gather| switch (index) {64                0 => gather.input,65                1 => gather.indices,66                else => null,67            },68            .scatter_add => |scatter_add| switch (index) {69                0 => scatter_add.input,70                1 => scatter_add.indices,71                2 => scatter_add.updates,72                else => null,73            },74            .sparse_cross_entropy => |sparse_cross_entropy| switch (index) {75                0 => sparse_cross_entropy.logits,76                1 => sparse_cross_entropy.targets,77                else => null,78            },79            .dot_general => |dot| switch (index) {80                0 => dot.lhs,81                1 => dot.rhs,82                else => null,83            },84            .compare => |compare| switch (index) {85                0 => compare.lhs,86                1 => compare.rhs,87                else => null,88            },89            .select => |select| switch (index) {90                0 => select.pred,91                1 => select.on_true,92                2 => select.on_false,93                else => null,94            },95            .custom_call => |custom| if (index < custom.operands.len) custom.operands[index] else null,96            .scan => |scan| if (index < scan.inits.len) scan.inits[index] else null,97            .projection => |projection| if (index == 0) projection.source else null,98        };99    }100};101102pub fn apply(allocator: std.mem.Allocator, source: *const program_mod.Program, pass: anytype) !program_mod.Program {103    var builder = try trace.Builder.init(allocator, source.name);104    errdefer builder.deinit();105106    const graph = interpret.Graph{ .builder = &builder };107    return try interpret.run(allocator, source, semantics(source, graph, pass));108}109110pub fn semantics(source: *const program_mod.Program, next: anytype, pass: anytype) interpret.Layer(trace.Value, @TypeOf(next), Rewrite(@TypeOf(pass))) {111    return interpret.layer(trace.Value, next, Rewrite(@TypeOf(pass)){112        .source = source,113        .pass = pass,114    });115}116117pub fn Rewrite(comptime Pass: type) type {118    return struct {119        source: *const program_mod.Program,120        pass: Pass,121122        pub fn bind(self: *@This(), layer: anytype) !trace.Value {123            var ctx = Context{124                .source = self.source,125                .builder = layer.builderHandle(),126                .op = layer.op,127                .args = layer.args,128            };129            return (try dispatch(&self.pass, &ctx, layer.op)) orelse layer.default();130        }131    };132}133134fn dispatch(pass: anytype, ctx: *Context, op: *const program_mod.Operation) !?trace.Value {135    const Pass = @TypeOf(pass.*);136    return switch (op.kind) {137        .binary => |binary| switch (binary.op) {138            .add => if (comptime @hasDecl(Pass, "add"))139                pass.add(ctx)140            else141                null,142            .sub => if (comptime @hasDecl(Pass, "sub"))143                pass.sub(ctx)144            else145                null,146            .mul => if (comptime @hasDecl(Pass, "mul"))147                pass.mul(ctx)148            else149                null,150            .div => if (comptime @hasDecl(Pass, "div"))151                pass.div(ctx)152            else153                null,154            .max => if (comptime @hasDecl(Pass, "max"))155                pass.max(ctx)156            else157                null,158            .min => if (comptime @hasDecl(Pass, "min"))159                pass.min(ctx)160            else161                null,162            .pow => if (comptime @hasDecl(Pass, "pow"))163                pass.pow(ctx)164            else165                null,166        },167        .unary => |unary| switch (unary.op) {168            .neg => if (comptime @hasDecl(Pass, "neg"))169                pass.neg(ctx)170            else171                null,172            .abs => if (comptime @hasDecl(Pass, "abs"))173                pass.abs(ctx)174            else175                null,176            .exp => if (comptime @hasDecl(Pass, "exp"))177                pass.exp(ctx)178            else179                null,180            .log => if (comptime @hasDecl(Pass, "log"))181                pass.log(ctx)182            else183                null,184            .sqrt => if (comptime @hasDecl(Pass, "sqrt"))185                pass.sqrt(ctx)186            else187                null,188            .tanh => if (comptime @hasDecl(Pass, "tanh"))189                pass.tanh(ctx)190            else191                null,192            .sin => if (comptime @hasDecl(Pass, "sin"))193                pass.sin(ctx)194            else195                null,196            .cos => if (comptime @hasDecl(Pass, "cos"))197                pass.cos(ctx)198            else199                null,200            .tan => if (comptime @hasDecl(Pass, "tan"))201                pass.tan(ctx)202            else203                null,204        },205        .reduce => if (comptime @hasDecl(Pass, "reduce"))206            pass.reduce(ctx)207        else208            null,209        .gather => if (comptime @hasDecl(Pass, "gather"))210            pass.gather(ctx)211        else212            null,213        .scatter_add => if (comptime @hasDecl(Pass, "scatterAdd"))214            pass.scatterAdd(ctx)215        else216            null,217        .sparse_cross_entropy => if (comptime @hasDecl(Pass, "sparseCrossEntropy"))218            pass.sparseCrossEntropy(ctx)219        else220            null,221        .dot_general => if (comptime @hasDecl(Pass, "dotGeneral"))222            pass.dotGeneral(ctx)223        else224            null,225        .compare => if (comptime @hasDecl(Pass, "compare"))226            pass.compare(ctx)227        else228            null,229        .select => if (comptime @hasDecl(Pass, "select"))230            pass.select(ctx)231        else232            null,233        .custom_call => if (comptime @hasDecl(Pass, "customCall"))234            pass.customCall(ctx)235        else236            null,237        .broadcast => if (comptime @hasDecl(Pass, "broadcast"))238            pass.broadcast(ctx)239        else240            null,241        .broadcast_in_dim => if (comptime @hasDecl(Pass, "broadcastInDim"))242            pass.broadcastInDim(ctx)243        else244            null,245        .reshape => if (comptime @hasDecl(Pass, "reshape"))246            pass.reshape(ctx)247        else248            null,249        .transpose => if (comptime @hasDecl(Pass, "transpose"))250            pass.transpose(ctx)251        else252            null,253        .iota => if (comptime @hasDecl(Pass, "iota"))254            pass.iota(ctx)255        else256            null,257        .parameter => if (comptime @hasDecl(Pass, "parameter"))258            pass.parameter(ctx)259        else260            null,261        .scan => if (comptime @hasDecl(Pass, "scan"))262            pass.scan(ctx)263        else264            null,265        .projection => if (comptime @hasDecl(Pass, "projection"))266            pass.projection(ctx)267        else268            null,269        .constant => null,270    };271}272273test "tensor transform defaults to structural copy" {274    var builder = try trace.Builder.init(std.testing.allocator, "copy");275    defer builder.deinit();276277    const x = try builder.input(.f32, .{ .lane = 4 });278    const y = try builder.input(.f32, .{ .lane = 4 });279    const out = try (try x.add(y)).tanh();280    var source = try builder.finish(&.{out});281    defer source.deinit();282283    var rewritten = try apply(std.testing.allocator, &source, struct {}{});284    defer rewritten.deinit();285286    try std.testing.expectEqual(source.valueCount(), rewritten.valueCount());287    try std.testing.expectEqual(source.operationCount(), rewritten.operationCount());288    try std.testing.expect(types.sameDims(source.typeOf(source.outputs[0]).dims, rewritten.typeOf(rewritten.outputs[0]).dims));289}290291const DropAddZero = struct {292    pub fn add(_: *@This(), ctx: *Context) !?trace.Value {293        const lhs = ctx.operand(0);294        const rhs = ctx.operand(1);295        if (ctx.isZero(1)) return lhs.value;296        if (ctx.isZero(0)) return rhs.value;297        return null;298    }299};300301test "tensor transform lets users rewrite one primitive" {302    var builder = try trace.Builder.init(std.testing.allocator, "drop_add_zero");303    defer builder.deinit();304305    const x = try builder.input(.f32, .{ .lane = 4 });306    const zero = try builder.full(.f32, .{ .lane = 4 }, 0.0);307    const out = try x.add(zero);308    var source = try builder.finish(&.{out});309    defer source.deinit();310311    var rewritten = try apply(std.testing.allocator, &source, DropAddZero{});312    defer rewritten.deinit();313314    try std.testing.expectEqual(@as(usize, 2), rewritten.valueCount());315    try std.testing.expectEqual(@as(u32, 0), rewritten.outputs[0].index);316}317318const BindSecondParameter = struct {319    pub fn parameter(_: *@This(), ctx: *Context) !?trace.Value {320        const parameter_info = ctx.op.kind.parameter;321        if (parameter_info.index == 1) return try ctx.builder.fullFloat(ctx.op.result, 1.0);322        return null;323    }324};325326test "tensor transform lets users bind parameters" {327    var builder = try trace.Builder.init(std.testing.allocator, "bind_parameter");328    defer builder.deinit();329330    const x = try builder.input(.f32, .{ .lane = 4 });331    const scale = try builder.input(.f32, .{ .lane = 4 });332    const out = try x.mul(scale);333    var source = try builder.finish(&.{out});334    defer source.deinit();335336    var rewritten = try apply(std.testing.allocator, &source, BindSecondParameter{});337    defer rewritten.deinit();338339    try std.testing.expectEqual(@as(usize, 1), rewritten.parameters.len);340    try types.expectExtents(&.{4}, rewritten.typeOf(rewritten.outputs[0]));341}