Eigen select

This commit is contained in:
Charles Schlosser 2024-04-19 17:52:34 +00:00 committed by Rasmus Munk Larsen
parent 5635d37f46
commit 1550c99541

View File

@ -840,6 +840,27 @@ struct ternary_evaluator<CwiseTernaryOp<TernaryOp, Arg1, Arg2, Arg3>, IndexBased
Data m_d;
};
// specialization for expresions like (a < b).select(c, d) to enable full vectorization
template <typename Arg1, typename Arg2, typename Scalar, typename CmpLhsType, typename CmpRhsType, ComparisonName cmp>
struct evaluator<CwiseTernaryOp<scalar_boolean_select_op<Scalar, Scalar, bool>, Arg1, Arg2,
CwiseBinaryOp<scalar_cmp_op<Scalar, Scalar, cmp, false>, CmpLhsType, CmpRhsType>>>
: public ternary_evaluator<
CwiseTernaryOp<scalar_boolean_select_op<Scalar, Scalar, Scalar>, Arg1, Arg2,
CwiseBinaryOp<scalar_cmp_op<Scalar, Scalar, cmp, true>, CmpLhsType, CmpRhsType>>> {
using DummyTernaryOp = scalar_boolean_select_op<Scalar, Scalar, bool>;
using DummyArg3 = CwiseBinaryOp<scalar_cmp_op<Scalar, Scalar, cmp, false>, CmpLhsType, CmpRhsType>;
using DummyXprType = CwiseTernaryOp<DummyTernaryOp, Arg1, Arg2, DummyArg3>;
using TernaryOp = scalar_boolean_select_op<Scalar, Scalar, Scalar>;
using Arg3 = CwiseBinaryOp<scalar_cmp_op<Scalar, Scalar, cmp, true>, CmpLhsType, CmpRhsType>;
using XprType = CwiseTernaryOp<TernaryOp, Arg1, Arg2, Arg3>;
using Base = ternary_evaluator<XprType>;
EIGEN_DEVICE_FUNC explicit evaluator(const DummyXprType& xpr)
: Base(XprType(xpr.arg1(), xpr.arg2(), Arg3(xpr.arg3().lhs(), xpr.arg3().rhs()))) {}
};
// -------------------- CwiseBinaryOp --------------------
// this is a binary expression