Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix pool gpu kernel #5446

Merged
merged 6 commits into from
Jul 9, 2021
Merged

fix pool gpu kernel #5446

Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e166bf8
fix pool gpu kernel
simonJJJ Jul 9, 2021
2e9b427
Merge branch 'master' into fix_pooling_kernel
simonJJJ Jul 9, 2021
6b75b2c
Merge branch 'master' into fix_pooling_kernel
oneflow-ci-bot Jul 9, 2021
cd83fc3
Merge branch 'master' into fix_pooling_kernel
oneflow-ci-bot Jul 9, 2021
cda450b
Merge branch 'master' into fix_pooling_kernel
oneflow-ci-bot Jul 9, 2021
df60db1
Merge branch 'master' into fix_pooling_kernel
oneflow-ci-bot Jul 9, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
12 changes: 6 additions & 6 deletions oneflow/user/kernels/pool_gpu_kernel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,9 @@ class CudnnPoolDesc final {

class GPUPoolOpKernelState final : public user_op::OpKernelState {
public:
GPUPoolOpKernelState(const int32_t dim, const std::string& pooling_type, const Shape& x_shape,
const Shape& y_shape, const std::string& data_format, const DataType& dtype,
const Params3D& params_3d)
GPUPoolOpKernelState(const int32_t dim, const std::string& pooling_type, const ShapeView& x_shape,
const ShapeView& y_shape, const std::string& data_format,
const DataType& dtype, const Params3D& params_3d)
: dim_(dim), pooling_type_(pooling_type) {
Reset(dim, pooling_type, x_shape, y_shape, data_format, dtype, params_3d);
}
Expand Down Expand Up @@ -82,7 +82,7 @@ class GPUPoolOpKernelState final : public user_op::OpKernelState {
const int32_t& dim, const std::string& pooling_type, user_op::KernelComputeContext* ctx) {
if (pooling_type != "MAX" && pooling_type != "AVG") { UNIMPLEMENTED(); }
const user_op::TensorDesc* x_desc = ctx->TensorDesc4ArgNameAndIndex("x", 0);
const Shape& x_shape = x_desc->shape();
const ShapeView& x_shape = ctx->Tensor4ArgNameAndIndex("x", 0)->shape();
const std::string& data_format = ctx->Attr<std::string>("data_format");
const std::string& padding = ctx->Attr<std::string>("padding");
const auto& padding_before = ctx->Attr<std::vector<int32_t>>("padding_before");
Expand All @@ -92,8 +92,8 @@ class GPUPoolOpKernelState final : public user_op::OpKernelState {
const bool ceil_mode = ctx->Attr<bool>("ceil_mode");
const Params3D params_3d(dim, x_shape, data_format, padding, padding_before, padding_after,
pool_size, strides, ceil_mode);
const Shape y_shape = ctx->TensorDesc4ArgNameAndIndex("y", 0)->shape();
const DataType dtype = x_desc->data_type();
const ShapeView& y_shape = ctx->Tensor4ArgNameAndIndex("y", 0)->shape();
const DataType dtype = ctx->Tensor4ArgNameAndIndex("x", 0)->data_type();
return std::make_shared<GPUPoolOpKernelState>(dim, pooling_type, x_shape, y_shape, data_format,
dtype, params_3d);
}
Expand Down