[libcxx-commits] [libcxx] [libc++] Add an ABI flag to optimize mersenne_twister_engine (PR #206423)

[Louis Dionne via libcxx-commits]

================
@@ -203,17 +206,43 @@ class mersenne_twister_engine {
           return;
       __x_[0] = result_type(1) << (__w - 1);
     }
+#ifdef _LIBCPP_ABI_VECTORIZED_MERSENNE_TWISTER_ENGINE
+    __update_all_states();
+#endif
+  }
+
+  void __update_state(size_t __i, size_t __k) {
+    const size_t __j         = (__i + 1) % __n;
+    const result_type __mask = __r == _Dt ? result_type(~0) : (result_type(1) << __r) - result_type(1);
+    const result_type __yp   = (__x_[__i] & ~__mask) | (__x_[__j] & __mask);
+    __x_[__i]                = __x_[__k] ^ __rshift<1>(__yp) ^ (__a * (__yp & 1));
+  }
+
+  void __update_state(size_t __i) { __update_state(__i, (__i + __m) % __n); }
+
+  void __update_all_states() {
+    size_t __i = 0;
+    // This is split into two loops to help the compiler vectorize the code
+    for (; __i != (__n - __m); ++__i)
+      __update_state(__i);
+    for (size_t __j = 0; __i != __n; ++__i, ++__j)
+      __update_state(__i, __j);
   }
 
   // generating functions
   [[__nodiscard__]] _LIBCPP_HIDE_FROM_ABI result_type operator()() {
-    const size_t __j         = (__i_ + 1) % __n;
-    const result_type __mask = __r == _Dt ? result_type(~0) : (result_type(1) << __r) - result_type(1);
-    const result_type __yp   = (__x_[__i_] & ~__mask) | (__x_[__j] & __mask);
-    const size_t __k         = (__i_ + __m) % __n;
-    __x_[__i_]               = __x_[__k] ^ __rshift<1>(__yp) ^ (__a * (__yp & 1));
-    result_type __z          = __x_[__i_] ^ (__rshift<__u>(__x_[__i_]) & __d);
-    __i_                     = __j;
+#ifdef _LIBCPP_ABI_VECTORIZED_MERSENNE_TWISTER_ENGINE
----------------
ldionne wrote:

This is technically an ABI break, but it's very pedantic. We'd need a mersenne twister to be passed through an ABI boundary compiled with a different version of libc++, and the `operator()` would have to be called on both sides of the ABI boundary.

I think introducing as an ABI break is the right call, but we should discuss whether this is something we can take unconditionally (as a separate patch).

https://github.com/llvm/llvm-project/pull/206423