make map order dependent

openvinotoolkit · Dec 14, 2020 · 148fb2b · 148fb2b
1 parent 6871e72
commit 148fb2b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 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
@@ -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
@@ -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}}});
         }
     }