From 148fb2be823204efec8426dd9bab207774c64972 Mon Sep 17 00:00:00 2001
From: George Zlobin <george.zlobin@intel.com>
Date: Mon, 14 Dec 2020 16:18:26 +0300
Subject: [PATCH] make map order dependent

---
 .../vpu/middleend/allocator/allocator.hpp     |  2 +-
 .../src/middleend/allocator/allocator.cpp     | 19 +++++++++++++++----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/inference-engine/src/vpu/graph_transformer/include/vpu/middleend/allocator/allocator.hpp b/inference-engine/src/vpu/graph_transformer/include/vpu/middleend/allocator/allocator.hpp
index 17429282384013..923aa38ae1873b 100644
--- a/inference-engine/src/vpu/graph_transformer/include/vpu/middleend/allocator/allocator.hpp
+++ b/inference-engine/src/vpu/graph_transformer/include/vpu/middleend/allocator/allocator.hpp
@@ -126,7 +126,7 @@ class Allocator final {
 
     DataMap<allocator::MemChunk*> _memChunksPerData;
 
-    std::map<DimValues, int> _offsets;
+    std::map<DimVector, std::map<DimValues, int>> _offsets;
 
     int _blobMemOffset = 0;
     int _inputMemOffset = 0;
diff --git a/inference-engine/src/vpu/graph_transformer/src/middleend/allocator/allocator.cpp b/inference-engine/src/vpu/graph_transformer/src/middleend/allocator/allocator.cpp
index 6556b1064f92af..ef07912e778335 100644
--- a/inference-engine/src/vpu/graph_transformer/src/middleend/allocator/allocator.cpp
+++ b/inference-engine/src/vpu/graph_transformer/src/middleend/allocator/allocator.cpp
@@ -316,15 +316,26 @@ ShapeLocation Allocator::allocateShape(const Data& data) {
     } else {
         // Static allocation
         shapeLocation.dimsLocation = Location::Blob;
+
         // Prevert allocation of same shapes multiple times
+        auto dimOrder = data->desc().dimsOrder().toPermutation();
         auto dimValues = data->desc().dims();
-        auto itr = _offsets.find(dimValues);
-        if (itr != _offsets.end()) {
-            shapeLocation.dimsOffset = itr->second;
+        // Find map of <DimValues, OffsetValue> for current DimOrder
+        auto itr_dim = _offsets.find(dimOrder);
+        if (itr_dim != _offsets.end()) {
+            // Try to find already existing Offset
+            auto itr_dims = itr_dim->second.find(dimValues);
+            if (itr_dims != itr_dim->second.end()) {
+                shapeLocation.dimsOffset = itr_dims->second;
+            } else {
+                shapeLocation.dimsOffset = _blobMemOffset;
+                itr_dim->second.insert({dimValues, shapeLocation.dimsOffset});
+                _blobMemOffset += dimsByteSize;
+            }
         } else {
             shapeLocation.dimsOffset = _blobMemOffset;
-            _offsets.insert({dimValues, shapeLocation.dimsOffset});
             _blobMemOffset += dimsByteSize;
+            _offsets.insert({dimOrder, std::map<DimValues, int>{{dimValues, shapeLocation.dimsOffset}}});
         }
     }