tiny.accy.kernel.library.histogram

1pub const Histogram = struct {2    bins: u64,3    count: u64,4    lo: f32 = 0,5    width: f32 = 1,6    dtype: DType = .f32,7    binning: HistogramBinningPolicy = .lower_inclusive_upper_exclusive,8    variant: HistogramVariant = .direct,9    threads: u32 = 256,10    bin_axis: []const u8 = "b",11    element_axis: []const u8 = "n",12};

1pub const HistogramVariant = enum {2    direct,3    shared_bins,4};

1pub const HistogramResolvedSchedule = struct {2    variant: HistogramVariant,3    threads: u32,4};

1pub const policy_parameter_name = "histogram_binning_policy";

1pub fn histogramBinForValue(instance: Histogram, value: f32) ?usize {2    return switch (instance.binning) {3        .lower_inclusive_upper_exclusive => {4            const relative = (value - instance.lo) / instance.width;5            const bins_float: f32 = @floatFromInt(instance.bins);6            if (!(relative >= 0) or !(relative < bins_float)) return null;7            return @intFromFloat(relative);8        },9    };10}

1pub fn histogramBinningPolicyCode(policy: HistogramBinningPolicy) u64 {2    return binning.code(policy);3}

1pub fn histogramBinningPolicyFromCode(value: u64) ?HistogramBinningPolicy {2    return binning.fromCode(value);3}

1pub fn histogramDTypeSupported(dtype: DType) bool {2    return dtype == .f32;3}

1pub fn histogramInstanceValid(instance: Histogram) bool {2    if (!histogramDTypeSupported(instance.dtype)) return false;3    if (instance.bins == 0 or instance.count == 0) return false;4    if (!(instance.width > 0)) return false;5    if (instance.variant == .shared_bins and instance.bins > histogram_shared_bins_cap) return false;6    return instance.threads != 0;7}

1pub const histogram_family_version: u32 = 2;

1pub const histogram_shared_bins_cap: u64 = 4096;

1pub fn histogramFamilyEntryName(allocator: std.mem.Allocator, instance: Histogram) ![]u8 {2    return switch (instance.variant) {3        .direct => std.fmt.allocPrint(4            allocator,5            "accy_kernel_histogram_histogram_family_{d}_{s}",6            .{ instance.threads, instance.dtype.name() },7        ),8        .shared_bins => std.fmt.allocPrint(9            allocator,10            "accy_kernel_histogram_histogram_family_shared{d}_{d}_{s}",11            .{ instance.bins, instance.threads, instance.dtype.name() },12        ),13    };14}

1pub fn histogramFamilyTarget(allocator: std.mem.Allocator, instance: Histogram) ![]u8 {2    return switch (instance.variant) {3        .direct => std.fmt.allocPrint(4            allocator,5            "accy.kernel.histogram.histogram_family_{d}_{s}",6            .{ instance.threads, instance.dtype.name() },7        ),8        .shared_bins => std.fmt.allocPrint(9            allocator,10            "accy.kernel.histogram.histogram_family_shared{d}_{d}_{s}",11            .{ instance.bins, instance.threads, instance.dtype.name() },12        ),13    };14}

1pub fn histogramRuntimeArguments(instance: Histogram) ![4]choir_abi.ScalarArgument {2    return .{3        .{ .u32 = try runtimeExtentArgument(instance.bins) },4        .{ .u32 = try runtimeExtentArgument(instance.count) },5        .{ .f32 = instance.lo },6        .{ .f32 = instance.width },7    };8}

1pub fn createHistogramFamilyArtifact(2    allocator: std.mem.Allocator,3    handle: kernel.BackendHandle,4    instance: Histogram,5    options: entry.ArtifactOptions,6) !kernel.OwnedKernelCallArtifact {7    if (!family.histogramInstanceValid(instance)) return error.InvalidKernelLibraryEntry;8    const target = try family.histogramFamilyTarget(allocator, instance);9    defer allocator.free(target);10    const entry_name = try family.histogramFamilyEntryName(allocator, instance);11    defer allocator.free(entry_name);12    const family_fingerprint = options.shape_family_fingerprint orelse try family.histogramFamilyFingerprint(allocator, instance);13    const shape_profile_dimensions = family.histogramShapeProfileDimensions(instance);14    const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{15        .name = "histogram",16        .fingerprint = family_fingerprint,17        .dimensions = shape_profile_dimensions[0..],18    };1920    var graph = switch (instance.dtype) {21        .f32 => try body.HistogramRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),22        else => return error.UnsupportedDType,23    };24    defer graph.deinit();25    return kernel.createKernelCallArtifact(allocator, handle, &graph, .{26        .target = target,27        .version = family.histogram_family_version,28        .format = options.format,29        .kernel_plan = options.kernel_plan,30        .element_count_argument = options.element_count_argument,31        .shape_family_fingerprint = family_fingerprint,32        .shape_profile = shape_profile,33        .launch = options.launch orelse try histogramDerivedLaunch(instance),34        .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 4 else options.runtime_scalar_argument_count,35        .static_arguments = options.static_arguments,36    });37}

