chore: import upstream snapshot with attribution

paddle/phi/api/ext/dispatch.h

@@ -0,0 +1,70 @@
+/* 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. */
+
+#pragma once
+
+#include "paddle/phi/core/visit_type.h"
+
+namespace paddle {
+
+// Note: Keep this file only for compatibility with custom operators
+
+///////// Floating Dispatch Marco ///////////
+
+#define PD_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
+  PD_VISIT_FLOATING_TYPES(TYPE, NAME, __VA_ARGS__)
+
+#define PD_DISPATCH_FLOATING_AND_HALF_TYPES(TYPE, NAME, ...) \
+  PD_VISIT_FLOATING_AND_HALF_TYPES(TYPE, NAME, __VA_ARGS__)
+
+///////// Integral Dispatch Marco ///////////
+
+#define PD_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
+  PD_VISIT_INTEGRAL_TYPES(TYPE, NAME, __VA_ARGS__)
+
+///////// Complex Dispatch Marco ///////////
+
+#define PD_DISPATCH_COMPLEX_TYPES(TYPE, NAME, ...) \
+  PD_VISIT_COMPLEX_TYPES(TYPE, NAME, __VA_ARGS__)
+
+///////// Floating and Integral Dispatch Marco ///////////
+
+#define PD_DISPATCH_FLOATING_AND_INTEGRAL_TYPES(TYPE, NAME, ...) \
+  PD_VISIT_FLOATING_AND_INTEGRAL_TYPES(TYPE, NAME, __VA_ARGS__)
+
+///////// Floating and Complex Dispatch Marco ///////////
+
+#define PD_DISPATCH_FLOATING_AND_COMPLEX_TYPES(TYPE, NAME, ...) \
+  PD_VISIT_FLOATING_AND_COMPLEX_TYPES(TYPE, NAME, __VA_ARGS__)
+
+///////// Floating and Complex and other type Dispatch Marco ///////////
+
+#define PD_DISPATCH_FLOATING_AND_COMPLEX_AND_1_TYPE( \
+    SPECIFIED_TYPE, TYPE, NAME, ...)                 \
+  PD_VISIT_FLOATING_AND_COMPLEX_AND_1_TYPE(          \
+      SPECIFIED_TYPE, TYPE, NAME, __VA_ARGS__)
+
+///////// Floating and Complex and 2 other type Dispatch Marco ///////////
+
+#define PD_DISPATCH_FLOATING_AND_COMPLEX_AND_2_TYPES(  \
+    SPECIFIED_TYPE1, SPECIFIED_TYPE2, TYPE, NAME, ...) \
+  PD_VISIT_FLOATING_AND_COMPLEX_AND_2_TYPES(           \
+      SPECIFIED_TYPE1, SPECIFIED_TYPE2, TYPE, NAME, __VA_ARGS__)
+
+///////// Floating, Integral and Complex Dispatch Marco ///////////
+
+#define PD_DISPATCH_FLOATING_AND_INTEGRAL_AND_COMPLEX_TYPES(TYPE, NAME, ...) \
+  PD_VISIT_FLOATING_AND_INTEGRAL_AND_COMPLEX_TYPES(TYPE, NAME, __VA_ARGS__)
+
+}  // namespace paddle

paddle/phi/api/ext/native_meta_tensor.h

@@ -0,0 +1,36 @@
+// Copyright (c) 2025 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.
+
+#pragma once
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/core/ddim.h"
+
+namespace phi {
+class PADDLE_API NativeMetaTensor {
+ public:
+  NativeMetaTensor() = default;
+  NativeMetaTensor(phi::DataType dtype, phi::DDim dims)
+      : dims_(dims), dtype_(dtype) {}
+  DDim dims() const;
+  DataType dtype() const;
+  void set_dims(const DDim& dims);
+  void set_dtype(DataType dtype);
+
+ private:
+  phi::DDim dims_;
+  phi::DataType dtype_{phi::DataType::FLOAT32};
+};
+}  // namespace phi

