tiny.choir.backends.gpu.cpu.lowering

1const std = @import("std");2const abi = @import("choir_abi");3const choir = @import("../../../root.zig");4const gpu = @import("../../../dialects/gpu/root.zig");56const ir = choir.ir;7const dialects = choir.dialects;8const ArithDialect = dialects.ArithDialect;9const BuiltinDialect = dialects.BuiltinDialect;10const FuncDialect = dialects.FuncDialect;11const GpuDialect = gpu.GpuDialect;12const MemrefDialect = dialects.MemrefDialect;13const ScfDialect = dialects.ScfDialect;14const Dimension = gpu.Dimension;1516pub const LowerOptions = struct {17    entry_name: []const u8,18    vector_width: ?u32 = null,19};2021pub fn lowerKernelToHostLoop(22    allocator: std.mem.Allocator,23    kernel_module: *ir.Operation,24    options: LowerOptions,25) abi.Error!*ir.Operation {26    const source = ir.inspection.functionDefinitionByName(kernel_module, options.entry_name) orelse return error.InvalidArtifact;27    const source_func = FuncDialect.FuncOp{ .op = source };28    if (source_func.getNumResults() != 0) return error.UnsupportedOperation;2930    const host_module = BuiltinDialect.ModuleOp.create(kernel_module.getContext(), source.location) catch |err| return loweringError(err);31    errdefer host_module.op.erase();3233    const source_args = source_func.getArguments();34    const index_type = ArithDialect.getIndexType(kernel_module.getContext()) catch |err| return loweringError(err);35    const buffer_arg_count = try countBufferArguments(kernel_module.getContext(), source_args);36    const host_arg_types = allocator.alloc(ir.Type, source_args.len + abi.launch_shape_argument_count) catch return error.OutOfMemory;37    defer allocator.free(host_arg_types);38    var buffer_arg_index: usize = 0;39    var scalar_arg_index: usize = buffer_arg_count;40    for (source_args) |argument| {41        const index = if (try isBufferType(kernel_module.getContext(), argument.type)) index: {42            const current = buffer_arg_index;43            buffer_arg_index += 1;44            break :index current;45        } else index: {46            const current = scalar_arg_index;47            scalar_arg_index += 1;48            break :index current;49        };50        host_arg_types[index] = cpuBoundaryType(kernel_module.getContext(), argument.type) catch |err| return loweringError(err);51    }52    for (0..abi.launch_shape_argument_count) |index| {53        host_arg_types[source_args.len + index] = index_type;54    }5556    const host_func = FuncDialect.FuncOp.create(57        kernel_module.getContext(),58        source.location,59        options.entry_name,60        host_arg_types,61        &.{},62    ) catch |err| return loweringError(err);63    host_module.getBodyBlock().addOperation(host_func.op) catch |err| return loweringError(err);6465    const entry = host_func.getEntryBlock();66    const host_args = host_func.getArguments();67    const mapped_source_args = allocator.alloc(*ir.Value, source_args.len) catch return error.OutOfMemory;68    defer allocator.free(mapped_source_args);69    buffer_arg_index = 0;70    scalar_arg_index = buffer_arg_count;71    for (source_args, 0..) |source_arg, source_index| {72        const index = if (try isBufferType(kernel_module.getContext(), source_arg.type)) index: {73            const current = buffer_arg_index;74            buffer_arg_index += 1;75            break :index current;76        } else index: {77            const current = scalar_arg_index;78            scalar_arg_index += 1;79            break :index current;80        };81        const boundary_value = try cpuBoundaryValue(kernel_module.getContext(), source.location, entry, source_arg.type, host_args[index]);82        mapped_source_args[source_index] = boundary_value;83    }8485    const source_block = source_func.getEntryBlock();86    const vectorization = try vectorizationForBlock(allocator, kernel_module.getContext(), source_block, options.vector_width);8788    const zero = createIndexConstant(kernel_module.getContext(), source.location, 0) catch |err| return loweringError(err);89    entry.addOperation(zero.op) catch |err| return loweringError(err);90    const one = createIndexConstant(kernel_module.getContext(), source.location, 1) catch |err| return loweringError(err);91    entry.addOperation(one.op) catch |err| return loweringError(err);9293    const upper = host_args[source_args.len];94    const grid_dims = [3]*ir.Value{95        host_args[source_args.len + 1],96        host_args[source_args.len + 2],97        host_args[source_args.len + 3],98    };99    const threadgroup_dims = [3]*ir.Value{100        host_args[source_args.len + 4],101        host_args[source_args.len + 5],102        host_args[source_args.len + 6],103    };104    const gpu_usage = collectGpuIndexUsage(source_block);105106    if (vectorization) |plan| {107        const vector_step = createIndexConstant(kernel_module.getContext(), source.location, plan.width) catch |err| return loweringError(err);108        entry.addOperation(vector_step.op) catch |err| return loweringError(err);109        const remainder = ArithDialect.RemOp.create(kernel_module.getContext(), source.location, upper, vector_step.getResult()) catch |err| return loweringError(err);110        entry.addOperation(remainder.op) catch |err| return loweringError(err);111        const vector_upper = ArithDialect.SubOp.create(kernel_module.getContext(), source.location, upper, remainder.getResult()) catch |err| return loweringError(err);112        entry.addOperation(vector_upper.op) catch |err| return loweringError(err);113114        const vector_loop = ScfDialect.ForOp.create(115            kernel_module.getContext(),116            source.location,117            zero.getResult(),118            vector_upper.getResult(),119            vector_step.getResult(),120            &.{},121            &.{},122        ) catch |err| return loweringError(err);123        entry.addOperation(vector_loop.op) catch |err| return loweringError(err);124        try cloneKernelBlockIntoLoop(125            allocator,126            kernel_module.getContext(),127            source.location,128            source_block,129            source_args,130            mapped_source_args,131            vector_loop,132            grid_dims,133            threadgroup_dims,134            gpu_usage,135            plan,136        );137138        const tail_loop = ScfDialect.ForOp.create(139            kernel_module.getContext(),140            source.location,141            vector_upper.getResult(),142            upper,143            one.getResult(),144            &.{},145            &.{},146        ) catch |err| return loweringError(err);147        entry.addOperation(tail_loop.op) catch |err| return loweringError(err);148        try cloneKernelBlockIntoLoop(149            allocator,150            kernel_module.getContext(),151            source.location,152            source_block,153            source_args,154            mapped_source_args,155            tail_loop,156            grid_dims,157            threadgroup_dims,158            gpu_usage,159            null,160        );161    } else {162        const loop = ScfDialect.ForOp.create(163            kernel_module.getContext(),164            source.location,165            zero.getResult(),166            upper,167            one.getResult(),168            &.{},169            &.{},170        ) catch |err| return loweringError(err);171        entry.addOperation(loop.op) catch |err| return loweringError(err);172        try cloneKernelBlockIntoLoop(173            allocator,174            kernel_module.getContext(),175            source.location,176            source_block,177            source_args,178            mapped_source_args,179            loop,180            grid_dims,181            threadgroup_dims,182            gpu_usage,183            null,184        );185    }186187    const ret = FuncDialect.ReturnOp.create(kernel_module.getContext(), source.location, &.{}) catch |err| return loweringError(err);188    entry.addOperation(ret.op) catch |err| return loweringError(err);189190    ir.verifyOperation(host_module.op, ir.verify.default_options) catch return error.CompilationFailed;191    return host_module.op;192}193194fn cloneKernelBlockIntoLoop(195    allocator: std.mem.Allocator,196    ctx: *ir.Context,197    loc: ir.Location,198    source_block: *ir.Block,199    source_args: []*ir.Value,200    mapped_source_args: []*ir.Value,201    loop: ScfDialect.ForOp,202    grid_dims: [3]*ir.Value,203    threadgroup_dims: [3]*ir.Value,204    gpu_usage: GpuIndexUsage,205    vectorization: ?VectorizationPlan,206) abi.Error!void {207    var mapping = ir.Mapping.init(allocator);208    defer mapping.deinit();209    try seedSourceArgumentMapping(&mapping, source_args, mapped_source_args);210211    const loop_body = loop.getBodyBlock();212    mapping.mapBlock(source_block, loop_body) catch return error.OutOfMemory;213    const gpu_values = try createGpuIndexValues(214        ctx,215        loc,216        loop_body,217        loop.getInductionVar(),218        grid_dims,219        threadgroup_dims,220        gpu_usage,221    );222    if (vectorization) |plan| {223        try cloneBlockIntoVectorHostLoop(224            source_block,225            loop_body,226            &mapping,227            gpu_values,228            plan,229            true,230        );231    } else {232        try cloneBlockIntoHostLoop(233            source_block,234            loop_body,235            &mapping,236            gpu_values,237            true,238        );239    }240}241242fn seedSourceArgumentMapping(243    mapping: *ir.Mapping,244    source_args: []*ir.Value,245    mapped_source_args: []*ir.Value,246) abi.Error!void {247    if (source_args.len != mapped_source_args.len) return error.InvalidArtifact;248    for (source_args, mapped_source_args) |source_arg, mapped_source_arg| {249        mapping.mapValue(source_arg, mapped_source_arg) catch return error.OutOfMemory;250    }251}252253fn createIndexConstant(ctx: *ir.Context, loc: ir.Location, value: i64) !ArithDialect.ConstantOp {254    const index_type = try ArithDialect.getIndexType(ctx);255    return try ArithDialect.ConstantOp.createInt(ctx, loc, index_type, value);256}257258fn cpuBoundaryType(ctx: *ir.Context, typ: ir.Type) !ir.Type {259    if (try ctx.getTypeParamPayload(typ, MemrefDialect.MemrefTypePayload)) |payload| {260        const element_type = payload.element_type orelse return error.InvalidMemrefElementType;261        const attrs = MemrefDialect.MemrefTypeAttrs{262            .alignment = payload.alignment,263            .exclusive = payload.exclusive,264            .indexing = payload.indexing,265        };266        if (payload.size) |size| {267            return try MemrefDialect.getMemrefType1DWithAttrs(ctx, size, element_type, .host, attrs);268        }269        return try MemrefDialect.getMemrefTypeDynamicWithAttrs(ctx, element_type, .host, attrs);270    }271    return switch (cpuScalarBoundaryKind(typ) orelse return typ) {272        .f32 => try ArithDialect.getScalarType(ctx, .i32),273        .f64 => try ArithDialect.getScalarType(ctx, .i64),274        else => typ,275    };276}277278fn cpuBoundaryValue(279    ctx: *ir.Context,280    loc: ir.Location,281    entry: *ir.Block,282    source_type: ir.Type,283    boundary_value: *ir.Value,284) abi.Error!*ir.Value {285    switch (cpuScalarBoundaryKind(source_type) orelse return boundary_value) {286        .f32, .f64 => {287            const cast = ArithDialect.BitcastOp.create(ctx, loc, boundary_value, source_type) catch |err| return loweringError(err);288            entry.addOperation(cast.op) catch |err| return loweringError(err);289            return cast.getResult();290        },291        else => return boundary_value,292    }293}294295fn cpuScalarBoundaryKind(typ: ir.Type) ?dialects.arith.ScalarKind {296    const type_name = typ.getDialectTypeName() orelse return null;297    return dialects.arith.scalarKindFromTypeName(type_name);298}299300fn countBufferArguments(ctx: *ir.Context, args: []*ir.Value) abi.Error!usize {301    var count: usize = 0;302    for (args) |argument| {303        if (try isBufferType(ctx, argument.type)) count += 1;304    }305    return count;306}307308fn isBufferType(ctx: *ir.Context, typ: ir.Type) abi.Error!bool {309    return (ctx.getTypeParamPayload(typ, MemrefDialect.MemrefTypePayload) catch |err| return loweringError(err)) != null;310}311312const VectorizationPlan = struct {313    width: u32,314};315316fn vectorizationForBlock(317    allocator: std.mem.Allocator,318    ctx: *ir.Context,319    block: *ir.Block,320    width: ?u32,321) abi.Error!?VectorizationPlan {322    const requested = width orelse return null;323    if (requested < 2) return null;324    if (!try canVectorizeBlock(allocator, ctx, block, requested)) return null;325    return .{ .width = requested };326}327328fn canVectorizeBlock(329    allocator: std.mem.Allocator,330    ctx: *ir.Context,331    block: *ir.Block,332    width: u32,333) abi.Error!bool {334    var index_values = std.AutoHashMapUnmanaged(*const ir.Value, void){};335    defer index_values.deinit(allocator);336    var vector_values = std.AutoHashMapUnmanaged(*const ir.Value, void){};337    defer vector_values.deinit(allocator);338    var scalar_values = std.AutoHashMapUnmanaged(*const ir.Value, void){};339    defer scalar_values.deinit(allocator);340    var saw_load = false;341    var saw_store = false;342343    for (block.arguments.items) |argument| {344        if ((try packedCpuVectorTypeForScalarType(ctx, argument.type, width)) != null) {345            scalar_values.put(allocator, argument, {}) catch return error.OutOfMemory;346        }347    }348349    var iter = block.getOperations();350    while (iter.next()) |op| {351        if (std.mem.eql(u8, op.name.name, FuncDialect.ReturnOp.operation_name)) continue;352353        if (std.mem.eql(u8, op.name.name, GpuDialect.GlobalIdxOp.operation_name)) {354            if ((gpuDimension(op) orelse return false) != .x) return false;355            const result = op.getResult(0) orelse return false;356            index_values.put(allocator, result, {}) catch return error.OutOfMemory;357            continue;358        }359360        if (std.mem.eql(u8, op.name.getDialectNamespace(), GpuDialect.name)) return false;361        if (op.successors.items.len != 0 or op.regions.items.len != 0) return false;362363        if (std.mem.eql(u8, op.name.name, MemrefDialect.LoadOp.operation_name)) {364            const load = MemrefDialect.LoadOp{ .op = op };365            if (!index_values.contains(load.getIndex())) return false;366            _ = try packedCpuVectorTypeForScalarType(ctx, load.getResult().type, width) orelse return false;367            vector_values.put(allocator, load.getResult(), {}) catch return error.OutOfMemory;368            saw_load = true;369            continue;370        }371372        if (std.mem.eql(u8, op.name.name, MemrefDialect.StoreOp.operation_name)) {373            const store = MemrefDialect.StoreOp{ .op = op };374            if (!index_values.contains(store.getIndex())) return false;375            if (!vector_values.contains(store.getValue())) return false;376            saw_store = true;377            continue;378        }379380        if (std.mem.eql(u8, op.name.name, ArithDialect.ConstantOp.operation_name)) {381            const result = op.getResult(0) orelse return false;382            if ((try packedCpuVectorTypeForScalarType(ctx, result.type, width)) == null) return false;383            scalar_values.put(allocator, result, {}) catch return error.OutOfMemory;384            continue;385        }386387        if (isVectorizableBinaryOp(op)) {388            const lhs = op.getOperand(0) orelse return false;389            const rhs = op.getOperand(1) orelse return false;390            const result = op.getResult(0) orelse return false;391            _ = try packedCpuVectorTypeForScalarType(ctx, result.type, width) orelse return false;392            const lhs_vector = vector_values.contains(lhs);393            const rhs_vector = vector_values.contains(rhs);394            if (!lhs_vector and !scalar_values.contains(lhs)) return false;395            if (!rhs_vector and !scalar_values.contains(rhs)) return false;396            if (lhs_vector or rhs_vector) {397                vector_values.put(allocator, result, {}) catch return error.OutOfMemory;398            } else {399                scalar_values.put(allocator, result, {}) catch return error.OutOfMemory;400            }401            continue;402        }403404        if (isVectorizableUnaryOp(op)) {405            const input = op.getOperand(0) orelse return false;406            const result = op.getResult(0) orelse return false;407            _ = try packedCpuVectorTypeForScalarType(ctx, result.type, width) orelse return false;408            if (vector_values.contains(input)) {409                vector_values.put(allocator, result, {}) catch return error.OutOfMemory;410            } else if (scalar_values.contains(input)) {411                scalar_values.put(allocator, result, {}) catch return error.OutOfMemory;412            } else {413                return false;414            }415            continue;416        }417418        return false;419    }420421    return saw_load and saw_store;422}423424fn cloneBlockIntoHostLoop(425    source: *ir.Block,426    dest: *ir.Block,427    mapping: *ir.Mapping,428    gpu_values: GpuIndexValues,429    function_body: bool,430) abi.Error!void {431    var iter = source.getOperations();432    while (iter.next()) |op| {433        if (std.mem.eql(u8, op.name.name, FuncDialect.ReturnOp.operation_name)) {434            if (!function_body) return error.UnsupportedOperation;435            const yield_op = ScfDialect.YieldOp.create(op.getContext(), op.location, &.{}) catch |err| return loweringError(err);436            dest.addOperation(yield_op.op) catch |err| return loweringError(err);437            continue;438        }439440        if (std.mem.eql(u8, op.name.name, GpuDialect.GlobalIdxOp.operation_name)) {441            const global_index = gpu_values.global.value(gpuDimension(op) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;442            const result = op.getResult(0) orelse return error.UnsupportedOperation;443            mapping.mapValue(result, global_index) catch return error.OutOfMemory;444            continue;445        }446447        if (std.mem.eql(u8, op.name.name, GpuDialect.ThreadIdxOp.operation_name)) {448            const thread_index = gpu_values.thread.value(gpuDimension(op) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;449            const result = op.getResult(0) orelse return error.UnsupportedOperation;450            mapping.mapValue(result, thread_index) catch return error.OutOfMemory;451            continue;452        }453454        if (std.mem.eql(u8, op.name.name, GpuDialect.BlockIdxOp.operation_name)) {455            const block_index = gpu_values.block.value(gpuDimension(op) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;456            const result = op.getResult(0) orelse return error.UnsupportedOperation;457            mapping.mapValue(result, block_index) catch return error.OutOfMemory;458            continue;459        }460461        if (std.mem.eql(u8, op.name.name, GpuDialect.BlockDimOp.operation_name)) {462            const block_dim = gpu_values.threadgroupValue(gpuDimension(op) orelse return error.UnsupportedOperation);463            const result = op.getResult(0) orelse return error.UnsupportedOperation;464            mapping.mapValue(result, block_dim) catch return error.OutOfMemory;465            continue;466        }467468        if (std.mem.eql(u8, op.name.name, GpuDialect.GridDimOp.operation_name)) {469            const grid_dim = gpu_values.gridValue(gpuDimension(op) orelse return error.UnsupportedOperation);470            const result = op.getResult(0) orelse return error.UnsupportedOperation;471            mapping.mapValue(result, grid_dim) catch return error.OutOfMemory;472            continue;473        }474475        if (std.mem.eql(u8, op.name.getDialectNamespace(), GpuDialect.name)) {476            return error.UnsupportedOperation;477        }478479        if (op.successors.items.len != 0) return error.UnsupportedOperation;480481        const has_regions = op.regions.items.len != 0;482        const cloned = op.cloneWithoutRegionsMapped(mapping, .{ .clone_operands = !has_regions }) catch |err| return loweringError(err);483        errdefer cloned.erase();484        dest.addOperation(cloned) catch |err| return loweringError(err);485        for (op.regions.items, 0..) |*source_region, region_index| {486            const dest_region = &cloned.regions.items[region_index];487            var source_block = source_region.blocks.head;488            while (source_block) |block| : (source_block = block.next) {489                const cloned_block = dest_region.addBlock() catch |err| return loweringError(err);490                mapping.mapBlock(block, cloned_block) catch return error.OutOfMemory;491                cloned_block.arguments.ensureTotalCapacity(cloned_block.allocator, block.arguments.items.len) catch return error.OutOfMemory;492                for (block.arguments.items) |argument| {493                    const cloned_argument = cloned_block.addArgument(argument.type, .unknown) catch |err| return loweringError(err);494                    mapping.mapValue(argument, cloned_argument) catch return error.OutOfMemory;495                }496            }497498            source_block = source_region.blocks.head;499            while (source_block) |block| : (source_block = block.next) {500                const cloned_block = mapping.lookupBlock(block) orelse return error.UnsupportedOperation;501                try cloneBlockIntoHostLoop(block, cloned_block, mapping, gpu_values, false);502            }503        }504        if (has_regions) try replaceOperandsMapped(op, cloned, mapping);505    }506}507508fn cloneBlockIntoVectorHostLoop(509    source: *ir.Block,510    dest: *ir.Block,511    mapping: *ir.Mapping,512    gpu_values: GpuIndexValues,513    plan: VectorizationPlan,514    function_body: bool,515) abi.Error!void {516    var iter = source.getOperations();517    while (iter.next()) |op| {518        if (std.mem.eql(u8, op.name.name, FuncDialect.ReturnOp.operation_name)) {519            if (!function_body) return error.UnsupportedOperation;520            const yield_op = ScfDialect.YieldOp.create(op.getContext(), op.location, &.{}) catch |err| return loweringError(err);521            dest.addOperation(yield_op.op) catch |err| return loweringError(err);522            continue;523        }524525        if (std.mem.eql(u8, op.name.name, GpuDialect.GlobalIdxOp.operation_name)) {526            if ((gpuDimension(op) orelse return error.UnsupportedOperation) != .x) return error.UnsupportedOperation;527            const global_index = gpu_values.global.x orelse return error.UnsupportedOperation;528            const result = op.getResult(0) orelse return error.UnsupportedOperation;529            mapping.mapValue(result, global_index) catch return error.OutOfMemory;530            continue;531        }532533        if (std.mem.eql(u8, op.name.name, MemrefDialect.LoadOp.operation_name)) {534            try cloneVectorLoad(op, dest, mapping, plan.width);535            continue;536        }537538        if (std.mem.eql(u8, op.name.name, MemrefDialect.StoreOp.operation_name)) {539            try cloneVectorStore(op, dest, mapping);540            continue;541        }542543        if (std.mem.eql(u8, op.name.name, ArithDialect.ConstantOp.operation_name)) {544            try cloneScalarOperationIntoVectorLoop(op, dest, mapping);545            continue;546        }547548        if (isVectorizableBinaryOp(op)) {549            try cloneVectorBinaryOp(op, dest, mapping, plan.width);550            continue;551        }552553        if (isVectorizableUnaryOp(op)) {554            try cloneVectorUnaryOp(op, dest, mapping, plan.width);555            continue;556        }557558        return error.UnsupportedOperation;559    }560}561562fn cloneVectorLoad(563    op: *ir.Operation,564    dest: *ir.Block,565    mapping: *ir.Mapping,566    width: u32,567) abi.Error!void {568    const load = MemrefDialect.LoadOp{ .op = op };569    const memref = mapping.lookupValue(load.getMemref()) orelse return error.UnsupportedOperation;570    const index = mapping.lookupValue(load.getIndex()) orelse return error.UnsupportedOperation;571    const result_type = try packedCpuVectorTypeForScalarType(op.getContext(), load.getResult().type, width) orelse return error.UnsupportedOperation;572    var cloned = MemrefDialect.LoadOp.create(op.getContext(), op.location, memref, index, result_type) catch |err| return loweringError(err);573    dest.addOperation(cloned.op) catch |err| return loweringError(err);574    mapping.mapValue(load.getResult(), cloned.getResult()) catch return error.OutOfMemory;575}576577fn cloneVectorStore(578    op: *ir.Operation,579    dest: *ir.Block,580    mapping: *ir.Mapping,581) abi.Error!void {582    const store = MemrefDialect.StoreOp{ .op = op };583    const value = mapping.lookupValue(store.getValue()) orelse return error.UnsupportedOperation;584    const memref = mapping.lookupValue(store.getMemref()) orelse return error.UnsupportedOperation;585    const index = mapping.lookupValue(store.getIndex()) orelse return error.UnsupportedOperation;586    const cloned = MemrefDialect.StoreOp.create(op.getContext(), op.location, value, memref, index) catch |err| return loweringError(err);587    dest.addOperation(cloned.op) catch |err| return loweringError(err);588}589590fn cloneVectorBinaryOp(591    op: *ir.Operation,592    dest: *ir.Block,593    mapping: *ir.Mapping,594    width: u32,595) abi.Error!void {596    const lhs = mapping.lookupValue(op.getOperand(0) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;597    const rhs = mapping.lookupValue(op.getOperand(1) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;598    const source_result = op.getResult(0) orelse return error.UnsupportedOperation;599    const result_type = try packedCpuVectorTypeForScalarType(op.getContext(), source_result.type, width) orelse return error.UnsupportedOperation;600    if (!lhs.type.eql(result_type) and !rhs.type.eql(result_type)) {601        try cloneScalarOperationIntoVectorLoop(op, dest, mapping);602        return;603    }604605    const vector_lhs = try vectorOperandForMappedValue(op, dest, lhs, result_type, width);606    const vector_rhs = try vectorOperandForMappedValue(op, dest, rhs, result_type, width);607608    var state = ir.Operation.State.init(op.name.name, op.location);609    state.addOperands(&.{ vector_lhs, vector_rhs });610    state.addTypes(&.{result_type});611    const cloned = op.getContext().createOperation(state) catch |err| return loweringError(err);612    errdefer cloned.erase();613    dest.addOperation(cloned) catch |err| return loweringError(err);614    const result = cloned.getResult(0) orelse return error.UnsupportedOperation;615    mapping.mapValue(source_result, result) catch return error.OutOfMemory;616}617618fn cloneVectorUnaryOp(619    op: *ir.Operation,620    dest: *ir.Block,621    mapping: *ir.Mapping,622    width: u32,623) abi.Error!void {624    const input = mapping.lookupValue(op.getOperand(0) orelse return error.UnsupportedOperation) orelse return error.UnsupportedOperation;625    const source_result = op.getResult(0) orelse return error.UnsupportedOperation;626    const result_type = try packedCpuVectorTypeForScalarType(op.getContext(), source_result.type, width) orelse return error.UnsupportedOperation;627    if (!input.type.eql(result_type)) {628        try cloneScalarOperationIntoVectorLoop(op, dest, mapping);629        return;630    }631632    var state = ir.Operation.State.init(op.name.name, op.location);633    state.addOperands(&.{input});634    state.addTypes(&.{result_type});635    const cloned = op.getContext().createOperation(state) catch |err| return loweringError(err);636    errdefer cloned.erase();637    dest.addOperation(cloned) catch |err| return loweringError(err);638    const result = cloned.getResult(0) orelse return error.UnsupportedOperation;639    mapping.mapValue(source_result, result) catch return error.OutOfMemory;640}641642fn cloneScalarOperationIntoVectorLoop(643    op: *ir.Operation,644    dest: *ir.Block,645    mapping: *ir.Mapping,646) abi.Error!void {647    const cloned = op.cloneWithoutRegionsMapped(mapping, .{ .clone_operands = true }) catch |err| return loweringError(err);648    errdefer cloned.erase();649    dest.addOperation(cloned) catch |err| return loweringError(err);650}651652fn vectorOperandForMappedValue(653    op: *ir.Operation,654    dest: *ir.Block,655    mapped: *ir.Value,656    result_type: ir.Type,657    width: u32,658) abi.Error!*ir.Value {659    if (mapped.type.eql(result_type)) return mapped;660    const vector_type = try packedCpuVectorTypeForScalarType(op.getContext(), mapped.type, width) orelse return error.UnsupportedOperation;661    if (!vector_type.eql(result_type)) return error.UnsupportedOperation;662    const splat = ArithDialect.SplatOp.create(op.getContext(), op.location, mapped, result_type) catch |err| return loweringError(err);663    dest.addOperation(splat.op) catch |err| return loweringError(err);664    return splat.getResult();665}666667fn isVectorizableBinaryOp(op: *const ir.Operation) bool {668    return std.mem.eql(u8, op.name.name, ArithDialect.AddOp.operation_name) or669        std.mem.eql(u8, op.name.name, ArithDialect.SubOp.operation_name) or670        std.mem.eql(u8, op.name.name, ArithDialect.MulOp.operation_name) or671        std.mem.eql(u8, op.name.name, ArithDialect.DivOp.operation_name) or672        std.mem.eql(u8, op.name.name, ArithDialect.MinOp.operation_name) or673        std.mem.eql(u8, op.name.name, ArithDialect.MaxOp.operation_name) or674        std.mem.eql(u8, op.name.name, ArithDialect.AndOp.operation_name) or675        std.mem.eql(u8, op.name.name, ArithDialect.OrOp.operation_name) or676        std.mem.eql(u8, op.name.name, ArithDialect.XorOp.operation_name) or677        std.mem.eql(u8, op.name.name, ArithDialect.ShlOp.operation_name) or678        std.mem.eql(u8, op.name.name, ArithDialect.ShrOp.operation_name) or679        std.mem.eql(u8, op.name.name, ArithDialect.UshrOp.operation_name);680}681682fn isVectorizableUnaryOp(op: *const ir.Operation) bool {683    return std.mem.eql(u8, op.name.name, ArithDialect.NegOp.operation_name) or684        std.mem.eql(u8, op.name.name, ArithDialect.NotOp.operation_name) or685        std.mem.eql(u8, op.name.name, ArithDialect.PopCountOp.operation_name);686}687688fn packedCpuVectorTypeForScalarType(ctx: *ir.Context, scalar_type: ir.Type, width: u32) abi.Error!?ir.Type {689    const scalar_name = scalar_type.getDialectTypeName() orelse return null;690    const scalar_kind = dialects.arith.scalarKindFromTypeName(scalar_name) orelse return null;691    switch (scalar_kind) {692        .f32, .i32, .u32 => if (width != 4) return null,693        .f64, .i64, .u64 => if (width != 2) return null,694        else => return null,695    }696    const vector_name = dialects.arith.vectorTypeNameForElement(width, scalar_name) orelse return null;697    return ctx.getDialectTypeFromName(vector_name) catch |err| return loweringError(err);698}699700fn replaceOperandsMapped(source: *ir.Operation, dest: *ir.Operation, mapping: *ir.Mapping) abi.Error!void {701    var operands: []*ir.Value = &.{};702    defer if (operands.len != 0) dest.allocator.free(operands);703    if (source.operand_values.len != 0) {704        operands = dest.allocator.alloc(*ir.Value, source.operand_values.len) catch return error.OutOfMemory;705        for (source.operand_values, 0..) |operand, index| {706            operands[index] = mapping.lookupOrDefaultValue(operand);707        }708    }709    dest.replaceOperands(operands) catch |err| return loweringError(err);710}711712const DimensionUsage = struct {713    x: bool = false,714    y: bool = false,715    z: bool = false,716717    fn add(self: *DimensionUsage, dim: Dimension) void {718        switch (dim) {719            .x => self.x = true,720            .y => self.y = true,721            .z => self.z = true,722        }723    }724725    fn addAll(self: *DimensionUsage, other: DimensionUsage) void {726        self.x = self.x or other.x;727        self.y = self.y or other.y;728        self.z = self.z or other.z;729    }730731    fn contains(self: DimensionUsage, dim: Dimension) bool {732        return switch (dim) {733            .x => self.x,734            .y => self.y,735            .z => self.z,736        };737    }738};739740const GpuIndexUsage = struct {741    global: DimensionUsage = .{},742    thread: DimensionUsage = .{},743    block: DimensionUsage = .{},744    block_dim: DimensionUsage = .{},745    grid_dim: DimensionUsage = .{},746747    fn merge(self: *GpuIndexUsage, other: GpuIndexUsage) void {748        self.global.addAll(other.global);749        self.thread.addAll(other.thread);750        self.block.addAll(other.block);751        self.block_dim.addAll(other.block_dim);752        self.grid_dim.addAll(other.grid_dim);753    }754755    fn derived(self: GpuIndexUsage) DimensionUsage {756        var result = DimensionUsage{};757        result.addAll(self.global);758        result.addAll(self.thread);759        result.addAll(self.block);760        return result;761    }762};763764const DimensionValues = struct {765    x: ?*ir.Value = null,766    y: ?*ir.Value = null,767    z: ?*ir.Value = null,768769    fn value(self: DimensionValues, dim: Dimension) ?*ir.Value {770        return switch (dim) {771            .x => self.x,772            .y => self.y,773            .z => self.z,774        };775    }776777    fn set(self: *DimensionValues, dim: Dimension, new_value: *ir.Value) void {778        switch (dim) {779            .x => self.x = new_value,780            .y => self.y = new_value,781            .z => self.z = new_value,782        }783    }784};785786const GpuIndexValues = struct {787    grid: [3]*ir.Value,788    threadgroup: [3]*ir.Value,789    global: DimensionValues = .{},790    thread: DimensionValues = .{},791    block: DimensionValues = .{},792793    fn gridValue(self: GpuIndexValues, dim: Dimension) *ir.Value {794        return self.grid[dimensionIndex(dim)];795    }796797    fn threadgroupValue(self: GpuIndexValues, dim: Dimension) *ir.Value {798        return self.threadgroup[dimensionIndex(dim)];799    }800};801802fn collectGpuIndexUsage(block: *ir.Block) GpuIndexUsage {803    var usage = GpuIndexUsage{};804    var iter = block.getOperations();805    while (iter.next()) |op| {806        if (std.mem.eql(u8, op.name.name, GpuDialect.GlobalIdxOp.operation_name)) {807            if (gpuDimension(op)) |dim| usage.global.add(dim);808        } else if (std.mem.eql(u8, op.name.name, GpuDialect.ThreadIdxOp.operation_name)) {809            if (gpuDimension(op)) |dim| usage.thread.add(dim);810        } else if (std.mem.eql(u8, op.name.name, GpuDialect.BlockIdxOp.operation_name)) {811            if (gpuDimension(op)) |dim| usage.block.add(dim);812        } else if (std.mem.eql(u8, op.name.name, GpuDialect.BlockDimOp.operation_name)) {813            if (gpuDimension(op)) |dim| usage.block_dim.add(dim);814        } else if (std.mem.eql(u8, op.name.name, GpuDialect.GridDimOp.operation_name)) {815            if (gpuDimension(op)) |dim| usage.grid_dim.add(dim);816        }817        for (op.regions.items) |*region| {818            var block_iter = region.getBlocks();819            while (block_iter.next()) |child_block| {820                usage.merge(collectGpuIndexUsage(child_block));821            }822        }823    }824    return usage;825}826827fn createGpuIndexValues(828    ctx: *ir.Context,829    loc: ir.Location,830    dest: *ir.Block,831    linear: *ir.Value,832    grid: [3]*ir.Value,833    threadgroup: [3]*ir.Value,834    usage: GpuIndexUsage,835) abi.Error!GpuIndexValues {836    var values = GpuIndexValues{837        .grid = grid,838        .threadgroup = threadgroup,839    };840    const derived = usage.derived();841    var y_linear: ?*ir.Value = null;842    var x_extent: ?*ir.Value = null;843    var y_extent: ?*ir.Value = null;844845    if (derived.x or derived.y or derived.z) {846        x_extent = try createProduct(ctx, loc, dest, grid[0], threadgroup[0]);847    }848    if (derived.y or derived.z) {849        y_extent = try createProduct(ctx, loc, dest, grid[1], threadgroup[1]);850    }851852    if (derived.x) {853        const x = ArithDialect.RemOp.create(ctx, loc, linear, x_extent.?) catch |err| return loweringError(err);854        dest.addOperation(x.op) catch |err| return loweringError(err);855        values.global.x = x.getResult();856        try createThreadBlockValues(ctx, loc, dest, &values, .x, x.getResult());857    }858859    if (derived.y or derived.z) {860        const div = ArithDialect.DivOp.create(ctx, loc, linear, x_extent.?) catch |err| return loweringError(err);861        dest.addOperation(div.op) catch |err| return loweringError(err);862        y_linear = div.getResult();863    }864865    if (derived.y) {866        const y = ArithDialect.RemOp.create(ctx, loc, y_linear.?, y_extent.?) catch |err| return loweringError(err);867        dest.addOperation(y.op) catch |err| return loweringError(err);868        values.global.y = y.getResult();869        try createThreadBlockValues(ctx, loc, dest, &values, .y, y.getResult());870    }871872    if (derived.z) {873        const xy_extent = try createProduct(ctx, loc, dest, x_extent.?, y_extent.?);874        const z = ArithDialect.DivOp.create(ctx, loc, linear, xy_extent) catch |err| return loweringError(err);875        dest.addOperation(z.op) catch |err| return loweringError(err);876        values.global.z = z.getResult();877        try createThreadBlockValues(ctx, loc, dest, &values, .z, z.getResult());878    }879880    return values;881}882883fn createThreadBlockValues(884    ctx: *ir.Context,885    loc: ir.Location,886    dest: *ir.Block,887    values: *GpuIndexValues,888    dim: Dimension,889    global_value: *ir.Value,890) abi.Error!void {891    const thread_extent = values.threadgroupValue(dim);892    const thread = ArithDialect.RemOp.create(ctx, loc, global_value, thread_extent) catch |err| return loweringError(err);893    dest.addOperation(thread.op) catch |err| return loweringError(err);894    values.thread.set(dim, thread.getResult());895896    const block = ArithDialect.DivOp.create(ctx, loc, global_value, thread_extent) catch |err| return loweringError(err);897    dest.addOperation(block.op) catch |err| return loweringError(err);898    values.block.set(dim, block.getResult());899}900901fn createProduct(902    ctx: *ir.Context,903    loc: ir.Location,904    dest: *ir.Block,905    lhs: *ir.Value,906    rhs: *ir.Value,907) abi.Error!*ir.Value {908    const product = ArithDialect.MulOp.create(ctx, loc, lhs, rhs) catch |err| return loweringError(err);909    dest.addOperation(product.op) catch |err| return loweringError(err);910    return product.getResult();911}912913fn dimensionIndex(dim: Dimension) usize {914    return switch (dim) {915        .x => 0,916        .y => 1,917        .z => 2,918    };919}920921fn gpuDimension(op: *const ir.Operation) ?Dimension {922    const attr = op.getAttrAs(ir.Attribute.DialectAttr, "dim") orelse return null;923    return Dimension.fromString(attr.payload);924}925926fn loweringError(err: anyerror) abi.Error {927    return switch (err) {928        error.OutOfMemory => error.OutOfMemory,929        else => error.CompilationFailed,930    };931}

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