Существует ли эффективный способ создания итератора волнового фронта? (Не связано с физикой.)

Question

// https://devblogs.nvidia.com/parallelforall/optimizing-recurrent-neural-networks-cudnn-5/ 
// The last optimization from the above post. This is a sorting based implementation. 
// I can't think of anything more elegant that a state machine at the moment 

let wavefront_order = 
    [| 
    for x=0 to 49 do 
     for y=0 to 9 do 
      for z=0 to 9 do 
       yield (x,y,z) 
    |] 
    |> Array.sortBy (fun (x,y,z) -> x+y+z) 

Фактически, работа с повторяющимся NN без сортировки по расстоянию от источника будет работать, но будет не так эффективна. Используя несколько потоков, планировщик Cuda будет засорен, если я не буду запускать ядра в порядке волнового фронта.

Что мне нужно выше для NN, или пример, находящийся в F #, является чисто контекстом.

Я знаю, как это сделать для двух измерений, но для 3D и более, все, что я имею в виду, немного хаки. Мне было интересно, есть ли какой-то изящный способ проведения итерации волнового фронта, аналогичный описанному выше методу сортировки для произвольных размеров размеров и количества измерений?

Answer 1

let x_len = 49 
let y_len = 9 
let z_len = 9 

let presort = 
    [| 
    for sum=0 to x_len+y_len+z_len do 
     for x=0 to min sum x_len do 
      for y=0 to min (sum-x) y_len do 
       let z = sum-x-y 
       if z <= z_len then yield (x,y,z) 
       |] 

Это прекрасно работает, это именно то, что я хотел. Спасибо, Джон Палмер.