Module Ir.Indexing

type symbol =
  1. | Symbol of Base.int
val compare_symbol : symbol -> symbol -> Base.int
val equal_symbol : symbol -> symbol -> Base.bool
val symbol_of_sexp : Sexplib0.Sexp.t -> symbol
val sexp_of_symbol : symbol -> Sexplib0.Sexp.t
val hash_fold_symbol : Ppx_hash_lib.Std.Hash.state -> symbol -> Ppx_hash_lib.Std.Hash.state
val hash_symbol : symbol -> Ppx_hash_lib.Std.Hash.hash_value
val symbol : Base.int -> symbol
val is_symbol : symbol -> bool
val symbol_val : symbol -> Base.int Stdlib.Option.t
module Variants_of_symbol : sig ... end
val unique_id : int Base.ref
val get_symbol : unit -> symbol
module CompareSymbol : sig ... end
module Symbol : sig ... end
val symbol_ident : symbol -> Base.String.t
type 'a environment = 'a Base.Map.M(Symbol).t
val environment_of_sexp : 'a. (Sexplib0.Sexp.t -> 'a) -> Sexplib0.Sexp.t -> 'a environment
val sexp_of_environment : 'a. ('a -> Sexplib0.Sexp.t) -> 'a environment -> Sexplib0.Sexp.t
val empty_env : 'a environment
type static_symbol = {
  1. static_symbol : symbol;
  2. mutable static_range : Base.int Base.option;
  3. mutable used_as_extent : Base.bool;
    (*

    The symbol is used as a symbolic dimension (gh-490): its bound value is a size in [0, static_range] (inclusive -- buffers are sized static_range), rather than an index in [0, static_range). Set by Row.get_sym_dim.

    *)
  4. mutable used_as_slice : Base.bool;
    (*

    The symbol is used as a batch-slice index (@|): its bound value indexes an axis of size static_range, so the strict [0, static_range) validation must apply. Set by the Batch_slice shape logic; mutually exclusive with used_as_extent (a single binding cannot be validated both inclusively and strictly).

    *)
}
val compare_static_symbol : static_symbol -> static_symbol -> Base.int
val equal_static_symbol : static_symbol -> static_symbol -> Base.bool
val static_symbol_of_sexp : Sexplib0.Sexp.t -> static_symbol
val sexp_of_static_symbol : static_symbol -> Sexplib0.Sexp.t
val hash_fold_static_symbol : Ppx_hash_lib.Std.Hash.state -> static_symbol -> Ppx_hash_lib.Std.Hash.state
val hash_static_symbol : static_symbol -> Ppx_hash_lib.Std.Hash.hash_value
type 'a bindings =
  1. | Empty
  2. | Bind of static_symbol * (Base.int -> 'a) bindings
val sexp_of_bindings : 'a. ('a -> Sexplib0.Sexp.t) -> 'a bindings -> Sexplib0.Sexp.t
val bound_symbols : 'a bindings -> static_symbol Base.List.t
type ('r, 'idcs, 'p1, 'p2) variadic =
  1. | Result of 'r
  2. | Param_idx of Base.int Base.ref * (Base.int -> 'r, Base.int -> 'idcs, 'p1, 'p2) variadic
  3. | Param_1 of 'p1 Base.option Base.ref * ('p1 -> 'r, 'idcs, 'p1, 'p2) variadic
  4. | Param_2 of 'p2 Base.option Base.ref * ('p2 -> 'r, 'idcs, 'p1, 'p2) variadic
  5. | Param_2f : ('p2f -> 'p2) * 'p2f Base.option Base.ref * ('p2 -> 'r, 'idcs, 'p1, 'p2) variadic -> ('r, 'idcs, 'p1, 'p2) variadic

Helps lowering the bindings.

