chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,37 @@
|
||||
// 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 <ATen/Device.h>
|
||||
#include <ATen/Functions.h>
|
||||
#include <ATen/Tensor.h>
|
||||
#include <ATen/TensorIndexing.h>
|
||||
#include <ATen/Utils.h>
|
||||
#include <c10/core/Device.h>
|
||||
#include <c10/core/DeviceType.h>
|
||||
#include <c10/core/MemoryFormat.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/SymIntArrayRef.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <c10/util/ArrayRef.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <c10/util/OptionalArrayRef.h>
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#endif
|
||||
@@ -0,0 +1,49 @@
|
||||
// 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#include <ATen/AccumulateType.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
c10::ScalarType toAccumulateType(c10::ScalarType type, c10::DeviceType device) {
|
||||
switch (type) {
|
||||
#define DEFINE_CASE(scalar_t, TypeNum) \
|
||||
case ScalarType::TypeNum: \
|
||||
switch (device) { \
|
||||
case DeviceType::CUDA: \
|
||||
return CppTypeToScalarType< \
|
||||
at::acc_type_device<scalar_t, c10::DeviceType::CUDA>>::value; \
|
||||
default: \
|
||||
return CppTypeToScalarType< \
|
||||
at::acc_type_device<scalar_t, c10::DeviceType::CPU>>::value; \
|
||||
}
|
||||
|
||||
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_EXCEPT_COMPLEX_HALF_F8NZ(DEFINE_CASE)
|
||||
#undef DEFINE_CASE
|
||||
|
||||
default:
|
||||
TORCH_INTERNAL_ASSERT(false, "Unrecognized ScalarType: ", type);
|
||||
}
|
||||
}
|
||||
|
||||
c10::ScalarType toAccumulateType(c10::ScalarType type, bool is_cuda) {
|
||||
return is_cuda ? toAccumulateType(type, c10::DeviceType::CUDA)
|
||||
: toAccumulateType(type, c10::DeviceType::CPU);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,115 @@
|
||||
// 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <c10/core/DeviceType.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/util/BFloat16.h>
|
||||
#include <c10/util/Float8_e4m3fn.h>
|
||||
// #include <c10/util/Float8_e4m3fnuz.h>
|
||||
#include <c10/util/Float8_e5m2.h>
|
||||
// #include <c10/util/Float8_e5m2fnuz.h>
|
||||
#include <c10/util/Half.h>
|
||||
|
||||
#if defined(__CUDACC__)
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#elif defined(__HIPCC__)
|
||||
#include <hip/hip_fp16.h>
|
||||
#include <hip/hip_runtime.h>
|
||||
#endif
|
||||
|
||||
namespace at {
|
||||
|
||||
template <typename T, c10::DeviceType D>
|
||||
struct AccumulateTypeDevice {};
|
||||
|
||||
template <typename T, bool>
|
||||
struct AccumulateType {};
|
||||
|
||||
template <typename T>
|
||||
struct AccumulateType<T, false> {
|
||||
using type = typename AccumulateTypeDevice<T, c10::DeviceType::CPU>::type;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct AccumulateType<T, true> {
|
||||
using type = typename AccumulateTypeDevice<T, c10::DeviceType::CUDA>::type;
|
||||
};
|
||||
|
||||
template <typename T, c10::DeviceType device>
|
||||
using acc_type_device = typename AccumulateTypeDevice<T, device>::type;
|
||||
|
||||
template <typename T, bool is_cuda>
|
||||
using acc_type = typename AccumulateType<T, is_cuda>::type;
|
||||
|
||||
#define ACC_TYPE(t, acc_t, device_type) \
|
||||
template <> \
|
||||
struct AccumulateTypeDevice<t, device_type> { \
|
||||
using type = acc_t; \
|
||||
};
|
||||
|
||||
#define CUDA_ACC_TYPE(t, acc_t) ACC_TYPE(t, acc_t, c10::DeviceType::CUDA)
|
||||
#define CPU_ACC_TYPE(t, acc_t) ACC_TYPE(t, acc_t, c10::DeviceType::CPU)
|
||||
|
||||
#if defined(__CUDACC__) || defined(__HIPCC__)
|
||||
CUDA_ACC_TYPE(half, float)
|
||||
#endif
|
||||
CUDA_ACC_TYPE(BFloat16, float)
|
||||
CUDA_ACC_TYPE(Half, float)
|
||||
CUDA_ACC_TYPE(Float8_e5m2, float)
|
||||
CUDA_ACC_TYPE(Float8_e4m3fn, float)
|
||||
// CUDA_ACC_TYPE(Float8_e5m2fnuz, float)
|
||||
// CUDA_ACC_TYPE(Float8_e4m3fnuz, float)
|
||||
CUDA_ACC_TYPE(float, float)
|
||||
CUDA_ACC_TYPE(double, double)
|
||||
CUDA_ACC_TYPE(int8_t, int64_t)
|
||||
CUDA_ACC_TYPE(uint8_t, int64_t)
|
||||
CUDA_ACC_TYPE(char, int64_t)
|
||||
CUDA_ACC_TYPE(int16_t, int64_t)
|
||||
CUDA_ACC_TYPE(int32_t, int64_t)
|
||||
CUDA_ACC_TYPE(int64_t, int64_t)
|
||||
CUDA_ACC_TYPE(bool, bool)
|
||||
CUDA_ACC_TYPE(c10::complex<Half>, c10::complex<float>)
|
||||
CUDA_ACC_TYPE(c10::complex<float>, c10::complex<float>)
|
||||
CUDA_ACC_TYPE(c10::complex<double>, c10::complex<double>)
|
||||
|
||||
CPU_ACC_TYPE(BFloat16, float)
|
||||
CPU_ACC_TYPE(Half, float)
|
||||
CPU_ACC_TYPE(Float8_e5m2, float)
|
||||
CPU_ACC_TYPE(Float8_e4m3fn, float)
|
||||
// CPU_ACC_TYPE(Float8_e5m2fnuz, float)
|
||||
// CPU_ACC_TYPE(Float8_e4m3fnuz, float)
|
||||
CPU_ACC_TYPE(float, double)
|
||||
CPU_ACC_TYPE(double, double)
|
||||
CPU_ACC_TYPE(int8_t, int64_t)
|
||||
CPU_ACC_TYPE(uint8_t, int64_t)
|
||||
CPU_ACC_TYPE(char, int64_t)
|
||||
CPU_ACC_TYPE(int16_t, int64_t)
|
||||
CPU_ACC_TYPE(int32_t, int64_t)
|
||||
CPU_ACC_TYPE(int64_t, int64_t)
|
||||
CPU_ACC_TYPE(bool, bool)
|
||||
CPU_ACC_TYPE(c10::complex<Half>, c10::complex<float>)
|
||||
CPU_ACC_TYPE(c10::complex<float>, c10::complex<double>)
|
||||
CPU_ACC_TYPE(c10::complex<double>, c10::complex<double>)
|
||||
|
||||
c10::ScalarType toAccumulateType(c10::ScalarType type, c10::DeviceType device);
|
||||
c10::ScalarType toAccumulateType(c10::ScalarType type, bool is_cuda);
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,16 @@
|
||||
// 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 <c10/core/Device.h>
|
||||
@@ -0,0 +1,35 @@
|
||||
// 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 <ATen/Tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/util/Optional.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::optional<Device> device_of(const Tensor& t) {
|
||||
if (t.defined()) {
|
||||
return t.device();
|
||||
} else {
|
||||
return std::nullopt;
|
||||
}
|
||||
}
|
||||
|
||||
inline std::optional<Device> device_of(const std::optional<Tensor>& t) {
|
||||
return t.has_value() ? device_of(t.value()) : std::nullopt;
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,82 @@
|
||||
// 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 <ATen/ops/_local_scalar_dense.h>
|
||||
#include <ATen/ops/_nnz.h>
|
||||
#include <ATen/ops/_values.h>
|
||||
#include <ATen/ops/abs.h>
|
||||
#include <ATen/ops/all.h>
|
||||
#include <ATen/ops/allclose.h>
|
||||
#include <ATen/ops/any.h>
|
||||
#include <ATen/ops/arange.h>
|
||||
#include <ATen/ops/as_strided.h>
|
||||
#include <ATen/ops/cat.h>
|
||||
#include <ATen/ops/chunk.h>
|
||||
#include <ATen/ops/clamp.h>
|
||||
#include <ATen/ops/coalesce.h>
|
||||
#include <ATen/ops/detach.h>
|
||||
#include <ATen/ops/dsplit.h>
|
||||
#include <ATen/ops/empty.h>
|
||||
#include <ATen/ops/empty_like.h>
|
||||
#include <ATen/ops/empty_strided.h>
|
||||
#include <ATen/ops/equal.h>
|
||||
#include <ATen/ops/expand.h>
|
||||
#include <ATen/ops/expand_as.h>
|
||||
#include <ATen/ops/eye.h>
|
||||
#include <ATen/ops/flatten.h>
|
||||
#include <ATen/ops/from_blob.h>
|
||||
#include <ATen/ops/full.h>
|
||||
#include <ATen/ops/hsplit.h>
|
||||
#include <ATen/ops/index.h>
|
||||
#include <ATen/ops/index_put.h>
|
||||
#include <ATen/ops/is_coalesced.h>
|
||||
#include <ATen/ops/item.h>
|
||||
#include <ATen/ops/masked_select.h>
|
||||
#include <ATen/ops/narrow.h>
|
||||
#include <ATen/ops/narrow_copy.h>
|
||||
#include <ATen/ops/new_empty.h>
|
||||
#include <ATen/ops/new_full.h>
|
||||
#include <ATen/ops/new_ones.h>
|
||||
#include <ATen/ops/new_zeros.h>
|
||||
#include <ATen/ops/ones.h>
|
||||
#include <ATen/ops/permute.h>
|
||||
#include <ATen/ops/reciprocal.h>
|
||||
#include <ATen/ops/record_stream.h>
|
||||
#include <ATen/ops/rename.h>
|
||||
#include <ATen/ops/reshape.h>
|
||||
#include <ATen/ops/resize.h>
|
||||
#include <ATen/ops/select.h>
|
||||
#include <ATen/ops/slice.h>
|
||||
#include <ATen/ops/sparse_coo_tensor.h>
|
||||
#include <ATen/ops/sparse_csr_tensor.h>
|
||||
#include <ATen/ops/split.h>
|
||||
#include <ATen/ops/split_with_sizes.h>
|
||||
#include <ATen/ops/squeeze.h>
|
||||
#include <ATen/ops/std.h>
|
||||
#include <ATen/ops/sum.h>
|
||||
#include <ATen/ops/t.h>
|
||||
#include <ATen/ops/tensor_split.h>
|
||||
#include <ATen/ops/to.h>
|
||||
#include <ATen/ops/transpose.h>
|
||||
#include <ATen/ops/unflatten.h>
|
||||
#include <ATen/ops/unsafe_split.h>
|
||||
#include <ATen/ops/unsafe_split_with_sizes.h>
|
||||
#include <ATen/ops/unsqueeze.h>
|
||||
#include <ATen/ops/view.h>
|
||||
#include <ATen/ops/view_as.h>
|
||||
#include <ATen/ops/vsplit.h>
|
||||
#include <ATen/ops/zeros.h>
|
||||
#include <ATen/ops/zeros_like.h>
|
||||
@@ -0,0 +1,69 @@
|
||||
// Copyright (c) 2026 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 <c10/core/ScalarType.h>
|
||||
#include <c10/util/BFloat16.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <c10/util/Float8_e4m3fn.h>
|
||||
#include <c10/util/Float8_e5m2.h>
|
||||
#include <c10/util/Half.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
// For FP16 or BFloat16 inputs, ops should perform internal math in FP32.
|
||||
template <typename scalar_t>
|
||||
struct OpMathType {
|
||||
using type = scalar_t;
|
||||
};
|
||||
template <>
|
||||
struct OpMathType<at::Half> {
|
||||
using type = float;
|
||||
};
|
||||
template <>
|
||||
struct OpMathType<at::BFloat16> {
|
||||
using type = float;
|
||||
};
|
||||
template <>
|
||||
struct OpMathType<at::Float8_e5m2> {
|
||||
using type = float;
|
||||
};
|
||||
template <>
|
||||
struct OpMathType<at::Float8_e4m3fn> {
|
||||
using type = float;
|
||||
};
|
||||
template <>
|
||||
struct OpMathType<c10::complex<Half>> {
|
||||
using type = c10::complex<float>;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using opmath_type = typename OpMathType<T>::type;
|
||||
|
||||
inline c10::ScalarType toOpMathType(const c10::ScalarType type) {
|
||||
switch (type) {
|
||||
#define DEFINE_CASE(scalar_t, TypeNum) \
|
||||
case ScalarType::TypeNum: \
|
||||
return CppTypeToScalarType<at::opmath_type<scalar_t>>::value;
|
||||
|
||||
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_CASE)
|
||||
#undef DEFINE_CASE
|
||||
|
||||
default:
|
||||
TORCH_INTERNAL_ASSERT(false, "Unrecognized ScalarType: ", type);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,17 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
@@ -0,0 +1,124 @@
|
||||
// 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 <c10/core/SymInt.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <limits>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <stdexcept>
|
||||
#include <type_traits>
|
||||
#include <utility>
|
||||
|
||||
namespace at {
|
||||
class Tensor;
|
||||
}
|
||||
|
||||
namespace at::indexing {
|
||||
|
||||
constexpr int64_t INDEX_MIN = std::numeric_limits<int64_t>::min();
|
||||
constexpr int64_t INDEX_MAX = std::numeric_limits<int64_t>::max();
|
||||
|
||||
enum class TensorIndexType { None, Ellipsis, SymInt, Boolean, Slice, Tensor };
|
||||
|
||||
constexpr std::nullopt_t None = std::nullopt;
|
||||
|
||||
struct EllipsisIndexType final {
|
||||
EllipsisIndexType() = default;
|
||||
};
|
||||
|
||||
const EllipsisIndexType Ellipsis = EllipsisIndexType();
|
||||
|
||||
struct Slice final {
|
||||
public:
|
||||
Slice(std::optional<c10::SymInt> start_index = std::nullopt,
|
||||
std::optional<c10::SymInt> stop_index = std::nullopt,
|
||||
std::optional<c10::SymInt> step_index = std::nullopt) {
|
||||
if (!step_index.has_value()) {
|
||||
step_ = c10::SymInt(1);
|
||||
} else {
|
||||
step_ = std::move(step_index).value();
|
||||
}
|
||||
|
||||
if (!start_index.has_value()) {
|
||||
start_ = c10::SymInt(step_ < 0 ? INDEX_MAX : 0);
|
||||
} else {
|
||||
start_ = std::move(start_index).value();
|
||||
}
|
||||
|
||||
if (!stop_index.has_value()) {
|
||||
stop_ = c10::SymInt(step_ < 0 ? INDEX_MIN : INDEX_MAX);
|
||||
} else {
|
||||
stop_ = std::move(stop_index).value();
|
||||
}
|
||||
}
|
||||
|
||||
inline c10::SymInt start() const { return start_; }
|
||||
inline c10::SymInt stop() const { return stop_; }
|
||||
inline c10::SymInt step() const { return step_; }
|
||||
|
||||
private:
|
||||
c10::SymInt start_;
|
||||
c10::SymInt stop_;
|
||||
c10::SymInt step_;
|
||||
};
|
||||
|
||||
struct TensorIndex final {
|
||||
TensorIndex(std::nullopt_t /*unused*/) : type_(TensorIndexType::None) {}
|
||||
|
||||
TensorIndex(at::indexing::EllipsisIndexType /*unused*/)
|
||||
: type_(TensorIndexType::Ellipsis) {}
|
||||
TensorIndex(const char* str) : TensorIndex(at::indexing::Ellipsis) {
|
||||
if (std::strcmp(str, "...") != 0) {
|
||||
throw std::invalid_argument(
|
||||
"Expected \"...\" to represent an ellipsis index.");
|
||||
}
|
||||
}
|
||||
|
||||
TensorIndex(c10::SymInt integer)
|
||||
: integer_(std::move(integer)), type_(TensorIndexType::SymInt) {}
|
||||
TensorIndex(int64_t integer) : TensorIndex(c10::SymInt(integer)) {}
|
||||
TensorIndex(int integer) : TensorIndex(c10::SymInt(integer)) {}
|
||||
|
||||
template <class T, class = std::enable_if_t<std::is_same_v<bool, T>>>
|
||||
TensorIndex(T boolean) : boolean_(boolean), type_(TensorIndexType::Boolean) {}
|
||||
|
||||
TensorIndex(Slice slice)
|
||||
: slice_(std::move(slice)), type_(TensorIndexType::Slice) {}
|
||||
|
||||
TensorIndex(const at::Tensor& tensor);
|
||||
|
||||
inline bool is_none() const { return type_ == TensorIndexType::None; }
|
||||
inline bool is_ellipsis() const { return type_ == TensorIndexType::Ellipsis; }
|
||||
inline bool is_integer() const { return type_ == TensorIndexType::SymInt; }
|
||||
inline c10::SymInt integer() const { return integer_; }
|
||||
inline bool is_boolean() const { return type_ == TensorIndexType::Boolean; }
|
||||
inline bool boolean() const { return boolean_; }
|
||||
inline bool is_slice() const { return type_ == TensorIndexType::Slice; }
|
||||
inline const Slice& slice() const { return slice_; }
|
||||
inline bool is_tensor() const { return type_ == TensorIndexType::Tensor; }
|
||||
const at::Tensor& tensor() const;
|
||||
|
||||
private:
|
||||
c10::SymInt integer_ = 0;
|
||||
bool boolean_ = false;
|
||||
Slice slice_;
|
||||
std::shared_ptr<at::Tensor> tensor_;
|
||||
TensorIndexType type_;
|
||||
};
|
||||
|
||||
} // namespace at::indexing
|
||||
@@ -0,0 +1,97 @@
|
||||
// Copyright (c) 2026 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 <ATen/Utils.h>
|
||||
|
||||
#include <ATen/ops/empty.h>
|
||||
#include <ATen/ops/to.h>
|
||||
#include <c10/core/Layout.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/util/ArrayRef.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <c10/util/accumulate.h>
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#include "paddle/common/macros.h"
|
||||
#include "paddle/phi/api/include/sparse_api.h"
|
||||
#include "paddle/phi/api/include/tensor.h"
|
||||
|
||||
namespace at {
|
||||
namespace detail {
|
||||
|
||||
template <typename T>
|
||||
Tensor tensor_cpu(ArrayRef<T> values, const TensorOptions& options) {
|
||||
constexpr auto native_scalar_type = c10::CppTypeToScalarType<T>::value;
|
||||
auto result = at::empty(values.size(), options.dtype(native_scalar_type));
|
||||
PD_CHECK(result.is_contiguous());
|
||||
std::copy(values.begin(), values.end(), result.template data_ptr<T>());
|
||||
if (options.dtype() != native_scalar_type) {
|
||||
return result.to(at::TensorOptions().dtype(options.dtype()));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Tensor tensor_backend(ArrayRef<T> values, const TensorOptions& options) {
|
||||
auto cpu_tensor =
|
||||
tensor_cpu(values, options.device(c10::Device(c10::DeviceType::CPU)));
|
||||
return cpu_tensor.to(options.device());
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Tensor tensor_complex_cpu(ArrayRef<T> values, const TensorOptions& options) {
|
||||
constexpr auto native_scalar_type = c10::CppTypeToScalarType<T>::value;
|
||||
auto result = at::empty(values.size(), options.dtype(native_scalar_type));
|
||||
PD_CHECK(result.is_contiguous());
|
||||
std::copy(values.begin(), values.end(), result.template data_ptr<T>());
|
||||
if (options.dtype() != native_scalar_type) {
|
||||
return result.to(at::TensorOptions().dtype(options.dtype()));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Tensor tensor_complex_backend(ArrayRef<T> values,
|
||||
const TensorOptions& options) {
|
||||
auto cpu_tensor = tensor_complex_cpu(
|
||||
values, options.device(c10::Device(c10::DeviceType::CPU)));
|
||||
return cpu_tensor.to(options.device());
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
#define TENSOR(T, _1) \
|
||||
PADDLE_API Tensor tensor(ArrayRef<T> values, const TensorOptions& options) { \
|
||||
if (options.device().type() != c10::DeviceType::CPU) { \
|
||||
return at::detail::tensor_backend(values, options); \
|
||||
} else { \
|
||||
return at::detail::tensor_cpu(values, options); \
|
||||
} \
|
||||
}
|
||||
AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR)
|
||||
#undef TENSOR
|
||||
|
||||
#define TENSOR(T, _1) \
|
||||
PADDLE_API Tensor tensor(ArrayRef<T> values, const TensorOptions& options) { \
|
||||
if (options.device().type() != c10::DeviceType::CPU) { \
|
||||
return at::detail::tensor_complex_backend(values, options); \
|
||||
} else { \
|
||||
return at::detail::tensor_complex_cpu(values, options); \
|
||||
} \
|
||||
}
|
||||
AT_FORALL_COMPLEX_TYPES(TENSOR)
|
||||
#undef TENSOR
|
||||
|
||||
} // namespace at
|
||||
@@ -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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
namespace detail {
|
||||
|
||||
template <typename T>
|
||||
Tensor tensor_cpu(ArrayRef<T> values, const TensorOptions& options);
|
||||
|
||||
template <typename T>
|
||||
Tensor tensor_backend(ArrayRef<T> values, const TensorOptions& options);
|
||||
|
||||
template <typename T>
|
||||
Tensor tensor_complex_cpu(ArrayRef<T> values, const TensorOptions& options);
|
||||
|
||||
template <typename T>
|
||||
Tensor tensor_complex_backend(ArrayRef<T> values, const TensorOptions& options);
|
||||
|
||||
} // namespace detail
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,197 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <c10/core/Device.h>
|
||||
#include <c10/core/DispatchKeySet.h>
|
||||
#include <c10/core/GeneratorImpl.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <c10/util/intrusive_ptr.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <deque>
|
||||
#include <memory>
|
||||
#include <mutex> // NOLINT(build/c++11)
|
||||
#include <optional>
|
||||
#include <utility>
|
||||
|
||||
/**
|
||||
* Note [Generator]
|
||||
* ~~~~~~~~~~~~~~~~
|
||||
* A Pseudo Random Number Generator (PRNG) is an engine that uses an algorithm
|
||||
* to generate a seemingly random sequence of numbers, that may be later be
|
||||
* used in creating a random distribution. Such an engine almost always
|
||||
* maintains a state and requires a seed to start off the creation of random
|
||||
* numbers. Often times, users have found it beneficial to be able to
|
||||
* explicitly create, retain, and destroy PRNG states and also be able to
|
||||
* have control over the seed value.
|
||||
*
|
||||
* A Generator in ATen gives users the ability to read, write and modify a
|
||||
* PRNG engine. For instance, it does so by letting users seed a PRNG engine,
|
||||
* fork the state of the engine, etc.
|
||||
*
|
||||
* By default, there is one generator per device, and a device's generator is
|
||||
* lazily created. A user can use the torch.Generator() api to create their
|
||||
* own generator.
|
||||
*
|
||||
* This implementation wraps Paddle's phi::Generator via c10::GeneratorImpl.
|
||||
*/
|
||||
|
||||
/**
|
||||
* Note [Acquire lock when using random generators]
|
||||
* ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
* Generator and its derived classes are NOT thread-safe. Please use the
|
||||
* public mutex_ when using any methods from these classes, except for the
|
||||
* read-only methods.
|
||||
*/
|
||||
|
||||
// Forward declare PyObject if not already available.
|
||||
#ifndef PyObject
|
||||
struct _object;
|
||||
using PyObject = _object;
|
||||
#endif
|
||||
|
||||
namespace at {
|
||||
|
||||
using c10::Device;
|
||||
using c10::DispatchKeySet;
|
||||
|
||||
class Tensor;
|
||||
|
||||
struct Generator {
|
||||
Generator() = default;
|
||||
|
||||
explicit Generator(c10::intrusive_ptr<c10::GeneratorImpl> gen_impl)
|
||||
: impl_(std::move(gen_impl)) {
|
||||
TORCH_CHECK(impl_.get(), "GeneratorImpl with nullptr is not supported");
|
||||
}
|
||||
|
||||
bool operator==(const Generator& rhs) const {
|
||||
return this->impl_ == rhs.impl_;
|
||||
}
|
||||
|
||||
bool operator!=(const Generator& rhs) const { return !((*this) == rhs); }
|
||||
|
||||
bool defined() const { return static_cast<bool>(impl_); }
|
||||
|
||||
c10::GeneratorImpl* unsafeGetGeneratorImpl() const { return impl_.get(); }
|
||||
|
||||
c10::GeneratorImpl* unsafeReleaseGeneratorImpl() { return impl_.release(); }
|
||||
|
||||
const c10::intrusive_ptr<c10::GeneratorImpl>& getIntrusivePtr() const {
|
||||
return impl_;
|
||||
}
|
||||
|
||||
void set_current_seed(uint64_t seed) { impl_->set_current_seed(seed); }
|
||||
|
||||
/// Sets the offset of Generator state to the desired offset.
|
||||
/// Supported for Philox based Generators (CUDA / MPS).
|
||||
void set_offset(uint64_t offset) { impl_->set_offset(offset); }
|
||||
|
||||
/// Returns the offset of Generator state.
|
||||
/// Supported for Philox based Generators (CUDA / MPS).
|
||||
uint64_t get_offset() const { return impl_->get_offset(); }
|
||||
|
||||
uint64_t current_seed() const { return impl_->current_seed(); }
|
||||
|
||||
uint64_t seed() { return impl_->seed(); }
|
||||
|
||||
// ----- state transfer (not inlined to break header cycles) ----------------
|
||||
// These methods mirror PyTorch's set_state / get_state which operate on
|
||||
// serialised byte tensors. In the Paddle compat layer we provide a simpler
|
||||
// state-copy semantic through graphsafe_set_state / graphsafe_get_state.
|
||||
|
||||
/// Copy the full PRNG state from another Generator.
|
||||
void graphsafe_set_state(const Generator& src) {
|
||||
TORCH_CHECK(src.defined(), "Source generator is not defined");
|
||||
TORCH_CHECK(defined(), "Target generator is not defined");
|
||||
auto src_state = src.impl_->paddle_generator()->GetState();
|
||||
impl_->paddle_generator()->SetState(src_state);
|
||||
}
|
||||
|
||||
/// Obtain a Generator whose state is a snapshot (clone) of this one.
|
||||
Generator graphsafe_get_state() const {
|
||||
TORCH_CHECK(defined(), "Generator is not defined");
|
||||
return clone();
|
||||
}
|
||||
|
||||
std::mutex& mutex() { return impl_->mutex_; }
|
||||
|
||||
DispatchKeySet key_set() const { return impl_->key_set(); }
|
||||
|
||||
Device device() const { return impl_->device(); }
|
||||
|
||||
inline void set_pyobj(PyObject* pyobj) const noexcept {
|
||||
impl_->set_pyobj(pyobj);
|
||||
}
|
||||
|
||||
inline PyObject* pyobj() const noexcept { return impl_->pyobj(); }
|
||||
|
||||
template <typename T>
|
||||
T* get() const {
|
||||
return static_cast<T*>(impl_.get());
|
||||
}
|
||||
|
||||
Generator clone() const { return Generator(impl_->clone()); }
|
||||
|
||||
/// Access the underlying Paddle phi::Generator (convenience).
|
||||
std::shared_ptr<phi::Generator> paddle_generator() const {
|
||||
return impl_->paddle_generator();
|
||||
}
|
||||
|
||||
private:
|
||||
c10::intrusive_ptr<c10::GeneratorImpl> impl_;
|
||||
};
|
||||
|
||||
template <class Impl, class... Args>
|
||||
Generator make_generator(Args&&... args) {
|
||||
return Generator(c10::make_intrusive<Impl>(std::forward<Args>(args)...));
|
||||
}
|
||||
|
||||
/**
|
||||
* Utility function to static cast input Generator to
|
||||
* the backend generator type (CPUGeneratorImpl / CUDAGeneratorImpl etc.)
|
||||
*/
|
||||
template <typename T>
|
||||
inline T* check_generator(std::optional<Generator> gen) {
|
||||
TORCH_CHECK(gen.has_value(), "Expected Generator but received nullopt");
|
||||
TORCH_CHECK(gen->defined(),
|
||||
"Generator with undefined implementation is not allowed");
|
||||
TORCH_CHECK(
|
||||
T::device_type() == gen->device().type(),
|
||||
"Expected a generator for ",
|
||||
phi::AllocationTypeStr(c10::DeviceTypeToPhi(T::device_type())),
|
||||
" but found one for ",
|
||||
phi::AllocationTypeStr(c10::DeviceTypeToPhi(gen->device().type())));
|
||||
return gen->get<T>();
|
||||
}
|
||||
|
||||
/**
|
||||
* Utility function used in tensor implementations, which supplies the
|
||||
* default generator to tensors if an input generator is not supplied.
|
||||
* The input Generator is also static-cast to the backend generator type.
|
||||
*/
|
||||
template <typename T>
|
||||
inline T* get_generator_or_default(const std::optional<Generator>& gen,
|
||||
const Generator& default_gen) {
|
||||
return gen.has_value() && gen->defined() ? check_generator<T>(gen)
|
||||
: check_generator<T>(default_gen);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,15 @@
|
||||
// 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 <c10/core/Scalar.h>
|
||||
@@ -0,0 +1,17 @@
|
||||
// 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 <ATen/core/TensorBody.h>
|
||||
@@ -0,0 +1,104 @@
|
||||
// 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/headeronly/core/TensorAccessor.h>
|
||||
|
||||
#include <c10/macros/Macros.h>
|
||||
#include <c10/util/ArrayRef.h>
|
||||
#include <c10/util/Exception.h>
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <type_traits>
|
||||
|
||||
namespace at {
|
||||
|
||||
using torch::headeronly::DefaultPtrTraits;
|
||||
#if defined(__CUDACC__) || defined(__HIPCC__)
|
||||
using torch::headeronly::RestrictPtrTraits;
|
||||
#endif
|
||||
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits,
|
||||
typename index_t = int64_t>
|
||||
using TensorAccessorBase = torch::headeronly::detail::
|
||||
TensorAccessorBase<c10::IntArrayRef, T, N, PtrTraits, index_t>;
|
||||
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits,
|
||||
typename index_t = int64_t>
|
||||
using TensorAccessor = torch::headeronly::detail::
|
||||
TensorAccessor<c10::IntArrayRef, T, N, PtrTraits, index_t>;
|
||||
|
||||
namespace detail {
|
||||
|
||||
template <size_t N, typename index_t>
|
||||
struct IndexBoundsCheck {
|
||||
explicit IndexBoundsCheck(index_t i) {
|
||||
TORCH_CHECK(0 <= i && i < index_t{N},
|
||||
"Index ",
|
||||
i,
|
||||
" is not within bounds of a tensor of dimension ",
|
||||
N);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace detail
|
||||
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits,
|
||||
typename index_t = int64_t>
|
||||
using GenericPackedTensorAccessorBase =
|
||||
torch::headeronly::detail::GenericPackedTensorAccessorBase<
|
||||
detail::IndexBoundsCheck<N, index_t>,
|
||||
T,
|
||||
N,
|
||||
PtrTraits,
|
||||
index_t>;
|
||||
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits,
|
||||
typename index_t = int64_t>
|
||||
using GenericPackedTensorAccessor =
|
||||
torch::headeronly::detail::GenericPackedTensorAccessor<
|
||||
TensorAccessor<T, N - 1, PtrTraits, index_t>,
|
||||
detail::IndexBoundsCheck<N, index_t>,
|
||||
T,
|
||||
N,
|
||||
PtrTraits,
|
||||
index_t>;
|
||||
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits>
|
||||
using PackedTensorAccessor32 =
|
||||
GenericPackedTensorAccessor<T, N, PtrTraits, int32_t>;
|
||||
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits>
|
||||
using PackedTensorAccessor64 =
|
||||
GenericPackedTensorAccessor<T, N, PtrTraits, int64_t>;
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,578 @@
|
||||
// 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 <ATen/core/TensorAccessor.h>
|
||||
#include <c10/core/Device.h>
|
||||
#include <c10/core/Layout.h>
|
||||
#include <c10/core/MemoryFormat.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/Storage.h>
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <c10/core/SymIntArrayRef.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/int_array_ref_conversion.h>
|
||||
#include <utils/scalar_type_conversion.h>
|
||||
#include <algorithm>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <type_traits>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
#include "paddle/common/layout.h"
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/api/include/tensor.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
|
||||
namespace at {
|
||||
using PaddleTensor = paddle::Tensor;
|
||||
|
||||
class PADDLE_API TensorBase {
|
||||
public:
|
||||
TensorBase() = default;
|
||||
explicit TensorBase(const PaddleTensor& tensor) : tensor_(tensor) {
|
||||
InitStorage();
|
||||
}
|
||||
TensorBase(const TensorBase&) = default;
|
||||
TensorBase(TensorBase&&) noexcept = default;
|
||||
~TensorBase() noexcept = default;
|
||||
|
||||
#if defined(_MSC_VER)
|
||||
TensorBase& operator=(const TensorBase& x) & {
|
||||
tensor_ = x.tensor_;
|
||||
storage_ = x.storage_;
|
||||
return *this;
|
||||
}
|
||||
TensorBase& operator=(TensorBase&& x) & noexcept {
|
||||
tensor_ = std::move(x.tensor_);
|
||||
storage_ = std::move(x.storage_);
|
||||
return *this;
|
||||
}
|
||||
#else
|
||||
TensorBase& operator=(const TensorBase& x) & = default;
|
||||
TensorBase& operator=(TensorBase&& x) & noexcept = default;
|
||||
#endif
|
||||
|
||||
TensorBase& operator=(const TensorBase&) && = delete;
|
||||
TensorBase& operator=(TensorBase&&) && noexcept = delete;
|
||||
|
||||
bool is_same(const TensorBase& other) const noexcept {
|
||||
return tensor_.impl().get() == other.tensor_.impl().get();
|
||||
}
|
||||
size_t use_count() const noexcept { return tensor_.impl().use_count(); }
|
||||
size_t weak_use_count() const noexcept {
|
||||
// PyTorch exposes an internal self weak-reference on live TensorImpls, so
|
||||
// the observable weak count starts at 1 even without user-created refs.
|
||||
return tensor_.defined() ? 1 : 0;
|
||||
}
|
||||
|
||||
void print() const {
|
||||
if (defined()) {
|
||||
std::cerr << '[' << toString() << ' ' << sizes() << ']' << '\n';
|
||||
} else {
|
||||
std::cerr << "[UndefinedTensor]" << '\n';
|
||||
}
|
||||
}
|
||||
|
||||
std::string toString() const {
|
||||
if (!tensor_.defined()) {
|
||||
return "UndefinedType";
|
||||
}
|
||||
|
||||
std::string backend_str;
|
||||
const auto& place = tensor_.place();
|
||||
|
||||
if (phi::is_cpu_place(place)) {
|
||||
backend_str = "CPU";
|
||||
} else if (phi::is_gpu_place(place)) {
|
||||
backend_str = "CUDA";
|
||||
} else {
|
||||
backend_str = "Undefined";
|
||||
}
|
||||
|
||||
std::string scalar_type_str = at::toString(scalar_type());
|
||||
|
||||
return backend_str + scalar_type_str + "Type";
|
||||
}
|
||||
|
||||
// Returns the tensor's current data pointer. Storage mutations flow through
|
||||
// the compat holder view, so tensor.data_ptr() stays aligned with storage()
|
||||
// while preserving tensor-specific offsets for views.
|
||||
void* data_ptr() const {
|
||||
if (!tensor_.defined()) {
|
||||
return nullptr;
|
||||
}
|
||||
return const_cast<void*>(tensor_.data());
|
||||
}
|
||||
template <typename T>
|
||||
T* data_ptr() const {
|
||||
return static_cast<T*>(data_ptr());
|
||||
}
|
||||
|
||||
const void* const_data_ptr() const { return data_ptr(); }
|
||||
|
||||
template <typename T>
|
||||
const T* const_data_ptr() const;
|
||||
|
||||
void* mutable_data_ptr() const { return data_ptr(); }
|
||||
|
||||
template <typename T>
|
||||
T* mutable_data_ptr() const;
|
||||
|
||||
int64_t stride(int64_t dim) const {
|
||||
if (dim < 0) {
|
||||
dim += tensor_.strides().size();
|
||||
}
|
||||
return tensor_.strides()[static_cast<int>(dim)];
|
||||
}
|
||||
|
||||
c10::SymInt sym_stride(int64_t dim) const {
|
||||
return static_cast<c10::SymInt>(stride(dim));
|
||||
}
|
||||
|
||||
c10::IntArrayRef strides() const {
|
||||
return compat::_PD_PhiDDimToIntArrayRef(tensor_.strides());
|
||||
}
|
||||
|
||||
c10::SymIntArrayRef sym_strides() const {
|
||||
return c10::SymIntArrayRef(
|
||||
reinterpret_cast<const c10::SymInt*>(strides().data()),
|
||||
strides().size());
|
||||
}
|
||||
|
||||
int64_t size(int64_t dim) const {
|
||||
if (dim < 0) {
|
||||
dim += tensor_.dims().size();
|
||||
}
|
||||
return tensor_.dims()[static_cast<int>(dim)];
|
||||
}
|
||||
|
||||
c10::SymInt sym_size(int64_t dim) const {
|
||||
return static_cast<c10::SymInt>(size(dim));
|
||||
}
|
||||
|
||||
c10::IntArrayRef sizes() const {
|
||||
return compat::_PD_PhiDDimToIntArrayRef(tensor_.dims());
|
||||
}
|
||||
|
||||
c10::SymIntArrayRef sym_sizes() const {
|
||||
return c10::SymIntArrayRef(
|
||||
reinterpret_cast<const c10::SymInt*>(sizes().data()), sizes().size());
|
||||
}
|
||||
|
||||
int64_t numel() const { return tensor_.numel(); }
|
||||
|
||||
c10::SymInt sym_numel() const { return static_cast<c10::SymInt>(numel()); }
|
||||
|
||||
caffe2::TypeMeta dtype() const {
|
||||
return caffe2::TypeMeta::fromScalarType(
|
||||
compat::_PD_PhiDataTypeToAtenScalarType(tensor_.dtype()));
|
||||
}
|
||||
|
||||
c10::Device device() const { return c10::Device(tensor_.place()); }
|
||||
c10::DeviceIndex get_device() const {
|
||||
return c10::Device(tensor_.place()).index();
|
||||
}
|
||||
|
||||
int64_t dim() const { return tensor_.dims().size(); }
|
||||
int64_t ndimension() const { return dim(); }
|
||||
|
||||
at::TensorBase contiguous(
|
||||
c10::MemoryFormat memory_format = c10::MemoryFormat::Contiguous) const {
|
||||
PD_CHECK(memory_format == c10::MemoryFormat::Contiguous,
|
||||
"`MemoryFormat` other than Contiguous");
|
||||
|
||||
return TensorBase(tensor_.contiguous());
|
||||
}
|
||||
|
||||
bool is_contiguous(
|
||||
at::MemoryFormat memory_format = at::MemoryFormat::Contiguous) const {
|
||||
PD_CHECK(memory_format == c10::MemoryFormat::Contiguous,
|
||||
"`MemoryFormat` other than Contiguous");
|
||||
|
||||
return tensor_.is_contiguous();
|
||||
}
|
||||
|
||||
bool is_non_overlapping_and_dense() const {
|
||||
if (numel() <= 1) {
|
||||
return true;
|
||||
}
|
||||
if (tensor_.is_contiguous()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
// For non-contiguous tensors, verify sorted strides form a valid dense
|
||||
// layout without gaps or overlaps.
|
||||
auto sizes_vec = sizes();
|
||||
auto strides_vec = strides();
|
||||
int64_t ndim = dim();
|
||||
|
||||
std::vector<int64_t> perm(ndim);
|
||||
for (int64_t i = 0; i < ndim; ++i) {
|
||||
perm[i] = i;
|
||||
}
|
||||
std::sort(perm.begin(), perm.end(), [&](int64_t a, int64_t b) {
|
||||
return strides_vec[a] < strides_vec[b];
|
||||
});
|
||||
|
||||
int64_t expected_stride = 1;
|
||||
for (int64_t i = 0; i < ndim; ++i) {
|
||||
int64_t dim_idx = perm[i];
|
||||
if (sizes_vec[dim_idx] == 0) {
|
||||
return true;
|
||||
}
|
||||
if (sizes_vec[dim_idx] == 1) {
|
||||
continue;
|
||||
}
|
||||
if (strides_vec[dim_idx] != expected_stride) {
|
||||
return false;
|
||||
}
|
||||
expected_stride *= sizes_vec[dim_idx];
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool is_contiguous_or_false(
|
||||
at::MemoryFormat memory_format = at::MemoryFormat::Contiguous) const {
|
||||
PD_CHECK(memory_format == c10::MemoryFormat::Contiguous,
|
||||
"`MemoryFormat` other than Contiguous");
|
||||
|
||||
return tensor_.is_contiguous();
|
||||
}
|
||||
|
||||
c10::ScalarType scalar_type() const {
|
||||
return compat::_PD_PhiDataTypeToAtenScalarType(tensor_.dtype());
|
||||
}
|
||||
|
||||
bool has_names() const {
|
||||
// In PyTorch, has_names() is used to check if any dimension has names.
|
||||
// In Paddle, we don't support named dimension yet, so always return false.
|
||||
return false;
|
||||
}
|
||||
|
||||
TensorOptions options() const {
|
||||
return TensorOptions().dtype(dtype()).device(device()).layout(layout());
|
||||
}
|
||||
|
||||
const TensorBase& fill_(const at::Scalar& scalar) const {
|
||||
paddle::experimental::fill_(const_cast<PaddleTensor&>(tensor_), scalar);
|
||||
return *this;
|
||||
}
|
||||
|
||||
const TensorBase& zero_() const {
|
||||
paddle::experimental::fill_(const_cast<PaddleTensor&>(tensor_), 0.0);
|
||||
return *this;
|
||||
}
|
||||
|
||||
at::TensorBase to(
|
||||
at::TensorOptions options = {},
|
||||
bool non_blocking = false,
|
||||
bool copy = false,
|
||||
std::optional<at::MemoryFormat> memory_format = std::nullopt) const {
|
||||
if (options.device_opt().has_value()) {
|
||||
PADDLE_THROW(common::errors::Unimplemented(
|
||||
"The `to` method with device option is not supported yet."));
|
||||
}
|
||||
if (memory_format.has_value()) {
|
||||
PADDLE_THROW(common::errors::Unimplemented(
|
||||
"The `to` method with memory_format option is not supported yet."));
|
||||
}
|
||||
return TensorBase(paddle::experimental::cast(
|
||||
tensor_, compat::_PD_AtenScalarTypeToPhiDataType(options.dtype())));
|
||||
}
|
||||
|
||||
bool is_complex() const { return at::isComplexType(this->scalar_type()); }
|
||||
|
||||
bool is_floating_point() const {
|
||||
return at::isFloatingType(this->scalar_type());
|
||||
}
|
||||
|
||||
bool is_signed() const { return at::isSignedType(this->scalar_type()); }
|
||||
|
||||
bool is_cpu() const { return phi::is_cpu_place(tensor_.place()); }
|
||||
bool is_cuda() const { return phi::is_gpu_place(tensor_.place()); }
|
||||
|
||||
bool is_sparse() const {
|
||||
return tensor_.is_sparse_coo_tensor() || tensor_.is_sparse_csr_tensor();
|
||||
}
|
||||
|
||||
bool is_sparse_csr() const { return tensor_.is_sparse_csr_tensor(); }
|
||||
|
||||
inline size_t nbytes() const {
|
||||
PD_CHECK(
|
||||
((tensor_.layout() != common::DataLayout::SPARSE_COO) &&
|
||||
(tensor_.layout() != common::DataLayout::SPARSE_CSR)),
|
||||
"nbytes is not defined for sparse tensors. If you want the size of "
|
||||
"the constituent "
|
||||
"tensors, add the nbytes of the indices and values. If you want the "
|
||||
"size of the "
|
||||
"equivalent dense tensor, multiply numel() by element_size()");
|
||||
return tensor_.numel() * SizeOf(tensor_.dtype());
|
||||
}
|
||||
|
||||
size_t itemsize() const { return SizeOf(tensor_.dtype()); }
|
||||
|
||||
int64_t element_size() const {
|
||||
return static_cast<int64_t>(SizeOf(tensor_.dtype()));
|
||||
}
|
||||
|
||||
bool defined() const { return tensor_.defined(); }
|
||||
|
||||
Layout layout() const {
|
||||
switch (tensor_.layout()) {
|
||||
case common::DataLayout::STRIDED:
|
||||
case common::DataLayout::NCHW:
|
||||
case common::DataLayout::NHWC:
|
||||
case common::DataLayout::NCDHW:
|
||||
case common::DataLayout::NDHWC:
|
||||
return c10::kStrided;
|
||||
case common::DataLayout::SPARSE_COO:
|
||||
return c10::kSparse;
|
||||
case common::DataLayout::SPARSE_CSR:
|
||||
return c10::kSparseCsr;
|
||||
case common::DataLayout::ONEDNN:
|
||||
return c10::kMkldnn;
|
||||
default:
|
||||
return c10::kStrided;
|
||||
}
|
||||
}
|
||||
|
||||
void reset() {
|
||||
tensor_.reset();
|
||||
storage_.reset();
|
||||
}
|
||||
|
||||
int64_t storage_offset() const {
|
||||
// Paddle DenseTensor stores offset in meta_.offset (in bytes)
|
||||
// We need to convert to element offset
|
||||
auto dense_tensor =
|
||||
std::dynamic_pointer_cast<phi::DenseTensor>(tensor_.impl());
|
||||
if (dense_tensor) {
|
||||
size_t byte_offset = dense_tensor->meta().offset;
|
||||
size_t element_size = SizeOf(tensor_.dtype());
|
||||
return element_size > 0 ? static_cast<int64_t>(byte_offset / element_size)
|
||||
: 0;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
c10::SymInt sym_storage_offset() const {
|
||||
return c10::SymInt(storage_offset());
|
||||
}
|
||||
|
||||
bool has_storage() const {
|
||||
SyncStorageFromTensor();
|
||||
return tensor_.defined() && storage_ && storage_->valid();
|
||||
}
|
||||
|
||||
// Returns a Storage handle backed by the shared StorageImpl for this tensor.
|
||||
const c10::Storage& storage() const {
|
||||
SyncStorageFromTensor();
|
||||
static const c10::Storage kEmptyStorage;
|
||||
return storage_ ? *storage_ : kEmptyStorage;
|
||||
}
|
||||
|
||||
bool is_alias_of(const at::TensorBase& other) const {
|
||||
return this->storage().is_alias_of(other.storage());
|
||||
}
|
||||
|
||||
private:
|
||||
template <typename DenseT>
|
||||
static auto MaybeResetHolder(DenseT* dense,
|
||||
const std::shared_ptr<phi::Allocation>& holder,
|
||||
int)
|
||||
-> decltype(dense->ResetHolder(holder), void()) {
|
||||
dense->ResetHolder(holder);
|
||||
}
|
||||
|
||||
static void MaybeResetHolder(phi::DenseTensor* dense,
|
||||
const std::shared_ptr<phi::Allocation>& holder,
|
||||
long) { // NOLINT
|
||||
TORCH_CHECK(dense != nullptr, "DenseTensor must not be null");
|
||||
|
||||
// External custom-kernel builds do not expose ResetHolder(), but Holder()
|
||||
// still returns the live holder reference used by DenseTensor.
|
||||
if (dense->numel() == 0) {
|
||||
auto& meta = const_cast<phi::DenseTensorMeta&>(dense->meta());
|
||||
meta.offset = 0;
|
||||
const_cast<std::shared_ptr<phi::Allocation>&>(dense->Holder()) = holder;
|
||||
return;
|
||||
}
|
||||
|
||||
if (dense->Holder() && dense->meta().is_contiguous()) {
|
||||
TORCH_CHECK(holder != nullptr, "Holder must not be null.");
|
||||
const auto required_bytes =
|
||||
dense->numel() * static_cast<int64_t>(phi::SizeOf(dense->dtype())) +
|
||||
static_cast<int64_t>(dense->meta().offset);
|
||||
TORCH_CHECK(required_bytes <= static_cast<int64_t>(holder->size()),
|
||||
"The size of Holder is not enough to store the Tensor.");
|
||||
}
|
||||
const_cast<std::shared_ptr<phi::Allocation>&>(dense->Holder()) = holder;
|
||||
}
|
||||
|
||||
void InitStorage() { SyncStorageFromTensor(); }
|
||||
|
||||
static std::shared_ptr<c10::Storage> GetOrCreateCanonicalStorage(
|
||||
c10::Storage&& live_storage) {
|
||||
auto impl = live_storage.get_impl();
|
||||
if (!impl) {
|
||||
return std::make_shared<c10::Storage>(std::move(live_storage));
|
||||
}
|
||||
|
||||
static std::mutex registry_mu;
|
||||
static std::unordered_map<c10::StorageImpl*, std::weak_ptr<c10::Storage>>
|
||||
registry;
|
||||
|
||||
std::lock_guard<std::mutex> guard(registry_mu);
|
||||
auto it = registry.find(impl.get());
|
||||
if (it != registry.end()) {
|
||||
if (auto cached = it->second.lock()) {
|
||||
return cached;
|
||||
}
|
||||
registry.erase(it);
|
||||
}
|
||||
|
||||
auto created = std::make_shared<c10::Storage>(std::move(live_storage));
|
||||
registry.emplace(impl.get(), created);
|
||||
return created;
|
||||
}
|
||||
|
||||
void SyncStorageFromTensor() const {
|
||||
auto dense = std::dynamic_pointer_cast<phi::DenseTensor>(tensor_.impl());
|
||||
if (!dense) {
|
||||
storage_.reset();
|
||||
return;
|
||||
}
|
||||
|
||||
auto holder = dense->Holder();
|
||||
if (!holder) {
|
||||
storage_.reset();
|
||||
return;
|
||||
}
|
||||
|
||||
c10::Storage live_storage = c10::Storage::createTensorStorage(holder);
|
||||
auto compat_holder = live_storage.ensureTensorHolder();
|
||||
if (holder != compat_holder) {
|
||||
MaybeResetHolder(dense.get(), compat_holder, 0);
|
||||
}
|
||||
|
||||
if (!storage_ || storage_->get_impl() != live_storage.get_impl()) {
|
||||
storage_ = GetOrCreateCanonicalStorage(std::move(live_storage));
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
// Return a `TensorAccessor` for CPU `Tensor`s. You have to specify scalar
|
||||
// type and
|
||||
// dimension.
|
||||
template <typename T, size_t N>
|
||||
TensorAccessor<T, N> accessor() const& {
|
||||
static_assert(
|
||||
N > 0,
|
||||
"accessor is used for indexing tensor, for scalars use *data_ptr<T>()");
|
||||
TORCH_CHECK(dim() == N,
|
||||
"TensorAccessor expected ",
|
||||
N,
|
||||
" dims but tensor has ",
|
||||
dim());
|
||||
T* ptr = nullptr;
|
||||
if constexpr (std::is_const_v<T>) {
|
||||
ptr = const_data_ptr<T>();
|
||||
} else {
|
||||
ptr = mutable_data_ptr<T>();
|
||||
}
|
||||
return TensorAccessor<T, N>(ptr, sizes().data(), strides().data());
|
||||
}
|
||||
template <typename T, size_t N>
|
||||
TensorAccessor<T, N> accessor() && = delete;
|
||||
|
||||
// Return a `GenericPackedTensorAccessor` for CUDA `Tensor`s. You have to
|
||||
// specify scalar type and dimension. You can optionally specify
|
||||
// RestrictPtrTraits as a template parameter to cast the data pointer to a
|
||||
// __restrict__ pointer. In order to use this, your CUDA kernel has to take a
|
||||
// corresponding GenericPackedTensorAccessor as an argument.
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits,
|
||||
typename index_t = int64_t>
|
||||
GenericPackedTensorAccessor<T, N, PtrTraits, index_t>
|
||||
generic_packed_accessor() const& {
|
||||
static_assert(
|
||||
N > 0,
|
||||
"accessor is used for indexing tensor, for scalars use *data_ptr<T>()");
|
||||
TORCH_CHECK(dim() == N,
|
||||
"TensorAccessor expected ",
|
||||
N,
|
||||
" dims but tensor has ",
|
||||
dim());
|
||||
T* ptr = nullptr;
|
||||
if constexpr (std::is_const_v<T>) {
|
||||
ptr = const_data_ptr<T>();
|
||||
} else {
|
||||
ptr = mutable_data_ptr<T>();
|
||||
}
|
||||
return GenericPackedTensorAccessor<T, N, PtrTraits, index_t>(
|
||||
static_cast<typename PtrTraits<T>::PtrType>(ptr),
|
||||
sizes().data(),
|
||||
strides().data());
|
||||
}
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits,
|
||||
typename index_t = int64_t>
|
||||
GenericPackedTensorAccessor<T, N> generic_packed_accessor() && = delete;
|
||||
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits>
|
||||
PackedTensorAccessor32<T, N, PtrTraits> packed_accessor32() const& {
|
||||
TORCH_CHECK(
|
||||
numel() <= static_cast<int64_t>(std::numeric_limits<int32_t>::max()),
|
||||
"numel needs to be smaller than int32_t max; otherwise, please use "
|
||||
"packed_accessor64");
|
||||
return generic_packed_accessor<T, N, PtrTraits, int32_t>();
|
||||
}
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits>
|
||||
PackedTensorAccessor32<T, N, PtrTraits> packed_accessor32() && = delete;
|
||||
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits>
|
||||
PackedTensorAccessor64<T, N, PtrTraits> packed_accessor64() const& {
|
||||
return generic_packed_accessor<T, N, PtrTraits, int64_t>();
|
||||
}
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits>
|
||||
PackedTensorAccessor64<T, N, PtrTraits> packed_accessor64() && = delete;
|
||||
|
||||
const PaddleTensor& _PD_GetInner() const& { return tensor_; }
|
||||
PaddleTensor& _PD_GetInner() & { return tensor_; }
|
||||
PaddleTensor&& _PD_GetInner() && { return std::move(tensor_); }
|
||||
|
||||
protected:
|
||||
PaddleTensor tensor_;
|
||||
// Cache a canonical Storage object shared by wrappers that reference the
|
||||
// same StorageImpl. This prevents independently-constructed wrappers around
|
||||
// one tensor impl from inflating Storage::use_count().
|
||||
mutable std::shared_ptr<c10::Storage> storage_;
|
||||
};
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,846 @@
|
||||
// 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 <ATen/TensorIndexing.h>
|
||||
#include <ATen/core/TensorBase.h>
|
||||
#include <c10/core/Backend.h>
|
||||
#include <c10/core/List.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/Stream.h>
|
||||
#include <c10/core/SymIntArrayRef.h>
|
||||
#include <c10/util/OptionalArrayRef.h>
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/api/include/tensor.h"
|
||||
#include "paddle/phi/common/int_array.h"
|
||||
#include "paddle/phi/common/scalar.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/memory/malloc.h"
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
#include <hip/hip_runtime.h>
|
||||
#elif defined(PADDLE_WITH_CUDA)
|
||||
#include <cuda_runtime_api.h>
|
||||
#endif
|
||||
|
||||
#include <limits>
|
||||
#include <optional>
|
||||
#include <type_traits>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
#include "paddle/common/ddim.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
|
||||
namespace at {
|
||||
class Tensor;
|
||||
|
||||
// Type aliases for ATen compatibility
|
||||
using Scalar = c10::Scalar;
|
||||
using TensorOptions = c10::TensorOptions;
|
||||
using MemoryFormat = c10::MemoryFormat;
|
||||
using IntArrayRef = c10::IntArrayRef;
|
||||
using OptionalIntArrayRef = c10::OptionalIntArrayRef;
|
||||
using ScalarType = c10::ScalarType;
|
||||
using TensorList = c10::ArrayRef<Tensor>;
|
||||
using ITensorListRef = c10::ArrayRef<Tensor>;
|
||||
} // namespace at
|
||||
|
||||
namespace at { // NOLINT(build/namespaces)
|
||||
using PaddleTensor = paddle::Tensor;
|
||||
using PaddlePlace = phi::Place;
|
||||
|
||||
// Stub for DimnameList (not supported in Paddle)
|
||||
using DimnameList = c10::ArrayRef<std::string>;
|
||||
|
||||
using Stream = c10::Stream;
|
||||
|
||||
class Tensor : public TensorBase {
|
||||
public:
|
||||
Tensor() = default;
|
||||
Tensor(const PaddleTensor& tensor) : TensorBase(tensor){}; // NOLINT
|
||||
Tensor(const Tensor& tensor) = default;
|
||||
Tensor(Tensor&& tensor) = default;
|
||||
|
||||
// Implicitly move-constructible from TensorBase, but must be explicit to
|
||||
// increase refcount
|
||||
explicit Tensor(const TensorBase& base) : TensorBase(base) {} // NOLINT
|
||||
/*implicit*/ Tensor(TensorBase&& base) // NOLINT
|
||||
: TensorBase(std::move(base)) {}
|
||||
|
||||
Tensor& operator=(const PaddleTensor& x) & noexcept {
|
||||
tensor_ = x;
|
||||
return *this;
|
||||
}
|
||||
Tensor& operator=(const TensorBase& x) & noexcept {
|
||||
const PaddleTensor& inner = x._PD_GetInner();
|
||||
tensor_ = inner;
|
||||
return *this;
|
||||
}
|
||||
Tensor& operator=(PaddleTensor&& x) & noexcept {
|
||||
tensor_ = std::move(x);
|
||||
return *this;
|
||||
}
|
||||
Tensor& operator=(TensorBase&& x) & noexcept {
|
||||
tensor_ = std::move(x)._PD_GetInner();
|
||||
return *this;
|
||||
}
|
||||
|
||||
Tensor& operator=(const Tensor& x) & noexcept {
|
||||
return operator=(static_cast<const TensorBase&>(x));
|
||||
}
|
||||
Tensor& operator=(Tensor&& x) & noexcept {
|
||||
return operator=(static_cast<TensorBase&&>(x));
|
||||
}
|
||||
Tensor& operator=(const Scalar& v) && {
|
||||
fill_(v);
|
||||
return *this;
|
||||
}
|
||||
Tensor& operator=(const Tensor& rhs) && {
|
||||
copy_(rhs);
|
||||
return *this;
|
||||
}
|
||||
Tensor& operator=(Tensor&& rhs) && {
|
||||
copy_(rhs);
|
||||
return *this;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T* data() const {
|
||||
return data_ptr<T>();
|
||||
}
|
||||
|
||||
Tensor toBackend(c10::Backend b) const {
|
||||
switch (b) {
|
||||
case c10::Backend::CPU:
|
||||
return tensor_.copy_to(PaddlePlace(phi::AllocationType::CPU), true);
|
||||
case c10::Backend::CUDA:
|
||||
return tensor_.copy_to(paddle::DefaultGPUPlace(), true);
|
||||
case c10::Backend::XPU:
|
||||
return tensor_.copy_to(paddle::DefaultXPUPlace(), true);
|
||||
case c10::Backend::IPU:
|
||||
return tensor_.copy_to(PaddlePlace(phi::IPUPlace()), true);
|
||||
default:
|
||||
PD_CHECK(false, "Unsupported backend");
|
||||
}
|
||||
return tensor_;
|
||||
}
|
||||
|
||||
Tensor cpu() const {
|
||||
PaddlePlace place(phi::AllocationType::CPU);
|
||||
return tensor_.copy_to(place, true);
|
||||
}
|
||||
|
||||
Tensor cuda() const {
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
auto place = paddle::DefaultGPUPlace();
|
||||
return tensor_.copy_to(place, true);
|
||||
#elif defined(PADDLE_WITH_XPU)
|
||||
return tensor_.copy_to(paddle::DefaultXPUPlace(), true);
|
||||
#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
|
||||
return tensor_.copy_to(paddle::DefaultCustomPlace(), true);
|
||||
#else
|
||||
PD_THROW(
|
||||
"cuda() is not supported: no GPU/XPU/Custom device enabled "
|
||||
"in this build.");
|
||||
#endif
|
||||
}
|
||||
|
||||
using TensorBase::stride;
|
||||
|
||||
using TensorBase::size;
|
||||
|
||||
at::Tensor to(
|
||||
at::TensorOptions options = {},
|
||||
bool non_blocking = false,
|
||||
bool copy = false,
|
||||
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const;
|
||||
at::Tensor to(::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory,
|
||||
bool non_blocking,
|
||||
bool copy,
|
||||
::std::optional<at::MemoryFormat> memory_format) const;
|
||||
at::Tensor to(
|
||||
at::Device device,
|
||||
at::ScalarType dtype,
|
||||
bool non_blocking = false,
|
||||
bool copy = false,
|
||||
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const;
|
||||
at::Tensor to(
|
||||
at::ScalarType dtype,
|
||||
bool non_blocking = false,
|
||||
bool copy = false,
|
||||
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const;
|
||||
at::Tensor to(
|
||||
const at::Tensor& other,
|
||||
bool non_blocking = false,
|
||||
bool copy = false,
|
||||
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const;
|
||||
|
||||
Tensor meta() const {
|
||||
PD_THROW("`meta()` is not supported in this Paddle build.");
|
||||
}
|
||||
|
||||
at::Scalar item() const;
|
||||
|
||||
template <typename T>
|
||||
T item() const;
|
||||
|
||||
bool equal(const at::Tensor& other) const;
|
||||
|
||||
// Clamp functions
|
||||
at::Tensor clamp(
|
||||
const ::std::optional<at::Scalar>& min,
|
||||
const ::std::optional<at::Scalar>& max = ::std::nullopt) const;
|
||||
|
||||
at::Tensor clamp(const ::std::optional<at::Tensor>& min = {},
|
||||
const ::std::optional<at::Tensor>& max = {}) const;
|
||||
|
||||
at::Tensor& clamp_(
|
||||
const ::std::optional<at::Scalar>& min,
|
||||
const ::std::optional<at::Scalar>& max = ::std::nullopt) const;
|
||||
|
||||
at::Tensor& clamp_(const ::std::optional<at::Tensor>& min = {},
|
||||
const ::std::optional<at::Tensor>& max = {}) const;
|
||||
|
||||
at::Tensor clamp_max(const at::Scalar& max) const;
|
||||
at::Tensor clamp_max(const at::Tensor& max) const;
|
||||
at::Tensor& clamp_max_(const at::Scalar& max) const;
|
||||
at::Tensor& clamp_max_(const at::Tensor& max) const;
|
||||
|
||||
at::Tensor clamp_min(const at::Scalar& min) const;
|
||||
at::Tensor clamp_min(const at::Tensor& min) const;
|
||||
at::Tensor& clamp_min_(const at::Scalar& min) const;
|
||||
at::Tensor& clamp_min_(const at::Tensor& min) const;
|
||||
|
||||
// as_strided: Create a tensor view with custom size, stride, and
|
||||
// storage_offset
|
||||
at::Tensor as_strided(
|
||||
at::IntArrayRef size,
|
||||
at::IntArrayRef stride,
|
||||
::std::optional<int64_t> storage_offset = ::std::nullopt) const;
|
||||
|
||||
// as_strided_: Inplace version
|
||||
const at::Tensor& as_strided_(
|
||||
at::IntArrayRef size,
|
||||
at::IntArrayRef stride,
|
||||
::std::optional<int64_t> storage_offset = ::std::nullopt) const;
|
||||
|
||||
// as_strided_scatter: Scatter src into a strided view
|
||||
at::Tensor as_strided_scatter(
|
||||
const at::Tensor& src,
|
||||
at::IntArrayRef size,
|
||||
at::IntArrayRef stride,
|
||||
::std::optional<int64_t> storage_offset = ::std::nullopt) const;
|
||||
|
||||
// Standard deviation functions
|
||||
Tensor std(bool unbiased) const;
|
||||
Tensor std(at::OptionalIntArrayRef dim,
|
||||
bool unbiased,
|
||||
bool keepdim = false) const;
|
||||
Tensor std(at::OptionalIntArrayRef dim = ::std::nullopt,
|
||||
const ::std::optional<at::Scalar>& correction = ::std::nullopt,
|
||||
bool keepdim = false) const;
|
||||
Tensor std(int dim) const { return std(at::IntArrayRef{dim}); }
|
||||
|
||||
Tensor tensor_data() const {
|
||||
PaddleTensor result;
|
||||
if (tensor_.initialized()) {
|
||||
auto src_impl = tensor_.impl();
|
||||
auto* src_tensor =
|
||||
std::dynamic_pointer_cast<phi::DenseTensor>(src_impl).get();
|
||||
if (src_tensor && src_tensor->meta().is_contiguous()) {
|
||||
result.set_impl(std::make_shared<phi::DenseTensor>());
|
||||
auto* dst_tensor =
|
||||
std::dynamic_pointer_cast<phi::DenseTensor>(result.impl()).get();
|
||||
dst_tensor->ShareDataWith(*src_tensor);
|
||||
} else {
|
||||
result = paddle::experimental::assign(tensor_);
|
||||
}
|
||||
}
|
||||
// For uninitialized tensor, return an uninitialized tensor (no assign
|
||||
// needed)
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
Tensor variable_data() const {
|
||||
PaddleTensor result;
|
||||
if (tensor_.initialized()) {
|
||||
auto src_impl = tensor_.impl();
|
||||
auto* src_tensor =
|
||||
std::dynamic_pointer_cast<phi::DenseTensor>(src_impl).get();
|
||||
if (src_tensor && src_tensor->meta().is_contiguous()) {
|
||||
result.set_impl(std::make_shared<phi::DenseTensor>());
|
||||
auto* dst_tensor =
|
||||
std::dynamic_pointer_cast<phi::DenseTensor>(result.impl()).get();
|
||||
dst_tensor->ShareDataWith(*src_tensor);
|
||||
} else {
|
||||
result = paddle::experimental::assign(tensor_);
|
||||
}
|
||||
}
|
||||
// For uninitialized tensor, return an uninitialized tensor (no assign
|
||||
// needed)
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
// index: Get values at specified tensor indices
|
||||
at::Tensor index(const c10::List<::std::optional<at::Tensor>>& indices) const;
|
||||
|
||||
// index_put_: Set values at specified indices in-place
|
||||
at::Tensor& index_put_(const c10::List<::std::optional<at::Tensor>>& indices,
|
||||
const at::Tensor& values,
|
||||
bool accumulate = false) const;
|
||||
|
||||
// index_put: Non-inplace version of index_put_
|
||||
at::Tensor index_put(const c10::List<::std::optional<at::Tensor>>& indices,
|
||||
const at::Tensor& values,
|
||||
bool accumulate = false) const;
|
||||
|
||||
Tensor toType(ScalarType t) const {
|
||||
return Tensor(paddle::experimental::cast(
|
||||
tensor_, compat::_PD_AtenScalarTypeToPhiDataType(t)));
|
||||
}
|
||||
|
||||
at::Tensor contiguous(
|
||||
c10::MemoryFormat memory_format = c10::MemoryFormat::Contiguous) const {
|
||||
PD_CHECK(memory_format == c10::MemoryFormat::Contiguous,
|
||||
"`MemoryFormat` other than Contiguous");
|
||||
|
||||
return tensor_.contiguous();
|
||||
}
|
||||
|
||||
at::Tensor flatten(int64_t start_dim = 0, int64_t end_dim = -1) const;
|
||||
at::Tensor unflatten(int64_t dim, at::IntArrayRef sizes) const;
|
||||
at::Tensor unflatten_symint(int64_t dim, c10::SymIntArrayRef sizes) const;
|
||||
|
||||
Tensor& fill_(const at::Scalar& value) const {
|
||||
paddle::experimental::fill_(const_cast<PaddleTensor&>(tensor_), value);
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
Tensor& zero_() const {
|
||||
paddle::experimental::fill_(const_cast<PaddleTensor&>(tensor_), 0.0);
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
bool is_pinned(::std::optional<c10::Device> device = ::std::nullopt) const {
|
||||
if (device.has_value()) {
|
||||
phi::enforce::ThrowWarnInternal(
|
||||
"The argument 'device' of Tensor.is_pinned() is deprecated. "
|
||||
"Please do not pass this argument.");
|
||||
}
|
||||
|
||||
const PaddlePlace place = tensor_.place();
|
||||
const bool is_gpu_pinned = phi::is_cuda_pinned_place(place);
|
||||
const bool is_xpu_pinned = phi::is_xpu_pinned_place(place);
|
||||
|
||||
// Keep parity with PyTorch behavior: only host tensors are pinnable.
|
||||
if (!(phi::is_cpu_place(place) || is_gpu_pinned || is_xpu_pinned)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!device.has_value()) {
|
||||
return is_gpu_pinned || is_xpu_pinned;
|
||||
}
|
||||
|
||||
const auto device_type = device.value().type();
|
||||
if (device_type == c10::DeviceType::CUDA) {
|
||||
return is_gpu_pinned;
|
||||
}
|
||||
if (device_type == c10::DeviceType::XPU) {
|
||||
return is_xpu_pinned;
|
||||
}
|
||||
// CPU and non-accelerator devices are not valid pinned backends.
|
||||
return false;
|
||||
}
|
||||
|
||||
Tensor pin_memory(
|
||||
::std::optional<c10::Device> device = ::std::nullopt) const {
|
||||
if (device.has_value()) {
|
||||
phi::enforce::ThrowWarnInternal(
|
||||
"The argument 'device' of Tensor.pin_memory() is deprecated. "
|
||||
"Please do not pass this argument.");
|
||||
}
|
||||
|
||||
if (is_pinned(device)) {
|
||||
return *this;
|
||||
}
|
||||
|
||||
const PaddlePlace current_place = tensor_.place();
|
||||
if (!phi::is_cpu_place(current_place)) {
|
||||
PD_THROW("cannot pin '" + this->toString() +
|
||||
"', only dense CPU tensors can be pinned");
|
||||
}
|
||||
|
||||
PaddlePlace pinned_place;
|
||||
|
||||
if (device.has_value()) {
|
||||
const auto device_type = device.value().type();
|
||||
if (device_type == c10::DeviceType::CUDA) {
|
||||
pinned_place = phi::Place(phi::GPUPinnedPlace());
|
||||
} else if (device_type == c10::DeviceType::XPU) {
|
||||
pinned_place = phi::Place(phi::XPUPinnedPlace());
|
||||
} else {
|
||||
PD_THROW("pin_memory device type must be an accelerator (GPU/XPU)");
|
||||
}
|
||||
} else {
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
pinned_place = phi::Place(phi::GPUPinnedPlace());
|
||||
#elif defined(PADDLE_WITH_XPU)
|
||||
pinned_place = phi::Place(phi::XPUPinnedPlace());
|
||||
#else
|
||||
PD_THROW("pin_memory is not supported: no GPU/XPU backend enabled");
|
||||
#endif
|
||||
}
|
||||
|
||||
return tensor_.copy_to(pinned_place, true);
|
||||
}
|
||||
|
||||
at::Tensor narrow_copy(int64_t dim, int64_t start, int64_t length) const;
|
||||
at::Tensor narrow_copy_symint(int64_t dim,
|
||||
c10::SymInt start,
|
||||
c10::SymInt length) const;
|
||||
|
||||
at::Tensor narrow(int64_t dim, int64_t start, int64_t length) const;
|
||||
at::Tensor narrow_symint(int64_t dim,
|
||||
c10::SymInt start,
|
||||
c10::SymInt length) const;
|
||||
at::Tensor narrow(int64_t dim, const at::Tensor& start, int64_t length) const;
|
||||
at::Tensor narrow_symint(int64_t dim,
|
||||
const at::Tensor& start,
|
||||
c10::SymInt length) const;
|
||||
|
||||
at::Tensor reshape(at::IntArrayRef shape) const;
|
||||
|
||||
at::Tensor transpose(int64_t dim0, int64_t dim1) const;
|
||||
at::Tensor& transpose_(int64_t dim0, int64_t dim1) const;
|
||||
|
||||
at::Tensor permute(at::IntArrayRef dims) const;
|
||||
|
||||
at::Tensor reciprocal() const;
|
||||
at::Tensor& reciprocal_() const;
|
||||
|
||||
at::Tensor detach() const;
|
||||
at::Tensor& detach_() const;
|
||||
|
||||
at::Tensor select(int64_t dim, int64_t index) const;
|
||||
at::Tensor select_symint(int64_t dim, c10::SymInt index) const;
|
||||
|
||||
at::Tensor& copy_(const at::Tensor& src, bool non_blocking = false) const {
|
||||
const_cast<PaddleTensor&>(tensor_).copy_(
|
||||
src._PD_GetInner(), tensor_.place(), /*blocking=*/!non_blocking);
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
at::Tensor view(at::IntArrayRef size) const;
|
||||
at::Tensor view(at::ScalarType dtype) const;
|
||||
|
||||
at::Tensor squeeze() const;
|
||||
at::Tensor squeeze(int64_t dim) const;
|
||||
at::Tensor squeeze(at::IntArrayRef dim) const;
|
||||
at::Tensor& squeeze_() const;
|
||||
at::Tensor& squeeze_(int64_t dim) const;
|
||||
at::Tensor& squeeze_(at::IntArrayRef dim) const;
|
||||
|
||||
at::Tensor unsqueeze(int64_t dim) const;
|
||||
at::Tensor& unsqueeze_(int64_t dim) const;
|
||||
|
||||
at::Tensor sum(::std::optional<at::ScalarType> dtype = ::std::nullopt) const;
|
||||
at::Tensor sum(at::OptionalIntArrayRef dim,
|
||||
bool keepdim = false,
|
||||
::std::optional<at::ScalarType> dtype = ::std::nullopt) const;
|
||||
|
||||
at::Tensor t() const;
|
||||
at::Tensor& t_() const;
|
||||
|
||||
at::Tensor view_as(const at::Tensor& other) const;
|
||||
|
||||
at::Tensor coalesce() const;
|
||||
bool is_coalesced() const;
|
||||
|
||||
int64_t _nnz() const;
|
||||
at::Tensor _values() const;
|
||||
|
||||
bool is_variable() const noexcept { return true; }
|
||||
|
||||
at::Tensor index_select(int64_t dim, const at::Tensor& index) const {
|
||||
return Tensor(
|
||||
paddle::experimental::index_select(tensor_, index._PD_GetInner(), dim));
|
||||
}
|
||||
|
||||
at::Tensor masked_select(const at::Tensor& mask) const;
|
||||
|
||||
std::vector<at::Tensor> tensor_split(int64_t sections, int64_t dim) const;
|
||||
std::vector<at::Tensor> tensor_split_symint(c10::SymInt sections,
|
||||
int64_t dim) const;
|
||||
std::vector<at::Tensor> tensor_split(at::IntArrayRef indices,
|
||||
int64_t dim) const;
|
||||
std::vector<at::Tensor> tensor_split_symint(c10::SymIntArrayRef indices,
|
||||
int64_t dim) const;
|
||||
std::vector<at::Tensor> tensor_split(
|
||||
const at::Tensor& tensor_indices_or_sections, int64_t dim) const;
|
||||
|
||||
std::vector<at::Tensor> split(int64_t split_size, int64_t dim) const;
|
||||
std::vector<at::Tensor> split_symint(c10::SymInt split_size,
|
||||
int64_t dim) const;
|
||||
std::vector<at::Tensor> split(at::IntArrayRef split_sizes, int64_t dim) const;
|
||||
std::vector<at::Tensor> split_symint(c10::SymIntArrayRef split_sizes,
|
||||
int64_t dim) const;
|
||||
|
||||
std::vector<at::Tensor> unsafe_split(int64_t split_size,
|
||||
int64_t dim = 0) const;
|
||||
std::vector<at::Tensor> unsafe_split_symint(c10::SymInt split_size,
|
||||
int64_t dim = 0) const;
|
||||
|
||||
std::vector<at::Tensor> split_with_sizes(at::IntArrayRef split_sizes,
|
||||
int64_t dim = 0) const;
|
||||
std::vector<at::Tensor> split_with_sizes_symint(
|
||||
c10::SymIntArrayRef split_sizes, int64_t dim = 0) const;
|
||||
|
||||
std::vector<at::Tensor> unsafe_split_with_sizes(at::IntArrayRef split_sizes,
|
||||
int64_t dim = 0) const;
|
||||
std::vector<at::Tensor> unsafe_split_with_sizes_symint(
|
||||
c10::SymIntArrayRef split_sizes, int64_t dim = 0) const;
|
||||
|
||||
std::vector<at::Tensor> hsplit(int64_t sections) const;
|
||||
std::vector<at::Tensor> hsplit(at::IntArrayRef indices) const;
|
||||
|
||||
std::vector<at::Tensor> vsplit(int64_t sections) const;
|
||||
std::vector<at::Tensor> vsplit(at::IntArrayRef indices) const;
|
||||
|
||||
std::vector<at::Tensor> dsplit(int64_t sections) const;
|
||||
std::vector<at::Tensor> dsplit(at::IntArrayRef indices) const;
|
||||
|
||||
at::Tensor bitwise_right_shift(const Scalar& other) const {
|
||||
return Tensor(paddle::experimental::bitwise_right_shift(
|
||||
tensor_, paddle::experimental::full({}, other, other.dtype())));
|
||||
}
|
||||
|
||||
at::Tensor slice(int64_t dim = 0,
|
||||
::std::optional<int64_t> start = ::std::nullopt,
|
||||
::std::optional<int64_t> end = ::std::nullopt,
|
||||
int64_t step = 1) const;
|
||||
|
||||
at::Tensor index(ArrayRef<at::indexing::TensorIndex> indices) const;
|
||||
inline at::Tensor index(
|
||||
std::initializer_list<at::indexing::TensorIndex> indices) const {
|
||||
return index(ArrayRef<at::indexing::TensorIndex>(indices));
|
||||
}
|
||||
Tensor& index_put_(ArrayRef<at::indexing::TensorIndex> indices,
|
||||
Tensor const& rhs);
|
||||
Tensor& index_put_(ArrayRef<at::indexing::TensorIndex> indices,
|
||||
const Scalar& v);
|
||||
Tensor& index_put_(std::initializer_list<at::indexing::TensorIndex> indices,
|
||||
Tensor const& rhs);
|
||||
Tensor& index_put_(std::initializer_list<at::indexing::TensorIndex> indices,
|
||||
const Scalar& v);
|
||||
|
||||
at::Tensor& floor_divide_(const at::Scalar& other) const {
|
||||
paddle::experimental::floor_divide_(
|
||||
const_cast<PaddleTensor&>(tensor_),
|
||||
paddle::experimental::full({}, other, other.dtype()));
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
// Paddle Tensor has no storage_offset, so we add it here, and it is always
|
||||
// 0.
|
||||
// int64_t storage_offset() const { return storage_offset_; }
|
||||
|
||||
inline Tensor clone(
|
||||
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const {
|
||||
(void)memory_format;
|
||||
PaddleTensor cloned_tensor = paddle::experimental::assign(tensor_);
|
||||
return Tensor(cloned_tensor);
|
||||
}
|
||||
|
||||
// all: Check if all elements are true (non-zero)
|
||||
at::Tensor all() const;
|
||||
at::Tensor all(int64_t dim, bool keepdim = false) const;
|
||||
at::Tensor all(at::OptionalIntArrayRef dim, bool keepdim = false) const;
|
||||
|
||||
// allclose: Check if two tensors are close to each other
|
||||
bool allclose(const at::Tensor& other,
|
||||
double rtol = 1e-05,
|
||||
double atol = 1e-08,
|
||||
bool equal_nan = false) const;
|
||||
|
||||
at::Tensor abs() const;
|
||||
|
||||
at::Tensor& abs_() const;
|
||||
|
||||
at::Tensor absolute() const { return abs(); }
|
||||
|
||||
at::Tensor& absolute_() const { return abs_(); }
|
||||
|
||||
Tensor operator[](int64_t index) const {
|
||||
// Use as_strided to create a view (shares storage with original tensor)
|
||||
// This allows fill_ to modify the original tensor
|
||||
int64_t numel = tensor_.numel();
|
||||
if (numel == 0) {
|
||||
PD_THROW("operator[]: cannot index empty tensor");
|
||||
}
|
||||
|
||||
// Handle negative index
|
||||
if (index < 0) {
|
||||
index += tensor_.dims()[0];
|
||||
}
|
||||
|
||||
// Check bounds
|
||||
if (index < 0 || index >= tensor_.dims()[0]) {
|
||||
PD_THROW("operator[]: index ",
|
||||
index,
|
||||
" out of range for tensor of size ",
|
||||
tensor_.dims(),
|
||||
" at dimension 0");
|
||||
}
|
||||
|
||||
// For 1D tensor: create a scalar view (0-dim tensor) with proper offset
|
||||
// For multi-D tensor: create a view of the row at index
|
||||
std::vector<int64_t> new_sizes;
|
||||
std::vector<int64_t> new_strides;
|
||||
|
||||
auto dims = tensor_.dims();
|
||||
auto stride = tensor_.strides();
|
||||
|
||||
// Skip the first dimension (dim 0)
|
||||
for (int i = 1; i < dims.size(); ++i) {
|
||||
new_sizes.push_back(dims[i]);
|
||||
new_strides.push_back(stride[i]);
|
||||
}
|
||||
|
||||
// Calculate storage offset
|
||||
int64_t storage_offset = index * stride[0];
|
||||
|
||||
return as_strided(c10::IntArrayRef(new_sizes),
|
||||
c10::IntArrayRef(new_strides),
|
||||
storage_offset);
|
||||
}
|
||||
|
||||
void record_stream(at::Stream s) const;
|
||||
|
||||
Tensor var(int dim) const { return var(at::IntArrayRef{dim}, true, false); }
|
||||
|
||||
Tensor var(bool unbiased) const {
|
||||
std::vector<int64_t> empty_dims;
|
||||
double correction = unbiased ? 1.0 : 0.0;
|
||||
return var_impl(empty_dims, correction, false);
|
||||
}
|
||||
|
||||
Tensor var(at::OptionalIntArrayRef dim,
|
||||
bool unbiased,
|
||||
bool keepdim = false) const {
|
||||
// Convert unbiased to correction: unbiased=True means correction=1
|
||||
double correction = unbiased ? 1.0 : 0.0;
|
||||
std::vector<int64_t> dims_vec;
|
||||
if (dim.has_value() && dim.value().size() > 0) {
|
||||
dims_vec.assign(dim.value().begin(), dim.value().end());
|
||||
}
|
||||
return var_impl(dims_vec, correction, keepdim);
|
||||
}
|
||||
|
||||
Tensor var(at::OptionalIntArrayRef dim = ::std::nullopt,
|
||||
const ::std::optional<at::Scalar>& correction = ::std::nullopt,
|
||||
bool keepdim = false) const {
|
||||
double correction_value = 1.0;
|
||||
if (correction.has_value()) {
|
||||
const at::Scalar& scalar = correction.value();
|
||||
correction_value = scalar.to<double>();
|
||||
}
|
||||
std::vector<int64_t> dims_vec;
|
||||
if (dim.has_value() && dim.value().size() > 0) {
|
||||
dims_vec.assign(dim.value().begin(), dim.value().end());
|
||||
}
|
||||
return var_impl(dims_vec, correction_value, keepdim);
|
||||
}
|
||||
|
||||
private:
|
||||
Tensor var_impl(const std::vector<int64_t>& dims_vec,
|
||||
double correction_value,
|
||||
bool keepdim) const {
|
||||
phi::IntArray dims_int_array(dims_vec);
|
||||
|
||||
PaddleTensor mean_tensor;
|
||||
if (dims_vec.empty()) {
|
||||
mean_tensor = paddle::experimental::mean(
|
||||
tensor_, phi::IntArray(std::vector<int64_t>{}), true);
|
||||
} else {
|
||||
mean_tensor = paddle::experimental::mean(tensor_, dims_int_array, true);
|
||||
}
|
||||
|
||||
PaddleTensor diff = paddle::experimental::subtract(tensor_, mean_tensor);
|
||||
PaddleTensor diff_squared = paddle::experimental::multiply(diff, diff);
|
||||
|
||||
PaddleTensor sum_squared_diff;
|
||||
if (dims_vec.empty()) {
|
||||
sum_squared_diff =
|
||||
paddle::experimental::sum(diff_squared,
|
||||
phi::IntArray(std::vector<int64_t>{}),
|
||||
diff_squared.dtype(),
|
||||
keepdim);
|
||||
} else {
|
||||
sum_squared_diff = paddle::experimental::sum(
|
||||
diff_squared, dims_int_array, diff_squared.dtype(), keepdim);
|
||||
}
|
||||
|
||||
int64_t n = tensor_.numel();
|
||||
if (!dims_vec.empty()) {
|
||||
n = 1;
|
||||
for (int64_t d : dims_vec) {
|
||||
int64_t dim_idx = d < 0 ? d + tensor_.dims().size() : d;
|
||||
if (dim_idx >= 0 &&
|
||||
dim_idx < static_cast<int64_t>(tensor_.dims().size())) {
|
||||
n *= tensor_.dims()[dim_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
double corrected_n = static_cast<double>(n) - correction_value;
|
||||
if (corrected_n <= 0.0) {
|
||||
corrected_n = static_cast<double>(n);
|
||||
}
|
||||
|
||||
std::vector<int64_t> result_shape_vec;
|
||||
for (int64_t i = 0; i < sum_squared_diff.dims().size(); ++i) {
|
||||
result_shape_vec.push_back(sum_squared_diff.dims()[i]);
|
||||
}
|
||||
PaddleTensor correction_scalar =
|
||||
paddle::experimental::full(phi::IntArray(result_shape_vec),
|
||||
phi::Scalar(corrected_n),
|
||||
sum_squared_diff.dtype(),
|
||||
sum_squared_diff.place());
|
||||
PaddleTensor result =
|
||||
paddle::experimental::divide(sum_squared_diff, correction_scalar);
|
||||
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
public:
|
||||
// Deprecated packed_accessor for compatibility with PyTorch
|
||||
// Use packed_accessor32 or packed_accessor64 instead
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits,
|
||||
typename index_t = int64_t>
|
||||
[[deprecated(
|
||||
"packed_accessor is deprecated, use packed_accessor32 or "
|
||||
"packed_accessor64 instead")]] GenericPackedTensorAccessor<T,
|
||||
N,
|
||||
PtrTraits,
|
||||
index_t>
|
||||
packed_accessor() const& {
|
||||
return this->template generic_packed_accessor<T, N, PtrTraits, index_t>();
|
||||
}
|
||||
|
||||
template <typename T,
|
||||
size_t N,
|
||||
template <typename U> class PtrTraits = DefaultPtrTraits,
|
||||
typename index_t = int64_t>
|
||||
[[deprecated(
|
||||
"packed_accessor is deprecated, use packed_accessor32 or "
|
||||
"packed_accessor64 instead")]] GenericPackedTensorAccessor<T,
|
||||
N,
|
||||
PtrTraits,
|
||||
index_t>
|
||||
packed_accessor() && = delete;
|
||||
|
||||
template <typename T>
|
||||
using hook_return_void_t =
|
||||
std::enable_if_t<std::is_void_v<std::invoke_result_t<T&, Tensor>>,
|
||||
unsigned>;
|
||||
template <typename T>
|
||||
using hook_return_var_t =
|
||||
std::enable_if_t<std::is_same_v<std::invoke_result_t<T&, Tensor>, Tensor>,
|
||||
unsigned>;
|
||||
|
||||
// register_hook - throws exception for Paddle compatibility
|
||||
// Paddle does not support gradient hooks
|
||||
template <typename T>
|
||||
hook_return_void_t<T> register_hook(T&&) const {
|
||||
throw std::runtime_error(
|
||||
"register_hook is not supported in Paddle, this is an ATen "
|
||||
"compatibility API that is not available");
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
hook_return_var_t<T> register_hook(T&&) const {
|
||||
throw std::runtime_error(
|
||||
"register_hook is not supported in Paddle, this is an ATen "
|
||||
"compatibility API that is not available");
|
||||
}
|
||||
|
||||
// any - returns true if any element is non-zero
|
||||
Tensor any(int64_t dim, bool keepdim = false) const;
|
||||
Tensor any(at::OptionalIntArrayRef dim, bool keepdim = false) const;
|
||||
Tensor any() const;
|
||||
|
||||
// chunk - splits tensor into chunks
|
||||
std::vector<Tensor> chunk(int64_t chunks, int64_t dim = 0) const;
|
||||
|
||||
// rename - stub for Paddle (Dimname not supported)
|
||||
Tensor rename(::std::optional<at::DimnameList> names) const;
|
||||
|
||||
// new_empty - creates uninitialized tensor with same dtype/device
|
||||
Tensor new_empty(at::IntArrayRef size, at::TensorOptions options = {}) const;
|
||||
Tensor new_empty(at::IntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) const;
|
||||
|
||||
// new_full - creates tensor filled with fill_value
|
||||
Tensor new_full(at::IntArrayRef size,
|
||||
const at::Scalar& fill_value,
|
||||
at::TensorOptions options = {}) const;
|
||||
Tensor new_full(at::IntArrayRef size,
|
||||
const at::Scalar& fill_value,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) const;
|
||||
|
||||
// new_zeros - creates zero tensor
|
||||
Tensor new_zeros(at::IntArrayRef size, at::TensorOptions options = {}) const;
|
||||
Tensor new_zeros(at::IntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) const;
|
||||
|
||||
// new_ones - creates tensor filled with ones
|
||||
Tensor new_ones(at::IntArrayRef size, at::TensorOptions options = {}) const;
|
||||
Tensor new_ones(at::IntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) const;
|
||||
|
||||
// resize_ - in-place resize
|
||||
const Tensor& resize_(
|
||||
at::IntArrayRef size,
|
||||
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) const;
|
||||
|
||||
// expand - expands tensor to new size
|
||||
Tensor expand(at::IntArrayRef size, bool implicit = false) const;
|
||||
|
||||
// expand_as - expands to same size as another tensor
|
||||
Tensor expand_as(const Tensor& other) const;
|
||||
|
||||
PaddleTensor _PD_GetInner() const { return tensor_; }
|
||||
PaddleTensor& _PD_GetInner() { return tensor_; }
|
||||
}; // NOLINT(readability/braces)
|
||||
} // namespace at
|
||||
@@ -0,0 +1,66 @@
|
||||
// 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#include <ATen/core/TensorBase.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <string_view>
|
||||
|
||||
namespace at {
|
||||
|
||||
void check_type(const TensorBase& tensor,
|
||||
ScalarType type,
|
||||
std::string_view type_name) {
|
||||
PD_CHECK(tensor.scalar_type() == type,
|
||||
"expected scalar type ",
|
||||
type_name,
|
||||
" but found ",
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(tensor.scalar_type()));
|
||||
}
|
||||
|
||||
#define DEFINE_CAST(T, name) \
|
||||
template <> \
|
||||
PADDLE_API const T* TensorBase::const_data_ptr() const { \
|
||||
check_type(*this, ScalarType::name, #name); \
|
||||
return const_cast<T*>(tensor_.data<T>()); \
|
||||
} \
|
||||
\
|
||||
template <> \
|
||||
PADDLE_API const T* TensorBase::const_data_ptr<const T>() const { \
|
||||
check_type(*this, ScalarType::name, #name); \
|
||||
return const_cast<T*>(tensor_.data<std::remove_const_t<T>>()); \
|
||||
} \
|
||||
\
|
||||
template <> \
|
||||
PADDLE_API T* TensorBase::mutable_data_ptr() const { \
|
||||
check_type(*this, ScalarType::name, #name); \
|
||||
return const_cast<PaddleTensor&>(tensor_).data<T>(); \
|
||||
} \
|
||||
\
|
||||
template <> \
|
||||
PADDLE_API T* TensorBase::data_ptr() const { \
|
||||
return const_cast<T*>(tensor_.data<T>()); \
|
||||
}
|
||||
|
||||
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_CAST) // missing half and float16
|
||||
// AT_FORALL_QINT_TYPES(DEFINE_CAST) // missing qint
|
||||
DEFINE_CAST(uint16_t, UInt16)
|
||||
DEFINE_CAST(uint32_t, UInt32)
|
||||
DEFINE_CAST(uint64_t, UInt64)
|
||||
#undef DEFINE_CAST
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,156 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ostream>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <unordered_set>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace c10 {
|
||||
/**
|
||||
* class AliasInfo
|
||||
*
|
||||
* Data structure to hold aliasing information for an `Argument`. They can be
|
||||
* nested to represent aliasing information on contained types.
|
||||
*
|
||||
* There is a `beforeSet` which describes the aliasing information before the
|
||||
* operator executes, and an `afterSet` that describes aliasing info
|
||||
* after execution.
|
||||
*/
|
||||
class AliasInfo {
|
||||
public:
|
||||
AliasInfo() = default;
|
||||
AliasInfo(bool is_write,
|
||||
const std::set<std::string>& before_qual_strings,
|
||||
const std::set<std::string>& after_qual_strings)
|
||||
: isWrite_(is_write) {
|
||||
for (const auto& s : before_qual_strings) {
|
||||
beforeSets_.insert(s);
|
||||
}
|
||||
for (const auto& s : after_qual_strings) {
|
||||
afterSets_.insert(s);
|
||||
}
|
||||
}
|
||||
|
||||
bool isWrite() const { return isWrite_; }
|
||||
|
||||
const std::unordered_set<std::string>& beforeSets() const {
|
||||
return beforeSets_;
|
||||
}
|
||||
|
||||
const std::unordered_set<std::string>& afterSets() const {
|
||||
return afterSets_;
|
||||
}
|
||||
|
||||
// the alias info for the contained types of the type
|
||||
// e.g. if this is an annotation on List[T], `sets` refers to
|
||||
// the alias sets that the list may be in
|
||||
// while containedTypes()[0] refers to the sets that members of the list
|
||||
// may be in
|
||||
void addContainedType(AliasInfo aliasInfo) {
|
||||
containedTypes_.push_back(std::move(aliasInfo));
|
||||
}
|
||||
const std::vector<AliasInfo>& containedTypes() const {
|
||||
return containedTypes_;
|
||||
}
|
||||
|
||||
private:
|
||||
std::unordered_set<std::string> beforeSets_;
|
||||
std::unordered_set<std::string> afterSets_;
|
||||
std::vector<AliasInfo> containedTypes_;
|
||||
bool isWrite_ = false;
|
||||
};
|
||||
|
||||
inline bool operator==(const AliasInfo& lhs, const AliasInfo& rhs) {
|
||||
return lhs.isWrite() == rhs.isWrite() &&
|
||||
lhs.beforeSets() == rhs.beforeSets() &&
|
||||
lhs.afterSets() == rhs.afterSets() &&
|
||||
lhs.containedTypes() == rhs.containedTypes();
|
||||
}
|
||||
|
||||
// this does match the way things are represented in the schema
|
||||
inline std::ostream& operator<<(std::ostream& out, const AliasInfo& aliasInfo) {
|
||||
out << '(';
|
||||
bool first = true;
|
||||
for (const auto& set : aliasInfo.beforeSets()) {
|
||||
if (first) {
|
||||
first = false;
|
||||
} else {
|
||||
out << '|';
|
||||
}
|
||||
out << set;
|
||||
}
|
||||
if (aliasInfo.isWrite()) {
|
||||
out << '!';
|
||||
}
|
||||
if (aliasInfo.beforeSets() != aliasInfo.afterSets()) {
|
||||
out << " -> ";
|
||||
first = true;
|
||||
for (const auto& set : aliasInfo.afterSets()) {
|
||||
if (first) {
|
||||
first = false;
|
||||
} else {
|
||||
out << '|';
|
||||
}
|
||||
out << set;
|
||||
}
|
||||
}
|
||||
out << ')';
|
||||
return out;
|
||||
}
|
||||
} // namespace c10
|
||||
|
||||
inline std::size_t hash_combine(std::size_t lhs, std::size_t rhs) {
|
||||
lhs ^= rhs + 0x9e3779b9 + (lhs << 6) + (lhs >> 2);
|
||||
return lhs;
|
||||
}
|
||||
|
||||
namespace std {
|
||||
template <>
|
||||
struct hash<c10::AliasInfo> {
|
||||
size_t operator()(const c10::AliasInfo& aliasInfo) const {
|
||||
auto hash = std::hash<bool>()(aliasInfo.isWrite());
|
||||
|
||||
// NOTE: for unordered_set hashes, we couldn't use hash_combine
|
||||
// because hash_combine is order dependent. Instead, we choose to
|
||||
// use XOR as the combining function as XOR is commutative.
|
||||
size_t before_set_hash_seed = 0;
|
||||
for (auto& e : aliasInfo.beforeSets()) {
|
||||
auto symbol_hash = std::hash<std::string>()(e);
|
||||
before_set_hash_seed = before_set_hash_seed ^ symbol_hash;
|
||||
}
|
||||
size_t after_set_hash_seed = 0;
|
||||
for (auto& e : aliasInfo.afterSets()) {
|
||||
auto symbol_hash = std::hash<std::string>()(e);
|
||||
after_set_hash_seed = after_set_hash_seed ^ symbol_hash;
|
||||
}
|
||||
|
||||
hash = hash_combine(hash, before_set_hash_seed);
|
||||
hash = hash_combine(hash, after_set_hash_seed);
|
||||
for (auto& e : aliasInfo.containedTypes()) {
|
||||
auto contained_type_hash = std::hash<c10::AliasInfo>()(e);
|
||||
hash = hash_combine(hash, contained_type_hash);
|
||||
}
|
||||
return hash;
|
||||
}
|
||||
};
|
||||
} // namespace std
|
||||
@@ -0,0 +1,201 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#include "ATen/core/function_schema.h"
|
||||
|
||||
namespace c10 {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr char kWildcardAliasSet[] = "*";
|
||||
|
||||
const char* schemaArgTypeName(SchemaArgType type) {
|
||||
if (type == SchemaArgType::input) {
|
||||
return "input";
|
||||
}
|
||||
if (type == SchemaArgType::output) {
|
||||
return "output";
|
||||
}
|
||||
return "unknown";
|
||||
}
|
||||
|
||||
bool aliasSetsMayOverlap(const std::unordered_set<std::string>& lhs,
|
||||
const std::unordered_set<std::string>& rhs) {
|
||||
if (lhs.empty() || rhs.empty()) {
|
||||
return false;
|
||||
}
|
||||
if (lhs.count(kWildcardAliasSet) > 0 || rhs.count(kWildcardAliasSet) > 0) {
|
||||
return true;
|
||||
}
|
||||
for (const auto& set : lhs) {
|
||||
if (rhs.count(set) > 0) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
const Argument& getSchemaArgumentOrThrow(const FunctionSchema& schema,
|
||||
const SchemaArgument& argument) {
|
||||
const auto& args = schema.getCorrectList(argument);
|
||||
TORCH_CHECK(argument.index < args.size(),
|
||||
"Schema ",
|
||||
schemaArgTypeName(argument.type),
|
||||
" index ",
|
||||
argument.index,
|
||||
" is out of bounds for size ",
|
||||
args.size());
|
||||
return args.at(argument.index);
|
||||
}
|
||||
|
||||
bool aliasInfoMayContainAlias(const AliasInfo& lhs,
|
||||
const AliasInfo& rhs,
|
||||
bool bidirectional) {
|
||||
if (aliasSetsMayOverlap(lhs.afterSets(), rhs.afterSets())) {
|
||||
return true;
|
||||
}
|
||||
|
||||
for (const auto& child : lhs.containedTypes()) {
|
||||
if (aliasInfoMayContainAlias(child, rhs, /*bidirectional=*/true)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
if (!bidirectional) {
|
||||
return false;
|
||||
}
|
||||
for (const auto& child : rhs.containedTypes()) {
|
||||
if (aliasInfoMayContainAlias(lhs, child, /*bidirectional=*/true)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
std::ostream& operator<<(std::ostream& out, const Argument& arg) {
|
||||
out << arg.type()->str() << " " << arg.name();
|
||||
if (arg.default_value()) {
|
||||
out << " = " << arg.default_value();
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
std::ostream& operator<<(std::ostream& out, const FunctionSchema& schema) {
|
||||
out << "(";
|
||||
bool first = true;
|
||||
for (const auto& arg : schema.arguments()) {
|
||||
if (!first) {
|
||||
out << ", ";
|
||||
}
|
||||
out << arg;
|
||||
first = false;
|
||||
}
|
||||
if (schema.is_vararg()) {
|
||||
if (!first) {
|
||||
out << ", ";
|
||||
}
|
||||
out << "...";
|
||||
}
|
||||
out << ")";
|
||||
|
||||
out << " -> ";
|
||||
|
||||
if (schema.returns().size() == 1) {
|
||||
out << schema.returns()[0];
|
||||
} else {
|
||||
out << "(";
|
||||
first = true;
|
||||
for (const auto& ret : schema.returns()) {
|
||||
if (!first) {
|
||||
out << ", ";
|
||||
}
|
||||
out << ret;
|
||||
first = false;
|
||||
}
|
||||
out << ")";
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
std::optional<int> FunctionSchema::argumentIndexWithName(
|
||||
const std::string& name) const {
|
||||
for (size_t i = 0; i < arguments_.size(); ++i) {
|
||||
if (arguments_[i].name() == name) {
|
||||
return static_cast<int>(i);
|
||||
}
|
||||
}
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
const std::vector<Argument>& FunctionSchema::getCorrectList(
|
||||
const SchemaArgument& argument) const {
|
||||
if (argument.type == SchemaArgType::input) {
|
||||
return arguments();
|
||||
}
|
||||
if (argument.type == SchemaArgType::output) {
|
||||
return returns();
|
||||
}
|
||||
TORCH_INTERNAL_ASSERT(false, "Could not match argument type");
|
||||
}
|
||||
|
||||
bool FunctionSchema::is_aliasing(const SchemaArgument& argument) const {
|
||||
const auto& arg = getSchemaArgumentOrThrow(*this, argument);
|
||||
return arg.alias_info() != nullptr;
|
||||
}
|
||||
|
||||
bool FunctionSchema::is_mutable(const SchemaArgument& argument) const {
|
||||
const auto& arg = getSchemaArgumentOrThrow(*this, argument);
|
||||
return arg.alias_info() != nullptr && arg.alias_info()->isWrite();
|
||||
}
|
||||
|
||||
bool FunctionSchema::is_mutable(const std::string& name) const {
|
||||
const auto index = argumentIndexWithName(name);
|
||||
TORCH_CHECK(
|
||||
index.has_value(), "Tried to test mutability of nonexistent name ", name);
|
||||
return is_mutable({SchemaArgType::input, static_cast<size_t>(*index)});
|
||||
}
|
||||
|
||||
bool FunctionSchema::may_alias(const SchemaArgument& lhs,
|
||||
const SchemaArgument& rhs) const {
|
||||
const auto& lhs_arg = getSchemaArgumentOrThrow(*this, lhs);
|
||||
const auto& rhs_arg = getSchemaArgumentOrThrow(*this, rhs);
|
||||
const auto* lhs_alias = lhs_arg.alias_info();
|
||||
const auto* rhs_alias = rhs_arg.alias_info();
|
||||
if (lhs_alias == nullptr || rhs_alias == nullptr) {
|
||||
return false;
|
||||
}
|
||||
return aliasSetsMayOverlap(lhs_alias->afterSets(), rhs_alias->afterSets());
|
||||
}
|
||||
|
||||
bool FunctionSchema::may_contain_alias(const SchemaArgument& lhs,
|
||||
const SchemaArgument& rhs,
|
||||
bool bidirectional) const {
|
||||
const auto& lhs_arg = getSchemaArgumentOrThrow(*this, lhs);
|
||||
const auto& rhs_arg = getSchemaArgumentOrThrow(*this, rhs);
|
||||
const auto* lhs_alias = lhs_arg.alias_info();
|
||||
const auto* rhs_alias = rhs_arg.alias_info();
|
||||
if (lhs_alias == nullptr || rhs_alias == nullptr) {
|
||||
return false;
|
||||
}
|
||||
return aliasInfoMayContainAlias(*lhs_alias, *rhs_alias, bidirectional);
|
||||
}
|
||||
|
||||
} // namespace c10
|
||||
@@ -0,0 +1,261 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
#include <ATen/core/alias_info.h>
|
||||
#include <ATen/core/ivalue.h>
|
||||
#include <ATen/core/jit_type.h>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include "paddle/common/macros.h" // For macro PADDLE_API
|
||||
|
||||
namespace c10 {
|
||||
|
||||
struct Argument;
|
||||
struct FunctionSchema;
|
||||
enum class SchemaArgType;
|
||||
struct SchemaArgument;
|
||||
|
||||
enum class SchemaArgType {
|
||||
input,
|
||||
output,
|
||||
};
|
||||
|
||||
struct SchemaArgument {
|
||||
SchemaArgType type;
|
||||
size_t index;
|
||||
};
|
||||
|
||||
struct PADDLE_API Argument {
|
||||
Argument(std::string name = "",
|
||||
const TypePtr& type = nullptr,
|
||||
std::optional<int32_t> N = std::nullopt,
|
||||
std::optional<torch::IValue> default_value = std::nullopt,
|
||||
bool kwarg_only = false,
|
||||
std::optional<c10::AliasInfo> alias_info = std::nullopt)
|
||||
: Argument(std::move(name),
|
||||
type,
|
||||
type,
|
||||
N,
|
||||
std::move(default_value),
|
||||
kwarg_only,
|
||||
std::move(alias_info)) {}
|
||||
|
||||
Argument(std::string name,
|
||||
TypePtr fake_type,
|
||||
TypePtr real_type,
|
||||
std::optional<int32_t> N = std::nullopt,
|
||||
std::optional<torch::IValue> default_value = std::nullopt,
|
||||
bool kwarg_only = false,
|
||||
std::optional<c10::AliasInfo> alias_info = std::nullopt)
|
||||
: name_(std::move(name)),
|
||||
type_(fake_type ? std::move(fake_type) : TensorType::get()),
|
||||
real_type_(real_type ? std::move(real_type) : type_),
|
||||
N_(N),
|
||||
default_value_(std::move(default_value)),
|
||||
alias_info_(alias_info ? std::make_unique<c10::AliasInfo>(
|
||||
std::move(*alias_info))
|
||||
: nullptr),
|
||||
kwarg_only_(kwarg_only) {
|
||||
// this is an softly-enforced invariant for out arguments.
|
||||
bool is_alias = alias_info_ != nullptr && alias_info_->isWrite();
|
||||
is_out_ = kwarg_only_ && is_alias;
|
||||
}
|
||||
|
||||
Argument(Argument&& rhs) noexcept = default;
|
||||
|
||||
Argument(const Argument& rhs)
|
||||
: name_(rhs.name_),
|
||||
type_(rhs.type_),
|
||||
real_type_(rhs.real_type_),
|
||||
N_(rhs.N_),
|
||||
default_value_(rhs.default_value_),
|
||||
alias_info_(rhs.alias_info_
|
||||
? std::make_unique<c10::AliasInfo>(*rhs.alias_info_)
|
||||
: nullptr),
|
||||
kwarg_only_(rhs.kwarg_only_),
|
||||
is_out_(rhs.is_out_) {}
|
||||
|
||||
Argument& operator=(Argument&& rhs) = default;
|
||||
|
||||
Argument& operator=(const Argument& rhs) {
|
||||
if (this != &rhs) {
|
||||
name_ = rhs.name_;
|
||||
type_ = rhs.type_;
|
||||
real_type_ = rhs.real_type_;
|
||||
N_ = rhs.N_;
|
||||
default_value_ = rhs.default_value_;
|
||||
alias_info_ = rhs.alias_info_
|
||||
? std::make_unique<c10::AliasInfo>(*rhs.alias_info_)
|
||||
: nullptr;
|
||||
kwarg_only_ = rhs.kwarg_only_;
|
||||
is_out_ = rhs.is_out_;
|
||||
}
|
||||
return *this;
|
||||
}
|
||||
~Argument() = default;
|
||||
|
||||
const std::string& name() const { return name_; }
|
||||
const TypePtr& type() const { return type_; }
|
||||
// if type() is non-null, this is guaranteed to be non-null (if no real
|
||||
// type was provided, this takes on type()'s value)
|
||||
const TypePtr& real_type() const { return real_type_; }
|
||||
const std::optional<int32_t>& N() const { return N_; }
|
||||
const std::optional<torch::IValue>& default_value() const {
|
||||
return default_value_;
|
||||
}
|
||||
bool kwarg_only() const { return kwarg_only_; }
|
||||
|
||||
bool is_out() const { return is_out_; }
|
||||
|
||||
[[nodiscard]] const c10::AliasInfo* alias_info() const {
|
||||
return alias_info_.get();
|
||||
}
|
||||
|
||||
bool is_inferred_type() const {
|
||||
bool is_inferred_type = false;
|
||||
TORCH_INTERNAL_ASSERT(type_);
|
||||
if (auto pt = type_->cast<TensorType>()) {
|
||||
if (pt->isInferredType()) {
|
||||
is_inferred_type = true;
|
||||
}
|
||||
}
|
||||
return is_inferred_type;
|
||||
}
|
||||
|
||||
std::string formatTypeMismatchMsg(const std::string& actual_type) const {
|
||||
std::string inferred_type_hint;
|
||||
if (is_inferred_type()) {
|
||||
inferred_type_hint = "Inferred '";
|
||||
inferred_type_hint += name();
|
||||
inferred_type_hint += "' to be of type 'Tensor' ";
|
||||
inferred_type_hint +=
|
||||
"because it was not annotated with an explicit type.\n";
|
||||
}
|
||||
std::string message;
|
||||
message += "Expected a value of type '";
|
||||
message += type()->repr_str();
|
||||
message += "' for argument '";
|
||||
message += name();
|
||||
message += "' but instead found type '";
|
||||
message += actual_type;
|
||||
message += "'.\n";
|
||||
message += inferred_type_hint;
|
||||
return message;
|
||||
}
|
||||
|
||||
Argument cloneWithType(const TypePtr& new_type) const {
|
||||
return Argument(name_,
|
||||
new_type,
|
||||
N_,
|
||||
default_value_,
|
||||
kwarg_only_,
|
||||
alias_info_ ? std::optional<c10::AliasInfo>(*alias_info_)
|
||||
: std::nullopt);
|
||||
}
|
||||
|
||||
friend PADDLE_API std::ostream& operator<<(std::ostream& out,
|
||||
const Argument& arg);
|
||||
|
||||
private:
|
||||
std::string name_;
|
||||
TypePtr type_;
|
||||
TypePtr real_type_; // this is ScalarType, not int, e.g.
|
||||
// for list types, an optional statically known length for the list
|
||||
// e.g. for int[3]: type = ListType::ofInts(), N = 3
|
||||
// If present, this will allow scalars to be broadcast to this length to
|
||||
// become a list.
|
||||
std::optional<int32_t> N_;
|
||||
|
||||
std::optional<torch::IValue> default_value_;
|
||||
// c10::AliasInfo is huge, so let's only allocate memory for it if
|
||||
// necessary (which it isn't during schema parsing on startup, to
|
||||
// give a pertinent example).
|
||||
std::unique_ptr<c10::AliasInfo> alias_info_;
|
||||
// is this only specifiable as a keyword argument?
|
||||
bool kwarg_only_;
|
||||
// marks if the argument is out variant of the schema
|
||||
bool is_out_;
|
||||
};
|
||||
|
||||
struct PADDLE_API FunctionSchema {
|
||||
FunctionSchema(std::vector<Argument> arguments,
|
||||
std::vector<Argument> returns,
|
||||
bool is_vararg = false,
|
||||
bool is_varret = false)
|
||||
: arguments_(std::move(arguments)),
|
||||
returns_(std::move(returns)),
|
||||
is_vararg_(is_vararg),
|
||||
is_varret_(is_varret) {
|
||||
checkSchema();
|
||||
}
|
||||
|
||||
const std::vector<Argument>& arguments() const { return arguments_; }
|
||||
|
||||
void checkSchema() const {
|
||||
bool seen_default_arg = false;
|
||||
for (const auto& arg : arguments()) {
|
||||
if (arg.default_value()) {
|
||||
seen_default_arg = true;
|
||||
} else {
|
||||
TORCH_INTERNAL_ASSERT(!seen_default_arg || arg.kwarg_only(),
|
||||
"Non-default positional argument follows default "
|
||||
"argument. Parameter ",
|
||||
arg.name(),
|
||||
" in ",
|
||||
*this);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const std::vector<Argument>& returns() const { return returns_; }
|
||||
|
||||
bool is_vararg() const { return is_vararg_; }
|
||||
|
||||
bool is_varret() const { return is_varret_; }
|
||||
|
||||
std::optional<int> argumentIndexWithName(const std::string& name) const;
|
||||
const std::vector<Argument>& getCorrectList(
|
||||
const SchemaArgument& argument) const;
|
||||
bool is_aliasing(const SchemaArgument& argument) const;
|
||||
bool is_mutable(const SchemaArgument& argument) const;
|
||||
bool is_mutable(const std::string& name) const;
|
||||
bool may_alias(const SchemaArgument& lhs, const SchemaArgument& rhs) const;
|
||||
bool may_contain_alias(const SchemaArgument& lhs,
|
||||
const SchemaArgument& rhs,
|
||||
bool bidirectional = true) const;
|
||||
|
||||
friend PADDLE_API std::ostream& operator<<(std::ostream& out,
|
||||
const FunctionSchema& schema);
|
||||
|
||||
private:
|
||||
std::vector<Argument> arguments_;
|
||||
std::vector<Argument> returns_;
|
||||
// if true then this schema takes an arbitrary number of additional arguments
|
||||
// after the argument specified in arguments
|
||||
// currently this is used primarily to represent 'primitive' operators whose
|
||||
// arguments are not checked by schema
|
||||
bool is_vararg_;
|
||||
bool is_varret_;
|
||||
};
|
||||
|
||||
PADDLE_API std::ostream& operator<<(std::ostream& out, const Argument& arg);
|
||||
PADDLE_API std::ostream& operator<<(std::ostream& out,
|
||||
const FunctionSchema& schema);
|
||||
|
||||
} // namespace c10
|
||||
@@ -0,0 +1,72 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <c10/util/ArrayRef.h>
|
||||
#include <vector>
|
||||
|
||||
namespace c10 {
|
||||
|
||||
// The passed in function must take T by value (T), or by
|
||||
// const reference (const T&); taking T by non-const reference
|
||||
// will result in an error like:
|
||||
//
|
||||
// error: no type named 'type' in 'class std::invoke_result<foobar::__lambda,
|
||||
// T>'
|
||||
//
|
||||
// No explicit template parameters are required.
|
||||
|
||||
// Overload for explicit function and ArrayRef
|
||||
template <class F, class T>
|
||||
inline auto fmap(const T& inputs, const F& fn)
|
||||
-> std::vector<decltype(fn(*inputs.begin()))> {
|
||||
std::vector<decltype(fn(*inputs.begin()))> r;
|
||||
r.reserve(inputs.size());
|
||||
for (const auto& input : inputs) r.push_back(fn(input));
|
||||
return r;
|
||||
}
|
||||
|
||||
// C++ forbids taking an address of a constructor, so here's a workaround...
|
||||
// Overload for constructor (R) application
|
||||
template <typename R, typename T>
|
||||
inline std::vector<R> fmap(const T& inputs) {
|
||||
std::vector<R> r;
|
||||
r.reserve(inputs.size());
|
||||
for (auto& input : inputs) r.push_back(R(input));
|
||||
return r;
|
||||
}
|
||||
|
||||
template <typename F, typename T>
|
||||
inline std::vector<T> filter(at::ArrayRef<T> inputs, const F& fn) {
|
||||
std::vector<T> r;
|
||||
r.reserve(inputs.size());
|
||||
for (auto& input : inputs) {
|
||||
if (fn(input)) {
|
||||
r.push_back(input);
|
||||
}
|
||||
}
|
||||
return r;
|
||||
}
|
||||
|
||||
template <typename F, typename T>
|
||||
inline std::vector<T> filter(const std::vector<T>& inputs, const F& fn) {
|
||||
return filter<F, T>(static_cast<at::ArrayRef<T>>(inputs), fn);
|
||||
}
|
||||
|
||||
} // namespace c10
|
||||
@@ -0,0 +1,686 @@
|
||||
// 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <cstddef>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <tuple>
|
||||
#include <type_traits>
|
||||
#include <utility>
|
||||
#include <variant>
|
||||
#include <vector>
|
||||
|
||||
namespace torch {
|
||||
|
||||
class CustomClassHolder {
|
||||
public:
|
||||
virtual ~CustomClassHolder() = default;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class intrusive_ptr {
|
||||
public:
|
||||
using element_type = T;
|
||||
using pointer = T*;
|
||||
|
||||
intrusive_ptr() : ptr_(nullptr) {}
|
||||
intrusive_ptr(T* ptr) : ptr_(std::shared_ptr<T>(ptr)) {} // NOLINT
|
||||
intrusive_ptr(std::shared_ptr<T> ptr) : ptr_(ptr) {} // NOLINT
|
||||
|
||||
template <typename... Args>
|
||||
static intrusive_ptr<T> make(Args&&... args) {
|
||||
return intrusive_ptr<T>(std::make_shared<T>(std::forward<Args>(args)...));
|
||||
}
|
||||
|
||||
T* get() const { return ptr_.get(); }
|
||||
T& operator*() const { return *ptr_; }
|
||||
T* operator->() const { return ptr_.get(); }
|
||||
|
||||
// For IValue
|
||||
std::shared_ptr<T> get_shared() const { return ptr_; }
|
||||
|
||||
explicit operator bool() const { return ptr_ != nullptr; }
|
||||
|
||||
private:
|
||||
std::shared_ptr<T> ptr_;
|
||||
};
|
||||
|
||||
template <typename T, typename... Args>
|
||||
intrusive_ptr<T> make_intrusive(Args&&... args) {
|
||||
return intrusive_ptr<T>::make(std::forward<Args>(args)...);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct _fake_type {};
|
||||
|
||||
enum class TypeTag {
|
||||
None = 0,
|
||||
Bool,
|
||||
Int,
|
||||
Double,
|
||||
String,
|
||||
Tensor,
|
||||
GenericList,
|
||||
CustomClass,
|
||||
Tuple
|
||||
};
|
||||
|
||||
class IValue; // Forward declaration
|
||||
|
||||
// Forward declaration of generic_to template function
|
||||
template <typename T>
|
||||
T generic_to(const IValue& ivalue, _fake_type<T>);
|
||||
|
||||
using GenericList = std::vector<IValue>;
|
||||
|
||||
// Separate tuple wrapper to avoid ambiguity with GenericList
|
||||
struct GenericTuple {
|
||||
std::vector<IValue> elements;
|
||||
|
||||
GenericTuple();
|
||||
GenericTuple(std::vector<IValue> elems); // NOLINT
|
||||
|
||||
size_t size() const;
|
||||
IValue& operator[](size_t idx);
|
||||
const IValue& operator[](size_t idx) const;
|
||||
};
|
||||
|
||||
class IValue {
|
||||
private:
|
||||
struct CustomClassWrapper {
|
||||
std::shared_ptr<CustomClassHolder> ptr;
|
||||
std::string class_name;
|
||||
|
||||
CustomClassWrapper(std::shared_ptr<CustomClassHolder> p,
|
||||
const std::string& name)
|
||||
: ptr(std::move(p)), class_name(name) {}
|
||||
};
|
||||
|
||||
public:
|
||||
IValue() : tag_(TypeTag::None), value_(std::monostate{}) {}
|
||||
|
||||
IValue(bool val) : tag_(TypeTag::Bool), value_(val) {} // NOLINT
|
||||
IValue(int val) // NOLINT
|
||||
: tag_(TypeTag::Int), value_(static_cast<int64_t>(val)) {}
|
||||
IValue(int64_t val) : tag_(TypeTag::Int), value_(val) {} // NOLINT
|
||||
IValue(double val) : tag_(TypeTag::Double), value_(val) {} // NOLINT
|
||||
IValue(const std::string& val) // NOLINT
|
||||
: tag_(TypeTag::String), value_(val) {}
|
||||
IValue(std::string&& val) // NOLINT
|
||||
: tag_(TypeTag::String), value_(std::move(val)) {}
|
||||
IValue(const char* val) // NOLINT
|
||||
: tag_(TypeTag::String), value_(std::string(val)) {}
|
||||
IValue(at::Tensor val) : tag_(TypeTag::Tensor), value_(val) {} // NOLINT
|
||||
IValue(ScalarType val) // NOLINT
|
||||
: tag_(TypeTag::Int),
|
||||
value_(static_cast<int64_t>(
|
||||
static_cast<std::underlying_type_t<ScalarType>>(val))) {}
|
||||
template <typename T>
|
||||
IValue(intrusive_ptr<T> ptr) // NOLINT
|
||||
: tag_(TypeTag::CustomClass),
|
||||
value_(CustomClassWrapper{ptr.get_shared(), typeid(T).name()}) {}
|
||||
|
||||
template <typename T,
|
||||
typename = std::enable_if_t<std::is_constructible_v<IValue, T>>>
|
||||
IValue(const std::vector<T>& vec) // NOLINT
|
||||
: tag_(TypeTag::GenericList) {
|
||||
GenericList generic_list;
|
||||
generic_list.reserve(vec.size());
|
||||
for (const auto& item : vec) {
|
||||
generic_list.emplace_back(IValue(item));
|
||||
}
|
||||
value_ = std::move(generic_list);
|
||||
}
|
||||
|
||||
template <typename T,
|
||||
typename = std::enable_if_t<std::is_constructible_v<IValue, T>>>
|
||||
IValue(std::vector<T>&& vec) // NOLINT
|
||||
: tag_(TypeTag::GenericList) {
|
||||
GenericList generic_list;
|
||||
generic_list.reserve(vec.size());
|
||||
for (auto&& item : vec) {
|
||||
generic_list.emplace_back(IValue(std::move(item)));
|
||||
}
|
||||
value_ = std::move(generic_list);
|
||||
}
|
||||
|
||||
template <typename T,
|
||||
typename = std::enable_if_t<std::is_constructible_v<IValue, T>>>
|
||||
IValue(ArrayRef<T> arr) : IValue(arr.vec()) {} // NOLINT
|
||||
|
||||
template <typename T>
|
||||
IValue(const std::optional<T>& opt) { // NOLINT
|
||||
if (opt.has_value()) {
|
||||
*this = IValue(*opt);
|
||||
} else {
|
||||
tag_ = TypeTag::None;
|
||||
value_ = std::monostate{};
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
IValue(std::optional<T>&& opt) { // NOLINT
|
||||
if (opt.has_value()) {
|
||||
*this = IValue(std::move(*opt));
|
||||
} else {
|
||||
tag_ = TypeTag::None;
|
||||
value_ = std::monostate{};
|
||||
}
|
||||
}
|
||||
|
||||
// Variadic template constructor for tuple of any number of tensors or
|
||||
// IValue-convertible types
|
||||
template <typename... Args>
|
||||
IValue(const std::tuple<Args...>& tuple_val) // NOLINT
|
||||
: tag_(TypeTag::Tuple) {
|
||||
static_assert(sizeof...(Args) > 0, "Tuple must have at least one element");
|
||||
std::vector<IValue> elements;
|
||||
elements.reserve(sizeof...(Args));
|
||||
tuple_to_ivalue_vector(
|
||||
tuple_val, elements, std::index_sequence_for<Args...>{});
|
||||
value_ = GenericTuple(std::move(elements));
|
||||
}
|
||||
|
||||
// Helper function to convert tuple elements to IValue vector using index
|
||||
// sequence
|
||||
template <typename Tuple, std::size_t... I>
|
||||
void tuple_to_ivalue_vector(const Tuple& tuple_val,
|
||||
std::vector<IValue>& elements, // NOLINT
|
||||
std::index_sequence<I...>) {
|
||||
(elements.emplace_back(std::get<I>(tuple_val)), ...);
|
||||
}
|
||||
|
||||
IValue(const IValue& other) = default;
|
||||
IValue(IValue&& other) = default;
|
||||
IValue& operator=(const IValue& other) = default;
|
||||
IValue& operator=(IValue&& other) = default;
|
||||
|
||||
bool is_none() const { return tag_ == TypeTag::None; }
|
||||
bool is_bool() const { return tag_ == TypeTag::Bool; }
|
||||
bool is_int() const { return tag_ == TypeTag::Int; }
|
||||
bool is_double() const { return tag_ == TypeTag::Double; }
|
||||
bool is_string() const { return tag_ == TypeTag::String; }
|
||||
bool is_list() const { return tag_ == TypeTag::GenericList; }
|
||||
bool is_tensor() const { return tag_ == TypeTag::Tensor; }
|
||||
bool is_custom_class() const { return tag_ == TypeTag::CustomClass; }
|
||||
bool is_tuple() const { return tag_ == TypeTag::Tuple; }
|
||||
|
||||
bool isNone() const { return is_none(); }
|
||||
bool isBool() const { return is_bool(); }
|
||||
bool isInt() const { return is_int(); }
|
||||
bool isDouble() const { return is_double(); }
|
||||
bool isString() const { return is_string(); }
|
||||
bool isList() const { return is_list(); }
|
||||
bool isTensor() const { return is_tensor(); }
|
||||
bool isCustomClass() const { return is_custom_class(); }
|
||||
bool isTuple() const { return is_tuple(); }
|
||||
|
||||
bool to_bool() const {
|
||||
if (!is_bool()) throw std::runtime_error("Not a bool");
|
||||
return std::get<bool>(value_);
|
||||
}
|
||||
|
||||
int64_t to_int() const {
|
||||
if (!is_int()) throw std::runtime_error("Not an int");
|
||||
return std::get<int64_t>(value_);
|
||||
}
|
||||
|
||||
double to_double() const {
|
||||
if (!is_double()) throw std::runtime_error("Not a double");
|
||||
return std::get<double>(value_);
|
||||
}
|
||||
|
||||
const std::string& to_string() const {
|
||||
if (!is_string()) throw std::runtime_error("Not a string");
|
||||
return std::get<std::string>(value_);
|
||||
}
|
||||
|
||||
const std::string_view to_string_view() const {
|
||||
if (!is_string()) throw std::runtime_error("Not a string");
|
||||
const auto& str = std::get<std::string>(value_);
|
||||
return std::string_view(str.data(), str.size());
|
||||
}
|
||||
|
||||
const GenericList& to_list() const {
|
||||
if (!is_list()) throw std::runtime_error("Not a list");
|
||||
return std::get<GenericList>(value_);
|
||||
}
|
||||
|
||||
GenericList& to_list() {
|
||||
if (!is_list()) throw std::runtime_error("Not a list");
|
||||
return std::get<GenericList>(value_);
|
||||
}
|
||||
|
||||
at::Tensor to_tensor() const {
|
||||
if (!is_tensor()) throw std::runtime_error("Not a tensor");
|
||||
return std::get<at::Tensor>(value_);
|
||||
}
|
||||
|
||||
const GenericTuple& to_tuple() const {
|
||||
if (!is_tuple()) throw std::runtime_error("Not a tuple");
|
||||
return std::get<GenericTuple>(value_);
|
||||
}
|
||||
|
||||
GenericTuple& to_tuple() {
|
||||
if (!is_tuple()) throw std::runtime_error("Not a tuple");
|
||||
return std::get<GenericTuple>(value_);
|
||||
}
|
||||
|
||||
at::ScalarType to_scalar_type() const {
|
||||
if (!is_int()) throw std::runtime_error("Not an int");
|
||||
return static_cast<at::ScalarType>(std::get<int64_t>(value_));
|
||||
}
|
||||
|
||||
bool toBool() const { return to_bool(); }
|
||||
int64_t toInt() const { return to_int(); }
|
||||
double toDouble() const { return to_double(); }
|
||||
const std::string& toStringRef() const { return to_string(); }
|
||||
std::string_view toStringView() const { return to_string_view(); }
|
||||
at::Tensor toTensor() const { return to_tensor(); }
|
||||
at::ScalarType toScalarType() const { return to_scalar_type(); }
|
||||
|
||||
std::string tagKind() const {
|
||||
switch (tag_) {
|
||||
case TypeTag::None:
|
||||
return "None";
|
||||
case TypeTag::Bool:
|
||||
return "Bool";
|
||||
case TypeTag::Int:
|
||||
return "Int";
|
||||
case TypeTag::Double:
|
||||
return "Double";
|
||||
case TypeTag::String:
|
||||
return "String";
|
||||
case TypeTag::Tensor:
|
||||
return "Tensor";
|
||||
case TypeTag::GenericList:
|
||||
return "GenericList";
|
||||
case TypeTag::CustomClass:
|
||||
return "CustomClass";
|
||||
case TypeTag::Tuple:
|
||||
return "Tuple";
|
||||
default:
|
||||
return "InvalidTag";
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
intrusive_ptr<T> to_custom_class() const {
|
||||
if (!is_custom_class()) throw std::runtime_error("Not a custom class");
|
||||
const auto& wrapper = std::get<CustomClassWrapper>(value_);
|
||||
auto casted = std::dynamic_pointer_cast<T>(wrapper.ptr);
|
||||
if (!casted) {
|
||||
throw std::runtime_error("Cannot cast custom class to requested type");
|
||||
}
|
||||
return intrusive_ptr<T>(casted);
|
||||
}
|
||||
|
||||
private:
|
||||
template <typename T>
|
||||
struct is_intrusive_ptr : std::false_type {};
|
||||
|
||||
template <typename T>
|
||||
struct is_intrusive_ptr<intrusive_ptr<T>> : std::true_type {};
|
||||
|
||||
template <typename T>
|
||||
static constexpr bool is_intrusive_ptr_v = is_intrusive_ptr<T>::value;
|
||||
|
||||
public:
|
||||
bool try_to_bool(bool& out) const { // NOLINT
|
||||
if (is_bool()) {
|
||||
out = std::get<bool>(value_);
|
||||
return true;
|
||||
} else if (is_int()) {
|
||||
out = (std::get<int64_t>(value_) != 0);
|
||||
return true;
|
||||
} else if (is_double()) {
|
||||
out = (std::get<double>(value_) != 0.0);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool try_to_int(int& out) const { // NOLINT
|
||||
if (is_int()) {
|
||||
out = static_cast<int>(std::get<int64_t>(value_));
|
||||
return true;
|
||||
} else if (is_double()) {
|
||||
double val = std::get<double>(value_);
|
||||
if (val != static_cast<int>(val)) {
|
||||
std::cout << "Warning: Converting double(" << val
|
||||
<< ") to int (precision loss)" << std::endl;
|
||||
}
|
||||
out = static_cast<int>(val);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool try_to_double(double& out) const { // NOLINT
|
||||
if (is_double()) {
|
||||
out = std::get<double>(value_);
|
||||
return true;
|
||||
} else if (is_int()) {
|
||||
out = static_cast<double>(std::get<int64_t>(value_));
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool try_to_string(std::string& out) const { // NOLINT
|
||||
if (is_string()) {
|
||||
out = std::get<std::string>(value_);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool try_to_tensor(at::Tensor& out) const { // NOLINT
|
||||
if (is_tensor()) {
|
||||
out = std::get<at::Tensor>(value_);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool try_to_scalar_type(at::ScalarType& out) const { // NOLINT
|
||||
if (is_int()) {
|
||||
out = static_cast<at::ScalarType>(std::get<int64_t>(value_));
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool try_to_optional_type(std::optional<T>& out) const { // NOLINT
|
||||
if (is_none()) {
|
||||
out = std::nullopt;
|
||||
return true;
|
||||
} else {
|
||||
T value;
|
||||
if (try_convert_to<T>(value)) {
|
||||
out = value;
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool try_to_custom_class(std::shared_ptr<CustomClassHolder>& out, // NOLINT
|
||||
const std::string& expected_class_name) const {
|
||||
if (is_custom_class()) {
|
||||
const auto& wrapper = std::get<CustomClassWrapper>(value_);
|
||||
if (wrapper.class_name == expected_class_name) {
|
||||
out = wrapper.ptr;
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool try_convert_to(T& out) const { // NOLINT
|
||||
// Remove reference and cv-qualifiers from T
|
||||
using BaseType = std::remove_cv_t<std::remove_reference_t<T>>;
|
||||
|
||||
if constexpr (std::is_same_v<BaseType, bool>) {
|
||||
return try_to_bool(const_cast<bool&>(reinterpret_cast<const bool&>(out)));
|
||||
} else if constexpr (std::is_same_v<BaseType, int>) {
|
||||
return try_to_int(const_cast<int&>(reinterpret_cast<const int&>(out)));
|
||||
} else if constexpr (std::is_same_v<BaseType, double>) {
|
||||
return try_to_double(
|
||||
const_cast<double&>(reinterpret_cast<const double&>(out)));
|
||||
} else if constexpr (std::is_same_v<BaseType, std::string>) {
|
||||
return try_to_string(
|
||||
const_cast<std::string&>(reinterpret_cast<const std::string&>(out)));
|
||||
} else if constexpr (std::is_same_v<BaseType, at::Tensor>) {
|
||||
return try_to_tensor(
|
||||
const_cast<at::Tensor&>(reinterpret_cast<const at::Tensor&>(out)));
|
||||
} else if constexpr (std::is_same_v<BaseType, at::ScalarType>) {
|
||||
return try_to_scalar_type(const_cast<at::ScalarType&>(
|
||||
reinterpret_cast<const at::ScalarType&>(out)));
|
||||
} else {
|
||||
try {
|
||||
// Handle const types by removing const and using const_cast
|
||||
using NonConstType = std::remove_const_t<T>;
|
||||
NonConstType temp = this->to<BaseType>();
|
||||
const_cast<NonConstType&>(out) = std::move(temp);
|
||||
return true;
|
||||
} catch (const std::exception&) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::string get_custom_class_name() const {
|
||||
if (!is_custom_class()) throw std::runtime_error("Not a custom class");
|
||||
const auto& wrapper = std::get<CustomClassWrapper>(value_);
|
||||
return wrapper.class_name;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T to() && {
|
||||
return generic_to(std::move(*this), _fake_type<T>{});
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T to() const& {
|
||||
return generic_to(*this, _fake_type<T>{});
|
||||
}
|
||||
|
||||
std::string type_string() const {
|
||||
switch (tag_) {
|
||||
case TypeTag::None:
|
||||
return "None";
|
||||
case TypeTag::Bool:
|
||||
return "Bool";
|
||||
case TypeTag::Int:
|
||||
return "Int";
|
||||
case TypeTag::Double:
|
||||
return "Double";
|
||||
case TypeTag::String:
|
||||
return "String";
|
||||
case TypeTag::Tensor:
|
||||
return "Tensor";
|
||||
case TypeTag::GenericList:
|
||||
return "List";
|
||||
case TypeTag::Tuple:
|
||||
return "Tuple";
|
||||
case TypeTag::CustomClass:
|
||||
return "CustomClass(" + get_custom_class_name() + ")";
|
||||
default:
|
||||
return "Unknown";
|
||||
}
|
||||
}
|
||||
|
||||
std::string to_repr() const {
|
||||
switch (tag_) {
|
||||
case TypeTag::None:
|
||||
return "None";
|
||||
case TypeTag::Bool:
|
||||
return std::get<bool>(value_) ? "true" : "false";
|
||||
case TypeTag::Int:
|
||||
return std::to_string(std::get<int64_t>(value_));
|
||||
case TypeTag::Double:
|
||||
return std::to_string(std::get<double>(value_));
|
||||
case TypeTag::String:
|
||||
return "\"" + std::get<std::string>(value_) + "\"";
|
||||
case TypeTag::Tensor: {
|
||||
const auto& tensor = std::get<at::Tensor>(value_);
|
||||
return "Tensor(" + std::to_string(tensor.numel()) + " elements)";
|
||||
}
|
||||
case TypeTag::GenericList: {
|
||||
const auto& list = std::get<GenericList>(value_);
|
||||
std::string result = "[";
|
||||
for (size_t i = 0; i < list.size(); ++i) {
|
||||
if (i > 0) result += ", ";
|
||||
result += list[i].to_repr();
|
||||
}
|
||||
result += "]";
|
||||
return result;
|
||||
}
|
||||
case TypeTag::Tuple: {
|
||||
const auto& tuple = std::get<GenericTuple>(value_);
|
||||
std::string result = "(";
|
||||
for (size_t i = 0; i < tuple.size(); ++i) {
|
||||
if (i > 0) result += ", ";
|
||||
result += tuple[i].to_repr();
|
||||
}
|
||||
if (tuple.size() == 1) result += ","; // Single element tuple
|
||||
result += ")";
|
||||
return result;
|
||||
}
|
||||
case TypeTag::CustomClass: {
|
||||
const auto& wrapper = std::get<CustomClassWrapper>(value_);
|
||||
return "CustomClass(" + wrapper.class_name + ")";
|
||||
}
|
||||
default:
|
||||
return "Unknown";
|
||||
}
|
||||
}
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const IValue& val) {
|
||||
return os << val.to_repr();
|
||||
}
|
||||
|
||||
private:
|
||||
TypeTag tag_;
|
||||
std::variant<std::monostate,
|
||||
bool,
|
||||
int64_t,
|
||||
double,
|
||||
std::string,
|
||||
at::Tensor,
|
||||
GenericList,
|
||||
CustomClassWrapper,
|
||||
GenericTuple>
|
||||
value_;
|
||||
template <typename T>
|
||||
friend T generic_to(const IValue& ivalue, _fake_type<T>);
|
||||
};
|
||||
|
||||
inline GenericTuple::GenericTuple() = default;
|
||||
inline GenericTuple::GenericTuple(std::vector<IValue> elems) // NOLINT
|
||||
: elements(std::move(elems)) {}
|
||||
|
||||
inline size_t GenericTuple::size() const { return elements.size(); }
|
||||
inline IValue& GenericTuple::operator[](size_t idx) { return elements[idx]; }
|
||||
inline const IValue& GenericTuple::operator[](size_t idx) const {
|
||||
return elements[idx];
|
||||
}
|
||||
|
||||
template <>
|
||||
inline bool generic_to(const IValue& ivalue, _fake_type<bool>) {
|
||||
return ivalue.to_bool();
|
||||
}
|
||||
|
||||
template <>
|
||||
inline int generic_to(const IValue& ivalue, _fake_type<int>) {
|
||||
return static_cast<int>(ivalue.to_int());
|
||||
}
|
||||
|
||||
template <>
|
||||
inline int64_t generic_to(const IValue& ivalue, _fake_type<int64_t>) {
|
||||
return ivalue.to_int();
|
||||
}
|
||||
|
||||
template <>
|
||||
inline double generic_to(const IValue& ivalue, _fake_type<double>) {
|
||||
return ivalue.to_double();
|
||||
}
|
||||
|
||||
template <>
|
||||
inline std::string generic_to(const IValue& ivalue, _fake_type<std::string>) {
|
||||
return ivalue.to_string();
|
||||
}
|
||||
|
||||
template <>
|
||||
inline std::string_view generic_to(const IValue& ivalue,
|
||||
_fake_type<std::string_view>) {
|
||||
return ivalue.to_string_view();
|
||||
}
|
||||
|
||||
template <>
|
||||
inline at::Tensor generic_to(const IValue& ivalue, _fake_type<at::Tensor>) {
|
||||
return ivalue.to_tensor();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::vector<T> generic_to(const IValue& ivalue, _fake_type<std::vector<T>>) {
|
||||
auto list = ivalue.to_list();
|
||||
std::vector<T> result;
|
||||
result.reserve(list.size());
|
||||
for (const auto& item : list) {
|
||||
result.push_back(item.to<T>());
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// Helper for converting IValue tuple to std::tuple using index sequence
|
||||
template <typename Tuple, std::size_t... I>
|
||||
Tuple ivalue_to_tuple_impl(const IValue& ivalue, std::index_sequence<I...>) {
|
||||
const auto& generic_tuple = ivalue.to_tuple();
|
||||
if (generic_tuple.size() != sizeof...(I)) {
|
||||
throw std::runtime_error("Tuple size mismatch: expected " +
|
||||
std::to_string(sizeof...(I)) + " but got " +
|
||||
std::to_string(generic_tuple.size()));
|
||||
}
|
||||
// Use std::get<I> with index instead of type to avoid ambiguity
|
||||
// when tuple contains multiple elements of the same type
|
||||
return Tuple{generic_tuple[I].to<std::tuple_element_t<I, Tuple>>()...};
|
||||
}
|
||||
|
||||
// Generic conversion from IValue to std::tuple
|
||||
template <typename... Args>
|
||||
std::tuple<Args...> generic_to(const IValue& ivalue,
|
||||
_fake_type<std::tuple<Args...>>) {
|
||||
return ivalue_to_tuple_impl<std::tuple<Args...>>(
|
||||
ivalue, std::index_sequence_for<Args...>{});
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
ArrayRef<T> generic_to(const IValue& ivalue, _fake_type<ArrayRef<T>>) {
|
||||
static thread_local std::vector<T> temp_storage;
|
||||
temp_storage = ivalue.to<std::vector<T>>();
|
||||
return ArrayRef<T>(temp_storage);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::optional<T> generic_to(const IValue& ivalue,
|
||||
_fake_type<std::optional<T>>) {
|
||||
if (ivalue.is_none()) {
|
||||
return std::nullopt;
|
||||
}
|
||||
return std::optional<T>(ivalue.to<T>());
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
intrusive_ptr<T> generic_to(const IValue& ivalue,
|
||||
_fake_type<intrusive_ptr<T>>) {
|
||||
return ivalue.to_custom_class<T>();
|
||||
}
|
||||
|
||||
} // namespace torch
|
||||
|
||||
namespace c10 {
|
||||
using IValue = ::torch::IValue;
|
||||
}
|
||||
@@ -0,0 +1,385 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/core/jit_type_base.h>
|
||||
|
||||
#include <c10/util/Exception.h>
|
||||
|
||||
#include <memory>
|
||||
#include <ostream>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace c10 {
|
||||
|
||||
inline bool operator!=(const Type& lhs, const Type& rhs) {
|
||||
return !(lhs == rhs);
|
||||
}
|
||||
|
||||
namespace detail {
|
||||
|
||||
// Lightweight runtime types used only by the compat schema parser.
|
||||
class SchemaAtomicType final : public SharedType {
|
||||
public:
|
||||
static std::shared_ptr<SchemaAtomicType> create(TypeKind kind,
|
||||
std::string repr) {
|
||||
return std::shared_ptr<SchemaAtomicType>(
|
||||
new SchemaAtomicType(kind, std::move(repr)));
|
||||
}
|
||||
|
||||
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
|
||||
|
||||
std::string str() const override { return repr_; }
|
||||
|
||||
private:
|
||||
SchemaAtomicType(TypeKind kind, std::string repr)
|
||||
: SharedType(kind), repr_(std::move(repr)) {}
|
||||
|
||||
std::string annotation_str_impl(
|
||||
[[maybe_unused]] const TypePrinter& printer = nullptr) const override {
|
||||
return repr_;
|
||||
}
|
||||
|
||||
std::string repr_;
|
||||
};
|
||||
|
||||
class SchemaOptionalType final : public SharedType {
|
||||
public:
|
||||
static const TypeKind Kind = TypeKind::OptionalType;
|
||||
|
||||
static std::shared_ptr<SchemaOptionalType> create(TypePtr elem) {
|
||||
return std::shared_ptr<SchemaOptionalType>(
|
||||
new SchemaOptionalType(std::move(elem)));
|
||||
}
|
||||
|
||||
bool equals(const Type& rhs) const override {
|
||||
if (rhs.kind() != kind()) {
|
||||
return false;
|
||||
}
|
||||
const auto contained = rhs.containedTypes();
|
||||
return contained.size() == 1 && *elem_.front() == *contained.front();
|
||||
}
|
||||
|
||||
std::string str() const override { return elem_.front()->str() + "?"; }
|
||||
|
||||
at::ArrayRef<TypePtr> containedTypes() const override { return elem_; }
|
||||
|
||||
TypePtr createWithContained(
|
||||
std::vector<TypePtr> contained_types) const override {
|
||||
TORCH_CHECK(contained_types.size() == 1,
|
||||
"Optional type expects exactly one contained type");
|
||||
return create(std::move(contained_types.front()));
|
||||
}
|
||||
|
||||
private:
|
||||
explicit SchemaOptionalType(TypePtr elem)
|
||||
: SharedType(Kind), elem_{std::move(elem)} {}
|
||||
|
||||
std::string annotation_str_impl(
|
||||
const TypePrinter& printer = nullptr) const override {
|
||||
return "Optional[" + elem_.front()->annotation_str(printer) + "]";
|
||||
}
|
||||
|
||||
std::vector<TypePtr> elem_;
|
||||
};
|
||||
|
||||
class SchemaTupleType final : public SharedType {
|
||||
public:
|
||||
static const TypeKind Kind = TypeKind::TupleType;
|
||||
|
||||
static std::shared_ptr<SchemaTupleType> create(
|
||||
std::vector<TypePtr> elements) {
|
||||
return std::shared_ptr<SchemaTupleType>(
|
||||
new SchemaTupleType(std::move(elements)));
|
||||
}
|
||||
|
||||
bool equals(const Type& rhs) const override {
|
||||
if (rhs.kind() != kind()) {
|
||||
return false;
|
||||
}
|
||||
const auto rhs_elems = rhs.containedTypes();
|
||||
if (rhs_elems.size() != elements_.size()) {
|
||||
return false;
|
||||
}
|
||||
for (size_t i = 0; i < elements_.size(); ++i) {
|
||||
if (*elements_[i] != *rhs_elems[i]) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
std::string str() const override {
|
||||
std::stringstream ss;
|
||||
ss << "(";
|
||||
for (size_t i = 0; i < elements_.size(); ++i) {
|
||||
if (i > 0) {
|
||||
ss << ", ";
|
||||
}
|
||||
ss << elements_[i]->str();
|
||||
}
|
||||
ss << ")";
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
at::ArrayRef<TypePtr> containedTypes() const override { return elements_; }
|
||||
|
||||
TypePtr createWithContained(
|
||||
std::vector<TypePtr> contained_types) const override {
|
||||
return create(std::move(contained_types));
|
||||
}
|
||||
|
||||
private:
|
||||
explicit SchemaTupleType(std::vector<TypePtr> elements)
|
||||
: SharedType(Kind), elements_(std::move(elements)) {}
|
||||
|
||||
std::string annotation_str_impl(
|
||||
const TypePrinter& printer = nullptr) const override {
|
||||
std::stringstream ss;
|
||||
ss << "Tuple[";
|
||||
for (size_t i = 0; i < elements_.size(); ++i) {
|
||||
if (i > 0) {
|
||||
ss << ", ";
|
||||
}
|
||||
ss << elements_[i]->annotation_str(printer);
|
||||
}
|
||||
ss << "]";
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
std::vector<TypePtr> elements_;
|
||||
};
|
||||
|
||||
} // namespace detail
|
||||
|
||||
inline TypePtr makeSchemaAtomicType(TypeKind kind, std::string repr) {
|
||||
return detail::SchemaAtomicType::create(kind, std::move(repr));
|
||||
}
|
||||
|
||||
inline TypePtr makeSchemaOptionalType(TypePtr elem) {
|
||||
return detail::SchemaOptionalType::create(std::move(elem));
|
||||
}
|
||||
|
||||
inline TypePtr makeSchemaTupleType(std::vector<TypePtr> elements) {
|
||||
return detail::SchemaTupleType::create(std::move(elements));
|
||||
}
|
||||
|
||||
struct TensorType;
|
||||
using TensorTypePtr = SingletonTypePtr<TensorType>;
|
||||
struct PADDLE_API TensorType : public Type {
|
||||
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
|
||||
|
||||
std::string str() const override { return "Tensor"; }
|
||||
|
||||
bool isInferredType() const { return is_inferred_; }
|
||||
|
||||
static const TypeKind Kind = TypeKind::TensorType;
|
||||
|
||||
static TensorTypePtr get() {
|
||||
static TensorType value(/*inferred=*/false);
|
||||
return TensorTypePtr(&value);
|
||||
}
|
||||
|
||||
static TensorTypePtr getInferred() {
|
||||
static TensorType value(/*inferred=*/true);
|
||||
return TensorTypePtr(&value);
|
||||
}
|
||||
|
||||
private:
|
||||
explicit TensorType(bool inferred)
|
||||
: Type(TypeKind::TensorType), is_inferred_(inferred) {}
|
||||
|
||||
bool is_inferred_;
|
||||
};
|
||||
|
||||
struct NumberType;
|
||||
using NumberTypePtr = SingletonTypePtr<NumberType>;
|
||||
struct PADDLE_API NumberType : public Type {
|
||||
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
|
||||
|
||||
bool isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const override {
|
||||
return rhs.kind() == TypeKind::NumberType ||
|
||||
Type::isSubtypeOfExt(rhs, why_not);
|
||||
}
|
||||
|
||||
std::string str() const override { return "Scalar"; }
|
||||
|
||||
static const TypeKind Kind = TypeKind::NumberType;
|
||||
|
||||
static NumberTypePtr get() {
|
||||
static NumberType value;
|
||||
return NumberTypePtr(&value);
|
||||
}
|
||||
|
||||
protected:
|
||||
explicit NumberType(TypeKind kind = TypeKind::NumberType) : Type(kind) {}
|
||||
|
||||
std::string annotation_str_impl(
|
||||
[[maybe_unused]] const TypePrinter& printer = nullptr) const override {
|
||||
return "number";
|
||||
}
|
||||
};
|
||||
|
||||
struct FloatType;
|
||||
using FloatTypePtr = SingletonTypePtr<FloatType>;
|
||||
struct PADDLE_API FloatType : public NumberType {
|
||||
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
|
||||
|
||||
std::string str() const override { return "float"; }
|
||||
|
||||
bool isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const override {
|
||||
return rhs.kind() == TypeKind::NumberType ||
|
||||
Type::isSubtypeOfExt(rhs, why_not);
|
||||
}
|
||||
|
||||
static const TypeKind Kind = TypeKind::FloatType;
|
||||
|
||||
static FloatTypePtr get() {
|
||||
static FloatType value;
|
||||
return FloatTypePtr(&value);
|
||||
}
|
||||
|
||||
private:
|
||||
FloatType() : NumberType(TypeKind::FloatType) {}
|
||||
|
||||
std::string annotation_str_impl(
|
||||
[[maybe_unused]] const TypePrinter& printer = nullptr) const override {
|
||||
return "float";
|
||||
}
|
||||
};
|
||||
|
||||
struct IntType;
|
||||
using IntTypePtr = SingletonTypePtr<IntType>;
|
||||
struct PADDLE_API IntType : public NumberType {
|
||||
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
|
||||
|
||||
std::string str() const override { return "int"; }
|
||||
|
||||
bool isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const override {
|
||||
return rhs.kind() == TypeKind::NumberType ||
|
||||
Type::isSubtypeOfExt(rhs, why_not);
|
||||
}
|
||||
|
||||
static const TypeKind Kind = TypeKind::IntType;
|
||||
|
||||
static IntTypePtr get() {
|
||||
static IntType value;
|
||||
return IntTypePtr(&value);
|
||||
}
|
||||
|
||||
private:
|
||||
IntType() : NumberType(TypeKind::IntType) {}
|
||||
|
||||
std::string annotation_str_impl(
|
||||
[[maybe_unused]] const TypePrinter& printer = nullptr) const override {
|
||||
return "int";
|
||||
}
|
||||
};
|
||||
|
||||
struct BoolType;
|
||||
using BoolTypePtr = SingletonTypePtr<BoolType>;
|
||||
struct PADDLE_API BoolType : public Type {
|
||||
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
|
||||
|
||||
std::string str() const override { return "bool"; }
|
||||
|
||||
static const TypeKind Kind = TypeKind::BoolType;
|
||||
|
||||
static BoolTypePtr get() {
|
||||
static BoolType value;
|
||||
return BoolTypePtr(&value);
|
||||
}
|
||||
|
||||
private:
|
||||
BoolType() : Type(TypeKind::BoolType) {}
|
||||
};
|
||||
|
||||
struct StringType;
|
||||
using StringTypePtr = SingletonTypePtr<StringType>;
|
||||
struct PADDLE_API StringType : public Type {
|
||||
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
|
||||
|
||||
std::string str() const override { return annotation_str(); }
|
||||
|
||||
std::string annotation_str_impl(
|
||||
[[maybe_unused]] const TypePrinter& printer = nullptr) const override {
|
||||
return "str";
|
||||
}
|
||||
|
||||
static const TypeKind Kind = TypeKind::StringType;
|
||||
|
||||
static StringTypePtr get() {
|
||||
static StringType value;
|
||||
return StringTypePtr(&value);
|
||||
}
|
||||
|
||||
private:
|
||||
StringType() : Type(TypeKind::StringType) {}
|
||||
};
|
||||
|
||||
struct NoneType;
|
||||
using NoneTypePtr = SingletonTypePtr<NoneType>;
|
||||
struct PADDLE_API NoneType : public Type {
|
||||
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
|
||||
|
||||
std::string str() const override { return "NoneType"; }
|
||||
|
||||
bool isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const override {
|
||||
return rhs.kind() == TypeKind::OptionalType ||
|
||||
Type::isSubtypeOfExt(rhs, why_not);
|
||||
}
|
||||
|
||||
static const TypeKind Kind = TypeKind::NoneType;
|
||||
|
||||
static NoneTypePtr get() {
|
||||
static NoneType value;
|
||||
return NoneTypePtr(&value);
|
||||
}
|
||||
|
||||
private:
|
||||
NoneType() : Type(TypeKind::NoneType) {}
|
||||
};
|
||||
|
||||
struct DeviceObjType;
|
||||
using DeviceObjTypePtr = SingletonTypePtr<DeviceObjType>;
|
||||
struct PADDLE_API DeviceObjType : public Type {
|
||||
bool equals(const Type& rhs) const override { return rhs.kind() == kind(); }
|
||||
|
||||
std::string str() const override { return "Device"; }
|
||||
|
||||
static const TypeKind Kind = TypeKind::DeviceObjType;
|
||||
|
||||
static DeviceObjTypePtr get() {
|
||||
static DeviceObjType value;
|
||||
return DeviceObjTypePtr(&value);
|
||||
}
|
||||
|
||||
private:
|
||||
DeviceObjType() : Type(TypeKind::DeviceObjType) {}
|
||||
};
|
||||
|
||||
inline std::ostream& operator<<(std::ostream& out, const Type& t) {
|
||||
out << t.str();
|
||||
return out;
|
||||
}
|
||||
|
||||
} // namespace c10
|
||||
@@ -0,0 +1,337 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <ostream>
|
||||
#include <string>
|
||||
#include <type_traits>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "ATen/core/type_ptr.h"
|
||||
#include "c10/util/ArrayRef.h"
|
||||
#include "c10/util/Exception.h"
|
||||
|
||||
namespace c10 {
|
||||
|
||||
#define C10_FORALL_TYPES(_) \
|
||||
_(TensorType) \
|
||||
_(StringType) \
|
||||
_(IntType) \
|
||||
_(FloatType) \
|
||||
_(BoolType) \
|
||||
_(NoneType) \
|
||||
_(TupleType) \
|
||||
_(NumberType) \
|
||||
_(OptionalType) \
|
||||
_(UnionType) \
|
||||
_(DeviceObjType) \
|
||||
_(DynamicType)
|
||||
|
||||
enum class TypeKind {
|
||||
#define DEFINE_TYPE(T) T,
|
||||
C10_FORALL_TYPES(DEFINE_TYPE)
|
||||
#undef DEFINE_TYPE
|
||||
};
|
||||
|
||||
struct Type;
|
||||
struct SharedType;
|
||||
using TypePrinter = std::function<std::optional<std::string>(const Type&)>;
|
||||
|
||||
namespace detail {
|
||||
template <typename T>
|
||||
struct IsSingletonType : std::false_type {};
|
||||
} // namespace detail
|
||||
#define TORCH_DECLARE_SINGLETON(Type) \
|
||||
struct Type; \
|
||||
namespace detail { \
|
||||
template <> \
|
||||
struct IsSingletonType<Type> : std::true_type {}; \
|
||||
}
|
||||
|
||||
TORCH_DECLARE_SINGLETON(NumberType)
|
||||
TORCH_DECLARE_SINGLETON(TensorType)
|
||||
TORCH_DECLARE_SINGLETON(StringType)
|
||||
TORCH_DECLARE_SINGLETON(IntType)
|
||||
TORCH_DECLARE_SINGLETON(FloatType)
|
||||
TORCH_DECLARE_SINGLETON(BoolType)
|
||||
TORCH_DECLARE_SINGLETON(NoneType)
|
||||
TORCH_DECLARE_SINGLETON(TupleType)
|
||||
TORCH_DECLARE_SINGLETON(OptionalType)
|
||||
TORCH_DECLARE_SINGLETON(DeviceObjType)
|
||||
|
||||
namespace detail {
|
||||
template <typename T, typename Enable = void>
|
||||
struct CastReturnType {
|
||||
using type = std::shared_ptr<T>;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct CastReturnType<T, std::enable_if_t<IsSingletonType<T>::value>> {
|
||||
using type = SingletonTypePtr<T>;
|
||||
};
|
||||
|
||||
template <typename T, typename Enable = void>
|
||||
struct CastConstReturnType {
|
||||
using type = std::shared_ptr<const T>;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct CastConstReturnType<T, std::enable_if_t<IsSingletonType<T>::value>> {
|
||||
using type = SingletonTypePtr<const T>;
|
||||
};
|
||||
|
||||
} // namespace detail
|
||||
|
||||
struct PADDLE_API Type {
|
||||
friend PADDLE_API bool operator==(const Type& lhs, const Type& rhs);
|
||||
|
||||
private:
|
||||
TypeKind kind_;
|
||||
|
||||
protected:
|
||||
explicit Type(TypeKind kind) : kind_(kind) {}
|
||||
|
||||
Type(const Type&) = default;
|
||||
Type& operator=(const Type&) = default;
|
||||
Type(Type&&) noexcept = default;
|
||||
Type& operator=(Type&&) noexcept = default;
|
||||
|
||||
virtual std::string annotation_str_impl(
|
||||
const TypePrinter& /*printer*/) const {
|
||||
return str();
|
||||
}
|
||||
virtual bool equals(const Type& rhs) const = 0;
|
||||
virtual bool symmetric() const { return true; }
|
||||
|
||||
public:
|
||||
template <typename T>
|
||||
class SingletonOrSharedTypePtr {
|
||||
public:
|
||||
using element_type = typename std::shared_ptr<T>::element_type;
|
||||
|
||||
SingletonOrSharedTypePtr() = default;
|
||||
|
||||
SingletonOrSharedTypePtr(std::shared_ptr<T> x) // NOLINT(runtime/explicit)
|
||||
: repr_(std::move(x)) {}
|
||||
|
||||
template <typename U,
|
||||
std::enable_if_t<std::is_convertible_v<U*, T*>, bool> = true>
|
||||
SingletonOrSharedTypePtr(std::shared_ptr<U> x) // NOLINT(runtime/explicit)
|
||||
: repr_(std::move(x)) {}
|
||||
|
||||
SingletonOrSharedTypePtr(std::nullptr_t) // NOLINT(runtime/explicit)
|
||||
: repr_(nullptr) {}
|
||||
|
||||
SingletonOrSharedTypePtr(SingletonTypePtr<T> p) // NOLINT(runtime/explicit)
|
||||
: repr_(makeSingletonSharedPtr(p.get())) {}
|
||||
|
||||
template <typename U,
|
||||
std::enable_if_t<std::is_convertible_v<U*, T*>, bool> = true>
|
||||
SingletonOrSharedTypePtr(SingletonTypePtr<U> p) // NOLINT(runtime/explicit)
|
||||
: repr_(makeSingletonSharedPtr(static_cast<T*>(p.get()))) {}
|
||||
|
||||
SingletonOrSharedTypePtr(const SingletonOrSharedTypePtr&) = default;
|
||||
SingletonOrSharedTypePtr(SingletonOrSharedTypePtr&&) noexcept = default;
|
||||
SingletonOrSharedTypePtr& operator=(const SingletonOrSharedTypePtr&) =
|
||||
default;
|
||||
SingletonOrSharedTypePtr& operator=(SingletonOrSharedTypePtr&&) noexcept =
|
||||
default;
|
||||
~SingletonOrSharedTypePtr() = default;
|
||||
|
||||
T* get() const { return repr_.get(); }
|
||||
|
||||
operator bool() const { return repr_ != nullptr; }
|
||||
|
||||
bool operator==(std::nullptr_t) const { return repr_ == nullptr; }
|
||||
|
||||
bool operator!=(std::nullptr_t) const { return repr_ != nullptr; }
|
||||
|
||||
template <typename U = T,
|
||||
std::enable_if_t<!std::is_same_v<std::remove_const_t<U>, void>,
|
||||
bool> = true>
|
||||
U& operator*() const {
|
||||
return *get();
|
||||
}
|
||||
|
||||
T* operator->() const { return get(); }
|
||||
|
||||
private:
|
||||
static std::shared_ptr<T> makeSingletonSharedPtr(T* ptr) {
|
||||
return std::shared_ptr<T>(std::shared_ptr<T>(), ptr);
|
||||
}
|
||||
|
||||
std::shared_ptr<T> repr_;
|
||||
};
|
||||
|
||||
using TypePtr = SingletonOrSharedTypePtr<Type>;
|
||||
|
||||
virtual bool isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const;
|
||||
bool isSubtypeOf(const Type& rhs) const {
|
||||
return isSubtypeOfExt(rhs, nullptr);
|
||||
}
|
||||
|
||||
// Compatibility shims to accommodate existing code that passes shared_ptrs
|
||||
// around. Ideally, we would just delete this, but it should be harmless.
|
||||
template <typename T>
|
||||
std::enable_if_t<std::is_base_of_v<Type, T>, bool> isSubtypeOf(
|
||||
const std::shared_ptr<T>& rhs) const {
|
||||
return isSubtypeOf(*rhs);
|
||||
}
|
||||
|
||||
virtual std::string str() const = 0;
|
||||
|
||||
std::string annotation_str(const TypePrinter& printer) const {
|
||||
if (printer) {
|
||||
if (auto renamed = printer(*this)) {
|
||||
return *renamed;
|
||||
}
|
||||
}
|
||||
return annotation_str_impl(printer);
|
||||
}
|
||||
std::string annotation_str() const { return annotation_str(nullptr); }
|
||||
|
||||
virtual std::string repr_str() const { return annotation_str(); }
|
||||
|
||||
TypeKind kind() const { return kind_; }
|
||||
|
||||
template <typename T,
|
||||
std::enable_if_t<!detail::IsSingletonType<T>::value, bool> = true>
|
||||
typename detail::CastReturnType<T>::type cast() {
|
||||
if (auto* typed = dynamic_cast<T*>(this)) {
|
||||
return std::static_pointer_cast<T>(typed->shared_from_this());
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
template <typename T,
|
||||
std::enable_if_t<detail::IsSingletonType<T>::value, bool> = true>
|
||||
typename detail::CastReturnType<T>::type cast() {
|
||||
if (auto* typed = dynamic_cast<T*>(this)) {
|
||||
return typename detail::CastReturnType<T>::type(typed);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
template <typename T,
|
||||
std::enable_if_t<!detail::IsSingletonType<T>::value, bool> = true>
|
||||
typename detail::CastConstReturnType<T>::type cast() const {
|
||||
if (auto* typed = dynamic_cast<const T*>(this)) {
|
||||
return std::static_pointer_cast<const T>(typed->shared_from_this());
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
template <typename T,
|
||||
std::enable_if_t<detail::IsSingletonType<T>::value, bool> = true>
|
||||
typename detail::CastConstReturnType<T>::type cast() const {
|
||||
if (auto* typed = dynamic_cast<const T*>(this)) {
|
||||
return typename detail::CastConstReturnType<T>::type(typed);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
virtual ~Type() = default;
|
||||
virtual at::ArrayRef<TypePtr> containedTypes() const { return {}; }
|
||||
virtual TypePtr createWithContained(
|
||||
std::vector<TypePtr> /*contained_types*/) const {
|
||||
TORCH_CHECK(
|
||||
false,
|
||||
"type with contained types did not overload createWithContained: ",
|
||||
str());
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using SingletonOrSharedTypePtr = Type::SingletonOrSharedTypePtr<T>;
|
||||
|
||||
template <typename T, typename U>
|
||||
bool operator==(const SingletonOrSharedTypePtr<T>& x,
|
||||
const SingletonOrSharedTypePtr<U>& y) {
|
||||
return static_cast<const void*>(x.get()) == static_cast<const void*>(y.get());
|
||||
}
|
||||
|
||||
template <typename T, typename U>
|
||||
bool operator==(const SingletonOrSharedTypePtr<T>& x,
|
||||
const SingletonTypePtr<U>& y) {
|
||||
return static_cast<const void*>(x.get()) == static_cast<const void*>(y.get());
|
||||
}
|
||||
|
||||
template <typename T, typename U>
|
||||
bool operator==(const SingletonTypePtr<T>& x,
|
||||
const SingletonOrSharedTypePtr<U>& y) {
|
||||
return static_cast<const void*>(x.get()) == static_cast<const void*>(y.get());
|
||||
}
|
||||
|
||||
template <typename T, typename U>
|
||||
bool operator!=(const SingletonOrSharedTypePtr<T>& x,
|
||||
const SingletonOrSharedTypePtr<U>& y) {
|
||||
return !(x == y);
|
||||
}
|
||||
|
||||
template <typename T, typename U>
|
||||
bool operator!=(const SingletonOrSharedTypePtr<T>& x,
|
||||
const SingletonTypePtr<U>& y) {
|
||||
return !(x == y);
|
||||
}
|
||||
|
||||
template <typename T, typename U>
|
||||
bool operator!=(const SingletonTypePtr<T>& x,
|
||||
const SingletonOrSharedTypePtr<U>& y) {
|
||||
return !(x == y);
|
||||
}
|
||||
|
||||
using TypePtr = SingletonOrSharedTypePtr<Type>;
|
||||
|
||||
// Base class for Types that are guaranteed to be owned by std::shared_ptr.
|
||||
struct PADDLE_API SharedType : public Type,
|
||||
public std::enable_shared_from_this<SharedType> {
|
||||
using Type::Type;
|
||||
};
|
||||
|
||||
inline bool operator==(const Type& lhs, const Type& rhs) {
|
||||
if (!rhs.symmetric()) {
|
||||
return rhs.equals(lhs);
|
||||
}
|
||||
return lhs.equals(rhs);
|
||||
}
|
||||
|
||||
inline bool Type::isSubtypeOfExt(const Type& rhs, std::ostream* why_not) const {
|
||||
if (*this == rhs) {
|
||||
return true;
|
||||
}
|
||||
if (rhs.kind() == TypeKind::OptionalType) {
|
||||
for (const auto& inner : rhs.containedTypes()) {
|
||||
if (*this == *inner) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
} // namespace c10
|
||||
|
||||
namespace std {
|
||||
template <typename T>
|
||||
struct hash<c10::SingletonOrSharedTypePtr<T>> {
|
||||
size_t operator()(const c10::SingletonOrSharedTypePtr<T>& x) const {
|
||||
return std::hash<T*>()(x.get());
|
||||
}
|
||||
};
|
||||
} // namespace std
|
||||
@@ -0,0 +1,70 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <type_traits>
|
||||
|
||||
#include "c10/util/Exception.h"
|
||||
|
||||
namespace c10 {
|
||||
|
||||
// Compatibility wrapper around a raw pointer so that existing code
|
||||
// written to deal with a shared_ptr can keep working.
|
||||
template <typename T>
|
||||
class SingletonTypePtr {
|
||||
public:
|
||||
SingletonTypePtr(T* p) : repr_(p) {} // NOLINT(runtime/explicit)
|
||||
|
||||
// We need this to satisfy Pybind11, but it shouldn't be hit.
|
||||
explicit SingletonTypePtr(std::shared_ptr<T> /*unused*/) {
|
||||
TORCH_CHECK(false);
|
||||
}
|
||||
|
||||
using element_type = typename std::shared_ptr<T>::element_type;
|
||||
|
||||
template <typename U = T,
|
||||
std::enable_if_t<!std::is_same_v<std::remove_const_t<U>, void>,
|
||||
bool> = true>
|
||||
T& operator*() const {
|
||||
return *repr_;
|
||||
}
|
||||
|
||||
T* get() const { return repr_; }
|
||||
|
||||
T* operator->() const { return repr_; }
|
||||
|
||||
operator bool() const { return repr_ != nullptr; }
|
||||
|
||||
private:
|
||||
T* repr_{nullptr};
|
||||
};
|
||||
|
||||
template <typename T, typename U>
|
||||
bool operator==(SingletonTypePtr<T> lhs, SingletonTypePtr<U> rhs) {
|
||||
return static_cast<const void*>(lhs.get()) ==
|
||||
static_cast<const void*>(rhs.get());
|
||||
}
|
||||
|
||||
template <typename T, typename U>
|
||||
bool operator!=(SingletonTypePtr<T> lhs, SingletonTypePtr<U> rhs) {
|
||||
return !(lhs == rhs);
|
||||
}
|
||||
|
||||
} // namespace c10
|
||||
@@ -0,0 +1,201 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA)
|
||||
|
||||
#include <ATen/cuda/CUDABlas.h>
|
||||
|
||||
#include "paddle/phi/backends/dynload/cublas.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
|
||||
namespace at::cuda::blas {
|
||||
|
||||
namespace {
|
||||
|
||||
inline cublasOperation_t to_cublas_op(char trans) {
|
||||
switch (trans) {
|
||||
case 'T':
|
||||
case 't':
|
||||
return CUBLAS_OP_T;
|
||||
case 'N':
|
||||
case 'n':
|
||||
return CUBLAS_OP_N;
|
||||
case 'C':
|
||||
case 'c':
|
||||
return CUBLAS_OP_C;
|
||||
default:
|
||||
PADDLE_THROW(common::errors::InvalidArgument(
|
||||
"at::cuda::blas::gemm: invalid transpose character '%c'", trans));
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
/* ───────────── gemm<double> ───────────── */
|
||||
template <>
|
||||
void gemm<double>(CUDABLAS_GEMM_ARGTYPES(double)) {
|
||||
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
|
||||
cublasOperation_t opa = to_cublas_op(transa);
|
||||
cublasOperation_t opb = to_cublas_op(transb);
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cublasDgemm(handle,
|
||||
opa,
|
||||
opb,
|
||||
static_cast<int>(m),
|
||||
static_cast<int>(n),
|
||||
static_cast<int>(k),
|
||||
&alpha,
|
||||
a,
|
||||
static_cast<int>(lda),
|
||||
b,
|
||||
static_cast<int>(ldb),
|
||||
&beta,
|
||||
c,
|
||||
static_cast<int>(ldc)));
|
||||
}
|
||||
|
||||
/* ───────────── gemm<float> ───────────── */
|
||||
template <>
|
||||
void gemm<float>(CUDABLAS_GEMM_ARGTYPES(float)) {
|
||||
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
|
||||
cublasOperation_t opa = to_cublas_op(transa);
|
||||
cublasOperation_t opb = to_cublas_op(transb);
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cublasSgemm(handle,
|
||||
opa,
|
||||
opb,
|
||||
static_cast<int>(m),
|
||||
static_cast<int>(n),
|
||||
static_cast<int>(k),
|
||||
&alpha,
|
||||
a,
|
||||
static_cast<int>(lda),
|
||||
b,
|
||||
static_cast<int>(ldb),
|
||||
&beta,
|
||||
c,
|
||||
static_cast<int>(ldc)));
|
||||
}
|
||||
|
||||
/* ───────────── gemm<c10::complex<double>> ───────────── */
|
||||
template <>
|
||||
void gemm<c10::complex<double>>(CUDABLAS_GEMM_ARGTYPES(c10::complex<double>)) {
|
||||
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
|
||||
cublasOperation_t opa = to_cublas_op(transa);
|
||||
cublasOperation_t opb = to_cublas_op(transb);
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cublasZgemm(
|
||||
handle,
|
||||
opa,
|
||||
opb,
|
||||
static_cast<int>(m),
|
||||
static_cast<int>(n),
|
||||
static_cast<int>(k),
|
||||
reinterpret_cast<const cuDoubleComplex *>(&alpha),
|
||||
reinterpret_cast<const cuDoubleComplex *>(a),
|
||||
static_cast<int>(lda),
|
||||
reinterpret_cast<const cuDoubleComplex *>(b),
|
||||
static_cast<int>(ldb),
|
||||
reinterpret_cast<const cuDoubleComplex *>(&beta),
|
||||
reinterpret_cast<cuDoubleComplex *>(c),
|
||||
static_cast<int>(ldc)));
|
||||
}
|
||||
|
||||
/* ───────────── gemm<c10::complex<float>> ───────────── */
|
||||
template <>
|
||||
void gemm<c10::complex<float>>(CUDABLAS_GEMM_ARGTYPES(c10::complex<float>)) {
|
||||
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
|
||||
cublasOperation_t opa = to_cublas_op(transa);
|
||||
cublasOperation_t opb = to_cublas_op(transb);
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cublasCgemm(
|
||||
handle,
|
||||
opa,
|
||||
opb,
|
||||
static_cast<int>(m),
|
||||
static_cast<int>(n),
|
||||
static_cast<int>(k),
|
||||
reinterpret_cast<const cuFloatComplex *>(&alpha),
|
||||
reinterpret_cast<const cuFloatComplex *>(a),
|
||||
static_cast<int>(lda),
|
||||
reinterpret_cast<const cuFloatComplex *>(b),
|
||||
static_cast<int>(ldb),
|
||||
reinterpret_cast<const cuFloatComplex *>(&beta),
|
||||
reinterpret_cast<cuFloatComplex *>(c),
|
||||
static_cast<int>(ldc)));
|
||||
}
|
||||
|
||||
/* ───────────── gemm<at::Half> ───────────── */
|
||||
template <>
|
||||
void gemm<at::Half>(CUDABLAS_GEMM_ARGTYPES(at::Half)) {
|
||||
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
|
||||
cublasOperation_t opa = to_cublas_op(transa);
|
||||
cublasOperation_t opb = to_cublas_op(transb);
|
||||
|
||||
// Use cublasGemmEx with FP32 compute for Half inputs
|
||||
float alpha_f = alpha;
|
||||
float beta_f = beta;
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
phi::dynload::cublasGemmEx(handle,
|
||||
opa,
|
||||
opb,
|
||||
static_cast<int>(m),
|
||||
static_cast<int>(n),
|
||||
static_cast<int>(k),
|
||||
&alpha_f,
|
||||
a,
|
||||
CUDA_R_16F,
|
||||
static_cast<int>(lda),
|
||||
b,
|
||||
CUDA_R_16F,
|
||||
static_cast<int>(ldb),
|
||||
&beta_f,
|
||||
c,
|
||||
CUDA_R_16F,
|
||||
static_cast<int>(ldc),
|
||||
CUDA_R_32F,
|
||||
CUBLAS_GEMM_DEFAULT_TENSOR_OP));
|
||||
}
|
||||
|
||||
/* ───────────── gemm<at::BFloat16> ───────────── */
|
||||
template <>
|
||||
void gemm<at::BFloat16>(CUDABLAS_GEMM_ARGTYPES(at::BFloat16)) {
|
||||
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
|
||||
cublasOperation_t opa = to_cublas_op(transa);
|
||||
cublasOperation_t opb = to_cublas_op(transb);
|
||||
|
||||
// Use cublasGemmEx with FP32 compute for BFloat16 inputs
|
||||
float alpha_f = alpha;
|
||||
float beta_f = beta;
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
phi::dynload::cublasGemmEx(handle,
|
||||
opa,
|
||||
opb,
|
||||
static_cast<int>(m),
|
||||
static_cast<int>(n),
|
||||
static_cast<int>(k),
|
||||
&alpha_f,
|
||||
a,
|
||||
CUDA_R_16BF,
|
||||
static_cast<int>(lda),
|
||||
b,
|
||||
CUDA_R_16BF,
|
||||
static_cast<int>(ldb),
|
||||
&beta_f,
|
||||
c,
|
||||
CUDA_R_16BF,
|
||||
static_cast<int>(ldc),
|
||||
CUDA_R_32F,
|
||||
CUBLAS_GEMM_DEFAULT_TENSOR_OP));
|
||||
}
|
||||
|
||||
} // namespace at::cuda::blas
|
||||
|
||||
#endif // PADDLE_WITH_CUDA
|
||||
@@ -0,0 +1,73 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
/*
|
||||
Provides a subset of CUDA BLAS functions as templates:
|
||||
|
||||
gemm<Dtype>(transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c,
|
||||
ldc)
|
||||
|
||||
where Dtype is double, float, c10::complex<double>, c10::complex<float>,
|
||||
at::Half or at::BFloat16. The functions are available in at::cuda::blas
|
||||
namespace.
|
||||
*/
|
||||
|
||||
#include <ATen/OpMathType.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
|
||||
#include "paddle/common/macros.h"
|
||||
|
||||
namespace at::cuda::blas {
|
||||
|
||||
/* LEVEL 3 BLAS FUNCTIONS */
|
||||
|
||||
#define CUDABLAS_GEMM_ARGTYPES(Dtype) \
|
||||
CUDABLAS_GEMM_ARGTYPES_AND_C_DTYPE(Dtype, Dtype)
|
||||
|
||||
#define CUDABLAS_GEMM_ARGTYPES_AND_C_DTYPE(Dtype, C_Dtype) \
|
||||
char transa, char transb, int64_t m, int64_t n, int64_t k, \
|
||||
at::opmath_type<Dtype> alpha, const Dtype *a, int64_t lda, \
|
||||
const Dtype *b, int64_t ldb, at::opmath_type<Dtype> beta, C_Dtype *c, \
|
||||
int64_t ldc
|
||||
|
||||
#define CUDABLAS_GEMM_ARGS(Dtype) \
|
||||
transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc
|
||||
|
||||
template <typename Dtype, typename C_Dtype = Dtype>
|
||||
inline void gemm(CUDABLAS_GEMM_ARGTYPES_AND_C_DTYPE(Dtype, C_Dtype)) {
|
||||
static_assert(false && sizeof(Dtype),
|
||||
"at::cuda::blas::gemm: not implemented");
|
||||
}
|
||||
|
||||
template <>
|
||||
PADDLE_API void gemm<double>(CUDABLAS_GEMM_ARGTYPES(double));
|
||||
template <>
|
||||
PADDLE_API void gemm<float>(CUDABLAS_GEMM_ARGTYPES(float));
|
||||
template <>
|
||||
PADDLE_API void gemm<c10::complex<double>>(
|
||||
CUDABLAS_GEMM_ARGTYPES(c10::complex<double>));
|
||||
template <>
|
||||
PADDLE_API void gemm<c10::complex<float>>(
|
||||
CUDABLAS_GEMM_ARGTYPES(c10::complex<float>));
|
||||
template <>
|
||||
PADDLE_API void gemm<at::Half>(CUDABLAS_GEMM_ARGTYPES(at::Half));
|
||||
template <>
|
||||
PADDLE_API void gemm<at::BFloat16>(CUDABLAS_GEMM_ARGTYPES(at::BFloat16));
|
||||
|
||||
} // namespace at::cuda::blas
|
||||
@@ -0,0 +1,214 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
|
||||
#include <c10/core/Allocator.h>
|
||||
#include <mutex>
|
||||
|
||||
#include "paddle/phi/backends/context_pool.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_info.h"
|
||||
#include "paddle/phi/core/memory/allocation/allocator_facade.h"
|
||||
|
||||
namespace at::cuda {
|
||||
|
||||
namespace {
|
||||
|
||||
inline void ensureDeviceContextPoolInitialized() {
|
||||
static std::once_flag init_pool_once;
|
||||
std::call_once(init_pool_once, []() {
|
||||
if (phi::DeviceContextPool::IsInitialized()) {
|
||||
return;
|
||||
}
|
||||
|
||||
std::vector<phi::Place> places;
|
||||
int gpu_count = phi::backends::gpu::GetGPUDeviceCount();
|
||||
for (int device = 0; device < gpu_count; ++device) {
|
||||
places.emplace_back(phi::GPUPlace(device));
|
||||
}
|
||||
places.emplace_back(phi::CPUPlace());
|
||||
places.emplace_back(phi::GPUPinnedPlace());
|
||||
phi::DeviceContextPool::Init(places);
|
||||
});
|
||||
}
|
||||
|
||||
/// Returns the GPUContext for the current device.
|
||||
inline phi::GPUContext* getCurrentGPUContext() {
|
||||
ensureDeviceContextPoolInitialized();
|
||||
int device_id = phi::backends::gpu::GetCurrentDeviceId();
|
||||
return static_cast<phi::GPUContext*>(
|
||||
phi::DeviceContextPool::Instance().Get(phi::GPUPlace(device_id)));
|
||||
}
|
||||
|
||||
/// Frees a phi::Allocation that was released with .release() during allocate().
|
||||
static void deletePaddleCUDAAllocation(void* p) {
|
||||
delete static_cast<phi::Allocation*>(p);
|
||||
}
|
||||
|
||||
/// Adapter class that wraps Paddle's AllocatorFacade as a c10::Allocator.
|
||||
/// This provides a bridge between Paddle's allocation interface and PyTorch's
|
||||
/// c10::Allocator interface for the CUDA compatibility layer.
|
||||
class PaddleCUDAAllocatorAdapter : public c10::Allocator {
|
||||
public:
|
||||
c10::DataPtr allocate(size_t n) override {
|
||||
int device_id = phi::backends::gpu::GetCurrentDeviceId();
|
||||
if (n == 0) {
|
||||
// Return a DataPtr that carries the current CUDA device without
|
||||
// allocating any memory. Callers that probe device identity via
|
||||
// DataPtr::device() (e.g. zero-byte tensor construction) will therefore
|
||||
// observe the correct CUDA device rather than a default CPU device.
|
||||
// NOTE: For HIP/ROCm builds, PyTorch's compatibility layer still
|
||||
// exposes DeviceType::CUDA (kCUDA) rather than a separate HIP device
|
||||
// type, so we follow the same convention here.
|
||||
return c10::DataPtr(nullptr,
|
||||
nullptr,
|
||||
nullptr,
|
||||
c10::Device(c10::DeviceType::CUDA, device_id));
|
||||
}
|
||||
auto* alloc = paddle::memory::allocation::AllocatorFacade::Instance()
|
||||
.GetAllocator(phi::GPUPlace(device_id))
|
||||
.get();
|
||||
auto phi_alloc = alloc->Allocate(n);
|
||||
void* ptr = phi_alloc->ptr();
|
||||
phi::Place place = phi_alloc->place();
|
||||
// Transfer ownership of phi_alloc to the DataPtr's context.
|
||||
auto* raw_alloc = phi_alloc.release();
|
||||
return c10::DataPtr(
|
||||
ptr, raw_alloc, deletePaddleCUDAAllocation, c10::Device(place));
|
||||
}
|
||||
|
||||
void copy_data(void* dst, const void* src, size_t n) const override {
|
||||
if (n == 0) return;
|
||||
// Use GPU device-to-device copy. std::memcpy is not valid for device
|
||||
// memory; callers such as c10::Allocator::clone() rely on this method to
|
||||
// perform correct D2D copies on CUDA/HIP memory.
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(hipMemcpy(dst, src, n, hipMemcpyDeviceToDevice));
|
||||
#else
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
cudaMemcpy(dst, src, n, cudaMemcpyDeviceToDevice));
|
||||
#endif
|
||||
}
|
||||
|
||||
c10::DeleterFnPtr raw_deleter() const override {
|
||||
// allocate() returns data=device_ptr, context=phi::Allocation*, so
|
||||
// get() != get_context() and the raw_allocate/raw_deallocate API is
|
||||
// unsafe for this allocator. Returning nullptr signals that.
|
||||
return nullptr;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
CUDAContextDeviceProp* getCurrentDeviceProperties() {
|
||||
int device = phi::backends::gpu::GetCurrentDeviceId();
|
||||
return getDeviceProperties(device);
|
||||
}
|
||||
|
||||
int warp_size() { return getCurrentDeviceProperties()->warpSize; }
|
||||
|
||||
CUDAContextDeviceProp* getDeviceProperties(c10::DeviceIndex device) {
|
||||
return const_cast<CUDAContextDeviceProp*>(
|
||||
&phi::backends::gpu::GetDeviceProperties(device));
|
||||
}
|
||||
|
||||
bool canDeviceAccessPeer(c10::DeviceIndex device,
|
||||
c10::DeviceIndex peer_device) {
|
||||
int can_access = 0;
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
hipDeviceCanAccessPeer(&can_access, device, peer_device);
|
||||
#else
|
||||
cudaDeviceCanAccessPeer(&can_access, device, peer_device);
|
||||
#endif
|
||||
return can_access != 0;
|
||||
}
|
||||
|
||||
/* Handles */
|
||||
|
||||
CUDAContextSparseHandle getCurrentCUDASparseHandle() {
|
||||
return getCurrentGPUContext()->cusparse_handle();
|
||||
}
|
||||
|
||||
CUDAContextBlasHandle getCurrentCUDABlasHandle() {
|
||||
return getCurrentGPUContext()->cublas_handle();
|
||||
}
|
||||
|
||||
CUDAContextBlasLtHandle getCurrentCUDABlasLtHandle() {
|
||||
return getCurrentGPUContext()->cublaslt_handle();
|
||||
}
|
||||
|
||||
void clearCublasWorkspaces() {
|
||||
// Workspaces are owned and managed by phi::GPUContext; no explicit
|
||||
// cleanup is required here.
|
||||
}
|
||||
|
||||
WorkspaceMapWithMutex& cublas_handle_stream_to_workspace() {
|
||||
static WorkspaceMapWithMutex workspace_map;
|
||||
return workspace_map;
|
||||
}
|
||||
|
||||
WorkspaceMapWithMutex& cublaslt_handle_stream_to_workspace() {
|
||||
static WorkspaceMapWithMutex workspace_map;
|
||||
return workspace_map;
|
||||
}
|
||||
|
||||
// Default workspace size consistent with PyTorch's chosen default (32 MiB).
|
||||
static constexpr size_t kDefaultWorkspaceSize = 32UL * 1024UL * 1024UL;
|
||||
|
||||
size_t getChosenWorkspaceSize() { return kDefaultWorkspaceSize; }
|
||||
|
||||
size_t getCUDABlasLtWorkspaceSize() {
|
||||
// Probe the context with the default size and return what was actually
|
||||
// allocated.
|
||||
auto [ptr, size] =
|
||||
getCurrentGPUContext()->cublaslt_workspace(kDefaultWorkspaceSize);
|
||||
(void)ptr;
|
||||
return size;
|
||||
}
|
||||
|
||||
void* getCUDABlasLtWorkspace() {
|
||||
return getCurrentGPUContext()
|
||||
->cublaslt_workspace(kDefaultWorkspaceSize)
|
||||
.first;
|
||||
}
|
||||
|
||||
CUDAContextSolverHandle getCurrentCUDASolverDnHandle() {
|
||||
return getCurrentGPUContext()->cusolver_dn_handle();
|
||||
}
|
||||
|
||||
#if defined(USE_CUDSS)
|
||||
cudssHandle_t getCurrentCudssHandle() {
|
||||
// cudss is not yet integrated into phi::GPUContext; not implemented.
|
||||
PADDLE_THROW(
|
||||
common::errors::Unimplemented("getCurrentCudssHandle() is not "
|
||||
"implemented in the Paddle compat layer."));
|
||||
return nullptr;
|
||||
}
|
||||
#endif // USE_CUDSS
|
||||
|
||||
c10::Allocator* getCUDADeviceAllocator() {
|
||||
static PaddleCUDAAllocatorAdapter adapter;
|
||||
return &adapter;
|
||||
}
|
||||
|
||||
} // namespace at::cuda
|
||||
|
||||
#endif // PADDLE_WITH_CUDA || PADDLE_WITH_HIP
|
||||
@@ -0,0 +1,26 @@
|
||||
// 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/cuda/CUDAContextLight.h>
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <ATen/cuda/Exceptions.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#endif
|
||||
@@ -0,0 +1,136 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
// Light-weight version of CUDAContext.h with fewer transitive includes
|
||||
|
||||
// cublasLT was introduced in CUDA 10.1 but we enable only for 11.1 that also
|
||||
// added bf16 support
|
||||
|
||||
#if defined(PADDLE_WITH_HIP)
|
||||
#include <hip/hip_runtime.h>
|
||||
#elif defined(PADDLE_WITH_CUDA)
|
||||
#if defined(USE_CUDSS)
|
||||
#include <cudss.h>
|
||||
#endif
|
||||
#include <driver_types.h>
|
||||
#endif
|
||||
|
||||
#include <c10/core/Allocator.h>
|
||||
#include <c10/cuda/CUDAFunctions.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <map>
|
||||
#include <shared_mutex>
|
||||
#include <tuple>
|
||||
|
||||
#include "paddle/common/macros.h"
|
||||
#include "paddle/phi/backends/gpu/forwards.h"
|
||||
|
||||
namespace c10 {
|
||||
struct Allocator;
|
||||
}
|
||||
|
||||
namespace at::cuda {
|
||||
|
||||
#if defined(PADDLE_WITH_HIP)
|
||||
using CUDAContextDeviceProp = phi::gpuDeviceProp;
|
||||
using CUDAContextSparseHandle = phi::sparseHandle_t;
|
||||
using CUDAContextBlasHandle = phi::blasHandle_t;
|
||||
using CUDAContextBlasLtHandle = phi::blasLtHandle_t;
|
||||
using CUDAContextSolverHandle = phi::solverHandle_t;
|
||||
#elif defined(PADDLE_WITH_CUDA)
|
||||
using CUDAContextDeviceProp = cudaDeviceProp;
|
||||
using CUDAContextSparseHandle = cusparseHandle_t;
|
||||
using CUDAContextBlasHandle = cublasHandle_t;
|
||||
using CUDAContextBlasLtHandle = cublasLtHandle_t;
|
||||
using CUDAContextSolverHandle = cusolverDnHandle_t;
|
||||
#endif
|
||||
|
||||
/*
|
||||
A common CUDA interface for ATen.
|
||||
|
||||
This interface is distinct from CUDAHooks, which defines an interface that links
|
||||
to both CPU-only and CUDA builds. That interface is intended for runtime
|
||||
dispatch and should be used from files that are included in both CPU-only and
|
||||
CUDA builds.
|
||||
|
||||
CUDAContext, on the other hand, should be preferred by files only included in
|
||||
CUDA builds. It is intended to expose CUDA functionality in a consistent
|
||||
manner.
|
||||
|
||||
This means there is some overlap between the CUDAContext and CUDAHooks, but
|
||||
the choice of which to use is simple: use CUDAContext when in a CUDA-only file,
|
||||
use CUDAHooks otherwise.
|
||||
|
||||
Note that CUDAContext simply defines an interface with no associated class.
|
||||
It is expected that the modules whose functions compose this interface will
|
||||
manage their own state. There is only a single CUDA context/state.
|
||||
*/
|
||||
|
||||
/**
|
||||
* DEPRECATED: use device_count() instead
|
||||
*/
|
||||
inline int64_t getNumGPUs() { return c10::cuda::device_count(); }
|
||||
|
||||
/**
|
||||
* CUDA is available if we compiled with CUDA, and there are one or more
|
||||
* devices. If we compiled with CUDA but there is a driver problem, etc.,
|
||||
* this function will report CUDA is not available (rather than raise an error.)
|
||||
*/
|
||||
inline bool is_available() { return c10::cuda::device_count() > 0; }
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
PADDLE_API CUDAContextDeviceProp* getCurrentDeviceProperties();
|
||||
|
||||
PADDLE_API int warp_size();
|
||||
|
||||
PADDLE_API CUDAContextDeviceProp* getDeviceProperties(c10::DeviceIndex device);
|
||||
|
||||
PADDLE_API bool canDeviceAccessPeer(c10::DeviceIndex device,
|
||||
c10::DeviceIndex peer_device);
|
||||
|
||||
/* Handles */
|
||||
PADDLE_API CUDAContextSparseHandle getCurrentCUDASparseHandle();
|
||||
PADDLE_API CUDAContextBlasHandle getCurrentCUDABlasHandle();
|
||||
PADDLE_API CUDAContextBlasLtHandle getCurrentCUDABlasLtHandle();
|
||||
|
||||
PADDLE_API void clearCublasWorkspaces();
|
||||
struct WorkspaceMapWithMutex {
|
||||
std::map<std::tuple<void*, void*>, at::DataPtr> map;
|
||||
std::shared_mutex mutex;
|
||||
};
|
||||
|
||||
PADDLE_API WorkspaceMapWithMutex& cublas_handle_stream_to_workspace();
|
||||
PADDLE_API WorkspaceMapWithMutex& cublaslt_handle_stream_to_workspace();
|
||||
PADDLE_API size_t getChosenWorkspaceSize();
|
||||
PADDLE_API size_t getCUDABlasLtWorkspaceSize();
|
||||
PADDLE_API void* getCUDABlasLtWorkspace();
|
||||
|
||||
PADDLE_API CUDAContextSolverHandle getCurrentCUDASolverDnHandle();
|
||||
|
||||
#if defined(USE_CUDSS)
|
||||
PADDLE_API cudssHandle_t getCurrentCudssHandle();
|
||||
#endif
|
||||
|
||||
// Get the CUDA device allocator for the current device.
|
||||
// Returns a pointer to a c10::Allocator that allocates GPU memory.
|
||||
PADDLE_API c10::Allocator* getCUDADeviceAllocator();
|
||||
#endif
|
||||
|
||||
} // namespace at::cuda
|
||||
@@ -0,0 +1,164 @@
|
||||
// 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <c10/core/ScalarType.h>
|
||||
|
||||
#if defined(PADDLE_WITH_HIP)
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <hip/library_types.h>
|
||||
#elif defined(PADDLE_WITH_CUDA)
|
||||
#include <cuda.h>
|
||||
#include <library_types.h>
|
||||
#endif
|
||||
|
||||
namespace at::cuda {
|
||||
|
||||
#if defined(PADDLE_WITH_HIP)
|
||||
using cudaDataType = hipDataType;
|
||||
#define CUDA_R_16F HIP_R_16F
|
||||
#define CUDA_R_32F HIP_R_32F
|
||||
#define CUDA_R_64F HIP_R_64F
|
||||
#define CUDA_C_16F HIP_C_16F
|
||||
#define CUDA_C_32F HIP_C_32F
|
||||
#define CUDA_C_64F HIP_C_64F
|
||||
#define CUDA_R_8U HIP_R_8U
|
||||
#define CUDA_R_8I HIP_R_8I
|
||||
#define CUDA_R_32I HIP_R_32I
|
||||
#define CUDA_R_16I HIP_R_16I
|
||||
#define CUDA_R_64I HIP_R_64I
|
||||
#define CUDA_R_16BF HIP_R_16BF
|
||||
#define CUDA_R_8F_E4M3 HIP_R_8F_E4M3
|
||||
#define CUDA_R_8F_E5M2 HIP_R_8F_E5M2
|
||||
#elif defined(PADDLE_WITH_CUDA)
|
||||
using cudaDataType = cudaDataType;
|
||||
#endif
|
||||
|
||||
template <typename scalar_t>
|
||||
cudaDataType getCudaDataType() {
|
||||
static_assert(false && sizeof(scalar_t),
|
||||
"Cannot convert type to cudaDataType.");
|
||||
return {};
|
||||
}
|
||||
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<at::Half>() {
|
||||
return CUDA_R_16F;
|
||||
}
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<float>() {
|
||||
return CUDA_R_32F;
|
||||
}
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<double>() {
|
||||
return CUDA_R_64F;
|
||||
}
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<c10::complex<c10::Half>>() {
|
||||
return CUDA_C_16F;
|
||||
}
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<c10::complex<float>>() {
|
||||
return CUDA_C_32F;
|
||||
}
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<c10::complex<double>>() {
|
||||
return CUDA_C_64F;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<uint8_t>() {
|
||||
return CUDA_R_8U;
|
||||
}
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<int8_t>() {
|
||||
return CUDA_R_8I;
|
||||
}
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<int>() {
|
||||
return CUDA_R_32I;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<int16_t>() {
|
||||
return CUDA_R_16I;
|
||||
}
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<int64_t>() {
|
||||
return CUDA_R_64I;
|
||||
}
|
||||
template <>
|
||||
inline cudaDataType getCudaDataType<at::BFloat16>() {
|
||||
return CUDA_R_16BF;
|
||||
}
|
||||
|
||||
inline cudaDataType ScalarTypeToCudaDataType(
|
||||
const c10::ScalarType& scalar_type) {
|
||||
switch (scalar_type) {
|
||||
case c10::ScalarType::Byte:
|
||||
return CUDA_R_8U;
|
||||
case c10::ScalarType::Char:
|
||||
return CUDA_R_8I;
|
||||
case c10::ScalarType::Int:
|
||||
return CUDA_R_32I;
|
||||
case c10::ScalarType::Half:
|
||||
return CUDA_R_16F;
|
||||
case c10::ScalarType::Float:
|
||||
return CUDA_R_32F;
|
||||
case c10::ScalarType::Double:
|
||||
return CUDA_R_64F;
|
||||
// case c10::ScalarType::ComplexHalf:
|
||||
// return CUDA_C_16F;
|
||||
case c10::ScalarType::ComplexFloat:
|
||||
return CUDA_C_32F;
|
||||
case c10::ScalarType::ComplexDouble:
|
||||
return CUDA_C_64F;
|
||||
case c10::ScalarType::Short:
|
||||
return CUDA_R_16I;
|
||||
case c10::ScalarType::Long:
|
||||
return CUDA_R_64I;
|
||||
case c10::ScalarType::BFloat16:
|
||||
return CUDA_R_16BF;
|
||||
#if defined(PADDLE_WITH_HIP)
|
||||
case c10::ScalarType::Float8_e4m3fn:
|
||||
return CUDA_R_8F_E4M3;
|
||||
case c10::ScalarType::Float8_e5m2:
|
||||
return CUDA_R_8F_E5M2;
|
||||
case c10::ScalarType::Float8_e4m3fnuz:
|
||||
return HIP_R_8F_E4M3_FNUZ;
|
||||
case c10::ScalarType::Float8_e5m2fnuz:
|
||||
return HIP_R_8F_E5M2_FNUZ;
|
||||
#elif !defined(USE_ROCM) || ROCM_VERSION >= 60300
|
||||
case c10::ScalarType::Float8_e4m3fn:
|
||||
return CUDA_R_8F_E4M3;
|
||||
case c10::ScalarType::Float8_e5m2:
|
||||
return CUDA_R_8F_E5M2;
|
||||
#endif
|
||||
// #if (defined(CUDA_VERSION) && CUDA_VERSION >= 12080) ||
|
||||
// (defined(USE_ROCM) && ROCM_VERSION >= 70000)
|
||||
// case c10::ScalarType::Float4_e2m1fn_x2:
|
||||
// return CUDA_R_4F_E2M1;
|
||||
// #endif
|
||||
default:
|
||||
TORCH_INTERNAL_ASSERT(
|
||||
false, "Cannot convert ScalarType ", scalar_type, " to cudaDataType.")
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace at::cuda
|
||||
@@ -0,0 +1,204 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#if defined(PADDLE_WITH_HIP)
|
||||
#include <hip/hip_runtime.h>
|
||||
#elif defined(PADDLE_WITH_CUDA)
|
||||
#include <cuda_runtime_api.h>
|
||||
#endif
|
||||
|
||||
#include <c10/core/Device.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
|
||||
namespace at::cuda {
|
||||
|
||||
/**
|
||||
* CUDAEvent is a movable, non-copyable wrapper around CUDA events.
|
||||
* Provides compatibility with PyTorch's CUDAEvent API.
|
||||
*/
|
||||
struct CUDAEvent {
|
||||
CUDAEvent() noexcept = default;
|
||||
|
||||
explicit CUDAEvent(unsigned int flags) noexcept : flags_(flags) {}
|
||||
|
||||
~CUDAEvent() {
|
||||
if (is_created_) {
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
hipEventDestroy(event_);
|
||||
#else
|
||||
cudaEventDestroy(event_);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
CUDAEvent(const CUDAEvent&) = delete;
|
||||
CUDAEvent& operator=(const CUDAEvent&) = delete;
|
||||
|
||||
CUDAEvent(CUDAEvent&& other) noexcept { moveHelper(std::move(other)); }
|
||||
CUDAEvent& operator=(CUDAEvent&& other) noexcept {
|
||||
if (this != &other) {
|
||||
moveHelper(std::move(other));
|
||||
}
|
||||
return *this;
|
||||
}
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
operator hipEvent_t() const { return event(); }
|
||||
|
||||
hipEvent_t event() const { return event_; }
|
||||
#else
|
||||
operator cudaEvent_t() const { return event(); }
|
||||
|
||||
cudaEvent_t event() const { return event_; }
|
||||
#endif
|
||||
|
||||
bool isCreated() const { return is_created_; }
|
||||
|
||||
c10::DeviceIndex device_index() const { return device_index_; }
|
||||
|
||||
bool query() const {
|
||||
if (!is_created_) return true;
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
hipError_t err = hipEventQuery(event_);
|
||||
if (err == hipSuccess) return true;
|
||||
if (err != hipErrorNotReady) C10_CUDA_CHECK(err);
|
||||
#else
|
||||
cudaError_t err = cudaEventQuery(event_);
|
||||
if (err == cudaSuccess) return true;
|
||||
if (err != cudaErrorNotReady) C10_CUDA_CHECK(err);
|
||||
#endif
|
||||
return false;
|
||||
}
|
||||
|
||||
void record() { record(getCurrentCUDAStream()); }
|
||||
|
||||
void record(const CUDAStream& stream) {
|
||||
if (!is_created_) {
|
||||
createEvent(stream.unwrap().device_index());
|
||||
}
|
||||
TORCH_CHECK(device_index_ == stream.unwrap().device_index(),
|
||||
"Event device ",
|
||||
device_index_,
|
||||
" does not match recording stream's device ",
|
||||
stream.unwrap().device_index(),
|
||||
".");
|
||||
c10::cuda::CUDAGuard guard(device_index_);
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
C10_CUDA_CHECK(hipEventRecord(event_, stream.stream()));
|
||||
#else
|
||||
C10_CUDA_CHECK(cudaEventRecord(event_, stream.stream()));
|
||||
#endif
|
||||
}
|
||||
|
||||
void recordOnce(const CUDAStream& stream) {
|
||||
if (!was_recorded_) {
|
||||
record(stream);
|
||||
was_recorded_ = true;
|
||||
}
|
||||
}
|
||||
|
||||
void block(const CUDAStream& stream) {
|
||||
if (is_created_) {
|
||||
c10::cuda::CUDAGuard guard(stream.unwrap().device_index());
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
C10_CUDA_CHECK(hipStreamWaitEvent(stream.stream(), event_, 0));
|
||||
#else
|
||||
C10_CUDA_CHECK(cudaStreamWaitEvent(stream.stream(), event_, 0));
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
void synchronize() const {
|
||||
if (is_created_) {
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
C10_CUDA_CHECK(hipEventSynchronize(event_));
|
||||
#else
|
||||
C10_CUDA_CHECK(cudaEventSynchronize(event_));
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
float elapsed_time(const CUDAEvent& other) const {
|
||||
TORCH_CHECK(
|
||||
is_created_ && other.isCreated(),
|
||||
"Both events must be recorded before calculating elapsed time.");
|
||||
TORCH_CHECK(
|
||||
query() && other.query(),
|
||||
"Both events must be completed before calculating elapsed time.");
|
||||
float time_ms = 0;
|
||||
c10::cuda::CUDAGuard guard(device_index_);
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
C10_CUDA_CHECK(hipEventElapsedTime(&time_ms, event_, other.event_));
|
||||
#else
|
||||
C10_CUDA_CHECK(cudaEventElapsedTime(&time_ms, event_, other.event_));
|
||||
#endif
|
||||
return time_ms;
|
||||
}
|
||||
|
||||
private:
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
unsigned int flags_ = hipEventDisableTiming;
|
||||
#else
|
||||
unsigned int flags_ = cudaEventDisableTiming;
|
||||
#endif
|
||||
bool is_created_ = false;
|
||||
bool was_recorded_ = false;
|
||||
c10::DeviceIndex device_index_ = -1;
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
hipEvent_t event_{};
|
||||
#else
|
||||
cudaEvent_t event_{};
|
||||
#endif
|
||||
|
||||
void createEvent(c10::DeviceIndex device_index) {
|
||||
device_index_ = device_index;
|
||||
c10::cuda::CUDAGuard guard(device_index_);
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
C10_CUDA_CHECK(hipEventCreateWithFlags(&event_, flags_));
|
||||
#else
|
||||
C10_CUDA_CHECK(cudaEventCreateWithFlags(&event_, flags_));
|
||||
#endif
|
||||
is_created_ = true;
|
||||
}
|
||||
|
||||
void moveHelper(CUDAEvent&& other) {
|
||||
flags_ = other.flags_;
|
||||
is_created_ = std::exchange(other.is_created_, false);
|
||||
was_recorded_ = other.was_recorded_;
|
||||
device_index_ = other.device_index_;
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
event_ = std::exchange(other.event_, hipEvent_t{});
|
||||
#else
|
||||
event_ = std::exchange(other.event_, cudaEvent_t{});
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace at::cuda
|
||||
|
||||
namespace torch {
|
||||
using at::cuda::CUDAEvent;
|
||||
using at::cuda::CUDAStream;
|
||||
} // namespace torch
|
||||
@@ -0,0 +1,156 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/core/Generator.h>
|
||||
#include <ATen/cuda/PhiloxCudaState.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_info.h"
|
||||
#include "paddle/phi/core/generator.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// Forward declaration
|
||||
struct CUDAGeneratorImpl;
|
||||
|
||||
inline DeviceIndex resolve_device_index(DeviceIndex idx) {
|
||||
if (idx < 0) {
|
||||
return static_cast<DeviceIndex>(phi::backends::gpu::GetCurrentDeviceId());
|
||||
}
|
||||
return idx;
|
||||
}
|
||||
|
||||
struct CUDAGeneratorImpl : public c10::GeneratorImpl {
|
||||
explicit CUDAGeneratorImpl(
|
||||
DeviceIndex device_index = -1, // NOLINT(runtime/int)
|
||||
bool use_default_gen = true)
|
||||
: c10::GeneratorImpl(
|
||||
c10::Device(c10::kCUDA, resolve_device_index(device_index)),
|
||||
use_default_gen
|
||||
? get_default_paddle_gen(resolve_device_index(device_index))
|
||||
: create_new_paddle_gen(resolve_device_index(device_index))),
|
||||
philox_offset_per_thread_(0) {}
|
||||
|
||||
~CUDAGeneratorImpl() override = default;
|
||||
|
||||
void set_current_seed(uint64_t seed) override {
|
||||
gen_->SetCurrentSeed(seed);
|
||||
philox_offset_per_thread_ = 0;
|
||||
}
|
||||
|
||||
uint64_t current_seed() const override { return gen_->GetCurrentSeed(); }
|
||||
|
||||
uint64_t seed() override {
|
||||
auto s = gen_->Seed();
|
||||
philox_offset_per_thread_ = 0;
|
||||
return s;
|
||||
}
|
||||
|
||||
void set_offset(uint64_t offset) override {
|
||||
philox_offset_per_thread_ = offset;
|
||||
}
|
||||
|
||||
uint64_t get_offset() const override { return philox_offset_per_thread_; }
|
||||
|
||||
void set_philox_offset_per_thread(uint64_t offset) {
|
||||
philox_offset_per_thread_ = offset;
|
||||
}
|
||||
|
||||
uint64_t philox_offset_per_thread() const {
|
||||
return philox_offset_per_thread_;
|
||||
}
|
||||
|
||||
PhiloxCudaState philox_cuda_state(uint64_t increment) {
|
||||
PhiloxCudaState state(gen_->GetCurrentSeed(), philox_offset_per_thread_);
|
||||
philox_offset_per_thread_ += increment;
|
||||
return state;
|
||||
}
|
||||
|
||||
std::pair<uint64_t, uint64_t> philox_engine_inputs(uint64_t increment) {
|
||||
uint64_t offset = philox_offset_per_thread_;
|
||||
philox_offset_per_thread_ += increment;
|
||||
return {gen_->GetCurrentSeed(), offset};
|
||||
}
|
||||
|
||||
c10::intrusive_ptr<c10::GeneratorImpl> clone() const override {
|
||||
auto new_gen = std::make_shared<phi::Generator>(gen_->GetCurrentSeed());
|
||||
auto state = gen_->GetState();
|
||||
new_gen->SetState(state);
|
||||
|
||||
auto impl = c10::make_intrusive<CUDAGeneratorImpl>(
|
||||
static_cast<DeviceIndex>(device_.index()));
|
||||
impl->gen_ = new_gen;
|
||||
impl->philox_offset_per_thread_ = philox_offset_per_thread_;
|
||||
return impl;
|
||||
}
|
||||
|
||||
static c10::DeviceType device_type() { return c10::kCUDA; }
|
||||
|
||||
private:
|
||||
uint64_t philox_offset_per_thread_;
|
||||
|
||||
static std::shared_ptr<phi::Generator> get_default_paddle_gen(
|
||||
DeviceIndex device_index) {
|
||||
return phi::DefaultCUDAGenerator(static_cast<int64_t>(device_index));
|
||||
}
|
||||
|
||||
static std::shared_ptr<phi::Generator> create_new_paddle_gen(
|
||||
DeviceIndex /*device_index*/) {
|
||||
return std::make_shared<phi::Generator>();
|
||||
}
|
||||
};
|
||||
|
||||
namespace cuda {
|
||||
namespace detail {
|
||||
|
||||
inline const Generator& getDefaultCUDAGenerator(DeviceIndex device_index = -1) {
|
||||
auto idx = resolve_device_index(device_index);
|
||||
static std::vector<Generator> generators;
|
||||
static std::once_flag init_flag;
|
||||
static int64_t num_devices = 0;
|
||||
|
||||
std::call_once(init_flag, []() {
|
||||
num_devices = phi::backends::gpu::GetGPUDeviceCount();
|
||||
generators.reserve(num_devices);
|
||||
for (int64_t i = 0; i < num_devices; ++i) {
|
||||
generators.emplace_back(c10::make_intrusive<CUDAGeneratorImpl>(
|
||||
static_cast<DeviceIndex>(i), /*use_default_gen=*/true));
|
||||
}
|
||||
});
|
||||
|
||||
TORCH_CHECK(idx < static_cast<DeviceIndex>(num_devices),
|
||||
"CUDA device index out of range: ",
|
||||
idx);
|
||||
return generators[static_cast<size_t>(idx)];
|
||||
}
|
||||
|
||||
inline Generator createCUDAGenerator(DeviceIndex device_index = -1) {
|
||||
return Generator(c10::make_intrusive<CUDAGeneratorImpl>(
|
||||
device_index, /*use_default_gen=*/false));
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
} // namespace cuda
|
||||
} // namespace at
|
||||
@@ -0,0 +1,46 @@
|
||||
// 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.
|
||||
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at::detail {
|
||||
|
||||
at::Tensor empty_cuda(IntArrayRef size,
|
||||
ScalarType dtype,
|
||||
std::optional<Device> device_opt,
|
||||
std::optional<c10::MemoryFormat> memory_format_opt) {
|
||||
PD_CHECK(!(memory_format_opt.has_value() &&
|
||||
memory_format_opt.value() != c10::MemoryFormat::Contiguous),
|
||||
"`MemoryFormat` other than Contiguous is not supported now.");
|
||||
return paddle::experimental::empty(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(dtype),
|
||||
device_opt && device_opt->has_index() ? phi::GPUPlace(device_opt->index())
|
||||
: paddle::DefaultGPUPlace());
|
||||
}
|
||||
|
||||
at::Tensor empty_cuda(IntArrayRef size, const TensorOptions &options) {
|
||||
auto place = options.has_device() && options.device().has_index()
|
||||
? phi::GPUPlace(options.device().index())
|
||||
: paddle::DefaultGPUPlace();
|
||||
return paddle::experimental::empty(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype_opt().value()),
|
||||
place);
|
||||
}
|
||||
|
||||
} // namespace at::detail
|
||||
@@ -0,0 +1,32 @@
|
||||
// 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 <ATen/core/TensorBody.h>
|
||||
|
||||
#include "paddle/common/macros.h"
|
||||
|
||||
namespace at::detail {
|
||||
|
||||
using at::Tensor;
|
||||
PADDLE_API at::Tensor empty_cuda(
|
||||
IntArrayRef size,
|
||||
ScalarType dtype,
|
||||
std::optional<Device> device_opt,
|
||||
std::optional<c10::MemoryFormat> memory_format_opt);
|
||||
|
||||
PADDLE_API at::Tensor empty_cuda(IntArrayRef size,
|
||||
const TensorOptions &options);
|
||||
|
||||
} // namespace at::detail
|
||||
@@ -0,0 +1,16 @@
|
||||
// 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 <c10/util/Exception.h>
|
||||
@@ -0,0 +1,56 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
namespace at {
|
||||
|
||||
struct PhiloxCudaState {
|
||||
PhiloxCudaState() = default;
|
||||
// Called if graph capture is not underway
|
||||
PhiloxCudaState(uint64_t seed, uint64_t offset) {
|
||||
seed_.val = seed;
|
||||
offset_.val = offset;
|
||||
}
|
||||
// Called if graph capture is underway
|
||||
PhiloxCudaState(int64_t* seed,
|
||||
int64_t* offset_extragraph,
|
||||
uint64_t offset_intragraph) {
|
||||
seed_.ptr = seed;
|
||||
offset_.ptr = offset_extragraph;
|
||||
offset_intragraph_ = offset_intragraph;
|
||||
captured_ = true;
|
||||
}
|
||||
|
||||
// Public members, directly accessible by at::cuda::philox::unpack.
|
||||
// If we made them private with getters/setters, the getters/setters
|
||||
// would have to be __device__, and we can't declare __device__ in ATen.
|
||||
union Payload {
|
||||
uint64_t val;
|
||||
int64_t* ptr;
|
||||
};
|
||||
|
||||
Payload seed_{};
|
||||
Payload offset_{};
|
||||
uint64_t offset_intragraph_ = 0;
|
||||
bool captured_ = false;
|
||||
};
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,49 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/cuda/PhiloxCudaState.h>
|
||||
#include <tuple>
|
||||
|
||||
namespace at::cuda::philox {
|
||||
|
||||
// In-kernel call to retrieve philox seed and offset from a PhiloxCudaState
|
||||
// instance whether that instance was created with graph capture underway or
|
||||
// not. See Note [CUDA Graph-safe RNG states].
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
__host__ __device__ __forceinline__ std::tuple<uint64_t, uint64_t> unpack(
|
||||
at::PhiloxCudaState arg) {
|
||||
#else
|
||||
inline std::tuple<uint64_t, uint64_t> unpack(at::PhiloxCudaState arg) {
|
||||
#endif
|
||||
if (arg.captured_) {
|
||||
// static_cast avoids "warning: invalid narrowing conversion from "long" to
|
||||
// "unsigned long".
|
||||
// *(arg.offset_.ptr) is a broadcast load of a single int64_t to the entire
|
||||
// kernel. For most threads' reads it will hit in cache, so it shouldn't
|
||||
// hurt performance.
|
||||
return std::make_tuple(
|
||||
static_cast<uint64_t>(*arg.seed_.ptr),
|
||||
static_cast<uint64_t>(*(arg.offset_.ptr) + arg.offset_intragraph_));
|
||||
} else {
|
||||
return std::make_tuple(arg.seed_.val, arg.offset_.val);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace at::cuda::philox
|
||||
@@ -0,0 +1,75 @@
|
||||
// Copyright (c) 2026 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 <ATen/AccumulateType.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <limits>
|
||||
|
||||
namespace at::native {
|
||||
|
||||
inline void arange_check_bounds(const c10::Scalar& start,
|
||||
const c10::Scalar& end,
|
||||
const c10::Scalar& step) {
|
||||
// use double precision for validation to avoid precision issues
|
||||
double dstart = start.to<double>();
|
||||
double dend = end.to<double>();
|
||||
double dstep = step.to<double>();
|
||||
|
||||
TORCH_CHECK(dstep > 0 || dstep < 0, "step must be nonzero");
|
||||
TORCH_CHECK(std::isfinite(dstart) && std::isfinite(dend),
|
||||
"unsupported range: ",
|
||||
dstart,
|
||||
" -> ",
|
||||
dend);
|
||||
TORCH_CHECK(
|
||||
((dstep > 0) && (dend >= dstart)) || ((dstep < 0) && (dend <= dstart)),
|
||||
"upper bound and lower bound inconsistent with step sign");
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
int64_t compute_arange_size(const Scalar& start,
|
||||
const Scalar& end,
|
||||
const Scalar& step) {
|
||||
arange_check_bounds(start, end, step);
|
||||
|
||||
// we use double precision for (start - end) / step
|
||||
// to compute size_d for consistency across devices.
|
||||
// The problem with using accscalar_t is that accscalar_t might be float32 on
|
||||
// gpu for a float32 scalar_t, but double on cpu for the same, and the
|
||||
// effective output size starts differing on CPU vs GPU because of precision
|
||||
// issues, which we dont want. the corner-case we do want to take into account
|
||||
// is int64_t, which has higher precision than double
|
||||
double size_d;
|
||||
if constexpr (std::is_same_v<scalar_t, int64_t>) {
|
||||
using accscalar_t = at::acc_type<scalar_t, false>;
|
||||
auto xstart = start.to<accscalar_t>();
|
||||
auto xend = end.to<accscalar_t>();
|
||||
auto xstep = step.to<accscalar_t>();
|
||||
int64_t sgn = (xstep > 0) - (xstep < 0);
|
||||
size_d = std::ceil((xend - xstart + xstep - sgn) / xstep);
|
||||
} else {
|
||||
size_d =
|
||||
std::ceil((end.to<double>() - start.to<double>()) / step.to<double>());
|
||||
}
|
||||
|
||||
TORCH_CHECK(size_d >= 0 && size_d <= static_cast<double>(
|
||||
std::numeric_limits<int64_t>::max()),
|
||||
"invalid size, possible overflow?");
|
||||
|
||||
return static_cast<int64_t>(size_d);
|
||||
}
|
||||
|
||||
} // namespace at::native
|
||||
@@ -0,0 +1,19 @@
|
||||
// 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
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
@@ -0,0 +1,77 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
|
||||
#include "paddle/phi/api/include/tensor.h"
|
||||
#include "paddle/phi/common/data_type.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
/// Extracts a scalar value from a single-element dense tensor.
|
||||
/// Mirrors PyTorch's at::_local_scalar_dense: copies the tensor to CPU if
|
||||
/// needed, then reads the first element according to its dtype.
|
||||
inline at::Scalar _local_scalar_dense(const at::Tensor& self) {
|
||||
PD_CHECK(self.numel() > 0, "_local_scalar_dense: Empty tensor not supported");
|
||||
|
||||
// Move to CPU if necessary (for compatibility with PyTorch behavior)
|
||||
const PaddleTensor& inner = self._PD_GetInner();
|
||||
PaddleTensor cpu_tensor = inner;
|
||||
if (!phi::is_cpu_place(inner.place())) {
|
||||
PaddlePlace place(phi::AllocationType::CPU);
|
||||
cpu_tensor = inner.copy_to(place, /*blocking=*/true);
|
||||
}
|
||||
|
||||
auto dtype = cpu_tensor.dtype();
|
||||
switch (dtype) {
|
||||
case phi::DataType::FLOAT32:
|
||||
return at::Scalar(*(cpu_tensor.data<float>()));
|
||||
case phi::DataType::FLOAT64:
|
||||
return at::Scalar(*(cpu_tensor.data<double>()));
|
||||
case phi::DataType::FLOAT16:
|
||||
return at::Scalar(
|
||||
static_cast<float>(*(cpu_tensor.data<phi::dtype::float16>())));
|
||||
case phi::DataType::BFLOAT16:
|
||||
return at::Scalar(
|
||||
static_cast<float>(*(cpu_tensor.data<phi::dtype::bfloat16>())));
|
||||
case phi::DataType::INT8:
|
||||
return at::Scalar(*(cpu_tensor.data<int8_t>()));
|
||||
case phi::DataType::INT16:
|
||||
return at::Scalar(*(cpu_tensor.data<int16_t>()));
|
||||
case phi::DataType::INT32:
|
||||
return at::Scalar(*(cpu_tensor.data<int32_t>()));
|
||||
case phi::DataType::INT64:
|
||||
return at::Scalar(*(cpu_tensor.data<int64_t>()));
|
||||
case phi::DataType::UINT8:
|
||||
return at::Scalar(*(cpu_tensor.data<uint8_t>()));
|
||||
case phi::DataType::BOOL:
|
||||
return at::Scalar(*(cpu_tensor.data<bool>()));
|
||||
case phi::DataType::COMPLEX64:
|
||||
return at::Scalar(*(cpu_tensor.data<phi::dtype::complex<float>>()));
|
||||
case phi::DataType::COMPLEX128:
|
||||
return at::Scalar(*(cpu_tensor.data<phi::dtype::complex<double>>()));
|
||||
default:
|
||||
PD_THROW("_local_scalar_dense: Unsupported data type");
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,44 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include "paddle/phi/core/sparse_coo_tensor.h"
|
||||
#include "paddle/phi/core/sparse_csr_tensor.h"
|
||||
|
||||
namespace at {} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline int64_t Tensor::_nnz() const {
|
||||
PD_CHECK(this->is_sparse(),
|
||||
"_nnz expected sparse tensor layout but got ",
|
||||
layout());
|
||||
if (tensor_.layout() == common::DataLayout::SPARSE_COO) {
|
||||
auto sparse_coo_tensor =
|
||||
std::dynamic_pointer_cast<phi::SparseCooTensor>(tensor_.impl());
|
||||
PD_CHECK(sparse_coo_tensor != nullptr,
|
||||
"_nnz: failed to cast tensor impl to SparseCooTensor");
|
||||
return sparse_coo_tensor->nnz();
|
||||
} else {
|
||||
auto sparse_csr_tensor =
|
||||
std::dynamic_pointer_cast<phi::SparseCsrTensor>(tensor_.impl());
|
||||
PD_CHECK(sparse_csr_tensor != nullptr,
|
||||
"_nnz: failed to cast tensor impl to SparseCsrTensor");
|
||||
return sparse_csr_tensor->nnz();
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,41 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include "paddle/phi/core/sparse_coo_tensor.h"
|
||||
#include "paddle/phi/core/sparse_csr_tensor.h"
|
||||
|
||||
namespace at {} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::_values() const {
|
||||
PD_CHECK(this->is_sparse(),
|
||||
"_values expected sparse tensor layout but got ",
|
||||
layout());
|
||||
if (tensor_.layout() == common::DataLayout::SPARSE_COO) {
|
||||
auto sparse_coo_tensor =
|
||||
std::dynamic_pointer_cast<phi::SparseCooTensor>(tensor_.impl());
|
||||
PD_CHECK(sparse_coo_tensor != nullptr,
|
||||
"_values: failed to cast tensor impl to SparseCooTensor");
|
||||
return paddle::Tensor(
|
||||
std::make_shared<phi::DenseTensor>(sparse_coo_tensor->values()));
|
||||
} else {
|
||||
PD_THROW("_values is not implemented for SparseCsr tensors");
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,57 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor abs(const at::Tensor& self) {
|
||||
if (!self.is_contiguous()) {
|
||||
phi::enforce::ThrowWarnInternal(
|
||||
"at::abs: input tensor is non-contiguous. PyTorch and Paddle handle "
|
||||
"non-contiguous tensors differently, which may produce logically "
|
||||
"incorrect results even though the code is syntactically valid. "
|
||||
"See https://github.com/PaddlePaddle/Paddle/pull/78099 for details.");
|
||||
}
|
||||
return paddle::experimental::abs(self._PD_GetInner());
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::abs() const { return at::abs(*this); }
|
||||
|
||||
inline at::Tensor& Tensor::abs_() const {
|
||||
if (!is_contiguous()) {
|
||||
phi::enforce::ThrowWarnInternal(
|
||||
"Tensor::abs_: tensor is non-contiguous. PyTorch and Paddle handle "
|
||||
"non-contiguous tensors differently, which may produce logically "
|
||||
"incorrect results even though the code is syntactically valid. "
|
||||
"See https://github.com/PaddlePaddle/Paddle/pull/78099 for details.");
|
||||
}
|
||||
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
|
||||
paddle::experimental::abs_(inner);
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,63 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/util/OptionalArrayRef.h>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// all: Check if all elements are true (non-zero)
|
||||
// Version 1: all() - check all elements in the tensor
|
||||
inline at::Tensor all(const at::Tensor& self) {
|
||||
return paddle::experimental::all(self._PD_GetInner(), {}, false);
|
||||
}
|
||||
|
||||
// Version 2: all(dim, keepdim) - check along a specific dimension
|
||||
inline at::Tensor all(const at::Tensor& self,
|
||||
int64_t dim,
|
||||
bool keepdim = false) {
|
||||
return paddle::experimental::all(self._PD_GetInner(), {dim}, keepdim);
|
||||
}
|
||||
|
||||
// Version 3: all(dim, keepdim) - check along optional dimensions
|
||||
inline at::Tensor all(const at::Tensor& self,
|
||||
at::OptionalIntArrayRef dim,
|
||||
bool keepdim = false) {
|
||||
std::vector<int64_t> axis_vec;
|
||||
if (dim.has_value()) {
|
||||
axis_vec.assign(dim.value().begin(), dim.value().end());
|
||||
}
|
||||
return paddle::experimental::all(self._PD_GetInner(), axis_vec, keepdim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
// Tensor member function implementations
|
||||
inline at::Tensor Tensor::all() const { return at::all(*this); }
|
||||
|
||||
inline at::Tensor Tensor::all(int64_t dim, bool keepdim) const {
|
||||
return at::all(*this, dim, keepdim);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::all(at::OptionalIntArrayRef dim, bool keepdim) const {
|
||||
return at::all(*this, dim, keepdim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,57 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/data_type.h"
|
||||
#include "paddle/phi/common/scalar.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// allclose: Check if two tensors are close to each other
|
||||
inline bool allclose(const at::Tensor& self,
|
||||
const at::Tensor& other,
|
||||
double rtol = 1e-05,
|
||||
double atol = 1e-08,
|
||||
bool equal_nan = false) {
|
||||
// Paddle's allclose returns a Tensor, but PyTorch's allclose returns bool.
|
||||
// The allclose kernel always sets output dtype to phi::DataType::BOOL via
|
||||
// kernel->OutputAt(0).SetDataType(phi::DataType::BOOL), so we read BOOL
|
||||
// directly.
|
||||
PaddleTensor result = paddle::experimental::allclose(self._PD_GetInner(),
|
||||
other._PD_GetInner(),
|
||||
phi::Scalar(rtol),
|
||||
phi::Scalar(atol),
|
||||
equal_nan);
|
||||
auto* result_tensor =
|
||||
std::dynamic_pointer_cast<phi::DenseTensor>(result.impl()).get();
|
||||
return *result_tensor->data<bool>();
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
// Tensor member function implementation
|
||||
inline bool Tensor::allclose(const at::Tensor& other,
|
||||
double rtol,
|
||||
double atol,
|
||||
bool equal_nan) const {
|
||||
return at::allclose(*this, other, rtol, atol, equal_nan);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,61 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/util/OptionalArrayRef.h>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// any - free functions
|
||||
inline Tensor any(const Tensor& self, int64_t dim, bool keepdim = false) {
|
||||
return paddle::experimental::any(self._PD_GetInner(), {dim}, keepdim);
|
||||
}
|
||||
|
||||
inline Tensor any(const Tensor& self,
|
||||
at::OptionalIntArrayRef dim,
|
||||
bool keepdim = false) {
|
||||
std::vector<int64_t> dims_vec;
|
||||
if (dim.has_value() && dim.value().size() > 0) {
|
||||
dims_vec.assign(dim.value().begin(), dim.value().end());
|
||||
}
|
||||
return paddle::experimental::any(self._PD_GetInner(), dims_vec, keepdim);
|
||||
}
|
||||
|
||||
inline Tensor any(const Tensor& self) {
|
||||
return paddle::experimental::any(self._PD_GetInner());
|
||||
}
|
||||
|
||||
// any - member function implementations
|
||||
inline Tensor Tensor::any(int64_t dim, bool keepdim) const {
|
||||
return paddle::experimental::any(_PD_GetInner(), {dim}, keepdim);
|
||||
}
|
||||
|
||||
inline Tensor Tensor::any(at::OptionalIntArrayRef dim, bool keepdim) const {
|
||||
std::vector<int64_t> dims_vec;
|
||||
if (dim.has_value() && dim.value().size() > 0) {
|
||||
dims_vec.assign(dim.value().begin(), dim.value().end());
|
||||
}
|
||||
return paddle::experimental::any(_PD_GetInner(), dims_vec, keepdim);
|
||||
}
|
||||
|
||||
inline Tensor Tensor::any() const {
|
||||
return paddle::experimental::any(_PD_GetInner());
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,157 @@
|
||||
// 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <ATen/native/RangeUtils.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/pinned_place.h>
|
||||
#include <optional>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
namespace detail {
|
||||
|
||||
inline bool _PD_IsIntegralArangeScalar(const at::Scalar& scalar) {
|
||||
switch (scalar.dtype()) {
|
||||
case phi::DataType::BOOL:
|
||||
case phi::DataType::UINT8:
|
||||
case phi::DataType::INT8:
|
||||
case phi::DataType::UINT16:
|
||||
case phi::DataType::INT16:
|
||||
case phi::DataType::UINT32:
|
||||
case phi::DataType::INT32:
|
||||
case phi::DataType::UINT64:
|
||||
case phi::DataType::INT64:
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
inline at::ScalarType _PD_ResolveArangeDtype(const at::Scalar& start,
|
||||
const at::Scalar& end,
|
||||
const at::Scalar& step,
|
||||
const at::TensorOptions& options) {
|
||||
if (options.has_dtype()) {
|
||||
return options.dtype().toScalarType();
|
||||
}
|
||||
if (_PD_IsIntegralArangeScalar(start) && _PD_IsIntegralArangeScalar(end) &&
|
||||
_PD_IsIntegralArangeScalar(step)) {
|
||||
return at::kLong;
|
||||
}
|
||||
return c10::get_default_dtype_as_scalartype();
|
||||
}
|
||||
|
||||
inline paddle::Tensor _PD_MakeArangeScalarTensor(const at::Scalar& scalar,
|
||||
phi::DataType dtype) {
|
||||
return paddle::experimental::full({}, scalar, dtype, phi::CPUPlace());
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
inline at::Tensor arange(const at::Scalar& start,
|
||||
const at::Scalar& end,
|
||||
const at::Scalar& step,
|
||||
at::TensorOptions options = {}) {
|
||||
// Match PyTorch: step must be non-zero and consistent with (end - start).
|
||||
at::native::arange_check_bounds(start, end, step);
|
||||
auto dtype = detail::_PD_ResolveArangeDtype(start, end, step, options);
|
||||
auto pd_dtype = compat::_PD_AtenScalarTypeToPhiDataType(dtype);
|
||||
if (options.pinned_memory()) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !options.device().is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
phi::Place base_place = options._PD_GetPlace();
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
|
||||
auto dense = paddle::experimental::arange(
|
||||
detail::_PD_MakeArangeScalarTensor(start, pd_dtype),
|
||||
detail::_PD_MakeArangeScalarTensor(end, pd_dtype),
|
||||
detail::_PD_MakeArangeScalarTensor(step, pd_dtype),
|
||||
pd_dtype,
|
||||
phi::CPUPlace());
|
||||
return dense.copy_to(pinned_place, /*blocking=*/true);
|
||||
}
|
||||
return paddle::experimental::arange(
|
||||
detail::_PD_MakeArangeScalarTensor(start, pd_dtype),
|
||||
detail::_PD_MakeArangeScalarTensor(end, pd_dtype),
|
||||
detail::_PD_MakeArangeScalarTensor(step, pd_dtype),
|
||||
pd_dtype,
|
||||
options._PD_GetPlace());
|
||||
}
|
||||
|
||||
inline at::Tensor arange(const at::Scalar& end,
|
||||
at::TensorOptions options = {}) {
|
||||
return arange(/*start=*/0, end, /*step=*/1, options);
|
||||
}
|
||||
|
||||
inline at::Tensor arange(const at::Scalar& end,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) {
|
||||
auto options = at::TensorOptions()
|
||||
.dtype(dtype)
|
||||
.layout(layout)
|
||||
.device(device)
|
||||
.pinned_memory(pin_memory);
|
||||
return arange(/*start=*/0, end, /*step=*/1, options);
|
||||
}
|
||||
|
||||
inline at::Tensor arange(const at::Scalar& start,
|
||||
const at::Scalar& end,
|
||||
at::TensorOptions options = {}) {
|
||||
return arange(start, end, /*step=*/1, options);
|
||||
}
|
||||
|
||||
inline at::Tensor arange(const at::Scalar& start,
|
||||
const at::Scalar& end,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) {
|
||||
auto options = at::TensorOptions()
|
||||
.dtype(dtype)
|
||||
.layout(layout)
|
||||
.device(device)
|
||||
.pinned_memory(pin_memory);
|
||||
return arange(start, end, /*step=*/1, options);
|
||||
}
|
||||
|
||||
inline at::Tensor arange(const at::Scalar& start,
|
||||
const at::Scalar& end,
|
||||
const at::Scalar& step,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) {
|
||||
auto options = at::TensorOptions()
|
||||
.dtype(dtype)
|
||||
.layout(layout)
|
||||
.device(device)
|
||||
.pinned_memory(pin_memory);
|
||||
return arange(start, end, step, options);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,113 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/util/ArrayRef.h>
|
||||
#include <optional>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/common/ddim.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// as_strided: Create a tensor view with custom size, stride, and storage_offset
|
||||
inline at::Tensor Tensor::as_strided(
|
||||
at::IntArrayRef size,
|
||||
at::IntArrayRef stride,
|
||||
::std::optional<int64_t> storage_offset) const {
|
||||
// Materialize the compat StorageHolderView before creating the view so
|
||||
// aliasing tensors share one StorageImpl and observe later resize_ growth.
|
||||
(void)this->storage();
|
||||
auto src_impl = tensor_.impl();
|
||||
auto* src_tensor =
|
||||
std::dynamic_pointer_cast<phi::DenseTensor>(src_impl).get();
|
||||
if (!src_tensor) {
|
||||
PD_THROW("as_strided: tensor must be a DenseTensor");
|
||||
}
|
||||
// Create new meta with desired shape and strides first
|
||||
std::vector<int64_t> size_vec(size.begin(), size.end());
|
||||
std::vector<int64_t> stride_vec(stride.begin(), stride.end());
|
||||
|
||||
// Create new DenseTensor with correct meta, then share data
|
||||
// We need to create a temporary DenseTensor with the right meta
|
||||
// because ShareDataWith copies the source meta which we don't want
|
||||
auto new_tensor = std::make_shared<phi::DenseTensor>();
|
||||
|
||||
// First, set up the holder by sharing data (this copies src meta, we'll
|
||||
// override)
|
||||
new_tensor->ShareDataWith(*src_tensor);
|
||||
|
||||
// Now create the correct meta with new shape/strides
|
||||
phi::DenseTensorMeta meta(src_tensor->dtype(),
|
||||
common::make_ddim(size_vec),
|
||||
common::make_ddim(stride_vec));
|
||||
// Calculate offset in bytes
|
||||
int64_t offset = storage_offset.has_value() ? storage_offset.value() : 0;
|
||||
meta.offset = src_tensor->meta().offset +
|
||||
static_cast<size_t>(offset) * phi::SizeOf(src_tensor->dtype());
|
||||
new_tensor->set_meta(meta);
|
||||
PaddleTensor result;
|
||||
result.set_impl(new_tensor);
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
// as_strided_: Inplace version
|
||||
inline const at::Tensor& Tensor::as_strided_(
|
||||
at::IntArrayRef size,
|
||||
at::IntArrayRef stride,
|
||||
::std::optional<int64_t> storage_offset) const {
|
||||
// Keep inplace metadata-only view rewrites attached to the same compat
|
||||
// storage as the original tensor.
|
||||
(void)this->storage();
|
||||
auto src_impl = tensor_.impl();
|
||||
auto* src_tensor =
|
||||
std::dynamic_pointer_cast<phi::DenseTensor>(src_impl).get();
|
||||
if (!src_tensor) {
|
||||
PD_THROW("as_strided_: tensor must be a DenseTensor");
|
||||
}
|
||||
std::vector<int64_t> size_vec(size.begin(), size.end());
|
||||
std::vector<int64_t> stride_vec(stride.begin(), stride.end());
|
||||
// Use set_meta instead of Resize + set_strides to avoid contiguous check
|
||||
phi::DenseTensorMeta meta(src_tensor->dtype(),
|
||||
common::make_ddim(size_vec),
|
||||
common::make_ddim(stride_vec));
|
||||
meta.layout = src_tensor->layout();
|
||||
int64_t offset = storage_offset.has_value() ? storage_offset.value() : 0;
|
||||
meta.offset = src_tensor->meta().offset +
|
||||
static_cast<size_t>(offset) * phi::SizeOf(src_tensor->dtype());
|
||||
src_tensor->set_meta(meta);
|
||||
return *this;
|
||||
}
|
||||
|
||||
// as_strided_scatter: Scatter src into a strided view
|
||||
// Returns a new tensor (copy of self) with the strided window filled by src.
|
||||
// The original tensor is NOT modified.
|
||||
inline at::Tensor Tensor::as_strided_scatter(
|
||||
const at::Tensor& src,
|
||||
at::IntArrayRef size,
|
||||
at::IntArrayRef stride,
|
||||
::std::optional<int64_t> storage_offset) const {
|
||||
// Clone self to an independent copy so the original tensor is left unchanged
|
||||
PaddleTensor self_copy = tensor_.copy_to(tensor_.place(), /*blocking=*/true);
|
||||
at::Tensor copy_tensor(self_copy);
|
||||
at::Tensor strided_view =
|
||||
copy_tensor.as_strided(size, stride, storage_offset);
|
||||
strided_view.copy_(src);
|
||||
return copy_tensor;
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,34 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <optional>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor cat(const at::ITensorListRef& tensors, int64_t dim = 0) {
|
||||
std::vector<paddle::Tensor> pd_tensors;
|
||||
pd_tensors.reserve(tensors.size());
|
||||
for (const auto& t : tensors) {
|
||||
pd_tensors.push_back(t._PD_GetInner());
|
||||
}
|
||||
return paddle::experimental::concat(pd_tensors, dim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,98 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// chunk - splits tensor into chunks
|
||||
inline std::vector<Tensor> chunk(const Tensor& self,
|
||||
int64_t chunks,
|
||||
int64_t dim = 0) {
|
||||
if (chunks <= 0) {
|
||||
PD_THROW("chunk expects chunks to be greater than 0, got ", chunks);
|
||||
}
|
||||
|
||||
std::vector<Tensor> result;
|
||||
paddle::Tensor pd_tensor = self._PD_GetInner();
|
||||
|
||||
int64_t rank = static_cast<int64_t>(pd_tensor.dims().size());
|
||||
if (rank == 0) {
|
||||
PD_THROW("chunk expects at least a 1-dimensional tensor");
|
||||
}
|
||||
|
||||
int64_t original_dim = dim;
|
||||
if (dim < 0) {
|
||||
dim += rank;
|
||||
}
|
||||
if (dim < 0 || dim >= rank) {
|
||||
PD_THROW("Dimension out of range (expected to be in range of [",
|
||||
-rank,
|
||||
", ",
|
||||
rank - 1,
|
||||
"], but got ",
|
||||
original_dim,
|
||||
")");
|
||||
}
|
||||
|
||||
int64_t dim_size = pd_tensor.dims()[dim];
|
||||
|
||||
if (dim_size == 0) {
|
||||
for (int64_t i = 0; i < chunks; ++i) {
|
||||
auto chunk_tensor =
|
||||
paddle::experimental::slice(pd_tensor, {dim}, {0}, {0}, {1}, {});
|
||||
result.push_back(Tensor(chunk_tensor));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// PyTorch returns at most 'dim_size' non-empty chunks when chunks > dim_size
|
||||
if (chunks > dim_size) {
|
||||
chunks = dim_size;
|
||||
}
|
||||
|
||||
int64_t chunk_size = (dim_size + chunks - 1) / chunks;
|
||||
int64_t remaining = dim_size;
|
||||
|
||||
for (int64_t i = 0; i < chunks && remaining > 0; ++i) {
|
||||
int64_t current_chunk_size = std::min(chunk_size, remaining);
|
||||
auto chunk_tensor =
|
||||
paddle::experimental::slice(pd_tensor,
|
||||
{dim},
|
||||
{i * chunk_size},
|
||||
{i * chunk_size + current_chunk_size},
|
||||
{1},
|
||||
{});
|
||||
result.push_back(Tensor(chunk_tensor));
|
||||
remaining -= current_chunk_size;
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
// Member function: Tensor::chunk
|
||||
inline std::vector<Tensor> Tensor::chunk(int64_t chunks, int64_t dim) const {
|
||||
return at::chunk(*this, chunks, dim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,212 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <ATen/core/TensorBase.h>
|
||||
#include <ATen/ops/full.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <limits>
|
||||
#include <optional>
|
||||
#include "paddle/phi/api/include/tensor.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// Helper function implementations
|
||||
namespace detail {
|
||||
inline at::Scalar get_default_min_value(c10::ScalarType dtype) {
|
||||
switch (dtype) {
|
||||
case c10::ScalarType::Byte:
|
||||
return at::Scalar(static_cast<uint8_t>(0));
|
||||
case c10::ScalarType::Char:
|
||||
return at::Scalar(std::numeric_limits<int8_t>::lowest());
|
||||
case c10::ScalarType::Short:
|
||||
return at::Scalar(std::numeric_limits<int16_t>::lowest());
|
||||
case c10::ScalarType::Int:
|
||||
return at::Scalar(std::numeric_limits<int32_t>::lowest());
|
||||
case c10::ScalarType::Long:
|
||||
return at::Scalar(std::numeric_limits<int64_t>::lowest());
|
||||
case c10::ScalarType::UInt16:
|
||||
return at::Scalar(static_cast<uint16_t>(0));
|
||||
case c10::ScalarType::UInt32:
|
||||
return at::Scalar(static_cast<uint32_t>(0));
|
||||
case c10::ScalarType::UInt64:
|
||||
return at::Scalar(static_cast<uint64_t>(0));
|
||||
case c10::ScalarType::Half:
|
||||
return at::Scalar(-std::numeric_limits<float>::infinity());
|
||||
case c10::ScalarType::Float:
|
||||
return at::Scalar(-std::numeric_limits<float>::infinity());
|
||||
case c10::ScalarType::Double:
|
||||
return at::Scalar(-std::numeric_limits<double>::infinity());
|
||||
case c10::ScalarType::BFloat16:
|
||||
return at::Scalar(-std::numeric_limits<float>::infinity());
|
||||
case c10::ScalarType::Bool:
|
||||
return at::Scalar(false);
|
||||
default:
|
||||
return at::Scalar(-std::numeric_limits<double>::infinity());
|
||||
}
|
||||
}
|
||||
|
||||
inline at::Scalar get_default_max_value(c10::ScalarType dtype) {
|
||||
switch (dtype) {
|
||||
case c10::ScalarType::Byte:
|
||||
return at::Scalar(std::numeric_limits<uint8_t>::max());
|
||||
case c10::ScalarType::Char:
|
||||
return at::Scalar(std::numeric_limits<int8_t>::max());
|
||||
case c10::ScalarType::Short:
|
||||
return at::Scalar(std::numeric_limits<int16_t>::max());
|
||||
case c10::ScalarType::Int:
|
||||
return at::Scalar(std::numeric_limits<int32_t>::max());
|
||||
case c10::ScalarType::Long:
|
||||
return at::Scalar(std::numeric_limits<int64_t>::max());
|
||||
case c10::ScalarType::UInt16:
|
||||
return at::Scalar(std::numeric_limits<uint16_t>::max());
|
||||
case c10::ScalarType::UInt32:
|
||||
return at::Scalar(std::numeric_limits<uint32_t>::max());
|
||||
case c10::ScalarType::UInt64:
|
||||
return at::Scalar(std::numeric_limits<uint64_t>::max());
|
||||
case c10::ScalarType::Half:
|
||||
return at::Scalar(std::numeric_limits<float>::infinity());
|
||||
case c10::ScalarType::Float:
|
||||
return at::Scalar(std::numeric_limits<float>::infinity());
|
||||
case c10::ScalarType::Double:
|
||||
return at::Scalar(std::numeric_limits<double>::infinity());
|
||||
case c10::ScalarType::BFloat16:
|
||||
return at::Scalar(std::numeric_limits<float>::infinity());
|
||||
case c10::ScalarType::Bool:
|
||||
return at::Scalar(true);
|
||||
default:
|
||||
return at::Scalar(std::numeric_limits<double>::infinity());
|
||||
}
|
||||
}
|
||||
} // namespace detail
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::clamp(const ::std::optional<at::Scalar>& min,
|
||||
const ::std::optional<at::Scalar>& max) const {
|
||||
// Handle cases where min or max is nullopt - don't apply that bound
|
||||
if (min.has_value() && !max.has_value()) {
|
||||
// Only min is specified - use clamp_min
|
||||
return clamp_min(min.value());
|
||||
} else if (!min.has_value() && max.has_value()) {
|
||||
// Only max is specified - use clamp_max
|
||||
return clamp_max(max.value());
|
||||
} else if (!min.has_value() && !max.has_value()) {
|
||||
// Neither specified - return copy of tensor
|
||||
return *this;
|
||||
}
|
||||
// Both specified - apply full clamp
|
||||
return Tensor(paddle::experimental::clip(tensor_, min.value(), max.value()));
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::clamp(const ::std::optional<at::Tensor>& min,
|
||||
const ::std::optional<at::Tensor>& max) const {
|
||||
PaddleTensor result = tensor_;
|
||||
if (min.has_value()) {
|
||||
result = paddle::experimental::maximum(result, min.value()._PD_GetInner());
|
||||
}
|
||||
if (max.has_value()) {
|
||||
result = paddle::experimental::minimum(result, max.value()._PD_GetInner());
|
||||
}
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::clamp_(
|
||||
const ::std::optional<at::Scalar>& min,
|
||||
const ::std::optional<at::Scalar>& max) const {
|
||||
// Handle cases where min or max is nullopt - don't apply that bound
|
||||
if (min.has_value() && !max.has_value()) {
|
||||
// Only min is specified - use clamp_min_
|
||||
return clamp_min_(min.value());
|
||||
} else if (!min.has_value() && max.has_value()) {
|
||||
// Only max is specified - use clamp_max_
|
||||
return clamp_max_(max.value());
|
||||
} else if (!min.has_value() && !max.has_value()) {
|
||||
// Neither specified - nothing to do
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
// Both specified - apply full clamp
|
||||
paddle::experimental::clip_(
|
||||
const_cast<PaddleTensor&>(tensor_), min.value(), max.value());
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::clamp_(
|
||||
const ::std::optional<at::Tensor>& min,
|
||||
const ::std::optional<at::Tensor>& max) const {
|
||||
if (min.has_value()) {
|
||||
PaddleTensor temp =
|
||||
paddle::experimental::maximum(tensor_, min.value()._PD_GetInner());
|
||||
const_cast<PaddleTensor&>(tensor_) = temp;
|
||||
}
|
||||
if (max.has_value()) {
|
||||
PaddleTensor temp =
|
||||
paddle::experimental::minimum(tensor_, max.value()._PD_GetInner());
|
||||
const_cast<PaddleTensor&>(tensor_) = temp;
|
||||
}
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::clamp_max(const at::Scalar& max) const {
|
||||
// Create a tensor with the same shape filled with the max value
|
||||
at::Tensor max_tensor = at::full(tensor_.shape(), max, {});
|
||||
return clamp_max(max_tensor);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::clamp_max(const at::Tensor& max) const {
|
||||
return Tensor(paddle::experimental::minimum(tensor_, max._PD_GetInner()));
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::clamp_max_(const at::Scalar& max) const {
|
||||
// Create a tensor with the same shape filled with the max value
|
||||
at::Tensor max_tensor = at::full(tensor_.shape(), max, {});
|
||||
return clamp_max_(max_tensor);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::clamp_max_(const at::Tensor& max) const {
|
||||
PaddleTensor temp =
|
||||
paddle::experimental::minimum(tensor_, max._PD_GetInner());
|
||||
const_cast<PaddleTensor&>(tensor_) = temp;
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::clamp_min(const at::Scalar& min) const {
|
||||
// Create a tensor with the same shape filled with the min value
|
||||
at::Tensor min_tensor = at::full(tensor_.shape(), min, {});
|
||||
return clamp_min(min_tensor);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::clamp_min(const at::Tensor& min) const {
|
||||
return Tensor(paddle::experimental::maximum(tensor_, min._PD_GetInner()));
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::clamp_min_(const at::Scalar& min) const {
|
||||
// Create a tensor with the same shape filled with the min value
|
||||
at::Tensor min_tensor = at::full(tensor_.shape(), min, {});
|
||||
return clamp_min_(min_tensor);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::clamp_min_(const at::Tensor& min) const {
|
||||
PaddleTensor temp =
|
||||
paddle::experimental::maximum(tensor_, min._PD_GetInner());
|
||||
const_cast<PaddleTensor&>(tensor_) = temp;
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,34 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include "paddle/phi/api/include/sparse_api.h"
|
||||
|
||||
namespace at {} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::coalesce() const {
|
||||
PD_CHECK(layout() == kSparse,
|
||||
"coalesce expected sparse coordinate tensor layout but got ",
|
||||
layout());
|
||||
if (is_coalesced()) {
|
||||
return *this;
|
||||
}
|
||||
return paddle::experimental::sparse::coalesce(tensor_);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,43 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor detach(const at::Tensor& self) {
|
||||
// Create a new Tensor that shares data but has no autograd history
|
||||
auto inner = self._PD_GetInner();
|
||||
PaddleTensor detached_tensor(inner.impl());
|
||||
detached_tensor.set_name(inner.name());
|
||||
detached_tensor.set_autograd_meta(nullptr);
|
||||
return Tensor(detached_tensor);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::detach() const { return at::detach(*this); }
|
||||
|
||||
inline at::Tensor& Tensor::detach_() const {
|
||||
// In-place version: clear autograd meta of current tensor
|
||||
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
|
||||
inner.set_autograd_meta(nullptr);
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,44 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <ATen/ops/tensor_split.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::vector<at::Tensor> dsplit(const at::Tensor& self,
|
||||
int64_t sections) {
|
||||
return tensor_split(self, sections, 2);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> dsplit(const at::Tensor& self,
|
||||
at::IntArrayRef indices) {
|
||||
return tensor_split(self, indices, 2);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::dsplit(int64_t sections) const {
|
||||
return at::dsplit(*this, sections);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::dsplit(at::IntArrayRef indices) const {
|
||||
return at::dsplit(*this, indices);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,75 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/dense_sparse_conversion.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor empty(
|
||||
at::IntArrayRef size,
|
||||
at::TensorOptions options = {},
|
||||
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) {
|
||||
PD_CHECK(!(memory_format.has_value() &&
|
||||
memory_format.value() != c10::MemoryFormat::Contiguous),
|
||||
"`MemoryFormat` other than Contiguous is not supported now.");
|
||||
if (options.pinned_memory()) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !options.device().is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
auto dense = paddle::experimental::empty(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
phi::CPUPlace());
|
||||
dense = dense.copy_to(phi::GPUPinnedPlace(), /*blocking=*/true);
|
||||
return compat::_PD_ConvertToSparseIfNeeded(dense, options.layout());
|
||||
}
|
||||
auto dense = paddle::experimental::empty(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
options._PD_GetPlace());
|
||||
return compat::_PD_ConvertToSparseIfNeeded(dense, options.layout());
|
||||
}
|
||||
|
||||
inline at::Tensor empty(at::IntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory,
|
||||
::std::optional<at::MemoryFormat> memory_format) {
|
||||
PD_CHECK(!(memory_format.has_value() &&
|
||||
memory_format.value() != c10::MemoryFormat::Contiguous),
|
||||
"`MemoryFormat` other than Contiguous is not supported now.");
|
||||
auto options =
|
||||
at::TensorOptions()
|
||||
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
|
||||
.layout(layout)
|
||||
.device(device.value_or(at::kCPU))
|
||||
.pinned_memory(pin_memory);
|
||||
return empty(size, options, memory_format);
|
||||
}
|
||||
|
||||
#define empty_symint empty // SymIntArrayRef is same as IntArrayRef
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,81 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/dense_sparse_conversion.h>
|
||||
#include <utils/pinned_place.h>
|
||||
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor empty_like(
|
||||
const at::Tensor& self,
|
||||
at::TensorOptions options = {},
|
||||
::std::optional<at::MemoryFormat> memory_format = ::std::nullopt) {
|
||||
PD_CHECK(!(memory_format.has_value() &&
|
||||
memory_format.value() != c10::MemoryFormat::Contiguous),
|
||||
"`MemoryFormat` other than Contiguous is not supported now.");
|
||||
|
||||
auto dtype = options.dtype_opt().value_or(self.dtype());
|
||||
paddle::Tensor dense;
|
||||
if (options.pinned_memory()) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !options.device().is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
auto dense_cpu = paddle::experimental::empty_like(
|
||||
self._PD_GetInner(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(dtype),
|
||||
phi::CPUPlace());
|
||||
phi::Place base_place = options._PD_GetPlace();
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
|
||||
dense = dense_cpu.copy_to(pinned_place, /*blocking=*/true);
|
||||
} else {
|
||||
auto place = options.device_opt().value_or(self.device());
|
||||
dense = paddle::experimental::empty_like(
|
||||
self._PD_GetInner(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(dtype),
|
||||
place._PD_GetInner());
|
||||
}
|
||||
return compat::_PD_ConvertToSparseIfNeeded(dense, options.layout());
|
||||
}
|
||||
|
||||
inline at::Tensor empty_like(const at::Tensor& self,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory,
|
||||
::std::optional<at::MemoryFormat> memory_format) {
|
||||
PD_CHECK(!(memory_format.has_value() &&
|
||||
memory_format.value() != c10::MemoryFormat::Contiguous),
|
||||
"`MemoryFormat` other than Contiguous is not supported now.");
|
||||
|
||||
auto options = at::TensorOptions()
|
||||
.dtype(dtype)
|
||||
.layout(layout)
|
||||
.device(device)
|
||||
.pinned_memory(pin_memory);
|
||||
return empty_like(self, options, memory_format);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,41 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <c10/util/ArrayRef.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor empty_strided(at::IntArrayRef size,
|
||||
at::IntArrayRef stride,
|
||||
at::TensorOptions options = {}) {
|
||||
auto empty_tensor = paddle::experimental::empty(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
options._PD_GetPlace());
|
||||
|
||||
return paddle::experimental::as_strided(
|
||||
empty_tensor,
|
||||
std::vector<int64_t>(size.begin(), size.end()),
|
||||
std::vector<int64_t>(stride.begin(), stride.end()));
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,60 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <ATen/ops/item.h>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
inline bool equal(const at::Tensor& self, const at::Tensor& other) {
|
||||
PD_CHECK(self.defined(),
|
||||
"Expected a proper Tensor but got None (or an undefined Tensor in "
|
||||
"C++)");
|
||||
PD_CHECK(other.defined(),
|
||||
"Expected a proper Tensor but got None (or an undefined Tensor in "
|
||||
"C++)");
|
||||
PD_CHECK(self.device() == other.device(),
|
||||
"Cannot compare two tensors on "
|
||||
"different devices. Got: ",
|
||||
self.device(),
|
||||
" and ",
|
||||
other.device());
|
||||
if (self.sizes() != other.sizes()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
auto lhs = self._PD_GetInner();
|
||||
auto rhs = other._PD_GetInner();
|
||||
if (self.scalar_type() != other.scalar_type()) {
|
||||
rhs = paddle::experimental::cast(
|
||||
rhs, compat::_PD_AtenScalarTypeToPhiDataType(self.scalar_type()));
|
||||
}
|
||||
|
||||
auto result = paddle::experimental::equal_all(lhs, rhs);
|
||||
return at::Tensor(std::move(result)).item<bool>();
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline bool Tensor::equal(const at::Tensor& other) const {
|
||||
return at::equal(*this, other);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,137 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// expand - expands tensor to new size
|
||||
// PyTorch's expand works by right-aligning dimensions and broadcasting
|
||||
// dimensions with size 1 to the target size
|
||||
// Unlike Paddle's expand_v2, PyTorch allows non-singleton dimensions to be
|
||||
// preserved when they match the corresponding target dimension
|
||||
inline Tensor expand(const Tensor& self,
|
||||
at::IntArrayRef size,
|
||||
bool implicit = false) {
|
||||
// implicit parameter is used by PyTorch's vmap for internal optimization.
|
||||
// It doesn't affect the actual expand operation, so we can safely ignore it.
|
||||
|
||||
paddle::Tensor pd_tensor = self._PD_GetInner();
|
||||
|
||||
// Target sizes - convert to vector
|
||||
std::vector<int64_t> target_size_vec(size.begin(), size.end());
|
||||
auto target_rank = target_size_vec.size();
|
||||
auto input_dims = pd_tensor.dims();
|
||||
auto input_rank = static_cast<size_t>(input_dims.size());
|
||||
|
||||
// PyTorch's expand uses right-alignment semantics:
|
||||
// - For 1D tensor expand to 2D: {3}.expand({3,4}) treats input as {3,1},
|
||||
// expands to {3,4}
|
||||
// - Non-singleton dimensions are preserved, singleton dimensions (1) can
|
||||
// expand
|
||||
//
|
||||
// For example:
|
||||
// {3}.expand({3, 4}) -> input {3} becomes {3, 1} implicitly
|
||||
// then expand: dim 0: 3 stays 3, dim 1: 1 -> 4 -> result {3, 4}
|
||||
|
||||
if (input_rank < target_rank) {
|
||||
// Add leading 1s to right-align with target shape (PyTorch behavior)
|
||||
// Input {1, 2}, target {2, 3, 2} -> reshape to {1, 1, 2}
|
||||
std::vector<int64_t> reshape_vec(target_rank, 1);
|
||||
for (size_t i = 0; i < input_rank; ++i) {
|
||||
reshape_vec[target_rank - input_rank + i] = input_dims[i];
|
||||
}
|
||||
|
||||
// Check if Paddle's expand can handle this right-aligned shape
|
||||
// Paddle allows: input[i] == 1 (can expand), or input[i] == target[i]
|
||||
// (match)
|
||||
bool can_use_paddle_expand = true;
|
||||
size_t fail_dim = 0;
|
||||
for (size_t i = 0; i < target_rank; ++i) {
|
||||
bool dim_can_expand = (reshape_vec[i] == 1);
|
||||
bool dim_is_matching = (reshape_vec[i] == target_size_vec[i]);
|
||||
if (!dim_can_expand && !dim_is_matching) {
|
||||
can_use_paddle_expand = false;
|
||||
fail_dim = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (can_use_paddle_expand) {
|
||||
// Reshape to right-aligned shape, then expand
|
||||
paddle::Tensor reshaped =
|
||||
paddle::experimental::reshape(pd_tensor, phi::IntArray(reshape_vec));
|
||||
paddle::Tensor result = paddle::experimental::expand(
|
||||
reshaped, phi::IntArray(target_size_vec));
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
PD_THROW("expand(): the expanded size of the tensor (",
|
||||
target_size_vec[fail_dim],
|
||||
") must match the existing size (",
|
||||
reshape_vec[fail_dim],
|
||||
") at non-singleton dimension ",
|
||||
fail_dim,
|
||||
".");
|
||||
} else if (input_rank == target_rank) {
|
||||
// Same rank - check if we can use expand directly
|
||||
bool can_use_paddle_expand = true;
|
||||
size_t fail_dim = 0;
|
||||
for (size_t i = 0; i < target_rank; ++i) {
|
||||
auto in_size = input_dims[i];
|
||||
auto target_size = target_size_vec[i];
|
||||
if (in_size != 1 && in_size != target_size) {
|
||||
can_use_paddle_expand = false;
|
||||
fail_dim = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (can_use_paddle_expand) {
|
||||
paddle::Tensor result = paddle::experimental::expand(
|
||||
pd_tensor, phi::IntArray(target_size_vec));
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
PD_THROW("expand(): the expanded size of the tensor (",
|
||||
target_size_vec[fail_dim],
|
||||
") must match the existing size (",
|
||||
input_dims[fail_dim],
|
||||
") at non-singleton dimension ",
|
||||
fail_dim,
|
||||
".");
|
||||
} else {
|
||||
PD_THROW("expand(): the number of sizes provided (",
|
||||
target_rank,
|
||||
") must be greater or equal to the number of dimensions in the "
|
||||
"tensor (",
|
||||
input_rank,
|
||||
").");
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
// Member function: Tensor::expand
|
||||
inline Tensor Tensor::expand(at::IntArrayRef size, bool implicit) const {
|
||||
return at::expand(*this, size, implicit);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,26 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <ATen/ops/expand.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline Tensor Tensor::expand_as(const Tensor& other) const {
|
||||
return at::expand(*this, other.sizes());
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,106 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/pinned_place.h>
|
||||
#include <optional>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// eye(n) — n×n identity matrix
|
||||
inline at::Tensor eye(int64_t n, at::TensorOptions options = {}) {
|
||||
if (options.pinned_memory()) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !options.device().is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
phi::Place base_place = options._PD_GetPlace();
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
|
||||
auto dense = paddle::experimental::eye(
|
||||
n,
|
||||
/*num_columns=*/-1,
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
phi::CPUPlace());
|
||||
return dense.copy_to(pinned_place, /*blocking=*/true);
|
||||
}
|
||||
return paddle::experimental::eye(
|
||||
n,
|
||||
/*num_columns=*/-1,
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
options._PD_GetPlace());
|
||||
}
|
||||
|
||||
// eye(n, m) — n×m identity-like matrix
|
||||
inline at::Tensor eye(int64_t n, int64_t m, at::TensorOptions options = {}) {
|
||||
if (options.pinned_memory()) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !options.device().is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
phi::Place base_place = options._PD_GetPlace();
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
|
||||
auto dense = paddle::experimental::eye(
|
||||
n,
|
||||
m,
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
phi::CPUPlace());
|
||||
return dense.copy_to(pinned_place, /*blocking=*/true);
|
||||
}
|
||||
return paddle::experimental::eye(
|
||||
n,
|
||||
m,
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
options._PD_GetPlace());
|
||||
}
|
||||
|
||||
// eye(n, dtype, layout, device, pin_memory)
|
||||
inline at::Tensor eye(int64_t n,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) {
|
||||
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
|
||||
auto options =
|
||||
at::TensorOptions()
|
||||
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
|
||||
.device(device.value_or(at::kCPU))
|
||||
.pinned_memory(pin_memory);
|
||||
return eye(n, options);
|
||||
}
|
||||
|
||||
// eye(n, m, dtype, layout, device, pin_memory)
|
||||
inline at::Tensor eye(int64_t n,
|
||||
int64_t m,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) {
|
||||
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
|
||||
auto options =
|
||||
at::TensorOptions()
|
||||
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
|
||||
.device(device.value_or(at::kCPU))
|
||||
.pinned_memory(pin_memory);
|
||||
return eye(n, m, options);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,37 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor flatten(const at::Tensor& self,
|
||||
int64_t start_dim = 0,
|
||||
int64_t end_dim = -1) {
|
||||
return Tensor(paddle::experimental::flatten(self._PD_GetInner(),
|
||||
static_cast<int>(start_dim),
|
||||
static_cast<int>(end_dim)));
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::flatten(int64_t start_dim, int64_t end_dim) const {
|
||||
return at::flatten(*this, start_dim, end_dim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,219 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
|
||||
#include "paddle/phi/api/include/tensor_utils.h"
|
||||
namespace at {
|
||||
|
||||
namespace detail {
|
||||
|
||||
inline void noopDelete(void* /*unused*/) {}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
class TensorMaker {
|
||||
friend TensorMaker for_blob(void* data, IntArrayRef sizes) noexcept;
|
||||
|
||||
public:
|
||||
using ContextDeleter = DeleterFnPtr;
|
||||
|
||||
TensorMaker& strides(OptionalIntArrayRef value) noexcept {
|
||||
strides_ = value;
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
TensorMaker& storage_offset(std::optional<int64_t> value) noexcept {
|
||||
storage_offset_ = value;
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
TensorMaker& deleter(std::function<void(void*)> value) noexcept {
|
||||
deleter_ = std::move(value);
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
TensorMaker& context(void* value, ContextDeleter deleter = nullptr) noexcept {
|
||||
ctx_ = std::unique_ptr<void, ContextDeleter>{
|
||||
value, deleter != nullptr ? deleter : detail::noopDelete};
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
TensorMaker& target_device(std::optional<Device> value) noexcept {
|
||||
device_ = value;
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
TensorMaker& options(TensorOptions value) noexcept {
|
||||
opts_ = value;
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
TensorMaker& resizeable_storage() noexcept {
|
||||
resizeable_ = true;
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
Tensor make_tensor() {
|
||||
PD_CHECK(!deleter_ || !ctx_,
|
||||
"The deleter and context arguments are mutually exclusive.");
|
||||
|
||||
PD_CHECK(!storage_offset_.has_value() || storage_offset_.value() == 0,
|
||||
"storage_offset` should be zero.");
|
||||
|
||||
if (device_.has_value() && opts_.has_device() &&
|
||||
opts_.device().has_index()) {
|
||||
PD_CHECK(opts_.device() == *device_,
|
||||
"Specified device ",
|
||||
opts_.device(),
|
||||
" does not match device of data ",
|
||||
*device_);
|
||||
}
|
||||
|
||||
phi::Place pd_place;
|
||||
if (device_.has_value()) {
|
||||
pd_place = device_->_PD_GetInner();
|
||||
} else if (opts_.has_device() && opts_.device().has_index()) {
|
||||
pd_place = opts_.device()._PD_GetInner();
|
||||
} else {
|
||||
pd_place = phi::Place(); // UNDEFINED → auto-detect inside from_blob
|
||||
}
|
||||
|
||||
// Build paddle deleter: prefer explicit deleter_, then wrap ctx_ so its
|
||||
// lifetime is tied to the tensor allocation.
|
||||
paddle::Deleter pd_deleter = nullptr;
|
||||
if (deleter_) {
|
||||
pd_deleter = deleter_;
|
||||
} else if (ctx_) {
|
||||
// shared_ptr takes ownership of the context and calls its deleter when
|
||||
// the last copy (held in the lambda) is destroyed.
|
||||
auto shared_ctx =
|
||||
std::shared_ptr<void>(ctx_.release(), ctx_.get_deleter());
|
||||
pd_deleter = [shared_ctx](void* /*data*/) {};
|
||||
}
|
||||
|
||||
if (strides_.has_value()) {
|
||||
return paddle::from_blob(
|
||||
data_,
|
||||
sizes_._PD_ToPaddleIntArray(),
|
||||
strides_.value()._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(opts_.dtype()),
|
||||
phi::DataLayout::NCHW,
|
||||
pd_place,
|
||||
pd_deleter);
|
||||
} else {
|
||||
return paddle::from_blob(
|
||||
data_,
|
||||
sizes_._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(opts_.dtype()),
|
||||
phi::DataLayout::NCHW,
|
||||
pd_place,
|
||||
pd_deleter);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
explicit TensorMaker(void* data, IntArrayRef sizes) noexcept
|
||||
: data_{data}, sizes_{sizes} {}
|
||||
|
||||
std::size_t computeStorageSize() const noexcept;
|
||||
|
||||
DataPtr makeDataPtrFromDeleter() noexcept;
|
||||
|
||||
DataPtr makeDataPtrFromContext() noexcept;
|
||||
|
||||
IntArrayRef makeTempSizes() const noexcept;
|
||||
|
||||
void* data_;
|
||||
IntArrayRef sizes_;
|
||||
OptionalIntArrayRef strides_;
|
||||
std::optional<int64_t> storage_offset_;
|
||||
std::function<void(void*)> deleter_;
|
||||
std::unique_ptr<void, ContextDeleter> ctx_{nullptr, detail::noopDelete};
|
||||
std::optional<Device> device_;
|
||||
TensorOptions opts_;
|
||||
bool resizeable_{};
|
||||
};
|
||||
|
||||
inline TensorMaker for_blob(void* data, IntArrayRef sizes) noexcept {
|
||||
return TensorMaker{data, sizes};
|
||||
}
|
||||
|
||||
inline Tensor from_blob(
|
||||
void* data,
|
||||
IntArrayRef sizes,
|
||||
IntArrayRef strides,
|
||||
const std::function<void(void*)>& deleter,
|
||||
const TensorOptions& options = {},
|
||||
const std::optional<Device> target_device = std::nullopt) {
|
||||
return for_blob(data, sizes)
|
||||
.strides(strides)
|
||||
.deleter(deleter)
|
||||
.options(options)
|
||||
.target_device(target_device)
|
||||
.make_tensor();
|
||||
}
|
||||
|
||||
inline Tensor from_blob(
|
||||
void* data,
|
||||
IntArrayRef sizes,
|
||||
IntArrayRef strides,
|
||||
int64_t storage_offset,
|
||||
const std::function<void(void*)>& deleter,
|
||||
const TensorOptions& options = {},
|
||||
const std::optional<Device> target_device = std::nullopt) {
|
||||
return for_blob(data, sizes)
|
||||
.strides(strides)
|
||||
.storage_offset(storage_offset)
|
||||
.deleter(deleter)
|
||||
.options(options)
|
||||
.target_device(target_device)
|
||||
.make_tensor();
|
||||
}
|
||||
|
||||
inline Tensor from_blob(
|
||||
void* data,
|
||||
IntArrayRef sizes,
|
||||
std::function<void(void*)> deleter,
|
||||
const TensorOptions& options = {},
|
||||
const std::optional<Device> target_device = std::nullopt) {
|
||||
return for_blob(data, sizes)
|
||||
.deleter(std::move(deleter))
|
||||
.options(options)
|
||||
.target_device(target_device)
|
||||
.make_tensor();
|
||||
}
|
||||
|
||||
inline Tensor from_blob(void* data,
|
||||
IntArrayRef sizes,
|
||||
IntArrayRef strides,
|
||||
const TensorOptions& options = {}) {
|
||||
return for_blob(data, sizes).strides(strides).options(options).make_tensor();
|
||||
}
|
||||
|
||||
inline Tensor from_blob(void* data,
|
||||
IntArrayRef sizes,
|
||||
const TensorOptions& options = {}) {
|
||||
return for_blob(data, sizes).options(options).make_tensor();
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,104 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/SymIntArrayRef.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/pinned_place.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor full(at::IntArrayRef size,
|
||||
const at::Scalar& fill_value,
|
||||
at::TensorOptions options = {}) {
|
||||
if (options.pinned_memory()) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !options.device().is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
phi::Place base_place = options._PD_GetPlace();
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
|
||||
auto dense = paddle::experimental::full(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
fill_value,
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
phi::CPUPlace());
|
||||
return dense.copy_to(pinned_place, /*blocking=*/true);
|
||||
}
|
||||
return paddle::experimental::full(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
fill_value,
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
options._PD_GetPlace());
|
||||
}
|
||||
|
||||
inline at::Tensor full(at::IntArrayRef size,
|
||||
const at::Scalar& fill_value,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) {
|
||||
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
|
||||
auto options =
|
||||
at::TensorOptions()
|
||||
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
|
||||
.device(device.value_or(at::kCPU))
|
||||
.pinned_memory(pin_memory);
|
||||
return full(size, fill_value, options);
|
||||
}
|
||||
|
||||
inline at::Tensor full_symint(c10::SymIntArrayRef size,
|
||||
const at::Scalar& fill_value,
|
||||
at::TensorOptions options = {}) {
|
||||
if (options.pinned_memory()) {
|
||||
phi::Place base_place = options._PD_GetPlace();
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
|
||||
auto dense = paddle::experimental::full(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
fill_value,
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
phi::CPUPlace());
|
||||
return dense.copy_to(pinned_place, /*blocking=*/true);
|
||||
}
|
||||
return paddle::experimental::full(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
fill_value,
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
options._PD_GetPlace());
|
||||
}
|
||||
|
||||
inline at::Tensor full_symint(c10::SymIntArrayRef size,
|
||||
const at::Scalar& fill_value,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) {
|
||||
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
|
||||
auto options =
|
||||
at::TensorOptions()
|
||||
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
|
||||
.device(device.value_or(at::kCPU))
|
||||
.pinned_memory(pin_memory);
|
||||
return full_symint(size, fill_value, options);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,47 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <ATen/ops/tensor_split.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::vector<at::Tensor> hsplit(const at::Tensor& self,
|
||||
int64_t sections) {
|
||||
// For 1D tensors, split along dim 0; otherwise split along dim 1
|
||||
int64_t dim = self._PD_GetInner().dims().size() == 1 ? 0 : 1;
|
||||
return at::tensor_split(self, sections, dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> hsplit(const at::Tensor& self,
|
||||
at::IntArrayRef indices) {
|
||||
int64_t dim = self._PD_GetInner().dims().size() == 1 ? 0 : 1;
|
||||
return at::tensor_split(self, indices, dim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::hsplit(int64_t sections) const {
|
||||
return at::hsplit(*this, sections);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::hsplit(at::IntArrayRef indices) const {
|
||||
return at::hsplit(*this, indices);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,203 @@
|
||||
// Copyright (c) 2026 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 <ATen/TensorIndexing.h>
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <c10/core/List.h>
|
||||
|
||||
namespace at::indexing {
|
||||
|
||||
inline TensorIndex::TensorIndex(const at::Tensor& tensor)
|
||||
: tensor_(std::make_shared<at::Tensor>(tensor)),
|
||||
type_(TensorIndexType::Tensor) {}
|
||||
|
||||
inline const at::Tensor& TensorIndex::tensor() const { return *tensor_; }
|
||||
|
||||
} // namespace at::indexing
|
||||
|
||||
namespace at::detail {
|
||||
|
||||
inline bool _PD_is_full_slice(const at::indexing::Slice& slice) {
|
||||
return static_cast<int64_t>(slice.start()) == 0 &&
|
||||
static_cast<int64_t>(slice.stop()) == at::indexing::INDEX_MAX &&
|
||||
static_cast<int64_t>(slice.step()) == 1;
|
||||
}
|
||||
|
||||
inline at::Tensor _PD_apply_tensor_index(
|
||||
const at::Tensor& self, ArrayRef<at::indexing::TensorIndex> indices) {
|
||||
int64_t output_dim = 0;
|
||||
int tensor_index_count = 0;
|
||||
at::Tensor result = self;
|
||||
|
||||
for (const auto& index : indices) {
|
||||
if (index.is_tensor()) {
|
||||
++tensor_index_count;
|
||||
PD_CHECK(tensor_index_count == 1,
|
||||
"Multiple tensor indices mixed with None/Slice are not "
|
||||
"supported yet.");
|
||||
result = paddle::experimental::index_select(
|
||||
result._PD_GetInner(), index.tensor()._PD_GetInner(), output_dim);
|
||||
++output_dim;
|
||||
} else if (index.is_none()) {
|
||||
result =
|
||||
paddle::experimental::unsqueeze(result._PD_GetInner(), {output_dim});
|
||||
++output_dim;
|
||||
} else if (index.is_slice()) {
|
||||
const auto& slice = index.slice();
|
||||
PD_CHECK(_PD_is_full_slice(slice),
|
||||
"Only full Slice() is supported when mixed with tensor/None "
|
||||
"indices.");
|
||||
++output_dim;
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
inline at::Tensor _PD_index_tensor_indices(
|
||||
const at::Tensor& self, ArrayRef<at::indexing::TensorIndex> indices) {
|
||||
if (indices.size() == 0) {
|
||||
PD_THROW("index() cannot be called with an empty index list");
|
||||
}
|
||||
|
||||
bool has_slice = false;
|
||||
bool has_tensor_like = false;
|
||||
for (const auto& index : indices) {
|
||||
has_slice = has_slice || index.is_slice();
|
||||
has_tensor_like = has_tensor_like || index.is_tensor() || index.is_none();
|
||||
PD_CHECK(!index.is_ellipsis(), "Ellipsis index is not supported yet.");
|
||||
PD_CHECK(!index.is_integer(), "Integer index is not supported yet.");
|
||||
PD_CHECK(!index.is_boolean(), "Boolean index is not supported yet.");
|
||||
}
|
||||
|
||||
if (has_slice && !has_tensor_like) {
|
||||
std::vector<int64_t> axes;
|
||||
std::vector<int64_t> starts;
|
||||
std::vector<int64_t> ends;
|
||||
std::vector<int64_t> strides;
|
||||
axes.reserve(indices.size());
|
||||
starts.reserve(indices.size());
|
||||
ends.reserve(indices.size());
|
||||
strides.reserve(indices.size());
|
||||
|
||||
int64_t dim = 0;
|
||||
for (const auto& index : indices) {
|
||||
const auto& slice = index.slice();
|
||||
axes.push_back(dim++);
|
||||
starts.push_back(static_cast<int64_t>(slice.start()));
|
||||
ends.push_back(static_cast<int64_t>(slice.stop()));
|
||||
strides.push_back(static_cast<int64_t>(slice.step()));
|
||||
}
|
||||
|
||||
return paddle::experimental::slice(
|
||||
self._PD_GetInner(), axes, starts, ends, strides, {});
|
||||
}
|
||||
|
||||
if (has_slice) {
|
||||
return _PD_apply_tensor_index(self, indices);
|
||||
}
|
||||
|
||||
c10::List<::std::optional<at::Tensor>> tensor_indices;
|
||||
for (const auto& index : indices) {
|
||||
if (index.is_none()) {
|
||||
tensor_indices.push_back(::std::nullopt);
|
||||
} else if (index.is_tensor()) {
|
||||
tensor_indices.push_back(index.tensor());
|
||||
}
|
||||
}
|
||||
return self.index(tensor_indices);
|
||||
}
|
||||
|
||||
} // namespace at::detail
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor index(const at::Tensor& self,
|
||||
const c10::List<::std::optional<at::Tensor>>& indices) {
|
||||
if (indices.empty()) {
|
||||
return self;
|
||||
}
|
||||
|
||||
bool all_none = true;
|
||||
for (const auto& idx : indices) {
|
||||
if (idx.has_value()) {
|
||||
all_none = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (all_none) {
|
||||
return self;
|
||||
}
|
||||
|
||||
std::vector<paddle::Tensor> pd_indices;
|
||||
std::vector<bool> has_index(indices.size(), false);
|
||||
pd_indices.reserve(indices.size());
|
||||
|
||||
for (size_t i = 0; i < indices.size(); ++i) {
|
||||
if (indices[i].has_value()) {
|
||||
pd_indices.push_back(indices[i].value()._PD_GetInner());
|
||||
has_index[i] = true;
|
||||
} else {
|
||||
pd_indices.push_back(paddle::Tensor());
|
||||
}
|
||||
}
|
||||
|
||||
int non_none_count = 0;
|
||||
size_t first_non_none = 0;
|
||||
for (size_t i = 0; i < has_index.size(); ++i) {
|
||||
if (has_index[i]) {
|
||||
non_none_count++;
|
||||
first_non_none = i;
|
||||
}
|
||||
}
|
||||
|
||||
if (non_none_count == 1 && first_non_none == 0) {
|
||||
return paddle::experimental::index_select(
|
||||
self._PD_GetInner(), pd_indices[0], 0);
|
||||
}
|
||||
|
||||
if (non_none_count == static_cast<int>(indices.size())) {
|
||||
auto stacked_indices = paddle::experimental::stack(pd_indices, -1);
|
||||
return paddle::experimental::gather_nd(self._PD_GetInner(),
|
||||
stacked_indices);
|
||||
}
|
||||
|
||||
auto self_dims = self._PD_GetInner().dims();
|
||||
int self_rank = self_dims.size();
|
||||
at::Tensor result = self;
|
||||
|
||||
for (size_t i = 0; i < indices.size() && i < static_cast<size_t>(self_rank);
|
||||
++i) {
|
||||
if (has_index[i]) {
|
||||
paddle::Tensor pd_result = result._PD_GetInner();
|
||||
result = paddle::experimental::index_select(
|
||||
pd_result, pd_indices[i], static_cast<int>(i));
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::index(
|
||||
ArrayRef<at::indexing::TensorIndex> indices) const {
|
||||
return at::detail::_PD_index_tensor_indices(*this, indices);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,205 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <ATen/ops/index.h>
|
||||
#include <c10/core/List.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
|
||||
namespace at::detail {
|
||||
|
||||
inline std::vector<at::Tensor> _PD_convert_indices_list(
|
||||
const c10::List<::std::optional<at::Tensor>>& indices) {
|
||||
std::vector<at::Tensor> result;
|
||||
result.reserve(indices.size());
|
||||
for (const auto& idx : indices) {
|
||||
if (idx.has_value()) {
|
||||
result.push_back(idx.value());
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
inline c10::List<::std::optional<at::Tensor>> _PD_convert_tensor_index_list(
|
||||
ArrayRef<at::indexing::TensorIndex> indices) {
|
||||
c10::List<::std::optional<at::Tensor>> result;
|
||||
for (const auto& index : indices) {
|
||||
PD_CHECK(!index.is_ellipsis(), "Ellipsis index is not supported yet.");
|
||||
PD_CHECK(!index.is_integer(), "Integer index is not supported yet.");
|
||||
PD_CHECK(!index.is_boolean(), "Boolean index is not supported yet.");
|
||||
if (index.is_slice()) {
|
||||
PD_CHECK(_PD_is_full_slice(index.slice()),
|
||||
"Only full Slice() is supported in index_put_ TensorIndex "
|
||||
"paths.");
|
||||
} else if (index.is_tensor()) {
|
||||
result.push_back(index.tensor());
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
inline at::Tensor _PD_squeeze_newaxis_value(
|
||||
const at::Tensor& values, ArrayRef<at::indexing::TensorIndex> indices) {
|
||||
std::vector<int64_t> value_shape(values.sizes().begin(),
|
||||
values.sizes().end());
|
||||
size_t value_dim = 0;
|
||||
bool changed = false;
|
||||
|
||||
for (const auto& index : indices) {
|
||||
if (index.is_none()) {
|
||||
if (!value_shape.empty()) {
|
||||
PD_CHECK(value_dim < value_shape.size(),
|
||||
"index_put_ value rank is too small for None index.");
|
||||
PD_CHECK(value_shape[value_dim] == 1,
|
||||
"index_put_ expected value dimension inserted by None to "
|
||||
"have size 1, but got ",
|
||||
value_shape[value_dim],
|
||||
".");
|
||||
value_shape.erase(value_shape.begin() + value_dim);
|
||||
changed = true;
|
||||
}
|
||||
} else if (index.is_tensor()) {
|
||||
if (!value_shape.empty()) {
|
||||
++value_dim;
|
||||
}
|
||||
} else if (index.is_slice()) {
|
||||
PD_CHECK(_PD_is_full_slice(index.slice()),
|
||||
"Only full Slice() is supported in index_put_ TensorIndex "
|
||||
"paths.");
|
||||
if (!value_shape.empty()) {
|
||||
++value_dim;
|
||||
}
|
||||
} else {
|
||||
PD_CHECK(!index.is_ellipsis(), "Ellipsis index is not supported yet.");
|
||||
PD_CHECK(!index.is_integer(), "Integer index is not supported yet.");
|
||||
PD_CHECK(!index.is_boolean(), "Boolean index is not supported yet.");
|
||||
}
|
||||
}
|
||||
|
||||
if (!changed) {
|
||||
return values;
|
||||
}
|
||||
return paddle::experimental::reshape(values._PD_GetInner(),
|
||||
phi::IntArray(value_shape));
|
||||
}
|
||||
|
||||
} // namespace at::detail
|
||||
|
||||
namespace at {
|
||||
|
||||
// index_put_: Set values at specified indices (in-place)
|
||||
inline at::Tensor& index_put_(
|
||||
at::Tensor& self, // NOLINT(runtime/references)
|
||||
const c10::List<::std::optional<at::Tensor>>& indices,
|
||||
const at::Tensor& values,
|
||||
bool accumulate = false) {
|
||||
std::vector<paddle::Tensor> pd_indices;
|
||||
pd_indices.reserve(indices.size());
|
||||
for (const auto& idx : indices) {
|
||||
if (idx.has_value()) {
|
||||
pd_indices.push_back(idx.value()._PD_GetInner());
|
||||
}
|
||||
}
|
||||
|
||||
paddle::experimental::index_put_(
|
||||
self._PD_GetInner(), pd_indices, values._PD_GetInner(), accumulate);
|
||||
return self;
|
||||
}
|
||||
|
||||
// index_put: Non-inplace version
|
||||
inline at::Tensor index_put(
|
||||
const at::Tensor& self,
|
||||
const c10::List<::std::optional<at::Tensor>>& indices,
|
||||
const at::Tensor& values,
|
||||
bool accumulate = false) {
|
||||
std::vector<paddle::Tensor> pd_indices;
|
||||
pd_indices.reserve(indices.size());
|
||||
for (const auto& idx : indices) {
|
||||
if (idx.has_value()) {
|
||||
pd_indices.push_back(idx.value()._PD_GetInner());
|
||||
}
|
||||
}
|
||||
|
||||
return paddle::experimental::index_put(
|
||||
self._PD_GetInner(), pd_indices, values._PD_GetInner(), accumulate);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::index(
|
||||
const c10::List<::std::optional<at::Tensor>>& indices) const {
|
||||
return at::index(*this, indices);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::index_put_(
|
||||
const c10::List<::std::optional<at::Tensor>>& indices,
|
||||
const at::Tensor& values,
|
||||
bool accumulate) const {
|
||||
return at::index_put_(
|
||||
const_cast<at::Tensor&>(*this), indices, values, accumulate);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::index_put_(
|
||||
ArrayRef<at::indexing::TensorIndex> indices, Tensor const& rhs) {
|
||||
auto tensor_indices = detail::_PD_convert_tensor_index_list(indices);
|
||||
at::Tensor values = detail::_PD_squeeze_newaxis_value(rhs, indices);
|
||||
if (tensor_indices.empty()) {
|
||||
return copy_(values);
|
||||
}
|
||||
return index_put_(tensor_indices, values);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::index_put_(
|
||||
ArrayRef<at::indexing::TensorIndex> indices, const Scalar& v) {
|
||||
auto tensor_indices = detail::_PD_convert_tensor_index_list(indices);
|
||||
if (tensor_indices.empty()) {
|
||||
std::vector<int64_t> value_shape(this->sizes().begin(),
|
||||
this->sizes().end());
|
||||
auto scalar_tensor =
|
||||
at::Tensor(paddle::experimental::full(phi::IntArray(value_shape),
|
||||
phi::Scalar(v.to<double>()),
|
||||
this->_PD_GetInner().dtype()));
|
||||
return copy_(scalar_tensor);
|
||||
}
|
||||
auto scalar_tensor = at::Tensor(paddle::experimental::full(
|
||||
{}, phi::Scalar(v.to<double>()), this->_PD_GetInner().dtype()));
|
||||
return index_put_(indices, scalar_tensor);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::index_put_(
|
||||
std::initializer_list<at::indexing::TensorIndex> indices,
|
||||
Tensor const& rhs) {
|
||||
return index_put_(ArrayRef<at::indexing::TensorIndex>(indices), rhs);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::index_put_(
|
||||
std::initializer_list<at::indexing::TensorIndex> indices, const Scalar& v) {
|
||||
return index_put_(ArrayRef<at::indexing::TensorIndex>(indices), v);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::index_put(
|
||||
const c10::List<::std::optional<at::Tensor>>& indices,
|
||||
const at::Tensor& values,
|
||||
bool accumulate) const {
|
||||
return at::index_put(*this, indices, values, accumulate);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,33 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include "paddle/phi/core/sparse_coo_tensor.h"
|
||||
|
||||
namespace at {} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline bool Tensor::is_coalesced() const {
|
||||
PD_CHECK(tensor_.layout() == common::DataLayout::SPARSE_COO,
|
||||
"is_coalesced expected sparse coordinate tensor layout but got ",
|
||||
layout());
|
||||
auto sparse_coo_tensor =
|
||||
std::dynamic_pointer_cast<phi::SparseCooTensor>(tensor_.impl());
|
||||
return sparse_coo_tensor->coalesced();
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,44 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <ATen/ops/_local_scalar_dense.h>
|
||||
|
||||
namespace at {} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Scalar Tensor::item() const {
|
||||
auto numel = this->sym_numel();
|
||||
PD_CHECK(numel == 1,
|
||||
"a Tensor with ",
|
||||
numel,
|
||||
" elements cannot be converted to Scalar");
|
||||
if (this->is_sparse()) {
|
||||
if (this->_nnz() == 0) return Scalar(0);
|
||||
if (this->is_coalesced()) return at::_local_scalar_dense(this->_values());
|
||||
return at::_local_scalar_dense(this->_values().sum());
|
||||
} else {
|
||||
return _local_scalar_dense(*this);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T Tensor::item() const {
|
||||
return item().to<T>();
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,35 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor masked_select(const at::Tensor& self,
|
||||
const at::Tensor& mask) {
|
||||
return Tensor(paddle::experimental::masked_select(self._PD_GetInner(),
|
||||
mask._PD_GetInner()));
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::masked_select(const at::Tensor& mask) const {
|
||||
return at::masked_select(*this, mask);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,115 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor narrow(const at::Tensor& self,
|
||||
int64_t dim,
|
||||
int64_t start,
|
||||
int64_t length) {
|
||||
// Bounds checks matching PyTorch behavior
|
||||
PD_CHECK(self.dim() > 0, "narrow() cannot be applied to a 0-dim tensor.");
|
||||
PD_CHECK(length >= 0, "narrow(): length must be non-negative.");
|
||||
|
||||
// Normalize negative dim
|
||||
int64_t ndim = self.dim();
|
||||
if (dim < 0) dim += ndim;
|
||||
PD_CHECK(dim >= 0 && dim < ndim,
|
||||
"start out of range (expected to be in range of [",
|
||||
-ndim,
|
||||
", ",
|
||||
ndim - 1,
|
||||
"], but got ",
|
||||
dim,
|
||||
")");
|
||||
|
||||
int64_t cur_size = self.sizes()[dim];
|
||||
|
||||
// Wrap negative start (matching PyTorch: only wrap when start != cur_size)
|
||||
if (start < 0) {
|
||||
start = start + cur_size;
|
||||
}
|
||||
PD_CHECK(start <= cur_size - length,
|
||||
"start (",
|
||||
start,
|
||||
") + length (",
|
||||
length,
|
||||
") exceeds dimension size (",
|
||||
cur_size,
|
||||
").");
|
||||
|
||||
// Use slice to implement narrow: narrow(dim, start, length) is equivalent
|
||||
// to slice(dim, start, start + length)
|
||||
return Tensor(paddle::experimental::slice(
|
||||
self._PD_GetInner(), {dim}, {start}, {start + length}, {1}, {}));
|
||||
}
|
||||
|
||||
inline at::Tensor narrow_symint(const at::Tensor& self,
|
||||
int64_t dim,
|
||||
c10::SymInt start,
|
||||
c10::SymInt length) {
|
||||
return narrow(self, dim, start, length);
|
||||
}
|
||||
|
||||
inline at::Tensor narrow(const at::Tensor& self,
|
||||
int64_t dim,
|
||||
const at::Tensor& start,
|
||||
int64_t length) {
|
||||
// Extract scalar value from start tensor
|
||||
PD_CHECK(start.numel() == 1,
|
||||
"start must be a 0-dim tensor or 1-element tensor");
|
||||
int64_t start_val = start.item<int64_t>();
|
||||
return narrow(self, dim, start_val, length);
|
||||
}
|
||||
|
||||
inline at::Tensor narrow_symint(const at::Tensor& self,
|
||||
int64_t dim,
|
||||
const at::Tensor& start,
|
||||
c10::SymInt length) {
|
||||
return narrow(self, dim, start, length);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::narrow(int64_t dim,
|
||||
int64_t start,
|
||||
int64_t length) const {
|
||||
return at::narrow(*this, dim, start, length);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::narrow_symint(int64_t dim,
|
||||
c10::SymInt start,
|
||||
c10::SymInt length) const {
|
||||
return at::narrow_symint(*this, dim, start, length);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::narrow(int64_t dim,
|
||||
const at::Tensor& start,
|
||||
int64_t length) const {
|
||||
return at::narrow(*this, dim, start, length);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::narrow_symint(int64_t dim,
|
||||
const at::Tensor& start,
|
||||
c10::SymInt length) const {
|
||||
return at::narrow_symint(*this, dim, start, length);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,53 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <ATen/ops/narrow.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor narrow_copy(const at::Tensor& self,
|
||||
int64_t dim,
|
||||
int64_t start,
|
||||
int64_t length) {
|
||||
// narrow_copy returns a copy of the narrowed tensor
|
||||
return narrow(self, dim, start, length).clone();
|
||||
}
|
||||
|
||||
inline at::Tensor narrow_copy_symint(const at::Tensor& self,
|
||||
int64_t dim,
|
||||
c10::SymInt start,
|
||||
c10::SymInt length) {
|
||||
return narrow_copy(self, dim, start, length);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::narrow_copy(int64_t dim,
|
||||
int64_t start,
|
||||
int64_t length) const {
|
||||
return at::narrow_copy(*this, dim, start, length);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::narrow_copy_symint(int64_t dim,
|
||||
c10::SymInt start,
|
||||
c10::SymInt length) const {
|
||||
return at::narrow_copy_symint(*this, dim, start, length);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,68 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/pinned_place.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// Member function: Tensor::new_empty
|
||||
inline Tensor Tensor::new_empty(at::IntArrayRef size,
|
||||
at::TensorOptions options) const {
|
||||
caffe2::TypeMeta actual_dtype = options.dtype_opt().value_or(dtype());
|
||||
auto actual_device = options.device_opt().value_or(device());
|
||||
auto actual_pin_memory = options.pinned_memory();
|
||||
|
||||
auto pd_dtype = compat::_PD_AtenScalarTypeToPhiDataType(actual_dtype);
|
||||
auto pd_place = actual_device._PD_GetInner();
|
||||
|
||||
paddle::Tensor result;
|
||||
if (actual_pin_memory) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !actual_device.is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(pd_place);
|
||||
auto dense_cpu = paddle::experimental::empty(
|
||||
size._PD_ToPaddleIntArray(), pd_dtype, phi::CPUPlace());
|
||||
result = dense_cpu.copy_to(pinned_place, /*blocking=*/true);
|
||||
} else {
|
||||
result = paddle::experimental::empty(
|
||||
size._PD_ToPaddleIntArray(), pd_dtype, pd_place);
|
||||
}
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
inline Tensor Tensor::new_empty(at::IntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout>,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) const {
|
||||
auto options = at::TensorOptions()
|
||||
.dtype(dtype.value_or(this->scalar_type()))
|
||||
.device(device.value_or(this->device()))
|
||||
.pinned_memory(pin_memory);
|
||||
return new_empty(size, options);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,70 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/pinned_place.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// Member function: Tensor::new_full
|
||||
inline Tensor Tensor::new_full(at::IntArrayRef size,
|
||||
const at::Scalar& fill_value,
|
||||
at::TensorOptions options) const {
|
||||
caffe2::TypeMeta actual_dtype = options.dtype_opt().value_or(dtype());
|
||||
auto actual_device = options.device_opt().value_or(device());
|
||||
auto actual_pin_memory = options.pinned_memory();
|
||||
|
||||
auto pd_dtype = compat::_PD_AtenScalarTypeToPhiDataType(actual_dtype);
|
||||
auto pd_place = actual_device._PD_GetInner();
|
||||
|
||||
paddle::Tensor result;
|
||||
if (actual_pin_memory) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !actual_device.is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(pd_place);
|
||||
auto dense_cpu = paddle::experimental::full(
|
||||
size._PD_ToPaddleIntArray(), fill_value, pd_dtype, phi::CPUPlace());
|
||||
result = dense_cpu.copy_to(pinned_place, /*blocking=*/true);
|
||||
} else {
|
||||
result = paddle::experimental::full(
|
||||
size._PD_ToPaddleIntArray(), fill_value, pd_dtype, pd_place);
|
||||
}
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
inline Tensor Tensor::new_full(at::IntArrayRef size,
|
||||
const at::Scalar& fill_value,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout>,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) const {
|
||||
auto options = at::TensorOptions()
|
||||
.dtype(dtype.value_or(this->scalar_type()))
|
||||
.device(device.value_or(this->device()))
|
||||
.pinned_memory(pin_memory);
|
||||
return new_full(size, fill_value, options);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,68 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/pinned_place.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// Member function: Tensor::new_ones
|
||||
inline Tensor Tensor::new_ones(at::IntArrayRef size,
|
||||
at::TensorOptions options) const {
|
||||
caffe2::TypeMeta actual_dtype = options.dtype_opt().value_or(dtype());
|
||||
auto actual_device = options.device_opt().value_or(device());
|
||||
auto actual_pin_memory = options.pinned_memory();
|
||||
|
||||
auto pd_dtype = compat::_PD_AtenScalarTypeToPhiDataType(actual_dtype);
|
||||
auto pd_place = actual_device._PD_GetInner();
|
||||
|
||||
paddle::Tensor result;
|
||||
if (actual_pin_memory) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !actual_device.is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(pd_place);
|
||||
auto dense_cpu = paddle::experimental::ones(
|
||||
size._PD_ToPaddleIntArray(), pd_dtype, phi::CPUPlace());
|
||||
result = dense_cpu.copy_to(pinned_place, /*blocking=*/true);
|
||||
} else {
|
||||
result = paddle::experimental::ones(
|
||||
size._PD_ToPaddleIntArray(), pd_dtype, pd_place);
|
||||
}
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
inline Tensor Tensor::new_ones(at::IntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout>,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) const {
|
||||
auto options = at::TensorOptions()
|
||||
.dtype(dtype.value_or(this->scalar_type()))
|
||||
.device(device.value_or(this->device()))
|
||||
.pinned_memory(pin_memory);
|
||||
return new_ones(size, options);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,69 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/SymIntArrayRef.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/pinned_place.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// Member function: Tensor::new_zeros
|
||||
inline Tensor Tensor::new_zeros(at::IntArrayRef size,
|
||||
at::TensorOptions options) const {
|
||||
caffe2::TypeMeta actual_dtype = options.dtype_opt().value_or(dtype());
|
||||
auto actual_device = options.device_opt().value_or(device());
|
||||
auto actual_pin_memory = options.pinned_memory();
|
||||
|
||||
auto pd_dtype = compat::_PD_AtenScalarTypeToPhiDataType(actual_dtype);
|
||||
auto pd_place = actual_device._PD_GetInner();
|
||||
|
||||
paddle::Tensor result;
|
||||
if (actual_pin_memory) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !actual_device.is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(pd_place);
|
||||
auto dense_cpu = paddle::experimental::zeros(
|
||||
size._PD_ToPaddleIntArray(), pd_dtype, phi::CPUPlace());
|
||||
result = dense_cpu.copy_to(pinned_place, /*blocking=*/true);
|
||||
} else {
|
||||
result = paddle::experimental::zeros(
|
||||
size._PD_ToPaddleIntArray(), pd_dtype, pd_place);
|
||||
}
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
inline Tensor Tensor::new_zeros(at::IntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout>,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) const {
|
||||
auto options = at::TensorOptions()
|
||||
.dtype(dtype.value_or(this->scalar_type()))
|
||||
.device(device.value_or(this->device()))
|
||||
.pinned_memory(pin_memory);
|
||||
return new_zeros(size, options);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,99 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/pinned_place.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor ones(at::IntArrayRef size, at::TensorOptions options = {}) {
|
||||
if (options.pinned_memory()) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !options.device().is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
phi::Place base_place = options._PD_GetPlace();
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
|
||||
auto dense = paddle::experimental::ones(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
phi::CPUPlace());
|
||||
return dense.copy_to(pinned_place, /*blocking=*/true);
|
||||
}
|
||||
return paddle::experimental::ones(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
options._PD_GetPlace());
|
||||
}
|
||||
|
||||
inline at::Tensor ones(at::IntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) {
|
||||
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
|
||||
auto options =
|
||||
at::TensorOptions()
|
||||
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
|
||||
.device(device.value_or(at::kCPU))
|
||||
.pinned_memory(pin_memory);
|
||||
return ones(size, options);
|
||||
}
|
||||
|
||||
inline at::Tensor ones_symint(c10::SymIntArrayRef size,
|
||||
at::TensorOptions options = {}) {
|
||||
if (options.pinned_memory()) {
|
||||
// Pinning memory is only supported for CPU tensors
|
||||
if (options.has_device() && !options.device().is_cpu()) {
|
||||
PD_THROW(
|
||||
"pin_memory=true requires device to be CPU, but got non-CPU device");
|
||||
}
|
||||
phi::Place base_place = options._PD_GetPlace();
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
|
||||
auto dense = paddle::experimental::ones(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
phi::CPUPlace());
|
||||
return dense.copy_to(pinned_place, /*blocking=*/true);
|
||||
}
|
||||
return paddle::experimental::ones(
|
||||
size._PD_ToPaddleIntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype()),
|
||||
options._PD_GetPlace());
|
||||
}
|
||||
|
||||
inline at::Tensor ones_symint(c10::SymIntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) {
|
||||
PD_CHECK(!layout.has_value(), "`layout` is not supported now.");
|
||||
auto options =
|
||||
at::TensorOptions()
|
||||
.dtype(dtype.value_or(c10::get_default_dtype_as_scalartype()))
|
||||
.device(device.value_or(at::kCPU))
|
||||
.pinned_memory(pin_memory);
|
||||
return ones_symint(size, options);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,37 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor permute(const at::Tensor& self, at::IntArrayRef dims) {
|
||||
std::vector<int> perm(dims.size());
|
||||
for (size_t i = 0; i < dims.size(); i++) {
|
||||
perm[i] = static_cast<int>(dims[i]);
|
||||
}
|
||||
return paddle::experimental::transpose(self._PD_GetInner(), perm);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::permute(at::IntArrayRef dims) const {
|
||||
return at::permute(*this, dims);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,37 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor reciprocal(const at::Tensor& self) {
|
||||
return Tensor(paddle::experimental::reciprocal(self._PD_GetInner()));
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::reciprocal() const { return at::reciprocal(*this); }
|
||||
|
||||
inline at::Tensor& Tensor::reciprocal_() const {
|
||||
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
|
||||
paddle::experimental::reciprocal_(inner);
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,52 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/Device.h>
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#endif
|
||||
|
||||
namespace at {
|
||||
inline void Tensor::record_stream(at::Stream s) const {
|
||||
auto dense_tensor =
|
||||
std::dynamic_pointer_cast<phi::DenseTensor>(tensor_.impl());
|
||||
PD_CHECK(dense_tensor != nullptr,
|
||||
"record_stream only supports DenseTensor, but got a non-dense "
|
||||
"tensor implementation.");
|
||||
PD_CHECK(dense_tensor->place().GetType() != phi::AllocationType::CPU,
|
||||
"record_stream is not supported for CPU tensors.");
|
||||
#if (defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)) && \
|
||||
!defined(PADDLE_WITH_CUSTOM_DEVICE)
|
||||
paddle::memory::RecordStream(
|
||||
dense_tensor->Holder(), reinterpret_cast<gpuStream_t>(s.native_handle()));
|
||||
#elif defined(PADDLE_WITH_XPU)
|
||||
paddle::memory::RecordStream(dense_tensor->Holder(),
|
||||
reinterpret_cast<XPUStream>(s.native_handle()));
|
||||
#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
|
||||
paddle::memory::RecordStream(
|
||||
dense_tensor->Holder(),
|
||||
reinterpret_cast<phi::stream::stream_t>(s.native_handle()));
|
||||
#else
|
||||
(void)s;
|
||||
(void)dense_tensor;
|
||||
PD_THROW(
|
||||
"record_stream is not supported: no GPU/XPU/Custom device enabled "
|
||||
"in this build.");
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,26 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
// Member function: Tensor::rename
|
||||
inline Tensor Tensor::rename(::std::optional<at::DimnameList>) const {
|
||||
return *this;
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,44 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor reshape(const at::Tensor& self, at::IntArrayRef shape) {
|
||||
return paddle::experimental::reshape(self._PD_GetInner(),
|
||||
shape._PD_ToPaddleIntArray());
|
||||
}
|
||||
|
||||
inline at::Tensor reshape_symint(const at::Tensor& self,
|
||||
c10::SymIntArrayRef shape) {
|
||||
return paddle::experimental::reshape(self._PD_GetInner(),
|
||||
shape._PD_ToPaddleIntArray());
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::reshape(at::IntArrayRef shape) const {
|
||||
return at::reshape(*this, shape);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,119 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <limits>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/ddim.h"
|
||||
#include "paddle/phi/core/memory/malloc.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
namespace detail {
|
||||
|
||||
inline int64_t ResizeCheckedNumel(at::IntArrayRef size) {
|
||||
int64_t numel = 1;
|
||||
for (const auto dim : size) {
|
||||
TORCH_CHECK(dim >= 0,
|
||||
"Trying to create tensor with negative dimension ",
|
||||
dim,
|
||||
": ",
|
||||
size);
|
||||
if (dim == 0) {
|
||||
numel = 0;
|
||||
continue;
|
||||
}
|
||||
TORCH_CHECK(numel <= std::numeric_limits<int64_t>::max() / dim,
|
||||
"resize_ size is too large, possible overflow for size ",
|
||||
size);
|
||||
numel *= dim;
|
||||
}
|
||||
return numel;
|
||||
}
|
||||
|
||||
inline size_t ResizeCheckedStorageBytes(int64_t numel,
|
||||
size_t itemsize,
|
||||
size_t storage_offset_bytes) {
|
||||
const auto numel_size = static_cast<size_t>(numel);
|
||||
TORCH_CHECK(
|
||||
itemsize == 0 || numel_size <= (std::numeric_limits<size_t>::max() -
|
||||
storage_offset_bytes) /
|
||||
itemsize,
|
||||
"resize_ size is too large in bytes");
|
||||
return storage_offset_bytes + numel_size * itemsize;
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
// resize_ - operate on the underlying DenseTensor directly so we preserve
|
||||
// storage semantics across shrink/grow round-trips. When growth exceeds the
|
||||
// current capacity, expand the shared storage itself so aliasing views keep
|
||||
// their storage offset and existing storage contents stay intact.
|
||||
inline const at::Tensor& Tensor::resize_(
|
||||
at::IntArrayRef size,
|
||||
::std::optional<at::MemoryFormat> memory_format) const {
|
||||
// Keep old compat behavior for memory_format in this split PR.
|
||||
// TODO(youge325): add real ChannelsLast/ChannelsLast3d restride support
|
||||
// later.
|
||||
(void)memory_format;
|
||||
|
||||
std::vector<int64_t> dims(size.begin(), size.end());
|
||||
int64_t new_numel = detail::ResizeCheckedNumel(size);
|
||||
auto dense_tensor =
|
||||
std::dynamic_pointer_cast<phi::DenseTensor>(tensor_.impl());
|
||||
TORCH_CHECK(dense_tensor != nullptr,
|
||||
"resize_ only supports DenseTensor, but got a non-dense tensor");
|
||||
TORCH_CHECK(tensor_.defined(),
|
||||
"resize_ is not allowed on an undefined tensor");
|
||||
|
||||
const size_t itemsize = phi::SizeOf(dense_tensor->dtype());
|
||||
const size_t new_storage_bytes = detail::ResizeCheckedStorageBytes(
|
||||
new_numel, itemsize, dense_tensor->meta().offset);
|
||||
const size_t current_storage_bytes =
|
||||
dense_tensor->Holder() == nullptr ? 0 : dense_tensor->Holder()->size();
|
||||
|
||||
if (new_storage_bytes <= current_storage_bytes || new_numel == 0) {
|
||||
dense_tensor->Resize(dims);
|
||||
return *this;
|
||||
}
|
||||
|
||||
// Sync through the compat Storage path first so the DenseTensor holder is a
|
||||
// live StorageHolderView backed by shared StorageImpl.
|
||||
auto storage = this->storage();
|
||||
const auto old_holder = dense_tensor->Holder();
|
||||
TORCH_CHECK(old_holder != nullptr,
|
||||
"resize_ cannot grow a tensor without allocated storage");
|
||||
const phi::Place place = old_holder->place();
|
||||
auto new_holder = paddle::memory::AllocShared(place, new_storage_bytes);
|
||||
TORCH_CHECK(new_holder != nullptr, "resize_ failed to allocate storage");
|
||||
const size_t copy_bytes = std::min(old_holder->size(), new_storage_bytes);
|
||||
if (copy_bytes > 0 && old_holder->ptr() != nullptr &&
|
||||
old_holder->ptr() != new_holder->ptr()) {
|
||||
phi::memory_utils::Copy(
|
||||
place, new_holder->ptr(), place, old_holder->ptr(), copy_bytes);
|
||||
}
|
||||
storage.set_data_ptr_noswap(std::move(new_holder));
|
||||
dense_tensor->Resize(phi::make_ddim(dims));
|
||||
return *this;
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,80 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor select(const at::Tensor& self, int64_t dim, int64_t index) {
|
||||
// Normalize dim to positive value for error messages
|
||||
int64_t orig_dim = dim;
|
||||
if (dim < 0) {
|
||||
dim += self.dim();
|
||||
}
|
||||
// Check dim is valid
|
||||
if (dim < 0 || dim >= self.dim()) {
|
||||
PD_CHECK(false,
|
||||
"select(): index ",
|
||||
orig_dim,
|
||||
" out of range for tensor of size ",
|
||||
self.sizes(),
|
||||
" at dimension ",
|
||||
orig_dim);
|
||||
}
|
||||
// Handle negative index
|
||||
int64_t orig_index = index;
|
||||
if (index < 0) {
|
||||
index = self.size(dim) + index;
|
||||
}
|
||||
// Check index is valid
|
||||
if (index < 0 || index >= self.size(dim)) {
|
||||
PD_CHECK(false,
|
||||
"select(): index ",
|
||||
orig_index,
|
||||
" out of range for tensor of size ",
|
||||
self.sizes(),
|
||||
" at dimension ",
|
||||
orig_dim < 0 ? orig_dim + self.dim() : orig_dim);
|
||||
}
|
||||
|
||||
return Tensor(
|
||||
paddle::experimental::slice(self._PD_GetInner(),
|
||||
/*axes=*/{static_cast<int>(dim)},
|
||||
/*starts=*/{index},
|
||||
/*ends=*/{index + 1},
|
||||
/*infer_flags=*/{1},
|
||||
/*decrease_axis=*/{static_cast<int>(dim)}));
|
||||
}
|
||||
|
||||
inline at::Tensor select_symint(const at::Tensor& self,
|
||||
int64_t dim,
|
||||
c10::SymInt index) {
|
||||
return select(self, dim, static_cast<int64_t>(index));
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::select(int64_t dim, int64_t index) const {
|
||||
return at::select(*this, dim, index);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::select_symint(int64_t dim, c10::SymInt index) const {
|
||||
return at::select_symint(*this, dim, index);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,51 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor slice(const at::Tensor& self,
|
||||
int64_t dim = 0,
|
||||
::std::optional<int64_t> start = ::std::nullopt,
|
||||
::std::optional<int64_t> end = ::std::nullopt,
|
||||
int64_t step = 1) {
|
||||
// Materialize the compat StorageHolderView before creating the slice so the
|
||||
// base tensor and its views observe the same shared storage during resize_.
|
||||
(void)self.storage();
|
||||
return paddle::experimental::slice(
|
||||
self._PD_GetInner(),
|
||||
{dim},
|
||||
start.has_value() ? IntArrayRef(start.value())._PD_ToPaddleIntArray()
|
||||
: IntArrayRef()._PD_ToPaddleIntArray(),
|
||||
end.has_value() ? IntArrayRef(end.value())._PD_ToPaddleIntArray()
|
||||
: IntArrayRef()._PD_ToPaddleIntArray(),
|
||||
{1},
|
||||
{});
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::slice(int64_t dim,
|
||||
::std::optional<int64_t> start,
|
||||
::std::optional<int64_t> end,
|
||||
int64_t step) const {
|
||||
return at::slice(*this, dim, start, end, step);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,119 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/pinned_place.h>
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/api/include/sparse_api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
#include "paddle/phi/core/sparse_coo_tensor.h"
|
||||
|
||||
namespace at::detail {
|
||||
|
||||
inline std::vector<int64_t> _PD_infer_sparse_coo_size(
|
||||
const at::Tensor& indices) {
|
||||
auto host_indices = indices.cpu().to(at::kLong);
|
||||
int64_t sparse_dim = host_indices.dim() > 0 ? host_indices.size(0) : 0;
|
||||
int64_t nnz = host_indices.dim() > 1 ? host_indices.size(1) : 0;
|
||||
|
||||
std::vector<int64_t> inferred_size(static_cast<size_t>(sparse_dim), 0);
|
||||
const int64_t* data = host_indices.const_data_ptr<int64_t>();
|
||||
for (int64_t dim = 0; dim < sparse_dim; ++dim) {
|
||||
for (int64_t i = 0; i < nnz; ++i) {
|
||||
inferred_size[static_cast<size_t>(dim)] = std::max(
|
||||
inferred_size[static_cast<size_t>(dim)], data[dim * nnz + i] + 1);
|
||||
}
|
||||
}
|
||||
return inferred_size;
|
||||
}
|
||||
|
||||
inline void _PD_set_sparse_coo_coalesced(at::Tensor* tensor,
|
||||
::std::optional<bool> is_coalesced) {
|
||||
if (!is_coalesced.has_value()) {
|
||||
return;
|
||||
}
|
||||
auto sparse_tensor = std::dynamic_pointer_cast<phi::SparseCooTensor>(
|
||||
tensor->_PD_GetInner().impl());
|
||||
PD_CHECK(sparse_tensor,
|
||||
"Expected SparseCooTensor result from sparse_coo_tensor.");
|
||||
sparse_tensor->SetCoalesced(is_coalesced.value());
|
||||
}
|
||||
|
||||
} // namespace at::detail
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor sparse_coo_tensor(
|
||||
const at::Tensor& indices,
|
||||
const at::Tensor& values,
|
||||
at::IntArrayRef size,
|
||||
at::TensorOptions options = {},
|
||||
::std::optional<bool> is_coalesced = ::std::nullopt) {
|
||||
paddle::Tensor idx = indices._PD_GetInner();
|
||||
paddle::Tensor vals = values._PD_GetInner();
|
||||
|
||||
// PyTorch ignores dtype mismatch between values and TensorOptions in
|
||||
// sparse_coo_tensor; the resulting sparse tensor uses values' original dtype.
|
||||
// Do not cast or throw here.
|
||||
|
||||
if (options.pinned_memory()) {
|
||||
phi::Place base_place = options._PD_GetPlace();
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
|
||||
idx = idx.copy_to(pinned_place, /*blocking=*/true);
|
||||
vals = vals.copy_to(pinned_place, /*blocking=*/true);
|
||||
}
|
||||
|
||||
// PyTorch: sparse_coo_tensor(indices, values, size)
|
||||
// Paddle: sparse_coo_tensor(values, indices, shape)
|
||||
at::Tensor result = paddle::experimental::sparse::sparse_coo_tensor(
|
||||
vals, idx, std::vector<int64_t>(size.begin(), size.end()));
|
||||
detail::_PD_set_sparse_coo_coalesced(&result, is_coalesced);
|
||||
return result;
|
||||
}
|
||||
|
||||
inline at::Tensor sparse_coo_tensor(const at::Tensor& indices,
|
||||
const at::Tensor& values,
|
||||
at::IntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory,
|
||||
::std::optional<bool> is_coalesced) {
|
||||
PD_CHECK(!layout.has_value() || layout.value() == c10::kSparse,
|
||||
"`layout` must be Sparse for sparse_coo_tensor.");
|
||||
auto options =
|
||||
at::TensorOptions().dtype(dtype).device(device).pinned_memory(pin_memory);
|
||||
return sparse_coo_tensor(indices, values, size, options, is_coalesced);
|
||||
}
|
||||
|
||||
inline at::Tensor sparse_coo_tensor(
|
||||
const at::Tensor& indices,
|
||||
const at::Tensor& values,
|
||||
at::TensorOptions options = {},
|
||||
::std::optional<bool> is_coalesced = ::std::nullopt) {
|
||||
return sparse_coo_tensor(indices,
|
||||
values,
|
||||
detail::_PD_infer_sparse_coo_size(indices),
|
||||
options,
|
||||
is_coalesced);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,137 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/pinned_place.h>
|
||||
#include <optional>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/api/include/tensor.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/sparse_csr_tensor.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor sparse_csr_tensor(const at::Tensor& crow_indices,
|
||||
const at::Tensor& col_indices,
|
||||
const at::Tensor& values,
|
||||
at::IntArrayRef size,
|
||||
at::TensorOptions options) {
|
||||
paddle::Tensor crows = crow_indices._PD_GetInner();
|
||||
paddle::Tensor cols = col_indices._PD_GetInner();
|
||||
paddle::Tensor vals = values._PD_GetInner();
|
||||
|
||||
// PyTorch ignores dtype mismatch between values and TensorOptions in
|
||||
// sparse_csr_tensor; the resulting sparse tensor uses values' original dtype.
|
||||
// Do not cast or throw here.
|
||||
|
||||
if (options.pinned_memory()) {
|
||||
phi::Place base_place = options._PD_GetPlace();
|
||||
phi::Place pinned_place = compat::_PD_GetCreatePinnedPlace(base_place);
|
||||
crows = crows.copy_to(pinned_place, /*blocking=*/true);
|
||||
cols = cols.copy_to(pinned_place, /*blocking=*/true);
|
||||
vals = vals.copy_to(pinned_place, /*blocking=*/true);
|
||||
}
|
||||
|
||||
// Get the underlying DenseTensors
|
||||
auto* dense_crows = dynamic_cast<phi::DenseTensor*>(crows.impl().get());
|
||||
auto* dense_cols = dynamic_cast<phi::DenseTensor*>(cols.impl().get());
|
||||
auto* dense_values = dynamic_cast<phi::DenseTensor*>(vals.impl().get());
|
||||
|
||||
PD_CHECK(dense_crows != nullptr,
|
||||
"crow_indices must be a dense tensor for sparse_csr_tensor.");
|
||||
PD_CHECK(dense_cols != nullptr,
|
||||
"col_indices must be a dense tensor for sparse_csr_tensor.");
|
||||
PD_CHECK(dense_values != nullptr,
|
||||
"values must be a dense tensor for sparse_csr_tensor.");
|
||||
|
||||
// Create the SparseCsrTensor
|
||||
std::shared_ptr<phi::SparseCsrTensor> csr_tensor =
|
||||
std::make_shared<phi::SparseCsrTensor>(
|
||||
*dense_crows,
|
||||
*dense_cols,
|
||||
*dense_values,
|
||||
common::make_ddim(std::vector<int64_t>(size.begin(), size.end())));
|
||||
|
||||
// Wrap in a Paddle Tensor
|
||||
paddle::Tensor result;
|
||||
result.set_impl(csr_tensor);
|
||||
return result;
|
||||
}
|
||||
|
||||
inline at::Tensor sparse_csr_tensor(const at::Tensor& crow_indices,
|
||||
const at::Tensor& col_indices,
|
||||
const at::Tensor& values,
|
||||
at::IntArrayRef size,
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory) {
|
||||
PD_CHECK(!layout.has_value() || layout.value() == c10::kSparseCsr,
|
||||
"`layout` must be SparseCsr for sparse_csr_tensor.");
|
||||
auto options =
|
||||
at::TensorOptions().dtype(dtype).device(device).pinned_memory(pin_memory);
|
||||
return sparse_csr_tensor(crow_indices, col_indices, values, size, options);
|
||||
}
|
||||
|
||||
inline at::Tensor sparse_csr_tensor(const at::Tensor& crow_indices,
|
||||
const at::Tensor& col_indices,
|
||||
const at::Tensor& values,
|
||||
at::TensorOptions options) {
|
||||
// Infer size from crow_indices and col_indices:
|
||||
// nrows = crow_indices.size(0) - 1
|
||||
// ncols = max(col_indices) + 1
|
||||
int64_t nrows = crow_indices.size(0) - 1;
|
||||
int64_t ncols = 0;
|
||||
|
||||
if (col_indices.numel() > 0) {
|
||||
auto* dense_cols = dynamic_cast<phi::DenseTensor*>(
|
||||
col_indices._PD_GetInner().impl().get());
|
||||
PD_CHECK(dense_cols != nullptr,
|
||||
"col_indices must be a dense tensor for sparse_csr_tensor.");
|
||||
PD_CHECK(
|
||||
dense_cols->place().GetType() == phi::AllocationType::CPU,
|
||||
"sparse_csr_tensor without explicit size only supports CPU "
|
||||
"col_indices for automatic size inference. Please provide the size "
|
||||
"parameter explicitly for non-CPU tensors.");
|
||||
|
||||
int64_t n = dense_cols->numel();
|
||||
if (dense_cols->dtype() == phi::DataType::INT64) {
|
||||
const int64_t* data = dense_cols->data<int64_t>();
|
||||
for (int64_t i = 0; i < n; ++i) {
|
||||
if (data[i] + 1 > ncols) ncols = data[i] + 1;
|
||||
}
|
||||
} else if (dense_cols->dtype() == phi::DataType::INT32) {
|
||||
const int32_t* data = dense_cols->data<int32_t>();
|
||||
for (int64_t i = 0; i < n; ++i) {
|
||||
int64_t val = static_cast<int64_t>(data[i]) + 1;
|
||||
if (val > ncols) ncols = val;
|
||||
}
|
||||
} else {
|
||||
PD_CHECK(false,
|
||||
"col_indices must have dtype int32 or int64 for automatic "
|
||||
"size inference in sparse_csr_tensor.");
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int64_t> size_vec = {nrows, ncols};
|
||||
return sparse_csr_tensor(
|
||||
crow_indices, col_indices, values, at::IntArrayRef(size_vec), options);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,93 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::vector<at::Tensor> split(const at::Tensor& self,
|
||||
int64_t split_size,
|
||||
int64_t dim = 0) {
|
||||
// Calculate number of splits based on split_size
|
||||
int64_t dim_size = self._PD_GetInner().dims()[dim];
|
||||
std::vector<int64_t> split_sizes;
|
||||
for (int64_t i = 0; i < dim_size; i += split_size) {
|
||||
split_sizes.push_back(std::min(split_size, dim_size - i));
|
||||
}
|
||||
auto outputs =
|
||||
paddle::experimental::split(self._PD_GetInner(), split_sizes, dim);
|
||||
std::vector<at::Tensor> at_tensors;
|
||||
at_tensors.reserve(outputs.size());
|
||||
for (const auto& paddle_tensor : outputs) {
|
||||
at_tensors.emplace_back(paddle_tensor);
|
||||
}
|
||||
return at_tensors;
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> split_symint(const at::Tensor& self,
|
||||
c10::SymInt split_size,
|
||||
int64_t dim = 0) {
|
||||
return split(self, static_cast<int64_t>(split_size), dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> split(const at::Tensor& self,
|
||||
at::IntArrayRef split_sizes,
|
||||
int64_t dim = 0) {
|
||||
auto outputs = paddle::experimental::split(
|
||||
self._PD_GetInner(), split_sizes._PD_ToPaddleIntArray(), dim);
|
||||
std::vector<at::Tensor> at_tensors;
|
||||
at_tensors.reserve(outputs.size());
|
||||
for (const auto& paddle_tensor : outputs) {
|
||||
at_tensors.emplace_back(paddle_tensor);
|
||||
}
|
||||
return at_tensors;
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> split_symint(const at::Tensor& self,
|
||||
c10::SymIntArrayRef split_sizes,
|
||||
int64_t dim = 0) {
|
||||
return split(
|
||||
self,
|
||||
at::IntArrayRef(reinterpret_cast<const int64_t*>(split_sizes.data()),
|
||||
split_sizes.size()),
|
||||
dim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::split(int64_t split_size,
|
||||
int64_t dim = 0) const {
|
||||
return at::split(*this, split_size, dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::split_symint(c10::SymInt split_size,
|
||||
int64_t dim = 0) const {
|
||||
return at::split_symint(*this, split_size, dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::split(at::IntArrayRef split_sizes,
|
||||
int64_t dim = 0) const {
|
||||
return at::split(*this, split_sizes, dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::split_symint(
|
||||
c10::SymIntArrayRef split_sizes, int64_t dim = 0) const {
|
||||
return at::split_symint(*this, split_sizes, dim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,47 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <ATen/ops/split.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::vector<at::Tensor> split_with_sizes(const at::Tensor& self,
|
||||
at::IntArrayRef split_sizes,
|
||||
int64_t dim = 0) {
|
||||
return at::split(self, split_sizes, dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> split_with_sizes_symint(
|
||||
const at::Tensor& self, c10::SymIntArrayRef split_sizes, int64_t dim = 0) {
|
||||
return at::split_symint(self, split_sizes, dim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::split_with_sizes(
|
||||
at::IntArrayRef split_sizes, int64_t dim) const {
|
||||
return at::split_with_sizes(*this, split_sizes, dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::split_with_sizes_symint(
|
||||
c10::SymIntArrayRef split_sizes, int64_t dim) const {
|
||||
return at::split_with_sizes_symint(*this, split_sizes, dim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,66 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor squeeze(const at::Tensor& self) {
|
||||
return paddle::experimental::squeeze(self._PD_GetInner(), {});
|
||||
}
|
||||
|
||||
inline at::Tensor squeeze(const at::Tensor& self, int64_t dim) {
|
||||
return paddle::experimental::squeeze(self._PD_GetInner(), {dim});
|
||||
}
|
||||
|
||||
inline at::Tensor squeeze(const at::Tensor& self, at::IntArrayRef dim) {
|
||||
return paddle::experimental::squeeze(self._PD_GetInner(),
|
||||
dim._PD_ToPaddleIntArray());
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::squeeze() const { return at::squeeze(*this); }
|
||||
|
||||
inline at::Tensor Tensor::squeeze(int64_t dim) const {
|
||||
return at::squeeze(*this, dim);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::squeeze(at::IntArrayRef dim) const {
|
||||
return at::squeeze(*this, dim);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::squeeze_() const {
|
||||
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
|
||||
paddle::experimental::squeeze_(inner, {});
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::squeeze_(int64_t dim) const {
|
||||
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
|
||||
paddle::experimental::squeeze_(inner, {dim});
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::squeeze_(at::IntArrayRef dim) const {
|
||||
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
|
||||
paddle::experimental::squeeze_(inner, dim._PD_ToPaddleIntArray());
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,146 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/util/ArrayRef.h>
|
||||
#include <c10/util/OptionalArrayRef.h>
|
||||
#include <optional>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/int_array.h"
|
||||
#include "paddle/phi/common/scalar.h"
|
||||
|
||||
namespace at::detail {
|
||||
|
||||
// Internal implementation for std (standard deviation = sqrt(variance))
|
||||
inline Tensor _PD_std_impl(const Tensor& self,
|
||||
const std::vector<int64_t>& dims_vec,
|
||||
double correction_value,
|
||||
bool keepdim) {
|
||||
// Validate dimensions before processing
|
||||
int64_t ndim = self.dim();
|
||||
for (int64_t d : dims_vec) {
|
||||
int64_t dim_idx = d < 0 ? d + ndim : d;
|
||||
if (dim_idx < 0 || dim_idx >= ndim) {
|
||||
PD_CHECK(false,
|
||||
"Dimension out of range (expected to be in range of [",
|
||||
-ndim,
|
||||
", ",
|
||||
ndim - 1,
|
||||
"], but got ",
|
||||
d,
|
||||
")");
|
||||
}
|
||||
}
|
||||
phi::IntArray dims_int_array(dims_vec);
|
||||
paddle::Tensor tensor = self._PD_GetInner();
|
||||
|
||||
paddle::Tensor mean_tensor;
|
||||
if (dims_vec.empty()) {
|
||||
mean_tensor = paddle::experimental::mean(
|
||||
tensor, phi::IntArray(std::vector<int64_t>{}), true);
|
||||
} else {
|
||||
mean_tensor = paddle::experimental::mean(tensor, dims_int_array, true);
|
||||
}
|
||||
|
||||
paddle::Tensor diff = paddle::experimental::subtract(tensor, mean_tensor);
|
||||
paddle::Tensor diff_squared = paddle::experimental::multiply(diff, diff);
|
||||
|
||||
paddle::Tensor sum_squared_diff;
|
||||
if (dims_vec.empty()) {
|
||||
sum_squared_diff =
|
||||
paddle::experimental::sum(diff_squared,
|
||||
phi::IntArray(std::vector<int64_t>{}),
|
||||
diff_squared.dtype(),
|
||||
keepdim);
|
||||
} else {
|
||||
sum_squared_diff = paddle::experimental::sum(
|
||||
diff_squared, dims_int_array, diff_squared.dtype(), keepdim);
|
||||
}
|
||||
|
||||
int64_t n = tensor.numel();
|
||||
if (!dims_vec.empty()) {
|
||||
n = 1;
|
||||
for (int64_t d : dims_vec) {
|
||||
int64_t dim_idx = d < 0 ? d + tensor.dims().size() : d;
|
||||
if (dim_idx >= 0 &&
|
||||
dim_idx < static_cast<int64_t>(tensor.dims().size())) {
|
||||
n *= tensor.dims()[dim_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
double corrected_n = static_cast<double>(n) - correction_value;
|
||||
if (corrected_n <= 0.0) {
|
||||
corrected_n = static_cast<double>(n);
|
||||
}
|
||||
|
||||
std::vector<int64_t> result_shape_vec;
|
||||
for (int64_t i = 0; i < sum_squared_diff.dims().size(); ++i) {
|
||||
result_shape_vec.push_back(sum_squared_diff.dims()[i]);
|
||||
}
|
||||
paddle::Tensor correction_scalar =
|
||||
paddle::experimental::full(phi::IntArray(result_shape_vec),
|
||||
phi::Scalar(corrected_n),
|
||||
sum_squared_diff.dtype(),
|
||||
sum_squared_diff.place());
|
||||
paddle::Tensor variance =
|
||||
paddle::experimental::divide(sum_squared_diff, correction_scalar);
|
||||
|
||||
paddle::Tensor result = paddle::experimental::sqrt(variance);
|
||||
|
||||
return Tensor(result);
|
||||
}
|
||||
|
||||
} // namespace at::detail
|
||||
|
||||
namespace at {
|
||||
|
||||
inline Tensor Tensor::std(bool unbiased) const {
|
||||
std::vector<int64_t> empty_dims;
|
||||
double correction = unbiased ? 1.0 : 0.0;
|
||||
return detail::_PD_std_impl(*this, empty_dims, correction, false);
|
||||
}
|
||||
|
||||
inline Tensor Tensor::std(at::OptionalIntArrayRef dim,
|
||||
bool unbiased,
|
||||
bool keepdim) const {
|
||||
double correction = unbiased ? 1.0 : 0.0;
|
||||
std::vector<int64_t> dims_vec;
|
||||
if (dim.has_value() && dim.value().size() > 0) {
|
||||
dims_vec.assign(dim.value().begin(), dim.value().end());
|
||||
}
|
||||
return detail::_PD_std_impl(*this, dims_vec, correction, keepdim);
|
||||
}
|
||||
|
||||
inline Tensor Tensor::std(at::OptionalIntArrayRef dim,
|
||||
const ::std::optional<at::Scalar>& correction,
|
||||
bool keepdim) const {
|
||||
double correction_value = 1.0;
|
||||
if (correction.has_value()) {
|
||||
const at::Scalar& scalar = correction.value();
|
||||
correction_value = scalar.to<double>();
|
||||
}
|
||||
std::vector<int64_t> dims_vec;
|
||||
if (dim.has_value() && dim.value().size() > 0) {
|
||||
dims_vec.assign(dim.value().begin(), dim.value().end());
|
||||
}
|
||||
return detail::_PD_std_impl(*this, dims_vec, correction_value, keepdim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,110 @@
|
||||
// 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <c10/util/OptionalArrayRef.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor sum(const at::Tensor& self,
|
||||
::std::optional<at::ScalarType> dtype = ::std::nullopt) {
|
||||
// Match PyTorch promotion: integer inputs -> int64; others -> keep input
|
||||
// dtype.
|
||||
at::ScalarType resolved_dtype;
|
||||
if (dtype.has_value()) {
|
||||
resolved_dtype = dtype.value();
|
||||
} else {
|
||||
at::ScalarType input_dtype = self.scalar_type();
|
||||
resolved_dtype = at::isIntegralType(input_dtype, /*includeBool=*/true)
|
||||
? at::kLong
|
||||
: input_dtype;
|
||||
}
|
||||
return paddle::experimental::sum(
|
||||
self._PD_GetInner(),
|
||||
{},
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(resolved_dtype),
|
||||
/*keepdim=*/false);
|
||||
}
|
||||
|
||||
inline at::Tensor sum(const at::Tensor& self,
|
||||
at::OptionalIntArrayRef dim,
|
||||
bool keepdim = false,
|
||||
::std::optional<at::ScalarType> dtype = ::std::nullopt) {
|
||||
// Match PyTorch promotion: integer inputs -> int64; others -> keep input
|
||||
// dtype.
|
||||
at::ScalarType resolved_dtype;
|
||||
if (dtype.has_value()) {
|
||||
resolved_dtype = dtype.value();
|
||||
} else {
|
||||
at::ScalarType input_dtype = self.scalar_type();
|
||||
resolved_dtype = at::isIntegralType(input_dtype, /*includeBool=*/true)
|
||||
? at::kLong
|
||||
: input_dtype;
|
||||
}
|
||||
return paddle::experimental::sum(
|
||||
self._PD_GetInner(),
|
||||
dim.has_value() ? dim.value()._PD_ToPaddleIntArray()
|
||||
: paddle::experimental::IntArray(),
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(resolved_dtype),
|
||||
keepdim);
|
||||
}
|
||||
|
||||
inline at::Tensor& sum_out(
|
||||
at::Tensor&
|
||||
out, // NOLINT: intentional non-const reference for output parameter
|
||||
const at::Tensor& self,
|
||||
at::OptionalIntArrayRef dim,
|
||||
bool keepdim = false,
|
||||
::std::optional<at::ScalarType> dtype = ::std::nullopt) {
|
||||
auto res = sum(self, dim, keepdim, dtype);
|
||||
paddle::experimental::assign_out_(res._PD_GetInner(), out._PD_GetInner());
|
||||
return out;
|
||||
}
|
||||
|
||||
inline at::Tensor& sum_out(
|
||||
at::Tensor&
|
||||
out, // NOLINT: intentional non-const reference for output parameter
|
||||
const at::Tensor& self,
|
||||
::std::optional<at::ScalarType> dtype = ::std::nullopt) {
|
||||
auto res = sum(self, dtype);
|
||||
paddle::experimental::assign_out_(res._PD_GetInner(), out._PD_GetInner());
|
||||
return out;
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::sum(::std::optional<at::ScalarType> dtype) const {
|
||||
return at::sum(*this, dtype);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::sum(at::OptionalIntArrayRef dim,
|
||||
bool keepdim,
|
||||
::std::optional<at::ScalarType> dtype) const {
|
||||
return at::sum(*this, dim, keepdim, dtype);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,35 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor t(const Tensor& self) {
|
||||
return self.transpose(0, self.dim() < 2 ? 0 : 1);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::t() const { return at::t(*this); }
|
||||
|
||||
inline at::Tensor& Tensor::t_() const {
|
||||
return transpose_(0, dim() < 2 ? 0 : 1);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,47 @@
|
||||
// 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
|
||||
#include "paddle/common/macros.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
#define TENSOR(T, S) \
|
||||
PADDLE_API Tensor tensor(ArrayRef<T> values, const TensorOptions& options); \
|
||||
inline Tensor tensor(std::initializer_list<T> values, \
|
||||
const TensorOptions& options) { \
|
||||
return at::tensor(ArrayRef<T>(values), options); \
|
||||
} \
|
||||
inline Tensor tensor(T value, const TensorOptions& options) { \
|
||||
return at::tensor(ArrayRef<T>(value), options); \
|
||||
} \
|
||||
inline Tensor tensor(ArrayRef<T> values) { \
|
||||
return at::tensor(std::move(values), at::dtype(k##S)); \
|
||||
} \
|
||||
inline Tensor tensor(std::initializer_list<T> values) { \
|
||||
return at::tensor(ArrayRef<T>(values)); \
|
||||
} \
|
||||
inline Tensor tensor(T value) { return at::tensor(ArrayRef<T>(value)); }
|
||||
AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR)
|
||||
AT_FORALL_COMPLEX_TYPES(TENSOR)
|
||||
#undef TENSOR
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,207 @@
|
||||
// Copyright (c) 2026 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.
|
||||
|
||||
// The file has been adapted from pytorch project
|
||||
// Licensed under BSD-style license -
|
||||
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::vector<at::Tensor> tensor_split(const at::Tensor& self,
|
||||
int64_t sections,
|
||||
int64_t dim = 0) {
|
||||
// Follow PyTorch's tensor_split_sections_symint implementation
|
||||
PD_CHECK(self._PD_GetInner().dims().size() > 0,
|
||||
"tensor_split expected at least a 1-dimensional tensor, but got a "
|
||||
"tensor with ",
|
||||
self._PD_GetInner().dims().size(),
|
||||
" dims");
|
||||
|
||||
PD_CHECK(
|
||||
sections > 0, "number of sections must be larger than 0, got ", sections);
|
||||
|
||||
int64_t dim_size = self._PD_GetInner().dims()[dim];
|
||||
|
||||
// Calculate split sizes: first (dim_size % sections) chunks get size
|
||||
// (dim_size / sections + 1), remaining chunks get size (dim_size / sections)
|
||||
auto min_split_size = dim_size / sections;
|
||||
auto num_splits_one_extra = dim_size % sections;
|
||||
|
||||
std::vector<int64_t> split_sizes;
|
||||
split_sizes.reserve(sections);
|
||||
|
||||
for (int64_t split_idx = 0; split_idx < sections; ++split_idx) {
|
||||
auto split_size = (split_idx < num_splits_one_extra) ? (min_split_size + 1)
|
||||
: min_split_size;
|
||||
split_sizes.push_back(split_size);
|
||||
}
|
||||
|
||||
// Use split with calculated sizes
|
||||
auto outputs =
|
||||
paddle::experimental::split(self._PD_GetInner(), split_sizes, dim);
|
||||
|
||||
std::vector<at::Tensor> at_tensors;
|
||||
at_tensors.reserve(outputs.size());
|
||||
for (const auto& paddle_tensor : outputs) {
|
||||
at_tensors.emplace_back(paddle_tensor);
|
||||
}
|
||||
return at_tensors;
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> tensor_split_symint(const at::Tensor& self,
|
||||
c10::SymInt sections,
|
||||
int64_t dim = 0) {
|
||||
return tensor_split(self, static_cast<int64_t>(sections), dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> tensor_split(const at::Tensor& self,
|
||||
at::IntArrayRef indices,
|
||||
int64_t dim = 0) {
|
||||
// Follow PyTorch's _tensor_split_indices implementation
|
||||
// indices are split positions, not sizes
|
||||
PD_CHECK(self._PD_GetInner().dims().size() > 0,
|
||||
"tensor_split expected at least a 1-dimensional tensor, but got a "
|
||||
"tensor with ",
|
||||
self._PD_GetInner().dims().size(),
|
||||
" dims");
|
||||
|
||||
int64_t num_indices = indices.size();
|
||||
int64_t dim_size = self._PD_GetInner().dims()[dim];
|
||||
|
||||
// Convert indices (positions) to sizes
|
||||
std::vector<int64_t> split_sizes;
|
||||
split_sizes.reserve(num_indices + 1);
|
||||
|
||||
int64_t start_idx = 0;
|
||||
for (int64_t i = 0; i < num_indices; ++i) {
|
||||
int64_t end_idx = indices[i];
|
||||
// Handle negative indices
|
||||
if (end_idx < 0) {
|
||||
end_idx += dim_size;
|
||||
}
|
||||
// Clamp to valid range
|
||||
end_idx = std::max(start_idx, std::min(end_idx, dim_size));
|
||||
split_sizes.push_back(end_idx - start_idx);
|
||||
start_idx = end_idx;
|
||||
}
|
||||
// Add the last segment
|
||||
split_sizes.push_back(dim_size - start_idx);
|
||||
|
||||
// Use split with calculated sizes
|
||||
auto outputs =
|
||||
paddle::experimental::split(self._PD_GetInner(), split_sizes, dim);
|
||||
|
||||
std::vector<at::Tensor> at_tensors;
|
||||
at_tensors.reserve(outputs.size());
|
||||
for (const auto& paddle_tensor : outputs) {
|
||||
at_tensors.emplace_back(paddle_tensor);
|
||||
}
|
||||
return at_tensors;
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> tensor_split_symint(const at::Tensor& self,
|
||||
c10::SymIntArrayRef indices,
|
||||
int64_t dim = 0) {
|
||||
return tensor_split(
|
||||
self,
|
||||
at::IntArrayRef(reinterpret_cast<const int64_t*>(indices.data()),
|
||||
indices.size()),
|
||||
dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> tensor_split(
|
||||
const at::Tensor& self,
|
||||
const at::Tensor& tensor_indices_or_sections,
|
||||
int64_t dim = 0) {
|
||||
// Follow PyTorch's validation and implementation
|
||||
PD_CHECK(self._PD_GetInner().dims().size() > 0,
|
||||
"tensor_split expected at least a 1-dimensional tensor, but got a "
|
||||
"tensor with ",
|
||||
self._PD_GetInner().dims().size(),
|
||||
" dims");
|
||||
|
||||
auto split_device = tensor_indices_or_sections.device();
|
||||
PD_CHECK(split_device.is_cpu(),
|
||||
"tensor_split expected tensor_indices_or_sections to be on cpu, but "
|
||||
"it's on ",
|
||||
split_device);
|
||||
|
||||
auto split_dtype = tensor_indices_or_sections.scalar_type();
|
||||
PD_CHECK(split_dtype == at::kLong,
|
||||
"tensor_split expected tensor_indices_or_sections to have dtype of "
|
||||
"long, but got ",
|
||||
split_dtype);
|
||||
|
||||
auto split_dim = tensor_indices_or_sections.dim();
|
||||
PD_CHECK(split_dim == 1 || split_dim == 0,
|
||||
"tensor_split expected tensor_indices_or_sections to be a "
|
||||
"zero-dimensional or one-dimensional tensor, but got a tensor with ",
|
||||
split_dim,
|
||||
" dims");
|
||||
|
||||
if (split_dim == 0) {
|
||||
// 0-dimensional tensor: treat as sections
|
||||
int64_t sections = tensor_indices_or_sections.item<int64_t>();
|
||||
return tensor_split(self, sections, dim);
|
||||
} else {
|
||||
// 1-dimensional tensor: treat as indices
|
||||
// Need to handle non-contiguous tensors properly
|
||||
const PaddleTensor& paddle_tensor =
|
||||
tensor_indices_or_sections._PD_GetInner();
|
||||
const int64_t* indices_data = paddle_tensor.data<int64_t>();
|
||||
auto stride = tensor_indices_or_sections.stride(0);
|
||||
auto numel = tensor_indices_or_sections.numel();
|
||||
std::vector<int64_t> indices(numel);
|
||||
for (int64_t offset = 0; offset < numel; ++offset) {
|
||||
// indices tensor could be non-contiguous
|
||||
indices[offset] = *(indices_data + offset * stride);
|
||||
}
|
||||
return tensor_split(self, at::IntArrayRef(indices), dim);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::tensor_split(int64_t sections,
|
||||
int64_t dim = 0) const {
|
||||
return at::tensor_split(*this, sections, dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::tensor_split_symint(
|
||||
c10::SymInt sections, int64_t dim = 0) const {
|
||||
return at::tensor_split_symint(*this, sections, dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::tensor_split(at::IntArrayRef indices,
|
||||
int64_t dim = 0) const {
|
||||
return at::tensor_split(*this, indices, dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::tensor_split_symint(
|
||||
c10::SymIntArrayRef indices, int64_t dim = 0) const {
|
||||
return at::tensor_split_symint(*this, indices, dim);
|
||||
}
|
||||
|
||||
inline std::vector<at::Tensor> Tensor::tensor_split(
|
||||
const at::Tensor& tensor_indices_or_sections, int64_t dim = 0) const {
|
||||
return at::tensor_split(*this, tensor_indices_or_sections, dim);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,133 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/Device.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <utils/scalar_type_conversion.h>
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
#include "paddle/phi/common/place.h"
|
||||
|
||||
namespace at {
|
||||
|
||||
// Overload 1: to(TensorOptions, non_blocking, copy, memory_format)
|
||||
inline at::Tensor Tensor::to(
|
||||
at::TensorOptions options,
|
||||
bool non_blocking,
|
||||
bool copy,
|
||||
::std::optional<at::MemoryFormat> memory_format) const {
|
||||
// Handle device transfer
|
||||
PaddleTensor result = tensor_;
|
||||
bool materialized_copy = false;
|
||||
|
||||
if (options.has_device()) {
|
||||
const c10::Device& dev = options.device();
|
||||
phi::Place place;
|
||||
switch (dev.type()) {
|
||||
case c10::DeviceType::CPU:
|
||||
case c10::DeviceType::CUDA:
|
||||
case c10::DeviceType::XPU:
|
||||
case c10::DeviceType::IPU:
|
||||
case c10::DeviceType::CUSTOM:
|
||||
place = dev._PD_GetInner();
|
||||
break;
|
||||
default:
|
||||
PD_THROW("Unsupported device type: ", dev.type());
|
||||
break;
|
||||
}
|
||||
if (place != tensor_.place()) {
|
||||
result = result.copy_to(place, /*blocking=*/!non_blocking);
|
||||
materialized_copy = true;
|
||||
}
|
||||
}
|
||||
|
||||
// Handle dtype cast
|
||||
if (options.has_dtype()) {
|
||||
auto target_dtype =
|
||||
compat::_PD_AtenScalarTypeToPhiDataType(options.dtype());
|
||||
if (target_dtype != result.dtype()) {
|
||||
result = paddle::experimental::cast(result, target_dtype);
|
||||
materialized_copy = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (copy && !materialized_copy) {
|
||||
result = paddle::experimental::assign(result);
|
||||
}
|
||||
|
||||
return at::Tensor(result);
|
||||
}
|
||||
|
||||
// Overload 2: to(optional<ScalarType>, optional<Layout>, optional<Device>,
|
||||
// optional<bool> pin_memory, non_blocking, copy, memory_format)
|
||||
inline at::Tensor Tensor::to(
|
||||
::std::optional<at::ScalarType> dtype,
|
||||
::std::optional<at::Layout> layout,
|
||||
::std::optional<at::Device> device,
|
||||
::std::optional<bool> pin_memory,
|
||||
bool non_blocking,
|
||||
bool copy,
|
||||
::std::optional<at::MemoryFormat> memory_format) const {
|
||||
at::TensorOptions options;
|
||||
if (dtype.has_value()) {
|
||||
options = options.dtype(dtype.value());
|
||||
}
|
||||
if (device.has_value()) {
|
||||
options = options.device(device.value());
|
||||
}
|
||||
if (pin_memory.has_value() && pin_memory.value()) {
|
||||
options = options.pinned_memory(true);
|
||||
}
|
||||
return to(options, non_blocking, copy, memory_format);
|
||||
}
|
||||
|
||||
// Overload 3: to(Device, ScalarType, non_blocking, copy, memory_format)
|
||||
inline at::Tensor Tensor::to(
|
||||
at::Device device,
|
||||
at::ScalarType dtype,
|
||||
bool non_blocking,
|
||||
bool copy,
|
||||
::std::optional<at::MemoryFormat> memory_format) const {
|
||||
at::TensorOptions options = at::TensorOptions().device(device).dtype(dtype);
|
||||
return to(options, non_blocking, copy, memory_format);
|
||||
}
|
||||
|
||||
// Overload 4: to(ScalarType, non_blocking, copy, memory_format)
|
||||
inline at::Tensor Tensor::to(
|
||||
at::ScalarType dtype,
|
||||
bool non_blocking,
|
||||
bool copy,
|
||||
::std::optional<at::MemoryFormat> memory_format) const {
|
||||
auto target_dtype = compat::_PD_AtenScalarTypeToPhiDataType(dtype);
|
||||
if (!copy && target_dtype == tensor_.dtype()) {
|
||||
return *this;
|
||||
}
|
||||
return at::Tensor(paddle::experimental::cast(tensor_, target_dtype));
|
||||
}
|
||||
|
||||
// Overload 5: to(const Tensor& other, non_blocking, copy, memory_format)
|
||||
inline at::Tensor Tensor::to(
|
||||
const at::Tensor& other,
|
||||
bool non_blocking,
|
||||
bool copy,
|
||||
::std::optional<at::MemoryFormat> memory_format) const {
|
||||
at::TensorOptions options =
|
||||
at::TensorOptions().device(other.device()).dtype(other.scalar_type());
|
||||
return to(options, non_blocking, copy, memory_format);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,88 @@
|
||||
// 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 <ATen/core/Tensor.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <limits>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/phi/api/include/api.h"
|
||||
|
||||
namespace at::detail {
|
||||
|
||||
inline int _PD_normalize_transpose_dim(int64_t dim,
|
||||
int64_t ndim,
|
||||
const char* name) {
|
||||
int64_t normalized = dim;
|
||||
if (normalized < 0) {
|
||||
normalized += ndim;
|
||||
}
|
||||
|
||||
PD_CHECK(normalized >= 0 && normalized < ndim, name, " out of range");
|
||||
PD_CHECK(normalized <= static_cast<int64_t>(std::numeric_limits<int>::max()),
|
||||
name,
|
||||
" out of int range");
|
||||
return static_cast<int>(normalized);
|
||||
}
|
||||
|
||||
inline std::vector<int> _PD_make_transpose_perm(int64_t ndim, int d0, int d1) {
|
||||
PD_CHECK(ndim <= static_cast<int64_t>(std::numeric_limits<int>::max()),
|
||||
"tensor rank out of int range");
|
||||
|
||||
std::vector<int> perm(static_cast<size_t>(ndim));
|
||||
for (int64_t i = 0; i < ndim; ++i) {
|
||||
perm[static_cast<size_t>(i)] = static_cast<int>(i);
|
||||
}
|
||||
std::swap(perm[d0], perm[d1]);
|
||||
return perm;
|
||||
}
|
||||
|
||||
} // namespace at::detail
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor transpose(const at::Tensor& self,
|
||||
int64_t dim0,
|
||||
int64_t dim1) {
|
||||
int64_t ndim = self.dim();
|
||||
int d0 = at::detail::_PD_normalize_transpose_dim(dim0, ndim, "dim0");
|
||||
int d1 = at::detail::_PD_normalize_transpose_dim(dim1, ndim, "dim1");
|
||||
auto perm = at::detail::_PD_make_transpose_perm(ndim, d0, d1);
|
||||
|
||||
return paddle::experimental::transpose(self._PD_GetInner(), perm);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::transpose(int64_t dim0, int64_t dim1) const {
|
||||
return at::transpose(*this, dim0, dim1);
|
||||
}
|
||||
|
||||
inline at::Tensor& Tensor::transpose_(int64_t dim0, int64_t dim1) const {
|
||||
int64_t ndim = this->dim();
|
||||
int d0 = at::detail::_PD_normalize_transpose_dim(dim0, ndim, "dim0");
|
||||
int d1 = at::detail::_PD_normalize_transpose_dim(dim1, ndim, "dim1");
|
||||
auto perm = at::detail::_PD_make_transpose_perm(ndim, d0, d1);
|
||||
PaddleTensor& inner = const_cast<PaddleTensor&>(tensor_);
|
||||
paddle::experimental::transpose_(inner, perm);
|
||||
return const_cast<at::Tensor&>(*this);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
@@ -0,0 +1,63 @@
|
||||
// Copyright (c) 2026 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 <ATen/core/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor unflatten(const at::Tensor& self,
|
||||
int64_t dim,
|
||||
at::IntArrayRef sizes) {
|
||||
// Compute the new shape by replacing the dimension at 'dim' with 'sizes'
|
||||
int64_t ndim = self._PD_GetInner().dims().size();
|
||||
int64_t actual_dim = dim < 0 ? dim + ndim : dim;
|
||||
std::vector<int64_t> new_shape;
|
||||
for (int64_t i = 0; i < ndim; ++i) {
|
||||
if (i == actual_dim) {
|
||||
for (auto s : sizes) {
|
||||
new_shape.push_back(s);
|
||||
}
|
||||
} else {
|
||||
new_shape.push_back(self._PD_GetInner().dims()[i]);
|
||||
}
|
||||
}
|
||||
return Tensor(paddle::experimental::reshape(self._PD_GetInner(), new_shape));
|
||||
}
|
||||
|
||||
inline at::Tensor unflatten_symint(const at::Tensor& self,
|
||||
int64_t dim,
|
||||
c10::SymIntArrayRef sizes) {
|
||||
return unflatten(
|
||||
self,
|
||||
dim,
|
||||
at::IntArrayRef(reinterpret_cast<const int64_t*>(sizes.data()),
|
||||
sizes.size()));
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
|
||||
namespace at {
|
||||
|
||||
inline at::Tensor Tensor::unflatten(int64_t dim, at::IntArrayRef sizes) const {
|
||||
return at::unflatten(*this, dim, sizes);
|
||||
}
|
||||
|
||||
inline at::Tensor Tensor::unflatten_symint(int64_t dim,
|
||||
c10::SymIntArrayRef sizes) const {
|
||||
return at::unflatten_symint(*this, dim, sizes);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user