CNE: Fix handling of theta nodes without loop variables (phate#1519)

The subregions of structural nodes without arguments were never visited
in CNE's mark phase, leading ultimately to problems in the divert phase
as no congruence sets could be found for nodes in such subregions.

Closes phate#1482

Co-authored-by: Magnus Sjalander <work@sjalander.com>

Author: phate

jlm/llvm/opt/CommonNodeElimination.cpp

@@ -642,30 +642,38 @@ markSubregionsFromInputs(
 
   for (auto & subregion : node.Subregions())
   {
-    // create a partitioning of the region arguments
-    std::vector<size_t> partitions(subregion.narguments());
-    // Arguments that do not belong to any input are given partition keys higher than any real index
-    size_t nextUniquePartitionKey = context.numCongruenceSets();
-
-    for (const auto argument : subregion.Arguments())
+    if (subregion.narguments() == 0)
     {
-      if (const auto input = argument->input())
+      markRegion(subregion, context);
+    }
+    else
+    {
+      // create a partitioning of the region arguments
+      std::vector<size_t> partitions(subregion.narguments());
+      // Arguments that do not belong to any input are given partition keys higher than any real
+      // index
+      size_t nextUniquePartitionKey = context.numCongruenceSets();
+
+      for (const auto argument : subregion.Arguments())
       {
-        // If the argument corresponds to an input, use the partition key of the input
-        partitions[argument->index()] = context.getSetFor(*input->origin());
+        if (const auto input = argument->input())
+        {
+          // If the argument corresponds to an input, use the partition key of the input
+          partitions[argument->index()] = context.getSetFor(*input->origin());
+        }
+        else
+        {
+          // Otherwise make sure the argument is not partitioned with any other argument
+          partitions[argument->index()] = nextUniquePartitionKey++;
+        }
       }
-      else
+
+      if (partitionArguments(subregion, partitions, context))
       {
-        // Otherwise make sure the argument is not partitioned with any other argument
-        partitions[argument->index()] = nextUniquePartitionKey++;
+        anyChanges = true;
+        markRegion(subregion, context);
       }
     }
-
-    if (partitionArguments(subregion, partitions, context))
-    {
-      anyChanges = true;
-      markRegion(subregion, context);
-    }
   }
 
   return anyChanges;

tests/jlm/llvm/opt/CommonNodeEliminationTests.cpp

@@ -556,3 +556,44 @@ TEST(CommonNodeEliminationTests, test_phi)
       jlm::rvsdg::AssertGetOwnerNode<jlm::rvsdg::LambdaNode>(*f1).input(0)->origin(),
       jlm::rvsdg::AssertGetOwnerNode<jlm::rvsdg::LambdaNode>(*f2).input(0)->origin());
 }
+
+TEST(CommonNodeEliminationTests, EmptyTheta)
+{
+  using namespace jlm::llvm;
+  using namespace jlm::rvsdg;
+
+  // Arrange
+  auto valueType = TestType::createValueType();
+  auto controlType = ControlType::Create(2);
+
+  LlvmRvsdgModule rvsdgModule(jlm::util::FilePath(""), "", "");
+  auto & rvsdg = rvsdgModule.Rvsdg();
+
+  auto thetaNode = ThetaNode::create(&rvsdg.GetRootRegion());
+
+  auto node1 = TestOperation::createNode(thetaNode->subregion(), {}, { valueType });
+  auto node2 =
+      TestOperation::createNode(thetaNode->subregion(), { node1->output(0) }, { valueType });
+  auto node3 =
+      TestOperation::createNode(thetaNode->subregion(), { node1->output(0) }, { valueType });
+  auto node4 = TestOperation::createNode(
+      thetaNode->subregion(),
+      { node2->output(0), node3->output(0) },
+      { controlType });
+
+  thetaNode->set_predicate(node4->output(0));
+
+  view(rvsdg, stdout);
+
+  // Act
+  CommonNodeElimination commonNodeElimination;
+  commonNodeElimination.Run(rvsdgModule, statisticsCollector);
+
+  thetaNode->subregion()->prune(false);
+
+  view(rvsdg, stdout);
+
+  // Assert
+  // We expect that node2 and node3 are unified in the theta subregion
+  EXPECT_EQ(thetaNode->subregion()->numNodes(), 3);
+}