[libc] add CPU_COUNT macro and backing function

Add the macro CPU_COUNT as well as a backing function to implement the
functionality.

Reviewed By: sivachandra

Differential Revision: https://reviews.llvm.org/D135179

libc/config/linux/x86_64/entrypoints.txt

@@ -25,6 +25,7 @@ set(TARGET_LIBC_ENTRYPOINTS
     # sched.h entrypoints
     libc.src.sched.sched_getaffinity
     libc.src.sched.sched_setaffinity
+    libc.src.sched.sched_getcpucount
 
     # string.h entrypoints
     libc.src.string.bcmp

libc/config/linux/x86_64/headers.txt

@@ -6,6 +6,7 @@ set(TARGET_PUBLIC_HEADERS
     libc.include.inttypes
     libc.include.math
     libc.include.pthread
+    libc.include.sched
     libc.include.signal
     libc.include.stdio
     libc.include.stdlib

libc/include/CMakeLists.txt

@@ -206,6 +206,7 @@ add_gen_header(
   DEPENDS
     .llvm_libc_common_h
     .llvm-libc-types.cpu_set_t
+    .llvm-libc-macros.sched_macros
 )
 
 # TODO: Not all platforms will have a include/sys directory. Add the sys

libc/include/llvm-libc-macros/CMakeLists.txt

@@ -14,6 +14,14 @@ add_header(
     file-seek-macros.h
 )
 
+add_header(
+  sched_macros
+  HDR
+    sched-macros.h
+  DEPENDS
+    .linux.sched_macros
+)
+
 add_header(
   signal_macros
   HDR

libc/include/llvm-libc-macros/linux/CMakeLists.txt

@@ -4,6 +4,12 @@ add_header(
     fcntl-macros.h
 )
 
+add_header(
+  sched_macros
+  HDR
+    sched-macros.h
+)
+
 add_header(
   sys_mman_macros
   HDR

libc/include/llvm-libc-macros/linux/sched-macros.h

@@ -0,0 +1,15 @@
+//===-- Definition of macros from sched.h ---------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef __LLVM_LIBC_MACROS_LINUX_SCHED_MACROS_H
+#define __LLVM_LIBC_MACROS_LINUX_SCHED_MACROS_H
+
+#define CPU_COUNT_S(setsize, set) __sched_getcpucount(setsize, set)
+#define CPU_COUNT(set) CPU_COUNT_S(sizeof(cpu_set_t), set)
+
+#endif // __LLVM_LIBC_MACROS_LINUX_SCHED_MACROS_H

libc/include/llvm-libc-macros/sched-macros.h

@@ -0,0 +1,16 @@
+//===-- Macros defined in sched.h header file -----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef __LLVM_LIBC_MACROS_SCHED_MACROS_H
+#define __LLVM_LIBC_MACROS_SCHED_MACROS_H
+
+#ifdef __unix__
+#include "linux/sched-macros.h"
+#endif
+
+#endif // __LLVM_LIBC_MACROS_SCHED_MACROS_H

libc/include/sched.h.def

@@ -10,6 +10,7 @@
 #define LLVM_LIBC_SCHED_H
 
 #include <__llvm-libc-common.h>
+#include <llvm-libc-macros/sched-macros.h>
 
 %%public_api()
 

libc/spec/gnu_ext.td

@@ -1,3 +1,7 @@
+def CpuSetT : NamedType<"cpu_set_t">;
+def CpuSetPtr : PtrType<CpuSetT>;
+def ConstCpuSetPtr : ConstType<CpuSetPtr>;
+
 def GnuExtensions : StandardSpec<"GNUExtensions"> {
   NamedType CookieIOFunctionsT = NamedType<"cookie_io_functions_t">;
   HeaderSpec CType = HeaderSpec<
@@ -29,9 +33,6 @@ def GnuExtensions : StandardSpec<"GNUExtensions"> {
       ]
   >;
 
-  NamedType CpuSetT = NamedType<"cpu_set_t">;
-  PtrType CpuSetPtr = PtrType<CpuSetT>;
-  ConstType ConstCpuSetPtr = ConstType<CpuSetPtr>;
   HeaderSpec Sched = HeaderSpec<
       "sched.h",
       [], // Macros

libc/spec/llvm_libc_ext.td

@@ -32,9 +32,24 @@ def LLVMLibcExt : StandardSpec<"llvm_libc_ext"> {
           >,
       ]
   >;
+
+  HeaderSpec Sched = HeaderSpec<
+      "sched.h",
+      [], // Macros
+      [PidT, SizeTType, CpuSetT], // Types
+      [], // Enumerations
+      [
+        FunctionSpec<
+            "sched_getcpucount",
+            RetValSpec<IntType>,
+            [ArgSpec<SizeTType>, ArgSpec<ConstCpuSetPtr>]
+        >,
+      ]
+  >;
 
   let Headers = [
     String,
+    Sched,
     Assert,
   ];
 }

libc/src/sched/CMakeLists.txt

@@ -15,3 +15,10 @@ add_entrypoint_object(
   DEPENDS
     .${LIBC_TARGET_OS}.sched_setaffinity
 )
+
+add_entrypoint_object(
+  sched_getcpucount
+  ALIAS
+  DEPENDS
+    .${LIBC_TARGET_OS}.sched_getcpucount
+)

libc/src/sched/linux/CMakeLists.txt

@@ -23,3 +23,13 @@ add_entrypoint_object(
     libc.src.__support.OSUtil.osutil
     libc.src.errno.errno
 )
+
+add_entrypoint_object(
+  sched_getcpucount
+  SRCS
+    sched_getcpucount.cpp
+  HDRS
+    ../sched_getcpucount.h
+  DEPENDS
+    libc.include.sched
+)

libc/src/sched/linux/sched_getcpucount.cpp

@@ -0,0 +1,27 @@
+//===-- Implementation of sched_getcpucount -------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/sched/sched_getcpucount.h"
+
+#include "src/__support/common.h"
+
+#include <sched.h>
+#include <stddef.h>
+
+namespace __llvm_libc {
+
+LLVM_LIBC_FUNCTION(int, __sched_getcpucount,
+                   (size_t cpuset_size, const cpu_set_t *mask)) {
+  int result = 0;
+  for (size_t i = 0; i < cpuset_size / sizeof(long); ++i) {
+    result += __builtin_popcountl(mask->__mask[i]);
+  }
+  return result;
+}
+
+} // namespace __llvm_libc

libc/src/sched/sched_getcpucount.h

@@ -0,0 +1,23 @@
+//===-- Implementation header for sched_getcpucount -------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC_UNISTD_SCHED_GETCPUCOUNT_H
+#define LLVM_LIBC_SRC_UNISTD_SCHED_GETCPUCOUNT_H
+
+#include <sched.h>
+#include <stddef.h>
+
+namespace __llvm_libc {
+
+// This function is for internal use in the CPU_COUNT macro, but since that's a
+// macro and will be applied to client files, this must be a public entrypoint.
+int __sched_getcpucount(size_t cpuset_size, const cpu_set_t *mask);
+
+} // namespace __llvm_libc
+
+#endif // LLVM_LIBC_SRC_UNISTD_SCHED_GETCPUCOUNT_H

libc/test/src/sched/CMakeLists.txt

@@ -15,3 +15,23 @@ add_libc_unittest(
     libc.src.sched.sched_setaffinity
     libc.test.errno_setter_matcher
 )
+
+# Since this test depends on a macro defined in the sched header, it won't work
+# without fullbuild.
+if(LLVM_LIBC_FULL_BUILD)
+  add_libc_unittest(
+    cpu_count_test
+    SUITE
+      libc_sched_unittests
+    SRCS
+      cpu_count_test.cpp
+    DEPENDS
+      libc.include.errno
+      libc.include.sched
+      libc.include.sys_syscall
+      libc.src.__support.OSUtil.osutil
+      libc.src.sched.sched_getaffinity
+      libc.src.sched.sched_getcpucount
+      libc.test.errno_setter_matcher
+  )
+endif()

libc/test/src/sched/cpu_count_test.cpp

@@ -0,0 +1,32 @@
+//===-- Unittests for sched_getaffinity and sched_setaffinity -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/__support/OSUtil/syscall.h"
+#include "src/sched/sched_getaffinity.h"
+#include "src/sched/sched_getcpucount.h"
+#include "test/ErrnoSetterMatcher.h"
+
+#include <errno.h>
+#include <sched.h>
+#include <sys/syscall.h>
+
+TEST(LlvmLibcSchedAffinityTest, SmokeTest) {
+  cpu_set_t mask;
+  errno = 0;
+  using __llvm_libc::testing::ErrnoSetterMatcher::Succeeds;
+  pid_t tid = __llvm_libc::syscall_impl(SYS_gettid);
+  ASSERT_GT(tid, pid_t(0));
+  ASSERT_THAT(__llvm_libc::sched_getaffinity(tid, sizeof(cpu_set_t), &mask),
+              Succeeds(0));
+
+  // CPU_COUNT is a macro, but it expands to an LLVM-libc internal function that
+  // needs to be in the appropriate namespace for the test.
+  int num_cpus = __llvm_libc::CPU_COUNT(&mask);
+  ASSERT_GT(num_cpus, 0);
+  ASSERT_LE(num_cpus, int(sizeof(cpu_set_t) * sizeof(unsigned long)));
+}