xref: /aosp_15_r20/external/pytorch/aten/src/ATen/native/quantized/TensorCompare.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <ATen/ATen.h>
2 #include <ATen/CPUApplyUtils.h>
3 #include <ATen/ExpandUtils.h>
4 #include <ATen/NativeFunctions.h>
5 #include <ATen/native/ReduceOpsUtils.h>
6 #include <ATen/native/TensorCompare.h>
7 #include <c10/util/Exception.h>
8 
9 
10 namespace at::native {
11 
max_quantized_cpu(const Tensor & self)12 Tensor max_quantized_cpu(const Tensor& self) {
13   return std::get<0>(self.reshape({-1}).max(/*dim=*/0));
14 }
15 
max_quantized_unary_out(const Tensor & self,Tensor & out)16 Tensor& max_quantized_unary_out(const Tensor& self, Tensor& out) {
17   // TODO this implementation is inefficient for now.
18   TORCH_CHECK(self.device() == out.device());
19 
20   TORCH_CHECK(canCast(
21       typeMetaToScalarType(self.dtype()),
22       typeMetaToScalarType(out.dtype())));
23   Tensor temp = max_quantized_cpu(self);
24   at::native::resize_output(out, temp.sizes());
25   out.copy_(temp);
26   return out;
27 }
28 
min_quantized_cpu(const Tensor & self)29 Tensor min_quantized_cpu(const Tensor& self) {
30   return std::get<0>(self.reshape({-1}).min(/*dim=*/0));
31 }
32 
min_quantized_unary_out(const Tensor & self,Tensor & out)33 Tensor& min_quantized_unary_out(const Tensor& self, Tensor& out) {
34   // TODO this implementation is inefficient for now.
35   TORCH_CHECK(self.device() == out.device());
36 
37   TORCH_CHECK(canCast(
38       typeMetaToScalarType(self.dtype()),
39       typeMetaToScalarType(out.dtype())));
40   Tensor temp = min_quantized_cpu(self);
41   at::native::resize_output(out, temp.sizes());
42   out.copy_(temp);
43   return out;
44 }
45 
46 // TODO: move to TensorMath.cpp
47 
sort_quantized_cpu_stable(const Tensor & self,std::optional<bool> stable,int64_t dim,bool descending)48 std::tuple<Tensor, Tensor> sort_quantized_cpu_stable(
49     const Tensor& self,
50     std::optional<bool> stable,
51     int64_t dim,
52     bool descending) {
53   auto [sort_int, sort_indicies] =
54       at::sort(self.int_repr(), stable, dim, descending);
55   return std::forward_as_tuple(
56       at::_make_per_tensor_quantized_tensor(
57           sort_int, self.q_scale(), self.q_zero_point()),
58       sort_indicies);
59 }
60 
61 } // namespace at::native
62