tiny.accy.kernel.library.reduction

1const std = @import("std");2const gpu = @import("gpu");34const entry = @import("entry.zig");5const kernel = @import("../root.zig");67pub const OneDimensional = struct {8    extent: u64,9    threads: u32,10    input_axis: []const u8 = "i",11    output_axis: []const u8 = "out",12};1314fn foldUpper(comptime extent: u64) i64 {15    if (extent > @as(u64, @intCast(std.math.maxInt(i64)))) {16        @compileError("kernel library reduction extent overflows index fold range");17    }18    return @intCast(extent);19}2021fn sumSpecialization(comptime spec: OneDimensional) entry.Specialization {22    return .{23        .dtype = .f32,24        .operation = .{ .reduction = .sum },25        .equation = "i->",26        .inputs = &.{entry.shape1D(spec.input_axis, spec.extent)},27        .outputs = &.{entry.shapeScalar()},28        .reductions = &.{entry.reduction("sum", .sum, entry.shape1D(spec.input_axis, spec.extent))},29        .launch = entry.launch1D(1, spec.threads),30        .schedule = entry.threadBlocks1D(spec.output_axis, 1, spec.threads),31    };32}3334fn dotSpecialization(comptime spec: OneDimensional) entry.Specialization {35    return .{36        .dtype = .f32,37        .operation = .{ .reduction = .dot_product },38        .equation = "i,i->",39        .inputs = &.{40            entry.shape1D(spec.input_axis, spec.extent),41            entry.shape1D(spec.input_axis, spec.extent),42        },43        .outputs = &.{entry.shapeScalar()},44        .reductions = &.{entry.reduction("dot", .dot_product, entry.shape1D(spec.input_axis, spec.extent))},45        .launch = entry.launch1D(1, spec.threads),46        .schedule = entry.threadBlocks1D(spec.output_axis, 1, spec.threads),47    };48}4950fn sum_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {51    const zero = try inner.constantFloat(.f32, 0.0);52    const sum = try inner.foldRange(0, foldUpper(ctx.spec.extent), 1, zero, .{53        .src = ctx.args.param(.src),54    }, sum_step);55    try ctx.args.param(.dst).store(inner, sum, index);56}5758fn sum_step(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {59    const value = try ctx.src.load(fold_inner, offset);60    return fold_inner.add(acc, value.raw());61}6263fn sumProgram(comptime spec: OneDimensional) type {64    const Body = struct {65        fn run(k: anytype, args: anytype) !void {66            _ = try k.forEach1D(spec.output_axis, 1, .{ .spec = spec, .args = args }, sum_each);67        }68    };6970    return kernel.logical.Program(.{71        .name = std.fmt.comptimePrint("accy_kernel_reduction_sum{}x{}_f32", .{ spec.extent, spec.threads }),72        .parameters = .{73            .dst = kernel.dynamicBuffer(.f32),74            .src = kernel.dynamicBuffer(.f32),75        },76        .body = Body.run,77    }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));78}7980fn dot_each(inner: anytype, index: kernel.Index1D, ctx: anytype) !void {81    const zero = try inner.constantFloat(.f32, 0.0);82    const sum = try inner.foldRange(0, foldUpper(ctx.spec.extent), 1, zero, .{83        .lhs = ctx.args.param(.lhs),84        .rhs = ctx.args.param(.rhs),85    }, dot_step);86    try ctx.args.param(.dst).store(inner, sum, index);87}8889fn dot_step(fold_inner: anytype, offset: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {90    const lhs = try ctx.lhs.load(fold_inner, offset);91    const rhs = try ctx.rhs.load(fold_inner, offset);92    const product = try lhs.mul(fold_inner, rhs);93    return fold_inner.add(acc, product.raw());94}9596fn dotProgram(comptime spec: OneDimensional) type {97    const Body = struct {98        fn run(k: anytype, args: anytype) !void {99            _ = try k.forEach1D(spec.output_axis, 1, .{ .spec = spec, .args = args }, dot_each);100        }101    };102103    return kernel.logical.Program(.{104        .name = std.fmt.comptimePrint("accy_kernel_reduction_dot{}x{}_f32", .{ spec.extent, spec.threads }),105        .parameters = .{106            .dst = kernel.dynamicBuffer(.f32),107            .lhs = kernel.dynamicBuffer(.f32),108            .rhs = kernel.dynamicBuffer(.f32),109        },110        .body = Body.run,111    }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));112}113114pub fn sumF32(comptime spec: OneDimensional) type {115    return entry.Entry(sumProgram(spec), .{116        .target = std.fmt.comptimePrint("accy.kernel.reduction.sum{}x{}_f32", .{ spec.extent, spec.threads }),117        .layer = .logical,118        .category = .reduction,119        .specialization = sumSpecialization(spec),120    });121}122123pub fn dotF32(comptime spec: OneDimensional) type {124    return entry.Entry(dotProgram(spec), .{125        .target = std.fmt.comptimePrint("accy.kernel.reduction.dot{}x{}_f32", .{ spec.extent, spec.threads }),126        .layer = .logical,127        .category = .reduction,128        .specialization = dotSpecialization(spec),129    });130}131132pub const Sum8F32 = sumF32(.{ .extent = 8, .threads = 1 });133pub const Dot8F32 = dotF32(.{ .extent = 8, .threads = 1 });134135test "reduction sum entry runs on CPU" {136    var src = [_]f32{ 1.0, -2.0, 3.0, 4.0, 5.5, 0.5, -1.0, 8.0 };137    var dst = [_]f32{0.0};138139    try Sum8F32.runCpu(std.testing.allocator, Sum8F32.Limits.testing, &.{140        kernel.argumentBuffer(f32, dst[0..]),141        kernel.argumentBuffer(f32, src[0..]),142    });143    try std.testing.expectEqual(@as(f32, 19.0), dst[0]);144145    const launch_value = try Sum8F32.launch(std.testing.allocator, Sum8F32.Limits.testing);146    try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);147    try std.testing.expectEqual(@as(u32, 1), launch_value.block[0]);148}149150test "reduction dot entry runs on CPU" {151    var lhs = [_]f32{ 1.0, -2.0, 3.0, 4.0, 0.5, -1.5, 2.0, -3.0 };152    var rhs = [_]f32{ 2.0, 3.0, -1.0, 0.25, 8.0, -2.0, 0.5, 4.0 };153    var dst = [_]f32{0.0};154155    try Dot8F32.runCpu(std.testing.allocator, Dot8F32.Limits.testing, &.{156        kernel.argumentBuffer(f32, dst[0..]),157        kernel.argumentBuffer(f32, lhs[0..]),158        kernel.argumentBuffer(f32, rhs[0..]),159    });160    try std.testing.expectEqual(@as(f32, -10.0), dst[0]);161162    const launch_value = try Dot8F32.launch(std.testing.allocator, Dot8F32.Limits.testing);163    try std.testing.expectEqual(@as(u32, 1), launch_value.grid[0]);164    try std.testing.expectEqual(@as(u32, 1), launch_value.block[0]);165}166167test "reduction constructor creates independent shape-specialized entries" {168    const Sum16F32 = sumF32(.{ .extent = 16, .threads = 1 });169    const Dot16F32 = dotF32(.{ .extent = 16, .threads = 1 });170171    try std.testing.expectEqualStrings("accy.kernel.reduction.sum8x1_f32", Sum8F32.target);172    try std.testing.expectEqualStrings("accy.kernel.reduction.sum16x1_f32", Sum16F32.target);173    try std.testing.expectEqualStrings("accy.kernel.reduction.dot8x1_f32", Dot8F32.target);174    try std.testing.expectEqualStrings("accy.kernel.reduction.dot16x1_f32", Dot16F32.target);175    try std.testing.expectEqual(@as(u64, 8), Sum8F32.specialization.inputs[0].elementCount().?);176    try std.testing.expectEqualStrings("sum", Sum16F32.specialization.reductions[0].name);177    try std.testing.expectEqual(entry.ReductionOperator.sum, Sum16F32.specialization.reductions[0].operator);178    try std.testing.expect(Sum16F32.specialization.operationIs(.{ .reduction = .sum }));179    try std.testing.expectEqual(@as(u64, 16), Sum16F32.specialization.reductions[0].shape.elementCount().?);180    try std.testing.expectEqual(@as(u64, 1), Sum16F32.specialization.outputs[0].elementCount().?);181    try std.testing.expectEqualStrings("dot", Dot16F32.specialization.reductions[0].name);182    try std.testing.expectEqual(entry.ReductionOperator.dot_product, Dot16F32.specialization.reductions[0].operator);183    try std.testing.expect(Dot16F32.specialization.operationIs(.{ .reduction = .dot_product }));184    try std.testing.expectEqual(@as(u64, 16), Dot16F32.specialization.reductions[0].shape.elementCount().?);185    try std.testing.expect(Dot16F32.specialization.outputHasExtents(0, &.{}));186    try std.testing.expectEqualDeep(Sum16F32.specialization.launch.?, Sum16F32.specialization.schedule.?.launch());187    try std.testing.expectEqualDeep(Dot16F32.specialization.launch.?, Dot16F32.specialization.schedule.?.launch());188189    var snapshot = try Sum16F32.scheduleSnapshot(std.testing.allocator, Sum16F32.Limits.testing);190    defer snapshot.deinit(std.testing.allocator);191    try std.testing.expect(Sum16F32.specialization.schedule.?.matchesSnapshot(&snapshot));192193    var dot_snapshot = try Dot16F32.scheduleSnapshot(std.testing.allocator, Dot16F32.Limits.testing);194    defer dot_snapshot.deinit(std.testing.allocator);195    try std.testing.expect(Dot16F32.specialization.schedule.?.matchesSnapshot(&dot_snapshot));196197    var src = @as([16]f32, @splat(1.0));198    var dst = [_]f32{0.0};199    try Sum16F32.runCpu(std.testing.allocator, Sum16F32.Limits.testing, &.{200        kernel.argumentBuffer(f32, dst[0..]),201        kernel.argumentBuffer(f32, src[0..]),202    });203    try std.testing.expectEqual(@as(f32, 16.0), dst[0]);204205    var dot_dst = [_]f32{0.0};206    try Dot16F32.runCpu(std.testing.allocator, Dot16F32.Limits.testing, &.{207        kernel.argumentBuffer(f32, dot_dst[0..]),208        kernel.argumentBuffer(f32, src[0..]),209        kernel.argumentBuffer(f32, src[0..]),210    });211    try std.testing.expectEqual(@as(f32, 16.0), dot_dst[0]);212}213214test "reduction sum entry creates registry-ready artifact" {215    const allocator = std.testing.allocator;216    var state = gpu.recording.BackendState{217        .allocator = allocator,218        .kind = .cuda,219        .format = .cuda_ptx,220    };221222    var call_artifact = try Sum8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = Sum8F32.Limits.testing });223    defer call_artifact.deinit();224225    const artifact = call_artifact.registry().find(Sum8F32.target, Sum8F32.version, .cuda_ptx) orelse {226        return error.TestExpectedKernelCallArtifact;227    };228    try std.testing.expectEqualStrings(Sum8F32.name, artifact.entry_name);229    try std.testing.expectEqual(@as(u32, 2), artifact.argument_count);230    try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);231    switch (artifact.launch) {232        .fixed => |geometry| {233            try std.testing.expectEqual(Sum8F32.specialization.launch.?.grid[0], geometry.grid[0]);234            try std.testing.expectEqual(Sum8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);235        },236        else => return error.TestExpectedFixedLaunch,237    }238}239240test "reduction dot entry creates registry-ready artifact" {241    const allocator = std.testing.allocator;242    var state = gpu.recording.BackendState{243        .allocator = allocator,244        .kind = .cuda,245        .format = .cuda_ptx,246    };247248    var call_artifact = try Dot8F32.createKernelCallArtifact(allocator, state.handle(), .{ .limits = Dot8F32.Limits.testing });249    defer call_artifact.deinit();250251    const artifact = call_artifact.registry().find(Dot8F32.target, Dot8F32.version, .cuda_ptx) orelse {252        return error.TestExpectedKernelCallArtifact;253    };254    try std.testing.expectEqualStrings(Dot8F32.name, artifact.entry_name);255    try std.testing.expectEqual(@as(u32, 3), artifact.argument_count);256    try std.testing.expectEqual(gpu.DTypeSet.init(&.{.f32}).bits, artifact.required_dtypes.bits);257    switch (artifact.launch) {258        .fixed => |geometry| {259            try std.testing.expectEqual(Dot8F32.specialization.launch.?.grid[0], geometry.grid[0]);260            try std.testing.expectEqual(Dot8F32.specialization.launch.?.threadgroup[0], geometry.threadgroup[0]);261        },262        else => return error.TestExpectedFixedLaunch,263    }264}

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