1 #pragma once
2
3 #include <ATen/core/Tensor.h>
4 #include <ATen/native/TypeProperties.h>
5 #include <ATen/ScalarOps.h>
6
7 #ifndef AT_PER_OPERATOR_HEADERS
8 #include <ATen/NativeFunctions.h>
9 #else
10 #include <ATen/ops/result_type.h>
11 #endif
12
13 namespace at::native {
14
15 // original values given by raw_*. If an original value is not contiguous, will make a contiguous copy to
16 // the corresponding trimmed_* value. Additionally, if the dtypes of the boundary and input tensor do not
17 // match, will change them to be a common super type so comparisons are done between the same types.
18 // For any trimmed_* tensor, if its outgoing value matches what it was incoming (typically null), then the
19 // corresponding raw_* version should be used since it was already contiguous of the right type.
searchsorted_maybe_trim_input_tensors(Tensor & trimmed_input,Tensor & trimmed_boundaries,Tensor & trimmed_sorter,const Tensor & raw_input,const Tensor & raw_boundaries,const Tensor & raw_sorter)20 inline void searchsorted_maybe_trim_input_tensors(
21 Tensor& trimmed_input,
22 Tensor& trimmed_boundaries,
23 Tensor& trimmed_sorter,
24 const Tensor& raw_input,
25 const Tensor& raw_boundaries,
26 const Tensor& raw_sorter) {
27 bool in_is_contiguous = raw_input.is_contiguous();
28 bool bd_is_contiguous = raw_boundaries.is_contiguous();
29 bool sort_is_contiguous = raw_sorter.is_contiguous();
30
31 if (!in_is_contiguous) {
32 TORCH_WARN_ONCE("torch.searchsorted(): input value tensor is non-contiguous, this will lower the performance due "
33 "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous input value "
34 "tensor if possible. This message will only appear once per program.");
35 trimmed_input = raw_input.contiguous();
36 }
37 if (!bd_is_contiguous) {
38 TORCH_WARN_ONCE("torch.searchsorted(): boundary tensor is non-contiguous, this will lower the performance due "
39 "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous boundary "
40 "tensor if possible. This message will only appear once per program.");
41 trimmed_boundaries = raw_boundaries.contiguous();
42 }
43 if (!sort_is_contiguous) {
44 TORCH_WARN_ONCE("torch.searchsorted(): sorter tensor is non-contiguous, this will lower the performance due "
45 "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous sorter "
46 "tensor if possible. This message will only appear once per program.");
47 trimmed_sorter = raw_sorter.contiguous();
48 }
49 if (raw_input.dtype() != raw_boundaries.dtype()) {
50 at::native::ResultTypeState state = {};
51 state = at::native::update_result_type_state(raw_boundaries, state);
52 state = at::native::update_result_type_state(raw_input, state);
53 ScalarType common_stype = at::native::result_type(state);
54
55 TORCH_INTERNAL_ASSERT(common_stype != ScalarType::Undefined);
56 if (common_stype != raw_input.scalar_type()) {
57 trimmed_input = in_is_contiguous ? raw_input.to(common_stype) : trimmed_input.to(common_stype);
58 }
59 if (common_stype != raw_boundaries.scalar_type()) {
60 trimmed_boundaries = bd_is_contiguous ? raw_boundaries.to(common_stype) : trimmed_boundaries.to(common_stype);
61 }
62 }
63 }
64
65 /* unused but needed for internal jagged tensor class */
searchsorted_maybe_trim_input_tensors(Tensor & trimmed_input,Tensor & trimmed_boundaries,const Tensor & raw_input,const Tensor & raw_boundaries)66 inline void searchsorted_maybe_trim_input_tensors(
67 Tensor& trimmed_input,
68 Tensor& trimmed_boundaries,
69 const Tensor& raw_input,
70 const Tensor& raw_boundaries) {
71 Tensor trimmed_sorter;
72 Tensor raw_sorter;
73 return searchsorted_maybe_trim_input_tensors(
74 trimmed_input,
75 trimmed_boundaries,
76 trimmed_sorter,
77 raw_input,
78 raw_boundaries,
79 raw_sorter);
80 }
81
searchsorted_dims_matched_before_last_dim(const Tensor & boundaries,const Tensor & input)82 inline bool searchsorted_dims_matched_before_last_dim(const Tensor& boundaries, const Tensor& input) {
83 if (boundaries.dim() != input.dim()) {
84 return false;
85 }
86 const auto& dims_bd = boundaries.sizes();
87 const auto& dims_in = input.sizes();
88 for (int64_t dim = 0; dim + 1 < boundaries.dim(); ++dim) {
89 if (dims_bd[dim] != dims_in[dim]) {
90 return false;
91 }
92 }
93 return true;
94 }
95
searchsorted_scalar_tensor(const Scalar & scalar,const c10::Device & device)96 inline Tensor searchsorted_scalar_tensor(const Scalar& scalar, const c10::Device& device) {
97 auto tensor = c10::scalar_to_tensor(scalar, device);
98 // This is to adopt the scalar promotion rules defined in native/TypeProperties.h
99 // So we have the same type promotion rules as binary operations.
100 tensor.unsafeGetTensorImpl()->set_wrapped_number(true);
101 return tensor;
102 }
103
searchsorted_pre_check(const Tensor & boundaries,const Tensor & input,const Tensor & output,const bool out_int32,const bool right,const std::optional<c10::string_view> side_opt,const Tensor & sorter)104 inline void searchsorted_pre_check(
105 const Tensor& boundaries,
106 const Tensor& input,
107 const Tensor& output,
108 const bool out_int32,
109 const bool right,
110 const std::optional<c10::string_view> side_opt,
111 const Tensor& sorter) {
112 if (side_opt) {
113 const c10::string_view side = *side_opt;
114 TORCH_CHECK(side == "left" || side == "right", "torch.searchsorted(): side can only be 'left' or 'right' but ",
115 "got ", side);
116
117 // assume the user has not explicitly set (right=False, side="right")
118 TORCH_CHECK(!right || side == "right", "torch.searchsorted(): side and right can't be set to opposites, got side "
119 "of ", side, " while right was True");
120 }
121
122 TORCH_CHECK(boundaries.device() == input.device(), "torch.searchsorted(): boundaries and input value tensors ",
123 "should have same device type, but got boundaries tensor device type ", boundaries.device(), " and input value ",
124 "tensor device type ", input.device());
125
126 if (sorter.defined()) {
127 TORCH_CHECK(sorter.device() == boundaries.device(), "torch.searchsorted(): sorter and boundary tensors should ",
128 "have same device type, but got sorter tensor device type ", sorter.device(), " and input value tensor ",
129 "device type ", boundaries.device());
130
131 TORCH_CHECK(sorter.sizes() == boundaries.sizes(), "torch.searchsorted(): boundary and sorter must have the same "
132 "size, but got boundary tensor ", boundaries.sizes(), "and got sorter tensor ", sorter.sizes());
133
134 TORCH_CHECK(sorter.scalar_type() == ScalarType::Long, "torch.searchsorted(): sorter must be a tensor of long ",
135 "dtype but got dtype ", sorter.scalar_type());
136
137 if (sorter.numel() > 0) {
138 auto minmax = sorter.aminmax();
139 int64_t vmin = std::get<0>(minmax).item().toLong();
140 int64_t vmax = std::get<1>(minmax).item().toLong();
141 TORCH_CHECK(vmin >= 0 && vmax < sorter.sizes().back(), "torch.searchsorted(): sorter index out of range");
142 }
143 }
144
145 TORCH_CHECK(input.dim() > 0 || (input.dim() == 0 && input.numel() == 1 && boundaries.dim() == 1),
146 "torch.searchsorted(): input value can be a scalar only when boundaries tensor dimension is 1, but we got ",
147 "boundaries tensor dim(", boundaries.dim(), ") and input value's dim(", input.dim(), ") numel(",
148 input.numel(), ")");
149
150 TORCH_CHECK(boundaries.dim() != 0, "torch.searchsorted(): boundaries tensor should have positive dimension, but ",
151 "got 0 dimension");
152
153 TORCH_CHECK(boundaries.dim() == 1 || searchsorted_dims_matched_before_last_dim(boundaries, input),
154 "torch.searchsorted(): boundaries tensor should be 1 dimension or the first N-1 dimensions of boundaries tensor ",
155 "and input value tensor must match, but we got boundaries tensor ", boundaries.sizes(), " and input value tensor ",
156 input.sizes());
157
158 ScalarType output_dtype = output.scalar_type();
159 TORCH_CHECK(
160 (output_dtype == ScalarType::Long && !out_int32) ||
161 (output_dtype == ScalarType::Int && out_int32),
162 "torch.searchsorted(): output tensor's dtype is wrong, it can only be Int(int32) or Long(int64) depending on ",
163 "whether out_int32 flag is True, but we got output tensor's dtype ", output_dtype,
164 " and out_int32 flag is ", (out_int32 ? "True" : "False"));
165
166 if (out_int32) {
167 TORCH_CHECK(boundaries.sizes().back() < INT_MAX,
168 "torch.searchsorted(): the size of boundaries' last dimension should be less than ", INT_MAX, ", but we got ",
169 boundaries.sizes().back());
170 }
171 }
172
173 } // namespace at::native
174