API reference — reverse of a graph

API reference — reverse of a graph#

The transposition behind eagle.diff: given a graph Description and the planes each node’s vjp reads and writes, eagle::diff::transpose returns the reverse pass (the order of the derivative steps, where each step’s incoming adjoint comes from, which adjoints are summed and in what order, and which forward planes must be copied before they are overwritten). eagle::diff::orderedSum is the fixed-order elementwise sum the joins use. Header-only and CUDA-free; the guide is Gradients through a graph: eagle.diff.

Namespace `eagle::diff`#

namespace diff#

Typedefs

using idx_t = std::int64_t#

using gidx = std::int64_t#: Signed index of nodes, frames and edges (-1 = none).

Functions

template<class Done, class Checkpoint, class Launch> idx_t recordForward(idx_t maxSteps, idx_t window, idx_t maxK, Done &&done, Checkpoint &&checkpoint, Launch &&launch)#

The recording forward: launches of at most maxK steps that never cross a window boundary, a checkpoint before every launch that starts a window, until done() or maxSteps. Returns the steps run.

Parameters:

maxSteps – step cap for the whole run.
window – steps between checkpoints.
maxK – steps one launch may take (1 for a single-step entry).
done – true once every sample has finished (checked before each launch).
checkpoint – checkpoint(k): save the state at the start of window k.
launch – launch(j, k): run steps j .. j+k-1 on every sample.

inline void gatherColumns(void *dst, idx_t dstLd, const void *src, idx_t srcLd, const idx_t *perm, idx_t count, idx_t rows, std::size_t bytes)#: dst[r*dstLd + s] = src[r*srcLd + perm[s]] for every row r < rows and s < count: count columns of a sample-last plane gathered into another sample order (and plane width). Element size bytes (1, 2, 4 or 8).

inline std::vector<idx_t> stopOrder(const idx_t *steps, idx_t n)#: The reverse’s sample order: samples sorted by stop step, longest first, ties in sample order. order[s] is the sample at reverse slot s.

inline std::vector<idx_t> liveCounts(const idx_t *steps, idx_t n)#: live[j] = how many samples take step j (stop > j), for j below the longest stop; steps in any order.

template<class Gather, class Recompute, class StepBack> void reverseWindows(idx_t window, const std::vector<idx_t> &live, Gather &&gather, Recompute &&recompute, StepBack &&stepBack)#

Run the reverse, window by window, last window first.

Parameters:

window – the tape’s window W (rows are step mod W).
live – live[j] from liveCounts: its size is the longest stop.
gather – gather(k, count): load window k’s start state for the first count reverse slots.
recompute – recompute(c, len): run the window’s steps c .. c+len-1 forward (tape rows from 0), each over its own live prefix; one fused launch or len single steps, as recorded.
stepBack – stepBack(j, count): the adjoint through step j.

template<class T> void accumulatePrefix(T *acc, const T *g, idx_t count)#: acc[s] += g[s] for s < count: one step’s share of a gradient that sums over steps (a parameter, a read-only input), in step order.

template<class T> void fillPrefix(T *plane, T value, idx_t count)#: plane[s] = value for s < count: the step index a time-indexed step’s vjp reads, the same for every sample still running.

inline std::vector<Snapshot> snapshots(const std::vector<NodeVjp> &vjps)#

The copies the reverse needs: for each plane a node reads, the first node at or after it that writes the plane (itself, when in place), deduped.

Parameters:: vjps – one entry per forward node, in node order.

inline ReversePlan transpose(const graph::Description &forward, const std::vector<NodeVjp> &vjps, const std::vector<std::string> &seeded)#

Reverse forward, given each node’s vjp planes and the planes the caller seeds (the gradient of the objective with respect to their final value); an unseeded plane contributes nothing of its own.

Throws:: std::invalid_argument – for a graph that is already a reverse (second order), an opaque node, or a node with no entry in vjps; the message names the node.

template<class T> inline void orderedSum(T *dst, const std::vector<const T*> &src, std::ptrdiff_t n)#: dst = ((src[0] + src[1]) + src[2]) + … elementwise, over n elements, in exactly the order of src: the same bits for every thread count. dst may alias src[0].