paddle/phi/api/ext/spmd_infer.h

@@ -0,0 +1,165 @@
+/* Copyright (c) 2023 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. */
+
+#pragma once
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+
+namespace paddle {
+
+using CustomSpmdInferTensorArg =
+    paddle::variant<phi::distributed::DistMetaTensor,
+                    std::vector<phi::distributed::DistMetaTensor>>;
+
+using CustomSpmdInferAttrArg = paddle::any;
+template <typename T>
+struct SpmdInferHelperTypeEnd {};
+
+#define PD_INFER_SPMD_CHECK_INPUTS_SIZE_GT(inputs, size_bound)                 \
+  PADDLE_ENFORCE_GT(inputs.size(),                                             \
+                    size_bound,                                                \
+                    common::errors::PreconditionNotMet(                        \
+                        "The size of %s must be greater than %d, but got %d.", \
+                        #inputs,                                               \
+                        size_bound,                                            \
+                        inputs.size()))
+
+#define PD_INFER_SPMD_CHECK_INPUTS_SIZE_EQ(inputs, size_bound)             \
+  PADDLE_ENFORCE_EQ(inputs.size(),                                         \
+                    size_bound,                                            \
+                    common::errors::PreconditionNotMet(                    \
+                        "The size of %s must be equal to %d, but got %d.", \
+                        #inputs,                                           \
+                        size_bound,                                        \
+                        inputs.size()))
+
+#define PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(attr_type)              \
+  template <typename... Tail>                                              \
+  struct SpmdInferHelper<attr_type, Tail...> {                             \
+    template <int in_idx, int attr_idx, typename... PreviousArgs>          \
+    static phi::distributed::SpmdInfo InferSpmd(                           \
+        const std::vector<CustomSpmdInferTensorArg>& inputs,               \
+        const std::vector<CustomSpmdInferAttrArg>& attrs,                  \
+        const PreviousArgs&... pargs) {                                    \
+      try {                                                                \
+        PD_INFER_SPMD_CHECK_INPUTS_SIZE_GT(attrs, attr_idx);               \
+        attr_type arg = paddle::any_cast<attr_type>(attrs[attr_idx]);      \
+        return SpmdInferHelper<Tail...>::template InferSpmd<in_idx,        \
+                                                            attr_idx + 1>( \
+            inputs, attrs, pargs..., arg);                                 \
+      } catch (paddle::bad_any_cast&) {                                    \
+        PD_THROW(                                                          \
+            "Attribute cast error in custom operator SpmdInferFunc "       \
+            "function. "                                                   \
+            "Expected " #attr_type                                         \
+            " value. SpmdInferFunc's attribute list must be exactly "      \
+            "same "                                                        \
+            "as "                                                          \
+            "Forward "                                                     \
+            "KernelFn's attribute list except std::vector<int64_t> "       \
+            "attribute.");                                                 \
+      }                                                                    \
+    }                                                                      \
+  }
+
+template <typename F, F f>
+struct SpmdInferImpl;
+
+template <typename... Args, phi::distributed::SpmdInfo (*impl_fn)(Args...)>
+struct SpmdInferImpl<phi::distributed::SpmdInfo (*)(Args...), impl_fn> {
+  static phi::distributed::SpmdInfo InferSpmd(
+      const std::vector<CustomSpmdInferTensorArg>& inputs,
+      const std::vector<CustomSpmdInferAttrArg>& attrs) {
+    return SpmdInferHelper<Args..., SpmdInferHelperTypeEnd<int>>::
+        template InferSpmd<0, 0>(inputs, attrs);
+  }
+
+ private:
+  template <typename... RemainingArgs>
+  struct SpmdInferHelper;
+
+  // Handle args for general tensor input case
+  template <typename... Tail>
+  struct SpmdInferHelper<const phi::distributed::DistMetaTensor&, Tail...> {
+    template <int in_idx, int attr_idx, typename... PreviousArgs>
+    static phi::distributed::SpmdInfo InferSpmd(
+        const std::vector<CustomSpmdInferTensorArg>& inputs,
+        const std::vector<CustomSpmdInferAttrArg>& attrs,
+        PreviousArgs&... pargs) {
+      static_assert(attr_idx == 0, "attributes must come after tensor inputs");
+      PD_INFER_SPMD_CHECK_INPUTS_SIZE_GT(inputs, in_idx);
+      auto& arg =
+          PADDLE_GET_CONST(phi::distributed::DistMetaTensor, inputs[in_idx]);
+      return SpmdInferHelper<Tail...>::template InferSpmd<in_idx + 1, attr_idx>(
+          inputs, attrs, pargs..., arg);
+    }
+  };
+
+  // Handle args for vector of Tensor input case
+  template <typename... Tail>
+  struct SpmdInferHelper<const std::vector<phi::distributed::DistMetaTensor>&,
+                         Tail...> {
+    template <int in_idx, int attr_idx, typename... PreviousArgs>
+    static phi::distributed::SpmdInfo InferSpmd(
+        const std::vector<CustomSpmdInferTensorArg>& inputs,
+        const std::vector<CustomSpmdInferAttrArg>& attrs,
+        PreviousArgs&... pargs) {
+      static_assert(attr_idx == 0, "attributes must come after tensor inputs");
+      PD_INFER_SPMD_CHECK_INPUTS_SIZE_GT(inputs, in_idx);
+      auto& arg = PADDLE_GET_CONST(
+          std::vector<phi::distributed::DistMetaTensor>, inputs[in_idx]);
+      return SpmdInferHelper<Tail...>::template InferSpmd<in_idx + 1, attr_idx>(
+          inputs, attrs, pargs..., arg);
+    }
+  };
+
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(bool);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(int);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(float);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(int64_t);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(const std::string&);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(const std::vector<int>&);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(const std::vector<float>&);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(const std::vector<std::string>&);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(const std::vector<int64_t>&);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(const bool&);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(const int&);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(const float&);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(const int64_t&);
+
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(std::string);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(std::vector<int>);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(std::vector<float>);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(std::vector<std::string>);
+  PD_SPECIALIZE_SpmdInferHelper_FOR_AttrType(const std::vector<int64_t>);
+
+  // end: base template
+  template <typename T>
+  struct SpmdInferHelper<SpmdInferHelperTypeEnd<T>> {
+    template <int in_idx, int attr_idx, typename... PreviousArgs>
+    static phi::distributed::SpmdInfo InferSpmd(
+        const std::vector<CustomSpmdInferTensorArg>& inputs,
+        const std::vector<CustomSpmdInferAttrArg>& attrs,
+        PreviousArgs&... pargs) {
+      PD_INFER_SPMD_CHECK_INPUTS_SIZE_EQ(inputs, in_idx);
+      PD_INFER_SPMD_CHECK_INPUTS_SIZE_EQ(attrs, attr_idx);
+      return impl_fn(pargs...);
+    }
+  };
+};
+
+#define PD_INFER_SPMD_RULE(...) \
+  ::paddle::SpmdInferImpl<decltype(&__VA_ARGS__), &__VA_ARGS__>::InferSpmd
+
+}  // namespace paddle

