chore: import upstream snapshot with attribution

paddle/phi/kernels/xpu/batch_norm_kernel.cc

@@ -0,0 +1,175 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/batch_norm_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/norm_utils.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void BatchNormKernel(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const DenseTensor& mean,
+                     const DenseTensor& variance,
+                     const optional<DenseTensor>& scale,
+                     const optional<DenseTensor>& bias,
+                     bool is_test,
+                     float momentum,
+                     float epsilon,
+                     const std::string& data_layout_str,
+                     bool use_global_stats,
+                     bool trainable_statistics,
+                     DenseTensor* y,
+                     DenseTensor* mean_out,
+                     DenseTensor* variance_out,
+                     DenseTensor* saved_mean,
+                     DenseTensor* saved_variance,
+                     DenseTensor* reserve_space) {
+  if (x.numel() == 0) {
+    dev_ctx.template Alloc<T>(y);
+    if (mean_out) dev_ctx.template Alloc<T>(mean_out);
+    if (variance_out) dev_ctx.template Alloc<T>(variance_out);
+    if (saved_mean) dev_ctx.template Alloc<T>(saved_mean);
+    if (saved_variance) dev_ctx.template Alloc<T>(saved_variance);
+    if (reserve_space) {
+      reserve_space->Resize({0});
+      dev_ctx.template Alloc<T>(reserve_space);
+    }
+    return;
+  }
+  using XPUType = typename XPUTypeTrait<T>::Type;
+  bool test_mode = is_test && (!trainable_statistics);
+  bool global_stats = test_mode || use_global_stats;
+  const auto data_layout = StringToDataLayout(data_layout_str);
+  PADDLE_ENFORCE_EQ(data_layout_str == "NCHW" || data_layout_str == "NHWC",
+                    true,
+                    common::errors::InvalidArgument(
+                        "The 'data_layout' attribute must be NCHW or NHWC. "
+                        "But received 'data_layout' is [%s].",
+                        data_layout_str));
+
+  const auto& x_dims = x.dims();
+  PADDLE_ENFORCE_EQ(
+      x_dims.size() >= 2 && x_dims.size() <= 5,
+      true,
+      common::errors::InvalidArgument(
+          "The size of input's dimensions should be between 2 and 5"
+          "But received: the size of input's dimensions is [%d]",
+          x_dims.size()));
+
+  int64_t N = -1, C = -1, H = -1, W = -1, D = -1;
+  funcs::ExtractNCWHD(x_dims, data_layout, &N, &C, &H, &W, &D);
+  N = (N == 0) ? 1 : N;
+  C = (C == 0) ? 1 : C;
+  H = (H == 0) ? 1 : H;
+  W = (W == 0) ? 1 : W;
+  D = (D == 0) ? 1 : D;
+
+  W = W * D;
+
+  auto* Scale = scale.get_ptr();
+  auto* Bias = bias.get_ptr();
+
+  DenseTensor new_scale;
+  DenseTensor new_bias;
+
+  if (Scale) {
+    new_scale = scale.get();
+  } else {
+    new_scale = phi::Full<T, Context>(dev_ctx, {C}, static_cast<T>(1));
+  }
+
+  if (Bias) {
+    new_bias = bias.get();
+  } else {
+    new_bias = phi::Full<T, Context>(dev_ctx, {C}, static_cast<T>(0));
+  }
+
+  const auto* x_data = reinterpret_cast<const XPUType*>(x.data<T>());
+  const auto* scale_data = new_scale.data<float>();
+  const auto* bias_data = new_bias.data<float>();
+
+  // alloc memory
+  auto* y_data = reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(y));
+  dev_ctx.template Alloc<float>(mean_out);
+  dev_ctx.template Alloc<float>(variance_out);
+  dev_ctx.template Alloc<float>(saved_mean);
+  dev_ctx.template Alloc<float>(saved_variance);
+  PADDLE_ENFORCE_LE(
+      x_dims.size(),
+      5,
+      common::errors::InvalidArgument(
+          "The size of input X's dimensions should be less than 6."
+          "But received: the size of input X's dimensions is [%d]",
+          x_dims.size()));
+
+  bool is_nchw = data_layout_str == "NCHW";
+
+  if (!global_stats) {
+    auto* mean_out_data = mean_out->data<float>();
+    auto* variance_out_data = variance_out->data<float>();
+    auto* saved_mean_data = saved_mean->data<float>();
+    auto* saved_variance_data = saved_variance->data<float>();
+
+    int r = xpu::batch_norm<XPUType>(dev_ctx.x_context(),
+                                     x_data,
+                                     y_data,
+                                     N,
+                                     C,
+                                     H,
+                                     W,
+                                     epsilon,
+                                     momentum,
+                                     scale_data,
+                                     bias_data,
+                                     saved_mean_data,
+                                     saved_variance_data,
+                                     mean_out_data,
+                                     variance_out_data,
+                                     is_nchw);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "batch_norm");
+  } else {
+    const auto* mean_data = mean.data<float>();
+    const auto* variance_data = variance.data<float>();
+    int r = xpu::batch_norm_infer<XPUType>(dev_ctx.x_context(),
+                                           x_data,
+                                           y_data,
+                                           N,
+                                           C,
+                                           H,
+                                           W,
+                                           epsilon,
+                                           scale_data,
+                                           bias_data,
+                                           mean_data,
+                                           variance_data,
+                                           is_nchw);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "batch_norm_infer");
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    batch_norm, XPU, ALL_LAYOUT, phi::BatchNormKernel, float, phi::float16) {
+  kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+  kernel->OutputAt(5).SetDataType(phi::DataType::UINT8);
+}