matrix.hpp

Go to the documentation of this file.

/// @file matrix.hpp
/// @brief Matrix operations
#pragma once
 
#include "core/vector.hpp"
#include "core/policy.hpp"
#include "core/small_matrix.hpp"
 
namespace num {
 
/// @brief Dense row-major matrix with optional GPU storage
class Matrix {
public:
    Matrix() : rows_(0), cols_(0), data_(nullptr) {}
    Matrix(idx rows, idx cols);
    Matrix(idx rows, idx cols, real val);
    ~Matrix();
 
    Matrix(const Matrix&);
    Matrix(Matrix&&) noexcept;
    Matrix& operator=(const Matrix&);
    Matrix& operator=(Matrix&&) noexcept;
 
    constexpr idx rows() const noexcept { return rows_; }
    constexpr idx cols() const noexcept { return cols_; }
    constexpr idx size() const noexcept { return rows_ * cols_; }
 
    real* data() { return data_.get(); }
    const real* data() const { return data_.get(); }
    real& operator()(idx i, idx j) { return data_[i * cols_ + j]; }
    real operator()(idx i, idx j) const { return data_[i * cols_ + j]; }
 
    void to_gpu();
    void to_cpu();
    real* gpu_data() { return d_data_; }
    const real* gpu_data() const { return d_data_; }
    bool on_gpu() const { return d_data_ != nullptr; }
 
private:
    idx rows_, cols_;
    std::unique_ptr<real[]> data_;
    real* d_data_ = nullptr;
};
 
/// @brief y = A * x
void matvec(const Matrix& A, const Vector& x, Vector& y, Backend b = default_backend);
 
/// @brief C = A * B
void matmul(const Matrix& A, const Matrix& B, Matrix& C, Backend b = default_backend);
 
/// @brief C = alpha*A + beta*B
void matadd(real alpha, const Matrix& A, real beta, const Matrix& B, Matrix& C,
            Backend b = default_backend);
 
// -- Named variants for benchmarking and pedagogy ----------------------------
// These expose the progression of optimisation techniques.  For production
// code, prefer matmul(A, B, C) or matmul(A, B, C, num::simd) etc.
 
/// @brief C = A * B  (cache-blocked)
///
/// Divides A, B, C into BLOCKxBLOCK tiles so the working set fits in L2 cache.
/// @param block_size  Tile edge length (default 64).
void matmul_blocked(const Matrix& A, const Matrix& B, Matrix& C,
                    idx block_size = 64);
 
/// @brief C = A * B  (register-blocked)
///
/// Extends cache blocking with a REGxREG register tile inside each cache tile.
/// @param block_size  Cache tile edge (default 64).
/// @param reg_size    Register tile edge (default 4).
void matmul_register_blocked(const Matrix& A, const Matrix& B, Matrix& C,
                              idx block_size = 64, idx reg_size = 4);
 
/// @brief C = A * B  (SIMD-accelerated)
///
/// Dispatches at compile time: AVX-256 + FMA on x86, NEON on AArch64,
/// falls back to matmul_blocked if neither is available.
void matmul_simd(const Matrix& A, const Matrix& B, Matrix& C,
                 idx block_size = 64);
 
/// @brief y = A * x  (SIMD-accelerated)
void matvec_simd(const Matrix& A, const Vector& x, Vector& y);
 
} // namespace num