[MLIR] Added llvm.indirectbr

[shraiysh]

Summary

This patch adds llvm.indirect_br.
LLVM Language reference suggests that repetition of labels, although allowed, is not useful. The same behaviour is implemented here. Also, the order of labels does not matter.

Syntax:
llvm.indirect_br <address> [ label1 label2 ... ] : <type>

Example:
LLVM IR
indirectbr i8* %address, [label %bb1, label %bb2, label %bb2]
MLIR LLVM Dialect
llvm.indirect_br %arg0 [^bb1, ^bb2] : !llvm<"i8*">

[rriddle]

Thanks for the patch shraiysh! I added a few comments.

Also, could you please re-upload with more context?
https://llvm.org/docs/Phabricator.html#requesting-a-review-via-the-web-interface

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

853

Do we need a map here? Can you just pass a ArrayRef<std::pair<>> instead?

856

nit: We can just do: result.addOperands(address).

857

This looks like C++17, LLVM/MLIR compiles with C++14.

861

*Op classes are value-typed, so they shouldn't be passed by reference or pointer.

864

nit: Can you just implement this inline with:

interleaveComma(llvm::seq<int>(0, op.getNumSuccessors()), p, [&](int i) {
  p.printSuccessorAndUseList(op, i);
});

[shraiysh]

Hi @rriddle ! Thanks for the comments. I have updated the patch with more context now. Also, I have addressed all comments except one. Please let me know if you think I should change it to ArrayRef (I mentioned my reason in the inline comment).

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

853

My reason for going with map here was that with ArrayRef<std::pair<>>, duplication of labels will have to be handled here (avoid entering the same label twice). Please let me know if you still want me to use ArrayRef<std::pair<>>.

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

475

Why choose a slightly different name (indirect_br vs indirectbr)? It may be slightly more readable, but is likely to cause issues due to this discrepancy.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

875

If LLVM supports repeated block labels, why should we restrict it? I'm not against it fundamentally, but I'd like a better argument then "it's not useful". For example, it creates a possibility to mention the same block with different arguments.

877

I find it unusual to have successors after the type of the argument, have you considered otherwise? Similarly, putting the labels in brackets looks unusual, we usually put block arguments in brackets.

885

const auto &

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

654

auto *?

Generally, I find auto appropriate when the type is obvious from the code (which isn't the case here) or when the type is obnoxiously long (which isn't the case here either).

656

Nit: operator[] on map default-initializes the value if it is not yet in the map, so you should be able to do failed(processBranchArgs(indBr, succ, succs[blocks[succ]]) directly and avoid a temporary. I've seen another opportunity like this above.

mlir/test/Dialect/LLVMIR/roundtrip.mlir

291

Please avoid matching on SSA names (https://mlir.llvm.org/getting_started/TestingGuide/).

293

Please drop the autogenerated "pred" comment from the test

304

The {{[23]}} regexp is a sign that you produce non-deterministic IR (due to the order in DenseMap), which is arguably a bad thing for testability. Consider using SetVector or another container with a deterministic order instead.

306

Please also add a test with blocks taking arguments.

Can we branch to blocks that take different arguments / different number of arguments, btw?

mlir/test/Target/llvmir.mlir

1181

How do you intend to handle LLVM's requirement that the argument of indirectbr must be produced by blockaddress ?

[shraiysh]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

877

I have shifted the type to the end of instruction.
About the block labels in brackets, please let me know if you want me to eliminate the braces. I figured square braces would be more llvm friendly.

mlir/test/Dialect/LLVMIR/roundtrip.mlir

306

I think we can branch to blocks with different arguments. Different blocks can have different PHI nodes in the LLVM IR and hence there could be different number of arguments for them and indirectbr in llvm should work when that happens.

mlir/test/Target/llvmir.mlir

1181

I tried running the following code through llvm compiler, and it seems to accept it. So, I think there is no check for this in LLVM IR. I am not sure how we could check that here too. If there is a way for it, or if I misunderstood something, please let me know.

test.ll

define void @indirectbrTest(i8* %address, i32* %sink) #0 {

entry:
  %x = alloca i32
  store i32 0, i32* %x
  indirectbr i32* %x, [label %bb1]

bb1:
  store volatile i32 0, i32* %sink
  ret void
}

I compiled this using clang++ -c test.ll. So, it looks like there are no checks for the argument.

[ftynse]

mlir/test/Dialect/LLVMIR/roundtrip.mlir

304

You seem to have removed DenseMap completely, instead of replacing it with SetVector as suggested. Now you have a different problem, which I hinted to in another comment. You can now create an indirect branch to the same block with different values llvm.indirectbr %condtion [^bb1(%0: i32), ^bb1(%1: i32)] , but LLVM IR has no way of expressing this directly (a PHI node can only have one "source" value per predecessor block).

mlir/test/Target/llvmir.mlir

1181

It does work because the value comes from the function argument and LLVM cannot reasonably know how the function will get called :) It should complain if you allocate an i8* and pass it to indirectbr. It was more a forward-looking question, without being able to take the address of a block, indirectbr has restricted applicability.

