WIP

smirnov-alexey · Nov 22, 2024 · 3d94b6d · 3d94b6d
1 parent 5698947
commit 3d94b6d
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 11 deletions.
diff --git a/src/plugins/intel_npu/src/plugin/npuw/lazy_tensor.cpp b/src/plugins/intel_npu/src/plugin/npuw/lazy_tensor.cpp
@@ -31,6 +31,7 @@ struct Concat {
 struct Unpack {
     LazyTensor w, z, s;
     ov::element::Type type;
+    ov::Shape shape;
 };
 struct Permute {
     LazyTensor tensor;
@@ -109,6 +110,9 @@ std::size_t LazyTensorImpl::get_hash() const {
                               seed ^= op.z.get_hash() + 0x9e3779b9;
                               seed ^= op.s.get_hash() + 0x9e3779b9;
                               seed ^= op.type.hash() + 0x9e3779b9;
+                              for (const auto& dim : op.shape) {
+                                  seed ^= std::hash<std::size_t>()(dim) + 0x9e3779b9;
+                              }
                           }},
                m_transform);
 
@@ -141,7 +145,7 @@ LazyTensorImpl::LazyTensorImpl(Transform&& t) {
                               m_transform = op::Concat{op.tensors, op.axis};
                           },
                           [this](const op::Unpack& op) {
-                              m_transform = op::Unpack{op.w, op.z, op.s, op.type};
+                              m_transform = op::Unpack{op.w, op.z, op.s, op.type, op.shape};
                           }},
                t);
 
@@ -163,7 +167,8 @@ bool LazyTensorImpl::operator==(const LazyTensorImpl& other) const {
                               meta_diff = (op.axis != std::get<op::Concat>(other.m_transform).axis);
                           },
                           [&](const op::Unpack& op) {
-                              meta_diff = (op.type != std::get<op::Unpack>(other.m_transform).type);
+                              meta_diff = (op.type != std::get<op::Unpack>(other.m_transform).type ||
+                                           op.shape != std::get<op::Unpack>(other.m_transform).shape);
                           }},
                m_transform);
 
@@ -200,30 +205,30 @@ ov::Tensor LazyTensorImpl::eval() const {
                               for (const auto& lt : op.tensors) {
                                   to_concat.push_back(lt.eval());
                               }
-                              result = ov::npuw::util::concat(to_concat, op.axis);
+                              result = std::move(ov::npuw::util::concat(to_concat, op.axis));
                           },
                           [&](const op::Unpack& op) {
                               const auto& gti = ov::get_tensor_impl;
                               const auto& tw = op.w.eval();
                               const auto& tz = op.z.eval();
                               const auto& ts = op.s.eval();
                               NPUW_ASSERT(tw);
-                              ov::Tensor dst(op.type, tw.get_shape());
+                              ov::Tensor dst(op.type, op.shape);
                               if (tw && tz && ts) {
                                   ov::npuw::util::unpack(gti(tw), gti(tz), gti(ts), gti(dst));
                               } else if (tw && ts) {
                                   ov::npuw::util::unpack(gti(tw), gti(ts), gti(dst));
                               } else {
                                   NPUW_ASSERT(false && "Unsupported combination");
                               }
-                              result = dst;
+                              result = std::move(dst);
                           },
                           [&](const op::Permute& op) {
-                              result = ov::npuw::util::permute(op.tensor.eval(), op.axes);
+                              result = std::move(ov::npuw::util::permute(op.tensor.eval(), op.axes));
                           },
                           [&](const op::Convert& op) {
                               NPUW_ASSERT(ov::element::f16 == op.type);
-                              result = ov::npuw::util::to_f16(op.tensor.eval());
+                              result = std::move(ov::npuw::util::to_f16(op.tensor.eval()));
                           }},
                m_transform);
 
@@ -236,8 +241,12 @@ LazyTensor::LazyTensor(const std::shared_ptr<ov::op::v0::Constant>& const_ptr)
     : m_impl(std::make_shared<LazyTensorImpl>(op::Const{const_ptr})) {}
 LazyTensor::LazyTensor(const std::vector<LazyTensor>& to_concat, const std::size_t axis)
     : m_impl(std::make_shared<LazyTensorImpl>(op::Concat{to_concat, axis})) {}
-LazyTensor::LazyTensor(const LazyTensor& cw, const LazyTensor& cz, const LazyTensor& cs, const ov::element::Type& type)
-    : m_impl(std::make_shared<LazyTensorImpl>(op::Unpack{cw, cz, cs, type})) {}
+LazyTensor::LazyTensor(const LazyTensor& cw,
+                       const LazyTensor& cz,
+                       const LazyTensor& cs,
+                       const ov::element::Type& type,
+                       const ov::Shape& shape)
+    : m_impl(std::make_shared<LazyTensorImpl>(op::Unpack{cw, cz, cs, type, shape})) {}
 
 LazyTensor LazyTensor::permute(const std::vector<std::size_t>& axes) {
     auto new_lt = std::make_shared<LazyTensorImpl>();

diff --git a/src/plugins/intel_npu/src/plugin/npuw/lazy_tensor.hpp b/src/plugins/intel_npu/src/plugin/npuw/lazy_tensor.hpp
@@ -30,7 +30,8 @@ class LazyTensor {
     LazyTensor(const LazyTensor& cw,
                const LazyTensor& cz,
                const LazyTensor& cs,
-               const ov::element::Type& type);  // construct from unpack
+               const ov::element::Type& type,
+               const ov::Shape& shape);  // construct from unpack
 
     LazyTensor permute(const std::vector<std::size_t>& axes);
     LazyTensor convert(const ov::element::Type& type);

diff --git a/src/plugins/intel_npu/src/plugin/npuw/partitioning/partitioning.cpp b/src/plugins/intel_npu/src/plugin/npuw/partitioning/partitioning.cpp
@@ -1867,7 +1867,8 @@ void Partitioner::optimize(const std::string& func_name) {
                 LazyTensor cw = funcall._lazy_closure[w_idx - f._param_offset];
                 LazyTensor cz = z_idx != -1 ? funcall._lazy_closure[z_idx - f._param_offset] : LazyTensor(ov::Tensor());
                 LazyTensor cs = funcall._lazy_closure[s_idx - f._param_offset];
-                funcall._lazy_closure.push_back(LazyTensor(cw, cz, cs, p.first->get_element_type()));
+                funcall._lazy_closure.push_back(
+                    LazyTensor(cw, cz, cs, p.first->get_element_type(), p.first->get_shape()));
                 // Some of the tensors might be in closure - preserve it's 1:1 idx mapping with _lazy_closure
                 funcall._closure.push_back(ov::Tensor());
             });