[MPSInductor] Implement prod reduction (#146396)

malfet · pytorchmergebot · commit 5d81bc3696d6 · 2025-02-04T14:08:04.000Z
Mostly reusing `sum` reduction logic Pull Request resolved: #146396 Approved by: https://github.com/dcci ghstack dependencies: #146369, #146370, #146380, #146389
diff --git a/c10/metal/reduction_utils.h b/c10/metal/reduction_utils.h
@@ -18,6 +18,18 @@ opmath_t<T> threadgroup_sum(threadgroup T* data, unsigned size) {
   return rc;
 }
 
+template <typename T>
+opmath_t<T> threadgroup_prod(threadgroup T* data, unsigned size) {
+  opmath_t<T> rc = 1;
+  // TODO: This should be moved to the callee
+  ::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
+  // TODO: Use `simd_shuffle_down`
+  for (auto idx = 0; idx < size; ++idx) {
+    rc *= data[idx];
+  }
+  return rc;
+}
+
 template <typename T>
 T threadgroup_max(threadgroup T* data, unsigned size) {
   // TODO: This should be moved to the callee
diff --git a/test/inductor/test_mps_basic.py b/test/inductor/test_mps_basic.py
@@ -154,6 +154,7 @@ def inc_(x):
     "test_min_max_reduction_nan",
     "test_nan_to_num",
     "test_pow2",
+    "test_prod",
     "test_randint_int64_mod",
     "test_randn_generator",
     "test_remainder",
diff --git a/torch/_inductor/codegen/mps.py b/torch/_inductor/codegen/mps.py
@@ -428,12 +428,12 @@ def reduction(
             """
             )
             return acc
-        if reduction_type == "sum":
+        if reduction_type in ["prod", "sum"]:
             acc_buf = self._new_accvar(src_dtype, reduction_dim.numel)
             self.body.splice(f"{acc_buf}[{reduction_dim.name}] = {value};")
             return self.cse.generate(
                 self.body,
-                f"c10::metal::threadgroup_sum({acc_buf}, {reduction_dim.numel})",
+                f"c10::metal::threadgroup_{reduction_type}({acc_buf}, {reduction_dim.numel})",
                 dtype=DTYPE_TO_COMPUTATION_DTYPE[dtype],
             )
         if reduction_type in ["max", "min"]:
