From c02aa15438459a2b148c1c84267fff3e926c2632 Mon Sep 17 00:00:00 2001
From: Guillaume Chatelet <gchatelet@google.com>
Date: Thu, 4 Nov 2021 12:14:31 +0000
Subject: [PATCH] [libc][NFC] Allow memset (and bzero) to be inlined

This allows shipping individual functions without also having to provide
memset or bzero at the expense of bigger functions.

Similar to D113097.

Differential Revision: https://reviews.llvm.org/D113108
---
 libc/src/string/CMakeLists.txt                |  5 +-
 libc/src/string/aarch64/memset.cpp            | 49 -------------
 libc/src/string/bzero.cpp                     |  4 +-
 ...emset_utils.h => memset_implementations.h} | 69 ++++++++++++++++---
 libc/src/string/memset.cpp                    |  6 +-
 5 files changed, 67 insertions(+), 66 deletions(-)
 delete mode 100644 libc/src/string/aarch64/memset.cpp
 rename libc/src/string/memory_utils/{memset_utils.h => memset_implementations.h} (54%)

diff --git a/libc/src/string/CMakeLists.txt b/libc/src/string/CMakeLists.txt
index 463181854e8e..5a130fc6ecc7 100644
--- a/libc/src/string/CMakeLists.txt
+++ b/libc/src/string/CMakeLists.txt
@@ -389,7 +389,7 @@ endif()
 
 function(add_memset memset_name)
   add_implementation(memset ${memset_name}
-    SRCS ${MEMSET_SRC}
+    SRCS ${LIBC_SOURCE_DIR}/src/string/memset.cpp
     HDRS ${LIBC_SOURCE_DIR}/src/string/memset.h
     DEPENDS
       .memory_utils.memory_utils
@@ -401,7 +401,6 @@ function(add_memset memset_name)
 endfunction()
 
 if(${LIBC_TARGET_ARCHITECTURE_IS_X86})
-  set(MEMSET_SRC ${LIBC_SOURCE_DIR}/src/string/memset.cpp)
   add_memset(memset_x86_64_opt_sse2   COMPILE_OPTIONS -march=k8             REQUIRE SSE2)
   add_memset(memset_x86_64_opt_sse4   COMPILE_OPTIONS -march=nehalem        REQUIRE SSE4_2)
   add_memset(memset_x86_64_opt_avx2   COMPILE_OPTIONS -march=haswell        REQUIRE AVX2)
@@ -409,12 +408,10 @@ if(${LIBC_TARGET_ARCHITECTURE_IS_X86})
   add_memset(memset_opt_host          COMPILE_OPTIONS ${LIBC_COMPILE_OPTIONS_NATIVE})
   add_memset(memset)
 elseif(${LIBC_TARGET_ARCHITECTURE_IS_AARCH64})
-  set(MEMSET_SRC ${LIBC_SOURCE_DIR}/src/string/aarch64/memset.cpp)
   add_memset(memset_opt_host          COMPILE_OPTIONS ${LIBC_COMPILE_OPTIONS_NATIVE}
                                       COMPILE_OPTIONS "SHELL:-mllvm --tail-merge-threshold=0")
   add_memset(memset                   COMPILE_OPTIONS "SHELL:-mllvm --tail-merge-threshold=0")
 else()
-  set(MEMSET_SRC ${LIBC_SOURCE_DIR}/src/string/memset.cpp)
   add_memset(memset_opt_host          COMPILE_OPTIONS ${LIBC_COMPILE_OPTIONS_NATIVE})
   add_memset(memset)
 endif()
diff --git a/libc/src/string/aarch64/memset.cpp b/libc/src/string/aarch64/memset.cpp
deleted file mode 100644
index fa66ffe1cc99..000000000000
--- a/libc/src/string/aarch64/memset.cpp
+++ /dev/null
@@ -1,49 +0,0 @@
-//===-- Implementation of memset ------------------------------------------===//
-//
-// 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/string/memset.h"
-#include "src/__support/common.h"
-#include "src/string/memory_utils/memset_utils.h"
-
-namespace __llvm_libc {
-
-using namespace __llvm_libc::aarch64_memset;
-
-inline static void AArch64Memset(char *dst, int value, size_t count) {
-  if (count == 0)
-    return;
-  if (count <= 3) {
-    SplatSet<_1>(dst, value);
-    if (count > 1)
-      SplatSet<Tail<_2>>(dst, value, count);
-    return;
-  }
-  if (count <= 8)
-    return SplatSet<HeadTail<_4>>(dst, value, count);
-  if (count <= 16)
-    return SplatSet<HeadTail<_8>>(dst, value, count);
-  if (count <= 32)
-    return SplatSet<HeadTail<_16>>(dst, value, count);
-  if (count <= 96) {
-    SplatSet<_32>(dst, value);
-    if (count <= 64)
-      return SplatSet<Tail<_32>>(dst, value, count);
-    SplatSet<Skip<32>::Then<_32>>(dst, value);
-    SplatSet<Tail<_32>>(dst, value, count);
-    return;
-  }
-  if (count < 448 || value != 0 || !AArch64ZVA(dst, count))
-    return SplatSet<Align<_16, Arg::_1>::Then<Loop<_64>>>(dst, value, count);
-}
-
-LLVM_LIBC_FUNCTION(void *, memset, (void *dst, int value, size_t count)) {
-  AArch64Memset((char *)dst, value, count);
-  return dst;
-}
-
-} // namespace __llvm_libc
diff --git a/libc/src/string/bzero.cpp b/libc/src/string/bzero.cpp
index 3c76ef64eed5..c57c922f6eff 100644
--- a/libc/src/string/bzero.cpp
+++ b/libc/src/string/bzero.cpp
@@ -8,12 +8,12 @@
 
 #include "src/string/bzero.h"
 #include "src/__support/common.h"
-#include "src/string/memory_utils/memset_utils.h"
+#include "src/string/memory_utils/memset_implementations.h"
 
 namespace __llvm_libc {
 
 LLVM_LIBC_FUNCTION(void, bzero, (void *ptr, size_t count)) {
-  GeneralPurposeMemset(reinterpret_cast<char *>(ptr), 0, count);
+  inline_memset(reinterpret_cast<char *>(ptr), 0, count);
 }
 
 } // namespace __llvm_libc
diff --git a/libc/src/string/memory_utils/memset_utils.h b/libc/src/string/memory_utils/memset_implementations.h
similarity index 54%
rename from libc/src/string/memory_utils/memset_utils.h
rename to libc/src/string/memory_utils/memset_implementations.h
index 666d64943479..e34b13a872fc 100644
--- a/libc/src/string/memory_utils/memset_utils.h
+++ b/libc/src/string/memory_utils/memset_implementations.h
@@ -1,4 +1,4 @@
-//===-- Memset utils --------------------------------------------*- C++ -*-===//
+//===-- Implementation of memset and bzero --------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -6,8 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef LLVM_LIBC_SRC_STRING_MEMORY_UTILS_MEMSET_UTILS_H
-#define LLVM_LIBC_SRC_STRING_MEMORY_UTILS_MEMSET_UTILS_H
+#ifndef LLVM_LIBC_SRC_STRING_MEMORY_UTILS_MEMSET_IMPLEMENTATIONS_H
+#define LLVM_LIBC_SRC_STRING_MEMORY_UTILS_MEMSET_IMPLEMENTATIONS_H
 
 #include "src/__support/architectures.h"
 #include "src/string/memory_utils/elements.h"
@@ -48,13 +48,65 @@ namespace __llvm_libc {
 // advance. SetAlignedBlocks<64> may waste up to 63 Bytes, SetAlignedBlocks<32>
 // may waste up to 31 Bytes. Benchmarks showed that SetAlignedBlocks<64> was not
 // superior for sizes that mattered.
-inline static void GeneralPurposeMemset(char *dst, unsigned char value,
-                                        size_t count) {
+inline static void inline_memset(char *dst, unsigned char value, size_t count) {
 #if defined(LLVM_LIBC_ARCH_X86)
-  using namespace ::__llvm_libc::x86;
+  /////////////////////////////////////////////////////////////////////////////
+  // LLVM_LIBC_ARCH_X86
+  /////////////////////////////////////////////////////////////////////////////
+  using namespace __llvm_libc::x86;
+  if (count == 0)
+    return;
+  if (count == 1)
+    return SplatSet<_1>(dst, value);
+  if (count == 2)
+    return SplatSet<_2>(dst, value);
+  if (count == 3)
+    return SplatSet<_3>(dst, value);
+  if (count <= 8)
+    return SplatSet<HeadTail<_4>>(dst, value, count);
+  if (count <= 16)
+    return SplatSet<HeadTail<_8>>(dst, value, count);
+  if (count <= 32)
+    return SplatSet<HeadTail<_16>>(dst, value, count);
+  if (count <= 64)
+    return SplatSet<HeadTail<_32>>(dst, value, count);
+  if (count <= 128)
+    return SplatSet<HeadTail<_64>>(dst, value, count);
+  return SplatSet<Align<_32, Arg::Dst>::Then<Loop<_32>>>(dst, value, count);
+#elif defined(LLVM_LIBC_ARCH_AARCH64)
+  /////////////////////////////////////////////////////////////////////////////
+  // LLVM_LIBC_ARCH_AARCH64
+  /////////////////////////////////////////////////////////////////////////////
+  using namespace __llvm_libc::aarch64_memset;
+  if (count == 0)
+    return;
+  if (count <= 3) {
+    SplatSet<_1>(dst, value);
+    if (count > 1)
+      SplatSet<Tail<_2>>(dst, value, count);
+    return;
+  }
+  if (count <= 8)
+    return SplatSet<HeadTail<_4>>(dst, value, count);
+  if (count <= 16)
+    return SplatSet<HeadTail<_8>>(dst, value, count);
+  if (count <= 32)
+    return SplatSet<HeadTail<_16>>(dst, value, count);
+  if (count <= 96) {
+    SplatSet<_32>(dst, value);
+    if (count <= 64)
+      return SplatSet<Tail<_32>>(dst, value, count);
+    SplatSet<Skip<32>::Then<_32>>(dst, value);
+    SplatSet<Tail<_32>>(dst, value, count);
+    return;
+  }
+  if (count < 448 || value != 0 || !AArch64ZVA(dst, count))
+    return SplatSet<Align<_16, Arg::_1>::Then<Loop<_64>>>(dst, value, count);
 #else
+  /////////////////////////////////////////////////////////////////////////////
+  // Default
+  /////////////////////////////////////////////////////////////////////////////
   using namespace ::__llvm_libc::scalar;
-#endif
 
   if (count == 0)
     return;
@@ -75,8 +127,9 @@ inline static void GeneralPurposeMemset(char *dst, unsigned char value,
   if (count <= 128)
     return SplatSet<HeadTail<_64>>(dst, value, count);
   return SplatSet<Align<_32, Arg::Dst>::Then<Loop<_32>>>(dst, value, count);
+#endif
 }
 
 } // namespace __llvm_libc
 
-#endif // LLVM_LIBC_SRC_STRING_MEMORY_UTILS_MEMSET_UTILS_H
+#endif // LLVM_LIBC_SRC_STRING_MEMORY_UTILS_MEMSET_IMPLEMENTATIONS_H
diff --git a/libc/src/string/memset.cpp b/libc/src/string/memset.cpp
index 945aeda234e6..549c0742dec7 100644
--- a/libc/src/string/memset.cpp
+++ b/libc/src/string/memset.cpp
@@ -8,13 +8,13 @@
 
 #include "src/string/memset.h"
 #include "src/__support/common.h"
-#include "src/string/memory_utils/memset_utils.h"
+#include "src/string/memory_utils/memset_implementations.h"
 
 namespace __llvm_libc {
 
 LLVM_LIBC_FUNCTION(void *, memset, (void *dst, int value, size_t count)) {
-  GeneralPurposeMemset(reinterpret_cast<char *>(dst),
-                       static_cast<unsigned char>(value), count);
+  inline_memset(reinterpret_cast<char *>(dst),
+                static_cast<unsigned char>(value), count);
   return dst;
 }