Efficient memory ordering for DiscreteDP

[oyamad]

(This is to summarize my comments #109 (comment), #109 (comment), #109 (comment).)

• n: number of states
• m: number of actions (Assume it is constant over the states)

Consider bellman_operator, where we compute

• vals = R + beta * Q * v and
• Tv = s_wise_max(vals).

Let me denote R + beta * Q * v by U(s, a) as a function of (s, a), to distinguish it from a Julia array; let me denote the probabilities by q(sa, s'), where sa is a state-action pair, while s' is a state tomorrow.

To compute s_wise_max, the data for U should be ordered in memory as

U(1, 1), ..., U(1, m), U(2, 1), ..., U(2, m), ..., U(n, 1), ..., U(n, m)

i.e., action should change faster.

So the same applies to R: action should change faster.

To compute Q * v, the data for q should be ordered in memory as

q(sa_1, 1), ..., q(sa_1, n), ..., q(sa_L, 1), ..., q(sa_L, n)

(where the state-action pairs are listed as sa_1, ..., sa_L)

i.e., state-tomorrow should change faster.

Given 1. and 2., for Q, state-tomorrow should change first, action second, and state-today third.

[sglyon]

This was very helpful, thank you for summarizing here.

From the very last statement, it sounds like in Julia we need to order Q(s', a, s) and in python we would need to order it Q(s, a, s'). Do you agree?

[oyamad]

Yes, that's my opinion.

[oyamad]

With #386 and #388 merged, the actionable content of this issue is implemented, and the remaining parts were measured and deliberately declined. Summary of the resolution:

The analysis in this thread was correct: the two hot kernels impose a joint memory order — state-tomorrow fastest, action second, state-today third — which is what NumPy's row-major Q[s, a, s'] gives QuantEcon.py for free, and what Julia's column-major layout inverts.

Implemented

• Dense product form: the expected values have always been computed as a single gemv over reshape(Q, n*m, n), which is near-optimal in column-major despite the layout (each summand column is contiguous). PERF: Optimize DiscreteDP kernels and solver loops #386 fixed the kernels that did not respect the layout: the s_wise_max! argmax kernel now traverses vals in column-major order, and RQ_sigma builds Q_sigma with a direct loop fill (54x). With those, the dense case needs no storage change.
• Sparse sa form: PERF: Store sparse sa-formulation Q transposed internally #388 stores Q transposed internally (states-tomorrow × sa-pairs), so each pair's transition probabilities occupy a contiguous CSC column, with ddp.Q a lazy Transpose view preserving the documented (L, n) shape and values. Q * v dispatches to the transpose matvec kernel and the RQ_sigma gather becomes a column gather (64x); on the benchmark suite, sparse MPFI 6.5x, VFI 1.7x, Bellman operator 1.5x.

Measured and declined

• Permuting the dense product-form storage to Q[s', a, s] (the pure-layout answer discussed above): the reshape-gemv path plus the PERF: Optimize DiscreteDP kernels and solver loops #386 kernel fixes already realize the layout benefits, and changing the user-facing convention would break parity with QuantEcon.py and the lectures for no measured gain. The user-facing constructor conventions remain identical to QuantEcon.py; the optimization is purely internal.
• Transposing the dense sa-form Q: measured a full-copy cost at construction (20 ms and 2x transient memory at L=50000, n=500) against ~10% solve-time gains — not worth it.

Outcome: with the benchmark suite added in #383, QuantEcon.jl now matches or beats QuantEcon.py on every benchmarked DiscreteDP operation (details in #388). #118 is resolved on the same grounds and can be closed together with this issue.

🤖 Generated with Claude Code (Claude Fable 5)