TriangularMatrixVector.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009 Gael Guennebaud <gael.guennebaud@inria.fr>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0
 
#ifndef EIGEN_TRIANGULARMATRIXVECTOR_H
#define EIGEN_TRIANGULARMATRIXVECTOR_H
 
// IWYU pragma: private
#include "../InternalHeaderCheck.h"
 
namespace Eigen {
 
namespace internal {
 
template <typename Index, int Mode, typename LhsScalar, bool ConjLhs, typename RhsScalar, bool ConjRhs,
          int StorageOrder, int Version = Specialized>
struct triangular_matrix_vector_product;
 
template <typename Index, int Mode, typename LhsScalar, bool ConjLhs, typename RhsScalar, bool ConjRhs, int Version>
struct triangular_matrix_vector_product<Index, Mode, LhsScalar, ConjLhs, RhsScalar, ConjRhs, ColMajor, Version> {
  using ResScalar = typename ScalarBinaryOpTraits<LhsScalar, RhsScalar>::ReturnType;
  static constexpr bool IsLower = (Mode & Lower) == Lower;
  static constexpr bool HasUnitDiag = (Mode & UnitDiag) == UnitDiag;
  static constexpr bool HasZeroDiag = (Mode & ZeroDiag) == ZeroDiag;
  static EIGEN_DONT_INLINE void run(Index _rows, Index _cols, const LhsScalar* lhs_, Index lhsStride,
                                    const RhsScalar* rhs_, Index rhsIncr, ResScalar* res_, Index resIncr,
                                    const RhsScalar& alpha);
};
 
template <typename Index, int Mode, typename LhsScalar, bool ConjLhs, typename RhsScalar, bool ConjRhs, int Version>
EIGEN_DONT_INLINE void triangular_matrix_vector_product<Index, Mode, LhsScalar, ConjLhs, RhsScalar, ConjRhs, ColMajor,
                                                        Version>::run(Index _rows, Index _cols, const LhsScalar* lhs_,
                                                                      Index lhsStride, const RhsScalar* rhs_,
                                                                      Index rhsIncr, ResScalar* res_, Index resIncr,
                                                                      const RhsScalar& alpha) {
  static const Index PanelWidth = EIGEN_TUNE_TRIANGULAR_PANEL_WIDTH;
  Index size = (std::min)(_rows, _cols);
  Index rows = IsLower ? _rows : (std::min)(_rows, _cols);
  Index cols = IsLower ? (std::min)(_rows, _cols) : _cols;
 
  using LhsMapper = const_blas_data_mapper<LhsScalar, Index, ColMajor>;
  using RhsMapper = const_blas_data_mapper<RhsScalar, Index, RowMajor>;
 
  conj_if<ConjLhs> cjl;
  conj_if<ConjRhs> cjr;
 
  for (Index pi = 0; pi < size; pi += PanelWidth) {
    Index actualPanelWidth = (std::min)(PanelWidth, size - pi);
 
    // Process the triangular panel using raw pointer operations with 2-column batching
    // to eliminate expression template overhead and share result loads/stores.
    EIGEN_IF_CONSTEXPR (IsLower) {
      Index k = 0;
      for (; k + 1 < actualPanelWidth; k += 2) {
        Index i0 = pi + k;
        Index i1 = i0 + 1;
        ResScalar s0 = alpha * cjr(rhs_[i0 * rhsIncr]);
        ResScalar s1 = alpha * cjr(rhs_[i1 * rhsIncr]);
        const LhsScalar* EIGEN_RESTRICT c0 = lhs_ + i0 * lhsStride;
        const LhsScalar* EIGEN_RESTRICT c1 = lhs_ + i1 * lhsStride;
 
        // Diagonal of column 0
        EIGEN_IF_CONSTEXPR (!(HasUnitDiag || HasZeroDiag)) res_[i0] += s0 * cjl(c0[i0]);
        // Row i1: contribution from column 0 + diagonal of column 1
        {
          ResScalar r1 = s0 * cjl(c0[i1]);
          EIGEN_IF_CONSTEXPR (!(HasUnitDiag || HasZeroDiag)) r1 += s1 * cjl(c1[i1]);
          res_[i1] += r1;
        }
        // Shared rows where both columns contribute
        Index panelEnd = pi + actualPanelWidth;
        for (Index j = i1 + 1; j < panelEnd; ++j) res_[j] += s0 * cjl(c0[j]) + s1 * cjl(c1[j]);
 
        EIGEN_IF_CONSTEXPR (HasUnitDiag) {
          res_[i0] += s0;
          res_[i1] += s1;
        }
      }
      if (k < actualPanelWidth) {
        Index i = pi + k;
        ResScalar s = alpha * cjr(rhs_[i * rhsIncr]);
        const LhsScalar* EIGEN_RESTRICT c = lhs_ + i * lhsStride;
        EIGEN_IF_CONSTEXPR (!(HasUnitDiag || HasZeroDiag)) res_[i] += s * cjl(c[i]);
        EIGEN_IF_CONSTEXPR (HasUnitDiag) res_[i] += s;
      }
    } else {
      // Upper triangular: process 2 columns at a time
      Index k = 0;
      for (; k + 1 < actualPanelWidth; k += 2) {
        Index i0 = pi + k;
        Index i1 = i0 + 1;
        ResScalar s0 = alpha * cjr(rhs_[i0 * rhsIncr]);
        ResScalar s1 = alpha * cjr(rhs_[i1 * rhsIncr]);
        const LhsScalar* EIGEN_RESTRICT c0 = lhs_ + i0 * lhsStride;
        const LhsScalar* EIGEN_RESTRICT c1 = lhs_ + i1 * lhsStride;
 
        // Shared rows before the diagonal block
        for (Index j = pi; j < i0; ++j) res_[j] += s0 * cjl(c0[j]) + s1 * cjl(c1[j]);
 
        // Row i0: diagonal of col0 + contribution from col1
        {
          ResScalar r0 = s1 * cjl(c1[i0]);
          EIGEN_IF_CONSTEXPR (!(HasUnitDiag || HasZeroDiag)) r0 += s0 * cjl(c0[i0]);
          res_[i0] += r0;
        }
        // Diagonal of column 1
        EIGEN_IF_CONSTEXPR (!(HasUnitDiag || HasZeroDiag)) res_[i1] += s1 * cjl(c1[i1]);
 
        EIGEN_IF_CONSTEXPR (HasUnitDiag) {
          res_[i0] += s0;
          res_[i1] += s1;
        }
      }
      if (k < actualPanelWidth) {
        Index i = pi + k;
        ResScalar s = alpha * cjr(rhs_[i * rhsIncr]);
        const LhsScalar* EIGEN_RESTRICT c = lhs_ + i * lhsStride;
        for (Index j = pi; j < i; ++j) res_[j] += s * cjl(c[j]);
        EIGEN_IF_CONSTEXPR (!(HasUnitDiag || HasZeroDiag)) res_[i] += s * cjl(c[i]);
        EIGEN_IF_CONSTEXPR (HasUnitDiag) res_[i] += s;
      }
    }
 
    // Rectangular part: delegate to optimized GEMV
    Index r = IsLower ? rows - pi - actualPanelWidth : pi;
    if (r > 0) {
      Index s = IsLower ? pi + actualPanelWidth : 0;
      general_matrix_vector_product<Index, LhsScalar, LhsMapper, ColMajor, ConjLhs, RhsScalar, RhsMapper, ConjRhs,
                                    BuiltIn>::run(r, actualPanelWidth, LhsMapper(&lhs_[pi * lhsStride + s], lhsStride),
                                                  RhsMapper(&rhs_[pi * rhsIncr], rhsIncr), &res_[s], resIncr, alpha);
    }
  }
  EIGEN_IF_CONSTEXPR (!IsLower) {
    if (cols > size) {
      general_matrix_vector_product<Index, LhsScalar, LhsMapper, ColMajor, ConjLhs, RhsScalar, RhsMapper, ConjRhs>::run(
          rows, cols - size, LhsMapper(&lhs_[size * lhsStride], lhsStride), RhsMapper(&rhs_[size * rhsIncr], rhsIncr),
          res_, resIncr, alpha);
    }
  }
}
 
template <typename Index, int Mode, typename LhsScalar, bool ConjLhs, typename RhsScalar, bool ConjRhs, int Version>
struct triangular_matrix_vector_product<Index, Mode, LhsScalar, ConjLhs, RhsScalar, ConjRhs, RowMajor, Version> {
  using ResScalar = typename ScalarBinaryOpTraits<LhsScalar, RhsScalar>::ReturnType;
  static constexpr bool IsLower = (Mode & Lower) == Lower;
  static constexpr bool HasUnitDiag = (Mode & UnitDiag) == UnitDiag;
  static constexpr bool HasZeroDiag = (Mode & ZeroDiag) == ZeroDiag;
  static EIGEN_DONT_INLINE void run(Index _rows, Index _cols, const LhsScalar* lhs_, Index lhsStride,
                                    const RhsScalar* rhs_, Index rhsIncr, ResScalar* res_, Index resIncr,
                                    const ResScalar& alpha);
};
 
template <typename Index, int Mode, typename LhsScalar, bool ConjLhs, typename RhsScalar, bool ConjRhs, int Version>
EIGEN_DONT_INLINE void triangular_matrix_vector_product<Index, Mode, LhsScalar, ConjLhs, RhsScalar, ConjRhs, RowMajor,
                                                        Version>::run(Index _rows, Index _cols, const LhsScalar* lhs_,
                                                                      Index lhsStride, const RhsScalar* rhs_,
                                                                      Index rhsIncr, ResScalar* res_, Index resIncr,
                                                                      const ResScalar& alpha) {
  static const Index PanelWidth = EIGEN_TUNE_TRIANGULAR_PANEL_WIDTH;
  Index diagSize = (std::min)(_rows, _cols);
  Index rows = IsLower ? _rows : diagSize;
  Index cols = IsLower ? diagSize : _cols;
 
  using LhsMapper = const_blas_data_mapper<LhsScalar, Index, RowMajor>;
  using RhsMapper = const_blas_data_mapper<RhsScalar, Index, RowMajor>;
 
  conj_if<ConjLhs> cjl;
  conj_if<ConjRhs> cjr;
 
  for (Index pi = 0; pi < diagSize; pi += PanelWidth) {
    Index actualPanelWidth = (std::min)(PanelWidth, diagSize - pi);
 
    // Process the triangular panel using raw dot products to eliminate
    // the cwiseProduct().sum() expression template overhead.
    for (Index k = 0; k < actualPanelWidth; ++k) {
      Index i = pi + k;
      const LhsScalar* EIGEN_RESTRICT row_i = lhs_ + i * lhsStride;
      ResScalar dot = ResScalar(0);
 
      EIGEN_IF_CONSTEXPR (IsLower) {
        Index s = pi;
        Index len = (HasUnitDiag || HasZeroDiag) ? k : k + 1;
        for (Index j = 0; j < len; ++j) dot += cjl(row_i[s + j]) * cjr(rhs_[s + j]);
      } else {
        Index s = (HasUnitDiag || HasZeroDiag) ? i + 1 : i;
        Index len = pi + actualPanelWidth - s;
        for (Index j = 0; j < len; ++j) dot += cjl(row_i[s + j]) * cjr(rhs_[s + j]);
      }
      res_[i * resIncr] += alpha * dot;
      EIGEN_IF_CONSTEXPR (HasUnitDiag) res_[i * resIncr] += alpha * cjr(rhs_[i]);
    }
 
    // Rectangular part: delegate to optimized GEMV
    Index r = IsLower ? pi : cols - pi - actualPanelWidth;
    if (r > 0) {
      Index s = IsLower ? 0 : pi + actualPanelWidth;
      general_matrix_vector_product<Index, LhsScalar, LhsMapper, RowMajor, ConjLhs, RhsScalar, RhsMapper, ConjRhs,
                                    BuiltIn>::run(actualPanelWidth, r, LhsMapper(&lhs_[pi * lhsStride + s], lhsStride),
                                                  RhsMapper(&rhs_[s], rhsIncr), &res_[pi * resIncr], resIncr, alpha);
    }
  }
  EIGEN_IF_CONSTEXPR (IsLower) {
    if (rows > diagSize) {
      general_matrix_vector_product<Index, LhsScalar, LhsMapper, RowMajor, ConjLhs, RhsScalar, RhsMapper, ConjRhs>::run(
          rows - diagSize, cols, LhsMapper(&lhs_[diagSize * lhsStride], lhsStride), RhsMapper(rhs_, rhsIncr),
          &res_[diagSize * resIncr], resIncr, alpha);
    }
  }
}
 
/***************************************************************************
 * Wrapper to product_triangular_vector
 ***************************************************************************/
 
template <int Mode, int StorageOrder>
struct trmv_selector;
 
}  // end namespace internal
 
namespace internal {
 
template <int Mode, typename Lhs, typename Rhs>
struct triangular_product_impl<Mode, true, Lhs, false, Rhs, true> {
  template <typename Dest>
  static void run(Dest& dst, const Lhs& lhs, const Rhs& rhs, const typename Dest::Scalar& alpha) {
    eigen_assert(dst.rows() == lhs.rows() && dst.cols() == rhs.cols());
 
    internal::trmv_selector<Mode, (int(internal::traits<Lhs>::Flags) & RowMajorBit) ? RowMajor : ColMajor>::run(
        lhs, rhs, dst, alpha);
  }
};
 
template <int Mode, typename Lhs, typename Rhs>
struct triangular_product_impl<Mode, false, Lhs, true, Rhs, false> {
  template <typename Dest>
  static void run(Dest& dst, const Lhs& lhs, const Rhs& rhs, const typename Dest::Scalar& alpha) {
    eigen_assert(dst.rows() == lhs.rows() && dst.cols() == rhs.cols());
 
    Transpose<Dest> dstT(dst);
    internal::trmv_selector<(Mode & (UnitDiag | ZeroDiag)) | ((Mode & Lower) ? Upper : Lower),
                            (int(internal::traits<Rhs>::Flags) & RowMajorBit) ? ColMajor
                                                                              : RowMajor>::run(rhs.transpose(),
                                                                                               lhs.transpose(), dstT,
                                                                                               alpha);
  }
};
 
}  // end namespace internal
 
namespace internal {
 
template <int Mode, typename Lhs, typename Rhs, typename Dest, typename LhsScalar>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void trmv_correct_unit_diagonal(const Lhs& lhs, const Rhs& rhs, Dest& dest,
                                                                      const LhsScalar& lhs_alpha) {
  EIGEN_IF_CONSTEXPR ((Mode & UnitDiag) == UnitDiag) {
    if (!numext::is_exactly_one(lhs_alpha)) {
      Index diagSize = (std::min)(lhs.rows(), lhs.cols());
      dest.head(diagSize) -= (lhs_alpha - LhsScalar(1)) * rhs.head(diagSize);
    }
  }
}
 
template <int Mode>
struct trmv_selector<Mode, ColMajor> {
  template <typename Lhs, typename Rhs, typename Dest>
  static void run(const Lhs& lhs, const Rhs& rhs, Dest& dest, const typename Dest::Scalar& alpha) {
    using LhsScalar = typename Lhs::Scalar;
    using RhsScalar = typename Rhs::Scalar;
    using ResScalar = typename Dest::Scalar;
 
    using LhsBlasTraits = internal::blas_traits<Lhs>;
    using ActualLhsType = typename LhsBlasTraits::DirectLinearAccessType;
    using RhsBlasTraits = internal::blas_traits<Rhs>;
    using ActualRhsType = typename RhsBlasTraits::DirectLinearAccessType;
    add_const_on_value_type_t<ActualLhsType> actualLhs = LhsBlasTraits::extract(lhs);
    add_const_on_value_type_t<ActualRhsType> actualRhs = RhsBlasTraits::extract(rhs);
 
    LhsScalar lhs_alpha = LhsBlasTraits::extractScalarFactor(lhs);
    RhsScalar rhs_alpha = RhsBlasTraits::extractScalarFactor(rhs);
    ResScalar actualAlpha = alpha * lhs_alpha * rhs_alpha;
 
    // FIXME find a way to allow an inner stride on the result if packet_traits<Scalar>::size==1
    // On the other hand, it is good for the cache to pack the vector anyways...
    constexpr bool EvalToDestAtCompileTime = Dest::InnerStrideAtCompileTime == 1;
    constexpr bool ComplexByReal = (NumTraits<LhsScalar>::IsComplex) && (!NumTraits<RhsScalar>::IsComplex);
    constexpr bool MightCannotUseDest = (Dest::InnerStrideAtCompileTime != 1) || ComplexByReal;
 
    gemv_static_vector_if<ResScalar, Dest::SizeAtCompileTime, Dest::MaxSizeAtCompileTime, MightCannotUseDest>
        static_dest;
 
    gemv_destination_policy<RhsScalar, ResScalar, EvalToDestAtCompileTime, ComplexByReal> destPolicy(actualAlpha);
 
    ei_declare_aligned_stack_constructed_variable(ResScalar, actualDestPtr, dest.size(),
                                                  destPolicy.eval_to_dest() ? dest.data() : static_dest.data());
 
    destPolicy.prepare(dest, actualDestPtr);
 
    internal::triangular_matrix_vector_product<Index, Mode, LhsScalar, LhsBlasTraits::NeedToConjugate, RhsScalar,
                                               RhsBlasTraits::NeedToConjugate,
                                               ColMajor>::run(actualLhs.rows(), actualLhs.cols(), actualLhs.data(),
                                                              actualLhs.outerStride(), actualRhs.data(),
                                                              actualRhs.innerStride(), actualDestPtr, 1,
                                                              destPolicy.compatible_alpha());
 
    destPolicy.copy_back(dest, actualDestPtr);
 
    trmv_correct_unit_diagonal<Mode>(lhs, rhs, dest, lhs_alpha);
  }
};
 
template <int Mode>
struct trmv_selector<Mode, RowMajor> {
  template <typename Lhs, typename Rhs, typename Dest>
  static void run(const Lhs& lhs, const Rhs& rhs, Dest& dest, const typename Dest::Scalar& alpha) {
    using LhsScalar = typename Lhs::Scalar;
    using RhsScalar = typename Rhs::Scalar;
    using ResScalar = typename Dest::Scalar;
 
    using LhsBlasTraits = internal::blas_traits<Lhs>;
    using ActualLhsType = typename LhsBlasTraits::DirectLinearAccessType;
    using RhsBlasTraits = internal::blas_traits<Rhs>;
    using ActualRhsType = typename RhsBlasTraits::DirectLinearAccessType;
    using ActualRhsTypeCleaned = internal::remove_all_t<ActualRhsType>;
 
    std::add_const_t<ActualLhsType> actualLhs = LhsBlasTraits::extract(lhs);
    std::add_const_t<ActualRhsType> actualRhs = RhsBlasTraits::extract(rhs);
 
    LhsScalar lhs_alpha = LhsBlasTraits::extractScalarFactor(lhs);
    RhsScalar rhs_alpha = RhsBlasTraits::extractScalarFactor(rhs);
    ResScalar actualAlpha = alpha * lhs_alpha * rhs_alpha;
 
    constexpr bool DirectlyUseRhs = ActualRhsTypeCleaned::InnerStrideAtCompileTime == 1;
 
    gemv_static_vector_if<RhsScalar, ActualRhsTypeCleaned::SizeAtCompileTime,
                          ActualRhsTypeCleaned::MaxSizeAtCompileTime, !DirectlyUseRhs>
        static_rhs;
 
    ei_declare_aligned_stack_constructed_variable(
        RhsScalar, actualRhsPtr, actualRhs.size(),
        DirectlyUseRhs ? const_cast<RhsScalar*>(actualRhs.data()) : static_rhs.data());
 
    gemv_prepare_rhs<DirectlyUseRhs>(actualRhs, actualRhsPtr);
 
    internal::triangular_matrix_vector_product<Index, Mode, LhsScalar, LhsBlasTraits::NeedToConjugate, RhsScalar,
                                               RhsBlasTraits::NeedToConjugate, RowMajor>::run(actualLhs.rows(),
                                                                                              actualLhs.cols(),
                                                                                              actualLhs.data(),
                                                                                              actualLhs.outerStride(),
                                                                                              actualRhsPtr, 1,
                                                                                              dest.data(),
                                                                                              dest.innerStride(),
                                                                                              actualAlpha);
 
    trmv_correct_unit_diagonal<Mode>(lhs, rhs, dest, lhs_alpha);
  }
};
 
}  // end namespace internal
 
}  // end namespace Eigen
 
#endif  // EIGEN_TRIANGULARMATRIXVECTOR_H