inline idx_t ceilSqrt(idx_t s)#: The smallest r with r*r >= s (exact for every int64 input).

inline idx_t defaultWindow(idx_t S)#: The default window for runs of up to S steps: clamp(ceil(sqrt(S)), 16, 1024).

inline WindowPlan planWindows(idx_t steps, idx_t window)#: A checked WindowPlan: steps >= 0 and window >= 1.

inline idx_t chunkSamples(idx_t budgetBytes, idx_t perSampleBytes, idx_t n, idx_t leaf)#

Samples per reverse pass under a memory budget.

Parameters:

budgetBytes – the bytes a pass may hold.
perSampleBytes – what one sample costs (its checkpoints and tape).
n – the batch.
leaf – the canonical reduction’s leaf block (a power of two).

Returns:

n when the whole batch fits; otherwise the largest leaf * 2^q that fits, so each pass is a whole subtree of the canonical sum and the combined result is bitwise the unchunked one; 0 when not even one leaf block (or, with n smaller, one sample) fits.

Variables

gidx kSeed = -1#: An adjoint source that is the caller’s seed for the plane.

gidx kZero = -2#: An adjoint source that is zero: nothing flows into the plane.

gidx kLive = -1#: A read that comes from the forward’s own plane, not a snapshot.

struct Join#: #include <Transpose.h>

Reverse node target sums the adjoints of plane produced by sources, elementwise, in exactly this order.

struct NodeVjp#: #include <Transpose.h>

What the reverse needs to know about one forward node: the planes its vjp reads (the node’s inputs) and the planes it writes (whose adjoints flow into its vjp). Plane names are keys that are unique in the graph: two nodes touching the same memory use the same key.

class RecordSlot#

#include <Record.h>

The single live-record slot of one recording forward. Thread-safe.

Public Types

enum class Begin : int#

Why a recording could not begin.

Values:

enumerator Ok#

enumerator Held#

enumerator Live#

Public Functions

inline Begin begin()#: Start a recording: refused while a reverse holds the last record (Held) or while it is still live (Live); otherwise the generation moves on and older records are stale from here.

inline void abandon()#: A recording that failed part-way leaves nothing live.

inline bool hold(std::uint64_t g)#: Hold the record of generation g for one read; false (and nothing held) when it is stale.

inline bool release(std::uint64_t g, bool consumed)#: End one read of generation g; consumed (a finished reverse or a keep) frees the slot for the next recording. Returns false when the generation moved while it was held.

struct ReversePlan#

#include <Transpose.h>

The reverse of a graph. Reverse nodes run in id order.

Public Members

graph::Description reversed#: vjp steps (kind “vjp”, or “broadcast” for a reduce) and summing joins (kind “sum”); RAW edges carry adjoint planes.

std::vector<gidx> forward#: Per reverse node, the forward node it reverses; -1 for a join.

std::vector<std::vector<gidx>> seedOf#: Per reverse node, per NodeVjp::writes entry: the reverse node whose adjoint it consumes, kSeed or kZero. Empty for a join.

std::vector<std::vector<gidx>> readOf#: Per reverse node, per NodeVjp::reads entry: the index of the snapshot holding the value the node read, or kLive.

std::vector<std::pair<std::string, gidx>> gradients#: Per graph input (a plane read before anything writes it), in first-use order: where its gradient comes from (a reverse node, or kZero).

struct Snapshot#: #include <Transpose.h>

Copy forward plane plane just before forward node node runs.

struct WindowPlan#

#include <Window.h>

Windows of window steps over a run of up to steps steps, aligned to multiples of window from step 0.

Public Functions

inline idx_t count() const#: How many windows cover steps (0 for an empty run).

inline idx_t start(idx_t k) const#: The first step of window k.

inline idx_t length(idx_t k) const#: How many steps window k holds (the last may be short).

inline idx_t windowOf(idx_t j) const#: The window holding step j.

inline idx_t row(idx_t j) const#: The tape row step j is stored in.

API reference — reverse of a graph

Contents

API reference — reverse of a graph#

Namespace eagle::diff#

Namespace `eagle::diff`#