mirrors / paddlepaddle--paddle

项目文件夹

paddlepaddle--paddle
文件
main
T
wehub-resource-sync e25996e7db chore: import upstream snapshot with attribution
2026-07-13 12:40:42 +08:00

6231 行
202 KiB
C++
原始文件 永久链接 Blame 文件历史

此文件含有模棱两可的 Unicode 字符
此文件含有可能会与其他字符混淆的 Unicode 字符。 如果您是想特意这样的，可以安全地忽略该警告。 使用 Escape 按钮显示他们。
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <glog/logging.h>
#include <algorithm>
#include <cmath>
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
#ifndef _USE_MATH_DEFINES
#define _USE_MATH_DEFINES
#endif
#include <type_traits>
#ifdef PADDLE_WITH_SLEEF
#include <sleef.h>
#endif
#include "paddle/phi/common/amp_type_traits.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/phi/kernels/funcs/eigen/common.h"
#include "paddle/phi/kernels/funcs/eigen/extensions.h"
#ifdef PADDLE_WITH_XPU_KP
#define __forceinline__ __inline__
#endif
namespace phi {
namespace funcs {
enum ActBwdOpFwdDeps {
kNoDeps = 0x00, // Do not need any forward input/output
kDepX = 0x01, // Only need forward input X
kDepOut = 0x02, // Only need forward output Out
};
template <typename T, typename AttrT = float>
struct BaseActivationFunctor {
using ELEMENT_TYPE = T;
using AttrPair = std::vector<std::pair<const char*, AttrT*>>;
AttrPair GetAttrs() { return AttrPair(); }
};
template <typename T>
struct Sine {
HOSTDEVICE T operator()(const T& val) const { return sin(val); }
};
// Specialized Sine for float using Sleef (matches PyTorch's u35 precision)
template <>
struct Sine<float> {
HOSTDEVICE float operator()(const float& val) const {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
return sin(val);
#elif defined(PADDLE_WITH_SLEEF)
return Sleef_sinf1_u35(val);
#else
return sin(val);
#endif
}
};
// Specialized Sine for double using Sleef (matches PyTorch's u10 precision)
template <>
struct Sine<double> {
HOSTDEVICE double operator()(const double& val) const {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
return sin(val);
#elif defined(PADDLE_WITH_SLEEF)
return Sleef_sind1_u10(val);
#else
return sin(val);
#endif
}
};
template <>
struct Sine<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
return float16(sin(static_cast<float>(val)));
#elif defined(PADDLE_WITH_SLEEF)
return float16(Sleef_sinf1_u35(static_cast<float>(val)));
#else
return float16(sin(static_cast<float>(val)));
#endif
}
};
template <>
struct Sine<bfloat16> {
HOSTDEVICE bfloat16 operator()(const bfloat16& val) const {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
return bfloat16(sin(static_cast<float>(val)));
#elif defined(PADDLE_WITH_SLEEF)
return bfloat16(Sleef_sinf1_u35(static_cast<float>(val)));
#else
return bfloat16(sin(static_cast<float>(val)));
#endif
}
};
template <typename T>
struct Cosine {
HOSTDEVICE T operator()(const T& val) const { return cos(val); }
};
// Specialized Cosine for float using Sleef (matches PyTorch's u35 precision)
template <>
struct Cosine<float> {
HOSTDEVICE float operator()(const float& val) const {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
return cos(val);
#elif defined(PADDLE_WITH_SLEEF)
return Sleef_cosf1_u35(val);
#else
return cos(val);
#endif
}
};
// Specialized Cosine for double using Sleef (matches PyTorch's u10 precision)
template <>
struct Cosine<double> {
HOSTDEVICE double operator()(const double& val) const {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
return cos(val);
#elif defined(PADDLE_WITH_SLEEF)
return Sleef_cosd1_u10(val);
#else
return cos(val);
#endif
}
};
template <>
struct Cosine<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
return float16(cos(static_cast<float>(val)));
#elif defined(PADDLE_WITH_SLEEF)
return float16(Sleef_cosf1_u35(static_cast<float>(val)));
#else
return float16(cos(static_cast<float>(val)));
#endif
}
};
template <>
struct Cosine<bfloat16> {
HOSTDEVICE bfloat16 operator()(const bfloat16& val) const {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
return bfloat16(cos(static_cast<float>(val)));
#elif defined(PADDLE_WITH_SLEEF)
return bfloat16(Sleef_cosf1_u35(static_cast<float>(val)));
#else
return bfloat16(cos(static_cast<float>(val)));
#endif
}
};
template <typename T>
using ComplexType = phi::dtype::complex<T>;
// T is phi::complex64 or phi::complex128
template <typename T>
struct Conj {
HOSTDEVICE ComplexType<T> operator()(const ComplexType<T>& val) const {
return ComplexType<T>(val.real, -val.imag);
}
};
// T is phi::complex64 or phi::complex128
template <typename T>
struct Real {
HOSTDEVICE ComplexType<T> operator()(const ComplexType<T>& val) const {
return ComplexType<T>(val.real);
}
};
// sine'(x) = cos(x)
template <typename T>
struct SinGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * x.unaryExpr(Cosine<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct SinGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * x.unaryExpr(Cosine<ComplexType<T>>()).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
// sine(x) = sin(x)
template <typename T>
struct SinFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
// Note(GGBond8488): Since Eigen3.3, Behavior like {A = (B * A).cwiseAbs()}
// will give wrong result, details see
// http://eigen.tuxfamily.org/dox/group__TopicAliasing.html
out.device(d) = x.unaryExpr(Sine<T>()).eval();
}
};
// sine''(x) = -sin(x)
template <typename T>
struct SinDoubleGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* dOut,
const DenseTensor* ddX,
DenseTensor* dX,
DenseTensor* ddOut) const {
auto* d = dev.eigen_device();
auto d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "d2d1x", "SinDoubleGrad"));
auto x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(X, "Input", "x", "SinDoubleGrad"));
// calculate d2x first, so d2d1y can inplace d2d1x
if (dX) {
auto d2x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dX, "Output", "d2x", "SinDoubleGrad"));
if (dOut) {
auto d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Output", "d1y", "SinDoubleGrad"));
d2x.device(*d) = -d2d1x * x.unaryExpr(Sine<T>()) * d1y;
} else {
d2x.device(*d) = -d2d1x * x.unaryExpr(Sine<T>()) * static_cast<T>(0);
}
}
// calculate d2d1y
if (ddOut) {
auto d2d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "d2d1y", "SinDoubleGrad"));
d2d1y.device(*d) = d2d1x * x.unaryExpr(Cosine<T>());
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
// 1st reverse grad
// y = sin(x)
// x --> y
// d1x = d1y * cos(x)
//
// 2nd reverse grad
// x, d1y --> d1x
// d2x = -sin(x) * d1y * d2d1x
// d2d1y = cos(x) * d2d1x
//
// 3rd reverse grad
// x, d1y, d2d1x --> d2x, d2d1y
// d3x = -cos(x) * d1y * d2d1x * d3d2x - sin(x) * d2d1x * d3d2d1y
// d3d1y = -sin(x) * d2d1x * d3d2x
// d3d2d1x = -sin(x) * d1y * d3d2x + cos(x) * d3d2d1y
template <typename T>
struct SinTripleGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* ddX,
const DenseTensor* dOut,
const DenseTensor* d_DDOut,
const DenseTensor* d_dx_New,
DenseTensor* d_d_Out,
DenseTensor* d_x_New,
DenseTensor* d_DDx) const {
auto* d = dev.eigen_device();
auto x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(X, "Input", "x", "SinTripleGrad"));
auto d3d2x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_dx_New, "Input", "d3d2x", "SinTripleGrad"));
if (d_x_New) {
auto d3x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_x_New, "Output", "d3x", "SinTripleGrad"));
if (dOut && ddX && d_DDOut) {
auto d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "d2d1x", "SinTripleGrad"));
auto d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "d1y", "SinTripleGrad"));
auto d3d2d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "d3d2d1y", "SinTripleGrad"));
d3x.device(*d) = -x.unaryExpr(Cosine<T>()) * d1y * d2d1x * d3d2x -
x.unaryExpr(Sine<T>()) * d2d1x * d3d2d1y;
} else if (!dOut && ddX && d_DDOut) {
auto d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "d2d1x", "SinTripleGrad"));
auto d3d2d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "d3d2d1y", "SinTripleGrad"));
d3x.device(*d) = -x.unaryExpr(Sine<T>()) * d2d1x * d3d2d1y;
} else if (dOut && ddX && !d_DDOut) {
auto d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "d2d1x", "SinTripleGrad"));
auto d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "d1y", "SinTripleGrad"));
d3x.device(*d) = -x.unaryExpr(Cosine<T>()) * d1y * d2d1x * d3d2x;
} else {
d3x.device(*d) = x * static_cast<T>(0);
}
}
if (d_d_Out) {
auto d3d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_d_Out, "Output", "d3d1y", "SinTripleGrad"));
if (ddX) {
auto d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "d2d1x", "SinTripleGrad"));
d3d1y.device(*d) = -x.unaryExpr(Sine<T>()) * d2d1x * d3d2x;
} else {
d3d1y.device(*d) = static_cast<T>(0) * x;
}
}
if (d_DDx) {
auto d3d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDx, "Output", "d3d2d1x", "SinTripleGrad"));
if (dOut && d_DDOut) {
auto d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "d1y", "SinTripleGrad"));
auto d3d2d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "d3d2d1y", "SinTripleGrad"));
d3d2d1x.device(*d) = -x.unaryExpr(Sine<T>()) * d1y * d3d2x +
x.unaryExpr(Cosine<T>()) * d3d2d1y;
} else if (dOut && !d_DDOut) {
auto d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "d1y", "SinTripleGrad"));
d3d2d1x.device(*d) = -x.unaryExpr(Sine<T>()) * d1y * d3d2x;
} else if (!dOut && d_DDOut) {
auto d3d2d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "d3d2d1y", "SinTripleGrad"));
d3d2d1x.device(*d) = x.unaryExpr(Cosine<T>()) * d3d2d1y;
} else {
d3d2d1x.device(*d) = x * static_cast<T>(0);
}
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
// reciprocal(x) = 1 / x
template <typename T>
struct ReciprocalFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = static_cast<T>(1) / x;
}
};
template <typename T>
struct Reciprocal {
HOSTDEVICE ComplexType<T> operator()(const ComplexType<T>& val) const {
auto both_inf = [](T real, T imag) {
return (std::isinf(real) && std::isinf(imag));
};
auto either_inf = [](T real, T imag) {
return std::isinf(real) || std::isinf(imag);
};
auto either_nan = [](T real, T imag) {
return std::isnan(real) || std::isnan(imag);
};
if (either_nan(val.real, val.imag) || both_inf(val.real, val.imag)) {
// If either is Nan or both are infinite, return {nan, nan}
return ComplexType<T>(std::numeric_limits<T>::quiet_NaN(),
std::numeric_limits<T>::quiet_NaN());
} else if (either_inf(val.real, val.imag)) {
// If either is Inf, return {0, 0}
return ComplexType<T>{static_cast<T>(0), static_cast<T>(0)};
}
return static_cast<ComplexType<T>>(1.0) / val;
}
};
template <typename T>
struct ReciprocalFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Reciprocal<T>());
}
};
template <typename T>
struct ReciprocalGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = dout * static_cast<T>(-1) * out * out;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct ReciprocalGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = dout * static_cast<ComplexType<T>>(-1) *
(out * out).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
// 1st reverse grad
// y = cos(x)
// x --> y
// d1x = d1y * -sin(x)
//
// 2nd reverse grad
// x, d1y --> d1x
// d2x = -cos(x) * d1y * d2d1x
// d2d1y = -sin(x) * d2d1x
//
// 3rd reverse grad
// x, d1y, d2d1x --> d2x, d2d1y
// d3x = sin(x) * d1y * d2d1x * d3d2x - cos(x) * d2d1x * d3d2d1y
// d3d1y = -cos(x) * d2d1x * d3d2x
// d3d2d1x = -cos(x) * d1y * d3d2x - sin(x) * d3d2d1y
// cosine'(x) = -sin(x)
template <typename T>
struct CosGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = -dout * x.unaryExpr(Sine<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct CosGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
-dout * x.unaryExpr(Sine<ComplexType<T>>()).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
// cos''(x) = -cos(x)
template <typename T>
struct CosDoubleGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* dOut,
const DenseTensor* ddX,
DenseTensor* dX,
DenseTensor* ddOut) const {
auto* device = dev.eigen_device();
auto d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "d2d1x", "CosDoubleGrad"));
auto x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(X, "Input", "x", "CosDoubleGrad"));
// calculate d2x first, so d2d1y can inplace d2d1x
if (dX) {
auto d2x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dX, "Output", "d2x", "CosDoubleGrad"));
auto d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Output", "d1y", "CosDoubleGrad"));
d2x.device(*device) = -d2d1x * x.unaryExpr(Cosine<T>()) * d1y;
}
if (ddOut) {
// calculate d2d1y
auto d2d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "d2d1y", "CosDoubleGrad"));
d2d1y.device(*device) = -d2d1x * x.unaryExpr(Sine<T>());
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct CosTripleGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* ddX,
const DenseTensor* dOut,
const DenseTensor* d_DDOut,
const DenseTensor* d_dx_New,
DenseTensor* d_d_Out,
DenseTensor* d_x_New,
DenseTensor* d_DDx) const {
auto* d = dev.eigen_device();
auto x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(X, "Input", "x", "CosTripleGrad"));
auto d3d2x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_dx_New, "Input", "d3d2x", "CosTripleGrad"));
if (d_x_New) {
auto d3x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_x_New, "Output", "d3x", "CosTripleGrad"));
if (dOut && ddX && d_DDOut) {
auto d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "d2d1x", "CosTripleGrad"));
auto d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "d1y", "CosTripleGrad"));
auto d3d2d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "d3d2d1y", "CosTripleGrad"));
d3x.device(*d) = x.unaryExpr(Sine<T>()) * d1y * d2d1x * d3d2x -
x.unaryExpr(Cosine<T>()) * d2d1x * d3d2d1y;
} else if (dOut && ddX && !d_DDOut) {
auto d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "d2d1x", "CosTripleGrad"));
auto d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "d1y", "CosTripleGrad"));
d3x.device(*d) = x.unaryExpr(Sine<T>()) * d1y * d2d1x * d3d2x;
} else if (!dOut && ddX && d_DDOut) {
auto d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "d2d1x", "CosTripleGrad"));
auto d3d2d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "d3d2d1y", "CosTripleGrad"));
d3x.device(*d) = -x.unaryExpr(Cosine<T>()) * d2d1x * d3d2d1y;
} else {
d3x.device(*d) = static_cast<T>(0) * x;
}
}
if (d_d_Out) {
auto d3d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_d_Out, "Output", "d3d1y", "CosTripleGrad"));
if (ddX) {
auto d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "d2d1x", "CosTripleGrad"));
d3d1y.device(*d) = -x.unaryExpr(Cosine<T>()) * d2d1x * d3d2x;
} else {
d3d1y.device(*d) = static_cast<T>(0) * x;
}
}
if (d_DDx) {
auto d3d2d1x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDx, "Output", "d3d2d1x", "CosTripleGrad"));
if (dOut && d_DDOut) {
auto d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "d1y", "CosTripleGrad"));
auto d3d2d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "d3d2d1y", "CosTripleGrad"));
d3d2d1x.device(*d) = -x.unaryExpr(Cosine<T>()) * d1y * d3d2x -
x.unaryExpr(Sine<T>()) * d3d2d1y;
} else if (!dOut && d_DDOut) {
auto d3d2d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "d3d2d1y", "CosTripleGrad"));
d3d2d1x.device(*d) = -x.unaryExpr(Sine<T>()) * d3d2d1y;
} else if (dOut && !d_DDOut) {
auto d1y = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "d1y", "CosTripleGrad"));
d3d2d1x.device(*d) = -x.unaryExpr(Cosine<T>()) * d1y * d3d2x;
} else {
d3d2d1x.device(*d) = static_cast<T>(0) * x;
}
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
// cosine(x) = cos(x)
template <typename T>
struct CosFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Cosine<T>()).eval();
}
};
template <typename T>
struct LogitFunctor {
template <typename Device, typename X, typename Out, typename P>
void operator()(Device d, X x, Out out, P p, double eps) const {
// logit(x) = ln(x/(1-x))
auto tmp_x =
(x.cwiseMin(static_cast<T>(1.0 - eps))).cwiseMax(static_cast<T>(eps));
if (!eps) {
out.device(d) = (x < static_cast<T>(0.0) || x > static_cast<T>(1.0))
.select(p.constant(static_cast<T>(NAN)),
(tmp_x / (static_cast<T>(1) - tmp_x)).log());
} else {
out.device(d) = (tmp_x / (static_cast<T>(1) - tmp_x)).log();
}
}
};
// mish(x) = x * tanh(softplus(x))
// softplus(x) = x, if x > threshold
// = ln(1 + exp(x)), otherwise
template <typename T>
struct MishFunctor : public BaseActivationFunctor<T> {
float threshold;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
auto sp = (x > static_cast<T>(threshold)) // NOLINT
.select(x, (static_cast<T>(1) + x.exp()).log());
out.device(d) = x * sp.tanh();
}
};
// dx = dout * (tanh(sp) + x * (1 - tanh(sp) ** 2) * (1 - exp(-sp)))
// sp = softplus(x)
template <typename T>
struct MishGradFunctor : public BaseActivationFunctor<T> {
float threshold;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto sp = (x > static_cast<T>(threshold)) // NOLINT
.select(x, (static_cast<T>(1) + x.exp()).log());
auto gsp = static_cast<T>(1) - (-sp).exp();
auto tsp = sp.tanh();
dx.device(d) = dout * (tsp + x * (static_cast<T>(1) - tsp * tsp) * gsp);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct STanhFunctor : public BaseActivationFunctor<T> {
float scale_a;
float scale_b;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"scale_a", &scale_a}, {"scale_b", &scale_b}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = static_cast<T>(scale_b) *
(static_cast<T>(scale_a) * x).tanh(); // NOLINT
}
};
template <typename T>
struct STanhGradFunctor : public BaseActivationFunctor<T> {
float scale_a;
float scale_b;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"scale_a", &scale_a}, {"scale_b", &scale_b}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto a = static_cast<T>(scale_a); // NOLINT
auto b = static_cast<T>(scale_b);
auto temp = (a * x).tanh() * (a * x).tanh();
dx.device(d) = dout * a * b * (static_cast<T>(1) - temp);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct STanhGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
float scale_a;
float scale_b;
typename BaseActivationFunctor<ComplexType<T>>::AttrPair GetAttrs() {
return {{"scale_a", &scale_a}, {"scale_b", &scale_b}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto a = static_cast<ComplexType<T>>(scale_a); // NOLINT
auto b = static_cast<ComplexType<T>>(scale_b);
auto temp = (a * x).tanh() * (a * x).tanh();
dx.device(d) =
dout *
(a * b * (static_cast<ComplexType<T>>(1) - temp)).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct Tangent {
HOSTDEVICE T operator()(const T& val) const { return tan(val); }
};
template <>
struct Tangent<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(tan(static_cast<float>(val)));
}
};
// Tangent'(x) = -Tangent(x)
template <typename T>
struct TanGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout / x.unaryExpr(Cosine<T>()).square();
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct TanGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
// auto dx_ =
// static_cast<ComplexType<T>>(x.unaryExpr(Cosine<T>()).square());
// ComplexType<T> dx_conj_(dx_.real, -dx_.imag);
// dx.device(d) = dout / dx_conj_;
dx.device(d) =
dout /
x.unaryExpr(Cosine<ComplexType<T>>()).square().unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
// square(x) = x^2
template <typename T>
struct SquareFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.square();
}
};
template <typename T>
struct SquareGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * static_cast<T>(2) * x;
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct SquareGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * static_cast<ComplexType<T>>(2) * x.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// sqrt(x) = x^(1/2)
template <typename T>
struct SqrtFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.sqrt();
}
};
template <typename T>
struct SqrtGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = static_cast<T>(0.5) * dout / out;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct SqrtGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) =
dout * (static_cast<ComplexType<T>>(0.5) / out).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
// rsqrt(x) = x^(-1/2)
template <typename T>
struct RsqrtFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.rsqrt();
}
};
template <typename T>
struct RsqrtGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = static_cast<T>(-0.5) * dout * out * out * out;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
// // For numerical stability, using the following formula instead of
// softplus(x) =
// // log(1 + exp(x))
// // softplus(x) = log(1 + exp(beta * x)) / beta when beta * x <=
// threshold(beta =
// // 1, threshold = 20 by default), otherwise x
template <typename T>
struct SoftplusFunctor : public BaseActivationFunctor<T> {
using AttrPair = std::vector<std::pair<const char*, double*>>;
double beta;
double threshold;
typename SoftplusFunctor<T>::AttrPair GetAttrs() {
return {{"beta", &beta}, {"threshold", &threshold}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
auto x_beta = static_cast<T>(beta) * x; // NOLINT
out.device(d) = (x_beta > static_cast<T>(threshold))
.select(x,
(static_cast<T>(1) + x_beta.exp()).log() /
static_cast<T>(beta));
}
};
template <typename T>
struct SoftplusFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
float beta;
float threshold;
typename BaseActivationFunctor<ComplexType<T>>::AttrPair GetAttrs() {
return {{"beta", &beta}, {"threshold", &threshold}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
auto x_beta = static_cast<ComplexType<T>>(beta) * x;
out.device(d) =
(x_beta > static_cast<ComplexType<T>>(threshold))
.select(x,
(static_cast<ComplexType<T>>(1) + x_beta.exp()).log() /
static_cast<ComplexType<T>>(beta));
}
};
// For numerical stability, using the following formula instead of
// d(softplus(x))/dx = 1 / (1 + exp(-x))
// d(softplus(x))/dx = 1 / (1 + exp(-beta * x)) when beta * x <= threshold(beta
// = 1, threshold = 20 by default), otherwise x
template <typename T>
struct SoftplusGradFunctor : public BaseActivationFunctor<T> {
using AttrPair = std::vector<std::pair<const char*, double*>>;
double beta;
double threshold;
typename SoftplusGradFunctor<T>::AttrPair GetAttrs() {
return {{"beta", &beta}, {"threshold", &threshold}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto x_beta = static_cast<T>(beta) * x; // NOLINT
dx.device(d) =
(x_beta > static_cast<T>(threshold))
.select(dout, dout / (static_cast<T>(1) + (-x_beta).exp()));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct SoftplusGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
using AttrPair = std::vector<std::pair<const char*, double*>>;
double beta;
double threshold;
typename SoftplusGradFunctor<ComplexType<T>>::AttrPair GetAttrs() {
return {{"beta", &beta}, {"threshold", &threshold}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto x_beta = static_cast<ComplexType<T>>(beta) * x; // NOLINT
dx.device(d) =
(x_beta > static_cast<ComplexType<T>>(threshold))
.select(dout,
dout / (static_cast<ComplexType<T>>(1) + (-x_beta).exp())
.unaryExpr(Conj<T>()));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct SoftplusDoubleGradFunctor : public BaseActivationFunctor<T> {
using AttrPair = std::vector<std::pair<const char*, double*>>;
double beta;
double threshold;
typename SoftplusDoubleGradFunctor<T>::AttrPair GetAttrs() {
return {{"beta", &beta}, {"threshold", &threshold}};
}
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* dOut,
const DenseTensor* ddX,
DenseTensor* dX,
DenseTensor* ddOut) const {
auto* d = dev.eigen_device();
auto x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(X, "Input", "X", "SoftplusDoubleGrad"));
auto x_beta = static_cast<T>(beta) * x; // NOLINT
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "SoftplusDoubleGrad"));
if (dX) {
auto dx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dX, "Output", "DX", "SoftplusDoubleGrad"));
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Output", "DOut", "SoftplusDoubleGrad"));
// ddx * dout * beta * exp(x_beta) / (exp(x_beta) + 1) ^ 2, if x_beta
// <= threshold
// 0, if x_beta > threshold
dx.device(*d) =
(x_beta > static_cast<T>(threshold))
.select(x.constant(static_cast<T>(0)),
ddx * dout * static_cast<T>(beta) * x_beta.exp() /
(x_beta.exp() + static_cast<T>(1))
.pow(static_cast<T>(2)));
}
if (ddOut) {
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "SoftplusDoubleGrad"));
// ddx / (1 + exp(-x_beta)), if x_beta <= threshold
// ddx, if x_beta > threshold
ddout.device(*d) =
(x_beta > static_cast<T>(threshold))
.select(ddx, ddx / (static_cast<T>(1) + (-x_beta).exp()));
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// Tangent(x) = tan(x)
template <typename T>
struct TanFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
// Note(GGBond8488): Since Eigen3.3, Behavior like {A = (B * A).cwiseAbs()}
// will give wrong result, details see
// http://eigen.tuxfamily.org/dox/group__TopicAliasing.html
out.device(d) = x.unaryExpr(Tangent<T>()).eval();
}
};
template <typename T>
struct Sinh {
HOSTDEVICE T operator()(const T& val) const { return sinh(val); }
};
template <>
struct Sinh<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(sinhf(static_cast<float>(val)));
}
};
template <typename T>
struct Cosh {
HOSTDEVICE T operator()(const T& val) const { return cosh(val); }
};
template <>
struct Cosh<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(coshf(static_cast<float>(val)));
}
};
// sinh(x) = sinh(x)
template <typename T>
struct SinhFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Sinh<T>()).eval();
}
};
// cosh(x) = cosh(x)
template <typename T>
struct CoshFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Cosh<T>()).eval();
}
};
// sinh'(x) = cosh(x)
template <typename T>
struct SinhGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * x.unaryExpr(Cosh<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct SinhGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * x.unaryExpr(Cosh<ComplexType<T>>()).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
// cosh'(x) = sinh(x)
template <typename T>
struct CoshGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * x.unaryExpr(Sinh<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct CoshGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * x.unaryExpr(Sinh<ComplexType<T>>()).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct Acos {
HOSTDEVICE T operator()(const T& val) const { return acos(val); }
};
template <>
struct Acos<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(acos(static_cast<float>(val)));
}
};
// Acos(x) = acos(x)
template <typename T>
struct AcosFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Acos<T>()).eval();
}
};
// acos'(x) = -1/sqrt(1-x^2)
template <typename T>
struct AcosGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
-dout * static_cast<T>(1) / (static_cast<T>(1) - x.square()).sqrt();
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct AcosGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
-dout * (static_cast<ComplexType<T>>(1) /
(static_cast<ComplexType<T>>(1) - x.square()).sqrt())
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct Asin {
HOSTDEVICE T operator()(const T& val) const { return asin(val); }
};
template <>
struct Asin<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(asin(static_cast<float>(val)));
}
};
// Asin(x) = asin(x)
template <typename T>
struct AsinFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Asin<T>()).eval();
}
};
// asin'(x) = 1/sqrt(1-x^2)
template <typename T>
struct AsinGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * static_cast<T>(1) / (static_cast<T>(1) - x.square()).sqrt();
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct AsinGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * (static_cast<ComplexType<T>>(1) /
(static_cast<ComplexType<T>>(1) - x.square()).sqrt())
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct Atan {
HOSTDEVICE T operator()(const T& val) const { return atan(val); }
};
template <>
struct Atan<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(atan(static_cast<float>(val)));
}
};
// Atan(x) = atan(x)
template <typename T>
struct AtanFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Atan<T>()).eval();
}
};
// atan'(x) = 1 / (1 + x^2)
template <typename T>
struct AtanGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * static_cast<T>(1) / (static_cast<T>(1) + x.square());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct AtanGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * (static_cast<ComplexType<T>>(1) /
(static_cast<ComplexType<T>>(1) + x.square()))
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct LogitGradFunctor {
template <typename Device, typename X, typename dOut, typename dX, typename P>
void operator()(Device d, X x, dOut dout, dX dx, P p, double eps) const {
// logit(x)' = 1/(x*(1-x))
if (!eps) {
dx.device(d) = (x < static_cast<T>(0.0) || x > static_cast<T>(1.0))
.select(p.constant(static_cast<T>(NAN)),
dout * (static_cast<T>(1) /
((static_cast<T>(1) - x) * x)));
} else {
dx.device(d) = (x < static_cast<T>(eps) || x > static_cast<T>(1.0 - eps))
.select(p.constant(static_cast<T>(0)),
dout * (static_cast<T>(1) /
((static_cast<T>(1) - x) * x)));
}
}
};
template <typename T>
struct Acosh {
HOSTDEVICE T operator()(const T& val) const { return acosh(val); }
};
template <>
struct Acosh<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(acosh(static_cast<float>(val)));
}
};
// Acosh(x) = acosh(x)
template <typename T>
struct AcoshFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Acosh<T>()).eval();
}
};
// acosh'(x) = 1/sqrt(x^2 - 1)
template <typename T>
struct AcoshGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * static_cast<T>(1) / (x * x - static_cast<T>(1)).sqrt();
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct AcoshGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * (static_cast<ComplexType<T>>(1) /
(-static_cast<ComplexType<T>>(1) + x.square()).sqrt())
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct Asinh {
HOSTDEVICE T operator()(const T& val) const { return asinh(val); }
};
template <>
struct Asinh<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(asinh(static_cast<float>(val)));
}
};
// Asinh(x) = asinh(x)
template <typename T>
struct AsinhFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Asinh<T>()).eval();
}
};
// asinh'(x) = 1/sqrt(x^2 + 1)
template <typename T>
struct AsinhGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * static_cast<T>(1) / (x.square() + static_cast<T>(1)).sqrt();
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct AsinhGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * (static_cast<ComplexType<T>>(1) /
(x.square() + static_cast<ComplexType<T>>(1)).sqrt())
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct Atanh {
HOSTDEVICE T operator()(const T& val) const { return atanh(val); }
};
template <>
struct Atanh<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(atanh(static_cast<float>(val)));
}
};
// Atanh(x) = atanh(x)
template <typename T>
struct AtanhFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Atanh<T>()).eval();
}
};
// atanh'(x) = 1/(1 - x^2)
template <typename T>
struct AtanhGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * static_cast<T>(1) / (static_cast<T>(1) - x.square());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
template <typename T>
struct AtanhGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * (static_cast<ComplexType<T>>(1) /
(static_cast<ComplexType<T>>(1) - x.square()))
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepX; }
};
// exp functor
// exp(x) = e^x
template <typename T>
struct ExpFunctor : public BaseActivationFunctor<T> {
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.template cast<U>().exp();
}
};
template <typename T>
struct ExpGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = dout * out;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct ExpGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = dout * out.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct Expm1 {};
template <typename T>
struct Expm1<ComplexType<T>> {
HOSTDEVICE ComplexType<T> operator()(const ComplexType<T>& val) const {
return exp(val) - static_cast<ComplexType<T>>(1);
}
};
// expm1(x) = e^x - 1
template <typename T>
struct Expm1Functor : public BaseActivationFunctor<T> {
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.template cast<U>().expm1();
}
};
template <typename T>
struct Expm1Functor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr(Expm1<ComplexType<T>>()).eval();
}
};
template <typename T>
struct Expm1GradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = dout * out + dout;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct Expm1GradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = dout * out.unaryExpr(Conj<T>()) + dout;
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepOut; }
};
// relu(x) = max(x, 0)
template <typename T>
struct ReluCPUFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr([] HOSTDEVICE(T v) {
return v > static_cast<T>(0) ? v : static_cast<T>(0);
});
}
};
template <typename T>
struct ReluCUDAFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.cwiseMax(static_cast<T>(0));
}
};
template <typename T>
struct ReluGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = dout * (out > static_cast<T>(0)).template cast<T>();
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct ReluGradGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X UNUSED,
const DenseTensor* Out,
const DenseTensor* ddX,
DenseTensor* ddOut,
DenseTensor* dOut UNUSED,
DenseTensor* dX UNUSED) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "ReluGradGrad"));
auto out = EigenVector<T>::Flatten(
GET_DATA_SAFELY(Out, "Output", "Out", "ReluGradGrad"));
if (ddOut) {
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "ReluGradGrad"));
ddout.device(*d) = ddx * (out > static_cast<T>(0)).template cast<T>();
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
// tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))
template <typename T>
struct TanhFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.tanh();
}
};
template <typename T>
struct TanhGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = dout * (static_cast<T>(1) - out * out);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct TanhGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
// auto dx_ = static_cast<ComplexType<T>>(1) - out * out;
// ComplexType<T> dx_conj_(dx_.real, -dx_.imag);
// dx.device(d) = dout * dx_conj_;
dx.device(d) =
dout *
(static_cast<ComplexType<T>>(1) - out * out).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct TanhGradGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* Out,
const DenseTensor* ddX,
const DenseTensor* dOut,
DenseTensor* dOutNew,
DenseTensor* ddOut) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "TanhGradGrad"));
auto out = EigenVector<T>::Flatten(
GET_DATA_SAFELY(Out, "Input", "Out", "TanhGradGrad"));
// tanh grad grad : ddout = (1 - out^2) * ddx, dout = - (dout_old * 2 * out
// * ddx)
if (dOutNew) {
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "DOut", "TanhGradGrad"));
auto dout_new = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOutNew, "Output", "DOutNew", "TanhGradGrad"));
dout_new.device(*d) =
static_cast<T>(-1) * dout * static_cast<T>(2) * out * ddx;
}
if (ddOut) {
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "TanhGradGrad"));
ddout.device(*d) = (static_cast<T>(1) - out * out) * ddx;
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
/*
Out
DOut D_Dout
DDx -> TanhTripleGrad -> D_DDx
D_DDout d_OutNew
D_Dout_new
D_Dout = (-2) * Out * DDx * D_Dout_new
D_DDx = (1-Out^2)*D_DDout + (-2) * Out * DOut * D_Dout_new
D_OutNew = (-2) * Out * DDx * D_DDout + (-2) * DOut * DDx * D_Dout_new
Out, DDX, DOut, D_DDOut, D_DOut_New // input
D_OutNew, D_DOut, D_DDx // output
*/
template <typename T>
struct TanhTripleGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* Out,
const DenseTensor* ddX,
const DenseTensor* dOut,
const DenseTensor* d_DDOut,
const DenseTensor* d_dOut_New,
DenseTensor* d_d_Out,
DenseTensor* d_Out_New,
DenseTensor* d_DDx) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "TanhTripleGrad"));
auto out = EigenVector<T>::Flatten(
GET_DATA_SAFELY(Out, "Input", "Out", "TanhTripleGrad"));
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "DOut", "TanhTripleGrad"));
if (d_Out_New) {
auto d_OutNew = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_Out_New, "Output", "D_OutNew", "TanhTripleGrad"));
if (d_DDOut && d_dOut_New) {
auto d_ddOut = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "D_DDOut", "TanhTripleGrad"));
auto d_dOutNew = EigenVector<T>::Flatten(GET_DATA_SAFELY(
d_dOut_New, "Input", "D_DOut_New", "TanhTripleGrad"));
d_OutNew.device(*d) = (static_cast<T>(-2) * out * ddx * d_ddOut) -
(static_cast<T>(2) * dout * ddx * d_dOutNew);
} else if (d_DDOut && !d_dOut_New) {
auto d_ddOut = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "D_DDOut", "TanhTripleGrad"));
d_OutNew.device(*d) = (static_cast<T>(-2) * out * ddx * d_ddOut);
} else if (!d_DDOut && d_dOut_New) {
auto d_dOutNew = EigenVector<T>::Flatten(GET_DATA_SAFELY(
d_dOut_New, "Input", "D_DOut_New", "TanhTripleGrad"));
d_OutNew.device(*d) = -(static_cast<T>(2) * dout * ddx * d_dOutNew);
} else {
d_OutNew.device(*d) = static_cast<T>(0) * out;
}
}
if (d_d_Out) {
auto d_dOut = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_d_Out, "Output", "D_DOut", "TanhTripleGrad"));
if (d_dOut_New) {
auto d_dOutNew = EigenVector<T>::Flatten(GET_DATA_SAFELY(
d_dOut_New, "Input", "D_DOut_New", "TanhTripleGrad"));
d_dOut.device(*d) = static_cast<T>(-2) * out * ddx * d_dOutNew;
} else {
d_dOut.device(*d) = static_cast<T>(0) * out;
}
}
if (d_DDx) {
auto d_ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDx, "Output", "D_DDx", "TanhTripleGrad"));
if (d_DDOut && d_dOut_New) {
auto d_ddOut = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "D_DDOut", "TanhTripleGrad"));
auto d_dOutNew = EigenVector<T>::Flatten(GET_DATA_SAFELY(
d_dOut_New, "Input", "D_DOut_New", "TanhTripleGrad"));
d_ddx.device(*d) = (static_cast<T>(1) - (out * out)) * d_ddOut -
static_cast<T>(2) * out * dout * d_dOutNew;
} else if (d_DDOut && !d_dOut_New) {
auto d_ddOut = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "D_DDOut", "TanhTripleGrad"));
d_ddx.device(*d) = (static_cast<T>(1) - (out * out)) * d_ddOut;
} else if (!d_DDOut && d_dOut_New) {
auto d_dOutNew = EigenVector<T>::Flatten(GET_DATA_SAFELY(
d_dOut_New, "Input", "D_DOut_New", "TanhTripleGrad"));
d_ddx.device(*d) = -static_cast<T>(2) * out * dout * d_dOutNew;
} else {
d_ddx.device(*d) = static_cast<T>(0) * ddx;
}
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct HardTanhFunctor : public BaseActivationFunctor<T> {
float t_min;
float t_max;
// NOTE: Explicit hides the `BaseActivationFunctor<T>::GetAttrs`
// not polymorphism for speed.
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"t_min", &t_min}, {"t_max", &t_max}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.cwiseMax(static_cast<T>(t_min))
.cwiseMin(static_cast<T>(t_max)); // NOLINT
}
};
template <typename T>
struct HardTanhGradFunctor : public BaseActivationFunctor<T> {
float t_min;
float t_max;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"t_min", &t_min}, {"t_max", &t_max}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * ((x > static_cast<T>(t_min)) *
(x < static_cast<T>(t_max))) // NOLINT
.template cast<T>();
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct LeakyReluFunctor : public BaseActivationFunctor<T, double> {
double alpha;
typename BaseActivationFunctor<T, double>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
if (alpha < 1.f) { // NOLINT
out.device(d) = x.cwiseMax(static_cast<T>(alpha) * x);
} else {
out.device(d) = x.cwiseMin(static_cast<T>(alpha) * x);
}
}
};
template <typename T>
struct LeakyReluGradFunctor : public BaseActivationFunctor<T, double> {
double alpha;
typename BaseActivationFunctor<T, double>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto temp1 = static_cast<T>(alpha) *
(x < static_cast<T>(0)).template cast<T>(); // NOLINT
auto temp2 = (x >= static_cast<T>(0)).template cast<T>();
dx.device(d) = dout * (temp1 + temp2).template cast<T>();
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct LeakyReluGradGradFunctor : public BaseActivationFunctor<T, double> {
double alpha;
typename BaseActivationFunctor<T, double>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* Out UNUSED,
const DenseTensor* ddX,
DenseTensor* ddOut,
DenseTensor* dOut UNUSED,
DenseTensor* dX UNUSED) const {
if (ddOut) {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "LeakyReluGradGrad"));
auto x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(X, "Input", "X", "LeakyReluGradGrad"));
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DOut", "LeakyReluGradGrad"));
ddout.device(*d) = ddx * ((x > static_cast<T>(0)).template cast<T>() +
static_cast<T>(alpha) *
(x <= static_cast<T>(0)).template cast<T>())
.template cast<T>();
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct ThresholdedReluFunctor : public BaseActivationFunctor<T> {
float threshold;
float value;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}, {"value", &value}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
auto th = static_cast<T>(threshold); // NOLINT
out.device(d) = (x > th).template cast<T>() * x +
(x <= th).template cast<T>() * static_cast<T>(value);
}
};
template <typename T>
struct ThresholdedReluGradFunctor : public BaseActivationFunctor<T> {
float threshold;
float value;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}, {"value", &value}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto th = static_cast<T>(threshold); // NOLINT
dx.device(d) = dout * (x > th).template cast<T>();
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// relu6(x) = min(max(0, x), 6)
template <typename T>
struct Relu6Functor : public BaseActivationFunctor<T> {
float threshold;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.cwiseMax(static_cast<T>(0))
.cwiseMin(static_cast<T>(threshold)); // NOLINT
}
};
template <typename T>
struct Relu6GradFunctor : public BaseActivationFunctor<T> {
typename BaseActivationFunctor<T>::AttrPair GetAttrs() { return {{}}; }
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
float threshold = 6;
dx.device(d) =
dout * ((out > static_cast<T>(0)) * (out < static_cast<T>(threshold)))
.template cast<T>();
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
// tanhshrink(x) = x - tanh(x)
// where tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))
template <typename T>
struct TanhShrinkFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x - x.tanh();
}
};
template <typename T>
struct TanhShrinkGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * (x.tanh() * x.tanh());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// tanhshrink(x) = x - tanh(x)
// where tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))
template <typename T>
struct HardShrinkFunctor : public BaseActivationFunctor<T> {
float threshold;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
auto temp1 = x <= static_cast<T>(threshold * -1.f); // NOLINT
auto temp2 = x >= static_cast<T>(threshold);
out.device(d) = x * (temp1 || temp2).template cast<T>();
}
};
template <typename T>
struct HardShrinkGradFunctor : public BaseActivationFunctor<T> {
float threshold;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto temp1 = x <= static_cast<T>(threshold * -1.f); // NOLINT
auto temp2 = x >= static_cast<T>(threshold);
dx.device(d) = dout * (temp1 || temp2).template cast<T>();
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// softshrink(x) = x - lambda, if x > lambda; x + lambda, if x < -lambda; 0
// otherwise
template <typename T>
struct SoftShrinkFunctor : public BaseActivationFunctor<T> {
float lambda;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"lambda", &lambda}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
auto lambdaT = static_cast<T>(lambda); // NOLINT
auto temp1 = (x > lambdaT).template cast<T>();
auto temp2 = (x < -lambdaT).template cast<T>();
out.device(d) = temp1 * (x - lambdaT) + temp2 * (x + lambdaT);
}
};
template <typename T>
struct SoftShrinkGradFunctor : public BaseActivationFunctor<T> {
float lambda;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"lambda", &lambda}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto lambdaT = static_cast<T>(lambda); // NOLINT
auto temp1 = (x > lambdaT).template cast<T>();
auto temp2 = (x < -lambdaT).template cast<T>();
dx.device(d) = dout * (temp1 + temp2).template cast<T>();
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct ELUFunctor : public BaseActivationFunctor<T> {
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) =
(x < static_cast<T>(0))
.select(static_cast<T>(alpha) * (x.exp() - static_cast<T>(1)),
x); // NOLINT
}
};
template <typename T>
struct ELUGradFunctor : public BaseActivationFunctor<T> {
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
// case 1: alpha >= 0
// dx = dout, if out > 0
// dx = dout * (out + alpha), if out <= 0
dx.device(d) = (out > static_cast<T>(0))
.select(dout, dout * (out + static_cast<T>(alpha)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct ELUGradNegativeAlphaFunctor : public BaseActivationFunctor<T> {
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
// case 2: alpha < 0
// dx = dout, if x > 0
// dx = dout * (out + alpha), if x <=0
dx.device(d) = (x > static_cast<T>(0))
.select(dout, dout * static_cast<T>(alpha) * x.exp());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct ELUGradGradFunctor : public BaseActivationFunctor<T> {
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* ddX,
DenseTensor* ddOut,
const DenseTensor* dOut,
DenseTensor* dX) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "ELUGradGrad"));
auto x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(X, "Input", "X", "ELUGradGrad"));
if (dX) {
auto dx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dX, "Output", "DX", "ELUGradGrad"));
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Output", "DOut", "ELUGradGrad"));
dx.device(*d) = ddx * dout * static_cast<T>(alpha) * x.exp() *
(x <= static_cast<T>(0)).template cast<T>();
}
if (ddOut) {
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "ELUGradGrad"));
ddout.device(*d) = ddx * ((x > static_cast<T>(0)).template cast<T>() +
static_cast<T>(alpha) * x.exp() *
(x <= static_cast<T>(0)).template cast<T>())
.template cast<T>();
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// silu(x) = x / (1 + exp(-x))
template <typename T>
struct SiluFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
auto temp = static_cast<T>(1) / (static_cast<T>(1) + (-x).exp());
out.device(d) = x * temp;
}
};
// silu'(x) = (1 / (1 + e^{-x})) * (1 + out * e^{-x}))
template <typename T>
struct SiluGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto temp1 = static_cast<T>(1) + (-x).exp(); // 1+e^(-x)
auto temp2 = x * (-x).exp(); // x*e^(-x)
dx.device(d) = dout * ((static_cast<T>(1) / temp1) *
(static_cast<T>(1) + (temp2 / temp1)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct SiluGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto temp1 = static_cast<ComplexType<T>>(1) + (-x).exp(); // 1+e^(-x)
auto temp2 = x * (-x).exp(); // x*e^(-x)
dx.device(d) = dout * ((static_cast<ComplexType<T>>(1) / temp1) *
(static_cast<ComplexType<T>>(1) + (temp2 / temp1)))
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct SoftsignFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x / (static_cast<T>(1) + x.abs());
}
};
// d(softsign(x))/dx = 1 / (1 + |x|)^2
// Taken from https://en.wikipedia.org/wiki/Activation_function
template <typename T>
struct SoftsignGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * (static_cast<T>(1) / (static_cast<T>(1) + x.abs()).square());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct SoftsignGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
auto temp = (-x / (one + x.abs()).square()).unaryExpr(Real<T>());
dx.device(d) = dout * (one / (one + x.abs()) + temp * x / x.abs());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// sigmoid(x) = 1 / (1 + exp(-x))
template <typename T>
struct SigmoidFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = static_cast<T>(1) / (static_cast<T>(1) + (-x).exp());
}
};
template <typename T>
struct SigmoidGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = dout * out * (static_cast<T>(1) - out);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct SigmoidGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
ComplexType<T> one = static_cast<ComplexType<T>>(1);
dx.device(d) = dout * (out * (one - out)).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
/*
Out
DOut -> SigmoidGradGrad -> DOutNew
DDX DDOut
DDOut = (1-Out)*Out*DDX
DOutNew = (1-2*Out)*DOut*DDX
*/
template <typename T>
struct SigmoidGradGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* Out,
const DenseTensor* ddX,
const DenseTensor* dOut,
DenseTensor* dOutNew,
DenseTensor* ddOut) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "SigmoidGradGrad"));
auto out = EigenVector<T>::Flatten(
GET_DATA_SAFELY(Out, "Input", "Out", "SigmoidGradGrad"));
if (dOutNew) {
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "DOut", "SigmoidGradGrad"));
auto dout_new = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOutNew, "Output", "DOutNew", "SigmoidGradGrad"));
dout_new.device(*d) =
(static_cast<T>(1) - static_cast<T>(2) * out) * dout * ddx;
}
if (ddOut) {
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "SigmoidGradGrad"));
ddout.device(*d) = (static_cast<T>(1) - out) * out * ddx;
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
/*
Out
DOut D_Dout
DDx -> SigmoidTripleGrad -> D_DDx
D_DDout d_OutNew
D_Dout_new
D_Dout = (1-2*Out)*DDx*D_Dout_new
D_DDx = (1-Out)*Out*D_DDout + (1-2*Out)*DOut*D_Dout_new
D_OutNew = (DDx-2*Out*DDx)*D_DDout - 2*DOut*DDx*D_Dout_new
Out, DDX, DOut, D_DDOut, D_DOut_New // input
D_OutNew, D_DOut, D_DDx // output
*/
template <typename T>
struct SigmoidTripleGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* Out,
const DenseTensor* ddX,
const DenseTensor* dOut,
const DenseTensor* d_DDOut,
const DenseTensor* d_dOut_New,
DenseTensor* d_d_Out,
DenseTensor* d_Out_New,
DenseTensor* d_DDx) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "SigmoidTripleGrad"));
auto out = EigenVector<T>::Flatten(
GET_DATA_SAFELY(Out, "Input", "Out", "SigmoidTripleGrad"));
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Input", "DOut", "SigmoidTripleGrad"));
auto d_dOutNew = EigenVector<T>::Flatten(GET_DATA_SAFELY(
d_dOut_New, "Input", "D_DOut_New", "SigmoidTripleGrad"));
if (d_Out_New) {
auto d_OutNew = EigenVector<T>::Flatten(GET_DATA_SAFELY(
d_Out_New, "Output", "D_OutNew", "SigmoidTripleGrad"));
d_OutNew.device(*d) = -static_cast<T>(2) * dout * ddx * d_dOutNew;
if (d_DDOut) {
auto d_ddOut = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "D_DDOut", "SigmoidTripleGrad"));
d_OutNew.device(*d) =
(ddx - static_cast<T>(2) * out * ddx) * d_ddOut + d_OutNew;
}
}
if (d_d_Out) {
auto d_dOut = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_d_Out, "Output", "D_DOut", "SigmoidTripleGrad"));
d_dOut.device(*d) =
(static_cast<T>(1) - static_cast<T>(2) * out) * ddx * d_dOutNew;
}
if (d_DDx) {
auto d_ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDx, "Output", "D_DDx", "SigmoidTripleGrad"));
d_ddx.device(*d) =
(static_cast<T>(1) - static_cast<T>(2) * out) * dout * d_dOutNew;
if (d_DDOut) {
auto d_ddOut = EigenVector<T>::Flatten(
GET_DATA_SAFELY(d_DDOut, "Input", "D_DDOut", "SigmoidTripleGrad"));
d_ddx.device(*d) = d_ddx + (static_cast<T>(1) - out) * out * d_ddOut;
}
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
// Originally: logsigmoid(x) = -log (1 + exp(-x))
// For numerical stability, we can use the log-sum-exp trick:
// https://hips.seas.harvard.edu/blog/2013/01/09/computing-log-sum-exp/
// We can rewrite the above equation as:
// out = -log( exp(0) + exp(-x)) [since exp(0) = 1]
// = -log( exp(max(-x, 0) - max(-x, 0)) + exp(-x + max(-x, 0) - max(-x, 0)))
// = -log( exp(max(-x, 0)) * exp(-max(-x, 0)) - exp(max(-x, 0)) * exp(-x -
// max(-x, 0)))
// = -log( exp(max(-x, 0)) * (exp(-max(-x, 0)) + exp(-x - max(-x, 0))))
// = -log( exp(max(-x, 0)) - log(exp(-max(-x, 0)) + exp(-x - max(-x, 0)))
//
// Hence, logsigmoid(x) = - (max(-x, 0) + log(exp(-max(-x, 0))
// + exp(-x - max(-x, 0))))
template <typename T>
struct LogSigmoidFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
auto temp = (-x).cwiseMax(static_cast<T>(0)); // temp = max(-x, 0)
out.device(d) = -temp - (((-temp).exp() + (-x - temp).exp()).log());
}
};
// Specialized implementation for complex numbers
template <typename T>
struct LogSigmoidFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
// For complex numbers, use log σ(x) = -log(1 + exp(-x))
ComplexType<T> one = ComplexType<T>(T(1), T(0));
// Cache exp(-x) to avoid redundant computation
auto exp_neg_x = (-x).exp();
out.device(d) = -(one + exp_neg_x).log();
}
};
// Originally: f' = exp(-x) / (1 + exp(-x))
// For numerical stability: f' = exp(-x - max(-x, 0)) / (exp(-max(-x, 0)) +
// exp(-x - max(-x, 0)))
template <typename T>
struct LogSigmoidGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto temp = (-x).cwiseMax(static_cast<T>(0)); // temp = max(-x, 0)
dx.device(d) =
dout * ((-x - temp).exp() / ((-temp).exp() + (-x - temp).exp()));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct LogSigmoidGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
// For complex numbers, use the direct formula:
// d/dx log(1/(1+exp(-x))) = exp(-x)/(1+exp(-x))
ComplexType<T> one = ComplexType<T>(T(1), T(0));
// Cache exp(-x) to avoid redundant computation
auto exp_neg_x = (-x).exp();
dx.device(d) = dout * (exp_neg_x / (one + exp_neg_x)).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct HardSigmoidFunctor : public BaseActivationFunctor<T> {
float slope;
float offset;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"slope", &slope}, {"offset", &offset}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
auto temp = x * static_cast<T>(slope) + static_cast<T>(offset); // NOLINT
out.device(d) =
temp.cwiseMax(static_cast<T>(0)).cwiseMin(static_cast<T>(1));
}
};
template <typename T>
struct HardSigmoidGradFunctor : public BaseActivationFunctor<T> {
float slope;
float offset;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"slope", &slope}, {"offset", &offset}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
dx.device(d) = dout * // NOLINT
((out > static_cast<T>(0)) * (out < static_cast<T>(1)))
.template cast<T>() *
static_cast<T>(slope);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct Log {
HOSTDEVICE T operator()(const T& val) const { return std::log(val); }
};
template <typename T>
struct Log<ComplexType<T>> {
HOSTDEVICE ComplexType<T> operator()(const ComplexType<T>& val) const {
return ComplexType<T>(std::log(std::complex<T>(val)));
}
};
template <>
struct Log<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(std::log(static_cast<float>(val)));
}
};
template <>
struct Log<bfloat16> {
HOSTDEVICE bfloat16 operator()(const bfloat16& val) const {
return bfloat16(std::log(static_cast<float>(val)));
}
};
// log(x) = natural logarithm of x
template <typename T>
struct LogFunctor : public BaseActivationFunctor<T> {
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.template cast<U>().unaryExpr(Log<U>()).eval();
}
};
template <typename T>
struct LogGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * (static_cast<T>(1) / x);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct LogGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * (static_cast<ComplexType<T>>(1) / x).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct Log2 {
HOSTDEVICE T operator()(const T& val) const { return std::log2(val); }
};
template <typename T>
struct Log2<ComplexType<T>> {
HOSTDEVICE ComplexType<T> operator()(const ComplexType<T>& val) const {
return ComplexType<T>(std::log(std::complex<T>(val)) /
std::log(std::complex<T>(2)));
}
};
template <>
struct Log2<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(std::log2(static_cast<float>(val)));
}
};
template <>
struct Log2<bfloat16> {
HOSTDEVICE bfloat16 operator()(const bfloat16& val) const {
return bfloat16(std::log2(static_cast<float>(val)));
}
};
// log2(x) = logarithm to the base 2 of the elements of x
template <typename T>
struct Log2Functor : public BaseActivationFunctor<T> {
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.template cast<U>().unaryExpr(Log2<U>()).eval();
}
};
// the gradient of log2(x) is 1/(x*ln(2))
template <typename T>
struct Log2GradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * static_cast<T>(1) / (x * static_cast<T>(log(2)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct Log2GradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * (static_cast<ComplexType<T>>(1) /
(x * static_cast<ComplexType<T>>(log(2))))
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct Log10 {
HOSTDEVICE T operator()(const T& val) const { return std::log10(val); }
};
template <typename T>
struct Log10<ComplexType<T>> {
HOSTDEVICE ComplexType<T> operator()(const ComplexType<T>& val) const {
return ComplexType<T>(std::log10(std::complex<T>(val)));
}
};
template <>
struct Log10<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(std::log10(static_cast<float>(val)));
}
};
template <>
struct Log10<bfloat16> {
HOSTDEVICE bfloat16 operator()(const bfloat16& val) const {
return bfloat16(std::log10(static_cast<float>(val)));
}
};
// log10(x) = logarithm to the base 10 of the elements of x
template <typename T>
struct Log10Functor : public BaseActivationFunctor<T> {
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.template cast<U>().unaryExpr(Log10<U>()).eval();
}
};
// the gradient of log10(x) is 1/(x*ln(10))
// PyTorch formula: grad / (self * 2.3025850929940456)
template <typename T>
struct Log10GradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
// Use PyTorch's exact constant (ln(10) to 16 significant digits) and
// matching evaluation order: dout / (x * ln10), i.e., multiply x by the
// constant first, then divide. This avoids runtime log(10) computation
// and aligns CPU/GPU paths with PyTorch's backward for bit-exact results.
T log_ten = static_cast<T>(2.3025850929940456);
dx.device(d) = dout / (x * log_ten);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct Log10GradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * (static_cast<ComplexType<T>>(1) /
(x * static_cast<ComplexType<T>>(log(10))))
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct Log1p {
HOSTDEVICE T operator()(const T& val) const { return std::log1p(val); }
};
template <typename T>
struct Log1p<ComplexType<T>> {
HOSTDEVICE ComplexType<T> operator()(const ComplexType<T>& val) const {
return ComplexType<T>(std::log(std::complex<T>(1) + std::complex<T>(val)));
}
};
template <>
struct Log1p<float16> {
HOSTDEVICE float16 operator()(const float16& val) const {
return float16(std::log1p(static_cast<float>(val)));
}
};
template <>
struct Log1p<bfloat16> {
HOSTDEVICE bfloat16 operator()(const bfloat16& val) const {
return bfloat16(std::log1p(static_cast<float>(val)));
}
};
// log1p(x) = natural logarithm of x+1
template <typename T>
struct Log1pFunctor : public BaseActivationFunctor<T> {
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.template cast<U>().unaryExpr(Log1p<U>()).eval();
}
};
template <typename T>
struct Log1pGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * (static_cast<T>(1) / (x + static_cast<T>(1)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct Log1pGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * (static_cast<ComplexType<T>>(1) /
(x + static_cast<ComplexType<T>>(1)))
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct LogGradGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* ddX,
DenseTensor* ddOut,
const DenseTensor* dOut,
DenseTensor* dX) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "LogGradGrad"));
auto x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(X, "Input", "X", "LogGradGrad"));
// ddout = ddx / x; dx = -(dout / x) * (ddx / x)
// calculate dx first, so ddout can inplace ddx
if (dX) {
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Output", "DOut", "LogGradGrad"));
auto dx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dX, "Output", "DX", "LogGradGrad"));
dx.device(*d) = dout * static_cast<T>(-1) * ddx / (x * x);
}
if (ddOut) {
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "LogGradGrad"));
ddout.device(*d) = ddx * static_cast<T>(1) / x;
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct LogGradGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* ddX,
DenseTensor* ddOut,
const DenseTensor* dOut,
DenseTensor* dX) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<ComplexType<T>>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "LogGradGrad"));
auto x = EigenVector<ComplexType<T>>::Flatten(
GET_DATA_SAFELY(X, "Input", "X", "LogGradGrad"));
// ddout = ddx / x; dx = -(dout / x) * (ddx / x)
// calculate dx first, so ddout can inplace ddx
if (dX) {
auto dout = EigenVector<ComplexType<T>>::Flatten(
GET_DATA_SAFELY(dOut, "Output", "DOut", "LogGradGrad"));
auto dx = EigenVector<ComplexType<T>>::Flatten(
GET_DATA_SAFELY(dX, "Output", "DX", "LogGradGrad"));
dx.device(*d) = dout * static_cast<ComplexType<T>>(-1) * ddx /
(x * x).unaryExpr(Conj<T>());
}
if (ddOut) {
auto ddout = EigenVector<ComplexType<T>>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "LogGradGrad"));
ddout.device(*d) =
ddx * static_cast<ComplexType<T>>(1) / x.unaryExpr(Conj<T>());
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// HardSwish = min(max(0, x+3), 6) * x / 6
template <typename T>
struct HardSwishFunctor : public BaseActivationFunctor<T> {
float threshold;
float scale;
float offset;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}, {"scale", &scale}, {"offset", &offset}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = (x + static_cast<T>(offset)) // NOLINT
.cwiseMax(static_cast<T>(0))
.cwiseMin(static_cast<T>(threshold)) *
x / static_cast<T>(scale);
}
};
template <typename T>
struct HardSwishGradFunctor : public BaseActivationFunctor<T> {
float threshold;
float scale;
float offset;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}, {"scale", &scale}, {"offset", &offset}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto tmp =
((x + static_cast<T>(offset)) < static_cast<T>(threshold)) // NOLINT
.template cast<T>();
dx.device(d) =
dout *
(((x + static_cast<T>(offset)) > static_cast<T>(0)).template cast<T>() *
(static_cast<T>(2) * x + static_cast<T>(offset)) /
static_cast<T>(scale) * tmp +
static_cast<T>(1) * (static_cast<T>(1) - tmp));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct HardSwishGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
float threshold;
float scale;
float offset;
typename BaseActivationFunctor<ComplexType<T>>::AttrPair GetAttrs() {
return {{"threshold", &threshold}, {"scale", &scale}, {"offset", &offset}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto offset_t = static_cast<ComplexType<T>>(offset);
auto threshold_t = static_cast<ComplexType<T>>(threshold);
auto one = static_cast<ComplexType<T>>(1);
auto zero = static_cast<ComplexType<T>>(0);
auto two = static_cast<ComplexType<T>>(2);
auto scale_t = static_cast<ComplexType<T>>(scale);
auto tmp1 = ((x + offset_t) < threshold_t) // NOLINT
.template cast<ComplexType<T>>();
auto tmp2 = ((x + offset_t) > zero).template cast<ComplexType<T>>();
// dx = 0, when x <= -offset
// dout , when x >= threshold - offset
// dout * (2 * x / scale + offset / scale), otherwise
// threshold = scale = 6, offset = 3 by default
dx.device(d) = dout * (tmp2 * (two * x + offset_t) / scale_t * tmp1 +
one * (one - tmp1))
.unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct SwishFunctor : public BaseActivationFunctor<T> {
float beta;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"beta", &beta}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) =
x / (static_cast<T>(1) + (static_cast<T>(-beta) * x).exp()); // NOLINT
}
};
template <typename T>
struct SwishGradFunctor : public BaseActivationFunctor<T> {
typename BaseActivationFunctor<T>::AttrPair GetAttrs() { return {{}}; }
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out fake_out UNUSED, dOut dout, dX dx) const {
float beta = 1.0;
auto temp1 = static_cast<T>(1) /
(static_cast<T>(1) + (static_cast<T>(-beta) * x).exp());
auto out = x * temp1;
auto temp2 = temp1 * (static_cast<T>(1) - (static_cast<T>(beta) * out));
dx.device(d) = dout * ((static_cast<T>(beta) * out) + temp2);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// FIXME(qijun) https://github.com/PaddlePaddle/Paddle/issues/5198
template <typename T>
struct PowFunctor : public BaseActivationFunctor<T> {
float factor;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"factor", &factor}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.pow(static_cast<T>(factor)); // NOLINT
}
};
template <typename T>
struct PowFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
float factor;
typename BaseActivationFunctor<ComplexType<T>>::AttrPair GetAttrs() {
return {{"factor", &factor}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.pow(static_cast<ComplexType<T>>(factor)); // NOLINT
}
};
template <typename T>
struct PowGradFunctor : public BaseActivationFunctor<T> {
float factor;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"factor", &factor}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) = dout * static_cast<T>(factor) *
x.pow(static_cast<T>(factor) - static_cast<T>(1));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct PowGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
float factor;
typename BaseActivationFunctor<ComplexType<T>>::AttrPair GetAttrs() {
return {{"factor", &factor}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
dx.device(d) =
dout * static_cast<ComplexType<T>>(factor) *
x.pow(static_cast<ComplexType<T>>(factor - 1)).unaryExpr(Conj<T>());
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// floor(x) = flooring(x)
template <typename T>
struct FloorFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
if constexpr ((std::is_same<T, uint8_t>::value) ||
(std::is_same<T, int8_t>::value) ||
(std::is_same<T, uint16_t>::value) ||
(std::is_same<T, int16_t>::value) ||
(std::is_same<T, int>::value) ||
(std::is_same<T, int64_t>::value)) {
out.device(d) = x;
} else {
out.device(d) = x.floor();
}
}
};
// rint(x) = [x]
template <typename T, typename Enable = void>
struct RintFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x.unaryExpr([](const T& val) {
return (std::isnan(val) || std::isinf(val)) ? val : std::rint(val);
});
}
};
template <typename T>
struct RintFunctor<T, std::enable_if_t<std::is_integral_v<T>>>
: public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x;
}
};
// round(x) = [x]
template <typename T, typename Enable = void>
struct RoundFunctor : public BaseActivationFunctor<T> {
int decimals;
std::vector<std::pair<const char*, int*>> GetAttrs() {
return {{"decimals", &decimals}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
if (decimals == 0) {
out.device(d) = x.unaryExpr([](const T& val) {
return (std::isnan(val) || std::isinf(val)) ? val : std::rint(val);
});
} else if (decimals > 0) {
auto ten_pow_decimals = static_cast<T>(std::pow(10, decimals));
out.device(d) = x.unaryExpr([ten_pow_decimals](const T& val) {
return (std::isnan(val) || std::isinf(val))
? val
: std::rint(val * ten_pow_decimals) / ten_pow_decimals;
});
} else {
auto ten_pow_decimals = static_cast<T>(std::pow(10, -decimals));
out.device(d) = x.unaryExpr([ten_pow_decimals](const T& val) {
return (std::isnan(val) || std::isinf(val))
? val
: std::rint(val / ten_pow_decimals) * ten_pow_decimals;
});
}
}
};
template <typename T>
struct RoundFunctor<T, std::enable_if_t<std::is_integral_v<T>>>
: public BaseActivationFunctor<T> {
int decimals;
std::vector<std::pair<const char*, int*>> GetAttrs() {
return {{"decimals", &decimals}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = x;
}
};
template <typename T>
struct RoundFunctor<phi::dtype::complex<T>>
: public BaseActivationFunctor<phi::dtype::complex<T>> {
int decimals;
std::vector<std::pair<const char*, int*>> GetAttrs() {
return {{"decimals", &decimals}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
using ComplexT = phi::dtype::complex<T>;
if (decimals == 0) {
out.device(d) = x.unaryExpr([](const ComplexT& c) {
T real = std::isnan(c.real) || std::isinf(c.real) ? c.real
: std::rint(c.real);
T imag = std::isnan(c.imag) || std::isinf(c.imag) ? c.imag
: std::rint(c.imag);
return ComplexT(real, imag);
});
} else if (decimals > 0) {
auto ten_pow_decimals = static_cast<T>(std::pow(10, decimals));
out.device(d) = x.unaryExpr([ten_pow_decimals](const ComplexT& c) {
T real = std::isnan(c.real) || std::isinf(c.real)
? c.real
: std::rint(c.real * ten_pow_decimals) / ten_pow_decimals;
T imag = std::isnan(c.imag) || std::isinf(c.imag)
? c.imag
: std::rint(c.imag * ten_pow_decimals) / ten_pow_decimals;
return ComplexT(real, imag);
});
} else {
auto ten_pow_decimals = static_cast<T>(std::pow(10, -decimals));
out.device(d) = x.unaryExpr([ten_pow_decimals](const ComplexT& c) {
T real = std::isnan(c.real) || std::isinf(c.real)
? c.real
: std::rint(c.real / ten_pow_decimals) * ten_pow_decimals;
T imag = std::isnan(c.imag) || std::isinf(c.imag)
? c.imag
: std::rint(c.imag / ten_pow_decimals) * ten_pow_decimals;
return ComplexT(real, imag);
});
}
}
};
// ceil(x) = ceiling(x)
template <typename T>
struct CeilFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
if constexpr ((std::is_same<T, uint8_t>::value) ||
(std::is_same<T, int8_t>::value) ||
(std::is_same<T, uint16_t>::value) ||
(std::is_same<T, int16_t>::value) ||
(std::is_same<T, int>::value) ||
(std::is_same<T, int64_t>::value)) {
out.device(d) = x;
} else {
out.device(d) = x.ceil();
}
}
};
template <typename T>
struct NegativeFunctor : public BaseActivationFunctor<T> {
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) = -x;
}
};
template <typename T>
struct ZeroGradFunctor : public BaseActivationFunctor<T> {
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(
Device d, X x UNUSED, Out out, dOut dout UNUSED, dX dx) const {
dx.device(d) = static_cast<T>(0) * out;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kNoDeps;
}
};
template <typename T>
struct SqrtGradGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* Out,
const DenseTensor* dX,
const DenseTensor* ddX,
DenseTensor* dOut,
DenseTensor* ddOut) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "SqrtGradGrad"));
auto out = EigenVector<T>::Flatten(
GET_DATA_SAFELY(Out, "Output", "Out", "SqrtGradGrad"));
// sqrt GradGrad: ddy = 0.5 * ddx / y, dy = -1 * dx * ddx
// calculate dy first, so ddy can inplace ddx
if (dOut) {
auto dx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dX, "Output", "DX", "SqrtGradGrad"));
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Output", "DOut", "SqrtGradGrad"));
dout.device(*d) = dx * ddx * static_cast<T>(-1) / out;
}
if (ddOut) {
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "SqrtGradGrad"));
ddout.device(*d) = ddx * static_cast<T>(0.5) / out;
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct RsqrtGradGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* Out,
const DenseTensor* dX,
const DenseTensor* ddX,
DenseTensor* dOut,
DenseTensor* ddOut) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "RsqrtGradGrad"));
auto out = EigenVector<T>::Flatten(
GET_DATA_SAFELY(Out, "Output", "Out", "RsqrtGradGrad"));
// rsqrt GradGrad: ddy = -0.5 * ddx * y * y * y, dy = (3/y) * dx * ddx
if (dOut) {
auto dx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dX, "Output", "DX", "RsqrtGradGrad"));
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Output", "DOut", "RsqrtGradGrad"));
dout.device(*d) = (static_cast<T>(3.0) / out) * dx * ddx;
}
if (ddOut) {
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "RsqrtGradGrad"));
ddout.device(*d) = ddx * static_cast<T>(-0.5) * out * out * out;
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CELUFunctor : public BaseActivationFunctor<T> {
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
template <typename Device, typename X, typename Out>
void operator()(Device d, X x, Out out) const {
out.device(d) =
(x < static_cast<T>(0))
.select(static_cast<T>(alpha) * ((x / static_cast<T>(alpha)).exp() -
static_cast<T>(1)), // NOLINT
x);
}
};
template <typename T>
struct CELUGradFunctor : public BaseActivationFunctor<T> {
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
template <typename Device,
typename X,
typename Out,
typename dOut,
typename dX>
void operator()(Device d, X x, Out out UNUSED, dOut dout, dX dx) const {
auto temp_a_pos = static_cast<T>(alpha > 0); // NOLINT
auto temp_a_neg = static_cast<T>(alpha <= 0);
auto temp_x_pos = (x > static_cast<T>(0)).template cast<T>();
auto temp_x_neg = (x <= static_cast<T>(0)).template cast<T>();
// dx = dout, if alpha > 0 and x > 0
// dx = dout * (x/alpha).exp(), if alpha > 0 and x <= 0
// dx = dout , if alpha < 0 and x > 0
// dx = dout * (x/alpha).exp(), if alpha < 0 and x <=0
dx.device(d) =
dout * temp_a_pos * temp_x_pos +
dout * (x / static_cast<T>(alpha)).exp() * temp_a_pos * temp_x_neg +
dout * temp_a_neg * temp_x_pos +
dout * (x / static_cast<T>(alpha)).exp() * temp_a_neg * temp_x_neg;
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CELUGradGradFunctor : public BaseActivationFunctor<T> {
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* dOut,
const DenseTensor* ddX,
DenseTensor* dX,
DenseTensor* ddOut) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "CELUGradGrad"));
auto x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(X, "Input", "X", "CELUGradGrad"));
if (dX) {
auto dx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dX, "Output", "DX", "CELUGradGrad"));
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Output", "DOut", "CELUGradGrad"));
dx.device(*d) = ddx * dout / static_cast<T>(alpha) *
(x / static_cast<T>(alpha)).exp() *
(x <= static_cast<T>(0)).template cast<T>();
}
if (ddOut) {
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "CELUGradGrad"));
ddout.device(*d) = ddx * ((x > static_cast<T>(0)).template cast<T>() +
(x / static_cast<T>(alpha)).exp() *
(x <= static_cast<T>(0)).template cast<T>())
.template cast<T>();
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct SquareGradGradFunctor : public BaseActivationFunctor<T> {
template <typename Device>
void operator()(const Device& dev,
const DenseTensor* X,
const DenseTensor* dOut,
const DenseTensor* ddX,
DenseTensor* dX,
DenseTensor* ddOut) const {
auto* d = dev.eigen_device();
auto ddx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddX, "Input", "DDX", "SquareGradGrad"));
auto x = EigenVector<T>::Flatten(
GET_DATA_SAFELY(X, "Input", "X", "SquareGradGrad"));
// square GradGrad: ddy=2x*ddx, dx=2dy*ddx
// calculate dx first, so ddy can inplace ddx
if (dX) {
auto dx = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dX, "Output", "DX", "SquareGradGrad"));
auto dout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(dOut, "Output", "DOut", "SquareGradGrad"));
dx.device(*d) = ddx * static_cast<T>(2) * dout;
}
if (ddOut) {
auto ddout = EigenVector<T>::Flatten(
GET_DATA_SAFELY(ddOut, "Output", "DDOut", "SquareGradGrad"));
ddout.device(*d) = ddx * static_cast<T>(2) * x;
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
#if defined(__NVCC__) || defined(__HIPCC__) || defined(__xpu__)
template <typename T>
struct CudaLogitFunctor : public BaseActivationFunctor<T> {
using AttrPair = std::vector<std::pair<const char*, double*>>;
using MT = typename MPTypeTrait<T>::Type;
MT zero = static_cast<MT>(0.0f);
MT one = static_cast<MT>(1.0f);
double eps;
typename CudaLogitFunctor<T>::AttrPair GetAttrs() { return {{"eps", &eps}}; }
// logit(x) = ln(x/(1-x))
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
MT y = min(x, (one - static_cast<MT>(eps)));
y = max(y, static_cast<MT>(eps));
if (!eps) {
y = x < zero || x > one ? static_cast<T>(NAN) : log(y / (one - y));
} else {
y = log(y / (one - y));
}
return static_cast<T>(y);
}
};
template <typename T>
struct CudaLogitGradFunctor : public BaseActivationFunctor<T> {
using AttrPair = std::vector<std::pair<const char*, double*>>;
using MT = typename MPTypeTrait<T>::Type;
double eps;
MT zero = static_cast<MT>(0.0f);
MT one = static_cast<MT>(1.0f);
typename CudaLogitGradFunctor<T>::AttrPair GetAttrs() {
return {{"eps", &eps}};
}
// logit(x)' = 1/(x*(1-x))
__device__ __forceinline__ T operator()(const T dout, const T arg_x) const {
MT x = static_cast<MT>(arg_x);
MT dx;
if (!eps) {
dx = (x < zero || x > one) ? static_cast<T>(NAN)
: (static_cast<MT>(dout) / (x * (one - x)));
} else {
dx = (x < static_cast<MT>(eps) || x > one - static_cast<MT>(eps))
? zero
: (static_cast<MT>(dout) / (x * (one - x)));
}
return static_cast<T>(dx);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaReluFunctor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
// relu(x) = max(x, 0)
__device__ __forceinline__ T operator()(const T x) const {
return x < zero ? zero : x;
}
};
template <typename T>
struct CudaReluGradFunctor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
// dx = dout * (out > 0)
__device__ __forceinline__ T operator()(const T dout, const T out) const {
return out > zero ? dout : zero;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaCosFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// cos(x) = cos(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(cos(x));
}
};
template <typename T>
struct CudaCosGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// dx = dout * (-sin(x))
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
if constexpr (std::is_same<T, phi::float16>::value ||
std::is_same<T, phi::bfloat16>::value) {
return static_cast<T>(-arg_dout * static_cast<T>(sin(x)));
} else {
return static_cast<T>(-dout * sin(x));
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaCosGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// dx = dout * (-sin(x))
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(-dout * conj(sin(x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaExpFunctor : public BaseActivationFunctor<T> {
// exp(x) = expf(x)
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
__device__ __forceinline__ U operator()(const T x) const {
return static_cast<U>(expf(static_cast<float>(x)));
}
};
template <>
struct CudaExpFunctor<double> : public BaseActivationFunctor<double> {
// exp(x) = exp(x)
__device__ __forceinline__ double operator()(const double x) const {
return exp(x);
}
};
template <typename T>
struct CudaExpFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// exp(x) = exp(x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(exp(x));
}
};
template <typename T>
struct CudaSeluFunctor : public BaseActivationFunctor<T> {
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"scale", &scale}, {"alpha", &alpha}};
}
__device__ __forceinline__ T operator()(const T x) const {
using MT =
typename std::conditional<(sizeof(T) > sizeof(float)), T, float>::type;
MT res = static_cast<MT>(x);
if (x <= zero) {
res = alpha * expf(res) - alpha;
}
res *= scale;
return static_cast<T>(res);
}
private:
float scale;
float alpha;
T zero = static_cast<T>(0.0f);
};
template <>
struct CudaSeluFunctor<double> : public BaseActivationFunctor<double> {
typename BaseActivationFunctor<double>::AttrPair GetAttrs() {
return {{"scale", &scale}, {"alpha", &alpha}};
}
__device__ __forceinline__ double operator()(const double x) const {
double res = x;
double alpha_cast = static_cast<double>(alpha);
double scale_cast = static_cast<double>(scale);
if (res <= zero) {
res = alpha_cast * exp(res) - alpha_cast;
}
res *= scale_cast;
return res;
}
private:
float scale;
float alpha;
double zero = static_cast<double>(0.0f);
};
template <typename T>
struct CudaSquareFunctor : public BaseActivationFunctor<T> {
// square(x) = x * x
__device__ __forceinline__ T operator()(const T x) const { return x * x; }
};
template <typename T>
struct CudaSquareGradFunctor : public BaseActivationFunctor<T> {
T two = static_cast<T>(2.0f);
// dx = dout * 2 * x
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return dout * two * x;
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSquareGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> two = static_cast<ComplexType<T>>(2.0f);
// dx = dout * 2 * x
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout * two * conj(x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaRsquareFunctor : public BaseActivationFunctor<T> {
// square(x) = 1 / (x * x)
T one = static_cast<T>(1.0f);
__device__ __forceinline__ T operator()(const T x) const {
return one / (x * x);
}
};
template <typename T>
struct CudaExpGradFunctor : public BaseActivationFunctor<T> {
// dx = dout * out
__device__ __forceinline__ T operator()(const T dout, const T out) const {
return dout * out;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaExpGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// dx = dout * exp(x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> out) const {
return static_cast<ComplexType<T>>(dout * conj(out));
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaReciprocalFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
__device__ __forceinline__ T operator()(const T x) const {
return static_cast<T>(one / static_cast<MT>(x));
}
};
template <typename T>
struct CudaReciprocalFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> x) const {
auto both_inf = [](T real, T imag) {
return (::isinf(real) && ::isinf(imag));
};
auto either_inf = [](T real, T imag) {
return ::isinf(real) || ::isinf(imag);
};
auto either_nan = [](T real, T imag) {
return ::isnan(real) || ::isnan(imag);
};
if (either_nan(x.real, x.imag) || both_inf(x.real, x.imag)) {
// If either is Nan or both are infinite, return {nan, nan}
if constexpr (std::is_same<T, float>::value) {
return ComplexType<T>(nanf(""), nanf(""));
} else if constexpr (std::is_same<T, double>::value) {
return ComplexType<T>(nan(""), nan(""));
}
} else if (either_inf(x.real, x.imag)) {
// If either is Inf, return {0, 0}
return ComplexType<T>(static_cast<T>(0), static_cast<T>(0));
}
return static_cast<ComplexType<T>>(1.0) / x;
}
};
template <typename T>
struct CudaReciprocalGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// dx = -dout * out^2
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_out) const {
MT dout = static_cast<MT>(arg_dout);
MT out = static_cast<MT>(arg_out);
return static_cast<T>(-dout * static_cast<MT>(static_cast<T>(out * out)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaReciprocalGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// dx = -dout * out^2
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> out) const {
return -dout * conj(out * out);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
// for pow(x, -1)
template <typename T>
struct CudaReciprocalGradDepXFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
// dx = -dout * out^2
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
return static_cast<T>(-dout * (one / (x * x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaReciprocalGradDepXFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = -dout * out^2
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return -dout * conj(one / (x * x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaExpm1Functor : public BaseActivationFunctor<T> {
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
// expm1(x) = expm1f(x)
__device__ __forceinline__ U operator()(const T x) const {
return static_cast<U>(::expm1f(static_cast<float>(x)));
}
};
template <>
struct CudaExpm1Functor<double> : public BaseActivationFunctor<double> {
// expm1(x) = expm1(x)
__device__ __forceinline__ double operator()(const double x) const {
return ::expm1(x);
}
};
template <typename T>
__device__ __forceinline__ ComplexType<T> local_expm1(const ComplexType<T>& z) {
T x = z.real;
T y = z.imag;
T a = std::sin(y / 2);
T er = std::expm1(x) * std::cos(y) - T(2) * a * a;
T ei = std::exp(x) * std::sin(y);
return {er, ei};
}
template <typename T>
struct CudaExpm1Functor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(local_expm1(x));
}
};
template <typename T>
struct CudaExpm1GradFunctor : public BaseActivationFunctor<T> {
T one = static_cast<T>(1.0f);
// dx = dout * out
__device__ __forceinline__ T operator()(const T dout, const T out) const {
return dout * (out + one);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaExpm1GradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = dout * exp(x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> out) const {
return static_cast<ComplexType<T>>(dout * (conj(out) + one));
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaSinFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// sin(x) = sin(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(sin(x));
}
};
template <typename T>
struct CudaSinGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// dx = dout * cos(x)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
if constexpr (std::is_same<T, phi::float16>::value ||
std::is_same<T, phi::bfloat16>::value) {
return static_cast<T>(arg_dout * static_cast<T>(cos(x)));
} else {
return static_cast<T>(dout * cos(x));
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSinGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// dx = dout * cos(x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout * conj(cos(x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaTanFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// tan(x) = tan(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(tan(x));
}
};
template <typename T>
struct CudaTanGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
T one = static_cast<T>(1.0f);
// dx = dout *(1 + tan(x)^2)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
if constexpr (std::is_same<T, double>::value) {
double td = ::tan(x);
double tsq = __dmul_rn(td, td);
double y = __dadd_rn(tsq, 1.0);
return static_cast<T>(dout * y);
} else if constexpr (std::is_same<T, float>::value) {
float tf = ::tanf(x);
float tsq = __fmul_rn(tf, tf);
float y = __fadd_rn(tsq, 1.0f);
return static_cast<T>(dout * y);
} else if constexpr (std::is_same<T, phi::float16>::value) {
__half tf = __float2half_rn(::tanf(x));
__half tmp_half = __hmul(tf, tf);
return arg_dout * (one + static_cast<T>(__half2float(tmp_half)));
} else {
return static_cast<T>(dout *
(static_cast<MT>(1.0f) + ::tan(x) * ::tan(x)));
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaTanGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// dx = dout *(1 + tan(x)^2)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
return static_cast<ComplexType<T>>(dout * conj(tan(x) * tan(x) + one));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAsinFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// asin(x) = asin(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(asin(x));
}
};
template <typename T>
struct CudaAsinGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
// dx = dout / sqrt(1 - x^2)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
return static_cast<T>(dout / sqrt(one - x * x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAsinGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = dout / sqrt(1 - x^2)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout / conj(sqrt(one - x * x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAcosFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// acos(x) = acos(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(acos(x));
}
};
template <typename T>
struct CudaAcosGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
// dx = -dout / sqrt(1 - x^2)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
return static_cast<T>(-dout / sqrt(one - x * x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAcosGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = -dout / sqrt(1 - x^2)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(-dout / conj(sqrt(one - x * x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaCoshFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// cosh(x) = cosh(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(cosh(x));
}
};
template <typename T>
struct CudaCoshGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// dx = dout * sinh(x)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
return static_cast<T>(dout * sinh(x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaCoshGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// dx = dout * sinh(x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout * conj(sinh(x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSinhFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// sinh(x) = sinh(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(sinh(x));
}
};
template <typename T>
struct CudaSinhGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// dx = dout * cosh(x)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
return static_cast<T>(dout * cosh(x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSinhGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// dx = dout * cosh(x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout * conj(cosh(x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAcoshFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// Acosh(x) = acosh(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(acosh(x));
}
};
template <typename T>
struct CudaAcoshGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
// dx = dout * 1 / sqrt(x^2 - 1)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
return static_cast<T>(dout * one / sqrt(x * x - one));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAcoshGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = dout * 1 / sqrt(x^2 - 1)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout * conj(one / sqrt(x * x - one)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAsinhFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// Asinh(x) = asinh(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(asinh(x));
}
};
template <typename T>
struct CudaAsinhGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
// dx = dout * 1/sqrt(x^2 + 1)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
return static_cast<T>(dout * one / sqrt(x * x + one));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAsinhGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = dout * 1/sqrt(x^2 + 1)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout * conj(one / sqrt(x * x + one)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAtanhFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// Atanh(x) = atanh(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(atanh(x));
}
};
template <typename T>
struct CudaSTanhFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
float scale_a;
float scale_b;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"scale_a", &scale_a}, {"scale_b", &scale_b}};
}
// stanh(x) = b * tanh(a * x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
MT a = static_cast<MT>(scale_a);
MT b = static_cast<MT>(scale_b);
return static_cast<T>(b * tanh(a * x));
}
};
template <typename T>
struct CudaSTanhGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
float scale_a;
float scale_b;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"scale_a", &scale_a}, {"scale_b", &scale_b}};
}
// dx = dout * a * b * (1 - tanh(a * x) * tanh(a * x))
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
MT a = static_cast<MT>(scale_a);
MT b = static_cast<MT>(scale_b);
MT temp = tanh(a * x);
return static_cast<T>(dout * a * b * (one - temp * temp));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSTanhGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
float scale_a;
float scale_b;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"scale_a", &scale_a}, {"scale_b", &scale_b}};
}
// dx = dout * a * b * (1 - tanh(a * x) * tanh(a * x))
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_dout, const ComplexType<T> arg_x) const {
ComplexType<T> dout = static_cast<ComplexType<T>>(arg_dout);
ComplexType<T> x = static_cast<ComplexType<T>>(arg_x);
ComplexType<T> a = static_cast<ComplexType<T>>(scale_a);
ComplexType<T> b = static_cast<ComplexType<T>>(scale_b);
ComplexType<T> temp = tanh(a * x);
return static_cast<ComplexType<T>>(dout *
conj(a * b * (one - temp * temp)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
__device__ __forceinline__ T log1p_local(T x) {
return log1p(x);
}
template <typename T>
__device__ __forceinline__ ComplexType<T> log1p_local(ComplexType<T> x) {
return log(ComplexType<T>{1.} + exp(x));
}
template <typename T>
struct CudaSoftplusFunctor : public BaseActivationFunctor<T> {
using AttrPair = std::vector<std::pair<const char*, double*>>;
using MT = typename MPTypeTrait<T>::Type;
double beta;
double threshold;
typename CudaSoftplusFunctor<T>::AttrPair GetAttrs() {
return {{"beta", &beta}, {"threshold", &threshold}};
}
// softplus(x) = beta * x > threshold ? x : log(1 + exp(beta * x)) / beta
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
MT b = static_cast<MT>(beta);
MT t = static_cast<MT>(threshold);
return static_cast<T>((x * b) > t ? x : (log1p_local(exp(x * b))) / b);
}
};
template <typename T>
struct CudaSoftplusFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
using MT = typename MPTypeTrait<ComplexType<T>>::Type;
MT one = static_cast<MT>(1.0f);
float beta;
float threshold;
typename BaseActivationFunctor<ComplexType<T>>::AttrPair GetAttrs() {
return {{"beta", &beta}, {"threshold", &threshold}};
}
// softplus(x) = beta * x > threshold ? x : log(1 + exp(beta * x)) / beta
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_x) const {
MT x = static_cast<MT>(arg_x);
MT b = static_cast<MT>(beta);
MT t = static_cast<MT>(threshold);
MT x_beta = x * static_cast<MT>(beta);
return static_cast<ComplexType<T>>(x_beta > t ? x
: log(one + exp(x_beta)) / b);
}
};
template <typename T>
struct CudaSoftplusGradFunctor : public BaseActivationFunctor<T> {
using AttrPair = std::vector<std::pair<const char*, double*>>;
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
double beta;
double threshold;
typename CudaSoftplusGradFunctor<T>::AttrPair GetAttrs() {
return {{"beta", &beta}, {"threshold", &threshold}};
}
// dx = x * beta > threshold ? dout : dout / (1 + exp(-beta * x))
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
MT b = static_cast<MT>(beta);
MT t = static_cast<MT>(threshold);
MT z = std::exp(x * b);
return (x * b) > t ? arg_dout : static_cast<T>(dout * z / (z + one));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSoftplusGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
using AttrPair = std::vector<std::pair<const char*, double*>>;
using MT = typename MPTypeTrait<ComplexType<T>>::Type;
MT one = static_cast<MT>(1.0f);
double beta;
double threshold;
typename CudaSoftplusGradFunctor<ComplexType<T>>::AttrPair GetAttrs() {
return {{"beta", &beta}, {"threshold", &threshold}};
}
// dx = x * beta > threshold ? dout : dout / (1 + exp(-beta * x))
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_dout, const ComplexType<T> arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
MT b = static_cast<MT>(beta);
MT t = static_cast<MT>(threshold);
MT z = exp(x * b);
return (x * b) > t
? dout
: static_cast<ComplexType<T>>(dout * conj(z / (z + one)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAtanhGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
// dx = dout * 1/(1- x^2)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
return static_cast<T>(dout * one / (one - x * x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAtanhGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = dout * 1/(1- x^2)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout * conj(one / (one - x * x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSqrtFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// sqrt(x) = sqrt(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(sqrt(x));
}
};
template <typename T>
struct CudaSqrtGradFunctor : public BaseActivationFunctor<T> {
T two = static_cast<T>(2);
// dx = dout / (2 * out)
__device__ __forceinline__ T operator()(const T dout, const T out) const {
return dout / (two * out);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaSqrtGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one_half = static_cast<ComplexType<T>>(0.5f);
// dx = dout * 0.5 / out
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> out) const {
return dout * conj(one_half / out);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
// for pow(x, 0.5)
template <typename T>
struct CudaSqrtGradDepXFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one_half = static_cast<MT>(0.5f);
// dx = dout * (0.5 * rsqrt(x))
__device__ __forceinline__ T operator()(const T dout, const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return dout * static_cast<T>(one_half * rsqrt(x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSqrtGradDepXFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one_half = static_cast<ComplexType<T>>(0.5f);
// dx = dout * conj(0.5 * rsqrt(x))
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return dout * conj(one_half / sqrt(x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaRsqrtFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// rsqrt(x) = rsqrt(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(rsqrt(x));
}
};
template <typename T>
struct CudaRsqrtFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// rsqrt(x) = 1 / sqrt(x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_x) const {
return one / sqrt(arg_x);
}
};
template <typename T, bool Compatible = false>
struct CudaRsqrtGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT minus_one_half = static_cast<MT>(-0.5f);
// dx = -0.5 * dout * out^3
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_out) const {
if constexpr (Compatible) {
T t1 = static_cast<T>(-0.5f) * arg_dout;
T cube = arg_out * arg_out * arg_out;
return t1 * cube;
} else {
MT dout = static_cast<MT>(arg_dout);
MT out = static_cast<MT>(arg_out);
return static_cast<T>(minus_one_half * dout * (out * out * out));
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaAtanFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// atan(x) = atan(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(atan(x));
}
};
template <typename T>
struct CudaAtanGradFunctor : public BaseActivationFunctor<T> {
T one = static_cast<T>(1.0f);
// dx = dout / (1 + x^2)
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return dout / (one + x * x);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaAtanGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = dout / (1 + x^2)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return dout / conj(one + x * x);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaTanhFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// tanh(x) = tanh(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(tanh(x));
}
};
template <typename T>
struct CudaTanhGradFunctor : public BaseActivationFunctor<T> {
T one = static_cast<T>(1.0f);
// dx = dout * (1 - out^2)
__device__ __forceinline__ T operator()(const T dout, const T out) const {
if constexpr (std::is_same<T, phi::float16>::value) {
__half out_half = __float2half_rn(static_cast<float>(out));
__half tmp_half = __hmul(out_half, out_half);
return dout * (one - static_cast<T>(__half2float(tmp_half)));
} else {
return dout * (one - out * out);
}
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaTanhGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = dout * (1 - out^2)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> out) const {
return dout * conj(one - out * out);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaHardTanhFunctor : public BaseActivationFunctor<T> {
float t_min;
float t_max;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"t_min", &t_min}, {"t_max", &t_max}};
}
// brelu(x) = min(max(x, t_min), t_max)
__device__ __forceinline__ T operator()(const T x) const {
T t_min_cast = static_cast<T>(t_min);
T t_max_cast = static_cast<T>(t_max);
T temp_max = x > t_min_cast ? x : t_min_cast;
T temp_min = temp_max < t_max_cast ? temp_max : t_max_cast;
return temp_min;
}
};
template <typename T>
struct CudaMishFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
float threshold;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}};
}
// mish(x) = x * tanh(softplus(x))
// softplus(x) = x, if x > threshold
// = ln(1 + exp(x)), otherwise
// Inputs: args[0], the input x
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
MT sp = (x > static_cast<MT>(threshold)) ? x : log(one + exp(x));
return static_cast<T>(x * tanh(sp));
}
};
template <typename T>
struct CudaMishGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
float threshold;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}};
}
// dx = dout * (tanh(sp) + x * (1 - tanh(sp) ** 2) * (1 - exp(-sp)))
// sp = softplus(x)
// Inputs: args[0], the input dout
// args[1], the input x
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
MT sp = (x > static_cast<MT>(threshold)) ? x : log(one + exp(x));
MT gsp = (x > static_cast<MT>(threshold)) ? one : one / (one + exp(-x));
MT tsp = tanh(sp);
return static_cast<T>(dout * (tsp + x * (one - tsp * tsp) * gsp));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaHardTanhGradFunctor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
float t_min;
float t_max;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"t_min", &t_min}, {"t_max", &t_max}};
}
// dx = (x > t_min && x < t_max) ? dout : 0
__device__ __forceinline__ T operator()(const T dout, const T x) const {
T t_min_cast = static_cast<T>(t_min);
T t_max_cast = static_cast<T>(t_max);
return (x > t_min_cast && x < t_max_cast) ? dout : zero;
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaThresholdedReluFunctor : public BaseActivationFunctor<T> {
float threshold;
float value;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}, {"value", &value}};
}
// thresholded_relu(x, threshold, value) = x > threshold ? x : value
__device__ __forceinline__ T operator()(const T x) const {
return x > static_cast<T>(threshold) ? x : static_cast<T>(value);
}
};
template <typename T>
struct CudaThresholdedReluGradFunctor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
float threshold;
float value;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}, {"value", &value}};
}
// dx = x > threshold ? dout : 0
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return x > static_cast<T>(threshold) ? dout : zero;
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaRelu6Functor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
float threshold;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}};
}
// relu6(x) = min(max(0, x), 6)
__device__ __forceinline__ T operator()(const T x) const {
T t = static_cast<T>(threshold);
return x <= zero ? zero : (x < t ? x : t);
}
};
template <typename T>
struct CudaRelu6GradFunctor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
typename BaseActivationFunctor<T>::AttrPair GetAttrs() { return {{}}; }
// dx = (out > 0 && out < t) ? dout : 0
__device__ __forceinline__ T operator()(const T dout, const T out) const {
float threshold = 6;
T t = static_cast<T>(threshold);
return (out > zero && out < t) ? dout : zero;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaLeakyReluFunctor : public BaseActivationFunctor<T, double> {
using MT = typename MPTypeTrait<T>::Type;
T zero = static_cast<T>(0.0f);
double alpha;
typename BaseActivationFunctor<T, double>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
// leakyrelu(x) = x > 0 ? x : alpha * x
__device__ __forceinline__ T operator()(const T x) const {
return x > zero
? x
: static_cast<T>(static_cast<MT>(alpha) * static_cast<MT>(x));
}
};
template <typename T>
struct CudaLeakyReluGradFunctor : public BaseActivationFunctor<T, double> {
using MT = typename MPTypeTrait<T>::Type;
T zero = static_cast<T>(0.0f);
double alpha;
typename BaseActivationFunctor<T, double>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
// dx = dout * (x > 0 ? 1 : alpha)
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return x > zero
? dout
: static_cast<T>(static_cast<MT>(alpha) * static_cast<MT>(dout));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSoftShrinkFunctor : public BaseActivationFunctor<T> {
float lambda;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"lambda", &lambda}};
}
// softshrink(x) = x - lambda, if x > lambda;
// x + lambda, if x < -lambda;
// 0, otherwise.
__device__ __forceinline__ T operator()(const T x) const {
T l = static_cast<T>(lambda);
T temp1 = static_cast<T>(x > l);
T temp2 = static_cast<T>(x < -l);
return temp1 * (x - l) + temp2 * (x + l);
}
};
template <typename T>
struct CudaSoftShrinkGradFunctor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
float lambda;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"lambda", &lambda}};
}
// dx = dout, if x > lambda or x < -lambda else 0
__device__ __forceinline__ T operator()(const T dout, const T x) const {
T l = static_cast<T>(lambda);
return (x >= -l && x <= l) ? zero : dout;
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaTanhShrinkFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
bool compatible = false;
// tanhshrink(x) = x - tanh(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
if (compatible) {
// Match PyTorch: tanh truncated to native dtype T before subtraction
T tanh_val = static_cast<T>(tanh(x));
return static_cast<T>(x - static_cast<MT>(tanh_val));
}
return static_cast<T>(x - tanh(x));
}
};
template <typename T>
struct CudaTanhShrinkGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
bool compatible = false;
// dx = dout * tanh(x)^2
// PyTorch decomposes tanhshrink as x - tanh(x), so backward is:
// tanh_grad_input = -grad * (1 - tanh_out_T * tanh_out_T)
// dx = grad + tanh_grad_input
// PyTorch's tanh backward for fp16/bf16 computes out*out in native dtype
// using __hmul (not promoted to fp32). We must use explicit __half ops
// to avoid implicit float promotion through operator float().
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
if (compatible) {
// Match PyTorch decomposed backward for tanhshrink = x - tanh(x):
// tanh_grad = grad * (1 - tanh_out^2) -- CudaTanhGradFunctor pattern
// dx = grad - tanh_grad
// tanh output is stored at native dtype T.
T tanh_val = static_cast<T>(tanh(x));
if constexpr (std::is_same<T, phi::float16>::value) {
// Match PyTorch: tanh backward computes in native fp16 (scalar_t),
// using __hmul for multiplication. Each intermediate truncated to fp16.
// PyTorch's tanh_backward: a * (scalar_t{1.} - b * b)
// NVCC may fuse __hsub(one, __hmul(t,t)) into HFMA2, but PyTorch
// does NOT fuse these for fp16. Use volatile to prevent FMA fusion
// for the t_sq computation, matching PyTorch's non-fused behavior.
__half t_half = __float2half_rn(static_cast<float>(tanh_val));
volatile __half t_sq_half = __hmul(t_half, t_half);
__half one_half = __float2half_rn(1.0f);
__half one_minus_t_sq = __hsub(one_half, t_sq_half);
__half dout_half = __float2half_rn(static_cast<float>(arg_dout));
volatile __half tanh_grad_half = __hmul(dout_half, one_minus_t_sq);
__half result_half = __hsub(dout_half, tanh_grad_half);
return static_cast<T>(__half2float(result_half));
} else if constexpr (std::is_same<T, phi::dtype::bfloat16>::value) {
// Match PyTorch: tanh backward computes in native bf16 (scalar_t),
// not promoted to opmath_t. Compute each step at T precision.
T one = static_cast<T>(1.0f);
T t_sq = static_cast<T>(static_cast<float>(tanh_val) *
static_cast<float>(tanh_val));
T one_minus_t_sq =
static_cast<T>(static_cast<float>(one) - static_cast<float>(t_sq));
T tanh_grad = static_cast<T>(static_cast<float>(arg_dout) *
static_cast<float>(one_minus_t_sq));
return static_cast<T>(static_cast<float>(arg_dout) -
static_cast<float>(tanh_grad));
} else if constexpr (std::is_same<T, float>::value) {
// For float32: T == MT == float.
// PyTorch decomposes tanhshrink backward into two SEPARATE kernels:
// Kernel 1 (tanh_backward): (-dout) * (1.0f - t*t)
// Kernel 2 (add): dout + tanh_backward_result
// Within Kernel 1, NVCC fuses (1.0f - t*t) into fma(-t, t, 1),
// so we ALLOW FMA here. The multiply dout*one_minus_t_sq is a
// separate fmul instruction in PyTorch's kernel.
// Between kernels, no FMA fusion occurs (global memory barrier).
//
// Bug: volatile on tanh_grad does NOT prevent NVCC from fusing
// dout * one_minus_t_sq and dout - tanh_grad into a single
// fmaf(-dout, one_minus_t_sq, dout). This causes 1-ULP errors
// when dout != 1.0 (e.g., .mean() backward where dout = 1/N).
// Fix: use __fmul_rn to force a non-FMA rounded multiply,
// which emits a mul.rn.f32 instruction that NVCC cannot fuse.
float t = static_cast<float>(tanh_val);
float one_minus_t_sq = 1.0f - t * t; // FMA allowed: fma(-t,t,1)
float tanh_grad = __fmul_rn(dout, one_minus_t_sq); // non-FMA mul
return dout - tanh_grad;
} else {
// For float64: T == MT == double.
// Same decomposition as float32. NVCC fuses (1 - t*t) via FMA.
// Use __dmul_rn to prevent FMA fusion of multiply+subtract,
// matching PyTorch's separate-kernel behavior.
double t = static_cast<double>(tanh_val);
double one_minus_t_sq = 1.0 - t * t; // FMA allowed: fma(-t,t,1)
double tanh_grad = __dmul_rn(dout, one_minus_t_sq); // non-FMA mul
return dout - tanh_grad;
}
}
return static_cast<T>(dout * tanh(x) * tanh(x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaHardShrinkFunctor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
float threshold;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}};
}
// hadrshrink(x) = (x > -threshold && x < threshold) ? 0 : x
__device__ __forceinline__ T operator()(const T x) const {
T t = static_cast<T>(threshold);
return (x >= -t && x <= t) ? zero : x;
}
};
template <typename T>
struct CudaHardShrinkGradFunctor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
float threshold;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}};
}
// dx = (x > -threshold && x < threshold) ? 0 : dout
__device__ __forceinline__ T operator()(const T dout, const T x) const {
T t = static_cast<T>(threshold);
return (x >= -t && x <= t) ? zero : dout;
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaELUFunctor : public BaseActivationFunctor<T> {
using CT = typename MPTypeTrait<T>::Type;
CT zero = static_cast<CT>(0.0f);
CT one = static_cast<CT>(1.0f);
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
// elu(x) = x, if x > 0
// elu(x) = alpha * (e^x - 1), if x <= 0
__device__ __forceinline__ T operator()(const T arg_x) const {
CT x = static_cast<CT>(arg_x);
CT temp = static_cast<CT>(alpha) * (exp(x) - one);
CT res = x > zero ? x : temp;
return static_cast<T>(res);
}
};
template <typename T>
struct CudaELUGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT zero = static_cast<MT>(0.0f);
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
// case 1: alpha >= 0
// dx = dout, if out > 0
// dx = dout * (out + alpha), if out <= 0
__device__ __forceinline__ T operator()(T arg_dout, T arg_out) const {
MT dout = static_cast<MT>(arg_dout);
MT out = static_cast<MT>(arg_out);
MT a = static_cast<MT>(alpha);
MT out_pos = static_cast<MT>(out > zero);
MT out_neg = static_cast<MT>(out <= zero);
return static_cast<T>(dout * (out_pos + out_neg * (out + a)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaELUGradNegativeAlphaFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT zero = static_cast<MT>(0.0f);
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
// case 2: alpha < 0
// dx = dout, if x > 0
// dx = dout * (out + alpha), if x <=0
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_out,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT out = static_cast<MT>(arg_out);
MT x = static_cast<MT>(arg_x);
MT a = static_cast<MT>(alpha);
MT x_pos = static_cast<MT>(x > zero);
MT x_neg = static_cast<MT>(x <= zero);
return static_cast<T>(dout * (x_pos + x_neg * (out + a)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSiluFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
// silu(x) = x / (1 + exp(-x))
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(x / (one + exp(-x)));
}
};
template <typename T>
struct CudaSiluGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
// dx = dout * (1 + exp(-x) + x * exp(-x) / (1 + exp(-x))^2)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
MT temp = one / (one + exp(-x));
return static_cast<T>(dout * temp * (one + x * (one - temp)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSiluGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = dout * (1 + exp(-x) + x * exp(-x) / (1 + exp(-x))^2)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_dout, const ComplexType<T> arg_x) const {
ComplexType<T> dout = static_cast<ComplexType<T>>(arg_dout);
ComplexType<T> x = static_cast<ComplexType<T>>(arg_x);
ComplexType<T> temp = one / (one + exp(-x));
return dout * conj(temp * (one + x * (one - temp)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSoftsignFunctor : public BaseActivationFunctor<T> {
T one = static_cast<T>(1.0f);
// softsign(x) = x / (1 + abs(x))
__device__ __forceinline__ T operator()(const T x) const {
// Using abs directly will cause namespace conflict
return x / (one + (x > -x ? x : -x));
}
};
template <typename T>
struct CudaSoftsignFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
using Complex = ComplexType<T>;
Complex one = static_cast<Complex>(1.0f);
__device__ __forceinline__ Complex operator()(const Complex x) const {
return x / (one + static_cast<Complex>(abs(x)));
}
};
template <typename T>
struct CudaSoftsignGradFunctor : public BaseActivationFunctor<T> {
T one = static_cast<T>(1.0f);
// dx = dout / (1 + abs(x))^2
__device__ __forceinline__ T operator()(const T dout, const T x) const {
// Using abs directly will cause namespace conflict
T temp = one + (x > -x ? x : -x);
return dout / (temp * temp);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSoftsignGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
using Complex = ComplexType<T>;
Complex one = static_cast<Complex>(1.0f);
__device__ __forceinline__ Complex operator()(const Complex dout,
const Complex x) const {
Complex abs_x = static_cast<Complex>(abs(x));
Complex abs_x_plus = one + abs_x;
Complex temp = static_cast<Complex>((-x / (abs_x_plus * abs_x_plus)).real);
return dout * (one / abs_x_plus + temp * x / abs_x);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSigmoidFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
// sigmoid(x) = 1 / (1 + exp(-x))
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<T>(one / (one + exp(-x)));
}
};
template <typename T>
struct CudaSigmoidGradFunctor : public BaseActivationFunctor<T> {
T one = static_cast<T>(1.0f);
// dx = dout * out * (1 - out)
__device__ __forceinline__ T operator()(const T dout, const T out) const {
return dout * (one - out) * out;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaSigmoidGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
using Complex = ComplexType<T>;
Complex one = Complex(1.0f);
// dx = dout * out * (1 - out)
__device__ __forceinline__ Complex operator()(const Complex dout,
const Complex out) const {
Complex y = out * (one - out);
return dout * conj(y);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
struct CudaLogSigmoidFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT zero = static_cast<MT>(0.0f);
// logsigmoid(x) = log(1 / (1 + exp(-x)))
// Use the numerically stable:
// log_sigmoid(x) = min(0, x) - log1p(exp(-abs(x)))
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
MT min0 = (x < zero) ? x : zero;
MT abs_x = abs(x);
return static_cast<T>(min0 - log1p_local(exp(-abs_x)));
}
};
// Specialized CUDA implementation for complex numbers
template <typename T>
struct CudaLogSigmoidFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = ComplexType<T>(T(1), T(0));
// For complex numbers, use log σ(x) = -log(1 + exp(-x))
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_x) const {
ComplexType<T> x = static_cast<ComplexType<T>>(arg_x);
// LogSigmoid formula: log σ(x) = -log(1 + exp(-x))
return -log(one + exp(-x));
}
};
template <typename T>
struct CudaLogSigmoidGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT zero = static_cast<MT>(0.0f);
MT one = static_cast<MT>(1.0f);
// dx = dout * exp(-x) / (1 + exp(-x))
// Use stable backward:
// grad = dout * (max_deriv - sign * (z / (1 + z)))
// where z = exp(-abs(x)), max_deriv = (x < 0) ? 1 : 0, sign = (x < 0) ? 1 :
// -1
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
// in_negative, max_deriv, sign
const bool in_negative = (x < zero);
const MT max_deriv = in_negative ? one : zero;
const MT sign = in_negative ? one : -one;
MT z = exp(-abs(x));
return static_cast<T>(dout * (max_deriv - sign * (z / (one + z))));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaLogSigmoidGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = ComplexType<T>(T(1), T(0));
// For complex numbers, gradient of log σ(x) is σ(-x) = exp(-x)/(1+exp(-x))
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_dout, const ComplexType<T> arg_x) const {
ComplexType<T> dout = static_cast<ComplexType<T>>(arg_dout);
ComplexType<T> x = static_cast<ComplexType<T>>(arg_x);
// Gradient of log σ(x) is σ(-x) = exp(-x)/(1+exp(-x))
auto exp_neg_x = exp(-x); // Cache exp(-x) to avoid redundant computation
return dout * conj(exp_neg_x / (one + exp_neg_x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaHardSigmoidFunctor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
T one = static_cast<T>(1.0f);
float slope;
float offset;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"slope", &slope}, {"offset", &offset}};
}
// hard_sigmoid(x) = 0, when x <= -3
// 1, when x >= 3
// x * slope + offset, otherwise
__device__ __forceinline__ T operator()(const T x) const {
T temp = x * static_cast<T>(slope) + static_cast<T>(offset);
T temp_max = temp > zero ? temp : zero;
T temp_min = temp_max < one ? temp_max : one;
return temp_min;
}
};
template <typename T>
struct CudaHardSigmoidGradFunctor : public BaseActivationFunctor<T> {
T zero = static_cast<T>(0.0f);
T one = static_cast<T>(1.0f);
float slope;
float offset;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"slope", &slope}, {"offset", &offset}};
}
// dx = (out > 0 && out < 1) ? dout * slope : 0
__device__ __forceinline__ T operator()(const T dout, const T out) const {
return (out > zero && out < one) ? dout * static_cast<T>(slope) : zero;
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kDepOut;
}
};
template <typename T>
__device__ __forceinline__
std::conditional_t<std::is_integral<T>::value, float, T>
log_local(T x) {
static_assert(!std::is_same<T, double>::value,
"this template must be used with float or less precise type");
return static_cast<std::conditional_t<std::is_integral<T>::value, float, T>>(
::log(static_cast<double>(x)));
}
template <>
__device__ __forceinline__ float log_local<float>(float x) {
return ::log(x);
}
template <>
__device__ __forceinline__ double log_local<double>(double x) {
return ::log(x);
}
template <typename T>
struct CudaLogFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
// log(x) = log(x)
__device__ __forceinline__ U operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<U>(log_local(x));
}
};
template <typename T>
struct CudaLogFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// log(x) = log(x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_x) const {
return static_cast<ComplexType<T>>(log(arg_x));
}
};
template <typename T>
struct CudaLogGradFunctor : public BaseActivationFunctor<T> {
// dx = dout / x
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return dout / x;
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaLogGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// dx = dout / conj(x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return dout / conj(x);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaLog1pFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
// log1p(x) = log(1 + x)
__device__ __forceinline__ U operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<U>(log_local(one + x));
}
};
template <typename T>
struct CudaLog1pFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// log1p(x) = log(1 + x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_x) const {
return static_cast<ComplexType<T>>(
log(static_cast<ComplexType<T>>(1) + arg_x));
}
};
template <typename T>
struct CudaLog1pGradFunctor : public BaseActivationFunctor<T> {
T one = static_cast<T>(1.0f);
// dx = dout / (1 + x)
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return dout / (one + x);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaLog1pGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
// dx = dout / conj(1 + x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return dout / conj(one + x);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
__device__ __forceinline__
std::conditional_t<std::is_integral<T>::value, float, T>
log2_local(T x) {
static_assert(!std::is_same<T, double>::value,
"this template must be used with float or less precise type");
#if defined(__CUDA_ARCH__) || defined(__HIP_ARCH__)
// use __logf fast approximation for peak bandwidth
return __log2f(x);
#else
return ::log2(x);
#endif
}
template <>
__device__ __forceinline__ double log2_local<double>(double x) {
return ::log2(x);
}
template <typename T>
struct CudaLog2Functor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
// log2(x) = log2(x)
__device__ __forceinline__ U operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return static_cast<U>(log2_local(x));
}
};
template <typename T>
struct CudaLog2Functor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// log2(x) = log(x)/log(2)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_x) const {
return static_cast<ComplexType<T>>(log(arg_x) /
static_cast<ComplexType<T>>(log(2.0f)));
}
};
template <typename T>
struct CudaLog2GradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
T log_two = static_cast<T>(log(static_cast<MT>(2.0f)));
// dx = dout / (x * log(2))
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return dout / (x * log_two);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaLog2GradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// dx = dout / conj(x * log(2))
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return dout / conj(x * static_cast<ComplexType<T>>(log(2.0f)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
__device__ __forceinline__
std::conditional_t<std::is_integral<T>::value, float, T>
log10_local(T x) {
static_assert(!std::is_same<T, double>::value,
"this template must be used with float or less precise type");
return ::log10(x);
}
template <>
__device__ __forceinline__ double log10_local(double x) {
return ::log10(x);
}
template <typename T>
struct CudaLog10Functor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
using U = typename std::conditional_t<std::is_integral<T>::value, float, T>;
// log10(x) = log10(x)
__device__ __forceinline__ U operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
// Cast to floating-point before log10_local to avoid calling
// host-only ::log10(int) on Windows NVCC when MT is integral
using FPType = std::conditional_t<std::is_integral<MT>::value, float, MT>;
return static_cast<U>(log10_local(static_cast<FPType>(x)));
}
};
template <typename T>
struct CudaLog10Functor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// log10(x) = log(x)/log(10)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> arg_x) const {
return static_cast<ComplexType<T>>(log(arg_x) /
static_cast<ComplexType<T>>(log(10.0)));
}
};
template <typename T>
struct CudaLog10GradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// ln(10) = 2.30258509299404568402... (M_LN10)
// Using PyTorch's exact 16-digit literal from derivatives.yaml for bit-exact
// alignment: grad / (self * 2.3025850929940456)
T log_ten = static_cast<T>(2.3025850929940456);
// dx = dout / (x * ln(10))
// PyTorch computes: grad / (self * 2.3025850929940456)
// i.e., multiply x by ln(10) first, then divide grad by the product.
// This matches PyTorch's evaluation order exactly: one multiplication
// followed by one division, rather than two sequential divisions.
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return dout / (x * log_ten);
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaLog10GradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
// dx = dout / conj(x * log(10))
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return dout / conj(x * static_cast<ComplexType<T>>(log(10.0)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaSwishFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
float beta = 1.0;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"beta", &beta}};
}
// swish(x) = x / (1 + exp(-beta * x))
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
MT b = static_cast<MT>(beta);
return static_cast<T>(x / (one + exp(-b * x)));
}
};
template <typename T>
struct CudaSwishGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT one = static_cast<MT>(1.0f);
typename BaseActivationFunctor<T>::AttrPair GetAttrs() { return {{}}; }
// dx = dout * (1 + exp(-b * x) + b * x * exp(-b * x) / (1 + exp(-b * x))^2)
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
float beta = 1.0;
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
MT b = static_cast<MT>(beta);
MT temp1 = one / (one + exp(-b * x));
MT out = x * temp1;
MT temp2 = b * out;
MT temp3 = temp1 * (one - temp2);
return static_cast<T>(dout * (temp2 + temp3));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaHardSwishFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
const MT zero = static_cast<MT>(0.0f);
float threshold;
float scale;
float offset;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}, {"scale", &scale}, {"offset", &offset}};
}
// hard_swish(x) = 0, when x <= -offset
// x , when x >= threshold - offset
// x * (x + offset) / scale, otherwise
// threshold = scale = 6, offset = 3 by default
__device__ __forceinline__ T operator()(const T x) const {
const MT x_t = static_cast<MT>(x);
const MT x_offset_t = x_t + static_cast<MT>(offset);
const MT temp_max = (x_offset_t >= zero) ? x_offset_t : zero;
const MT threshold_t = static_cast<MT>(threshold);
const MT temp_min = (temp_max < threshold_t) ? temp_max : threshold_t;
return static_cast<T>(temp_min * x_t / static_cast<MT>(scale));
}
};
template <typename T>
struct CudaHardSwishGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
const MT zero = static_cast<MT>(0.0f);
const MT one = static_cast<MT>(1.0f);
const MT two = static_cast<MT>(2.0f);
float threshold;
float scale;
float offset;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"threshold", &threshold}, {"scale", &scale}, {"offset", &offset}};
}
// dx = 0, when x <= -offset
// dout , when x >= threshold - offset
// dout * (2 * x / scale + offset / scale), otherwise
// threshold = scale = 6, offset = 3 by default
__device__ __forceinline__ T operator()(const T dout, const T x) const {
const MT dout_t = static_cast<MT>(dout);
const MT x_t = static_cast<MT>(x);
const MT offset_t = static_cast<MT>(offset);
const MT scale_t = static_cast<MT>(scale);
const MT temp1 = static_cast<MT>(x_t + offset_t > zero);
const MT temp2 =
static_cast<MT>(x_t + offset_t < static_cast<MT>(threshold));
return static_cast<T>(
dout_t *
(temp1 * temp2 * (two * x_t + offset_t) / scale_t + one - temp2));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaHardSwishGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
const ComplexType<T> zero = static_cast<ComplexType<T>>(0.0f);
const ComplexType<T> one = static_cast<ComplexType<T>>(1.0f);
const ComplexType<T> two = static_cast<ComplexType<T>>(2.0f);
float threshold;
float scale;
float offset;
typename BaseActivationFunctor<ComplexType<T>>::AttrPair GetAttrs() {
return {{"threshold", &threshold}, {"scale", &scale}, {"offset", &offset}};
}
// dx = 0, when x <= -offset
// dout , when x >= threshold - offset
// dout * (2 * x / scale + offset / scale), otherwise
// threshold = scale = 6, offset = 3 by default
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
const ComplexType<T> dout_t = static_cast<ComplexType<T>>(dout);
const ComplexType<T> x_t = static_cast<ComplexType<T>>(x);
const ComplexType<T> offset_t = static_cast<ComplexType<T>>(offset);
const ComplexType<T> scale_t = static_cast<ComplexType<T>>(scale);
const ComplexType<T> temp1 =
static_cast<ComplexType<T>>(x_t + offset_t > zero);
const ComplexType<T> temp2 = static_cast<ComplexType<T>>(
x_t + offset_t < static_cast<ComplexType<T>>(threshold));
return static_cast<ComplexType<T>>(
dout_t *
conj(temp1 * temp2 * (two * x_t + offset_t) / scale_t + one - temp2));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaCeilFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// ceil(x) = ceil(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
if constexpr ((std::is_same<T, uint8_t>::value) ||
(std::is_same<T, int8_t>::value) ||
(std::is_same<T, uint16_t>::value) ||
(std::is_same<T, int16_t>::value) ||
(std::is_same<T, int>::value) ||
(std::is_same<T, int64_t>::value)) {
return static_cast<T>(x);
} else {
return static_cast<T>(ceil(x));
}
}
};
template <typename T>
__device__ __forceinline__
typename std::enable_if<std::is_integral<T>::value, int64_t>::type
compute_pow(const T a, const double b) {
// TODO(wujionghao): A potential speed improvement is supporting different
// types in C++.
// On CUDAPlace, pow(3, 1) calls pow(float, float), and
// it will return a float number like 2.99... , which floor to 2
// when cast to int by default and it is wrong.
// Use llrint to cast it to the nearest integer, which is 3.
return llrint(pow(static_cast<double>(a), b));
}
template <typename T, typename MT>
__device__ __forceinline__
typename std::enable_if<!std::is_integral<T>::value, MT>::type
compute_pow(const T a, const MT b) {
return pow(static_cast<MT>(a), b);
}
template <typename T, typename MT>
__device__ __forceinline__
typename std::enable_if<!std::is_integral<T>::value, ComplexType<MT>>::type
compute_pow(const ComplexType<T> a, const ComplexType<MT> b) {
return pow(static_cast<ComplexType<MT>>(a), b);
}
template <typename T>
struct BaseCudaPowFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT factor;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"factor", &factor}};
}
void SetFactor(double factor) { this->factor = static_cast<MT>(factor); }
};
template <typename T>
struct BaseCudaPowGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT factor;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"factor", &factor}};
}
void SetFactor(double factor) { this->factor = static_cast<MT>(factor); }
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaPowFunctor : public BaseCudaPowFunctor<T> {
__device__ __forceinline__ T operator()(const T x) const {
return static_cast<T>(compute_pow(x, this->factor));
}
};
template <typename T>
struct CudaPowGradFunctor : public BaseCudaPowGradFunctor<T> {
// dx = dout * n * pow(x, n - 1)
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return dout *
static_cast<T>(this->factor * compute_pow(x, this->factor - 1));
}
};
template <typename T>
struct CudaPowGradFunctor<ComplexType<T>>
: public BaseCudaPowGradFunctor<ComplexType<T>> {
using MT = typename MPTypeTrait<ComplexType<T>>::Type;
MT one = static_cast<MT>(1.0f);
// dx = dout * (4 * (x*x*x))
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return dout * static_cast<ComplexType<T>>(
conj(this->factor * compute_pow(x, this->factor - one)));
}
};
template <typename T>
struct CudaCubeFunctor : public BaseActivationFunctor<T> {
// cube(x) = x * x * x
__device__ __forceinline__ T operator()(const T x) const { return x * x * x; }
};
template <typename T>
struct CudaCubeGradFunctor : public BaseActivationFunctor<T> {
T three = static_cast<T>(3.0f);
// dx = dout * 3 * x * x
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return dout * (three * (x * x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaCubeGradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> three = static_cast<ComplexType<T>>(3.0f);
// dx = dout * conj(3 * x * x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout * conj(three * (x * x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaPow4GradFunctor : public BaseActivationFunctor<T> {
T four = static_cast<T>(4.0f);
// dx = dout * 4 * x * x * x
__device__ __forceinline__ T operator()(const T dout, const T x) const {
return dout * (four * (x * x * x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaPow4GradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> four = static_cast<ComplexType<T>>(4.0f);
// dx = dout * conj(4 * x * x * x)
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout * conj(four * (x * x * x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
// for pow(x, 1.5)
template <typename T>
struct CudaPow1p5GradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT f1p5 = static_cast<T>(1.5f);
// dx = dout * 1.5 * sqrt(x)
__device__ __forceinline__ T operator()(const T dout, const T arg_x) const {
MT x = static_cast<MT>(arg_x);
return dout * static_cast<T>(f1p5 * sqrt(x));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaPow1p5GradFunctor<ComplexType<T>>
: public BaseActivationFunctor<ComplexType<T>> {
ComplexType<T> f1p5 = static_cast<ComplexType<T>>(1.5f);
// dx = dout * conj(1.5 * sqrt(x))
__device__ __forceinline__ ComplexType<T> operator()(
const ComplexType<T> dout, const ComplexType<T> x) const {
return static_cast<ComplexType<T>>(dout * conj(f1p5 * sqrt(x)));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
template <typename T>
struct CudaFloorFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// floor(x) = floor(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
if constexpr ((std::is_same<T, uint8_t>::value) ||
(std::is_same<T, int8_t>::value) ||
(std::is_same<T, uint16_t>::value) ||
(std::is_same<T, int16_t>::value) ||
(std::is_same<T, int>::value) ||
(std::is_same<T, int64_t>::value)) {
return static_cast<T>(x);
} else {
return static_cast<T>(floor(x));
}
}
};
template <typename T, typename Enable = void>
struct CudaRintFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
// rint(x) = rint(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
if (isnan(x) || isinf(x)) return arg_x;
return static_cast<T>(std::rint(x));
}
};
template <typename T>
struct CudaRintFunctor<T, std::enable_if_t<std::is_integral_v<T>>>
: public BaseActivationFunctor<T> {
// rint(x) = x
__device__ __forceinline__ T operator()(const T arg_x) const { return arg_x; }
};
template <typename T, typename Enable = void>
struct CudaRoundFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
int decimals;
std::vector<std::pair<const char*, int*>> GetAttrs() {
return {{"decimals", &decimals}};
}
// round(x) = round(x)
__device__ __forceinline__ T operator()(const T arg_x) const {
MT x = static_cast<MT>(arg_x);
if (isnan(x) || isinf(x)) return arg_x;
if (decimals == 0) {
return static_cast<T>(std::rint(x));
} else if (decimals > 0) {
MT ten_pow_decimals = pow(static_cast<MT>(10), static_cast<MT>(decimals));
return static_cast<T>(rint(x * static_cast<MT>(ten_pow_decimals)) /
ten_pow_decimals);
} else {
MT ten_pow_decimals =
pow(static_cast<MT>(10), static_cast<MT>(-decimals));
return static_cast<T>(rint(x / static_cast<MT>(ten_pow_decimals)) *
ten_pow_decimals);
}
}
};
template <typename T>
struct CudaRoundFunctor<T, std::enable_if_t<std::is_integral_v<T>>>
: public BaseActivationFunctor<T> {
int decimals;
std::vector<std::pair<const char*, int*>> GetAttrs() {
return {{"decimals", &decimals}};
}
// round(x) = round(x)
__device__ __forceinline__ T operator()(const T arg_x) const { return arg_x; }
};
template <typename T>
struct CudaRoundFunctor<phi::dtype::complex<T>>
: public BaseActivationFunctor<phi::dtype::complex<T>> {
using MT = typename MPTypeTrait<T>::Type;
int decimals;
std::vector<std::pair<const char*, int*>> GetAttrs() {
return {{"decimals", &decimals}};
}
__device__ __forceinline__ phi::dtype::complex<T> operator()(
const phi::dtype::complex<T> arg_x) const {
MT real_part = static_cast<MT>(arg_x.real);
MT imag_part = static_cast<MT>(arg_x.imag);
bool real_special = isnan(real_part) || isinf(real_part);
bool imag_special = isnan(imag_part) || isinf(imag_part);
MT real, imag;
if (decimals == 0) {
real = real_special ? real_part : rint(real_part);
imag = imag_special ? imag_part : rint(imag_part);
} else if (decimals > 0) {
MT ten_pow_decimals = pow(static_cast<MT>(10), static_cast<MT>(decimals));
real = real_special
? real_part
: rint(real_part * ten_pow_decimals) / ten_pow_decimals;
imag = imag_special
? imag_part
: rint(imag_part * ten_pow_decimals) / ten_pow_decimals;
} else {
MT ten_pow_decimals =
pow(static_cast<MT>(10), static_cast<MT>(-decimals));
real = real_special
? real_part
: rint(real_part / ten_pow_decimals) * ten_pow_decimals;
imag = imag_special
? imag_part
: rint(imag_part / ten_pow_decimals) * ten_pow_decimals;
}
return phi::dtype::complex<T>(static_cast<T>(real), static_cast<T>(imag));
}
};
// GradFunctor for ceil, floor and round
template <typename T>
struct CudaZeroGradFunctor : public BaseActivationFunctor<T> {
__device__ __forceinline__ T operator()(const T x) const {
return static_cast<T>(0.0f);
}
static constexpr ActBwdOpFwdDeps FwdDeps() {
return ActBwdOpFwdDeps::kNoDeps;
}
};
template <typename T>
struct CudaCELUFunctor : public BaseActivationFunctor<T> {
using CT = typename MPTypeTrait<T>::Type;
CT zero = static_cast<CT>(0.0f);
CT one = static_cast<CT>(1.0f);
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
// celu(x) = max(0, x) + min(0, alpha * (exp(x/alpha) - 1))
__device__ __forceinline__ T operator()(const T arg_x) const {
CT x = static_cast<CT>(arg_x);
CT temp = static_cast<CT>(alpha) * (exp(x / static_cast<CT>(alpha)) - one);
CT res = (x > zero ? x : zero) + (temp > zero ? zero : temp);
return static_cast<T>(res);
}
};
template <typename T>
struct CudaCELUGradFunctor : public BaseActivationFunctor<T> {
using MT = typename MPTypeTrait<T>::Type;
MT zero = static_cast<MT>(0.0f);
MT one = static_cast<MT>(1.0f);
float alpha;
typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
return {{"alpha", &alpha}};
}
// dx = dout, if alpha > 0 and x > 0
// dx = dout * (x/alpha).exp(), if alpha > 0 and x <= 0
// dx = dout , if alpha < 0 and x > 0
// dx = dout * (x/alpha).exp(), if alpha < 0 and x <=0
__device__ __forceinline__ T operator()(const T arg_dout,
const T arg_x) const {
MT dout = static_cast<MT>(arg_dout);
MT x = static_cast<MT>(arg_x);
MT a = static_cast<MT>(alpha);
MT temp_a_pos = static_cast<MT>(alpha > 0.0f);
MT temp_a_neg = static_cast<MT>(alpha <= 0.0f);
MT temp_x_pos = static_cast<MT>(x > zero);
MT temp_x_neg = static_cast<MT>(x <= zero);
return static_cast<T>(
dout *
(temp_a_pos * temp_x_pos + temp_a_pos * temp_x_neg * exp(x / a) +
temp_a_neg * temp_x_pos + exp(x / a) * temp_a_neg * temp_x_neg));
}
static constexpr ActBwdOpFwdDeps FwdDeps() { return ActBwdOpFwdDeps::kDepX; }
};
#endif
template <typename T>
struct SwiGLUFunctor {
using MT = typename MPTypeTrait<T>::Type;
HOSTDEVICE T operator()(T x, T y) const {
MT mp_x = static_cast<MT>(x);
MT mp_y = static_cast<MT>(y);
MT one = static_cast<MT>(1);
return static_cast<T>(mp_y * mp_x / (one + exp(-mp_x)));
}
};
template <typename T, bool HasDX = true, bool HasDY = true>
struct SwiGLUGradFunctor {
using MT = typename MPTypeTrait<T>::Type;
HOSTDEVICE void operator()(T x, T y, T dz, T* dx, T* dy) const {
MT one = static_cast<MT>(1);
MT mp_x = static_cast<MT>(x);
MT mp_dz = static_cast<MT>(dz);
MT sigmoid = one / (one + exp(-mp_x));
MT tmp = mp_x * sigmoid;
if (HasDX) {
MT mp_y = static_cast<MT>(y);
*dx = static_cast<T>(mp_dz * mp_y * sigmoid * (one + mp_x - tmp));
}
if (HasDY) {
*dy = static_cast<T>(mp_dz * tmp);
}
}
};
} // namespace funcs
} // namespace phi
在新工单中引用 查看 Git Blame 复制永久链接