[Flint] adding multiple registers/blocks support

src/intel/FlintCommandBuffer.zig

@@ -332,9 +332,9 @@ pub fn dispatchBase(interface: *Interface, base_group_x: u32, base_group_y: u32,
     const self: *Self = @alignCast(@fieldParentPtr("interface", interface));
     if (group_count_x == 0 or group_count_y == 0 or group_count_z == 0)
         return;
-    if (base_group_x != 0 or base_group_y != 0 or base_group_z != 0 or
-        group_count_x != 1 or group_count_y != 1 or group_count_z != 1)
+    if (base_group_x != 0 or base_group_y != 0 or base_group_z != 0)
         return VkError.FeatureNotPresent;
+    const group_count: [3]u32 = .{ group_count_x, group_count_y, group_count_z };
 
     const pipeline = self.bound_compute_pipeline orelse return VkError.ValidationFailed;
     const artifact = pipeline.computeArtifact() orelse return VkError.FeatureNotPresent;
@@ -364,6 +364,7 @@ pub fn dispatchBase(interface: *Interface, base_group_x: u32, base_group_y: u32,
         sizes[resource.binding_table_index] = range.size;
     }
 
+    std.log.scoped(.FlintDispatch).debug("dispatch {d}x{d}x{d}, kernel {d} bytes, state buffers {any}, existing batch {d} dwords", .{ group_count_x, group_count_y, group_count_z, kernel.len, sizes[0..artifact.resources.bindings.len], self.batch.items.len });
     const old_engine = self.engine;
     try self.requireEngine(.render);
     const old_batch_len = self.batch.items.len;
@@ -386,7 +387,7 @@ pub fn dispatchBase(interface: *Interface, base_group_x: u32, base_group_y: u32,
     errdefer if (state_owned) state.deinit(&device.kmd, self.interface.owner.io());
 
     const mapped = try state.map(&device.kmd, self.interface.owner.io(), 0, gen9_dispatch.page_size);
-    const state_layout = gen9_dispatch.writeState(mapped, kernel, sizes[0..artifact.resources.bindings.len]) catch |err| switch (err) {
+    const state_layout = gen9_dispatch.writeState(mapped, kernel, sizes[0..artifact.resources.bindings.len], group_count) catch |err| switch (err) {
         error.StateTooLarge,
         error.UnsupportedBufferSize,
         error.EmptyBuffer,
@@ -465,7 +466,7 @@ pub fn dispatchBase(interface: *Interface, base_group_x: u32, base_group_y: u32,
     try self.emitSlice(&gen9_dispatch.pipeControl(gen9_dispatch.pipe_control.cs_stall));
     try self.emitSlice(&gen9_dispatch.mediaVfeState());
     try self.emitSlice(&gen9_dispatch.interfaceDescriptorLoad(state_layout.interface_descriptor_offset));
-    try self.emitSlice(&gen9_dispatch.gpgpuWalker(.{ 1, 1, 1 }, 1));
+    try self.emitSlice(&gen9_dispatch.gpgpuWalker(group_count, 1));
     try self.emitSlice(&gen9_dispatch.mediaStateFlush);
     try self.emitSlice(&gen9_dispatch.pipeControl(gen9_dispatch.pipe_control.cs_stall |
         gen9_dispatch.pipe_control.dc_flush));

src/intel/FlintPipeline.zig

@@ -320,6 +320,6 @@ test "Flint pipeline: lower common compute IR" {
     const text = try compiler.printer.allocPrint(std.testing.allocator, program);
     defer std.testing.allocator.free(text);
     try std.testing.expect(std.mem.indexOf(u8, text, "load_global_invocation_id") == null);
-    try std.testing.expect(std.mem.indexOf(u8, text, "mov r1:u32, 0:u32") != null);
+    try std.testing.expect(std.mem.indexOf(u8, text, "r0:u32[byte=4, broadcast]") != null);
     try std.testing.expect(std.mem.indexOf(u8, text, "surface_message write bti(0)") != null);
 }

src/intel/compiler/ir/instruction.zig

@@ -9,6 +9,8 @@ pub const LoadGlobalInvocationId = struct {
     component: u8,
 };
 
+pub const LoadNumWorkgroups = LoadGlobalInvocationId;
+
 pub const BufferReference = union(enum) {
     logical: ids.StorageBufferId,
     binding_table: u8,
@@ -113,6 +115,7 @@ pub const Math = struct {
 
 pub const Operation = union(enum) {
     load_global_invocation_id: LoadGlobalInvocationId,
+    load_num_workgroups: LoadNumWorkgroups,
     load_buffer: LoadBuffer,
     store_buffer: StoreBuffer,
     array_length: ArrayLength,

src/intel/compiler/ir/printer.zig

@@ -111,8 +111,8 @@ fn writeInstruction(program: *const program_ir.Program, writer: *std.Io.Writer,
 
 fn writeOperation(program: *const program_ir.Program, writer: *std.Io.Writer, execution_size: device.ExecutionSize, operation: inst_ir.Operation) !void {
     switch (operation) {
-        .load_global_invocation_id => |op| {
-            try writer.writeAll("load_global_invocation_id ");
+        .load_global_invocation_id, .load_num_workgroups => |op| {
+            try writer.print("{s} ", .{@tagName(operation)});
             try writeDestination(program, writer, execution_size, op.destination);
             try writer.print(", component({d})", .{op.component});
         },

src/intel/compiler/ir/validator.zig

@@ -22,6 +22,7 @@ pub const Error = error{
     InvalidLaneCount,
     InvalidMath,
     InvalidMessage,
+    InvalidNumWorkgroups,
     InvalidParallelCopyDestination,
     InvalidPayloadLayout,
     InvalidPhysicalRegister,
@@ -140,6 +141,13 @@ fn validateInstruction(program: *const program_ir.Program, inst: instruction.Ins
             if (op.component >= 3 or op.destination.type != .u32)
                 return Error.InvalidGlobalInvocationId;
         },
+        .load_num_workgroups => |op| {
+            if (program.properties.system_values_lowered)
+                return Error.UnloweredSystemValue;
+            try validateDestination(program, op.destination);
+            if (op.component >= 3 or op.destination.type != .u32)
+                return Error.InvalidNumWorkgroups;
+        },
         .load_buffer => |op| {
             if (program.properties.messages_lowered)
                 return Error.UnloweredMessage;
@@ -507,6 +515,21 @@ test "[ir] validator checks compute system values and resources" {
     try std.testing.expectError(Error.UnloweredSystemValue, validate(&program));
     program.properties.system_values_lowered = false;
 
+    program.instructions.getMut(system_value_id).?.operation = .{ .load_num_workgroups = .{
+        .destination = .{ .register = .{ .virtual = register }, .type = .u32 },
+        .component = 2,
+    } };
+    try validate(&program);
+    program.instructions.getMut(system_value_id).?.operation.load_num_workgroups.component = 3;
+    try std.testing.expectError(Error.InvalidNumWorkgroups, validate(&program));
+    program.instructions.getMut(system_value_id).?.operation.load_num_workgroups.component = 0;
+    program.instructions.getMut(system_value_id).?.operation.load_num_workgroups.destination.type = .i32;
+    try std.testing.expectError(Error.InvalidNumWorkgroups, validate(&program));
+    program.instructions.getMut(system_value_id).?.operation.load_num_workgroups.destination.type = .u32;
+    program.properties.system_values_lowered = true;
+    try std.testing.expectError(Error.UnloweredSystemValue, validate(&program));
+    program.properties.system_values_lowered = false;
+
     program.instructions.getMut(buffer_load_id).?.operation.load_buffer.buffer = .{ .logical = ids.StorageBufferId.fromIndex(99) };
     try std.testing.expectError(Error.InvalidStorageBuffer, validate(&program));
     program.instructions.getMut(buffer_load_id).?.operation.load_buffer.buffer = .{ .logical = buffer };

src/intel/compiler/lower/common_ir.zig

@@ -764,7 +764,7 @@ const LoweringState = struct {
         switch (variable.semantic) {
             .builtin => |builtin| switch (builtin) {
                 .global_invocation_id => try self.lowerGlobalInvocationId(block_id, result_id, variable),
-                .num_workgroups => try self.lowerNumWorkgroups(result_id),
+                .num_workgroups => try self.lowerNumWorkgroups(block_id, result_id),
                 .workgroup_size => base.unsupported("workgroup size builtin is not yet supported in Flint", .{}),
                 else => return Error.UnsupportedOperation,
             },
@@ -799,26 +799,22 @@ const LoweringState = struct {
         }
     }
 
-    fn lowerNumWorkgroups(self: *LoweringState, result_id: shader_ir.id.ValueId) Error!void {
+    fn lowerNumWorkgroups(self: *LoweringState, block_id: ids.BlockId, result_id: shader_ir.id.ValueId) Error!void {
         const result_value = self.lowerer.module.values.get(result_id) orelse return Error.InvalidModule;
         const lowered_type = try self.lowerType(result_value.type);
 
         if (lowered_type.element_type != .u32 or lowered_type.component_count != 3)
             return Error.UnsupportedType;
 
-        const vec = try self.storage.alloc(operand.Source, 3);
-
-        for (self.lowerer.module.execution_modes.workgroup_size.?, vec) |value, *component| {
-            component.* = .{
-                .register = .{
-                    .immediate = .{ .u32 = value },
+        const result_components = try self.addRegisterLocation(result_id, .temporary);
+        for (result_components, 0..) |result_component, component_index| {
+            try self.appendInstruction(block_id, null, .{
+                .load_num_workgroups = .{
+                    .destination = try destinationFromSource(result_component),
+                    .component = @intCast(component_index),
                 },
-                .type = .u32,
-                .region = operand.Region.broadcast(),
-            };
+            });
         }
-
-        try self.putLocation(result_id, .{ .components = vec });
     }
 
     fn lowerStoreInterface(self: *LoweringState, block_id: ids.BlockId, result: ?shader_ir.id.ValueId, operation: shader_ir.instruction.StoreInterface) Error!void {
@@ -1397,6 +1393,42 @@ test "[ir] Lower: global invocation ID" {
     }, &.{});
 }
 
+test "[ir] Lower: number of workgroups is a runtime system value" {
+    const source =
+        \\shader compute @main
+        \\{
+        \\    @group_count: vec3[u32] = input[builtin(num_workgroups)]
+        \\    @destination: vec3[u32] = storage_buffer[set(0), binding(0)]
+        \\    %offset: constant u32 = 0
+        \\    fn @main() -> void
+        \\    {
+        \\        .entry():
+        \\            branch .load()
+        \\        .load():
+        \\            %count: vec3[u32] = load_interface @group_count
+        \\            store_buffer @destination, %offset, %count
+        \\            return
+        \\    }
+        \\}
+    ;
+
+    try expectLoweredFragments(source, &.{
+        "%count_x: vgrf u32[8], class(temporary)",
+        "%count_y: vgrf u32[8], class(temporary)",
+        "%count_z: vgrf u32[8], class(temporary)",
+        ".load:\n    [simd8] load_num_workgroups %count_x:u32, component(0)",
+        "[simd8] load_num_workgroups %count_y:u32, component(1)",
+        "[simd8] load_num_workgroups %count_z:u32, component(2)",
+        "[simd8] store_buffer @destination, 0:u32, %count_x:u32",
+        "[simd8] store_buffer @destination, 0:u32, offset(4), %count_y:u32",
+        "[simd8] store_buffer @destination, 0:u32, offset(8), %count_z:u32",
+    }, &.{
+        "store_buffer @destination, 0:u32, 1:u32",
+        "store_buffer @destination, 0:u32, offset(4), 1:u32",
+        "store_buffer @destination, 0:u32, offset(8), 1:u32",
+    });
+}
+
 test "[ir] Lower: vector storage-buffer operations" {
     const source =
         \\shader compute @main

src/intel/compiler/targets/gen9/compute/dispatch.zig

@@ -2,6 +2,8 @@ const std = @import("std");
 
 pub const max_storage_surfaces: usize = 4;
 pub const max_surfaces: usize = max_storage_surfaces + 1;
+pub const num_workgroups_offset: u32 = max_storage_surfaces * @sizeOf(u32);
+const size_table_size: u32 = num_workgroups_offset + 3 * @sizeOf(u32);
 pub const page_size: usize = 4096;
 pub const surface_state_size: usize = 64;
 pub const interface_descriptor_size: usize = 32;
@@ -29,7 +31,7 @@ pub const StateLayout = struct {
     interface_descriptor_offset: u32,
 };
 
-pub fn writeState(destination: []u8, kernel: []const u8, buffer_sizes: []const u64) Error!StateLayout {
+pub fn writeState(destination: []u8, kernel: []const u8, buffer_sizes: []const u64, group_count: [3]u32) Error!StateLayout {
     if (buffer_sizes.len > max_storage_surfaces)
         return Error.TooManySurfaces;
 
@@ -69,7 +71,7 @@ pub fn writeState(destination: []u8, kernel: []const u8, buffer_sizes: []const u
 
     cursor = alignForward(cursor, @alignOf(u32));
     layout.size_table_offset = @intCast(cursor);
-    cursor += @max(buffer_sizes.len, 1) * @sizeOf(u32);
+    cursor += size_table_size;
 
     cursor = alignForward(cursor, 64);
     layout.interface_descriptor_offset = @intCast(cursor);
@@ -86,7 +88,10 @@ pub fn writeState(destination: []u8, kernel: []const u8, buffer_sizes: []const u
         putU32(destination, layout.binding_table_offset + @as(u32, @intCast(index * @sizeOf(u32))), layout.surface_offsets[index]);
         putU32(destination, layout.size_table_offset + @as(u32, @intCast(index * @sizeOf(u32))), @intCast(size));
     }
-    _ = try encodeRawBufferSurface(destination, layout.surface_offsets[size_table_surface], @max(buffer_sizes.len, 1) * @sizeOf(u32));
+    for (group_count, 0..) |count, component| {
+        putU32(destination, layout.size_table_offset + num_workgroups_offset + @as(u32, @intCast(component * @sizeOf(u32))), count);
+    }
+    _ = try encodeRawBufferSurface(destination, layout.surface_offsets[size_table_surface], size_table_size);
     putU32(destination, layout.binding_table_offset + @as(u32, @intCast(size_table_surface * @sizeOf(u32))), layout.surface_offsets[size_table_surface]);
 
     const idd = layout.interface_descriptor_offset;
@@ -185,9 +190,45 @@ fn putU32(destination: []u8, offset: u32, value: u32) void {
     std.mem.writeInt(u32, destination[offset..][0..@sizeOf(u32)], value, .little);
 }
 
+test "[gen9] dispatch: fixed size-table ABI includes workgroup counts" {
+    const buffer_sizes = [_]u64{ 4096, 8192, 16384, 32768 };
+    const group_count: [3]u32 = .{ 7, 11, 13 };
+    try std.testing.expectEqual(@as(u32, 16), num_workgroups_offset);
+    try std.testing.expectEqual(@as(u32, 28), size_table_size);
+
+    for (0..max_storage_surfaces + 1) |buffer_count| {
+        var state: [page_size]u8 = undefined;
+        const layout = try writeState(&state, &.{ 0xaa, 0xbb }, buffer_sizes[0..buffer_count], group_count);
+        for (0..max_storage_surfaces) |index| {
+            const actual = std.mem.readInt(u32, state[layout.size_table_offset + index * @sizeOf(u32) ..][0..4], .little);
+            const expected: u32 = if (index < buffer_count) @intCast(buffer_sizes[index]) else 0;
+            try std.testing.expectEqual(expected, actual);
+        }
+        for (group_count, 0..) |expected, component| {
+            const actual = std.mem.readInt(u32, state[layout.size_table_offset + num_workgroups_offset + component * @sizeOf(u32) ..][0..4], .little);
+            try std.testing.expectEqual(expected, actual);
+        }
+        const surface_offset = layout.surface_offsets[buffer_count];
+        try std.testing.expectEqual(@as(u32, 27), std.mem.readInt(u32, state[surface_offset + 8 ..][0..4], .little));
+        try std.testing.expectEqual(@as(u32, 0), std.mem.readInt(u32, state[surface_offset + 12 ..][0..4], .little));
+        try std.testing.expect(layout.size_table_offset + size_table_size <= layout.interface_descriptor_offset);
+    }
+}
+
+test "[gen9] dispatch: walker preserves multidimensional group counts with one active lane" {
+    const words = gpgpuWalker(.{ 7, 11, 13 }, 1);
+    try std.testing.expectEqual(@as(u32, 7), words[7]);
+    try std.testing.expectEqual(@as(u32, 11), words[10]);
+    try std.testing.expectEqual(@as(u32, 13), words[12]);
+    try std.testing.expectEqual(@as(u32, 1), words[13]);
+    try std.testing.expectEqual(@as(u32, 0), words[4]);
+    for ([_]usize{ 5, 8, 11 }) |index|
+        try std.testing.expectEqual(@as(u32, 0), words[index]);
+}
+
 test "[gen9] dispatch: interface descriptor exposes internal size-table surface" {
     var state: [page_size]u8 = undefined;
-    const layout = try writeState(&state, &.{ 0xaa, 0xbb }, &.{ 4096, 8192 });
+    const layout = try writeState(&state, &.{ 0xaa, 0xbb }, &.{ 4096, 8192 }, .{ 1, 1, 1 });
 
     try std.testing.expectEqual(@as(u8, 3), layout.surface_count);
 

src/intel/compiler/targets/gen9/compute/eu_encoder.zig

@@ -149,7 +149,7 @@ fn encodeJumpWithPredicate(displacement_bytes: i32, predicate: ?operand.Predicat
             .physical => |physical| physical,
             .virtual => return Error.UnsupportedOperand,
         };
-        if (flag.register != 0 or flag.subregister > 1)
+        if (flag.register > 1 or flag.subregister > 1)
             return Error.InvalidRegister;
 
         encoded.setBits(19, 16, 1); // Normal predicate control.
@@ -197,7 +197,7 @@ pub fn encodeCompare(execution_size: device.ExecutionSize, compare: ir_instructi
         .physical => |value| value,
         .virtual => return Error.UnsupportedOperand,
     };
-    if (flag.register != 0 or flag.subregister > 1)
+    if (flag.register > 1 or flag.subregister > 1)
         return Error.InvalidRegister;
 
     var encoded = try instructionHeader(.cmp, execution_size);
@@ -477,6 +477,39 @@ test "[gen9] EU encoder: encode unsigned less-than comparison" {
     try std.testing.expectEqual(@as(u64, 1), encoded.bits(32, 32));
 }
 
+test "[gen9] EU encoder: encode all flag halves and reject out of range flags" {
+    const binary = testBinary(.add);
+    var compare: ir_instruction.Compare = .{
+        .opcode = .equal,
+        .destination = .{ .physical = .{} },
+        .lhs = binary.lhs,
+        .rhs = binary.rhs,
+    };
+    for (0..2) |register| {
+        for (0..2) |subregister| {
+            const flag: operand.FlagRef = .{ .physical = .{
+                .register = @intCast(register),
+                .subregister = @intCast(subregister),
+            } };
+            compare.destination = flag;
+            const encoded_compare = try encodeCompare(.simd8, compare);
+            const encoded_jump = try encodePredicatedJump(16, .{ .flag = flag });
+            for ([_]EncodedInstruction{ encoded_compare, encoded_jump }) |encoded| {
+                try std.testing.expectEqual(@as(u64, @intCast(register)), encoded.bits(33, 33));
+                try std.testing.expectEqual(@as(u64, @intCast(subregister)), encoded.bits(32, 32));
+            }
+        }
+    }
+    for ([_]operand.PhysicalFlag{
+        .{ .register = 2 },
+        .{ .register = 1, .subregister = 2 },
+    }) |physical| {
+        compare.destination = .{ .physical = physical };
+        try std.testing.expectError(Error.InvalidRegister, encodeCompare(.simd8, compare));
+        try std.testing.expectError(Error.InvalidRegister, encodePredicatedJump(16, .{ .flag = .{ .physical = physical } }));
+    }
+}
+
 test "[gen9] EU encoder: encode predicated jump" {
     const encoded = try encodePredicatedJump(-32, .{
         .flag = .{ .physical = .{ .register = 0, .subregister = 1 } },

src/intel/compiler/targets/gen9/compute/kernel_encoder.zig

@@ -20,7 +20,11 @@ const JumpFixup = struct {
 pub fn encode(allocator: std.mem.Allocator, program: *program_ir.Program) Error![]u8 {
     if (!program.properties.registers_allocated)
         return Error.InvalidProgram;
-    if (program.program_data.total_grf_count > eu.eot_payload_grf)
+
+    if (program.program_data.total_grf_count > program.device_info.grf_count)
+        return Error.InvalidProgram;
+
+    if (program.device_info.grf_count <= eu.eot_payload_grf)
         return Error.EotRegisterUnavailable;
 
     const entry_id = program.entry_block orelse return Error.InvalidProgram;
@@ -81,7 +85,7 @@ pub fn encode(allocator: std.mem.Allocator, program: *program_ir.Program) Error!
                 const instructions = try eu.encodeEndThread(header);
                 for (instructions) |encoded|
                     try appendInstruction(allocator, &kernel, encoded);
-                program.program_data.total_grf_count = eu.eot_payload_grf + 1;
+                program.program_data.total_grf_count = @max(program.program_data.total_grf_count, eu.eot_payload_grf + 1);
             },
             .@"unreachable" => return Error.UnsupportedControlFlow,
         }
@@ -145,6 +149,38 @@ fn appendInstruction(allocator: std.mem.Allocator, kernel: *std.ArrayList(u8), i
     try kernel.appendSlice(allocator, &bytes);
 }
 
+test "[gen9] kernel encoder: use upper GRFs and encode repeatedly" {
+    const device_info = @import("../../../device.zig").DeviceInfo{
+        .generation = .gen9,
+        .platform = .skylake,
+        .pci_device_id = 0x1912,
+        .grf_count = 128,
+    };
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, device_info, .simd8);
+    defer program.deinit();
+    const entry = try program.addBlock("entry");
+    _ = try program.appendInstruction(entry, .simd8, null, .{ .move = .{
+        .destination = .{ .register = .{ .physical_grf = .{ .number = 127 } }, .type = .u32 },
+        .source = .{ .register = .{ .physical_grf = .{ .number = 112 } }, .type = .u32, .region = @import("../../../ir/operand.zig").Region.contiguous(.simd8) },
+    } });
+    try program.setTerminator(entry, .end_thread);
+    program.payload.header_grf = .{ .number = 0 };
+    program.properties.registers_allocated = true;
+    program.program_data.total_grf_count = 128;
+
+    const first = try encode(std.testing.allocator, &program);
+    defer std.testing.allocator.free(first);
+    try std.testing.expectEqual(@as(u16, 128), program.program_data.total_grf_count);
+    const second = try encode(std.testing.allocator, &program);
+    defer std.testing.allocator.free(second);
+    try std.testing.expectEqualSlices(u8, first, second);
+    try std.testing.expectEqual(@as(u64, 127), (std.mem.readInt(u64, first[0..8], .little) >> 53) & 0xff);
+
+    program.device_info.grf_count = 112;
+    program.program_data.total_grf_count = 112;
+    try std.testing.expectError(Error.EotRegisterUnavailable, encode(std.testing.allocator, &program));
+}
+
 test "[gen9] kernel encoder: patch unconditional jump between blocks" {
     const device = @import("../../../device.zig");
 

src/intel/compiler/targets/gen9/compute/pipeline.zig

@@ -78,7 +78,7 @@ pub fn compile(allocator: std.mem.Allocator, module: *shader_ir.module.Module, d
     errdefer program.deinit();
 
     try abi.run(&program);
-    try system_values.run(&program);
+
     try block_arguments.run(allocator, &program);
     try parallel_copies.run(allocator, &program);
 
@@ -86,6 +86,7 @@ pub fn compile(allocator: std.mem.Allocator, module: *shader_ir.module.Module, d
     errdefer resources.deinit(allocator);
 
     try resource_lowering.run(&program, &resources);
+    try system_values.run(&program, resources.bindings.len);
     try array_length_lowering.run(&program, &resources);
     try message_lowering.run(&program);
     try message_addresses.run(&program);

src/intel/compiler/targets/gen9/compute/regions.zig

@@ -43,7 +43,22 @@ fn legalizeSource(source: *operand.Source, execution_size: @import("../../../dev
     const byte_offset = source.region.byte_offset;
     source.region = switch (source.register) {
         .immediate => operand.Region.broadcast(),
+        // Physical payload operands already describe the hardware ABI, including
+        // scalar header fields that must be broadcast rather than read as vectors
+        .physical_grf => source.region,
         else => operand.Region.contiguous(execution_size),
     };
     source.region.byte_offset = byte_offset;
 }
+
+test "[gen9] regions: preserve scalar physical payload reads" {
+    const std = @import("std");
+    var source: operand.Source = .{
+        .register = .{ .physical_grf = .{ .number = 0, .byte_offset = 28 } },
+        .type = .u32,
+        .region = operand.Region.broadcast(),
+    };
+    legalizeSource(&source, .simd8);
+    try std.testing.expectEqual(operand.Region.broadcast(), source.region);
+    try std.testing.expectEqual(@as(u8, 28), source.register.physical_grf.byte_offset);
+}

src/intel/compiler/targets/gen9/compute/system_values.zig

@@ -3,16 +3,20 @@ const std = @import("std");
 const operand = @import("../../../ir/operand.zig");
 const program_ir = @import("../../../ir/program.zig");
 const validator = @import("../../../ir/validator.zig");
+const dispatch = @import("dispatch.zig");
 
 pub const Error = error{InvalidProgram};
 
-pub fn run(program: *program_ir.Program) Error!void {
+pub fn run(program: *program_ir.Program, storage_surface_count: usize) Error!void {
     validator.validate(program) catch return Error.InvalidProgram;
     if (program.properties.system_values_lowered)
         return;
 
-    // Flint dispatch currently accepts only one invocation in one workgroup at
-    // base group zero, so every component of GlobalInvocationId is zero.
+    if (storage_surface_count > dispatch.max_storage_surfaces)
+        return Error.InvalidProgram;
+
+    // Dispatch currently supports one invocation per workgroup. Its global ID
+    // equals the group ID delivered in the hardware thread header
     if (!std.mem.eql(u32, &program.workgroup_size, &.{ 1, 1, 1 }))
         return;
 
@@ -21,7 +25,16 @@ pub fn run(program: *program_ir.Program) Error!void {
         inst.operation = switch (inst.operation) {
             .load_global_invocation_id => |op| .{ .move = .{
                 .destination = op.destination,
-                .source = zero(),
+                .source = groupId(op.component),
+            } },
+            .load_num_workgroups => |op| .{ .load_buffer = .{
+                .destination = op.destination,
+                .buffer = .{ .binding_table = @intCast(storage_surface_count) },
+                .byte_offset = .{
+                    .register = .{ .immediate = .{ .u32 = dispatch.num_workgroups_offset + @as(u32, op.component) * @sizeOf(u32) } },
+                    .type = .u32,
+                    .region = operand.Region.broadcast(),
+                },
             } },
             else => inst.operation,
         };
@@ -31,9 +44,11 @@ pub fn run(program: *program_ir.Program) Error!void {
     validator.validate(program) catch return Error.InvalidProgram;
 }
 
-fn zero() operand.Source {
+fn groupId(component: u8) operand.Source {
+    // Gen9 GPGPU thread payload: group X in r0.1, Y in r0.6, Z in r0.7
+    const dwords = [_]u8{ 1, 6, 7 };
     return .{
-        .register = .{ .immediate = .{ .u32 = 0 } },
+        .register = .{ .physical_grf = .{ .number = 0, .byte_offset = dwords[component] * @sizeOf(u32) } },
         .type = .u32,
         .region = operand.Region.broadcast(),
     };
@@ -68,11 +83,11 @@ test "[gen9] system values: lower global invocation ID for single invocation" {
     try builder.setTerminator(entry, .end_thread);
     try builder.setEntryBlock(entry);
 
-    try run(&program);
+    try run(&program, 0);
 
     try std.testing.expect(program.properties.system_values_lowered);
     const move = program.instructions.get(load).?.operation.move;
-    try std.testing.expectEqual(@as(u32, 0), move.source.register.immediate.u32);
+    try std.testing.expectEqual(operand.PhysicalGrf{ .number = 0, .byte_offset = 28 }, move.source.register.physical_grf);
 }
 
 test "[gen9] system values: preserve IDs for unsupported workgroup sizes" {
@@ -90,6 +105,6 @@ test "[gen9] system values: preserve IDs for unsupported workgroup sizes" {
     try program.setTerminator(entry, .end_thread);
     try program.setEntryBlock(entry);
 
-    try run(&program);
+    try run(&program, 0);
     try std.testing.expect(!program.properties.system_values_lowered);
 }

src/intel/compiler/targets/gen9/flag_allocation.zig

@@ -14,7 +14,7 @@ pub const Error = std.mem.Allocator.Error || error{
     OutOfFlagRegisters,
 };
 
-const physical_flag_count = 2;
+const physical_flag_count = 4;
 
 pub fn run(allocator: std.mem.Allocator, program: *program_ir.Program) Error!void {
     if (!program.properties.block_parameters_lowered)
@@ -35,31 +35,167 @@ pub fn run(allocator: std.mem.Allocator, program: *program_ir.Program) Error!voi
     var occupied: [physical_flag_count]bool = @splat(false);
     try visitProgramFlags(program, allocations, &occupied, false);
 
-    for (allocations) |*allocation| {
-        const marker = allocation.* orelse continue;
-        if (marker.subregister != std.math.maxInt(u8))
-            return Error.InvalidProgram;
-
-        const subregister = std.mem.indexOfScalar(bool, &occupied, false) orelse return Error.OutOfFlagRegisters;
-
-        allocation.* = .{
-            .register = 0,
-            .subregister = @intCast(subregister),
-        };
-        occupied[subregister] = true;
-    }
+    try allocateLiveFlags(allocator, program, allocations, occupied);
 
     try visitProgramFlags(program, allocations, &occupied, true);
     program.properties.flags_allocated = true;
     validator.validate(program) catch return Error.InvalidProgram;
 }
 
-fn visitProgramFlags(
-    program: *program_ir.Program,
-    allocations: []?operand.PhysicalFlag,
-    occupied: *[physical_flag_count]bool,
-    rewrite: bool,
-) Error!void {
+fn addUse(live: []bool, flag: operand.FlagRef) void {
+    switch (flag) {
+        .virtual => |id| live[id.index()] = true,
+        .physical => {},
+    }
+}
+
+fn interfere(graph: []bool, count: usize, a: usize, b: usize) void {
+    if (a == b) return;
+    graph[a * count + b] = true;
+    graph[b * count + a] = true;
+}
+
+fn addLiveInterference(graph: []bool, live: []const bool) void {
+    for (live, 0..) |a_live, a| {
+        if (!a_live) continue;
+        for (live[0..a], 0..) |b_live, b| {
+            if (b_live) interfere(graph, live.len, a, b);
+        }
+    }
+}
+
+fn mergeSuccessor(live: []bool, live_in: []const bool, edge: instruction.Edge) void {
+    const successor = live_in[edge.target.index() * live.len ..][0..live.len];
+    for (live, successor) |*value, incoming| value.* = value.* or incoming;
+    for (edge.arguments) |argument| switch (argument) {
+        .source => {},
+        .predicate => |value| switch (value) {
+            .constant => {},
+            .dynamic => |predicate| addUse(live, predicate.flag),
+        },
+    };
+}
+
+fn scanBlock(program: *const program_ir.Program, block: instruction.Block, live_in: []const bool, live: []bool, graph: ?[]bool) void {
+    @memset(live, false);
+    switch (block.terminator.?) {
+        .jump => |edge| mergeSuccessor(live, live_in, edge),
+        .conditional_branch => |branch| {
+            mergeSuccessor(live, live_in, branch.true_edge);
+            mergeSuccessor(live, live_in, branch.false_edge);
+            addUse(live, branch.predicate.flag);
+        },
+        .end_thread, .@"unreachable" => {},
+    }
+    if (graph) |edges| addLiveInterference(edges, live);
+    var index = block.instructions.items.len;
+    while (index > 0) {
+        index -= 1;
+        const inst = program.instructions.get(block.instructions.items[index]).?;
+        if (inst.operation == .compare) {
+            switch (inst.operation.compare.destination) {
+                .virtual => |destination| {
+                    if (graph) |edges| {
+                        for (live, 0..) |is_live, other| {
+                            if (is_live) interfere(edges, live.len, destination.index(), other);
+                        }
+                        // Keep a compare's predicate distinct from its destination.
+                        if (inst.predicate) |predicate| switch (predicate.flag) {
+                            .virtual => |source| interfere(edges, live.len, destination.index(), source.index()),
+                            .physical => {},
+                        };
+                    }
+                    // A predicated write preserves the old value on inactive lanes.
+                    if (inst.predicate == null) live[destination.index()] = false;
+                },
+                .physical => {},
+            }
+        }
+        if (inst.predicate) |predicate| addUse(live, predicate.flag);
+        if (graph) |edges| addLiveInterference(edges, live);
+    }
+}
+
+fn allocateLiveFlags(allocator: std.mem.Allocator, program: *const program_ir.Program, allocations: []?operand.PhysicalFlag, occupied: [physical_flag_count]bool) Error!void {
+    const count = allocations.len;
+    const live_in = try allocator.alloc(bool, program.blocks.entries.items.len * count);
+    defer allocator.free(live_in);
+    @memset(live_in, false);
+    const live = try allocator.alloc(bool, count);
+    defer allocator.free(live);
+
+    var changed = true;
+    while (changed) {
+        changed = false;
+        for (program.blocks.entries.items, 0..) |entry, block_index| {
+            const block = entry orelse continue;
+            scanBlock(program, block, live_in, live, null);
+            const incoming = live_in[block_index * count ..][0..count];
+            if (!std.mem.eql(bool, incoming, live)) {
+                @memcpy(incoming, live);
+                changed = true;
+            }
+        }
+    }
+
+    const graph = try allocator.alloc(bool, count * count);
+    defer allocator.free(graph);
+    @memset(graph, false);
+    for (program.blocks.entries.items) |entry| {
+        const block = entry orelse continue;
+        scanBlock(program, block, live_in, live, graph);
+    }
+
+    var available: [physical_flag_count]u8 = undefined;
+    var available_count: usize = 0;
+    for (occupied, 0..) |reserved, slot| {
+        if (reserved) continue;
+        available[available_count] = @intCast(slot);
+        available_count += 1;
+    }
+
+    const colors = try allocator.alloc(?u8, count);
+    defer allocator.free(colors);
+    @memset(colors, null);
+
+    while (true) {
+        var selected: ?usize = null;
+        var best_saturation: usize = 0;
+        var best_degree: usize = 0;
+        var selected_used: [physical_flag_count]bool = @splat(false);
+        for (allocations, 0..) |allocation, candidate| {
+            if (allocation == null or colors[candidate] != null) continue;
+            var used: [physical_flag_count]bool = @splat(false);
+            var degree: usize = 0;
+            for (graph[candidate * count ..][0..count], 0..) |adjacent, other| {
+                if (!adjacent) continue;
+                degree += 1;
+                if (colors[other]) |color| used[color] = true;
+            }
+            const saturation = std.mem.count(bool, &used, &.{true});
+            if (selected == null or saturation > best_saturation or
+                (saturation == best_saturation and degree > best_degree))
+            {
+                selected = candidate;
+                best_saturation = saturation;
+                best_degree = degree;
+                selected_used = used;
+            }
+        }
+        const current = selected orelse break;
+        const color = std.mem.indexOfScalar(bool, selected_used[0..available_count], false) orelse return Error.OutOfFlagRegisters;
+        colors[current] = @intCast(color);
+    }
+
+    for (allocations, colors) |*allocation, color| {
+        if (color) |value| allocation.* = .{
+            .register = available[value] / 2,
+            .subregister = available[value] % 2,
+        };
+    }
+}
+
+fn visitProgramFlags(program: *program_ir.Program, allocations: []?operand.PhysicalFlag, occupied: *[physical_flag_count]bool, rewrite: bool) Error!void {
     for (program.instructions.entries.items, 0..) |entry, instruction_index| {
         _ = entry orelse continue;
         const inst = program.instructions.getMut(ids.InstructionId.fromIndex(instruction_index)) orelse
@@ -140,18 +276,34 @@ fn visitFlagRef(
             }
 
             const physical = allocations[virtual.index()] orelse return Error.InvalidProgram;
-            if (physical.subregister >= physical_flag_count)
+            if (physical.register > 1 or physical.subregister > 1)
                 return Error.InvalidProgram;
             flag.* = .{ .physical = physical };
         },
         .physical => |physical| {
-            if (physical.register != 0 or physical.subregister >= physical_flag_count)
+            if (physical.register > 1 or physical.subregister > 1)
                 return Error.InvalidProgram;
-            occupied[physical.subregister] = true;
+            occupied[physical.register * 2 + physical.subregister] = true;
         },
     }
 }
 
+fn defineFlag(program: *program_ir.Program, block: ids.BlockId, flag: operand.FlagRef) !ids.InstructionId {
+    return program.appendInstruction(block, .simd8, null, .{ .compare = .{
+        .opcode = .equal,
+        .destination = flag,
+        .lhs = immediateU32(0),
+        .rhs = immediateU32(0),
+    } });
+}
+
+fn useFlag(program: *program_ir.Program, block: ids.BlockId, flag: operand.FlagRef) !ids.InstructionId {
+    return program.appendInstruction(block, .simd8, .{ .flag = flag }, .{ .move = .{
+        .destination = .{ .register = .{ .physical_grf = .{ .number = 10 } }, .type = .u32 },
+        .source = immediateU32(0),
+    } });
+}
+
 const test_device = @import("../../device.zig").DeviceInfo{
     .generation = .gen9,
     .platform = .skylake,
@@ -227,13 +379,15 @@ test "[gen9] flag allocation: rewrite compares and predicates" {
     try std.testing.expectEqual(@as(u8, 1), branch.predicate.flag.physical.subregister);
 }
 
-test "[gen9] flag allocation: report exhaustion without rewriting" {
+test "[gen9] flag allocation: fifth live flag reports exhaustion without rewriting" {
     var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
     defer program.deinit();
 
     const first = try program.addVirtualFlag(.{});
     const second = try program.addVirtualFlag(.{});
     const third = try program.addVirtualFlag(.{});
+    const fourth = try program.addVirtualFlag(.{});
+    const fifth = try program.addVirtualFlag(.{});
     const entry = try program.addBlock("entry");
 
     const first_compare = try program.appendInstruction(entry, .simd8, null, .{ .compare = .{
@@ -248,6 +402,13 @@ test "[gen9] flag allocation: report exhaustion without rewriting" {
         .lhs = immediateU32(1),
         .rhs = immediateU32(1),
     } });
+    _ = try defineFlag(&program, entry, .{ .virtual = fourth });
+    _ = try defineFlag(&program, entry, .{ .virtual = fifth });
+    _ = try useFlag(&program, entry, .{ .virtual = fourth });
+    _ = try useFlag(&program, entry, .{ .virtual = fifth });
+    _ = try useFlag(&program, entry, .{ .virtual = first });
+    _ = try useFlag(&program, entry, .{ .virtual = second });
+    _ = try useFlag(&program, entry, .{ .virtual = third });
     try program.setTerminator(entry, .end_thread);
     markPrerequisites(&program);
 
@@ -255,3 +416,157 @@ test "[gen9] flag allocation: report exhaustion without rewriting" {
     try std.testing.expect(!program.properties.flags_allocated);
     try std.testing.expectEqual(first, program.instructions.get(first_compare).?.operation.compare.destination.virtual);
 }
+
+test "[gen9] flag allocation: all four flag halves simultaneously live" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+    defer program.deinit();
+    const entry = try program.addBlock("entry");
+    var flags: [4]operand.FlagRef = undefined;
+    var definitions: [4]ids.InstructionId = undefined;
+    var uses: [4]ids.InstructionId = undefined;
+    for (&flags, &definitions) |*flag, *definition| {
+        flag.* = .{ .virtual = try program.addVirtualFlag(.{}) };
+        definition.* = try defineFlag(&program, entry, flag.*);
+    }
+    for (flags, &uses) |flag, *use| use.* = try useFlag(&program, entry, flag);
+    try program.setTerminator(entry, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    var seen: [4]bool = @splat(false);
+    for (definitions, uses) |definition, use| {
+        const flag = program.instructions.get(definition).?.operation.compare.destination.physical;
+        try std.testing.expect(flag.register <= 1 and flag.subregister <= 1);
+        const slot = flag.register * 2 + flag.subregister;
+        try std.testing.expect(!seen[slot]);
+        seen[slot] = true;
+        try std.testing.expectEqual(flag, program.instructions.get(use).?.predicate.?.flag.physical);
+    }
+    try std.testing.expectEqual([_]bool{ true, true, true, true }, seen);
+}
+
+test "[gen9] flag allocation: reuse sequential lifetimes" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+    defer program.deinit();
+    const entry = try program.addBlock("entry");
+    var definitions: [12]ids.InstructionId = undefined;
+    for (&definitions) |*definition| {
+        const flag: operand.FlagRef = .{ .virtual = try program.addVirtualFlag(.{}) };
+        definition.* = try defineFlag(&program, entry, flag);
+        _ = try useFlag(&program, entry, flag);
+    }
+    try program.setTerminator(entry, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    for (definitions) |definition| {
+        try std.testing.expectEqual(@as(u8, 0), program.instructions.get(definition).?.operation.compare.destination.physical.subregister);
+    }
+    try run(std.testing.allocator, &program);
+}
+
+test "[gen9] flag allocation: cross block live values and mutually exclusive lifetimes" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+    defer program.deinit();
+    const entry = try program.addBlock("entry");
+    const left = try program.addBlock("left");
+    const right = try program.addBlock("right");
+    const merge = try program.addBlock("merge");
+    const carried: operand.FlagRef = .{ .virtual = try program.addVirtualFlag(.{}) };
+    const definition = try defineFlag(&program, entry, carried);
+    try program.setTerminator(entry, .{ .conditional_branch = .{
+        .predicate = .{ .flag = carried },
+        .true_edge = .{ .target = left, .arguments = &.{} },
+        .false_edge = .{ .target = right, .arguments = &.{} },
+    } });
+    var locals: [2]ids.InstructionId = undefined;
+    for ([_]ids.BlockId{ left, right }, &locals) |block, *local| {
+        const flag: operand.FlagRef = .{ .virtual = try program.addVirtualFlag(.{}) };
+        local.* = try defineFlag(&program, block, flag);
+        _ = try useFlag(&program, block, flag);
+        try program.setTerminator(block, .{ .jump = .{ .target = merge, .arguments = &.{} } });
+    }
+    const use = try useFlag(&program, merge, carried);
+    try program.setTerminator(merge, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    const physical = program.instructions.get(definition).?.operation.compare.destination.physical;
+    try std.testing.expectEqual(physical, program.instructions.get(use).?.predicate.?.flag.physical);
+    for (locals) |local| {
+        try std.testing.expect(physical.subregister != program.instructions.get(local).?.operation.compare.destination.physical.subregister);
+    }
+}
+
+test "[gen9] flag allocation: loop backedge preserves live flags" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+    defer program.deinit();
+    const entry = try program.addBlock("entry");
+    const header = try program.addBlock("header");
+    const body = try program.addBlock("body");
+    const exit = try program.addBlock("exit");
+    const carried: operand.FlagRef = .{ .virtual = try program.addVirtualFlag(.{}) };
+    const definition = try defineFlag(&program, entry, carried);
+    try program.setTerminator(entry, .{ .jump = .{ .target = header, .arguments = &.{} } });
+    try program.setTerminator(header, .{ .conditional_branch = .{
+        .predicate = .{ .flag = carried },
+        .true_edge = .{ .target = body, .arguments = &.{} },
+        .false_edge = .{ .target = exit, .arguments = &.{} },
+    } });
+    var locals: [3]ids.InstructionId = undefined;
+    for (&locals) |*local| {
+        const flag: operand.FlagRef = .{ .virtual = try program.addVirtualFlag(.{}) };
+        local.* = try defineFlag(&program, body, flag);
+        _ = try useFlag(&program, body, flag);
+    }
+    try program.setTerminator(body, .{ .jump = .{ .target = header, .arguments = &.{} } });
+    try program.setTerminator(exit, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    const physical = program.instructions.get(definition).?.operation.compare.destination.physical;
+    for (locals) |local| {
+        try std.testing.expect(physical.subregister != program.instructions.get(local).?.operation.compare.destination.physical.subregister);
+    }
+}
+
+test "[gen9] flag allocation: predicated definitions preserve old destination lanes" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+    defer program.deinit();
+    const entry = try program.addBlock("entry");
+    const carried: operand.FlagRef = .{ .virtual = try program.addVirtualFlag(.{}) };
+    const temporary: operand.FlagRef = .{ .virtual = try program.addVirtualFlag(.{}) };
+    const condition: operand.FlagRef = .{ .virtual = try program.addVirtualFlag(.{}) };
+    const original = try defineFlag(&program, entry, carried);
+    const clobber = try defineFlag(&program, entry, temporary);
+    _ = try defineFlag(&program, entry, condition);
+    _ = try program.appendInstruction(entry, .simd8, .{ .flag = condition }, .{ .compare = .{
+        .opcode = .equal,
+        .destination = carried,
+        .lhs = immediateU32(1),
+        .rhs = immediateU32(0),
+    } });
+    _ = try useFlag(&program, entry, carried);
+    try program.setTerminator(entry, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    try std.testing.expect(program.instructions.get(original).?.operation.compare.destination.physical.subregister !=
+        program.instructions.get(clobber).?.operation.compare.destination.physical.subregister);
+}
+
+test "[gen9] flag allocation: reserve explicit physical flags while reusing virtual slots" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+    defer program.deinit();
+    const entry = try program.addBlock("entry");
+    _ = try defineFlag(&program, entry, .{ .physical = .{ .register = 0, .subregister = 0 } });
+    _ = try defineFlag(&program, entry, .{ .physical = .{ .register = 0, .subregister = 1 } });
+    _ = try defineFlag(&program, entry, .{ .physical = .{ .register = 1, .subregister = 0 } });
+    var definitions: [3]ids.InstructionId = undefined;
+    for (&definitions) |*definition| {
+        const flag: operand.FlagRef = .{ .virtual = try program.addVirtualFlag(.{}) };
+        definition.* = try defineFlag(&program, entry, flag);
+        _ = try useFlag(&program, entry, flag);
+    }
+    try program.setTerminator(entry, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    for (definitions) |definition| {
+        try std.testing.expectEqual(operand.PhysicalFlag{ .register = 1, .subregister = 1 }, program.instructions.get(definition).?.operation.compare.destination.physical);
+    }
+}

src/intel/compiler/targets/gen9/gen9.zig

@@ -152,6 +152,241 @@ test "[gen9] target: lower 256 KiB SSBO copy loop" {
     try std.testing.expectEqual(@as(usize, 4), store_count);
 }
 
+test "[gen9] target: reuse registers across conditional blocks and a carrying loop" {
+    var module = try shader_ir.parser.parseString(std.testing.allocator,
+        \\shader compute @main
+        \\{
+        \\    @storage: u32 = storage_buffer[set(0), binding(0)]
+        \\    %offset: constant u32 = bits(0x0)
+        \\    %zero: constant i32 = bits(0x0)
+        \\    %one: constant i32 = bits(0x1)
+        \\    %limit: constant i32 = bits(0x8)
+        \\    fn @main() -> void
+        \\    {
+        \\        .entry():
+        \\            %input: u32 = load_buffer @storage, %offset
+        \\            %seed: i32 = bitcast %input
+        \\            %retained: i32 = integer_add %seed, %one
+        \\            branch .test1(%seed)
+        \\        .test1(%a: i32):
+        \\            %c1: bool = cmp_signed_less %a, %limit
+        \\            conditional_branch %c1, .add1(), .test2(%a)
+        \\        .add1():
+        \\            %b: i32 = integer_add %a, %one
+        \\            branch .test2(%b)
+        \\        .test2(%c: i32):
+        \\            %c2: bool = cmp_signed_less %c, %limit
+        \\            conditional_branch %c2, .add2(), .test3(%c)
+        \\        .add2():
+        \\            %d: i32 = integer_add %c, %one
+        \\            branch .test3(%d)
+        \\        .test3(%e: i32):
+        \\            %c3: bool = cmp_signed_less %e, %limit
+        \\            conditional_branch %c3, .add3(), .test4(%e)
+        \\        .add3():
+        \\            %f: i32 = integer_add %e, %one
+        \\            branch .test4(%f)
+        \\        .test4(%g: i32):
+        \\            %c4: bool = cmp_signed_less %g, %limit
+        \\            conditional_branch %c4, .add4(), .test5(%g)
+        \\        .add4():
+        \\            %h: i32 = integer_add %g, %one
+        \\            branch .test5(%h)
+        \\        .test5(%j: i32):
+        \\            %c5: bool = cmp_signed_less %j, %limit
+        \\            conditional_branch %c5, .add5(), .header(%zero, %j)
+        \\        .add5():
+        \\            %k: i32 = integer_add %j, %one
+        \\            branch .header(%zero, %k)
+        \\        .header(%index: i32, %sum: i32):
+        \\            %in_bounds: bool = cmp_signed_less %index, %limit
+        \\            conditional_branch %in_bounds, .body(), .exit()
+        \\        .body():
+        \\            %next_sum: i32 = integer_add %sum, %index
+        \\            %next_index: i32 = integer_add %index, %one
+        \\            branch .header(%next_index, %next_sum)
+        \\        .exit():
+        \\            %result: i32 = integer_add %sum, %retained
+        \\            %output: u32 = bitcast %result
+        \\            store_buffer @storage, %offset, %output
+        \\            return
+        \\    }
+        \\}
+    );
+    defer module.deinit();
+    module.execution_modes.workgroup_size = .{ 1, 1, 1 };
+
+    var artifact = try compileCompute(std.testing.allocator, &module, .{
+        .generation = .gen9,
+        .platform = .skylake,
+        .pci_device_id = 0x1912,
+        .grf_count = 128,
+    }, .{});
+    defer artifact.deinit(std.testing.allocator);
+    try std.testing.expect(artifact.kernel != null);
+    const program = &artifact.program;
+    try std.testing.expect(program.properties.registers_allocated);
+    try std.testing.expect(program.properties.flags_allocated);
+
+    const PhysicalGrf = @import("../../ir/operand.zig").PhysicalGrf;
+    var retained: ?PhysicalGrf = null;
+    var last_rhs: ?PhysicalGrf = null;
+    var destinations = std.StaticBitSet(128).initEmpty();
+    var binary_count: usize = 0;
+    var compare_count: usize = 0;
+    var store_count: usize = 0;
+    for (program.instructions.entries.items) |entry| {
+        const inst = entry orelse continue;
+        switch (inst.operation) {
+            .binary => |op| {
+                const destination = op.destination.register.physical_grf;
+                if (retained) |live| {
+                    // The entry value survives every conditional and the loop until the final add.
+                    try std.testing.expect(destination.number != live.number);
+                } else {
+                    retained = destination;
+                }
+                destinations.set(destination.number);
+                binary_count += 1;
+                last_rhs = if (op.rhs.register == .physical_grf) op.rhs.register.physical_grf else null;
+            },
+            .compare => compare_count += 1,
+            .surface_message => |op| {
+                if (op.kind == .write) store_count += 1;
+            },
+            else => {},
+        }
+    }
+    try std.testing.expectEqual(@as(usize, 6), compare_count);
+    try std.testing.expectEqual(@as(usize, 9), binary_count);
+    try std.testing.expectEqual(@as(usize, 1), store_count);
+    try std.testing.expect(retained != null and last_rhs != null);
+    try std.testing.expectEqual(retained.?, last_rhs.?);
+    // Count arithmetic GRFs before the encoder's EOT header copy, not the total
+    // high-water mark: EOT reserves r112 even when shader temporaries reuse GRFs.
+    try std.testing.expect(destinations.count() < binary_count);
+    try std.testing.expect(program.program_data.total_grf_count >= 113);
+}
+
+test "[gen9] target: CTS multiple invocations invert copy" {
+    var module = try shader_ir.parser.parseString(std.testing.allocator,
+        \\shader compute @main
+        \\{
+        \\    @source: runtime_array[u32] = storage_buffer[set(0), binding(0)]
+        \\    @destination: runtime_array[u32] = storage_buffer[set(0), binding(1)]
+        \\    @global_id: vec3[u32] = input[builtin(global_invocation_id)]
+        \\    @group_count: vec3[u32] = input[builtin(num_workgroups)]
+        \\    %zero: constant u32 = 0
+        \\    %one: constant u32 = 1
+        \\    %stride: constant u32 = 4
+        \\    fn @main() -> void
+        \\    {
+        \\        .entry():
+        \\            %id: vec3[u32] = load_interface @global_id
+        \\            %x: u32 = composite_extract %id[0]
+        \\            %y: u32 = composite_extract %id[1]
+        \\            %z: u32 = composite_extract %id[2]
+        \\            %count: vec3[u32] = load_interface @group_count
+        \\            %nx: u32 = composite_extract %count[0]
+        \\            %ny: u32 = composite_extract %count[1]
+        \\            %nz: u32 = composite_extract %count[2]
+        \\            %zy: u32 = integer_multiply %z, %ny
+        \\            %row: u32 = integer_add %zy, %y
+        \\            %row_start: u32 = integer_multiply %row, %nx
+        \\            %linear_id: u32 = integer_add %row_start, %x
+        \\            %nxy: u32 = integer_multiply %nx, %ny
+        \\            %invocations: u32 = integer_multiply %nxy, %nz
+        \\            %length: u32 = array_length @source, %zero, stride 4
+        \\            %per_invocation: u32 = unsigned_divide %length, %invocations
+        \\            %begin: u32 = integer_multiply %linear_id, %per_invocation
+        \\            %end: u32 = integer_add %begin, %per_invocation
+        \\            branch .header(%begin)
+        \\        .header(%index: u32):
+        \\            %in_bounds: bool = cmp_unsigned_less %index, %end
+        \\            conditional_branch %in_bounds, .body(), .exit()
+        \\        .body():
+        \\            %offset: u32 = integer_multiply %index, %stride
+        \\            %value: u32 = load_buffer @source, %offset
+        \\            %inverted: u32 = bitwise_not %value
+        \\            store_buffer @destination, %offset, %inverted
+        \\            branch .continue()
+        \\        .continue():
+        \\            %next: u32 = integer_add %index, %one
+        \\            branch .header(%next)
+        \\        .exit():
+        \\            return
+        \\    }
+        \\}
+    );
+    defer module.deinit();
+    module.execution_modes.workgroup_size = .{ 1, 1, 1 };
+
+    var artifact = try compileCompute(std.testing.allocator, &module, .{
+        .generation = .gen9,
+        .platform = .skylake,
+        .pci_device_id = 0x1912,
+        .grf_count = 128,
+    }, .{});
+    defer artifact.deinit(std.testing.allocator);
+    try std.testing.expect(artifact.kernel != null);
+    const program = &artifact.program;
+    try std.testing.expect(program.properties.system_values_lowered);
+    try std.testing.expect(program.properties.resources_lowered);
+    try std.testing.expect(program.properties.registers_allocated);
+    try std.testing.expectEqual(@as(usize, 2), artifact.resources.bindings.len);
+
+    const header_offsets = [_]u8{ 4, 24, 28 };
+    var header_components: usize = 0;
+    var internal_reads: usize = 0;
+    var source_reads: usize = 0;
+    var destination_writes: usize = 0;
+    // Inspect block order: message payload moves are inserted before their sends.
+    for (program.blocks.entries.items) |block_entry| {
+        const block = block_entry orelse continue;
+        for (block.instructions.items, 0..) |instruction_id, index| {
+            const inst = program.instructions.get(instruction_id).?;
+            switch (inst.operation) {
+                .load_global_invocation_id, .load_num_workgroups => return error.UnloweredSystemValue,
+                .array_length => return error.UnloweredArrayLength,
+                .move => |op| {
+                    if (op.source.register == .physical_grf and op.source.register.physical_grf.number == 0) {
+                        // Ignore the whole-header EOT copy, if present.
+                        if (op.source.register.physical_grf.byte_offset == 0) continue;
+                        try std.testing.expect(header_components < header_offsets.len);
+                        try std.testing.expectEqual(header_offsets[header_components], op.source.register.physical_grf.byte_offset);
+
+                        header_components += 1;
+                    }
+                },
+                .surface_message => |op| {
+                    if (op.binding_table == artifact.resources.bindings.len) {
+                        try std.testing.expectEqual(.read, op.kind);
+                        try std.testing.expect(index > 0);
+                        const payload = program.instructions.get(block.instructions.items[index - 1]).?.operation.move;
+                        try std.testing.expectEqual(op.payload.base, payload.destination.register);
+                        // Three NumWorkgroups loads, then source buffer byte size.
+                        const offsets = [_]u32{ 16, 20, 24, 0 };
+                        try std.testing.expect(internal_reads < offsets.len);
+                        try std.testing.expectEqual(offsets[internal_reads], payload.source.register.immediate.u32);
+                        internal_reads += 1;
+                    } else if (op.kind == .read) {
+                        try std.testing.expectEqual(@as(u8, 0), op.binding_table);
+                        source_reads += 1;
+                    } else {
+                        try std.testing.expectEqual(@as(u8, 1), op.binding_table);
+                        destination_writes += 1;
+                    }
+                },
+                else => {},
+            }
+        }
+    }
+    try std.testing.expectEqual(@as(usize, 3), header_components);
+    try std.testing.expectEqual(@as(usize, 4), internal_reads);
+    try std.testing.expectEqual(@as(usize, 1), source_reads);
+    try std.testing.expectEqual(@as(usize, 1), destination_writes);
+}
+
 test "[gen9] target: encode runtime array length" {
     var module = try shader_ir.parser.parseString(std.testing.allocator,
         \\shader compute @main

src/intel/compiler/targets/gen9/register_allocation.zig

@@ -15,117 +15,348 @@ pub const Error = std.mem.Allocator.Error || error{
 pub fn run(allocator: std.mem.Allocator, program: *program_ir.Program) Error!void {
     if (!program.properties.block_parameters_lowered)
         return Error.BlockParametersNotLowered;
+
     if (!program.properties.parallel_copies_lowered)
         return Error.ParallelCopiesNotLowered;
+
     if (program.properties.registers_allocated)
         return;
 
-    const grf_size = program.device_info.grf_size_bytes;
-    if (grf_size == 0)
+    const grf_size: usize = program.device_info.grf_size_bytes;
+    if (grf_size == 0 or grf_size > 256 or !std.math.isPowerOfTwo(grf_size))
         return Error.InvalidProgram;
 
-    const allocations = try allocator.alloc(?operand.PhysicalGrf, program.virtual_registers.entries.items.len);
-    defer allocator.free(allocations);
+    const capacity = @as(usize, program.device_info.grf_count) * grf_size;
+    const count = program.virtual_registers.entries.items.len;
+    const block_count = program.blocks.entries.items.len;
+
+    var arena = std.heap.ArenaAllocator.init(allocator);
+    defer arena.deinit();
+
+    const scratch = arena.allocator();
+    const allocations = try scratch.alloc(?operand.PhysicalGrf, count);
     @memset(allocations, null);
 
-    var next_byte: usize = @as(usize, program.program_data.payload_grf_count) * grf_size;
-    next_byte = try reserveExistingPhysicalRegisters(program, next_byte, grf_size);
-    const capacity: usize = @as(usize, program.device_info.grf_count) * grf_size;
+    const fixed = try scratch.alloc(bool, capacity);
+    @memset(fixed, false);
+
+    const alignments = try scratch.alloc(usize, count);
+    for (program.virtual_registers.entries.items, 0..) |entry, index| {
+        const register = entry orelse continue;
+        if (register.size_bytes == 0 or !std.math.isPowerOfTwo(register.alignment_bytes))
+            return Error.InvalidProgram;
+        alignments[index] = register.alignment_bytes;
+    }
+
+    var analysis = Analysis{ .program = program, .fixed = fixed, .alignments = alignments, .grf_size = grf_size };
+    const payload_end = @as(usize, program.program_data.payload_grf_count) * grf_size;
+
+    if (payload_end > capacity)
+        return Error.OutOfRegisters;
+
+    @memset(fixed[0..payload_end], true);
+
+    if (program.payload.header_grf) |header|
+        try analysis.reserve(header, 0, grf_size);
+
+    const accesses = try scratch.alloc(Access, program.instructions.entries.items.len);
+    for (program.instructions.entries.items, 0..) |entry, index| {
+        accesses[index] = if (entry) |inst| try analysis.instructionAccess(inst) else .{};
+    }
+
+    const matrix_size = std.math.mul(usize, block_count, count) catch return Error.OutOfMemory;
+    const live_in = try scratch.alloc(bool, matrix_size);
+    @memset(live_in, false);
+
+    const live = try scratch.alloc(bool, count);
+    const graph_size = std.math.mul(usize, count, count) catch return Error.OutOfMemory;
+    var graph = try std.DynamicBitSetUnmanaged.initEmpty(scratch, graph_size);
+
+    var changed = true;
+    while (changed) {
+        changed = false;
+        var block_index = block_count;
+        while (block_index != 0) {
+            block_index -= 1;
+            const block = program.blocks.entries.items[block_index] orelse continue;
+            try successorLive(program, block, live_in, live);
+            var position = block.instructions.items.len;
+
+            while (position != 0) {
+                position -= 1;
+                const id = block.instructions.items[position];
+                const inst = program.instructions.get(id) orelse return Error.InvalidProgram;
+
+                if (inst.parent_block.index() != block_index)
+                    return Error.InvalidProgram;
+
+                accesses[id.index()].transfer(live);
+            }
+
+            const input = live_in[block_index * count ..][0..count];
+
+            if (!std.mem.eql(bool, input, live)) {
+                @memcpy(input, live);
+                changed = true;
+            }
+        }
+    }
+
+    for (program.blocks.entries.items) |entry| {
+        const block = entry orelse continue;
+        try successorLive(program, block, live_in, live);
+        addClique(&graph, live);
+        var position = block.instructions.items.len;
+        while (position != 0) {
+            position -= 1;
+            const access = accesses[block.instructions.items[position].index()];
+
+            for (access.uses[0..access.use_count]) |use|
+                live[use] = true;
+
+            if (access.definition) |definition|
+                live[definition] = true;
+
+            addClique(&graph, live);
+            access.transfer(live);
+        }
+    }
+
+    var high_water = payload_end;
+    for (fixed, 0..) |reserved, byte| {
+        if (reserved)
+            high_water = byte + 1;
+    }
 
     for (program.virtual_registers.entries.items, 0..) |entry, index| {
         const register = entry orelse continue;
-        const start = std.mem.alignForward(usize, next_byte, register.alignment_bytes);
-        const end = std.math.add(usize, start, register.size_bytes) catch return Error.OutOfRegisters;
-        if (end > capacity)
+        const size: usize = register.size_bytes;
+
+        if (size > capacity)
             return Error.OutOfRegisters;
 
-        allocations[index] = .{
-            .number = @intCast(start / grf_size),
-            .byte_offset = @intCast(start % grf_size),
-        };
-        next_byte = end;
+        var start: usize = 0;
+        while (true) : (start += alignments[index]) {
+            if (start > capacity - size)
+                return Error.OutOfRegisters;
+
+            const end = start + size;
+            if (std.mem.indexOfScalar(bool, fixed[start..end], true) != null)
+                continue;
+
+            var conflict = false;
+            for (allocations[0..index], 0..) |allocated, other| {
+                const physical = allocated orelse continue;
+
+                if (!graph.isSet(index * count + other))
+                    continue;
+
+                const other_start = @as(usize, physical.number) * grf_size + physical.byte_offset;
+                const other_end = other_start + program.virtual_registers.entries.items[other].?.size_bytes;
+
+                if (start < other_end and other_start < end) {
+                    conflict = true;
+                    break;
+                }
+            }
+
+            if (conflict)
+                continue;
+            allocations[index] = .{ .number = @intCast(start / grf_size), .byte_offset = @intCast(start % grf_size) };
+            high_water = @max(high_water, end);
+
+            break;
+        }
     }
 
     try rewriteProgram(program, allocations);
-    program.program_data.total_grf_count = @intCast(std.math.divCeil(usize, next_byte, grf_size) catch return Error.InvalidProgram);
+    program.program_data.total_grf_count = @intCast(std.math.divCeil(usize, high_water, grf_size) catch return Error.InvalidProgram);
     program.properties.registers_allocated = true;
 }
 
-fn reserveExistingPhysicalRegisters(program: *const program_ir.Program, initial: usize, grf_size: usize) Error!usize {
-    var next_byte = initial;
-    if (program.payload.header_grf) |header|
-        reservePhysical(&next_byte, header, grf_size);
+const Access = struct {
+    uses: [3]usize = undefined,
+    use_count: usize = 0,
+    definition: ?usize = null,
+    full_overwrite: bool = false,
 
-    for (program.instructions.entries.items) |entry| {
-        const inst = entry orelse continue;
+    fn use(self: *Access, index: usize) void {
+        self.uses[self.use_count] = index;
+        self.use_count += 1;
+    }
+
+    fn transfer(self: Access, live: []bool) void {
+        if (self.definition) |definition|
+            live[definition] = !self.full_overwrite;
+
+        for (self.uses[0..self.use_count]) |index|
+            live[index] = true;
+    }
+};
+
+const Analysis = struct {
+    program: *const program_ir.Program,
+    fixed: []bool,
+    alignments: []usize,
+    grf_size: usize,
+
+    fn reserve(self: *Analysis, physical: operand.PhysicalGrf, offset: usize, size: usize) Error!void {
+        if (physical.byte_offset >= self.grf_size)
+            return Error.InvalidProgram;
+
+        const start = @as(usize, physical.number) * self.grf_size + physical.byte_offset + offset;
+
+        if (start > self.fixed.len or size > self.fixed.len - start)
+            return Error.InvalidProgram;
+
+        @memset(self.fixed[start..][0..size], true);
+    }
+
+    fn register(self: *Analysis, ref: operand.RegisterRef, offset: usize, size: usize) Error!?usize {
+        switch (ref) {
+            .virtual => |id| {
+                const value = self.program.virtual_registers.get(id) orelse return Error.InvalidProgram;
+
+                if (offset > value.size_bytes or size > value.size_bytes - offset)
+                    return Error.InvalidProgram;
+
+                return id.index();
+            },
+            .physical_grf => |physical| try self.reserve(physical, offset, size),
+            else => {},
+        }
+        return null;
+    }
+
+    fn read(self: *Analysis, access: *Access, value: operand.Source, lanes: usize) Error!void {
+        const region = value.region;
+
+        if (region.width == 0)
+            return Error.InvalidProgram;
+
+        var last: usize = 0;
+        for (0..lanes) |lane| {
+            last = @max(last, (lane / region.width) * region.vertical_stride + (lane % region.width) * region.horizontal_stride);
+        }
+
+        if (try self.register(value.register, region.byte_offset, (last + 1) * value.type.sizeBytes())) |index|
+            access.use(index);
+    }
+
+    fn write(self: *Analysis, access: *Access, value: operand.Destination, lanes: usize, predicated: bool) Error!void {
+        const size = ((lanes - 1) * value.region.horizontal_stride + 1) * value.type.sizeBytes();
+
+        if (try self.register(value.register, value.region.byte_offset, size)) |index| {
+            access.definition = index;
+            access.full_overwrite = !predicated and value.region.byte_offset == 0 and
+                (lanes == 1 or value.region.horizontal_stride == 1) and
+                size == self.program.virtual_registers.entries.items[index].?.size_bytes;
+        }
+    }
+
+    fn span(self: *Analysis, access: *Access, value: operand.RegisterSpan, destination: bool, predicated: bool) Error!void {
+        if (value.register_count == 0)
+            return Error.InvalidProgram;
+
+        const size = @as(usize, value.register_count) * self.grf_size;
+        if (try self.register(value.base, 0, size)) |index| {
+            self.alignments[index] = @max(self.alignments[index], self.grf_size);
+            if (destination) {
+                access.definition = index;
+                access.full_overwrite = !predicated and size == self.program.virtual_registers.entries.items[index].?.size_bytes;
+            } else access.use(index);
+        }
+    }
+
+    fn instructionAccess(self: *Analysis, inst: instruction.Instruction) Error!Access {
+        var access = Access{};
+        const lanes: usize = @intFromEnum(inst.execution_size);
+        const predicated = inst.predicate != null;
         switch (inst.operation) {
-            .load_global_invocation_id => |op| reserveRegister(&next_byte, op.destination.register, grf_size),
-            .load_buffer => |op| {
-                reserveRegister(&next_byte, op.destination.register, grf_size);
-                reserveRegister(&next_byte, op.byte_offset.register, grf_size);
+            .load_global_invocation_id, .load_num_workgroups => |op| try self.write(&access, op.destination, lanes, predicated),
+            inline .load_buffer, .array_length => |op| {
+                try self.write(&access, op.destination, lanes, predicated);
+                try self.read(&access, op.byte_offset, lanes);
             },
             .store_buffer => |op| {
-                reserveRegister(&next_byte, op.byte_offset.register, grf_size);
-                reserveRegister(&next_byte, op.source.register, grf_size);
-            },
-            .array_length => |op| {
-                reserveRegister(&next_byte, op.destination.register, grf_size);
-                reserveRegister(&next_byte, op.byte_offset.register, grf_size);
+                try self.read(&access, op.byte_offset, lanes);
+                try self.read(&access, op.source, lanes);
             },
             .surface_read => |op| {
-                reserveRegister(&next_byte, op.destination.register, grf_size);
-                reserveRegister(&next_byte, op.address.register, grf_size);
+                try self.write(&access, op.destination, lanes, predicated);
+                try self.read(&access, op.address, lanes);
             },
             .surface_write => |op| {
-                reserveRegister(&next_byte, op.address.register, grf_size);
-                reserveRegister(&next_byte, op.data.register, grf_size);
+                try self.read(&access, op.address, lanes);
+                try self.read(&access, op.data, lanes);
             },
             .surface_message => |op| {
-                reserveRegister(&next_byte, op.payload.base, grf_size);
-                if (op.response) |response|
-                    reserveRegister(&next_byte, response.base, grf_size);
+                try self.span(&access, op.payload, false, predicated);
+                if (op.response) |response| try self.span(&access, response, true, predicated);
             },
             .move => |op| {
-                reserveRegister(&next_byte, op.destination.register, grf_size);
-                reserveRegister(&next_byte, op.source.register, grf_size);
+                try self.write(&access, op.destination, lanes, predicated);
+                try self.read(&access, op.source, lanes);
             },
-            .binary => |op| {
-                reserveRegister(&next_byte, op.destination.register, grf_size);
-                reserveRegister(&next_byte, op.lhs.register, grf_size);
-                reserveRegister(&next_byte, op.rhs.register, grf_size);
+            inline .binary, .math => |op| {
+                try self.write(&access, op.destination, lanes, predicated);
+                try self.read(&access, op.lhs, lanes);
+                try self.read(&access, op.rhs, lanes);
             },
             .compare => |op| {
-                reserveRegister(&next_byte, op.lhs.register, grf_size);
-                reserveRegister(&next_byte, op.rhs.register, grf_size);
-            },
-            .math => |op| {
-                reserveRegister(&next_byte, op.destination.register, grf_size);
-                reserveRegister(&next_byte, op.lhs.register, grf_size);
-                reserveRegister(&next_byte, op.rhs.register, grf_size);
+                try self.read(&access, op.lhs, lanes);
+                try self.read(&access, op.rhs, lanes);
             },
             .parallel_copy => return Error.ParallelCopiesNotLowered,
         }
+        return access;
     }
-    return next_byte;
+};
+
+fn mergeEdge(program: *const program_ir.Program, edge: instruction.Edge, live_in: []const bool, live: []bool) Error!void {
+    if (!program.blocks.isLive(edge.target))
+        return Error.InvalidProgram;
+
+    if (edge.arguments.len != 0)
+        return Error.BlockParametersNotLowered;
+
+    const successor = live_in[edge.target.index() * live.len ..][0..live.len];
+    for (live, successor) |*value, incoming|
+        value.* = value.* or incoming;
 }
 
-fn reserveRegister(next_byte: *usize, register: operand.RegisterRef, grf_size: usize) void {
-    switch (register) {
-        .physical_grf => |physical| reservePhysical(next_byte, physical, grf_size),
-        else => {},
+fn successorLive(program: *const program_ir.Program, block: instruction.Block, live_in: []const bool, live: []bool) Error!void {
+    if (block.parameters.items.len != 0)
+        return Error.BlockParametersNotLowered;
+
+    @memset(live, false);
+    switch (block.terminator orelse return Error.InvalidProgram) {
+        .jump => |edge| try mergeEdge(program, edge, live_in, live),
+        .conditional_branch => |branch| {
+            try mergeEdge(program, branch.true_edge, live_in, live);
+            try mergeEdge(program, branch.false_edge, live_in, live);
+        },
+        .end_thread, .@"unreachable" => {},
     }
 }
 
-fn reservePhysical(next_byte: *usize, physical: operand.PhysicalGrf, grf_size: usize) void {
-    const end = (@as(usize, physical.number) + 1) * grf_size;
-    next_byte.* = @max(next_byte.*, end);
+fn addClique(graph: *std.DynamicBitSetUnmanaged, live: []const bool) void {
+    for (live, 0..) |active, index| {
+        if (!active) continue;
+        for (live[0..index], 0..) |other_active, other| {
+            if (!other_active)
+                continue;
+            graph.set(index * live.len + other);
+            graph.set(other * live.len + index);
+        }
+    }
 }
 
 fn rewriteProgram(program: *program_ir.Program, allocations: []const ?operand.PhysicalGrf) Error!void {
     for (program.instructions.entries.items) |*entry| {
         const inst = if (entry.*) |*value| value else continue;
         switch (inst.operation) {
-            .load_global_invocation_id => |*op| try rewriteDestination(program, &op.destination, allocations),
+            .load_global_invocation_id, .load_num_workgroups => |*op| try rewriteDestination(program, &op.destination, allocations),
             .load_buffer => |*op| {
                 try rewriteDestination(program, &op.destination, allocations);
                 try rewriteSource(program, &op.byte_offset, allocations);
@@ -175,9 +406,12 @@ fn rewriteProgram(program: *program_ir.Program, allocations: []const ?operand.Ph
 
     for (program.blocks.entries.items) |*entry| {
         const block = if (entry.*) |*value| value else continue;
+
         if (block.parameters.items.len != 0)
             return Error.BlockParametersNotLowered;
+
         const terminator = if (block.terminator) |*value| value else return Error.InvalidProgram;
+
         switch (terminator.*) {
             .jump => |*edge| try rewriteEdge(program, edge, allocations),
             .conditional_branch => |*branch| {
@@ -209,8 +443,10 @@ fn rewriteRegister(program: *const program_ir.Program, register: *operand.Regist
         .virtual => |value| value,
         else => return,
     };
+
     if (!program.virtual_registers.isLive(virtual) or virtual.index() >= allocations.len)
         return Error.InvalidProgram;
+
     const physical = allocations[virtual.index()] orelse return Error.InvalidProgram;
     register.* = .{ .physical_grf = physical };
 }
@@ -269,6 +505,271 @@ test "[gen9] register allocation: assign non-overlapping physical GRFs" {
     try std.testing.expect(program.properties.registers_allocated);
 }
 
+fn defineRegister(program: *program_ir.Program, block: ids.BlockId, register: ids.VirtualRegisterId) !ids.InstructionId {
+    return program.appendInstruction(block, .simd8, null, .{ .move = .{
+        .destination = .{ .register = .{ .virtual = register }, .type = .u32 },
+        .source = .{ .register = .{ .immediate = .{ .u32 = 0 } }, .type = .u32, .region = operand.Region.scalar() },
+    } });
+}
+
+fn useRegister(program: *program_ir.Program, block: ids.BlockId, register: ids.VirtualRegisterId) !void {
+    _ = try program.appendInstruction(block, .simd8, null, .{ .move = .{
+        .destination = .{ .register = .null, .type = .u32 },
+        .source = source(register),
+    } });
+}
+
+fn assigned(program: *const program_ir.Program, id: ids.InstructionId) operand.PhysicalGrf {
+    return program.instructions.get(id).?.operation.move.destination.register.physical_grf;
+}
+
+const test_predicate = operand.Predicate{ .flag = .{ .physical = .{} } };
+
+test "[gen9] register allocation: reuse disjoint lifetimes across blocks" {
+    var limited_device = test_device;
+    limited_device.grf_count = 1;
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, limited_device, .simd8);
+    defer program.deinit();
+    const first = try addRegister(&program, 32, 32);
+    const second = try addRegister(&program, 32, 32);
+    const entry = try program.addBlock("entry");
+    const next = try program.addBlock("next");
+    const a = try defineRegister(&program, entry, first);
+    try useRegister(&program, entry, first);
+    try program.setTerminator(entry, .{ .jump = .{ .target = next, .arguments = &.{} } });
+    const b = try defineRegister(&program, next, second);
+    try useRegister(&program, next, second);
+    try program.setTerminator(next, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    try std.testing.expectEqual(assigned(&program, a), assigned(&program, b));
+    try std.testing.expectEqual(@as(u16, 1), program.program_data.total_grf_count);
+}
+
+test "[gen9] register allocation: more virtual registers than GRFs with disjoint lifetimes" {
+    var device = test_device;
+    device.grf_count = 1;
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, device, .simd8);
+    defer program.deinit();
+    const entry = try program.addBlock("entry");
+    var definitions: [256]ids.InstructionId = undefined;
+    for (&definitions) |*definition| {
+        const register = try addRegister(&program, 32, 32);
+        definition.* = try defineRegister(&program, entry, register);
+        try useRegister(&program, entry, register);
+    }
+    try program.setTerminator(entry, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    for (definitions) |definition|
+        try std.testing.expectEqual(operand.PhysicalGrf{ .number = 0 }, assigned(&program, definition));
+    try std.testing.expectEqual(@as(u16, 1), program.program_data.total_grf_count);
+}
+
+test "[gen9] register allocation: branch union and loop backedge liveness" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+    defer program.deinit();
+    const carried = try addRegister(&program, 32, 32);
+    const temporary = try addRegister(&program, 32, 32);
+    const exit_value = try addRegister(&program, 32, 32);
+    const entry = try program.addBlock("entry");
+    const loop = try program.addBlock("loop");
+    const body = try program.addBlock("body");
+    const exit = try program.addBlock("exit");
+    const a = try defineRegister(&program, entry, carried);
+    const c = try defineRegister(&program, entry, exit_value);
+    try program.setTerminator(entry, .{ .jump = .{ .target = loop, .arguments = &.{} } });
+    try useRegister(&program, loop, carried);
+    try program.setTerminator(loop, .{ .conditional_branch = .{
+        .predicate = test_predicate,
+        .true_edge = .{ .target = body, .arguments = &.{} },
+        .false_edge = .{ .target = exit, .arguments = &.{} },
+    } });
+    const b = try defineRegister(&program, body, temporary);
+    try useRegister(&program, body, temporary);
+    try program.setTerminator(body, .{ .jump = .{ .target = loop, .arguments = &.{} } });
+    try useRegister(&program, exit, exit_value);
+    try program.setTerminator(exit, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    try std.testing.expect(assigned(&program, a).number != assigned(&program, b).number);
+    try std.testing.expect(assigned(&program, c).number != assigned(&program, b).number);
+    try std.testing.expect(assigned(&program, a).number != assigned(&program, c).number);
+    try std.testing.expectEqual(@as(u16, 3), program.program_data.total_grf_count);
+}
+
+test "[gen9] register allocation: only full unpredicated writes kill" {
+    for (0..5) |mode| {
+        var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+        defer program.deinit();
+        const value = try addRegister(&program, 32, 32);
+        const temporary = try addRegister(&program, 32, 32);
+        const entry = try program.addBlock("entry");
+        const a = try defineRegister(&program, entry, value);
+        try useRegister(&program, entry, value);
+        const b = try defineRegister(&program, entry, temporary);
+        try useRegister(&program, entry, temporary);
+        _ = try program.appendInstruction(entry, if (mode == 0 or mode == 4) .simd8 else .simd4, if (mode == 4) test_predicate else null, .{ .move = .{
+            .destination = .{
+                .register = .{ .virtual = value },
+                .type = .u32,
+                .region = .{ .byte_offset = if (mode == 1) 16 else 0, .horizontal_stride = if (mode == 2) 2 else 1 },
+            },
+            .source = .{ .register = .{ .immediate = .{ .u32 = 1 } }, .type = .u32, .region = operand.Region.scalar() },
+        } });
+        try useRegister(&program, entry, value);
+        try program.setTerminator(entry, .end_thread);
+        markPrerequisites(&program);
+        try run(std.testing.allocator, &program);
+        try std.testing.expectEqual(mode == 0, assigned(&program, a).number == assigned(&program, b).number);
+    }
+}
+
+test "[gen9] register allocation: high fixed SEND spans leave usable holes" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+    defer program.deinit();
+    program.program_data.payload_grf_count = 1;
+    program.payload.header_grf = .{ .number = 0 };
+    const large = try addRegister(&program, 119 * 32, 32);
+    const small = try addRegister(&program, 3 * 32, 32);
+    const entry = try program.addBlock("entry");
+    const a = try defineRegister(&program, entry, large);
+    const b = try defineRegister(&program, entry, small);
+    _ = try program.appendInstruction(entry, .simd8, null, .{ .surface_message = .{
+        .kind = .read,
+        .binding_table = 0,
+        .payload = .{ .base = .{ .physical_grf = .{ .number = 120 } }, .register_count = 2 },
+        .response = .{ .base = .{ .physical_grf = .{ .number = 125 } }, .register_count = 3 },
+        .data_type = .u32,
+    } });
+    try useRegister(&program, entry, large);
+    try useRegister(&program, entry, small);
+    try program.setTerminator(entry, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    try std.testing.expectEqual(@as(u16, 1), assigned(&program, a).number);
+    try std.testing.expectEqual(@as(u16, 122), assigned(&program, b).number);
+    try std.testing.expectEqual(@as(u16, 128), program.program_data.total_grf_count);
+}
+
+test "[gen9] register allocation: full GRF capacity and simultaneous exhaustion" {
+    for ([_]usize{ 128, 129 }) |count| {
+        var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+        defer program.deinit();
+        const entry = try program.addBlock("entry");
+        var registers: [129]ids.VirtualRegisterId = undefined;
+        var definitions: [129]ids.InstructionId = undefined;
+        for (0..count) |index| {
+            registers[index] = try addRegister(&program, 32, 32);
+            definitions[index] = try defineRegister(&program, entry, registers[index]);
+        }
+        for (registers[0..count]) |register| try useRegister(&program, entry, register);
+        try program.setTerminator(entry, .end_thread);
+        markPrerequisites(&program);
+        if (count == 129) {
+            try std.testing.expectError(Error.OutOfRegisters, run(std.testing.allocator, &program));
+            try std.testing.expect(!program.properties.registers_allocated);
+            try std.testing.expect(program.instructions.get(definitions[0]).?.operation.move.destination.register == .virtual);
+        } else {
+            try run(std.testing.allocator, &program);
+            for (definitions[0..count], 0..) |definition, index|
+                try std.testing.expectEqual(@as(u16, @intCast(index)), assigned(&program, definition).number);
+            try std.testing.expectEqual(@as(u16, 128), program.program_data.total_grf_count);
+        }
+    }
+}
+
+test "[gen9] register allocation: fixed regions reserve bytes across GRF boundaries" {
+    var device = test_device;
+    device.grf_count = 3;
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, device, .simd8);
+    defer program.deinit();
+    program.program_data.payload_grf_count = 1;
+    const first = try addRegister(&program, 28, 4);
+    const second = try addRegister(&program, 4, 4);
+    const entry = try program.addBlock("entry");
+    var definitions: [2]ids.InstructionId = undefined;
+    for ([_]ids.VirtualRegisterId{ first, second }, 0..) |register, index| {
+        definitions[index] = try program.appendInstruction(entry, .simd1, null, .{ .move = .{
+            .destination = .{ .register = .{ .virtual = register }, .type = .u32 },
+            .source = .{ .register = .{ .immediate = .{ .u32 = 0 } }, .type = .u32, .region = operand.Region.scalar() },
+        } });
+    }
+    _ = try program.appendInstruction(entry, .simd8, null, .{ .move = .{
+        .destination = .{ .register = .null, .type = .u32 },
+        .source = .{
+            .register = .{ .physical_grf = .{ .number = 1, .byte_offset = 4 } },
+            .type = .u32,
+            .region = .{ .byte_offset = 24, .vertical_stride = 8, .width = 8, .horizontal_stride = 1 },
+        },
+    } });
+    for ([_]ids.VirtualRegisterId{ first, second }) |register| {
+        _ = try program.appendInstruction(entry, .simd1, null, .{ .move = .{
+            .destination = .{ .register = .null, .type = .u32 },
+            .source = .{ .register = .{ .virtual = register }, .type = .u32, .region = operand.Region.scalar() },
+        } });
+    }
+    try program.setTerminator(entry, .end_thread);
+    markPrerequisites(&program);
+    try run(std.testing.allocator, &program);
+    try std.testing.expectEqual(operand.PhysicalGrf{ .number = 1 }, assigned(&program, definitions[0]));
+    try std.testing.expectEqual(operand.PhysicalGrf{ .number = 2, .byte_offset = 28 }, assigned(&program, definitions[1]));
+    try std.testing.expectEqual(@as(u16, 3), program.program_data.total_grf_count);
+}
+
+test "[gen9] register allocation: SEND response overwrite and payload interference" {
+    for ([_]bool{ false, true }) |partial| {
+        var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+        defer program.deinit();
+        const response = try addRegister(&program, 64, 4);
+        const temporary = try addRegister(&program, 32, 32);
+        const payload = try addRegister(&program, 32, 4);
+        const entry = try program.addBlock("entry");
+        const a = try defineRegister(&program, entry, response);
+        try useRegister(&program, entry, response);
+        const b = try defineRegister(&program, entry, temporary);
+        try useRegister(&program, entry, temporary);
+        const p = try defineRegister(&program, entry, payload);
+        const send = try program.appendInstruction(entry, .simd8, null, .{ .surface_message = .{
+            .kind = .read,
+            .binding_table = 0,
+            .payload = .{ .base = .{ .virtual = payload }, .register_count = 1 },
+            .response = .{ .base = .{ .virtual = response }, .register_count = if (partial) 1 else 2 },
+            .data_type = .u32,
+        } });
+        try useRegister(&program, entry, response);
+        try program.setTerminator(entry, .end_thread);
+        markPrerequisites(&program);
+        try run(std.testing.allocator, &program);
+        try std.testing.expectEqual(!partial, assigned(&program, a).number == assigned(&program, b).number);
+        try std.testing.expect(assigned(&program, p).number >= assigned(&program, a).number + 2);
+        const message = program.instructions.get(send).?.operation.surface_message;
+        try std.testing.expectEqual(assigned(&program, a), message.response.?.base.physical_grf);
+        try std.testing.expectEqual(assigned(&program, p), message.payload.base.physical_grf);
+        try std.testing.expectEqual(@as(u8, 0), message.payload.base.physical_grf.byte_offset);
+    }
+}
+
+test "[gen9] register allocation: reject fixed spans beyond capacity before rewriting" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device, .simd8);
+    defer program.deinit();
+    const register = try addRegister(&program, 32, 32);
+    const entry = try program.addBlock("entry");
+    const definition = try defineRegister(&program, entry, register);
+    _ = try program.appendInstruction(entry, .simd8, null, .{ .surface_message = .{
+        .kind = .read,
+        .binding_table = 0,
+        .payload = .{ .base = .{ .physical_grf = .{ .number = 0 } }, .register_count = 1 },
+        .response = .{ .base = .{ .physical_grf = .{ .number = 127 } }, .register_count = 2 },
+        .data_type = .u32,
+    } });
+    try program.setTerminator(entry, .end_thread);
+    markPrerequisites(&program);
+    try std.testing.expectError(Error.InvalidProgram, run(std.testing.allocator, &program));
+    try std.testing.expect(!program.properties.registers_allocated);
+    try std.testing.expect(program.instructions.get(definition).?.operation.move.destination.register == .virtual);
+}
+
 test "[gen9] register allocation: report GRF exhaustion" {
     var limited_device = test_device;
     limited_device.grf_count = 2;

src/intel/compiler/targets/gen9/validator.zig

@@ -53,7 +53,7 @@ fn validateInstruction(inst: instruction.Instruction) Error!void {
     if (inst.predicate) |predicate|
         try validateFlag(predicate.flag);
     switch (inst.operation) {
-        .load_global_invocation_id => |op| try validateDestination(op.destination),
+        .load_global_invocation_id, .load_num_workgroups => |op| try validateDestination(op.destination),
         .load_buffer => |op| {
             try validateBufferReference(op.buffer);
             try validateDestination(op.destination);
@@ -172,11 +172,21 @@ fn validateEdge(edge: instruction.Edge) Error!void {
 fn validateFlag(flag: operand.FlagRef) Error!void {
     switch (flag) {
         .virtual => {},
-        .physical => |physical| if (physical.register != 0 or physical.subregister > 1)
+        .physical => |physical| if (physical.register > 1 or physical.subregister > 1)
             return Error.InvalidPhysicalFlag,
     }
 }
 
+test "[gen9] validator: accept all four flag halves and reject out of range flags" {
+    for (0..2) |register| {
+        for (0..2) |subregister| {
+            try validateFlag(.{ .physical = .{ .register = @intCast(register), .subregister = @intCast(subregister) } });
+        }
+    }
+    try std.testing.expectError(Error.InvalidPhysicalFlag, validateFlag(.{ .physical = .{ .register = 2 } }));
+    try std.testing.expectError(Error.InvalidPhysicalFlag, validateFlag(.{ .physical = .{ .register = 1, .subregister = 2 } }));
+}
+
 fn validatePayload(program: *const program_ir.Program) Error!void {
     if (program.payload.header_grf) |header| {
         if (header.number != 0 or header.byte_offset != 0)

src/intel/i915/kmd.zig

@@ -123,19 +123,7 @@ pub const Device = struct {
                 group_index = groups.items.len - 1;
             }
 
-            const domain: u32 = switch (relocation.domain) {
-                .none => 0,
-                .render => _i915.gem_domain_render,
-                .instruction => _i915.gem_domain_instruction,
-            };
-            groups.items[group_index].entries.append(allocator, .{
-                .target_handle = relocation.target_handle,
-                .delta = relocation.delta,
-                .offset = relocation.offset,
-                .presumed_offset = 0,
-                .read_domains = if (relocation.read) domain else 0,
-                .write_domain = if (relocation.write) domain else 0,
-            }) catch return VkError.OutOfHostMemory;
+            groups.items[group_index].entries.append(allocator, relocationEntry(relocation)) catch return VkError.OutOfHostMemory;
         }
 
         var objects = std.ArrayList(_i915.ExecObject2).empty;
@@ -202,6 +190,24 @@ pub const Device = struct {
     }
 };
 
+fn relocationEntry(relocation: common_kmd.Relocation) _i915.RelocationEntry {
+    const domain: u32 = switch (relocation.domain) {
+        .none => 0,
+        .render => _i915.gem_domain_render,
+        .instruction => _i915.gem_domain_instruction,
+    };
+    return .{
+        .target_handle = relocation.target_handle,
+        .delta = relocation.delta,
+        .offset = relocation.offset,
+        // GPU address zero is valid. These locations have not been patched yet,
+        // so never let i915 skip the initial relocation, including its delta.
+        .presumed_offset = std.math.maxInt(u64),
+        .read_domains = if (relocation.read) domain else 0,
+        .write_domain = if (relocation.write) domain else 0,
+    };
+}
+
 pub const Memory = struct {
     handle: u32,
     size: vk.DeviceSize,
@@ -292,3 +298,21 @@ pub const Memory = struct {
         ) catch return VkError.DeviceLost;
     }
 };
+
+test "[i915] initial relocations force patching even at GPU address zero" {
+    const entry = relocationEntry(.{
+        .source_handle = 4,
+        .target_handle = 4,
+        .offset = 64,
+        .delta = 1920,
+        .read = true,
+        .write = false,
+        .domain = .render,
+    });
+    try std.testing.expectEqual(std.math.maxInt(u64), entry.presumed_offset);
+    try std.testing.expectEqual(@as(u32, 1920), entry.delta);
+    try std.testing.expectEqual(@as(u64, 64), entry.offset);
+    try std.testing.expectEqual(@as(u32, 4), entry.target_handle);
+    try std.testing.expectEqual(_i915.gem_domain_render, entry.read_domains);
+    try std.testing.expectEqual(@as(u32, 0), entry.write_domain);
+}