tiny.choir.backends.gpu.lowering

1const std = @import("std");2const choir = @import("../../root.zig");3const spirv = @import("spirv/root.zig");4const nvptx = @import("nvptx/root.zig");56pub const gpu_to_spirv_pipeline_name = "gpu-to-spirv-pipeline";7pub const gpu_to_spirv_pipeline_description = "Lower the GPU dialect to the SPIR-V target dialect";89pub const gpu_to_nvptx_pipeline_name = "gpu-to-nvptx-pipeline";10pub const gpu_to_nvptx_pipeline_description = "Lower the GPU dialect to the NVPTX target dialect";1112pub const TargetLowering = enum {13    spirv,14    nvptx,15};1617const Stage = struct {18    lowering: TargetLowering,19    target_spec: choir.backends.TargetSpec,20    pass_registration: choir.passes.PassRegistration,21    pipeline_registration: choir.passes.PipelineRegistration,22};2324fn buildGpuToSpirvPipeline(pm: *choir.passes.OpPassManager) anyerror!void {25    try pm.addPass(spirv.createGpuToSpirvPass());26}2728fn buildGpuToNvptxPipeline(pm: *choir.passes.OpPassManager) anyerror!void {29    try pm.addPass(nvptx.createGpuToNvptxPass());30}3132pub const gpu_to_spirv_pipeline_registration = choir.passes.PipelineRegistration{33    .name = gpu_to_spirv_pipeline_name,34    .description = gpu_to_spirv_pipeline_description,35    .build = buildGpuToSpirvPipeline,36};3738pub const gpu_to_nvptx_pipeline_registration = choir.passes.PipelineRegistration{39    .name = gpu_to_nvptx_pipeline_name,40    .description = gpu_to_nvptx_pipeline_description,41    .build = buildGpuToNvptxPipeline,42};4344fn targetSpecWithPipeline(45    spec: choir.backends.TargetSpec,46    pipeline_registration: choir.passes.PipelineRegistration,47) choir.backends.TargetSpec {48    var target_spec = spec;49    target_spec.pipeline_name = pipeline_registration.name;50    target_spec.pipeline_description = pipeline_registration.description;51    return target_spec;52}5354const stages = [_]Stage{55    .{56        .lowering = .spirv,57        .target_spec = targetSpecWithPipeline(spirv.target_spec, gpu_to_spirv_pipeline_registration),58        .pass_registration = spirv.pass_registration,59        .pipeline_registration = gpu_to_spirv_pipeline_registration,60    },61    .{62        .lowering = .nvptx,63        .target_spec = targetSpecWithPipeline(nvptx.target_spec, gpu_to_nvptx_pipeline_registration),64        .pass_registration = nvptx.pass_registration,65        .pipeline_registration = gpu_to_nvptx_pipeline_registration,66    },67};6869fn passRegistrations() [stages.len]choir.passes.PassRegistration {70    comptime {71        var registrations: [stages.len]choir.passes.PassRegistration = undefined;72        for (stages, 0..) |stage, index| {73            registrations[index] = stage.pass_registration;74        }75        return registrations;76    }77}7879fn pipelineRegistrations() [stages.len]choir.passes.PipelineRegistration {80    comptime {81        var registrations: [stages.len]choir.passes.PipelineRegistration = undefined;82        for (stages, 0..) |stage, index| {83            registrations[index] = stage.pipeline_registration;84        }85        return registrations;86    }87}8889fn stageFor(lowering: TargetLowering) Stage {90    return switch (lowering) {91        .spirv => stages[0],92        .nvptx => stages[1],93    };94}9596fn specs() [stages.len]choir.backends.TargetSpec {97    comptime {98        var registrations: [stages.len]choir.backends.TargetSpec = undefined;99        for (stages, 0..) |stage, index| {100            registrations[index] = stage.target_spec;101        }102        return registrations;103    }104}105106pub const pass_registrations = passRegistrations();107pub const pipeline_registrations = pipelineRegistrations();108pub const target_specs = specs();109pub const spirv_target_spec = target_specs[0];110pub const nvptx_target_spec = target_specs[1];111pub const target_lowering_stage_count = stages.len;112pub const target_spec_count = target_specs.len;113pub const target_lowering_pass_count = pass_registrations.len;114pub const target_lowering_pipeline_count = pipeline_registrations.len;115116pub fn targetLoweringPassName(index: usize) []const u8 {117    return stages[index].pass_registration.name;118}119120pub fn targetLoweringPipelineName(index: usize) []const u8 {121    return stages[index].pipeline_registration.name;122}123124pub fn targetLoweringTargetDialectName(index: usize) []const u8 {125    return stages[index].target_spec.target_dialect_name;126}127128pub fn targetLoweringSpec(index: usize) choir.backends.TargetSpec {129    return target_specs[index];130}131132pub fn addGpuToSpirvPipeline(pm: *choir.passes.PassManager) !void {133    try gpu_to_spirv_pipeline_registration.addTo(&pm.root);134}135136pub fn addGpuToNvptxPipeline(pm: *choir.passes.PassManager) !void {137    try gpu_to_nvptx_pipeline_registration.addTo(&pm.root);138}139140pub fn addTargetLoweringPipeline(pm: *choir.passes.PassManager, lowering: TargetLowering) !void {141    try stageFor(lowering).pipeline_registration.addTo(&pm.root);142}143144test "target lowering stage table defines registries" {145    const testing = std.testing;146147    try testing.expectEqual(stages.len, target_lowering_stage_count);148    try testing.expectEqual(stages.len, target_spec_count);149    try testing.expectEqual(stages.len, target_lowering_pass_count);150    try testing.expectEqual(stages.len, target_lowering_pipeline_count);151    try testing.expectEqual(stages.len, target_specs.len);152    try testing.expectEqual(stages.len, pass_registrations.len);153    try testing.expectEqual(stages.len, pipeline_registrations.len);154155    inline for (stages, 0..) |stage, index| {156        try testing.expectEqual(stage.lowering, stageFor(stage.lowering).lowering);157        const target_spec = targetLoweringSpec(index);158        try testing.expectEqualStrings(stage.target_spec.name, target_spec.name);159        try testing.expectEqualStrings(stage.target_spec.target_dialect_name, targetLoweringTargetDialectName(index));160        try testing.expectEqualStrings(stage.target_spec.target_dialect_name, target_spec.target_dialect_name);161        try testing.expectEqualStrings(stage.target_spec.pass_name, stage.pass_registration.name);162        try testing.expectEqualStrings(stage.target_spec.pass_description, stage.pass_registration.description);163        try testing.expectEqualStrings(stage.target_spec.pipeline_name, stage.pipeline_registration.name);164        try testing.expectEqualStrings(stage.target_spec.pipeline_description, stage.pipeline_registration.description);165        try testing.expect(stage.target_spec.legalizesDialect(stage.target_spec.target_dialect_name));166        try testing.expectEqualStrings(stage.pass_registration.name, targetLoweringPassName(index));167        try testing.expectEqualStrings(stage.pass_registration.name, pass_registrations[index].name);168        try testing.expectEqualStrings(stage.pass_registration.description, pass_registrations[index].description);169        try testing.expectEqualStrings(stage.pass_registration.pass.name, pass_registrations[index].pass.name);170        try testing.expectEqualStrings(stage.pipeline_registration.name, targetLoweringPipelineName(index));171        try testing.expectEqualStrings(stage.pipeline_registration.name, pipeline_registrations[index].name);172        try testing.expectEqualStrings(stage.pipeline_registration.description, pipeline_registrations[index].description);173    }174}

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