tiny.accy.kernel

1pub const Options = struct {2    target: []const u8,3    version: u32 = 1,4    format: ?gpu.ArtifactFormat = null,5    kernel_plan: plan_mod.Options = .{},6    element_count_argument: artifact_product.ElementCountArgument = .none,7    shape_family_fingerprint: ?u64 = null,8    shape_profile: ?artifact_product.KernelCallShapeProfile = null,9    launch: ?artifact_product.KernelCallLaunch = null,10    runtime_scalar_argument_count: u32 = 0,11    static_arguments: []const choir_abi.ScalarArgument = &.{},12};

1pub const OwnedArtifact = struct {2    allocator: std.mem.Allocator,3    artifacts: [1]artifact_product.KernelCallArtifact,45    pub fn entry(self: *const OwnedArtifact) artifact_product.KernelCallArtifact {6        return self.artifacts[0];7    }89    pub fn registry(self: *const OwnedArtifact) artifact_product.KernelCallRegistry {10        return .{ .entries = self.artifacts[0..] };11    }1213    pub fn fingerprint(self: *const OwnedArtifact) u64 {14        const registry_value = self.registry();15        return artifact_product.kernelCallRegistryFingerprint(&registry_value);16    }1718    pub fn productStamp(self: *const OwnedArtifact) choir.product.incremental.ProductStamp {19        return choir.product.incremental.productStamp(product_name, self.fingerprint());20    }2122    pub fn deinit(self: *OwnedArtifact) void {23        const allocator = self.allocator;24        const artifact = &self.artifacts[0];25        allocator.free(@constCast(artifact.target));26        allocator.free(@constCast(artifact.entry_name));27        if (artifact.shape_profile) |profile| artifact_product.deinitKernelCallShapeProfile(allocator, profile);28        if (artifact.static_arguments.len != 0) allocator.free(@constCast(artifact.static_arguments));29        deinitCompilePayload(allocator, artifact.payload);30        self.* = undefined;31    }32};

1pub const RawLimits = struct {2    context: choir.ir.Context.Limits,3    parameters: usize,4    kernel_name_bytes: usize,5    temporary_values: usize,6    temporary_types: usize,78    pub const standard: RawLimits = .{9        .context = choir.ir.Context.Limits.standard,10        .parameters = 64,11        .kernel_name_bytes = 512,12        .temporary_values = 64,13        .temporary_types = 64,14    };1516    pub const testing: RawLimits = .{17        .context = choir.ir.Context.Limits.testing,18        .parameters = 256,19        .kernel_name_bytes = 2048,20        .temporary_values = 256,21        .temporary_types = 256,22    };2324    pub const borrowed_standard: RawLimits = .{25        .context = borrowedContextLimits(false),26        .parameters = 64,27        .kernel_name_bytes = 512,28        .temporary_values = 64,29        .temporary_types = 64,30    };3132    pub const borrowed_testing: RawLimits = .{33        .context = borrowedContextLimits(true),34        .parameters = 256,35        .kernel_name_bytes = 2048,36        .temporary_values = 256,37        .temporary_types = 256,38    };39};

