tiny.accy.preparation.canonicalization

1const std = @import("std");2const alloc_fixed = @import("alloc_fixed");3const choir = @import("choir");4const accy_choir = @import("../choir/root.zig");5const dialect_mod = accy_choir.dialect;67const ir = choir.ir;8const rewrite = ir.rewrite;9const passes = choir.passes;10const work = passes.pass.work;11const AccyDialect = dialect_mod.AccyDialect;1213pub const canonicalization_pass_name = "accy-choir-canonicalize";14pub const canonicalization_pass_description =15    "Canonicalize identity Accy Choir tensor operations";16const RewriteFn = *const fn (*ir.Operation, *rewrite.PatternRewriter) rewrite.PatternResult;1718pub const CanonicalizationPatternEntry = struct {19    spec: rewrite.RewritePatternSpec,20    rewrite: RewriteFn,21};2223pub const canonicalization_pattern_entries = [_]CanonicalizationPatternEntry{24    .{ .spec = .{ .name = "accy-reshape-identity", .root_op_name = AccyDialect.ReshapeOp.operation_name, .products = .none }, .rewrite = rewriteReshape },25    .{ .spec = .{ .name = "accy-transpose-identity", .root_op_name = AccyDialect.TransposeOp.operation_name, .products = .none }, .rewrite = rewriteTranspose },26    .{ .spec = .{ .name = "accy-broadcast-identity", .root_op_name = AccyDialect.BroadcastOp.operation_name, .products = .none }, .rewrite = rewriteBroadcast },27    .{ .spec = .{ .name = "accy-broadcast-in-dim-identity", .root_op_name = AccyDialect.BroadcastInDimOp.operation_name, .products = .none }, .rewrite = rewriteBroadcastInDim },28    .{ .spec = .{ .name = "accy-slice-identity", .root_op_name = AccyDialect.SliceOp.operation_name, .products = .none }, .rewrite = rewriteSlice },29    .{ .spec = .{ .name = "accy-concatenate-single-input", .root_op_name = AccyDialect.ConcatenateOp.operation_name, .products = .none }, .rewrite = rewriteConcatenate },30    .{ .spec = .{ .name = "accy-dot-flatten-shared-rhs", .root_op_name = AccyDialect.DotGeneralOp.operation_name, .products = .none }, .rewrite = rewriteDotGeneral },31    .{ .spec = .{ .name = "accy-cumsum-attach-scratch", .root_op_name = AccyDialect.CumsumOp.operation_name, .products = .none }, .rewrite = rewriteCumsum },32};3334pub const cumsum_block_tile = 8192;3536const greedy_config: passes.GreedyRewriteConfig = .{};3738const AccyCanonicalizationPass = passes.CanonicalizationPass(.{39    .name = canonicalization_pass_name,40    .description = canonicalization_pass_description,41    .populate_patterns = populateCanonicalizationPatterns,42    .greedy_config = greedy_config,43});4445pub fn canonicalizationPass() passes.Pass {46    var result = AccyCanonicalizationPass.init();47    result.work_contract = .{48        .identity = .{ .name = canonicalization_pass_name, .version = 1 },49        .estimate = canonicalizationWork,50    };51    return result;52}5354const CanonicalizationWorkFacts = struct {55    type_key_bytes: u64 = 0,56    list_bytes: u64 = 0,57    broadcasts: u64 = 0,58    dots: u64 = 0,59    scans: u64 = 0,6061    fn visit(self: *CanonicalizationWorkFacts, op: *ir.Operation) !ir.WalkResult {62        const namespace = op.name.getDialectNamespace();63        if (!std.mem.eql(u8, namespace, "accy") and !std.mem.eql(u8, namespace, "func") and64            !std.mem.eql(u8, namespace, "builtin")) return error.UnsupportedDialect;65        if (op.hasInterface(ir.interfaces.FoldOpInterface)) return error.MissingWorkContract;66        const name = op.name.name;67        if (std.mem.eql(u8, name, AccyDialect.BroadcastOp.operation_name) or68            std.mem.eql(u8, name, AccyDialect.BroadcastInDimOp.operation_name))69        {70            self.broadcasts = try work.add(self.broadcasts, 1);71        }72        if (std.mem.eql(u8, name, AccyDialect.DotGeneralOp.operation_name)) {73            self.dots = try work.add(self.dots, 1);74        }75        if (std.mem.eql(u8, name, AccyDialect.CumsumOp.operation_name)) {76            self.scans = try work.add(self.scans, 1);77        }78        for (op.results.items) |*value| self.typeKey(value.type);79        for (op.getOperandValues()) |value| self.typeKey(value.type);80        var attributes = op.getAttrs();81        while (attributes.next()) |entry| {82            if (entry.value.cast(ir.Attribute.DialectAttr)) |attr| {83                self.list_bytes = @max(self.list_bytes, attr.payload.len);84            }85        }86        return .advance;87    }8889    fn typeKey(self: *CanonicalizationWorkFacts, typ: ir.Type) void {90        if (typ.getDialectParamKey()) |key| self.type_key_bytes = @max(self.type_key_bytes, key.len);91    }92};9394/// The compile chain calls this before the pass runs, to charge the costs a pass declares before it95/// runs against the caller's limits as the work bound, so compilation can refuse the pass on them.96/// The rules this pass admits remove operations that change nothing (reshape, transpose, broadcast,97/// `broadcast_in_dim`, slice, and a concatenate with one input), rewrite a plain broadcast as one98/// `broadcast_in_dim`, merge a chain of `broadcast_in_dim` into one, split one batched dot into a99/// reshape, a dot and a reshape, and attach one scratch value to a cumulative sum. No rule makes a100/// chain longer, so the operations the pass can create are bounded by counts taken from the input:101/// at most three per dot, two per cumulative sum, and a quadratic term in the number of broadcasts.102/// The bound multiplies those counts by the rewriter's iteration limit.103fn canonicalizationWork(input: work.Input) !work.Bounds {104    const counts = try work.Census.inspect(input.operation);105    var facts: CanonicalizationWorkFacts = .{};106    _ = try input.operation.walk(.{ .order = .pre_order }, &facts, CanonicalizationWorkFacts.visit);107    const patterns = try canonicalizationPatternStorage(input.operation.context);108    const broadcasts = try work.multiply(facts.broadcasts, try work.add(facts.broadcasts, 1));109    const dot_and_scan = try work.add(try work.multiply(3, facts.dots), try work.multiply(2, facts.scans));110    const generated = try work.add(broadcasts, dot_and_scan);111    const traversed = try work.add(counts.operations, try work.multiply(2, generated));112    const attempts = try work.multiply(greedy_config.max_iterations, traversed);113    const dimensions = try work.multiply(try work.add(facts.type_key_bytes, 128), @sizeOf(i64));114    const scratch = try work.multiply(12, try work.add(115        try work.add(dimensions, facts.list_bytes),116        @alignOf(i64),117    ));118    const queues = try work.multiply(2, try work.arrayListGrowth(*ir.Operation, traversed));119    const iterations = try work.multiply(greedy_config.max_iterations, queues);120    const workspace = try work.add(patterns.bytes, try work.add(121        iterations,122        try work.multiply(attempts, scratch),123    ));124    const units = try work.add(try work.add(counts.atoms, counts.input_bytes), generated);125    const scans = try work.multiply(try work.add(patterns.patterns, 64), try work.multiply(attempts, try work.add(units, 1)));126    return .{127        .work = .{128            .input_bytes = try work.add(counts.input_bytes, patterns.name_bytes),129            .output_bytes = input.operation.context.capacity.storage_bytes,130            .structural_visits = try work.add(scans, patterns.visits),131            .rewrite_attempts = try work.multiply(attempts, canonicalization_pattern_entries.len),132            .allocation_capacity = workspace,133        },134        .workspace = workspace,135    };136}137138fn canonicalizationPatternStorage(139    context: *ir.Context,140) !passes.canonicalization.PatternPopulationBounds {141    const Interface = rewrite.DialectCanonicalizationInterface;142    const known = choir.dialects.arith.canonicalization_patterns[0..];143    var interfaces = context.dialect_registry.interfaces.iterator();144    while (interfaces.next()) |entry| {145        for (entry.value_ptr.items) |interface| {146            if (interface.id != Interface.id) continue;147            const patterns = Interface.fromOpaque(interface.vtable).patterns;148            if (patterns.len != 0 and (patterns.len != known.len or patterns.ptr != known.ptr)) {149                return error.MissingWorkContract;150            }151        }152    }153    const extra = comptime blk: {154        var specs: [canonicalization_pattern_entries.len]rewrite.RewritePatternSpec = undefined;155        for (canonicalization_pattern_entries, 0..) |entry, index| specs[index] = entry.spec;156        break :blk specs;157    };158    return passes.canonicalization.patternPopulationBounds(context, &extra);159}160161pub fn canonicalizeModule(162    allocator: std.mem.Allocator,163    choir_module: *ir.Operation,164    ctx: *ir.Context,165) !usize {166    var pm = passes.PassManager.init(allocator);167    defer pm.deinit();168    pm.enableVerifier();169    try pm.addPass(canonicalizationPass());170    if (pm.run(choir_module, ctx) == .failure) return error.CanonicalizationFailed;171    return pm.stats.passes_modified;172}173174pub fn populateCanonicalizationPatterns(patterns: *rewrite.RewritePatternSet) !void {175    for (canonicalization_pattern_entries) |entry| {176        try patterns.add(rewrite.RewritePattern.init(entry.spec, entry.rewrite));177    }178}179180fn rewriteResult(result: anyerror!bool) rewrite.PatternResult {181    return if (result catch return .failure) .success else .failure;182}183184fn rewriteReshape(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {185    return rewriteResult(canonicalizeReshape(rewriter.allocator, op, rewriter));186}187188fn rewriteTranspose(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {189    return rewriteResult(canonicalizeTranspose(rewriter.allocator, op, rewriter));190}191192fn rewriteBroadcast(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {193    return rewriteResult(canonicalizeBroadcast(rewriter.allocator, op, rewriter));194}195196fn rewriteBroadcastInDim(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {197    return rewriteResult(canonicalizeBroadcastInDim(rewriter.allocator, op, rewriter));198}199200fn rewriteSlice(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {201    return rewriteResult(canonicalizeSlice(rewriter.allocator, op, rewriter));202}203204fn rewriteConcatenate(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {205    return rewriteResult(canonicalizeConcatenate(rewriter.allocator, op, rewriter));206}207208fn broadcastBatchedSource(rhs: *ir.Value) ?*ir.Value {209    const def_any = rhs.getDefiningOp() orelse return null;210    const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));211    if (!std.mem.eql(u8, def_op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) return null;212    if (def_op.getNumOperands() != 1) return null;213    var dims_buffer: [8]i64 = undefined;214    const attr = def_op.getAttrAs(ir.Attribute.DialectAttr, "broadcast_dims") orelse return null;215    if (attr.payload.len != 2 * @sizeOf(i64)) return null;216    @memcpy(std.mem.sliceAsBytes(dims_buffer[0..2]), attr.payload[0 .. 2 * @sizeOf(i64)]);217    if (dims_buffer[0] != 1 or dims_buffer[1] != 2) return null;218    return def_op.getOperand(0);219}220221fn rewriteDotGeneral(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {222    return rewriteResult(canonicalizeDotGeneral(rewriter.allocator, op, rewriter));223}224225fn rewriteCumsum(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {226    return rewriteResult(canonicalizeCumsum(rewriter.allocator, op, rewriter));227}228229fn canonicalizeCumsum(230    allocator: std.mem.Allocator,231    op: *ir.Operation,232    rewriter: *rewrite.PatternRewriter,233) !bool {234    if (op.getNumOperands() != 1 or op.getNumResults() != 1) return false;235    const input = op.getOperand(0) orelse return false;236    const result = op.getResult(0) orelse return false;237238    const shape = try tensorShapeAlloc(allocator, input.type) orelse return false;239    defer allocator.free(shape);240    if (shape.len != 1) return false;241    const total = shape[0];242    if (total <= cumsum_block_tile) return false;243244    const blocks = @divTrunc(total + cumsum_block_tile - 1, cumsum_block_tile);245    const words = 1 + 3 * blocks;246247    const ctx = rewriter.ir_ctx;248    const scratch_type = try dialect_mod.accyTensorType(ctx, .u32, &.{words});249    const scratch = try AccyDialect.ScratchOp.create(ctx, op.location, scratch_type, words);250    _ = try rewriter.insert(scratch.op);251252    const axis_attr = op.getAttrAs(ir.Attribute.IntegerAttr, "axis") orelse return false;253    const replacement = try AccyDialect.CumsumOp.createWithScratch(ctx, op.location, input, scratch.getResult(), result.type, axis_attr.getValue());254    _ = try rewriter.insert(replacement.op);255    try rewriter.replaceOpWithOperation(op, replacement.op);256    return true;257}258259fn canonicalizeDotGeneral(260    allocator: std.mem.Allocator,261    op: *ir.Operation,262    rewriter: *rewrite.PatternRewriter,263) !bool {264    if (op.getNumOperands() != 2 or op.getNumResults() != 1) return false;265    const lhs = op.getOperand(0) orelse return false;266    var rhs = op.getOperand(1) orelse return false;267    const result = op.getResult(0) orelse return false;268269    const lhs_shape = try tensorShapeAlloc(allocator, lhs.type) orelse return false;270    defer allocator.free(lhs_shape);271    const result_shape = try tensorShapeAlloc(allocator, result.type) orelse return false;272    defer allocator.free(result_shape);273    if (lhs_shape.len != 3 or result_shape.len != 3) return false;274275    const lhs_contract = try readI64ListAttrAlloc(allocator, op, "lhs_contract", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("lhs_contract")) orelse return false;276    defer allocator.free(lhs_contract);277    const rhs_contract = try readI64ListAttrAlloc(allocator, op, "rhs_contract", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("rhs_contract")) orelse return false;278    defer allocator.free(rhs_contract);279    const lhs_batch = try readI64ListAttrAlloc(allocator, op, "lhs_batch", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("lhs_batch")) orelse return false;280    defer allocator.free(lhs_batch);281    const rhs_batch = try readI64ListAttrAlloc(allocator, op, "rhs_batch", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("rhs_batch")) orelse return false;282    defer allocator.free(rhs_batch);283284    if (lhs_contract.len != 1 or rhs_contract.len != 1) return false;285    if (lhs_contract[0] != 2) return false;286287    if (rhs_batch.len == 0) {288        if (rhs_contract[0] != 0) return false;289        if (lhs_batch.len > 1) return false;290        if (lhs_batch.len == 1 and lhs_batch[0] != 0) return false;291    } else {292        if (lhs_batch.len != 1 or rhs_batch.len != 1) return false;293        if (lhs_batch[0] != 0 or rhs_batch[0] != 0) return false;294        if (rhs_contract[0] != 1) return false;295        rhs = broadcastBatchedSource(rhs) orelse return false;296    }297298    const rhs_shape = try tensorShapeAlloc(allocator, rhs.type) orelse return false;299    defer allocator.free(rhs_shape);300    if (rhs_shape.len != 2) return false;301302    const batch = lhs_shape[0];303    const m = lhs_shape[1];304    const k = lhs_shape[2];305    const n = rhs_shape[1];306    if (rhs_shape[0] != k) return false;307    if (result_shape[0] != batch or result_shape[1] != m or result_shape[2] != n) return false;308    const flat_m = std.math.mul(i64, batch, m) catch return false;309310    var dtype_arena_buffer: [256]u8 = undefined;311    var dtype_arena = alloc_fixed.FixedBuffer.init(dtype_arena_buffer[0..]);312    const lhs_decoded = dialect_mod.decodeTensorType(dtype_arena.allocator(), lhs.type) catch return false;313314    const ctx = rewriter.ir_ctx;315    const flat_lhs_type = try dialect_mod.accyTensorType(ctx, lhs_decoded.dtype, &.{ flat_m, k });316    const flat_out_type = try dialect_mod.accyTensorType(ctx, lhs_decoded.dtype, &.{ flat_m, n });317318    const flat_lhs = try dialect_mod.AccyDialect.ReshapeOp.create(ctx, op.location, lhs, flat_lhs_type, &.{ flat_m, k });319    _ = try rewriter.insert(flat_lhs.op);320    const flat_dot = try dialect_mod.AccyDialect.DotGeneralOp.create(ctx, op.location, flat_lhs.getResult(), rhs, flat_out_type, &.{}, &.{}, &.{1}, &.{0});321    _ = try rewriter.insert(flat_dot.op);322    const restored = try dialect_mod.AccyDialect.ReshapeOp.create(ctx, op.location, flat_dot.getResult(), result.type, result_shape);323    _ = try rewriter.insert(restored.op);324    try rewriter.replaceOpWithOperation(op, restored.op);325    return true;326}327328fn canonicalizeReshape(329    allocator: std.mem.Allocator,330    op: *ir.Operation,331    rewriter: *rewrite.PatternRewriter,332) !bool {333    const input = identityUnaryInput(op) orelse return false;334    const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;335    defer allocator.free(input_shape);336    const new_shape = try readI64ListAttrAlloc(allocator, op, "new_shape", dialect_mod.AccyDialect.ReshapeOp.dialectAttrName("new_shape")) orelse return false;337    defer allocator.free(new_shape);338    if (!std.mem.eql(i64, input_shape, new_shape)) return false;339    return replaceWithInput(op, input, rewriter);340}341342fn canonicalizeTranspose(343    allocator: std.mem.Allocator,344    op: *ir.Operation,345    rewriter: *rewrite.PatternRewriter,346) !bool {347    const input = identityUnaryInput(op) orelse return false;348    const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;349    defer allocator.free(input_shape);350    const permutation = try readI64ListAttrAlloc(allocator, op, "permutation", dialect_mod.AccyDialect.TransposeOp.dialectAttrName("permutation")) orelse return false;351    defer allocator.free(permutation);352    if (permutation.len != input_shape.len) return false;353    if (!isIdentityPermutation(permutation)) return false;354    return replaceWithInput(op, input, rewriter);355}356357fn canonicalizeBroadcast(358    allocator: std.mem.Allocator,359    op: *ir.Operation,360    rewriter: *rewrite.PatternRewriter,361) !bool {362    if (op.getNumOperands() != 1 or op.getNumResults() != 1) return false;363    const input = op.getOperand(0) orelse return false;364    const result = op.getResult(0) orelse return false;365    const sizes = try readI64ListAttrAlloc(allocator, op, "sizes", dialect_mod.AccyDialect.BroadcastOp.dialectAttrName("sizes")) orelse return false;366    defer allocator.free(sizes);367    if (sizes.len == 0) {368        if (!result.type.eql(input.type)) return false;369        return replaceWithInput(op, input, rewriter);370    }371372    const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;373    defer allocator.free(input_shape);374375    const result_shape = try allocator.alloc(i64, sizes.len + input_shape.len);376    defer allocator.free(result_shape);377    @memcpy(result_shape[0..sizes.len], sizes);378    @memcpy(result_shape[sizes.len..], input_shape);379380    const broadcast_dims = try allocator.alloc(i64, input_shape.len);381    defer allocator.free(broadcast_dims);382    for (broadcast_dims, 0..) |*dim, index| dim.* = @intCast(sizes.len + index);383384    const replacement = try dialect_mod.AccyDialect.BroadcastInDimOp.create(385        rewriter.ir_ctx,386        op.location,387        input,388        result.type,389        broadcast_dims,390        result_shape,391    );392    _ = try rewriter.insert(replacement.op);393    try rewriter.replaceOpWithOperation(op, replacement.op);394    return true;395}396397fn canonicalizeBroadcastInDim(398    allocator: std.mem.Allocator,399    op: *ir.Operation,400    rewriter: *rewrite.PatternRewriter,401) !bool {402    if (try collapseChainedBroadcastInDim(allocator, op, rewriter)) return true;403    const input = identityUnaryInput(op) orelse return false;404    const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;405    defer allocator.free(input_shape);406    const broadcast_dims = try readI64ListAttrAlloc(allocator, op, "broadcast_dims", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("broadcast_dims")) orelse return false;407    defer allocator.free(broadcast_dims);408    const result_shape = try readI64ListAttrAlloc(allocator, op, "result_shape", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("result_shape")) orelse return false;409    defer allocator.free(result_shape);410    if (broadcast_dims.len != input_shape.len) return false;411    if (!isIdentityPermutation(broadcast_dims)) return false;412    if (!std.mem.eql(i64, input_shape, result_shape)) return false;413    return replaceWithInput(op, input, rewriter);414}415416fn collapseChainedBroadcastInDim(417    allocator: std.mem.Allocator,418    op: *ir.Operation,419    rewriter: *rewrite.PatternRewriter,420) !bool {421    if (op.getNumOperands() != 1 or op.getNumResults() != 1) return false;422    const mid = op.getOperand(0) orelse return false;423    const result = op.getResult(0) orelse return false;424    const inner_any = mid.getDefiningOp() orelse return false;425    const inner: *ir.Operation = @ptrCast(@alignCast(inner_any));426    if (!std.mem.eql(u8, inner.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) return false;427    if (inner.getNumOperands() != 1) return false;428    const source = inner.getOperand(0) orelse return false;429430    const outer_dims = try readI64ListAttrAlloc(allocator, op, "broadcast_dims", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("broadcast_dims")) orelse return false;431    defer allocator.free(outer_dims);432    const outer_shape = try readI64ListAttrAlloc(allocator, op, "result_shape", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("result_shape")) orelse return false;433    defer allocator.free(outer_shape);434    const inner_dims = try readI64ListAttrAlloc(allocator, inner, "broadcast_dims", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("broadcast_dims")) orelse return false;435    defer allocator.free(inner_dims);436437    const composed = try allocator.alloc(i64, inner_dims.len);438    defer allocator.free(composed);439    for (inner_dims, 0..) |mid_dim, index| {440        if (mid_dim < 0 or mid_dim >= outer_dims.len) return false;441        composed[index] = outer_dims[@intCast(mid_dim)];442    }443444    const replacement = try dialect_mod.AccyDialect.BroadcastInDimOp.create(445        rewriter.ir_ctx,446        op.location,447        source,448        result.type,449        composed,450        outer_shape,451    );452    _ = try rewriter.insert(replacement.op);453    try rewriter.replaceOpWithOperation(op, replacement.op);454    return true;455}456457fn canonicalizeSlice(458    allocator: std.mem.Allocator,459    op: *ir.Operation,460    rewriter: *rewrite.PatternRewriter,461) !bool {462    const input = identityUnaryInput(op) orelse return false;463    const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;464    defer allocator.free(input_shape);465    const starts = try readI64ListAttrAlloc(allocator, op, "starts", dialect_mod.AccyDialect.SliceOp.dialectAttrName("starts")) orelse return false;466    defer allocator.free(starts);467    const limits = try readI64ListAttrAlloc(allocator, op, "limits", dialect_mod.AccyDialect.SliceOp.dialectAttrName("limits")) orelse return false;468    defer allocator.free(limits);469    const strides = try readI64ListAttrAlloc(allocator, op, "strides", dialect_mod.AccyDialect.SliceOp.dialectAttrName("strides")) orelse return false;470    defer allocator.free(strides);471472    if (starts.len != input_shape.len or473        limits.len != input_shape.len or474        strides.len != input_shape.len)475    {476        return false;477    }478    for (starts, limits, strides, input_shape) |start, limit, stride, dim| {479        if (start != 0 or limit != dim or stride != 1) return false;480    }481    return replaceWithInput(op, input, rewriter);482}483484fn canonicalizeConcatenate(485    allocator: std.mem.Allocator,486    op: *ir.Operation,487    rewriter: *rewrite.PatternRewriter,488) !bool {489    if (op.getNumOperands() != 1 or op.getNumResults() != 1) return false;490    const input = op.getOperand(0) orelse return false;491    const result = op.getResult(0) orelse return false;492    if (!result.type.eql(input.type)) return false;493    const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;494    defer allocator.free(input_shape);495    const dimension = readIntegerAttr(op, "dimension") orelse return false;496    if (dimension < 0 or dimension >= @as(i64, @intCast(input_shape.len))) return false;497    try rewriter.replaceOpWithValue(op, input);498    return true;499}500501fn identityUnaryInput(op: *ir.Operation) ?*ir.Value {502    if (op.getNumOperands() != 1 or op.getNumResults() != 1) return null;503    const input = op.getOperand(0) orelse return null;504    const result = op.getResult(0) orelse return null;505    if (!result.type.eql(input.type)) return null;506    return input;507}508509fn replaceWithInput(510    op: *ir.Operation,511    input: *ir.Value,512    rewriter: *rewrite.PatternRewriter,513) !bool {514    try rewriter.replaceOpWithValue(op, input);515    return true;516}517518pub fn readDialectPayload(op: *const ir.Operation, attr_name: []const u8, dialect_attr_name: []const u8) ?[]const u8 {519    const attr = op.getAttr(attr_name) orelse return null;520    if (!std.mem.eql(u8, attr.abstract.name, dialect_attr_name)) return null;521    const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return null;522    return dialect_attr.payload;523}524525pub fn readI64ListAttrAlloc(526    allocator: std.mem.Allocator,527    op: *const ir.Operation,528    attr_name: []const u8,529    dialect_attr_name: []const u8,530) !?[]i64 {531    const payload = readDialectPayload(op, attr_name, dialect_attr_name) orelse return null;532    if (payload.len % @sizeOf(i64) != 0) return null;533    const values = try allocator.alloc(i64, payload.len / @sizeOf(i64));534    for (values, 0..) |*value, i| {535        const start = i * @sizeOf(i64);536        @memcpy(std.mem.asBytes(value), payload[start..][0..@sizeOf(i64)]);537    }538    return values;539}540541fn readIntegerAttr(op: *const ir.Operation, attr_name: []const u8) ?i64 {542    const attr = op.getAttr(attr_name) orelse return null;543    if (!std.mem.eql(u8, attr.abstract.name, "builtin.integer")) return null;544    const int_attr = attr.cast(ir.Attribute.IntegerAttr) orelse return null;545    return int_attr.getValue();546}547548fn tensorShapeAlloc(allocator: std.mem.Allocator, typ: ir.Type) !?[]i64 {549    const type_name = typ.getDialectTypeName() orelse return null;550    if (!std.mem.eql(u8, type_name, dialect_mod.tensor_type_name)) return null;551    const key = typ.getDialectParamKey() orelse return null;552    const comma = std.mem.indexOfScalar(u8, key, ',') orelse return null;553    const dims_text = key[comma + 1 ..];554    if (dims_text.len == 0) return try allocator.alloc(i64, 0);555556    var dim_count: usize = 1;557    for (dims_text) |ch| {558        if (ch == 'x') dim_count += 1;559    }560561    const dims = try allocator.alloc(i64, dim_count);562    errdefer allocator.free(dims);563    var iter = std.mem.splitScalar(u8, dims_text, 'x');564    var index: usize = 0;565    while (iter.next()) |part| {566        if (part.len == 0 or index >= dim_count) {567            allocator.free(dims);568            return null;569        }570        const dim = std.fmt.parseInt(i64, part, 10) catch {571            allocator.free(dims);572            return null;573        };574        if (dim < 0) {575            allocator.free(dims);576            return null;577        }578        dims[index] = dim;579        index += 1;580    }581    if (index != dim_count) {582        allocator.free(dims);583        return null;584    }585    return dims;586}587588fn isIdentityPermutation(values: []const i64) bool {589    for (values, 0..) |value, i| {590        if (value != @as(i64, @intCast(i))) return false;591    }592    return true;593}594595const testing = std.testing;596const semantic = accy_choir.semantic;597598fn readSymbolName(func: *ir.Operation) ?[]const u8 {599    return ir.SymbolTable.getSymbolName(func);600}601602fn findOpNamedInBlock(block: *ir.Block, name: []const u8) ?*ir.Operation {603    var iter = block.operations.head;604    while (iter) |op_ptr| {605        const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));606        if (std.mem.eql(u8, op.name.name, name)) return op;607        iter = op.next_op;608    }609    return null;610}611612test "canonicalization patterns publish Choir rewrite specs" {613    const allocator = testing.allocator;614615    var patterns = rewrite.RewritePatternSet.init(allocator);616    defer patterns.deinit();617618    try populateCanonicalizationPatterns(&patterns);619620    try testing.expectEqual(canonicalization_pattern_entries.len, patterns.patterns.items.len);621    for (canonicalization_pattern_entries, patterns.patterns.items) |entry, pattern| {622        try testing.expectEqualStrings(entry.spec.name, pattern.spec.name);623        try testing.expectEqualStrings(entry.spec.root_op_name, pattern.spec.root_op_name);624        try testing.expectEqual(entry.spec.benefit, pattern.spec.benefit);625        try testing.expectEqual(entry.spec.kind, pattern.spec.kind);626        switch (entry.spec.products) {627            .none => {},628            else => return error.TestExpectedNoProducts,629        }630    }631    try testing.expectEqualStrings(AccyDialect.ReshapeOp.operation_name, canonicalization_pattern_entries[0].spec.root_op_name);632}633634test "canonicalizeModule removes identity shape operations" {635    const allocator = testing.allocator;636637    var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);638    defer builder.deinit();639    const f32_2x3 = try builder.tensor(.f32, &.{ 2, 3 });640    var fb = try builder.beginFunction("identity_shapes", &.{f32_2x3}, &.{f32_2x3});641    const r0 = try fb.reshape(fb.parameter(0), f32_2x3, &.{ 2, 3 });642    const r1 = try fb.transpose(r0, f32_2x3, &.{ 0, 1 });643    const r2 = try fb.broadcast(r1, f32_2x3, &.{});644    const r3 = try fb.broadcastInDim(r2, f32_2x3, &.{ 2, 3 }, &.{ 0, 1 });645    const r4 = try fb.slice(r3, f32_2x3, &.{ 0, 0 }, &.{ 2, 3 }, &.{ 1, 1 });646    const r5 = try fb.concatenate(&.{r4}, f32_2x3, 0);647    try fb.return_(&.{r5});648    try fb.finish();649    const module = try builder.finish();650    defer module.deinit();651652    const choir_mod = module.choir_module;653    const ctx = module.context();654    try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.ReshapeOp.operation_name));655    try testing.expectEqual(@as(usize, 1), try canonicalizeModule(allocator, choir_mod, ctx));656    try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.ReshapeOp.operation_name));657    try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.TransposeOp.operation_name));658    try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.BroadcastOp.operation_name));659    try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name));660    try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.SliceOp.operation_name));661    try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.ConcatenateOp.operation_name));662663    const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;664    const func = ir.inspection.functionByNameInBlock(module_body, "identity_shapes") orelse return error.TestExpectedFunc;665    const entry = func.getRegion(0).?.getEntryBlock().?;666    const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;667    try testing.expectEqual(entry.getArgument(0).?, ret.getOperand(0).?);668}669670test "canonicalization pass preserves analyses when it makes no changes" {671    const allocator = testing.allocator;672673    var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);674    defer builder.deinit();675    const f32_4 = try builder.tensor(.f32, &.{4});676    var fb = try builder.beginFunction("canonicalize_add4", &.{ f32_4, f32_4 }, &.{f32_4});677    const sum = try fb.add(fb.parameter(0), fb.parameter(1));678    try fb.return_(&.{sum});679    try fb.finish();680    const module = try builder.finish();681    defer module.deinit();682683    const choir_mod = module.choir_module;684    const ctx = module.context();685    var pm = passes.PassManager.init(allocator);686    defer pm.deinit();687    try pm.addPass(canonicalizationPass());688689    try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));690    try testing.expectEqual(@as(u64, 1), pm.stats.pass_runs);691    try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);692}693694test "canonicalizeModule flattens shared-rhs batched dots" {695    const allocator = testing.allocator;696697    var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);698    defer builder.deinit();699    const lhs_ty = try builder.tensor(.f32, &.{ 8, 2, 4 });700    const rhs_ty = try builder.tensor(.f32, &.{ 4, 3 });701    const out_ty = try builder.tensor(.f32, &.{ 8, 2, 3 });702    var fb = try builder.beginFunction("flatten_dot", &.{ lhs_ty, rhs_ty }, &.{out_ty});703    const product = try fb.dotGeneral(fb.parameter(0), fb.parameter(1), out_ty, &.{2}, &.{0}, &.{}, &.{});704    try fb.return_(&.{product});705    try fb.finish();706    const module = try builder.finish();707    defer module.deinit();708709    const choir_mod = module.choir_module;710    const ctx = module.context();711    _ = try canonicalizeModule(allocator, choir_mod, ctx);712    try testing.expectEqual(@as(usize, 2), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.ReshapeOp.operation_name));713    try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.DotGeneralOp.operation_name));714}715716test "canonicalization accounting reserves worker storage and closes actual identity rewrites" {717    const module = try accountingIdentityModule(testing.allocator);718    defer module.deinit();719    try checkCanonicalizationAccounting(module, 6);720}721722fn checkCanonicalizationAccounting(module: *semantic.SemanticModule, rewrites: u64) !void {723    const revision = choir.product.revision;724    const allocator = testing.allocator;725    const selected = canonicalizationPass();726    const declared = try selected.work_contract.?.estimate(.{ .operation = module.choir_module });727    const ledger = try revision.AccountingV1.create(allocator, .{728        .allowance = .uniform(std.math.maxInt(u64)),729        .workspace = declared.workspace,730        .events = 16,731    }, &.{.{ .name = canonicalization_pass_name, .version = 1 }});732    defer ledger.destroy();733    var observed = testing.FailingAllocator.init(allocator, .{});734    var cache = try passes.AnalysisCache.initAccounted(observed.allocator(), null, ledger, .{}, 8);735    defer cache.deinit();736    var manager = passes.PassManager.init(observed.allocator());737    defer manager.deinit();738    try manager.addPass(selected);739    const before = observed.allocated_bytes;740    try testing.expectEqual(.success, manager.runWithAnalysisCache(741        module.choir_module,742        module.context(),743        &cache,744        .{ .max_threads = 1 },745    ));746    const worker_bytes = observed.allocated_bytes - before;747    try ledger.producersComplete();748    const receipt = ledger.view();749    try testing.expectEqual(1, receipt.executed.counters.pass_runs);750    try testing.expectEqual(@intFromBool(rewrites != 0), receipt.executed.counters.passes_modified);751    try testing.expectEqual(rewrites, receipt.executed.counters.successful_rewrites);752    const iterations: u64 = if (rewrites == 0) 1 else 2;753    try testing.expectEqual(iterations, receipt.executed.counters.rewrite_iterations);754    var pass_events: u32 = 0;755    for (receipt.events) |event| {756        if (event.phase != .pass) continue;757        pass_events += 1;758        try testing.expectEqual(declared.workspace, event.workspace);759        try testing.expect(event.workspace >= worker_bytes);760        try testing.expect(event.charged.allocation_capacity >= worker_bytes);761        try testing.expect(event.charged.rewrite_attempts >= rewrites);762    }763    try testing.expectEqual(1, pass_events);764    try module.verify();765}766767fn accountingIdentityModule(allocator: std.mem.Allocator) !*semantic.SemanticModule {768    var builder = try semantic.Builder.init(allocator, .testing);769    defer builder.deinit();770    const typ = try builder.tensor(.f32, &.{ 2, 3 });771    var function = try builder.beginFunction("accounted_identity", &.{typ}, &.{typ});772    const reshape = try function.reshape(function.parameter(0), typ, &.{ 2, 3 });773    const transpose = try function.transpose(reshape, typ, &.{ 0, 1 });774    const broadcast = try function.broadcast(transpose, typ, &.{});775    const in_dim = try function.broadcastInDim(broadcast, typ, &.{ 2, 3 }, &.{ 0, 1 });776    const slice = try function.slice(in_dim, typ, &.{ 0, 0 }, &.{ 2, 3 }, &.{ 1, 1 });777    const concat = try function.concatenate(&.{slice}, typ, 0);778    try function.return_(&.{concat});779    try function.finish();780    const module = try builder.finish();781    return module;782}783784test "canonicalization accounting refuses workspace and rewrite budgets before mutation" {785    inline for (.{ false, true }) |exhaust_workspace| {786        try checkCanonicalizationRefusal(exhaust_workspace);787    }788}789790fn checkCanonicalizationRefusal(exhaust_workspace: bool) !void {791    const revision = choir.product.revision;792    const allocator = testing.allocator;793    const module = try accountingIdentityModule(allocator);794    defer module.deinit();795    const selected = canonicalizationPass();796    const declared = try selected.work_contract.?.estimate(.{ .operation = module.choir_module });797    try testing.expect(declared.workspace > 0);798    try testing.expect(declared.work.rewrite_attempts > 0);799    var allowance = revision.WorkVector.uniform(std.math.maxInt(u64));800    if (!exhaust_workspace) allowance.rewrite_attempts = declared.work.rewrite_attempts - 1;801    const ledger = try revision.AccountingV1.create(allocator, .{802        .allowance = allowance,803        .workspace = declared.workspace - @intFromBool(exhaust_workspace),804        .events = 16,805    }, &.{.{ .name = canonicalization_pass_name, .version = 1 }});806    defer ledger.destroy();807    var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);808    defer cache.deinit();809    var manager = passes.PassManager.init(allocator);810    defer manager.deinit();811    try manager.addPass(selected);812    try testing.expectEqual(.failure, manager.runWithAnalysisCache(813        module.choir_module,814        module.context(),815        &cache,816        .{ .max_threads = 1 },817    ));818    const receipt = ledger.view();819    try testing.expectEqual(.exhausted, receipt.outcome);820    try testing.expectEqual(0, receipt.executed.counters.pass_runs);821    try testing.expectEqual(0, receipt.executed.counters.passes_modified);822    try testing.expectEqual(0, receipt.executed.counters.successful_rewrites);823    try testing.expectEqual(0, receipt.executed.counters.rewrite_iterations);824    try testing.expectEqual(1, ir.inspection.countOperationsNamed(825        module.choir_module,826        AccyDialect.ReshapeOp.operation_name,827    ));828    try testing.expectError(error.WorkExhausted, ledger.producersComplete());829    try module.verify();830}831832test "canonicalization accounting covers generated dots scans and broadcast chains" {833    const allocator = testing.allocator;834    const dot = try accountingDotModule(allocator);835    defer dot.deinit();836    try checkCanonicalizationAccounting(dot, 1);837    const scan = try accountingScanModule(allocator);838    defer scan.deinit();839    try checkCanonicalizationAccounting(scan, 1);840    for ([_]u32{ 0, 1, 4, 16 }) |depth| {841        const chain = try accountingBroadcastModule(allocator, depth);842        defer chain.deinit();843        try checkCanonicalizationAccounting(chain, if (depth == 0) 0 else 2 * depth - 1);844        try testing.expectEqual(@intFromBool(depth > 0), ir.inspection.countOperationsNamed(845            chain.choir_module,846            AccyDialect.BroadcastInDimOp.operation_name,847        ));848    }849}850851fn accountingDotModule(allocator: std.mem.Allocator) !*semantic.SemanticModule {852    var builder = try semantic.Builder.init(allocator, .testing);853    defer builder.deinit();854    const lhs = try builder.tensor(.f32, &.{ 8, 2, 4 });855    const rhs = try builder.tensor(.f32, &.{ 4, 3 });856    const output = try builder.tensor(.f32, &.{ 8, 2, 3 });857    var function = try builder.beginFunction("accounted_dot", &.{ lhs, rhs }, &.{output});858    const value = try function.dotGeneral(859        function.parameter(0),860        function.parameter(1),861        output,862        &.{2},863        &.{0},864        &.{},865        &.{},866    );867    try function.return_(&.{value});868    try function.finish();869    return builder.finish();870}871872fn accountingScanModule(allocator: std.mem.Allocator) !*semantic.SemanticModule {873    var builder = try semantic.Builder.init(allocator, .testing);874    defer builder.deinit();875    const typ = try builder.tensor(.f32, &.{16384});876    var function = try builder.beginFunction("accounted_scan", &.{typ}, &.{typ});877    const value = try function.cumsum(function.parameter(0), typ, 0);878    try function.return_(&.{value});879    try function.finish();880    return builder.finish();881}882883fn accountingBroadcastModule(allocator: std.mem.Allocator, depth: u32) !*semantic.SemanticModule {884    std.debug.assert(depth <= 16);885    var builder = try semantic.Builder.init(allocator, .testing);886    defer builder.deinit();887    var dimensions: [17]i64 = @splat(2);888    dimensions[depth] = 1;889    const typ = try builder.tensor(.f32, &.{1});890    const output = try builder.tensor(.f32, dimensions[0 .. depth + 1]);891    var function = try builder.beginFunction("accounted_broadcast", &.{typ}, &.{output});892    var value = function.parameter(0);893    for (0..depth) |index| {894        const result = try builder.tensor(.f32, dimensions[depth - index - 1 .. depth + 1]);895        value = try function.broadcast(value, result, &.{2});896    }897    try function.return_(&.{value});898    try function.finish();899    return builder.finish();900}901902test "canonicalization accounting refuses an unmodeled registered callback" {903    const allocator = testing.allocator;904    const module = try accountingIdentityModule(allocator);905    defer module.deinit();906    const extra = comptime [_]rewrite.RewritePattern{rewrite.RewritePattern.init(.{907        .name = "unmodeled-reshape",908        .root_op_name = AccyDialect.ReshapeOp.operation_name,909        .benefit = 20,910    }, unmodeledCanonicalization)};911    const registered = comptime rewrite.DialectCanonicalizationInterface.entryFor("accy", &extra);912    try module.context().registerDialectInterface("accy", registered);913    const ledger = try choir.product.revision.AccountingV1.create(allocator, .{914        .allowance = .uniform(std.math.maxInt(u64)),915        .workspace = 128 * 1024 * 1024,916        .events = 16,917    }, &.{.{ .name = canonicalization_pass_name, .version = 1 }});918    defer ledger.destroy();919    var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);920    defer cache.deinit();921    var manager = passes.PassManager.init(allocator);922    defer manager.deinit();923    try manager.addPass(canonicalizationPass());924    try testing.expectEqual(.failure, manager.runWithAnalysisCache(925        module.choir_module,926        module.context(),927        &cache,928        .{ .max_threads = 1 },929    ));930    try testing.expectEqual(.rejected, ledger.view().outcome);931    try testing.expect(ledger.view().missing_work_contract);932    try testing.expectEqual(0, ledger.view().executed.counters.pass_runs);933    try testing.expectEqual(0, ledger.view().executed.counters.successful_rewrites);934    try testing.expectEqual(1, ir.inspection.countOperationsNamed(935        module.choir_module,936        AccyDialect.ReshapeOp.operation_name,937    ));938}939940fn unmodeledCanonicalization(_: *ir.Operation, _: *rewrite.PatternRewriter) rewrite.PatternResult {941    @panic("an unmodeled callback reached accounted work");942}

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