phy: optimize precoding for PDSCH DM-RS

fix NEON compilation

phy: fix compilation

Author: xavierarteaga

include/srsran/phy/generic_functions/precoding/channel_precoder.h

@@ -34,7 +34,7 @@ class channel_precoder
   /// of RE per layer of the input buffer.
   /// \remark An assertion is triggered if the precoding matrix dimensions do not match the number of layers of the
   /// input buffer and the number of antenna ports of the output buffer.
-  virtual void apply_precoding(re_buffer_writer<>&            output,
+  virtual void apply_precoding(re_buffer_writer<cbf16_t>&     output,
                                const re_buffer_reader<>&      input,
                                const precoding_weight_matrix& precoding) const = 0;
 

include/srsran/phy/support/resource_grid_writer.h

@@ -85,7 +85,7 @@ class resource_grid_writer : public resource_grid_base
   /// \param[in] symbols Symbols to be written into the resource grid.
   /// \note The RE positions given \c k_init, the number of elements in \c symbols and the \c stride shall be within the
   /// resource grid number of subcarriers.
-  virtual void put(unsigned port, unsigned l, unsigned k_init, unsigned stride, span<const cf_t> symbols) = 0;
+  virtual void put(unsigned port, unsigned l, unsigned k_init, unsigned stride, span<const cbf16_t> symbols) = 0;
 
   /// \brief Gets a read-write view of an OFDM symbol for a given port.
   ///

lib/phy/generic_functions/precoding/channel_precoder_avx2.cpp

@@ -43,9 +43,26 @@ simd_cf_interleaved operator*(const simd_cf_interleaved& re, const simd_cf_t& we
   return _mm256_fmaddsub_ps(re, weight.re, _mm256_mul_ps(_mm256_shuffle_ps(re, re, 0xb1), weight.im));
 }
 
+inline __m128i ps_to_cbf16(simd_cf_interleaved in)
+{
+  const __m256i bias = _mm256_set1_epi32(0x7fff);
+  const __m256i one  = _mm256_set1_epi32(0x1);
+
+  __m256i a_i32 = _mm256_castps_si256(in);
+
+  // Round to nearest even.
+  a_i32 = _mm256_add_epi32(a_i32, _mm256_add_epi32(bias, _mm256_and_si256(_mm256_srli_epi32(a_i32, 16), one)));
+
+  // Shift right 16 bits.
+  a_i32 = _mm256_srai_epi32(a_i32, 16);
+
+  // Pack both parts in 32-bit registers.
+  return _mm_packs_epi32(_mm256_extractf128_si256(a_i32, 0), _mm256_extractf128_si256(a_i32, 1));
+}
+
 } // namespace
 
-void channel_precoder_avx2::apply_precoding_port(span<cf_t>                port_re,
+void channel_precoder_avx2::apply_precoding_port(span<cbf16_t>             port_re,
                                                  const re_buffer_reader<>& input_re,
                                                  span<const cf_t>          port_weights) const
 {
@@ -84,15 +101,16 @@ void channel_precoder_avx2::apply_precoding_port(span<cf_t>                port_
     }
 
     // Store.
-    _mm256_storeu_ps(reinterpret_cast<float*>(&port_re[i_re]), re_out);
+    _mm_storeu_si128(reinterpret_cast<__m128i*>(&port_re[i_re]), ps_to_cbf16(re_out));
   }
 
   for (; i_re != nof_re; ++i_re) {
-    port_re[i_re] = layer_re_view_list[0][i_re] * port_weights[0];
+    cf_t sum = layer_re_view_list[0][i_re] * port_weights[0];
 
     for (unsigned i_layer = 1; i_layer != nof_layers; ++i_layer) {
-      port_re[i_re] += layer_re_view_list[i_layer][i_re] * port_weights[i_layer];
+      sum += layer_re_view_list[i_layer][i_re] * port_weights[i_layer];
     }
+    port_re[i_re] = sum;
   }
 }
 
@@ -181,23 +199,6 @@ static inline void layer4_map_and_ci8_to_cf(simd_cf_interleaved& out_l0,
   from_ci8_to_cf(out_l0, out_l1, out_l2, out_l3, tmp);
 }
 
