Generalized Loop-fusion / mapReduce with functional syntax

[mratsim]

Related to #31

It would be nice to do

loopFusion:
  zip(t1, t2).map(f).fold(init, g) # without temporaries.

See https://github.com/alehander42/zero-functional

[mratsim]

From https://github.com/mratsim/Arraymancer/blob/a59a982e8a424326a2899d4e33eb0f9ff4b9674b/benchmarks/euler_tensor_optim.nim

The usage should transform this code:

proc f(T: Tensor[float]): Tensor[float] =
  a_dz2 * (T[_, 0..^3] - 2.0 * T[_, 1..^2] + T[_, 2..^1])

To this

proc f(T: Tensor[float]): Tensor[float] =
  loopFusion:
    a_dz2 * (T[_, 0..^3] - 2.0 * T[_, 1..^2] + T[_, 2..^1])

[mratsim]

This can follow the syntax of the new loop-fusion package https://github.com/numforge/loop-fusion

import loopfusion

block: # Simple
  let a = @[1, 2, 3]
  let b = @[11, 12, 13]
  let c = @[10, 10, 10]

  forEach x in a, y in b, z in c:
    echo (x + y) * z

  # 120
  # 140
  # 160

block: # With index
  let a = @[1, 2, 3]
  let b = @[11, 12, 13]
  let c = @[10, 10, 10]
  var d: seq[int] = @[]

  forEach i, x in a, y in b, z in c:
    d.add i + x + y + z

  doAssert d == @[22, 25, 28]

block: # With mutation
  var a = @[1, 2, 3]
  let b = @[11, 12, 13]
  let c = @[10, 10, 10]

  forEach x in var a, y in b, z in c:
    x += y * z

  doAssert a == @[111, 122, 133]

block: # With mutation, index and multiple statements
  var a = @[1, 2, 3]
  let b = @[11, 12, 13]
  let c = @[10, 10, 10]

  forEach i, x in var a, y in b, z in c:
    let tmp = i * (y - z)
    x += tmp

  doAssert a == @[1, 4, 9]

block: # With iteration on seq of different types
  let a = @[1, 2, 3]
  let b = @[false, true, true]

  forEach integer in a, boolean in b:
    if boolean:
      echo integer

block: # With an expression
  let a = @[1, 2, 3]
  let b = @[4, 5, 6]


  let c = forEach(x in a, y in b):
    x + y

  doAssert c == @[5, 7, 9]


block: # With arrays + seq, mutation, index and multiple statements
  var a = [1, 2, 3]
  let b = [11, 12, 13]
  let c = @[10, 10, 10]

  forEach i, x in var a, y in b, z in c:
    let tmp = i * (y - z)
    x += tmp

  doAssert a == [1, 4, 9]

Even though you can do reduction with forEach, a specific forEachReduceOMP is needed for OpenMP parallel updates to deal with multithreading specificities.

[mratsim]

Unfortunately, this bug prevents adapting loopfusion to custom container types :/ nim-lang/Nim#7737

[mratsim]

Implemented in laser:

• https://github.com/numforge/laser/blob/e9bbb1055517c6e76649b67737d7b8807af22eb2/examples/ex05_tensor_parallel_reduction.nim#L19-L20
• https://github.com/numforge/laser/blob/e9bbb1055517c6e76649b67737d7b8807af22eb2/benchmarks/loop_iteration/iter_bench_prod.nim#L89-L90