[Flint] adding compute resource layout lowering

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

@@ -1,5 +1,9 @@
 const std = @import("std");
 
+pub const resource_layout = @import("resource_layout.zig");
+pub const resource_lowering = @import("resource_lowering.zig");
+pub const ResourceLayout = resource_layout.Layout;
+
 pub const Error = error{UnsupportedWorkgroupSize};
 
 pub fn validateWorkgroupSize(size: [3]u32) Error!void {

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

@@ -0,0 +1,141 @@
+const std = @import("std");
+const ids = @import("../../../ir/id.zig");
+const program_ir = @import("../../../ir/program.zig");
+
+pub const max_storage_buffers: usize = 4;
+
+pub const Error = std.mem.Allocator.Error || error{
+    TooManyStorageBuffers,
+};
+
+pub const Binding = struct {
+    set: u32,
+    binding: u32,
+    binding_table_index: u8,
+};
+
+const Candidate = struct {
+    resource: ids.StorageBufferId,
+    set: u32,
+    binding: u32,
+};
+
+pub const Layout = struct {
+    bindings: []Binding,
+    resource_indices: []?u8,
+
+    pub fn init(allocator: std.mem.Allocator, program: *const program_ir.Program) Error!Layout {
+        var candidates: std.ArrayList(Candidate) = .empty;
+        defer candidates.deinit(allocator);
+
+        for (program.storage_buffers.entries.items, 0..) |entry, index| {
+            const buffer = entry orelse continue;
+            try candidates.append(allocator, .{
+                .resource = ids.StorageBufferId.fromIndex(index),
+                .set = buffer.set,
+                .binding = buffer.binding,
+            });
+        }
+        std.mem.sort(Candidate, candidates.items, {}, lessThan);
+
+        var unique_count: usize = 0;
+        for (candidates.items, 0..) |candidate, index| {
+            if (index == 0 or candidate.set != candidates.items[index - 1].set or candidate.binding != candidates.items[index - 1].binding)
+                unique_count += 1;
+        }
+        if (unique_count > max_storage_buffers)
+            return Error.TooManyStorageBuffers;
+
+        const bindings = try allocator.alloc(Binding, unique_count);
+        errdefer allocator.free(bindings);
+        const resource_indices = try allocator.alloc(?u8, program.storage_buffers.entries.items.len);
+        errdefer allocator.free(resource_indices);
+        @memset(resource_indices, null);
+
+        var binding_index: usize = 0;
+        for (candidates.items, 0..) |candidate, index| {
+            if (index == 0 or candidate.set != candidates.items[index - 1].set or candidate.binding != candidates.items[index - 1].binding) {
+                bindings[binding_index] = .{
+                    .set = candidate.set,
+                    .binding = candidate.binding,
+                    .binding_table_index = @intCast(binding_index),
+                };
+                binding_index += 1;
+            }
+            resource_indices[candidate.resource.index()] = @intCast(binding_index - 1);
+        }
+        std.debug.assert(binding_index == bindings.len);
+
+        return .{
+            .bindings = bindings,
+            .resource_indices = resource_indices,
+        };
+    }
+
+    pub fn deinit(self: *Layout, allocator: std.mem.Allocator) void {
+        allocator.free(self.bindings);
+        allocator.free(self.resource_indices);
+        self.* = undefined;
+    }
+
+    pub fn bindingTableIndex(self: *const Layout, resource: ids.StorageBufferId) ?u8 {
+        if (resource.index() >= self.resource_indices.len)
+            return null;
+        return self.resource_indices[resource.index()];
+    }
+};
+
+fn lessThan(_: void, lhs: Candidate, rhs: Candidate) bool {
+    if (lhs.set != rhs.set)
+        return lhs.set < rhs.set;
+    if (lhs.binding != rhs.binding)
+        return lhs.binding < rhs.binding;
+    return lhs.resource.index() < rhs.resource.index();
+}
+
+test "[gen9] compute resource layout: assign stable binding-table indices" {
+    const device = @import("../../../device.zig");
+
+    const device_info: device.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 third = try program.addStorageBuffer(.{ .set = 2, .binding = 7 });
+    const first = try program.addStorageBuffer(.{ .set = 0, .binding = 3 });
+    const alias = try program.addStorageBuffer(.{ .set = 0, .binding = 3 });
+    const second = try program.addStorageBuffer(.{ .set = 1, .binding = 0 });
+
+    var layout = try Layout.init(std.testing.allocator, &program);
+    defer layout.deinit(std.testing.allocator);
+
+    try std.testing.expectEqual(@as(usize, 3), layout.bindings.len);
+    try std.testing.expectEqual(Binding{ .set = 0, .binding = 3, .binding_table_index = 0 }, layout.bindings[0]);
+    try std.testing.expectEqual(Binding{ .set = 1, .binding = 0, .binding_table_index = 1 }, layout.bindings[1]);
+    try std.testing.expectEqual(Binding{ .set = 2, .binding = 7, .binding_table_index = 2 }, layout.bindings[2]);
+    try std.testing.expectEqual(@as(?u8, 0), layout.bindingTableIndex(first));
+    try std.testing.expectEqual(@as(?u8, 0), layout.bindingTableIndex(alias));
+    try std.testing.expectEqual(@as(?u8, 1), layout.bindingTableIndex(second));
+    try std.testing.expectEqual(@as(?u8, 2), layout.bindingTableIndex(third));
+}
+
+test "[gen9] compute resource layout: enforce advertised storage-buffer limit" {
+    const device = @import("../../../device.zig");
+
+    const device_info: device.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();
+    for (0..max_storage_buffers + 1) |binding|
+        _ = try program.addStorageBuffer(.{ .set = 0, .binding = @intCast(binding) });
+
+    try std.testing.expectError(Error.TooManyStorageBuffers, Layout.init(std.testing.allocator, &program));
+}

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

@@ -0,0 +1,127 @@
+const instruction = @import("../../../ir/instruction.zig");
+const program_ir = @import("../../../ir/program.zig");
+const validator = @import("../../../ir/validator.zig");
+const resource_layout = @import("resource_layout.zig");
+
+pub const Error = error{
+    InvalidProgram,
+    InvalidResourceLayout,
+};
+
+pub fn run(program: *program_ir.Program, layout: *const resource_layout.Layout) Error!void {
+    validator.validate(program) catch return Error.InvalidProgram;
+    if (program.properties.resources_lowered)
+        return;
+    if (layout.resource_indices.len != program.storage_buffers.entries.items.len)
+        return Error.InvalidResourceLayout;
+
+    for (program.blocks.entries.items) |block_entry| {
+        const block = block_entry orelse continue;
+        for (block.instructions.items) |instruction_id| {
+            const inst = program.instructions.get(instruction_id) orelse return Error.InvalidProgram;
+            const reference = bufferReference(inst.operation) orelse continue;
+            const resource = switch (reference) {
+                .logical => |value| value,
+                .binding_table => return Error.InvalidProgram,
+            };
+            const binding_table_index = layout.bindingTableIndex(resource) orelse return Error.InvalidResourceLayout;
+            if (binding_table_index >= layout.bindings.len)
+                return Error.InvalidResourceLayout;
+            const buffer = program.storage_buffers.get(resource) orelse return Error.InvalidProgram;
+            const binding = layout.bindings[binding_table_index];
+            if (binding.binding_table_index != binding_table_index or binding.set != buffer.set or binding.binding != buffer.binding)
+                return Error.InvalidResourceLayout;
+        }
+    }
+
+    for (program.blocks.entries.items) |block_entry| {
+        const block = block_entry orelse continue;
+        for (block.instructions.items) |instruction_id| {
+            const inst = program.instructions.getMut(instruction_id) orelse unreachable;
+            const reference = bufferReferenceMut(&inst.operation) orelse continue;
+            const resource = reference.logical;
+            const binding_table_index = layout.bindingTableIndex(resource).?;
+            reference.* = .{ .binding_table = binding_table_index };
+        }
+    }
+
+    program.properties.resources_lowered = true;
+    validator.validate(program) catch return Error.InvalidProgram;
+}
+
+fn bufferReference(operation: instruction.Operation) ?instruction.BufferReference {
+    return switch (operation) {
+        .load_buffer => |op| op.buffer,
+        .store_buffer => |op| op.buffer,
+        else => null,
+    };
+}
+
+fn bufferReferenceMut(operation: *instruction.Operation) ?*instruction.BufferReference {
+    return switch (operation.*) {
+        .load_buffer => |*op| &op.buffer,
+        .store_buffer => |*op| &op.buffer,
+        else => null,
+    };
+}
+
+test "[gen9] compute resource lowering: resolve logical buffers" {
+    const std = @import("std");
+    const Builder = @import("../../../ir/Builder.zig");
+    const device = @import("../../../device.zig");
+    const operand = @import("../../../ir/operand.zig");
+    const printer = @import("../../../ir/printer.zig");
+
+    const device_info: device.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();
+    var builder = Builder.init(&program);
+
+    const value = try builder.addVirtualRegister(.{
+        .size_bytes = 32,
+        .alignment_bytes = 32,
+        .element_type = .u32,
+        .lane_count = 8,
+        .class = .temporary,
+    });
+    const buffer = try builder.addStorageBuffer(.{ .set = 1, .binding = 3, .name = "storage" });
+    const entry = try builder.addBlock("entry");
+    const store_id = try builder.appendInstruction(entry, .simd8, null, .{
+        .store_buffer = .{
+            .buffer = .{ .logical = buffer },
+            .byte_offset = .{
+                .register = .{ .immediate = .{ .u32 = 0 } },
+                .type = .u32,
+                .region = operand.Region.broadcast(),
+            },
+            .source = .{
+                .register = .{ .virtual = value },
+                .type = .u32,
+                .region = operand.Region.contiguous(.simd8),
+            },
+        },
+    });
+    try builder.setTerminator(entry, .end_thread);
+
+    var layout = try resource_layout.Layout.init(std.testing.allocator, &program);
+    defer layout.deinit(std.testing.allocator);
+    layout.bindings[0].binding = 4;
+    try std.testing.expectError(Error.InvalidResourceLayout, run(&program, &layout));
+    try std.testing.expect(!program.properties.resources_lowered);
+    try std.testing.expect(program.instructions.get(store_id).?.operation.store_buffer.buffer == .logical);
+    layout.bindings[0].binding = 3;
+
+    try run(&program, &layout);
+    try validator.validate(&program);
+
+    try std.testing.expect(program.properties.resources_lowered);
+    try std.testing.expectEqual(@as(u8, 0), program.instructions.get(store_id).?.operation.store_buffer.buffer.binding_table);
+    const text = try printer.allocPrint(std.testing.allocator, &program);
+    defer std.testing.allocator.free(text);
+    try std.testing.expect(std.mem.indexOf(u8, text, "store_buffer bti(0), 0:u32") != null);
+}

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

@@ -8,6 +8,8 @@ pub const compute = @import("compute/compute.zig");
 pub const validator = @import("validator.zig");
 
 pub const Options = common_ir.Options;
+pub const ResourceLoweringError = compute.resource_lowering.Error;
+
 pub const Error = common_ir.Error || compute.Error || error{
     UnsupportedGeneration,
     UnsupportedStage,
@@ -38,6 +40,11 @@ pub fn lower(
     return program;
 }
 
+pub fn lowerComputeResources(program: *program_ir.Program, layout: *const compute.ResourceLayout) ResourceLoweringError!void {
+    try compute.resource_lowering.run(program, layout);
+    validator.validate(program) catch return ResourceLoweringError.InvalidProgram;
+}
+
 test "[gen9] target: reject unsupported target configurations" {
     var module = try shader_ir.parser.parseString(std.testing.allocator,
         \\shader compute @main

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

@@ -12,6 +12,7 @@ pub const Error = shared.Error || compute.Error || error{
     UnsupportedExecutionSize,
     UnsupportedDataType,
     InvalidPhysicalFlag,
+    InvalidBindingTableIndex,
     InvalidPayloadLayout,
 };
 
@@ -54,10 +55,12 @@ fn validateInstruction(inst: instruction.Instruction) Error!void {
     switch (inst.operation) {
         .load_global_invocation_id => |op| try validateDestination(op.destination),
         .load_buffer => |op| {
+            try validateBufferReference(op.buffer);
             try validateDestination(op.destination);
             try validateSource(op.byte_offset);
         },
         .store_buffer => |op| {
+            try validateBufferReference(op.buffer);
             try validateSource(op.byte_offset);
             try validateSource(op.source);
         },
@@ -88,6 +91,14 @@ fn validateInstruction(inst: instruction.Instruction) Error!void {
     }
 }
 
+fn validateBufferReference(reference: instruction.BufferReference) Error!void {
+    switch (reference) {
+        .logical => {},
+        .binding_table => |index| if (index >= compute.resource_layout.max_storage_buffers)
+            return Error.InvalidBindingTableIndex,
+    }
+}
+
 fn validateSource(source: operand.Source) Error!void {
     try validateType(source.type);
     switch (source.register) {