paddle/phi/api/ext/tensor_compat.h

@@ -0,0 +1,163 @@
+/* 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. */
+
+#pragma once
+
+#include "paddle/phi/api/include/api.h"
+#include "paddle/phi/api/include/tensor.h"
+
+// Note(chenweihang): In order to be compatible with the original custom
+// operator Tensor interface, only available to external users, the file
+// cannot be included in paddle
+
+namespace paddle {
+// using several Tensor initialize functions in paddle namespace
+using experimental::abs;
+using experimental::acos;
+using experimental::acosh;
+using experimental::add;
+using experimental::addmm;
+using experimental::allclose;
+using experimental::argsort;
+using experimental::asin;
+using experimental::asinh;
+using experimental::atan;
+using experimental::atan2;
+using experimental::atanh;
+using experimental::bernoulli;
+using experimental::ceil;
+using experimental::cholesky;
+using experimental::cholesky_solve;
+using experimental::clip;
+using experimental::concat;
+using experimental::conj;
+using experimental::cos;
+using experimental::cosh;
+using experimental::cross;
+using experimental::det;
+using experimental::diag;
+using experimental::diagonal;
+using experimental::digamma;
+using experimental::dist;
+using experimental::divide;
+using experimental::dot;
+using experimental::elu;
+using experimental::empty;
+using experimental::empty_like;
+using experimental::equal;
+using experimental::equal_all;
+using experimental::erf;
+using experimental::erfinv;
+using experimental::exp;
+using experimental::expand;
+using experimental::expm1;
+using experimental::flatten;
+using experimental::flip;
+using experimental::floor;
+using experimental::floor_divide;
+using experimental::fmax;
+using experimental::fmin;
+using experimental::frame;
+using experimental::full;
+using experimental::gather;
+using experimental::gather_nd;
+using experimental::gelu;
+using experimental::greater_equal;
+using experimental::greater_than;
+using experimental::gumbel_softmax;
+using experimental::hardswish;
+using experimental::hardtanh;
+using experimental::imag;
+using experimental::increment;
+using experimental::index_sample;
+using experimental::is_empty;
+using experimental::isclose;
+using experimental::isfinite;
+using experimental::isinf;
+using experimental::isnan;
+using experimental::kron;
+using experimental::kthvalue;
+using experimental::label_smooth;
+using experimental::lerp;
+using experimental::less_equal;
+using experimental::less_than;
+using experimental::lgamma;
+using experimental::log;
+using experimental::log10;
+using experimental::log1p;
+using experimental::log2;
+using experimental::logit;
+using experimental::masked_select;
+using experimental::matmul;
+using experimental::matrix_power;
+using experimental::maximum;
+using experimental::maxout;
+using experimental::meshgrid;
+using experimental::minimum;
+using experimental::mode;
+using experimental::multi_dot;
+using experimental::multinomial;
+using experimental::multiply;
+using experimental::mv;
+using experimental::nll_loss;
+using experimental::not_equal;
+using experimental::npu_identity;
+using experimental::one_hot;
+using experimental::ones;
+using experimental::overlap_add;
+using experimental::pixel_shuffle;
+using experimental::poisson;
+using experimental::put_along_axis;
+using experimental::qr;
+using experimental::real;
+using experimental::reciprocal;
+using experimental::relu;
+using experimental::relu6;
+using experimental::remainder;
+using experimental::reshape;
+using experimental::roll;
+using experimental::round;
+using experimental::rsqrt;
+using experimental::scatter;
+using experimental::scatter_nd_add;
+using experimental::selu;
+using experimental::send_u_recv;
+using experimental::send_ue_recv;
+using experimental::send_uv;
+using experimental::sigmoid;
+using experimental::sign;
+using experimental::silu;
+using experimental::sin;
+using experimental::sinh;
+using experimental::split;
+using experimental::sqrt;
+using experimental::square;
+using experimental::stack;
+using experimental::standard_gamma;
+using experimental::strided_slice;
+using experimental::subtract;
+using experimental::swish;
+using experimental::tanh;
+using experimental::thresholded_relu;
+using experimental::tile;
+using experimental::trace;
+using experimental::triangular_solve;
+using experimental::tril;
+using experimental::unbind;
+using experimental::unique;
+using experimental::unsqueeze;
+using experimental::where;
+using experimental::zeros;
+
+}  // namespace paddle