lib/choir/src/backends/gpu/lowering.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const choir = @import("../../root.zig");
  3 const spirv = @import("spirv/root.zig");
  4 const nvptx = @import("nvptx/root.zig");
  5 
  6 pub const gpu_to_spirv_pipeline_name = "gpu-to-spirv-pipeline";
  7 pub const gpu_to_spirv_pipeline_description = "Lower the GPU dialect to the SPIR-V target dialect";
  8 
  9 pub const gpu_to_nvptx_pipeline_name = "gpu-to-nvptx-pipeline";
 10 pub const gpu_to_nvptx_pipeline_description = "Lower the GPU dialect to the NVPTX target dialect";
 11 
 12 pub const TargetLowering = enum {
 13     spirv,
 14     nvptx,
 15 };
 16 
 17 const Stage = struct {
 18     lowering: TargetLowering,
 19     target_spec: choir.backends.TargetSpec,
 20     pass_registration: choir.passes.PassRegistration,
 21     pipeline_registration: choir.passes.PipelineRegistration,
 22 };
 23 
 24 fn buildGpuToSpirvPipeline(pm: *choir.passes.OpPassManager) anyerror!void {
 25     try pm.addPass(spirv.createGpuToSpirvPass());
 26 }
 27 
 28 fn buildGpuToNvptxPipeline(pm: *choir.passes.OpPassManager) anyerror!void {
 29     try pm.addPass(nvptx.createGpuToNvptxPass());
 30 }
 31 
 32 pub 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 };
 37 
 38 pub 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 };
 43 
 44 fn 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 }
 53 
 54 const 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 };
 68 
 69 fn 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 }
 78 
 79 fn 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 }
 88 
 89 fn stageFor(lowering: TargetLowering) Stage {
 90     return switch (lowering) {
 91         .spirv => stages[0],
 92         .nvptx => stages[1],
 93     };
 94 }
 95 
 96 fn 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 }
105 
106 pub const pass_registrations = passRegistrations();
107 pub const pipeline_registrations = pipelineRegistrations();
108 pub const target_specs = specs();
109 pub const spirv_target_spec = target_specs[0];
110 pub const nvptx_target_spec = target_specs[1];
111 pub const target_lowering_stage_count = stages.len;
112 pub const target_spec_count = target_specs.len;
113 pub const target_lowering_pass_count = pass_registrations.len;
114 pub const target_lowering_pipeline_count = pipeline_registrations.len;
115 
116 pub fn targetLoweringPassName(index: usize) []const u8 {
117     return stages[index].pass_registration.name;
118 }
119 
120 pub fn targetLoweringPipelineName(index: usize) []const u8 {
121     return stages[index].pipeline_registration.name;
122 }
123 
124 pub fn targetLoweringTargetDialectName(index: usize) []const u8 {
125     return stages[index].target_spec.target_dialect_name;
126 }
127 
128 pub fn targetLoweringSpec(index: usize) choir.backends.TargetSpec {
129     return target_specs[index];
130 }
131 
132 pub fn addGpuToSpirvPipeline(pm: *choir.passes.PassManager) !void {
133     try gpu_to_spirv_pipeline_registration.addTo(&pm.root);
134 }
135 
136 pub fn addGpuToNvptxPipeline(pm: *choir.passes.PassManager) !void {
137     try gpu_to_nvptx_pipeline_registration.addTo(&pm.root);
138 }
139 
140 pub fn addTargetLoweringPipeline(pm: *choir.passes.PassManager, lowering: TargetLowering) !void {
141     try stageFor(lowering).pipeline_registration.addTo(&pm.root);
142 }
143 
144 test "target lowering stage table defines registries" {
145     const testing = std.testing;
146 
147     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);
154 
155     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 }