PJRT_Api Function-Pointer Table Reconstruction

All addresses on this page apply to libtpu 0.0.40 (cp314, manylinux_2_31_x86_64), build-id 89edbbe81c5b328a958fe628a9f2207d. The PJRT C-API surface is v0.103. Other libtpu releases pin different minors and will renumber the late-addition slots.

Abstract

libtpu.so is a PJRT plugin: it exports one symbol, GetPjrtApi, and everything a host framework (JAX, PyTorch/XLA) ever calls reaches the TPU through the single PJRT_Api struct that symbol returns. That struct is a C ABI vtable — a flat array of 140 qword slots: five scalar header fields followed by 135 function pointers, one per PJRT_* C-API entry point. The host reads api->PJRT_LoadedExecutable_Execute(args) and the call lands in libtpu's text section. There is no C++ inheritance, no dispatch object, no per-version trampoline table — just this one struct, populated once, immutable for process lifetime.

This page is the backbone reference for the PJRT section: the complete, ordered, slot-by-slot reconstruction of that table. It owns the struct header decode (struct_size, extension_start, the embedded PJRT_Api_Version), the full 140-slot map (every slot's field name, the libtpu implementation symbol it points at, and that symbol's virtual address), and the populated-vs-injected map (which five slots are TPU-specialized and which 130 are stock XLA wrappers). The slot ordering matches public xla/pjrt/c/pjrt_c_api.h v0.103 exactly, including the late-addition cluster appended in feature-addition order beyond the original v0.40 surface.

The vtable is assembled by one function — pjrt::CreatePjrtApi @ 0xf874160 — which does nothing but 140 stores into a caller-supplied buffer. That makes reconstruction unusually clean: the slot index is the array subscript in the initializer, and the field name is the assigned pjrt::PJRT_* symbol. Everything below is anchored to those stores. Per-area behavior (how each implementation actually works) lives on the area pages linked at the bottom; the extension linked list dangling off extension_start lives on extension-chain.md. This page is the index they all hang from.

For reimplementation, the contract is:

• The struct layout — 1120 bytes = 140 × 8; five scalar header slots (0..4), then 135 function pointers (5..139), in the exact order below.
• The version pin — PJRT_Api_Version = {struct_size=24, priv=NULL, major=0, minor=103}, encoded as the qword 0x6700000000 at offset +0x20.
• The slot-to-symbol map — every populated slot resolves to a text-section symbol; no slot in this build is left null.
• The five injection points — the only slots not compile-time-fixed: passed in as parameters to CreatePjrtApi by its caller.
• The per-slot backward-compat guard — every wrapper's first act is ActualStructSizeIsGreaterOrEqual, the mechanism that lets an older-minor host call this v0.103 plugin.

Slot count per area

Population Path and Storage

Purpose

The table is not a static initializer in a PROGBITS section — it is built lazily, once, on the first GetPjrtApi call, into zero-filled .lbss. Understanding the population path is the difference between a reimplementation that returns a correct table and one that races or returns garbage.

Entry Point

dlsym("GetPjrtApi")
  GetPjrtApi @ 0xE6A83A0                       ── thunk; tail-calls the builder
    └─ pjrt::tpu_plugin::GetTpuPjrtApi @ 0xE6AA440
         ── 16 __cxa_guard blocks build the .bss extensions (newest→oldest)
         ── final __cxa_guard wraps:
         └─ pjrt::CreatePjrtApi @ 0xF874160     ── 140 stores into &pjrt_api
              returns &pjrt_api = 0x227BA840

The exported symbol is a pure thunk. Its entire decompiled body is return pjrt::tpu_plugin::GetTpuPjrtApi(a1); (confirmed at 0xE6A83A0, marked thunk by IDA). GetTpuPjrtApi is the orchestrator: it runs one Itanium-ABI __cxa_guard-protected constructor per extension to build the extension chain in .bss, then a final guard around CreatePjrtApi, which writes every slot. The function returns &pjrt_api, the address of a function-local static.

Algorithm

function GetTpuPjrtApi():                       // 0xE6AA440
    // one-shot init of the extension chain (covered on extension-chain.md);
    // each extension is a __cxa_guard-protected function-local static in .bss
    build_extensions_if_needed();               // 16 guards, newest→oldest
    if (guard_for(pjrt_api) not yet set):        // _ZGV... byte == 0
        CreatePjrtApi(&pjrt_api,                  // 0xF874160 — writes all 140 slots
                      Client_Create_impl,         // -> slot 15  (a2)
                      ExecuteContext_Create_impl,  // -> slot 103 (a3)
                      TopologyDescription_Create_impl, // -> slot 87 (a4)
                      Plugin_Initialize_impl,      // -> slot 8   (a5)
                      extension_start,             // -> slot 1   (a6)
                      Plugin_Attributes_Xla_impl); // -> slot 9   (a7)
        mark_guard_set(pjrt_api);
    return &pjrt_api;                             // 0x227BA840

