ConvolutionCostState

All addresses on this page apply to libtpu.so from the libtpu-0.0.40-cp314 wheel (BuildID md5 89edbbe81c5b328a958fe628a9f2207d). The binary is not stripped — every symbol below is a demangled C++ name. Section map: .text/.rodata VMA == file offset; .data.rel.ro VMA − 0x200000 == file offset. All addresses are VMA.

Abstract

ConvCostState is the running state the cost model builds once per convolution before it prices the op. It is the conv-shaped analogue of an LLVM MachineFunction's analysis cache: extract the conv-like HLO once, decode its ConvolutionDimensionNumbers, kernel Shape, and Window, map each logical dimension to its physical (layout-mapped) position, run the Target's per-axis chunk grid over the three operand shapes, and stash the result so the per-kernel pricing leaf can read it directly. The struct is built by CostModel::GetConvolutionCostState @0x130a4b20 and is the const ConvCostState& argument carried through the whole conv-cost emitter chain.

There are in fact two distinct structs, and conflating them is the central trap. ConvCostState is the large persistent state (1180+ bytes — protos at +0x158/+0x1e8/+0x328, six layout-mapped dim sizes at +0x350..+0x378, a set of per-dim vector masks at +0x380..+0x458, the 21-byte ConvolutionLoweringStrategy dtype/pack flags at +0x470..+0x484, and three inlined_vector<long,6> chunk-count grids at +0x8/+0x78/+0xe8). ConvState is the small per-call dim-product struct (+0x0..+0x78) that CostModel::RecordConvolutionCycles @0x130b6ce0 builds inside the deposit, on the stack, from ConvCostState's chunk-count vectors divided by the Target geometry. The pricing leaf CostModel::RecordConvKernelCycles @0x130caf20 reads dtype/pack flags from ConvCostState and dim products from ConvState, multiplies them into an op count, and deposits cycles into the Matmul / Matpush / Xlu resource slots.

This page documents the field offsets of both structs and the throughput bridge that turns the conv's op count into cycles. The bridge is the central detail: the matmul-rate multiplier the conv deposit needs is VfCycleTable::GetCyclesForThroughput(CT 0), which is not a constant — it is a call into MxuLatencyTable::GetResourceUsage(instr 212, res 3), and the res argument is a resource-index remap (res 3 → array index 15 = MatmulAccA; res 11 → array index 0 = MatpushPushPort), not a cycle seed. So the VfCycleTable throughput integer and the MxuLatencyTable reservation cycle of mxu-latency-overview are the same number viewed two ways.

For reimplementation, the contract is:

• The ConvCostState field map: the three copied protos, the six logical→physical dim sizes, the per-dim vector masks, the ConvolutionLoweringStrategy flag bytes (read only for opcode 0x2b=convolution), and the three chunk-count vectors.
• The ConvState field map: the dim products at +0x8..+0x78, the four shape-case flag bytes at +0x0..+0x6, and the Product(chunk_vector)/SublaneCount/ChunksPerTile formula that fills each.
• The Matmul/Matpush/Xlu deposit: op count × thru(CT) × format-rate (0.5/0.25) ÷ Target+0x4ac (matmul rate) or Target+0x4b0 (Xlu rate) ÷ ElementPackingFactor.
• The VfCycleTable → MxuLatencyTable bridge and its res→index remap.
• That Target+0x4ac/+0x4b0 are Target vector-ISA rate fields, not ConvCostState-local derate fields.

ConvCostState — The Persistent State

Purpose

GetConvolutionCostState @0x130a4b20 has the signature bool GetConvolutionCostState(const HloInstruction* inst, ConvCostState* out). It extracts the conv-like HLO (fusion_util::ExtractConvLikeHlo @0x1d6aa140; a null result is a LogMessageFatal on conv_like_hlo != nullptr at cost_model.cc:3374), then fills the struct top-to-bottom: copy the three shape/window protos, run ChunkCountsWithTmp over the three operand shapes, map each logical conv dimension to its physical position via LayoutUtil::MakeLogicalToPhysical, initialise the per-dim vector masks, classify each dim (batch / outer / inner) into the three mask vectors, and finally — only for a real convolution opcode — copy the ConvolutionLoweringStrategy flag block.

Field Map

