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

Co-authored-by: QingshuChen <Qingshu.Chen714@gamil.com>

Author: QingshuChen

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",

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>,