function CreatePjrtApi(a1, a2, a3, a4, a5, a6, a7):  // 0xF874160
    *a1     = 1120;                 // slot 0  struct_size
    a1[1]   = a6;                   // slot 1  extension_start  (injected: chain head)
    a1[2]   = 24;                   // slot 2  pjrt_api_version.struct_size
    a1[3]   = 0;                    // slot 3  pjrt_api_version.priv
    a1[4]   = 0x6700000000LL;       // slot 4  {major=0, minor=103}
    a1[5]   = pjrt::PJRT_Error_Destroy;   // slot 5  ... 130 fixed wrappers ...
    a1[8]   = a5;                   // slot 8  Plugin_Initialize    (injected)
    a1[9]   = a7;                   // slot 9  Plugin_Attributes    (injected)
    a1[15]  = a2;                   // slot 15 Client_Create        (injected)
    a1[87]  = a4;                   // slot 87 TopologyDescription_Create (injected)
    a1[103] = a3;                   // slot 103 ExecuteContext_Create (injected)
    // ... all other slots assigned a compile-time-fixed pjrt::PJRT_* symbol ...
    a1[139] = pjrt::PJRT_Executable_ParameterMemoryKinds;  // slot 139
    return a1;

NOTE — CreatePjrtApi contains no branches and no loop: it is 140 straight-line stores. The slot index is the C-array subscript, which is why the reconstruction below can be exact rather than inferred. The five injected slots (a2..a7) are the only values the function does not hard-code; everything else is a pjrt::PJRT_* relocation baked at link time. This was confirmed by reading the initializer in full at 0xF874160.

Considerations

The struct lives in .lbss (section [47], NOBITS, large-model BSS), at the function-local static _ZZN4pjrt10tpu_plugin13GetTpuPjrtApiEvE8pjrt_api @ 0x227BA840. Two consequences for a reimplementer and for any tool that inspects the binary statically:

• Static disassembly cannot show the populated table. The 1120 bytes are zero on disk and zero-filled at load; the function pointers only exist after the first GetPjrtApi call runs CreatePjrtApi. The slot map below is reconstructed from the stores in the initializer, not from a data dump — there is nothing to dump until runtime. Runtime confirmation would require reading /proc/self/mem at 0x227BA840 inside a live JAX worker after the first call.
• Concurrency is __cxa_guard, not a lock the reader holds. First callers serialize through the Itanium ABI guard (one thread runs the constructor, the rest block on its mutex). After the first call the struct is immutable, so steady-state readers take no lock. A reimplementation that rebuilds the table per call, or that omits the one-shot guard, diverges from this contract.

Note — the struct lives in .lbss at the concrete VA 0x227BA840, not in .bss/.data.rel.ro; any &stru_… + _GLOBAL_OFFSET_TABLE_ form in a disassembler is a PIC-relocation artifact, not a real address. Slot 9 is not a TPU-plugin override — it is fed from CreatePjrtApi's a7 parameter = pjrt::PJRT_Plugin_Attributes_Xla, a stock XLA implementation.

Struct Header (Slots 0..4)

Purpose

The first five qwords are not function pointers. They are the ABI handshake: total size (for forward/backward compat), the extension-chain head, and an embedded PJRT_Api_Version substruct. A host reads these before touching any function pointer.

Layout

struct PJRT_Api_Version {       // embedded at PJRT_Api +0x010, 24 bytes
    size_t struct_size;          // = 24
    void*  priv;                 // = NULL
    int    major_version;        // = 0
    int    minor_version;        // = 103 (0x67)
};

The version qword decodes by little-endian byte layout: bytes at +0x20 are 00 00 00 00 67 00 00 00, i.e. major = low 32 bits = 0, minor = high 32 bits = 0x67 = 103. struct_size = 1120 lets a host compiled against a different minor know exactly how many trailing slots it may safely read.

QUIRK — extension_start (slot 1) points at the newest extension, not the oldest. The chain is built in __cxa_guard order and linked head-insert, so walking .next goes newest→oldest and terminates at the profiler extension in .data. A host that assumes "type 0 first" or any type ordering is wrong; it must iterate to NULL. The chain itself is documented on extension-chain.md.

The 140-Slot Map

Purpose

The complete ordered table. Slots 0..4 are the scalar header above. Slots 5..139 are function pointers; each cell below names the field (per pjrt_c_api.h v0.103) and the libtpu symbol the slot points at, with that symbol's virtual address. All 135 function-pointer slots are populated in this build — there is no null/unimplemented slot. Confidence is CERTAIN for every slot whose store and target symbol were both read in the decompiled initializer and confirmed against the symbol table; the 12 spot-confirmed addresses are flagged.

Slot ordering note

The order is grouped by area (Error/Plugin/Event/Client/DeviceDescription/Device/Memory/Executable/LoadedExecutable/Buffer/CopyStream/Topology/Compile) through slot 94, then a flat late-addition cluster (slots 95..139) appended in feature-addition order across minors 0.40→0.103. The late cluster is not re-grouped by area — Buffer_*, Client_*, Executable_*, and Event_* entries are interleaved in the order XLA added them. The "Area" column tags each late slot back to its conceptual home; the area deep-dive pages own the behavior.

