diff options
Diffstat (limited to 'third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Tridiag.hs')
-rw-r--r-- | third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Tridiag.hs | 16 |
1 files changed, 16 insertions, 0 deletions
diff --git a/third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Tridiag.hs b/third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Tridiag.hs new file mode 100644 index 000000000000..7668deace132 --- /dev/null +++ b/third_party/bazel/rules_haskell/examples/vector/benchmarks/Algo/Tridiag.hs @@ -0,0 +1,16 @@ +module Algo.Tridiag ( tridiag ) where + +import Data.Vector.Unboxed as V + +tridiag :: (Vector Double, Vector Double, Vector Double, Vector Double) + -> Vector Double +{-# NOINLINE tridiag #-} +tridiag (as,bs,cs,ds) = V.prescanr' (\(c,d) x' -> d - c*x') 0 + $ V.prescanl' modify (0,0) + $ V.zip (V.zip as bs) (V.zip cs ds) + where + modify (c',d') ((a,b),(c,d)) = + let id = 1 / (b - c'*a) + in + id `seq` (c*id, (d-d'*a)*id) + |