1pub const HistogramRuntimeFamilyF32 = histogramRuntimeFamily(.f32);

1pub fn histogramFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: Histogram) !u64 {2    var family = try histogramShapeFamily(backing_allocator, instance);3    defer family.deinit();4    return shape.fingerprint(family);5}

1pub fn histogramShapeFamily(backing_allocator: std.mem.Allocator, instance: Histogram) !shape.Family {2    var builder = try shape.Builder.init(backing_allocator, "histogram");3    errdefer builder.deinit();45    const bin = try builder.symbol(instance.bin_axis);6    const element = try builder.symbol(instance.element_axis);7    const bin_expr = try builder.symbolExpression(bin);8    const element_expr = try builder.symbolExpression(element);910    _ = try builder.tensor("dst", &.{bin_expr});11    _ = try builder.tensor("data", &.{element_expr});12    try builder.assumeBounds(bin_expr, histogramRuntimeExtentBounds());13    try builder.assumeBounds(element_expr, histogramRuntimeExtentBounds());1415    return builder.finish();16}

1pub fn histogramShapeProfileDimensions(instance: Histogram) [2]artifact.KernelCallShapeProfileDimension {2    const bounds = histogramRuntimeExtentBounds();3    return .{4        .{ .name = instance.bin_axis, .runtime_scalar_argument_index = 0, .bounds = bounds },5        .{ .name = instance.element_axis, .runtime_scalar_argument_index = 1, .bounds = bounds },6    };7}

1pub fn histogramFamilySpecialization(backing_allocator: std.mem.Allocator, instance: Histogram) !entry.OwnedSpecialization {2    var owned = entry.OwnedSpecialization.init(backing_allocator);3    errdefer owned.deinit();4    const lifetime_allocator = owned.allocator();56    const inputs = try lifetime_allocator.alloc(entry.Shape, 1);7    inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.count);89    const outputs = try lifetime_allocator.alloc(entry.Shape, 1);10    outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.bin_axis, instance.bins);1112    const static_parameters = try lifetime_allocator.alloc(entry.StaticParameter, 1);13    static_parameters[0] = try entry.runtimeStaticParameter(14        lifetime_allocator,15        model.histogram_binning_policy_parameter,16        model.histogramBinningPolicyCode(instance.binning),17    );1819    owned.value = .{20        .dtype = instance.dtype,21        .operation = .{ .indexing = .histogram },22        .inputs = inputs,23        .outputs = outputs,24        .static_parameters = static_parameters,25        .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, instance.element_axis, instance.count, instance.threads),26        .structure = @tagName(instance.variant),27    };28    owned.value.launch = owned.value.schedule.?.launch();29    var family = try shape.histogramShapeFamily(backing_allocator, instance);30    errdefer family.deinit();31    try owned.takeShapeFamily(&family);32    return owned;33}

1pub fn histogramInstanceFromSpecialization(specialization: entry.Specialization) ?Histogram {2    if (!specialization.scheduleMatchesLaunch()) return null;3    if (!specialization.operationIs(.{ .indexing = .histogram })) return null;4    const dtype = specialization.dtype orelse return null;5    if (!model.histogramDTypeSupported(dtype)) return null;6    if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;7    if (specialization.reductions.len != 0) return null;8    if (specialization.static_parameters.len != 1) return null;9    const binning = model.histogramBinningPolicyFromCode(10        specialization.staticParameterValue(model.histogram_binning_policy_parameter) orelse return null,11    ) orelse return null;12    const data = specialization.inputs[0];13    const bins_shape = specialization.outputs[0];14    if (data.axes.len != 1 or bins_shape.axes.len != 1) return null;15    const launch = specialization.launch orelse return null;16    if (launch.threadgroup[0] == 0) return null;17    const structure = specialization.structure orelse return null;18    const variant = std.meta.stringToEnum(HistogramVariant, structure) orelse return null;19    const instance = Histogram{20        .bins = bins_shape.axes[0].extent,21        .count = data.axes[0].extent,22        .dtype = dtype,23        .binning = binning,24        .variant = variant,25        .threads = launch.threadgroup[0],26        .bin_axis = bins_shape.axes[0].name,27        .element_axis = data.axes[0].name,28    };29    if (!model.histogramInstanceValid(instance)) return null;30    return instance;31}

