Add LoadNonVolatileOperation::NormalizeLoadMemoryStateJoin() reduction

jlm/llvm/ir/operators/Load.cpp

@@ -317,53 +317,64 @@ perform_multiple_origin_reduction(
   return results;
 }
 
+template<class TMemoryStateMergeOrJoinOperation>
 std::optional<std::vector<rvsdg::Output *>>
-LoadNonVolatileOperation::NormalizeLoadMemoryStateMerge(
+NormalizeMemoryStateMergeOrJoin(
     const LoadNonVolatileOperation & operation,
     const std::vector<rvsdg::Output *> & operands)
 {
+  static_assert(
+      std::is_same_v<TMemoryStateMergeOrJoinOperation, MemoryStateMergeOperation>
+          || std::is_same_v<TMemoryStateMergeOrJoinOperation, MemoryStateJoinOperation>,
+      "Template parameter T must be a MemoryStateMergeOperation or a MemoryStateJoinOperation!");
+
   auto & address = *operands[0];
   const auto oldLoadMemoryStates = std::vector(std::next(operands.begin()), operands.end());
 
-  bool foundMemoryStateMergeOperation = false;
+  bool foundMergeOrJoinOperation = false;
   std::vector<rvsdg::Output *> newLoadMemoryStates;
   for (const auto memoryState : oldLoadMemoryStates)
   {
-    auto [memoryStateMergeNode, memoryStateMergeOperation] =
-        rvsdg::TryGetSimpleNodeAndOptionalOp<MemoryStateMergeOperation>(*memoryState);
-    if (memoryStateMergeOperation)
+    auto [memoryStateNode, memoryStateOperation] =
+        rvsdg::TryGetSimpleNodeAndOptionalOp<TMemoryStateMergeOrJoinOperation>(*memoryState);
+    if (memoryStateOperation)
     {
-      foundMemoryStateMergeOperation = true;
-      auto memoryStateMergeOperands = rvsdg::operands(memoryStateMergeNode);
+      foundMergeOrJoinOperation = true;
+      auto memoryStateOpOperands = rvsdg::operands(memoryStateNode);
       newLoadMemoryStates.insert(
           newLoadMemoryStates.end(),
-          memoryStateMergeOperands.begin(),
-          memoryStateMergeOperands.end());
+          memoryStateOpOperands.begin(),
+          memoryStateOpOperands.end());
     }
     else
     {
       newLoadMemoryStates.push_back(memoryState);
     }
   }
-  if (!foundMemoryStateMergeOperation)
+  if (!foundMergeOrJoinOperation)
     return std::nullopt;
 
-  auto & newLoadNode =
-      CreateNode(address, newLoadMemoryStates, operation.GetLoadedType(), operation.GetAlignment());
+  auto & newLoadNode = LoadNonVolatileOperation::CreateNode(
+      address,
+      newLoadMemoryStates,
+      operation.GetLoadedType(),
+      operation.GetAlignment());
 
   size_t newMemoryStateResultIndex = 1;
   std::vector<rvsdg::Output *> results;
-  results.push_back(&LoadedValueOutput(newLoadNode));
+  results.push_back(&LoadOperation::LoadedValueOutput(newLoadNode));
   for (auto & oldMemoryStateOperand : oldLoadMemoryStates)
   {
-    auto [memoryStateMergeNode, memoryStateMergeOperation] =
-        rvsdg::TryGetSimpleNodeAndOptionalOp<MemoryStateMergeOperation>(*oldMemoryStateOperand);
-    if (memoryStateMergeOperation)
+    auto [memoryStateNode, memoryStateOperation] =
+        rvsdg::TryGetSimpleNodeAndOptionalOp<TMemoryStateMergeOrJoinOperation>(
+            *oldMemoryStateOperand);
+    if (memoryStateOperation)
     {
-      size_t numMemoryStates = memoryStateMergeNode->ninputs();
-      auto memoryStateMergeOperands =
+      size_t numMemoryStates = memoryStateNode->ninputs();
+      auto memoryStateOpOperands =
           rvsdg::Outputs(newLoadNode, newMemoryStateResultIndex, numMemoryStates);
-      const auto result = MemoryStateMergeOperation::CreateNode(memoryStateMergeOperands).output(0);
+      const auto result =
+          TMemoryStateMergeOrJoinOperation::CreateNode(memoryStateOpOperands).output(0);
       results.push_back(result);
       newMemoryStateResultIndex += numMemoryStates;
     }
@@ -379,6 +390,22 @@ LoadNonVolatileOperation::NormalizeLoadMemoryStateMerge(
   return results;
 }
 
+std::optional<std::vector<rvsdg::Output *>>
+LoadNonVolatileOperation::NormalizeLoadMemoryStateMerge(
+    const LoadNonVolatileOperation & operation,
+    const std::vector<rvsdg::Output *> & operands)
+{
+  return NormalizeMemoryStateMergeOrJoin<MemoryStateMergeOperation>(operation, operands);
+}
+
+std::optional<std::vector<rvsdg::Output *>>
+LoadNonVolatileOperation::NormalizeLoadMemoryStateJoin(
+    const LoadNonVolatileOperation & operation,
+    const std::vector<rvsdg::Output *> & operands)
+{
+  return NormalizeMemoryStateMergeOrJoin<MemoryStateJoinOperation>(operation, operands);
+}
+
 std::optional<std::vector<rvsdg::Output *>>
 LoadNonVolatileOperation::NormalizeLoadStore(
     const LoadNonVolatileOperation & operation,

jlm/llvm/ir/operators/Load.hpp

@@ -289,6 +289,23 @@ class LoadNonVolatileOperation final : public LoadOperation
       const LoadNonVolatileOperation & operation,
       const std::vector<rvsdg::Output *> & operands);
 
+  /**
+   * \brief Swaps a \ref MemoryStateJoinOperation and a \ref LoadNonVolatileOperation
+   *
+   * sx1 = MemoryStateJoinOperation si1 ... siM
+   * v sl1 = LoadNonVolatileOperation a sx1
+   * =>
+   * v sl1 ... slM = LoadNonVolatileOperation a si1 ... siM
+   * sx1 = MemoryStateJoinOperation sl1 ... SlM
+   *
+   * @return If the normalization could be applied, then the results of the \ref
+   * LoadNonVolatileOperation after the transformation. Otherwise, std::nullopt.
+   */
+  static std::optional<std::vector<rvsdg::Output *>>
+  NormalizeLoadMemoryStateJoin(
+      const LoadNonVolatileOperation & operation,
+      const std::vector<rvsdg::Output *> & operands);
+
   /**
    * \brief If the producer of a load's address is an alloca operation, then we can remove all
    * state edges originating from other alloca operations.

jlm/llvm/opt/reduction.cpp

@@ -253,6 +253,7 @@ NodeReduction::NormalizeLoadNode(
 {
   static std::vector<rvsdg::NodeNormalization<LoadNonVolatileOperation>> loadNodeNormalizations(
       { LoadNonVolatileOperation::NormalizeLoadMemoryStateMerge,
+        LoadNonVolatileOperation::NormalizeLoadMemoryStateJoin,
         LoadNonVolatileOperation::NormalizeLoadStore,
         LoadNonVolatileOperation::NormalizeLoadAlloca,
         LoadNonVolatileOperation::NormalizeDuplicateStates,

tests/jlm/llvm/ir/operators/LoadTests.cpp

@@ -18,7 +18,7 @@
 #include <jlm/rvsdg/view.hpp>
 
 static void
-OperationEquality()
+LoadNonVolatileOperationEquality()
 {
   using namespace jlm::llvm;
 
@@ -42,8 +42,8 @@ OperationEquality()
 }
 
 JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadNonVolatileTests-OperationEquality",
-    OperationEquality)
+    "jlm/llvm/ir/operators/LoadTests-LoadNonVolatileOperationEquality",
+    LoadNonVolatileOperationEquality)
 
 static void
 TestCopy()
@@ -75,7 +75,7 @@ TestCopy()
       == jlm::util::AssertedCast<jlm::rvsdg::SimpleNode>(copiedNode)->GetOperation());
 }
 
-JLM_UNIT_TEST_REGISTER("jlm/llvm/ir/operators/LoadNonVolatileTests-Copy", TestCopy)
+JLM_UNIT_TEST_REGISTER("jlm/llvm/ir/operators/LoadTests-Copy", TestCopy)
 
 static void
 TestLoadAllocaReduction()
@@ -116,11 +116,11 @@ TestLoadAllocaReduction()
 }
 
 JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadNonVolatileTests-LoadAllocaReduction",
+    "jlm/llvm/ir/operators/LoadTests-LoadAllocaReduction",
     TestLoadAllocaReduction)
 
 static void
-LoadMuxReduction_Success()
+LoadMemoryStateMergeReduction_Success()
 {
   using namespace jlm::llvm;
 
@@ -196,11 +196,11 @@ LoadMuxReduction_Success()
 }
 
 JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadNonVolatileTests-LoadMuxReduction_Success",
-    LoadMuxReduction_Success)
+    "jlm/llvm/ir/operators/LoadTests-LoadMemoryStateMergeReduction_Success",
+    LoadMemoryStateMergeReduction_Success)
 
 static void
-LoadMuxReduction_LoadWithoutStates()
+LoadMemoryStateMergeReduction_LoadWithoutStates()
 {
   using namespace jlm::llvm;
 
@@ -234,8 +234,87 @@ LoadMuxReduction_LoadWithoutStates()
 }
 
 JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadNonVolatileTests-LoadMuxReduction_LoadWithoutStates",
-    LoadMuxReduction_LoadWithoutStates)
+    "jlm/llvm/ir/operators/LoadTests-LoadMemoryStateMergeReduction_LoadWithoutStates",
+    LoadMemoryStateMergeReduction_LoadWithoutStates)
+
+static void
+LoadMemoryStateJoinReduction_Success()
+{
+  using namespace jlm::llvm;
+  using namespace jlm::rvsdg;
+
+  // Arrange
+  const auto memoryStateType = MemoryStateType::Create();
+  const auto pointerType = PointerType::Create();
+  const auto bitstringType = BitType::Create(32);
+
+  Graph graph;
+  const auto address = &jlm::rvsdg::GraphImport::Create(graph, pointerType, "address");
+  auto s1 = &jlm::rvsdg::GraphImport::Create(graph, memoryStateType, "state1");
+  auto s2 = &jlm::rvsdg::GraphImport::Create(graph, memoryStateType, "state2");
+  auto s3 = &jlm::rvsdg::GraphImport::Create(graph, memoryStateType, "state3");
+
+  auto mergeResult1 = MemoryStateJoinOperation::CreateNode({ s1, s2, s3 }).output(0);
+  auto mergeResult2 = MemoryStateJoinOperation::CreateNode({ s1, s2, s3 }).output(0);
+  auto & loadNode = LoadNonVolatileOperation::CreateNode(
+      *address,
+      { mergeResult1, mergeResult2 },
+      bitstringType,
+      4);
+
+  auto & ex1 = GraphExport::Create(*loadNode.output(0), "l");
+  auto & ex2 = GraphExport::Create(*loadNode.output(1), "s1");
+  auto & ex3 = GraphExport::Create(*loadNode.output(2), "s2");
+
+  view(&graph.GetRootRegion(), stdout);
+
+  // Act
+  const auto success = jlm::rvsdg::ReduceNode<LoadNonVolatileOperation>(
+      LoadNonVolatileOperation::NormalizeLoadMemoryStateJoin,
+      loadNode);
+  graph.PruneNodes();
+
+  view(&graph.GetRootRegion(), stdout);
+
+  // Assert
+  assert(success);
+  const auto reducedLoadNode = jlm::rvsdg::TryGetOwnerNode<Node>(*ex1.origin());
+  assert(is<LoadNonVolatileOperation>(reducedLoadNode));
+  assert(reducedLoadNode->ninputs() == 7);
+  assert(reducedLoadNode->input(0)->origin() == address);
+  assert(reducedLoadNode->input(1)->origin() == s1);
+  assert(reducedLoadNode->input(2)->origin() == s2);
+  assert(reducedLoadNode->input(3)->origin() == s3);
+  assert(reducedLoadNode->input(4)->origin() == s1);
+  assert(reducedLoadNode->input(5)->origin() == s2);
+  assert(reducedLoadNode->input(6)->origin() == s3);
+
+  {
+    const auto merge = jlm::rvsdg::TryGetOwnerNode<Node>(*ex2.origin());
+    assert(is<MemoryStateJoinOperation>(merge));
+    assert(merge->ninputs() == 3);
+    for (size_t n = 0; n < merge->ninputs(); n++)
+    {
+      const auto expectedLoadNode = jlm::rvsdg::TryGetOwnerNode<Node>(*merge->input(n)->origin());
+      assert(expectedLoadNode == reducedLoadNode);
+    }
+  }
+
+  {
+    const auto merge = jlm::rvsdg::TryGetOwnerNode<Node>(*ex3.origin());
+    assert(is<MemoryStateJoinOperation>(merge));
+    assert(merge->ninputs() == 3);
+    for (size_t n = 0; n < merge->ninputs(); n++)
+    {
+      const auto expectedLoadNode = jlm::rvsdg::TryGetOwnerNode<Node>(*merge->input(n)->origin());
+      assert(expectedLoadNode == reducedLoadNode);
+    }
+  }
+}
+
+JLM_UNIT_TEST_REGISTER(
+    "jlm/llvm/ir/operators/LoadTests-LoadMemoryStateJoinReduction_Success",
+    LoadMemoryStateJoinReduction_Success)
 
 static void
 TestDuplicateStateReduction()
@@ -287,7 +366,7 @@ TestDuplicateStateReduction()
 }
 
 JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadNonVolatileTests-DuplicateStateReduction",
+    "jlm/llvm/ir/operators/LoadTests-DuplicateStateReduction",
     TestDuplicateStateReduction)
 
 static void
@@ -339,7 +418,7 @@ TestLoadStoreStateReduction()
 }
 
 JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadNonVolatileTests-LoadStoreStateReduction",
+    "jlm/llvm/ir/operators/LoadTests-LoadStoreStateReduction",
     TestLoadStoreStateReduction)
 
 static void
@@ -380,7 +459,7 @@ TestLoadStoreReduction_Success()
 }
 
 JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadNonVolatileTests-LoadStoreReduction_Success",
+    "jlm/llvm/ir/operators/LoadTests-LoadStoreReduction_Success",
     TestLoadStoreReduction_Success)
 
 /**
@@ -435,7 +514,7 @@ LoadStoreReduction_DifferentValueOperandType()
 }
 
 JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadNonVolatileTests-LoadStoreReduction_DifferentValueOperandType",
+    "jlm/llvm/ir/operators/LoadTests-LoadStoreReduction_DifferentValueOperandType",
     LoadStoreReduction_DifferentValueOperandType)
 
 static void
@@ -510,9 +589,7 @@ TestLoadLoadReduction()
   }
 }
 
-JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadNonVolatileTests-LoadLoadReduction",
-    TestLoadLoadReduction)
+JLM_UNIT_TEST_REGISTER("jlm/llvm/ir/operators/LoadTests-LoadLoadReduction", TestLoadLoadReduction)
 
 static void
 IOBarrierAllocaAddressNormalization()
@@ -675,7 +752,7 @@ LoadVolatileOperationEquality()
 }
 
 JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadVolatileTests-OperationEquality",
+    "jlm/llvm/ir/operators/LoadTests-LoadVolatileOperationEquality",
     LoadVolatileOperationEquality)
 
 static void
@@ -697,7 +774,7 @@ OperationCopy()
   assert(*copiedOperation == operation);
 }
 
-JLM_UNIT_TEST_REGISTER("jlm/llvm/ir/operators/LoadVolatileTests-OperationCopy", OperationCopy)
+JLM_UNIT_TEST_REGISTER("jlm/llvm/ir/operators/LoadTests-OperationCopy", OperationCopy)
 
 static void
 OperationAccessors()
@@ -721,9 +798,7 @@ OperationAccessors()
   assert(operation.nresults() == numMemoryStates + 2);   // [loadedValue, ioState, memoryStates]
 }
 
-JLM_UNIT_TEST_REGISTER(
-    "jlm/llvm/ir/operators/LoadVolatileTests-OperationAccessors",
-    OperationAccessors)
+JLM_UNIT_TEST_REGISTER("jlm/llvm/ir/operators/LoadTests-OperationAccessors", OperationAccessors)
 
 static void
 NodeCopy()
@@ -765,4 +840,4 @@ NodeCopy()
   assert(*copiedOperation->GetLoadedType() == *valueType);
 }
 
-JLM_UNIT_TEST_REGISTER("jlm/llvm/ir/operators/LoadVolatileTests-NodeCopy", NodeCopy)
+JLM_UNIT_TEST_REGISTER("jlm/llvm/ir/operators/LoadTests-NodeCopy", NodeCopy)