|
|
|
@ -72,25 +72,6 @@ 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>
|
|
|
|
|
struct StackGradFunctor {
|
|
|
|
|
HOSTDEVICE StackGradFunctor(const VecDxType &dx, const T *dy, int n, int post)
|
|
|
|
@ -110,14 +91,6 @@ struct StackGradFunctor {
|
|
|
|
|
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>
|
|
|
|
|
static inline void StackGradFunctorForRange(const DeviceContext &ctx,
|
|
|
|
|
const VecDxType &dx, const T *dy,
|
|
|
|
|