From 4166764d2752a03493ed975071d4722ac9131208 Mon Sep 17 00:00:00 2001
From: Penporn Koanantakool <penporn@google.com>
Date: Mon, 30 Sep 2024 08:31:34 -0700
Subject: [PATCH] [tsl] Add functions to check if a CPU is x86/aarch64.

PiperOrigin-RevId: 680587338
---
 tsl/platform/BUILD            | 11 +++++++++++
 tsl/platform/cpu_info.h       | 19 ++++++++++++++++++
 tsl/platform/cpu_info_test.cc | 36 +++++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+)
 create mode 100644 tsl/platform/cpu_info_test.cc

diff --git a/tsl/platform/BUILD b/tsl/platform/BUILD
index c31b1d6d5..e20b9af9d 100644
--- a/tsl/platform/BUILD
+++ b/tsl/platform/BUILD
@@ -100,6 +100,17 @@ cc_library(
     ],
 )
 
+tsl_cc_test(
+    name = "cpu_info_test",
+    size = "small",
+    srcs = ["cpu_info_test.cc"],
+    deps = [
+        ":platform_port",
+        ":test",
+        ":test_main",
+    ],
+)
+
 cc_library(
     name = "criticality",
     compatible_with = get_compatible_with_portable(),
diff --git a/tsl/platform/cpu_info.h b/tsl/platform/cpu_info.h
index 68506b1d3..c8d3903ff 100644
--- a/tsl/platform/cpu_info.h
+++ b/tsl/platform/cpu_info.h
@@ -21,6 +21,7 @@ limitations under the License.
 // TODO(ahentz): This is not strictly required here but, for historical
 // reasons, many people depend on cpu_info.h in order to use kLittleEndian.
 #include "tsl/platform/byte_order.h"
+#include "tsl/platform/platform.h"
 
 #if defined(_MSC_VER)
 // included so __cpuidex function is available for GETCPUID on Windows
@@ -150,6 +151,24 @@ bool TestAarch64CPU(Aarch64CPU cpu);
 // Checks CPU registers to return hardware capabilities.
 bool TestCPUFeature(CPUFeature feature);
 
+// Checks whether the current processor is x86.
+constexpr bool IsX86CPU() {
+#ifdef PLATFORM_IS_X86
+  return true;
+#else
+  return false;
+#endif
+}
+
+// Checks whether the current processor is aarch64.
+constexpr bool IsAarch64CPU() {
+#if defined(PLATFORM_IS_ARM64) && !defined(__APPLE__) && !defined(__OpenBSD__)
+  return true;
+#else
+  return false;
+#endif
+}
+
 // Returns CPU Vendor string (i.e. 'GenuineIntel', 'AuthenticAMD', etc.)
 std::string CPUVendorIDString();
 
diff --git a/tsl/platform/cpu_info_test.cc b/tsl/platform/cpu_info_test.cc
new file mode 100644
index 000000000..dbef5a57f
--- /dev/null
+++ b/tsl/platform/cpu_info_test.cc
@@ -0,0 +1,36 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tsl/platform/cpu_info.h"
+
+#include "tsl/platform/test.h"
+
+namespace tsl {
+
+TEST(CPUInfo, CommonX86CPU) {
+  // CPUs from 1999 onwards support SSE.
+  if (port::TestCPUFeature(port::CPUFeature::SSE)) {
+    EXPECT_TRUE(port::IsX86CPU());
+  }
+}
+
+TEST(CPUInfo, Aarch64NeoverseV1CPU) {
+  if (port::TestAarch64CPU(port::Aarch64CPU::ARM_NEOVERSE_V1)) {
+    EXPECT_TRUE(port::IsAarch64CPU());
+  }
+}
+
+}  // namespace tsl