[ftynse]

See https://llvm.discourse.group/t/duplicate-successors-in-llvm-dialect/640 for repeated-successor discussion.

[ftynse]

The discussion seems to have converged to allowing repeated successors with different arguments and doing the legalization as a pre-pass to translation. Let's update this diff to reflect that.

[shraiysh]

Hi @ftynse ! I will update this diff accordingly. I'll figure out how to handle BlockAddresses first.

[ftynse]

This looks abandoned. Ping me if it is not the case.

[shraiysh]

Hi, Yes. I'm sorry, but I'm a little busy with other commitments right now. I'll work on this later.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_CondBrOp : LLVM_TerminatorOp<"cond_br", []> {
   let verifier = [{
     if (getNumSuccessors() != 2)
       return emitOpError("expected exactly two successors");
     return success();
   }];
   let parser = [{ return parseCondBrOp(parser, result); }];
   let printer = [{ printCondBrOp(p, *this); }];
 }
+def LLVM_IndirectBrOp : LLVM_TerminatorOp<"indirect_br", []> {
+  let arguments = (ins LLVM_Type:$address);
+  let builders = [OpBuilder<
+    "Builder *b, OperationState &result, Value address,"
+    "DenseMap<Block*, SmallVector<Value, 4>> successors",
+    [{
+      result.addOperands(address);
+      for (auto successor : successors)
+        result.addSuccessor(successor.getFirst(), successor.getSecond());
+    }]>];
+  let printer = [{ printIndirectBrOp(p, *this); }];
+  let parser = [{ return parseIndirectBrOp(parser, result); }];
+}
 def LLVM_ReturnOp : LLVM_TerminatorOp<"return", []> {
   string llvmBuilder = [{
     if ($_numOperands != 0)
       builder.CreateRet($args[0]);
     else
       builder.CreateRetVoid();
   }];
 

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

if (parser.parseOperand(condition) || parser.parseComma() ||
       parser.resolveOperand(condition, i1Type, result.operands))
     return failure();
 
   result.addSuccessor(trueDest, trueOperands);
   result.addSuccessor(falseDest, falseOperands);
   return success();
 }
 
