Eliminate allocs in vector-based TimeDependentSums (#171)

* Eliminate allocs in vector-based tdops * avoid allocs with bigger tuples
qojulia · Jul 9, 2024 · b7edca6 · b7edca6 · amilsted · Jul 9, 2024
1 parent 7c112ac
commit b7edca6
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "QuantumOpticsBase"
 uuid = "4f57444f-1401-5e15-980d-4471b28d5678"
-version = "0.5.0"
+version = "0.5.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/time_dependent_operator.jl b/src/time_dependent_operator.jl
@@ -160,7 +160,8 @@ function set_time!(o::TimeDependentSum, t::Number)
         o.current_time = t
         update_static_coefficients!(static_operator(o), coefficients(o), t)
     end
-    set_time!.(suboperators(o), t)
+    # foreach is type-stable for tuples and concrete-typed vectors
+    foreach(o -> set_time!(o, t), suboperators(o))
     o
 end
 
@@ -291,13 +292,15 @@ end
 
 @inline eval_coefficients(coeffs::Tuple, t::Number) = map(c->eval_coefficient(c, t), coeffs)
 
-# This is the performance-critical implementation.
-# To avoid allocations in most cases, we model this on map(f, t::Tuple).
-@inline eval_coefficients(::Type{T}, coeffs::Tuple{Any,}, t::Number) where T          = (T(eval_coefficient(coeffs[1], t)),)
-@inline eval_coefficients(::Type{T}, coeffs::Tuple{Any, Any}, t::Number) where T      = (T(eval_coefficient(coeffs[1], t)), T(eval_coefficient(coeffs[2], t)))
-@inline eval_coefficients(::Type{T}, coeffs::Tuple{Any, Any, Any}, t::Number) where T = (T(eval_coefficient(coeffs[1], t)), T(eval_coefficient(coeffs[2], t)), T(eval_coefficient(coeffs[3], t)))
-@inline eval_coefficients(::Type{T}, coeffs::Tuple, t::Number) where T                = (T(eval_coefficient(coeffs[1], t)), eval_coefficients(T, Base.tail(coeffs), t)...)
+# This is the performance-critical implementation. map(f, ::Tuple) avoids allocs in most cases
+@inline eval_coefficients(::Type{T}, coeffs::Tuple, t::Number) where T = map(c->T(eval_coefficient(c, t)), coeffs)
 
+# Now just using map here instead. Maybe restore this in case of regressions.
+## To avoid allocations in most cases, we model this on map(f, t::Tuple).
+#@inline eval_coefficients(::Type{T}, coeffs::Tuple{Any,}, t::Number) where T          = (T(eval_coefficient(coeffs[1], t)),)
+#@inline eval_coefficients(::Type{T}, coeffs::Tuple{Any, Any}, t::Number) where T      = (T(eval_coefficient(coeffs[1], t)), T(eval_coefficient(coeffs[2], t)))
+#@inline eval_coefficients(::Type{T}, coeffs::Tuple{Any, Any, Any}, t::Number) where T = (T(eval_coefficient(coeffs[1], t)), T(eval_coefficient(coeffs[2], t)), T(eval_coefficient(coeffs[3], t)))
+#@inline eval_coefficients(::Type{T}, coeffs::Tuple, t::Number) where T                = (T(eval_coefficient(coeffs[1], t)), eval_coefficients(T, Base.tail(coeffs), t)...)
 
 _timeshift_coeff(coeff, t0) = (@inline shifted_coeff(t) = coeff(t-t0))
 _timeshift_coeff(coeff::Number, _) = coeff