Task: inventory_replenishment_dp
Category: Operations research
Difficulty: 3
Tags: dynamic programming,inventory control,multi-dimensional state space,stochastic optimization
Source: See \
Description: Sim 2010, "Robust Approximation to Multiperiod Inventory Management," Operations Research: https://doi.org/10.1287/opre.1090.0746 & Multi-period stochastic inventory-replenishment dynamic program over inventory states and demand outcomes.

Prompt template (YAML):
task_description: |
  A finite-horizon dynamic programming problem for multi-period inventory replenishment with autoregressive demand.
  The state space is 4-dimensional, tracking current inventory `I`, accumulated backlog `B`, and two statistics of past demand shocks: sum `Phi` and sum of squares `Psi`.
  At each time step `t` (from `T-1` down to `0`), and for each state in the 4D grid, the algorithm must find the optimal adjustment quantity `x` that minimizes expected holding, backlogging, and adjustment costs by integrating over the distribution of the current demand shock `phi`.

  The state transitions are:
  - Realized demand: `xi = mu + nu * (phi + alpha * Phi)`
  - Next inventory: `I' = I + y + x - xi`
  - Next backlog: `B' = B + max(0, -I')`
  - Next Phi: `Phi' = Phi + phi`
  - Next Psi: `Psi' = Psi + phi^2`

  The immediate cost is:
  `cost = c_t[t] * |x| + h_t[t] * max(0, I') + b_t[t] * max(0, -I')`

  The expected cost for a given action `x` is the sum over all `phi` realizations of:
  `prob * (cost + V_{t+1}(I', B', Phi', Psi'))`
  where `V_{t+1}` is the value function at the next time step, evaluated using multi-linear interpolation on the 4D grid.
  The value `V_t(I, B, Phi, Psi)` is the minimum expected cost over all actions `x`.

  The boundary conditions for interpolation: if a value falls outside the grid, it should be clamped to the nearest grid point (i.e., the boundary).

interface: |
  void solution_compute(
      int N_I, int N_B, int N_Phi, int N_Psi, int N_x, int N_phi, int T,
      const float* c_t, const float* h_t, const float* b_t,
      float mu, float nu, float alpha, float y,
      const float* grid_I, const float* grid_B, const float* grid_Phi, const float* grid_Psi,
      const float* actions, const float* shocks, const float* shock_probs,
      float* V_out
  );

input_size_notes: |
  - `N_I`, `N_B`, `N_Phi`, `N_Psi` are the dimensions of the 4D state grid (up to ~32 each).
  - `N_x` is the number of discrete actions (up to ~32).
  - `N_phi` is the number of discrete demand shocks (up to ~32).
  - `T` is the time horizon (up to ~15).
  - The total number of states per time step is `N_I * N_B * N_Phi * N_Psi`, which can be over 1 million.
  - The total number of iterations per time step is `Total States * N_x * N_phi`, which can be over 1 billion.

output_contract: |
  - The output array `V_out` should contain the computed value function `V_0` for all states at the initial time step `t=0`.
  - The array should be flattened in row-major order: `(((iI * N_B + iB) * N_Phi + iPhi) * N_Psi + iPsi)`.
  - The values should be numerically accurate (within 0.01 relative/absolute tolerance).

algorithm_background: |
  The algorithm performs backward induction. Initialize `V_T` to all zeros.
  For each time step `t` from `T-1` down to `0`:
    For each state `(I, B, Phi, Psi)`:
      `min_cost = INF`
      For each action `x`:
        `expected_cost = 0`
        For each shock `phi` with probability `prob`:
          Compute next state `(I', B', Phi', Psi')`
          Compute immediate `cost`
          Interpolate `V_{t+1}` at `(I', B', Phi', Psi')` to get `future_cost`
          `expected_cost += prob * (cost + future_cost)`
        `min_cost = min(min_cost, expected_cost)`
      `V_t(I, B, Phi, Psi) = min_cost`

  Multi-linear interpolation in 4D requires finding the bounding hyper-box for the target point and computing a weighted average of the 16 corner values.

hints_l2: |
  - This problem is perfectly suited for GPU acceleration. Each state in the 4D grid can be evaluated independently by a single CUDA thread.
  - The outer loop over time steps `t` must be sequential, but the inner loops over the state space can be fully parallelized.
  - You will need two buffers for the value function: `V_curr` and `V_next`. Swap them after each time step.
  - The 4D interpolation requires 16 memory reads from `V_next`. Optimizing these reads is critical for performance.
  - Consider using CUDA textures or `__ldg()` for reading `V_next` to benefit from the read-only cache.
  - The action `x` and shock `phi` loops should be kept thread-local.

hints_l1: |
  - Allocate `V_curr` and `V_next` on the device using `cudaMalloc`.
  - Copy the grid arrays (`grid_I`, `grid_B`, `grid_Phi`, `grid_Psi`), `actions`, `shocks`, and `shock_probs` to the device.
  - Launch a kernel with a 1D or 3D grid that covers all `N_I * N_B * N_Phi * N_Psi` states.
  - Inside the kernel, compute the 4D indices `(iI, iB, iPhi, iPsi)` from the thread index.
  - Implement a device function for the 4D interpolation. It should first find the 1D indices and weights for each dimension (using binary search or linear search, clamping to boundaries), and then compute the weighted sum of the 16 corners.
  - After the time step loop, copy `V_curr` back to the host array `V_out`.


CPU reference implementation (C):
#include <stdlib.h>
#include <math.h>
#include <string.h>

static void get_interp_weights(float val, const float* grid, int N, int* idx, float* w) {
    if (val <= grid[0]) {
        *idx = 0;
        *w = 0.0f;
    } else if (val >= grid[N-1]) {
        *idx = N - 2;
        *w = 1.0f;
    } else {
        int low = 0, high = N - 1;
        while (low <= high) {
            int mid = low + (high - low) / 2;
            if (grid[mid] <= val) {
                low = mid + 1;
            } else {
                high = mid - 1;
            }
        }
        *idx = high;
        if (*idx < 0) *idx = 0;
        if (*idx >= N - 1) *idx = N - 2;
        *w = (val - grid[*idx]) / (grid[*idx+1] - grid[*idx]);
    }
}

void solution_compute(
    int N_I, int N_B, int N_Phi, int N_Psi, int N_x, int N_phi, int T,
    const float* c_t, const float* h_t, const float* b_t,
    float mu, float nu, float alpha, float y,
    const float* grid_I, const float* grid_B, const float* grid_Phi, const float* grid_Psi,
    const float* actions, const float* shocks, const float* shock_probs,
    float* V_out
) {
    int total_states = N_I * N_B * N_Phi * N_Psi;
    float* V_next = (float*)calloc(total_states, sizeof(float));
    float* V_curr = (float*)malloc(total_states * sizeof(float));

    for (int t = T - 1; t >= 0; t--) {
        for (int iI = 0; iI < N_I; iI++) {
            for (int iB = 0; iB < N_B; iB++) {
                for (int iPhi = 0; iPhi < N_Phi; iPhi++) {
                    for (int iPsi = 0; iPsi < N_Psi; iPsi++) {
                        float min_cost = 1e30f;

                        for (int ix = 0; ix < N_x; ix++) {
                            float x = actions[ix];
                            float expected_cost = 0.0f;

                            for (int iphi = 0; iphi < N_phi; iphi++) {
                                float phi = shocks[iphi];
                                float prob = shock_probs[iphi];

                                float xi = mu + nu * (phi + alpha * grid_Phi[iPhi]);
                                float I_next = grid_I[iI] + y + x - xi;
                                float B_next = grid_B[iB] + fmaxf(0.0f, -I_next);
                                float Phi_next = grid_Phi[iPhi] + phi;
                                float Psi_next = grid_Psi[iPsi] + phi * phi;

                                float cost = c_t[t] * fabsf(x) + h_t[t] * fmaxf(0.0f, I_next) + b_t[t] * fmaxf(0.0f, -I_next);

                                int idxI, idxB, idxPhi, idxPsi;
                                float wI, wB, wPhi, wPsi;
                                get_interp_weights(I_next, grid_I, N_I, &idxI, &wI);
                                get_interp_weights(B_next, grid_B, N_B, &idxB, &wB);
                                get_interp_weights(Phi_next, grid_Phi, N_Phi, &idxPhi, &wPhi);
                                get_interp_weights(Psi_next, grid_Psi, N_Psi, &idxPsi, &wPsi);

                                float future_cost = 0.0f;
                                for (int dI = 0; dI <= 1; dI++) {
                                    float cI = dI ? wI : (1.0f - wI);
                                    for (int dB = 0; dB <= 1; dB++) {
                                        float cB = dB ? wB : (1.0f - wB);
                                        for (int dPhi = 0; dPhi <= 1; dPhi++) {
                                            float cPhi = dPhi ? wPhi : (1.0f - wPhi);
                                            for (int dPsi = 0; dPsi <= 1; dPsi++) {
                                                float cPsi = dPsi ? wPsi : (1.0f - wPsi);
                                                
                                                int flat_idx = (((idxI + dI) * N_B + (idxB + dB)) * N_Phi + (idxPhi + dPhi)) * N_Psi + (idxPsi + dPsi);
                                                future_cost += cI * cB * cPhi * cPsi * V_next[flat_idx];
                                            }
                                        }
                                    }
                                }
                                expected_cost += prob * (cost + future_cost);
                            }
                            if (expected_cost < min_cost) {
                                min_cost = expected_cost;
                            }
                        }
                        int flat_idx = (((iI) * N_B + (iB)) * N_Phi + (iPhi)) * N_Psi + (iPsi);
                        V_curr[flat_idx] = min_cost;
                    }
                }
            }
        }
        memcpy(V_next, V_curr, total_states * sizeof(float));
    }
    memcpy(V_out, V_next, total_states * sizeof(float));
    free(V_next);
    free(V_curr);
}

void solution_free(void) {}