1pub fn createArtifact(2    allocator: std.mem.Allocator,3    handle: gpu.BackendHandle,4    program: *program_mod.Program,5    options: Options,6) !OwnedArtifact {7    if (options.target.len == 0 or options.version == 0) return error.InvalidArtifact;89    var authored_plan = try program.createCheckedPlan(allocator, options.kernel_plan);10    defer authored_plan.deinit();1112    const caps = try handle.queryCapabilities();13    const backend_kind = handle.backendKind() orelse caps.identity.backend;14    const format = options.format orelse artifact_product.defaultArtifactFormat(backend_kind) orelse {15        return error.UnsupportedOperation;16    };17    if (!caps.supportsArtifactFormat(format)) return error.UnsupportedArtifactFormat;18    const geometry = launchGeometry(&authored_plan);19    try caps.validateLaunchGeometry(geometry);2021    const compilation = try target.compileKernelForArtifactFormat(22        allocator,23        format,24        authored_plan.entry_name,25        program.kernelModule(),26        compileOptions(format, geometry),27    );28    errdefer deinitCompilePayload(allocator, compilation.payload);2930    const owned_target = try allocator.dupe(u8, options.target);31    errdefer allocator.free(owned_target);32    const owned_entry_name = try allocator.dupe(u8, authored_plan.entry_name);33    errdefer allocator.free(owned_entry_name);34    const static_arguments = try compileStaticArguments(allocator, format, geometry, options.static_arguments);35    errdefer if (static_arguments.len != 0) allocator.free(static_arguments);3637    const shape_family_fingerprint = options.shape_family_fingerprint orelse if (options.shape_profile) |profile| profile.fingerprint else null;38    if (options.shape_profile) |profile| {39        if (shape_family_fingerprint.? != profile.fingerprint) return error.InvalidArtifact;40        try profile.validate(options.runtime_scalar_argument_count);41    }42    const shape_profile = if (options.shape_profile) |profile|43        try artifact_product.duplicateKernelCallShapeProfile(allocator, profile)44    else45        null;46    errdefer if (shape_profile) |profile| artifact_product.deinitKernelCallShapeProfile(allocator, profile);4748    return .{49        .allocator = allocator,50        .artifacts = .{.{51            .target = owned_target,52            .version = options.version,53            .format = format,54            .entry_name = owned_entry_name,55            .argument_count = try compileArgumentCount(format, authored_plan.argument_count),56            .shape_family_fingerprint = shape_family_fingerprint,57            .shape_profile = shape_profile,58            .required_dtypes = requiredDTypesForParams(authored_plan.params),59            .required_features = gpu_codegen.featureRequirementsForModule(program.kernelModule()),60            .required_subgroup = gpu_codegen.subgroupRequirementsForModule(program.kernelModule()),61            .push_constants = compilation.push_constants,62            .payload = compilation.payload,63            .launch = options.launch orelse .{ .fixed = geometry },64            .element_count_argument = options.element_count_argument,65            .runtime_scalar_argument_count = options.runtime_scalar_argument_count,66            .static_arguments = static_arguments,67        }},68    };69}

1pub fn createBackendArtifactFromEntry(2    allocator: std.mem.Allocator,3    handle: gpu.BackendHandle,4    entry: artifact_product.KernelCallArtifact,5    diagnostic_id: ?[]const u8,6) !gpu.KernelArtifact {7    const caps = try handle.queryCapabilities();8    const scalar_argument_count = try entry.scalarArgumentCount();9    var artifact = try gpu.KernelArtifact.init(allocator, .{10        .backend = handle.backendKind() orelse caps.identity.backend,11        .format = entry.format,12        .entry_name = entry.entry_name,13        .argument_count = entry.argument_count,14        .scalar_argument_count = scalar_argument_count,15        .diagnostic_id = diagnostic_id,16        .interface = .{17            .features = entry.required_features,18            .subgroup = entry.required_subgroup,19            .push_constants = entry.push_constants,20        },21    });22    errdefer artifact.deinit();23    switch (entry.payload) {24        .text => |text| artifact.setBorrowedText(text),25        .bytes => |bytes| artifact.setBorrowedBytes(bytes),26        .words_u32 => |words| artifact.setBorrowedWords(words),27        .none => return error.InvalidArtifact,28    }29    return artifact;30}

1pub fn launchGeometryForEntry(2    entry: artifact_product.KernelCallArtifact,3    runtime_arguments: []const choir_abi.ScalarArgument,4) gpu.BackendError!choir_abi.LaunchGeometry {5    return switch (entry.launch) {6        .fixed => |geometry| geometry,7        .derived => |derived| try derived.geometry(runtime_arguments),8    };9}

