linCmtB(): thread-safe forward-mode AD Jacobian solves (#1126)

* linCmtB(): thread-safe forward-mode AD Jacobian solves

Forward-mode AD (the default linCmtB Jacobian path, sensType 3/30/auto)
is stack-local fvar with no shared Stan arena, so parallel subject solves
are safe.  Two fixes enable it:

- linCmtScaleInitPar/N/linCmtZeroJac read the current thread's __linCmtB
  slot (rx_get_thread) instead of the hardcoded [0] slot that another
  thread mutates and may resize during a parallel solve.
- The notThreadLinCmtB solve-time guard now single-cores only the
  reverse-mode AD (31) and finite-difference paths; the forward-AD path
  threads like any thread-safe method.

Validated: 80-subject linCmtB sensitivity solve is bitwise identical at
cores=1 vs cores=4; reverse-AD path stays serial and correct.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* linCmtB(): regression test + NEWS for thread-safe forward-AD solves

Add a parallel-vs-serial regression test (cores=1 == cores=2, bitwise)
to the linCmtB sensitivity suite, refresh the stale "not thread safe"
parse-time comment, and note the change in NEWS.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* linCmtB(): skip finite-difference H setup on the AD paths

setupLinH() ran shi21ForwardH() for the default AD sensType (3, and
auto=100 which remaps to 3) -- a base solve plus several per-parameter
probe solves used to size a finite-difference step.  The AD Jacobian
(linCmtFwdJac / reverse jacobian) never reads ind->linH, so those solves
were pure waste on every linCmt sensitivity setup.

Give linCmtSensIsAD() external linkage (its finite-difference set is
exactly the sensTypes linCmtB reads ind->linH for) and early-out of
setupLinH() for the AD paths with a constant linH fill.  The now-dead
case 3 shi21ForwardH branch is removed.

Validated: AD linCmtB vs ODE reference unchanged (3.5e-7 < 1e-5);
lincmt-solve (1623) and lincmt-solve-sens (9) pass.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* linCmtB(): address PR review -- shared sensType classifier, no narrowing

- Move the linCmt sensitivity-method classifiers into a shared header
  src/linCmtSensType.h so linCmt.cpp, par_solve.cpp and rxData.cpp use one
  definition instead of a cross-TU forward declaration.  Add
  linCmtSensForwardAdThreadSafe() (3/30/100) and document why auto (100) is
  safe -- setupLinH() remaps 100->3 and nothing resolves auto to reverse-mode
  AD or a finite-difference method -- and why reverse-mode AD (31) and the FD
  paths are excluded.  The thread gate in rxData.cpp now calls this predicate.
- Cast the __linCmtB pool size to int once per accessor (linCmtScaleInitPar/N,
  linCmtZeroJac) since rx_get_thread() takes an int, avoiding a narrowing
  conversion.
- Skip the thread-safety regression test unless getRxThreads() >= 2 so it does
  not pass vacuously on a single-core runner.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: mattfidler

NEWS.md

@@ -86,6 +86,11 @@
   inside `linCmt()` and mixed `linCmt()`+ODE models, with Jacobian handling of
   the dosing events; `odeToLin()` preserves and renames them when converting.
 
+- `linCmt()` sensitivity (`linCmtB`) solves now run in parallel across subjects
+  on the default forward-mode AD Jacobian path (`linCmtSensType="AD"`), which is
+  stack-local with no shared Stan arena.  The reverse-mode AD (`"ADr"`) and
+  finite-difference paths remain single threaded.
+
 - Inductive linearization and matrix exponentials rewritten with a more
   NONMEM-like interface (automatic ODE->syntax translation retained) and
   symbolic-differentiation gradients.

src/linCmt.cpp

@@ -8,6 +8,7 @@
 #include "timsort.h"
 #define SORT gfx::timsort
 #include "linCmt.h"
+#include "linCmtSensType.h"
 
 
 extern rx_solving_options op_global;
@@ -395,17 +396,26 @@ extern "C" double linCmtA(rx_solve *rx, int id,
 #undef yp
 }
 
+// These scaling/step-size helpers are read from the finite-difference H setup
+// (shi21ForwardH/gillForwardH) after ind_linCmtFH has populated THIS thread's
+// linCmtB slot, so they must read the same per-thread slot rx_get_thread()
+// selects -- not the hardcoded [0] slot, which another thread mutates (and may
+// resize) during a parallel linCmt solve.  rx_get_thread() takes an int, so the
+// size_t pool size is cast once to avoid a narrowing conversion.
 extern "C" double linCmtScaleInitPar(int which) {
-  return __linCmtB[0].lc.initPar(which);
+  int tid = rx_get_thread((int)__linCmtB.size());
+  return __linCmtB[tid].lc.initPar(which);
 }
 
 extern "C" double linCmtScaleInitN() {
-  Eigen::Matrix<double, Eigen::Dynamic, 1> theta = __linCmtB[0].lc.initPar();
+  int tid = rx_get_thread((int)__linCmtB.size());
+  Eigen::Matrix<double, Eigen::Dynamic, 1> theta = __linCmtB[tid].lc.initPar();
   return theta.size();
 }
 
 extern "C" int linCmtZeroJac(int i) {
-  return __linCmtB[0].lc.parDepV1(i);
+  int tid = rx_get_thread((int)__linCmtB.size());
+  return __linCmtB[tid].lc.parDepV1(i);
 }
 
 
@@ -484,20 +494,10 @@ extern "C" int linCmtZeroJac(int i) {
  * @author Matthew Fidler
  *
 */
-// True for the automatic-differentiation jacobian methods (forward-mode
-// fvar = 3/30, reverse-mode = 31, and the auto/default path), which all want
-// the unscaled thetaSens + passthrough trueTheta. The finite-difference
-// methods (1,2,4,5,6,7,10,20,40,50) keep the scaled path.
-static inline bool linCmtSensIsAD(int sensType) {
-  switch (sensType) {
-  case 1: case 2: case 4: case 5: case 6: case 7:
-  case 10: case 20: case 40: case 50:
-    return false;
-  default:
-    return true;
-  }
-}
-
+// linCmtSensIsAD() (the AD jacobian classifier; forward-mode fvar 3/30,
+// reverse-mode 31, auto 100 -> unscaled thetaSens + passthrough trueTheta; the
+// finite-difference methods keep the scaled path) lives in linCmtSensType.h so
+// linCmt.cpp, par_solve.cpp and rxData.cpp share one definition.
 extern "C" double linCmtB(rx_solve *rx, int id,
                           double _t, int linCmt,
                           int ncmt, int oral0,

src/linCmtSensType.h

@@ -0,0 +1,36 @@
+#ifndef __LINCMTSENSTYPE_H__
+#define __LINCMTSENSTYPE_H__
+// Single source of truth for classifying the linCmt() sensitivity (Jacobian)
+// method encoded in rx->sensType.  Shared by linCmt.cpp (thetaSens scaling),
+// par_solve.cpp (setupLinH step-size skip) and rxData.cpp (thread decision) so
+// the classification cannot drift between translation units.
+
+// True for the automatic-differentiation Jacobian methods: forward-mode fvar
+// (3/30), reverse-mode (31) and the auto default (100).  The finite-difference
+// methods (1,2,4,5,6,7,10,20,40,50) return false.  linCmtB reads ind->linH only
+// on the finite-difference methods, so these are exactly the AD methods that do
+// not need finite-difference step-size estimation.
+static inline int linCmtSensIsAD(int sensType) {
+  switch (sensType) {
+  case 1: case 2: case 4: case 5: case 6: case 7:
+  case 10: case 20: case 40: case 50:
+    return 0;
+  default:
+    return 1;
+  }
+}
+
+// True for the forward-mode AD Jacobian methods that are safe to solve across
+// threads.  3 and 30 are explicit forward-mode fvar -- stack-local with no
+// shared AD state.  100 (auto) is included because it always resolves to the
+// forward-mode path: setupLinH() remaps 100 -> 3 and it is the only runtime
+// mutation of rx->sensType (nothing resolves auto to reverse-mode or a
+// finite-difference method).  Reverse-mode AD (31) is deliberately excluded --
+// it uses Stan's shared ChainableStack.  The finite-difference methods are
+// excluded -- they share a first-subject scaling/step-size setup that is not
+// per-thread.
+static inline int linCmtSensForwardAdThreadSafe(int sensType) {
+  return (sensType == 3 || sensType == 30 || sensType == 100);
+}
+
+#endif // __LINCMTSENSTYPE_H__

src/par_solve.cpp

@@ -17,6 +17,7 @@
 #include "../inst/include/rxode2parseGetTime.h"
 #include "../inst/include/rxode2EventTranslate.h"
 #include "linCmtDiffConstant.h"
+#include "linCmtSensType.h"
 
 #define SORT gfx::timsort
 
@@ -6792,16 +6793,23 @@ void setupLinH(rx_solve *rx, int solveid,
   rx_solving_options_ind *ind = &(rx->subjects[neq[1]]);
   double *hh = ind->linH;
   if (ind->linCmtHparIndex == -3) return; // already setup
+  if (linCmtSensIsAD(rx->sensType)) {
+    // AD Jacobians (forward-mode fvar 3/30, reverse-mode 31) are exact and
+    // never read ind->linH, so skip the finite-difference step-size estimation
+    // (shi21ForwardH/gillForwardH) -- it would run a base solve plus several
+    // probe solves per parameter that the AD path discards.
+    std::fill_n(ind->linH, 7, rx->sensH);
+    ind->linCmtH = NA_REAL;
+    ind->linCmtHparIndex = -3;
+    return;
+  }
   switch (rx->sensType) {
   case 1: // forward; shi difference
     shi21ForwardH(rx, op, solveid, neq, dydt, u_inis);
     break;
   case 2: // central; shi
     shi21CentralH(rx, op, solveid, neq, dydt, u_inis);
     break;
-  case 3: // 3pt forward; shi
-    shi21ForwardH(rx, op, solveid, neq, dydt, u_inis);
-    break;
   case 4: // 5-point endpoint difference; shi
     shi21CentralH(rx, op, solveid, neq, dydt, u_inis);
     for (int i = 0; i < 7; i++) {

src/print_node.h

@@ -193,8 +193,9 @@ static inline int nodeFunLinCmtB(char *value) {
     sAppendN(&sbt,"linCmtB", 7);
     tb.linCmt=2;
 
-    // right now linCmtB isn't thread safe,
-    // the Jacobian can cause a null free in stan math currently.
+    // Mark the model as linCmtB-bearing.  The solve-time thread decision
+    // (rxData.cpp) threads the forward-mode AD Jacobian path and keeps the
+    // reverse-mode AD (stan shared arena) and finite-difference paths serial.
     if (tb.thread == threadSafe) {
       tb.thread = notThreadLinCmtB;
     }

src/rxData.cpp

@@ -57,6 +57,7 @@ extern "C" void ensureRworkPool(int nCores, int lrw, int liw);
 #include "rxThreadData.h"
 
 #include "threadSafeConstants.h"
+#include "linCmtSensType.h"
 //#include "seed.h"
 
 SEXP rxSaveState_();         // defined in rxSerialize.cpp
@@ -5937,6 +5938,13 @@ SEXP rxSolve_(const RObject &obj, const List &rxControl,
     } else {
       op->cores = asInt(rxControl[Rxc_cores], "cores");
       int thread = INTEGER(rxSolveDat->mv[RxMv_flags])[RxMvFlag_thread];
+      // linCmtB is thread safe on the forward-mode AD Jacobian path: each thread
+      // evaluates its own subject on its own __linCmtB[rx_get_thread()] slot with
+      // stack-local fvar and no shared AD arena.  linCmtSensForwardAdThreadSafe()
+      // (linCmtSensType.h) is the shared classifier -- it excludes reverse-mode
+      // AD (31, Stan's shared ChainableStack) and the finite-difference paths
+      // (first-subject scaling/step setup), which keep the single-core guard.
+      int linCmtBThreadSafe = linCmtSensForwardAdThreadSafe(rx->sensType);
       if (op->cores == 0) {
         switch (thread) {
         case threadSafeRepNumThread:
@@ -5957,8 +5965,13 @@ SEXP rxSolve_(const RObject &obj, const List &rxControl,
           rxSolveDat->throttle = false;
           break;
         case notThreadLinCmtB:
-          op->cores = 1;
-          rxSolveDat->throttle = false;
+          if (linCmtBThreadSafe) {
+            op->cores = getRxThreads(INT_MAX, false);
+            rxSolveDat->throttle = true;
+          } else {
+            op->cores = 1;
+            rxSolveDat->throttle = false;
+          }
           break;
         }
       } else {
@@ -5979,8 +5992,13 @@ SEXP rxSolve_(const RObject &obj, const List &rxControl,
           rxSolveDat->throttle = false;
           break;
         case notThreadLinCmtB:
-          op->cores = 1;
-          rxSolveDat->throttle = false;
+          if (linCmtBThreadSafe) {
+            // Thread safe (forward-mode AD); keep the user-requested core count.
+            rxSolveDat->throttle = true;
+          } else {
+            op->cores = 1;
+            rxSolveDat->throttle = false;
+          }
           break;
         }
       }

tests/testthat/test-lincmt-solve-sens.R

@@ -117,4 +117,35 @@ rxTest({
       expect_equal(outOde[[col]], outLin[[col]], tolerance=1e-5)
     })
   }
+
+  test_that("linCmtB forward-AD solves are thread safe (parallel == serial)", {
+    skip_on_cran()
+    # Without >= 2 available threads the cores = 2 request cannot exercise the
+    # parallel path, so the comparison would pass vacuously; skip instead.
+    # rxode2 exposes no per-solve "cores actually used" telemetry, so gate on
+    # thread availability (this model's linCmtB carries the forward-AD flag that
+    # the cores = 2 request now honors -- see rxData.cpp thread decision).
+    skip_if_not(rxode2::getRxThreads() >= 2L)
+    .nsub <- 60
+    set.seed(1)
+    .params <- data.frame(
+      id = seq_len(.nsub),
+      "THETA[1]" = log(4), "THETA[2]" = log(70), "THETA[3]" = log(1),
+      "THETA[4]" = 0.1,
+      "ETA[1]" = rnorm(.nsub, 0, 0.3),
+      "ETA[2]" = rnorm(.nsub, 0, 0.3),
+      "ETA[3]" = rnorm(.nsub, 0, 0.3),
+      check.names = FALSE)
+    .ev <- et(amt = 100, ii = 12, until = 48) |>
+      et(seq(0, 48, by = 1)) |>
+      et(id = seq_len(.nsub))
+    .s1 <- rxSolve(lin, params = .params, events = .ev, cores = 1,
+                   returnType = "data.frame")
+    .s2 <- rxSolve(lin, params = .params, events = .ev, cores = 2,
+                   returnType = "data.frame")
+    .cols <- grep("sens_rx_pred.*ETA|^rx_pred_$", names(.s1), value = TRUE)
+    for (.col in .cols) {
+      expect_identical(.s1[[.col]], .s2[[.col]])
+    }
+  })
 })