[c10/metal] Add a vectype variant for short/int/long (#145430)

dcci · pytorchmergebot · commit f56c638849ce · 2025-01-23T04:52:56.000Z
Some of the kernels (exp_complex/atan_complex) need the specialization. Pull Request resolved: #145430 Approved by: https://github.com/malfet, https://github.com/jansel
diff --git a/c10/metal/utils.h b/c10/metal/utils.h
@@ -29,6 +29,28 @@ struct vectypes<bfloat> {
   using type2 = bfloat2;
 };
 #endif
+
+template <>
+struct vectypes<short> {
+  using type4 = short4;
+  using type3 = short3;
+  using type2 = short2;
+};
+
+template <>
+struct vectypes<int> {
+  using type4 = int4;
+  using type3 = int3;
+  using type2 = int2;
+};
+
+template <>
+struct vectypes<long> {
+  using type4 = short4;
+  using type3 = short3;
+  using type2 = short2;
+};
+
 } // namespace detail
 
 template <typename T>
