/// <reference types="dist" />
import { Device } from "./device";
import { ATypedArray, Dtype } from "./dtype";
import { ExprCode, CompiledExpr, EvalEnv } from "./expr";
export type ShaderType = "u8" | "array<u8>" | "i32" | "array<i32>" | "u32" | "array<u32>" | "f32" | "array<f32>" | "vec3<i32>" | "vec3<u32>" | "vec3<f32>";
export interface KernelSpec {
    name: string;
    config: KernelConfigSpec[];
    workgroupSize: [ExprCode, ExprCode, ExprCode];
    parameters: KernelParamSpec[];
    workgroupCount: [ExprCode, ExprCode, ExprCode];
    inputs: KernelInputSpec[];
    outputs: KernelOutputSpec[];
    shader: string;
}
export interface KernelInputSpec {
    name: string;
    shaderType: ShaderType;
}
export interface KernelOutputSpec {
    name: string;
    shaderType: ShaderType;
    size: ExprCode;
}
export interface KernelParamSpec {
    name: string;
    shaderType: "u32" | "f32";
}
export interface KernelConfigSpec {
    name: string;
}
export type KernelConfigInput = {
    [name: string]: string | number;
};
export type KernelParamsInput = {
    [name: string]: number;
};
export type KernelConfig = (string | number)[];
export type KernelKey = string;
export declare abstract class Kernel {
    private _key;
    private _spec;
    private _config;
    private _device;
    private _workgroupCountXFunc;
    private _workgroupCountYFunc;
    private _workgroupCountZFunc;
    protected _outputSizeFuncs: CompiledExpr[];
    get key(): KernelKey;
    get spec(): KernelSpec;
    get config(): KernelConfig;
    get device(): Device;
    constructor(spec: KernelSpec, config: KernelConfig, device: Device);
    abstract run(inputs: (GPUBuffer | ATypedArray)[], parameters: KernelParamsInput, outputs?: (GPUBuffer | ATypedArray)[]): (GPUBuffer | ATypedArray)[];
    protected getRunEnv(parameters: KernelParamsInput): EvalEnv;
    getWorkgroupCounts(env: EvalEnv): [number, number, number];
}
export declare function getKernelConfig(spec: KernelSpec, config: KernelConfigInput): KernelConfig;
export declare function getKernelKey(spec: KernelSpec, config: KernelConfig): KernelKey;
export declare function getShaderTypeElementByteSize(shaderType: ShaderType): number;
export declare function shaderTypeToDtype(shaderType: ShaderType): Dtype;
export declare function getKernelShaderCode(spec: KernelSpec, config: KernelConfig): string;
export declare function getKernelJavaScriptCode(spec: KernelSpec, config: KernelConfig): string;