Sync C10 headers with updated PyTorch pin (pytorch#23126)

### Summary

Advances the pinned PyTorch revision from
`ea89e4e90dc68302e8ef5cba3ddbdaa9d50d9512` to
`2bb2a9c522e7f944bf98c02c9b85fd959fa0bf0b` and synchronizes the two
vendored C10 headers changed by that range. Keeping the pin and headers
in one change prevents `check-c10-sync` from observing mismatched
revisions.

Authored with assistance from OpenAI Codex.

### Test plan

`.ci/scripts/check_c10_sync.sh`

`lintrunner -a .ci/docker/ci_commit_pins/pytorch.txt
runtime/core/portable_type/c10/torch/headeronly/macros/Macros.h
runtime/core/portable_type/c10/torch/headeronly/util/overflows.h`

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Author: JakeStevens

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-ea89e4e90dc68302e8ef5cba3ddbdaa9d50d9512
+2bb2a9c522e7f944bf98c02c9b85fd959fa0bf0b

backends/arm/test/models/test_torch_functions.py

@@ -94,14 +94,7 @@ def forward(self, *args):
 test_parameters = {test[0]: test[1:] for test in module_tests}
 
 
-@parametrize(
-    "test_data",
-    test_parameters,
-    xfails={
-        "topk": "NotImplementedError: No registered serialization name for <class 'torch.return_types.topk'> found",
-        "sort": "NotImplementedError: No registered serialization name for <class 'torch.return_types.sort'> found",
-    },
-)
+@parametrize("test_data", test_parameters)
 def test_torch_functions_tosa_FP(test_data):
     module, inputs = test_data
     pipeline = TosaPipelineFP[input_t](
@@ -119,15 +112,7 @@ def test_torch_functions_tosa_FP(test_data):
             raise e
 
 
-@parametrize(
-    "test_data",
-    test_parameters,
-    xfails={
-        "topk": "NotImplementedError: No registered serialization name for <class 'torch.return_types.topk'> found",
-        "sort": "NotImplementedError: No registered serialization name for <class 'torch.return_types.sort'> found",
-    },
-    strict=True,
-)
+@parametrize("test_data", test_parameters)
 def test_torch_functions_tosa_INT(test_data):
     module, inputs = test_data
     pipeline = TosaPipelineINT[input_t](

runtime/core/portable_type/c10/torch/headeronly/macros/Macros.h

@@ -471,14 +471,15 @@ __host__ __device__
     (void)(printf(                                                    \
         "[CUDA_KERNEL_ASSERT] " __FILE__ ":" C10_STRINGIZE(           \
             __LINE__) ": %s: block: [%d,%d,%d], thread: [%d,%d,%d]: " \
-                      "Assertion failed: `" #cond "`: " msg "\n",     \
+                      "Assertion failed: `%s`: " msg "\n",            \
         __func__,                                                     \
         blockIdx.x,                                                   \
         blockIdx.y,                                                   \
         blockIdx.z,                                                   \
         threadIdx.x,                                                  \
         threadIdx.y,                                                  \
         threadIdx.z,                                                  \
+        #cond,                                                        \
         ##__VA_ARGS__));                                              \
     (void)(_wassert(                                                  \
                _CRT_WIDE(#cond),                                      \
@@ -554,11 +555,89 @@ __host__ __device__
     abort();                     \
   }
 #else
+#if defined(USE_ROCM) && defined(__HIP_DEVICE_COMPILE__)
+namespace torch::headeronly::detail {
+// Compile-time string fragment for CUDA_KERNEL_ASSERT message pieces.
+template <unsigned N>
+struct RoCmAssertLit {
+  char data[N]{};
+  constexpr RoCmAssertLit(const char (&s)[N]) {
+    for (unsigned i = 0; i < N; ++i) {
+      data[i] = s[i];
+    }
+  }
+  static constexpr unsigned size() {
+    return N;
+  }
+  constexpr unsigned percent() const {
+    unsigned count = 0;
+    for (unsigned i = 0; i < N - 1; ++i) {
+      if (data[i] == '%') {
+        ++count;
+      }
+    }
+    return count;
+  }
+};
+
+// OCKL is_last=1 treats the append as printf format; escape '%'. Size the
+// buffer as Out + NumPercent.
+template <RoCmAssertLit A, RoCmAssertLit B, RoCmAssertLit C>
+constexpr auto rocm_assert_concat() {
+  constexpr unsigned out = A.size() + B.size() + C.size() - 2;
+  constexpr unsigned pct = A.percent() + B.percent() + C.percent();
+  struct {
+    char data[out + pct];
+    unsigned length;
+  } msg{};
+  unsigned j = 0;
+  auto append = [&](const char* s, unsigned n) {
+    for (unsigned k = 0; k < n; ++k) {
+      if (s[k] == '%') {
+        msg.data[j++] = '%';
+      }
+      msg.data[j++] = s[k];
+    }
+  };
+  append(A.data, A.size() - 1);
+  append(B.data, B.size() - 1);
+  append(C.data, C.size());
+  msg.length = j;
+  return msg;
+}
+
+// __host__ __device__: hipcc's device pass still type-checks host-only call
+// sites (e.g. MemoryAccess.cuh LoadWithCast ctors). Host runtime uses #else
+// __assert_fail; this body is not emitted for host codegen.
+template <unsigned N>
+__host__ __device__ inline __attribute__((always_inline, flatten)) void
+rocm_assert_one_shot(const char (&msg)[N], unsigned length) {
+  auto d = __ockl_fprintf_stderr_begin();
+  __ockl_fprintf_append_string_n(d, msg, length, 1);
+  __builtin_trap();
+}
+} // namespace torch::headeronly::detail
+
+#define ROCM_ASSERT_LIT(s) ::torch::headeronly::detail::RoCmAssertLit(s)
+
+#define CUDA_KERNEL_ASSERT(cond)                                        \
+  if C10_UNLIKELY (!(cond)) {                                           \
+    static constexpr auto _rocm_assert_msg =                            \
+        ::torch::headeronly::detail::rocm_assert_concat<                \
+            ROCM_ASSERT_LIT(__FILE__ ":" C10_STRINGIZE(__LINE__) ": "), \
+            ROCM_ASSERT_LIT(__func__),                                  \
+            ROCM_ASSERT_LIT(": Device-side assertion `" #cond           \
+                            "' failed.\n")>();                          \
+    ::torch::headeronly::detail::rocm_assert_one_shot(                  \
+        _rocm_assert_msg.data, _rocm_assert_msg.length);                \
+  }
+#else
 #define CUDA_KERNEL_ASSERT(cond)                                         \
   if (C10_UNLIKELY(!(cond))) {                                           \
     __assert_fail(                                                       \
         #cond, __FILE__, static_cast<unsigned int>(__LINE__), __func__); \
   }
+#endif // defined(USE_ROCM) && defined(__HIP_DEVICE_COMPILE__)
 #define CUDA_KERNEL_ASSERT_MSG(cond, msg)                              \
   if (C10_UNLIKELY(!(cond))) {                                         \
     __assert_fail(                                                     \
@@ -569,14 +648,15 @@ __host__ __device__
     printf(                                                            \
         "[CUDA_KERNEL_ASSERT] " __FILE__ ":" C10_STRINGIZE(            \
             __LINE__) ": %s: block: [%d,%d,%d], thread: [%d,%d,%d]: "  \
-            "Assertion failed: `" #cond "`: " msg "\n",                \
+            "Assertion failed: `%s`: " msg "\n",                        \
         __func__,                                                      \
         blockIdx.x,                                                    \
         blockIdx.y,                                                    \
         blockIdx.z,                                                    \
         threadIdx.x,                                                   \
         threadIdx.y,                                                   \
         threadIdx.z,                                                   \
+        #cond,                                                         \
         ##__VA_ARGS__); \
     __assert_fail(                                                       \
         #cond, __FILE__, static_cast<unsigned int>(__LINE__), __func__); \

runtime/core/portable_type/c10/torch/headeronly/util/overflows.h

@@ -8,7 +8,7 @@
 #include <limits>
 #include <type_traits>
 
-HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+namespace c10 {
 // In some versions of MSVC, there will be a compiler error when building.
 // C4146: unary minus operator applied to unsigned type, result still unsigned
 // C4804: unsafe use of type 'bool' in operation
@@ -43,25 +43,26 @@ std::enable_if_t<std::is_same_v<From, bool>, bool> overflows(
 template <typename To, typename From>
 std::enable_if_t<std::is_integral_v<From> && !std::is_same_v<From, bool>, bool>
 overflows(From f, bool strict_unsigned = false) {
-  using limit = std::numeric_limits<typename scalar_value_type<To>::type>;
+  using limit = std::numeric_limits<
+      typename torch::headeronly::scalar_value_type<To>::type>;
   if constexpr (!limit::is_signed && std::numeric_limits<From>::is_signed) {
     // allow for negative numbers to wrap using two's complement arithmetic.
     // For example, with uint8, this allows for `a - b` to be treated as
     // `a + 255 * b`.
     if (!strict_unsigned) {
-      return greater_than_max<To>(f) ||
-          (is_negative(f) &&
+      return c10::greater_than_max<To>(f) ||
+          (c10::is_negative(f) &&
            -static_cast<uint64_t>(f) > static_cast<uint64_t>(limit::max()));
     }
   }
-  return less_than_lowest<To>(f) || greater_than_max<To>(f);
+  return c10::less_than_lowest<To>(f) || c10::greater_than_max<To>(f);
 }
 
 template <typename To, typename From>
 std::enable_if_t<std::is_floating_point_v<From>, bool> overflows(
     From f,
     bool strict_unsigned [[maybe_unused]] = false) {
-  using ToScalar = typename scalar_value_type<To>::type;
+  using ToScalar = typename torch::headeronly::scalar_value_type<To>::type;
   using limit = std::numeric_limits<ToScalar>;
   if (limit::has_infinity && std::isinf(static_cast<double>(f))) {
     return false;
@@ -93,27 +94,27 @@ C10_CLANG_DIAGNOSTIC_POP()
 #endif
 
 template <typename To, typename From>
-std::enable_if_t<is_complex<From>::value, bool> overflows(
+std::enable_if_t<torch::headeronly::is_complex<From>::value, bool> overflows(
     From f,
     bool strict_unsigned = false) {
   // casts from complex to real are considered to overflow if the
   // imaginary component is non-zero
-  if (!is_complex<To>::value && f.imag() != 0) {
+  if (!torch::headeronly::is_complex<To>::value && f.imag() != 0) {
     return true;
   }
   // Check for overflow componentwise
   // (Technically, the imag overflow check is guaranteed to be false
   // when !is_complex<To>, but any optimizer worth its salt will be
   // able to figure it out.)
-  return overflows<
-             typename scalar_value_type<To>::type,
+  return c10::overflows<
+             typename torch::headeronly::scalar_value_type<To>::type,
              typename From::value_type>(f.real(), strict_unsigned) ||
-      overflows<
-             typename scalar_value_type<To>::type,
+      c10::overflows<
+             typename torch::headeronly::scalar_value_type<To>::type,
              typename From::value_type>(f.imag(), strict_unsigned);
 }
-HIDDEN_NAMESPACE_END(torch, headeronly)
-
-namespace c10 {
-using torch::headeronly::overflows;
 } // namespace c10
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using c10::overflows;
+HIDDEN_NAMESPACE_END(torch, headeronly)