fix kunlun sharding v2 bug. Sharding v2 need inplace assign (#64209)

Co-authored-by: QingshuChen <Qingshu.Chen714@gamil.com>
PaddlePaddle · May 11, 2024 · bedd212 · bedd212
1 parent 32a4fe7
commit bedd212
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -64,12 +64,15 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT64,
                      phi::DataType::INT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::BOOL})},
       {"assign_value",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
                      phi::DataType::INT64,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::BOOL})},
       {"atan", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"atan_grad",

diff --git a/paddle/phi/kernels/assign_kernel.cc b/paddle/phi/kernels/assign_kernel.cc
@@ -196,6 +196,8 @@ PD_REGISTER_KERNEL(assign_value,
                    bool,
                    int,
                    float,
+                   phi::dtype::bfloat16,
+                   phi::dtype::float16,
                    double,
                    int64_t,
                    phi::dtype::complex<float>,