DOCSEARCHbyAlgolia
Submit Docs

Mooncake.jl

repository·main·Indexed 19 days ago

https://github.com/chalk-lab/mooncake.jl

A high-performance Automatic Differentiation (AD) package written in Julia. It provides optimized reverse-mode, forward-mode, and Hessian computations. A key feature is its caching mechanism, allowing users to prepare caches for functions and inputs to accelerate repeated evaluations of gradients and Hessians.

Tokens
33.7K
Snippets
87
Records
132
Agent score
63%

What's inside Mooncake.jl

  1. What is Mooncake.jl?

    main

    Mooncake.jl is a Julia package designed for performing reverse-mode algorithmic differentiation (AD).

    Unlike other AD systems such as Zygote, Diffractor, or ChainRules, Mooncake.jl explicitly supports in-place operations and mutation. This makes it suitable for complex workflows where data structures are modified during computation, but it also means users should understand its specific mathematical handling of mutated states to use it effectively.

  2. What is a CFGBlock and why is it used?

    main

    A CFGBlock is Mooncake's internal, builder-local basic-block representation used during the reverse-mode assembly process. It is not a public IR for general use, but a specialized format used within reverse_mode.jl to manage the complexity of AD.

    Structure:

    • A stable internal Mooncake ID (independent of compiler SSA or block numbers).
    • inst_ids: A vector of statement IDs.
    • insts: A vector of NewInstructions.

    Why it is necessary: Directly manipulating compiler IRCode during reverse-mode is brittle because AD requires inserting new blocks, splitting blocks, and rethreading control flow. CFGBlock provides a 'loose' construction format where block identity is stable and instructions can be inserted without immediate, complex SSA renumbering. The final structure is only converted back to IRCode once the assembly is complete.

  3. Understand Mooncake's Tangent Type and FData/RData Split

    main

    Mooncake associates every primal type (the original data) with a tangent type (the type storing derivatives). To optimize performance, Mooncake splits tangents into two components:

    1. fdata (Forward Data): Components typically identified by address (like arrays or mutable fields) that are carried along and updated in-place during the forward pass.
    2. rdata (Reverse Data): Value-identified components (like plain numbers) that are only needed during the reverse pass.

    For any tangent t, the relationship must hold: Mooncake.tangent(Mooncake.fdata(t), Mooncake.rdata(t)) must reconstruct the original t.

  4. The mathematical model for Julia functions in Mooncake

    main

    Because Julia allows in-place mutation and memory allocation, Mooncake models a function f as a mapping:

    f : X -> X × A

    Where:

    • X is the real finite Hilbert space of the arguments before execution.
    • A is the real finite Hilbert space of any newly allocated data that is externally visible after execution (e.g., return values or data stored in Refs).

    The output of the function is a 2-tuple: (updated_state_of_arguments, newly_allocated_data).

    Externally-Visible Effects

    Mooncake considers only two ways a function can communicate results to the outside world:

    1. Return Value: The explicit result of the function.
    2. Modification of Arguments: In-place mutation of arguments (e.g., x .+= 1).

    Note on Global State: Mooncake does not support communication via global mutable state. Functions relying on global mutable state are explicitly unsupported to ensure correctness and avoid awkward handling.

  5. How Mooncake.jl's Rule System works

    main

    Mooncake.jl uses a recursive approach to Automatic Differentiation (AD). It defines a single specification for differentiating a Julia callable, primarily implemented via reverse-mode AD.

    A rule r(f, x) for a function f(x) operates in two distinct phases:

    1. Forwards Pass: Executes the original function and performs bookkeeping required for the reverse pass.
    2. Reverse Pass: Undoes the computation from the forwards pass, "backpropagates" the gradient with respect to the function's output by applying the adjoint of the derivative, and writes the results to the appropriate locations.

    To handle Julia's mutation capabilities, Mooncake models functions as transitions between states rather than simple mathematical mappings.

    # Concept: Rule Phases
