import { DefaultMap, Enum, WeakValueMap } from '../helpers/helpers.ts';
import { type sint, UOp, type Variable } from '../ops.ts';
import { ProgramSpec, type Renderer, type TensorCore } from '../renderer/index.ts';
import { ShapeTracker } from '../shape/shapetracker.ts';
export declare class OptOps<Name extends string = string, Value extends number = number> extends Enum {
    private static VALUES;
    static values: () => OptOps<string, number>[];
    _id: bigint;
    constructor(name: Name);
    static readonly TC: OptOps<"TC", number>;
    static readonly UPCAST: OptOps<"UPCAST", number>;
    static readonly UNROLL: OptOps<"UNROLL", number>;
    static readonly LOCAL: OptOps<"LOCAL", number>;
    static readonly GROUP: OptOps<"GROUP", number>;
    static readonly GROUPTOP: OptOps<"GROUPTOP", number>;
    static readonly NOLOCALS: OptOps<"NOLOCALS", number>;
    static readonly PADTO: OptOps<"PADTO", number>;
    static readonly SWAP: OptOps<"SWAP", number>;
}
export declare class KernelOptError extends Error {
}
export declare const check: (cond: boolean, msg?: string) => void;
export declare class Opt {
    op: OptOps;
    axis?: number | undefined;
    amt?: number | undefined;
    _id: bigint;
    static cache: WeakValueMap<bigint, Opt>;
    constructor(op: OptOps, axis?: number | undefined, amt?: number | undefined);
    toString: () => string;
    real_axis: (k: Kernel) => number;
}
export declare class TensorCoreOptions {
    axes: number[];
    axes_exist: boolean[];
    axis_pads: [number, number][];
    constructor(axes: number[], // the location of the original N and M axes if still in the shape
    axes_exist: boolean[], // true if the original N and M axes are still in the shape
    axis_pads: [number, number][]);
    fix_axes: (removed_axis: number) => void;
}
export declare class Kernel {
    ast: UOp;
    opts: Renderer;
    reduceops: UOp[];
    vars: Variable[];
    bufs: UOp[];
    full_buf_index: number;
    sts: ShapeTracker[];
    applied_opts: Opt[];
    group_for_reduces: number;
    uops?: UOp[];
    upcasted: number;
    local_dims: number;
    tensor_core?: TensorCore;
    tensor_core_opts?: TensorCoreOptions;
    use_tensor_cores: number;
    dont_use_locals: boolean;
    constructor(ast: UOp, opts?: Renderer);
    copy: () => Kernel;
    get membufs(): UOp[];
    float4_axis: (i: number) => number[];
    upcasted_axis: (i: number) => [number, undefined | sint, boolean][];
    get first_reduce(): number;
    get first_upcast(): number;
    get reduceop(): UOp | undefined;
    get output_shape(): sint[];
    get full_shape(): sint[];
    get full_unupcasted_shape(): sint[];
    get shape_len(): number;
    get upcast_in_mid_reduce_axes(): number[];
    get global_dims(): number;
    colors: () => string[];
    colored_shape: (pad?: number, dense?: boolean) => string;
    reshape_and_permute: (new_shape_fxn?: (a: sint[]) => sint[], axis?: number[]) => void;
    upcast: () => void;
    shift_to: (axis: number, amount: sint, top?: boolean, insert_before?: number) => void;
    simplify_ones: () => boolean;
    simplify_merge_adjacent: () => void;
    _create_tc_opts: (reduceop: UOp, tc: TensorCore, axis: number, opt_level: number) => TensorCoreOptions | undefined;
    _apply_tc_opt: (use_tensor_cores: number, axis: number, opt_level: number) => boolean;
    /**
     * Attempts to apply a tensor core optimization to the kernel.  If one exists && applies properly, return true, otherwise return false.
     * Tensor cores are optimized instructions that matrix multiply-accumulate across a wave of threads: D(M, N) = A(M, K) * B(K, N) + C(M, N).
     *
     * Keyword arguments:
     * use_tensor_cores -- controls how tensor cores are applied (default 1)
     * 0: will disable any tensor core matching
     * 1: enable tensor cores
     * 2: apply tensor core shape but don't use UOp.WMMA
     * extra_opts -- additional Opt's to apply after the tensor core instead of the hand-coded additional Opt's (default undefined)
     * tc_opt -- controls which kinds of kernels may be eligible for tensor cores application (default 2 during BEAM, 0 otherwise)
     * 0: applies to only kernels with a single reduce axis && direct UOps.LOAD into Ops.MUL
     * 1: allows kernels with multiple reduce axes && also multiplication of UOps.CAST'd buffers
     * 2: allows kernels with M, N, K axes that are not multiples of the tensor core dimensions by applying padding those axes as needed
     */
    apply_tensor_cores: (use_tensor_cores?: number, extra_opts?: Opt[], axis?: number, tc_opt?: number) => boolean;
    apply_opt: (opt: Opt, append_opt?: boolean) => void;
    required_optimizations: () => Kernel;
    hand_coded_optimizations: () => Kernel;
    static kernel_cnt: DefaultMap<string, number>;
    private _name;
    get name(): string;
    fixup_ast: (op: UOp) => UOp;
    get_optimized_ast: () => UOp;
    linearize: () => Kernel;
    to_program: (name_override?: string) => ProgramSpec;
}
export declare const _assert_valid_uop: (uop: UOp, st: ShapeTracker, sts: Map<UOp, ShapeTracker>) => undefined;
export declare const verify_ast: (ast: UOp) => undefined;