tiny.accy.kernel.library.layout

1const std = @import("std");2const gpu = @import("gpu");34const entry = @import("entry.zig");5const kernel = @import("../root.zig");67fn indexUpper(comptime extent: u64) i64 {8    if (extent > @as(u64, @intCast(std.math.maxInt(i64)))) {9        @compileError("kernel library layout extent overflows index range");10    }11    return @intCast(extent);12}1314fn transposeSpecialization(comptime spec: entry.Matrix2D) entry.Specialization {15    return .{16        .dtype = .f32,17        .operation = .{ .layout = .transpose },18        .equation = "ij->ji",19        .inputs = &.{entry.shape2D(spec.row_axis, spec.rows, spec.col_axis, spec.cols)},20        .outputs = &.{entry.shape2D(spec.col_axis, spec.cols, spec.row_axis, spec.rows)},21        .launch = entry.launch2D(spec.cols, spec.rows, spec.threads.x, spec.threads.y),22        .schedule = entry.threadBlocks2D(spec.col_axis, spec.cols, spec.row_axis, spec.rows, spec.threads.x, spec.threads.y),23    };24}2526fn transposeInputIndex(inner: anytype, comptime spec: entry.Matrix2D, row: kernel.Value, col: kernel.Value) !kernel.Value {27    const col_stride = try inner.constantIndex(indexUpper(spec.cols));28    const row_offset = try inner.mul(row, col_stride);29    return inner.add(row_offset, col);30}3132fn transposeOutputIndex(inner: anytype, comptime spec: entry.Matrix2D, row: kernel.Value, col: kernel.Value) !kernel.Value {33    const row_stride = try inner.constantIndex(indexUpper(spec.rows));34    const col_offset = try inner.mul(col, row_stride);35    return inner.add(col_offset, row);36}3738fn transpose_each(inner: anytype, index: kernel.Index2D, ctx: anytype) !void {39    const input_index = try transposeInputIndex(inner, ctx.spec, index.y.index, index.x.index);40    const output_index = try transposeOutputIndex(inner, ctx.spec, index.y.index, index.x.index);41    const value = try ctx.args.param(.src).load(inner, input_index);42    try ctx.args.param(.dst).store(inner, value, output_index);43}4445fn transposeProgram(comptime spec: entry.Matrix2D) type {46    const Body = struct {47        fn run(k: anytype, args: anytype) !void {48            _ = try k.forEach2D(.{49                .x = kernel.logical.axis(spec.col_axis, spec.cols),50                .y = kernel.logical.axis(spec.row_axis, spec.rows),51            }, .{ .spec = spec, .args = args }, transpose_each);52        }53    };5455    return kernel.logical.Program(.{56        .name = std.fmt.comptimePrint(57            "accy_kernel_layout_transpose{}x{}_{}x{}_f32",58            .{ spec.rows, spec.cols, spec.threads.x, spec.threads.y },59        ),60        .parameters = .{61            .dst = kernel.dynamicBuffer(.f32),62            .src = kernel.dynamicBuffer(.f32),63        },64        .body = Body.run,65    }).withSchedule(kernel.logical.schedule.threadBlocks(.{66        .x = spec.threads.x,67        .y = spec.threads.y,68    }));69}7071pub fn transposeF32(comptime spec: entry.Matrix2D) type {72    return entry.Entry(transposeProgram(spec), .{73        .target = std.fmt.comptimePrint(74            "accy.kernel.layout.transpose{}x{}_{}x{}_f32",75            .{ spec.rows, spec.cols, spec.threads.x, spec.threads.y },76        ),77        .layer = .logical,78        .category = .layout,79        .specialization = transposeSpecialization(spec),80    });81}8283pub const Transpose8x16F32 = transposeF32(.{84    .rows = 8,85    .cols = 16,86    .threads = .{ .x = 8, .y = 4 },87});8889test "layout transpose entry runs on CPU and records schedule" {90    const Transpose2x3F32 = transposeF32(.{91        .rows = 2,92        .cols = 3,93        .threads = .{ .x = 2, .y = 2 },94    });9596    var src = [_]f32{97        1.0, 2.0, 3.0,98        4.0, 5.0, 6.0,99    };100    var dst = @as([6]f32, @splat(0.0));101102    try Transpose2x3F32.runCpu(std.testing.allocator, Transpose2x3F32.Limits.testing, &.{103        kernel.argumentBuffer(f32, dst[0..]),104        kernel.argumentBuffer(f32, src[0..]),105    });106    try std.testing.expectEqualSlices(f32, &.{ 1.0, 4.0, 2.0, 5.0, 3.0, 6.0 }, dst[0..]);107108    const launch_value = try Transpose2x3F32.launch(std.testing.allocator, Transpose2x3F32.Limits.testing);109    try std.testing.expectEqual(@as(u32, 2), launch_value.grid[0]);110    try std.testing.expectEqual(@as(u32, 1), launch_value.grid[1]);111    try std.testing.expectEqual(@as(u32, 2), launch_value.block[0]);112    try std.testing.expectEqual(@as(u32, 2), launch_value.block[1]);113    try std.testing.expect(Transpose2x3F32.specialization.operationIs(.{ .layout = .transpose }));114    try std.testing.expectEqualStrings("ij->ji", Transpose2x3F32.specialization.equation.?);115    try std.testing.expectEqual(@as(u64, 6), Transpose2x3F32.specialization.inputs[0].elementCount().?);116    try std.testing.expectEqual(@as(u64, 6), Transpose2x3F32.specialization.outputs[0].elementCount().?);117    try std.testing.expectEqualDeep(Transpose2x3F32.specialization.launch.?, Transpose2x3F32.specialization.schedule.?.launch());118119    var snapshot = try Transpose2x3F32.scheduleSnapshot(std.testing.allocator, Transpose2x3F32.Limits.testing);120    defer snapshot.deinit(std.testing.allocator);121    try std.testing.expect(Transpose2x3F32.specialization.schedule.?.matchesSnapshot(&snapshot));122}123124test "layout transpose entry creates registry-ready artifact" {125    const allocator = std.testing.allocator;126    var state = gpu.recording.BackendState{127        .allocator = allocator,128        .kind = .cuda,129        .format = .cuda_ptx,130    };131132    var call_artifact = try Transpose8x16F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = Transpose8x16F32.Limits.testing });133    defer call_artifact.deinit();134135    const artifact = call_artifact.registry().find(Transpose8x16F32.target, Transpose8x16F32.version, .cuda_ptx) orelse {136        return error.TestExpectedKernelCallArtifact;137    };138    try std.testing.expectEqualStrings(Transpose8x16F32.name, artifact.entry_name);139    try std.testing.expectEqual(@as(u32, 2), artifact.argument_count);140    try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);141    switch (artifact.launch) {142        .fixed => |geometry| {143            try std.testing.expectEqual(Transpose8x16F32.specialization.launch.?.grid[0], geometry.grid[0]);144            try std.testing.expectEqual(Transpose8x16F32.specialization.launch.?.grid[1], geometry.grid[1]);145            try std.testing.expectEqual(Transpose8x16F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);146            try std.testing.expectEqual(Transpose8x16F32.specialization.launch.?.threadgroup[1], geometry.threadgroup[1]);147        },148        else => return error.TestExpectedFixedLaunch,149    }150}

1pub const layout = @import("layout.zig");