[stream_executor] NFC: Optimize kernel arguments packing

Remove unused packed arguments members that are mostly leftovers from passing shared memory arguments to OpenCL kernels

name                          old cpu/op   new cpu/op   delta
BM_PackDeviceMemoryArgs/4     26.8ns ± 3%  21.8ns ± 8%  -18.63%  (p=0.000 n=20+20)
BM_PackDeviceMemoryArgs/8     33.0ns ± 3%  30.9ns ± 4%   -6.21%  (p=0.000 n=20+19)
BM_PackDeviceMemoryArgs/32     105ns ± 3%    80ns ± 3%  -23.97%  (p=0.000 n=20+18)
BM_PackDeviceMemoryArgs/64     201ns ± 4%   156ns ± 4%  -22.43%  (p=0.000 n=19+20)
BM_PackDeviceMemoryArgs/128    520ns ± 4%   362ns ± 4%  -30.35%  (p=0.000 n=19+20)
BM_PackDeviceMemoryArgs/256    783ns ± 7%   586ns ± 4%  -25.21%  (p=0.000 n=20+20)
BM_PackDeviceMemoryArgs/512   1.58µs ± 4%  1.14µs ± 4%  -27.71%  (p=0.000 n=20+20)
BM_PackDeviceMemoryArgs/1024  3.48µs ± 4%  2.39µs ± 4%  -31.34%  (p=0.000 n=20+20)

PiperOrigin-RevId: 580343747

Author: ezhulenev

xla/stream_executor/kernel.h

@@ -70,6 +70,7 @@ limitations under the License.
 #define XLA_STREAM_EXECUTOR_KERNEL_H_
 
 #include <array>
+#include <cstddef>
 #include <cstdint>
 #include <cstring>
 #include <memory>
@@ -376,7 +377,6 @@ class KernelArgsPackedArray : public KernelArgsPackedArrayBase {
     std::memcpy(generic_arg_storage, &arg, sizeof(T));
 
     argument_addresses_[number_of_argument_addresses_] = generic_arg_storage;
-    argument_sizes_[number_of_argument_addresses_] = sizeof(arg);
     ++number_of_argument_addresses_;
   }
 
@@ -386,7 +386,6 @@ class KernelArgsPackedArray : public KernelArgsPackedArrayBase {
         &device_memory_opaque_pointers_[number_of_argument_addresses_];
     *copy_ptr = arg.opaque();
     argument_addresses_[number_of_argument_addresses_] = copy_ptr;
-    argument_sizes_[number_of_argument_addresses_] = sizeof(void *);
     ++number_of_argument_addresses_;
   }
 
@@ -395,17 +394,13 @@ class KernelArgsPackedArray : public KernelArgsPackedArrayBase {
   // The only significant information about a shared argument is its size, so
   // that is the only parameter in this function.
   void add_shared_bytes(size_t number_of_bytes) {
-    shared_memory_indices_[number_of_shared_memory_arguments_] =
-        number_of_argument_addresses_ + number_of_shared_memory_arguments_;
-    shared_memory_bytes_[number_of_shared_memory_arguments_] = number_of_bytes;
-    ++number_of_shared_memory_arguments_;
     total_shared_memory_bytes_ += number_of_bytes;
   }
 
   // Gets the number of arguments added so far, including shared memory
   // arguments.
   size_t number_of_arguments() const override {
-    return number_of_argument_addresses_ + number_of_shared_memory_arguments_;
+    return number_of_argument_addresses_ + (total_shared_memory_bytes_ > 0);
   }
 
   // Gets the total number of shared memory bytes added so far.
@@ -427,28 +422,15 @@ class KernelArgsPackedArray : public KernelArgsPackedArrayBase {
   std::array<const void *, kNumArgs> argument_addresses_;
 
   // Storage for arguments of templated type.
-  alignas(kMaxGenericArgSize)
+  alignas(std::max_align_t)
       std::array<char, kNumArgs * kMaxGenericArgSize> generic_arguments_;
 
-  // Sizes for non-shared-memory arguments.
-  std::array<size_t, kNumArgs> argument_sizes_;
-
-  // Size in bytes for each shared memory argument.
-  std::array<size_t, kNumArgs> shared_memory_bytes_;
-
-  // Indices in the arguments array for shared memory arguments.
-  std::array<size_t, kNumArgs> shared_memory_indices_;
-
   // Total of all shared memory sizes.
   size_t total_shared_memory_bytes_ = 0;
 
   // Number of significant entries in argument_addresses_ and argument_sizes_.
   size_t number_of_argument_addresses_ = 0;
 
-  // Number of significant entries in shared_memory_bytes_ and
-  // shared_memory_indices_.
-  size_t number_of_shared_memory_arguments_ = 0;
-
   // The number of generic arguments that have been added to generic_arguments_.
   size_t number_of_generic_arguments_ = 0;
 };