a3 is the ConvCostState*. Offsets are byte offsets from the struct base; the decompiler renders them as (char*)a3 + N (proto copies) or *((_QWORD*)a3 + N/8) (scalar/pointer stores).

NOTE — The three chunk-count inlined_vector<long,6> members sit at +0x8/+0x78/+0xe8 on a 0x70 stride (112 B per member, including its header): the three Storage<long,6>::Assign calls target (__int64*)a3 + 1, +15, +29 immediately after the three ChunkCountsWithTmp calls. The first vector is fed from the operand shape (v128), the second from the output shape, the third from the kernel shape — easy to mis-bind if the operand vector is overlooked.

The ConvolutionLoweringStrategy Flag Block

The +0x470..+0x484 region is the 21-byte ConvolutionLoweringStrategy of mxu-latency-overview's sibling. It is copied as one vmovups [r15+0x470] of the 16-byte head plus a qword tail at +0x47d, only when the conv-like HLO's opcode byte ([conv+12]) is 0x2b (convolution). RecordConvKernelCycles reads it byte-by-byte as the dtype / packing decision. The byte offsets within the struct (confirmed by the byte reads in @0x130caf20):

These bytes OR-combine (@0x130caf38..af86) to decide whether TransferSizeUtil::ElementPackingFactor @0x1d6b03e0 is consulted (called only if any flag is set) versus the rbx = 1 / 2 default. The flag names are inferred from the ConvolutionLoweringStrategy proto field order and the gate semantics; the offsets and read sites are byte-confirmed.

ConvState — The Per-Call Dim-Product Struct

Purpose

ConvState is built on the stack (at rbp-0x1d0) inside RecordConvolutionCycles (4-RV) @0x130b6ce0 and passed by const& to RecordConvKernelCycles. It is not persisted. Each dim field is a Product(chunk_dim_sizes) over one of ConvCostState's chunk-count vectors, divided by a Target geometry constant (SublaneCount() = 8, LaneCount() = 128, ChunksPerTile() = 16), with paired remainder / divisibility flags. The three xla::Product calls (@0x130b6f1d/6f43/6f69) reduce the operand / kernel / output chunk vectors; the four leading bytes are the conv-shape-case flags the deposit branches on.

Field Map

a3 (in RecordConvKernelCycles) is the ConvState*. The decompiler renders fields as a3[k] (8-byte stride) or a3[byte].

The pair (+0x60, +0x68) plus the [+0x3]^1 XOR-select (@0x130cafc3) is the "which-feature-dim-first" choice: the deposit reads a3[16*a3[3] + 104] (=[+0x68]-region) and a3[16*(a3[3]^1) + 96] (=[+0x60]-region). The op count is a3[3] · a3[2] · a3[4] · a3[1] (the [+0x18]·[+0x10]·[+0x20]·[+0x8] block) × [+0x50] × feat0 × feat1, then a branch on the flag bytes [+0x0], [+0x2], [+0x3] folds in [+0x48] / [+0x70] / [+0x78] (the four conv-shape cases at @0x130cb011/0c0/0ee/113), finally divided by ElementPackingFactor and ChunksPerTile.

NOTE — the semantic role of each ConvState dim offset (which is output-feature vs input-feature vs spatial vs batch vs kernel-window) follows the imul chain and the ChunkCountsWithTmp source order; this build carries no DWARF, so the dimension names are an interpretation while the offsets and the formula shape are taken straight from the disassembly. The authoritative source is GetConvIterationCounts @0x130b3c80.

The Throughput Bridge — VfCycleTable → MxuLatencyTable

Why a Bridge Exists

RecordConvKernelCycles needs a throughput (cycles-per-op) for the matmul and the matpush, not the per-op latency of the base Performance grid. That throughput is the per-resource occupancy from mxu-latency-overview. VfCycleTable::GetCyclesForThroughput(CycleTable::Instruction) @0x1c89e2c0 is the adapter: it maps each 32-entry CycleTable::Instruction (CT) ordinal to the right call into the MXU occupancy table. The conv deposit reaches it through the Target's CycleTable vtable — the (*(…+16LL))(this, CT, …) indirect calls in @0x130caf20 (e.g. CT 0 at @0x130cb..., CT 1 at @0x130cb375).

The CT → (instr, res, transpose) Table

