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 0000000000..7668deace1
--- /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)
+