Integrate missing streaming

src/serinv/algs/pobts.py

@@ -4,6 +4,7 @@
 from serinv import (
     ArrayLike,
     _get_module_from_array,
+    _get_module_from_str,
 )
 
 
@@ -41,8 +42,11 @@
     else:
         # Natural arrowhead
         if device_streaming:
-            raise NotImplementedError(
-                "Streaming is not implemented for the natural arrowhead."
+            _pobts_streaming(
+                L_diagonal_blocks,
+                L_lower_diagonal_blocks,
+                B,
+                trans,
             )
         else:
             _pobts(
@@ -163,3 +167,222 @@
             )
     else:
         raise ValueError(f"Invalid transpose argument: {trans}.")
+
+
+def _pobts_streaming(
+    L_diagonal_blocks: ArrayLike,
+    L_lower_diagonal_blocks: ArrayLike,
+    B: ArrayLike,
+    trans: str,
+):
+    arr_module, _ = _get_module_from_array(arr=L_diagonal_blocks)
+    if arr_module.__name__ != "numpy":
+        raise TypeError(
+            "Host<->Device streaming only works when host-arrays are given."
+        )
+
+    cp, cu_la = _get_module_from_str(module_str="cupy")
+
+    # Vars
+    diag_blocksize = L_diagonal_blocks.shape[1]
+    n_diag_blocks = L_diagonal_blocks.shape[0]
+
+    # Streams
+    compute_stream = cp.cuda.Stream(non_blocking=True)
+    h2d_stream = cp.cuda.Stream(non_blocking=True)
+    d2h_stream = cp.cuda.Stream(non_blocking=True)
+
+    # Device Buffers
+    # B Buffers
+    B_shape = B[0:diag_blocksize]
+    B_d = cp.empty((2, *B_shape.shape), dtype=B_shape.dtype)
+    B_previous_d = cp.empty((2, *B_shape.shape), dtype=B_shape.dtype)
+    del B_shape
+
+    # L Buffers
+    L_diagonal_blocks_d = cp.empty(
+        (2, *L_diagonal_blocks.shape[1:]), dtype=L_diagonal_blocks.dtype
+    )
+    L_lower_diagonal_blocks_d = cp.empty(
+        (2, *L_diagonal_blocks.shape[1:]), dtype=L_diagonal_blocks.dtype
+    )
+
+    # Events
+    compute_B_events = [cp.cuda.Event(), cp.cuda.Event()]
+    h2d_events = [cp.cuda.Event(), cp.cuda.Event()]
+    d2h_events = [cp.cuda.Event(), cp.cuda.Event()]
+
+    if trans == "N":
+        # ----- Forward substitution -----
+
+        # --- H2D: transfers ---
+        B_d[0].set(arr=B[:diag_blocksize], stream=h2d_stream)
+        L_diagonal_blocks_d[0].set(arr=L_diagonal_blocks[0], stream=h2d_stream)
+
+        h2d_events[1].record(stream=h2d_stream)
+
+        if n_diag_blocks > 1:
+            B_d[1].set(arr=B[diag_blocksize : (2 * diag_blocksize)], stream=h2d_stream)
+            L_diagonal_blocks_d[1].set(arr=L_diagonal_blocks[1], stream=h2d_stream)
+            L_lower_diagonal_blocks_d[1].set(
+                arr=L_lower_diagonal_blocks[0], stream=h2d_stream
+            )
+
+            h2d_events[0].record(stream=h2d_stream)
+
+        with compute_stream:
+            # Solve first B block
+            compute_stream.wait_event(h2d_events[1])
+
+            B_previous_d[0] = cu_la.solve_triangular(
+                L_diagonal_blocks_d[0],
+                B_d[0],
+                lower=True,
+            )
+
+            compute_B_events[0].record(stream=compute_stream)
+
+        for i in range(1, n_diag_blocks):
+
+            if i + 1 < n_diag_blocks:
+                # Pass next blocks
+                h2d_stream.wait_event(compute_B_events[(i + 1) % 2])
+
+                B_d[(i + 1) % 2].set(
+                    arr=B[(i + 1) * diag_blocksize : (i + 2) * diag_blocksize],
+                    stream=h2d_stream,
+                )
+                L_diagonal_blocks_d[(i + 1) % 2].set(
+                    arr=L_diagonal_blocks[i + 1], stream=h2d_stream
+                )
+                L_lower_diagonal_blocks_d[(i + 1) % 2].set(
+                    arr=L_lower_diagonal_blocks[i], stream=h2d_stream
+                )
+
+                h2d_events[i % 2].record(stream=h2d_stream)
+
+            with compute_stream:
+                # X_{i} = L_{i,i}^{-T} (Y_{i} - L_{i+1,i}^{T} X_{i+1}) - L_{ndb+1,i}^T X_{ndb+1}
+                compute_stream.wait_event(h2d_events[(i + 1) % 2])
+                compute_stream.wait_event(d2h_events[(i + 1) % 2])
+
+                B_previous_d[i % 2] = cu_la.solve_triangular(
+                    L_diagonal_blocks_d[i % 2],
+                    B_d[i % 2]
+                    - L_lower_diagonal_blocks_d[i % 2] @ B_previous_d[(i + 1) % 2],
+                    lower=True,
+                )
+
+                compute_B_events[i % 2].record(compute_stream)
+
+            # Pass previous B block back
+            d2h_stream.wait_event(compute_B_events[(i - 1) % 2])
+
+            B_previous_d[(i + 1) % 2].get(
+                out=B[(i - 1) * diag_blocksize : i * diag_blocksize],
+                stream=d2h_stream,
+                blocking=False,
+            )
+
+            d2h_events[i % 2].record(stream=d2h_stream)
+
+        # Pass last B block back
+        d2h_stream.wait_event(compute_B_events[(n_diag_blocks + 1) % 2])
+
+        B_previous_d[(n_diag_blocks + 1) % 2].get(
+            out=B[-diag_blocksize:], stream=d2h_stream, blocking=False
+        )
+
+    elif trans == "T" or trans == "C":
+        # ----- Backward substitution -----
+
+        # --- H2D: transfers ---
+        B_d[(n_diag_blocks - 1) % 2].set(arr=B[-diag_blocksize:], stream=h2d_stream)
+        L_diagonal_blocks_d[(n_diag_blocks - 1) % 2].set(
+            arr=L_diagonal_blocks[-1], stream=h2d_stream
+        )
+
+        h2d_events[(n_diag_blocks - 1) % 2].record(stream=h2d_stream)
+
+        if n_diag_blocks > 1:
+
+            B_d[n_diag_blocks % 2].set(
+                arr=B[-(2 * diag_blocksize) : -diag_blocksize], stream=h2d_stream
+            )
+            L_diagonal_blocks_d[n_diag_blocks % 2].set(
+                arr=L_diagonal_blocks[-2], stream=h2d_stream
+            )
+            L_lower_diagonal_blocks_d[n_diag_blocks % 2].set(
+                arr=L_lower_diagonal_blocks[-1], stream=h2d_stream
+            )
+
+            h2d_events[(n_diag_blocks - 1) % 2].record(stream=h2d_stream)
+
+        with compute_stream:
+            # X_{ndb} = L_{ndb,ndb}^{-T} (Y_{ndb} - L_{ndb+1,ndb}^{T} X_{ndb+1})
+            compute_stream.wait_event(h2d_events[(n_diag_blocks - 1) % 2])
+
+            B_previous_d[(n_diag_blocks - 1) % 2] = cu_la.solve_triangular(
+                L_diagonal_blocks_d[(n_diag_blocks - 1) % 2],
+                B_d[(n_diag_blocks - 1) % 2],
+                lower=True,
+                trans="C",
+            )
+
+            compute_B_events[(n_diag_blocks - 1) % 2].record(stream=compute_stream)
+
+        for i in range(n_diag_blocks - 2, -1, -1):
+
+            if i > 0:
+                # pass next blocks
+                h2d_stream.wait_event(compute_B_events[(i - 1) % 2])
+
+                B_d[(i - 1) % 2].set(
+                    arr=B[(i - 1) * diag_blocksize : i * diag_blocksize],
+                    stream=h2d_stream,
+                )
+                L_diagonal_blocks_d[(i - 1) % 2].set(
+                    arr=L_diagonal_blocks[i - 1], stream=h2d_stream
+                )
+                L_lower_diagonal_blocks_d[(i - 1) % 2].set(
+                    arr=L_lower_diagonal_blocks[i - 1], stream=h2d_stream
+                )
+
+                h2d_events[i % 2].record(stream=h2d_stream)
+
+            with compute_stream:
+                # X_{i} = L_{i,i}^{-T} (Y_{i} - L_{i+1,i}^{T} X_{i+1}) - L_{ndb+1,i}^T X_{ndb+1}
+                compute_stream.wait_event(h2d_events[(i - 1) % 2])
+                compute_stream.wait_event(d2h_events[(i - 1) % 2])
+
+                B_previous_d[i % 2] = cu_la.solve_triangular(
+                    L_diagonal_blocks_d[i % 2],
+                    B_d[i % 2]
+                    - L_lower_diagonal_blocks_d[i % 2].conj().T
+                    @ B_previous_d[(i - 1) % 2],
+                    lower=True,
+                    trans="C",
+                )
+
+                compute_B_events[i % 2].record(compute_stream)
+
+            # Pass previous B block back
+            d2h_stream.wait_event(compute_B_events[(i - 1) % 2])
+
+            B_previous_d[(i - 1) % 2].get(
+                out=B[(i + 1) * diag_blocksize : (i + 2) * diag_blocksize],
+                stream=d2h_stream,
+                blocking=False,
+            )
+
+            d2h_events[i % 2].record(stream=d2h_stream)
+
+        # Pass last B block back
+        d2h_stream.wait_event(compute_B_events[0])
+
+        B_previous_d[0].get(out=B[:diag_blocksize], stream=d2h_stream, blocking=False)
+
+    else:
+        raise ValueError(f"Invalid transpose argument: {trans}.")
+
+    cp.cuda.Device().synchronize()

src/serinv/wrappers/ddbtars.py

@@ -15,7 +15,6 @@
     import cupyx as cpx
     import cupy as cp
 
-
 def allocate_ddbtars(
     A_diagonal_blocks: ArrayLike,
     A_lower_diagonal_blocks: ArrayLike,

src/serinv/wrappers/pddbtasci.py

@@ -44,7 +44,7 @@ def pddbtasci(
         The arrow tip block of the block tridiagonal with arrowhead matrix.
     comm : MPI.Comm
         The MPI communicator. Default is MPI.COMM_WORLD.
-        
+
     Keyword Arguments
     -----------------
     rhs : dict

src/serinv/wrappers/pddbtsc.py

@@ -192,6 +192,4 @@
         quadratic=quadratic,
     )
 
-    comm.Barrier()
-
-    return elapsed
+    comm.Barrier()

src/serinv/wrappers/pddbtsci.py

@@ -34,7 +34,7 @@ def pddbtsci(
         The upper diagonal blocks of the block tridiagonal matrix.
     comm : MPI.Comm
         The MPI communicator. Default is MPI.COMM_WORLD.
-        
+
     Keyword Arguments
     -----------------
     rhs : dict

tests/conftest.py

@@ -1,6 +1,5 @@
 # Copyright 2023-2025 ETH Zurich. All rights reserved.
 # Global pytest fixtures for the Serinv tests.
-
 import pytest
 
 from serinv import backend_flags
@@ -15,7 +14,6 @@
         ]
     )
 
-
 DTYPE = [
     pytest.param("float64", id="float64"),
     pytest.param("complex128", id="complex128"),

tests/tests_algs/regular/tests_bt/test_pobtf.py

@@ -3,6 +3,8 @@
 import numpy as np
 import pytest
 
+from ....conftest import ARRAY_TYPE
+
 from serinv import backend_flags, _get_module_from_array
 from ....testing_utils import bt_dense_to_arrays, dd_bt, symmetrize
 
@@ -11,6 +13,16 @@
 if backend_flags["cupy_avail"]:
     import cupyx as cpx
 
+    ARRAY_TYPE.extend(
+        [
+            pytest.param("streaming", id="streaming"),
+        ]
+    )
+
+    @pytest.fixture(params=ARRAY_TYPE, autouse=True)
+    def array_type(request: pytest.FixtureRequest) -> str:
+        return request.param
+
 
 @pytest.mark.mpi_skip()
 def test_pobtf(

tests/tests_algs/regular/tests_bt/test_pobts.py

@@ -3,11 +3,26 @@
 import numpy as np
 import pytest
 
-from serinv import _get_module_from_array
+from ....conftest import ARRAY_TYPE
+
+from serinv import backend_flags, _get_module_from_array
 from ....testing_utils import bt_dense_to_arrays, dd_bt, symmetrize, rhs
 
 from serinv.algs import pobtf, pobts
 
+if backend_flags["cupy_avail"]:
+    import cupyx as cpx
+
+    ARRAY_TYPE.extend(
+        [
+            pytest.param("streaming", id="streaming"),
+        ]
+    )
+
+    @pytest.fixture(params=ARRAY_TYPE, autouse=True)
+    def array_type(request: pytest.FixtureRequest) -> str:
+        return request.param
+
 
 @pytest.mark.mpi_skip()
 @pytest.mark.parametrize("n_rhs", [1, 2, 3])
@@ -18,6 +33,7 @@ def test_pobts(
     array_type: str,
     dtype: np.dtype,
 ):
+
     A = dd_bt(
         diagonal_blocksize,
         n_diag_blocks,
@@ -47,9 +63,22 @@ def test_pobts(
         _,
     ) = bt_dense_to_arrays(A, diagonal_blocksize, n_diag_blocks)
 
+    if backend_flags["cupy_avail"] and array_type == "streaming":
+        A_diagonal_blocks_pinned = cpx.zeros_like_pinned(A_diagonal_blocks)
+        A_diagonal_blocks_pinned[:, :, :] = A_diagonal_blocks[:, :, :]
+        A_lower_diagonal_blocks_pinned = cpx.zeros_like_pinned(A_lower_diagonal_blocks)
+        A_lower_diagonal_blocks_pinned[:, :, :] = A_lower_diagonal_blocks[:, :, :]
+        B_pinned = cpx.zeros_like_pinned(B)
+        B_pinned[:, :] = B[:, :]
+
+        A_diagonal_blocks = A_diagonal_blocks_pinned
+        A_lower_diagonal_blocks = A_lower_diagonal_blocks_pinned
+        B = B_pinned
+
     pobtf(
         A_diagonal_blocks,
         A_lower_diagonal_blocks,
+        device_streaming=True if array_type == "streaming" else False,
     )
 
     # Forward solve: Y=L^{-1}B
@@ -58,6 +87,7 @@ def test_pobts(
         A_lower_diagonal_blocks,
         B,
         trans="N",
+        device_streaming=True if array_type == "streaming" else False,
     )
 
     # Backward solve: X=L^{-T}Y
@@ -66,6 +96,7 @@ def test_pobts(
         A_lower_diagonal_blocks,
         B,
         trans="C",
+        device_streaming=True if array_type == "streaming" else False,
     )
 
     assert xp.allclose(B, X_ref)