The decompile of @0x1c89e2c0 (a switch(a2) over CT, default → return 1) confirms three call shapes: matmul/matpush CTs call MxuLatencyTable::GetResourceUsage; mxres/Xlu CTs call ViperfishPerformance::GetResourceUsage(…, res 14) + 1; CT 10/16 are a LogMessageFatal("Unsupported PushGainsS4.") at cycle_table.cc:682.

The transpose flag is the 4th GetResourceUsage arg; the matpush key pre-transforms (267 direct, 271 → latch_mode ^ 0xB, 277 → latch_mode | 0x14) are applied inside GetResourceUsage (see mxu-latency-overview).

The Throughput Bridge — res Is a Resource-Index Remap

MxuLatencyTable::GetResourceUsage(instr, res, transpose) @0x1c8ae5c0 does not index its array<int,19> reservation vector with res directly. It first remaps:

function GetResourceUsage(instr, res, transpose):       // VF @0x1c8ae5c0
    if   res == 3:  idx = 15        // → MatmulAccA
    elif res == 11: idx = 0         //  → MatpushPushPort
    else: return InvalidArgument("Unsupported kind of resource")   // mxu_latency_table_vf.cc:547

    // build the modifier key from the opcode, then find:
    switch (instr):
        case 212: key = MatmulModifier{format=1}; map = this+0x20
        case 218: key = MatmulModifier{format=2}; map = this+0x20
        case 230: key = MatmulModifier{format=6}; map = this+0x20
        case 267: key = MatpushKey(latch_mode);          map = this+0x00
        case 271: key = MatpushKey(latch_mode ^ 0xB);    map = this+0x00
        case 277: key = MatpushKey(latch_mode | 0x14);   map = this+0x00
        default:  LogFatal("Unsupported opcode")          // vf.cc:578

    entry = map.find(key)                 // miss → out_of_range
    array<int,19> v = entry.value         // vmovups: matmul reads [rdx+8], matpush reads [rdx+4]
    CHECK(idx < 0x13)                      // == 19
    return (float) v[idx]                  // the hold-cycle count

So the res argument 3/11 is a named-sub-resource selector, remapped to the concrete array index (3 → 15 = MatmulAccA, 11 → 0 = MatpushPushPort), and that indexes the per-modifier reservation array. The matmul CTs (0, 1, 4) therefore read array[15]; the matpush CTs (5, 6, 9, 11, 12, 15) read array[0]. The VfCycleTable throughput integer is the MxuLatencyTable reservation cycle — the two cost tables are the same numbers.

NOTE — the non-bf16 entries follow the per-format reservation groups; the live MatmulModifier/MatpushModifier maps per MatmulDataFormat give the exact values. The {2,1,1} / {4,3,2} / {8,7,6} reservation triplet's first element (the MatpushPushPort value) is exactly thru(CT matpush); the other two (matrix-staging-register holds) are read by separate GetResourceUsage calls in the MXU scheduler, not by this bridge.

Per-Generation Bridge

The bridge is per-generation. The Ghostlite helper GlcCycleTable::GetCyclesForThroughputHelper @0x1c89ed20 is structurally identical but uses the Ghostlite instruction set — matmul CTs (0, 1, 4) use instrs 292/298/310 (0x124/0x12a/0x136); matpush CTs (5, 6, 9, 10) use 347/349/356/358 (0x15b/0x15d/0x164/0x166); the GL MxuLatencyTable array<int,11> remaps res 4 → idx 3 and res 9 → idx 9. GfcCycleTable::GetCyclesForThroughputHelper @0x1c89f400 is the TPU7x variant with the same shape (array<int,11>). A reimplementation must read the (instr, res) pair from the per-gen helper.

The Deposit — Putting It Together

RecordConvKernelCycles @0x130caf20 computes one op count from ConvState and deposits cycles into three slots (ResourceVector::Acc(slot, cycles)):

The 0.5 (.rodata 0xa2df5c8) is the half-rate the matmul multiplies by when the format flag v166 is clear; the 0.25 (.rodata 0xa2df6d0) replaces it only when ConvCostState+0x480 (generate_x8_packed_vlatches) == 1 and Target+0x398 == 2 (the topology gate at @0x130cb468/cb471). The matmul-rate divisor is read as vcvtsi2sd xmm1, [rax+0x4ac] (Target+0x4ac); the Xlu divisor as [rax+0x4b0] reached through the CycleTable's Target pointer.

