[pull] master from phate:master

jlm/hls/backend/rvsdg2rhls/mem-sep.cpp

@@ -184,7 +184,6 @@ MemoryStateSeparation::separateMemoryStates(const rvsdg::LambdaNode & lambdaNode
   const auto lambdaSubregion = lambdaNode.subregion();
   auto & memoryStateArgument = llvm::GetMemoryStateRegionArgument(lambdaNode);
 
-  auto & state_user = *memoryStateArgument.Users().begin();
   auto tracedPointerNodesVector = TracePointerArguments(&lambdaNode);
   for (auto & tp : tracedPointerNodesVector)
   {
@@ -210,27 +209,41 @@ MemoryStateSeparation::separateMemoryStates(const rvsdg::LambdaNode & lambdaNode
   }
 
   const size_t numMemoryStates = tracedPointerNodesVector.size() + 1;
+
+  // Assign memory node ids incrementally, used by both the split and merge
   std::vector<llvm::MemoryNodeId> memoryNodeIds;
   for (size_t i = 0; i < numMemoryStates; ++i)
   {
     memoryNodeIds.push_back(i);
   }
+
   auto & lambdaEntrySplitNode = llvm::LambdaEntryMemoryStateSplitOperation::CreateNode(
       memoryStateArgument,
       numMemoryStates,
       memoryNodeIds);
   auto memoryStates = outputs(&lambdaEntrySplitNode);
 
-  auto state_result = &llvm::GetMemoryStateRegionResult(lambdaNode);
   // handle existing state edge - TODO: remove entirely?
+  // The memory state chain that was previously between the state argument and state result,
+  // should now be attached as a chain between the final state on the split and merge, respectively
   auto common_edge = memoryStates.back();
+  // Divert everything except the lambdaEntrySplit itself, to the lambdaEntrySplit
+  memoryStateArgument.divertUsersWhere(
+      *common_edge,
+      [&](const rvsdg::Input & input)
+      {
+        return &input != lambdaEntrySplitNode.input(0);
+      });
+
+  auto state_result = &llvm::GetMemoryStateRegionResult(lambdaNode);
+  memoryStates.back() = state_result->origin();
+
+  auto & lambdaExitMergeNode = llvm::LambdaExitMemoryStateMergeOperation::CreateNode(
+      *lambdaSubregion,
+      memoryStates,
+      memoryNodeIds);
   memoryStates.pop_back();
-  state_user.divert_to(common_edge);
-  memoryStates.push_back(state_result->origin());
-  auto & merged_state =
-      llvm::LambdaExitMemoryStateMergeOperation::Create(*lambdaSubregion, memoryStates);
-  memoryStates.pop_back();
-  state_result->divert_to(&merged_state);
+  state_result->divert_to(lambdaExitMergeNode.output(0));
 
   for (auto tp : tracedPointerNodesVector)
   {

jlm/llvm/ir/operators/MemoryStateOperations.hpp

@@ -448,19 +448,6 @@ class LambdaExitMemoryStateMergeOperation final : public MemoryStateOperation
              : rvsdg::CreateOpNode<LambdaExitMemoryStateMergeOperation>(operands, memoryNodeIds);
   }
 
-  // FIXME: Deprecated, needs to be removed
-  static rvsdg::Output &
-  Create(rvsdg::Region & region, const std::vector<rvsdg::Output *> & operands)
-  {
-    std::vector<MemoryNodeId> memoryNodeIds;
-    for (size_t i = 0; i < operands.size(); ++i)
-    {
-      memoryNodeIds.push_back(i);
-    }
-
-    return *CreateNode(region, operands, std::move(memoryNodeIds)).output(0);
-  }
-
 private:
   util::BijectiveMap<MemoryNodeId, size_t> MemoryNodeIdToIndex_{};
 };

jlm/mlir/frontend/MlirToJlmConverter.cpp

@@ -730,8 +730,7 @@ MlirToJlmConverter::ConvertOperation(
         LambdaEntryMemstateSplitOp.getNumResults(),
         std::move(memoryNodeIds)));
   }