# 1. Forwards Pass: f(x) -> (updated_args, new_data)
# 2. Reverse Pass: backpropagate gradients using the adjoint

  6. Understand the interaction between `@mooncake_overlay` and `@is_primitive`

    main

    Mooncake provides two ways to intercept differentiation: @mooncake_overlay (swaps the function body) and @is_primitive (dispatches to a hand-written rule). These do not compose and can lead to silent errors or runtime TypeErrors.

    Key Limitations:

    1. Shadowing: Marking a signature as a primitive (@is_primitive) shadows any overlay reachable through it. The rule fires, but the overlay is ignored.
    2. Return Type Mismatch: Mooncake infers the return type from the original method. If an overlay changes the return type, a manual rule returning the overlaid type will trigger a TypeError because it disagrees with Mooncake's inferred type.
    3. Inert Overlays: An overlay inside a primitive's body will never run because the rule executes the primal through ordinary Julia dispatch, which ignores Mooncake overlays.

    Best Practice: Apply only one of @mooncake_overlay or @is_primitive to a given signature, and ensure no overlay you rely on is hidden behind a primitive's rule.

  7. How Mooncake handles recursive function calls

    main
    Mooncake handles recursion by delaying code generation for generic function calls until the first time they are actually executed. This mechanism is encapsulated in the Mooncake.LazyDerivedRule abstraction. This approach allows the system to manage complex dependency graphs and recursive transformations without triggering immediate, potentially invalid, code generation.

  8. Understand the CFGBlock data structure and ID system

    main

    The CFGBlock representation differs from standard IRCode in three key ways:

    1. Implicit Control Flow: The CFG is not stored as a separate field. Instead, it is derived from the order of blocks and their terminators. To explicitly compute the graph, use Mooncake.control_flow_graph, Mooncake._compute_cfg_successors, or Mooncake._compute_cfg_predecessors.
    2. Instruction IDs: All SSA values (e.g., %1, %2) are replaced with unique Mooncake.ID objects. Each instruction in a block has a corresponding ID stored in the inst_ids field of the CFGBlock.
    3. Block IDs: Basic blocks are identified by a unique id field rather than their position in the IR.

    This ID-based system ensures that the 'name' of a block or instruction remains stable even when new blocks or instructions are inserted during transformations.

    # Accessing IDs in a CFGBlock
blocks[3].inst_ids      # Vector{ID} of instruction IDs
blocks[3].id            # The unique ID of the block itself

  9. Understand Control Flow and Basic Blocks in Julia IR

    main

    Julia's IR is organized into Basic Blocks. A basic block is a sequence of statements that execute sequentially without interruption. Control flow is managed by terminators at the end of a block, such as goto, goto if not, or return.

    • Terminators: Determine if the execution jumps to another block or returns.
    • Control Flow Graph (CFG): A data structure that maps the relationships between blocks. You can access it via the .cfg field of the IR object. It tracks successors (blocks that can run next) and predecessors (blocks that could have run immediately prior).
    # Example: Inspecting the CFG structure
function bar(x)
    if x > 0
        return x
    else
        return 5x
    end
end

ir = Base.code_ircode_by_type(Tuple{typeof(bar), Float64})[1][1]
println(ir.cfg)

  11. Understand the Mooncake reverse-mode rule structure

    main

    In Mooncake.jl, a reverse-mode rule for a function is designed to perform a forward pass and return an adjoint function (the reverse pass). For a function $f(x, y)$, a Mooncake-style rule rr(f, x, y) typically returns a tuple containing the function result and an adjoint function adj_f(db). This adjoint function takes the gradient of the output (db) and returns the gradients for all inputs (e.g., dx, dy) along with any necessary reverse-mode data (rdata).

    Key components of the rule execution:

    1. Forward-pass: The rule replaces calls to functions with calls to their respective rules.
    2. Reverse-pass: The adjoints are run in reverse order of the forward pass. If a variable is used multiple times, its adjoint contributions are added together.
    function f(x, y)
    a = g(x)
    b = h(a, y)
    return b
end

# A correct reverse-mode rule implementation:
function rr(f, x, y)
    a, adj_g = rr(g, x)
    b, adj_h = rr(h, a, y)
    function adj_f(db)
        _, da, dy = adj_h(db)
        _, dx = adj_g(da)
        return NoRData(), dx, dy
    end
    return b, adj_f
end

  12. Differentiating CUDA kernels in Mooncake.jl

    main

    Mooncake.jl supports differentiating CUDA kernels generally, provided a suitable rule exists. However, there are specific constraints:

    Unsupported Kernels:

    • Kernels that surface as foreign calls (e.g., those generated via KernelAbstractions.jl) are not supported.
    • Workaround: Provide a custom rule, potentially using another AD tool to assist.

    Second-order AD (HVP / Hessian) Constraints:

    • Supported: Array-level operations whose rules do not launch custom per-element kernels (e.g., sum(x), dot, matrix multiplication).
    • Unsupported: Operations that map a Julia function over array elements inside a GPU kernel (e.g., broadcasting, sum(f, x)-style reductions). These will raise an ArgumentError.
    • Note: Gradients and JVPs are unaffected by these second-order restrictions.