From 1fc046c3a8eb62690cd78a6da1b62463e9133f6d Mon Sep 17 00:00:00 2001
From: Eugene Brevdo <ebrevdo@google.com>
Date: Wed, 19 Dec 2018 00:09:25 -0800
Subject: [PATCH] [TF port] Disable tests for GetCurrentCpu() on iOS/OS X
 because MacOS has a __cpuid bug.

PiperOrigin-RevId: 226125575
---
 tensorflow/core/platform/port_test.cc  |  7 ++++---
 tensorflow/core/platform/posix/port.cc | 18 +++++++++++-------
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/platform/port_test.cc b/tensorflow/core/platform/port_test.cc
index 33c66a6f25a..0567130e8b9 100644
--- a/tensorflow/core/platform/port_test.cc
+++ b/tensorflow/core/platform/port_test.cc
@@ -35,10 +35,11 @@ TEST(Port, AlignedMalloc) {
 
 TEST(Port, GetCurrentCPU) {
   const int cpu = GetCurrentCPU();
-  // TODO(b/120919972): Re-enable this EXPECT_GE after fixing MacOS Kokoro
-  // failures.
-  // EXPECT_GE(cpu, 0);
+#if !defined(__APPLE__)
+  // GetCurrentCPU does not currently work on MacOS.
+  EXPECT_GE(cpu, 0);
   EXPECT_LT(cpu, NumTotalCPUs());
+#endif
 }
 
 TEST(ConditionVariable, WaitForMilliseconds_Timeout) {
diff --git a/tensorflow/core/platform/posix/port.cc b/tensorflow/core/platform/posix/port.cc
index 0fac8b1a889..ea6066ac7bd 100644
--- a/tensorflow/core/platform/posix/port.cc
+++ b/tensorflow/core/platform/posix/port.cc
@@ -29,7 +29,7 @@ limitations under the License.
 #include <sys/syscall.h>
 #endif
 
-#if !defined(__APPLE__) && (__x86_64__ || __i386__)
+#if (__x86_64__ || __i386__)
 #include <cpuid.h>
 #endif
 
@@ -78,20 +78,24 @@ int NumSchedulableCPUs() {
 
 int NumTotalCPUs() {
   int count = absl::base_internal::NumCPUs();
-  return (count == 0) ? kUnknownCPU : count;
+  return (count <= 0) ? kUnknownCPU : count;
 }
 
 int GetCurrentCPU() {
 #if defined(__linux__) && !defined(__ANDROID__)
   return sched_getcpu();
-#elif defined(__cpuid_count)
   // Attempt to use cpuid on all other platforms.  If that fails, perform a
   // syscall.
-  uint32_t eax, ebx, ecx, edx;
-  __cpuid_count(/*leaf=*/1, /*subleaf=*/0, eax, ebx, ecx, edx);
-  if ((edx & (1 << 9)) != 0) {
+#elif defined(__cpuid) && !defined(__APPLE__)
+  // TODO(b/120919972): __cpuid returns invalid APIC ids on OS X.
+  uint32_t eax = 0;
+  uint32_t ebx = 0;
+  uint32_t ecx = 0;
+  uint32_t edx = 0;
+  __cpuid(/*level=*/1, eax, ebx, ecx, edx);
+  if ((edx & /*bit_APIC=*/(1 << 9)) != 0) {
     // EBX bits 24-31 are APIC ID
-    return static_cast<unsigned int>(ebx >> 24);
+    return (ebx & 0xFF) >> 24;
   }
 #elif defined(__NR_getcpu)
   unsigned int cpu;