3D relaxation in Futhark

gistfile1.txt

let hood 'a (xs: []a) : [](a,a,a) =
  zip3 (rotate (-1) xs) xs (rotate 1 xs)

let hood_2d 'a (xss: [][]a) : [][]((a,a,a),
                                   (a,a,a),
                                   (a,a,a)) =
  map hood xss -- [][](a, a, a)
  |> hood -- []([](a, a, a), [](a, a, a), [](a, a, a))
  |> map (\(a,b,c) -> zip3 a b c)

let hood_3d 'a (xsss: [][][]a) : [][][](((a,a,a),(a,a,a),(a,a,a)),
                                        ((a,a,a),(a,a,a),(a,a,a)),
                                        ((a,a,a),(a,a,a),(a,a,a))) =
  map hood_2d xsss -- [][][]((a, a, a), (a, a, a), (a, a, a))
  |> hood -- []([][]((a, a, a), (a, a, a), (a, a, a)),
          --    [][]((a, a, a), (a, a, a), (a, a, a)),
          --    [][]((a, a, a), (a, a, a), (a, a, a)))
  |> map (\(a,b,c) -> map3 zip3 a b c)

let relax (xsss: [][][]f32) : [][][]f32 =
  let on_point h =
    (h.1.1.1 + h.1.1.2 + h.1.1.3 +
     h.1.2.1 + h.1.2.2 + h.1.2.3 +
     h.1.3.1 + h.1.3.2 + h.1.3.3 +
     h.2.1.1 + h.2.1.2 + h.2.1.3 +
     h.2.2.1 + h.2.2.2 + h.2.2.3 +
     h.2.3.1 + h.2.3.2 + h.2.3.3 +
     h.3.1.1 + h.3.1.2 + h.3.1.3 +
     h.3.2.1 + h.3.2.2 + h.3.2.3 +
     h.3.3.1 + h.3.3.2 + h.3.3.3) / 27
  in xsss |> hood_3d |> map (map (map on_point))

let main = relax