tiny.choir.backends.gpu.registration

1const std = @import("std");2const choir = @import("../../root.zig");3const gpu = @import("../../dialects/gpu/root.zig");4const spirv = @import("spirv/root.zig");5const nvptx = @import("nvptx/root.zig");6const lowering = @import("lowering.zig");78fn loadGpuDialect(ctx: *choir.ir.Context) !void {9    try choir.ir.dialects.loadDialectSpec(ctx, gpu.GpuDialect.spec);10}1112fn loadSpirvDialect(ctx: *choir.ir.Context) !void {13    try choir.ir.dialects.loadDialectSpec(ctx, spirv.SpirvDialect.spec);14}1516fn loadNvptxDialect(ctx: *choir.ir.Context) !void {17    try choir.ir.dialects.loadDialectSpec(ctx, nvptx.NvptxDialect.spec);18}1920pub const dialect_registrations = [_]choir.extensions.DialectRegistration{21    .{ .name = "gpu", .load = loadGpuDialect },22    .{ .name = "spirv", .load = loadSpirvDialect, .is_backend = true },23    .{ .name = "nvptx", .load = loadNvptxDialect, .is_backend = true },24};2526pub const package_extension = choir.extensions.PackageExtension{27    .name = "choir-gpu",28    .dialects = &dialect_registrations,29    .passes = &lowering.pass_registrations,30    .pipelines = &lowering.pipeline_registrations,31};3233pub fn registerTargetDialects(ctx: *choir.ir.Context) !void {34    for (dialect_registrations) |registration| {35        ctx.registerDialectLoader(registration.name, registration.load) catch |err| switch (err) {36            error.DuplicateDialectLoader => {},37            else => return err,38        };39        if (registration.is_backend) try ctx.registerBackendDialect(registration.name);40    }41}4243/// Load the native compilation dialect closure before an owned program Context activates.44pub fn prepareCompilationDialects(ctx: *choir.ir.Context) !void {45    try choir.dialects.registerChoirDialect(ctx);46    try registerTargetDialects(ctx);47    inline for (.{ "builtin", "arith", "memref", "scf", "func" }) |name| {48        _ = try ctx.getOrLoadDialect(name);49    }50    for (dialect_registrations) |registration| _ = try ctx.getOrLoadDialect(registration.name);51}5253test "registerTargetDialects is idempotent for kernel contexts" {54    const testing = std.testing;55    var ctx = try choir.ir.Context.init(testing.allocator, choir.ir.Context.Limits.testing);56    defer ctx.deinit(testing.allocator);5758    try registerTargetDialects(&ctx);59    try registerTargetDialects(&ctx);6061    _ = try ctx.getOrLoadDialect("gpu");62    _ = try ctx.getOrLoadDialect("spirv");63    _ = try ctx.getOrLoadDialect("nvptx");64    try testing.expect(ctx.isBackendDialect("spirv"));65    try testing.expect(ctx.isBackendDialect("nvptx"));66}

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