Skip to content

ToT: strided arena kernels for complex inner cells; complex ToT scaling #573

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

Open
kshitij-05 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

ToT: strided arena kernels for complex inner cells; complex ToT scaling #573

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
41 changes: 25 additions & 16 deletions src/TiledArray/expressions/cont_engine.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1366,26 +1366,27 @@ class ContEngine : public BinaryEngine<Derived> {
// non-identity inner result perm is applied downstream and left
// to the per-cell path here.
// The strided kernel is specialized to view (arena) inner cells
// with double storage, and its static_assert requires that of
// ALL THREE operands (result, left, right). Gate on the same
// 3-operand predicate so a mixed-operand contraction (e.g. a
// view/double result with a non-view or non-double operand, or
// float/complex inner) stays on the generic per-cell path and
// never instantiates the double-view-only kernel (which would
// be a hard compile error rather than a graceful fallback).
// whose numeric type is one of the four BLAS gemm element types
// (is_strided_dgemm_numeric_v), shared by ALL THREE operands
// (result, left, right). Gate on the same 3-operand predicate
// so a mixed-operand contraction (e.g. a view result with a
// non-view operand, a non-BLAS inner numeric type, or a
// real-ToT x complex-ToT product) stays on the generic
// per-cell path and never instantiates the view-only kernel
// (which would be a hard compile error rather than a graceful
// fallback).
if constexpr (
TiledArray::is_tensor_view_v<result_tile_element_type> &&
TiledArray::is_tensor_view_v<left_tile_element_type> &&
TiledArray::is_tensor_view_v<right_tile_element_type> &&
std::is_same_v<
typename result_tile_element_type::numeric_type,
double> &&
TiledArray::detail::is_strided_dgemm_numeric_v<
typename result_tile_element_type::numeric_type> &&
std::is_same_v<
typename left_tile_element_type::numeric_type,
double> &&
typename result_tile_element_type::numeric_type> &&
std::is_same_v<
typename right_tile_element_type::numeric_type,
double>) {
typename result_tile_element_type::numeric_type>) {
if (contrreduce_op.gemm_helper().num_contract_ranks() == 0 &&
!bool(inner(this->perm_))) {
const scalar_type factor = this->factor_;
Expand All @@ -1401,7 +1402,9 @@ class ContEngine : public BinaryEngine<Derived> {
Cc, Lt, Rt, static_cast<std::size_t>(M),
static_cast<std::size_t>(N),
static_cast<std::size_t>(K), gh.left_op(),
gh.right_op(), double(factor));
gh.right_op(),
static_cast<typename result_tile_element_type::
numeric_type>(factor));
};
}
// ce+ce (hce+ce): inner CONTRACTION (num_contract_ranks() >=
Expand Down Expand Up @@ -1517,7 +1520,10 @@ class ContEngine : public BinaryEngine<Derived> {
Cc, Lt, Rt, static_cast<std::size_t>(Mo),
static_cast<std::size_t>(No),
static_cast<std::size_t>(Ko), gh.left_op(),
gh.right_op(), double(factor), left_inner_T);
gh.right_op(),
static_cast<typename result_tile_element_type::
numeric_type>(factor),
left_inner_T);
};
} else if (left_arm_ok) {
const scalar_type factor = this->factor_;
Expand All @@ -1537,11 +1543,14 @@ class ContEngine : public BinaryEngine<Derived> {
Cc, Lt, Rt, static_cast<std::size_t>(Mo),
static_cast<std::size_t>(No),
static_cast<std::size_t>(Ko), gh.left_op(),
gh.right_op(), double(factor), right_inner_T);
gh.right_op(),
static_cast<typename result_tile_element_type::
numeric_type>(factor),
right_inner_T);
};
}
// [strided-dgemm] install-decision instrumentation. For each
// ToT contraction reaching this double-view path, report
// ToT contraction reaching this view-cell path, report
// whether a strided-DGEMM regime (hce+e / hc+e / hce+ce)
// FIRED or the contraction REVERTED to the generic by-cell
// evaluation path (with the blocking guard). Gated by
Expand Down
Loading
Loading