Module Tolk.Search

actions () is the optimisations the search tries under the current settings, in order:

• splits of the axes 0 to 9 into Opt.target.Upcast lanes by 0, 2, 3, 4, 5 and 7;
• splits of the axes 0 to 7 into Opt.target.Local threads by 0, 2, 3, 4, 8, 13, 16 and 29, then from the top by 13, 16, 28, 29, 32, 49, 64 and 256;
• under Setting.beam_padto, pads of the axes 0 to 6 to a multiple of 32;
• a split of the axis 0 into 32 local threads;
• tensor cores on the axis 0 with the level tc_opt 0, then on the axes 0 to 8 with the level 2, which admits every kernel a tensor core fits, each with the first core that fits and the level use_tc of Setting.use_tc;
• swaps of each two of the axes 0 to 4.

get_kernel_actions ~max_up k is the kernels that one action makes of k: (a, k') for an action a of actions () and k' a copy of k it was applied to, in the order of actions. k is left as it is. An action is left out when:

• it does not apply (Postrange.Scheduler.apply_opt);
• it is no tensor core and its axis is not one of k's, or it splits a whole axis by its size while actions () splits it by 0 alike;
• k' has more upcast lanes than max_up (default Setting.beam_upcast_max), not counting a tensor core's product for each of its threads, or more warp and local threads than Setting.beam_local_max. Under Setting.beam_log_surpass_max, these are reported on standard output.

Raises Invalid_argument as Postrange.Scheduler.apply_opt does, and if whether a kernel has too many lanes or threads depends on the value of a variable.

beam_search ~link ~time ~allow_test_size amt k is k, or a copy of it, with the optimisations that a beam search of width amt finds fastest. k's kernel is the sink that Postrange.apply_opts optimises.

link call is the call call (Op.t.Call) of a compiled program of k's kernel made ready to run on a device of k's renderer's target. Its arguments hold storage parameters in their slots, each naming the target's device (Helpers.Target.t), which link places on a device of the target, and constant values for positional scalar parameters; an unused slot before the last argument holds a byte. Each scalar of k's kernel (Shape.variables) is sampled at the middle of its bounds, (vmin + vmax) / 2 rounded down. Launch sizes are evaluated with these samples, and the kernel keeps its parameters.

time ~vars p is the time in seconds of one run of the program p links, from cold caches where the device can, with each variable bound to its value in vars, the samples of the named variables. The search links each program it times once. A program whose link or time raises Failure when it is first timed is dropped; the search raises any other exception. The search limits neither a compilation nor a run: a candidate whose compilation or run hangs hangs the search, one more reason to run it offline.

A kernel is compiled for k's renderer named "test". It is first linearized (Codegen.linearize) and dropped if it has Setting.beam_uops_max instructions or more, unless that is 0 or less. It is then completed (Codegen.compile), its binary read from or kept in its compiler's disk cache (Renderer.Compiler.compile_cached); its program is not kept past the search. A kernel whose compilation raises Failure is dropped, and so is one whose compilation raises another exception, unless Setting.beam_strict_mode holds and the exception is raised.

A candidate is first timed up to three times, stopping once its least time exceeds an early stop. Its samples are the times of this first timing; the times a round takes again, below, are not samples.

The beam holds up to amt kernels, and starts as k, untimed, of an infinite time. Each round:

1. The candidates are the kernels get_kernel_actions makes of the beam's kernels, in order. They are compiled on the domains Worker.iter spreads them over while the calling domain times those compiled, in the order their compilations end.
2. In that order, a candidate is dropped if its binary was timed before in this search, or if its program's estimated operations (Ops.estimates, 0 if unknown) are more than 1000 times the fewest of this round's programs so far. Each other one is timed with an early stop of three times the least sample of the beam's first kernel.
3. The fastest candidate is the one of least sample, the first timed of ties. The first round compares k's infinite time with the fastest's least sample. In a later round, the beam's first kernel and the fastest are timed again in turn, three times each, starting with the beam's first, and the round compares the least of each one's three times; a time that raises Failure counts as infinite.
4. The search ends if no candidate was timed, if the fastest took less than Setting.beam_min_progress microseconds, or if it beats the beam's first kernel by less than that, or not at all. It then keeps the fastest candidate alone if it is faster than the beam's first kernel. Otherwise the beam becomes the amt candidates of least sample, ties in the order they were timed, and a new round starts.

The result is the beam's first kernel.

A round decides whether the search goes on, and what it keeps, by comparing only times taken in that round. A device's clock follows its load, power and temperature, as a GPU's boost and idle states and a host's turbo do, so times taken rounds apart may have run at different clocks; the beam's first kernel and the fastest, timed in turn, run at nearly the same. A clock that moves within a round still shifts that round's samples against each other, and the early stop, three times the beam's first kernel's least sample from the round that kept it, can cut a candidate's samples short. A time that holds its device's clocks removes both, and measures the device at the held clocks.

With allow_test_size (default Setting.beam_estimate), a program launching more than 65536 workgroups is timed launching fewer: of its global sizes, the last one greater than 16 is halved until their product is at most 65536, and its time is scaled up by the ratio of the products.

While the setting Setting.cachelevel is positive, each search keeps the optimisations it finds in the Helpers.Diskcache table "beam_search", keyed by what the search is a function of but the times it measures: k's kernel (Ops.key), amt, allow_test_size, the name and target of k's renderer and the table of its compiler's binaries (Renderer.Compiler.cachekey), the settings that shape compilation (Setting.shaping), among them those that pick the candidates (Setting.use_tc and Setting.beam_padto) and how it measures and stops (Setting.beam_uops_max, Setting.beam_upcast_max, Setting.beam_local_max, Setting.beam_min_progress and Setting.beam_estimate), and the sources of this library. Unless Setting.ignore_beam_cache holds, a search whose key is kept compiles, links and times nothing, and applies the optimisations kept beyond as many as k has to a copy of k. A kept result is what an earlier search measured fastest: another search may measure otherwise. A key whose entry is damaged on disk, or cannot be read, is searched again, as one never kept.

When the setting Setting.debug is 2 or more, the progress of the search is printed on standard output, and Setting.beam_debug prints more.

Raises Invalid_argument if amt is not positive or as get_kernel_actions does, and Failure if the kept optimisations are malformed.