[MPS] Add aminmax op (pytorch#101691)

Pull Request resolved: pytorch#101691 Approved by: https://github.com/malfet

Author: qqaatw

aten/src/ATen/native/Histogram.cpp

@@ -16,8 +16,6 @@
 #include <ATen/ops/_histogramdd_from_bin_tensors.h>
 #include <ATen/ops/_histogramdd_from_bin_tensors_native.h>
 #include <ATen/ops/aminmax.h>
-#include <ATen/ops/amin.h>
-#include <ATen/ops/amax.h>
 #include <ATen/ops/empty.h>
 #include <ATen/ops/histc_native.h>
 #include <ATen/ops/histogram_native.h>
@@ -196,9 +194,8 @@ select_outer_bin_edges(const Tensor& input, c10::optional<c10::ArrayRef<double>>
  // non-empty input
  AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "histogramdd", [&]() {
  if (input.is_mps()) {
- // aminmax has not been implemented on mps.
- Tensor min = at::amin(input, 0);
- Tensor max = at::amax(input, 0);
+ Tensor min, max;
+ std::tie(min, max) = at::aminmax(input, 0);
 
  for (const auto i : c10::irange(N)) {
  leftmost_edges[i] = min[i].item().to<scalar_t>();
@@ -239,18 +236,9 @@ std::pair<double, double> histc_select_outer_bin_edges(const Tensor& input,
  double rightmost_edge = max.to<double>();
 
  if (leftmost_edge == rightmost_edge && input.numel() > 0) {
- if (input.is_mps()) {
- // aminmax has not been implemented on mps.
- Tensor min = at::amin(input);
- Tensor max = at::amax(input);
-
- leftmost_edge = min.item<double>();
- rightmost_edge = max.item<double>();
- } else {
- auto extrema = aminmax(input);
- leftmost_edge = std::get<0>(extrema).item<double>();
- rightmost_edge = std::get<1>(extrema).item<double>();
- }
+ auto extrema = aminmax(input);
+ leftmost_edge = std::get<0>(extrema).item<double>();
+ rightmost_edge = std::get<1>(extrema).item<double>();
  }
 
  if (leftmost_edge == rightmost_edge) {

aten/src/ATen/native/mps/operations/ReduceOps.mm

@@ -15,6 +15,7 @@
 #include <ATen/ops/all_native.h>
 #include <ATen/ops/amax_native.h>
 #include <ATen/ops/amin_native.h>
+#include <ATen/ops/aminmax_native.h>
 #include <ATen/ops/any_native.h>
 #include <ATen/ops/argmax_native.h>
 #include <ATen/ops/argmin_native.h>
@@ -989,6 +990,24 @@ Tensor trace_mps(const Tensor& self) {
  mps::reduction_out_mps(input_t, dim, keepdim, c10::nullopt, output_t, mps::MPSReductionType::AMIN, "amin_out_mps");
 }
 
+TORCH_IMPL_FUNC(aminmax_out_mps)
+(const Tensor& input_t, c10::optional<int64_t> dim_opt, bool keepdim, const Tensor& min_t, const Tensor& max_t) {
+ mps::reduction_out_mps(input_t,
+ dim_opt.has_value() ? OptionalIntArrayRef({*dim_opt}) : c10::nullopt,
+ keepdim,
+ c10::nullopt,
+ min_t,
+ mps::MPSReductionType::AMIN,
+ "aminmax_out_mps_min");
+ mps::reduction_out_mps(input_t,
+ dim_opt.has_value() ? OptionalIntArrayRef({*dim_opt}) : c10::nullopt,
+ keepdim,
+ c10::nullopt,
+ max_t,
+ mps::MPSReductionType::AMAX,
+ "aminmax_out_mps_max");
+}
+
 Tensor prod_mps(const Tensor& self, c10::optional<ScalarType> opt_dtype) {
  std::vector<int64_t> dims(self.dim());
  std::iota(dims.begin(), dims.end(), 0);

aten/src/ATen/native/native_functions.yaml

@@ -3571,6 +3571,7 @@
  structured: True
  dispatch:
  CPU, CUDA: aminmax_out
+ MPS: aminmax_out_mps
 
 - func: _compute_linear_combination(Tensor input, Tensor coefficients) -> Tensor
  dispatch:

test/test_mps.py

@@ -77,6 +77,7 @@ def mps_ops_grad_modifier(ops):
  'cdist': [torch.float32],
  'masked.scatter': [torch.float16, torch.float32],
  'index_fill': [torch.float16, torch.float32], # missing `aten::_unique`.
+ 'aminmax': [torch.float32],
 
  # Correctness issues
  'atanh': [torch.float32],
@@ -394,7 +395,6 @@ def mps_ops_modifier(ops):
  'rounddecimals_3': None,
  'rounddecimals_0': None,
  '__rsub__': None,
- 'aminmax': None,
  'angle': None,
  'bucketize': None,
  'cauchy_': None,