type unit_bindings = (Base.unit -> Base.unit) bindings
val sexp_of_unit_bindings : unit_bindings -> Sexplib0.Sexp.t
type lowered_bindings = (static_symbol, Base.int Base.ref) Base.List.Assoc.t
val sexp_of_lowered_bindings : lowered_bindings -> Sexplib0.Sexp.t
val apply : 'r 'idcs 'p1 'p2. ('r, 'idcs, 'p1, 'p2) variadic -> 'r

apply run_variadic () applies the parameters in reverse order to how they appear in the run_variadic list.

val lowered_bindings : 'a bindings -> ('b, 'a, 'p1, 'p2) variadic -> (static_symbol * Base.int Base.ref) Base.List.t
val find_exn : lowered_bindings -> static_symbol -> Base.int Base.ref
val get_static_symbol : ?static_range:Base.int -> (Base.int -> 'a) bindings -> static_symbol * 'a bindings
val validate_bound_value : ?width64:Base.bool -> static_symbol -> Base.int -> unit

Validates a launch-parameter value against its declared range (minimum 0 is implicit) and, when width64 is false, against the int32 index width (docs/proposals/signed-index-precision.md: bind-time validation mirroring tinygrad's bind assert; one host compare per launch). The per-routine node-extent contract is NOT sound for launch parameters -- value-embedded parameters (step counters) take runtime values unrelated to any touched node's extent -- so narrowing an unbounded parameter requires declaring (and thereby bind-validating) a range.

val validate_lowered_bindings : ?width64:Base.bool -> lowered_bindings -> unit
val dims_to_string : ?with_axis_numbers:bool -> Base.Int.t Base.Array.t -> Base.String.t

Dimensions to string, "x"-separated, e.g. 1x2x3 for batch dims 1, input dims 3, output dims 2. Outputs "-" for empty dimensions.

type axis_index =
  1. | Fixed_idx of Base.int
    (*

    A fixed position along an axis

    *)
  2. | Iterator of symbol
    (*

    A simple iterator symbol

    *)
  3. | Affine of {
    1. symbols : (Base.int * symbol) Base.list;
    2. offset : Base.int;
    }
    (*

    An affine combination of symbols with coefficients and an offset. Represents:

    Σ(coeff_i * symbol_i) + offset

    For convolutions: symbols = [(stride, i1); (dilation, i2)] and offset = ~-padding. Note: for readability, we use Fixed_idx and Iterator as separate variants and require Affine to not be ambiguous: symbols should be longer than 1 or have a coefficient different from 1 and 0.

    *)
  4. | Sub_axis
    (*

    This axis belongs to an adjacent multi-axis index.

    *)
  5. | Concat of symbol Base.list
    (*

    This axis is formed by concatenating multiple axes, each represented by an iterator symbol. Concat indices are eliminated during lowering.

    *)
val compare_axis_index : axis_index -> axis_index -> Base.int
val equal_axis_index : axis_index -> axis_index -> Base.bool
val axis_index_of_sexp : Sexplib0.Sexp.t -> axis_index
val sexp_of_axis_index : axis_index -> Sexplib0.Sexp.t
type str_osym_map = (Base.string, symbol Base.option, Base.String.comparator_witness) Base.Map.t
val sexp_of_str_osym_map : str_osym_map -> Base.Sexp.t
type projections_debug = {
  1. spec : Base.string;
  2. derived_for : Base.Sexp.t;
  3. trace : (Base.string * Base.int) Base.list;
}
val projections_debug_of_sexp : Sexplib0.Sexp.t -> projections_debug
val sexp_of_projections_debug : projections_debug -> Sexplib0.Sexp.t
val unique_debug_id : unit -> int
type projections = {
  1. product_space : Base.int Base.list Base.array;
    (*

    The product space dimensions that an operation should parallelize (map-reduce) over. Singletons except for concatenation dimensions.

    *)
  2. lhs_dims : Base.int Base.array;
    (*

    The dimensions of the LHS array.

    *)
  3. rhs_dims : Base.int Base.array Base.array;
    (*

    The dimensions of the RHS arrays, needed for deriving projections from other projections.

    *)
  4. product_iterators : symbol Base.list Base.array;
    (*

    The product space iterators (concatentation of the relevant batch, output, input axes) for iterating over the product_space axes, where same axes are at same array indices. These may be shared; lowering creates fresh symbols for loop indices. Singletons except for concatenation dimensions. A concatenation dimension is a connected component of symbols participating in concatenated axes.

    *)
  5. project_lhs : axis_index Base.array;
    (*

    A projection that takes an product_space-bound index and produces an index into the result of an operation.

    *)
  6. project_rhs : axis_index Base.array Base.array;
    (*

    project_rhs.(i) Produces an index into the i+1th argument of an operation.

    *)
  7. extent_syms : (symbol * static_symbol) Base.list;
    (*

    gh-490 symbolic extents: maps a product iterator to the static symbol whose bound value is the axis's runtime extent. The corresponding product_space entry is the maximum extent (the symbol's declared range); lowering guards the loop body with iterator < value when the symbol is among the routine's bindings.

    *)
  8. debug_info : projections_debug;
}

All the information relevant for code generation.

val compare_projections : projections -> projections -> Base.int
val equal_projections : projections -> projections -> Base.bool
val projections_of_sexp : Sexplib0.Sexp.t -> projections
val sexp_of_projections : projections -> Sexplib0.Sexp.t
val iterated : int -> bool
val opt_symbol : int -> symbol option
val opt_iterator : symbol option -> axis_index
val coalesce_affine_terms : (Base__Int.t * Symbol.t) Base.List.t -> (Base__Int.t * Symbol.t) Base.List.t

Coalesce an affine symbols list into canonical (coeff, symbol) terms: repeated symbols are summed and zero-coefficient terms dropped. Order is unspecified.

val affine_injective : symbol_range:(Symbol.t -> Base__Int.t) -> axis_index Base.array -> Base.bool
val identity_projections : ?debug_info:projections_debug -> ?derived_for:string -> lhs_dims:Base.int Base.Array.t -> unit -> projections

Projections for a pointwise unary operator. Provide only one of debug_info or derived_for.

val prod_dims : projections -> Base.int Base.array

The extents of the product-space components of p: one entry per component, in product_space order. A concatenation component's extent is the sum of its segment extents. Non-iterated (dimension-1) axes do not appear.

val prod_project_for : projections -> dims:Base.int Base.array -> axis_index Base.array

The identity projection of a tensor laid out over the full product space of p (e.g. the (output x kernel) pair space of a windowed reduction): the tensor's axes are the product components, possibly interleaved with non-iterated (dimension-1) axes, which the product space omits. dims is the tensor's dimensions; each dimension-1 axis maps to Fixed_idx 0 and every other axis consumes the first not-yet-consumed product component of equal extent. Matching by extent rather than strictly by position tolerates a proxy-shape layout whose axis order deviates from the product order (e.g. a reduced-over row variable placed in the result's output row while the result keeps input axes): any pairing is correct as long as it is consistent, and this is a pure function of p and dims, so every access to the tensor uses the same pairing. A leftover on either side means the product-space proxy shape (gh-512) got out of sync with the derived projections, and raises rather than corrupting memory.

val reflect_projection : dims:Base.int Base.array -> projection:axis_index Base.array -> axis_index
type variable_ref = {
  1. ref_label : Base.string;
  2. mutable solved_dim : Base.int Base.option;
  3. mutable solved_sym : static_symbol Base.option;
    (*

    The dimension is a symbolic extent (gh-490) rather than a concrete integer; where a concrete value is required (e.g. Embed_dim), it materializes at the symbol's declared range (the maximum extent). At most one of solved_dim, solved_sym is set.

    *)
}
val sexp_of_variable_ref : variable_ref -> Sexplib0.Sexp.t
val equal_variable_ref : variable_ref -> variable_ref -> Base.bool
module Doc_helpers : sig ... end