GOTCHA — Target+0x4ac (matmul rate) and Target+0x4b0 (Xlu rate) are Target fields, not ConvCostState fields. They are chip-wide vector-ISA throughput-rate parameters populated in Target::Init from TpuSequencerParts::vector_isa() @0x20b31840: vector_isa[+0xc] (8 bytes) → Target+0x4ac & +0x4b0; vector_isa[+0x14] → Target+0x4a8. The conv deposit reads them through *(_QWORD*)a1 (the CostModel's Target*) and through the CycleTable's Target*, never out of ConvCostState — they are not a conv-local derate.

Because Matmul, Matpush, and Xlu (R[0]/R[1]/R[2]) are in the plain-MAX group of MaxResourceCycles (the MXU pipes overlap — see resource-enum), a compute-bound conv's bundle cost is ≈ max(Matmul, Matpush, Xlu), not their sum.

Worked Example — bf16 Conv on Viperfish

thru(CT 0)  = VfCycleTable::GetCyclesForThroughput(0)
            = MxuLatencyTable::GetResourceUsage(instr=212, res=3, false)
            = MatmulModifier{format=1 bf16}, array[remap(res=3)=15]
            = array[15] = MatmulAccA = 8

R[1] Matmul  = op_count × 8 × 0.5 ÷ Target[+0x4ac] ÷ ElementPackingFactor[bf16=1]

thru(CT 5)   = GetResourceUsage(instr=267, res=11, false) = array[0] = MatpushPushPort = 2
R[0] Matpush = matpush_tile_count × 2

thru(CT 28)  = VfPerf.GetResourceUsage(287, res=14) + 1            (the Xlu mxres-read rate)
R[2] Xlu     = thru(CT 28) × ChunksPerTile × remainder ÷ Target[+0x4b0]

bundle cost ≈ max(R[0], R[1], R[2])                                (MXU pipes overlap)

For an int8-x8 conv the matmul format selects group 4 → array[15] = 32 (matmul), array[0] = 8 (matpush) — 4× the bf16 rate, the four-byte-plane x8 latch sequence.

Sibling — Reduce-Window Cost Reuses ConvState

CostModel::RecordReduceWindowCycles @0x130b5ec0 is the pooling sibling. It takes three Windows directly (not a ConvCostState) and builds the same stack ConvState at rbp-0x1d0 by the same Product / SublaneCount / ChunksPerTile division. A trivial path: when CostModel+0x14 is false it deposits zero into R[7] VectorLoad, R[3] VectorAlu0, R[4] VectorAlu1, R[2] Xlu and returns (Acc(a9, 7/3/4/2) @line 782..788). Otherwise it dispatches on GetReduceWindowType @0x1454d4a0 (0 → RecordLaneReduceWindowCycles, 1 → RecordSublaneReduceWindowCycles, 2 → RecordMajorReduceWindowCycles), the Lane and Sublane paths gated on !HasBaseDilation. The lane leaf (@0x130c97e0) deposits R[7] VectorLoad (op_volume = ConvState[+0x20]·[+0x18]·[+0x70]), an R[5] VectorAluAny per-element combine (thru(CT 22)), the UpdateCostBasedOnReductionFunction combiner cost (count = op_volume × (ConvState[+0x68] − 1)), and an R[2] Xlu cross-lane drain (thru(CT 27) ÷ Target+0x4b0). The full pooling cost is on reduce-window-pooling-cost.

Function Map

Related Components

Cross-References

• MXU Latency Overview — the MxuLatencyTable reservation model, the MatmulModifier/MatpushModifier keys, and the array[idx] read this bridge surfaces
• VfCycleTable — the 32-entry CT → (instr, res) dump and the full GetCyclesForThroughput per-gen bridge
• WindowDescription Byte-Cost — the conv/DMA byte + throughput primitive
• Reduce-Window / Pooling Cost — RecordReduceWindowCycles, the pooling sibling that reuses ConvState
• Resource Enum (23-slot) — the ResourceVector slots (Matmul/Matpush/Xlu) the conv deposit writes and the MaxResourceCycles overlap rule
• MXU Slot — the LLO MXU instructions (matmul / matpush / matres) the conv deposit prices