[INTEL_HPU] Add set value OP (#1852)

Signed-off-by: Fei Wang <[email protected]>

Author: feiwan1

backends/intel_hpu/kernels/funcs.h

@@ -206,6 +206,34 @@ inline paddle::Tensor copy_tensor_wrapper(const phi::CustomContext* dev_ctx,
   return paddle::Tensor(dst_dt);
 }
 
+/**
+ * CPU -> INTEL_HPU
+ */
+template <typename T>
+inline void TensorFromVector(const phi::CustomContext& ctx,
+                             const std::vector<T>& src,
+                             const phi::CustomContext& dev_ctx,
+                             phi::DenseTensor* dst) {
+  auto dst_place = dev_ctx.GetPlace();
+  C_Device_st device{dst_place.GetDeviceId()};
+  auto src_ptr = static_cast<const void*>(src.data());
+  dst->Resize({static_cast<int64_t>(src.size())});
+  auto dst_ptr = static_cast<void*>(dev_ctx.template Alloc<T>(dst));
+  auto size = src.size() * sizeof(T);
+  if (UNLIKELY(size == 0)) return;
+
+  if (dst_place.GetType() == phi::AllocationType::CUSTOM) {
+    AsyncMemCpyH2D(&device,
+                   static_cast<C_Stream>(dev_ctx.stream()),
+                   dst_ptr,
+                   src_ptr,
+                   size);
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "TensorFromVector on %s is not supported.", dst_place));
+  }
+}
+
 inline int CanonicalAxis(const int axis, const int rank) {
   if (axis < 0) {
     return axis + rank;

backends/intel_hpu/kernels/set_value_kernel.cc

@@ -233,53 +233,49 @@ void SetTensorValueKernel(const Context& dev_ctx,
   runner.Run(reinterpret_cast<C_Stream>(dev_ctx.stream()), tensors);
 }
 
-// template <typename T, typename Context>
-// void SetValueKernel(const Context& dev_ctx,
-//                        const phi::DenseTensor& x,
-//                        const phi::IntArray& starts,
-//                        const phi::IntArray& ends,
-//                        const phi::IntArray& steps,
-//                        const std::vector<int64_t>& axes,
-//                        const std::vector<int64_t>& decrease_axes,
-//                        const std::vector<int64_t>& none_axes,
-//                        const std::vector<int64_t>& shape,
-//                        const std::vector<phi::Scalar>& values,
-//                        phi::DenseTensor* out) {
-//   std::vector<T> assgin_values;
-//   assgin_values.reserve(values.size());
-//   for (const auto& val : values) {
-//     assgin_values.push_back(val.to<T>());
-//   }
-//   phi::DenseTensor value_tensor;
-//   value_tensor.Resize(phi::make_ddim(shape));
-//   custom_kernel::TensorFromVector(
-//       dev_ctx, assgin_values, dev_ctx, &value_tensor);
-//   value_tensor.Resize(phi::make_ddim(shape));
-
-//   custom_kernel::SetTensorValueKernel<T, Context>(dev_ctx,
-//                                                      x,
-//                                                      value_tensor,
-//                                                      starts,
-//                                                      ends,
-//                                                      steps,
-//                                                      axes,
-//                                                      decrease_axes,
-//                                                      none_axes,
-//                                                      out);
-// }
+template <typename T, typename Context>
+void SetValueKernel(const Context& dev_ctx,
+                    const phi::DenseTensor& x,
+                    const phi::IntArray& starts,
+                    const phi::IntArray& ends,
+                    const phi::IntArray& steps,
+                    const std::vector<int64_t>& axes,
+                    const std::vector<int64_t>& decrease_axes,
+                    const std::vector<int64_t>& none_axes,
+                    const std::vector<int64_t>& shape,
+                    const std::vector<phi::Scalar>& values,
+                    phi::DenseTensor* out) {
+  std::vector<T> assign_values;
+  assign_values.reserve(values.size());
+  for (const auto& val : values) {
+    assign_values.push_back(val.to<T>());
+  }
+  phi::DenseTensor value_tensor;
+  value_tensor.Resize(phi::make_ddim(shape));
+  TensorFromVector(dev_ctx, assign_values, dev_ctx, &value_tensor);
+  value_tensor.Resize(phi::make_ddim(shape));
 
-//
+  custom_kernel::SetTensorValueKernel<T, Context>(dev_ctx,
+                                                  x,
+                                                  value_tensor,
+                                                  starts,
+                                                  ends,
+                                                  steps,
+                                                  axes,
+                                                  decrease_axes,
+                                                  none_axes,
+                                                  out);
+}
 
 }  // namespace custom_kernel
 
-// PD_REGISTER_PLUGIN_KERNEL(set_value,
-//                           intel_hpu,
-//                           ALL_LAYOUT,
-//                           custom_kernel::SetValueKernel,
-//                           float,
-//                           phi::dtype::float16,
-//                           phi::dtype::bfloat16) {
-// }
+PD_REGISTER_PLUGIN_KERNEL(set_value,
+                          intel_hpu,
+                          ALL_LAYOUT,
+                          custom_kernel::SetValueKernel,
+                          float,
+                          phi::dtype::float16,
+                          phi::dtype::bfloat16) {}
 
 PD_REGISTER_PLUGIN_KERNEL(set_value_with_tensor,
                           intel_hpu,