chore: import upstream snapshot with attribution

paddle/phi/kernels/cpu/lrn_kernel.cc

@@ -0,0 +1,122 @@
+// Copyright (c) 2024 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/impl/lrn_kernel_impl.h"
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "paddle/phi/backends/onednn/onednn_helper.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace phi {
+
+template <typename T>
+struct LRNFunctor<CPUContext, T> {
+  void operator()(const CPUContext& dev_ctx,
+                  const DenseTensor& input,
+                  DenseTensor* out,
+                  DenseTensor* mid,
+                  int64_t N,
+                  int64_t C,
+                  int64_t H,
+                  int64_t W,
+                  int n,
+                  T k,
+                  T alpha,
+                  T beta,
+                  const DataLayout data_layout) {
+    auto blas = funcs::GetBlas<CPUContext, T>(dev_ctx);
+    funcs::Transpose<CPUContext, T, 4> transpose;
+    DenseTensor in_transpose, mid_transpose, out_transpose;
+    // if channel_last, transpose to channel_first
+    if (data_layout == DataLayout::NHWC) {
+      auto in_dims = input.dims();
+      std::vector<int64_t> shape(
+          {in_dims[0], in_dims[3], in_dims[1], in_dims[2]});
+      in_transpose.Resize(shape);
+      mid_transpose.Resize(shape);
+      out_transpose.Resize(shape);
+      dev_ctx.Alloc<T>(&in_transpose);
+      dev_ctx.Alloc<T>(&mid_transpose);
+      dev_ctx.Alloc<T>(&out_transpose);
+      std::vector<int> axis = {0, 3, 1, 2};
+      transpose(dev_ctx, input, &in_transpose, axis);
+    } else {
+      in_transpose = input;
+      mid_transpose = *mid;
+      out_transpose = *out;
+      mid_transpose.Resize(mid->dims());
+      out_transpose.Resize(out->dims());
+      dev_ctx.Alloc<T>(&mid_transpose);
+      dev_ctx.Alloc<T>(&out_transpose);
+    }
+
+    const T* idata = in_transpose.data<T>();
+    T* odata = out_transpose.data<T>();
+    T* mdata = mid_transpose.data<T>();
+
+    DenseTensor squared;
+    squared.Resize({1, C + n - 1, H, W});
+    T* sdata = dev_ctx.Alloc<T>(&squared);
+    std::memset(sdata, 0, sizeof(T) * squared.numel());
+    for (int64_t i = 0; i < mid->numel(); ++i) {
+      mdata[i] = k;
+    }
+    int64_t img_size = H * W;
+    int64_t fea_size = C * img_size;
+    int pre_pad = (n - 1) / 2;
+    // compute batches one by one
+    for (int64_t i = 0; i < N; ++i) {
+      blas.VSQUARE(fea_size, idata + i * fea_size, sdata + pre_pad * img_size);
+      // init the first channel of mid
+      for (int c = 0; c < n; ++c) {
+        blas.AXPY(img_size, alpha, sdata + c * img_size, mdata + i * fea_size);
+      }
+      for (int64_t c = 1; c < C; ++c) {
+        // copy previous scale
+        int64_t mid_offset = i * fea_size + c * img_size;
+        std::memcpy(mdata + mid_offset,
+                    mdata + mid_offset - img_size,
+                    img_size * sizeof(T));
+        // add last
+        blas.AXPY(img_size,
+                  alpha,
+                  sdata + (c + n - 1) * img_size,
+                  mdata + mid_offset);
+        // sub rest
+        blas.AXPY(
+            img_size, -alpha, sdata + (c - 1) * img_size, mdata + mid_offset);
+      }
+    }
+    // compute the final output
+    blas.VPOW(mid->numel(), mdata, -beta, odata);
+    blas.VMUL(mid->numel(), odata, idata, odata);
+
+    // if channel_last, transpose the output(NCHW) to channel_last
+    if (data_layout == DataLayout::NHWC) {
+      std::vector<int> axis = {0, 2, 3, 1};
+      transpose(dev_ctx, mid_transpose, mid, axis);
+      transpose(dev_ctx, out_transpose, out, axis);
+    }
+  }
+};
+template struct LRNFunctor<CPUContext, float>;
+template struct LRNFunctor<CPUContext, double>;
+}  // namespace phi
+
+PD_REGISTER_KERNEL(lrn, CPU, ALL_LAYOUT, phi::LRNKernel, float) {}