Add memory descriptor to the OutputConfig of a Subgraph

maxnick · Oct 25, 2024 · 0524249 · 0524249
1 parent 3d0af3b
commit 0524249
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 13 deletions.
diff --git a/src/plugins/intel_cpu/src/nodes/composite.cpp b/src/plugins/intel_cpu/src/nodes/composite.cpp
@@ -50,7 +50,7 @@ void Composite::selectOptimalPrimitiveDescriptor() {
     }
 
     std::vector<Input::OutputConfig> graphOutputConfig;
-    for (size_t i = 0; i < getParentEdges().size(); i++) {
+    for (size_t i = 0; i < outputShapes.size(); i++) {
         graphOutputConfig.emplace_back(node::Input::OutputConfig{true, true});
     }
 

diff --git a/src/plugins/intel_cpu/src/nodes/input.cpp b/src/plugins/intel_cpu/src/nodes/input.cpp
@@ -430,6 +430,7 @@ Input::Input(const std::shared_ptr<ov::Node>& op,
              const GraphContext::CPtr context,
              OutputConfig config)
     : Input(op, context) {
+    extMemDesc = config.desc;
     m_useParentMemoryDescForOutput = config.useParentMemoryDescForOutput;
     m_isInPlace = config.inPlace;
 }

diff --git a/src/plugins/intel_cpu/src/nodes/input.h b/src/plugins/intel_cpu/src/nodes/input.h
@@ -19,9 +19,17 @@ class Input : public Node {
     };
 
     struct OutputConfig {
+        OutputConfig() = default;
+        OutputConfig(bool useParentMemoryDesc_, bool inPlace_)
+            : useParentMemoryDescForOutput(useParentMemoryDesc_),
+              inPlace(inPlace_) {}
+
+        OutputConfig(MemoryDescPtr desc_, bool inPlace_) : desc(std::move(desc_)), inPlace(inPlace_) {}
+
         // @todo better to use memory desc with any layout and undefined precision
-        bool useParentMemoryDescForOutput;
-        bool inPlace;
+        MemoryDescPtr desc = nullptr;
+        bool useParentMemoryDescForOutput = false;
+        bool inPlace = false;
     };
 
     Input(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context);

diff --git a/src/plugins/intel_cpu/src/nodes/lora.cpp b/src/plugins/intel_cpu/src/nodes/lora.cpp
@@ -54,27 +54,26 @@ void LoRA::selectOptimalPrimitiveDescriptor() {
     }
 
     std::vector<Input::OutputConfig> graphOutputConfig;
-    for (size_t i = 0; i < getParentEdges().size(); i++) {
-        graphOutputConfig.emplace_back(node::Input::OutputConfig{true, true});
-    }
+    // enforce the same memory descriptor on the output as on the input to allow inPlace memory
+    graphOutputConfig.emplace_back(node::Input::OutputConfig{inConfs.front().getMemDesc(), true});
 
     // configure the inner graph to get the information about output memory descriptors
     m_graph.Init(m_body, context, graphInputConfig, graphOutputConfig);
 
     // for the output descriptors, use the configuration of the graph's output nodes
     auto outputDescriptors = m_graph.getOutputMemoryDescriptors();
 
-    std::vector<PortConfig> outConfs;
     const auto& desc = outputDescriptors.front();
 
-    outConfs.emplace_back(desc);
-
-    if (desc->isCompatible(*(inConfs.front().getMemDesc()))) {
-        outConfs.at(0).inPlace(0); // use the memory from the first input inPlace
-    } else {
-        THROW_CPU_NODE_ERR("Unexpected input/output descriptor mismatch"); //FIXME: add enforcing the output descriptor
+    if (!desc->isCompatible(*(inConfs.front().getMemDesc()))) {
+        // just a sanity check
+        THROW_CPU_NODE_ERR("Unexpected input/output descriptor mismatch");
     }
 
+    std::vector<PortConfig> outConfs;
+    outConfs.emplace_back(desc);
+    outConfs.front().inPlace(0); // use the memory from the first input inPlace
+
     const NodeConfig config(inConfs, outConfs);
 
     supportedPrimitiveDescriptors.clear();