[Mlir-commits] [mlir] bea16e7 - [mlir] Fix invalid assertion in ModuleTranslation.cpp

[Alex Zinenko]

Author: Alex Zinenko
Date: 2022-01-14T12:56:33+01:00
New Revision: bea16e72a757fc72b9237098606ead945377170f

URL: https://github.com/llvm/llvm-project/commit/bea16e72a757fc72b9237098606ead945377170f
DIFF: https://github.com/llvm/llvm-project/commit/bea16e72a757fc72b9237098606ead945377170f.diff

LOG: [mlir] Fix invalid assertion in ModuleTranslation.cpp

LLVM dialect supports terminators with repeated successor blocks that take
different operands. This cannot be directly expressed in LLVM IR though since
it uses the number of the predecessor block to differentiate values in its PHI
nodes. Therefore, the translation to LLVM IR inserts dummy blocks to forward
arguments in case of repeated succesors with arguments. The insertion works
correctly. However, when connecting PHI nodes to their source values, the
assertion of the insertion having worked correctly was incorrect: it would only
trigger if repeated blocks were adjacent in the successor list (not guaranteed
by anything) and would not check if the successors have operands (no need for
dummy blocks in absence of operands since no PHIs are being created). Change
the assertion to only trigger in case of duplicate successors with operands,
and don't expect them to be adjacent.

Reviewed By: wsmoses

Differential Revision: https://reviews.llvm.org/D117214

Added: 
    

Modified: 
    mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
    mlir/test/Target/LLVMIR/llvmir.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 57733a68fc24c..f347b20667607 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -362,11 +362,19 @@ static Value getPHISourceValue(Block *current, Block *pred,
   if (isa<LLVM::BrOp>(terminator))
     return terminator.getOperand(index);
 
-  SuccessorRange successors = terminator.getSuccessors();
-  assert(std::adjacent_find(successors.begin(), successors.end()) ==
-             successors.end() &&
-         "successors with arguments in LLVM branches must be 
diff erent blocks");
-  (void)successors;
+#ifndef NDEBUG
+  llvm::SmallPtrSet<Block *, 4> seenSuccessors;
+  for (unsigned i = 0, e = terminator.getNumSuccessors(); i < e; ++i) {
+    Block *successor = terminator.getSuccessor(i);
+    auto branch = cast<BranchOpInterface>(terminator);
+    Optional<OperandRange> successorOperands = branch.getSuccessorOperands(i);
+    assert(
+        (!seenSuccessors.contains(successor) ||
+         (successorOperands && successorOperands->empty())) &&
+        "successors with arguments in LLVM branches must be 
diff erent blocks");
+    seenSuccessors.insert(successor);
+  }
+#endif
 
   // For instructions that branch based on a condition value, we need to take
   // the operands for the branch that was taken.

diff  --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index 69e881228d14b..644d480404043 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -1766,3 +1766,99 @@ module {
 // CHECK-DAG: ![[SCOPES13]] = !{![[SCOPE1]], ![[SCOPE3]]}
 // CHECK-DAG: ![[SCOPES23]] = !{![[SCOPE2]], ![[SCOPE3]]}
 
+
+// -----
+
+// It is okay to have repeated successors if they have no arguments.
+
+// CHECK-LABEL: @duplicate_block_in_switch
+// CHECK-SAME: float %[[FIRST:.*]],
+// CHECK-SAME: float %[[SECOND:.*]])
+// CHECK:   switch i32 %{{.*}}, label %[[DEFAULT:.*]] [
+// CHECK:     i32 105, label %[[DUPLICATE:.*]]
+// CHECK:     i32 108, label %[[BLOCK:.*]]
+// CHECK:     i32 106, label %[[DUPLICATE]]
+// CHECK:   ]
+
+// CHECK: [[DEFAULT]]:
+// CHECK:   phi float [ %[[FIRST]], %{{.*}} ]
+// CHECK:   call void @bar
+
+// CHECK: [[DUPLICATE]]:
+// CHECK:   call void @baz
+
+// CHECK: [[BLOCK]]:
+// CHECK:   phi float [ %[[SECOND]], %{{.*}} ]
+// CHECK:   call void @qux
+
+llvm.func @duplicate_block_in_switch(%cond : i32, %arg1: f32, %arg2: f32) {
+  llvm.switch %cond : i32, ^bb1(%arg1: f32) [
+    105: ^bb2,
+    108: ^bb3(%arg2: f32),
+    106: ^bb2
+  ]
+
+^bb1(%arg3: f32):
+  llvm.call @bar(%arg3): (f32) -> ()
+  llvm.return
+
+^bb2:
+  llvm.call @baz() : () -> ()
+  llvm.return
+
+^bb3(%arg4: f32):
+  llvm.call @qux(%arg4) : (f32) -> ()
+  llvm.return
+}
+
+// If there are repeated successors with arguments, a new block must be created
+// for repeated successors to ensure PHI can disambiguate values based on the
+// predecessor they come from.
+
+// CHECK-LABEL: @duplicate_block_with_args_in_switch
+// CHECK-SAME: float %[[FIRST:.*]],
+// CHECK-SAME: float %[[SECOND:.*]])
+// CHECK:   switch i32 %{{.*}}, label %[[DEFAULT:.*]] [
+// CHECK:     i32 106, label %[[DUPLICATE:.*]]
+// CHECK:     i32 105, label %[[BLOCK:.*]]
+// CHECK:     i32 108, label %[[DEDUPLICATED:.*]]
+// CHECK:   ]
+
+// CHECK: [[DEFAULT]]:
+// CHECK:   phi float [ %[[FIRST]], %{{.*}} ]
+// CHECK:   call void @bar
+
+// CHECK: [[BLOCK]]:
+// CHECK:   call void @baz
+
+// CHECK: [[DUPLICATE]]:
+// CHECK:   phi float [ %[[PHI:.*]], %[[DEDUPLICATED]] ], [ %[[FIRST]], %{{.*}} ]
+// CHECK:   call void @qux
+
+// CHECK: [[DEDUPLICATED]]:
+// CHECK:   %[[PHI]] = phi float [ %[[SECOND]], %{{.*}} ]
+// CHECK:   br label %[[DUPLICATE]]
+
+llvm.func @duplicate_block_with_args_in_switch(%cond : i32, %arg1: f32, %arg2: f32) {
+  llvm.switch %cond : i32, ^bb1(%arg1: f32) [
+    106: ^bb3(%arg1: f32),
+    105: ^bb2,
+    108: ^bb3(%arg2: f32)
+  ]
+
+^bb1(%arg3: f32):
+  llvm.call @bar(%arg3): (f32) -> ()
+  llvm.return
+
+^bb2:
+  llvm.call @baz() : () -> ()
+  llvm.return
+
+^bb3(%arg4: f32):
+  llvm.call @qux(%arg4) : (f32) -> ()
+  llvm.return
+}
+
+llvm.func @bar(f32)
+llvm.func @baz()
+llvm.func @qux(f32)