提交 d91740ac 编写于 作者: Y Yihua Xu

Revert "Remove the remnant code (test=develop)"

This reverts commit be506703.
上级 be506703
...@@ -72,6 +72,25 @@ class StackOpMaker : public framework::OpProtoAndCheckerMaker { ...@@ -72,6 +72,25 @@ class StackOpMaker : public framework::OpProtoAndCheckerMaker {
} }
}; };
template <typename VecXType, typename T>
struct StackFunctor {
HOSTDEVICE StackFunctor(const VecXType &x, T *y, int n, int post)
: x_(x), y_(y), n_(n), post_(post) {}
HOSTDEVICE void operator()(int idx) {
int i = idx / (n_ * post_);
int which_x = idx / post_ - i * n_;
int x_index = i * post_ + idx % post_;
y_[idx] = x_[which_x][x_index];
}
private:
VecXType x_;
T *y_;
int n_;
int post_;
};
template <typename VecDxType, typename T> template <typename VecDxType, typename T>
struct StackGradFunctor { struct StackGradFunctor {
HOSTDEVICE StackGradFunctor(const VecDxType &dx, const T *dy, int n, int post) HOSTDEVICE StackGradFunctor(const VecDxType &dx, const T *dy, int n, int post)
...@@ -91,6 +110,14 @@ struct StackGradFunctor { ...@@ -91,6 +110,14 @@ struct StackGradFunctor {
int post_; int post_;
}; };
template <typename DeviceContext, typename VecXType, typename T>
static inline void StackFunctorForRange(const DeviceContext &ctx,
const VecXType &x, T *y, int total_num,
int n, int post) {
platform::ForRange<DeviceContext> for_range(ctx, total_num);
for_range(StackFunctor<VecXType, T>(x, y, n, post));
}
template <typename DeviceContext, typename VecDxType, typename T> template <typename DeviceContext, typename VecDxType, typename T>
static inline void StackGradFunctorForRange(const DeviceContext &ctx, static inline void StackGradFunctorForRange(const DeviceContext &ctx,
const VecDxType &dx, const T *dy, const VecDxType &dx, const T *dy,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册