-inline __m128i ps_to_cbf16(simd_cf_interleaved in)
-{
-  const __m256i bias = _mm256_set1_epi32(0x7fff);
-  const __m256i one  = _mm256_set1_epi32(0x1);
-
-  __m256i a_i32 = _mm256_castps_si256(in);
-
-  // Round to nearest even.
-  a_i32 = _mm256_add_epi32(a_i32, _mm256_add_epi32(bias, _mm256_and_si256(_mm256_srli_epi32(a_i32, 16), one)));
-
-  // Shift right 16 bits.
-  a_i32 = _mm256_srai_epi32(a_i32, 16);
-
-  // Pack both parts in 32-bit registers.
-  return _mm_packs_epi32(_mm256_extractf128_si256(a_i32, 0), _mm256_extractf128_si256(a_i32, 1));
-}
-
 void channel_precoder_avx2::apply_layer_map_and_precoding(re_buffer_writer<cbf16_t>&     output,
                                                           span<const ci8_t>              input,
                                                           const precoding_weight_matrix& precoding) const

lib/phy/generic_functions/precoding/channel_precoder_avx2.h

@@ -22,7 +22,7 @@ class channel_precoder_avx2 : public channel_precoder_impl
 {
 public:
   // See interface for documentation.
-  void apply_precoding_port(span<cf_t>                port_re,
+  void apply_precoding_port(span<cbf16_t>             port_re,
                             const re_buffer_reader<>& input_re,
                             span<const cf_t>          port_weights) const override;
 

lib/phy/generic_functions/precoding/channel_precoder_avx512.cpp

@@ -37,6 +37,27 @@ struct simd_cf_t {
 // Type to hold a set of complex numbers using an AVX512 register, with interleaved real and imaginary parts.
 using simd_cf_interleaved = __m512;
 
+inline __m256i ps_to_cbf16(simd_cf_interleaved in)
+{
+#if __AVX512BF16__
+  return (__m256i)_mm512_cvtneps_pbh(in);
+#else  // __AVX512BF16__
+  const __m512i bias = _mm512_set1_epi32(0x7fff);
+  const __m512i one  = _mm512_set1_epi32(0x1);
+
+  __m512i a_i32 = _mm512_castps_si512(in);
+
+  // Round to nearest even.
+  a_i32 = _mm512_add_epi32(a_i32, _mm512_add_epi32(bias, _mm512_and_si512(_mm512_srli_epi32(a_i32, 16), one)));
+
+  // Shift right 16 bits.
+  a_i32 = _mm512_srli_epi32(a_i32, 16);
+
+  // Pack both parts in 32-bit registers.
+  return _mm512_cvtepi32_epi16(a_i32);
+#endif // __AVX512BF16__
+}
+
 } // namespace
 
 // Multiplication operator for the precoding weights.
@@ -203,28 +224,7 @@ static inline void layer4_map_and_ci8_to_cf(simd_cf_interleaved& out0,
   from_ci8_to_cf(out0, out1, out2, out3, tmp);
 }
 
-inline __m256i ps_to_cbf16(simd_cf_interleaved in)
-{
-#if __AVX512BF16__
-  return (__m256i)_mm512_cvtneps_pbh(in);
-#else  // __AVX512BF16__
-  const __m512i bias = _mm512_set1_epi32(0x7fff);
-  const __m512i one  = _mm512_set1_epi32(0x1);
-
-  __m512i a_i32 = _mm512_castps_si512(in);
-
-  // Round to nearest even.
-  a_i32 = _mm512_add_epi32(a_i32, _mm512_add_epi32(bias, _mm512_and_si512(_mm512_srli_epi32(a_i32, 16), one)));
-
-  // Shift right 16 bits.
-  a_i32 = _mm512_srli_epi32(a_i32, 16);
-
-  // Pack both parts in 32-bit registers.
-  return _mm512_cvtepi32_epi16(a_i32);
-#endif // __AVX512BF16__
-}
-
-void channel_precoder_avx512::apply_precoding_port(span<cf_t>                port_re,
+void channel_precoder_avx512::apply_precoding_port(span<cbf16_t>             port_re,
                                                    const re_buffer_reader<>& input_re,
                                                    span<const cf_t>          port_weights) const
 {
@@ -263,15 +263,15 @@ void channel_precoder_avx512::apply_precoding_port(span<cf_t>                por
     }
 
     // Store.
-    _mm512_storeu_ps(reinterpret_cast<float*>(&port_re[i_re]), re_out);
+    _mm256_storeu_si256(reinterpret_cast<__m256i*>(&port_re[i_re]), ps_to_cbf16(re_out));
   }
 
   for (; i_re != nof_re; ++i_re) {
-    port_re[i_re] = layer_re_view_list[0][i_re] * port_weights[0];
-
+    cf_t sum = layer_re_view_list[0][i_re] * port_weights[0];
     for (unsigned i_layer = 1; i_layer != nof_layers; ++i_layer) {
-      port_re[i_re] += layer_re_view_list[i_layer][i_re] * port_weights[i_layer];
+      sum += layer_re_view_list[i_layer][i_re] * port_weights[i_layer];
     }
+    port_re[i_re] = sum;
   }
 }
 

lib/phy/generic_functions/precoding/channel_precoder_avx512.h

@@ -21,7 +21,7 @@ namespace srsran {
 class channel_precoder_avx512 : public channel_precoder_impl
 {
   // See interface for documentation.
-  void apply_precoding_port(span<cf_t>                port_re,
+  void apply_precoding_port(span<cbf16_t>             port_re,
                             const re_buffer_reader<>& input_re,
                             span<const cf_t>          port_weights) const override;
 

lib/phy/generic_functions/precoding/channel_precoder_generic.cpp

@@ -12,7 +12,7 @@
 
 using namespace srsran;
 
-void channel_precoder_generic::apply_precoding_port(span<cf_t>                port_re,
+void channel_precoder_generic::apply_precoding_port(span<cbf16_t>             port_re,
                                                     const re_buffer_reader<>& input_re,
                                                     span<const cf_t>          port_weights) const
 {
@@ -26,12 +26,13 @@ void channel_precoder_generic::apply_precoding_port(span<cf_t>                po
 
   for (unsigned i_re = 0; i_re != nof_re; ++i_re) {
     // Set the port RE to the contribution of the first layer.
-    port_re[i_re] = layer_re_view_list[0][i_re] * port_weights[0];
+    cf_t sum = layer_re_view_list[0][i_re] * port_weights[0];
 
     for (unsigned i_layer = 1; i_layer != nof_layers; ++i_layer) {
       // Accumulate the contributions of all other layers.
-      port_re[i_re] += layer_re_view_list[i_layer][i_re] * port_weights[i_layer];
+      sum += layer_re_view_list[i_layer][i_re] * port_weights[i_layer];
     }
+    port_re[i_re] = sum;
   }
 }
 

lib/phy/generic_functions/precoding/channel_precoder_generic.h

@@ -21,7 +21,7 @@ namespace srsran {
 class channel_precoder_generic : public channel_precoder_impl
 {
   // See interface for documentation.
-  void apply_precoding_port(span<cf_t>                port_re,
+  void apply_precoding_port(span<cbf16_t>             port_re,
                             const re_buffer_reader<>& input_re,
                             span<const cf_t>          port_weights) const override;
 

lib/phy/generic_functions/precoding/channel_precoder_impl.cpp

@@ -12,7 +12,7 @@
 
 using namespace srsran;
 
-void channel_precoder_impl::apply_precoding(re_buffer_writer<>&            output,
+void channel_precoder_impl::apply_precoding(re_buffer_writer<cbf16_t>&     output,
                                             const re_buffer_reader<>&      input,
                                             const precoding_weight_matrix& precoding) const
 {
@@ -46,7 +46,7 @@ void channel_precoder_impl::apply_precoding(re_buffer_writer<>&            outpu
 
   for (unsigned i_port = 0; i_port != nof_tx_ports; ++i_port) {
     // View of the output RE for a single antenna port.
-    span<cf_t> port_re_view = output.get_slice(i_port);
+    span<cbf16_t> port_re_view = output.get_slice(i_port);
 
     // View of the precoding weights applicable to a single antenna port, i.e., the coefficients applied to each
     // layer for the antenna port.

lib/phy/generic_functions/precoding/channel_precoder_impl.h

@@ -25,7 +25,7 @@ class channel_precoder_impl : public channel_precoder
   explicit channel_precoder_impl() = default;
 
   // See interface for documentation.
-  void apply_precoding(re_buffer_writer<>&            output,
+  void apply_precoding(re_buffer_writer<cbf16_t>&     output,
                        const re_buffer_reader<>&      input,
                        const precoding_weight_matrix& precoding) const override;
 
@@ -35,8 +35,9 @@ class channel_precoder_impl : public channel_precoder
   /// \param[out] port_re   View over the RE of a single antenna port.
   /// \param[in] input     Input symbols, indexed by RE and transmit layer.
   /// \param[in] precoding Precoding coefficients, indexed by layer.
-  virtual void
-  apply_precoding_port(span<cf_t> port_re, const re_buffer_reader<>& input_re, span<const cf_t> port_weights) const = 0;
+  virtual void apply_precoding_port(span<cbf16_t>             port_re,
+                                    const re_buffer_reader<>& input_re,
+                                    span<const cf_t>          port_weights) const = 0;
 };
 
 } // namespace srsran