Skip to content

[PHI decoupling] migrate transpose_op.cu.h and gpu_utils.h to phi #48286

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

Merged
merged 15 commits into from
Nov 30, 2022
Merged

[PHI decoupling] migrate transpose_op.cu.h and gpu_utils.h to phi #48286

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
47dc074
migrate transpose_op.cu.h and gpu_utils.h
Patrick-Star125 Nov 23, 2022
b159631
format code style
Patrick-Star125 Nov 23, 2022
9723302
Merge remote-tracking branch 'paddle/develop' into develop
Patrick-Star125 Nov 24, 2022
782ca36
fix some problems
Patrick-Star125 Nov 24, 2022
ff73bd0
format code
Patrick-Star125 Nov 24, 2022
c0238b4
Merge remote-tracking branch 'paddle/develop' into develop
Patrick-Star125 Nov 25, 2022
29216da
Merge remote-tracking branch 'paddle/develop' into develop
Patrick-Star125 Nov 25, 2022
7c40def
reset tranpose_op.cc
Patrick-Star125 Nov 25, 2022
cf22538
test commit
Patrick-Star125 Nov 25, 2022
0baef02
Merge remote-tracking branch 'paddle/develop' into develop
Patrick-Star125 Nov 25, 2022
2bd9b22
recover transpose_op.h
Patrick-Star125 Nov 25, 2022
16acf1f
Merge remote-tracking branch 'paddle/develop' into develop
Patrick-Star125 Nov 25, 2022
aa5481b
Merge remote-tracking branch 'paddle/develop' into develop
Patrick-Star125 Nov 28, 2022
c26e979
delete transpose_op.h
Patrick-Star125 Nov 28, 2022
3490d8a
adjust header files order in transpose_op.cc
Patrick-Star125 Nov 29, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions paddle/fluid/operators/fused/fmha_ref.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,12 @@ limitations under the License. */

#include "paddle/fluid/operators/dropout_impl.cu.h"
#include "paddle/fluid/operators/fused/fused_softmax_mask.cu.h"
#include "paddle/fluid/operators/transpose_op.cu.h"
#include "paddle/phi/kernels/funcs/broadcast_function.h"
#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
#include "paddle/phi/kernels/funcs/elementwise_base.h"
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
#include "paddle/phi/kernels/funcs/functors.h"
#include "paddle/phi/kernels/funcs/transpose_op.cu.h"
#include "paddle/phi/kernels/gpudnn/softmax_gpudnn.h"

namespace paddle {
Expand Down Expand Up @@ -98,7 +98,7 @@ class FMHARef {
// transpose with perm [2, 0, 3, 1, 4],
// output_shape: [3, bs, num_head, seq_len, head_dim]
std::vector<int> perm_1 = {2, 0, 3, 1, 4};
TransposeGPUKernelDriver<T>(
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, qkv_input_tensor, perm_1, transpose_2_out_tensor);
T* qkv_data = transpose_2_out_tensor->data<T>();
T* qk_out_data = qk_out_tensor->data<T>();
Expand Down Expand Up @@ -254,7 +254,7 @@ class FMHARef {
// transpose: [0, 2, 1, 3]
// output shape: [batch_size, seq_len, num_heads, head_dim]
std::vector<int> perm_3 = {0, 2, 1, 3};
TransposeGPUKernelDriver<T>(
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, *qktv_out_tensor, perm_3, fmha_out_tensor);
}

Expand Down Expand Up @@ -428,7 +428,7 @@ class FMHARef {
// transpose: [0, 2, 1, 3]
// output shape: [batch_size, seq_len, num_heads, head_dim]
std::vector<int> perm_3 = {0, 2, 1, 3};
TransposeGPUKernelDriver<T>(
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, *qktv_out_tensor, perm_3, fmha_out_tensor);
}

Expand Down Expand Up @@ -470,7 +470,7 @@ class FMHARef {

// transpose bw
std::vector<int> perm_3 = {0, 2, 1, 3};
TransposeGPUKernelDriver<T>(
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, fmha_out_grad_tensor, perm_3, qktv_out_grad_tensor);

// recall batchedgemm(nn) fw: softmax_out_data(x) * v_ptr(y) =
Expand Down Expand Up @@ -648,7 +648,7 @@ class FMHARef {

// transpose bw
std::vector<int> perm_1 = {1, 3, 0, 2, 4};
TransposeGPUKernelDriver<T>(
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, *transpose_2_out_grad_tensor, perm_1, qkv_input_grad_tensor);
}

Expand Down
27 changes: 16 additions & 11 deletions paddle/fluid/operators/fused/fused_gate_attention.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,11 @@ limitations under the License. */

#pragma once

#include "paddle/fluid/operators/transpose_op.cu.h"
#include "paddle/phi/kernels/funcs/broadcast_function.h"
#include "paddle/phi/kernels/funcs/elementwise_base.h"
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
#include "paddle/phi/kernels/funcs/reduce_function.h"
#include "paddle/phi/kernels/funcs/transpose_op.cu.h"
#include "paddle/phi/kernels/gpudnn/softmax_gpudnn.h"

namespace paddle {
Expand Down Expand Up @@ -626,9 +626,12 @@ class FMHAGateRef {
phi::DenseTensor* k_transpose_out,
phi::DenseTensor* v_transpose_out) {
std::vector<int> perm = {0, 1, 3, 2, 4};
TransposeGPUKernelDriver<T>(dev_ctx_, q_out, perm, q_transpose_out);
TransposeGPUKernelDriver<T>(dev_ctx_, k_out, perm, k_transpose_out);
TransposeGPUKernelDriver<T>(dev_ctx_, v_out, perm, v_transpose_out);
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, q_out, perm, q_transpose_out);
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, k_out, perm, k_transpose_out);
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, v_out, perm, v_transpose_out);
}

void ComputeQKVTransposeBackward(const phi::DenseTensor& q_transpose_out_grad,
Expand All @@ -638,11 +641,11 @@ class FMHAGateRef {
phi::DenseTensor* k_out_grad,
phi::DenseTensor* v_out_grad) {
std::vector<int> perm = {0, 1, 3, 2, 4};
TransposeGPUKernelDriver<T>(
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, q_transpose_out_grad, perm, q_out_grad);
TransposeGPUKernelDriver<T>(
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, k_transpose_out_grad, perm, k_out_grad);
TransposeGPUKernelDriver<T>(
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, v_transpose_out_grad, perm, v_out_grad);
}

Expand All @@ -651,14 +654,15 @@ class FMHAGateRef {
void ComputeQKVTransposeForward(const phi::DenseTensor& qkv_out,
phi::DenseTensor* qkv_transpose_out) {
std::vector<int> perm = {3, 0, 1, 4, 2, 5};
TransposeGPUKernelDriver<T>(dev_ctx_, qkv_out, perm, qkv_transpose_out);
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, qkv_out, perm, qkv_transpose_out);
}

void ComputeQKVTransposeBackward(
const phi::DenseTensor& qkv_transpose_out_grad,
phi::DenseTensor* qkv_out_grad) {
std::vector<int> perm = {1, 2, 4, 0, 3, 5};
TransposeGPUKernelDriver<T>(
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, qkv_transpose_out_grad, perm, qkv_out_grad);
}

Expand All @@ -667,13 +671,14 @@ class FMHAGateRef {
void ComputeQKTVTransposeForward(const phi::DenseTensor& qktv_out,
phi::DenseTensor* fmha_out) {
std::vector<int> perm = {0, 1, 3, 2, 4};
TransposeGPUKernelDriver<T>(dev_ctx_, qktv_out, perm, fmha_out);
phi::funcs::TransposeGPUKernelDriver<T>(dev_ctx_, qktv_out, perm, fmha_out);
}

void ComputeQKTVTransposeBackward(const phi::DenseTensor& fmha_out_grad,
phi::DenseTensor* qktv_out_grad) {
std::vector<int> perm = {0, 1, 3, 2, 4};
TransposeGPUKernelDriver<T>(dev_ctx_, fmha_out_grad, perm, qktv_out_grad);
phi::funcs::TransposeGPUKernelDriver<T>(
dev_ctx_, fmha_out_grad, perm, qktv_out_grad);
}

// qk_out = qk_out + nonbatched_bias + src_mask
Expand Down
22 changes: 11 additions & 11 deletions paddle/fluid/framework/gpu_utils.h → paddle/phi/kernels/funcs/gpu_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,11 @@

#include <array>

#include "paddle/fluid/platform/enforce.h"
#include "paddle/phi/core/enforce.h"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

把gpu_utils.h这个文件放到phi/backends/gpu目录下吧

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

#include "unsupported/Eigen/CXX11/Tensor"

namespace paddle {
namespace framework {
namespace phi {
namespace funcs {

template <typename T, int Size, T DefaultValue>
struct DeviceArray {
Expand Down Expand Up @@ -110,16 +110,16 @@ IntType CeilOrFloor(IntType x, IntType deviser) {
PADDLE_ENFORCE_GT(
deviser,
0,
platform::errors::InvalidArgument("deviser should be greater than 0, "
"but received is:%d",
deviser));
phi::errors::InvalidArgument("deviser should be greater than 0, "
"but received is:%d",
deviser));

PADDLE_ENFORCE_GT(
x,
0,
platform::errors::InvalidArgument("input should be greater than 0, "
"but received is:%d",
x));
phi::errors::InvalidArgument("input should be greater than 0, "
"but received is:%d",
x));

const IntType round_to_zero = x / deviser;
const IntType inte_result = round_to_zero * deviser;
Expand All @@ -140,5 +140,5 @@ IntType CeilOrFloor(IntType x, IntType deviser) {
}
}

} // namespace framework
} // namespace paddle
} // namespace funcs
} // namespace phi
Loading