[Flint] implement parallel copy lowering

src/intel/compiler/lower/lower.zig

@@ -1,2 +1,3 @@
 pub const block_arguments = @import("block_arguments.zig");
 pub const common_ir = @import("common_ir.zig");
+pub const parallel_copies = @import("parallel_copies.zig");

src/intel/compiler/lower/parallel_copies.zig

@@ -0,0 +1,466 @@
+const std = @import("std");
+const device = @import("../device.zig");
+const Builder = @import("../ir/Builder.zig");
+const ids = @import("../ir/id.zig");
+const instruction = @import("../ir/instruction.zig");
+const operand = @import("../ir/operand.zig");
+const program_ir = @import("../ir/program.zig");
+const pseudo = @import("../ir/pseudo.zig");
+const validator = @import("../ir/validator.zig");
+
+pub const Error = std.mem.Allocator.Error || error{
+    InvalidProgram,
+};
+
+const EmittedInstruction = struct {
+    predicate: ?operand.Predicate = null,
+    operation: instruction.Operation,
+};
+
+const FlagValue = union(enum) {
+    constant: bool,
+    snapshot: ids.VirtualRegisterId,
+};
+
+const FlagWrite = struct {
+    destination: ids.VirtualFlagId,
+    value: FlagValue,
+};
+
+pub fn run(allocator: std.mem.Allocator, program: *program_ir.Program) Error!void {
+    validator.validate(program) catch return error.InvalidProgram;
+    if (program.properties.parallel_copies_lowered)
+        return;
+
+    var builder = Builder.init(program);
+    for (program.blocks.entries.items, 0..) |entry, block_index| {
+        _ = entry orelse continue;
+        const block_id = ids.BlockId.fromIndex(block_index);
+        var instruction_index: usize = 0;
+
+        while (true) {
+            const block = program.blocks.get(block_id) orelse return error.InvalidProgram;
+            if (instruction_index >= block.instructions.items.len)
+                break;
+
+            const instruction_id = block.instructions.items[instruction_index];
+            const inst = program.instructions.get(instruction_id) orelse return error.InvalidProgram;
+            const parallel_copy = switch (inst.operation) {
+                .parallel_copy => |copy| copy,
+                else => {
+                    instruction_index += 1;
+                    continue;
+                },
+            };
+            if (inst.predicate != null)
+                return error.InvalidProgram;
+            const execution_size = inst.execution_size;
+
+            var emitted: std.ArrayList(EmittedInstruction) = .empty;
+            defer emitted.deinit(allocator);
+            try lowerParallelCopy(allocator, &builder, execution_size, parallel_copy, &emitted);
+
+            if (emitted.items.len == 0) {
+                const mutable_block = program.blocks.getMut(block_id) orelse return error.InvalidProgram;
+                const removed_id = mutable_block.instructions.orderedRemove(instruction_index);
+                if (removed_id != instruction_id or !program.instructions.remove(instruction_id))
+                    return error.InvalidProgram;
+                continue;
+            }
+
+            builder.replaceOperation(instruction_id, emitted.items[0].operation) catch |err|
+                return mapBuilderError(err);
+            const replacement = program.instructions.getMut(instruction_id) orelse return error.InvalidProgram;
+            replacement.predicate = emitted.items[0].predicate;
+
+            for (emitted.items[1..], 1..) |item, offset| {
+                _ = builder.insertInstruction(
+                    block_id,
+                    instruction_index + offset,
+                    execution_size,
+                    item.predicate,
+                    item.operation,
+                ) catch |err| return mapBuilderError(err);
+            }
+            instruction_index += emitted.items.len;
+        }
+    }
+
+    program.properties.parallel_copies_lowered = true;
+    validator.validate(program) catch return error.InvalidProgram;
+}
+
+fn lowerParallelCopy(
+    allocator: std.mem.Allocator,
+    builder: *Builder,
+    execution_size: device.ExecutionSize,
+    copy: pseudo.ParallelCopy,
+    emitted: *std.ArrayList(EmittedInstruction),
+) Error!void {
+    var pending_registers: std.ArrayList(pseudo.RegisterCopy) = .empty;
+    defer pending_registers.deinit(allocator);
+    for (copy.register_copies) |item| {
+        if (!isRegisterIdentity(item, execution_size))
+            try pending_registers.append(allocator, item);
+    }
+
+    var flag_writes: std.ArrayList(FlagWrite) = .empty;
+    defer flag_writes.deinit(allocator);
+    try snapshotFlagSources(allocator, builder, execution_size, copy.flag_copies, emitted, &flag_writes);
+    try scheduleRegisterCopies(allocator, builder, execution_size, &pending_registers, emitted);
+    try emitFlagWrites(allocator, execution_size, flag_writes.items, emitted);
+}
+
+fn scheduleRegisterCopies(
+    allocator: std.mem.Allocator,
+    builder: *Builder,
+    execution_size: device.ExecutionSize,
+    pending: *std.ArrayList(pseudo.RegisterCopy),
+    emitted: *std.ArrayList(EmittedInstruction),
+) Error!void {
+    while (pending.items.len != 0) {
+        if (findReadyCopy(pending.items)) |ready_index| {
+            const ready = pending.orderedRemove(ready_index);
+            try emitted.append(allocator, .{ .operation = .{ .move = .{
+                .destination = ready.destination,
+                .source = ready.source,
+            } } });
+            continue;
+        }
+
+        const cycle_copy = &pending.items[0];
+        const destination_id = destinationVirtualRegister(cycle_copy.destination) orelse
+            return error.InvalidProgram;
+        const destination_register = builder.program.virtual_registers.get(destination_id) orelse
+            return error.InvalidProgram;
+        const temporary = builder.addVirtualRegister(.{
+            .size_bytes = destination_register.size_bytes,
+            .alignment_bytes = destination_register.alignment_bytes,
+            .element_type = destination_register.element_type,
+            .lane_count = destination_register.lane_count,
+            .class = .temporary,
+            .spillable = destination_register.spillable,
+        }) catch |err| return mapBuilderError(err);
+
+        var temporary_destination = cycle_copy.destination;
+        temporary_destination.register = .{ .virtual = temporary };
+        try emitted.append(allocator, .{ .operation = .{ .move = .{
+            .destination = temporary_destination,
+            .source = cycle_copy.source,
+        } } });
+
+        cycle_copy.source = .{
+            .register = .{ .virtual = temporary },
+            .type = cycle_copy.source.type,
+            .region = operand.Region.contiguous(execution_size),
+        };
+    }
+}
+
+fn findReadyCopy(pending: []const pseudo.RegisterCopy) ?usize {
+    for (pending, 0..) |candidate, candidate_index| {
+        const destination_id = destinationVirtualRegister(candidate.destination) orelse continue;
+        var destination_is_source = false;
+        for (pending, 0..) |other, other_index| {
+            if (candidate_index == other_index)
+                continue;
+            switch (other.source.register) {
+                .virtual => |source_id| if (source_id == destination_id) {
+                    destination_is_source = true;
+                    break;
+                },
+                else => {},
+            }
+        }
+        if (!destination_is_source)
+            return candidate_index;
+    }
+    return null;
+}
+
+fn snapshotFlagSources(
+    allocator: std.mem.Allocator,
+    builder: *Builder,
+    execution_size: device.ExecutionSize,
+    copies: []const pseudo.FlagCopy,
+    emitted: *std.ArrayList(EmittedInstruction),
+    writes: *std.ArrayList(FlagWrite),
+) Error!void {
+    for (copies) |copy| {
+        if (isFlagIdentity(copy))
+            continue;
+
+        const value: FlagValue = switch (copy.source) {
+            .constant => |constant| .{ .constant = constant },
+            .dynamic => |predicate| value: {
+                const temporary = builder.addVirtualRegister(.{
+                    .size_bytes = @as(u32, @intFromEnum(execution_size)) * @sizeOf(u32),
+                    .alignment_bytes = builder.program.device_info.grf_size_bytes,
+                    .element_type = .u32,
+                    .lane_count = @intFromEnum(execution_size),
+                    .class = .temporary,
+                }) catch |err| return mapBuilderError(err);
+                const destination: operand.Destination = .{
+                    .register = .{ .virtual = temporary },
+                    .type = .u32,
+                };
+                try emitted.append(allocator, .{ .operation = .{ .move = .{
+                    .destination = destination,
+                    .source = immediateU32(0),
+                } } });
+                try emitted.append(allocator, .{
+                    .predicate = predicate,
+                    .operation = .{ .move = .{
+                        .destination = destination,
+                        .source = immediateU32(1),
+                    } },
+                });
+                break :value .{ .snapshot = temporary };
+            },
+        };
+        try writes.append(allocator, .{
+            .destination = copy.destination,
+            .value = value,
+        });
+    }
+}
+
+fn emitFlagWrites(
+    allocator: std.mem.Allocator,
+    execution_size: device.ExecutionSize,
+    writes: []const FlagWrite,
+    emitted: *std.ArrayList(EmittedInstruction),
+) Error!void {
+    for (writes) |write| {
+        const value = switch (write.value) {
+            .constant => |constant| immediateU32(@intFromBool(constant)),
+            .snapshot => |temporary| operand.Source{
+                .register = .{ .virtual = temporary },
+                .type = .u32,
+                .region = operand.Region.contiguous(execution_size),
+            },
+        };
+        try emitted.append(allocator, .{ .operation = .{ .compare = .{
+            .opcode = .not_equal,
+            .destination = .{ .virtual = write.destination },
+            .lhs = value,
+            .rhs = immediateU32(0),
+        } } });
+    }
+}
+
+fn isRegisterIdentity(copy: pseudo.RegisterCopy, execution_size: device.ExecutionSize) bool {
+    const destination_id = destinationVirtualRegister(copy.destination) orelse return false;
+    const source_id = switch (copy.source.register) {
+        .virtual => |id| id,
+        else => return false,
+    };
+    if (destination_id != source_id or copy.source.negate or copy.source.absolute)
+        return false;
+
+    const contiguous = operand.Region.contiguous(execution_size);
+    return copy.destination.type == copy.source.type and
+        copy.destination.region.byte_offset == contiguous.byte_offset and
+        copy.destination.region.horizontal_stride == 1 and
+        copy.source.region.byte_offset == contiguous.byte_offset and
+        copy.source.region.vertical_stride == contiguous.vertical_stride and
+        copy.source.region.width == contiguous.width and
+        copy.source.region.horizontal_stride == contiguous.horizontal_stride;
+}
+
+fn isFlagIdentity(copy: pseudo.FlagCopy) bool {
+    return switch (copy.source) {
+        .constant => false,
+        .dynamic => |predicate| !predicate.inverse and switch (predicate.flag) {
+            .virtual => |source| source == copy.destination,
+            .physical => false,
+        },
+    };
+}
+
+fn destinationVirtualRegister(destination: operand.Destination) ?ids.VirtualRegisterId {
+    return switch (destination.register) {
+        .virtual => |id| id,
+        else => null,
+    };
+}
+
+fn immediateU32(value: u32) operand.Source {
+    return .{
+        .register = .{ .immediate = .{ .u32 = value } },
+        .type = .u32,
+        .region = operand.Region.broadcast(),
+    };
+}
+
+fn mapBuilderError(err: anyerror) Error {
+    return switch (err) {
+        error.OutOfMemory => error.OutOfMemory,
+        else => error.InvalidProgram,
+    };
+}
+
+const test_device_info: device.DeviceInfo = .{
+    .generation = .gen9,
+    .platform = .skylake,
+    .pci_device_id = 0x1912,
+    .grf_count = 128,
+};
+
+fn addTestRegister(builder: *Builder) !ids.VirtualRegisterId {
+    return builder.addVirtualRegister(.{
+        .size_bytes = 32,
+        .alignment_bytes = 32,
+        .element_type = .u32,
+        .lane_count = 8,
+        .class = .temporary,
+    });
+}
+
+fn testDestination(register: ids.VirtualRegisterId) operand.Destination {
+    return .{
+        .register = .{ .virtual = register },
+        .type = .u32,
+    };
+}
+
+fn testSource(register: ids.VirtualRegisterId) operand.Source {
+    return .{
+        .register = .{ .virtual = register },
+        .type = .u32,
+        .region = operand.Region.contiguous(.simd8),
+    };
+}
+
+test "[intel] parallel copies: lower independent copies" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device_info, .simd8);
+    defer program.deinit();
+    var builder = Builder.init(&program);
+
+    const source_a = try addTestRegister(&builder);
+    const source_b = try addTestRegister(&builder);
+    const destination_a = try addTestRegister(&builder);
+    const destination_b = try addTestRegister(&builder);
+    const entry = try builder.addBlock("entry");
+    const copies = [_]pseudo.RegisterCopy{
+        .{ .destination = testDestination(destination_a), .source = testSource(source_a) },
+        .{ .destination = testDestination(destination_b), .source = testSource(source_b) },
+    };
+    _ = try builder.appendInstruction(entry, .simd8, null, .{ .parallel_copy = .{
+        .register_copies = &copies,
+        .flag_copies = &.{},
+    } });
+    try builder.setTerminator(entry, .end_thread);
+
+    try validator.validate(&program);
+    try run(std.testing.allocator, &program);
+    try validator.validate(&program);
+
+    try std.testing.expect(program.properties.parallel_copies_lowered);
+    const instructions = program.blocks.get(entry).?.instructions.items;
+    try std.testing.expectEqual(@as(usize, 2), instructions.len);
+    const first = program.instructions.get(instructions[0]).?;
+    const second = program.instructions.get(instructions[1]).?;
+    try std.testing.expect(first.operation == .move);
+    try std.testing.expect(second.operation == .move);
+    try std.testing.expectEqual(destination_a, first.operation.move.destination.register.virtual);
+    try std.testing.expectEqual(source_a, first.operation.move.source.register.virtual);
+    try std.testing.expectEqual(destination_b, second.operation.move.destination.register.virtual);
+    try std.testing.expectEqual(source_b, second.operation.move.source.register.virtual);
+}
+
+test "[intel] parallel copies: remove identity copies" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device_info, .simd8);
+    defer program.deinit();
+    var builder = Builder.init(&program);
+
+    const register = try addTestRegister(&builder);
+    const entry = try builder.addBlock("entry");
+    const copies = [_]pseudo.RegisterCopy{.{
+        .destination = testDestination(register),
+        .source = testSource(register),
+    }};
+    const copy_id = try builder.appendInstruction(entry, .simd8, null, .{ .parallel_copy = .{
+        .register_copies = &copies,
+        .flag_copies = &.{},
+    } });
+    try builder.setTerminator(entry, .end_thread);
+
+    try validator.validate(&program);
+    try run(std.testing.allocator, &program);
+    try validator.validate(&program);
+
+    try std.testing.expect(program.properties.parallel_copies_lowered);
+    try std.testing.expectEqual(@as(usize, 0), program.blocks.get(entry).?.instructions.items.len);
+    try std.testing.expect(program.instructions.get(copy_id) == null);
+}
+
+test "[intel] parallel copies: break a two-register cycle" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device_info, .simd8);
+    defer program.deinit();
+    var builder = Builder.init(&program);
+
+    const register_a = try addTestRegister(&builder);
+    const register_b = try addTestRegister(&builder);
+    const entry = try builder.addBlock("entry");
+    const copies = [_]pseudo.RegisterCopy{
+        .{ .destination = testDestination(register_a), .source = testSource(register_b) },
+        .{ .destination = testDestination(register_b), .source = testSource(register_a) },
+    };
+    _ = try builder.appendInstruction(entry, .simd8, null, .{ .parallel_copy = .{
+        .register_copies = &copies,
+        .flag_copies = &.{},
+    } });
+    try builder.setTerminator(entry, .end_thread);
+
+    try validator.validate(&program);
+    try run(std.testing.allocator, &program);
+    try validator.validate(&program);
+
+    const instructions = program.blocks.get(entry).?.instructions.items;
+    try std.testing.expectEqual(@as(usize, 3), instructions.len);
+    const snapshot = program.instructions.get(instructions[0]).?.operation.move;
+    const restore_b = program.instructions.get(instructions[1]).?.operation.move;
+    const restore_a = program.instructions.get(instructions[2]).?.operation.move;
+    const temporary = snapshot.destination.register.virtual;
+
+    try std.testing.expect(temporary != register_a and temporary != register_b);
+    try std.testing.expectEqual(register_b, snapshot.source.register.virtual);
+    try std.testing.expectEqual(register_b, restore_b.destination.register.virtual);
+    try std.testing.expectEqual(register_a, restore_b.source.register.virtual);
+    try std.testing.expectEqual(register_a, restore_a.destination.register.virtual);
+    try std.testing.expectEqual(temporary, restore_a.source.register.virtual);
+    try std.testing.expectEqual(operand.RegisterClass.temporary, program.virtual_registers.get(temporary).?.class);
+}
+
+test "[intel] parallel copies: snapshot flag cycles" {
+    var program = program_ir.Program.init(std.testing.allocator, .{ 1, 1, 1 }, test_device_info, .simd8);
+    defer program.deinit();
+    var builder = Builder.init(&program);
+
+    const flag_a = try builder.addVirtualFlag(.{});
+    const flag_b = try builder.addVirtualFlag(.{});
+    const entry = try builder.addBlock("entry");
+    const copies = [_]pseudo.FlagCopy{
+        .{ .destination = flag_a, .source = .{ .dynamic = .{ .flag = .{ .virtual = flag_b } } } },
+        .{ .destination = flag_b, .source = .{ .dynamic = .{ .flag = .{ .virtual = flag_a } } } },
+    };
+    _ = try builder.appendInstruction(entry, .simd8, null, .{ .parallel_copy = .{
+        .register_copies = &.{},
+        .flag_copies = &copies,
+    } });
+    try builder.setTerminator(entry, .end_thread);
+
+    try validator.validate(&program);
+    try run(std.testing.allocator, &program);
+    try validator.validate(&program);
+
+    const instructions = program.blocks.get(entry).?.instructions.items;
+    try std.testing.expectEqual(@as(usize, 6), instructions.len);
+    for (instructions[0..4]) |instruction_id|
+        try std.testing.expect(program.instructions.get(instruction_id).?.operation == .move);
+    try std.testing.expect(program.instructions.get(instructions[4]).?.operation == .compare);
+    try std.testing.expect(program.instructions.get(instructions[5]).?.operation == .compare);
+    try std.testing.expectEqual(flag_a, program.instructions.get(instructions[4]).?.operation.compare.destination.virtual);
+    try std.testing.expectEqual(flag_b, program.instructions.get(instructions[5]).?.operation.compare.destination.virtual);
+}

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

