thomas.cpp

Go to the documentation of this file.

#include "linalg/factorization/thomas.hpp"
#include "core/parallel/cuda_ops.hpp"
#include <stdexcept>
 
namespace num {
 
void thomas(const Vector& a, const Vector& b, const Vector& c,
            const Vector& d, Vector& x, Backend backend) {
    idx n = b.size();
    if (a.size() != n - 1 || c.size() != n - 1 || d.size() != n || x.size() != n)
        throw std::invalid_argument("Dimension mismatch in Thomas solver");
 
    if (backend == Backend::gpu) {
#ifdef NUMERICS_HAS_CUDA
        Vector ag = a; ag.to_gpu();
        Vector bg = b; bg.to_gpu();
        Vector cg = c; cg.to_gpu();
        Vector dg = d; dg.to_gpu();
        x = Vector(n);
        x.to_gpu();
        cuda::thomas_batched(ag.gpu_data(), bg.gpu_data(), cg.gpu_data(),
                             dg.gpu_data(), x.gpu_data(), n, 1);
        x.to_cpu();
        return;
#endif
    }
 
    Vector b_work = b;
    Vector d_work = d;
 
    for (idx i = 1; i < n; ++i) {
        real w = a[i - 1] / b_work[i - 1];
        b_work[i] -= w * c[i - 1];
        d_work[i] -= w * d_work[i - 1];
    }
 
    x[n - 1] = d_work[n - 1] / b_work[n - 1];
    for (idx i = n - 1; i > 0; --i)
        x[i - 1] = (d_work[i - 1] - c[i - 1] * x[i]) / b_work[i - 1];
}
 
} // namespace num