1pub fn compileFragment(2    allocator: std.mem.Allocator,3    handle: gpu.BackendHandle,4    graph: *program.Program,5    options: FragmentCompilerOptions,6) !*exec.CompiledFragment {7    const artifact_start = nowNs();8    const artifact_job = try artifact.createJob(9        allocator,10        handle,11        graph,12        artifactPlanOptions(options),13    );14    defer artifact_job.deinit();15    try options.instrumentation.record(.plan_create_backend_artifacts, artifact_start);16    return try exec.compileFragmentFromArtifactJob(allocator, artifact_job);17}

1pub fn createArtifact(2    allocator: std.mem.Allocator,3    handle: gpu.BackendHandle,4    graph: *program.Program,5    options: ArtifactOptions,6) !gpu.KernelArtifact {7    const compiled = try compileFragment(allocator, handle, graph, .{8        .artifact_format = options.artifact_format,9        .authored_kernel_diagnostic_id = options.authored_kernel_diagnostic_id,10        .instrumentation = options.instrumentation,11    });12    defer compiled.deinit();13    if (compiled.kernelCount() != 1) return error.InvalidArtifact;14    return try exec.candidate.copyKernelArtifactToAllocator(15        allocator,16        compiled.artifactPlan().kernels.items[0].artifact,17    );18}

1pub fn Family(comptime definition: anytype) type {2    return struct {3        pub const name = definition.name;4        pub const parameters = definition.parameters;5        pub const Layout: type = parameter.Standard(parameters);6        pub const Instance: type = definition.Instance;7        pub const Limits: type = kernel.Limits;89        pub fn arg(comptime name_value: anytype) usize {10            return Layout.index(name_value);11        }1213        pub fn schema() []const kernel.Param {14            return parameter.schema(parameters);15        }1617        pub fn build(allocator: std.mem.Allocator, limits: Limits, instance: Instance) !kernel.Graph {18            return execute.buildInstanceSource(definition, instance, allocator, limits);19        }2021        pub fn buildNamed(allocator: std.mem.Allocator, limits: Limits, graph_name: []const u8, instance: Instance) !kernel.Graph {22            return execute.buildNamedInstanceSource(definition, graph_name, instance, allocator, limits);23        }2425        pub fn interpret(allocator: std.mem.Allocator, limits: Limits, instance: Instance, initial: anytype) !execute.result(@TypeOf(initial)) {26            return execute.interpretInstanceSource(definition, instance, allocator, limits, initial);27        }2829        pub fn launch(allocator: std.mem.Allocator, limits: Limits, instance: Instance) !kernel.Launch {30            var graph = try build(allocator, limits, instance);31            defer graph.deinit();32            return graph.launch();33        }3435        pub fn scheduleSnapshot(allocator: std.mem.Allocator, limits: Limits, instance: Instance) !kernel.ScheduleSnapshot {36            var graph = try build(allocator, limits, instance);37            defer graph.deinit();38            return graph.scheduleSnapshot(allocator);39        }4041        pub fn createPlan(allocator: std.mem.Allocator, limits: Limits, instance: Instance, options: kernel.PlanOptions) !kernel.Plan {42            var graph = try build(allocator, limits, instance);43            defer graph.deinit();44            return graph.createPlan(allocator, options);45        }4647        pub fn createCheckedPlan(allocator: std.mem.Allocator, limits: Limits, instance: Instance, options: kernel.PlanOptions) !kernel.Plan {48            var graph = try build(allocator, limits, instance);49            defer graph.deinit();50            return graph.createCheckedPlan(allocator, options);51        }5253        pub fn compileFragment(54            allocator: std.mem.Allocator,55            limits: Limits,56            instance: Instance,57            handle: kernel.BackendHandle,58            options: executable.FragmentCompilerOptions,59        ) !*executable.CompiledFragment {60            var graph = try build(allocator, limits, instance);61            defer graph.deinit();62            return kernel_compile.compileFragment(allocator, handle, &graph, options);63        }6465        pub fn createKernelArtifact(66            allocator: std.mem.Allocator,67            limits: Limits,68            instance: Instance,69            handle: kernel.BackendHandle,70            options: executable.KernelCompilerOptions,71        ) !kernel.KernelArtifact {72            var graph = try build(allocator, limits, instance);73            defer graph.deinit();74            return kernel_compile.createArtifact(allocator, handle, &graph, options);75        }7677        pub fn createKernelCallArtifact(78            allocator: std.mem.Allocator,79            limits: Limits,80            instance: Instance,81            handle: kernel.BackendHandle,82            options: kernel.KernelCallArtifactOptions,83        ) !kernel.OwnedKernelCallArtifact {84            var graph = try build(allocator, limits, instance);85            defer graph.deinit();86            return kernel.createKernelCallArtifact(allocator, handle, &graph, options);87        }8889        pub fn runCpu(allocator: std.mem.Allocator, limits: Limits, instance: Instance, args: []const kernel.Argument) !void {90            var graph = try build(allocator, limits, instance);91            defer graph.deinit();92            try graph.runCpu(allocator, args);93        }9495        pub fn runCpuWithDiagnostic(96            allocator: std.mem.Allocator,97            limits: Limits,98            instance: Instance,99            args: []const kernel.Argument,100            diagnostic: *kernel.ExecutionDiagnostic,101        ) !void {102            var graph = try build(allocator, limits, instance);103            defer graph.deinit();104            try graph.runCpuWithDiagnostic(allocator, args, diagnostic);105        }106107        pub fn verify(allocator: std.mem.Allocator, limits: Limits, instance: Instance) !void {108            var graph = try build(allocator, limits, instance);109            defer graph.deinit();110            try graph.verify();111        }112    };113}

