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

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 const abi = @import("choir_abi");
 3 const choir = @import("../../root.zig");
 4 const gpu = @import("../../dialects/gpu/root.zig");
 5 
 6 const ir = choir.ir;
 7 
 8 pub fn requirementsForModule(module: *ir.Operation) abi.SubgroupRequirements {
 9     var requirements: abi.SubgroupRequirements = .{};
10     appendRequirements(&requirements, module);
11     return requirements;
12 }
13 
14 pub fn requirementsForOperation(op: *ir.Operation) abi.SubgroupRequirements {
15     const GpuDialect = gpu.GpuDialect;
16     if (isName(op.name.name, GpuDialect.LaneIdOp.operation_name)) return .{ .supported = true };
17     if (isName(op.name.name, GpuDialect.WarpIdOp.operation_name)) return .{ .supported = true };
18     if (isName(op.name.name, GpuDialect.ActiveMaskOp.operation_name)) return .{ .supported = true, .ballot = true };
19     if (isName(op.name.name, GpuDialect.SyncWarpOp.operation_name)) return .{ .supported = true };
20     if (isName(op.name.name, GpuDialect.ShflSyncOp.operation_name)) return .{ .supported = true, .shuffle = true };
21     if (isName(op.name.name, GpuDialect.AllSyncOp.operation_name)) return .{ .supported = true, .vote = true };
22     if (isName(op.name.name, GpuDialect.AnySyncOp.operation_name)) return .{ .supported = true, .vote = true };
23     if (isName(op.name.name, GpuDialect.BallotSyncOp.operation_name)) return .{ .supported = true, .ballot = true };
24     if (isName(op.name.name, GpuDialect.MatchAnyOp.operation_name)) return .{ .supported = true, .ballot = true, .vote = true };
25     if (isName(op.name.name, GpuDialect.MatchAllOp.operation_name)) return .{ .supported = true, .ballot = true, .vote = true };
26     if (isName(op.name.name, GpuDialect.WarpReduceOp.operation_name)) return .{ .supported = true, .arithmetic = true };
27     if (isName(op.name.name, GpuDialect.WarpScanOp.operation_name)) return .{ .supported = true, .arithmetic = true, .scan = true };
28     if (isName(op.name.name, GpuDialect.BarrierOp.operation_name)) {
29         return requirementsForScope((GpuDialect.BarrierOp{ .op = op }).getScope());
30     }
31     if (isName(op.name.name, GpuDialect.FenceOp.operation_name)) {
32         return requirementsForScope((GpuDialect.FenceOp{ .op = op }).getScope());
33     }
34     return .{};
35 }
36 
37 fn appendRequirements(requirements: *abi.SubgroupRequirements, op: *ir.Operation) void {
38     mergeRequirement(requirements, requirementsForOperation(op));
39     for (op.regions.items) |*region| {
40         var block_iter = region.getBlocks();
41         while (block_iter.next()) |block| {
42             var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
43             while (current) |current_op| {
44                 appendRequirements(requirements, current_op);
45                 current = current_op.next_op;
46             }
47         }
48     }
49 }
50 
51 fn requirementsForScope(scope: ?gpu.Scope) abi.SubgroupRequirements {
52     if (scope) |value| {
53         if (value == .warp) return .{ .supported = true };
54     }
55     return .{};
56 }
57 
58 fn mergeRequirement(requirements: *abi.SubgroupRequirements, required: abi.SubgroupRequirements) void {
59     requirements.* = requirements.*.merge(required);
60 }
61 
62 fn isName(actual: []const u8, expected: []const u8) bool {
63     return std.mem.eql(u8, actual, expected);
64 }
65 
66 test "active mask requires subgroup ballot support" {
67     var ctx = try ir.Context.init(
68         std.testing.allocator,
69         ir.Context.Limits.testing,
70     );
71     defer ctx.deinit(std.testing.allocator);
72 
73     const active = try gpu.GpuDialect.ActiveMaskOp.create(&ctx, ir.Location.getUnknown());
74     const requirements = requirementsForOperation(active.op);
75     try std.testing.expect(requirements.supported);
76     try std.testing.expect(requirements.ballot);
77 }