Skip to content

ToT: conj() on nested-tile contractions; real-plain x complex-ToT products via one real gemm #574

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
kshitij-05 wants to merge 1 commit into
base: kshitij/feature/strided-tot-kernels-complex
Choose a base branch
from
Draft

ToT: conj() on nested-tile contractions; real-plain x complex-ToT products via one real gemm #574

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 15 additions & 9 deletions src/TiledArray/expressions/cont_engine.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1171,7 +1171,9 @@ class ContEngine : public BinaryEngine<Derived> {
const auto* rp = right.data();
result_tile_element_type acc{0};
for (std::size_t j = 0; j < n; ++j) acc += lp[j] * rp[j];
return static_cast<result_tile_element_type>(factor) * acc;
return TiledArray::detail::elem_factor<result_tile_element_type>(
factor) *
acc;
};
this->element_nonreturn_op_ = [flat_dot](
result_tile_element_type& result,
Expand Down Expand Up @@ -1295,7 +1297,8 @@ class ContEngine : public BinaryEngine<Derived> {
Numeric acc{0};
for (std::size_t j = 0; j < n; ++j) acc += lp[j] * rp[j];
// result cell is pre-shaped [1] by the unit_range plan.
result.data()[0] += static_cast<Numeric>(factor) * acc;
result.data()[0] +=
TiledArray::detail::elem_factor<Numeric>(factor) * acc;
};
if (this->outer_product_uses_summa()) {
this->arena_plan_ =
Expand Down Expand Up @@ -1403,8 +1406,9 @@ class ContEngine : public BinaryEngine<Derived> {
static_cast<std::size_t>(N),
static_cast<std::size_t>(K), gh.left_op(),
gh.right_op(),
static_cast<typename result_tile_element_type::
numeric_type>(factor));
TiledArray::detail::elem_factor<
typename result_tile_element_type::
numeric_type>(factor));
};
}
// ce+ce (hce+ce): inner CONTRACTION (num_contract_ranks() >=
Expand Down Expand Up @@ -1521,8 +1525,9 @@ class ContEngine : public BinaryEngine<Derived> {
static_cast<std::size_t>(No),
static_cast<std::size_t>(Ko), gh.left_op(),
gh.right_op(),
static_cast<typename result_tile_element_type::
numeric_type>(factor),
TiledArray::detail::elem_factor<
typename result_tile_element_type::
numeric_type>(factor),
left_inner_T);
};
} else if (left_arm_ok) {
Expand All @@ -1544,8 +1549,9 @@ class ContEngine : public BinaryEngine<Derived> {
static_cast<std::size_t>(No),
static_cast<std::size_t>(Ko), gh.left_op(),
gh.right_op(),
static_cast<typename result_tile_element_type::
numeric_type>(factor),
TiledArray::detail::elem_factor<
typename result_tile_element_type::
numeric_type>(factor),
right_inner_T);
};
}
Expand Down Expand Up @@ -1702,7 +1708,7 @@ class ContEngine : public BinaryEngine<Derived> {
const auto* rp = right.data();
Numeric acc{0};
for (std::size_t j = 0; j < n; ++j) acc += lp[j] * rp[j];
acc *= static_cast<Numeric>(factor);
acc *= TiledArray::detail::elem_factor<Numeric>(factor);
if (TA::empty(result)) {
using R = typename result_tile_element_type::range_type;
TiledArray::container::svector<std::size_t> ext(
Expand Down
23 changes: 23 additions & 0 deletions src/TiledArray/tensor/complex.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,29 @@ TILEDARRAY_FORCE_INLINE L& operator*=(L& value,
return value;
}

template <typename T>
struct is_complex_conjugate : std::false_type {};
template <typename T>
struct is_complex_conjugate<ComplexConjugate<T>> : std::true_type {};
/// true if \c T is a ComplexConjugate<...> contraction factor
template <typename T>
inline constexpr bool is_complex_conjugate_v =
is_complex_conjugate<std::remove_cv_t<T>>::value;

/// The numeric multiplier to bake into a per-element (per-cell) multiply-add
/// op for a contraction with factor \c factor: the factor itself (converted
/// to \c Numeric) for a numeric factor, and \c Numeric(1) for a
/// ComplexConjugate<...> factor -- the conjugation (and, for
/// ComplexConjugate<Scalar>, the scale) of such a factor is applied to the
/// finished result by ContractReduce's finalization step, not per element.
template <typename Numeric, typename Scalar>
TILEDARRAY_FORCE_INLINE Numeric elem_factor(const Scalar& factor) {
if constexpr (is_complex_conjugate_v<Scalar>)
return Numeric(1);
else
return static_cast<Numeric>(factor);
}

template <typename T>
inline auto abs(const ComplexConjugate<T>& a) {
return std::abs(a.factor());
Expand Down
48 changes: 36 additions & 12 deletions src/TiledArray/tensor/tensor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3141,6 +3141,10 @@ class Tensor {
Tensor& gemm(const Tensor<As...>& A, const Tensor<Bs...>& B, const W alpha,
const math::GemmHelper& gemm_helper) {
numeric_type beta = 1;
// A ComplexConjugate<...> alpha (conj(A*B) at the expression level) is
// applied to the finished result by ContractReduce's finalization step;
// the gemm itself runs with alpha = 1 (see detail::elem_factor).
const numeric_type alpha_n = detail::elem_factor<numeric_type>(alpha);
if (this->empty()) {
*this =
Tensor(gemm_helper.make_result_range<range_type>(A.range_, B.range()),
Expand All @@ -3165,7 +3169,7 @@ class Tensor {
}
for (size_t i = 0; i < this->nbatch(); ++i) {
auto Ci = this->batch(i);
TiledArray::gemm(alpha, A.batch(i), B.batch(i),
TiledArray::gemm(alpha_n, A.batch(i), B.batch(i),
twostep ? numeric_type(0) : numeric_type(1), Ci,
gemm_helper);
}
Expand Down Expand Up @@ -3219,7 +3223,7 @@ class Tensor {
#else // TA_ENABLE_TILE_OPS_LOGGING
for (size_t i = 0; i < this->nbatch(); ++i) {
auto Ci = this->batch(i);
TiledArray::detail::gemm(alpha, A.batch(i), B.batch(i), beta, Ci,
TiledArray::detail::gemm(alpha_n, A.batch(i), B.batch(i), beta, Ci,
gemm_helper);
}
#endif // TA_ENABLE_TILE_OPS_LOGGING
Expand Down Expand Up @@ -3326,7 +3330,16 @@ class Tensor {
if constexpr (detail::is_numeric_v<V> && is_tensor_view_v<U> &&
is_tensor_view_v<value_type>) {
using Real = std::remove_cv_t<typename value_type::value_type>;
if constexpr (std::is_same_v<std::remove_cv_t<V>, Real>) {
using Vr = std::remove_cv_t<V>;
// Same element type: one gemm in that type. Real plain scalars (Vr)
// against complex<Vr> inner cells: the complex slabs are viewed as real
// matrices with the inner extent doubled (re,im interleaved), so one
// real gemm with alpha = beta = 1 accumulates both parts exactly.
constexpr bool same_type = std::is_same_v<Vr, Real>;
constexpr bool interleaved =
!same_type && std::is_same_v<std::complex<Vr>, Real>;
constexpr integer cw = interleaved ? 2 : 1; // reals per inner element
if constexpr (same_type || interleaved) {
if (gemm_helper.left_op() == TiledArray::math::blas::NoTranspose &&
gemm_helper.right_op() == TiledArray::math::blas::NoTranspose) {
// kernel-total timer: destroyed at `return *this;` below, so it
Expand Down Expand Up @@ -3426,18 +3439,20 @@ class Tensor {
detail::g_scale[0].gemm_flop.fetch_add(
2ull * static_cast<std::uint64_t>(K) *
static_cast<std::uint64_t>(N) *
static_cast<std::uint64_t>(A),
static_cast<std::uint64_t>(A) * cw,
std::memory_order_relaxed);
}
const integer Ai = static_cast<integer>(A);
detail::ScopedScaleTimer _scale_gt(detail::g_scale[0].gemm_ns);
TiledArray::math::blas::gemm(
TiledArray::math::blas::Transpose,
TiledArray::math::blas::NoTranspose,
/*M=*/N, /*N=*/Ai, /*K=*/K, Real(1),
/*M=*/N, /*N=*/Ai * cw, /*K=*/K, Vr(1),
/*A=*/right_data, /*lda=*/N,
/*B=*/lc0[0].data(), /*ldb=*/ldb, Real(1),
/*C=*/rc0[0].data(), /*ldc=*/ldc);
/*B=*/reinterpret_cast<const Vr*>(lc0[0].data()),
/*ldb=*/ldb * cw, Vr(1),
/*C=*/reinterpret_cast<Vr*>(rc0[0].data()),
/*ldc=*/ldc * cw);
} else { // per-cell AXPY fallback for this row
if (detail::scale_gemm_timing_enabled()) {
// classify fallback reason (re-scan; observation only, does
Expand Down Expand Up @@ -3506,7 +3521,14 @@ class Tensor {
if constexpr (detail::is_numeric_v<U> && is_tensor_view_v<V> &&
is_tensor_view_v<value_type>) {
using Real = std::remove_cv_t<typename value_type::value_type>;
if constexpr (std::is_same_v<std::remove_cv_t<U>, Real>) {
using Ur = std::remove_cv_t<U>;
// see the tot_x_t block: same type, or real plain x complex<Ur> inner
// cells via the re,im-interleaved real gemm
constexpr bool same_type = std::is_same_v<Ur, Real>;
constexpr bool interleaved =
!same_type && std::is_same_v<std::complex<Ur>, Real>;
constexpr integer cw = interleaved ? 2 : 1; // reals per inner element
if constexpr (same_type || interleaved) {
if (gemm_helper.left_op() == TiledArray::math::blas::NoTranspose &&
gemm_helper.right_op() == TiledArray::math::blas::NoTranspose) {
// kernel-total timer (see tot_x_t block); destroyed at `return`.
Expand Down Expand Up @@ -3590,18 +3612,20 @@ class Tensor {
detail::g_scale[1].gemm_flop.fetch_add(
2ull * static_cast<std::uint64_t>(M) *
static_cast<std::uint64_t>(K) *
static_cast<std::uint64_t>(A),
static_cast<std::uint64_t>(A) * cw,
std::memory_order_relaxed);
}
const integer Ai = static_cast<integer>(A);
detail::ScopedScaleTimer _scale_gt(detail::g_scale[1].gemm_ns);
TiledArray::math::blas::gemm(
TiledArray::math::blas::NoTranspose,
TiledArray::math::blas::NoTranspose,
/*M=*/M, /*N=*/Ai, /*K=*/K, Real(1),
/*M=*/M, /*N=*/Ai * cw, /*K=*/K, Ur(1),
/*A=*/left_data, /*lda=*/K,
/*B=*/right_data[n].data(), /*ldb=*/ldb, Real(1),
/*C=*/this_data[n].data(), /*ldc=*/ldc);
/*B=*/reinterpret_cast<const Ur*>(right_data[n].data()),
/*ldb=*/ldb * cw, Ur(1),
/*C=*/reinterpret_cast<Ur*>(this_data[n].data()),
/*ldc=*/ldc * cw);
} else { // per-cell AXPY fallback for this column
if (detail::scale_gemm_timing_enabled()) {
// classify fallback reason (re-scan; observation only) +
Expand Down
124 changes: 65 additions & 59 deletions src/TiledArray/tile_op/contract_reduce.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,57 @@ class ContractReduceBase {
strided_oprod_op() const {
return pimpl_->strided_oprod_op_;
}
/// Reduce two partial results: \c result += \c arg. Arena ToT partials
/// reduced from disjoint K-panel subsets can carry different inner-cell
/// sparsity, so their shapes are unioned before accumulating.
template <typename R>
void reduce_results(R& result, const R& arg) const {
if constexpr (
detail::is_contraction_arena_tot_v<
R, std::remove_cv_t<std::remove_reference_t<first_argument_type>>,
std::remove_cv_t<std::remove_reference_t<second_argument_type>>>) {
detail::arena_tot_add_to(result, arg);
} else {
using TiledArray::add_to;
add_to(result, arg);
}
}

/// Nested-tile accumulate: \c result += \c left * \c right through the
/// per-cell multiply-add op (via the arena plan and the strided
/// outer-product op when installed). Shared by the primary ContractReduce
/// and its ComplexConjugate specializations -- a ComplexConjugate factor is
/// applied to the finished result by those specializations' finalization
/// step, so the accumulate itself is identical.
template <typename R, typename L, typename Rt>
void accumulate_nested(R& result, const L& left, const Rt& right) const {
using TiledArray::empty;
using TiledArray::gemm;
TA_ASSERT(this->elem_muladd_op());
if constexpr (detail::is_contraction_arena_tot_v<
R, std::remove_cv_t<std::remove_reference_t<L>>,
std::remove_cv_t<std::remove_reference_t<Rt>>>) {
// The result tile is shaped from operand inner cells. A SUMMA
// reduction streams K-panels one at a time: the first panel sizes the
// result; a later panel of a contracted-dimension-sparse ToT operand
// can touch inner cells the first panel left null, so each subsequent
// panel extends the result to cover its own cells.
if (this->arena_plan().has_value()) {
if (empty(result))
result = this->arena_plan()->reserve_and_construct(
left, right, this->gemm_helper());
else
this->arena_plan()->grow_to_cover(result, left, right,
this->gemm_helper());
}
if (this->strided_oprod_op()) {
this->strided_oprod_op()(result, left, right, this->gemm_helper());
return;
}
}
gemm(result, left, right, this->gemm_helper(), this->elem_muladd_op());
}

void set_strided_oprod_op(
TiledArray::function_ref<typename Impl::strided_oprod_op_type> op) {
pimpl_->strided_oprod_op_ = op;
Expand Down Expand Up @@ -384,19 +435,7 @@ class ContractReduce : public ContractReduceBase<Result, Left, Right, Scalar> {
/// target
/// \param[in] arg The argument that will be added to \c result
void operator()(result_type& result, const result_type& arg) const {
if constexpr (
detail::is_contraction_arena_tot_v<
result_type,
std::remove_cv_t<std::remove_reference_t<first_argument_type>>,
std::remove_cv_t<std::remove_reference_t<second_argument_type>>>) {
// Two partial contraction results reduced from disjoint K-panel
// subsets can carry different inner-cell sparsity; union their shapes
// before accumulating.
detail::arena_tot_add_to(result, arg);
} else {
using TiledArray::add_to;
add_to(result, arg);
}
this->reduce_results(result, arg);
}

/// Contract a pair of tiles and add to a target tile
Expand All @@ -413,34 +452,7 @@ class ContractReduce : public ContractReduceBase<Result, Left, Right, Scalar> {
if (empty(left) || empty(right)) return;

if constexpr (!ContractReduceBase_::plain_tensors) {
TA_ASSERT(this->elem_muladd_op());
if constexpr (detail::is_contraction_arena_tot_v<
result_type,
std::remove_cv_t<
std::remove_reference_t<first_argument_type>>,
std::remove_cv_t<
std::remove_reference_t<second_argument_type>>>) {
// The result tile is shaped from operand inner cells. A SUMMA
// reduction streams K-panels one at a time: the first panel sizes the
// result; a later panel of a contracted-dimension-sparse ToT operand
// can touch inner cells the first panel left null, so each subsequent
// panel extends the result to cover its own cells.
if (this->arena_plan().has_value()) {
if (empty(result))
result = this->arena_plan()->reserve_and_construct(
left, right, this->gemm_helper());
else
this->arena_plan()->grow_to_cover(result, left, right,
this->gemm_helper());
}
if (this->strided_oprod_op()) {
this->strided_oprod_op()(result, left, right,
ContractReduceBase_::gemm_helper());
return;
}
}
gemm(result, left, right, ContractReduceBase_::gemm_helper(),
this->elem_muladd_op());
this->accumulate_nested(result, left, right);
} else { // plain tensors
TA_ASSERT(!this->elem_muladd_op());
if (empty(result))
Expand Down Expand Up @@ -476,10 +488,8 @@ class ContractReduce<Result, Left, Right,
typedef typename ContractReduceBase_::first_argument_type
first_argument_type; ///< The left tile type
typedef typename ContractReduceBase_::second_argument_type
second_argument_type; ///< The right tile type
typedef decltype(gemm(std::declval<Left>(), std::declval<Right>(), 1,
std::declval<math::GemmHelper>()))
result_type; ///< The result tile type.
second_argument_type; ///< The right tile type
typedef Result result_type; ///< The result tile type.
typedef TiledArray::detail::ComplexConjugate<void> scalar_type;

using typename ContractReduceBase_::elem_muladd_op_type;
Expand Down Expand Up @@ -555,8 +565,7 @@ class ContractReduce<Result, Left, Right,
/// target
/// \param[in] arg The argument that will be added to \c result
void operator()(result_type& result, const result_type& arg) const {
using TiledArray::add_to;
add_to(result, arg);
this->reduce_results(result, arg);
}

/// Contract a pair of tiles and add to a target tile
Expand All @@ -568,10 +577,10 @@ class ContractReduce<Result, Left, Right,
/// \param[in] right The right-hand tile to be contracted
void operator()(result_type& result, const first_argument_type& left,
const second_argument_type& right) const {
using TiledArray::empty;
if (empty(left) || empty(right)) return;
if constexpr (!ContractReduceBase_::plain_tensors) {
TA_ASSERT(this->elem_muladd_op());
// not yet implemented
abort();
this->accumulate_nested(result, left, right);
} else {
TA_ASSERT(!this->elem_muladd_op());
using TiledArray::empty;
Expand Down Expand Up @@ -608,10 +617,8 @@ class ContractReduce<Result, Left, Right,
typedef typename ContractReduceBase_::first_argument_type
first_argument_type; ///< The left tile type
typedef typename ContractReduceBase_::second_argument_type
second_argument_type; ///< The right tile type
typedef decltype(gemm(std::declval<Left>(), std::declval<Right>(), 1,
std::declval<math::GemmHelper>()))
result_type; ///< The result tile type.
second_argument_type; ///< The right tile type
typedef Result result_type; ///< The result tile type.
typedef TiledArray::detail::ComplexConjugate<Scalar> scalar_type;

using typename ContractReduceBase_::elem_muladd_op_type;
Expand Down Expand Up @@ -687,8 +694,7 @@ class ContractReduce<Result, Left, Right,
/// target
/// \param[in] arg The argument that will be added to \c result
void operator()(result_type& result, const result_type& arg) const {
using TiledArray::add_to;
add_to(result, arg);
this->reduce_results(result, arg);
}

/// Contract a pair of tiles and add to a target tile
Expand All @@ -700,10 +706,10 @@ class ContractReduce<Result, Left, Right,
/// \param[in] right The right-hand tile to be contracted
void operator()(result_type& result, const first_argument_type& left,
const second_argument_type& right) const {
using TiledArray::empty;
if (empty(left) || empty(right)) return;
if constexpr (!ContractReduceBase_::plain_tensors) {
TA_ASSERT(this->elem_muladd_op());
// not yet implemented
abort();
this->accumulate_nested(result, left, right);
} else {
TA_ASSERT(!this->elem_muladd_op());
using TiledArray::empty;
Expand Down
Loading
Loading