tiny.accy.preparation.call

1const std = @import("std");2const choir_abi = @import("choir_abi");3const choir = @import("choir");4const accy_choir = @import("../choir/root.zig");5const kernel_library = @import("../kernel/library/root.zig");67const ir = choir.ir;8const rewrite = ir.rewrite;9const dialects = choir.dialects;10const dialect_mod = accy_choir.dialect;11const semantic = accy_choir.semantic;1213pub const CatalogCallOptions = struct {14    has_side_effects: bool = false,15    operand_effects: ?[]const semantic.KernelOperandEffect = null,16    result_aliases: ?[]const ?usize = null,17    runtime_scalars: ?[]const dialect_mod.AccyDialect.KernelCallScalar = null,18};1920pub fn catalogCallScalars(21    comptime count: usize,22    arguments: [count]choir_abi.ScalarArgument,23) [count]dialect_mod.AccyDialect.KernelCallScalar {24    var scalars: [count]dialect_mod.AccyDialect.KernelCallScalar = undefined;25    for (arguments, 0..) |argument, index| {26        scalars[index] = switch (argument) {27            .i32 => |value| .{ .kind = .i32, .bits = @as(u32, @bitCast(value)) },28            .u32 => |value| .{ .kind = .u32, .bits = value },29            .i64 => |value| .{ .kind = .i64, .bits = @bitCast(value) },30            .u64 => |value| .{ .kind = .u64, .bits = value },31            .f32 => |value| .{ .kind = .f32, .bits = @as(u32, @bitCast(value)) },32            .f64 => |value| .{ .kind = .f64, .bits = @bitCast(value) },33        };34    }35    return scalars;36}3738pub const CatalogCallRequest = struct {39    descriptor: kernel_library.CatalogDescriptor,40    operands: []const *ir.Value,41    result_types: []const ir.Type,42    options: CatalogCallOptions = .{},43};4445pub fn insertCatalogCall(46    rewriter: *rewrite.PatternRewriter,47    request: CatalogCallRequest,48) !dialect_mod.AccyDialect.KernelCallOp {49    const operand_effects = try kernelOperandEffects(rewriter.allocator, request.operands.len, request.options.operand_effects);50    defer if (request.options.operand_effects == null) rewriter.allocator.free(operand_effects);5152    const result_aliases = try kernelResultAliases(rewriter.allocator, request.result_types.len, request.options.result_aliases);53    defer if (request.options.result_aliases == null) rewriter.allocator.free(result_aliases);5455    const metadata = request.descriptor.metadata;56    const call = try dialect_mod.AccyDialect.KernelCallOp.create(57        rewriter.ir_ctx,58        ir.Location.getUnknown(),59        request.operands,60        request.result_types,61        metadata.target,62        metadata.version,63        request.options.has_side_effects,64        operand_effects,65        result_aliases,66    );67    if (request.options.runtime_scalars) |scalars| {68        try dialect_mod.AccyDialect.setKernelCallRuntimeScalars(rewriter.ir_ctx, call.op, scalars);69    }70    _ = try rewriter.insert(call.op);71    return call;72}7374fn kernelOperandEffects(75    allocator: std.mem.Allocator,76    operand_count: usize,77    provided: ?[]const semantic.KernelOperandEffect,78) ![]const semantic.KernelOperandEffect {79    if (provided) |effects| {80        if (effects.len != operand_count) return error.InvalidKernelCallContract;81        return effects;82    }8384    const effects = try allocator.alloc(semantic.KernelOperandEffect, operand_count);85    @memset(effects, .read);86    return effects;87}8889fn kernelResultAliases(90    allocator: std.mem.Allocator,91    result_count: usize,92    provided: ?[]const ?usize,93) ![]const ?usize {94    if (provided) |aliases| {95        if (aliases.len != result_count) return error.InvalidKernelCallContract;96        return aliases;97    }9899    const aliases = try allocator.alloc(?usize, result_count);100    @memset(aliases, null);101    return aliases;102}103104fn findOpNamed(op: *ir.Operation, name: []const u8) ?*ir.Operation {105    if (std.mem.eql(u8, op.name.name, name)) return op;106    for (op.regions.items) |*region| {107        var block_iter = region.getBlocks();108        while (block_iter.next()) |block| {109            var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));110            while (current) |current_op| {111                if (findOpNamed(current_op, name)) |found| return found;112                current = current_op.next_op;113            }114        }115    }116    return null;117}118119const testing = std.testing;120121test "catalog call insertion emits descriptor target and default contract" {122    const allocator = testing.allocator;123124    var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);125    defer builder.deinit();126    const f32_8 = try builder.tensor(.f32, &.{8});127    var fb = try builder.beginFunction("catalog_call_insertion", &.{ f32_8, f32_8 }, &.{f32_8});128    const lhs = fb.parameter(0);129    const rhs = fb.parameter(1);130    const out = try fb.add(lhs, rhs);131    try fb.return_(&.{out});132    try fb.finish();133    const module = try builder.finish();134    defer module.deinit();135136    const return_op = findOpNamed(module.choir_module, dialects.FuncDialect.ReturnOp.operation_name) orelse {137        return error.TestExpectedReturn;138    };139    var rewriter = rewrite.PatternRewriter.init(allocator, module.context());140    defer rewriter.deinit();141    rewriter.setInsertionPointBefore(return_op);142143    const result_types = [_]ir.Type{f32_8};144    const descriptor = kernel_library.findEntry(145        kernel_library.elementwise.VectorAdd8F32.target,146        kernel_library.elementwise.VectorAdd8F32.version,147    ) orelse return error.TestExpectedCatalogDescriptor;148    const call = try insertCatalogCall(&rewriter, .{149        .descriptor = descriptor,150        .operands = &.{ lhs, rhs },151        .result_types = &result_types,152    });153154    try module.verify();155    const target_attr = call.op.getAttr("target") orelse return error.TestExpectedTarget;156    const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedTarget;157    try testing.expectEqualStrings(dialect_mod.AccyDialect.KernelCallOp.target_attr_name, target_attr.abstract.name);158    try testing.expectEqualStrings(kernel_library.elementwise.VectorAdd8F32.target, target.payload);159160    const version = call.op.getAttrAs(ir.Attribute.IntegerAttr, "version") orelse return error.TestExpectedVersion;161    try testing.expectEqual(@as(i64, kernel_library.elementwise.VectorAdd8F32.version), version.getValue());162163    const effects_attr = call.op.getAttr("operand_effects") orelse {164        return error.TestExpectedOperandEffects;165    };166    const effects = effects_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedOperandEffects;167    try testing.expectEqualStrings(dialect_mod.AccyDialect.KernelCallOp.operand_effects_attr_name, effects_attr.abstract.name);168    try testing.expectEqual(@as(usize, 2), effects.payload.len);169    try testing.expectEqual(@backingInt(semantic.KernelOperandEffect.read), effects.payload[0]);170    try testing.expectEqual(@backingInt(semantic.KernelOperandEffect.read), effects.payload[1]);171172    const aliases_attr = call.op.getAttr("result_aliases") orelse {173        return error.TestExpectedResultAliases;174    };175    const aliases = aliases_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedResultAliases;176    try testing.expectEqualStrings(dialect_mod.AccyDialect.KernelCallOp.result_aliases_attr_name, aliases_attr.abstract.name);177    try testing.expectEqual(@as(usize, @sizeOf(i64)), aliases.payload.len);178    const values = std.mem.bytesAsSlice(i64, aliases.payload);179    try testing.expectEqual(@as(i64, -1), values[0]);180}

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