wehub-resource-sync
76 lines
2.8 KiB
C++
76 lines
2.8 KiB
C++
// 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
|