lib/accy/src/kernel/library/layout.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3
4 const entry = @import("entry.zig");
5 const kernel = @import("../root.zig");
6
7 fn 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 }
13
14 fn 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 }
25
26 fn 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 }
31
32 fn 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 }
37
38 fn 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 }
44
45 fn 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 };
54
55 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 }
70
71 pub 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 }
82
83 pub const Transpose8x16F32 = transposeF32(.{
84 .rows = 8,
85 .cols = 16,
86 .threads = .{ .x = 8, .y = 4 },
87 });
88
89 test "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 });
95
96 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));
101
102 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..]);
107
108 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());
118
119 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 }
123
124 test "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 };
131
132 var call_artifact = try Transpose8x16F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = Transpose8x16F32.Limits.testing });
133 defer call_artifact.deinit();
134
135 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 }