+///===---------------------------------------------------------------------===//
+/// Printing/parsing for LLVM::IndirectBrOp.
+///===---------------------------------------------------------------------===//
+
+static void printIndirectBrOp(OpAsmPrinter &p, IndirectBrOp op) {
+  p << op.getOperationName() << ' ' << op.address() << " : "
+    << op.address().getType() << " [";
+  interleaveComma(llvm::seq<int>(0, op.getNumSuccessors()), p,
+                  [&](int i) { p.printSuccessorAndUseList(op, i); });
+  p << "]";
+}
+
+/// <operation> ::= `llvm.indirect_br` ssa-use `,`
+///                 `[` bb-id (`[` ssa-use-and-type-list `]`)?
+///                           (`,` bb-id (`[` ssa-use-and-type-list `]`)?)* `]`
+///                  attribute-dict?
+static ParseResult parseIndirectBrOp(OpAsmParser &parser,
+                                     OperationState &result) {
+  OpAsmParser::OperandType address;
+  Type ty;
+
+  if (parser.parseOperand(address) || parser.parseColonType(ty) ||
+      parser.resolveOperand(address, ty, result.operands) ||
+      parser.parseLSquare())
+    return failure();
+
+  // Successors are stored in a DenseMap before adding them to result to avoid
+  // duplication of labels.
+  DenseMap<Block *, SmallVector<Value, 4>> successors;
+
+  do {
+    Block *dest;
+    SmallVector<Value, 4> oper;
+    if (parser.parseSuccessorAndUseList(dest, oper))
+      return failure();
+    successors[dest] = oper;
+  } while (succeeded(parser.parseOptionalComma()));
+
+  for (auto successor : successors)
+    result.addSuccessor(successor.getFirst(), successor.getSecond());
+
+  if (parser.parseRSquare() || parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Printing/parsing for LLVM::ReturnOp.
 //===----------------------------------------------------------------------===//
 
 static void printReturnOp(OpAsmPrinter &p, ReturnOp &op) {
   p << op.getOperationName();
   p.printOptionalAttrDict(op.getAttrs());
   assert(op.getNumOperands() <= 1);

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

for (auto *succ : llvm::reverse(brInst->successors())) {
       SmallVector<Value, 4> blockArguments;
       if (failed(processBranchArgs(brInst, succ, blockArguments)))
         return failure();
       state.addSuccessor(blocks[succ], blockArguments);
     }
     b.createOperation(state);
     return success();
   }
+  case llvm::Instruction::IndirectBr: {
+    auto *indBr = cast<llvm::IndirectBrInst>(inst);
+    Value address = processValue(indBr->getAddress());
+    if (!address)
+      return failure();
+    DenseMap<Block *, SmallVector<Value, 4>> succs;
+    for (auto succ : indBr->successors()) {
+      SmallVector<Value, 4> blockArguments;
+      if (failed(processBranchArgs(indBr, succ, blockArguments)))
+        return failure();
+      succs[blocks[succ]] = blockArguments;
+    }
+    b.create<IndirectBrOp>(loc, address, succs);
+    return success();
+  }
   case llvm::Instruction::PHI: {
     LLVMType type = processType(inst->getType());
     if (!type)
       return failure();
     v = b.getInsertionBlock()->addArgument(type);
     return success();
   }
   case llvm::Instruction::Call: {

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

if (auto brOp = dyn_cast<LLVM::BrOp>(opInst)) {
     return success();
   }
   if (auto condbrOp = dyn_cast<LLVM::CondBrOp>(opInst)) {
     builder.CreateCondBr(valueMapping.lookup(condbrOp.getOperand(0)),
                          blockMapping[condbrOp.getSuccessor(0)],
                          blockMapping[condbrOp.getSuccessor(1)]);
     return success();
   }
+  if (auto indBrOp = dyn_cast<LLVM::IndirectBrOp>(opInst)) {
+    unsigned nSucc = indBrOp.getNumSuccessors();
+    llvm::IndirectBrInst *ibi =
+        builder.CreateIndirectBr(valueMapping.lookup(indBrOp.address()), nSucc);
+    for (unsigned idx = 0; idx < nSucc; idx++)
+      ibi->addDestination(blockMapping[indBrOp.getSuccessor(idx)]);
+    return success();
+  }
 
   // Emit addressof.  We need to look up the global value referenced by the
   // operation and store it in the MLIR-to-LLVM value mapping.  This does not
   // emit any LLVM instruction.
   if (auto addressOfOp = dyn_cast<LLVM::AddressOfOp>(opInst)) {
     LLVM::GlobalOp global = addressOfOp.getGlobal();
     // The verifier should not have allowed this.
     assert(global && "referencing an undefined global");

mlir/test/Dialect/LLVMIR/roundtrip.mlir

 ^bb3:
   llvm.invoke @bar(%8, %6, %4) to ^bb2 unwind ^bb1 : (!llvm<"i8*">, !llvm<"i8*">, !llvm<"i8*">) -> ()
 
 // CHECK-NEXT: ^bb4:
 // CHECK-NEXT:   llvm.return %[[a0]] : !llvm.i32
 ^bb4:
   llvm.return %0 : !llvm.i32
 }
+
+// CHECK-LABEL: @indirectBr
+llvm.func @indirectBr(%arg0: !llvm<"i8*">, %arg1: !llvm<"i32*">) {
+  %0 = llvm.mlir.constant(2 : i32) : !llvm.i32
+  %1 = llvm.mlir.constant(1 : i32) : !llvm.i32
+  %2 = llvm.mlir.constant(0 : i32) : !llvm.i32
+  // CHECK: llvm.indirect_br %arg0 : !llvm<"i8*"> [^bb{{[12]}}, ^bb{{[12]}}]
+  llvm.indirect_br %arg0 : !llvm<"i8*"> [^bb2, ^bb1]
+^bb1: // pred: ^bb0
+  llvm.store %2, %arg1 : !llvm<"i32*">
+  llvm.br ^bb4
+^bb2: // 2 preds: ^bb0, ^bb4
+  llvm.store %1, %arg1 : !llvm<"i32*">
+  llvm.br ^bb4
+^bb3: // pred: ^bb4
+  llvm.store %0, %arg1 : !llvm<"i32*">
+  llvm.br ^bb4
+// CHECK: ^bb4:
+^bb4: // 3 preds: ^bb1, ^bb2, ^bb3
+  // CHECK: llvm.indirect_br %arg0 : !llvm<"i8*"> [^bb{{[23]}}, ^bb{{[23]}}]
+  llvm.indirect_br %arg0 : !llvm<"i8*"> [^bb3, ^bb2, ^bb2]
+}

mlir/test/Target/import.ll

 ; CHECK: ^bb3:
   ; CHECK: %{{[0-9]+}} = llvm.invoke @bar(%[[a3]]) to ^bb2 unwind ^bb1 : (!llvm<"i8*">) -> !llvm<"i8*">
   %6 = invoke i8* @bar(i8* %1) to label %4 unwind label %2
 
 ; CHECK: ^bb4:
   ; CHECK: llvm.return %{{[0-9]+}} : !llvm.i32
   ret i32 0
 }
+
+; CHECK-LABEL: llvm.func @indirectbrTest(%arg0: !llvm<"i8*">, %arg1: !llvm<"i32*">)
+define void @indirectbrTest(i8* %address, i32* %sink) #0 {
+
+entry:
+  ; CHECK: llvm.indirect_br %arg0 : !llvm<"i8*"> [^bb{{[12]}}, ^bb{{[12]}}]
+  indirectbr i8* %address, [label %bb1, label %bb2]
+
+; ^bb1: // pred: ^bb0
+bb1:
+  store volatile i32 0, i32* %sink
+  br label %latch
+
+; CHECK: ^bb2:  // 2 preds: ^bb0, ^bb4
+bb2:
+  store volatile i32 1, i32* %sink
+  br label %latch
+
+; CHECK: ^bb3:  // pred: ^bb4
+bb3:
+  store volatile i32 2, i32* %sink
+  br label %latch
+
+; CHECK: ^bb4:  // 3 preds: ^bb1, ^bb2, ^bb3
+latch:
+  ; CHECK: llvm.indirect_br %arg0 : !llvm<"i8*"> [^bb{{[23]}}, ^bb{{[23]}}]
+  indirectbr i8* %address, [label %bb2, label %bb3, label %bb3]
+}

mlir/test/Target/llvmir.mlir

^bb2: // 2 preds: ^bb0, ^bb3
   llvm.return %5 : !llvm.i32
 
 // CHECK: [[final]]:
 // CHECK-NEXT: %{{[0-9]+}} = invoke i8* @bar(i8* %[[a1]])
 // CHECK-NEXT:          to label %[[normal]] unwind label %[[unwind]]
 ^bb3: // pred: ^bb1
   %8 = llvm.invoke @bar(%6) to ^bb2 unwind ^bb1 : (!llvm<"i8*">) -> !llvm<"i8*">
 }
