[PTen]Reshape Kernel Refactor #37164
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -15,7 +15,12 @@ limitations under the License. */ | |
| #include <string> | ||
|
|
||
| #include "paddle/fluid/framework/op_registry.h" | ||
| #include "paddle/fluid/framework/pten_utils.h" | ||
|
|
||
| // only can include the headers in paddle/pten/api dirs | ||
| #include "paddle/pten/api/lib/utils/tensor_utils.h" | ||
| #include "paddle/pten/include/core.h" | ||
| #include "paddle/pten/include/manipulation.h" | ||
| namespace paddle { | ||
| namespace framework { | ||
| class InferShapeContext; | ||
|
|
@@ -248,13 +253,6 @@ class ReshapeOp : public framework::OperatorWithKernel { | |
| auto input_data_type = | ||
| framework::OperatorWithKernel::IndicateVarDataType(ctx, "X"); | ||
|
|
||
| return framework::OpKernelType(input_data_type, ctx.GetPlace()); | ||
| } | ||
|
|
||
|
|
@@ -366,13 +364,6 @@ class ReshapeGradOp : public framework::OperatorWithKernel { | |
| auto input_data_type = | ||
| framework::OperatorWithKernel::IndicateVarDataType(ctx, "X"); | ||
|
|
||
| return framework::OpKernelType(input_data_type, ctx.GetPlace()); | ||
| } | ||
| }; | ||
|
|
@@ -382,42 +373,117 @@ class ReshapeKernel { | |
| void operator()(const framework::ExecutionContext &ctx) const { | ||
| auto *out = ctx.Output<framework::LoDTensor>("Out"); | ||
| auto *in = ctx.Input<framework::LoDTensor>("X"); | ||
| // framework::DDim out_dims = out->dims(); | ||
| auto pt_x = paddle::experimental::MakePtenDenseTensor(*in); | ||
|
|
||
| // we can't MakePtenDenseTensor by out, because reshape will realloc memory | ||
| // and this will throw error(can't realloc shared memory) in current | ||
| // DenseTensor | ||
| // design. So, codes below create a tmp densetensor for output. | ||
| // TODO(YuanRisheng) we can use MakePtenDenseTensor after #36916 merge. | ||
| const auto alloc = std::make_shared<paddle::experimental::DefaultAllocator>( | ||
|
There was a problem hiding this comment. #36916 合入之后，会支持realloc，这里需要修改写法，建议记个TODO There was a problem hiding this comment. 目前这样的写法可能会对模型性能造成影响，后面需要尽快修改 |
||
| paddle::platform::CPUPlace()); | ||
| pten::DenseTensorMeta meta{pten::TransToPtenDataType(in->type()), | ||
| in->dims(), | ||
| pten::TransToPtenDataLayout(in->layout())}; | ||
| auto pt_out_tmp = | ||
| std::make_shared<pten::DenseTensor>(alloc, std::move(meta)); | ||
| pten::DenseTensor *pt_out = nullptr; | ||
| if (in == out) { | ||
| pt_out = pt_x.get(); | ||
| } else { | ||
| pt_out = pt_out_tmp.get(); | ||
| } | ||
|
|
||
| auto list_new_shape_tensor = | ||
| ctx.MultiInput<framework::Tensor>("ShapeTensor"); | ||
| auto *shape_tensor = ctx.HasInput("Shape") | ||
| ? ctx.Input<framework::LoDTensor>("Shape") | ||
| : nullptr; | ||
| if (list_new_shape_tensor.size() > 0) { | ||
| // have shape tensor | ||
| std::vector<pten::DenseTensor> pt_vec_shape; | ||
| for (auto &tensor : list_new_shape_tensor) { | ||
| if (platform::is_gpu_place(tensor->place()) || | ||
| platform::is_xpu_place(tensor->place())) { | ||
| framework::Tensor temp; | ||
| TensorCopySync(*tensor, platform::CPUPlace(), &temp); | ||
| pt_vec_shape.push_back( | ||
| std::move(*(paddle::experimental::MakePtenDenseTensor(temp)))); | ||
| } else { | ||
| pt_vec_shape.push_back( | ||
| std::move(*(paddle::experimental::MakePtenDenseTensor(*tensor)))); | ||
| } | ||
| } | ||
| if (platform::is_cpu_place(ctx.GetPlace())) { | ||
| auto &dev_ctx = ctx.device_context<platform::CPUDeviceContext>(); | ||
| pten::ReshapeFromVectorDT(dev_ctx, *pt_x.get(), pt_vec_shape, pt_out); | ||
| } | ||
| #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) | ||
| if (platform::is_gpu_place(ctx.GetPlace())) { | ||
| auto &dev_ctx = ctx.device_context<platform::CUDADeviceContext>(); | ||
| pten::ReshapeFromVectorDT(dev_ctx, *pt_x.get(), pt_vec_shape, pt_out); | ||
| } | ||
| #endif | ||
| #ifdef PADDLE_WITH_XPU | ||
| if (platform::is_xpu_place(ctx.GetPlace())) { | ||
| auto &dev_ctx = ctx.device_context<platform::XPUDeviceContext>(); | ||
| pten::ReshapeFromVectorDT(dev_ctx, *pt_x.get(), pt_vec_shape, pt_out); | ||
| } | ||
| #endif | ||
| } else if (shape_tensor) { | ||
| std::unique_ptr<pten::DenseTensor> pt_shape; | ||
| if (platform::is_gpu_place(shape_tensor->place()) || | ||
| platform::is_xpu_place(shape_tensor->place())) { | ||
| framework::Tensor temp; | ||
| TensorCopySync(*shape_tensor, platform::CPUPlace(), &temp); | ||
| pt_shape = paddle::experimental::MakePtenDenseTensor(temp); | ||
| } else { | ||
| pt_shape = paddle::experimental::MakePtenDenseTensor(*shape_tensor); | ||
| } | ||
|
|
||
| if (platform::is_cpu_place(ctx.GetPlace())) { | ||
| auto &dev_ctx = ctx.device_context<platform::CPUDeviceContext>(); | ||
| pten::ReshapeFromDT(dev_ctx, *pt_x.get(), *pt_shape.get(), pt_out); | ||
| } | ||
| #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) | ||
| if (platform::is_gpu_place(ctx.GetPlace())) { | ||
| auto &dev_ctx = ctx.device_context<platform::CUDADeviceContext>(); | ||
| pten::ReshapeFromDT(dev_ctx, *pt_x.get(), *pt_shape.get(), pt_out); | ||
| } | ||
| #endif | ||
| #ifdef PADDLE_WITH_XPU | ||
| if (platform::is_xpu_place(ctx.GetPlace())) { | ||
| auto &dev_ctx = ctx.device_context<platform::XPUDeviceContext>(); | ||
| pten::ReshapeFromDT(dev_ctx, *pt_x.get(), *pt_shape.get(), pt_out); | ||
| } | ||
| #endif | ||
| } else { | ||
| auto &shape_vec = ctx.Attr<std::vector<int>>("shape"); | ||
| if (platform::is_cpu_place(ctx.GetPlace())) { | ||
| auto &dev_ctx = ctx.device_context<platform::CPUDeviceContext>(); | ||
| pten::ReshapeFromVectorVal(dev_ctx, *pt_x.get(), shape_vec, pt_out); | ||
| } | ||
| #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) | ||
| if (platform::is_gpu_place(ctx.GetPlace())) { | ||
| auto &dev_ctx = ctx.device_context<platform::CUDADeviceContext>(); | ||
| pten::ReshapeFromVectorVal(dev_ctx, *pt_x.get(), shape_vec, pt_out); | ||
| } | ||
| #endif | ||
| #ifdef PADDLE_WITH_XPU | ||
| if (platform::is_xpu_place(ctx.GetPlace())) { | ||
| auto &dev_ctx = ctx.device_context<platform::XPUDeviceContext>(); | ||
| pten::ReshapeFromVectorVal(dev_ctx, *pt_x.get(), shape_vec, pt_out); | ||
| } | ||
| #endif | ||
| } | ||
| // non-inplace need move all result from pt_out to out, inplace need set | ||
| // result dims. | ||
| if (in != out) { | ||
| paddle::experimental::MovesStorage(pt_out, static_cast<Tensor *>(out)); | ||
| } else { | ||
| out->Resize(pt_out->dims()); | ||
| } | ||
| } | ||
| }; | ||
|
|
||
|
|
@@ -479,6 +545,21 @@ class Reshape2Op : public ReshapeOp { | |
|
|
||
| ReshapeOp::InferShape(ctx); | ||
| } | ||
|
|
||
| framework::KernelSignature GetExpectedPtenKernelArgs( | ||
| const framework::ExecutionContext &ctx) const override { | ||
| auto multi_inputs = ctx.MultiInput<framework::Tensor>("ShapeTensor"); | ||
| if (multi_inputs.size() > 0) { | ||
| return framework::KernelSignature( | ||
| "reshape2.mulhost.mid", {"X", "ShapeTensor"}, {}, {"XShape", "Out"}); | ||
| } else if (ctx.HasInput("Shape")) { | ||
| return framework::KernelSignature("reshape2.host.mid", {"X", "Shape"}, {}, | ||
| {"XShape", "Out"}); | ||
| } else { | ||
| return framework::KernelSignature("reshape2.mid", {"X"}, {"shape"}, | ||
| {"XShape", "Out"}); | ||
| } | ||
| } | ||
| }; | ||
|
|
||
| class Reshape2OpMaker : public ReshapeOpMaker { | ||
|
|
@@ -557,13 +638,6 @@ class Reshape2GradOp : public framework::OperatorWithKernel { | |
| auto input_data_type = framework::OperatorWithKernel::IndicateVarDataType( | ||
| ctx, framework::GradVarName("Out")); | ||
|
|
||
| return framework::OpKernelType(input_data_type, ctx.GetPlace()); | ||
| } | ||
|
|
||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
shape需要使用vector<int64_t>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
reshape maker定义中使用的vector，使用vector<int64_t>会有兼容问题