tiny.accy.preparation.fusion

1pub const ClaimKind = enum {2    recompute,3    slot,4};

1pub const FusionCluster = struct {2    ops: []*ir.Operation,3    kind: FusionClusterKind = .elementwise,45    pub fn len(self: FusionCluster) usize {6        return self.ops.len;7    }89    pub fn root(self: FusionCluster) ?*ir.Operation {10        if (self.ops.len == 0) return null;11        return self.ops[self.ops.len - 1];12    }1314    fn deinit(self: *FusionCluster, allocator: std.mem.Allocator) void {15        allocator.free(self.ops);16        self.* = undefined;17    }18};

1pub const FusionPlanAnalysis = struct {2    allocator: std.mem.Allocator,3    clusters: std.ArrayListUnmanaged(FusionCluster),4    elided: std.ArrayListUnmanaged(*ir.Operation),5    fused_op_count: usize = 0,6    max_cluster_len: usize = 0,78    pub fn init(allocator: std.mem.Allocator) FusionPlanAnalysis {9        return .{10            .allocator = allocator,11            .clusters = .empty,12            .elided = .empty,13        };14    }1516    pub fn deinit(self: *FusionPlanAnalysis) void {17        for (self.clusters.items) |*cluster| {18            cluster.deinit(self.allocator);19        }20        self.clusters.deinit(self.allocator);21        self.elided.deinit(self.allocator);22        self.* = undefined;23    }2425    pub fn clusterCount(self: FusionPlanAnalysis) usize {26        return self.clusters.items.len;27    }2829    fn addCluster(self: *FusionPlanAnalysis, ops: []const *ir.Operation) !void {30        try self.addClusterOfKind(ops, .elementwise);31    }3233    fn addClusterOfKind(self: *FusionPlanAnalysis, ops: []const *ir.Operation, kind: FusionClusterKind) !void {34        if (ops.len < 2) return;35        try self.addOwnedCluster(ops, kind);36    }3738    fn addReductionCluster(self: *FusionPlanAnalysis, ops: []const *ir.Operation) !void {39        if (ops.len == 0) return;40        try self.addOwnedCluster(ops, .reduction_input);41    }4243    fn addOwnedCluster(self: *FusionPlanAnalysis, ops: []const *ir.Operation, kind: FusionClusterKind) !void {44        const owned = try self.allocator.alloc(*ir.Operation, ops.len);45        @memcpy(owned, ops);46        errdefer self.allocator.free(owned);47        try self.clusters.append(self.allocator, .{ .ops = owned, .kind = kind });48        self.fused_op_count += owned.len;49        self.max_cluster_len = @max(self.max_cluster_len, owned.len);50    }51};

1pub const ClaimMap = std.AutoHashMap(*ir.Operation, ClaimKind);

1pub fn broadcastLeafSource(operand: *ir.Value) ?*ir.Value {2    const def_any = operand.getDefiningOp() orelse return null;3    const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));4    if (!isName(def_op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) return null;5    const operands = def_op.getOperandValues();6    if (operands.len != 1) return null;7    if (operands[0].getDefiningOp() != null) return null;8    return operands[0];9}

1pub const flash_attention_dim = 64;

1pub const flash_attention_min_seq = 128;

1pub fn fusionPlanningPass() passes.Pass {2    return .{3        .name = fusion_planning_pass_name,4        .description = fusion_planning_pass_description,5        .run_fn = runFusionPlanningPass,6        .work_contract = .{7            .identity = .{ .name = fusion_planning_pass_name, .version = 1 },8            .estimate = fusionPassWork,9        },10    };11}

1pub const fusion_plan_analysis_descriptor = passes.AnalysisDescriptor{2    .id = passes.analysisId(fusion_plan_analysis_name),3    .name = fusion_plan_analysis_name,4    .work_contract = .{5        .identity = .{ .name = fusion_plan_analysis_name, .version = 1 },6        .estimate = fusionAnalysisWork,7    },8};

1pub const fusion_plan_analysis_name = "accy-choir-fusion-plan";

1pub const fusion_planning_pass_description =2    "Plan block-local Accy Choir elementwise fusion clusters";

1pub const fusion_planning_pass_name = "accy-choir-plan-fusion";

1pub fn getFusionPlanAnalysis(2    pass_ctx: *passes.PassContext,3    op: *ir.Operation,4) !*FusionPlanAnalysis {5    const ptr = try pass_ctx.getAnalysis(6        op,7        &fusion_plan_analysis_descriptor,8        computeFusionPlanAnalysis,9        cleanupFusionPlanAnalysis,10    );11    return @ptrCast(@alignCast(ptr));12}

1pub fn isFlatGatherOp(op: *ir.Operation) bool {2    if (!isName(op.name.name, dialect_mod.AccyDialect.GatherOp.operation_name)) return false;3    if (op.getNumResults() != 1) return false;4    const operands = op.getOperandValues();5    if (operands.len != 2) return false;6    const attr = op.getAttr("axis") orelse return false;7    const int_attr = attr.cast(ir.Attribute.IntegerAttr) orelse return false;8    if (int_attr.getValue() != 0) return false;9    var arena_buffer: [256]u8 = undefined;10    var arena = alloc_fixed.FixedBuffer.init(arena_buffer[0..]);11    const src_type = dialect_mod.decodeTensorType(arena.allocator(), operands[0].type) catch return false;12    return src_type.dims.len == 1;13}

