[libc] add CPU_COUNT macro and backing function

[michaelrj]

Summary

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

[sivachandra]

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

13

These macros get activated only in the full build mode. So, do we really need a separate helper like CPU_COUNT_S? The internal function can be llvm libc only (as in full build only) and can take a single argument for the set?

libc/spec/gnu_ext.td

55

This should go into https://github.com/llvm/llvm-project/blob/main/libc/spec/llvm_libc_ext.td. Also, the function is really an internal function - it's name should have the __ prefix.

libc/src/sched/linux/sched_getcpucount.cpp

32

We can use __builtin_popcountl here as this function will only be used with LLVM libc's definition of cpu_set_t.

[michaelrj]

move function definition into llvm_libc_ext.td
rename function to __sched_getcpucount

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

13

CPU_COUNT_S is another macro that linux defines, and given that that exists it makes more sense to have one internal function that takes the same arguments and just chain the macros like this.

libc/spec/gnu_ext.td

55

in that case I should probably also move __llvm_libc_syscall into llvm_libc_ext.td in a followup patch.

[sivachandra]

libc/spec/gnu_ext.td

55

Yes, I missed that.

[michaelrj]

fix the non-full build and add a comment about the namespaced macro

libc/config/linux/x86_64/entrypoints.txt

set(TARGET_LIBC_ENTRYPOINTS
     # fcntl.h entrypoints
     libc.src.fcntl.creat
     libc.src.fcntl.open
     libc.src.fcntl.openat
 
     # 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.src.string.bzero
     libc.src.string.memccpy
     libc.src.string.memchr
     libc.src.string.memcmp
     libc.src.string.memcpy

libc/config/linux/x86_64/headers.txt

 set(TARGET_PUBLIC_HEADERS
     libc.include.assert_h
     libc.include.ctype
     libc.include.errno
     libc.include.fenv
     libc.include.inttypes
     libc.include.math
     libc.include.pthread
+    libc.include.sched
     libc.include.signal
     libc.include.stdio
     libc.include.stdlib
     libc.include.string
     libc.include.sys_mman
     libc.include.sys_syscall
     libc.include.threads
     libc.include.time
     libc.include.unistd
 )

libc/include/CMakeLists.txt

 
 add_gen_header(
   sched
   DEF_FILE sched.h.def
   GEN_HDR sched.h
   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
 # directory and the targets for sys/*.h files conditional to the OS requiring
 # them.
 file(MAKE_DIRECTORY "sys")
 
 add_gen_header(

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

 
 add_header(
   file_seek_macros
   HDR
     file-seek-macros.h
 )
 
 add_header(
+  sched_macros
+  HDR
+    sched-macros.h
+  DEPENDS
+    .linux.sched_macros
+)
+
+add_header(
   signal_macros
   HDR
     signal-macros.h
   DEPENDS
     .linux.signal_macros
 )
 
 add_header(

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

 add_header(
   fcntl_macros
   HDR
     fcntl-macros.h
 )
 
 add_header(
+  sched_macros
+  HDR
+    sched-macros.h
+)
+
+add_header(
   sys_mman_macros
   HDR
     sys-mman-macros.h
 )
 
 add_header(
   sys_stat_macros
   HDR

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

+//===-- 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

-//===-- C standard library header sched.h ---------------------------------===//
+//===-- 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_SCHED_H
-#define LLVM_LIBC_SCHED_H
+#ifndef __LLVM_LIBC_MACROS_SCHED_MACROS_H
+#define __LLVM_LIBC_MACROS_SCHED_MACROS_H
 
-#include <__llvm-libc-common.h>
+#ifdef __unix__
+#include "linux/sched-macros.h"
+#endif
 
-%%public_api()
-
-#endif // LLVM_LIBC_SCHED_H 
+#endif // __LLVM_LIBC_MACROS_SCHED_MACROS_H

libc/include/sched.h.def

 //===-- C standard library header 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_SCHED_H
 #define LLVM_LIBC_SCHED_H
 
 #include <__llvm-libc-common.h>
+#include <llvm-libc-macros/sched-macros.h>
 
 %%public_api()
 
 #endif // LLVM_LIBC_SCHED_H 

libc/spec/gnu_ext.td

HeaderSpec Sched = HeaderSpec<
             RetValSpec<IntType>,
             [ArgSpec<PidT>, ArgSpec<SizeTType>, ArgSpec<CpuSetPtr>]
         >,
         FunctionSpec<
             "sched_setaffinity",
             RetValSpec<IntType>,
             [ArgSpec<PidT>, ArgSpec<SizeTType>, ArgSpec<ConstCpuSetPtr>]
         >,
+        FunctionSpec<
+            "sched_getcpucount",
+            RetValSpec<IntType>,
+            [ArgSpec<SizeTType>, ArgSpec<ConstCpuSetPtr>]
+        >,
       ]
   >;
  
   HeaderSpec String = HeaderSpec<
       "string.h",
       [], // Macros
       [], // Types
       [], // Enumerations

libc/src/sched/CMakeLists.txt

 )
 
 add_entrypoint_object(
   sched_setaffinity
   ALIAS
   DEPENDS
     .${LIBC_TARGET_OS}.sched_setaffinity
 )
+
+add_entrypoint_object(
+  sched_getcpucount
+  ALIAS
+  DEPENDS
+    .${LIBC_TARGET_OS}.sched_getcpucount
+)

libc/src/sched/linux/CMakeLists.txt

add_entrypoint_object(
   HDRS
     ../sched_setaffinity.h
   DEPENDS
     libc.include.errno
     libc.include.sched
     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

+//===-- 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 {
+
+// counts the number of bits that are set in a given number.
+// TODO: replace with builtin_popcount
+size_t bit_count(unsigned long val) {
+  size_t result = 0;
+  for (; val != 0; val >>= 1) {
+    result += val & 1;
+  }
+  return result;
+}
+
+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 += bit_count(mask->__mask[i]);
+  }
+  return result;
+}
+
+} // namespace __llvm_libc

libc/src/sched/sched_getcpucount.h

+//===-- 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

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_setaffinity
     libc.test.errno_setter_matcher
 )
+
+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
+)

libc/test/src/sched/cpu_count_test.cpp

+//===-- 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));
+
+  int num_cpus = __llvm_libc::CPU_COUNT(&mask);
+  ASSERT_GT(num_cpus, 0);
+}