-  else if (
-      auto LambdaExitMemstateMergeOp =
+  if (auto LambdaExitMemstateMergeOp =
           ::mlir::dyn_cast<::mlir::rvsdg::LambdaExitMemoryStateMerge>(&mlirOperation))
   {
     auto memoryNodeIds =

tests/jlm/hls/backend/rvsdg2rhls/UnusedStateRemovalTests.cpp

@@ -326,11 +326,12 @@ TestInvariantMemoryState()
 
   // LambdaExitMemoryStateMerge node
   std::vector<jlm::rvsdg::Output *> outputs;
-  auto & memoryStateMerge = LambdaExitMemoryStateMergeOperation::Create(
+  auto & memoryStateMerge = LambdaExitMemoryStateMergeOperation::CreateNode(
       *lambda->subregion(),
-      { loadOutput[1], memoryStateSplitNode.output(1) });
+      { loadOutput[1], memoryStateSplitNode.output(1) },
+      { 0, 1 });
 
-  auto lambdaOutput = lambda->finalize({ &memoryStateMerge });
+  auto lambdaOutput = lambda->finalize({ memoryStateMerge.output(0) });
   jlm::rvsdg::GraphExport::Create(*lambdaOutput, "f");
 
   jlm::rvsdg::view(rvsdgModule->Rvsdg(), stdout);

tests/jlm/llvm/opt/InvariantValueRedirectionTests.cpp

@@ -266,12 +266,13 @@ TestCallWithMemoryStateNodes()
     auto gammaOutputMemoryState1 = gammaNode->AddExitVar(gammaInputMemoryState2.branchArgument);
     auto gammaOutputMemoryState2 = gammaNode->AddExitVar(gammaInputMemoryState1.branchArgument);
 
-    auto & lambdaExitMergeResult = LambdaExitMemoryStateMergeOperation::Create(
+    auto & lambdaExitMergeNode = LambdaExitMemoryStateMergeOperation::CreateNode(
         *lambdaNode->subregion(),
-        { gammaOutputMemoryState1.output, gammaOutputMemoryState2.output });
+        { gammaOutputMemoryState1.output, gammaOutputMemoryState2.output },
+        { 0, 1 });
 
-    lambdaOutputTest1 =
-        lambdaNode->finalize({ gammaOutputX.output, ioStateArgument, &lambdaExitMergeResult });
+    lambdaOutputTest1 = lambdaNode->finalize(
+        { gammaOutputX.output, ioStateArgument, lambdaExitMergeNode.output(0) });
   }
 
   jlm::rvsdg::Output * lambdaOutputTest2 = nullptr;
@@ -306,12 +307,14 @@ TestCallWithMemoryStateNodes()
         CallOperation::GetMemoryStateOutput(callNode),
         { 0, 1 });
 
-    auto & lambdaExitMergeResult = LambdaExitMemoryStateMergeOperation::Create(
+    auto & lambdaExitMergeNode = LambdaExitMemoryStateMergeOperation::CreateNode(
         *lambdaNode->subregion(),
-        outputs(&callExitSplitNode));
+        outputs(&callExitSplitNode),
+        { 0, 1 });
 
-    lambdaOutputTest2 = lambdaNode->finalize(
-        { callNode.output(0), &CallOperation::GetIOStateOutput(callNode), &lambdaExitMergeResult });
+    lambdaOutputTest2 = lambdaNode->finalize({ callNode.output(0),
+                                               &CallOperation::GetIOStateOutput(callNode),
+                                               lambdaExitMergeNode.output(0) });
     jlm::rvsdg::GraphExport::Create(*lambdaOutputTest2, "test2");
   }
 
@@ -375,12 +378,13 @@ TestCallWithMissingMemoryStateNodes()
         { memoryStateArgument },
         4);
 
-    auto & lambdaExitMergeResult = LambdaExitMemoryStateMergeOperation::Create(
+    auto & lambdaExitMergeNode = LambdaExitMemoryStateMergeOperation::CreateNode(
         *lambdaNode->subregion(),
-        { storeNode.output(0) });
+        { storeNode.output(0) },
+        { 0 });
 
-    lambdaOutputTest1 =
-        lambdaNode->finalize({ zeroNode.output(0), ioStateArgument, &lambdaExitMergeResult });
+    lambdaOutputTest1 = lambdaNode->finalize(
+        { zeroNode.output(0), ioStateArgument, lambdaExitMergeNode.output(0) });
   }
 
   Output * lambdaOutputTest2 = nullptr;
@@ -409,12 +413,14 @@ TestCallWithMissingMemoryStateNodes()
         CallOperation::GetMemoryStateOutput(callNode),
         { 0 });
 
-    auto & lambdaExitMergeResult = LambdaExitMemoryStateMergeOperation::Create(
+    auto & lambdaExitMergeNode = LambdaExitMemoryStateMergeOperation::CreateNode(
         *lambdaNode->subregion(),
-        outputs(&callExitSplitNode));
+        outputs(&callExitSplitNode),
+        { 0 });
 
-    lambdaOutputTest2 = lambdaNode->finalize(
-        { callNode.output(0), &CallOperation::GetIOStateOutput(callNode), &lambdaExitMergeResult });
+    lambdaOutputTest2 = lambdaNode->finalize({ callNode.output(0),
+                                               &CallOperation::GetIOStateOutput(callNode),
+                                               lambdaExitMergeNode.output(0) });
     GraphExport::Create(*lambdaOutputTest2, "test2");
   }