1pub fn isFusableElementwiseOp(op: *ir.Operation) bool {2    if (op.regions.items.len > 0) return false;3    if (op.getNumResults() != 1) return false;4    const name = op.name.name;5    return std.mem.eql(u8, name, dialect_mod.AccyDialect.AddOp.operation_name) or6        std.mem.eql(u8, name, dialect_mod.AccyDialect.SubOp.operation_name) or7        std.mem.eql(u8, name, dialect_mod.AccyDialect.MulOp.operation_name) or8        std.mem.eql(u8, name, dialect_mod.AccyDialect.DivOp.operation_name) or9        std.mem.eql(u8, name, dialect_mod.AccyDialect.MaxOp.operation_name) or10        std.mem.eql(u8, name, dialect_mod.AccyDialect.MinOp.operation_name) or11        std.mem.eql(u8, name, dialect_mod.AccyDialect.PowOp.operation_name) or12        std.mem.eql(u8, name, dialect_mod.AccyDialect.Atan2Op.operation_name) or13        std.mem.eql(u8, name, dialect_mod.AccyDialect.NegOp.operation_name) or14        std.mem.eql(u8, name, dialect_mod.AccyDialect.ExpOp.operation_name) or15        std.mem.eql(u8, name, dialect_mod.AccyDialect.LogOp.operation_name) or16        std.mem.eql(u8, name, dialect_mod.AccyDialect.TanhOp.operation_name) or17        std.mem.eql(u8, name, dialect_mod.AccyDialect.SqrtOp.operation_name) or18        std.mem.eql(u8, name, dialect_mod.AccyDialect.AbsOp.operation_name) or19        std.mem.eql(u8, name, dialect_mod.AccyDialect.SinOp.operation_name) or20        std.mem.eql(u8, name, dialect_mod.AccyDialect.CosOp.operation_name) or21        std.mem.eql(u8, name, dialect_mod.AccyDialect.TanOp.operation_name) or22        std.mem.eql(u8, name, dialect_mod.AccyDialect.FloorOp.operation_name) or23        std.mem.eql(u8, name, dialect_mod.AccyDialect.RoundOp.operation_name) or24        std.mem.eql(u8, name, dialect_mod.AccyDialect.TruncOp.operation_name) or25        std.mem.eql(u8, name, dialect_mod.AccyDialect.CompareOp.operation_name) or26        std.mem.eql(u8, name, dialect_mod.AccyDialect.ConvertOp.operation_name) or27        std.mem.eql(u8, name, dialect_mod.AccyDialect.SelectOp.operation_name);28}

1pub fn isSeeThroughShapeOp(op: *ir.Operation) bool {2    if (op.getNumResults() != 1) return false;3    if (isName(op.name.name, dialect_mod.AccyDialect.SliceOp.operation_name)) return true;4    if (isName(op.name.name, dialect_mod.AccyDialect.PadOp.operation_name)) {5        return padHasZeroInterior(op);6    }7    return false;8}

1pub const max_concat_reduction_group = 8;

1pub const max_dot_epilogue_ops = 6;

1pub const max_reduction_prologue_ops = 48;

1pub const max_row_pipeline_cols = 8192;

1pub const max_row_pipeline_ops = 24;

1pub const max_row_pipeline_reduces = 4;

1pub fn regionsAreOpaque(op: *ir.Operation) bool {2    return isName(op.name.name, dialect_mod.AccyDialect.IterateOp.operation_name);3}

1pub const row_pipeline_threads = 256;

1pub const FusionClusterKind = accy_choir.record.dispatch.FusionClusterKind;

1const pass = @import("pass.zig");23pub const fusion_plan_analysis_name = pass.fusion_plan_analysis_name;4pub const fusion_planning_pass_name = pass.fusion_planning_pass_name;5pub const fusion_planning_pass_description = pass.fusion_planning_pass_description;6pub const ClaimKind = pass.ClaimKind;7pub const ClaimMap = pass.ClaimMap;8pub const FusionClusterKind = pass.FusionClusterKind;9pub const FusionCluster = pass.FusionCluster;10pub const FusionPlanAnalysis = pass.FusionPlanAnalysis;11pub const fusion_plan_analysis_descriptor = pass.fusion_plan_analysis_descriptor;12pub const getFusionPlanAnalysis = pass.getFusionPlanAnalysis;13pub const fusionPlanningPass = pass.fusionPlanningPass;14pub const isFlatGatherOp = pass.isFlatGatherOp;15pub const regionsAreOpaque = pass.regionsAreOpaque;16pub const isSeeThroughShapeOp = pass.isSeeThroughShapeOp;17pub const flash_attention_dim = pass.flash_attention_dim;18pub const flash_attention_min_seq = pass.flash_attention_min_seq;19pub const row_pipeline_threads = pass.row_pipeline_threads;20pub const max_row_pipeline_ops = pass.max_row_pipeline_ops;21pub const max_row_pipeline_reduces = pass.max_row_pipeline_reduces;22pub const max_row_pipeline_cols = pass.max_row_pipeline_cols;23pub const max_dot_epilogue_ops = pass.max_dot_epilogue_ops;24pub const max_reduction_prologue_ops = pass.max_reduction_prologue_ops;25pub const max_concat_reduction_group = pass.max_concat_reduction_group;26pub const broadcastLeafSource = pass.broadcastLeafSource;27pub const isFusableElementwiseOp = pass.isFusableElementwiseOp;

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