WIP vectorized copy

Author: smirnov-alexey

src/plugins/intel_npu/src/plugin/CMakeLists.txt

@@ -50,7 +50,7 @@ cross_compiled_file(${TARGET_NAME}
         ARCH AVX2 ANY
                     npuw/util_xarch.cpp
         API         npuw/util_xarch.hpp
-        NAME        unpack_i4i8 unpack_u4i8 unpack_i4f16 unpack_i4f16_scale unpack_i4f16_z unpack_u4f16 unpack_u4f16_scale_zp unpack_u4f16_asymm_zp unpack_u4f16_z unpack_u4f32 unpack_i8f16 unpack_i8f16_scale unpack_u8f16 to_f16 copy_row_as_column transpose_i4 transpose_f16 transpose_f32
+        NAME        unpack_i4i8 unpack_u4i8 unpack_i4f16 unpack_i4f16_scale unpack_i4f16_z unpack_u4f16 unpack_u4f16_scale_zp unpack_u4f16_asymm_zp unpack_u4f16_z unpack_u4f32 unpack_i8f16 unpack_i8f16_scale unpack_u8f16 to_f16 copy_row_as_column transpose_i4 transpose_f16 transpose_f32 copy
         NAMESPACE   ov::npuw::util::XARCH
 )
 

src/plugins/intel_npu/src/plugin/npuw/util.cpp

@@ -654,6 +654,10 @@ ov::Tensor ov::npuw::util::transpose(const ov::Tensor& t) {
     return tnew;
 }
 
+void ov::npuw::util::copy(const ov::Tensor& src, ov::Tensor& dst) {
+    ov::npuw::util::XARCH::copy(src, dst);
+}
+
 ov::Tensor ov::npuw::util::permute(const ov::Tensor& t, const std::vector<std::size_t>& axes) {
     ov::Shape shape = t.get_shape();
     NPUW_ASSERT(shape.size() == 3);  // Yes, so far only transpose 3D tensors

src/plugins/intel_npu/src/plugin/npuw/util.hpp

@@ -18,6 +18,8 @@ namespace ov {
 namespace npuw {
 namespace util {
 
+void copy(const ov::Tensor& src, ov::Tensor& dst);
+
 bool is_set(const std::size_t sub_idx,
             const std::string& opt,
             const std::size_t real_idx = SIZE_MAX,

src/plugins/intel_npu/src/plugin/npuw/util_xarch.cpp

@@ -1492,6 +1492,31 @@ void ov::npuw::util::XARCH::copy_row_as_column(const ov::SoPtr<ov::ITensor>& fro
 #endif
 }
 
+void ov::npuw::util::XARCH::copy(const ov::Tensor& from, ov::Tensor& to) {
+#if defined(HAVE_AVX2)
+    constexpr uint32_t block_size = sizeof(__m256i) / sizeof(uint32_t);
+    size_t total_bytes = from.get_size() * from.get_element_type().size() / 4;
+    if (from.get_element_type() == ov::element::u4 || from.get_element_type() == ov::element::i4 ||
+        from.get_element_type() == ov::element::f4e2m1 || from.get_element_type() == ov::element::nf4) {
+        total_bytes = from.get_size() / 8;
+    }
+
+    const auto* pSrc = reinterpret_cast<uint32_t*>(from.data());
+    auto* pDst = reinterpret_cast<uint32_t*>(to.data());
+
+    size_t i = 0;
+    for (; i + block_size < total_bytes; i += block_size) {
+        __m256i input = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(pSrc + i));
+        _mm256_storeu_si256(reinterpret_cast<__m256i*>(pDst + i), input);
+    }
+    if (i < total_bytes) {
+        std::memcpy(pDst + i, pSrc + i, (total_bytes - i) * 4);
+    }
+#else
+    from.copy_to(to);
+#endif
+}
+
 void ov::npuw::util::XARCH::transpose_i4(const uint8_t* src, uint8_t* dst, size_t rows, size_t cols) {
 #if defined(HAVE_AVX2)
     size_t c_step = 8;

src/plugins/intel_npu/src/plugin/npuw/util_xarch.hpp

@@ -88,6 +88,8 @@ void transpose_i4(const uint8_t* src, uint8_t* dst, size_t rows, size_t cols);
 void transpose_f16(const uint16_t* src, uint16_t* dst, size_t rows, size_t cols);
 void transpose_f32(const float* src, float* dst, size_t rows, size_t cols);
 
+void copy(const ov::Tensor& src, ov::Tensor& dst);
+
 }  // namespace XARCH
 }  // namespace util
 }  // namespace npuw

src/plugins/intel_npu/src/plugin/npuw/weights_bank.cpp

@@ -115,7 +115,7 @@ void Bank::evaluate_and_allocate() {
 void Bank::evaluate_cpu(Bank::DeviceBank& device_bank, const std::vector<LazyTensor>& to_process) {
     // Note: not locking here. This is a private function, so Bank should handle the locks around it
     // as we lock in evaluate_and_allocate() now.
-    ov::parallel_for(to_process.size(), [&](std::size_t idx) {
+    ov::npuw::util::non_parallel_for(to_process.size(), [&](std::size_t idx) {
         const auto& lt = to_process[idx];
         auto iter_device_registered = device_bank.registered_tensors.find(lt);
         NPUW_ASSERT(iter_device_registered != device_bank.registered_tensors.end() &&
@@ -124,7 +124,8 @@ void Bank::evaluate_cpu(Bank::DeviceBank& device_bank, const std::vector<LazyTen
         auto t = lt.eval();
         device_bank.storage.at(uid).tensor = ov::Tensor(t.get_element_type(), t.get_shape());
         // Get ownership of the weights, might be a mmaped object during import
-        t.copy_to(device_bank.storage.at(uid).tensor);
+        // t.copy_to(device_bank.storage.at(uid).tensor);
+        ov::npuw::util::copy(t, device_bank.storage.at(uid).tensor);
         const_cast<LazyTensor&>(lt).detach();
     });
 }
@@ -172,7 +173,8 @@ void Bank::evaluate_and_allocate_on_device(Bank::DeviceBank& device_bank,
         auto& stored_tensor = device_bank.storage.at(allocated.uid);
 
         auto transformed = stored_tensor.lt.eval();
-        transformed.copy_to(allocated.allocated_tensor);
+        // transformed.copy_to(allocated.allocated_tensor);
+        ov::npuw::util::copy(transformed, allocated.allocated_tensor);
         stored_tensor.tensor = std::move(allocated.allocated_tensor);
 
         // Detach the evaluated LazyTensor from its memory here - when it is 100%