// Copyright (c) 2021 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/math_kernel.h" #include "paddle/phi/api/ext/dispatch.h" #include "paddle/phi/backends/cpu/cpu_context.h" #include "paddle/phi/common/scalar.h" #include "paddle/phi/core/kernel_registry.h" #include "paddle/phi/kernels/cpu/elementwise.h" #include "paddle/phi/kernels/cpu/reduce.h" #include "paddle/phi/kernels/funcs/elementwise_functor.h" #include "paddle/phi/kernels/funcs/reduce_functor.h" // See Note [ Why still include the fluid headers? ] #include "paddle/fluid/framework/eigen.h" #include "paddle/phi/common/bfloat16.h" #include "paddle/phi/common/complex.h" namespace phi { #define DEFINE_CPU_ELEMENTWISE_OP(name) \ template \ void name##RawKernel(const Context& dev_ctx, \ const DenseTensor& x, \ const DenseTensor& y, \ int axis, \ DenseTensor* out) { \ dev_ctx.template Alloc(out); \ if (x.dims() == y.dims()) { \ SameDimsElementwiseCompute>()( \ dev_ctx, x, y, out); \ } else { \ auto x_dims = x.dims(); \ auto y_dims = y.dims(); \ if (x_dims.size() >= y_dims.size()) { \ ElementwiseCompute, T>( \ dev_ctx, x, y, axis, funcs::name##Functor(), out); \ } else { \ ElementwiseCompute, T>( \ dev_ctx, x, y, axis, funcs::Inverse##name##Functor(), out); \ } \ } \ } template void MeanRawKernel(const Context& dev_ctx, const DenseTensor& x, const std::vector& dims, bool keep_dim, bool reduce_all, DenseTensor* out) { auto out_dtype = x.dtype(); phi::Reduce( dev_ctx, x, reduce_all, dims, keep_dim, out_dtype, out); } template void SumRawKernel(const Context& dev_ctx, const DenseTensor& x, const std::vector& dims, bool keep_dim, bool reduce_all, DataType out_dtype, DenseTensor* out) { phi::Reduce( dev_ctx, x, reduce_all, dims, keep_dim, out_dtype, out); } template void DivideRawKernel(const Context& dev_ctx, const DenseTensor& x, const DenseTensor& y, int axis, DenseTensor* out) { // allocate memory for out dev_ctx.template Alloc(out); if (x.dims() == y.dims() && std::is_floating_point::value) { SameDimsElementwiseCompute>()( dev_ctx, x, y, out); } else { auto x_dims = x.dims(); auto y_dims = y.dims(); if (x_dims.size() >= y_dims.size()) { ElementwiseCompute, T>( dev_ctx, x, y, axis, funcs::DivideFunctor(), out); } else { ElementwiseCompute, T>( dev_ctx, x, y, axis, funcs::InverseDivideFunctor(), out); } } } // Create the definition of Add DEFINE_CPU_ELEMENTWISE_OP(Add) // Create the definition of Subtract DEFINE_CPU_ELEMENTWISE_OP(Subtract) // Create the definition of Multiply DEFINE_CPU_ELEMENTWISE_OP(Multiply) } // namespace phi using complex64 = ::phi::dtype::complex; using complex128 = ::phi::dtype::complex; // NOTE(chenweihang): using bfloat16 will cause redefine with xpu bfloat16 // using bfloat16 = ::phi::dtype::bfloat16; PD_REGISTER_KERNEL(add_raw, CPU, ALL_LAYOUT, phi::AddRawKernel, float, double, int16_t, int, int64_t, complex64, complex128) {} PD_REGISTER_KERNEL(subtract_raw, CPU, ALL_LAYOUT, phi::SubtractRawKernel, float, double, int16_t, int, int64_t, complex64, complex128) {} PD_REGISTER_KERNEL(divide_raw, CPU, ALL_LAYOUT, phi::DivideRawKernel, float, double, int, int64_t, complex64, complex128) {} PD_REGISTER_KERNEL(multiply_raw, CPU, ALL_LAYOUT, phi::MultiplyRawKernel, float, double, int, int64_t, bool, complex64, complex128) {} PD_REGISTER_KERNEL(sum_raw, CPU, ALL_LAYOUT, phi::SumRawKernel, bool, float, double, phi::dtype::float16, int16_t, int, int64_t, complex64, complex128) { kernel->OutputAt(0).SetDataType(paddle::experimental::DataType::UNDEFINED); } PD_REGISTER_KERNEL( mean_raw, CPU, ALL_LAYOUT, phi::MeanRawKernel, float, double, bool) {}