fix set_constant error (#59905)

danleifeng · web-flow · commit a38ab8b9796e · 2023-12-15T11:49:23.000+08:00
diff --git a/paddle/fluid/framework/hogwild_worker.cc b/paddle/fluid/framework/hogwild_worker.cc
@@ -1131,7 +1131,7 @@ template <typename T>
 void HogwildWorker::SetZero(phi::DenseTensor *tensor,
                             const phi::DenseTensor &root_tensor) {
   tensor->mutable_data<T>(root_tensor.dims(), place_);
-  phi::funcs::set_constant(*dev_ctx_, tensor, static_cast<T>(0.0));
+  phi::funcs::set_constant(*dev_ctx_, tensor, 0.0);
 }
 
 void HogwildWorker::BindingDataFeedMemory() {
diff --git a/paddle/phi/kernels/funcs/math_function.cc b/paddle/phi/kernels/funcs/math_function.cc
@@ -143,71 +143,26 @@ DEFINE_CPU_TRANS_NORMAL(phi::dtype::complex<float>);
 DEFINE_CPU_TRANS_NORMAL(phi::dtype::complex<double>);
 
 struct TensorSetConstantCPU {
-  TensorSetConstantCPU(phi::DenseTensor* tensor, const void* value)
+  TensorSetConstantCPU(phi::DenseTensor* tensor, float value)
       : tensor_(tensor), value_(value) {}
   template <typename T>
   void apply() const {
     auto cpu = phi::CPUPlace();
     auto* begin = tensor_->mutable_data<T>(cpu);
-    const T* num_ptr = reinterpret_cast<const T*>(value_);
-    T num = *num_ptr;
-    std::fill(begin, begin + tensor_->numel(), num);
+    std::fill(begin, begin + tensor_->numel(), static_cast<T>(value_));
   }
   phi::DenseTensor* tensor_;
-  const void* value_;
+  float value_;
 };
 
-#ifdef PADDLE_WITH_XPU
-struct TensorSetConstantXPU {
-  TensorSetConstantXPU(const phi::DeviceContext& context,
-                       phi::DenseTensor* tensor,
-                       const void* value,
-                       phi::Place place)
-      : context_(context), tensor_(tensor), value_(value), place_(place) {}
-  template <typename T>
-  void apply() const {
-    auto* ctx = phi::DeviceContextPool::Instance().Get(place_);
-    auto data = ctx->Alloc<T>(tensor_);
-    const T* num = reinterpret_cast<const T*>(value_);
-    T num_value = static_cast<T>(*num);
-    int numel = tensor_->numel();
-    if (((std::is_same<T, float>::value) ||
-         (std::is_same<T, phi::dtype::float16>::value)) &&
-        (place_ == phi::XPUPlace())) {
-      using XPUType = typename XPUTypeTrait<T>::Type;
-      auto* dev_ctx = static_cast<phi::XPUContext*>(ctx);
-      int r = xpu::constant(dev_ctx->x_context(),
-                            reinterpret_cast<XPUType*>(data),
-                            numel,
-                            static_cast<XPUType>(num_value));
-      PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
-      dev_ctx->Wait();
-    } else {
-      std::unique_ptr<T[]> data_cpu(new T[numel]);
-      std::fill(
-          data_cpu.get(), data_cpu.get() + numel, static_cast<T>(num_value));
-      memory_utils::Copy(place_,
-                         data,
-                         phi::CPUPlace(),
-                         static_cast<void*>(data_cpu.get()),
-                         numel * sizeof(T));
-    }
-  }
-  const phi::DeviceContext& context_;
-  phi::DenseTensor* tensor_;
-  const void* value_;
-  phi::Place place_;
-};
-#endif
-
 template <>
 void set_constant_with_place<phi::XPUPlace>(const phi::DeviceContext& context,
                                             phi::DenseTensor* tensor,
-                                            const void* value) {
+                                            float value) {
 #ifdef PADDLE_WITH_XPU
   phi::VisitDataType(
       tensor->dtype(),
-      TensorSetConstantXPU(context, tensor, value, tensor->place()));
+      TensorSetConstantXPU<float>(tensor, value, tensor->place()));
 #else
   PADDLE_THROW(phi::errors::PreconditionNotMet("Not compiled with XPU!"));
 #endif
@@ -216,15 +171,13 @@ void set_constant_with_place<phi::XPUPlace>(const phi::DeviceContext& context,
 template <>
 void set_constant_with_place<phi::IPUPlace>(const phi::DeviceContext& context,
                                             phi::DenseTensor* tensor,
-                                            const void* value) {
+                                            float value) {
   PADDLE_THROW(phi::errors::Unimplemented("IPUPlace is not supported"));
 }
 
 template <>
 void set_constant_with_place<phi::CustomPlace>(
-    const phi::DeviceContext& context,
-    phi::DenseTensor* tensor,
-    const void* value) {
+    const phi::DeviceContext& context, phi::DenseTensor* tensor, float value) {
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
   auto kernel_result = phi::KernelFactory::Instance().SelectKernelOrThrowError(
       "full",
@@ -237,12 +190,10 @@ void set_constant_with_place<phi::CustomPlace>(
                                     const phi::Scalar&,
                                     DataType,
                                     phi::DenseTensor*);
-  const float* num_ptr = reinterpret_cast<const float*>(value);
-  float num = *num_ptr;
   auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
   (*kernel_fn)(context,
                phi::IntArray(common::vectorize(tensor->dims())),
-               phi::Scalar(num),
+               phi::Scalar(value),
                tensor->dtype(),
                tensor);
 #else
@@ -253,15 +204,13 @@ void set_constant_with_place<phi::CustomPlace>(
 template <>
 void set_constant_with_place<phi::CPUPlace>(const phi::DeviceContext& context,
                                             phi::DenseTensor* tensor,
-                                            const void* value) {
+                                            float value) {
   phi::VisitDataType(tensor->dtype(), TensorSetConstantCPU(tensor, value));
 }
 
 template <>
 void set_constant_with_place<phi::GPUPinnedPlace>(
-    const phi::DeviceContext& context,
-    phi::DenseTensor* tensor,
-    const void* value) {
+    const phi::DeviceContext& context, phi::DenseTensor* tensor, float value) {
   phi::VisitDataType(tensor->dtype(), TensorSetConstantCPU(tensor, value));
 }
 
@@ -270,7 +219,7 @@ struct TensorSetConstantWithPlace {
   using result_type = void;
   TensorSetConstantWithPlace(const phi::DeviceContext& context,
                              phi::DenseTensor* tensor,
-                             const void* value)
+                             float value)
       : context_(context), tensor_(tensor), value_(value) {}
 
   template <typename Place>
@@ -280,12 +229,12 @@ struct TensorSetConstantWithPlace {
 
   const phi::DeviceContext& context_;
   phi::DenseTensor* tensor_;
-  const void* value_;
+  float value_;
 };
 
 void set_constant(const phi::DeviceContext& context,
                   phi::DenseTensor* tensor,
-                  const void* value) {
+                  float value) {
   TensorSetConstantWithPlace func(context, tensor, value);
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
   if (context.GetPlace().GetType() == phi::AllocationType::CUSTOM) {
diff --git a/paddle/phi/kernels/funcs/math_function.cu b/paddle/phi/kernels/funcs/math_function.cu
@@ -336,34 +336,26 @@ DEFINE_GPU_TRANS_NORMAL(phi::dtype::complex<double>);
 struct TensorSetConstantGPU {
   TensorSetConstantGPU(const phi::DeviceContext& context,
                        phi::DenseTensor* tensor,
-                       const void* value)
+                       float value)
       : context_(context), tensor_(tensor), value_(value) {}
 
   template <typename T>
   void apply() const {
-    // SetConstant<phi::GPUContext, T> functor;
-    // functor(reinterpret_cast<const phi::GPUContext&>(context_),
-    //         tensor_,
-    //         static_cast<T>(value_));
-    int N = static_cast<int>(tensor_->numel());
-    if (N <= 0) {
-      return;
-    }
-    auto& ctx = reinterpret_cast<const phi::GPUContext&>(context_);
-    const T* num = reinterpret_cast<const T*>(value_);
-    FillConstantKernel<T><<<(N + 512 - 1) / 512, 512, 0, ctx.stream()>>>(
-        N, tensor_->mutable_data<T>(ctx.GetPlace()), static_cast<T>(*num));
+    SetConstant<phi::GPUContext, T> functor;
+    functor(reinterpret_cast<const phi::GPUContext&>(context_),
+            tensor_,
+            static_cast<T>(value_));
   }
 
   const phi::DeviceContext& context_;
   phi::DenseTensor* tensor_;
-  const void* value_;
+  float value_;
 };
 
 template <>
 void set_constant_with_place<phi::GPUPlace>(const phi::DeviceContext& context,
                                             phi::DenseTensor* tensor,
-                                            const void* value) {
+                                            float value) {
   phi::VisitDataType(tensor->dtype(),
                      TensorSetConstantGPU(context, tensor, value));
 }
diff --git a/paddle/phi/kernels/funcs/math_function.h b/paddle/phi/kernels/funcs/math_function.h
@@ -64,21 +64,23 @@ struct SetConstant {
                   T num);
 };
 
+#ifdef PADDLE_WITH_XPU
+template <typename T>
+struct SetConstant<phi::XPUContext, T> {
+  void operator()(const phi::XPUContext& context,
+                  phi::DenseTensor* tensor,
+                  T num);
+};
+#endif
+
 template <typename Place>
 void set_constant_with_place(const phi::DeviceContext& context,
                              phi::DenseTensor* tensor,
-                             const void* value);
-
-void set_constant(const phi::DeviceContext& context,
-                  phi::DenseTensor* tensor,
-                  const void* value);
+                             float value);
 
-template <typename T>
 void set_constant(const phi::DeviceContext& context,
                   phi::DenseTensor* tensor,
-                  const T value) {
-  set_constant(context, tensor, reinterpret_cast<const void*>(&value));
-}
+                  float value);
 
 template <typename DeviceContext, typename T>
 struct RowwiseAdd {
@@ -109,6 +111,65 @@ struct RowwiseMean {
                   phi::DenseTensor* vec);
 };
 
+#ifdef PADDLE_WITH_XPU
+template <typename U>
+struct TensorSetConstantXPU {
+  TensorSetConstantXPU(phi::DenseTensor* tensor, U value, phi::Place place)
+      : tensor_(tensor), value_(value), place_(place) {}
+  template <typename T>
+  void apply() const {
+    auto* ctx = phi::DeviceContextPool::Instance().Get(place_);
+    auto begin = ctx->Alloc<T>(tensor_);
+    int numel = tensor_->numel();
+    std::unique_ptr<T[]> data_cpu(new T[numel]);
+    std::fill(data_cpu.get(), data_cpu.get() + numel, static_cast<T>(value_));
+    memory_utils::Copy(place_,
+                       begin,
+                       phi::CPUPlace(),
+                       static_cast<void*>(data_cpu.get()),
+                       numel * sizeof(T));
+  }
+  phi::DenseTensor* tensor_;
+  U value_;
+  phi::Place place_;
+};
+
+template <>
+struct TensorSetConstantXPU<float> {
+  TensorSetConstantXPU(phi::DenseTensor* tensor, float value, phi::Place place)
+      : tensor_(tensor), value_(value), place_(place) {}
+  template <typename T>
+  void apply() const {
+    auto* ctx = phi::DeviceContextPool::Instance().Get(place_);
+    auto begin = ctx->Alloc<T>(tensor_);
+    int numel = tensor_->numel();
+    if (((std::is_same<T, float>::value) ||
+         (std::is_same<T, phi::dtype::float16>::value)) &&
+        (place_ == phi::XPUPlace())) {
+      using XPUType = typename XPUTypeTrait<T>::Type;
+      auto* dev_ctx = static_cast<phi::XPUContext*>(ctx);
+      int r = xpu::constant<XPUType>(dev_ctx->x_context(),
+                                     reinterpret_cast<XPUType*>(begin),
+                                     numel,
+                                     static_cast<XPUType>(value_));
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
+      dev_ctx->Wait();
+    } else {
+      std::unique_ptr<T[]> data_cpu(new T[numel]);
+      std::fill(data_cpu.get(), data_cpu.get() + numel, static_cast<T>(value_));
+      memory_utils::Copy(place_,
+                         begin,
+                         phi::CPUPlace(),
+                         static_cast<void*>(data_cpu.get()),
+                         numel * sizeof(T));
+    }
+  }
+  phi::DenseTensor* tensor_;
+  float value_;
+  phi::Place place_;
+};
+#endif
+
 template <typename Context, typename T>
 inline void TransCompute(const int dim,
                          const Context& dev_ctx,
diff --git a/paddle/phi/kernels/funcs/math_function_impl.h b/paddle/phi/kernels/funcs/math_function_impl.h
@@ -29,11 +29,20 @@ template <typename DeviceContext, typename T>
 void SetConstant<DeviceContext, T>::operator()(const DeviceContext& context,
                                                phi::DenseTensor* tensor,
                                                T num) {
-  // auto t = phi::EigenVector<T>::Flatten(*tensor);
-  // t.device(*context.eigen_device()) = t.constant(static_cast<T>(num));
-  set_constant(context, tensor, reinterpret_cast<const void*>(&num));
+  auto t = phi::EigenVector<T>::Flatten(*tensor);
+  t.device(*context.eigen_device()) = t.constant(static_cast<T>(num));
 }
 
+#ifdef PADDLE_WITH_XPU
+template <typename T>
+void SetConstant<phi::XPUContext, T>::operator()(const phi::XPUContext& context,
+                                                 phi::DenseTensor* tensor,
+                                                 T num) {
+  phi::VisitDataType(tensor->dtype(),
+                     TensorSetConstantXPU<T>(tensor, num, context.GetPlace()));
+}
+#endif
+
 template <typename DeviceContext, typename T, int Rank>
 void Transpose<DeviceContext, T, Rank>::operator()(
     const DeviceContext& context,
diff --git a/paddle/phi/kernels/legacy/cpu/one_hot_kernel.cc b/paddle/phi/kernels/legacy/cpu/one_hot_kernel.cc
@@ -38,7 +38,7 @@ struct OneHotV2OpFunctor {
     auto* p_in_data = in_->data<InT>();
     auto numel = in_->numel();
     auto* p_out_data = ctx_.template Alloc<OutT>(out_);
-    funcs::set_constant(ctx_, out_, static_cast<OutT>(0.0));
+    funcs::set_constant(ctx_, out_, 0.0);
 
     for (int i = 0; i < numel; ++i) {
       PADDLE_ENFORCE_GE(
diff --git a/paddle/phi/kernels/legacy/gpu/one_hot_kernel.cu b/paddle/phi/kernels/legacy/gpu/one_hot_kernel.cu
@@ -59,7 +59,7 @@ struct OneHotV2OpCUDAFunctor {
     auto numel = in_->numel();
     auto* p_out_data = ctx_.template Alloc<OutT>(out_);
     auto stream = ctx_.stream();
-    funcs::set_constant(ctx_, out_, static_cast<OutT>(0.0));
+    funcs::set_constant(ctx_, out_, 0.0);
 
     auto config = phi::backends::gpu::GetGpuLaunchConfig1D(ctx_, numel);
 

Original file line number	Diff line number	Diff line change
`@@ -1131,7 +1131,7 @@ template <typename T>`
`1131`	`1131`	`void HogwildWorker::SetZero(phi::DenseTensor *tensor,`
`1132`	`1132`	`const phi::DenseTensor &root_tensor) {`
`1133`	`1133`	`tensor->mutable_data<T>(root_tensor.dims(), place_);`
`1134`		`- phi::funcs::set_constant(*dev_ctx_, tensor, static_cast<T>(0.0));`
	`1134`	`+ phi::funcs::set_constant(*dev_ctx_, tensor, 0.0);`
`1135`	`1135`	`}`
`1136`	`1136`
`1137`	`1137`	`void HogwildWorker::BindingDataFeedMemory() {`