+
+// CHECK-LABEL: @indirectbr(i8* %0, i32* %1)
+llvm.func @indirectbr(%arg0: !llvm<"i8*">, %arg1: !llvm<"i32*">) {
+  %0 = llvm.mlir.constant(2 : i32) : !llvm.i32
+  %1 = llvm.mlir.constant(1 : i32) : !llvm.i32
+  %2 = llvm.mlir.constant(0 : i32) : !llvm.i32
+  // CHECK: indirectbr i8* %0, [label %{{[34]}}, label %{{[34]}}]
+  llvm.indirect_br %arg0 : !llvm<"i8*"> [^bb1, ^bb2]
+^bb1: // pred: ^bb0
+  llvm.store %2, %arg1 : !llvm<"i32*">
+  llvm.br ^bb4
+^bb2: // 2 preds: ^bb0, ^bb4
+  llvm.store %1, %arg1 : !llvm<"i32*">
+  llvm.br ^bb4
+^bb3: // pred: ^bb4
+  llvm.store %0, %arg1 : !llvm<"i32*">
+  llvm.br ^bb4
+^bb4: // 3 preds: ^bb1, ^bb2, ^bb3
+  // CHECK: indirectbr i8* %0, [label %{{[45]}}, label %{{[45]}}]
+  llvm.indirect_br %arg0 : !llvm<"i8*"> [^bb3, ^bb2, ^bb2]
+}