Error / Plugin / Event (slots 5..14)

Client (slots 15..27)

DeviceDescription / Device (slots 28..39)

Memory (slots 40..44)

Executable (slots 45..54)

LoadedExecutable (slots 55..62)

NOTE — slot 62 PJRT_LoadedExecutable_Fingerprint is the deprecated fingerprint entry; v0.103 hosts use slot 99 PJRT_Executable_Fingerprint instead. Both are populated; the deprecated one is kept for older-minor callers.

Buffer (slots 63..81)

CopyToDeviceStream (slots 82..86)

TopologyDescription / Compile (slots 87..94)

Late-addition cluster (slots 95..139, addition order)

These slots were appended across minors 0.40→0.103, in the order XLA added them — not regrouped by area. The Area column tags each back to its conceptual home.

Populated vs Injected Map

Purpose

Every one of the 135 function-pointer slots is populated; none is null in this build. But the slots split into two provenance classes, and the split is the heart of the plugin pattern: 130 slots are generic XLA C-API wrappers baked at link time, and 5 are injected at runtime by CreatePjrtApi's caller.

The five injection points

The only slots CreatePjrtApi does not hard-code are its a2..a7 parameters. GetTpuPjrtApi passes the TPU-specific implementations into exactly these positions; everything else is a fixed relocation.

The four tpu_plugin::-namespaced impls (slots 8, 15, 87, 103) are the genuine TPU specializations. Slot 9 is injected but resolves to the stock pjrt::PJRT_Plugin_Attributes_Xla — the caller passes the XLA implementation, not a TPU override. Slot 1 is the extension-chain head, also injected (see extension-chain.md). The remaining 130 function pointers are compile-time-fixed pjrt::PJRT_* symbols from XLA's pjrt_c_api_wrapper_impl.cc.

QUIRK — "TPU-specialized" is a much smaller set than a reimplementer might expect. Only four slots carry TPU-specific code; the other 131 function pointers are byte-for-byte the same generic XLA wrappers any PJRT plugin would ship. The TPU-ness lives almost entirely behind Client_Create (which builds the PjRtClient the generic wrappers then dispatch through) and in the extension chain — not in the per-call wrappers.

Backward-Compatibility Guard

Purpose

A host compiled against an older PJRT minor passes a smaller _Args struct than this v0.103 plugin expects. The plugin must accept it and read only the fields the caller actually provided. The mechanism is a per-slot size check at the top of every wrapper.

Algorithm

function PJRT_Error_Destroy(args):                  // 0xF85ECE0
    rc = pjrt::ActualStructSizeIsGreaterOrEqual(     // 0xF8A4EC0
            "PJRT_Error_Destroy_Args",                // API name (for the log line)
            23,                                       // min accepted struct_size
            24,                                       // current (v0.103) struct_size
            args->struct_size);                       // what the caller passed
    if (rc != ok):
        LogMessage(".../pjrt_c_api_wrapper_impl.cc", 545) << status;  // diagnostic
    if (args->struct_size < 0x18): return;            // too small: bail
    // ... read only fields within args->struct_size ...

Every wrapper's first instruction is this ActualStructSizeIsGreaterOrEqual call. The hard-coded (min, current) pair bounds the per-_Args versions the plugin accepts; for PJRT_Error_Destroy_Args it is (23, 24). Fields beyond the caller's struct_size are never read. This is what lets one v0.103 plugin serve a range of host minors without per-version trampolines — the size prefix on each args struct, not a version branch, carries the compatibility logic.

NOTE — the source-path string third_party/tensorflow/compiler/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc is baked into the diagnostic at 0xF85ECE0. It confirms the 130 fixed wrappers are XLA's generic pjrt_c_api_wrapper_impl.cc code, identical across PJRT plugins — the binary's own provenance string, not an external claim.

Hot-Path Slots

The slots a JAX/PyTorch-XLA step touches most. This ranking is a semantic estimate — no call-count instrumentation was run, so confidence is LOW on the ordering, though the slot identities are CERTAIN.

GOTCHA — the Top-20 ranking in the source findings is an ordering heuristic, not measured data. A reimplementer optimizing the dispatch path should treat slot 60 (Execute) and the Event slots (10..14) as certainly hot, but should not trust the precise rank order of the cooler slots without instrumentation.

Cross-References

• PJRT Overview — where this table sits in the plugin lifecycle; the GetPjrtApi → GetTpuPjrtApi → CreatePjrtApi chain in context
• Extension Chain — the linked list dangling off extension_start (slot 1); 17 extensions, type-IDs, per-extension layouts
• Client and Device — slots 15..44: Client_Create (the key TPU injection) plus device/memory queries
• Buffer and Memory — slots 63..81, 97, 105, 125, 130, 136: buffer lifecycle, transfers, DMA
• Executable Execution — slots 45..62, 122, 135 and slot 60 LoadedExecutable_Execute, the hot path
• Events and Async — slots 10..14, 131..132: the PJRT_Event model
• Callbacks — OnReady / async-tracking-event slots and the callback extension
• Collectives Communicator — the collectives extension surface reached via the chain, not via these slots