[Flint] adding multiple registers/blocks support

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;