github ↗
docs menu

Module Tolk.Opt

Kernel optimisations.

An optimisation reshapes the loops of a kernel. It names the axis it acts on by its position among the kernel's axes.

#type target =
  1. #| Local
  2. #| Upcast

The type for the roles a split gives the axis it makes, in the order of their axis types (Ops.Axis_type.compare): an axis across the threads of a workgroup, or an axis unrolled into the lanes of a thread, which a reduction's axis unrolls as its loop.

#type t =
  1. #| Tc of {
    1. #axis : int;
    2. #tc_select : int;
    3. #tc_opt : int;
    4. #use_tc : int;
    }
    (*

    Use a tensor core on the reduction at axis: tc_select picks the core (-1 the first that fits), tc_opt how far to relax its requirements, from 0 to 2, and use_tc how to use it, 1 or 2. tc_opt carries the level of Setting.tc_opt, or the search's 2, and use_tc that of Setting.use_tc.

    *)
  2. #| Split of {
    1. #axis : int;
    2. #amount : int;
    3. #target : target;
    4. #top : bool;
    }
    (*

    Split amount out of axis into a new axis for target, taken from the outer end if top. An amount of 0 takes the whole axis.

    *)
  3. #| Padto of {
    1. #axis : int;
    2. #amount : int;
    }
    (*

    Pad axis to a multiple of amount.

    *)
  4. #| Swap of {
    1. #axis : int;
    2. #with_axis : int;
    }
    (*

    Swap two global axes.

    *)

The type for kernel optimisations.

#val axis : t -> int

axis o is the axis o acts on.

#val equal : t -> t -> bool

equal o0 o1 is true iff o0 and o1 are the same optimisation.

#val compare : t -> t -> int

compare o0 o1 orders optimisations by kind (Tc, Split, Padto, Swap), then axis, then arguments, a split's target among them in the order of target.

#val pp : Stdlib.Format.formatter -> t -> unit

pp formats Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.UPCAST)).