leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

CUTLASS 3.5.0 (#1411)

Browse Source
This commit is contained in:
Vijay Thakkar
2024-03-19 17:51:04 -04:00
committed by GitHub
parent ffa34e7075
commit 629f4653c3
468 changed files with 48729 additions and 7252 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/algorithm/axpby.hpp
+13 -6
View File
@@ -33,6 +33,7 @@
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_predicate.hpp>
namespace cute
{
@@ -43,15 +44,17 @@ namespace cute
template <class Alpha,
class XEngine, class XLayout,
class Beta,
class YEngine, class YLayout>
class YEngine, class YLayout,
class PrdTensor = TrivialPredTensor>
CUTE_HOST_DEVICE
void
axpby(Alpha const& alpha,
Tensor<XEngine, XLayout> const& x,
Beta const& beta,
Tensor<YEngine, YLayout> && y)
Tensor<YEngine, YLayout> && y,
PrdTensor const& p = {})
{
return axpby(alpha, x, beta, y);
return axpby(alpha, x, beta, y, p);
}
//
@@ -60,13 +63,15 @@ axpby(Alpha const& alpha,
template <class Alpha,
class XEngine, class XLayout,
class Beta,
class YEngine, class YLayout>
class YEngine, class YLayout,
class PrdTensor = TrivialPredTensor>
CUTE_HOST_DEVICE
void
axpby(Alpha const& alpha,
Tensor<XEngine, XLayout> const& x,
Beta const& beta,
Tensor<YEngine, YLayout> & y)
Tensor<YEngine, YLayout> & y,
PrdTensor const& p = {})
{
auto isBetaZero = [&] () {
if constexpr (is_complex<Beta>::value) {
@@ -81,7 +86,9 @@ axpby(Alpha const& alpha,
CUTE_UNROLL
for (int i = 0; i < size(x); ++i) {
y(i) = (isBetaZero ? alpha * x(i) : alpha * x(i) + beta * y(i));
if (p(i)) {
y(i) = (isBetaZero ? alpha * x(i) : alpha * x(i) + beta * y(i));
}
}
}