[libc] CopyAlignedBlocks can now specify alignment on top of block size

[gchatelet]

Summary

This has been requested in D92236

[avieira]

Eh you beat me to it. Had a similar patch I've been measuring and testing. Your's is more complete though :)

LGTM.

[gchatelet]

In D94770#2501005, @avieira wrote:

Eh you beat me to it. Had a similar patch I've been measuring and testing. Your's is more complete though :)

LGTM.

Ha! My apologies. I should have coordinated with you before.
Thx for the review!

libc/src/string/memory_utils/memcpy_utils.h

 static void CopyBlockOverlap(char *__restrict dst, const char *__restrict src,
                              size_t count) {
   CopyBlock<kBlockSize>(dst, src);
   CopyLastBlock<kBlockSize>(dst, src, count);
 }
 
 // Copies `count` bytes by blocks of `kBlockSize` bytes.
 // Copies at the start and end of the buffer are unaligned.
-// Copies in the middle of the buffer are aligned to `kBlockSize`.
+// Copies in the middle of the buffer are aligned to `kAlignment`.
 //
 // e.g. with
 // [12345678123456781234567812345678]
-// [__XXXXXXXXXXXXXXXXXXXXXXXXXXX___]
-// [__XXXXXXXX______________________]
-// [________XXXXXXXX________________]
-// [________________XXXXXXXX________]
-// [_____________________XXXXXXXX___]
-//
-// Precondition: `count > 2 * kBlockSize` for efficiency.
-//               `count >= kBlockSize` for correctness.
-template <size_t kBlockSize>
+// [__XXXXXXXXXXXXXXXXXXXXXXXXXXXX___]
+// [__XXXX___________________________]
+// [_____XXXXXXXX____________________]
+// [_____________XXXXXXXX____________]
+// [_____________________XXXXXXXX____]
+// [______________________XXXXXXXX___]
+//
+// Precondition: `kAlignment <= kBlockSize`
+//               `count > 2 * kBlockSize` for efficiency.
+//               `count >= kAlignment` for correctness.
+template <size_t kBlockSize, size_t kAlignment = kBlockSize>
 static void CopyAlignedBlocks(char *__restrict dst, const char *__restrict src,
                               size_t count) {
-  CopyBlock<kBlockSize>(dst, src); // Copy first block
+  static_assert(is_power2(kAlignment), "kAlignment must be a power of two");
+  static_assert(is_power2(kBlockSize), "kBlockSize must be a power of two");
+  static_assert(kAlignment <= kBlockSize,
+                "kAlignment must be less or equal to block size");
+  CopyBlock<kAlignment>(dst, src); // Copy first block
 
   // Copy aligned blocks
-  const size_t ofla = offset_from_last_aligned<kBlockSize>(src);
+  const size_t ofla = offset_from_last_aligned<kAlignment>(src);
   const size_t limit = count + ofla - kBlockSize;
-  for (size_t offset = kBlockSize; offset < limit; offset += kBlockSize)
-    CopyBlock<kBlockSize>(dst - ofla + offset, src - ofla + offset);
+  for (size_t offset = kAlignment; offset < limit; offset += kBlockSize)
+    CopyBlock<kBlockSize>(dst - ofla + offset,
+                          assume_aligned<kAlignment>(src - ofla + offset));
 
   CopyLastBlock<kBlockSize>(dst, src, count); // Copy last block
 }
 
 } // namespace __llvm_libc
 
 #endif //  LIBC_SRC_STRING_MEMORY_UTILS_MEMCPY_UTILS_H

libc/src/string/memory_utils/utils.h

template <size_t alignment> intptr_t offset_to_next_aligned(const void *ptr) {
   return -reinterpret_cast<uintptr_t>(ptr) & (alignment - 1U);
 }
 
 // Returns the offset from `ptr` to the next cache line.
 static inline intptr_t offset_to_next_cache_line(const void *ptr) {
   return offset_to_next_aligned<LLVM_LIBC_CACHELINE_SIZE>(ptr);
 }
 
+template <size_t alignment, typename T> static T *assume_aligned(T *ptr) {
+  return reinterpret_cast<T *>(__builtin_assume_aligned(ptr, alignment));
+}
+
 } // namespace __llvm_libc
 
 #endif // LLVM_LIBC_SRC_MEMORY_UTILS_H

libc/test/src/string/memory_utils/memcpy_utils_test.cpp

TEST(MemcpyUtilsTest, CopyAlignedBlocks) {
   // + "000000001111"
   // = "011121111111"
   trace.Clear();
   CopyAlignedBlocks<4>(I(0), I(1), 11);
   EXPECT_STREQ(trace.Write(), "11121111111");
   EXPECT_STREQ(trace.Read(), "011121111111");
 }
 
-TEST(MemcpyUtilsTest, MaxReloads) {
+TEST(MemcpyUtilsTest, CopyAlignedBlocksWithAlignment) {
+  auto &trace = GetTrace();
+  // Source is aligned and multiple of alignment.
+  //   "11111111"
+  trace.Clear();
+  CopyAlignedBlocks<8, 4>(I(0), I(0), 8);
+  EXPECT_STREQ(trace.Write(), "22221111");
+  EXPECT_STREQ(trace.Read(), "22221111");
+
+  // Source is aligned and multiple of alignment.
+  //   "111111111"
+  trace.Clear();
+  CopyAlignedBlocks<8, 4>(I(0), I(0), 9);
+  EXPECT_STREQ(trace.Write(), "122211111");
+  EXPECT_STREQ(trace.Read(), "122211111");
+}
+
+TEST(MemcpyUtilsTest, CopyAlignedBlocksMaxReloads) {
   auto &trace = GetTrace();
   for (size_t alignment = 0; alignment < 32; ++alignment) {
     for (size_t count = 64; count < 768; ++count) {
       trace.Clear();
       // We should never reload more than twice when copying from count = 2x32.
       CopyAlignedBlocks<32>(I(alignment), I(0), count);
       const char *const written = trace.Write();
       // First bytes are untouched.
       for (size_t i = 0; i < alignment; ++i)
         EXPECT_EQ(written[i], '0');
       // Next bytes are loaded once or twice but no more.
       for (size_t i = alignment; i < count; ++i) {
         EXPECT_GE(written[i], '1');
         EXPECT_LE(written[i], '2');
       }
     }
   }
 }
 
+TEST(MemcpyUtilsTest, CopyAlignedBlocksWithAlignmentMaxReloads) {
+  auto &trace = GetTrace();
+  for (size_t alignment = 0; alignment < 32; ++alignment) {
+    for (size_t count = 64; count < 768; ++count) {
+      trace.Clear();
+      // We should never reload more than twice when copying from count = 2x32.
+      CopyAlignedBlocks<32, 16>(I(alignment), I(0), count);
+      const char *const written = trace.Write();
+      // First bytes are untouched.
+      for (size_t i = 0; i < alignment; ++i)
+        EXPECT_EQ(written[i], '0');
+      // Next bytes are loaded once or twice but no more.
+      for (size_t i = alignment; i < count; ++i) {
+        EXPECT_GE(written[i], '1');
+        EXPECT_LE(written[i], '2');
+      }
+    }
+  }
+}
+
 } // namespace __llvm_libc