NMODL is slow in translating complicated models

[JCGoran]

Overview of the issue

Not a bug per se, but warrants a discussion. For this particular mod file https://github.com/ModelDBRepository/138382/blob/3548e4565608f7f61ea51ead2432e3f40bdf279e/cdp3.mod, NMODL takes over 7 minutes to just generate the code:

$ time nmodl --neuron cdp3.mod
real    7m11.235s
user    7m9.113s
sys     0m0.507s

My guess is that this is due to the KINETIC block being rather complex, and the resulting C++ file has almost 40000 lines.

On the other hand, nocmodl translates the file much faster:

$ time nocmodl cdp3.mod
real    0m0.089s
user    0m0.071s
sys     0m0.018s

Note that, for this particular model, compilation actually fails when using the NMODL translator, with the error (this is technically unrelated to the slowness, but it happens to the same model so I'm mentioning it for completeness):

/include/Eigen/src/Core/DenseStorage.h: In instantiation of ‘void Eigen::internal::check_static_allocation_size() [with T = double; int Size = 35721]’:
/include/Eigen/src/Core/DenseStorage.h:51:41:   required from ‘Eigen::internal::plain_array<T, Size, MatrixOrArrayOptions, Alignment>::plain_array() [with T = double; int Size = 35721; int MatrixOrArrayOptions = 0; int Alignment = 0]
’
/include/Eigen/src/Core/DenseStorage.h:211:38:   required from ‘Eigen::DenseStorage<T, Size, _Rows, _Cols, _Options>::DenseStorage() [with T = double; int Size = 35721; int _Rows = 189; int _Cols = 189; int _Options = 0]’
/include/Eigen/src/Core/PlainObjectBase.h:476:45:   required from ‘Eigen::PlainObjectBase<Derived>::PlainObjectBase() [with Derived = Eigen::Matrix<double, 189, 189>]’
/include/Eigen/src/Core/Matrix.h:259:21:   required from ‘Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows, _MaxCols>::Matrix() [with _Scalar = double; int _Rows = 189; int _Cols = 189; int _Options = 0; int _MaxRows = 189; int _MaxCols = 189]’
138382/AnwarEtAl2010/x86_64/cdp3.cpp:303:33:   required from ‘int nmodl::newton::newton_solver(Eigen::Matrix<double, N, 1>&, FUNC, double, int) [with int N = 189; FUNC = neuron::functor_cdp3_0]’
138382/AnwarEtAl2010/x86_64/cdp3.cpp:39091:65:   required from here
/include/Eigen/src/Core/DenseStorage.h:33:40: error: static assertion failed: OBJECT_ALLOCATED_ON_STACK_IS_TOO_BIG
   33 |   EIGEN_STATIC_ASSERT(Size * sizeof(T) <= EIGEN_STACK_ALLOCATION_LIMIT, OBJECT_ALLOCATED_ON_STACK_IS_TOO_BIG);
/include/Eigen/src/Core/DenseStorage.h:33:3: note: in expansion of macro ‘EIGEN_STATIC_ASSERT’
   33 |   EIGEN_STATIC_ASSERT(Size * sizeof(T) <= EIGEN_STACK_ALLOCATION_LIMIT, OBJECT_ALLOCATED_ON_STACK_IS_TOO_BIG);
      |   ^~~~~~~~~~~~~~~~~~~
/include/Eigen/src/Core/DenseStorage.h:33:40: note: the comparison reduces to ‘(285768 <= 131072)’
   33 |   EIGEN_STATIC_ASSERT(Size * sizeof(T) <= EIGEN_STACK_ALLOCATION_LIMIT, OBJECT_ALLOCATED_ON_STACK_IS_TOO_BIG);
/include/Eigen/src/Core/DenseStorage.h:33:3: note: in expansion of macro ‘EIGEN_STATIC_ASSERT’

The key line being:

note: the comparison reduces to ‘(285768 <= 131072)’

Since sizeof(T) is 8 on my machine, it seems the default limit for the an N x N matrix is N = 128. We can get around this issue by adding:

#define EIGEN_STACK_ALLOCATION_LIMIT 0

somewhere (I added it by hand to the file, but I'm sure we can add it as a flag to the nrnivmodl makefile).

Expected result/behavior

NMODL should not take so long to just translate the whole file. My suspicion is that the sympy solver, being written in pure Python, is rather slow for large systems.

[JCGoran]

The following file showcases the problem (extracted from the above model) using just the solver, no NMODL components necessary (file is a bit large to paste in full):

slow.py

According to sympy/sympy#27886 (and any issues linked within), as well as https://stackoverflow.com/a/78321840, it seems the subs function is rather slow:

By replacing all of its calls with xreplace in solve_non_lin_system, the time drops down to 170 seconds, which is a 5x speed-up. The call graph points to the new bottlenecks:

The demangle_protected_identifiers is because I ran the profiler on this branch: https://github.com/neuronsimulator/nrn/tree/jelic/temp_attempt_modeldb

Though I am surprised that the regexing now takes 1/3 of the total runtime. In any case, removing the calls to simplify we can get it down to sub-100 seconds:

Can we speed-up the call to jacobian? I am not so sure, but a 10x speed-up is pretty good so far. I still need to validate the results though. In addition, I don't know how we could speed-up the demangling of identifiers, but if we can get rid of it (or at least make it negligible), we could get a 20x speed-up compared to the current setup. Still, sympy's jacobian would remain the bottleneck.

EDIT: to note, I tried running this with both SymPy 1.13.3 and latest master (commit a8299470d018f12be72f3d043ce19ff2e06fb5ae), but found no difference in speed.

---

For reference, I generated the profiles using the builtin cProfile module:

python -m cProfile -o <output_file> slow.py

which can then be viewed in the browser using snakeviz (installable with pip install snakeviz):

snakeviz <output_file>

[JCGoran]

To continue discussion of the Jacobian, I've opened up sympy/sympy#28445 and we get a 10-20% performance increase for the call to free_symbols, which was a rather large part of the function.py:1232(__new__) call. This brings the total runtime to 92 seconds:

We can almost entirely get rid of the performance bottleneck of the regex matching by using the following diff:

diff --git a/share/lib/python/neuron/nmodl/ode.py b/share/lib/python/neuron/nmodl/ode.py
index 8adfc7c51..4b42f0912 100644
--- a/share/lib/python/neuron/nmodl/ode.py
+++ b/share/lib/python/neuron/nmodl/ode.py
@@ -119,12 +119,18 @@ def demangle_protected_identifiers(eqs: Iterable[str]) -> list[str]:
     --------
     mangle_protected_identifiers : The inverse of this function.
     """
-    for var in forbidden_var:
-        r = re.compile(rf"\b{MANGLE_PREFIX}{var}\b")
-        f = var
-        eqs = [re.sub(r, f, x) for x in eqs]
+    # Build single regex with all variables
+    escaped_vars = "|".join(re.escape(f"{MANGLE_PREFIX}{var}") for var in forbidden_var)
+    pattern = re.compile(rf"\b({escaped_vars})\b")

-    return eqs
+    # Create replacement mapping
+    replacements = {f"{MANGLE_PREFIX}{var}": var for var in forbidden_var}
+
+    # Single pass replacement
+    def replace_func(match):
+        return replacements[match.group(0)]
+
+    return [pattern.sub(replace_func, x) for x in eqs]


 def _get_custom_functions(fcts):

What this does is it compiles a single regex pattern that looks like \b(_mangled_nmodl_identifierif|_mangled_nmodl_identifierelse|...)\b, and uses that for matching instead of compiling a pattern for each individual "forbidden" variable (if, else, def, etc.).

This brings the runtime down to 40 seconds:

which yields a 23x speed-up compared to the current implementation.

Note that the rest of the sympy stuff seems to involve Python builtins (mostly getattr, set.update, and isinstance):

which I am assuming are already rather optimized.