1pub fn Program(comptime definition: anytype) type {2    return Source(definition);3}

1const gpu = @import("gpu");2const model = @import("model/root.zig");3pub const builder = model.core.builder;4const call = @import("call.zig");5pub const domain = model.core.domain;6const dsl = @import("dsl/root.zig");7const compile = @import("compile/root.zig");8const interpret_mod = @import("interpret/root.zig");9const library_mod = @import("library/root.zig");10const limits = model.core;11const logical_mod = @import("logical/root.zig");12pub const oracle = @import("oracle/root.zig");13pub const plan = model.plan;14pub const program = @import("program/root.zig");15pub const schedule = model.core.schedule;16pub const typed = model.core.typed;17pub const vector = model.core.vector;18pub const view = model.core.view;19pub const artifact = compile.artifact;2021pub const Param = builder.Param;22pub const Buffer = builder.Buffer;23pub const Type = builder.Type;24pub const Kernel = builder.Kernel;25pub const RawBuilder = builder.Builder;26pub const RawStorage = builder.Storage;27pub const Builder = program.Builder;28pub const MmaShape = builder.MmaShape;29pub const Value = builder.Value;30pub const Block = builder.Block;31pub const If = builder.If;32pub const For = builder.For;33pub const ForScope = builder.ForScope;34pub const While = builder.While;35pub const WhileScope = builder.WhileScope;36pub const InsertionScope = builder.InsertionScope;37pub const VerificationDiagnostic = builder.VerificationDiagnostic;38pub const Plan = plan.Plan;39pub const PlanOptions = plan.Options;40pub const PlanError = plan.Error;41pub const PlanVersion = plan.plan_version;42pub const BackendHandle = gpu.BackendHandle;43pub const KernelArtifact = gpu.KernelArtifact;44pub const KernelCallArtifactOptions = call.Options;45pub const OwnedKernelCallArtifact = call.OwnedArtifact;46pub const Argument = oracle.Argument;47pub const Scalar = oracle.Scalar;48pub const ExecutionDiagnostic = oracle.ExecutionDiagnostic;49pub const Graph = program.Program;50pub const Program = dsl.Program;51pub const Family = dsl.Family;52pub const program_product_name = program.product_name;53pub const interpret = interpret_mod;54pub const library = library_mod;55pub const logical = logical_mod;56pub const DomainAxis = domain.DomainAxis;57pub const Domain2D = domain.Domain2D;58pub const Domain3D = domain.Domain3D;59pub const Index1D = domain.Index1D;60pub const VectorIndex1D = domain.VectorIndex1D;61pub const Index2D = domain.Index2D;62pub const Index3D = domain.Index3D;63pub const BufferView = view.BufferView;64pub const KeyBufferView = view.KeyBufferView;65pub const TypedValue = typed.Value;66pub const KeyValue = typed.KeyValue;67pub const TypedVec2 = typed.Vec2;68pub const TypedVec3 = typed.Vec3;69pub const Vec2 = vector.Vec2;70pub const Vec3 = vector.Vec3;71pub const Guard = program.Guard;72pub const Schedule = schedule.Schedule;73pub const ScheduleError = schedule.ScheduleError;74pub const Launch = schedule.Launch;75pub const Limits = limits.Limits;76pub const Capacity = program.Builder.Capacity;77pub const RawLimits = limits.RawLimits;78pub const RawCapacity = builder.Builder.Capacity;79pub const RawUsage = builder.Storage.Usage;80pub const RawExhaustion = builder.Storage.Exhaustion;81pub const ScheduleLimits = schedule.ScheduleLimits;82pub const ScheduleCapacity = schedule.ScheduleCapacity;83pub const ScheduleUsage = schedule.ScheduleUsage;84pub const ScheduleExhaustion = schedule.ScheduleExhaustion;85pub const ScheduleSnapshotLimits = schedule.SnapshotLimits;86pub const ScheduleSnapshotCapacity = schedule.SnapshotCapacity;87pub const ScheduleSnapshotUsage = schedule.SnapshotUsage;88pub const ScheduleSnapshotExhaustion = schedule.SnapshotExhaustion;89pub const Axis = schedule.Axis;90pub const AxisId = schedule.AxisId;91pub const AxisStep = schedule.AxisStep;92pub const BindTarget = schedule.BindTarget;93pub const ScheduleSnapshot = schedule.Snapshot;94pub const ScheduleSnapshotVersion = schedule.snapshot_version;95pub const Split = schedule.Split;96pub const Step = schedule.Step;97pub const scalar = builder.scalar;98pub const buffer = builder.buffer;99pub const dynamicBuffer = builder.dynamicBuffer;100pub const domainAxis = domain.domainAxis;101pub const createPlan = plan.create;102pub const createArtifactJob = artifact.createJob;103pub const createArtifactPlan = artifact.createPlan;104pub const createKernelCallArtifact = call.createArtifact;105pub const createBackendArtifactFromKernelCallEntry = call.createBackendArtifactFromEntry;106pub const kernelCallEntryLaunchGeometry = call.launchGeometryForEntry;107pub const compileFragment = compile.compileFragment;108pub const createKernelArtifact = compile.createArtifact;109pub const argumentBuffer = oracle.argumentBuffer;110pub const argumentBool = oracle.argumentBool;111pub const argumentI32 = oracle.argumentI32;112pub const argumentU32 = oracle.argumentU32;113pub const argumentI64 = oracle.argumentI64;114pub const argumentF32 = oracle.argumentF32;115pub const argumentF64 = oracle.argumentF64;116pub const Compare = builder.Compare;117pub const Dimension = builder.Dimension;118pub const AddressSpace = builder.AddressSpace;119pub const MemoryOrder = builder.MemoryOrder;120pub const Scope = builder.Scope;121pub const WarpOpKind = builder.WarpOpKind;

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