Investigate slowness and memory usage compared to Python

[mratsim]

Reported by @narimiran on IRC

import math
import arraymancer

const
  dz = 0.01
  z = 100
  spaceSteps = int(z / dz)
  timeSteps = 50000
  dt = 0.12 / timeSteps
  alpha = 2.0
  startingTemp = 30.0
  oscillations = 20.0

var
  Ts = startingTemp * ones[float](timeSteps, spaceSteps)


for j in 0 ..< timeSteps:
  Ts[j, 0] = startingTemp - oscillations * sin(2.0 * PI * j.float / timeSteps)

proc f(T: Tensor[float]): Tensor[float] =
  let d2T_dz2 = (T[_, 0..^3] - 2.0 * T[_, 1..^2] + T[_, 2..^1]) / dz^2
  return alpha * d2T_dz2

proc eulerSolve(Ts: var Tensor[float]) =
  for t in 0 ..< timeSteps-1:
    Ts[t+1, 1..^2] = Ts[t, 1..^2] + dt * f(Ts[t, _])
    Ts[t+1, ^1] = Ts[t+1, ^2]

Ts.eulerSolve()

[mratsim]

[mratsim]

Nim is spending 21 seconds on this, there is a big issue on assignments:

import math
import arraymancer

const
  dz = 0.01
  z = 100
  spaceSteps = int(z / dz)
  timeSteps = 50000
  dt = 0.12 / timeSteps
  alpha = 2.0
  startingTemp = 30.0
  oscillations = 20.0

var
  Ts = startingTemp * ones[float](timeSteps, spaceSteps)


for j in 0 ..< timeSteps:
  Ts[j, 0] = startingTemp - oscillations * sin(2.0 * PI * j.float / timeSteps)

[mratsim]

There might be an allocator bug because it's not even the loop that is slow:
This takes a whooping 17seconds half of it allocating memory

import math
import arraymancer

const
  dz = 0.01
  z = 100
  spaceSteps = int(z / dz)
  timeSteps = 50000
  dt = 0.12 / timeSteps
  alpha = 2.0
  startingTemp = 30.0
  oscillations = 20.0

var
  Ts = startingTemp * ones[float](timeSteps, spaceSteps)

[mratsim]

Raised upstream: nim-lang/Nim#7295

[narimiran]

@mratsim thank you for raising this and immediately investigating!!

[narimiran]

For a comparison, here is a version without Arraymancer, using Nim's arrays.

Notice timeSteps = 10000 - I cannot use 50000 with arrays, it raises an error during compilation.

import math

const
  dz = 0.01
  z = 100
  spaceSteps = int(z / dz)
  timeSteps = 10000
  dt = 0.12 / timeSteps
  alpha = 2.0
  startingTemp = 30.0
  oscillations = 20.0

var Ts: array[timeSteps, array[spaceSteps, float]]

for t in 0 ..< timeSteps:
  Ts[t][0] = startingTemp - oscillations * sin(2.0 * PI * t.float / timeSteps) 
  for s in 1 ..< spaceSteps:
    Ts[t][s] = startingTemp


proc f(T: array[spaceSteps, float]): array[spaceSteps-2, float] =
  for i in 0 .. spaceSteps-3:
    result[i] = alpha * (T[i] - 2.0 * T[i+1] + T[i+2]) / dz^2


proc eulerSolve(Ts: var array[timeSteps, array[spaceSteps, float]]) =
  for t in 0 ..< timeSteps-1:
    let
      k = f(Ts[t])
    for i in 1 .. spaceSteps-2:
      Ts[t+1][i] = Ts[t][i] + dt * k[i-1]
    Ts[t+1][spaceSteps-1] = Ts[t+1][spaceSteps-2] 

Ts.eulerSolve()

On my computer, this runs in 0.8 sec (Arraymancer needs 2.6 sec)

[narimiran]

Python version:

import numpy as np


dz = 0.01
z = 100
space_steps = int(z / dz)

time_steps = 50000
ts = np.linspace(0, 0.12, time_steps)
dt = 0.12 / time_steps
alpha = 2
starting_temp = 30
oscillations = 20

Ts = starting_temp * np.ones((time_steps, space_steps), dtype=np.float64)
Ts[:, 0] = starting_temp - oscillations * np.sin(2 * np.pi * ts / 12)


def f(T):
    d2T_dz2 = (T[:-2] - 2 * T[1:-1] + T[2:]) / dz**2
    return alpha * d2T_dz2


def euler_solve(Ts):
    for t in range(time_steps-1):
        Ts[t+1, 1:-1] = Ts[t, 1:-1] + dt * f(Ts[t])
        Ts[t+1, -1] = Ts[t+1, -2]
    return Ts


euler = euler_solve(Ts)

[narimiran]

On my computer, this runs in 0.8 sec (Arraymancer needs 2.6 sec)

With the updated Arraymancer, the timings for a problem with 10000 time steps are as follows:

nim+Arraymancer | 1.3 sec
nim with arrays | 0.7 sec
python+numpy    | 1.3 sec

[mratsim]

A couple updates on this:

• Template substitution will trigger proc parameters multiple times #193 has been merged, previously functions/computations may have been done twice if they were used as template inputs
• TRIOT: iterating over (variadic !) tensors without branching #164 I thought I could optimize items iteration further by implementing the TRIOT paper mentioned in the issue, (un?)fortunately the current iteration scheme is 20% faster than macros generated nested for loops

Next steps:

There shouldn't be any reason why Arraymancer has the same speed as Python here (there is no external libraries like BLAS involved I think).
I don't think there is a way to get as fast as arrays purely on the stack though.

So:

1. I'll investigate the remaining overhead and squash it (slicing? assignments with unnecessary copies?).