1pub fn histogramFamilyTuningKey(2    backing_allocator: std.mem.Allocator,3    device_fingerprint: u64,4    instance: Histogram,5) !tuning.FamilyTuningKey {6    const family_fingerprint = try shape.histogramFamilyFingerprint(backing_allocator, instance);7    const extents = histogramTuningExtents(instance);8    return tuning.FamilyTuningKey.init(9        device_fingerprint,10        family_fingerprint,11        entry.operationFingerprint(histogramTuningOperation(instance)),12        instance.dtype,13        model.histogram_family_version,14        extents[0..],15    ) orelse unreachable;16}

1pub fn histogramThreadCandidatesForCount(count: u64) geometry.Thread1DCandidates {2    return geometry.threadCandidatesForExtent(count, histogram_thread_caps);3}

1pub fn histogramThreadsForCount(count: u64) u32 {2    return geometry.threadsForExtent(count, histogram_thread_caps);3}

1pub fn histogramTuningExtents(instance: Histogram) [2]u64 {2    return .{ instance.bins, instance.count };3}

1pub fn histogramTuningOperation(instance: Histogram) entry.Operation {2    _ = instance;3    return .{ .indexing = .histogram };4}

1pub fn resolveHistogramSchedule(2    backing_allocator: std.mem.Allocator,3    reader: tuning.FamilyTuningReader,4    instance: Histogram,5) !?HistogramResolvedSchedule {6    const key = try histogramFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);7    const record = reader.table.find(key) orelse return null;8    const thread_candidates = histogramThreadCandidatesForCount(instance.count);9    const variants = [_]HistogramVariant{ .direct, .shared_bins };10    for (variants) |variant| {11        for (thread_candidates.slice()) |threads| {12            var candidate = instance;13            candidate.variant = variant;14            candidate.threads = threads;15            if (variant == .shared_bins and candidate.bins > model.histogram_shared_bins_cap) continue;16            const target = try naming.histogramFamilyTarget(backing_allocator, candidate);17            defer backing_allocator.free(target);18            if (std.mem.eql(u8, target, record.target)) {19                return .{ .variant = variant, .threads = threads };20            }21        }22    }23    return null;24}

1const model = @import("model.zig");2const naming = @import("naming.zig");3const tuning = @import("tuning.zig");4const runtime = @import("runtime/root.zig");5const shape = @import("shape.zig");6const specialization = @import("specialization.zig");78pub const HistogramBinningPolicy = model.HistogramBinningPolicy;9pub const HistogramVariant = model.HistogramVariant;10pub const Histogram = model.Histogram;11pub const histogram_family_version = model.histogram_family_version;12pub const histogram_shared_bins_cap = model.histogram_shared_bins_cap;13pub const histogram_binning_policy_parameter = model.histogram_binning_policy_parameter;14pub const histogramDTypeSupported = model.histogramDTypeSupported;15pub const histogramInstanceValid = model.histogramInstanceValid;16pub const histogramBinningPolicyCode = model.histogramBinningPolicyCode;17pub const histogramBinningPolicyFromCode = model.histogramBinningPolicyFromCode;18pub const histogramBinForValue = model.histogramBinForValue;19pub const histogramFamilyTarget = naming.histogramFamilyTarget;20pub const histogramFamilyEntryName = naming.histogramFamilyEntryName;21pub const HistogramResolvedSchedule = tuning.HistogramResolvedSchedule;22pub const histogramTuningExtents = tuning.histogramTuningExtents;23pub const histogramTuningOperation = tuning.histogramTuningOperation;24pub const histogramFamilyTuningKey = tuning.histogramFamilyTuningKey;25pub const resolveHistogramSchedule = tuning.resolveHistogramSchedule;26pub const histogramThreadsForCount = tuning.histogramThreadsForCount;27pub const histogramThreadCandidatesForCount = tuning.histogramThreadCandidatesForCount;28pub const histogramRuntimeArguments = runtime.histogramRuntimeArguments;29pub const HistogramRuntimeFamilyF32 = runtime.HistogramRuntimeFamilyF32;30pub const createHistogramFamilyArtifact = runtime.createHistogramFamilyArtifact;31pub const histogramShapeFamily = shape.histogramShapeFamily;32pub const histogramFamilyFingerprint = shape.histogramFamilyFingerprint;33pub const histogramShapeProfileDimensions = shape.histogramShapeProfileDimensions;34pub const histogramFamilySpecialization = specialization.histogramFamilySpecialization;35pub const histogramInstanceFromSpecialization = specialization.histogramInstanceFromSpecialization;

1pub const histogram = @import("histogram/family/root.zig");