@@ -3,6 +3,7 @@ const shader_ir = @import("shader_ir").ir;
 const device = @import("../../device.zig");
 const program_ir = @import("../../ir/program.zig");
 const common_ir = @import("../../lower/common_ir.zig");
+const parallel_copies = @import("../../lower/parallel_copies.zig");
 
 pub const compute = @import("compute/compute.zig");
 pub const validator = @import("validator.zig");
@@ -17,12 +18,7 @@ pub const Error = common_ir.Error || compute.Error || error{
     UnsupportedGrfSize,
 };
 
-pub fn lower(
+pub fn lower(allocator: std.mem.Allocator, module: *shader_ir.module.Module, device_info: device.DeviceInfo, options: Options) Error!program_ir.Program {
-    allocator: std.mem.Allocator,
-    module: *shader_ir.module.Module,
-    device_info: device.DeviceInfo,
-    options: Options,
-) Error!program_ir.Program {
     if (device_info.generation != .gen9)
         return Error.UnsupportedGeneration;
     if (module.stage != .compute)
@@ -36,6 +32,10 @@ pub fn lower(
 
     var program = try common_ir.lower(allocator, module, device_info, options);
     errdefer program.deinit();
+    parallel_copies.run(allocator, &program) catch |err| return switch (err) {
+        error.OutOfMemory => Error.OutOfMemory,
+        error.InvalidProgram => Error.InvalidLoweredProgram,
+    };
     validator.validate(&program) catch return Error.InvalidLoweredProgram;
     return program;
 }
@@ -78,3 +78,103 @@ test "[gen9] target: reject unsupported target configurations" {
     wide_grf.grf_size_bytes = 64;
     try std.testing.expectError(Error.UnsupportedGrfSize, lower(std.testing.allocator, &module, wide_grf, .{}));
 }
+
+test "[gen9] target: lower 256 KiB SSBO copy loop" {
+    const source =
+        \\shader compute @main
+        \\{
+        \\    @source: vec4[u32] = storage_buffer[set(0), binding(0)]
+        \\    @destination: vec4[u32] = storage_buffer[set(0), binding(1)]
+        \\    %zero: constant i32 = bits(0x0)
+        \\    %one: constant i32 = bits(0x1)
+        \\    %stride: constant i32 = bits(0x10)
+        \\    %element_count: constant i32 = bits(0x4000)
+        \\    fn @main() -> void
+        \\    {
+        \\        .entry():
+        \\            branch .header(%zero)
+        \\        .header(%index: i32):
+        \\            %in_bounds: bool = cmp_signed_less %index, %element_count
+        \\            conditional_branch %in_bounds, .body(), .exit()
+        \\        .body():
+        \\            %signed_offset: i32 = integer_multiply %index, %stride
+        \\            %offset: u32 = bitcast %signed_offset
+        \\            %value: vec4[u32] = load_buffer @source, %offset
+        \\            store_buffer @destination, %offset, %value
+        \\            branch .continue()
+        \\        .continue():
+        \\            %next: i32 = integer_add %index, %one
+        \\            branch .header(%next)
+        \\        .exit():
+        \\            return
+        \\    }
+        \\}
+    ;
+
+    var module = try shader_ir.parser.parseString(std.testing.allocator, source);
+    defer module.deinit();
+    module.execution_modes.workgroup_size = .{ 1, 1, 1 };
+
+    const gen9_device: device.DeviceInfo = .{
+        .generation = .gen9,
+        .platform = .skylake,
+        .pci_device_id = 0x1912,
+        .grf_count = 128,
+    };
+    var program = try lower(std.testing.allocator, &module, gen9_device, .{});
+    defer program.deinit();
+
+    try std.testing.expect(program.properties.common_ir_lowered);
+    try std.testing.expect(program.properties.block_parameters_lowered);
+    try std.testing.expect(program.properties.parallel_copies_lowered);
+
+    var resources = try compute.ResourceLayout.init(std.testing.allocator, &program);
+    defer resources.deinit(std.testing.allocator);
+    try std.testing.expectEqual(@as(usize, 2), resources.bindings.len);
+    try std.testing.expectEqual(compute.resource_layout.Binding{
+        .set = 0,
+        .binding = 0,
+        .binding_table_index = 0,
+    }, resources.bindings[0]);
+    try std.testing.expectEqual(compute.resource_layout.Binding{
+        .set = 0,
+        .binding = 1,
+        .binding_table_index = 1,
+    }, resources.bindings[1]);
+
+    try lowerComputeResources(&program, &resources);
+    try validator.validate(&program);
+
+    var load_offsets: [4]bool = @splat(false);
+    var store_offsets: [4]bool = @splat(false);
+    var load_count: usize = 0;
+    var store_count: usize = 0;
+    for (program.instructions.entries.items) |instruction_entry| {
+        const inst = instruction_entry orelse continue;
+        switch (inst.operation) {
+            .load_buffer => |operation| {
+                try std.testing.expectEqual(@as(u8, 0), operation.buffer.binding_table);
+                try std.testing.expect(operation.immediate_offset % @sizeOf(u32) == 0);
+                const component = operation.immediate_offset / @sizeOf(u32);
+                try std.testing.expect(component < load_offsets.len);
+                load_offsets[component] = true;
+                load_count += 1;
+            },
+            .store_buffer => |operation| {
+                try std.testing.expectEqual(@as(u8, 1), operation.buffer.binding_table);
+                try std.testing.expect(operation.immediate_offset % @sizeOf(u32) == 0);
+                const component = operation.immediate_offset / @sizeOf(u32);
+                try std.testing.expect(component < store_offsets.len);
+                store_offsets[component] = true;
+                store_count += 1;
+            },
+            .parallel_copy => return error.UnloweredParallelCopy,
+            else => {},
+        }
+    }
+
+    try std.testing.expectEqual(@as(usize, 4), load_count);
+    try std.testing.expectEqual(@as(usize, 4), store_count);
+    try std.testing.expectEqual([4]bool{ true, true, true, true }, load_offsets);
+    try std.testing.expectEqual([4]bool{ true, true, true, true }, store_offsets);
+}